💻 computer science

frax: Fast Robot Kinematics and Dynamics in JAX

Dit artikel introduceert frax, een hoogpresterende, pure-Python JAX-bibliotheek voor robotkinematica en -dynamica die lage latentie CPU-uitvoering en massale GPU-parallellisatie biedt, terwijl het automatische differentiatie ondersteunt voor optimalisatiegebaseerde besturing en leren.

Oorspronkelijke auteurs: Daniel Morton, Marco Pavone

Gepubliceerd 2026-05-18

📖 4 min leestijd☕ Koffiepauze-leesvoer

CC BY 4.0

Oorspronkelijke auteurs: Daniel Morton, Marco Pavone

Oorspronkelijk artikel gelicentieerd onder CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dit is een AI-gegenereerde uitleg van het onderstaande artikel. Het is niet geschreven of goedgekeurd door de auteurs. Raadpleeg het oorspronkelijke artikel voor technische nauwkeurigheid. Lees de volledige disclaimer

Stel je voor dat je probeert een robot te leren hoe hij zijn arm moet bewegen, moet lopen of een bal moet vangen. Om dit te doen, heb je een "hersenen" nodig die de fysica direct kan berekenen: hoe zwaar de arm is, hoe snel hij zwaait, en precies hoeveel kracht nodig is om hem te stoppen of van richting te laten veranderen. Dit noemen we robotdynamica.

Lange tijd waren de hulpmiddelen die ingenieurs voor deze berekening tot hun beschikking hadden, als het hebben van twee verschillende soorten rekenmachines:

De "Snelle" Rekenmachine: Deze was ongelooflijk snel, maar werkte alleen op een specifiek type computerchip (CPU). Als je probeerde deze op een moderne grafische kaart (GPU) te gebruiken, zou hij stukgaan.
De "Massaproductie"-Rekenmachine: Deze was geweldig om duizenden berekeningen tegelijkertijd uit te voeren (zoals het simuleren van 1.000 robots tegelijk), maar was traag en onhandig als je gewoon een enkele robot wilde aansturen.

Het artikel introduceert frax, een nieuw hulpmiddel dat fungeert als een universele vertaler voor robotfysica. Het is gebouwd met een technologie genaamd JAX, waardoor het geschreven kan worden in eenvoudige Python-code, maar draait met de rauwe snelheid van hoogwaardige C++-code.

Hier is hoe frax werkt, met behulp van alledaagse analogieën:

1. De "Assemblagelijn" versus de "Teamhuddle"

De meeste oude robotbibliotheken berekenen fysica als een enkele persoon die een huddle doet. Ze vragen: "Oké, gewricht 1 beweegt, dus gewricht 2 beweegt, waardoor gewricht 3 beweegt..." Ze doen dit stap voor stap, één voor één. Dit is snel voor één robot, maar als je het probeert voor 1.000 robots, raakt de lijn verstopt.

frax verandert het spel. In plaats van een huddle, fungeert het als een massale, gesynchroniseerde assemblagelijn. Het berekent de beweging van elk gewricht voor elke robot allemaal tegelijkertijd op exact hetzelfde moment.

De Truc: Het gebruikt een "masker" (een digitale filter) om uit te zoeken welke gewrichten welke anderen beïnvloeden, waardoor de computer alle wiskunde parallel kan uitvoeren.
Het Resultaat: Op één computer is het snel genoeg om een robot in real-time aan te sturen (duizenden keren per seconde). Op een krachtige grafische kaart kan het 100 miljoen robotbewegingen simuleren in één seconde.

2. De "Magische Spiegel" (Automatische Differentiatie)

Een van de moeilijkste delen van robotleren is uitzoeken hoe je de beweging van de robot kunt verbeteren. Meestal moet je complexe wiskundige formules handmatig schrijven om te zien hoe het veranderen van één klein ding het hele systeem beïnvloedt.

frax komt met een ingebouwde "Magische Spiegel". Omdat het is gebouwd op JAX, kan het automatisch kijken naar de beweging van de robot en je direct precies vertellen hoe je het moet aanpassen om betere resultaten te krijgen. Je hoeft geen wiskundig genie te zijn om deze formules af te leiden; de bibliotheek doet het zware werk voor je. Dit maakt het veel gemakkelijker om robots te ontwerpen die kunnen leren en zich kunnen aanpassen.

3. Snelheid en Flexibiliteit

De auteurs hebben frax getest op twee echte robots: een Franka Panda (een robotarm) en een Unitree G1 (een humanoïde robot).

Op een standaardcomputer (CPU): frax is zo snel dat het andere populaire Python-tools met 2 tot 5 keer verslaat. Het komt in de buurt van de snelheid van de allerbeste tools die zijn geschreven in complexe C++-code, maar zonder de hoofdpijn van het schrijven van C++.
Op een grafische kaart (GPU): Het presteert net zo goed als de beste tools die specifiek zijn ontworpen voor grafische kaarten, maar het start veel sneller op en compileert sneller, waardoor ingenieurs tijd besparen wanneer ze nieuwe ideeën testen.

De Conclusie

frax is een nieuwe, open-source bibliotheek die ingenieurs toelaat om robotbesturingssoftware te schrijven in eenvoudige Python. Het is snel genoeg om een enkele robot in real-time aan te sturen op een standaardlaptop, maar krachtig genoeg om duizenden robots tegelijkertijd te simuleren op een supercomputer. Het overbrugt de kloof tussen "gemakkelijk te gebruiken" en "extreem snel", waardoor het gemakkelijker wordt om robots te bouwen die kunnen bewegen, leren en veilig kunnen interageren met de echte wereld.

Technische Samenvatting: frax: Snelle Robotkinematica en Dynamica in JAX

Probleemstelling

Op het gebied van robotbesturing, planning en leren is er een kritieke behoefte aan library's voor dynamica van starre lichamen die tegelijkertijd hoogpresterend, gebruiksvriendelijk en compatibel zijn met diverse hardware-architecturen (CPU's, GPU's en TPUs). Bestaande oplossingen staan voor aanzienlijke afwegingen:

Op CPU geoptimaliseerde library's: Library's zoals Pinocchio en MuJoCo blinken uit in uitvoering met lage latentie op CPU's, maar kunnen niet direct op versnellers worden uitgevoerd. Hun Python-bindingen introduceren vaak overhead en bemoeilijken de integratie met workflows voor automatische differentiatie (AD).
Op GPU geoptimaliseerde library's: Library's zoals MJX en BRAX maken gebruik van GPU-parallelisme voor simulatie met hoge doorvoer, maar lijden vaak onder hoge overhead wanneer ze worden gebruikt voor besturing van een enkele robot op CPU's, waardoor ze ongeschikt zijn voor real-time besturingslussen met een frequentie in de kilohertz.
AD-compatibiliteit: Het integreren van op C++ gebaseerde dynamica-library's in op Python gebaseerde machine learning- of planning-workflows creëert vaak wrijving vanwege de moeilijkheid om automatische differentiatie over taalgrenzen toe te passen.

Er ontbreekt een unificerend framework dat meerdere architecturen target zonder in te leveren op prestaties of gebruiksgemak.

Methodologie

De auteurs introduceren frax, een pure-Python library gebouwd op JAX die een hoogpresterende interface biedt voor robotkinematica en -dynamica. De kernmethodologische innovatie is een volledig vectoriserende aanpak voor dynamica van starre lichamen, die afwijkt van de traditionele methoden voor recursieve boomdoorloping (zoals het Recursieve Newton-Euler-algoritme) die in de meeste C++ library's worden gebruikt.

Belangrijke technische componenten:

Vectoriserende formulering: In plaats van recursieve lussen codeert frax de kinematische boomstructuur met behulp van een vooroudermasker ( $U \in \{0, 1\}^{n \times n}$ $U \in {0, 1}^{n \times n}$ ). Dit maakt het mogelijk dynamische berekeningen uit te voeren via matrixvermenigvuldiging en array-broadcasting.
- Inverse dynamica (RNEA): Berekend door krachten en versnellingen op en neer door de boom op te tellen met behulp van het masker $U$ , waardoor sequentiële afhankelijkheden worden vermeden.
- Massamatrix (CRBA): Berekend door het accumuleren van samengestelde inertie te vectoriseren.
Afweging complexiteit: Deze aanpak verhoogt de theoretische rekencomplexiteit van algoritmen zoals RNEA bewust van $O(n)$ naar $O(n^2)$ en CRBA van $O(nd)$ naar $O(n^2)$ (waarbij $n$ het aantal vrijheidsgraden is). De auteurs betogen echter dat de prestatiewinst door fijnkorrelig parallelisme op moderne hardware (via XLA-compilatie) en het elimineren van lus-overschrijding de redundante berekeningen ruimschoots opwegen.
JAX-integratie: Door gebruik te maken van JAX ondersteunt frax native:
- Just-In-Time (JIT)-compilatie: Voor geoptimaliseerde uitvoering op CPU, GPU en TPU.
- Automatische differentiatie (AD): Het mogelijk maken van efficiënte gradiëntberekening voor op optimalisatie gebaseerde besturing zonder de kosten van het traceren door lange, uitgerolde lussen.
- Jacobian-vectorproducten (JVP's): Het mogelijk maken van efficiënte evaluatie van Lie-afgeleiden voor besturingsbeperkingen (bijvoorbeeld Control Barrier Functions) zonder volledige Jacobianen te berekenen.

Belangrijkste bijdragen

Unificerend architectuur: frax biedt één codebase die efficiënt draait op CPU's, GPU's en TPUs, en zo de kloof overbrugt tussen besturing met lage latentie en simulatie met hoge doorvoer.
Prestaties:
- Op CPU: Bereikt berekeningstijden in de lage microseconden die geschikt zijn voor besturingsfrequenties in de kilohertz, en presteert 2–5x beter dan Python-bindingen van Pinocchio en MuJoCo, en nadert de ruwe prestaties van geoptimaliseerde C++-implementaties.
- Op GPU/TPU: Schaleert tot duizenden parallelle instanties, met meer dan 100 miljoen dynamische evaluaties per seconde, en komt overeen met of overtreft de prestaties van dedicated GPU-simulators zoals MJX en BRAX.
Gebruiksgemak: Als pure-Python library vereenvoudigt het prototyping en integreert het naadloos met het Python-ecosysteem, waardoor de noodzaak van C++-bindingen wordt weggenomen terwijl hoge prestaties worden behouden.
Open source: De library wordt als open-source software uitgebracht, met ondersteuning voor specifieke robots zoals de Franka Panda en Unitree G1, inclusief afgestemde botsingsmodellen.

Resultaten

De auteurs hebben frax gevalideerd op een Franka Panda-manipulator en een Unitree G1-humanoid.

Latentie: Op CPU bereikte frax tijden voor inverse kinematica (differentiële IK) van 4,09 µs (Franka) en 6,71 µs (G1), en tijden voor inverse dynamica (Operationele Ruimte Besturing) van 7,03 µs (Franka) en 24,6 µs (G1).
Vergelijking:
- In vergelijking met Python-API's van Pinocchio en MuJoCo was frax 2–5x sneller door JIT-compilatie en compilerfusies.
- In vergelijking met C++-interfaces van Pinocchio/MuJoCo nadert frax hun ruwe prestaties, maar behoudt het Python-flexibiliteit en AD-mogelijkheden.
- In vergelijking met MJX en BRAX toonde frax vergelijkbare of iets snellere prestaties op GPU, maar aanzienlijk kortere JIT-compilatietijden (1–2 seconden versus 6–12 seconden), wat snellere iteratieve ontwikkeling mogelijk maakt.
Doorvoer: Op GPU demonstreerde de library doorvoersnelheden van meer dan 100 miljoen evaluaties per seconde bij grote batchgroottes (4096), wat wijst op geschiktheid voor massaal parallelle trainingsomgevingen.

Betekenis en claims

Het artikel claimt dat frax een belangrijke stap voorwaarts is naar het unificeren van workflows voor robotbesturing en -leren. De primaire betekenis ligt in het mogelijk maken van real-time besturing van een enkele robot op CPU's, terwijl tegelijkertijd massaal parallelle simulatie op versnellers wordt ondersteund binnen dezelfde codebase.

De auteurs benadrukken dat frax onderzoekers in staat stelt om:

Besturingsystemen te prototypen in Python met automatische differentiatie (bijvoorbeeld voor niet-lineaire trajectoptimalisatie of MPC) zonder in te leveren op de prestaties die nodig zijn voor hardware-implementatie.
Volledige Jacobiaan-berekeningen te vervangen door efficiënte JVP's voor veiligheidskritieke beperkingen.
Dezelfde logica te implementeren op verschillende hardware-backends (CPU voor uitvoering, GPU/TPU voor training) zonder algoritmen opnieuw te schrijven.

Het artikel houdt een bescheiden standpunt, met de erkenning dat frax momenteel minimaler is dan uitgebreide library's zoals Pinocchio (bijvoorbeeld zonder globale IK-oplossers en volledige simulatieomgevingen) en dat C++-implementaties mogelijk nog marginale snellere besturingsfrequenties bieden. Het positioneert frax echter als een flexibele, hoogpresterende basis die specifiek is ontworpen voor de snijvlak van moderne robotbesturing en machine learning.