Fix Articulation GPU issue (#44)

Author: yuecideng

embodichain/lab/sim/cfg.py

@@ -165,7 +165,6 @@ def attr(self) -> PhysicalAttr:
         attr.mass = self.mass
         attr.contact_offset = self.contact_offset
         attr.rest_offset = self.rest_offset
-        attr.enable_collision = self.enable_collision
         attr.dynamic_friction = self.dynamic_friction
         attr.static_friction = self.static_friction
         attr.angular_damping = self.angular_damping

embodichain/lab/sim/objects/articulation.py

@@ -511,9 +511,7 @@ def __init__(
         self.device = device
 
         # Store all indices for batch operations
-        self._all_indices = torch.arange(
-            len(entities), dtype=torch.int32, device=device
-        )
+        self._all_indices = torch.arange(len(entities), dtype=torch.int32).tolist()
 
         if device.type == "cuda":
             self._world.update(0.001)
@@ -799,6 +797,7 @@ def set_local_pose(
             # we should keep `pose_` life cycle to the end of the function.
             pose_ = torch.cat((quat, xyz), dim=-1)
             indices = self.body_data.gpu_indices[local_env_ids]
+            torch.cuda.synchronize(self.device)
             self._ps.gpu_apply_root_data(
                 data=pose_,
                 gpu_indices=indices,
@@ -978,6 +977,7 @@ def set_qpos(
             indices = self.body_data.gpu_indices[local_env_ids]
             qpos_set = self.body_data._qpos[local_env_ids]
             qpos_set[:, local_joint_ids] = qpos
+            torch.cuda.synchronize(self.device)
             self._ps.gpu_apply_joint_data(
                 data=qpos_set,
                 gpu_indices=indices,
@@ -1041,6 +1041,7 @@ def set_qvel(
                 self.body_data.qvel
                 qvel_set = self.body_data._qvel[local_env_ids]
                 qvel_set[:, joint_ids] = qvel
+            torch.cuda.synchronize(self.device)
             self._ps.gpu_apply_joint_data(
                 data=qvel_set,
                 gpu_indices=indices,
@@ -1081,6 +1082,7 @@ def set_qf(
                 self.body_data.qf
                 qf_set = self.body_data._qf[local_env_ids]
                 qf_set[:, joint_ids] = qf
+            torch.cuda.synchronize(self.device)
             self._ps.gpu_apply_joint_data(
                 data=qf_set,
                 gpu_indices=indices,
@@ -1161,11 +1163,13 @@ def clear_dynamics(self, env_ids: Sequence[int] | None = None) -> None:
                 (len(local_env_ids), self.dof), dtype=torch.float32, device=self.device
             )
             indices = self.body_data.gpu_indices[local_env_ids]
+            torch.cuda.synchronize(self.device)
             self._ps.gpu_apply_joint_data(
                 data=zeros,
                 gpu_indices=indices,
                 data_type=ArticulationGPUAPIWriteType.JOINT_VELOCITY,
             )
+            torch.cuda.synchronize(self.device)
             self._ps.gpu_apply_joint_data(
                 data=zeros,
                 gpu_indices=indices,

embodichain/lab/sim/objects/rigid_object.py

@@ -176,9 +176,7 @@ def __init__(
         self._world = dexsim.default_world()
         self._ps = self._world.get_physics_scene()
 
-        self._all_indices = torch.arange(
-            len(entities), dtype=torch.int32, device=device
-        )
+        self._all_indices = torch.arange(len(entities), dtype=torch.int32).tolist()
 
         # data for managing body data (only for dynamic and kinematic bodies) on GPU.
         self._data: RigidBodyData | None = None
@@ -200,6 +198,12 @@ def __init__(
         # set default collision filter
         self._set_default_collision_filter()
 
+        # TODO: Must be called after setting all attributes.
+        # May be improved in the future.
+        if cfg.attrs.enable_collision is False:
+            flag = torch.zeros(len(entities), dtype=torch.bool)
+            self.enable_collision(flag)
+
         # reserve flag for collision visible node existence
         self._has_collision_visible_node = False
 
@@ -614,6 +618,26 @@ def get_user_ids(self) -> torch.Tensor:
             device=self.device,
         )
 
+    def enable_collision(
+        self, enable: torch.Tensor, env_ids: Sequence[int] | None = None
+    ) -> None:
+        """Enable or disable collision for the rigid bodies.
+
+        Args:
+            enable (torch.Tensor): A tensor of shape (N,) representing whether to enable collision for each rigid body.
+            env_ids (Sequence[int] | None): Environment indices. If None, then all indices are used.
+        """
+        local_env_ids = self._all_indices if env_ids is None else env_ids
+
+        if len(local_env_ids) != len(enable):
+            logger.log_error(
+                f"Length of env_ids {len(local_env_ids)} does not match enable length {len(enable)}."
+            )
+
+        enable_list = enable.tolist()
+        for i, env_idx in enumerate(local_env_ids):
+            self._entities[env_idx].enable_collision(bool(enable_list[i]))
+
     def clear_dynamics(self, env_ids: Sequence[int] | None = None) -> None:
         """Clear the dynamics of the rigid bodies by resetting velocities and applying zero forces and torques.
 

embodichain/lab/sim/objects/rigid_object_group.py

@@ -187,12 +187,10 @@ def __init__(
         self._world = dexsim.default_world()
         self._ps = self._world.get_physics_scene()
 
-        self._all_indices = torch.arange(
-            len(entities), dtype=torch.int32, device=device
-        )
+        self._all_indices = torch.arange(len(entities), dtype=torch.int32).tolist()
         self._all_obj_indices = torch.arange(
-            len(entities[0]), dtype=torch.int32, device=device
-        )
+            len(entities[0]), dtype=torch.int32
+        ).tolist()
 
         # data for managing body data (only for dynamic and kinematic bodies) on GPU.
         self._data = RigidBodyGroupData(entities=entities, ps=self._ps, device=device)

embodichain/lab/sim/objects/soft_object.py

@@ -165,9 +165,7 @@ def __init__(
     ) -> None:
         self._world = dexsim.default_world()
         self._ps = self._world.get_physics_scene()
-        self._all_indices = torch.arange(
-            len(entities), dtype=torch.int32, device=device
-        )
+        self._all_indices = torch.arange(len(entities), dtype=torch.int32).tolist()
 
         self._data = SoftBodyData(entities=entities, ps=self._ps, device=device)