Add camera extrinsics randomization functor (#32)

Author: yuecideng

README.md

@@ -1,11 +1,11 @@
-
 # EmbodiChain
 
 ![teaser](assets/imgs/teaser.jpg)
 
-[![Version](https://img.shields.io/badge/version-0.0.1-blue.svg)](https://github.com/DexForce/EmbodiChain/releases)
 [![License](https://img.shields.io/github/license/DexForce/EmbodiChain)](LICENSE)
 [![GitHub Pages](https://img.shields.io/badge/GitHub%20Pages-docs-blue?logo=github&logoColor=white)](https://dexforce.github.io/EmbodiChain/introduction.html)
+[![Python](https://img.shields.io/badge/python-3.10-blue.svg)](https://docs.python.org/3/whatsnew/3.10.html)
+[![Version](https://img.shields.io/badge/version-0.0.1-blue.svg)](https://github.com/DexForce/EmbodiChain/releases)
 
 ---
 

configs/gym/cobotmagic.json

@@ -59,6 +59,7 @@
     },
     "sensor": [
         {
+            "uid": "camera_1",
             "sensor_type": "Camera",
             "width": 640,
             "height": 480,

embodichain/lab/gym/envs/managers/events.py

@@ -111,13 +111,13 @@ def __init__(self, cfg: FunctorCfg, env: EmbodiedEnv):
 
             # remove regular expression from patterns
             patterns = remove_regex_chars(patterns)
-            full_path = get_data_path(f"{folder_path}/")
+            self._full_path = get_data_path(f"{folder_path}/")
             self._asset_group_path = get_all_files_in_directory(
-                full_path, patterns=patterns
+                self._full_path, patterns=patterns
             )
         else:
-            full_path = get_data_path(folder_path)
-            self._asset_group_path = get_all_files_in_directory(full_path)
+            self._full_path = get_data_path(folder_path)
+            self._asset_group_path = get_all_files_in_directory(self._full_path)
 
     def __call__(
         self,

embodichain/lab/gym/envs/managers/observations.py

@@ -40,6 +40,9 @@ def get_rigid_object_pose(
 ) -> torch.Tensor:
     """Get the world poses of the rigid objects in the environment.
 
+    If the rigid object with the specified UID does not exist in the environment,
+    a zero tensor will be returned.
+
     Args:
         env: The environment instance.
         obs: The observation dictionary.
@@ -49,6 +52,9 @@ def get_rigid_object_pose(
         A tensor of shape (num_envs, 4, 4) representing the world poses of the rigid objects.
     """
 
+    if entity_cfg.uid not in env.sim.get_rigid_object_uid_list():
+        return torch.zeros((env.num_envs, 4, 4), dtype=torch.float32)
+
     obj = env.sim.get_rigid_object(entity_cfg.uid)
 
     return obj.get_local_pose(to_matrix=True)
@@ -130,7 +136,10 @@ def compute_semantic_mask(
 
     robot_mask = (mask_exp == robot_uids_exp).any(-1).squeeze_(-1)
 
-    foreground_assets = [env.sim.get_asset(uid) for uid in foreground_uids]
+    asset_uids = env.sim.asset_uids
+    foreground_assets = [
+        env.sim.get_asset(uid) for uid in foreground_uids if uid in asset_uids
+    ]
 
     # cat assets uid (num_envs, n) into dim 1
     foreground_uids = torch.cat(

embodichain/lab/gym/envs/managers/randomization/rendering.py

@@ -31,7 +31,11 @@
     VisualMaterialCfg,
 )
 from embodichain.utils.string import resolve_matching_names
-from embodichain.utils.math import sample_uniform
+from embodichain.utils.math import (
+    sample_uniform,
+    quat_from_euler_xyz,
+    euler_xyz_from_quat,
+)
 from embodichain.utils import logger
 from embodichain.data import get_data_path
 
@@ -40,12 +44,135 @@
 
 
 __all__ = [
+    "randomize_camera_extrinsics",
     "randomize_light",
     "randomize_camera_intrinsics",
     "randomize_visual_material",
 ]
 
 
+def randomize_camera_extrinsics(
+    env: EmbodiedEnv,
+    env_ids: Union[torch.Tensor, None],
+    entity_cfg: SceneEntityCfg,
+    pos_range: Optional[tuple[list[float], list[float]]] = None,
+    euler_range: Optional[tuple[list[float], list[float]]] = None,
+    eye_range: Optional[tuple[list[float], list[float]]] = None,
+    target_range: Optional[tuple[list[float], list[float]]] = None,
+    up_range: Optional[tuple[list[float], list[float]]] = None,
+) -> None:
+    """
+    Randomize camera extrinsic properties (position and orientation).
+
+    Behavior:
+    - If extrinsics config has a parent field (attach mode), pos_range/euler_range are used to perturb the initial pose (pos, quat),
+        and set_local_pose is called to attach the camera to the parent node. In this case, pose is related to parent.
+    - If extrinsics config uses eye/target/up (no parent), eye_range/target_range/up_range are used to perturb the initial eye, target, up vectors,
+        and look_at is called to set the camera orientation.
+
+    Args:
+        env: The environment instance.
+        env_ids: The environment IDs to apply the randomization.
+        entity_cfg (SceneEntityCfg): The configuration of the scene entity to randomize.
+        pos_range: Position perturbation range (attach mode).
+        euler_range: Euler angle perturbation range (attach mode).
+        eye_range: Eye position perturbation range (look_at mode).
+        target_range: Target position perturbation range (look_at mode).
+        up_range: Up vector perturbation range (look_at mode).
+    """
+    camera: Union[Camera, StereoCamera] = env.sim.get_sensor(entity_cfg.uid)
+    num_instance = len(env_ids)
+
+    extrinsics = camera.cfg.extrinsics
+
+    if extrinsics.parent is not None:
+        # If extrinsics has a parent field, use pos/euler perturbation and attach camera to parent node
+        init_pos = getattr(extrinsics, "pos", [0.0, 0.0, 0.0])
+        init_quat = getattr(extrinsics, "quat", [0.0, 0.0, 0.0, 1.0])
+        new_pose = torch.tensor(
+            [init_pos + init_quat], dtype=torch.float32, device=env.device
+        ).repeat(num_instance, 1)
+        if pos_range:
+            random_value = sample_uniform(
+                lower=torch.tensor(pos_range[0]),
+                upper=torch.tensor(pos_range[1]),
+                size=(num_instance, 3),
+            )
+            new_pose[:, :3] += random_value
+        if euler_range:
+            # 1. quat -> euler
+            init_quat_np = (
+                torch.tensor(init_quat, dtype=torch.float32, device=env.device)
+                .unsqueeze_(0)
+                .repeat(num_instance, 1)
+            )
+            init_euler = euler_xyz_from_quat(init_quat_np)
+            # 2. Sample perturbation for euler angles
+            random_value = sample_uniform(
+                lower=torch.tensor(euler_range[0]),
+                upper=torch.tensor(euler_range[1]),
+                size=(num_instance, 3),
+            )
+            # 3. Add perturbation to each environment and convert back to quaternion
+            new_quat = quat_from_euler_xyz(init_euler + random_value)
+            new_pose[:, 3:7] = new_quat
+
+        camera.set_local_pose(new_pose, env_ids=env_ids)
+
+    elif extrinsics.eye is not None:
+        # If extrinsics uses eye/target/up, use perturbation for look_at mode
+        init_eye = (
+            torch.tensor(extrinsics.eye, dtype=torch.float32, device=env.device)
+            .unsqueeze(0)
+            .repeat(num_instance, 1)
+        )
+        init_target = (
+            torch.tensor(extrinsics.target, dtype=torch.float32, device=env.device)
+            .unsqueeze(0)
+            .repeat(num_instance, 1)
+        )
+        init_up = (
+            torch.tensor(extrinsics.up, dtype=torch.float32, device=env.device)
+            .unsqueeze(0)
+            .repeat(num_instance, 1)
+        )
+
+        if eye_range:
+            eye_delta = sample_uniform(
+                lower=torch.tensor(eye_range[0]),
+                upper=torch.tensor(eye_range[1]),
+                size=(num_instance, 3),
+            )
+            new_eye = init_eye + eye_delta
+        else:
+            new_eye = init_eye
+
+        if target_range:
+            target_delta = sample_uniform(
+                lower=torch.tensor(target_range[0]),
+                upper=torch.tensor(target_range[1]),
+                size=(num_instance, 3),
+            )
+            new_target = init_target + target_delta
+        else:
+            new_target = init_target
+
+        if up_range:
+            up_delta = sample_uniform(
+                lower=torch.tensor(up_range[0]),
+                upper=torch.tensor(up_range[1]),
+                size=(num_instance, 3),
+            )
+            new_up = init_up + up_delta
+        else:
+            new_up = init_up
+
+        camera.look_at(new_eye, new_target, new_up, env_ids=env_ids)
+
+    else:
+        logger.log_error("Unsupported extrinsics format for camera randomization.")
+
+
 def randomize_light(
     env: EmbodiedEnv,
     env_ids: Union[torch.Tensor, None],

embodichain/lab/gym/envs/managers/randomization/spatial.py

@@ -121,6 +121,9 @@ def randomize_rigid_object_pose(
         relative_rotation (bool): Whether to randomize the rotation relative to the object's initial rotation. Default is False.
     """
 
+    if entity_cfg.uid not in env.sim.get_rigid_object_uid_list():
+        return
+
     rigid_object: RigidObject = env.sim.get_rigid_object(entity_cfg.uid)
     num_instance = len(env_ids)
 

embodichain/lab/sim/objects/gizmo.py

@@ -331,7 +331,7 @@ def _update_robot_ik(self, target_transform: torch.Tensor):
                     new_qpos = new_qpos.unsqueeze(0)  # Make it (1, dof) for set_qpos
 
                 # Update robot joint positions
-                self.target.set_qpos(qpos=new_qpos, joint_ids=current_joint_ids)
+                self.target.set_qpos(qpos=new_qpos[0], joint_ids=current_joint_ids)
                 return True
             else:
                 logger.log_warning("IK solution not found")

embodichain/lab/sim/objects/rigid_object.py

@@ -516,6 +516,23 @@ def get_visual_material_inst(
         ids = env_ids if env_ids is not None else range(self.num_instances)
         return [self._visual_material[i] for i in ids]
 
+    def share_visual_material_inst(self, mat_insts: List[VisualMaterialInst]) -> None:
+        """Share material instances for the rigid object.
+
+        Args:
+            mat_insts (List[VisualMaterialInst]): List of material instances to share.
+        """
+        if len(self._entities) != len(mat_insts):
+            logger.log_error(
+                f"Length of entities {len(self._entities)} does not match length of material instances {len(mat_insts)}."
+            )
+
+        for i, entity in enumerate(self._entities):
+            if mat_insts[i] is None:
+                continue
+            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:
         """
         Retrieve the body scale for specified environment instances.

embodichain/lab/sim/sim_manager.py

@@ -143,21 +143,12 @@ class SimulationManager:
 
     This class is used to manage the global simulation environment and simulated assets.
         - assets loading, creation, modification and deletion.
-            - assets include robots, fixed actors, dynamic actors and background.
+            - assets include rigid objects, soft objects, articulations, robots, sensors and lights.
         - manager the scenes and the simulation environment.
             - parallel scenes simulation on both CPU and GPU.
-            - sensors arrangement
-            - lighting and indirect lighting
-            - physics simulation parameters control
-        - ...
-
-    Note:
-        1. The arena is used as a standalone space for robots to simulate in. When :meth:`build_multiple_arenas` is called,
-             it will create multiple arenas in a grid pattern. Meanwhile, each simulation assets adding interface will
-             take an additional parameter `arena_index` to specify which arena to place the asset. The name of the asset to
-             be added will be appended with the arena index to avoid name conflict.
-        2. In GUI mode, the physics will be set to a fps (or a wait time for manual mode) for better visualization.
-
+            - create and setup the rendering related settings, eg. environment map, lighting, materials, etc.
+            - physics simulation management, eg. time step, manual update, etc.
+            - interactive control via gizmo and window callbacks events.
 
     Args:
         sim_config (SimulationManagerCfg, optional): simulation configuration. Defaults to SimulationManagerCfg().
@@ -199,11 +190,8 @@ def __init__(
         self._world.set_delta_time(sim_config.physics_dt)
         self._world.show_coordinate_axis(False)
 
-        if sys.platform == "linux":
-            dexsim.set_physics_config(**sim_config.physics_config.to_dexsim_args())
-            dexsim.set_physics_gpu_memory_config(
-                **sim_config.gpu_memory_config.to_dict()
-            )
+        dexsim.set_physics_config(**sim_config.physics_config.to_dexsim_args())
+        dexsim.set_physics_gpu_memory_config(**sim_config.gpu_memory_config.to_dict())
 
         self._is_initialized_gpu_physics = False
         self._ps = self._world.get_physics_scene()
@@ -216,9 +204,10 @@ def __init__(
         self._default_resources = SimResources()
 
         # set unique material path to accelerate material creation.
+        # TODO: This will be removed.
         if self.sim_config.enable_rt is False:
             self._env.set_unique_mat_path(
-                os.path.join(self._material_cache_dir, "dexsim_mat")
+                os.path.join(self._material_cache_dir, "default_mat")
             )
 
         # arena is used as a standalone space for robots to simulate in.
@@ -661,6 +650,8 @@ def get_asset(
             return self._rigid_objects[uid]
         if uid in self._rigid_object_groups:
             return self._rigid_object_groups[uid]
+        if uid in self._soft_objects:
+            return self._soft_objects[uid]
         if uid in self._articulations:
             return self._articulations[uid]
 
@@ -1261,34 +1252,6 @@ def remove_asset(self, uid: str) -> bool:
 
         return False
 
-    def get_asset(
-        self, uid: str
-    ) -> Optional[Union[RigidObject, Articulation, Robot, Light, BaseSensor]]:
-        """Get an asset by its UID.
-
-        The asset can be a rigid object, articulation or robot.
-
-        Args:
-            uid (str): The UID of the asset.
-        """
-        if uid in self._rigid_objects:
-            return self._rigid_objects[uid]
-
-        if uid in self._articulations:
-            return self._articulations[uid]
-
-        if uid in self._robots:
-            return self._robots[uid]
-
-        if uid in self._lights:
-            return self._lights[uid]
-
-        if uid in self._sensors:
-            return self._sensors[uid]
-
-        logger.log_warning(f"Asset {uid} not found.")
-        return None
-
     def draw_marker(
         self,
         cfg: MarkerCfg,

pyproject.toml

@@ -48,7 +48,6 @@ dependencies = [
   "cvxpy==1.4.0",
   "ortools",
   "prettytable",
-  "hdfdict@git+http://69.235.177.182:8081/externalrepo/hdfdict",
   "black==24.3.0",
   "h5py",
 ]