mani-skill-nightly 2025.6.14.814__py3-none-any.whl → 2025.6.14.2004__py3-none-any.whl

mani_skill/envs/tasks/digital_twins/so100_arm/grasp_cube.py

@@ -1,9 +1,11 @@
-from dataclasses import dataclass
-from typing import Any, Dict, Tuple
+from dataclasses import asdict, dataclass
+from typing import Any, Dict, Sequence, Union
 
+import dacite
 import numpy as np
 import sapien
 import torch
+from sapien.render import RenderBodyComponent
 from transforms3d.euler import euler2quat
 
 import mani_skill.envs.utils.randomization as randomization
@@ -25,9 +27,10 @@ from mani_skill.utils.structs.types import GPUMemoryConfig, SimConfig
 class SO100GraspCubeDomainRandomizationConfig:
     ### task agnostic domain randomizations, many of which you can copy over to your own tasks ###
     initial_qpos_noise_scale: float = 0.02
-    randomize_robot_color: bool = True
+    robot_color: Union[str, Sequence[float]] = (1, 1, 1)
+    """Color of the robot in RGB format in scale of 0 to 1 mapping to 0 to 255. If you want to randomize it just set this value to "random"."""
     randomize_lighting: bool = True
-    max_camera_offset: Tuple[float, float, float] = (0.025, 0.025, 0.025)
+    max_camera_offset: Sequence[float] = (0.025, 0.025, 0.025)
     """max camera offset from the base camera position in x, y, and z axes"""
     camera_target_noise: float = 1e-3
     """scale of noise added to the camera target position"""
@@ -37,12 +40,15 @@ class SO100GraspCubeDomainRandomizationConfig:
     """scale of noise added to the camera fov"""
 
     ### task-specific related domain randomizations that occur during scene loading ###
-    cube_half_size_range: Tuple[float, float] = (0.022 / 2, 0.028 / 2)
+    cube_half_size_range: Sequence[float] = (0.022 / 2, 0.028 / 2)
     cube_friction_mean: float = 0.3
     cube_friction_std: float = 0.05
-    cube_friction_bounds: Tuple[float, float] = (0.1, 0.5)
+    cube_friction_bounds: Sequence[float] = (0.1, 0.5)
     randomize_cube_color: bool = True
 
+    def dict(self):
+        return {k: v for k, v in asdict(self).items()}
+
 
 @register_env("SO100GraspCube-v1", max_episode_steps=64)
 class SO100GraspCubeEnv(BaseDigitalTwinEnv):
@@ -69,7 +75,9 @@ class SO100GraspCubeEnv(BaseDigitalTwinEnv):
         robot_uids="so100",
         control_mode="pd_joint_target_delta_pos",
         greenscreen_overlay_path=None,
-        domain_randomization_config=SO100GraspCubeDomainRandomizationConfig(),
+        domain_randomization_config: Union[
+            SO100GraspCubeDomainRandomizationConfig, dict
+        ] = SO100GraspCubeDomainRandomizationConfig(),
         domain_randomization=True,
         base_camera_settings=dict(
             fov=52 * np.pi / 180,
@@ -82,8 +90,18 @@ class SO100GraspCubeEnv(BaseDigitalTwinEnv):
     ):
         self.domain_randomization = domain_randomization
         """whether randomization is turned on or off."""
-        self.domain_randomization_config = domain_randomization_config
+        self.domain_randomization_config = SO100GraspCubeDomainRandomizationConfig()
         """domain randomization config"""
+        merged_domain_randomization_config = self.domain_randomization_config.dict()
+        if isinstance(domain_randomization_config, dict):
+            common.dict_merge(
+                merged_domain_randomization_config, domain_randomization_config
+            )
+            self.domain_randomization_config = dacite.from_dict(
+                data_class=SO100GraspCubeDomainRandomizationConfig,
+                data=domain_randomization_config,
+                config=dacite.Config(strict=True),
+            )
         self.base_camera_settings = base_camera_settings
         """what the camera fov, position and target are when domain randomization is off. DR is centered around these settings"""
 
@@ -142,7 +160,12 @@ class SO100GraspCubeEnv(BaseDigitalTwinEnv):
     def _load_agent(self, options: dict):
         # load the robot arm at this initial pose
         super()._load_agent(
-            options, sapien.Pose(p=[0, 0, 0], q=euler2quat(0, 0, np.pi / 2))
+            options,
+            sapien.Pose(p=[0, 0, 0], q=euler2quat(0, 0, np.pi / 2)),
+            build_separate=True
+            if self.domain_randomization
+            and self.domain_randomization_config.robot_color == "random"
+            else False,
         )
 
     def _load_lighting(self, options: dict):
@@ -251,6 +274,33 @@ class SO100GraspCubeEnv(BaseDigitalTwinEnv):
         builder.initial_pose = sapien.Pose()
         self.camera_mount = builder.build_kinematic("camera_mount")
 
+        # randomize or set a fixed robot color
+        for link in self.agent.robot.links:
+            for i, obj in enumerate(link._objs):
+                # modify the i-th object which is in parallel environment i
+                render_body_component: RenderBodyComponent = (
+                    obj.entity.find_component_by_type(RenderBodyComponent)
+                )
+                if render_body_component is not None:
+                    for render_shape in render_body_component.render_shapes:
+                        for part in render_shape.parts:
+                            if (
+                                self.domain_randomization
+                                and self.domain_randomization_config.robot_color
+                                == "random"
+                            ):
+                                part.material.set_base_color(
+                                    self._batched_episode_rng[i]
+                                    .uniform(low=0.0, high=1.0, size=(3,))
+                                    .tolist()
+                                    + [1]
+                                )
+                            else:
+                                part.material.set_base_color(
+                                    list(self.domain_randomization_config.robot_color)
+                                    + [1]
+                                )
+
     def sample_camera_poses(self, n: int):
         # a custom function to sample random camera poses
         # the way this works is we first sample "eyes", which are the camera positions

{mani_skill_nightly-2025.6.14.814.dist-info → mani_skill_nightly-2025.6.14.2004.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mani-skill-nightly
-Version: 2025.6.14.814
+Version: 2025.6.14.2004
 Summary: ManiSkill3: A Unified Benchmark for Generalizable Manipulation Skills
 Home-page: https://github.com/haosulab/ManiSkill
 Author: ManiSkill contributors

{mani_skill_nightly-2025.6.14.814.dist-info → mani_skill_nightly-2025.6.14.2004.dist-info}/RECORD

@@ -569,7 +569,7 @@ mani_skill/envs/tasks/digital_twins/bridge_dataset_eval/__init__.py,sha256=2okxy
 mani_skill/envs/tasks/digital_twins/bridge_dataset_eval/base_env.py,sha256=HFjZwcMOSNduDmyD-qo9zVeHqHOdXdx0H11japbTpEI,21997
 mani_skill/envs/tasks/digital_twins/bridge_dataset_eval/put_on_in_scene.py,sha256=kDKEX-e1Hq5kZiBs9xfEzU2jjx35eHWriAAYFJcBbgE,9419
 mani_skill/envs/tasks/digital_twins/so100_arm/__init__.py,sha256=uehHSCaHoZDB9awMtF81r4A-X4fM6NFmErzOsgzgSs4,42
-mani_skill/envs/tasks/digital_twins/so100_arm/grasp_cube.py,sha256=vJMNx6tmPVMP1ttiB0AMo-a5PhQR0y8GBsbu2X_GR3E,19397
+mani_skill/envs/tasks/digital_twins/so100_arm/grasp_cube.py,sha256=U0ldL-lDMipcjEUhrlWvn5Kzh_r1TkkjfdUxj9OGyX0,21875
 mani_skill/envs/tasks/drawing/__init__.py,sha256=eZ7CnzIWBOHUihvOKStSQK94eKSv5wchRdSodiDYBlw,72
 mani_skill/envs/tasks/drawing/draw.py,sha256=WlhxJjt0-DlkxC3t-o0M8BoOwdwWpM9reupFg5OqiZc,8146
 mani_skill/envs/tasks/drawing/draw_svg.py,sha256=UKnVl_tMfUNHmOg24Ny-wFynb5CaZC0uIHvW9sBxbyo,16206
@@ -824,8 +824,8 @@ mani_skill/vector/wrappers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJ
 mani_skill/vector/wrappers/gymnasium.py,sha256=0kNe_iquf8J3503obkbmCKNaYr5m3cPrC_hBZS0Nl9Y,7408
 mani_skill/vector/wrappers/sb3.py,sha256=SlXdiEPqcNHYMhJCzA29kBU6zK7DKTe1nc0L6Z3QQtY,4722
 mani_skill/viewer/__init__.py,sha256=srvDBsk4LQU75K2VIttrhiQ68p_ro7PSDqQRls2PY5c,1722
-mani_skill_nightly-2025.6.14.814.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-mani_skill_nightly-2025.6.14.814.dist-info/METADATA,sha256=5xnWX8qPJnA3aZ-J_IYKKIhiKte0AsNTrdTuoHUKS9A,9271
-mani_skill_nightly-2025.6.14.814.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-mani_skill_nightly-2025.6.14.814.dist-info/top_level.txt,sha256=bkBgOVl_MZMoQx2aRFsSFEYlZLxjWlip5vtJ39FB3jA,11
-mani_skill_nightly-2025.6.14.814.dist-info/RECORD,,
+mani_skill_nightly-2025.6.14.2004.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+mani_skill_nightly-2025.6.14.2004.dist-info/METADATA,sha256=apDEpbBIfYmUimW8Nsj_OP8RHhfx_UH4q8I8wCfb2jY,9272
+mani_skill_nightly-2025.6.14.2004.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+mani_skill_nightly-2025.6.14.2004.dist-info/top_level.txt,sha256=bkBgOVl_MZMoQx2aRFsSFEYlZLxjWlip5vtJ39FB3jA,11
+mani_skill_nightly-2025.6.14.2004.dist-info/RECORD,,