Résumé technique : frax : Cinématique et dynamique rapides pour robots en JAX

Énoncé du problème

Dans les domaines du contrôle, de la planification et de l'apprentissage des robots, il existe un besoin critique de bibliothèques de dynamique des corps rigides qui soient simultanément haute performance, faciles à utiliser et compatibles avec des architectures matérielles diverses (CPU, GPU et TPU). Les solutions existantes font face à des compromis significatifs :

• Bibliothèques optimisées pour CPU : Des bibliothèques comme Pinocchio et MuJoCo excellent dans l'exécution à faible latence sur CPU, mais ne peuvent pas être exécutées directement sur des accélérateurs. Leurs liaisons Python introduisent souvent une surcharge et compliquent l'intégration avec les flux de travail de différenciation automatique (AD).
• Bibliothèques optimisées pour GPU : Des bibliothèques telles que MJX et BRAX exploitent le parallélisme GPU pour des simulations à haut débit, mais souffrent souvent d'une surcharge élevée lorsqu'elles sont utilisées pour le contrôle d'un seul robot sur CPU, ce qui les rend inadaptées aux boucles de contrôle en temps réel à des taux de kilohertz.
• Compatibilité AD : L'intégration de bibliothèques de dynamique basées sur C++ dans des flux de travail d'apprentissage automatique ou de planification basés sur Python crée souvent des frictions en raison de la difficulté d'appliquer la différenciation automatique à travers les frontières linguistiques.

Il manque un cadre unifié visant plusieurs architectures sans compromettre les performances ni la facilité d'utilisation.

Méthodologie

Les auteurs présentent frax, une bibliothèque purement Python construite sur JAX qui fournit une interface haute performance pour la cinématique et la dynamique des robots. L'innovation méthodologique centrale est une approche entièrement vectorisée de la dynamique des corps rigides, s'écartant des méthodes traditionnelles de parcours récursif d'arbres (par exemple, l'algorithme récursif de Newton-Euler) utilisées dans la plupart des bibliothèques C++.

Composants techniques clés :

1. Formulation vectorisée : Au lieu de boucles récursives, frax encode la structure de l'arbre cinématique à l'aide d'un masque d'ancêtre ($U \in \{0, 1\}^{n \times n}$). Cela permet d'effectuer les calculs de dynamique via la multiplication matricielle et la diffusion de tableaux.
   • Dynamique inverse (RNEA) : Calculée en sommant les forces et les accélérations vers le haut et vers le bas de l'arbre en utilisant le masque $U$, évitant ainsi les dépendances séquentielles.
   • Matrice de masse (CRBA) : Calculée en vectorisant l'accumulation des inerties composites.
2. Compromis de complexité : Cette approche augmente intentionnellement la complexité computationnelle théorique d'algorithmes comme RNEA de $O(n)$ à $O(n^2)$ et CRBA de $O(nd)$à$O(n^2)$(où$n$ est le nombre de degrés de liberté). Cependant, les auteurs soutiennent que les gains de performance issus du parallélisme fin sur le matériel moderne (via la compilation XLA) et l'élimination de la surcharge des boucles l'emportent largement sur les calculs redondants.
3. Intégration JAX : En exploitant JAX, frax prend nativement en charge :
   • Compilation Juste-à-Temps (JIT) : Pour une exécution optimisée sur CPU, GPU et TPU.
   • Différenciation automatique (AD) : Permettant un calcul efficace des gradients pour le contrôle basé sur l'optimisation sans le coût du traçage à travers de longues boucles déroulées.
   • Produits Jacobien-Vecteur (JVP) : Permettant une évaluation efficace des dérivées de Lie pour les contraintes de contrôle (par exemple, les fonctions de barrière de contrôle) sans calculer les jacobiens complets.

Contributions clés

1. Architecture unifiée : frax fournit une seule base de code qui s'exécute efficacement sur CPU, GPU et TPU, comblant le fossé entre le contrôle à faible latence et la simulation à haut débit.
2. Performances :
   • Sur CPU : Atteint des temps de calcul de l'ordre de la microseconde basse, adaptés aux taux de contrôle en kilohertz, surpassant les liaisons Python de Pinocchio et MuJoCo de 2 à 5 fois et approchant les performances brutes des implémentations C++ optimisées.
   • Sur GPU/TPU : S'adapte à des milliers d'instances parallèles, atteignant plus de 100 millions d'évaluations de dynamique par seconde, égalant ou dépassant les performances de simulateurs GPU dédiés comme MJX et BRAX.
3. Facilité d'utilisation : En tant que bibliothèque purement Python, elle simplifie le prototypage et s'intègre parfaitement à l'écosystème Python, éliminant le besoin de liaisons C++ tout en maintenant des performances élevées.
4. Open Source : La bibliothèque est publiée sous forme de logiciel open source, avec un support pour des robots spécifiques comme le Franka Panda et le Unitree G1, incluant des modèles de collision ajustés.

Résultats

Les auteurs ont validé frax sur un manipulateur Franka Panda et un humanoïde Unitree G1.

• Latence : Sur CPU, frax a atteint des temps de cinématique inverse (IK différentielle) de 4,09 µs (Franka) et 6,71 µs (G1), et des temps de dynamique inverse (contrôle dans l'espace opérationnel) de 7,03 µs (Franka) et 24,6 µs (G1).
• Comparaison :
  • Par rapport aux API Python de Pinocchio et MuJoCo, frax était 2 à 5 fois plus rapide grâce à la compilation JIT et aux fusions de compilateur.
  • Par rapport aux interfaces C++ de Pinocchio/MuJoCo, frax approche leurs performances brutes tout en conservant la flexibilité Python et les capacités AD.
  • Par rapport à MJX et BRAX, frax a montré des performances comparables ou légèrement supérieures sur GPU, mais avec des temps de compilation JIT considérablement réduits (1 à 2 secondes contre 6 à 12 secondes), facilitant un développement itératif plus rapide.
• Débit : Sur GPU, la bibliothèque a démontré des taux de débit dépassant 100 millions d'évaluations par seconde pour de grandes tailles de lots (4096), indiquant une adéquation pour des environnements d'entraînement massivement parallèles.

Importance et revendications

L'article revendique que frax représente une étape significative vers l'unification des flux de travail de contrôle et d'apprentissage des robots. Son importance principale réside dans la possibilité de permettre un contrôle en temps réel d'un seul robot sur CPU tout en soutenant simultanément une simulation massivement parallèle sur accélérateurs au sein de la même base de code.

Les auteurs soulignent que frax permet aux chercheurs de :

• Prototyper des contrôleurs en Python avec différenciation automatique (par exemple, pour l'optimisation de trajectoire non linéaire ou le MPC) sans sacrifier les performances requises pour le déploiement matériel.
• Remplacer les calculs complets de jacobien par des JVP efficaces pour des contraintes critiques pour la sécurité.
• Déployer la même logique sur différents backends matériels (CPU pour l'exécution, GPU/TPU pour l'entraînement) sans réécrire les algorithmes.

L'article adopte une position modeste, reconnaissant que frax est actuellement plus minimal que des bibliothèques complètes comme Pinocchio (par exemple, l'absence de solveurs d'IK globaux et d'environnements de simulation complets) et que les implémentations C++ peuvent encore offrir des taux de contrôle légèrement plus rapides. Cependant, il positionne frax comme une fondation flexible et haute performance, conçue spécifiquement pour l'intersection du contrôle robotique moderne et de l'apprentissage automatique.