MuJoCo Warp documentation: Does MJWarp work with multiple GPUs?.

PiperOrigin-RevId: 836642204
Change-Id: I559642b2cb902a704257e7e0344f768b690678a1
This commit is contained in:
Taylor Howell
2025-11-25 06:25:49 -08:00
committed by Copybara-Service
parent d8712808b9
commit 893882171e
+26
View File
@@ -383,6 +383,32 @@ Warp's `automatic differentiation <https://nvidia.github.io/warp/modules/differe
functionality. Updates from the team related to enabling automatic differentiation for MJWarp are tracked in this
`GitHub issue <https://github.com/google-deepmind/mujoco_warp/issues/500>`__.
**Does MJWarp work with multiple GPUs?**
Yes. Warp's ``wp.ScopedDevice`` enables multi-GPU computation
.. code-block:: python
# create a graph for each device
graph = {}
for device in wp.get_cuda_devices():
with wp.ScopedDevice(device):
m = mjw.put_model(mjm)
d = mjw.make_data(mjm)
with wp.ScopedCapture(device) as capture:
mjw.step(m, d)
graph[device] = capture.graph
# launch a graph on each device
for device in wp.get_cuda_devices():
wp.capture_launch(graph[device])
Please see the
`Warp documentation <https://nvidia.github.io/modules/devices.html#example-using-wp-scopeddevice-with-multiple-gpus>`__
for details and
`mjlab distributed training <https://github.com/mujocolab/mjlab/tree/main/docs/api/distributed_training.md>`__ for a
reinforcement learning example.
Orientation representation
--------------------------