decorator to assert error when gpu is initialized for calling functio…

…ns user should not call.

mani_skill2/envs/sapien_env.py

@@ -856,7 +856,7 @@ def render_human(self):
 
         # TODO (stao): currently in GPU mode we cannot render all sub-scenes together in the GUI yet. So we have this
         # naive solution which shows whatever scene is selected by self._viewer_scene_idx
-        if physx.is_gpu_enabled() and self._scene._buffers_ready:
+        if physx.is_gpu_enabled() and self._scene._gpu_sim_initialized:
             # TODO (stao): This is the slow method, update objects via cpu
             for actor in self._scene.actors.values():
                 if actor.px_body_type == "static":

mani_skill2/envs/scene.py

@@ -41,7 +41,7 @@ class ManiSkillScene:
     def __init__(self, sub_scenes: List[sapien.Scene], debug_mode: bool = True):
         self.sub_scenes = sub_scenes
         self.px = self.sub_scenes[0].physx_system
-        self._buffers_ready = False
+        self._gpu_sim_initialized = False
         self.debug_mode = debug_mode
         super().__init__()
 
@@ -513,7 +513,7 @@ def _setup_gpu(self):
         )  # zero out all q velocities
         self.px.gpu_apply_articulation_qvel()
 
-        self._buffers_ready = True
+        self._gpu_sim_initialized = True
         self._gpu_fetch_all()
 
     def _gpu_apply_all(self):

mani_skill2/envs/tasks/stack_cube.py

@@ -72,7 +72,7 @@ def _initialize_actors(self):
             xy = torch.rand((self.num_envs, 2)) * 0.2 - 0.1
             region = [[-0.1, -0.2], [0.1, 0.2]]
             sampler = UniformPlacementSampler(bounds=region, batch_size=self.num_envs)
-            radius = (torch.linalg.norm(torch.Tensor([0.02, 0.02])) + 0.001).to(
+            radius = (torch.linalg.norm(torch.tensor([0.02, 0.02])) + 0.001).to(
                 self.device
             )
             cubeA_xy = xy + sampler.sample(radius, 100)

mani_skill2/utils/structs/actor.py

@@ -117,6 +117,9 @@ def get_state(self):
         else:
             vel = self.get_linear_velocity()  # [N, 3]
             ang_vel = self.get_angular_velocity()  # [N, 3]
+        import ipdb
+
+        ipdb.set_trace()
         return torch.hstack([pose.p, pose.q, vel, ang_vel])
 
     def set_state(self, state: Array):

mani_skill2/utils/structs/base.py

@@ -8,6 +8,7 @@
 import torch
 
 from mani_skill2.utils.sapien_utils import to_numpy, to_tensor
+from mani_skill2.utils.structs.decorators import before_gpu_init
 from mani_skill2.utils.structs.types import Array
 
 if TYPE_CHECKING:
@@ -51,6 +52,8 @@ def px(self):
 @dataclass
 class PhysxRigidBodyComponentStruct:
     # Reference to the data for this rigid body on the GPU
+    _scene: ManiSkillScene
+    """The ManiSkillScene object that manages the sub-scenes this dataclasses's objects are in"""
     _body_data_name: str
     _bodies: List[physx.PhysxRigidBodyComponent]
     _body_data_index_internal: slice = None
@@ -115,10 +118,11 @@ def set_mass(self, arg0: float) -> None:
     # def set_max_contact_impulse(self, impulse: float) -> None: ... # TODO (Stao)
     # def set_max_depenetraion_velocity(self, velocity: float) -> None: ... # TODO (Stao)
     @property
-    def angular_damping(self) -> float:
-        return self._bodies[0].angular_damping
+    def angular_damping(self) -> torch.Tensor:
+        return torch.tensor([body.angular_damping for body in self._bodies])
 
     @angular_damping.setter
+    @before_gpu_init
     def angular_damping(self, arg1: float) -> None:
         for rb in self._bodies:
             rb.angular_damping = arg1
@@ -131,8 +135,8 @@ def angular_velocity(self) -> torch.Tensor:
             return torch.from_numpy(self._bodies[0].angular_velocity[None, :])
 
     @property
-    def auto_compute_mass(self) -> bool:
-        return self._bodies[0].auto_compute_mass
+    def auto_compute_mass(self) -> torch.Tensor:
+        return torch.tensor([body.auto_compute_mass for body in self._bodies])
 
     # @property
     # def cmass_local_pose(self) -> sapien.pysapien.Pose:
@@ -143,10 +147,11 @@ def auto_compute_mass(self) -> bool:
     # def cmass_local_pose(self, arg1: sapien.pysapien.Pose) -> None:
     #     pass
     @property
-    def disable_gravity(self) -> bool:
-        return self._bodies[0].disable_gravity
+    def disable_gravity(self) -> torch.Tensor:
+        return torch.tensor([body.disable_gravity for body in self._bodies])
 
     @disable_gravity.setter
+    @before_gpu_init
     def disable_gravity(self, arg1: bool) -> None:
         for rb in self._bodies:
             rb.disable_gravity = arg1
@@ -160,10 +165,11 @@ def disable_gravity(self, arg1: bool) -> None:
     # def inertia(self, arg1: numpy.ndarray[numpy.float32, _Shape, _Shape[3]]) -> None:
     #     pass
     @property
-    def linear_damping(self) -> float:
-        return self._bodies[0].linear_damping
+    def linear_damping(self) -> torch.Tensor:
+        return torch.tensor([body.linear_damping for body in self._bodies])
 
     @linear_damping.setter
+    @before_gpu_init
     def linear_damping(self, arg1: float) -> None:
         for rb in self._bodies:
             rb.linear_damping = arg1
@@ -176,17 +182,18 @@ def linear_velocity(self) -> torch.Tensor:
             return torch.from_numpy(self._bodies[0].linear_velocity[None, :])
 
     @property
-    def mass(self) -> float:
-        return self._bodies[0].mass
+    def mass(self) -> torch.Tensor:
+        return torch.tensor([body.mass for body in self._bodies])
 
     @mass.setter
+    @before_gpu_init
     def mass(self, arg1: float) -> None:
         if physx.is_gpu_enabled():
             raise NotImplementedError(
                 "Setting mass is not supported on GPU sim at the moment."
             )
         else:
-            return self._bodies[0].mass
+            self._bodies[0].mass = arg1
 
     # @property
     # def max_contact_impulse(self) -> float:
@@ -285,6 +292,7 @@ def gpu_pose_index(self):
             )
 
     @property
+    @before_gpu_init
     def is_sleeping(self):
         if physx.is_gpu_enabled():
             return [b.is_sleeping for b in self._bodies]

mani_skill2/utils/structs/decorators.py

@@ -0,0 +1,13 @@
+def before_gpu_init(func):
+    """
+    decorator to throw an error if a function is called when gpu sim has been initialized already. Used for functions such as setting friction values which currently
+    cannot be changed once the gpu simulation has started.
+    """
+
+    def wrapper(self, *args, **kwargs):
+        assert (
+            self._scene._gpu_sim_initialized == False
+        ), f"{func} can only be called when the GPU simulation has not been initialized yet"
+        return func(self, *args, **kwargs)
+
+    return wrapper