Comment utiliser NVIDIA Warp et MjWarp pour accélérer les workflows de simulation et d’apprentissage en robotique

Comment utiliser NVIDIA Warp et MjWarp pour accélérer les workflows de simulation et d’apprentissage en robotique

La simulation robotique a un problème de débit. Une tâche riche en contacts, comme la manipulation dans la main ou la locomotion à pattes, oblige un moteur physique à résoudre des milliers de petits problèmes de contraintes par pas de contrôle, et un agent d’apprentissage par renforcement peut avoir besoin de dizaines de millions de ces pas avant qu’une politique converge. Exécuter cette boucle sur un cœur CPU — ou même sur un nombre modeste de cœurs CPU — transforme une idée de recherche en engagement de calcul de plusieurs jours.

NVIDIA Warp et MjWarp s’attaquent à ce goulot d’étranglement dans la même direction : ils déplacent la simulation elle-même sur le GPU, et ils vous permettent d’exécuter de nombreuses copies d’une scène en parallèle plutôt qu’une seule. Warp est la couche sous-jacente de programmation de kernels ; MjWarp est une implémentation accélérée sur GPU du moteur physique MuJoCo construite par-dessus. Cet article détaille ce que fait chaque élément, comment les installer et les vérifier, et comment les intégrer dans un workflow d’entraînement.

Ce que NVIDIA Warp apporte à la simulation robotique

Warp est un framework Python pour écrire du code de simulation et de calcul spatial haute performance qui s’exécute sur les GPU NVIDIA. Plutôt que de vous demander d’écrire du CUDA C++ brut, Warp vous permet de définir des kernels dans une syntaxe proche de Python et de les compiler juste-à-temps en code GPU. De simples fonctions Python décorées comme kernels sont lancées sur un grand nombre de threads, chaque thread étant identifié par un indice et opérant sur des éléments de tableau en parallèle.

Deux propriétés comptent le plus pour la robotique :

• Exécution par lots. Un lancement de kernel peut traiter des centaines de milliers d’éléments indépendants à la fois. Si l’état de votre robot est stocké dans des tableaux — positions, vitesses, couples articulaires — le lot entier avance en un seul lancement.
• Différentiabilité. Warp prend en charge le calcul de gradients à travers ses kernels, ce qui signifie que les pas de simulation peuvent participer à une boucle d’optimisation plutôt que d’être traités comme une boîte noire opaque.

Warp gère aussi la plomberie qui fait généralement dérailler les projets de simulation GPU : allocation de mémoire sur le dispositif, passage de tableaux typé et vérifié entre hôte et dispositif, et requêtes de maillages et volumes creux pour les charges de travail de type collision. Pour une équipe de robotique, cela signifie que les parties personnalisées d’un simulateur — un modèle de contact de pince, un actionneur spécifique à un domaine, une simulation de capteur — peuvent être écrites dans le même langage que le reste de la pile.

Pourquoi MjWarp compte pour les workflows MuJoCo

MuJoCo est le moteur physique par défaut pour l’apprentissage robotique depuis des années, en grande partie grâce à sa gestion précise des contacts et à sa vitesse sur CPU pour la simulation d’un seul environnement. MjWarp réimplémente ce moteur par-dessus Warp afin que la même classe de modèles s’exécute sur GPU avec de nombreuses instances parallèles.

La conséquence pratique est un changement de forme du workflow. Au lieu de faire avancer un environnement dans une boucle Python, vous instanciez un lot d’environnements — par exemple 1 024 ou 4 096 copies du même modèle XML — et vous les faites avancer ensemble. C’est exactement ce que veulent les algorithmes d’apprentissage par renforcement on-policy : ils collectent des rollouts auprès de nombreux agents parallèles, puis mettent à jour une politique à partir du lot agrégé. Quand la simulation et l’apprentissage vivent tous deux sur GPU, l’aller-retour de transfert de données qui les sépare normalement disparaît en grande partie.

Parce que MjWarp vise à préserver la sémantique de modélisation de MuJoCo, les fichiers de modèle MJCF existants restent le point de départ. Vous ne réécrivez pas la description de votre robot ; vous changez le backend d’exécution.

Prérequis

Avant d’installer quoi que ce soit, vérifiez les éléments suivants :

• Un GPU NVIDIA avec une architecture de calcul compatible CUDA. Warp compile les kernels pour le GPU présent sur la machine ; les architectures très anciennes peuvent ne pas être prises en charge.
• Un pilote NVIDIA récent correspondant à la version de CUDA que vous comptez utiliser.
• Python 3.9 ou plus récent, avec pip disponible. Un environnement virtuel est fortement recommandé pour que les versions des paquets n’entrent pas en conflit avec une installation MuJoCo ou PyTorch existante.
• Un environnement Linux pour l’expérience la plus fluide. Les configurations Windows et WSL fonctionnent, mais ont tendance à demander plus d’attention aux chemins des pilotes et de la boîte à outils.
• Le paquet Python MuJoCo, car MjWarp dépend des structures de modèle de MuJoCo et de l’analyse MJCF.

Notez que les noms de paquets, les chemins de modules et les versions minimales de cet écosystème évoluent rapidement. Considérez les commandes ci-dessous comme le schéma d’installation standard et vérifiez les noms actuels dans la documentation amont de Warp et MjWarp avant de figer quoi que ce soit en production.

Installation étape par étape

Commencez par créer et activer un environnement isolé. Cela garde les artefacts compilés de Warp et les versions de MuJoCo séparés de tout autre projet sur la machine.

python3 -m venv ~/warp-robotics
source ~/warp-robotics/bin/activate
python -m pip install --upgrade pip

Installez Warp lui-même. Le nom de distribution PyPI est warp-lang, et il tire l’environnement d’exécution et la chaîne de compilation nécessaires pour construire des kernels pour votre GPU.

pip install warp-lang

Installez MuJoCo. MjWarp s’appuie sur les abstractions de modèle et de données de MuJoCo ; il s’agit donc d’une dépendance requise, et non facultative.

pip install mujoco

Installez le paquet MjWarp. Vérifiez le nom exact de distribution dans le dépôt amont, car c’est le composant le plus susceptible d’être distribué sous un nom légèrement différent ou via une compilation depuis les sources plutôt que depuis PyPI.

pip install mujoco-warp

Si un paquet précompilé n’est pas disponible pour votre version de CUDA, l’alternative consiste à cloner le dépôt et à l’installer en mode éditable afin que les kernels Warp se compilent avec votre boîte à outils locale.

git clone <mjwarp-repository-url>
cd mujoco_warp
pip install -e .

Pour les travaux basés sur les gradients, associez la pile à un framework d’apprentissage profond compatible GPU. Les tenseurs Warp interopèrent avec les bibliothèques de tableaux qui exposent l’interface de tableau CUDA, et PyTorch est le compagnon le plus courant dans les pipelines d’apprentissage robotique.

pip install torch --index-url https://download.pytorch.org/whl/cu124

Ajustez le suffixe CUDA pour qu’il corresponde à votre boîte à outils installée. Une inadéquation ici est la cause la plus fréquente d’une pile qui s’importe proprement mais échoue au premier lancement de kernel.

Vérification de l’installation

Vérifiez toujours Warp avant d’écrire du code de simulation. L’initialisation de Warp affiche des diagnostics sur l’environnement d’exécution, le pilote CUDA et les dispositifs détectés.

import warp as wp

wp.init()
print(wp.get_devices())

Si cela affiche un dispositif CUDA avec un nom et une quantité de mémoire plausibles, le compilateur de kernels fonctionne. Ensuite, vérifiez que MuJoCo peut analyser un modèle et que MjWarp expose ses structures de données GPU.

import mujoco

model = mujoco.MjModel.from_xml_string("<mujoco><worldbody/></mujoco>")
print("nq:", model.nq, "nv:", model.nv)

Un échec à ce stade indique un problème d’installation de MuJoCo plutôt qu’un problème lié à Warp, ce qu’il est utile de savoir avant de déboguer quoi que ce soit de plus complexe.

Exemples d’utilisation

Exemple 1 : un kernel Warp minimal

Le plus petit programme Warp utile est un kernel qui fait avancer des tableaux d’état. Ce schéma sous-tend presque chaque composant de simulateur personnalisé que vous écrirez.

import numpy as np
import warp as wp

wp.init()

@wp.kernel
def integrate(position: wp.array(dtype=wp.vec3),
              velocity: wp.array(dtype=wp.vec3),
              dt: float):
    i = wp.tid()
    position[i] = position[i] + velocity[i] * dt

n = 100_000
pos = wp.array(np.zeros((n, 3), dtype=np.float32), dtype=wp.vec3, device="cuda:0")
vel = wp.array(np.ones((n, 3), dtype=np.float32), dtype=wp.vec3, device="cuda:0")

wp.launch(integrate, dim=n, inputs=[pos, vel, 0.01], device="cuda:0")
wp.synchronize()

L’idée clé est que 100 000 corps indépendants avancent en un seul lancement. Sur CPU, vous feriez une boucle, et la boucle dominerait le temps d’exécution.

Exemple 2 : faire avancer une scène MuJoCo par lots

Le workflow MjWarp remplace un unique MjData par un conteneur de données par lots. Les signatures exactes du constructeur et de la fonction step varient selon la version, donc consultez la référence de l’API — mais la forme du code reste cohérente.

import mujoco
import mujoco_warp as mjw

model = mujoco.MjModel.from_xml_path("humanoid.xml")

# Un objet de données contenant N copies indépendantes du même modèle
data = mjw.Data(model, nworld=1024, device="cuda:0")

for _ in range(1000):
    mjw.step(model, data)

# Les résultats sont des tableaux de forme (nworld, ...)
print(data.qpos.shape)

Mille pas de contrôle sur mille environnements s’exécutent désormais comme du travail GPU. Pour relire l’état en vue d’une mise à jour d’apprentissage, copiez uniquement les tenseurs dont vous avez besoin plutôt que toute la structure de données.

Exemple 3 : randomisation de domaine sur le lot

Comme chaque environnement est une tranche indépendante d’un tableau GPU, la randomisation est une opération en masse plutôt qu’une branche Python par environnement. Réinitialiser un sous-ensemble de mondes lorsqu’ils se terminent devient une écriture indexée.

import torch

# Masque de réinitialisation produit par votre logique d’environnement
reset_mask = torch.rand(1024, device="cuda") < 0.01

# Écrire les états initiaux uniquement dans les environnements qui en ont besoin
initial = torch.zeros((reset_mask.sum(), model.nq), device="cuda")
# ... répartir `initial` dans le qpos par lots aux indices masqués ...

C’est là que le gain d’ingénierie apparaît : le calendrier de randomisation, la logique de terminaison et la physique restent tous sur le dispositif, de sorte qu’aucun blocage de synchronisation n’interrompt la collecte des rollouts.

Exemple 4 : intégrer MjWarp dans une boucle d’entraînement de politique

La boucle d’entraînement elle-même ne change pas conceptuellement. Ce qui change, c’est le coût de l’étape de rollout.

for iteration in range(num_iterations):
    # Rollouts : simulation GPU par lots, aucun aller-retour vers l’hôte
    with torch.no_grad():
        obs = collect_rollouts(model, data, policy, horizon=64)

    # Mise à jour de l’apprentissage : aussi sur GPU
    loss = policy_update(policy, obs)
    loss.backward()
    optimizer.step()

Le goulot d’étranglement classique dans cette boucle est l’écart entre un simulateur CPU produisant des transitions et un apprenant GPU les consommant. Exécuter MjWarp aux côtés de la politique sur le même dispositif réduit cet écart. Pour les réseaux plus petits, le pas de physique peut cesser complètement d’être le facteur limitant.

Exemple 5 : utiliser Warp pour des capteurs personnalisés et des gradients

Là où MjWarp fournit le moteur principal, Warp fournit le point d’extension. Si votre robot a besoin d’un capteur de distance synthétique, d’un câble déformable ou d’un modèle de contact que MuJoCo ne fournit pas, vous l’écrivez comme un kernel Warp qui lit et écrit les mêmes tableaux d’état.

La différentiabilité est le second usage. Le contrôle basé sur l’optimisation et l’identification de système bénéficient tous deux du fait que le pas de simulation fournit des gradients plutôt que d’exiger des différences finies. Les kernels différentiables de Warp rendent cela possible, bien que la portée pratique dépende des opérations de votre pipeline qui sont différentiables et de celles qui ne le sont pas.

Notes pratiques sur les performances et l’ingénierie

Quelques habitudes séparent une pile qui passe à l’échelle d’une pile qui cale :

• Utilisez des lots agressifs. Le débit de simulation sur le GPU augmente avec le nombre de mondes parallèles jusqu’à saturer la bande passante mémoire ou l’occupation. Des lots sous-dimensionnés gaspillent le dispositif.
• Minimisez la synchronisation hôte-dispositif. Chaque copie vers le CPU force le GPU à vider sa file d’attente. Gardez les rollouts sur le dispositif et ne transférez que des statistiques agrégées à l’entraîneur.
• Réutilisez les allocations. Allouer des tableaux dans la boucle de step provoque des allocations répétées et de la fragmentation. Allouez les tampons une fois à la configuration.
• Capturez les séquences de lancement répétées. Pour une séquence fixe de kernels exécutée à chaque pas, la capture de graphe CUDA — que Warp prend en charge via ses utilitaires de capture — supprime le surcoût par lancement. Vérifiez l’API actuelle avant de vous y fier.
• Faites correspondre les dtypes. Float32 partout est généralement le bon choix par défaut ; mélanger les dtypes force silencieusement des conversions.

Considérez ces points comme des conseils d’ingénierie plutôt que comme des affirmations mesurées. Les accélérations réelles dépendent fortement du modèle, de la complexité des contacts, de la taille du lot et du GPU.

Limites et questions ouvertes

Deux mises en garde honnêtes ont leur place dans tout article sur cette pile.

Premièrement, ce n’est pas un rapport de benchmark. Les chiffres publiés pour la physique GPU sont notoirement sensibles au scénario, et quiconque cite un multiplicateur unique sans préciser la taille du lot, le modèle et le matériel vous dit très peu de choses.

Deuxièmement, MjWarp suit la sémantique de MuJoCo mais est une implémentation distincte avec son propre rythme de publication. La couverture des fonctionnalités, le comportement dans les cas limites et la stabilité de l’API peuvent différer du moteur CPU. Si votre workflow dépend d’une option de solveur inhabituelle ou d’une fonctionnalité MJCF rarement utilisée, validez-la par rapport à MuJoCo CPU avant d’y consacrer une exécution d’entraînement. Le déterminisme entre exécutions GPU est un autre aspect qu’il vaut la peine de tester explicitement pour votre propre modèle, car les réductions parallèles peuvent ne pas être bit à bit identiques au chemin CPU.

Conclusion

NVIDIA Warp et MjWarp s’attaquent au même goulot d’étranglement structurel de la recherche en robotique sous deux angles complémentaires. Warp vous donne un moyen d’écrire du code de simulation parallèle sur GPU — y compris de la physique personnalisée, des capteurs et des composants différentiables — dans quelque chose de proche de Python. MjWarp vous donne un moteur MuJoCo parallèle sur GPU afin que les modèles MJCF existants puissent s’exécuter en grands lots sans réécrire la description du robot.

Le parcours d’installation est court : créez un environnement, installez warp-lang, installez mujoco, installez le paquet MjWarp, vérifiez la liste des dispositifs, puis commencez avec une scène par lots plutôt qu’une scène unique. Le changement de workflow qui compte le plus est de traiter le lot comme l’unité de simulation, et non l’environnement individuel. Une fois que les rollouts, la randomisation et les mises à jour de politique vivent tous sur le même dispositif, l’aller-retour vers l’hôte qui limite traditionnellement le débit disparaît, et la simulation cesse d’être la contrainte de temps d’horloge de vos expériences.

Pour le tutoriel d’origine et les détails d’API les plus récents, consultez l’article de blog NVIDIA sur Hugging Face : https://huggingface.co/blog/nvidia/how-to-use-nvidia-warp-and-mjwarp. Comme les noms de paquets et les signatures de cet écosystème changent fréquemment, vérifiez les commandes d’installation et les chemins de modules dans la documentation amont avant de les figer dans un pipeline de production.