💻 computer science

frax: Fast Robot Kinematics and Dynamics in JAX

Questo articolo presenta frax, una libreria JAX in puro Python ad alte prestazioni per la cinematica e la dinamica dei robot, che offre esecuzione CPU a bassa latenza e parallelizzazione massiva su GPU, supportando al contempo la differenziazione automatica per il controllo e l'apprendimento basati sull'ottimizzazione.

Autori originali: Daniel Morton, Marco Pavone

Pubblicato 2026-05-18

📖 4 min di lettura☕ Lettura da pausa caffè

CC BY 4.0

Autori originali: Daniel Morton, Marco Pavone

Articolo originale sotto licenza CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Questa è una spiegazione generata dall'IA dell'articolo qui sotto. Non è stata scritta né approvata dagli autori. Per precisione tecnica, consulta l'articolo originale. Leggi il disclaimer completo

Immagina di dover insegnare a un robot come muovere il braccio, camminare o afferrare una palla. Per farlo, hai bisogno di un "cervello" in grado di calcolare istantaneamente la fisica: quanto pesa il braccio, quanto velocemente oscilla e esattamente quanta forza è necessaria per fermarlo o cambiarne la direzione. Questo è chiamato dinamica robotica.

Per lungo tempo, gli strumenti a disposizione degli ingegneri per questi calcoli erano come avere due diversi tipi di calcolatrici:

La calcolatrice "Veloce": Era incredibilmente veloce ma funzionava solo su un tipo specifico di chip per computer (CPU). Se provavi a usarla su una moderna scheda grafica (GPU), si rompeva.
La calcolatrice "Produzione di massa": Era ottima per eseguire migliaia di calcoli contemporaneamente (come simulare 1.000 robot allo stesso tempo), ma era lenta e ingombrante quando volevi controllare un singolo robot.

Il documento introduce frax, un nuovo strumento che agisce come un traduttore universale per la fisica dei robot. È costruito utilizzando una tecnologia chiamata JAX, che gli permette di essere scritto in codice Python semplice ma eseguito con la velocità grezza di codice C++ ad alte prestazioni.

Ecco come funziona frax, utilizzando alcune analogie quotidiane:

1. La "Linea di Assemblaggio" contro l'"Incontro di Squadra"

La maggior parte delle vecchie librerie per robot calcola la fisica come se fosse una singola persona che tiene un incontro. Chiedono: "Ok, l'articolazione 1 si muove, quindi l'articolazione 2 si muove, il che fa muovere l'articolazione 3..." Lo fanno passo dopo passo, uno dopo l'altro. Questo è veloce per un robot, ma se provi a farlo per 1.000 robot, la linea si intasa.

frax cambia le regole del gioco. Invece di un incontro, agisce come una massiccia linea di assemblaggio sincronizzata. Calcola il movimento di ogni articolazione per ogni robot tutti allo stesso identico momento.

Il trucco: Utilizza una "maschera" (un filtro digitale) per capire quali articolazioni influenzano quali altre, permettendo al computer di eseguire tutti i calcoli in parallelo.
Il risultato: Su un singolo computer, è abbastanza veloce da controllare un robot in tempo reale (migliaia di volte al secondo). Su una potente scheda grafica, può simulare 100 milioni di movimenti robotici in un singolo secondo.

2. Lo "Specchio Magico" (Differenziazione Automatica)

Una delle parti più difficili dell'apprendimento robotico è capire come migliorare il movimento del robot. Di solito, devi scrivere manualmente formule matematiche complesse per vedere come cambiare una piccola cosa influenzi l'intero sistema.

frax viene fornito con uno "Specchio Magico" integrato. Poiché è costruito su JAX, può automaticamente osservare il movimento del robot e dirti istantaneamente esattamente come modificarlo per ottenere risultati migliori. Non devi essere un genio della matematica per derivare queste formule; la libreria fa il lavoro pesante per te. Questo rende molto più facile progettare robot che possono imparare e adattarsi.

3. Velocità e Flessibilità

Gli autori hanno testato frax su due robot reali: un Franka Panda (un braccio robotico) e un Unitree G1 (un robot umanoide).

Su un computer standard (CPU): frax è così veloce da superare di 2 o 5 volte altri popolari strumenti Python. Si avvicina alla velocità dei migliori strumenti scritti in codice C++ complesso, ma senza il mal di testa di scrivere C++.
Su una scheda grafica (GPU): Si comporta esattamente come i migliori strumenti progettati specificamente per le schede grafica, ma si avvia e compila molto più velocemente, risparmiando tempo agli ingegneri quando testano nuove idee.

La Conclusione

frax è una nuova libreria open-source che permette agli ingegneri di scrivere software di controllo robotico in Python semplice. È abbastanza veloce da controllare un singolo robot in tempo reale su un laptop standard, ma abbastanza potente da simulare migliaia di robot contemporaneamente su un supercomputer. Colma il divario tra "facile da usare" ed "estremamente veloce", rendendo più facile costruire robot che possono muoversi, imparare e interagire con il mondo reale in modo sicuro.

Riepilogo Tecnico: frax: Cinematica e Dinamica Robotica Veloci in JAX

Enunciato del Problema

Nei campi del controllo, della pianificazione e dell'apprendimento robotici, esiste un bisogno critico di librerie per la dinamica dei corpi rigidi che siano simultaneamente ad alte prestazioni, facili da usare e compatibili con diverse architetture hardware (CPU, GPU e TPU). Le soluzioni esistenti presentano compromessi significativi:

Librerie Ottimizzate per CPU: Librerie come Pinocchio e MuJoCo eccellono nell'esecuzione a bassa latenza su CPU, ma non possono essere eseguite direttamente su acceleratori. I loro binding Python spesso introducono sovraccarichi e complicano l'integrazione con flussi di lavoro di differenziazione automatica (AD).
Librerie Ottimizzate per GPU: Librerie come MJX e BRAX sfruttano il parallelismo GPU per simulazioni ad alto throughput, ma spesso soffrono di un elevato sovraccarico quando utilizzate per il controllo di singoli robot su CPU, rendendole inadatte a loop di controllo in tempo reale a frequenza chilohertz.
Compatibilità AD: Integrare librerie di dinamica basate su C++ in flussi di lavoro di apprendimento automatico o pianificazione basati su Python crea spesso attriti a causa della difficoltà di applicare la differenziazione automatica attraverso i confini linguistici.

Manca un quadro unificato che miri a più architetture senza compromettere le prestazioni o la facilità d'uso.

Metodologia

Gli autori introducono frax, una libreria puramente Python costruita su JAX che fornisce un'interfaccia ad alte prestazioni per la cinematica e la dinamica robotica. L'innovazione metodologica fondamentale è un approccio completamente vettorizzato alla dinamica dei corpi rigidi, che si discosta dai tradizionali metodi di attraversamento ricorsivo degli alberi (ad esempio, l'Algoritmo Ricorsivo di Newton-Euler) utilizzati nella maggior parte delle librerie C++.

Componenti Tecnici Chiave:

Formulazione Vettorizzata: Invece di cicli ricorsivi, frax codifica la struttura dell'albero cinematico utilizzando una maschera degli antenati ( $U \in \{0, 1\}^{n \times n}$ $U \in {0, 1}^{n \times n}$ ). Ciò consente di eseguire i calcoli dinamici tramite moltiplicazione di matrici e trasmissione (broadcasting) di array.
- Dinamica Inversa (RNEA): Calcolata sommando forze e accelerazioni su e giù per l'albero utilizzando la maschera $U$ , evitando dipendenze sequenziali.
- Matrice di Massa (CRBA): Calcolata vettorizzando l'accumulo di inerzie composte.
Compromesso sulla Complessità: Questo approccio aumenta intenzionalmente la complessità computazionale teorica di algoritmi come RNEA da $O(n)$ a $O(n^2)$ e CRBA da $O(nd) $a$ O(n^2)$ (dove $n$ è il numero di gradi di libertà). Tuttavia, gli autori sostengono che i guadagni di prestazioni derivanti dal parallelismo fine-granulare sull'hardware moderno (tramite compilazione XLA) e l'eliminazione del sovraccarico dei cicli superano di gran lunga i calcoli ridondanti.
Integrazione JAX: Sfruttando JAX, frax supporta nativamente:
- Compilazione Just-In-Time (JIT): Per un'esecuzione ottimizzata su CPU, GPU e TPU.
- Differenziazione Automatica (AD): Abilitando il calcolo efficiente dei gradienti per il controllo basato sull'ottimizzazione senza il costo di tracciare lunghi cicli srotolati.
- Prodotti Jacobiano-Vettore (JVP): Consentendo una valutazione efficiente delle derivate di Lie per vincoli di controllo (ad esempio, Funzioni di Barriera di Controllo) senza calcolare i Jacobiani completi.

Contributi Chiave

Architettura Unificata: frax fornisce un'unica base di codice che esegue in modo efficiente su CPU, GPU e TPU, colmando il divario tra controllo a bassa latenza e simulazione ad alto throughput.
Prestazioni:
- Su CPU: Raggiunge tempi di calcolo nell'ordine dei microsecondi bassi, adatti a frequenze di controllo chilohertz, superando i binding Python di Pinocchio e MuJoCo di 2–5 volte e avvicinandosi alle prestazioni grezze delle implementazioni C++ ottimizzate.
- Su GPU/TPU: Si scala a migliaia di istanze parallele, raggiungendo oltre 100 milioni di valutazioni dinamiche al secondo, eguagliando o superando le prestazioni di simulatori GPU dedicati come MJX e BRAX.
Facilità d'Uso: Essendo una libreria puramente Python, semplifica la prototipazione e si integra perfettamente con l'ecosistema Python, eliminando la necessità di binding C++ mantenendo al contempo elevate prestazioni.
Open Source: La libreria è rilasciata come software open source, con supporto per robot specifici come Franka Panda e Unitree G1, inclusi modelli di collisione tarati.

Risultati

Gli autori hanno validato frax su un manipolatore Franka Panda e su un umanoide Unitree G1.

Latenza: Su CPU, frax ha raggiunto tempi di cinematica inversa (IK differenziale) di 4,09 µs (Franka) e 6,71 µs (G1), e tempi di dinamica inversa (Controllo nello Spazio Operativo) di 7,03 µs (Franka) e 24,6 µs (G1).
Confronto:
- Rispetto alle API Python di Pinocchio e MuJoCo, frax è stato 2–5 volte più veloce grazie alla compilazione JIT e alle fusioni del compilatore.
- Rispetto alle interfacce C++ di Pinocchio/MuJoCo, frax ne avvicina le prestazioni grezze ma mantiene la flessibilità Python e le capacità AD.
- Rispetto a MJX e BRAX, frax ha mostrato prestazioni comparabili o leggermente superiori su GPU, ma con tempi di compilazione JIT significativamente ridotti (1–2 secondi contro 6–12 secondi), facilitando uno sviluppo iterativo più rapido.
Throughput: Su GPU, la libreria ha dimostrato tassi di throughput superiori a 100 milioni di valutazioni al secondo con grandi dimensioni di batch (4096), indicando l'idoneità per ambienti di addestramento massivamente paralleli.

Significato e Affermazioni

Il documento afferma che frax rappresenta un passo significativo verso l'unificazione dei flussi di lavoro di controllo e apprendimento robotici. Il suo significato principale risiede nel consentire il controllo in tempo reale di singoli robot su CPU supportando simultaneamente la simulazione massivamente parallela su acceleratori all'interno della stessa base di codice.

Gli autori sottolineano che frax permette ai ricercatori di:

Prototipare controllori in Python con differenziazione automatica (ad esempio, per ottimizzazione di traiettorie non lineari o MPC) senza sacrificare le prestazioni richieste per il deployment sull'hardware.
Sostituire i calcoli completi dei Jacobiani con JVP efficienti per vincoli critici per la sicurezza.
Distribuire la stessa logica su diversi backend hardware (CPU per l'esecuzione, GPU/TPU per l'addestramento) senza riscrivere gli algoritmi.

Il documento mantiene una posizione modesta, riconoscendo che frax è attualmente più minimale rispetto a librerie complete come Pinocchio (ad esempio, manca di solver IK globali e ambienti di simulazione completi) e che le implementazioni C++ potrebbero ancora offrire frequenze di controllo marginalmente superiori. Tuttavia, posiziona frax come una fondazione flessibile e ad alte prestazioni progettata specificamente per l'intersezione tra il controllo robotico moderno e l'apprendimento automatico.