pi zero pytorch

pi-zéro-pytorch (wip)

Implémentation de π₀ l'architecture du modèle de fondation robotique proposée par Physical Intelligence

Le résumé de ce travail serait qu'il s'agit d'une transfusion simplifiée (Zhou et al.) avec l'influence de Stable Diffusion 3 (Esser et al.), principalement l'adoption de l'appariement de flux au lieu de la diffusion pour la génération de politiques, ainsi que la séparation des paramètres (Attention conjointe de mmDIT). Ils s'appuient sur un modèle de langage de vision pré-entraîné, PaliGemma 2B.

Appréciation

• Einops pour l'incroyable pack et déballage, largement utilisé ici pour gérer divers ensembles de jetons

• Flex Attention pour permettre un mélange facile d'attention autorégressive et bidirectionnelle

• @Wonder1905 pour la révision du code et l'identification des problèmes

• Toi? peut-être un doctorant qui souhaite contribuer à la dernière architecture SOTA pour le clonage comportemental ?

Installer

$ pip install pi-zero-pytorch

Usage

 import torch
from pi_zero_pytorch import π0

model = π0 (
    dim = 512 ,
    dim_action_input = 6 ,
    dim_joint_state = 12 ,
    num_tokens = 20_000
)

vision = torch . randn ( 1 , 1024 , 512 )
commands = torch . randint ( 0 , 20_000 , ( 1 , 1024 ))
joint_state = torch . randn ( 1 , 12 )
actions = torch . randn ( 1 , 32 , 6 )

loss , _ = model ( vision , commands , joint_state , actions )
loss . backward ()

# after much training

sampled_actions = model ( vision , commands , joint_state , trajectory_length = 32 ) # (1, 32, 6)

Contribuer

À la racine du projet, exécutez

$ pip install ' .[test] ' # or `uv pip install '.[test]'`

Ajoutez ensuite vos tests à tests/test_pi_zero.py et exécutez

$ pytest tests/

C'est ça

Citation

 @misc { Black2024 ,
    author  = { Kevin Black, Noah Brown, Danny Driess, Adnan Esmail, Michael Equi, Chelsea Finn, Niccolo Fusai, Lachy Groom, Karol Hausman, Brian Ichter, Szymon Jakubczak, Tim Jones, Liyiming Ke, Sergey Levine, Adrian Li-Bell, Mohith Mothukuri, Suraj Nair, Karl Pertsch, Lucy Xiaoyang Shi, James Tanner, Quan Vuong, Anna Walling, Haohuan Wang, Ury Zhilinsky } ,
    url     = { https://www.physicalintelligence.company/download/pi0.pdf }
}

 @inproceedings { Zhou2024ValueRL ,
    title   = { Value Residual Learning For Alleviating Attention Concentration In Transformers } ,
    author  = { Zhanchao Zhou and Tianyi Wu and Zhiyun Jiang and Zhenzhong Lan } ,
    year    = { 2024 } ,
    url     = { https://api.semanticscholar.org/CorpusID:273532030 }
}

 @inproceedings { Darcet2023VisionTN ,
    title   = { Vision Transformers Need Registers } ,
    author  = { Timoth'ee Darcet and Maxime Oquab and Julien Mairal and Piotr Bojanowski } ,
    year    = { 2023 } ,
    url     = { https://api.semanticscholar.org/CorpusID:263134283 }
}

 @article { Li2024ImmiscibleDA ,
    title   = { Immiscible Diffusion: Accelerating Diffusion Training with Noise Assignment } ,
    author  = { Yiheng Li and Heyang Jiang and Akio Kodaira and Masayoshi Tomizuka and Kurt Keutzer and Chenfeng Xu } ,
    journal = { ArXiv } ,
    year    = { 2024 } ,
    volume  = { abs/2406.12303 } ,
    url     = { https://api.semanticscholar.org/CorpusID:270562607 }
}

 @inproceedings { Sadat2024EliminatingOA ,
    title   = { Eliminating Oversaturation and Artifacts of High Guidance Scales in Diffusion Models } ,
    author  = { Seyedmorteza Sadat and Otmar Hilliges and Romann M. Weber } ,
    year    = { 2024 } ,
    url     = { https://api.semanticscholar.org/CorpusID:273098845 }
}

 @article { Bulatov2022RecurrentMT ,
    title   = { Recurrent Memory Transformer } ,
    author  = { Aydar Bulatov and Yuri Kuratov and Mikhail S. Burtsev } ,
    journal = { ArXiv } ,
    year    = { 2022 } ,
    volume  = { abs/2207.06881 } ,
    url     = { https://api.semanticscholar.org/CorpusID:250526424 }
}

 @inproceedings { Bessonov2023RecurrentAT ,
    title   = { Recurrent Action Transformer with Memory } ,
    author  = { A. B. Bessonov and Alexey Staroverov and Huzhenyu Zhang and Alexey K. Kovalev and D. Yudin and Aleksandr I. Panov } ,
    year    = { 2023 } ,
    url     = { https://api.semanticscholar.org/CorpusID:259188030 }
}

 @article { Zhu2024HyperConnections ,
    title   = { Hyper-Connections } ,
    author  = { Defa Zhu and Hongzhi Huang and Zihao Huang and Yutao Zeng and Yunyao Mao and Banggu Wu and Qiyang Min and Xun Zhou } ,
    journal = { ArXiv } ,
    year    = { 2024 } ,
    volume  = { abs/2409.19606 } ,
    url     = { https://api.semanticscholar.org/CorpusID:272987528 }
}

chère Alice