Fix Rigid Object/Group CUDA memory issue (#46)

Author: yuecideng

embodichain/lab/sim/objects/rigid_object.py

@@ -350,6 +350,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_rigid_body_data(
                 data=pose.clone(),
                 gpu_indices=indices,
@@ -458,6 +459,7 @@ def add_force_torque(
 
         else:
             indices = self.body_data.gpu_indices[local_env_ids]
+            torch.cuda.synchronize(self.device)
             if force is not None:
                 self._ps.gpu_apply_rigid_body_data(
                     data=force,
@@ -658,6 +660,7 @@ def clear_dynamics(self, env_ids: Sequence[int] | None = None) -> None:
                 (len(local_env_ids), 3), dtype=torch.float32, device=self.device
             )
             indices = self.body_data.gpu_indices[local_env_ids]
+            torch.cuda.synchronize(self.device)
             self._ps.gpu_apply_rigid_body_data(
                 data=zeros,
                 gpu_indices=indices,

embodichain/lab/sim/objects/rigid_object_group.py

@@ -362,6 +362,7 @@ def set_local_pose(
             indices = self.body_data.gpu_indices[local_env_ids][
                 :, local_obj_ids
             ].flatten()
+            torch.cuda.synchronize(self.device)
             self._ps.gpu_apply_rigid_body_data(
                 data=pose.clone(),
                 gpu_indices=indices,
@@ -433,6 +434,7 @@ def clear_dynamics(self, env_ids: Sequence[int] | None = None) -> None:
                 device=self.device,
             )
             indices = self.body_data.gpu_indices[local_env_ids].flatten()
+            torch.cuda.synchronize(self.device)
             self._ps.gpu_apply_rigid_body_data(
                 data=zeros,
                 gpu_indices=indices,