frax : cinématique et dynamique robotique rapide en JAX

Une équipe de chercheurs a publié frax, une bibliothèque open-source de cinématique et dynamique de corps rigides pour la robotique, construite sur JAX, le framework d'autodifférenciation de Google. Documentée dans l'article arXiv:2604.04310, la librairie cible une lacune précise : la plupart des bibliothèques existantes excellent soit en exécution CPU à faible latence, soit en débit GPU massif, mais rarement les deux à la fois. frax propose une interface unifiée en Python pur, opérationnelle sur CPU, GPU et TPU sans modification de code. Sur CPU, les temps de calcul descendent à quelques microsecondes, compatibles avec des boucles de contrôle à l'échelle du kilohertz, à parité avec des implémentations C++ optimisées. Sur GPU, la même implémentation vectorisée atteint plus de 100 millions d'évaluations de dynamique par seconde en parallélisant des milliers d'instances simultanément. Les performances ont été validées sur deux plateformes de référence : le bras manipulateur Franka Panda et l'humanoïde Unitree G1.

L'impact pour les équipes de robotique est double. Le support natif de l'autodifférenciation via JAX ouvre la voie à des méthodes d'optimisation basées sur les gradients directement dans la boucle de simulation, ce qui accélère le model-based reinforcement learning et le trajectory optimization sans nécessiter d'implémentations séparées. Le fait d'atteindre des temps de cycle CPU comparables au C++ depuis un code Python pur représente également un gain de friction considérable pour les intégrateurs, entre le stade prototype et le déploiement. La scalabilité GPU à 100 millions d'évaluations par seconde est particulièrement pertinente pour le sim-to-real, où l'entraînement massif en simulation parallèle est devenu la norme dans les pipelines de robot learning modernes.

frax s'inscrit dans un écosystème JAX en expansion pour la robotique, aux côtés de Brax (Google DeepMind) et MuJoCo XLA (mjx). Les auteurs ne prétendent pas remplacer ces outils, mais proposer une alternative axée sur la polyvalence multi-matériel et la simplicité d'usage. La bibliothèque C++ Pinocchio, développée par le Gepetto team du LAAS-CNRS à Toulouse en partenariat avec l'INRIA, reste une référence pour les systèmes embarqués temps réel, mais frax vise explicitement le segment recherche et apprentissage. La bibliothèque est disponible en open-source, et les prochaines étapes attendues incluent l'intégration avec des pipelines RL existants et l'extension à la dynamique de contact.