[BugFix] Fix multi-gpu setup documentation and future proof some code

docs/source/user_guide/getting_started/quickstart.md

@@ -161,19 +161,9 @@ This will then open up a GUI that looks like so:
 
 ### Additional GPU simulation/rendering customization
 
-Finally on servers with multiple GPUs you can directly pick which devices/backends to use for simulation and rendering.
+Finally on servers with multiple GPUs you can directly pick which devices/backends to use for simulation and rendering by setting the `CUDA_VISIBLE_DEVICES` environment variable. You can do this by e.g. running `export CUDA_VISIBLE_DEVICES=1` and then run the same code. While everything is labeled as device "cuda:0" it is actually using GPU device 1 now, which you can verify by running `nvidia-smi`.
 
-```python
-import gymnasium as gym
-import mani_skill.envs
-
-env = gym.make(
-    "PickCube-v1",
-    num_envs=16,
-    sim_backend="cuda:1", # selects the GPU with index 1
-    render_backend="cuda", # auto selects a GPU
-)
-```
+We currently do not properly support exposing multiple visible CUDA devices to a single process as it has some rendering bugs at the moment.
 
 ## Task Instantiation Options
 

mani_skill/agents/robots/allegro_hand/allegro.py

@@ -159,15 +159,17 @@ def tip_poses(self):
         """
         Get the tip pose for each of the finger, four fingers in total
         """
-        tip_poses = [vectorize_pose(link.pose) for link in self.tip_links]
+        tip_poses = [
+            vectorize_pose(link.pose, device=self.device) for link in self.tip_links
+        ]
         return torch.stack(tip_poses, dim=-2)
 
     @property
     def palm_pose(self):
         """
         Get the palm pose for allegro hand
         """
-        return vectorize_pose(self.palm_link.pose)
+        return vectorize_pose(self.palm_link.pose, device=self.device)
 
 
 @register_agent()

mani_skill/agents/robots/dclaw/dclaw.py

@@ -117,5 +117,7 @@ def tip_poses(self):
         """
         Get the tip pose for each of the finger, three fingers in total
         """
-        tip_poses = [vectorize_pose(link.pose) for link in self.tip_links]
+        tip_poses = [
+            vectorize_pose(link.pose, device=self.device) for link in self.tip_links
+        ]
         return torch.stack(tip_poses, dim=-2)

mani_skill/agents/robots/trifingerpro/trifingerpro.py

@@ -8,12 +8,8 @@
 from mani_skill.agents.base_agent import BaseAgent
 from mani_skill.agents.controllers import *
 from mani_skill.agents.registration import register_agent
-from mani_skill.agents.utils import (
-    get_active_joint_indices,
-)
-from mani_skill.utils.sapien_utils import (
-    get_objs_by_names,
-)
+from mani_skill.agents.utils import get_active_joint_indices
+from mani_skill.utils.sapien_utils import get_objs_by_names
 from mani_skill.utils.structs.pose import vectorize_pose
 
 
@@ -23,6 +19,7 @@ class TriFingerPro(BaseAgent):
     Modified from https://github.com/NVIDIA-Omniverse/IsaacGymEnvs/blob/main/isaacgymenvs/tasks/trifinger.py
 
     """
+
     uid = "trifingerpro"
     urdf_path = f"{PACKAGE_ASSET_DIR}/robots/trifinger/trifingerpro.urdf"
     urdf_config = dict(
@@ -66,8 +63,16 @@ def __init__(self, *args, **kwargs):
         self.joint_stiffness = 1e2
         self.joint_damping = 1e1
         self.joint_force_limit = 2e1
-        self.tip_link_names = ["finger_tip_link_0", "finger_tip_link_120", "finger_tip_link_240"]
-        self.root_joint_names = ["finger_base_to_upper_joint_0", "finger_base_to_upper_joint_120", "finger_base_to_upper_joint_240"]
+        self.tip_link_names = [
+            "finger_tip_link_0",
+            "finger_tip_link_120",
+            "finger_tip_link_240",
+        ]
+        self.root_joint_names = [
+            "finger_base_to_upper_joint_0",
+            "finger_base_to_upper_joint_120",
+            "finger_base_to_upper_joint_240",
+        ]
 
         super().__init__(*args, **kwargs)
 
@@ -164,7 +169,9 @@ def tip_poses(self):
         """
         Get the tip pose for each of the finger, three fingers in total
         """
-        tip_poses = [vectorize_pose(link.pose) for link in self.tip_links]
+        tip_poses = [
+            vectorize_pose(link.pose, device=self.device) for link in self.tip_links
+        ]
         return torch.stack(tip_poses, dim=-1)
 
     # @property

mani_skill/envs/sapien_env.py

@@ -2,7 +2,7 @@
 import gc
 import os
 from functools import cached_property
-from typing import Any, Dict, List, Sequence, Tuple, Union
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
 
 import dacite
 import gymnasium as gym
@@ -423,7 +423,7 @@ def elapsed_steps(self):
     def obs_mode(self):
         return self._obs_mode
 
-    def get_obs(self, info: Dict = None):
+    def get_obs(self, info: Optional[Dict] = None):
         """
         Return the current observation of the environment. User may call this directly to get the current observation
         as opposed to taking a step with actions in the environment.
@@ -581,7 +581,6 @@ def _reconfigure(self, options = dict()):
         self._setup_scene()
         self._load_agent(options)
         self._load_scene(options)
-
         self._load_lighting(options)
 
         if physx.is_gpu_enabled():

mani_skill/utils/common.py

@@ -3,7 +3,7 @@
 """
 
 from collections import defaultdict
-from typing import Dict, Sequence, Tuple, Union
+from typing import Dict, Optional, Sequence, Tuple, Union
 
 import gymnasium as gym
 import numpy as np
@@ -124,7 +124,9 @@ def index_dict_array(x1, idx: Union[int, slice], inplace=True):
 
 
 # TODO (stao): this code can be simplified
-def to_tensor(array: Union[torch.Tensor, np.array, Sequence], device: Device = None):
+def to_tensor(
+    array: Union[torch.Tensor, np.array, Sequence], device: Optional[Device] = None
+):
     """
     Maps any given sequence to a torch tensor on the CPU/GPU. If physx gpu is not enabled then we use CPU, otherwise GPU, unless specified
     by the device argument
@@ -149,7 +151,11 @@ def to_tensor(array: Union[torch.Tensor, np.array, Sequence], device: Device = N
         else:
             ret = torch.tensor(array)
         if device is None:
-            return ret.cuda()
+            if ret.device.type == "cpu":
+                # TODO (stao): note that .cuda does move a tensor to the torch.device context, it moves to the default
+                return ret.cuda()
+            # keep same device if already on GPU
+            return ret
         else:
             return ret.to(device)
     else:

mani_skill/utils/structs/actor.py

@@ -327,7 +327,7 @@ def pose(self) -> Pose:
     def pose(self, arg1: Union[Pose, sapien.Pose, Array]) -> None:
         if physx.is_gpu_enabled():
             if not isinstance(arg1, torch.Tensor):
-                arg1 = vectorize_pose(arg1)
+                arg1 = vectorize_pose(arg1, device=self.device)
             if self.hidden:
                 self.before_hide_pose[self.scene._reset_mask[self._scene_idxs]] = arg1
                 return

mani_skill/utils/structs/articulation.py

@@ -504,7 +504,7 @@ def set_root_pose(self, pose: sapien.Pose) -> None:
 
     @cached_property
     def dof(self) -> torch.tensor:
-        return torch.tensor([obj.dof for obj in self._objs])
+        return torch.tensor([obj.dof for obj in self._objs], device=self.device)
 
     # @property
     # def gpu_index(self) -> int:
@@ -580,10 +580,7 @@ def qlimits(self):
             ]
         )
         padded_qlimits = torch.from_numpy(padded_qlimits).float()
-        if physx.is_gpu_enabled():
-            return padded_qlimits.cuda()
-        else:
-            return padded_qlimits
+        return padded_qlimits.to(self.device)
 
     @property
     def qpos(self):

mani_skill/utils/structs/base.py

@@ -53,10 +53,7 @@ def __maniskill_hash__(self):
 
     @property
     def device(self):
-        if physx.is_gpu_enabled():
-            return torch.device("cuda")
-        else:
-            return torch.device("cpu")
+        return self.scene.device
 
     @property
     def _num_objs(self):
@@ -108,7 +105,7 @@ def _body_data_index(self):
         """a list of indexes of each GPU rigid body in the `px.cuda_rigid_body_data` buffer, one for each element in `self._objs`"""
         if self._body_data_index_internal is None:
             self._body_data_index_internal = torch.tensor(
-                [body.gpu_pose_index for body in self._bodies], device="cuda"
+                [body.gpu_pose_index for body in self._bodies], device=self.device
             )
         return self._body_data_index_internal
 

mani_skill/utils/structs/link.py

@@ -213,7 +213,7 @@ def pose(self) -> Pose:
     def pose(self, arg1: Union[Pose, sapien.Pose, Array]) -> None:
         if physx.is_gpu_enabled():
             if not isinstance(arg1, torch.Tensor):
-                arg1 = vectorize_pose(arg1)
+                arg1 = vectorize_pose(arg1, device=self.device)
             if self.scene.parallel_in_single_scene:
                 if len(arg1.shape) == 1:
                     arg1 = arg1.view(1, -1)

mani_skill/utils/structs/pose.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import List, Union
+from typing import List, Optional, Union
 
 import numpy as np
 import sapien
@@ -94,11 +94,16 @@ def create_from_pq(
 
     @classmethod
     def create(
-        cls, pose: Union[torch.Tensor, sapien.Pose, List[sapien.Pose], "Pose"]
+        cls,
+        pose: Union[torch.Tensor, sapien.Pose, List[sapien.Pose], "Pose"],
+        device: Optional[Device] = None,
     ) -> "Pose":
         if isinstance(pose, sapien.Pose):
             raw_pose = torch.hstack(
-                [common.to_tensor(pose.p), common.to_tensor(pose.q)]
+                [
+                    common.to_tensor(pose.p, device=device),
+                    common.to_tensor(pose.q, device=device),
+                ]
             )
             return cls(raw_pose=add_batch_dim(raw_pose))
         elif isinstance(pose, cls):
@@ -109,8 +114,8 @@ def create(
             for p in pose:
                 ps.append(p.p)
                 qs.append(p.q)
-            ps = common.to_tensor(ps)
-            qs = common.to_tensor(qs)
+            ps = common.to_tensor(ps, device=device)
+            qs = common.to_tensor(qs, device=device)
             return cls(raw_pose=torch.hstack([ps, qs]))
 
         else:
@@ -136,6 +141,11 @@ def shape(self):
     def device(self):
         return self.raw_pose.device
 
+    def to(self, device: Device):
+        if self.raw_pose.device == device:
+            return self
+        return Pose.create(self.raw_pose.to(device))
+
     # -------------------------------------------------------------------------- #
     # Functions from sapien.Pose
     # -------------------------------------------------------------------------- #
@@ -224,21 +234,26 @@ def q(self, arg1: torch.Tensor):
     #     pass
 
 
-def vectorize_pose(pose: Union[sapien.Pose, Pose, Array]) -> torch.Tensor:
+def vectorize_pose(
+    pose: Union[sapien.Pose, Pose, Array], device: Optional[Device] = None
+) -> torch.Tensor:
     """
     Maps several formats of Pose representation to the appropriate tensor representation
     """
     if isinstance(pose, sapien.Pose):
         if physx.is_gpu_enabled():
             return torch.concatenate(
-                [common.to_tensor(pose.p), common.to_tensor(pose.q)]
+                [
+                    common.to_tensor(pose.p, device=device),
+                    common.to_tensor(pose.q, device=device),
+                ]
             )
         else:
             return np.hstack([pose.p, pose.q])
     elif isinstance(pose, Pose):
         return pose.raw_pose
     else:
-        return common.to_tensor(pose)
+        return common.to_tensor(pose, device=device)
 
 
 def to_sapien_pose(pose: Union[torch.Tensor, sapien.Pose, Pose]) -> sapien.Pose: