frax

Biblioteca JAX para cinemática, dinámica y cinemática inversa de robots, rápidas y diferenciables

Licencia

MIT

Resumen

frax es una biblioteca de cinemática y dinámica de robots en Python puro, construida sobre JAX por el Autonomous Systems Lab de Stanford. A partir de un URDF calcula transformaciones, jacobianos, la matriz de masa, los términos de gravedad y de Coriolis y la dinámica directa e inversa, todo compilado JIT y diferenciable. Los controladores típicos se ejecutan en unos pocos microsegundos en una CPU, y el mismo código procesa por lotes miles de robots en una GPU o TPU. La versión 0.2.0 añade frax.ik, un conjunto de herramientas de cinemática inversa diferencial y retargeting basado en el solver ElastiQP.

frax, de Daniel Morton y Marco Pavone en el Autonomous Systems Lab de Stanford, es una biblioteca de cinemática y dinámica de cuerpo rígido para robots escrita en Python puro sobre JAX. Cargas un robot desde un URDF, o uno de los modelos incluidos como el Franka Panda o el Unitree G1, y le pides transformaciones de articulaciones, eslabones y marcos, jacobianos, la matriz de masa, los términos de gravedad y de Coriolis o la dinámica directa e inversa. Como cada término es una función de JAX, se puede compilar con JIT, ejecutar en CPU, GPU o TPU, vectorizar sobre muchos robots a la vez y derivar con diferenciación automática.

Los autores apuntan a dos tipos de trabajo. Para controlar un solo robot, una CPU calcula los términos de un controlador típico en unos pocos microsegundos, lo que el artículo sitúa en frecuencias de control de kilohercios; el README recomienda aquí la CPU incluso cuando hay una GPU. Para lotes, el mismo código alcanza más de 100 millones de evaluaciones de dinámica por segundo en una GPU. Los modelos de colisión basados en esferas, ya preparados para el Franka Panda/FR3 y el Unitree G1, permiten a un controlador añadir la evitación de autocolisiones y de obstáculos como funciones simples de la cinemática del robot.

La versión 0.2.0 (4 de noviembre de 2026) añade frax.ik, un conjunto de herramientas de cinemática inversa diferencial al estilo de mink y pink, pero diferenciable y basado en ElastiQP, un solver QP del mismo laboratorio que sigue siendo factible cuando las restricciones entran en conflicto. Gestiona control y teleoperación en línea a más de 1.000 Hz y retargeting de movimiento fuera de línea a más de 100.000 fotogramas por segundo: la documentación adapta todo el conjunto de datos de movimiento humano LAFAN1 a un Unitree G1 en unos cuatro segundos en una sola CPU. Se instala con pip install "frax[ik]".

Por qué importa

El código de control de robots suele tener que elegir entre una biblioteca de dinámica rápida en C++ como Pinocchio y un simulador en GPU con JAX. frax ofrece los propios términos de dinámica como funciones diferenciables de JAX que son rápidas en una sola CPU y también se procesan por lotes en aceleradores, de modo que el mismo modelo sirve para un controlador en tiempo real, para una optimización que necesita gradientes y para un gran trabajo de retargeting. Con frax.ik se suma además a mink y PyRoki como opción abierta para el retargeting de humanoides, a una velocidad que convierte en rutina procesar conjuntos de datos de movimiento completos.

Qué problema resuelve

Conseguir una cinemática y una dinámica de robots que sean lo bastante rápidas para el control a kilohercios, diferenciables para la optimización y paralelas para grandes lotes suele obligar a combinar varias bibliotecas. frax ofrece las tres cosas a partir de un único modelo del robot en Python.

Puntos clave

• Términos de cinemática y dinámica a partir de un URDF como funciones de JAX compiladas con JIT y diferenciables
• Términos de controlador en pocos microsegundos en una CPU; más de 100 millones de evaluaciones de dinámica por segundo por lotes en una GPU
• frax.ik en la 0.2.0: cinemática inversa diferencial y retargeting sobre el solver siempre factible ElastiQP, a más de 100.000 fotogramas por segundo fuera de línea
• Modelos de colisión basados en esferas para el Franka Panda/FR3 y el Unitree G1, para evitar autocolisiones y obstáculos
• Licencia MIT, en PyPI como frax, con un artículo en arXiv (2604.04310)

Para quién es

• Investigadores de control que diseñan controladores de cinemática inversa, dinámica inversa y filtros de seguridad en Python
• Grupos de aprendizaje robótico que adaptan datos de movimiento humano a humanoides como el Unitree G1
• Desarrolladores que necesitan dinámica de cuerpo rígido diferenciable o por lotes en GPU o TPU

Limitaciones

• Por ahora solo URDF; el soporte de MJCF está en la lista de pendientes de los autores
• La geometría de colisión debe ser de esferas, así que otros robots necesitan primero un modelo de colisión con esferas
• Sin cadenas cinemáticas cerradas, y solo articulaciones de revolución, prismáticas, fijas y de base flotante
• Las articulaciones de pinzas y otras fuera de la cadena deben marcarse como fixed en el URDF antes de usarlo
• El rendimiento depende de la versión y la configuración de JAX: el README recomienda JAX 0.4.30 para dinámica solo en CPU
• La interfaz aún cambia: la 0.2.0 cambió la convención de las articulaciones de base flotante y algunos nombres