MuJoCo Warp documentation: Does MJWarp work with multiple GPUs?.
PiperOrigin-RevId: 836642204 Change-Id: I559642b2cb902a704257e7e0344f768b690678a1
This commit is contained in:
committed by
Copybara-Service
parent
d8712808b9
commit
893882171e
@@ -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
|
||||
--------------------------
|
||||
|
||||
|
||||
Reference in New Issue
Block a user