diff --git a/embodichain/lab/sim/objects/articulation.py b/embodichain/lab/sim/objects/articulation.py index b20650a..f1325a6 100644 --- a/embodichain/lab/sim/objects/articulation.py +++ b/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() diff --git a/embodichain/lab/sim/objects/gizmo.py b/embodichain/lab/sim/objects/gizmo.py index 28053a7..b3f8dc3 100644 --- a/embodichain/lab/sim/objects/gizmo.py +++ b/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. diff --git a/embodichain/lab/sim/objects/rigid_object.py b/embodichain/lab/sim/objects/rigid_object.py index 047df41..f6fc48c 100644 --- a/embodichain/lab/sim/objects/rigid_object.py +++ b/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) diff --git a/embodichain/lab/sim/objects/rigid_object_group.py b/embodichain/lab/sim/objects/rigid_object_group.py index d510233..bcf1385 100644 --- a/embodichain/lab/sim/objects/rigid_object_group.py +++ b/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() diff --git a/embodichain/lab/sim/objects/robot.py b/embodichain/lab/sim/objects/robot.py index d53a04e..aebcd5c 100644 --- a/embodichain/lab/sim/objects/robot.py +++ b/embodichain/lab/sim/objects/robot.py @@ -586,6 +586,24 @@ def get_control_part_base_pose( link_name=root_link_name, env_ids=local_env_ids, to_matrix=to_matrix ) + def get_control_part_link_names(self, name: str | None = None) -> List[str]: + """Get the link names of the control part. + + Args: + name (str | None): The name of the control part. If None, return all link names. + Returns: + List[str]: link names of the control part. + """ + if name is None: + return self.link_names + if name in self._control_groups: + return self._control_groups[name].link_names + else: + logger.log_warning( + f"The control part '{name}' does not exist in the robot's control parts." + ) + return [] + def _extract_control_groups(self) -> Dict[str, ControlGroup]: r"""Extract control groups from the active joint names. @@ -648,5 +666,46 @@ def build_pk_serial_chain(self) -> None: """ self.pk_serial_chain = self.cfg.build_pk_serial_chain(device=self.device) + def set_physical_visible( + self, + visible: bool = True, + control_part: str | None = None, + rgba: Sequence[float] | None = None, + ): + """set collision of the robot or a specific control part. + + Args: + visible (bool, optional): is collision body visible. Defaults to True. + control_part (str | None, optional): control part to set visibility. Defaults to None. If None, all links are set. + 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.get_control_part_link_names(name=control_part) + + # 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: return super().destroy() diff --git a/embodichain/lab/sim/sim_manager.py b/embodichain/lab/sim/sim_manager.py index 86579ee..a11626d 100644 --- a/embodichain/lab/sim/sim_manager.py +++ b/embodichain/lab/sim/sim_manager.py @@ -1159,7 +1159,7 @@ def set_gizmo_visibility( """ gizmo = self.get_gizmo(uid, control_part) if gizmo is not None: - gizmo.set_visibility(visible) + gizmo.set_visible(visible) def add_sensor(self, sensor_cfg: SensorCfg) -> BaseSensor: """General interface to add a sensor to the scene and returns a handle. diff --git a/scripts/tutorials/sim/create_robot.py b/scripts/tutorials/sim/create_robot.py index 1c8012b..29a4606 100644 --- a/scripts/tutorials/sim/create_robot.py +++ b/scripts/tutorials/sim/create_robot.py @@ -46,7 +46,7 @@ def main(): description="Create and simulate a robot in SimulationManager" ) parser.add_argument( - "--num_envs", type=int, default=1, help="Number of environments to simulate" + "--num_envs", type=int, default=4, help="Number of environments to simulate" ) parser.add_argument( "--device", diff --git a/scripts/tutorials/sim/create_scene.py b/scripts/tutorials/sim/create_scene.py index 823586a..be353d2 100644 --- a/scripts/tutorials/sim/create_scene.py +++ b/scripts/tutorials/sim/create_scene.py @@ -24,8 +24,9 @@ from embodichain.lab.sim import SimulationManager, SimulationManagerCfg from embodichain.lab.sim.cfg import RigidBodyAttributesCfg -from embodichain.lab.sim.shapes import CubeCfg +from embodichain.lab.sim.shapes import CubeCfg, MeshCfg from embodichain.lab.sim.objects import RigidObject, RigidObjectCfg +from dexsim.utility.path import get_resources_data_path def main(): diff --git a/tests/sim/objects/test_articulation.py b/tests/sim/objects/test_articulation.py index b3ecae8..fe93e63 100644 --- a/tests/sim/objects/test_articulation.py +++ b/tests/sim/objects/test_articulation.py @@ -175,6 +175,15 @@ def test_remove_articulation(self): self.art.uid not in self.sim.asset_uids ), "FAIL: Articulation UID still present after removal" + def test_set_physical_visible(self): + self.art.set_physical_visible( + visible=True, + rgba=(0.1, 0.1, 0.9, 0.4), + ) + self.art.set_physical_visible(visible=False) + all_link_names = self.art.link_names + self.art.set_physical_visible(visible=True, link_names=all_link_names[:3]) + def teardown_method(self): """Clean up resources after each test method.""" self.sim.destroy() diff --git a/tests/sim/objects/test_rigid_object.py b/tests/sim/objects/test_rigid_object.py index 56cbe7a..86ed73d 100644 --- a/tests/sim/objects/test_rigid_object.py +++ b/tests/sim/objects/test_rigid_object.py @@ -67,6 +67,7 @@ def setup_simulation(self, sim_device): uid="table", shape=MeshCfg(fpath=table_path), body_type="static" ), ) + self.chair: RigidObject = self.sim.add_rigid_object( cfg=RigidObjectCfg( uid="chair", shape=MeshCfg(fpath=chair_path), body_type="kinematic" @@ -253,6 +254,18 @@ def test_remove(self): self.duck.uid not in self.sim.asset_uids ), "Duck UID still present after removal" + def test_set_physical_visible(self): + self.table.set_physical_visible( + visible=True, + rgba=(0.1, 0.1, 0.9, 0.4), + ) + self.table.set_physical_visible(visible=True) + self.table.set_physical_visible(visible=False) + + def test_set_visible(self): + self.table.set_visible(visible=True) + self.table.set_visible(visible=False) + def teardown_method(self): """Clean up resources after each test method.""" self.sim.destroy() diff --git a/tests/sim/objects/test_rigid_object_group.py b/tests/sim/objects/test_rigid_object_group.py index 8acff65..e9a818e 100644 --- a/tests/sim/objects/test_rigid_object_group.py +++ b/tests/sim/objects/test_rigid_object_group.py @@ -107,6 +107,14 @@ def test_remove(self): self.obj_group.uid not in self.sim.asset_uids ), "Object group UID still present after removal" + def test_set_physical_visible(self): + self.obj_group.set_physical_visible(visible=True) + self.obj_group.set_physical_visible(visible=False) + + def test_set_visible(self): + self.obj_group.set_visible(visible=True) + self.obj_group.set_visible(visible=False) + def teardown_method(self): """Clean up resources after each test method.""" self.sim.destroy() diff --git a/tests/sim/objects/test_robot.py b/tests/sim/objects/test_robot.py index 282abed..e5c28ca 100644 --- a/tests/sim/objects/test_robot.py +++ b/tests/sim/objects/test_robot.py @@ -245,6 +245,21 @@ def teardown_method(self): """Clean up resources after each test method.""" self.sim.destroy() + def test_set_physical_visible(self): + self.robot.set_physical_visible( + visible=True, + rgba=(0.1, 0.1, 0.9, 0.4), + control_part="left_arm", + ) + self.robot.set_physical_visible( + visible=True, + control_part="left_arm", + ) + self.robot.set_physical_visible( + visible=False, + control_part="left_arm", + ) + class TestRobotCPU(BaseRobotTest): def setup_method(self):