NVIDIA Warp und MJWarp: Roboter-Simulation im GPU-Maßstab

Für Reinforcement Learning in der Robotik ist weniger die Frage, wie schnell eine Simulationswelt läuft - sondern wie viele gleichzeitig. NVIDIA hat auf dem HuggingFace-Blog den zweiten Teil seiner Serie “State of Simulation for Physical AI” veröffentlicht und zeigt darin den Weg von einem einzelnen MuJoCo-Modell bis zu 2.048 parallelen Simulationen auf der GPU.

Die Bausteine

Die Basis ist NVIDIA Warp, ein Python-Framework für GPU-Kernel: Entwickler schreiben statisch getypte Kernel in einer Python-Teilmenge, die per JIT zu nativem CUDA kompiliert werden. Ein logischer Thread bearbeitet dabei einen Punkt, einen Contact oder eine ganze Welt - derselbe Code skaliert von wenigen Elementen bis zu Millionen. Warp-Kernel sind zudem differenzierbar (aufgezeichnete Vorwärts-Läufe werden per wp.Tape rückwärts abgespielt) und seit Version 1.15 optional deterministisch, was reproduzierbare Regressionstests ermöglicht.

Darauf setzt MuJoCo Warp (MJWarp) auf - eine Implementierung der MuJoCo-Physik-Pipeline auf Warp. Statt einzelne CPU-Welten zu advancieren, liegt ein Batch unabhängiger Zustände auf der GPU. Ein einziger Aufruf von mjw.step verarbeitet den kompletten Batch.

Latenz vs. Durchsatz

Der entscheidende Perspektivwechsel im Artikel: MJWarp optimiert nicht die Latenz eines einzelnen Simulationsschritts, sondern die Aggregat-Durchsatzrate (aggregate throughput) - die Gesamtzahl abgeschlossener World-Steps pro Sekunde. Für MPC oder Teleoperation eines einzelnen Roboters bleibt klassisches CPU-MuJoCo die richtige Wahl. Für RL-Workflows, die massiv Erfahrung sammeln müssen, zahlt sich das GPU-Batching aus.

Der Migrationspfad

Am Beispiel eines SO-101-Follower-Arms beschreibt der Beitrag den pragmatischen Ablauf: eine MuJoCo-Welt validieren, per mjw.put_model und mjw.put_data auf die GPU überführen, den Batch bilden und anschließend Parallelität und Korrektheit messen. Die Autoren betonen den Validierungsschritt zwischen jedem Stadium.

Als Faustregel geben sie mit: MuJoCo CPU für Einzelroboter, MJWarp oder mjlab für maximalen Physik-Durchsatz, MuJoCo Playground/MJX für JAX-basierte Training-Rezepte und Newton für Multi-Solver-Integrationen mit Isaac Lab - dazu folgt ein eigener Beitrag.

Mehr dazu: HuggingFace Blog