💻 computer science

frax: Fast Robot Kinematics and Dynamics in JAX

Dieser Beitrag stellt frax vor, eine hochleistungsfähige, reine Python-basierte JAX-Bibliothek für Robotik-Kinematik und -Dynamik, die eine CPU-Ausführung mit niedriger Latenz und massive GPU-Parallelisierung bietet und gleichzeitig automatische Differentiation für optimierungsbasierte Regelung und Lernen unterstützt.

Ursprüngliche Autoren: Daniel Morton, Marco Pavone

Veröffentlicht 2026-05-18

📖 4 Min. Lesezeit☕ Kaffeepausen-Lektüre

CC BY 4.0

Ursprüngliche Autoren: Daniel Morton, Marco Pavone

Originalarbeit lizenziert unter CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Dies ist eine KI-generierte Erklärung des untenstehenden Papers. Sie wurde nicht von den Autoren verfasst oder gebilligt. Für technische Genauigkeit konsultieren Sie das Originalpaper. Vollständigen Haftungsausschluss lesen

Stellen Sie sich vor, Sie versuchen einem Roboter beizubringen, wie er seinen Arm bewegt, geht oder einen Ball fängt. Dafür benötigen Sie ein „Gehirn", das die Physik sofort berechnen kann: wie schwer der Arm ist, wie schnell er schwingt und genau wie viel Kraft benötigt wird, um ihn zu stoppen oder seine Richtung zu ändern. Dies wird als Roboter-Dynamik bezeichnet.

Lange Zeit waren die Werkzeuge, die Ingenieuren für diese Berechnungen zur Verfügung standen, wie zwei verschiedene Arten von Taschenrechnern:

Der „Schnelle" Rechner: Er war unglaublich schnell, funktionierte aber nur auf einem bestimmten Typ von Computerchip (CPU). Versuchte man, ihn auf einer modernen Grafikkarte (GPU) zu verwenden, ging er kaputt.
Der „Massenproduktions"-Rechner: Er war hervorragend darin, Tausende von Berechnungen gleichzeitig durchzuführen (wie die Simulation von 1.000 Robotern zur gleichen Zeit), war aber langsam und sperrig, wenn man nur einen einzelnen Roboter steuern wollte.

Die Arbeit stellt frax vor, ein neues Werkzeug, das wie ein universeller Übersetzer für Roboterphysik fungiert. Es wurde mit einer Technologie namens JAX entwickelt, die es ermöglicht, es in einfachem Python-Code zu schreiben, aber mit der rohen Geschwindigkeit von hochleistungsfähigem C++-Code auszuführen.

So funktioniert frax, unter Verwendung einiger alltäglicher Analogien:

1. Die „Fließband" vs. das „Team-Huddle"

Die meisten alten Roboter-Bibliotheken berechnen die Physik wie eine einzelne Person, die ein Huddle macht. Sie fragen: „Okay, Gelenk 1 bewegt sich, also bewegt sich Gelenk 2, was Gelenk 3 bewegt..." Sie tun dies schrittweise, nacheinander. Dies ist schnell für einen Roboter, aber wenn man es für 1.000 Roboter versucht, staut sich die Schlange.

frax ändert das Spiel. Statt eines Huddles fungiert es wie eine massive, synchronisierte Fließband. Es berechnet die Bewegung jedes Gelenks für jeden Roboter alle zur exakt gleichen Zeit.

Der Trick: Es verwendet eine „Maske" (ein digitaler Filter), um herauszufinden, welche Gelenke welche anderen beeinflussen, wodurch der Computer alle Mathematik parallel durchführen kann.
Das Ergebnis: Auf einem einzelnen Computer ist es schnell genug, um einen Roboter in Echtzeit zu steuern (Tausende Male pro Sekunde). Auf einer leistungsstarken Grafikkarte kann es 100 Millionen Roboterbewegungen in einer einzigen Sekunde simulieren.

2. Der „Magische Spiegel" (Automatische Differentiation)

Einer der schwierigsten Teile des Roboterlernens besteht darin herauszufinden, wie man die Bewegung des Roboters verbessern kann. Normalerweise muss man komplexe mathematische Formeln manuell schreiben, um zu sehen, wie die Änderung einer winzigen Sache das gesamte System beeinflusst.

frax kommt mit einem eingebauten „Magischen Spiegel". Da es auf JAX basiert, kann es automatisch die Bewegung des Roboters betrachten und Ihnen sofort genau sagen, wie man sie anpassen muss, um bessere Ergebnisse zu erzielen. Sie müssen kein Mathe-Genie sein, um diese Formeln herzuleiten; die Bibliothek übernimmt die schwere Arbeit für Sie. Dies macht es viel einfacher, Roboter zu entwerfen, die lernen und sich anpassen können.

3. Geschwindigkeit und Flexibilität

Die Autoren testeten frax an zwei echten Robotern: einem Franka Panda (ein Roboterarm) und einem Unitree G1 (ein humanoider Roboter).

Auf einem Standardcomputer (CPU): frax ist so schnell, dass es andere beliebte Python-Tools um das 2- bis 5-fache schlägt. Es kommt der Geschwindigkeit der allerbesten Tools nahe, die in komplexem C++-Code geschrieben sind, aber ohne die Kopfschmerzen, C++ zu schreiben.
Auf einer Grafikkarte (GPU): Es leistet genauso gut wie die besten Tools, die speziell für Grafikkarten entwickelt wurden, startet und kompiliert jedoch viel schneller, was Ingenieuren Zeit spart, wenn sie neue Ideen testen.

Das Fazit

frax ist eine neue, quelloffene Bibliothek, die Ingenieuren ermöglicht, Robotersteuerungssoftware in einfachem Python zu schreiben. Es ist schnell genug, um einen einzelnen Roboter in Echtzeit auf einem Standard-Laptop zu steuern, aber leistungsstark genug, um Tausende von Robotern gleichzeitig auf einem Supercomputer zu simulieren. Es schließt die Lücke zwischen „einfach zu bedienen" und „extrem schnell" und macht es einfacher, Roboter zu bauen, die sich bewegen, lernen und sicher mit der realen Welt interagieren können.

Technische Zusammenfassung: frax: Schnelle Roboterkinematische und Dynamik in JAX

Problemstellung

In den Bereichen Robotiksteuerung, Planung und Lernen besteht ein kritischer Bedarf an Bibliotheken für die Starrkörperdynamik, die gleichzeitig hochleistungsfähig, einfach zu bedienen und mit verschiedenen Hardwarearchitekturen (CPUs, GPUs und TPUs) kompatibel sind. Bestehende Lösungen weisen erhebliche Kompromisse auf:

Für CPUs optimierte Bibliotheken: Bibliotheken wie Pinocchio und MuJoCo zeichnen sich durch eine geringe Latenz auf CPUs aus, können jedoch nicht direkt auf Beschleunigern ausgeführt werden. Ihre Python-Bindings führen oft zu Overhead und erschweren die Integration in Workflows für automatische Differentiation (AD).
Für GPUs optimierte Bibliotheken: Bibliotheken wie MJX und BRAX nutzen die GPU-Parallelität für simulations mit hohem Durchsatz, leiden jedoch oft unter hohem Overhead, wenn sie für die Steuerung einzelner Roboter auf CPUs verwendet werden, was sie für Echtzeit-Steuerkreise im Kilohertz-Bereich ungeeignet macht.
AD-Kompatibilität: Die Integration von auf C++ basierenden Dynamikbibliotheken in Python-basierte Machine-Learning- oder Planungsworkflows führt häufig zu Reibungsverlusten aufgrund der Schwierigkeit, automatische Differentiation über Sprachgrenzen hinweg anzuwenden.

Es fehlt ein einheitliches Framework, das mehrere Architekturen anspricht, ohne Leistung oder Benutzerfreundlichkeit zu beeinträchtigen.

Methodik

Die Autoren stellen frax vor, eine rein in Python geschriebene Bibliothek, die auf JAX basiert und eine hochleistungsfähige Schnittstelle für Roboterkinematische und Dynamik bietet. Die zentrale methodische Innovation ist ein vollständig vektorisierter Ansatz für die Starrkörperdynamik, der sich von den traditionellen rekursiven Baum-Durchlaufmethoden (z. B. der rekursive Newton-Euler-Algorithmus) unterscheidet, wie sie in den meisten C++-Bibliotheken verwendet werden.

Wichtige technische Komponenten:

Vektorisierte Formulierung: Anstelle rekursiver Schleifen kodiert frax die kinematische Baumstruktur mithilfe einer Ahnenmaske ( $U \in \{0, 1\}^{n \times n}$ $U \in {0, 1}^{n \times n}$ ). Dies ermöglicht die Berechnung der Dynamik durch Matrixmultiplikation und Array-Broadcasting.
- Inverse Dynamik (RNEA): Berechnet durch Summation von Kräften und Beschleunigungen auf- und abwärts im Baum unter Verwendung der Maske $U$ , wodurch sequenzielle Abhängigkeiten vermieden werden.
- Massenmatrix (CRBA): Berechnet durch Vektorisierung der Akkumulation zusammengesetzter Trägheiten.
Komplexitätskompromiss: Dieser Ansatz erhöht absichtlich die theoretische Rechenkomplexität von Algorithmen wie RNEA von $O(n)$ auf $O(n^2)$ und von CRBA von $O(nd)$ auf $O(n^2)$ (wobei $n$ die Freiheitsgrade sind). Die Autoren argumentieren jedoch, dass die Leistungsgewinne durch feingranulare Parallelität auf moderner Hardware (via XLA-Kompilierung) und die Beseitigung von Schleifen-Overhead die redundanten Berechnungen weitgehend kompensieren.
JAX-Integration: Durch die Nutzung von JAX unterstützt frax nativ:
- Just-In-Time (JIT)-Kompilierung: Für optimierte Ausführung auf CPU, GPU und TPU.
- Automatische Differentiation (AD): Ermöglicht effiziente Gradientenberechnung für optimierungsbasierte Steuerung ohne die Kosten des Durchlaufens langer, entrollter Schleifen.
- Jacobian-Vector-Produkte (JVPs): Erlaubt die effiziente Auswertung von Lie-Ableitungen für Steuerungsbeschränkungen (z. B. Control Barrier Functions) ohne die Berechnung vollständiger Jacobi-Matrizen.

Wichtige Beiträge

Einheitliche Architektur: frax bietet eine einzige Codebasis, die effizient auf CPUs, GPUs und TPUs läuft und die Lücke zwischen latenzarmer Steuerung und simulations mit hohem Durchsatz schließt.
Leistung:
- Auf CPU: Erreicht Berechnungszeiten im niedrigen Mikrosekundenbereich, die für Steuerungen im Kilohertz-Bereich geeignet sind, und übertrifft die Python-Bindings von Pinocchio und MuJoCo um das 2- bis 5-fache, wobei sie sich der Rohleistung optimierter C++-Implementierungen annähern.
- Auf GPU/TPU: Skaliert auf Tausende paralleler Instanzen und erreicht über 100 Millionen Dynamikbewertungen pro Sekunde, was die Leistung dedizierter GPU-Simulatoren wie MJX und BRAX erreicht oder übertrifft.
Benutzerfreundlichkeit: Als rein in Python geschriebene Bibliothek vereinfacht sie die Prototypenerstellung und integriert sich nahtlos in das Python-Ökosystem, wodurch der Bedarf an C++-Bindings entfällt, während die hohe Leistung erhalten bleibt.
Open Source: Die Bibliothek wird als Open-Source-Software veröffentlicht, mit Unterstützung für spezifische Roboter wie den Franka Panda und Unitree G1, einschließlich abgestimmter Kollisionsmodelle.

Ergebnisse

Die Autoren validierten frax an einem Franka Panda-Manipulator und einem Unitree G1-Humanoiden.

Latenz: Auf CPU erreichte frax Zeiten für inverse Kinematik (differenzielle IK) von 4,09 µs (Franka) und 6,71 µs (G1) sowie Zeiten für inverse Dynamik (Steuerung im Operationsraum) von 7,03 µs (Franka) und 24,6 µs (G1).
Vergleich:
- Im Vergleich zu den Python-APIs von Pinocchio und MuJoCo war frax 2–5-mal schneller aufgrund von JIT-Kompilierung und Compiler-Fusionen.
- Im Vergleich zu den C++-Schnittstellen von Pinocchio/MuJoCo nähert sich frax deren Rohleistung an, behält jedoch die Flexibilität von Python und AD-Fähigkeiten bei.
- Im Vergleich zu MJX und BRAX zeigte frax auf der GPU eine vergleichbare oder leicht schnellere Leistung, jedoch mit deutlich reduzierten JIT-Kompilierungszeiten (1–2 Sekunden vs. 6–12 Sekunden), was eine schnellere iterative Entwicklung ermöglicht.
Durchsatz: Auf GPU demonstrierte die Bibliothek Durchsatzraten von über 100 Millionen Bewertungen pro Sekunde bei großen Batch-Größen (4096), was ihre Eignung für massiv parallele Trainingsumgebungen unterstreicht.

Bedeutung und Behauptungen

Die Arbeit behauptet, dass frax einen bedeutenden Schritt hin zu einer Vereinheitlichung von Robotiksteuerungs- und Lernworkflows darstellt. Ihre primäre Bedeutung liegt darin, Echtzeit-Steuerung einzelner Roboter auf CPUs zu ermöglichen und gleichzeitig massiv parallele Simulation auf Beschleunigern innerhalb derselben Codebasis zu unterstützen.

Die Autoren betonen, dass frax Forschern ermöglicht:

Controller in Python mit automatischer Differentiation zu prototypisieren (z. B. für nichtlineare Trajektorienoptimierung oder MPC), ohne die für den Hardware-Einsatz erforderliche Leistung zu opfern.
Vollständige Jacobi-Berechnungen durch effiziente JVPs für sicherheitskritische Beschränkungen zu ersetzen.
Dieselbe Logik über verschiedene Hardware-Backends hinweg einzusetzen (CPU für die Ausführung, GPU/TPU für das Training), ohne Algorithmen neu schreiben zu müssen.

Die Arbeit nimmt eine bescheidene Haltung ein und räumt ein, dass frax derzeit minimaler ist als umfassende Bibliotheken wie Pinocchio (z. B. fehlen globale IK-Löser und vollständige Simulationsumgebungen) und dass C++-Implementierungen möglicherweise noch marginale schnellere Steuerungszeiten bieten. Dennoch positioniert sie frax als flexible, hochleistungsfähige Grundlage, die speziell für die Schnittstelle moderner Robotiksteuerung und Machine Learning entwickelt wurde.