frax

Libreria JAX per cinematica, dinamica e cinematica inversa dei robot, veloci e differenziabili

Licenza

MIT

Panoramica

frax è una libreria di cinematica e dinamica dei robot in puro Python, basata su JAX, dell'Autonomous Systems Lab di Stanford. A partire da un URDF calcola trasformazioni, jacobiani, matrice delle masse, termini di gravità e di Coriolis e dinamica diretta e inversa, tutto compilato JIT e differenziabile. I controller tipici girano in pochi microsecondi su una CPU, e lo stesso codice si esegue in batch su migliaia di robot su una GPU o una TPU. La versione 0.2.0 aggiunge frax.ik, un toolkit di cinematica inversa differenziale e di retargeting basato sul solver ElastiQP.

frax, di Daniel Morton e Marco Pavone dell'Autonomous Systems Lab di Stanford, è una libreria per la cinematica e la dinamica dei corpi rigidi dei robot scritta in puro Python su JAX. Si carica un robot da un URDF, o uno dei modelli integrati come il Franka Panda o l'Unitree G1, e gli si chiedono le trasformazioni di giunti, link e frame, i jacobiani, la matrice delle masse, i termini di gravità e di Coriolis, o la dinamica diretta e inversa. Poiché ogni termine è una funzione JAX, può essere compilato JIT, eseguito su CPU, GPU o TPU, vettorizzato su molti robot contemporaneamente e differenziato con la differenziazione automatica.

Gli autori puntano a due tipi di lavoro. Per il controllo di un singolo robot, una CPU calcola i termini di un controller tipico in pochi microsecondi, il che secondo l'articolo consente frequenze di controllo nell'ordine dei kilohertz; in questo caso il README consiglia la CPU anche quando è presente una GPU. Per i batch, lo stesso codice supera i 100 milioni di valutazioni della dinamica al secondo su una GPU. I modelli di collisione a sfere, già pronti per il Franka Panda/FR3 e l'Unitree G1, permettono a un controller di aggiungere l'autocollisione e l'evitamento degli ostacoli come semplici funzioni della cinematica del robot.

La versione 0.2.0 (4 nov 2026) aggiunge frax.ik, un toolkit di cinematica inversa differenziale nello spirito di mink e pink, ma differenziabile e basato su ElastiQP, un solver QP dello stesso laboratorio che resta ammissibile quando i vincoli sono in conflitto. Gestisce il controllo online e la teleoperazione a oltre 1.000 Hz e il retargeting del movimento offline a oltre 100.000 frame al secondo: la documentazione esegue il retargeting dell'intero dataset di movimento umano LAFAN1 su un Unitree G1 in circa quattro secondi su una sola CPU. L'installazione è pip install "frax[ik]".

Perché è importante

Il codice di controllo dei robot di solito sceglie tra una libreria di dinamica C++ veloce come Pinocchio e un simulatore GPU in JAX. frax fornisce i termini della dinamica stessi come funzioni JAX differenziabili, veloci su una singola CPU e eseguibili anche in batch sugli acceleratori, così lo stesso modello serve un controller in tempo reale, un'ottimizzazione che richiede gradienti e un grande lavoro di retargeting. Con frax.ik si affianca inoltre a mink e PyRoki come opzione aperta per il retargeting sugli umanoidi, a una velocità che rende ordinaria l'elaborazione di interi dataset di movimento.

Il problema che risolve

Ottenere una cinematica e una dinamica dei robot abbastanza veloci per il controllo nell'ordine dei kilohertz, differenziabili per l'ottimizzazione e parallele per grandi batch di solito significa combinare più librerie. frax fornisce tutte e tre le cose da un unico modello Python del robot.

Punti chiave

• Termini di cinematica e dinamica da un URDF come funzioni JAX compilate JIT e differenziabili
• Termini del controller in pochi microsecondi su una CPU; oltre 100 milioni di valutazioni della dinamica al secondo in batch su una GPU
• frax.ik nella 0.2.0: cinematica inversa differenziale e retargeting sul solver ElastiQP sempre ammissibile, oltre 100.000 frame al secondo offline
• Modelli di collisione a sfere per Franka Panda/FR3 e Unitree G1, per l'autocollisione e l'evitamento degli ostacoli
• Licenza MIT, su PyPI come frax, con un articolo su arXiv (2604.04310)

A chi si rivolge

• Ricercatori di controllo che progettano in Python controller di cinematica inversa, di dinamica inversa e con filtri di sicurezza
• Gruppi di apprendimento robotico che eseguono il retargeting di dati di movimento umano su umanoidi come l'Unitree G1
• Sviluppatori che necessitano di dinamica dei corpi rigidi differenziabile o in batch su GPU o TPU

Limiti

• Per ora solo URDF; il supporto MJCF è nell'elenco delle cose da fare degli autori
• La geometria di collisione deve essere a sfere, quindi per altri robot serve prima un modello di collisione a sfere
• Nessuna catena cinematica chiusa e solo giunti rotoidali, prismatici, fissi e flottanti
• I giunti della pinza e gli altri giunti fuori catena vanno impostati come fissi nell'URDF prima dell'uso
• Le prestazioni dipendono dalla versione e dalle impostazioni di JAX: il README consiglia JAX 0.4.30 per la dinamica solo su CPU
• L'interfaccia sta ancora cambiando: la 0.2.0 ha modificato la convenzione del giunto flottante e alcuni nomi