frax

Bibliothèque JAX de cinématique, dynamique et cinématique inverse de robots, rapide et différentiable

Licence

MIT

Aperçu

frax est une bibliothèque de cinématique et de dynamique des robots en Python pur, construite sur JAX par l'Autonomous Systems Lab de Stanford. À partir d'un URDF, elle calcule les transformations, les jacobiennes, la matrice de masse, les termes de gravité et de Coriolis ainsi que la dynamique directe et inverse, le tout compilé JIT et différentiable. Les contrôleurs typiques s'exécutent en quelques microsecondes sur un CPU, et le même code traite en lot des milliers de robots sur un GPU ou un TPU. La version 0.2.0 ajoute frax.ik, une boîte à outils de cinématique inverse différentielle et de retargeting reposant sur le solveur ElastiQP.

frax, de Daniel Morton et Marco Pavone à l'Autonomous Systems Lab de Stanford, est une bibliothèque de cinématique et de dynamique des corps rigides pour robots, écrite en Python pur sur JAX. Vous chargez un robot depuis un URDF, ou l'un des modèles intégrés comme le Franka Panda ou l'Unitree G1, puis vous lui demandez les transformations des articulations, des segments et des repères, les jacobiennes, la matrice de masse, les termes de gravité et de Coriolis, ou la dynamique directe et inverse. Chaque terme étant une fonction JAX, il peut être compilé JIT, exécuté sur CPU, GPU ou TPU, vectorisé sur de nombreux robots à la fois et dérivé par différentiation automatique.

Les auteurs visent deux types de travail. Pour commander un seul robot, un CPU calcule les termes d'un contrôleur typique en quelques microsecondes, ce que l'article situe à des fréquences de commande de l'ordre du kilohertz ; le README recommande ici le CPU même lorsqu'un GPU est disponible. En lot, le même code dépasse 100 millions d'évaluations de dynamique par seconde sur un GPU. Des modèles de collision à base de sphères, déjà prêts pour le Franka Panda/FR3 et l'Unitree G1, permettent à un contrôleur d'ajouter l'évitement des autocollisions et des obstacles sous forme de simples fonctions de la cinématique du robot.

La version 0.2.0 (4 novembre 2026) ajoute frax.ik, une boîte à outils de cinématique inverse différentielle dans l'esprit de mink et pink, mais différentiable et reposant sur ElastiQP, un solveur QP du même laboratoire qui reste réalisable lorsque les contraintes se contredisent. Elle gère la commande et la téléopération en ligne à plus de 1 000 Hz et le retargeting de mouvement hors ligne à plus de 100 000 images par seconde : la documentation transpose l'ensemble du jeu de données de mouvements humains LAFAN1 sur un Unitree G1 en environ quatre secondes sur un seul CPU. L'installation se fait avec pip install "frax[ik]".

Pourquoi c'est important

Le code de commande des robots doit généralement choisir entre une bibliothèque de dynamique rapide en C++ comme Pinocchio et un simulateur sur GPU en JAX. frax fournit les termes de dynamique eux-mêmes sous forme de fonctions JAX différentiables, rapides sur un seul CPU et capables de traitement en lot sur des accélérateurs : le même modèle sert ainsi à un contrôleur temps réel, à une optimisation qui a besoin de gradients et à un gros travail de retargeting. Avec frax.ik, il rejoint aussi mink et PyRoki comme option ouverte pour le retargeting d'humanoïdes, à une vitesse qui rend banal le traitement de jeux de données de mouvement entiers.

Le problème résolu

Obtenir une cinématique et une dynamique de robots assez rapides pour une commande au kilohertz, différentiables pour l'optimisation et parallèles pour de gros lots oblige généralement à combiner plusieurs bibliothèques. frax fournit les trois à partir d'un seul modèle Python du robot.

Points clés

• Termes de cinématique et de dynamique issus d'un URDF sous forme de fonctions JAX compilées JIT et différentiables
• Termes de contrôleur en quelques microsecondes sur un CPU ; plus de 100 millions d'évaluations de dynamique par seconde en lot sur un GPU
• frax.ik dans la 0.2.0 : cinématique inverse différentielle et retargeting sur le solveur toujours réalisable ElastiQP, à plus de 100 000 images par seconde hors ligne
• Modèles de collision à base de sphères pour le Franka Panda/FR3 et l'Unitree G1, pour l'évitement des autocollisions et des obstacles
• Licence MIT, sur PyPI sous le nom frax, avec un article sur arXiv (2604.04310)

Pour qui

• Chercheurs en commande qui conçoivent en Python des contrôleurs de cinématique inverse, de dynamique inverse et des filtres de sécurité
• Équipes d'apprentissage robotique qui transposent des données de mouvement humain sur des humanoïdes comme l'Unitree G1
• Développeurs qui ont besoin d'une dynamique des corps rigides différentiable ou en lot sur GPU ou TPU

Limites

• URDF uniquement pour l'instant ; la prise en charge de MJCF figure dans la liste des tâches des auteurs
• La géométrie de collision doit être faite de sphères, donc les autres robots ont d'abord besoin d'un modèle de collision en sphères
• Pas de chaînes cinématiques fermées, et uniquement des articulations rotoïdes, prismatiques, fixes et à base flottante
• Les articulations de pinces et autres hors de la chaîne doivent être déclarées fixed dans l'URDF avant utilisation
• Les performances dépendent de la version et des réglages de JAX : le README recommande JAX 0.4.30 pour la dynamique sur CPU seul
• L'interface évolue encore : la 0.2.0 a modifié la convention des articulations à base flottante et certains noms