support collision body visualization for rigid object and articulation (#27)

Co-authored-by: chenjian <[email protected]>
Co-authored-by: Yueci Deng <[email protected]>

Author: 3 people

embodichain/lab/sim/objects/articulation.py

@@ -579,6 +579,11 @@ def __init__(
         # set default collision filter
         self._set_default_collision_filter()
 
+        # flag for collision visible node existence
+        self._has_collision_visible_node_dict = dict()
+        for link_name in self.link_names:
+            self._has_collision_visible_node_dict[link_name] = False
+
     def __str__(self) -> str:
         parent_str = super().__str__()
         return parent_str + f" | dof: {self.dof} | num_links: {self.num_links}"
@@ -1535,6 +1540,47 @@ def get_visual_material_inst(
                 result.append(mat_dict)
         return result
 
+    def set_physical_visible(
+        self,
+        visible: bool = True,
+        link_names: List[str] | None = None,
+        rgba: Sequence[float] | None = None,
+    ):
+        """set collision
+
+        Args:
+            visible (bool, optional): is collision body visible. Defaults to True.
+            link_names (List[str] | None, optional): links to set visibility. Defaults to None.
+            rgba (Sequence[float] | None, optional): collision body visible rgba. It will be defined at the first time the function is called. Defaults to None.
+        """
+        rgba = rgba if rgba is not None else (0.8, 0.2, 0.2, 0.7)
+        if len(rgba) != 4:
+            logger.log_error(f"Invalid rgba {rgba}, should be a sequence of 4 floats.")
+        rgba = np.array(
+            [
+                rgba[0],
+                rgba[1],
+                rgba[2],
+                rgba[3],
+            ]
+        )
+        link_names = self.link_names if link_names is None else link_names
+
+        # create collision visible node if not exist
+        if visible:
+            for i, env_idx in enumerate(self._all_indices):
+                for link_name in link_names:
+                    if self._has_collision_visible_node_dict[link_name] is False:
+                        self._entities[env_idx].create_physical_visible_node(
+                            rgba, link_name
+                        )
+                        self._has_collision_visible_node_dict[link_name] = True
+
+        # set visibility
+        for i, env_idx in enumerate(self._all_indices):
+            for link_name in link_names:
+                self._entities[env_idx].set_physical_visible(visible, link_name)
+
     def destroy(self) -> None:
         env = self._world.get_env()
         arenas = env.get_all_arenas()

embodichain/lab/sim/objects/gizmo.py

@@ -424,7 +424,7 @@ def toggle_visibility(self) -> bool:
 
         return self._is_visible
 
-    def set_visibility(self, visible: bool):
+    def set_visible(self, visible: bool):
         """
         Set the visibility of the gizmo.
 

embodichain/lab/sim/objects/rigid_object.py

@@ -19,7 +19,7 @@
 import numpy as np
 
 from dataclasses import dataclass
-from typing import List, Sequence, Optional, Union
+from typing import List, Sequence, Union
 
 from dexsim.models import MeshObject
 from dexsim.types import RigidBodyGPUAPIReadType, RigidBodyGPUAPIWriteType
@@ -177,7 +177,7 @@ def __init__(
         )
 
         # data for managing body data (only for dynamic and kinematic bodies) on GPU.
-        self._data: Optional[RigidBodyData] = None
+        self._data: RigidBodyData | None = None
         if self.is_static is False:
             self._data = RigidBodyData(entities=entities, ps=self._ps, device=device)
 
@@ -196,6 +196,9 @@ def __init__(
         # set default collision filter
         self._set_default_collision_filter()
 
+        # reserve flag for collision visible node existence
+        self._has_collision_visible_node = False
+
     def __str__(self) -> str:
         parent_str = super().__str__()
         return (
@@ -204,7 +207,7 @@ def __str__(self) -> str:
         )
 
     @property
-    def body_data(self) -> Optional[RigidBodyData]:
+    def body_data(self) -> RigidBodyData | None:
         """Get the rigid body data manager for this rigid object.
 
         Returns:
@@ -266,7 +269,7 @@ def _set_default_collision_filter(self) -> None:
         self.set_collision_filter(collision_filter_data)
 
     def set_collision_filter(
-        self, filter_data: torch.Tensor, env_ids: Optional[Sequence[int]] = None
+        self, filter_data: torch.Tensor, env_ids: Sequence[int] | None = None
     ) -> None:
         """set collision filter data for the rigid object.
 
@@ -276,7 +279,7 @@ def set_collision_filter(
                 If 2nd element is 0, the object will collision with all other objects in world.
                 3rd and 4th elements are not used currently.
 
-            env_ids (Optional[Sequence[int]], optional): Environment indices. If None, then all indices are used. Defaults to None.
+            env_ids (Sequence[int] | None, optional): Environment indices. If None, then all indices are used. Defaults to None.
         """
         local_env_ids = self._all_indices if env_ids is None else env_ids
 
@@ -292,13 +295,13 @@ def set_collision_filter(
             )
 
     def set_local_pose(
-        self, pose: torch.Tensor, env_ids: Optional[Sequence[int]] = None
+        self, pose: torch.Tensor, env_ids: Sequence[int] | None = None
     ) -> None:
         """Set local pose of the rigid object.
 
         Args:
             pose (torch.Tensor): The local pose of the rigid object with shape (N, 7) or (N, 4, 4).
-            env_ids (Optional[Sequence[int]], optional): Environment indices. If None, then all indices are used.
+            env_ids (Sequence[int] | None, optional): Environment indices. If None, then all indices are used.
         """
         local_env_ids = self._all_indices if env_ids is None else env_ids
 
@@ -395,10 +398,10 @@ def get_local_pose_cpu(
 
     def add_force_torque(
         self,
-        force: Optional[torch.Tensor] = None,
-        torque: Optional[torch.Tensor] = None,
-        pos: Optional[torch.Tensor] = None,
-        env_ids: Optional[Sequence[int]] = None,
+        force: torch.Tensor | None = None,
+        torque: torch.Tensor | None = None,
+        pos: torch.Tensor | None = None,
+        env_ids: Sequence[int] | None = None,
     ) -> None:
         """Add force and/or torque to the rigid object.
 
@@ -409,10 +412,10 @@ def add_force_torque(
             - if not `pos` is specified, the force and torque are applied at the center of mass of the rigid body.
 
         Args:
-            force (Optional[torch.Tensor] = None): The force to add with shape (N, 3). Defaults to None.
-            torque (Optional[torch.Tensor], optional): The torque to add with shape (N, 3). Defaults to None.
-            pos (Optional[torch.Tensor], optional): The position to apply the force at with shape (N, 3). Defaults to None.
-            env_ids (Optional[Sequence[int]], optional): Environment indices. If None, then all indices are used.
+            force (torch.Tensor | None = None): The force to add with shape (N, 3). Defaults to None.
+            torque (torch.Tensor | None, optional): The torque to add with shape (N, 3). Defaults to None.
+            pos (torch.Tensor | None, optional): The position to apply the force at with shape (N, 3). Defaults to None.
+            env_ids (Sequence[int] | None, optional): Environment indices. If None, then all indices are used.
         """
         if force is None and torque is None:
             logger.log_warning(
@@ -463,13 +466,13 @@ def add_force_torque(
     def set_attrs(
         self,
         attrs: Union[RigidBodyAttributesCfg, List[RigidBodyAttributesCfg]],
-        env_ids: Optional[Sequence[int]] = None,
+        env_ids: Sequence[int] | None = None,
     ) -> None:
         """Set physical attributes for the rigid object.
 
         Args:
             attrs (Union[RigidBodyAttributesCfg, List[RigidBodyAttributesCfg]]): The physical attributes to set.
-            env_ids (Optional[Sequence[int]], optional): Environment indices. If None, then all indices are used.
+            env_ids (Sequence[int] | None, optional): Environment indices. If None, then all indices are used.
         """
         local_env_ids = self._all_indices if env_ids is None else env_ids
 
@@ -487,13 +490,13 @@ def set_attrs(
                 self._entities[env_idx].set_physical_attr(attrs[i].attr())
 
     def set_visual_material(
-        self, mat: VisualMaterial, env_ids: Optional[Sequence[int]] = None
+        self, mat: VisualMaterial, env_ids: Sequence[int] | None = None
     ) -> None:
         """Set visual material for the rigid object.
 
         Args:
             mat (VisualMaterial): The material to set.
-            env_ids (Optional[Sequence[int]], optional): Environment indices. If None, then all indices are used.
+            env_ids (Sequence[int] | None, optional): Environment indices. If None, then all indices are used.
         """
         local_env_ids = self._all_indices if env_ids is None else env_ids
 
@@ -503,12 +506,12 @@ def set_visual_material(
             self._visual_material[env_idx] = mat_inst
 
     def get_visual_material_inst(
-        self, env_ids: Optional[Sequence[int]] = None
+        self, env_ids: Sequence[int] | None = None
     ) -> List[VisualMaterialInst]:
         """Get material instances for the rigid object.
 
         Args:
-            env_ids (Optional[Sequence[int]], optional): Environment indices. If None, then all indices are used.
+            env_ids (Sequence[int] | None, optional): Environment indices. If None, then all indices are used.
 
         Returns:
             List[MaterialInst]: List of material instances.
@@ -533,12 +536,12 @@ def share_visual_material_inst(self, mat_insts: List[VisualMaterialInst]) -> Non
             entity.set_material(mat_insts[i].mat)
             self._visual_material[i] = mat_insts[i]
 
-    def get_body_scale(self, env_ids: Optional[Sequence[int]] = None) -> torch.Tensor:
+    def get_body_scale(self, env_ids: Sequence[int] | None = None) -> torch.Tensor:
         """
         Retrieve the body scale for specified environment instances.
 
         Args:
-            env_ids (Optional[Sequence[int]]): A sequence of environment instance IDs.
+            env_ids (Sequence[int] | None): A sequence of environment instance IDs.
                 If None, retrieves the body scale for all instances.
 
         Returns:
@@ -553,13 +556,13 @@ def get_body_scale(self, env_ids: Optional[Sequence[int]] = None) -> torch.Tenso
         )
 
     def set_body_scale(
-        self, scale: torch.Tensor, env_ids: Optional[Sequence[int]] = None
+        self, scale: torch.Tensor, env_ids: Sequence[int] | None = None
     ) -> None:
         """Set the scale of the rigid body.
 
         Args:
             scale (torch.Tensor): The scale to set with shape (N, 3).
-            env_ids (Optional[Sequence[int]], optional): Environment indices. If None, then all indices are used.
+            env_ids (Sequence[int] | None, optional): Environment indices. If None, then all indices are used.
         """
         local_env_ids = self._all_indices if env_ids is None else env_ids
 
@@ -575,12 +578,12 @@ def set_body_scale(
         else:
             logger.log_error(f"Setting body scale on GPU is not supported yet.")
 
-    def get_vertices(self, env_ids: Optional[Sequence[int]] = None) -> torch.Tensor:
+    def get_vertices(self, env_ids: Sequence[int] | None = None) -> torch.Tensor:
         """
         Retrieve the vertices of the rigid objects.
 
         Args:
-            env_ids (Optional[Sequence[int]]): A sequence of environment IDs for which to retrieve vertices.
+            env_ids (Sequence[int] | None): A sequence of environment IDs for which to retrieve vertices.
                                                 If None, retrieves vertices for all instances.
 
         Returns:
@@ -607,11 +610,11 @@ def get_user_ids(self) -> torch.Tensor:
             device=self.device,
         )
 
-    def clear_dynamics(self, env_ids: Optional[Sequence[int]] = None) -> None:
+    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.
 
         Args:
-            env_ids (Optional[Sequence[int]]): Environment indices. If None, then all indices are used.
+            env_ids (Sequence[int] | None): Environment indices. If None, then all indices are used.
         """
         if self.is_non_dynamic:
             return
@@ -648,7 +651,51 @@ def clear_dynamics(self, env_ids: Optional[Sequence[int]] = None) -> None:
                 data_type=RigidBodyGPUAPIWriteType.TORQUE,
             )
 
-    def reset(self, env_ids: Optional[Sequence[int]] = None) -> None:
+    def set_physical_visible(
+        self,
+        visible: bool = True,
+        rgba: Sequence[float] | None = None,
+    ):
+        """set collion render visibility
+
+        Args:
+            visible (bool, optional): is collision body visible. Defaults to True.
+            rgba (Sequence[float] | None, optional): collision body visible rgba. It will be defined at the first time the function is called. Defaults to None.
+        """
+        rgba = rgba if rgba is not None else (0.8, 0.2, 0.2, 0.7)
+        if len(rgba) != 4:
+            logger.log_error(f"Invalid rgba {rgba}, should be a sequence of 4 floats.")
+
+        # create collision visible node if not exist
+        if visible:
+            if not self._has_collision_visible_node:
+                for i, env_idx in enumerate(self._all_indices):
+                    self._entities[env_idx].create_physical_visible_node(
+                        np.array(
+                            [
+                                rgba[0],
+                                rgba[1],
+                                rgba[2],
+                                rgba[3],
+                            ]
+                        )
+                    )
+                self._has_collision_visible_node = True
+
+        # create collision visible node if not exist
+        for i, env_idx in enumerate(self._all_indices):
+            self._entities[env_idx].set_physical_visible(visible)
+
+    def set_visible(self, visible: bool = True) -> None:
+        """Set the visibility of the rigid object.
+
+        Args:
+            visible (bool, optional): Whether the rigid object is visible. Defaults to True.
+        """
+        for i, env_idx in enumerate(self._all_indices):
+            self._entities[env_idx].set_visible(visible)
+
+    def reset(self, env_ids: Sequence[int] | None = None) -> None:
         local_env_ids = self._all_indices if env_ids is None else env_ids
         num_instances = len(local_env_ids)
         self.set_attrs(self.cfg.attrs, env_ids=local_env_ids)

embodichain/lab/sim/objects/rigid_object_group.py

@@ -204,6 +204,10 @@ def __init__(
         # set default collision filter
         self._set_default_collision_filter()
 
+        # reserve flag for collision visible node existence
+        n_instances = len(self._entities[0])
+        self._has_collision_visible_node_list = [False] * n_instances
+
     def __str__(self) -> str:
         parent_str = super().__str__()
         return (
@@ -503,6 +507,53 @@ def reset(self, env_ids: Sequence[int] | None = None) -> None:
 
         self.clear_dynamics(env_ids=local_env_ids)
 
+    def set_physical_visible(
+        self,
+        visible: bool = True,
+        rgba: Sequence[float] | None = None,
+    ):
+        """set collion render visibility
+
+        Args:
+            visible (bool, optional): is collision body visible. Defaults to True.
+            rgba (Sequence[float] | None, optional): collision body visible rgba. It will be defined at the first time the function is called. Defaults to None.
+        """
+        rgba = rgba if rgba is not None else (0.8, 0.2, 0.2, 0.7)
+        if len(rgba) != 4:
+            logger.log_error(f"Invalid rgba {rgba}, should be a sequence of 4 floats.")
+
+        # create collision visible node if not exist
+        if visible:
+            for i, env_idx in enumerate(self._all_indices):
+                for intance_id, entity in enumerate(self._entities[env_idx]):
+                    if not self._has_collision_visible_node_list[intance_id]:
+                        entity.create_physical_visible_node(
+                            np.array(
+                                [
+                                    rgba[0],
+                                    rgba[1],
+                                    rgba[2],
+                                    rgba[3],
+                                ]
+                            )
+                        )
+                        self._has_collision_visible_node_list[intance_id] = True
+
+        # create collision visible node if not exist
+        for i, env_idx in enumerate(self._all_indices):
+            for entity in self._entities[env_idx]:
+                entity.set_physical_visible(visible)
+
+    def set_visible(self, visible: bool = True) -> None:
+        """Set the visibility of the rigid object group.
+
+        Args:
+            visible (bool, optional): Whether the rigid object group is visible. Defaults to True.
+        """
+        for i, env_idx in enumerate(self._all_indices):
+            for entity in self._entities[env_idx]:
+                entity.set_visible(visible)
+
     def destroy(self) -> None:
         env = self._world.get_env()
         arenas = env.get_all_arenas()