🔨 enable multi-GPU support using custom PPO

duburcqa · Apr 27, 2024 · 3cdc6a6 · 3cdc6a6
1 parent daeb266
commit 3cdc6a6
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/python/gym_jiminy/rllib/gym_jiminy/rllib/ppo.py b/python/gym_jiminy/rllib/gym_jiminy/rllib/ppo.py
@@ -155,6 +155,10 @@ def _compute_mirrored_value(value: torch.Tensor,
     """Compute mirrored value from observation space based on provided
     mirroring transformation.
     """
+    # Make sure value and mirror_mat are on the same device.
+    # This is needed for multi-GPU training.
+    mirror_mat = mirror_mat.to(value.device)
+
     def _update_flattened_slice(data: torch.Tensor,
                                 shape: Tuple[int, ...],
                                 mirror_mat: torch.Tensor) -> torch.Tensor: