Introduzione

MuJoCo classico fornisce una simulazione robotica veloce basata su CPU, ideale per sviluppare, testare e controllare robot con la possibilità di parallelizzare il campionamento sui core della CPU. Con l’aumento dei carichi di lavoro di apprendimento, la domanda si sposta da “quanto velocemente può girare un singolo mondo” a “quanti mondi possono girare simultaneamente”. L’accelerazione su GPU rende possibile avanzare molti mondi in grandi batch mantenendo i dati di simulazione e di apprendimento vicini al dispositivo.

MuJoCo Warp (MJWarp), costruito sopra NVIDIA Warp, trasforma i modelli MuJoCo compatibili in un regime di scala GPU. Nel presente articolo spostiamo un braccio SO‑101 follower da un tipico flusso di lavoro MuJoCo a un batch di fino a 2.048 ambienti MJWarp, illustrando la tecnologia e i passaggi di validazione che rendono possibile la transizione.

Che cos'è NVIDIA Warp

NVIDIA Warp è un framework Python per scrivere kernel ad alte prestazioni accelerati su GPU. Warp consente agli sviluppatori di authorare kernel tipizzati staticamente in Python, compilandoli per esecuzione su CPU o CUDA. Il primo lancio costruisce e memorizza nella cache un modulo nativo; i lanci successivi lo riutilizzano. Il linguaggio dei kernel è una sotto‑insieme orientata alle prestazioni di Python, mentre il Python “normale” rimane responsabile della configurazione, dell’allocazione e dell’orchestrazione dei lanci.

Un piccolo kernel orientato alla robotica sposta le posizioni dei punti sotto gravità. Un thread logico gestisce un punto, così lo stesso codice scala da due punti a milioni senza introdurre concetti di GPU nel flusso di controllo.

Due capacità aggiuntive di Warp sono importanti anche se non usate direttamente nel workflow SO‑101:

Differenziabilità: un wp.Tape registra i lanci di kernel forward all’interno del suo contesto e ne riproduce gli adjoint al contrario quando viene chiamato backward(). Questo permette di costruire geometrie differenziabili, CFD e fisica custom in Warp, incluse pipeline CAE per ottimizzazione di progetto.
Esecuzione deterministica: introdotta in Warp 1.15, consente di ottenere ordine riproducibile delle operazioni atomiche GPU, scambiando una parte di performance per risultati consistenti in test di regressione e validazione.

Queste caratteristiche sono intrinseche a Warp, non garanzie di differenziabilità o determinismo per un’intera distribuzione MJWarp.

Per provare Warp:

pip install warp-lang (≥ 1.15 per la determinismo GPU)
python -m warp.examples.browse
consultare i notebook tutorial

Migrazione del modello SO‑101 a MJWarp

Un simulatore robotico calcola ripetutamente “cosa succederà dopo”: dato lo stato corrente di giunti, velocità, controlli e contatti, avanza la scena di un piccolo intervallo di tempo. In questo contesto, un mondo è una copia indipendente di tale scena e del suo stato. Un mondo può contenere il braccio SO‑101 che cerca di afferrare un cubo; un altro può contenere lo stesso braccio che parte da una posa leggermente diversa.

MuJoCo e MJWarp possono eseguire lo stesso robot e lo stesso compito, ma organizzano il lavoro diversamente. MuJoCo è naturale per sviluppare e ispezionare uno o pochi mondi su CPU. MJWarp è un’implementazione in NVIDIA Warp della pipeline fisica di MuJoCo, che colloca il modello e un batch di stati indipendenti sulle GPU NVIDIA; una chiamata a mjw.step avanza l’intero batch.

Il valore di MJWarp non è necessariamente una simulazione più veloce per un singolo mondo, ma la capacità di avanzare centinaia o migliaia di mondi insieme, fornendo alla GPU sufficiente lavoro parallelo per migliorare la throughput aggregata, cioè il numero totale di passi di mondo completati al secondo. Questo è vantaggioso per il reinforcement learning e il campionamento su larga scala, dove la raccolta di esperienza è più importante della latenza di un singolo ambiente.

Configurazione dei parametri di batch

Per allocare risorse batch è necessario definire i seguenti parametri (riferirsi a “Batch sizes”):

CUDA graph capture: mjw.step genera molteplici lanci di kernel; catturare il grafo una volta e riprodurlo più volte riduce l’overhead di lancio.
Dimensioni di nconmax / naconmax / njmax: questi valori determinano la memoria e il lavoro richiesti; è consigliabile scalarli strettamente al carico previsto. Si può sintonizzare con mjwarp-testspeed --measure_alloc e monitorare eventuali overflow in mjwarp-viewer.

Altre considerazioni di tuning includono:

Testare i limiti di iterazione del solver senza alterare il comportamento del compito.
I settaggi di mesh e CCD possono aumentare l’uso di memoria; ridurre nccdmax e naccdmax quando il conteggio dei contatti misurato lo permette.
Il solver compatto di MJWarp utilizza il Newton constraint solver di MuJoCo e la modalità “sleeping”, non il framework Newton physics‑engine separato.

Configurazioni multi‑GPU o ottimizzazioni avanzate del solver sono fuori dallo scopo di questa guida; per approfondire si consulti la documentazione di MJWarp performance‑tuning.

Installazione e prova rapida

Per provare MJWarp eseguire:

pip install mujoco-warp
mjwarp-viewer path/to/scene.xml
seguire il tutorial Colab fornito nel repository ufficiale

Il file di scena è un MJCF standard contenente un braccio SO‑101, un tavolo e due cubi da impilare. Non è specifico di MJWarp, quindi può essere caricato sia da MuJoCo CPU sia da MJWarp senza modifiche.

Nel repository di accompagnamento la scena è generata automaticamente: resolvepickplace_scene() copia il braccio Menagerie