Add files using upload-large-folder tool

GR00T-WholeBodyControl/decoupled_wbc/dexmg/gr00trobocasa/robocasa/environments/locomanipulation/__init__.py

@@ -0,0 +1,13 @@
+"""
+Full list of loco-manipulation tasks.
+
+GroundOnly - ground only environments
+
+locomanip_pnp - factory environments, pick and place tasks:
+LMBottlePnP
+LMBoxPnP
+"""
+
+from .base import REGISTERED_LOCOMANIPULATION_ENVS
+
+ALL_LOCOMANIPULATION_ENVIRONMENTS = REGISTERED_LOCOMANIPULATION_ENVS.keys()

GR00T-WholeBodyControl/decoupled_wbc/dexmg/gr00trobocasa/robocasa/environments/locomanipulation/base.py

@@ -0,0 +1,1658 @@
+from copy import deepcopy
+import os
+from typing import Optional, Type
+import warnings
+import xml.etree.ElementTree as ET
+
+import mujoco
+import numpy as np
+import robosuite
+from robosuite.environments.base import EnvMeta
+from robosuite.environments.manipulation.manipulation_env import ManipulationEnv
+from robosuite.models.arenas import Arena
+from robosuite.models.tasks import ManipulationTask
+from robosuite.utils.mjcf_utils import array_to_string, find_elements, xml_path_completion
+from robosuite.utils.observables import Observable, sensor
+
+import robocasa
+from robocasa.models.objects.objects import MJCFObject
+from robocasa.models.scenes import GroundArena
+import robocasa.utils.camera_utils as CamUtils
+from robocasa.utils.dexmg_utils import DexMGConfigHelper
+from robocasa.utils.object_utils import check_obj_upright
+from robocasa.utils.visuals_utls import Gradient, randomize_materials_rgba
+
+REGISTERED_LOCOMANIPULATION_ENVS = {}
+
+
+def register_locomanipulation_env(target_class):
+    REGISTERED_LOCOMANIPULATION_ENVS[target_class.__name__] = target_class
+
+
+class LocoManipulationEnvMeta(EnvMeta):
+    """Metaclass for registering robocasa environments"""
+
+    def __new__(meta, name, bases, class_dict):
+        cls = super().__new__(meta, name, bases, class_dict)
+        register_locomanipulation_env(cls)
+        return cls
+
+
+class CameraPoseRandomizer:
+    @staticmethod
+    def randomize_cameras(
+        env: "LocoManipulationEnv",
+        cam_names: list[str],
+        pos_range: tuple[np.ndarray, np.ndarray],
+        euler_range: tuple[np.ndarray, np.ndarray],
+    ):
+        """
+        Randomize camera poses while maintaining their relative transforms.
+
+        Args:
+            env: The environment instance
+            cam_names: List of camera names to randomize together
+            pos_range: Tuple of (min_pos, max_pos) as 3D arrays for position randomization
+            euler_range: Tuple of (min_euler, max_euler) as 3D arrays for euler angle randomization (in radians)
+        """
+        if len(cam_names) == 0:
+            return
+
+        # Sample random transform offset
+        random_pos_offset = env.rng.uniform(pos_range[0], pos_range[1])
+        random_euler_offset = env.rng.uniform(euler_range[0], euler_range[1])
+
+        # Convert euler offset to quaternion
+        quat_offset = np.zeros(4, dtype=float)
+        mujoco.mju_euler2Quat(quat_offset, random_euler_offset, "xyz")
+
+        # Apply the same transform to all specified cameras
+        for cam_name in cam_names:
+            if cam_name not in env._cam_configs:
+                warnings.warn(f"Camera {cam_name} not found in camera configs. Skipping.")
+                continue
+
+            cam_config = env._cam_configs[cam_name]
+
+            # Get original position and quaternion
+            original_pos = np.array(cam_config["pos"], dtype=float)
+            original_quat = np.array(cam_config["quat"], dtype=float)
+
+            # Apply rotation offset to position (rotate position offset by the random rotation)
+            rotated_offset = np.zeros(3, dtype=float)
+            mujoco.mju_rotVecQuat(rotated_offset, random_pos_offset, original_quat)
+            new_pos = original_pos + rotated_offset
+
+            # Compose quaternions: new_quat = quat_offset * original_quat
+            new_quat = np.zeros(4, dtype=float)
+            mujoco.mju_mulQuat(new_quat, quat_offset, original_quat)
+
+            # Update camera config
+            cam_config["pos"] = new_pos.tolist()
+            cam_config["quat"] = new_quat.tolist()
+
+            # Update in simulation if already created
+            if hasattr(env, "sim") and env.sim is not None:
+                try:
+                    cam_id = env.sim.model.camera_name2id(cam_name)
+                    env.sim.model.cam_pos[cam_id] = new_pos
+                    env.sim.model.cam_quat[cam_id] = new_quat
+                except:
+                    # Camera might not be in the model yet
+                    pass
+
+
+class RobotPoseRandomizer:
+    @staticmethod
+    def set_pose(
+        env: "LocoManipulationEnv",
+        x_range: [tuple[float, float]],
+        y_range: [tuple[float, float]],
+        yaw_range: [tuple[float, float]],
+    ):
+        new_x = env.rng.uniform(*x_range)
+        new_y = env.rng.uniform(*y_range)
+        new_yaw = env.rng.uniform(*yaw_range)
+
+        if env.robots[0].name == "G1":
+            base_offset = env.ROBOT_POS_OFFSETS[env.robots[0].robot_model.__class__.__name__]
+            target_pos = np.array([new_x, new_y, base_offset[2]], dtype=float)
+            quat = np.zeros(4, dtype=float)
+            mujoco.mju_euler2Quat(quat, np.array([0.0, 0.0, new_yaw]), "xyz")
+            base_freejoint = f"{env.robots[0].robot_model.naming_prefix}base"
+            if base_freejoint in env.sim.model.joint_names:
+                env.sim.data.set_joint_qpos(base_freejoint, np.concatenate([target_pos, quat]))
+            else:
+                warnings.warn(f"Base joint {base_freejoint} not found in the model.")
+        else:
+            base_joint_pos = np.array([new_x, new_y, new_yaw])
+            base_joint_names = [
+                "mobilebase0_joint_mobile_forward",
+                "mobilebase0_joint_mobile_side",
+                "mobilebase0_joint_mobile_yaw",
+            ]
+            for i, base_joint_name in enumerate(base_joint_names):
+                if base_joint_name not in env.sim.model.joint_names:
+                    warnings.warn(
+                        f"Base joint {base_joint_name} not found in the model. "
+                        f"Skipping randomization of {base_joint_name}."
+                    )
+                else:
+                    env.sim.data.set_joint_qpos(base_joint_name, base_joint_pos[i])
+
+    @staticmethod
+    def set_arm(env: ManipulationEnv, elbow_qpos: float, shoulder_pitch_qpos: float):
+        """Helper function to reinitialize G1 robot arm configuration."""
+        robot = env.robots[0]
+        if "G1" not in robot.name:
+            # avoid reinitializing arm configuration for non-G1 robots
+            return
+
+        joint_names = robot.robot_joints
+        joint_pos_indices = robot._ref_joint_pos_indexes
+        for joint_name, pos_idx in zip(joint_names, joint_pos_indices):
+            if "elbow" in joint_name:
+                print(f"reinitializing G1 {joint_name} with idx {pos_idx} to {elbow_qpos}")
+                env.sim.data.qpos[pos_idx] = elbow_qpos
+            elif "shoulder_pitch" in joint_name:
+                print(f"reinitializing G1 {joint_name} with idx {pos_idx} to {shoulder_pitch_qpos}")
+                env.sim.data.qpos[pos_idx] = shoulder_pitch_qpos
+
+
+class LocoManipulationEnv(ManipulationEnv, metaclass=LocoManipulationEnvMeta):
+    """
+    Initialized a Base Ground Standing environment.
+    """
+
+    MUJOCO_ARENA_CLS: Type[Arena] = GroundArena
+
+    ROBOT_POS_OFFSETS: dict[str, list[float]] = {
+        "PandaOmron": [0, 0, 0],
+        "GR1FloatingBody": [0, 0, 0.97],
+        "GR1": [0, 0, 0.97],
+        "GR1FixedLowerBody": [0, 0, 0.97],
+        "GR1FixedLowerBodyInspireHands": [0, 0, 0.97],
+        "GR1FixedLowerBodyFourierHands": [0, 0, 0.97],
+        "GR1ArmsOnly": [0, 0, 0.97],
+        "GR1ArmsOnlyInspireHands": [0, 0, 0.97],
+        "GR1ArmsOnlyFourierHands": [0, 0, 0.97],
+        "GR1ArmsAndWaistFourierHands": [0, 0, 0.97],
+        "G1": [0, 0, 0.793],
+        "G1FixedBase": [0, 0, 0.793],
+        "G1FixedLowerBody": [0, 0, 0.793],
+        "G1ArmsOnly": [0, 0, 0.793],
+        "G1ArmsOnlyFloating": [0, 0, 0.793],
+        "G1FloatingBody": [0, 0, 0.793],
+        "G1FloatingBodyWithVertical": [0, 0, 0.793],
+    }
+
+    def __init__(
+        self,
+        translucent_robot: bool = False,
+        use_object_obs: bool = False,
+        randomize_cameras: bool = False,
+        *args,
+        **kwargs,
+    ):
+        self.mujoco_objects = []
+        self.randomize_cameras = randomize_cameras
+
+        super().__init__(
+            *args,
+            **kwargs,
+        )
+
+        self.translucent_robot = translucent_robot
+
+    def _load_model(self):
+        super()._load_model()
+
+        self.mujoco_arena = self.MUJOCO_ARENA_CLS()
+        self.mujoco_arena.set_origin([0, 0, 0])
+        self.set_cameras()
+
+        self.model = ManipulationTask(
+            mujoco_arena=self.mujoco_arena,
+            mujoco_robots=[robot.robot_model for robot in self.robots],
+            mujoco_objects=self.mujoco_objects,
+        )
+
+        robot_base_pos = self.ROBOT_POS_OFFSETS[self.robots[0].robot_model.__class__.__name__]
+        robot_model = self.robots[0].robot_model
+        robot_model.set_base_xpos(robot_base_pos)
+        # robot_model.set_base_ori(robot_base_ori)
+
+    def set_cameras(self):
+        """
+        Adds new tabletop-relevant cameras to the environment. Will randomize cameras if specified.
+        """
+
+        self._cam_configs = deepcopy(CamUtils.CAM_CONFIGS)
+
+        for robot in self.robots:
+            if hasattr(robot.robot_model, "get_camera_configs"):
+                self._cam_configs.update(robot.robot_model.get_camera_configs())
+
+        for cam_name, cam_cfg in self._cam_configs.items():
+            if cam_cfg.get("parent_body", None) is not None:
+                continue
+
+            self.mujoco_arena.set_camera(
+                camera_name=cam_name,
+                pos=cam_cfg["pos"],
+                quat=cam_cfg["quat"],
+                camera_attribs=cam_cfg.get("camera_attribs", None),
+            )
+
+        self.mujoco_arena.set_camera(
+            camera_name="egoview",
+            pos=[0.078, 0, 1.308],
+            quat=[0.66491268, 0.24112495, -0.24112507, -0.66453637],
+            camera_attribs=dict(fovy="90"),
+        )
+
+    def visualize(self, vis_settings):
+        """
+        In addition to super call, make the robot semi-transparent
+
+        Args:
+            vis_settings (dict): Visualization keywords mapped to T/F, determining whether that specific
+                component should be visualized. Should have "grippers" keyword as well as any other relevant
+                options specified.
+        """
+        # Run superclass method first
+        super().visualize(vis_settings=vis_settings)
+
+        visual_geom_names = []
+
+        for robot in self.robots:
+            robot_model = robot.robot_model
+            visual_geom_names += robot_model.visual_geoms
+
+        for name in visual_geom_names:
+            rgba = self.sim.model.geom_rgba[self.sim.model.geom_name2id(name)]
+            if self.translucent_robot:
+                rgba[-1] = 0.10
+            else:
+                rgba[-1] = 1.0
+
+    def reward(self, action=None):
+        """
+        Reward function for the task. The reward function is based on the task
+        and to be implemented in the subclasses. Returns 0 by default.
+
+        Returns:
+            float: Reward for the task
+        """
+        reward = 0
+        if self._check_success():
+            reward = 1.0
+        return reward
+
+    def _check_success(self):
+        """
+        Checks if the task has been successfully completed.
+        Success condition is based on the task and to be implemented in the
+        subclasses. Returns False by default.
+
+        Returns:
+            bool: True if the task is successfully completed, False otherwise
+        """
+        return False
+
+    def edit_model_xml(self, xml_str):
+        """
+        This function postprocesses the model.xml collected from a MuJoCo demonstration
+        for retrospective model changes.
+
+        Args:
+            xml_str (str): Mujoco sim demonstration XML file as string
+
+        Returns:
+            str: Post-processed xml file as string
+        """
+        xml_str = super().edit_model_xml(xml_str)
+
+        tree = ET.fromstring(xml_str)
+        root = tree
+        worldbody = root.find("worldbody")
+        actuator = root.find("actuator")
+        asset = root.find("asset")
+        meshes = asset.findall("mesh")
+        textures = asset.findall("texture")
+        all_elements = meshes + textures
+
+        robosuite_path_split = os.path.split(robosuite.__file__)[0].split("/")
+        robocasa_path_split = os.path.split(robocasa.__file__)[0].split("/")
+
+        # replace robocasa-specific asset paths
+        for elem in all_elements:
+            old_path = elem.get("file")
+            if old_path is None:
+                continue
+
+            old_path_split = old_path.split("/")
+            # maybe replace all paths to robosuite assets
+            if "models/assets" in old_path:
+                if "/robosuite/" in old_path:
+                    check_lst = [
+                        loc for loc, val in enumerate(old_path_split) if val == "robosuite"
+                    ]
+                    ind = max(check_lst)  # last occurrence index
+                    new_path_split = robosuite_path_split + old_path_split[ind + 1 :]
+                elif "/robocasa/" in old_path:
+                    check_lst = [loc for loc, val in enumerate(old_path_split) if val == "robocasa"]
+                    ind = max(check_lst)  # last occurrence index
+                    new_path_split = robocasa_path_split + old_path_split[ind + 1 :]
+                else:
+                    raise ValueError
+
+                new_path = "/".join(new_path_split)
+                elem.set("file", new_path)
+
+        # set cameras
+        for cam_name, cam_config in self._cam_configs.items():
+            parent_body = cam_config.get("parent_body", None)
+
+            cam_root = worldbody
+            if parent_body is not None:
+                cam_root = find_elements(root=worldbody, tags="body", attribs={"name": parent_body})
+                if cam_root is None:
+                    # camera config refers to body that doesnt exist on the robot
+                    continue
+
+            cam = find_elements(root=cam_root, tags="camera", attribs={"name": cam_name})
+
+            if cam is None:
+                old_cam = find_elements(root=worldbody, tags="camera", attribs={"name": cam_name})
+                if old_cam is not None:
+                    # old camera associated with different body
+                    continue
+
+                cam = ET.Element("camera")
+                cam.set("mode", "fixed")
+                cam.set("name", cam_name)
+                cam_root.append(cam)
+
+            cam.set("pos", array_to_string(cam_config["pos"]))
+            cam.set("quat", array_to_string(cam_config["quat"]))
+            for k, v in cam_config.get("camera_attribs", {}).items():
+                cam.set(k, v)
+
+        # replace base -> mobilebase (this is needed for old PandaOmron demos)
+        for elem in find_elements(
+            root=worldbody, tags=["geom", "site", "body", "joint"], return_first=False
+        ):
+            if elem.get("name") is None:
+                continue
+            if elem.get("name").startswith("base0_"):
+                old_name = elem.get("name")
+                new_name = "mobilebase0_" + old_name[6:]
+                elem.set("name", new_name)
+        for elem in find_elements(
+            root=actuator,
+            tags=["velocity", "position", "motor", "general"],
+            return_first=False,
+        ):
+            if elem.get("name") is None:
+                continue
+            if elem.get("name").startswith("base0_"):
+                old_name = elem.get("name")
+                new_name = "mobilebase0_" + old_name[6:]
+                elem.set("name", new_name)
+        for elem in find_elements(
+            root=actuator,
+            tags=["velocity", "position", "motor", "general"],
+            return_first=False,
+        ):
+            if elem.get("joint") is None:
+                continue
+            if elem.get("joint").startswith("base0_"):
+                old_joint = elem.get("joint")
+                new_joint = "mobilebase0_" + old_joint[6:]
+                elem.set("joint", new_joint)
+
+        # result = ET.tostring(root, encoding="utf8").decode("utf8")
+        result = ET.tostring(root).decode("utf8")
+
+        # # replace with generative textures
+        # if (self.generative_textures is not None) and (
+        #     self.generative_textures is not False
+        # ):
+        #     # sample textures
+        #     assert self.generative_textures == "100p"
+        #     self._curr_gen_fixtures = get_random_textures(self.rng)
+
+        #     cab_tex = self._curr_gen_fixtures["cab_tex"]
+        #     counter_tex = self._curr_gen_fixtures["counter_tex"]
+        #     wall_tex = self._curr_gen_fixtures["wall_tex"]
+        #     floor_tex = self._curr_gen_fixtures["floor_tex"]
+
+        #     result = replace_cab_textures(
+        #         self.rng, result, new_cab_texture_file=cab_tex
+        #     )
+        #     result = replace_counter_top_texture(
+        #         self.rng, result, new_counter_top_texture_file=counter_tex
+        #     )
+        #     result = replace_wall_texture(
+        #         self.rng, result, new_wall_texture_file=wall_tex
+        #     )
+        #     result = replace_floor_texture(
+        #         self.rng, result, new_floor_texture_file=floor_tex
+        #     )
+
+        return result
+
+    def _setup_references(self):
+        super()._setup_references()
+
+        self.obj_body_id = {}
+
+    def _randomize_robot_cameras(self):
+        """Randomize the poses of robot-mounted cameras while preserving their relative transforms."""
+        cam_names = ["robot0_oak_egoview", "robot0_oak_left_monoview", "robot0_oak_right_monoview"]
+
+        # Define randomization ranges
+        pos_range = (
+            np.array([-0.02, -0.02, -0.02]),  # min position offset [x, y, z] in meters
+            np.array([0.02, 0.02, 0.02]),  # max position offset [x, y, z] in meters
+        )
+        euler_range = (
+            np.array([-0.1, -0.1, -0.1]),  # min euler angles [roll, pitch, yaw] in radians
+            np.array([0.1, 0.1, 0.1]),  # max euler angles [roll, pitch, yaw] in radians
+        )
+
+        CameraPoseRandomizer.randomize_cameras(
+            env=self, cam_names=cam_names, pos_range=pos_range, euler_range=euler_range
+        )
+
+    def _reset_internal(self):
+        super()._reset_internal()
+
+        if self.randomize_cameras:
+            self._randomize_robot_cameras()
+
+    def _reset_observables(self):
+        if self.hard_reset:
+            self._observables = self._setup_observables()
+
+        # these sensors need a lot of computation, so we disable them by default for speed up simulation
+        disabled_sensors = [
+            "base_to_left_eef_pos",
+            "base_to_left_eef_quat",
+            "base_to_left_eef_quat_site",
+            "base_to_right_eef_pos",
+            "base_to_right_eef_quat",
+            "base_to_right_eef_quat_site",
+        ]
+        for name in disabled_sensors:
+            for robot in self.robots:
+                robot_name_prefix = robot.robot_model.naming_prefix
+                if f"{robot_name_prefix}{name}" in self._observables:
+                    self._observables[f"{robot_name_prefix}{name}"].set_enabled(False)
+                    self._observables[f"{robot_name_prefix}{name}"].set_active(False)
+
+    def get_state(self):
+        return {"states": self.sim.get_state().flatten()}
+
+
+class GroundOnly(LocoManipulationEnv):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+
+class PrimitiveBottle:
+    DEFAULT_RGB = [0.3, 0.7, 0.8]
+
+    def __init__(
+        self,
+        name="bottle",
+        radius: float = 0.03,
+        half_height: float = 0.075,
+        rgb: Optional[list[float]] = None,
+    ):
+        self.name = name
+        self.assets = [
+            ET.Element(
+                "texture",
+                type="2d",
+                name=f"{name}_tex",
+                builtin="flat",
+                rgb1=" ".join(map(str, self.DEFAULT_RGB if rgb is None else rgb)),
+                width="512",
+                height="512",
+            ),
+            ET.Element(
+                "material",
+                name=f"{name}_mat",
+                texture=f"{name}_tex",
+                texuniform="true",
+                reflectance="0.1",
+            ),
+        ]
+
+        self.body = ET.Element("body", name=f"{self.name}_body", pos="0.35 0 0.8")
+        bottle_vis_geom = ET.Element(
+            "geom",
+            name=f"{name}_vis",
+            pos="0 0 0",
+            size=f"{radius} {half_height}",
+            type="cylinder",
+            material=f"{name}_mat",
+            group="1",
+            conaffinity="0",
+            contype="0",
+        )
+        self.body.append(bottle_vis_geom)
+
+        # Cylinder collider approximation for stable contacts
+        self.contact_geoms = []
+        n_sides = 3
+        half_width = radius * np.tan(np.pi / n_sides / 2)
+        for i in range(n_sides):
+            coll_name = f"{self.name}_collider_{i}"
+            angle = np.pi / n_sides * i
+            quat = np.zeros(4)
+            euler = np.array([0, 0, angle])
+            mujoco.mju_euler2Quat(quat, euler, "xyz")
+            box_geom = ET.Element(
+                "geom",
+                name=coll_name,
+                type="box",
+                pos="0 0 0",
+                size=f"{radius} {half_width} {half_height}",
+                quat=" ".join(map(str, quat)),
+                solimp="0.998 0.998 0.001",
+                solref="0.001 2",
+                density="100",
+                friction="0.95 0.3 0.1",
+            )
+            self.body.append(box_geom)
+            self.contact_geoms.append(coll_name)
+
+        bottle_joint = ET.Element(
+            "joint",
+            name=f"{self.name}_joint",
+            type="free",
+            damping="0.0005",
+        )
+        self.body.append(bottle_joint)
+
+
+class PrimitiveFixture:
+    DEFAULT_RGB = [0.8, 0.8, 0.8]
+
+    def __init__(
+        self,
+        name: str,
+        pos: np.ndarray = np.array([0.0, 0.0, 0.8]),
+        half_size: np.ndarray = np.array([0.1, 0.1, 0.001]),
+        rgb: Optional[str] = None,
+    ):
+        """
+        A simple primitive fixture as a flat box.
+
+        Args:
+            half_size: Half-sizes in [x, y, z] directions. Default creates a 20cm x 20cm x 2mm box.
+        """
+        self.half_size = half_size
+
+        self.assets = [
+            ET.Element(
+                "texture",
+                type="2d",
+                name=f"{name}",
+                builtin="flat",
+                rgb1=" ".join(map(str, self.DEFAULT_RGB if rgb is None else rgb)),
+                width="512",
+                height="512",
+            ),
+            ET.Element(
+                "material",
+                name=f"{name}",
+                texture=f"{name}",
+                texuniform="true",
+                reflectance="0.05",  # Less reflective than bottle
+            ),
+        ]
+
+        self.body = ET.Element("body", name=f"{name}_body", pos=array_to_string(pos))
+
+        # Visual geometry
+        fixture_vis_geom = ET.Element(
+            "geom",
+            name=f"{name}_vis",
+            pos="0 0 0",
+            size=f"{half_size[0]} {half_size[1]} {half_size[2]}",
+            type="box",
+            material=f"{name}",
+            group="1",
+            conaffinity="0",
+            contype="0",
+        )
+        self.body.append(fixture_vis_geom)
+
+        # Collision geometry - just a single box since it's already a simple shape
+        self.contact_geoms = []
+        fixture_collider = ET.Element(
+            "geom",
+            name=f"{name}_collider",
+            type="box",
+            pos="0 0 0",
+            size=f"{half_size[0]} {half_size[1]} {half_size[2]}",
+            solimp="0.998 0.998 0.001",
+            solref="0.001 2",
+            density="100",
+            friction="0.6 0.01 0.001",  # Similar to add_fixture_body friction
+        )
+        self.body.append(fixture_collider)
+        self.contact_geoms.append("fixture_collider")
+
+
+class PnPBottle(LocoManipulationEnv, DexMGConfigHelper):
+    TABLE_GRADIENT: Gradient = Gradient(
+        np.array([0.68, 0.34, 0.07, 1.0]), np.array([1.0, 1.0, 1.0, 1.0])
+    )
+    DEFAULT_BOTTLE_POS: np.ndarray = np.array([0.4, 0, 0.77])
+    BOTTLE_POS_RANGE_X = (-0.08, 0.04)
+    BOTTLE_POS_RANGE_Y = (-0.08, 0.08)
+
+    def __init__(self, *args, **kwargs):
+        self.objects = {}
+        super().__init__(*args, **kwargs)
+
+    def _load_model(self):
+        self.mujoco_objects = [self._create_table("table_body", [0.5, 0, 0], [0, 0, np.pi / 2])]
+
+        super()._load_model()
+
+        self.bottle = self._create_bottle()
+
+    @staticmethod
+    def _create_table(name: str, position: list[float], euler: list[float]) -> MJCFObject:
+        table = MJCFObject(
+            name=name,
+            mjcf_path=xml_path_completion(
+                "objects/omniverse/locomanip/lab_table/model.xml", root=robocasa.models.assets_root
+            ),
+            scale=1.0,
+            solimp=(0.998, 0.998, 0.001),
+            solref=(0.001, 1),
+            density=10,
+            friction=(1, 1, 1),
+            static=True,
+        )
+        table.set_pos(position)
+        table.set_euler(euler)
+        return table
+
+    def _create_bottle(
+        self, name: str = "bottle", rgb: Optional[list[float]] = None
+    ) -> PrimitiveBottle:
+        bottle = PrimitiveBottle(name=name, radius=0.03, half_height=0.075, rgb=rgb)
+        self.model.asset.extend(bottle.assets)
+        self.model.worldbody.append(bottle.body)
+        self.objects[name] = {"name": f"{name}_body"}
+        return bottle
+
+    def _reset_internal(self):
+        """
+        Resets simulation internal configurations.
+        """
+        super()._reset_internal()
+
+        if not self.deterministic_reset:
+            self._randomize_bottle_placement()
+            self._randomize_table_texture()
+
+    def _randomize_bottle_placement(
+        self, name: str = "bottle", base_pos: Optional[np.ndarray] = None
+    ):
+        if not self.deterministic_reset:
+            bottle_joint = f"{name}_joint"
+            base_pos = self.DEFAULT_BOTTLE_POS if base_pos is None else base_pos
+
+            random_x = self.rng.uniform(*self.BOTTLE_POS_RANGE_X)
+            random_y = self.rng.uniform(*self.BOTTLE_POS_RANGE_Y)
+            new_pos = base_pos + np.array([random_x, random_y, 0])
+
+            current_qpos = self.sim.data.get_joint_qpos(bottle_joint)
+            new_qpos = current_qpos.copy()
+            new_qpos[:3] = new_pos
+
+            self.sim.data.set_joint_qpos(bottle_joint, new_qpos)
+
+    def _randomize_table_texture(self):
+        table = self.mujoco_objects[0]
+        randomize_materials_rgba(
+            rng=self.rng, mjcf_obj=table, gradient=self.TABLE_GRADIENT, linear=True
+        )
+
+    def _setup_references(self):
+        super()._setup_references()
+
+        self.obj_body_id = {}
+        for name, model in self.objects.items():
+            self.obj_body_id[name] = self.sim.model.body_name2id(model["name"])
+
+    def _check_success(self):
+        check_grasp = self._check_grasp(self.robots[0].gripper["right"], self.bottle.contact_geoms)
+
+        bottle_z = self.sim.data.body_xpos[self.obj_body_id["bottle"]][2]
+        table_z = self.mujoco_objects[0].top_offset[2]
+        check_bottle_in_air = bottle_z > table_z + 0.2
+        # check bottle and table collision
+        # check_bottle_in_air = not self.check_contact("bottle", "table")
+        return check_grasp and check_bottle_in_air
+
+    def get_object(self):
+        return dict(
+            bottle=dict(obj_name=self.objects["bottle"]["name"], obj_type="body"),
+        )
+
+    def get_subtask_term_signals(self):
+        signals = dict()
+        signals["grasp_bottle"] = int(
+            self._check_grasp(self.robots[0].gripper["right"], self.bottle.contact_geoms)
+        )
+        return signals
+
+    @staticmethod
+    def task_config():
+        task = DexMGConfigHelper.AttrDict()
+        task.task_spec_0.subtask_1 = dict(
+            object_ref="bottle",
+            subtask_term_signal=None,
+            subtask_term_offset_range=None,
+            selection_strategy="random",
+            selection_strategy_kwargs=None,
+            action_noise=0.05,
+            num_interpolation_steps=5,
+            num_fixed_steps=0,
+            apply_noise_during_interpolation=False,
+        )
+        task.task_spec_1.subtask_1 = dict(
+            object_ref=None,
+            subtask_term_signal=None,
+            subtask_term_offset_range=None,
+            selection_strategy="random",
+            selection_strategy_kwargs=None,
+            action_noise=0.05,
+            num_interpolation_steps=5,
+            num_fixed_steps=0,
+            apply_noise_during_interpolation=False,
+        )
+        return task.to_dict()
+
+
+def create_shelf(pos: list[float], euler: list[float]) -> MJCFObject:
+    shelf = MJCFObject(
+        name="shelf_body",
+        mjcf_path=xml_path_completion(
+            "objects/aigc/shelf/model.xml", root=robocasa.models.assets_root
+        ),
+        scale=[1.0, 1.0, 1.0],
+        solimp=(0.998, 0.998, 0.001),
+        solref=(0.001, 1),
+        density=10,
+        friction=(1, 1, 1),
+        static=True,
+    )
+    shelf.set_pos(pos)
+    shelf.set_euler(euler)
+    return shelf
+
+
+class PickBottleShelf(PnPBottle):
+    def _load_model(self):
+        # Create both the original table and the target table
+        self.mujoco_objects = [create_shelf(pos=[0.8, 0.4, 0], euler=[0, 0, np.pi / 2])]
+
+        LocoManipulationEnv._load_model(self)
+
+        self.bottle = self._create_bottle()
+
+    def _reset_internal(self):
+        """
+        Resets simulation internal configurations.
+        """
+        LocoManipulationEnv._reset_internal(self)
+
+        if not self.deterministic_reset:
+            # Base position on ground (z=0.075 is bottle radius)
+            # Level 2 of shelf
+            self._randomize_bottle_placement(base_pos=np.array([0.7, 0.4, 0.376660 + 0.075 + 0.02]))
+            self._randomize_table_texture()
+            RobotPoseRandomizer.set_arm(self, elbow_qpos=-0.5, shoulder_pitch_qpos=0.5)
+
+
+class PnPBottleHigh(PnPBottle):
+    def _load_model(self):
+        self.mujoco_objects = [self._create_table("table_body", [0.5, 0, 0.1], [0, 0, np.pi / 2])]
+
+        LocoManipulationEnv._load_model(self)
+
+        self.bottle = self._create_bottle()
+
+    def _reset_internal(self):
+        """
+        Resets simulation internal configurations.
+        """
+        LocoManipulationEnv._reset_internal(self)
+
+        # Randomize bottle position within +/- 0.1 range on x and y axes
+        if not self.deterministic_reset:
+            # Base position of the bottle
+            base_pos = np.array([0.4, 0, 0.875])
+
+            # Add random offset within +/- 0.1 range for x and y
+            random_x = np.random.uniform(-0.1, 0.1)
+            random_y = np.random.uniform(-0.1, 0.1)
+            # New randomized position (keep z constant)
+            new_pos = base_pos + np.array([random_x, random_y, 0])
+
+            # Set the bottle position using the free joint
+            # For free joints, qpos includes [x, y, z, qw, qx, qy, qz]
+            current_qpos = self.sim.data.get_joint_qpos("bottle_joint")
+            new_qpos = current_qpos.copy()
+            new_qpos[:3] = new_pos  # Update position (x, y, z)
+
+            self.sim.data.set_joint_qpos("bottle_joint", new_qpos)
+
+    def _setup_observables(self):
+        observables = super()._setup_observables()
+
+        @sensor(modality="object")
+        def obj_pos(obs_cache):
+            return self.sim.data.body_xpos[self.obj_body_id["bottle"]]
+
+        @sensor(modality="object")
+        def obj_quat(obs_cache):
+            return self.sim.data.body_xquat[self.obj_body_id["bottle"]]
+
+        @sensor(modality="object")
+        def obj_linear_vel(obs_cache):
+            return self.sim.data.get_body_xvelp("bottle_body")
+
+        @sensor(modality="object")
+        def obj_angular_vel(obs_cache):
+            return self.sim.data.get_body_xvelr("bottle_body")
+
+        sensors = [obj_pos, obj_quat, obj_linear_vel, obj_angular_vel]
+        names = [s.__name__ for s in sensors]
+
+        for name, s in zip(names, sensors):
+            observables[name] = Observable(
+                name=name,
+                sensor=s,
+                sampling_rate=self.control_freq,
+            )
+
+        return observables
+
+    def get_privileged_obs_keys(self):
+        return {
+            "obj_pos": (3,),
+            "obj_quat": (4,),
+            "obj_linear_vel": (3,),
+            "obj_angular_vel": (3,),
+        }
+
+
+class NavPickBottle(PnPBottle):
+    """
+    Pick-and-Place Bottle environment with robot position randomized at reset.
+    """
+
+    def _reset_internal(self):
+        super()._reset_internal()
+
+        if not self.deterministic_reset:
+            RobotPoseRandomizer.set_pose(self, (-0.3, -0.16), (-0.2, 0.2), (-np.pi / 6, np.pi / 6))
+
+
+class PnPBottleRandRobotPose(NavPickBottle):
+    pass
+
+
+class VisualReach(LocoManipulationEnv):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def _load_model(self):
+        super()._load_model()
+
+        self.create_visual_only_goal_cube()
+
+    def create_visual_only_goal_cube(self):
+        cube_tex = ET.Element(
+            "texture",
+            type="2d",
+            name="cube",
+            builtin="flat",
+            rgb1="1.0 0.0 0.0",
+            width="512",
+            height="512",
+        )
+        cube_mat = ET.Element(
+            "material",
+            name="cube",
+            texture="cube",
+            texuniform="true",
+            reflectance="0.1",
+        )
+        self.model.asset.append(cube_tex)
+        self.model.asset.append(cube_mat)
+
+        self.objects = {}
+        cube_body = ET.Element("body", name="cube_body", pos="0.4 0 0.875")
+
+        cube_vis_geom = ET.Element(
+            "geom",
+            name="cube_vis",
+            pos="0 0 0",
+            size="0.0375 0.0375 0.0375",
+            type="box",
+            material="cube",
+            group="1",
+            conaffinity="0",
+            contype="0",
+        )
+
+        cube_body.append(cube_vis_geom)
+        self.model.worldbody.append(cube_body)
+        self.objects["cube"] = {"name": "cube_body"}
+
+    def _setup_references(self):
+        super()._setup_references()
+
+        self.obj_body_id = {}
+        for name, model in self.objects.items():
+            self.obj_body_id[name] = self.sim.model.body_name2id(model["name"])
+
+    def _check_success(self):
+        # check_grasp = self._check_grasp(self.robots[0].gripper["right"], self.objects["bottle"])
+        # check_reach = self._check_reach(self.objects["bottle"])
+        return True
+
+    def _check_reach(self, obj_name):
+        raise NotImplementedError
+        # To be implemented by the subclass
+
+    def get_object(self):
+        return dict(
+            cube=dict(obj_name=self.objects["cube"].root_body, obj_type="body"),
+        )
+
+    def reset_obj_pos(self):
+        # reset object pos randomly around bottle_body pos="0.4 0 0.875"
+        init_pos = np.array([0.4, 0, 0.875])
+        random_x = np.random.uniform(-0.3, 0.15)
+        random_y = np.random.uniform(-0.15, 0.15)
+        random_z = np.random.uniform(-0.15, 0.30)
+        self.sim.model.body_pos[self.obj_body_id["cube"]] = init_pos + np.array(
+            [random_x, random_y, random_z]
+        )
+
+    def set_cameras(self):
+        super().set_cameras()
+        self.mujoco_arena.set_camera(
+            camera_name="egoview",
+            pos=[0.078, 0, 1.308],
+            quat=[0.66491268, 0.24112495, -0.24112507, -0.66453637],
+            camera_attribs=dict(fovy="90"),
+        )
+
+    def _setup_observables(self):
+        observables = super()._setup_observables()
+
+        @sensor(modality="object")
+        def obj_pos(obs_cache):
+            return self.sim.data.body_xpos[self.obj_body_id["cube"]]
+
+        @sensor(modality="object")
+        def obj_quat(obs_cache):
+            return self.sim.data.body_xquat[self.obj_body_id["cube"]]
+
+        @sensor(modality="object")
+        def obj_linear_vel(obs_cache):
+            return self.sim.data.get_body_xvelp("cube_body")
+
+        @sensor(modality="object")
+        def obj_angular_vel(obs_cache):
+            return self.sim.data.get_body_xvelr("cube_body")
+
+        sensors = [obj_pos, obj_quat, obj_linear_vel, obj_angular_vel]
+        names = [s.__name__ for s in sensors]
+
+        for name, s in zip(names, sensors):
+            observables[name] = Observable(
+                name=name,
+                sensor=s,
+                sampling_rate=self.control_freq,
+            )
+
+        return observables
+
+    def get_privileged_obs_keys(self):
+        return {
+            "obj_pos": (3,),
+            "obj_quat": (4,),
+            "obj_linear_vel": (3,),
+            "obj_angular_vel": (3,),
+        }
+
+
+class PnPBottleFixtureToFixture(PnPBottle):
+    """
+    Task: Robot picks up bottle and places it on a fixture.
+
+    Initialization: bottle rests on a source fixture.
+
+    Idea: by changing the location of the target fixture, we can change the data generation task layout for
+    these placement related tasks.
+    """
+
+    SEPARATION_THRESH_M: float = 0.0005  # 0.05 cm
+    DISTMAX_SCAN_M: float = 0.05  # 5 cm window for distance queries
+    _SRC_NAME = "start_fixture"
+    _TGT_NAME = "target_fixture"
+    _FIXTURE_HALF_SIZE = np.array([0.05, 0.05, 0.001])
+    _BOTTLE_HALF_HEIGHT = 0.075
+    _X_SRC_RANGE = (0.30, 0.55)
+    _X_TGT_RANGE = (0.30, 0.55)
+    _Y_SRC_RANGE = (-0.20, -0.05)
+    _Y_TGT_RANGE = (0.05, 0.20)
+    _SRC_FIXTURE_VISIBLE = False
+    _TGT_FIXTURE_VISIBLE = True
+
+    def _load_model(self):
+        self.mujoco_objects = [self._create_table("table_body", [0.5, 0, 0], [0, 0, np.pi / 2])]
+        LocoManipulationEnv._load_model(self)
+        self.bottle = self._create_bottle()
+        self._create_fixture(self._SRC_NAME, visible=self._SRC_FIXTURE_VISIBLE, rgb="1 0 0")
+        self._create_fixture(self._TGT_NAME, visible=self._TGT_FIXTURE_VISIBLE, rgb="0 1 0")
+        self._src_body = f"{self._SRC_NAME}_body"
+        self._tgt_body = f"{self._TGT_NAME}_body"
+        self._src_coll = f"{self._SRC_NAME}_collider"
+        self._tgt_coll = f"{self._TGT_NAME}_collider"
+
+    def _create_fixture(self, name: str, visible: bool, rgb: Optional[str] = None) -> None:
+        """Create a flat box fixture; add assets + body to the compiled model."""
+        fx = PrimitiveFixture(
+            name=name, pos=np.array([0.0, 0.0, 0.0]), half_size=self._FIXTURE_HALF_SIZE, rgb=rgb
+        )
+
+        # Make source fixture invisible (keep collision only)
+        if not visible:
+            # Find the visual geom and hide it
+            for child in list(fx.body):
+                if child.tag == "geom" and child.get("name") == f"{name}_vis":
+                    child.set("rgba", "0 0 0 0")  # invisible visual
+                    break
+
+        # Register assets + body into the scene graph
+        self.model.asset.extend(fx.assets)
+        self.model.worldbody.append(fx.body)
+
+    def _setup_references(self):
+        super()._setup_references()
+        # Table root body is "<name>_main" (same convention as target_table above)
+        self.table_body_id = self.sim.model.body_name2id("table_body_main")
+        self.src_fixture_id = self.sim.model.body_name2id(self._src_body)
+        self.tgt_fixture_id = self.sim.model.body_name2id(self._tgt_body)
+
+    def _check_success(self) -> bool:
+        """Bottle touches target fixture collider and is upright."""
+        bottle_on_target = self.check_contact(self.bottle.contact_geoms, [self._tgt_coll])
+        bottle_upright = check_obj_upright(self, "bottle", threshold=0.8, symmetric=True)
+        return bottle_on_target and bottle_upright
+
+    # --- runtime table height ---
+    def _table_top_z(self) -> float:
+        base_z = float(self.sim.data.body_xpos[self.table_body_id][2])
+        top_offset_z = float(self.mujoco_objects[0].top_offset[2])
+        return base_z + top_offset_z
+
+    def _reset_internal(self):
+        LocoManipulationEnv._reset_internal(self)
+
+        if not self.deterministic_reset:
+            # Sample fixture XY, compute Z from current table pose
+            x_src = self.rng.uniform(*self._X_SRC_RANGE)
+            y_src = self.rng.uniform(*self._Y_SRC_RANGE)
+
+            x_tgt = self.rng.uniform(*self._X_TGT_RANGE)
+            y_tgt = self.rng.uniform(*self._Y_TGT_RANGE)
+            # y_tgt = self.rng.uniform(*self._Y_TGT_RANGE)
+
+            z_top = self._table_top_z()  # dynamic table top
+            src_pos = np.array([x_src, y_src, z_top])
+            tgt_pos = np.array([x_tgt, y_tgt, z_top])
+
+            # Reset fixture body poses (static bodies): write to model; MuJoCo will use it after forward()
+            self.sim.model.body_pos[self.src_fixture_id] = src_pos
+            self.sim.model.body_pos[self.tgt_fixture_id] = tgt_pos
+
+            # Place bottle on source fixture: top of fixture + bottle half-height + tiny clearance
+            bottle_z = z_top + self._FIXTURE_HALF_SIZE[2] + self._BOTTLE_HALF_HEIGHT + 0.002
+            qpos = self.sim.data.get_joint_qpos("bottle_joint").copy()
+            qpos[:3] = np.array([x_src, y_src, bottle_z])
+            qpos[3:7] = np.array([1.0, 0.0, 0.0, 0.0])  # upright
+            self.sim.data.set_joint_qpos("bottle_joint", qpos)
+
+            self._randomize_table_texture()
+            RobotPoseRandomizer.set_pose(self, (-0.3, -0.16), (-0.2, 0.2), (-np.pi / 6, np.pi / 6))
+
+    # --- distance via MuJoCo ---
+    def _min_signed_distance_mj(self, geoms_a: list[str], geoms_b: list[str]) -> float:
+        model, data = self.sim.model, self.sim.data
+        dmin = np.inf
+        fromto = np.empty(6, dtype=np.float64)
+        a_ids = [model.geom_name2id(n) for n in geoms_a]
+        b_ids = [model.geom_name2id(n) for n in geoms_b]
+        for ga in a_ids:
+            for gb in b_ids:
+                dist = mujoco.mj_geomDistance(
+                    model._model, data._data, ga, gb, self.DISTMAX_SCAN_M + 0.01, fromto
+                )
+                dmin = min(dmin, float(dist))
+        return dmin
+
+    def get_subtask_term_signals(self) -> dict[str, int]:
+        """
+        1 iff (no contact between bottle and source fixture) AND
+             (min signed distance > DISTMAX_SCAN_M).
+        """
+        in_contact = self.check_contact(self.bottle.contact_geoms, [self._src_coll])
+        min_dist = self._min_signed_distance_mj(self.bottle.contact_geoms, [self._src_coll])
+        return {
+            "obj_off_source_fixture": int((not in_contact) and (min_dist > self.DISTMAX_SCAN_M))
+        }
+
+    def get_object(self) -> dict:
+        return dict(
+            bottle=dict(obj_name=self.objects["bottle"]["name"], obj_type="body"),
+            source_fixture=dict(obj_name=self._src_body, obj_type="body"),
+            target_fixture=dict(obj_name=self._tgt_body, obj_type="body"),
+        )
+
+    @staticmethod
+    def task_config() -> dict:
+        task = DexMGConfigHelper.AttrDict()
+        # Subtask 1: pick (leave source fixture)
+        task.task_spec_0.subtask_1 = dict(
+            object_ref="bottle",
+            subtask_term_signal="obj_off_source_fixture",
+            subtask_term_offset_range=(5, 10),
+            selection_strategy="random",
+            selection_strategy_kwargs=None,
+            action_noise=0.05,
+            num_interpolation_steps=5,
+            num_fixed_steps=0,
+            apply_noise_during_interpolation=False,
+        )
+        # Subtask 2: place on target fixture
+        task.task_spec_0.subtask_2 = dict(
+            object_ref="target_fixture",
+            subtask_term_signal=None,
+            subtask_term_offset_range=None,
+            selection_strategy="random",
+            selection_strategy_kwargs=None,
+            action_noise=0.05,
+            num_interpolation_steps=5,
+            num_fixed_steps=0,
+            apply_noise_during_interpolation=False,
+        )
+        # Default filler for task_spec_1, mirroring other tasks
+        task.task_spec_1.subtask_1 = dict(
+            object_ref=None,
+            subtask_term_signal=None,
+            subtask_term_offset_range=None,
+            selection_strategy="random",
+            selection_strategy_kwargs=None,
+            action_noise=0.05,
+            num_interpolation_steps=5,
+            num_fixed_steps=0,
+            apply_noise_during_interpolation=False,
+        )
+        return task.to_dict()
+
+
+class PnPBottleFixtureToFixtureSourceDemo(PnPBottleFixtureToFixture):
+    """
+    Environment for collecting source demo for PnPBottleFixtureToFixture tasks.
+    """
+
+    _X_SRC_RANGE = (0.375, 0.375)
+    _X_TGT_RANGE = (0.375, 0.375)
+    _Y_SRC_RANGE = (-0.15, -0.15)
+    _Y_TGT_RANGE = (0.1, 0.1)
+    _SRC_FIXTURE_VISIBLE = False
+    _TGT_FIXTURE_VISIBLE = True
+
+
+class PnPBottleShelfToTable(PnPBottleFixtureToFixture):
+    """
+    Task: Robot picks up bottle from a fixture on a shelf and places it on a fixture on a table.
+
+    Initialization: bottle rests on a source fixture on the shelf.
+    Target: place bottle on target fixture on the table.
+    """
+
+    # Adjust ranges for shelf-to-table layout
+    _X_SRC_RANGE = (-0.05, 0.05)  # Shelf position range
+    _X_TGT_RANGE = (-0.05 - 0.2, 0.05 - 0.2)  # Table position range
+    # TODO: could be better to have some 'center' specified here
+    _Y_SRC_RANGE = (-0.05, 0.05)  # Shelf position range
+    _Y_TGT_RANGE = (-0.05, 0.05)  # Table position range
+    _SRC_FIXTURE_VISIBLE = True
+    _TGT_FIXTURE_VISIBLE = True
+    _FIXTURE_HALF_SIZE = np.array([0.05, 0.05, 0.001])
+
+    # Shelf height constants (from PnPBottleShelf)
+    # _SHELF_HEIGHT = 0.386660  # Level 2 of shelf from the original shelf environment
+    _SHELF_HEIGHT = 0.753321 + 0.015  # Level 3 of shelf from the original shelf environment
+
+    def _load_model(self):
+        # Create both shelf and table
+        self.mujoco_objects = [
+            self._create_table("table_body", [0.5, 0.6, 0], [0, 0, np.pi / 2]),
+            create_shelf(pos=[0.8, -0.4, 0], euler=[0, 0, np.pi / 2]),
+        ]
+
+        LocoManipulationEnv._load_model(self)
+
+        self.bottle = self._create_bottle()
+        self._create_fixture(self._SRC_NAME, visible=self._SRC_FIXTURE_VISIBLE, rgb="1 0 0")
+        self._create_fixture(self._TGT_NAME, visible=self._TGT_FIXTURE_VISIBLE, rgb="0 1 0")
+        self._src_body = f"{self._SRC_NAME}_body"
+        self._tgt_body = f"{self._TGT_NAME}_body"
+        self._src_coll = f"{self._SRC_NAME}_collider"
+        self._tgt_coll = f"{self._TGT_NAME}_collider"
+
+    def _setup_references(self):
+        super()._setup_references()
+        # Add reference to shelf
+        self.shelf_body_id = self.sim.model.body_name2id("shelf_body_main")
+
+    def _shelf_top_z(self) -> float:
+        """Get the Z coordinate of the shelf top surface"""
+        # Use the same shelf height as in PnPBottleShelf
+        return self._SHELF_HEIGHT
+
+    def _shelf_xy(self) -> tuple[float, float]:
+        """Get the XY coordinates of the shelf"""
+        return self.sim.data.body_xpos[self.shelf_body_id][:2]
+
+    def _table_xy(self) -> tuple[float, float]:
+        """Get the XY coordinates of the table"""
+        return self.sim.data.body_xpos[self.table_body_id][:2]
+
+    def _reset_internal(self):
+        LocoManipulationEnv._reset_internal(self)
+
+        if not self.deterministic_reset:
+            # Sample fixture XY positions
+            x_src = self.rng.uniform(*self._X_SRC_RANGE)
+            y_src = self.rng.uniform(*self._Y_SRC_RANGE)
+
+            x_tgt = self.rng.uniform(*self._X_TGT_RANGE)
+            y_tgt = self.rng.uniform(*self._Y_TGT_RANGE)
+
+            # Source fixture on shelf
+            z_shelf = self._shelf_top_z()
+            x_shelf, y_shelf = self._shelf_xy()
+            src_pos = np.array([x_src, y_src, z_shelf])
+            src_pos += np.array([x_shelf, y_shelf, 0])
+
+            # table pos
+            # Target fixture on table
+            z_table = self._table_top_z()
+            x_table, y_table = self._table_xy()
+            tgt_pos = np.array([x_tgt, y_tgt, z_table])
+            tgt_pos += np.array([x_table, y_table, 0])
+
+            # Reset fixture body poses
+            self.sim.model.body_pos[self.src_fixture_id] = src_pos
+            self.sim.model.body_pos[self.tgt_fixture_id] = tgt_pos
+
+            # Place bottle on source fixture (shelf): top of fixture + bottle half-height + clearance
+            bottle_z = z_shelf + self._FIXTURE_HALF_SIZE[2] + self._BOTTLE_HALF_HEIGHT + 0.002
+            qpos = self.sim.data.get_joint_qpos("bottle_joint").copy()
+            qpos[:3] = np.array([src_pos[0], src_pos[1], bottle_z])
+            qpos[3:7] = np.array([1.0, 0.0, 0.0, 0.0])  # upright
+            self.sim.data.set_joint_qpos("bottle_joint", qpos)
+
+            self._randomize_table_texture()
+            RobotPoseRandomizer.set_pose(self, (-0.3, -0.16), (-0.2, 0.2), (-np.pi / 6, np.pi / 6))
+
+    def _randomize_table_texture(self):
+        """Randomize textures for the table (shelf texture is static)"""
+        # Only randomize the table texture (index 1), not the shelf
+        table = self.mujoco_objects[1]
+        randomize_materials_rgba(
+            rng=self.rng, mjcf_obj=table, gradient=self.TABLE_GRADIENT, linear=True
+        )
+
+
+class PnPBottleTableToTable(PnPBottle):
+    def _load_model(self):
+        # Create both the original table and the target table
+        self.mujoco_objects = [
+            self._create_table("table_body", [0.5, 0, 0], [0, 0, np.pi / 2]),
+            self._create_table("target_table_body", [0.5, 1.2, 0], [0, 0, np.pi / 2]),
+        ]
+
+        LocoManipulationEnv._load_model(self)
+
+        self.bottle = self._create_bottle()
+
+    def _setup_references(self):
+        super()._setup_references()
+
+        # Add reference to target table - note the _main suffix
+        self.target_table_body_id = self.sim.model.body_name2id("target_table_body_main")
+
+    def _check_success(self):
+        """Check if bottle is successfully placed on the target table"""
+        bottle_on_table = self.check_contact(self.bottle.contact_geoms, self.mujoco_objects[1])
+        bottle_is_upright = check_obj_upright(self, "bottle", threshold=0.8, symmetric=True)
+        return bottle_on_table and bottle_is_upright
+
+    def _randomize_table_texture(self):
+        """Randomize textures for both tables"""
+        # Randomize original table
+        original_table = self.mujoco_objects[0]
+        randomize_materials_rgba(
+            rng=self.rng, mjcf_obj=original_table, gradient=self.TABLE_GRADIENT, linear=True
+        )
+
+        # Randomize target table
+        target_table = self.mujoco_objects[1]
+        randomize_materials_rgba(
+            rng=self.rng, mjcf_obj=target_table, gradient=self.TABLE_GRADIENT, linear=True
+        )
+
+    def get_object(self):
+        return dict(
+            bottle=dict(obj_name=self.objects["bottle"]["name"], obj_type="body"),
+            target_table=dict(obj_name="target_table_body_main", obj_type="body"),
+        )
+
+    @staticmethod
+    def task_config():
+        task = DexMGConfigHelper.AttrDict()
+        task.task_spec_0.subtask_1 = dict(
+            object_ref="bottle",
+            subtask_term_signal="obj_off_table",
+            subtask_term_offset_range=(5, 10),
+            selection_strategy="random",
+            selection_strategy_kwargs=None,
+            action_noise=0.05,
+            num_interpolation_steps=5,
+            num_fixed_steps=0,
+            apply_noise_during_interpolation=False,
+        )
+        # Second subtask for placing on target table
+        task.task_spec_0.subtask_2 = dict(
+            object_ref="target_table",
+            subtask_term_signal=None,
+            subtask_term_offset_range=None,
+            selection_strategy="random",
+            selection_strategy_kwargs=None,
+            action_noise=0.05,
+            num_interpolation_steps=5,
+            num_fixed_steps=0,
+            apply_noise_during_interpolation=False,
+        )
+        task.task_spec_1.subtask_1 = dict(
+            object_ref=None,
+            subtask_term_signal=None,
+            subtask_term_offset_range=None,
+            selection_strategy="random",
+            selection_strategy_kwargs=None,
+            action_noise=0.05,
+            num_interpolation_steps=5,
+            num_fixed_steps=0,
+            apply_noise_during_interpolation=False,
+        )
+        return task.to_dict()
+
+    def get_subtask_term_signals(self):
+        """
+        Retrieve signals used to define subtask termination conditions.
+
+        Returns:
+            dict: Dictionary mapping signal names to their current values
+        """
+        signals = dict()
+
+        obj_z = self.sim.data.body_xpos[self.obj_body_id["bottle"]][2]
+        target_table_pos = self.sim.data.body_xpos[self.target_table_body_id]
+        target_table_z = target_table_pos[2] + self.mujoco_objects[1].top_offset[2]
+
+        th = 0.15
+        signals["obj_off_table"] = int(obj_z - target_table_z > th)
+
+        return signals
+
+
+class PickBottleGround(PnPBottle):
+    """
+    Pick-and-Place Bottle environment with bottle initialized on the ground.
+    """
+
+    def _reset_internal(self):
+        """
+        Resets simulation internal configurations.
+        """
+        LocoManipulationEnv._reset_internal(self)
+
+        if not self.deterministic_reset:
+            # Base position on ground (z=0.075 is bottle radius)
+            self._randomize_bottle_placement(base_pos=np.ndarray([0.4, 0, 0.075]))
+            self._randomize_table_texture()
+
+    def _randomize_table_texture(self):
+        pass
+
+    def _check_success(self):
+        check_grasp = self._check_grasp(self.robots[0].gripper["right"], "bottle")
+
+        bottle_z = self.sim.data.body_xpos[self.obj_body_id["bottle"]][2]
+        ground_z = 0
+        check_bottle_in_air = bottle_z > ground_z + 0.2
+        # check bottle and table collision
+        # check_bottle_in_air = not self.check_contact("bottle", "table")
+        return check_grasp and check_bottle_in_air
+
+    def _load_model(self):
+        self.mujoco_objects = []
+
+        super(PnPBottle, self)._load_model()
+        self._create_bottle()
+
+
+class PickBottles(PnPBottle):
+    BOTTLE_POS_RANGE_X = (-0.08, 0.04)
+    BOTTLE_POS_RANGE_Y = (-0.04, 0.04)
+
+    COLOURS: list[list[float]] = [[0.3, 0.7, 0.8], [0.8, 0.4, 0.3]]
+    BOTTLES_COUNT = 2
+    Y_OFFSET_STEP = 0.1
+
+    @staticmethod
+    def _get_bottle_names() -> list[str]:
+        return [f"bottle_{i}" for i in range(PickBottles.BOTTLES_COUNT)]
+
+    def _load_model(self):
+        self.mujoco_objects = [self._create_table("table_body", [0.5, 0, 0], [0, 0, np.pi / 2])]
+
+        LocoManipulationEnv._load_model(self)
+
+        self.bottles = self._create_bottles()
+
+    def _create_bottles(self) -> list[PrimitiveBottle]:
+        bottles = []
+        for i, name in enumerate(self._get_bottle_names()):
+            rgb = self.COLOURS[i % len(self.COLOURS)]
+            bottles.append(self._create_bottle(name=name, rgb=rgb))
+        return bottles
+
+    def _reset_internal(self):
+        LocoManipulationEnv._reset_internal(self)
+
+        n = len(self.bottles)
+        offsets = np.arange(n) - (n - 1) / 2.0
+        for i, bottle in enumerate(self.bottles):
+            self._randomize_bottle_placement(
+                name=bottle.name,
+                base_pos=self.DEFAULT_BOTTLE_POS
+                + np.array([0, self.Y_OFFSET_STEP * offsets[i], 0]),
+            )
+            self._randomize_table_texture()
+
+    def _check_success(self):
+        for bottle in self.bottles:
+            check_grasp = self._check_grasp(
+                self.robots[0].gripper["right"], bottle.contact_geoms
+            ) or self._check_grasp(self.robots[0].gripper["left"], bottle.contact_geoms)
+            bottle_z = self.sim.data.body_xpos[self.obj_body_id[bottle.name]][2]
+            table_z = self.mujoco_objects[0].top_offset[2]
+            check_bottle_in_air = bottle_z > table_z + 0.2
+            if check_grasp and check_bottle_in_air:
+                continue
+            return False
+        return True
+
+    def get_object(self):
+        result = {}
+        for bottle in self.bottles:
+            result[bottle.name] = dict(obj_name=self.objects[bottle.name]["name"], obj_type="body")
+        return result
+
+    def get_subtask_term_signals(self):
+        signals = dict()
+        for bottle in self.bottles:
+            signals[f"grasp_{bottle.name}"] = int(
+                self._check_grasp(self.robots[0].gripper["right"], bottle.contact_geoms)
+                or self._check_grasp(self.robots[0].gripper["left"], bottle.contact_geoms)
+            )
+        return signals
+
+    @staticmethod
+    def task_config():
+        task = DexMGConfigHelper.AttrDict()
+        bottle_names = PickBottles._get_bottle_names()
+        assert len(bottle_names) == 2
+        for i, name in enumerate(bottle_names):
+            subtask = dict(
+                object_ref=name,
+                subtask_term_signal=f"grasp_{name}",
+                subtask_term_offset_range=None,
+                selection_strategy="random",
+                selection_strategy_kwargs=None,
+                action_noise=0.05,
+                num_interpolation_steps=5,
+                num_fixed_steps=0,
+                apply_noise_during_interpolation=False,
+            )
+            spec_attr = f"task_spec_{i}"
+            setattr(getattr(task, spec_attr), "subtask_1", subtask)
+        return task.to_dict()
+
+
+class NavPickBottles(PickBottles):
+    """
+    PickBottles environment with robot position randomized further from table at reset.
+    """
+
+    def _reset_internal(self):
+        super()._reset_internal()
+
+        if not self.deterministic_reset:
+            RobotPoseRandomizer.set_pose(self, (-0.3, -0.16), (-0.2, 0.2), (-np.pi / 6, np.pi / 6))
+
+
+class PnPBottlesTableToTable(PickBottles):
+    def _load_model(self):
+        self.mujoco_objects = [
+            self._create_table("table_body", [0.5, 0, 0], [0, 0, np.pi / 2]),
+            self._create_table("target_table_body", [0.5, 1.2, 0], [0, 0, np.pi / 2]),
+        ]
+
+        LocoManipulationEnv._load_model(self)
+
+        self.bottles = self._create_bottles()
+
+    def _check_success(self):
+        """Check if bottles are successfully placed on the target table"""
+        for bottle in self.bottles:
+            bottle_on_table = self.check_contact(bottle.contact_geoms, self.mujoco_objects[1])
+            bottle_is_upright = check_obj_upright(self, bottle.name, threshold=0.8, symmetric=True)
+            if bottle_on_table and bottle_is_upright:
+                continue
+            return False
+        return True
+
+    def _setup_references(self):
+        super()._setup_references()
+
+        # Add reference to target table - note the _main suffix
+        self.target_table_body_id = self.sim.model.body_name2id("target_table_body_main")
+
+    def get_object(self):
+        result = super().get_object()
+        result["target_table"] = dict(obj_name="target_table_body_main", obj_type="body")
+        return result
+
+    @staticmethod
+    def task_config():
+        task = DexMGConfigHelper.AttrDict()
+
+        bottle_names = PickBottles._get_bottle_names()
+        assert len(bottle_names) == 2
+        for i, name in enumerate(bottle_names):
+
+            # pick subtask per arm
+            subtask = dict(
+                object_ref=name,
+                subtask_term_signal=f"{name}_off_table",
+                subtask_term_offset_range=None,
+                selection_strategy="random",
+                selection_strategy_kwargs=None,
+                action_noise=0.05,
+                num_interpolation_steps=5,
+                num_fixed_steps=0,
+                apply_noise_during_interpolation=False,
+            )
+            spec_attr = f"task_spec_{i}"
+            setattr(getattr(task, spec_attr), "subtask_1", subtask)
+
+            # place subtask per arm
+            subtask = dict(
+                object_ref="target_table",
+                subtask_term_signal=None,
+                subtask_term_offset_range=None,
+                selection_strategy="random",
+                selection_strategy_kwargs=None,
+                action_noise=0.05,
+                num_interpolation_steps=5,
+                num_fixed_steps=0,
+                apply_noise_during_interpolation=False,
+            )
+            spec_attr = f"task_spec_{i}"
+            setattr(getattr(task, spec_attr), "subtask_2", subtask)
+
+        return task.to_dict()
+
+    def get_subtask_term_signals(self):
+        signals = dict()
+        for bottle in self.bottles:
+            obj_z = self.sim.data.body_xpos[self.obj_body_id[bottle.name]][2]
+            target_table_pos = self.sim.data.body_xpos[self.target_table_body_id]
+            target_table_z = target_table_pos[2] + self.mujoco_objects[1].top_offset[2]
+            th = 0.15
+            signals[f"{bottle.name}_off_table"] = int(obj_z - target_table_z > th)
+        return signals

GR00T-WholeBodyControl/decoupled_wbc/dexmg/gr00trobocasa/robocasa/environments/locomanipulation/locomanip.py

@@ -0,0 +1,83 @@
+from abc import abstractmethod
+from typing import Optional
+
+from robocasa.environments.locomanipulation.base import LocoManipulationEnv
+from robocasa.models.scenes import GroundArena
+from robocasa.models.scenes.factory_arena import FactoryArena
+from robocasa.utils.scene.configs import SceneConfig, SceneScaleConfig
+from robocasa.utils.scene.scene import Scene, SceneObject
+from robocasa.utils.scene.success_criteria import SuccessCriteria
+
+
+class LMEnvBase(LocoManipulationEnv):
+    SCENE_SCALE = SceneScaleConfig()
+
+    def __init__(
+        self,
+        translucent_robot: bool = False,
+        use_object_obs: bool = False,
+        scene_scale: Optional[SceneScaleConfig] = None,
+        *args,
+        **kwargs,
+    ):
+        self.scene_scale = scene_scale or self.SCENE_SCALE
+        super().__init__(translucent_robot, use_object_obs, *args, **kwargs)
+
+    def _load_model(self):
+        self.scene = Scene(self, self._get_env_config(), self.scene_scale)
+        self.mujoco_objects = self.scene.mujoco_objects
+
+        super()._load_model()
+
+    def _reset_internal(self):
+        """
+        Resets simulation internal configurations.
+        """
+        super()._reset_internal()
+
+        if not self.deterministic_reset:
+            self.scene.reset()
+
+    def _setup_references(self):
+        super()._setup_references()
+
+        self.obj_body_id = {}
+        for obj in self.mujoco_objects:
+            self.obj_body_id[obj.name] = self.sim.model.body_name2id(obj.root_body)
+
+    def _get_env_config(self) -> SceneConfig:
+        return SceneConfig(
+            objects=self._get_objects(),
+            success=self._get_success_criteria(),
+            instruction=self._get_instruction(),
+        )
+
+    @abstractmethod
+    def _get_objects(self) -> list[SceneObject]:
+        raise NotImplementedError
+
+    @abstractmethod
+    def _get_success_criteria(self) -> SuccessCriteria:
+        raise NotImplementedError
+
+    @abstractmethod
+    def _get_instruction(self) -> str:
+        raise NotImplementedError
+
+    def _check_success(self):
+        return self.scene.success()
+
+    def get_ep_meta(self):
+        ep_meta = super().get_ep_meta()
+        ep_meta["lang"] = self.scene.instruction
+        return ep_meta
+
+
+# noinspection PyAbstractClass
+class LMSimpleEnv(LMEnvBase):
+    MUJOCO_ARENA_CLS = GroundArena
+
+
+# noinspection PyAbstractClass
+class LMFactoryEnv(LMEnvBase):
+    MUJOCO_ARENA_CLS = FactoryArena

GR00T-WholeBodyControl/decoupled_wbc/dexmg/gr00trobocasa/robocasa/environments/locomanipulation/locomanip_basic.py

@@ -0,0 +1,732 @@
+import numpy as np
+from robocasa.environments.locomanipulation.base import RobotPoseRandomizer
+from robocasa.environments.locomanipulation.locomanip import LMSimpleEnv
+from robocasa.utils.dexmg_utils import DexMGConfigHelper
+from robocasa.utils.scene.configs import (
+    ObjectConfig,
+    ReferenceConfig,
+    SamplingConfig,
+    SceneScaleConfig,
+)
+from robocasa.utils.scene.scene import SceneObject
+from robocasa.utils.scene.success_criteria import (
+    AllCriteria,
+    AnyCriteria,
+    IsGrasped,
+    IsInContact,
+    IsPositionInRange,
+    IsRobotInRange,
+    IsUpright,
+    NotCriteria,
+    SuccessCriteria,
+)
+from robocasa.utils.visuals_utls import Gradient, randomize_materials_rgba
+
+
+class LMPickBottle(LMSimpleEnv, DexMGConfigHelper):
+    SCENE_SCALE = SceneScaleConfig(planar_scale=1.0)
+
+    TABLE_GRADIENT: Gradient = Gradient(
+        np.array([0.68, 0.34, 0.07, 1.0]), np.array([1.0, 1.0, 1.0, 1.0])
+    )
+    LIFT_OFFSET = 0.1
+
+    def _get_objects(self) -> list[SceneObject]:
+        self.table = SceneObject(
+            ObjectConfig(
+                name="table",
+                mjcf_path="objects/omniverse/locomanip/lab_table/model.xml",
+                scale=1.0,
+                static=True,
+                sampler_config=SamplingConfig(
+                    x_range=np.array([-0.02, 0.02]),
+                    y_range=np.array([-0.02, 0.02]),
+                    reference_pos=np.array([0.5, 0, 0]),
+                    rotation=np.array([np.pi * 0.5, np.pi * 0.5]),
+                ),
+            )
+        )
+        self.bottle = SceneObject(
+            ObjectConfig(
+                name="bottle",
+                mjcf_path="objects/omniverse/locomanip/jug_a01/model.xml",
+                static=False,
+                scale=0.6,
+                sampler_config=SamplingConfig(
+                    x_range=np.array([-0.08, 0.04]),
+                    y_range=np.array([-0.08, 0.08]),
+                    rotation=np.array([-np.pi, np.pi]),
+                    reference_pos=np.array([0.4, 0, self.table.mj_obj.top_offset[2]]),
+                ),
+            )
+        )
+        return [self.table, self.bottle]
+
+    def _get_success_criteria(self) -> SuccessCriteria:
+        return AllCriteria(
+            IsGrasped(self.bottle, "right"),
+            IsPositionInRange(self.bottle, 2, self.table.mj_obj.top_offset[2] + self.LIFT_OFFSET),
+        )
+
+    def _get_instruction(self) -> str:
+        return "Pick up the bottle."
+
+    def get_object(self):
+        return dict(
+            bottle=dict(obj_name=self.bottle.mj_obj.root_body, obj_type="body"),
+        )
+
+    def get_subtask_term_signals(self):
+        signals = dict()
+        signals["grasp_bottle"] = int(
+            self._check_grasp(self.robots[0].gripper["right"], self.bottle.mj_obj.contact_geoms)
+        )
+        return signals
+
+    @staticmethod
+    def task_config():
+        task = DexMGConfigHelper.AttrDict()
+        task.task_spec_0.subtask_1 = dict(
+            object_ref="bottle",
+            subtask_term_signal=None,
+            subtask_term_offset_range=None,
+            selection_strategy="random",
+            selection_strategy_kwargs=None,
+            action_noise=0.05,
+            num_interpolation_steps=5,
+            num_fixed_steps=0,
+            apply_noise_during_interpolation=False,
+        )
+        task.task_spec_1.subtask_1 = dict(
+            object_ref=None,
+            subtask_term_signal=None,
+            subtask_term_offset_range=None,
+            selection_strategy="random",
+            selection_strategy_kwargs=None,
+            action_noise=0.05,
+            num_interpolation_steps=5,
+            num_fixed_steps=0,
+            apply_noise_during_interpolation=False,
+        )
+        return task.to_dict()
+
+    def _reset_internal(self):
+        super()._reset_internal()
+
+        if not self.deterministic_reset:
+            self._randomize_table_rgba()
+
+    def _randomize_table_rgba(self):
+        randomize_materials_rgba(
+            rng=self.rng, mjcf_obj=self.table.mj_obj, gradient=self.TABLE_GRADIENT, linear=True
+        )
+
+
+class LMPickBottleHigh(LMPickBottle):
+    TABLE_OFFSET = 0.1
+
+    def _get_objects(self) -> list[SceneObject]:
+        self.table = SceneObject(
+            ObjectConfig(
+                name="table",
+                mjcf_path="objects/omniverse/locomanip/lab_table/model.xml",
+                scale=1.0,
+                static=True,
+                sampler_config=SamplingConfig(
+                    x_range=np.array([-0.02, 0.02]),
+                    y_range=np.array([-0.02, 0.02]),
+                    reference_pos=np.array([0.5, 0, self.TABLE_OFFSET]),
+                    rotation=np.array([np.pi * 0.5, np.pi * 0.5]),
+                ),
+            )
+        )
+        self.bottle = SceneObject(
+            ObjectConfig(
+                name="bottle",
+                mjcf_path="objects/omniverse/locomanip/jug_a01/model.xml",
+                static=False,
+                scale=0.6,
+                sampler_config=SamplingConfig(
+                    x_range=np.array([-0.08, 0.04]),
+                    y_range=np.array([-0.08, 0.08]),
+                    rotation=np.array([-np.pi, np.pi]),
+                    reference_pos=np.array(
+                        [0.4, 0, self.TABLE_OFFSET + self.table.mj_obj.top_offset[2]]
+                    ),
+                    reference=ReferenceConfig(obj=self.table),
+                ),
+            )
+        )
+        return [self.table, self.bottle]
+
+
+class LMNavPickBottle(LMPickBottle):
+    def _reset_internal(self):
+        super()._reset_internal()
+
+        if not self.deterministic_reset:
+            RobotPoseRandomizer.set_pose(self, (-0.3, -0.16), (-0.2, 0.2), (-np.pi / 6, np.pi / 6))
+
+    def _get_instruction(self) -> str:
+        return "Walk forward and pick up the bottle from the table."
+
+
+class LMPickBottleGround(LMPickBottle):
+    def _get_objects(self) -> list[SceneObject]:
+        self.bottle = SceneObject(
+            ObjectConfig(
+                name="bottle",
+                mjcf_path="objects/omniverse/locomanip/jug_a01/model.xml",
+                static=False,
+                scale=0.6,
+                sampler_config=SamplingConfig(
+                    x_range=np.array([-0.08, 0.04]),
+                    y_range=np.array([-0.08, 0.08]),
+                    rotation=np.array([-np.pi, np.pi]),
+                    reference_pos=np.array(
+                        [0.4, 0, 0.075]
+                    ),  # Base position on ground (z=0.075 is bottle radius)
+                ),
+            )
+        )
+        return [self.bottle]
+
+    def _get_success_criteria(self) -> SuccessCriteria:
+        return AllCriteria(
+            IsGrasped(self.bottle, "right"),
+            IsPositionInRange(self.bottle, 2, self.LIFT_OFFSET, 10),
+        )
+
+    def _randomize_table_rgba(self):
+        pass
+
+
+class LMPnPBottle(LMPickBottle):
+    LIFT_OFFSET = 0.15
+
+    def _get_objects(self) -> list[SceneObject]:
+        super()._get_objects()
+        self.table_target = SceneObject(
+            ObjectConfig(
+                name="table_target",
+                mjcf_path="objects/omniverse/locomanip/lab_table/model.xml",
+                scale=1.0,
+                static=True,
+                sampler_config=SamplingConfig(
+                    x_range=np.array([-0.02, 0.02]),
+                    y_range=np.array([-0.02, 0.02]),
+                    reference_pos=np.array([0.5, 1.2, 0]),
+                    rotation=np.array([np.pi * 0.5, np.pi * 0.5]),
+                ),
+            )
+        )
+        return [self.table, self.table_target, self.bottle]
+
+    def _get_success_criteria(self) -> SuccessCriteria:
+        return AllCriteria(
+            IsUpright(self.bottle, symmetric=True), IsInContact(self.bottle, self.table_target)
+        )
+
+    def _get_instruction(self) -> str:
+        return "Pick up the bottle and place it on the other table."
+
+    def get_object(self):
+        return dict(
+            bottle=dict(obj_name=self.bottle.mj_obj.root_body, obj_type="body"),
+            target_table=dict(obj_name=self.table_target.mj_obj.root_body, obj_type="body"),
+        )
+
+    def get_subtask_term_signals(self):
+        obj_z = self.sim.data.body_xpos[self.obj_body_id(self.bottle.mj_obj.name)][2]
+        target_table_pos = self.sim.data.body_xpos[self.obj_body_id(self.table_target.mj_obj.name)]
+        target_table_z = target_table_pos[2] + self.table_target.mj_obj.top_offset[2]
+        return dict(obj_off_table=int(obj_z - target_table_z > self.LIFT_OFFSET))
+
+    @staticmethod
+    def task_config():
+        task = DexMGConfigHelper.AttrDict()
+        task.task_spec_0.subtask_1 = dict(
+            object_ref="bottle",
+            subtask_term_signal="obj_off_table",
+            subtask_term_offset_range=(5, 10),
+            selection_strategy="random",
+            selection_strategy_kwargs=None,
+            action_noise=0.05,
+            num_interpolation_steps=5,
+            num_fixed_steps=0,
+            apply_noise_during_interpolation=False,
+        )
+        # Second subtask for placing on target table
+        task.task_spec_0.subtask_2 = dict(
+            object_ref="target_table",
+            subtask_term_signal=None,
+            subtask_term_offset_range=None,
+            selection_strategy="random",
+            selection_strategy_kwargs=None,
+            action_noise=0.05,
+            num_interpolation_steps=5,
+            num_fixed_steps=0,
+            apply_noise_during_interpolation=False,
+        )
+        task.task_spec_1.subtask_1 = dict(
+            object_ref=None,
+            subtask_term_signal=None,
+            subtask_term_offset_range=None,
+            selection_strategy="random",
+            selection_strategy_kwargs=None,
+            action_noise=0.05,
+            num_interpolation_steps=5,
+            num_fixed_steps=0,
+            apply_noise_during_interpolation=False,
+        )
+        return task.to_dict()
+
+    def _randomize_table_rgba(self):
+        for table in [self.table_target, self.table]:
+            randomize_materials_rgba(
+                rng=self.rng, mjcf_obj=table.mj_obj, gradient=self.TABLE_GRADIENT, linear=True
+            )
+
+
+class LMPickMultipleBottles(LMPickBottle):
+    BOTTLE_COLOURS = [(0.3, 0.7, 0.8, 1.0), (0.8, 0.4, 0.3, 1.0)]
+    BOTTLES_COUNT = 2
+    Y_OFFSET_STEP = 0.1
+
+    def _get_objects(self) -> list[SceneObject]:
+        self.table = SceneObject(
+            ObjectConfig(
+                name="table",
+                mjcf_path="objects/omniverse/locomanip/lab_table/model.xml",
+                scale=1.0,
+                static=True,
+                sampler_config=SamplingConfig(
+                    x_range=np.array([-0.02, 0.02]),
+                    y_range=np.array([-0.02, 0.02]),
+                    reference_pos=np.array([0.5, 0, 0]),
+                    rotation=np.array([np.pi * 0.5, np.pi * 0.5]),
+                ),
+            )
+        )
+
+        self.bottles = []
+        offsets = np.arange(self.BOTTLES_COUNT) - (self.BOTTLES_COUNT - 1) / 2.0
+        for i in range(self.BOTTLES_COUNT):
+            reference_pos = np.array([0.4, 0, self.table.mj_obj.top_offset[2]])
+            reference_pos += np.array([0, self.Y_OFFSET_STEP * offsets[i], 0])
+            bottle = SceneObject(
+                ObjectConfig(
+                    name=f"bottle_{i}",
+                    mjcf_path="objects/omniverse/locomanip/jug_a01/model.xml",
+                    static=False,
+                    scale=0.6,
+                    sampler_config=SamplingConfig(
+                        x_range=np.array([-0.08, 0.04]),
+                        y_range=np.array([-0.04, 0.04]),
+                        rotation=np.array([-np.pi, np.pi]),
+                        reference_pos=reference_pos,
+                    ),
+                    rgba=self.BOTTLE_COLOURS[i % len(self.BOTTLE_COLOURS)],
+                )
+            )
+            self.bottles.append(bottle)
+        return [self.table, *self.bottles]
+
+    def _get_success_criteria(self) -> SuccessCriteria:
+        criteria = []
+        for bottle in self.bottles:
+            criteria.append(AnyCriteria(IsGrasped(bottle, "right"), IsGrasped(bottle, "left")))
+            criteria.append(
+                IsPositionInRange(bottle, 2, self.table.mj_obj.top_offset[2] + self.LIFT_OFFSET, 10)
+            )
+        return AllCriteria(*criteria)
+
+    def _get_instruction(self) -> str:
+        return "Pick up bottles."
+
+    def get_object(self):
+        return {
+            bottle.mj_obj.name: dict(obj_name=bottle.mj_obj.root_body, obj_type="body")
+            for bottle in self.bottles
+        }
+
+    def get_subtask_term_signals(self):
+        return {
+            f"grasp_{bottle.mj_obj.name}": int(
+                self._check_grasp(self.robots[0].gripper["right"], bottle.mj_obj)
+                or self._check_grasp(self.robots[0].gripper["left"], bottle.mj_obj)
+            )
+            for bottle in self.bottles
+        }
+
+    @staticmethod
+    def task_config():
+        task = DexMGConfigHelper.AttrDict()
+        for i in range(LMPickMultipleBottles.BOTTLES_COUNT):
+            subtask = dict(
+                object_ref=f"bottle_{i}",
+                subtask_term_signal=None,
+                subtask_term_offset_range=None,
+                selection_strategy="random",
+                selection_strategy_kwargs=None,
+                action_noise=0.05,
+                num_interpolation_steps=5,
+                num_fixed_steps=0,
+                apply_noise_during_interpolation=False,
+            )
+            setattr(task.task_spec_0, f"subtask_{i+1}", subtask)
+        task.task_spec_1.subtask_1 = dict(
+            object_ref=None,
+            subtask_term_signal=None,
+            subtask_term_offset_range=None,
+            selection_strategy="random",
+            selection_strategy_kwargs=None,
+            action_noise=0.05,
+            num_interpolation_steps=5,
+            num_fixed_steps=0,
+            apply_noise_during_interpolation=False,
+        )
+        return task.to_dict()
+
+
+class LMPnPMultipleBottles(LMPickMultipleBottles):
+    def _get_objects(self) -> list[SceneObject]:
+        super()._get_objects()
+        self.table_target = SceneObject(
+            ObjectConfig(
+                name="table_target",
+                mjcf_path="objects/omniverse/locomanip/lab_table/model.xml",
+                scale=1.0,
+                static=True,
+                sampler_config=SamplingConfig(
+                    x_range=np.array([-0.02, 0.02]),
+                    y_range=np.array([-0.02, 0.02]),
+                    reference_pos=np.array([0.5, 1.2, 0]),
+                    rotation=np.array([np.pi * 0.5, np.pi * 0.5]),
+                ),
+            )
+        )
+        return [self.table, self.table_target, *self.bottles]
+
+    def _get_success_criteria(self) -> SuccessCriteria:
+        criteria = [
+            AllCriteria(IsInContact(bottle, self.table_target), IsUpright(bottle, symmetric=True))
+            for bottle in self.bottles
+        ]
+        return AllCriteria(*criteria)
+
+    def _get_instruction(self) -> str:
+        return "Pick up bottles from one table and place it on the other."
+
+    @staticmethod
+    def task_config():
+        task = DexMGConfigHelper.AttrDict()
+        for i in range(LMPnPMultipleBottles.BOTTLES_COUNT):
+            bottle_name = f"bottle_{i}"
+            subtask = dict(
+                object_ref=bottle_name,
+                subtask_term_signal=f"{bottle_name}_off_table",
+                subtask_term_offset_range=None,
+                selection_strategy="random",
+                selection_strategy_kwargs=None,
+                action_noise=0.05,
+                num_interpolation_steps=5,
+                num_fixed_steps=0,
+                apply_noise_during_interpolation=False,
+            )
+            setattr(task.task_spec_0, f"subtask_{i+1}", subtask)
+        # Next subtask for placing on target table
+        task.task_spec_0.subtask_3 = dict(
+            object_ref="target_table",
+            subtask_term_signal=None,
+            subtask_term_offset_range=None,
+            selection_strategy="random",
+            selection_strategy_kwargs=None,
+            action_noise=0.05,
+            num_interpolation_steps=5,
+            num_fixed_steps=0,
+            apply_noise_during_interpolation=False,
+        )
+        task.task_spec_1.subtask_1 = dict(
+            object_ref=None,
+            subtask_term_signal=None,
+            subtask_term_offset_range=None,
+            selection_strategy="random",
+            selection_strategy_kwargs=None,
+            action_noise=0.05,
+            num_interpolation_steps=5,
+            num_fixed_steps=0,
+            apply_noise_during_interpolation=False,
+        )
+        return task.to_dict()
+
+    def get_subtask_term_signals(self):
+        signals = dict()
+        for bottle in self.bottles:
+            obj_z = self.sim.data.body_xpos[self.obj_body_id(self.bottle.mj_obj.name)][2]
+            target_table_pos = self.sim.data.body_xpos[
+                self.obj_body_id(self.table_target.mj_obj.name)
+            ]
+            target_table_z = target_table_pos[2] + self.table_target.mj_obj.top_offset[2]
+            signals[f"{bottle.mj_obj.name}_off_table"] = int(
+                obj_z - target_table_z > self.LIFT_OFFSET
+            )
+        return signals
+
+    def _randomize_table_rgba(self):
+        for table in [self.table_target, self.table]:
+            randomize_materials_rgba(
+                rng=self.rng, mjcf_obj=table.mj_obj, gradient=self.TABLE_GRADIENT, linear=True
+            )
+
+
+class LMPickBottleShelf(LMPickBottle):
+    def _get_objects(self) -> list[SceneObject]:
+        super()._get_objects()
+        self.shelf = SceneObject(
+            ObjectConfig(
+                name="shelf",
+                mjcf_path="objects/omniverse/locomanip/lab_shelf/model.xml",
+                static=True,
+                sampler_config=SamplingConfig(
+                    rotation=np.array([np.pi / 2, np.pi / 2]),
+                    reference_pos=np.array([0.9, 0, 0]),
+                ),
+            )
+        )
+        self.bottle = SceneObject(
+            ObjectConfig(
+                name="bottle",
+                mjcf_path="objects/omniverse/locomanip/jug_a01/model.xml",
+                static=False,
+                scale=0.6,
+                sampler_config=SamplingConfig(
+                    x_range=np.array([-0.14, -0.06]),
+                    y_range=np.array([-0.08, 0.08]),
+                    rotation=np.array([-np.pi, np.pi]),
+                    reference=ReferenceConfig(self.shelf, spawn_id=2),
+                ),
+            )
+        )
+        return [self.shelf, self.bottle]
+
+    def _get_success_criteria(self) -> SuccessCriteria:
+        return AllCriteria(
+            IsGrasped(self.bottle, "right"),
+            NotCriteria(IsInContact(self.bottle, self.shelf)),
+        )
+
+
+class LMNavPickBottleShelf(LMPickBottleShelf):
+    ROBOT_DISTANCE_THRESHOLD = 1.0
+
+    def _reset_internal(self):
+        super()._reset_internal()
+        if not self.deterministic_reset:
+            RobotPoseRandomizer.set_pose(self, (-0.1, 0.1), (-0.1, 0.1), (-np.pi / 6, np.pi / 6))
+
+    def _get_success_criteria(self) -> SuccessCriteria:
+        return AllCriteria(
+            NotCriteria(IsRobotInRange(self.shelf, self.ROBOT_DISTANCE_THRESHOLD, True)),
+            IsGrasped(self.bottle, "right"),
+            NotCriteria(IsInContact(self.bottle, self.shelf)),
+        )
+
+    def _get_instruction(self) -> str:
+        return "Pick up the bottle from the shelf and move backward away from it."
+
+
+class LMPickBottleShelfLow(LMPickBottleShelf):
+    def _get_objects(self) -> list[SceneObject]:
+        super()._get_objects()
+        self.bottle = SceneObject(
+            ObjectConfig(
+                name="bottle",
+                mjcf_path="objects/omniverse/locomanip/jug_a01/model.xml",
+                static=False,
+                scale=0.6,
+                sampler_config=SamplingConfig(
+                    x_range=np.array([-0.14, -0.06]),
+                    y_range=np.array([-0.08, 0.08]),
+                    rotation=np.array([-np.pi, np.pi]),
+                    reference=ReferenceConfig(self.shelf, spawn_id=1),
+                ),
+            )
+        )
+        return [self.shelf, self.bottle]
+
+
+class LMNavPickBottleShelfLow(LMNavPickBottleShelf):
+    def _get_objects(self) -> list[SceneObject]:
+        super()._get_objects()
+        self.bottle = SceneObject(
+            ObjectConfig(
+                name="bottle",
+                mjcf_path="objects/omniverse/locomanip/jug_a01/model.xml",
+                static=False,
+                scale=0.6,
+                sampler_config=SamplingConfig(
+                    x_range=np.array([-0.14, -0.06]),
+                    y_range=np.array([-0.08, 0.08]),
+                    rotation=np.array([-np.pi, np.pi]),
+                    reference=ReferenceConfig(self.shelf, spawn_id=1),
+                ),
+            )
+        )
+        return [self.shelf, self.bottle]
+
+
+class LMPnPBottleToPlate(LMPnPBottle):
+    def _get_objects(self) -> list[SceneObject]:
+        super()._get_objects()
+        self.plate = SceneObject(
+            ObjectConfig(
+                name="plate",
+                mjcf_path="objects/omniverse/locomanip/plate_1/model.xml",
+                scale=1.0,
+                static=True,
+                sampler_config=SamplingConfig(
+                    x_range=np.array([-0.2 - 0.08, -0.2 + 0.04]),
+                    y_range=np.array([-0.08, 0.08]),
+                    rotation=np.array([-np.pi, np.pi]),
+                    reference=ReferenceConfig(self.table_target),
+                ),
+            )
+        )
+        return [self.table, self.table_target, self.bottle, self.plate]
+
+    def _get_success_criteria(self) -> SuccessCriteria:
+        return AllCriteria(
+            IsUpright(self.bottle, symmetric=True), IsInContact(self.bottle, self.plate)
+        )
+
+    def _get_instruction(self) -> str:
+        return "Pick up the bottle and place it on the plate."
+
+    def get_object(self):
+        return dict(
+            bottle=dict(obj_name=self.bottle.mj_obj.root_body, obj_type="body"),
+            plate=dict(obj_name=self.plate.mj_obj.root_body, obj_type="body"),
+        )
+
+    def get_subtask_term_signals(self):
+        obj_z = self.sim.data.body_xpos[self.obj_body_id(self.bottle.mj_obj.name)][2]
+        target_table_pos = self.sim.data.body_xpos[self.obj_body_id(self.table_target.mj_obj.name)]
+        target_table_z = target_table_pos[2] + self.table_target.mj_obj.top_offset[2]
+        return dict(obj_off_table=int(obj_z - target_table_z > self.LIFT_OFFSET))
+
+    @staticmethod
+    def task_config():
+        task = DexMGConfigHelper.AttrDict()
+        task.task_spec_0.subtask_1 = dict(
+            object_ref="bottle",
+            subtask_term_signal="obj_off_table",
+            subtask_term_offset_range=(5, 10),
+            selection_strategy="random",
+            selection_strategy_kwargs=None,
+            action_noise=0.05,
+            num_interpolation_steps=5,
+            num_fixed_steps=0,
+            apply_noise_during_interpolation=False,
+        )
+        # Second subtask for placing on plate
+        task.task_spec_0.subtask_2 = dict(
+            object_ref="plate",
+            subtask_term_signal=None,
+            subtask_term_offset_range=None,
+            selection_strategy="random",
+            selection_strategy_kwargs=None,
+            action_noise=0.05,
+            num_interpolation_steps=5,
+            num_fixed_steps=0,
+            apply_noise_during_interpolation=False,
+        )
+        task.task_spec_1.subtask_1 = dict(
+            object_ref=None,
+            subtask_term_signal=None,
+            subtask_term_offset_range=None,
+            selection_strategy="random",
+            selection_strategy_kwargs=None,
+            action_noise=0.05,
+            num_interpolation_steps=5,
+            num_fixed_steps=0,
+            apply_noise_during_interpolation=False,
+        )
+        return task.to_dict()
+
+
+class LMPnPAppleToPlate(LMPnPBottleToPlate):
+    def _get_objects(self) -> list[SceneObject]:
+        super()._get_objects()
+        self.apple = SceneObject(
+            ObjectConfig(
+                name="apple",
+                mjcf_path="objects/omniverse/locomanip/apple_0/model.xml",
+                static=False,
+                scale=1.0,
+                sampler_config=SamplingConfig(
+                    x_range=np.array([-0.08, 0.04]),
+                    y_range=np.array([-0.08, 0.08]),
+                    rotation=np.array([-np.pi, np.pi]),
+                    reference_pos=np.array([0.4, 0, self.table.mj_obj.top_offset[2]]),
+                ),
+            )
+        )
+        return [self.table, self.table_target, self.apple, self.plate]
+
+    def _get_success_criteria(self) -> SuccessCriteria:
+        return IsInContact(self.apple, self.plate)
+
+    def _get_instruction(self) -> str:
+        return "pick up the apple, walk left and place the apple on the plate."
+
+    def get_object(self):
+        return dict(
+            apple=dict(obj_name=self.apple.mj_obj.root_body, obj_type="body"),
+            plate=dict(obj_name=self.plate.mj_obj.root_body, obj_type="body"),
+        )
+
+    def get_subtask_term_signals(self):
+        obj_z = self.sim.data.body_xpos[self.obj_body_id(self.apple.mj_obj.name)][2]
+        target_table_pos = self.sim.data.body_xpos[self.obj_body_id(self.table_target.mj_obj.name)]
+        target_table_z = target_table_pos[2] + self.table_target.mj_obj.top_offset[2]
+        return dict(obj_off_table=int(obj_z - target_table_z > self.LIFT_OFFSET))
+
+    @staticmethod
+    def task_config():
+        task = DexMGConfigHelper.AttrDict()
+        task.task_spec_0.subtask_1 = dict(
+            object_ref="apple",
+            subtask_term_signal="obj_off_table",
+            subtask_term_offset_range=(5, 10),
+            selection_strategy="random",
+            selection_strategy_kwargs=None,
+            action_noise=0.05,
+            num_interpolation_steps=5,
+            num_fixed_steps=0,
+            apply_noise_during_interpolation=False,
+        )
+        # Second subtask for placing on plate
+        task.task_spec_0.subtask_2 = dict(
+            object_ref="plate",
+            subtask_term_signal=None,
+            subtask_term_offset_range=None,
+            selection_strategy="random",
+            selection_strategy_kwargs=None,
+            action_noise=0.05,
+            num_interpolation_steps=5,
+            num_fixed_steps=0,
+            apply_noise_during_interpolation=False,
+        )
+        task.task_spec_1.subtask_1 = dict(
+            object_ref=None,
+            subtask_term_signal=None,
+            subtask_term_offset_range=None,
+            selection_strategy="random",
+            selection_strategy_kwargs=None,
+            action_noise=0.05,
+            num_interpolation_steps=5,
+            num_fixed_steps=0,
+            apply_noise_during_interpolation=False,
+        )
+        return task.to_dict()

GR00T-WholeBodyControl/decoupled_wbc/dexmg/gr00trobocasa/robocasa/environments/locomanipulation/locomanip_dc.py

@@ -0,0 +1,15 @@
+from robocasa import (
+    LMNavPickBottle,
+    LMPnPAppleToPlate,
+)
+from robocasa.models.scenes.lab_arena import LabArena
+
+
+class LabEnvMixin:
+    MUJOCO_ARENA_CLS = LabArena
+
+
+class LMNavPickBottleDC(LabEnvMixin, LMNavPickBottle): ...
+
+
+class LMPnPAppleToPlateDC(LabEnvMixin, LMPnPAppleToPlate): ...

GR00T-WholeBodyControl/decoupled_wbc/dexmg/gr00trobocasa/robocasa/environments/locomanipulation/locomanip_pnp.py

@@ -0,0 +1,99 @@
+import numpy as np
+from robocasa.environments.locomanipulation.locomanip import LMFactoryEnv
+from robocasa.utils.scene.configs import (
+    ObjectConfig,
+    ReferenceConfig,
+    SamplingConfig,
+    SceneHandedness,
+    SceneScaleConfig,
+)
+from robocasa.utils.scene.scene import SceneObject
+from robocasa.utils.scene.success_criteria import (
+    AllCriteria,
+    IsInContact,
+    IsUpright,
+    SuccessCriteria,
+)
+
+
+class LMBottlePnP(LMFactoryEnv):
+    SCENE_SCALE = SceneScaleConfig(planar_scale=(1, 1), handedness=SceneHandedness.RIGHT)
+
+    def _get_objects(self) -> list[SceneObject]:
+        self.table_target = SceneObject(
+            ObjectConfig(
+                name="table_target",
+                mjcf_path="objects/omniverse/locomanip/factory_ergo_table/model.xml",
+                static=True,
+                sampler_config=SamplingConfig(
+                    x_range=np.array([-0.02, 0.02]),
+                    y_range=np.array([-0.02, 0.02]),
+                    reference_pos=np.array([1.2, 0.8, 0]),
+                    rotation=np.array([np.pi, np.pi]),
+                ),
+            )
+        )
+        self.table_origin = SceneObject(
+            ObjectConfig(
+                name="table_origin",
+                mjcf_path="objects/omniverse/locomanip/factory_ergo_table/model.xml",
+                static=True,
+                sampler_config=SamplingConfig(
+                    x_range=np.array([-0.02, 0.02]),
+                    y_range=np.array([-0.02, 0.02]),
+                    reference_pos=np.array([1.2, -0.8, 0]),
+                    rotation=np.array([np.pi, np.pi]),
+                ),
+            )
+        )
+        self.bottle = SceneObject(
+            ObjectConfig(
+                name="obj",
+                mjcf_path="objects/omniverse/locomanip/jug_a01/model.xml",
+                static=False,
+                scale=0.6,
+                sampler_config=SamplingConfig(
+                    x_range=np.array([-0.4, -0.35]),
+                    y_range=np.array([-0.1, 0.1]),
+                    rotation=np.array([-np.pi, np.pi]),
+                    reference=ReferenceConfig(self.table_origin),
+                ),
+            )
+        )
+        return [self.table_origin, self.table_target, self.bottle]
+
+    def _get_success_criteria(self) -> SuccessCriteria:
+        return AllCriteria(IsInContact(self.bottle, self.table_target), IsUpright(self.bottle))
+
+    def _get_instruction(self) -> str:
+        return "Pick up the bottle from one table and place it on the other."
+
+
+class LMBoxPnP(LMBottlePnP):
+    SCENE_SCALE = SceneScaleConfig(planar_scale=(1, 1), handedness=SceneHandedness.RIGHT)
+
+    def _get_objects(self) -> list[SceneObject]:
+        super()._get_objects()
+        self.box = SceneObject(
+            ObjectConfig(
+                name="obj",
+                mjcf_path="objects/omniverse/locomanip/cardbox_a1/model.xml",
+                static=False,
+                scale=0.7,
+                density=1,
+                friction=(2, 1, 1),
+                sampler_config=SamplingConfig(
+                    x_range=np.array([-0.35, -0.3]),
+                    y_range=np.array([-0.1, 0.1]),
+                    rotation=np.array([np.pi * 0.9, np.pi * 1.1]),
+                    reference=ReferenceConfig(self.table_origin),
+                ),
+            )
+        )
+        return [self.table_origin, self.table_target, self.box]
+
+    def _get_success_criteria(self) -> SuccessCriteria:
+        return AllCriteria(IsInContact(self.box, self.table_target), IsUpright(self.box))
+
+    def _get_instruction(self) -> str:
+        return "Pick up the box from one table and place it on the other."

GR00T-WholeBodyControl/decoupled_wbc/dexmg/gr00trobocasa/robocasa/examples/third_party_controller/default_mink_ik_g1_gear_wbc.json

@@ -0,0 +1,113 @@
+{
+    "type": "HYBRID_WHOLE_BODY_MINK_IK",
+    "composite_controller_specific_configs": {
+        "ref_name": ["gripper0_right_grip_site", "gripper0_left_grip_site"],
+        "interpolation": null,
+        "actuation_part_names": ["torso", "left", "right"],
+        "external_part_names": ["legs"],
+        "max_dq": 4,
+        "ik_pseudo_inverse_damping": 5e-2,
+        "ik_integration_dt": 1e-1,
+        "ik_input_type": "absolute",
+        "ik_input_ref_frame": "base", 
+        "ik_input_rotation_repr": "axis_angle",
+        "verbose": false,
+        "ik_posture_weights": {
+            "robot0_waist_yaw_joint": 100.0,
+            "robot0_waist_roll_joint": 200.0,
+            "robot0_waist_pitch_joint": 400.0,
+            "robot0_left_shoulder_pitch_joint": 4.0,
+            "robot0_left_shoulder_roll_joint": 3.0,
+            "robot0_left_shoulder_yaw_joint": 2.0,
+            "robot0_left_elbow_joint": 1.0,
+            "robot0_right_shoulder_pitch_joint": 4.0,
+            "robot0_right_shoulder_roll_joint": 3.0,
+            "robot0_right_shoulder_yaw_joint": 2.0,
+            "robot0_right_elbow_joint": 1.0
+        },
+        "ik_hand_pos_cost": 10.0,
+        "ik_hand_ori_cost": 5,
+        "use_joint_angle_action_input": false
+    },
+    "body_parts": {
+        "legs": {
+            "type" : "JOINT_POSITION",
+            "input_max": 100,
+            "input_min": -100,
+            "input_type": "absolute",
+            "output_max": 100,
+            "output_min": -100,
+            "kd": [2, 2, 2, 4, 2, 2, 2, 2, 2, 4, 2, 2],
+            "kv": 0,
+            "kp": [150, 150, 150, 300, 40, 40, 150, 150, 150, 200, 40, 40],
+            "velocity_limits": [-1,1],
+            "kp_limits": [0, 1000],
+            "interpolation": null,
+            "ramp_ratio": 0.2,
+            "use_torque_compensation": false,
+            "desired_torque_as_acceleration": false
+        },
+        "arms": {
+            "left": {
+                "type" : "JOINT_POSITION",
+                "input_max": 100,
+                "input_min": -100,
+                "input_type": "absolute",
+                "output_max": 100,
+                "output_min": -100,
+                "kd": [5.0, 5.0, 2.0, 2.0, 2.0, 2.0, 2.0],
+                "kp": [100, 100, 40, 40, 20, 20, 20],
+                "velocity_limits": [-1,1],
+                "kp_limits": [0, 1000],
+                "interpolation": null,
+                "ramp_ratio": 0.2,
+                "gripper": {
+                    "type": "GRIP",
+                    "use_action_scaling": false
+                },
+                "use_torque_compensation": false,
+                "desired_torque_as_acceleration": false
+            },
+            "right": {
+                "type" : "JOINT_POSITION",
+                "input_max": 100,
+                "input_min": -100,
+                "input_type": "absolute",
+                "output_max": 100,
+                "output_min": -100,
+                "kd": [5.0, 5.0, 2.0, 2.0, 2.0, 2.0, 2.0],
+                "kp": [100, 100, 40, 40, 20, 20, 20],
+                "velocity_limits": [-1,1],
+                "kp_limits": [0, 1000],
+                "interpolation": null,
+                "ramp_ratio": 0.2,
+                "gripper": {
+                    "type": "GRIP",
+                    "use_action_scaling": false
+                },
+                "use_torque_compensation": false,
+                "desired_torque_as_acceleration": false
+            }
+        },
+        "torso": {
+            "type" : "JOINT_POSITION",
+            "input_max": 100,
+            "input_min": -100,
+            "input_type": "absolute",
+            "output_max": 100,
+            "output_min": -100,
+            "kd": 5.0,
+            "kp": 250.0,
+            "velocity_limits": [-1,1],
+            "kp_limits": [0, 1000],
+            "interpolation": null,
+            "ramp_ratio": 0.2,
+            "use_torque_compensation": false,
+            "desired_torque_as_acceleration": false
+        },
+        "base": {
+            "type": "JOINT_VELOCITY_AND_POSITION",
+            "interpolation": "null"
+        }
+    }
+}

GR00T-WholeBodyControl/decoupled_wbc/dexmg/gr00trobocasa/robocasa/examples/third_party_controller/default_mink_ik_g1_gear_wbc_gc.json

@@ -0,0 +1,117 @@
+{
+    "type": "HYBRID_WHOLE_BODY_MINK_IK",
+    "composite_controller_specific_configs": {
+        "ref_name": ["gripper0_right_grip_site", "gripper0_left_grip_site"],
+        "interpolation": null,
+        "actuation_part_names": ["torso", "left", "right"],
+        "external_part_names": ["legs"],
+        "max_dq": 4,
+        "ik_pseudo_inverse_damping": 5e-2,
+        "ik_integration_dt": 1e-1,
+        "ik_input_type": "absolute",
+        "ik_input_ref_frame": "base", 
+        "ik_input_rotation_repr": "axis_angle",
+        "verbose": false,
+        "ik_posture_weights": {
+            "robot0_waist_yaw_joint": 100.0,
+            "robot0_waist_roll_joint": 200.0,
+            "robot0_waist_pitch_joint": 400.0,
+            "robot0_left_shoulder_pitch_joint": 4.0,
+            "robot0_left_shoulder_roll_joint": 3.0,
+            "robot0_left_shoulder_yaw_joint": 2.0,
+            "robot0_left_elbow_joint": 1.0,
+            "robot0_right_shoulder_pitch_joint": 4.0,
+            "robot0_right_shoulder_roll_joint": 3.0,
+            "robot0_right_shoulder_yaw_joint": 2.0,
+            "robot0_right_elbow_joint": 1.0
+        },
+        "ik_hand_pos_cost": 10.0,
+        "ik_hand_ori_cost": 5,
+        "use_joint_angle_action_input": false
+    },
+    "body_parts": {
+        "legs": {
+            "type" : "JOINT_POSITION",
+            "input_max": 100,
+            "input_min": -100,
+            "input_type": "absolute",
+            "output_max": 100,
+            "output_min": -100,
+            "kd": [2, 2, 2, 4, 2, 2, 2, 2, 2, 4, 2, 2],
+            "kv": 0,
+            "kp": [150, 150, 150, 300, 40, 40, 150, 150, 150, 200, 40, 40],
+            "velocity_limits": [-1,1],
+            "kp_limits": [0, 1000],
+            "interpolation": null,
+            "ramp_ratio": 0.2,
+            "use_torque_compensation": false,
+            "use_external_torque_compensation": true,
+            "desired_torque_as_acceleration": false
+        },
+        "arms": {
+            "left": {
+                "type" : "JOINT_POSITION",
+                "input_max": 100,
+                "input_min": -100,
+                "input_type": "absolute",
+                "output_max": 100,
+                "output_min": -100,
+                "kd": [5.0, 5.0, 2.0, 2.0, 2.0, 2.0, 2.0],
+                "kp": [100, 100, 40, 40, 20, 20, 20],
+                "velocity_limits": [-1,1],
+                "kp_limits": [0, 1000],
+                "interpolation": null,
+                "ramp_ratio": 0.2,
+                "gripper": {
+                    "type": "GRIP",
+                    "use_action_scaling": false
+                },
+                "use_torque_compensation": false,
+                "use_external_torque_compensation": true,
+                "desired_torque_as_acceleration": false
+            },
+            "right": {
+                "type" : "JOINT_POSITION",
+                "input_max": 100,
+                "input_min": -100,
+                "input_type": "absolute",
+                "output_max": 100,
+                "output_min": -100,
+                "kd": [5.0, 5.0, 2.0, 2.0, 2.0, 2.0, 2.0],
+                "kp": [100, 100, 40, 40, 20, 20, 20],
+                "velocity_limits": [-1,1],
+                "kp_limits": [0, 1000],
+                "interpolation": null,
+                "ramp_ratio": 0.2,
+                "gripper": {
+                    "type": "GRIP",
+                    "use_action_scaling": false
+                },
+                "use_torque_compensation": false,
+                "use_external_torque_compensation": true,
+                "desired_torque_as_acceleration": false
+            }
+        },
+        "torso": {
+            "type" : "JOINT_POSITION",
+            "input_max": 100,
+            "input_min": -100,
+            "input_type": "absolute",
+            "output_max": 100,
+            "output_min": -100,
+            "kd": 5.0,
+            "kp": 250.0,
+            "velocity_limits": [-1,1],
+            "kp_limits": [0, 1000],
+            "interpolation": null,
+            "ramp_ratio": 0.2,
+            "use_torque_compensation": false,
+            "use_external_torque_compensation": true,
+            "desired_torque_as_acceleration": false
+        },
+        "base": {
+            "type": "JOINT_VELOCITY_AND_POSITION",
+            "interpolation": "null"
+        }
+    }
+}

GR00T-WholeBodyControl/decoupled_wbc/dexmg/gr00trobocasa/robocasa/models/objects/composite/__init__.py

@@ -0,0 +1,6 @@
+from .box_pattern_object import BoxPatternObject
+from .needle import NeedleObject
+from .ring_tripod import RingTripodObject
+from .bin import Bin
+from .lid import Lid
+from .pot_with_handles import PotWithHandlesObject

GR00T-WholeBodyControl/decoupled_wbc/dexmg/gr00trobocasa/robocasa/models/objects/composite/bin.py

@@ -0,0 +1,205 @@
+import numpy as np
+
+import robosuite.utils.transform_utils as T
+from robosuite.models.objects import CompositeObject
+from robosuite.utils.mjcf_utils import CustomMaterial, add_to_dict
+
+
+class Bin(CompositeObject):
+    """
+    Generates a four-walled bin container with an open top.
+    Args:
+        name (str): Name of this Bin object
+        bin_size (3-array): (x,y,z) full size of bin
+        wall_thickness (float): How thick to make walls of bin
+        transparent_walls (bool): If True, walls will be semi-translucent
+        friction (3-array or None): If specified, sets friction values for this bin. None results in default values
+        density (float): Density value to use for all geoms. Defaults to 1000
+        use_texture (bool): If true, geoms will be defined by realistic textures and rgba values will be ignored
+        rgba (4-array or None): If specified, sets rgba values for all geoms. None results in default values
+        material: If specified, use this material
+        upside_down (bool): if True, construct and initialize the Bin so the bottom geom is at the top
+    """
+
+    def __init__(
+        self,
+        name,
+        bin_size=(0.3, 0.3, 0.15),
+        wall_thickness=0.01,
+        transparent_walls=True,
+        friction=None,
+        density=1000.0,
+        use_texture=True,
+        rgba=(0.2, 0.1, 0.0, 1.0),
+        material=None,
+        upside_down=False,
+        add_second_base=False,
+        transparent_base=False,
+    ):
+        # Set name
+        self._name = name
+
+        # Set object attributes
+        self.bin_size = np.array(bin_size)
+        self.wall_thickness = wall_thickness
+        self.transparent_walls = transparent_walls
+        self.friction = friction if friction is None else np.array(friction)
+        self.density = density
+        self.use_texture = use_texture
+        self.rgba = rgba
+        self.bin_mat_name = "dark_wood_mat"
+
+        # if box should be constructed and initialized upside down
+        self.upside_down = upside_down
+
+        # if box should have a second base (so it will be a closed box)
+        self.add_second_base = add_second_base
+        if self.add_second_base:
+            assert not self.upside_down
+
+        # if base(s) should be transparent
+        self.transparent_base = transparent_base
+
+        self.has_material = material is not None
+        if self.has_material:
+            assert isinstance(material, CustomMaterial)
+            self.material = material
+            self.bin_mat_name = self.material.mat_attrib["name"]
+        else:
+            # default material
+            tex_attrib = {
+                "type": "cube",
+            }
+            mat_attrib = {
+                "texrepeat": "3 3",
+                "specular": "0.4",
+                "shininess": "0.1",
+            }
+            bin_mat = CustomMaterial(
+                texture="WoodDark",
+                tex_name="dark_wood",
+                mat_name=self.bin_mat_name,
+                tex_attrib=tex_attrib,
+                mat_attrib=mat_attrib,
+            )
+            self.material = bin_mat
+
+        # Element references
+        self._base_geom = "base"
+        if self.add_second_base:
+            self._second_base_geom = "base2"
+
+        # Other private attributes
+        self._important_sites = {}
+
+        # Create dictionary of values to create geoms for composite object and run super init
+        super().__init__(**self._get_geom_attrs())
+
+        # Define materials we want to use for this object
+        self.append_material(self.material)
+
+    def _get_geom_attrs(self):
+        """
+        Creates geom elements that will be passed to superclass CompositeObject constructor
+        Returns:
+            dict: args to be used by CompositeObject to generate geoms
+        """
+        # Initialize dict of obj args that we'll pass to the CompositeObject constructor
+        base_args = {
+            "total_size": self.bin_size / 2.0,
+            "name": self.name,
+            "locations_relative_to_center": True,
+            "obj_types": "all",
+            "density": self.density,
+        }
+        obj_args = {}
+
+        # Base(s)
+        base_geom_loc = (0, 0, -(self.bin_size[2] - self.wall_thickness) / 2)
+        if self.upside_down:
+            base_geom_loc = (
+                base_geom_loc[0],
+                base_geom_loc[1],
+                -1.0 * base_geom_loc[2],
+            )
+        if self.transparent_base:
+            base_rgba = (1.0, 1.0, 1.0, 0.3)
+            base_mat = None
+        else:
+            base_rgba = None if self.use_texture else self.rgba
+            base_mat = self.bin_mat_name if self.use_texture else None
+
+        base_geom_names = [self._base_geom]
+        base_geom_locs = [base_geom_loc]
+        if self.add_second_base:
+            base_geom_names.append(self._second_base_geom)
+            base_geom_locs.append((base_geom_loc[0], base_geom_loc[1], -1.0 * base_geom_loc[2]))
+
+        for base_g_name, base_g_loc in zip(base_geom_names, base_geom_locs):
+            add_to_dict(
+                dic=obj_args,
+                geom_types="box",
+                geom_locations=base_g_loc,
+                geom_quats=(1, 0, 0, 0),
+                geom_sizes=(
+                    np.array((self.bin_size[0], self.bin_size[1], self.wall_thickness))
+                    - np.array((self.wall_thickness, self.wall_thickness, 0))
+                )
+                / 2,
+                geom_names=base_g_name,
+                geom_rgbas=base_rgba,
+                geom_materials=base_mat,
+                geom_frictions=self.friction,
+            )
+
+        # Walls
+        x_vals = np.array(
+            [
+                0,
+                -(self.bin_size[0] - self.wall_thickness) / 2,
+                0,
+                (self.bin_size[0] - self.wall_thickness) / 2,
+            ]
+        )
+        y_vals = np.array(
+            [
+                -(self.bin_size[1] - self.wall_thickness) / 2,
+                0,
+                (self.bin_size[1] - self.wall_thickness) / 2,
+                0,
+            ]
+        )
+        w_vals = np.array([self.bin_size[0], self.bin_size[1], self.bin_size[0], self.bin_size[1]])
+        r_vals = np.array([np.pi / 2, 0, -np.pi / 2, np.pi])
+        if self.transparent_walls:
+            wall_rgba = (1.0, 1.0, 1.0, 0.3)
+            wall_mat = None
+        else:
+            wall_rgba = None if self.use_texture else self.rgba
+            wall_mat = self.bin_mat_name if self.use_texture else None
+        for i, (x, y, w, r) in enumerate(zip(x_vals, y_vals, w_vals, r_vals)):
+            add_to_dict(
+                dic=obj_args,
+                geom_types="box",
+                geom_locations=(x, y, 0),
+                geom_quats=T.convert_quat(T.axisangle2quat(np.array([0, 0, r])), to="wxyz"),
+                geom_sizes=(self.wall_thickness / 2, w / 2, self.bin_size[2] / 2),
+                geom_names=f"wall{i}",
+                geom_rgbas=wall_rgba,
+                geom_materials=wall_mat,
+                geom_frictions=self.friction,
+            )
+
+        # Add back in base args and site args
+        obj_args.update(base_args)
+
+        # Return this dict
+        return obj_args
+
+    @property
+    def base_geoms(self):
+        """
+        Returns:
+            list of str: geom names corresponding to bin base
+        """
+        return [self.correct_naming(self._base_geom)]

GR00T-WholeBodyControl/decoupled_wbc/dexmg/gr00trobocasa/robocasa/models/objects/composite/box_pattern_object.py

@@ -0,0 +1,124 @@
+import numpy as np
+
+from robosuite.models.objects import CompositeObject
+from robosuite.utils.mjcf_utils import add_to_dict
+import robosuite.utils.transform_utils as T
+
+
+class BoxPatternObject(CompositeObject):
+    """
+    Generates shapes by using a pattern of unit-size boxes.
+
+    Args:
+        name (str): Name of this Needle object
+    """
+
+    def __init__(
+        self,
+        name,
+        unit_size,
+        pattern,
+        rgba=None,
+        material=None,
+        density=100.0,
+        # solref=[0.02, 1.],
+        # solimp=[0.9, 0.95, 0.001],
+        friction=None,
+    ):
+        """
+        Args:
+            unit_size (3d array / list): size of each unit block in each dimension
+
+            pattern (3d array / list): array of normalized sizes specifying the
+                geometry of the shape. A "0" indicates the absence of a cube and
+                a "1" indicates the presence of a full unit block. The dimensions
+                correspond to z, x, and y respectively.
+        """
+        self._name = name
+        self.rgba = rgba
+        self.material = material
+        self.density = density
+        self.friction = friction
+
+        # number of blocks in z, x, and y
+        self.pattern = np.array(pattern)
+        self.nz, self.nx, self.ny = self.pattern.shape
+        self.unit_size = unit_size
+        self.total_size = [
+            self.nx * unit_size[0],
+            self.ny * unit_size[1],
+            self.nz * unit_size[2],
+        ]
+
+        # Other private attributes
+        self._important_sites = {}
+
+        # Create dictionary of values to create geoms for composite object and run super init
+        super().__init__(**self._get_geom_attrs())
+
+        # Define materials we want to use for this object
+        if self.material is not None:
+            self.append_material(self.material)
+
+    def _get_geom_attrs(self):
+        """
+        Creates geom elements that will be passed to superclass CompositeObject constructor
+
+        Returns:
+            dict: args to be used by CompositeObject to generate geoms
+        """
+        # Initialize dict of obj args that we'll pass to the CompositeObject constructor
+        base_args = {
+            "total_size": self.total_size,
+            "name": self.name,
+            "locations_relative_to_center": False,
+            "obj_types": "all",
+            "density": self.density,
+        }
+        obj_args = {}
+
+        geom_locations = []
+        geom_sizes = []
+        geom_names = []
+        nz, nx, ny = self.pattern.shape
+        for k in range(nz):
+            for i in range(nx):
+                for j in range(ny):
+                    if self.pattern[k, i, j] > 0:
+                        geom_sizes.append(
+                            [
+                                self.unit_size[0],
+                                self.unit_size[1],
+                                self.unit_size[2],
+                            ]
+                        )
+                        geom_locations.append(
+                            [
+                                i * 2.0 * self.unit_size[0],
+                                j * 2.0 * self.unit_size[1],
+                                k * 2.0 * self.unit_size[2],
+                            ]
+                        )
+                        geom_names.append("{}_{}_{}".format(k, i, j))
+
+        # geom_rgbas = [rgba for _ in geom_locations]
+        # geom_frictions = [friction for _ in geom_locations]
+        for i in range(len(geom_locations)):
+            add_to_dict(
+                dic=obj_args,
+                geom_types="box",
+                # needle geom needs to be offset from boundary in (x, z)
+                geom_locations=tuple(geom_locations[i]),
+                geom_quats=(1, 0, 0, 0),
+                geom_sizes=tuple(geom_sizes[i]),
+                geom_names=geom_names[i],
+                geom_rgbas=self.rgba,
+                geom_materials=self.material.name if self.material is not None else None,
+                geom_frictions=None,
+            )
+
+        # Add back in base args and site args
+        obj_args.update(base_args)
+
+        # Return this dict
+        return obj_args

GR00T-WholeBodyControl/decoupled_wbc/dexmg/gr00trobocasa/robocasa/models/objects/composite/lid.py

@@ -0,0 +1,136 @@
+import numpy as np
+
+import robosuite.utils.transform_utils as T
+from robosuite.models.objects import CompositeObject
+from robosuite.utils.mjcf_utils import CustomMaterial, add_to_dict
+
+
+class Lid(CompositeObject):
+    """
+    Generates a square lid with a simple handle.
+    Args:
+        name (str): Name of this Lid object
+        lid_size (3-array): (length, width, thickness) of lid
+        handle_size (3-array): (thickness, length, height) of handle
+        transparent (bool): If True, lid will be semi-translucent
+        friction (3-array or None): If specified, sets friction values for this lid. None results in default values
+        density (float): Density value to use for all geoms. Defaults to 1000
+        use_texture (bool): If true, geoms will be defined by realistic textures and rgba values will be ignored
+        rgba (4-array or None): If specified, sets rgba values for all geoms. None results in default values
+    """
+
+    def __init__(
+        self,
+        name,
+        lid_size=(0.3, 0.3, 0.01),
+        handle_size=(0.02, 0.08, 0.03),
+        transparent=True,
+        friction=None,
+        density=250.0,
+        use_texture=True,
+        rgba=(0.2, 0.1, 0.0, 1.0),
+    ):
+        # Set name
+        self._name = name
+
+        # Set object attributes
+        self.lid_size = np.array(lid_size)
+        self.handle_size = np.array(handle_size)
+        self.transparent = transparent
+        self.friction = friction if friction is None else np.array(friction)
+        self.density = density
+        self.use_texture = use_texture
+        self.rgba = rgba
+        self.lid_mat_name = "dark_wood_mat"
+
+        # Element references
+        self._handle_geom = "handle"
+
+        # Other private attributes
+        self._important_sites = {}
+
+        # Create dictionary of values to create geoms for composite object and run super init
+        super().__init__(**self._get_geom_attrs())
+
+        # Define materials we want to use for this object
+        tex_attrib = {
+            "type": "cube",
+        }
+        mat_attrib = {
+            "texrepeat": "3 3",
+            "specular": "0.4",
+            "shininess": "0.1",
+        }
+        lid_mat = CustomMaterial(
+            texture="WoodDark",
+            tex_name="dark_wood",
+            mat_name=self.lid_mat_name,
+            tex_attrib=tex_attrib,
+            mat_attrib=mat_attrib,
+        )
+        self.append_material(lid_mat)
+
+    def _get_geom_attrs(self):
+        """
+        Creates geom elements that will be passed to superclass CompositeObject constructor
+        Returns:
+            dict: args to be used by CompositeObject to generate geoms
+        """
+        full_height = self.lid_size[2] + self.handle_size[2]
+        full_size = np.array([self.lid_size[0], self.lid_size[1], full_height])
+        # Initialize dict of obj args that we'll pass to the CompositeObject constructor
+        base_args = {
+            "total_size": full_size / 2.0,
+            "name": self.name,
+            "locations_relative_to_center": True,
+            "obj_types": "all",
+        }
+        obj_args = {}
+
+        # Top
+        if self.transparent:
+            top_rgba = (1.0, 1.0, 1.0, 0.3)
+            top_mat = None
+        else:
+            top_rgba = None if self.use_texture else self.rgba
+            top_mat = self.lid_mat_name if self.use_texture else None
+        add_to_dict(
+            dic=obj_args,
+            geom_types="box",
+            geom_locations=(0, 0, (-full_size[2] + self.lid_size[2]) / 2),
+            geom_quats=(1, 0, 0, 0),
+            geom_sizes=np.array((full_size[0], full_size[1], self.lid_size[2])) / 2,
+            geom_names="top",
+            geom_rgbas=top_rgba,
+            geom_materials=top_mat,
+            geom_frictions=self.friction,
+            density=self.density,
+        )
+
+        # Handle
+        add_to_dict(
+            dic=obj_args,
+            geom_types="box",
+            geom_locations=(0, 0, (full_size[2] - self.handle_size[2]) / 2),
+            geom_quats=(1, 0, 0, 0),
+            geom_sizes=self.handle_size / 2,
+            geom_names=self._handle_geom,
+            geom_rgbas=None if self.use_texture else self.rgba,
+            geom_materials=self.lid_mat_name if self.use_texture else None,
+            geom_frictions=self.friction,
+            density=self.density * 2,
+        )
+
+        # Add back in base args and site args
+        obj_args.update(base_args)
+
+        # Return this dict
+        return obj_args
+
+    @property
+    def handle_geoms(self):
+        """
+        Returns:
+            list of str: geom names corresponding to lid handle
+        """
+        return [self.correct_naming(self._handle_geom)]

GR00T-WholeBodyControl/decoupled_wbc/dexmg/gr00trobocasa/robocasa/models/objects/composite/needle.py

@@ -0,0 +1,109 @@
+import numpy as np
+
+from robosuite.models.objects import CompositeObject
+from robosuite.utils.mjcf_utils import add_to_dict, CustomMaterial
+import robosuite.utils.transform_utils as T
+
+
+class NeedleObject(CompositeObject):
+    """
+    Generates a needle with a handle (used in Threading task)
+
+    Args:
+        name (str): Name of this Needle object
+    """
+
+    def __init__(
+        self,
+        name,
+    ):
+
+        ### TODO: make this object more general (with more args and configuration options) later ###
+
+        # Set object attributes
+        self._name = name
+        self.needle_mat_name = "darkwood_mat"
+
+        # Other private attributes
+        self._important_sites = {}
+
+        # Create dictionary of values to create geoms for composite object and run super init
+        super().__init__(**self._get_geom_attrs())
+
+        # Define materials we want to use for this object
+        tex_attrib = {
+            "type": "cube",
+        }
+        mat_attrib = {
+            "texrepeat": "1 1",
+            "specular": "0.4",
+            "shininess": "0.1",
+        }
+        needle_mat = CustomMaterial(
+            texture="WoodDark",
+            tex_name="darkwood",
+            mat_name="darkwood_mat",
+            tex_attrib=tex_attrib,
+            mat_attrib=mat_attrib,
+        )
+        self.append_material(needle_mat)
+
+    def _get_geom_attrs(self):
+        """
+        Creates geom elements that will be passed to superclass CompositeObject constructor
+
+        Returns:
+            dict: args to be used by CompositeObject to generate geoms
+        """
+        # Initialize dict of obj args that we'll pass to the CompositeObject constructor
+        base_args = {
+            "total_size": [0.02, 0.08, 0.02],
+            "name": self.name,
+            "locations_relative_to_center": False,
+            "obj_types": "all",
+            "density": 100.0,
+        }
+        obj_args = {}
+
+        # make a skinny needle object with a large handle
+        needle_size = [0.005, 0.06, 0.005]
+        handle_size = [0.02, 0.02, 0.02]
+
+        # Needle
+        add_to_dict(
+            dic=obj_args,
+            geom_types="box",
+            # needle geom needs to be offset from boundary in (x, z)
+            geom_locations=(
+                (handle_size[0] - needle_size[0]),
+                0.0,
+                (handle_size[2] - needle_size[2]),
+            ),
+            geom_quats=(1, 0, 0, 0),
+            geom_sizes=tuple(needle_size),
+            geom_names="needle",
+            geom_rgbas=None,
+            geom_materials=self.needle_mat_name,
+            # make the needle low friction to ensure easy insertion
+            geom_frictions=(0.3, 5e-3, 1e-4),
+        )
+
+        # Handle
+        add_to_dict(
+            dic=obj_args,
+            geom_types="box",
+            # handle geom needs to be offset in y
+            geom_locations=(0.0, 2.0 * needle_size[1], 0.0),
+            geom_quats=(1, 0, 0, 0),
+            geom_sizes=tuple(handle_size),
+            geom_names="handle",
+            geom_rgbas=None,
+            geom_materials=self.needle_mat_name,
+            geom_frictions=None,
+        )
+
+        # Add back in base args and site args
+        obj_args.update(base_args)
+
+        # Return this dict
+        return obj_args

GR00T-WholeBodyControl/decoupled_wbc/dexmg/gr00trobocasa/robocasa/models/objects/composite/pot_with_handles.py

@@ -0,0 +1,396 @@
+import numpy as np
+
+import robosuite.utils.transform_utils as T
+from robosuite.models.objects import CompositeObject
+from robosuite.utils.mjcf_utils import (
+    BLUE,
+    GREEN,
+    RED,
+    CustomMaterial,
+    add_to_dict,
+    array_to_string,
+)
+
+
+class PotWithHandlesObject(CompositeObject):
+    """
+    Generates the Pot object with side handles (used in TwoArmLift)
+
+    Args:
+        name (str): Name of this Pot object
+
+        body_half_size (3-array of float): If specified, defines the (x,y,z) half-dimensions of the main pot
+            body. Otherwise, defaults to [0.07, 0.07, 0.07]
+
+        handle_radius (float): Determines the pot handle radius
+
+        handle_length (float): Determines the pot handle length
+
+        handle_width (float): Determines the pot handle width
+
+        handle_friction (float): Friction value to use for pot handles. Defauls to 1.0
+
+        density (float): Density value to use for all geoms. Defaults to 1000
+
+        use_texture (bool): If true, geoms will be defined by realistic textures and rgba values will be ignored
+
+        rgba_body (4-array or None): If specified, sets pot body rgba values
+
+        rgba_handle_0 (4-array or None): If specified, sets handle 0 rgba values
+
+        rgba_handle_1 (4-array or None): If specified, sets handle 1 rgba values
+
+        solid_handle (bool): If true, uses a single geom to represent the handle
+
+        thickness (float): How thick to make the pot body walls
+    """
+
+    def __init__(
+        self,
+        name,
+        body_half_size=(0.07, 0.07, 0.07),
+        handle_radius=0.01,
+        handle_length=0.09,
+        handle_width=0.09,
+        handle_friction=1.0,
+        density=1000,
+        use_texture=True,
+        rgba_body=None,
+        rgba_handle_0=None,
+        rgba_handle_1=None,
+        solid_handle=False,
+        thickness=0.01,  # For body
+    ):
+        # Set name
+        self._name = name
+
+        # Set object attributes
+        self.body_half_size = np.array(body_half_size)
+        self.thickness = thickness
+        self.handle_radius = handle_radius
+        self.handle_length = handle_length
+        self.handle_width = handle_width
+        self.handle_friction = handle_friction
+        self.density = density
+        self.use_texture = use_texture
+        self.rgba_body = np.array(rgba_body) if rgba_body else RED
+        self.rgba_handle_0 = np.array(rgba_handle_0) if rgba_handle_0 else GREEN
+        self.rgba_handle_1 = np.array(rgba_handle_1) if rgba_handle_1 else BLUE
+        self.solid_handle = solid_handle
+
+        # Element references to be filled when generated
+        self._handle0_geoms = None
+        self._handle1_geoms = None
+        self.pot_base = None
+
+        # Other private attributes
+        self._important_sites = {}
+
+        # Create dictionary of values to create geoms for composite object and run super init
+        super().__init__(**self._get_geom_attrs())
+
+        # Define materials we want to use for this object
+        tex_attrib = {
+            "type": "cube",
+        }
+        mat_attrib = {
+            "texrepeat": "1 1",
+            "specular": "0.4",
+            "shininess": "0.1",
+        }
+        redwood = CustomMaterial(
+            texture="WoodRed",
+            tex_name="redwood",
+            mat_name="pot_mat",
+            tex_attrib=tex_attrib,
+            mat_attrib=mat_attrib,
+        )
+        greenwood = CustomMaterial(
+            texture="WoodGreen",
+            tex_name="greenwood",
+            mat_name="handle0_mat",
+            tex_attrib=tex_attrib,
+            mat_attrib=mat_attrib,
+        )
+        bluewood = CustomMaterial(
+            texture="WoodBlue",
+            tex_name="bluewood",
+            mat_name="handle1_mat",
+            tex_attrib=tex_attrib,
+            mat_attrib=mat_attrib,
+        )
+        self.append_material(redwood)
+        self.append_material(greenwood)
+        self.append_material(bluewood)
+
+    def _get_geom_attrs(self):
+        """
+        Creates geom elements that will be passed to superclass CompositeObject constructor
+
+        Returns:
+            dict: args to be used by CompositeObject to generate geoms
+        """
+        full_size = np.array(
+            (
+                self.body_half_size,
+                self.body_half_size + self.handle_length * 2,
+                self.body_half_size,
+            )
+        )
+        # Initialize dict of obj args that we'll pass to the CompositeObject constructor
+        base_args = {
+            "total_size": full_size / 2.0,
+            "name": self.name,
+            "locations_relative_to_center": True,
+            "obj_types": "all",
+        }
+        site_attrs = []
+        obj_args = {}
+
+        # Initialize geom lists
+        self._handle0_geoms = []
+        self._handle1_geoms = []
+
+        # Add main pot body
+        # Base geom
+        name = f"base"
+        self.pot_base = [name]
+        add_to_dict(
+            dic=obj_args,
+            geom_types="box",
+            geom_locations=(0, 0, -self.body_half_size[2] + self.thickness / 2),
+            geom_quats=(1, 0, 0, 0),
+            geom_sizes=np.array(
+                [self.body_half_size[0], self.body_half_size[1], self.thickness / 2]
+            ),
+            geom_names=name,
+            geom_rgbas=None if self.use_texture else self.rgba_body,
+            geom_materials="pot_mat" if self.use_texture else None,
+            geom_frictions=None,
+            density=self.density,
+        )
+
+        # Walls
+        x_off = np.array(
+            [
+                0,
+                -(self.body_half_size[0] - self.thickness / 2),
+                0,
+                self.body_half_size[0] - self.thickness / 2,
+            ]
+        )
+        y_off = np.array(
+            [
+                -(self.body_half_size[1] - self.thickness / 2),
+                0,
+                self.body_half_size[1] - self.thickness / 2,
+                0,
+            ]
+        )
+        w_vals = np.array(
+            [
+                self.body_half_size[0],
+                self.body_half_size[1],
+                self.body_half_size[0],
+                self.body_half_size[1],
+            ]
+        )
+        r_vals = np.array([np.pi / 2, 0, -np.pi / 2, np.pi])
+        for i, (x, y, w, r) in enumerate(zip(x_off, y_off, w_vals, r_vals)):
+            add_to_dict(
+                dic=obj_args,
+                geom_types="box",
+                geom_locations=(x, y, 0),
+                geom_quats=T.convert_quat(T.axisangle2quat(np.array([0, 0, r])), to="wxyz"),
+                geom_sizes=np.array([self.thickness / 2, w, self.body_half_size[2]]),
+                geom_names=f"body{i}",
+                geom_rgbas=None if self.use_texture else self.rgba_body,
+                geom_materials="pot_mat" if self.use_texture else None,
+                geom_frictions=None,
+                density=self.density,
+            )
+
+        # Add handles
+        main_bar_size = np.array(
+            [
+                self.handle_width / 2 + self.handle_radius,
+                self.handle_radius,
+                self.handle_radius,
+            ]
+        )
+        side_bar_size = np.array([self.handle_radius, self.handle_length / 2, self.handle_radius])
+        handle_z = self.body_half_size[2] - self.handle_radius
+        for i, (g_list, handle_side, rgba) in enumerate(
+            zip(
+                [self._handle0_geoms, self._handle1_geoms],
+                [1.0, -1.0],
+                [self.rgba_handle_0, self.rgba_handle_1],
+            )
+        ):
+            handle_center = np.array(
+                (
+                    0,
+                    handle_side * (self.body_half_size[1] + self.handle_length),
+                    handle_z,
+                )
+            )
+            # Solid handle case
+            if self.solid_handle:
+                handle_center = np.array(
+                    (
+                        0,
+                        handle_side * (self.body_half_size[1] + self.handle_length / 2),
+                        handle_z,
+                    )
+                )
+                name = f"handle{i}"
+                g_list.append(name)
+                add_to_dict(
+                    dic=obj_args,
+                    geom_types="box",
+                    geom_locations=handle_center,
+                    geom_quats=(1, 0, 0, 0),
+                    geom_sizes=np.array(
+                        [
+                            self.handle_width / 2,
+                            self.handle_length / 2,
+                            self.handle_radius,
+                        ]
+                    ),
+                    geom_names=name,
+                    geom_rgbas=None if self.use_texture else rgba,
+                    geom_materials=f"handle{i}_mat" if self.use_texture else None,
+                    geom_frictions=(self.handle_friction, 0.005, 0.0001),
+                    density=self.density,
+                )
+            # Hollow handle case
+            else:
+                # Center bar
+                name = f"handle{i}_c"
+                g_list.append(name)
+                add_to_dict(
+                    dic=obj_args,
+                    geom_types="box",
+                    geom_locations=handle_center,
+                    geom_quats=(1, 0, 0, 0),
+                    geom_sizes=main_bar_size,
+                    geom_names=name,
+                    geom_rgbas=None if self.use_texture else rgba,
+                    geom_materials=f"handle{i}_mat" if self.use_texture else None,
+                    geom_frictions=(self.handle_friction, 0.005, 0.0001),
+                    density=self.density,
+                )
+                # Side bars
+                for bar_side, suffix in zip([-1.0, 1.0], ["-", "+"]):
+                    name = f"handle{i}_{suffix}"
+                    g_list.append(name)
+                    add_to_dict(
+                        dic=obj_args,
+                        geom_types="box",
+                        geom_locations=(
+                            bar_side * self.handle_width / 2,
+                            handle_side * (self.body_half_size[1] + self.handle_length / 2),
+                            handle_z,
+                        ),
+                        geom_quats=(1, 0, 0, 0),
+                        geom_sizes=side_bar_size,
+                        geom_names=name,
+                        geom_rgbas=None if self.use_texture else rgba,
+                        geom_materials=f"handle{i}_mat" if self.use_texture else None,
+                        geom_frictions=(self.handle_friction, 0.005, 0.0001),
+                        density=self.density,
+                    )
+            # Add relevant site
+            handle_site = self.get_site_attrib_template()
+            handle_name = f"handle{i}"
+            handle_site.update(
+                {
+                    "name": handle_name,
+                    "pos": array_to_string(handle_center - handle_side * np.array([0, 0.005, 0])),
+                    "size": "0.005",
+                    "rgba": rgba,
+                }
+            )
+            site_attrs.append(handle_site)
+            # Add to important sites
+            self._important_sites[f"handle{i}"] = self.naming_prefix + handle_name
+
+        # Add pot body site
+        pot_site = self.get_site_attrib_template()
+        center_name = "center"
+        pot_site.update(
+            {
+                "name": center_name,
+                "size": "0.005",
+            }
+        )
+        site_attrs.append(pot_site)
+        # Add to important sites
+        self._important_sites["center"] = self.naming_prefix + center_name
+
+        # Add back in base args and site args
+        obj_args.update(base_args)
+        obj_args["sites"] = site_attrs  # All sites are part of main (top) body
+
+        # Return this dict
+        return obj_args
+
+    @property
+    def handle_distance(self):
+        """
+        Calculates how far apart the handles are
+
+        Returns:
+            float: handle distance
+        """
+        return self.body_half_size[1] * 2 + self.handle_length * 2
+
+    @property
+    def handle0_geoms(self):
+        """
+        Returns:
+            list of str: geom names corresponding to handle0 (green handle)
+        """
+        return self.correct_naming(self._handle0_geoms)
+
+    @property
+    def handle1_geoms(self):
+        """
+        Returns:
+            list of str: geom names corresponding to handle1 (blue handle)
+        """
+        return self.correct_naming(self._handle1_geoms)
+
+    @property
+    def handle_geoms(self):
+        """
+        Returns:
+            list of str: geom names corresponding to both handles
+        """
+        return self.handle0_geoms + self.handle1_geoms
+
+    @property
+    def important_sites(self):
+        """
+        Returns:
+            dict: In addition to any default sites for this object, also provides the following entries
+
+                :`'handle0'`: Name of handle0 location site
+                :`'handle1'`: Name of handle1 location site
+        """
+        # Get dict from super call and add to it
+        dic = super().important_sites
+        dic.update(self._important_sites)
+        return dic
+
+    @property
+    def bottom_offset(self):
+        return np.array([0, 0, -1 * self.body_half_size[2]])
+
+    @property
+    def top_offset(self):
+        return np.array([0, 0, self.body_half_size[2]])
+
+    @property
+    def horizontal_radius(self):
+        return np.sqrt(2) * (max(self.body_half_size) + self.handle_length)

GR00T-WholeBodyControl/decoupled_wbc/dexmg/gr00trobocasa/robocasa/models/objects/composite/ring_tripod.py

@@ -0,0 +1,194 @@
+import numpy as np
+
+from robosuite.models.objects import CompositeObject
+from robosuite.utils.mjcf_utils import add_to_dict
+from robosuite.utils.mjcf_utils import CustomMaterial
+import robosuite.utils.transform_utils as T
+
+
+class RingTripodObject(CompositeObject):
+    """
+    Generates a tripod base with a small ring for threading a needle through it (used in Threading task)
+
+    Args:
+        name (str): Name of this RingTripod object
+    """
+
+    def __init__(
+        self,
+        name,
+    ):
+
+        ### TODO: make this object more general (with more args and configuration options) later ###
+
+        # Set object attributes
+        self._name = name
+        self.tripod_mat_name = "lightwood_mat"
+
+        # Other private attributes
+        self._important_sites = {}
+
+        # Create dictionary of values to create geoms for composite object and run super init
+        super().__init__(**self._get_geom_attrs())
+
+        # Define materials we want to use for this object
+        tex_attrib = {
+            "type": "cube",
+        }
+        mat_attrib = {
+            "texrepeat": "1 1",
+            "specular": "0.4",
+            "shininess": "0.1",
+        }
+        tripod_mat = CustomMaterial(
+            texture="WoodLight",
+            tex_name="lightwood",
+            mat_name="lightwood_mat",
+            tex_attrib=tex_attrib,
+            mat_attrib=mat_attrib,
+        )
+        self.append_material(tripod_mat)
+
+    def _get_geom_attrs(self):
+        """
+        Creates geom elements that will be passed to superclass CompositeObject constructor
+
+        Returns:
+            dict: args to be used by CompositeObject to generate geoms
+        """
+        # Initialize dict of obj args that we'll pass to the CompositeObject constructor
+        total_size = (0.05, 0.05, 0.1)
+        base_args = {
+            "total_size": total_size,
+            "name": self.name,
+            "locations_relative_to_center": False,
+            "obj_types": "all",
+            "density": 100.0,
+            # NOTE: this lower value of solref allows the thin hole wall to avoid penetration through it
+            "solref": (0.02, 1.0),
+            "solimp": (0.9, 0.95, 0.001),
+        }
+        obj_args = {}
+
+        # pattern for threading ring
+        unit_size = [0.005, 0.002, 0.002]
+        pattern = np.ones((6, 1, 6))
+        for i in range(1, 5):
+            pattern[i][0][1:5] = np.zeros(4)
+        ring_size = [
+            unit_size[0] * pattern.shape[1],
+            unit_size[1] * pattern.shape[2],
+            unit_size[2] * pattern.shape[0],
+        ]
+        self.ring_size = np.array(ring_size)
+
+        # ring offset for where the ring is located relative to the (0, 0, 0) corner
+        ring_offset = [
+            total_size[0] - ring_size[0],
+            total_size[1] - ring_size[1],
+            2.0 * (total_size[2] - ring_size[2]),
+        ]
+
+        # RING-GEOMS: use the pattern to instantiate geoms corresponding to the threading ring
+        nz, nx, ny = pattern.shape
+        self.num_ring_geoms = 0
+        for k in range(nz):
+            for i in range(nx):
+                for j in range(ny):
+                    if pattern[k, i, j] > 0:
+                        add_to_dict(
+                            dic=obj_args,
+                            geom_types="box",
+                            # needle geom needs to be offset from boundary in (x, z)
+                            geom_locations=(
+                                (i * 2.0 * unit_size[0]) + ring_offset[0],
+                                (j * 2.0 * unit_size[1]) + ring_offset[1],
+                                (k * 2.0 * unit_size[2]) + ring_offset[2],
+                            ),
+                            geom_quats=(1, 0, 0, 0),
+                            geom_sizes=tuple(unit_size),
+                            geom_names="ring_{}".format(self.num_ring_geoms),
+                            geom_rgbas=None,
+                            geom_materials=self.tripod_mat_name,
+                            # make the ring low friction to ensure easy insertion
+                            geom_frictions=(0.3, 5e-3, 1e-4),
+                        )
+                        self.num_ring_geoms += 1
+
+        # TRIPOD-GEOMS: legs of the tripod
+        tripod_capsule_r = 0.01
+        tripod_capsule_h = 0.03
+        tripod_geom_locations = [
+            (0.0, 0.0, 0.0),
+            (0.0, 2.0 * total_size[1] - 2.0 * tripod_capsule_r, 0.0),
+            (
+                2.0 * total_size[0] - 2.0 * tripod_capsule_r,
+                total_size[1] - tripod_capsule_r,
+                0.0,
+            ),
+        ]
+        # rotate the legs to resemble a tripod
+        tripod_center = np.array([total_size[0], total_size[1], 0.0])
+        xy_offset = np.array([tripod_capsule_r, tripod_capsule_r, 0.0])
+        rotation_angle = -np.pi / 6.0  # 30 degrees
+        tripod_geom_quats = []
+        for i in range(3):
+            capsule_loc = np.array(tripod_geom_locations[i]) + xy_offset
+            capsule_loc[2] = 0.0  # only care about location in x-y plane
+            vec_to_center = tripod_center - capsule_loc
+            vec_to_center = vec_to_center / np.linalg.norm(vec_to_center)
+            # cross-product with z unit vector to get vector to rotate about
+            rot_vec = np.cross(vec_to_center, np.array([0.0, 0.0, 1.0]))
+            rot_quat = T.mat2quat(T.rotation_matrix(angle=rotation_angle, direction=rot_vec))
+            tripod_geom_quats.append(T.convert_quat(rot_quat, to="wxyz"))
+
+        for i in range(3):
+            add_to_dict(
+                dic=obj_args,
+                geom_types="capsule",
+                geom_locations=tripod_geom_locations[i],
+                geom_quats=tripod_geom_quats[i],
+                geom_sizes=(tripod_capsule_r, tripod_capsule_h),
+                geom_names="tripod_{}".format(i),
+                geom_rgbas=None,
+                geom_materials=self.tripod_mat_name,
+                geom_frictions=None,
+            )
+
+        # POST-GEOMS: mounted base + post
+        base_thickness = 0.005
+        post_size = 0.005
+        post_geom_sizes = [
+            (total_size[0], total_size[1], base_thickness),
+            (
+                post_size,
+                post_size,
+                total_size[2] - ring_size[2] - base_thickness - tripod_capsule_r - tripod_capsule_h,
+            ),
+        ]
+        post_geom_locations = [
+            (0.0, 0.0, 2.0 * (tripod_capsule_r + tripod_capsule_h)),
+            (
+                total_size[0] - post_size,
+                total_size[1] - post_size,
+                2.0 * (tripod_capsule_r + tripod_capsule_h + base_thickness),
+            ),
+        ]
+        for i in range(2):
+            add_to_dict(
+                dic=obj_args,
+                geom_types="box",
+                geom_locations=post_geom_locations[i],
+                geom_quats=(1, 0, 0, 0),
+                geom_sizes=post_geom_sizes[i],
+                geom_names="post_{}".format(i),
+                geom_rgbas=None,
+                geom_materials=self.tripod_mat_name,
+                geom_frictions=None,
+            )
+
+        # Add back in base args and site args
+        obj_args.update(base_args)
+
+        # Return this dict
+        return obj_args

GR00T-WholeBodyControl/decoupled_wbc/dexmg/gr00trobocasa/robocasa/models/objects/composite_body/bin_with_handles.py

@@ -0,0 +1,162 @@
+from robosuite.models.objects import CompositeBodyObject, BoxObject, Bin
+import numpy as np
+
+import robosuite.utils.transform_utils as T
+from robosuite.utils.mjcf_utils import array_to_string
+from robosuite.utils.mjcf_utils import RED, BLUE, CustomMaterial
+
+
+class BinWithHandles(CompositeBodyObject):
+    """
+    Bin with simple square handles on each side.
+    """
+
+    def __init__(
+        self,
+        name,
+        bin_size,
+        bin_wall_thickness,
+        bin_transparent_walls,
+        bin_upside_down,
+        center_handle_size,
+        adjacent_handle_size,
+        joints="default",
+        rgba=(0.2, 0.1, 0.0, 1.0),
+        material=None,
+        density=1000.0,
+        friction=None,
+    ):
+
+        # Object properties
+
+        # FULL size of bin
+        self.bin_size = list(bin_size)
+        self.bin_wall_thickness = bin_wall_thickness
+        self.bin_transparent_walls = bin_transparent_walls
+        self.bin_upside_down = bin_upside_down
+
+        # half-sizes of box geom used for center part of handle (which you grab)
+        self.center_handle_size = list(center_handle_size)
+
+        # half-sizes of box geoms used for adjacent parts of handle (not grabbed)
+        self.adjacent_handle_size = list(adjacent_handle_size)
+
+        # Create objects
+        objects = []
+        object_locations = []
+        object_quats = []
+        object_parents = []
+
+        # bin
+        self.bin = Bin(
+            name="bin",
+            bin_size=self.bin_size,
+            wall_thickness=self.bin_wall_thickness,
+            transparent_walls=self.bin_transparent_walls,
+            rgba=rgba,
+            material=material,
+            density=density,
+            friction=friction,
+            upside_down=bin_upside_down,
+        )
+        objects.append(self.bin)
+        object_locations.append([0.0, 0.0, 0.0])
+        object_quats.append([1.0, 0.0, 0.0, 0.0])
+        object_parents.append(None)
+
+        # handles on each side
+
+        left_handle_1_loc = [
+            0.0,
+            -(
+                self.bin_size[1] / 2.0
+                + 2.0 * self.adjacent_handle_size[1]
+                + self.center_handle_size[1]
+            ),
+            0.0,
+        ]
+        left_handle_1_size = self.center_handle_size
+
+        left_handle_2_loc = [
+            (self.center_handle_size[0] - self.adjacent_handle_size[0]),
+            -(self.bin_size[1] / 2.0 + self.adjacent_handle_size[1]),
+            0.0,
+        ]
+        left_handle_2_size = self.adjacent_handle_size
+
+        left_handle_3_loc = [
+            -(self.center_handle_size[0] - self.adjacent_handle_size[0]),
+            -(self.bin_size[1] / 2.0 + self.adjacent_handle_size[1]),
+            0.0,
+        ]
+        left_handle_3_size = self.adjacent_handle_size
+
+        right_handle_1_loc = [
+            0.0,
+            (
+                self.bin_size[1] / 2.0
+                + 2.0 * self.adjacent_handle_size[1]
+                + self.center_handle_size[1]
+            ),
+            0.0,
+        ]
+        right_handle_1_size = self.center_handle_size
+
+        right_handle_2_loc = [
+            (self.center_handle_size[0] - self.adjacent_handle_size[0]),
+            (self.bin_size[1] / 2.0 + self.adjacent_handle_size[1]),
+            0.0,
+        ]
+        right_handle_2_size = self.adjacent_handle_size
+
+        right_handle_3_loc = [
+            -(self.center_handle_size[0] - self.adjacent_handle_size[0]),
+            (self.bin_size[1] / 2.0 + self.adjacent_handle_size[1]),
+            0.0,
+        ]
+        right_handle_3_size = self.adjacent_handle_size
+
+        handle_locs = [
+            left_handle_1_loc,
+            left_handle_2_loc,
+            left_handle_3_loc,
+            right_handle_1_loc,
+            right_handle_2_loc,
+            right_handle_3_loc,
+        ]
+        handle_sizes = [
+            left_handle_1_size,
+            left_handle_2_size,
+            left_handle_3_size,
+            right_handle_1_size,
+            right_handle_2_size,
+            right_handle_3_size,
+        ]
+        handle_ind = 1
+        for b_loc, b_size in zip(handle_locs, handle_sizes):
+            this_handle = BoxObject(
+                name="handle_{}".format(handle_ind),
+                size=b_size,
+                rgba=rgba,
+                material=material,
+                density=density,
+                friction=friction,
+                joints=None,
+            )
+            objects.append(this_handle)
+            object_locations.append(b_loc)
+            object_quats.append([1.0, 0.0, 0.0, 0.0])
+            object_parents.append(None)
+            handle_ind += 1
+
+        # Run super init
+        super().__init__(
+            name=name,
+            objects=objects,
+            object_locations=object_locations,
+            object_quats=object_quats,
+            object_parents=object_parents,
+            joints=joints,
+            # total_size=body_total_size,
+            # locations_relative_to_corner=True,
+        )

GR00T-WholeBodyControl/decoupled_wbc/dexmg/gr00trobocasa/robocasa/models/objects/composite_body/coffee_machine.py

@@ -0,0 +1,244 @@
+from robosuite.models.objects import CompositeBodyObject, BoxObject
+from robocasa.models.objects.composite_body.cup import (
+    CupObject,
+)
+from robocasa.models.objects.xml_objects import (
+    CoffeeMachineBodyObject,
+    CoffeeMachineLidObject,
+    CoffeeMachineBaseObject,
+)
+import numpy as np
+
+from robosuite.utils.mjcf_utils import array_to_string
+from robosuite.utils.mjcf_utils import RED, BLUE, CustomMaterial
+
+
+class CoffeeMachineObject(CompositeBodyObject):
+    """
+    Coffee machine object with a lid fixed on a hinge joint.
+    """
+
+    def __init__(
+        self,
+        name,
+        add_cup=True,
+        pod_holder_friction=None,
+        joints="default",
+        density=1000.0,
+    ):
+
+        # pieces of the coffee machine
+        body = CoffeeMachineBodyObject(name="body")
+        body_size = body.get_bounding_box_half_size()
+        body_location = [0.0, 0.0, 0.0]
+
+        lid = CoffeeMachineLidObject(name="lid")
+        lid_size = self.lid_size = lid.get_bounding_box_half_size()
+        # add tolerance to allow lid to open fully
+        lid_location = [
+            body_size[0] - lid_size[0],
+            2.0 * body_size[1] + 0.01,
+            2.0 * (body_size[2] - lid_size[2]) + 0.005,
+        ]
+
+        # add in hinge joint to lid
+        hinge_pos = [0.0, -lid_size[1], 0.0]
+        hinge_joint = dict(
+            type="hinge",
+            axis="1 0 0",
+            pos=array_to_string(hinge_pos),
+            limited="true",
+            range="{} {}".format(0, 2.0 * np.pi / 3.0),
+            damping="0.005",
+        )
+        body_joints = dict(lid_main=[hinge_joint])  # note: "main" gets appended to body name
+        lid = CoffeeMachineLidObject(name="lid")
+
+        base = CoffeeMachineBaseObject(name="base")
+        base_size = base.get_bounding_box_half_size()
+        base_location = [body_size[0] - base_size[0], 2.0 * body_size[1], 0.0]
+
+        pod_holder_holder = BoxObject(
+            name="pod_holder_holder",
+            size=[
+                0.01,
+                # tolerance for having the lid stick out a little from the holder
+                0.9 * (lid_size[1] - lid_size[0]),
+                0.005,
+            ],
+            rgba=[0.839, 0.839, 0.839, 1],  # silver
+            joints=None,
+        )
+        pod_holder_holder_size = pod_holder_holder.get_bounding_box_half_size()
+        pod_holder_holder_location = [
+            body_size[0] - pod_holder_holder_size[0],
+            2.0 * body_size[1],
+            # put right underneath lid
+            2.0 * (body_size[2] - lid_size[2] - pod_holder_holder_size[2]),
+        ]
+
+        pod_holder = CupObject(
+            name="pod_holder",
+            outer_cup_radius=lid_size[0],
+            inner_cup_radius=0.028,
+            cup_height=0.028,
+            cup_ngeoms=64,  # 8,
+            cup_base_height=0.005,
+            cup_base_offset=0.002,
+            add_handle=False,
+            rgba=[1, 0, 0, 1],
+            density=density,
+            joints=None,
+            friction=pod_holder_friction,
+        )
+        pod_holder_size = self.pod_holder_size = pod_holder.get_bounding_box_half_size()
+        # pod_holder_size = self.pod_holder_size = np.array([0.0295, 0.0295, 0.028 ])
+        pod_holder_location = [
+            body_size[0] - pod_holder_size[0],
+            2.0 * (body_size[1] + pod_holder_holder_size[1]),
+            # put right underneath lid
+            2.0 * (body_size[2] - lid_size[2] - pod_holder_size[2]),
+        ]
+
+        total_size = [
+            body_size[0],
+            body_size[1] + base_size[1],
+            body_size[2],
+        ]
+
+        objects = [
+            body,
+            lid,
+            base,
+            pod_holder_holder,
+            pod_holder,
+        ]
+
+        object_locations = [
+            body_location,
+            lid_location,
+            base_location,
+            pod_holder_holder_location,
+            pod_holder_location,
+        ]
+
+        object_quats = [
+            [0.0, 0.0, 0.0, 1.0],  # z-rotate body and base by 180
+            [1.0, 0.0, 0.0, 0.0],
+            [0.0, 0.0, 0.0, 1.0],
+            [1.0, 0.0, 0.0, 0.0],
+            [1.0, 0.0, 0.0, 0.0],
+        ]
+
+        # add a rigidly mounted cup to the base
+        self.add_cup = add_cup
+        if self.add_cup:
+            cup = CupObject(
+                name="cupppp",
+                outer_cup_radius=0.03,
+                inner_cup_radius=0.025,
+                cup_height=0.025,
+                cup_ngeoms=64,  # 8,
+                cup_base_height=0.005,
+                cup_base_offset=0.005,
+                add_handle=True,
+                handle_outer_radius=0.015,
+                handle_inner_radius=0.010,
+                handle_thickness=0.003,
+                handle_ngeoms=64,
+                rgba=[0.839, 0.839, 0.839, 1],
+                density=1000.0,
+                joints=None,
+            )
+            cup_total_size = cup.get_bounding_box_half_size()
+            # cup_total_size = np.array([0.03 , 0.045, 0.025])
+            objects.append(cup)
+            object_locations.append(
+                [
+                    body_size[0] - cup_total_size[0],
+                    2.0 * (body_size[1] + pod_holder_holder_size[1])
+                    + pod_holder_size[1]
+                    - cup_total_size[1],
+                    2.0 * base_size[2],
+                ]
+            )
+            rot_angle = -np.pi / 2.0
+            object_quats.append([np.cos(rot_angle / 2), 0, 0, np.sin(rot_angle / 2)])
+
+        object_parents = [None] * len(objects)
+
+        """
+        Variables to compare:
+
+        objects
+        [<robosuite.models.objects.xml_objects.CoffeeMachineBodyObject at 0x10f0f7640>,
+         <robosuite.models.objects.xml_objects.CoffeeMachineLidObject at 0x10f0f7c10>,
+         <robosuite.models.objects.xml_objects.CoffeeMachineBaseObject at 0x10f0f7190>,
+         <robosuite.models.objects.generated_objects.BoxObject at 0x10f0f7d90>,
+         <robosuite.models.objects.generated_objects.CupObject at 0x10f0f71f0>,
+         <robosuite.models.objects.generated_objects.CupObject at 0x10f0f7070>]
+
+         object_locations
+         [[0.0, 0.0, 0.0],
+         [0.056999999999999995, 0.21100000000000002, 0.20700000000000002],
+         [0.04449999999999999, 0.201, 0.0],
+         [0.0765, 0.201, 0.192],
+         [0.056999999999999995, 0.22710000000000002, 0.14600000000000002],
+         [0.056499999999999995, 0.21160000000000007, 0.01]]
+
+        object_quats
+         [[0.0, 0.0, 0.0, 1.0],
+         [1.0, 0.0, 0.0, 0.0],
+         [0.0, 0.0, 0.0, 1.0],
+         [1.0, 0.0, 0.0, 0.0],
+         [1.0, 0.0, 0.0, 0.0],
+         [0.7071067811865476, 0, 0, -0.7071067811865475]]
+
+         body_size
+         array([0.0865, 0.1005, 0.1105])
+
+         lid_size
+         array([0.0295, 0.044 , 0.0095])
+
+         lid_location
+         [0.056999999999999995, 0.21100000000000002, 0.20700000000000002]
+
+         base_size
+         array([0.042, 0.05 , 0.005])
+
+         base_location
+         [0.04449999999999999, 0.201, 0.0]
+
+         pod_holder_holder_size
+         array([0.01   , 0.01305, 0.005  ])
+
+         pod_holder_holder_location
+         [0.0765, 0.201, 0.192]
+
+         pod_holder_size
+         array([0.0295, 0.0295, 0.028 ])
+
+         pod_holder_location
+         [0.056999999999999995, 0.22710000000000002, 0.14600000000000002]
+
+         total_size
+         [0.0865, 0.15050000000000002, 0.1105]
+
+         cup.total_size
+         array([0.03 , 0.045, 0.025])
+        """
+
+        # Run super init
+        super().__init__(
+            name=name,
+            objects=objects,
+            object_locations=object_locations,
+            object_quats=object_quats,
+            object_parents=object_parents,
+            body_joints=body_joints,  # make sure hinge joint is added
+            joints=joints,
+            # joints="default", # coffee machine can move
+            # joints=None, # coffee machine does not move
+            total_size=total_size,
+            locations_relative_to_corner=True,
+        )

GR00T-WholeBodyControl/decoupled_wbc/dexmg/gr00trobocasa/robocasa/models/objects/composite_body/inverse_stacked_cylinder.py

@@ -0,0 +1,139 @@
+from robosuite.models.objects import (
+    CompositeBodyObject,
+    BoxObject,
+    CylinderObject,
+    HollowCylinderObject,
+)
+import numpy as np
+
+import robosuite.utils.transform_utils as T
+from robosuite.utils.mjcf_utils import array_to_string
+from robosuite.utils.mjcf_utils import RED, BLUE, CustomMaterial
+
+
+class InverseStackedCylinderObject(CompositeBodyObject):
+    """
+    Inverse of stacked cylinder object, where the top piece is a hollow cylinder object
+    and the bottom piece is a cylinder object. Optionally add a square base for stability.
+    """
+
+    def __init__(
+        self,
+        name,
+        radius_1,
+        radius_2,
+        height_1,
+        height_2,
+        ngeoms=64,
+        joints="default",
+        rgba=None,
+        material=None,
+        density=100.0,
+        friction=None,
+        square_base_width=None,
+        square_base_height=None,
+    ):
+
+        # Object properties
+
+        # radius of first (bottom) cylinder and inner radius of second (top) hollow cylinder
+        self.r1 = radius_1
+        self.r2 = radius_2
+
+        # half-height of first (bottom) cylinder and second (top) hollow cylinder
+        self.h1 = height_1
+        self.h2 = height_2
+
+        # num geoms to approximate the hollow cylinder
+        self.ngeoms = ngeoms
+
+        # whether to add square base
+        self.add_square_base = (square_base_width is not None) and (square_base_height is not None)
+
+        # half-width and half-height for square base
+        self.square_base_width = square_base_width
+        self.square_base_height = square_base_height
+
+        # Create objects
+        objects = []
+        object_locations = []
+        object_quats = []
+        object_parents = []
+
+        # NOTE: we will place the object frame at the vertical center of the two stacked cylinders
+        z_center = (self.h1 + self.h2) / 2.0
+        c1_offset = self.h1 - z_center
+        c2_offset = 2.0 * self.h1 + self.h2 - z_center
+
+        # first (bottom) cylinder
+        self.cylinder_1 = CylinderObject(
+            name="cylinder_1",
+            size=[self.r1, self.h1],
+            rgba=rgba,
+            material=material,
+            density=density,
+            friction=friction,
+            solref=[0.02, 1.0],
+            # solimp=[0.998, 0.998, 0.001],
+            solimp=[0.9, 0.95, 0.001],
+            joints=None,
+        )
+        objects.append(self.cylinder_1)
+        object_locations.append([0.0, 0.0, c1_offset])
+        object_quats.append([1.0, 0.0, 0.0, 0.0])
+        object_parents.append(None)
+
+        # second (top) hollow cylinder
+        self.cylinder_2 = HollowCylinderObject(
+            name="cylinder_2",
+            outer_radius=self.r1,  # match radius of first cylinder
+            inner_radius=self.r2,
+            height=self.h2,
+            ngeoms=self.ngeoms,
+            rgba=rgba,
+            material=material,
+            density=density,
+            friction=friction,
+            # TODO: maybe tune solimp and try (0.998, 0.998, 0.001)
+            solref=[0.02, 1.0],
+            solimp=[0.9, 0.95, 0.001],
+            # solimp=(0.998, 0.998, 0.001),
+        )
+        objects.append(self.cylinder_2)
+        object_locations.append([0.0, 0.0, c2_offset])
+        object_quats.append([1.0, 0.0, 0.0, 0.0])
+        object_parents.append(None)
+
+        # # total size of object
+        # max_r = max(self.r1, self.r2)
+        # body_total_size = [max_r, max_r, self.h1 + self.h2]
+
+        if self.add_square_base:
+            # add square base underneath bottom cylinder
+            s1_offset = c1_offset - (self.square_base_height + self.h1)
+            self.square_base = BoxObject(
+                name="square_base",
+                size=[
+                    self.square_base_width,
+                    self.square_base_width,
+                    self.square_base_height,
+                ],
+                rgba=rgba,
+                material=material,
+            )
+            objects.append(self.square_base)
+            object_locations.append([0.0, 0.0, s1_offset])
+            object_quats.append([1.0, 0.0, 0.0, 0.0])
+            object_parents.append(None)
+
+        # Run super init
+        super().__init__(
+            name=name,
+            objects=objects,
+            object_locations=object_locations,
+            object_quats=object_quats,
+            object_parents=object_parents,
+            joints=joints,
+            # total_size=body_total_size,
+            # locations_relative_to_corner=True,
+        )

GR00T-WholeBodyControl/decoupled_wbc/dexmg/gr00trobocasa/robocasa/models/objects/composite_body/lightbulb.py

@@ -0,0 +1,153 @@
+from robosuite.models.objects import (
+    CompositeBodyObject,
+    BoxObject,
+    CylinderObject,
+    BallObject,
+)
+import numpy as np
+
+import robosuite.utils.transform_utils as T
+from robosuite.utils.mjcf_utils import array_to_string, new_site
+from robosuite.utils.mjcf_utils import RED, BLUE, CustomMaterial
+
+
+class BallObjectWithSite(BallObject):
+    """
+    A ball object with a inner site (used for the bulb).
+    """
+
+    def _get_object_subtree(self):
+        # tree = super()._get_object_subtree()
+        tree = self._get_object_subtree_(ob_type="sphere")
+        site_element_attr = self.get_site_attrib_template()
+
+        site_element_attr["pos"] = "0 0 0.1"
+        site_element_attr["name"] = "center_site"
+        site_element_attr["size"] = "{} {} {}".format(
+            2.0 * self.size[0], 2.0 * self.size[0], 2.0 * self.size[0]
+        )
+        site_element_attr["rgba"] = "1.0 0.0 0.0 1.0"
+        site_element_attr["group"] = "1"
+        # site_element_attr["rgba"] = "1.0 0.976 0.839 0.9"
+        # site_element_attr["rgba"] = "1.0 0.976 0.839 0.0"
+        tree.append(new_site(**site_element_attr))
+        return tree
+
+
+class LightbulbObject(CompositeBodyObject):
+    """
+    A simple lightbulb constructed out of a base of alternating radius cylinders
+    and a sphere on top.
+    """
+
+    def __init__(
+        self,
+        name,
+        radius_low,
+        radius_high,
+        cylinder_height,
+        num_cylinders,
+        sphere_radius,
+        joints="default",
+        density=100.0,
+        friction=None,
+    ):
+
+        # Object properties
+
+        # radii of alternating cylinders for base
+        self.radius_low = radius_low
+        self.radius_high = radius_high
+
+        # half-height of each cylinder
+        self.cylinder_height = cylinder_height
+
+        # number of cylinders for base
+        self.num_cylinders = num_cylinders
+
+        # radius of sphere at top
+        self.sphere_radius = sphere_radius
+
+        # toggle between translucent and yellow
+        self.translucent_rgba = (1.0, 1.0, 1.0, 0.3)
+        self.yellow_rgba = (1.0, 0.976, 0.839, 0.7)
+        # self.yellow_rgba = (0.0, 0.0, 0.0, 0.0)
+
+        # Create objects
+        objects = []
+        object_locations = []
+        object_quats = []
+        object_parents = []
+
+        # NOTE: we will place the object frame at the vertical center of all the stacked objects
+        self.z_center = (
+            (2.0 * self.cylinder_height) * self.num_cylinders + 2.0 * self.sphere_radius
+        ) / 2.0
+
+        metal = CustomMaterial(
+            texture="Metal",
+            tex_name="metal",
+            mat_name="MatMetal",
+            tex_attrib={"type": "cube"},
+            mat_attrib={"specular": "1", "shininess": "0.3", "rgba": "0.9 0.9 0.9 1"},
+        )
+
+        # we will define all objects relative to the bottom of the object, and then subtract the z_center value
+        for cylinder_ind in range(self.num_cylinders):
+            r = self.radius_low if ((cylinder_ind % 2) == 0) else self.radius_high
+            cyl_obj = CylinderObject(
+                name="cylinder_{}".format(cylinder_ind),
+                size=[r, self.cylinder_height],
+                rgba=None,
+                material=metal,
+                density=density,
+                friction=friction,
+                solref=[0.02, 1.0],
+                solimp=[0.9, 0.95, 0.001],
+                joints=None,
+            )
+            objects.append(cyl_obj)
+            z_cyl = (2.0 * cylinder_ind + 1) * self.cylinder_height
+            object_locations.append([0.0, 0.0, z_cyl])
+            object_quats.append([1.0, 0.0, 0.0, 0.0])
+            object_parents.append(None)
+
+        # then add translucent sphere at top
+        self.bulb = BallObject(
+            name="bulb",
+            size=[self.sphere_radius],
+            density=density,
+            friction=friction,
+            rgba=self.translucent_rgba,
+            material=None,
+            joints=None,
+        )
+        objects.append(self.bulb)
+        z_bulb = object_locations[-1][2] + self.cylinder_height + self.sphere_radius
+        object_locations.append([0.0, 0.0, z_bulb])
+        object_quats.append([1.0, 0.0, 0.0, 0.0])
+        object_parents.append(None)
+
+        # do frame conversion from bottom of object to z_center
+        object_locations = [[loc[0], loc[1], loc[2] - self.z_center] for loc in object_locations]
+
+        # add site that can be toggled to turn lightbulb on
+        sites = [
+            dict(
+                name="bulb_on",
+                pos=array_to_string(object_locations[-1]),
+                size="{}".format(0.95 * self.sphere_radius),
+                rgba=array_to_string(self.yellow_rgba),
+            )
+        ]
+
+        # Run super init
+        super().__init__(
+            name=name,
+            objects=objects,
+            object_locations=object_locations,
+            object_quats=object_quats,
+            object_parents=object_parents,
+            joints=joints,
+            sites=sites,
+        )

GR00T-WholeBodyControl/decoupled_wbc/dexmg/gr00trobocasa/robocasa/models/objects/composite_body/sliding_box.py

@@ -0,0 +1,132 @@
+import numpy as np
+
+from robosuite.models.objects import BoxObject, CompositeBodyObject, CylinderObject
+from robosuite.utils.mjcf_utils import BLUE, RED, CustomMaterial, array_to_string
+
+
+class SlidingBoxObject(CompositeBodyObject):
+    """
+    An example object that demonstrates the CompositeBodyObject functionality. This object consists of two cube bodies
+    joined together by a slide joint allowing one box to slide on top of the other.
+
+    Args:
+        name (str): Name of this object
+
+        box1_size (3-array): (L, W, H) half-sizes for the first box
+
+        box2_size (3-array): (L, W, H) half-sizes for the second box
+
+        use_texture (bool): set True if using wood textures for the blocks
+    """
+
+    def __init__(
+        self,
+        name,
+        box1_size=(0.1, 0.1, 0.02),
+        box2_size=(0.02, 0.02, 0.02),
+        use_texture=True,
+    ):
+        # Set box sizes
+        self.box1_size = np.array(box1_size)
+        self.box2_size = np.array(box2_size)
+
+        # Set box densities
+        self.box1_density = 10000.0
+        self.box2_density = 100.0
+
+        # Set texture attributes
+        self.use_texture = use_texture
+        self.box1_material = None
+        self.box2_material = None
+        self.box1_rgba = RED
+        self.box2_rgba = BLUE
+
+        # Define materials we want to use for this object
+        if self.use_texture:
+            # Remove RGBAs
+            self.box1_rgba = None
+            self.box2_rgba = None
+
+            # Set materials for each box
+            tex_attrib = {
+                "type": "cube",
+            }
+            mat_attrib = {
+                "texrepeat": "3 3",
+                "specular": "0.4",
+                "shininess": "0.1",
+            }
+            self.box1_material = CustomMaterial(
+                texture="WoodRed",
+                tex_name="box1_tex",
+                mat_name="box1_mat",
+                tex_attrib=tex_attrib,
+                mat_attrib=mat_attrib,
+            )
+            self.box2_material = CustomMaterial(
+                texture="WoodBlue",
+                tex_name="box2_tex",
+                mat_name="box2_mat",
+                tex_attrib=tex_attrib,
+                mat_attrib=mat_attrib,
+            )
+
+        # Create objects
+        objects = []
+        for i, (size, mat, rgba, density) in enumerate(
+            zip(
+                (self.box1_size, self.box2_size),
+                (self.box1_material, self.box2_material),
+                (self.box1_rgba, self.box2_rgba),
+                (self.box1_density, self.box2_density),
+            )
+        ):
+            objects.append(
+                BoxObject(
+                    name=f"box{i + 1}",
+                    size=size,
+                    rgba=rgba,
+                    material=mat,
+                )
+            )
+
+        # Define slide joint
+        rel_joint_pos = [0, 0, 0]  # at second box
+        joint_lim = self.box1_size[1] - self.box2_size[1]
+        slide_joint = {
+            "name": "box_slide",
+            "type": "slide",
+            "axis": "0 1 0",  # y-axis slide
+            "pos": array_to_string(rel_joint_pos),
+            "springref": "0",
+            "springdamper": "0.1 1.0",  # mass-spring system with 0.1 time constant, 1.0 damping ratio
+            "limited": "true",
+            "range": "{} {}".format(-joint_lim, joint_lim),
+        }
+
+        # Define positions -- second box should lie on top of first box
+        positions = [
+            np.zeros(3),  # First box is centered at top-level body anyways
+            np.array([0, 0, self.box1_size[2] + self.box2_size[2]]),
+        ]
+
+        quats = [
+            None,  # Default quaternion for box 1
+            None,  # Default quaternion for box 2
+        ]
+
+        # Define parents -- which body each is aligned to
+        parents = [
+            None,  # box 1 attached to top-level body
+            objects[0].root_body,  # box 2 attached to box 1
+        ]
+
+        # Run super init
+        super().__init__(
+            name=name,
+            objects=objects,
+            object_locations=positions,
+            object_quats=quats,
+            object_parents=parents,
+            body_joints={objects[1].root_body: [slide_joint]},
+        )

GR00T-WholeBodyControl/decoupled_wbc/dexmg/gr00trobocasa/robocasa/models/objects/composite_body/stacked_box.py

@@ -0,0 +1,111 @@
+from robosuite.models.objects import CompositeBodyObject, BoxObject, CylinderObject, Bin
+import numpy as np
+
+import robosuite.utils.transform_utils as T
+from robosuite.utils.mjcf_utils import array_to_string
+from robosuite.utils.mjcf_utils import RED, BLUE, CustomMaterial
+
+
+class StackedBoxObject(CompositeBodyObject):
+    """
+    Two boxes - one stacked on top of the other.
+    """
+
+    def __init__(
+        self,
+        name,
+        box_1_size,
+        box_2_size,
+        joints="default",
+        box_1_rgba=None,
+        box_2_rgba=None,
+        box_1_material=None,
+        box_2_material=None,
+        density=100.0,
+        friction=None,
+        make_box_2_transparent=False,
+    ):
+
+        # Object properties
+
+        # half-sizes of first (bottom) box
+        self.box_1_size = list(box_1_size)
+
+        # half-sizes of second (top) box
+        self.box_2_size = list(box_2_size)
+
+        # maybe make box 2 have transparent top and bottom walls
+        self.make_box_2_transparent = make_box_2_transparent
+
+        # Create objects
+        objects = []
+        object_locations = []
+        object_quats = []
+        object_parents = []
+
+        # NOTE: we will place the object frame at the vertical center of the two stacked boxes
+        z_center = (self.box_1_size[2] + self.box_2_size[2]) / 2.0
+        b1_offset = self.box_1_size[2] - z_center
+        b2_offset = 2.0 * self.box_1_size[2] + self.box_2_size[2] - z_center
+
+        # first (bottom) box
+        self.box_1 = BoxObject(
+            name="box_1",
+            size=self.box_1_size,
+            rgba=box_1_rgba,
+            material=box_1_material,
+            density=density,
+            friction=friction,
+            joints=None,
+        )
+        objects.append(self.box_1)
+        object_locations.append([0.0, 0.0, b1_offset])
+        object_quats.append([1.0, 0.0, 0.0, 0.0])
+        object_parents.append(None)
+
+        # second (top) box
+        if self.make_box_2_transparent:
+            self.box_2 = Bin(
+                name="box_2",
+                bin_size=(
+                    2.0 * self.box_2_size[0],
+                    2.0 * self.box_2_size[1],
+                    2.0 * self.box_2_size[2],
+                ),
+                wall_thickness=0.01,
+                transparent_walls=False,
+                friction=friction,
+                density=density,
+                use_texture=True,
+                rgba=box_2_rgba,
+                material=box_2_material,
+                upside_down=False,
+                add_second_base=True,
+                transparent_base=True,
+            )
+        else:
+            self.box_2 = BoxObject(
+                name="box_2",
+                size=self.box_2_size,
+                rgba=box_2_rgba,
+                material=box_2_material,
+                density=density,
+                friction=friction,
+                joints=None,
+            )
+        objects.append(self.box_2)
+        object_locations.append([0.0, 0.0, b2_offset])
+        object_quats.append([1.0, 0.0, 0.0, 0.0])
+        object_parents.append(None)
+
+        # Run super init
+        super().__init__(
+            name=name,
+            objects=objects,
+            object_locations=object_locations,
+            object_quats=object_quats,
+            object_parents=object_parents,
+            joints=joints,
+            # total_size=body_total_size,
+            # locations_relative_to_corner=True,
+        )

GR00T-WholeBodyControl/decoupled_wbc/dexmg/gr00trobocasa/robocasa/models/objects/composite_body/stacked_cylinder.py

@@ -0,0 +1,127 @@
+from robosuite.models.objects import CompositeBodyObject, BoxObject, CylinderObject
+import numpy as np
+
+import robosuite.utils.transform_utils as T
+from robosuite.utils.mjcf_utils import array_to_string
+from robosuite.utils.mjcf_utils import RED, BLUE, CustomMaterial
+
+
+class StackedCylinderObject(CompositeBodyObject):
+    """
+    Two cylinders - one stacked on top of the other.
+    Optionally add a square base for stability.
+    """
+
+    def __init__(
+        self,
+        name,
+        radius_1,
+        radius_2,
+        height_1,
+        height_2,
+        joints="default",
+        rgba=None,
+        material=None,
+        density=100.0,
+        friction=None,
+        square_base_width=None,
+        square_base_height=None,
+    ):
+
+        # Object properties
+
+        # radius of first (bottom) cylinder and second (top) cylinder
+        self.r1 = radius_1
+        self.r2 = radius_2
+
+        # half-height of first (bottom) cylinder and second (top) cylinder
+        self.h1 = height_1
+        self.h2 = height_2
+
+        # whether to add square base
+        self.add_square_base = (square_base_width is not None) and (square_base_height is not None)
+
+        # half-width and half-height for square base
+        self.square_base_width = square_base_width
+        self.square_base_height = square_base_height
+
+        # Create objects
+        objects = []
+        object_locations = []
+        object_quats = []
+        object_parents = []
+
+        # NOTE: we will place the object frame at the vertical center of the two stacked cylinders
+        z_center = (self.h1 + self.h2) / 2.0
+        c1_offset = self.h1 - z_center
+        c2_offset = 2.0 * self.h1 + self.h2 - z_center
+
+        # first (bottom) cylinder
+        self.cylinder_1 = CylinderObject(
+            name="cylinder_1",
+            size=[self.r1, self.h1],
+            rgba=rgba,
+            material=material,
+            density=density,
+            friction=friction,
+            solref=[0.02, 1.0],
+            solimp=[0.9, 0.95, 0.001],
+            # solimp=[0.998, 0.998, 0.001],
+            joints=None,
+        )
+        objects.append(self.cylinder_1)
+        object_locations.append([0.0, 0.0, c1_offset])
+        object_quats.append([1.0, 0.0, 0.0, 0.0])
+        object_parents.append(None)
+
+        # second (top) cylinder
+        self.cylinder_2 = CylinderObject(
+            name="cylinder_2",
+            size=[self.r2, self.h2],
+            rgba=rgba,
+            material=material,
+            density=density,
+            friction=friction,
+            solref=[0.02, 1.0],
+            # solimp=[0.998, 0.998, 0.001],
+            solimp=[0.9, 0.95, 0.001],
+            joints=None,
+        )
+        objects.append(self.cylinder_2)
+        object_locations.append([0.0, 0.0, c2_offset])
+        object_quats.append([1.0, 0.0, 0.0, 0.0])
+        object_parents.append(None)
+
+        # # total size of object
+        # max_r = max(self.r1, self.r2)
+        # body_total_size = [max_r, max_r, self.h1 + self.h2]
+
+        if self.add_square_base:
+            # add square base underneath bottom cylinder
+            s1_offset = c1_offset - (self.square_base_height + self.h1)
+            self.square_base = BoxObject(
+                name="square_base",
+                size=[
+                    self.square_base_width,
+                    self.square_base_width,
+                    self.square_base_height,
+                ],
+                rgba=rgba,
+                material=material,
+            )
+            objects.append(self.square_base)
+            object_locations.append([0.0, 0.0, s1_offset])
+            object_quats.append([1.0, 0.0, 0.0, 0.0])
+            object_parents.append(None)
+
+        # Run super init
+        super().__init__(
+            name=name,
+            objects=objects,
+            object_locations=object_locations,
+            object_quats=object_quats,
+            object_parents=object_parents,
+            joints=joints,
+            # total_size=body_total_size,
+            # locations_relative_to_corner=True,
+        )

GR00T-WholeBodyControl/decoupled_wbc/dexmg/gr00trobocasa/robocasa/models/objects/composite_body/stove_plug.py

@@ -0,0 +1,396 @@
+from robosuite.models.objects import CompositeBodyObject, BoxObject, CylinderObject
+import numpy as np
+
+import robosuite.utils.transform_utils as T
+from robosuite.utils.mjcf_utils import array_to_string
+from robosuite.utils.mjcf_utils import RED, BLUE, CustomMaterial
+from .inverse_stacked_cylinder import InverseStackedCylinderObject
+from .lightbulb import LightbulbObject
+
+import robosuite_task_zoo
+from robosuite_task_zoo.models.kitchen import StoveObject
+
+
+class StoveObjectNew(StoveObject):
+    """
+    Override some offsets for placement sampler.
+    """
+
+    @property
+    def bottom_offset(self):
+        # unused since we directly hardcode z
+        return np.array([0, 0, -0.02])
+
+    @property
+    def top_offset(self):
+        # unused since we directly hardcode z
+        return np.array([0, 0, 0.02])
+
+    @property
+    def horizontal_radius(self):
+        return 0.1
+
+
+class StovePlugObject(CompositeBodyObject):
+    """
+    Stove on wooden block with chain (plug) connected. Optionally replace the stove
+    with a lightbulb if @lightbulb_args is provided.
+    """
+
+    def __init__(
+        self,
+        name,
+        joints="default",
+        rgba=None,
+        material=None,
+        density=100.0,
+        friction=None,
+        stove_base_size=(0.12, 0.12, 0.01),
+        stove_z_half_size=0.025,
+        wire_box_geom_size=(0.005, 0.02, 0.005),
+        wire_box_geom_rgba=(0.0, 0.0, 0.0, 1.0),
+        num_box_geoms_left=5,
+        num_box_geoms_vert=8,
+        num_box_geoms_right=8,
+        merge_box_geoms=False,
+        merge_size=1,
+        cylinder_args=None,
+        lightbulb_args=None,
+    ):
+
+        # Object properties
+
+        # half sizes for stove base box object
+        # self.stove_base_size = (0.1, 0.1, 0.02)
+        self.stove_base_size = stove_base_size
+        self.stove_z_half_size = stove_z_half_size  # note: estimated approximately
+
+        # box geoms used for wire
+        self.wire_box_geom_size = wire_box_geom_size
+        self.wire_box_geom_rgba = wire_box_geom_rgba
+
+        # wire parameters - number of geoms to use for left, down, and right portions
+        self.num_box_geoms_left = num_box_geoms_left
+        self.num_box_geoms_vert = num_box_geoms_vert
+        self.num_box_geoms_right = num_box_geoms_right
+
+        # if true, merge the box geoms along each direction of the wire into a single box geom
+        self.merge_box_geoms = merge_box_geoms
+
+        # number of box geoms to use for each merged size (set to higher than 1 to merge the geoms into more
+        # than one box geom)
+        self.merge_size = merge_size
+
+        if cylinder_args is None:
+            # default cylinder args
+            cylinder_args = dict(
+                # bottom cylinder radius and half-height
+                radius_1=0.03,
+                height_1=0.01,
+                # top hollow cylinder inner radius and half-height
+                radius_2=0.025,
+                height_2=0.025,
+                # NOTE: reduce to 8 geoms if desired
+                ngeoms=64,
+                rgba=[0.839, 0.839, 0.839, 1],
+                density=1000.0,
+                # add square base
+                square_base_width=0.03,
+                square_base_height=0.005,
+            )
+        self.cylinder_args = dict(cylinder_args)
+        self.cylinder_args["joints"] = None
+
+        self.use_lightbulb = lightbulb_args is not None
+        if self.use_lightbulb:
+            self.lightbulb_args = dict(lightbulb_args)
+            self.lightbulb_args["joints"] = None
+
+        # materials
+        box_geom_material = CustomMaterial(
+            texture="Metal",
+            tex_name="metal",
+            mat_name="MatMetal",
+            tex_attrib={"type": "cube"},
+            mat_attrib={"specular": "1", "shininess": "0.3", "rgba": "0.9 0.9 0.9 1"},
+        )
+        stove_base_material = CustomMaterial(
+            texture="WoodLight",
+            tex_name="lightwood",
+            mat_name="lightwood_mat",
+            tex_attrib={"type": "cube"},
+            mat_attrib={"texrepeat": "1 1", "specular": "0.4", "shininess": "0.1"},
+        )
+
+        # params for ball joints used
+
+        # <joint name="J1" type="ball" pos="0 0 0" frictionloss="&friction;"/>
+        ball_joint_spec = {
+            "type": "ball",
+            "pos": "0 {} 0".format(self.wire_box_geom_size[1]),
+            "springref": "0",
+            "springdamper": "0.1 1.0",  # mass-spring system with 0.1 time constant, 1.0 damping ratio
+            "limited": "true",
+            "range": "0 {}".format(np.pi / 4),
+        }
+
+        # Create objects
+        objects = []
+        object_locations = []
+        object_quats = []
+        object_parents = []
+        object_joints = dict()
+
+        # NOTE: For absolute object locations (objects not defined relative to parent) we will use the stove base frame
+        #       as a frame of reference, and add an offset from the center of the object (approximated via full width)
+        #       to it.
+
+        # get an approximate x-y bounding box below, and place the center there, and define offset relative to stove base cente
+        approx_full_width = self.num_box_geoms_left * (2.0 * self.wire_box_geom_size[1]) + (
+            2.0 * self.stove_base_size[1]
+        )
+        approx_full_height = (self.num_box_geoms_vert + 2) * (2.0 * self.wire_box_geom_size[1])
+
+        stove_base_x_off = -(approx_full_height / 2.0) + self.stove_base_size[0]
+        stove_base_y_off = (approx_full_width / 2.0) - self.stove_base_size[1]
+
+        # base of stove
+        self.stove_base = BoxObject(
+            name="stove_base",
+            size=list(self.stove_base_size),
+            rgba=rgba,
+            material=stove_base_material,
+            joints=None,
+        )
+        objects.append(self.stove_base)
+        object_locations.append([stove_base_x_off, stove_base_y_off, 0.0])
+        object_quats.append([1.0, 0.0, 0.0, 0.0])
+        object_parents.append(None)
+
+        if self.use_lightbulb:
+            # lightbulb
+            self.lightbulb = LightbulbObject(
+                name="lightbulb",
+                **self.lightbulb_args,
+            )
+            objects.append(self.lightbulb)
+            object_locations.append(
+                [
+                    stove_base_x_off,
+                    stove_base_y_off,
+                    (self.stove_base_size[2] + self.lightbulb.z_center),
+                ]
+            )
+            object_quats.append([1.0, 0.0, 0.0, 0.0])
+            object_parents.append(None)
+        else:
+            # stove
+            self.stove = StoveObjectNew(
+                name="new_stove",
+                joints=None,
+            )
+            objects.append(self.stove)
+            object_locations.append(
+                [
+                    stove_base_x_off,
+                    stove_base_y_off,
+                    (self.stove_base_size[2] + self.stove_z_half_size),
+                ]
+            )
+            object_quats.append([1.0, 0.0, 0.0, 0.0])
+            object_parents.append(None)
+
+        # chain to the left of stove base
+        chain_ind = 0
+        left_chain_size = list(self.wire_box_geom_size)
+        num_geoms_iter = self.num_box_geoms_left - 1
+        if self.merge_box_geoms:
+            # only one big geom instead of chain of geoms
+            left_chain_size[1] *= self.num_box_geoms_left
+            left_chain_size[1] /= self.merge_size
+            # number of additional geoms to add
+            num_geoms_iter = self.merge_size - 1
+        left_chain_obj = BoxObject(
+            name="chain_{}".format(chain_ind),
+            size=list(left_chain_size),
+            rgba=list(self.wire_box_geom_rgba),
+            material=box_geom_material,
+            joints=None,
+        )
+        chain_ind += 1
+        objects.append(left_chain_obj)
+        object_locations.append(
+            [
+                stove_base_x_off - 0.75 * self.stove_base_size[0],
+                stove_base_y_off - (self.stove_base_size[1] + left_chain_size[1]),
+                0.0,
+            ]
+        )
+        # object_locations.append([0., -(self.stove_base_size[1] + left_chain_size[1]), 0.])
+        object_quats.append([1.0, 0.0, 0.0, 0.0])
+        object_parents.append(None)
+        if self.merge_box_geoms:
+            # add ball joint and make sure to place it at the edge of the geom
+            ball_joint = dict(ball_joint_spec)
+            ball_joint["name"] = "ball_joint_{}".format(chain_ind)
+            ball_joint["pos"] = "0 {} 0".format(left_chain_size[1])
+            object_joints[left_chain_obj.root_body] = [ball_joint]
+        for i in range(num_geoms_iter):
+            left_chain_obj = BoxObject(
+                name="chain_{}".format(chain_ind),
+                size=list(left_chain_size),
+                rgba=list(self.wire_box_geom_rgba),
+                material=box_geom_material,
+                joints=None,
+            )
+            chain_ind += 1
+            parent = objects[-1].root_body
+            objects.append(left_chain_obj)
+            object_locations.append([0.0, -2.0 * left_chain_size[1], 0.0])
+            object_quats.append([1.0, 0.0, 0.0, 0.0])
+            object_parents.append(parent)
+
+            # add ball joint and make sure to place it at the edge of the geom
+            ball_joint = dict(ball_joint_spec)
+            ball_joint["name"] = "ball_joint_{}".format(chain_ind)
+            ball_joint["pos"] = "0 {} 0".format(left_chain_size[1])
+            object_joints[left_chain_obj.root_body] = [ball_joint]
+
+        # add chain in downward direction
+        rot_quat = T.convert_quat(
+            T.axisangle2quat(np.array([0.0, 0.0, 1.0]) * (np.pi / 2.0)), to="wxyz"
+        )
+        vert_chain_size = list(self.wire_box_geom_size)
+        num_geoms_iter = self.num_box_geoms_vert - 1
+        if self.merge_box_geoms:
+            # only one big geom instead of chain of geoms
+            vert_chain_size[1] *= self.num_box_geoms_vert
+            vert_chain_size[1] /= self.merge_size
+            # number of additional geoms to add
+            num_geoms_iter = self.merge_size - 1
+        vert_chain_obj = BoxObject(
+            name="chain_{}".format(chain_ind),
+            size=list(vert_chain_size),
+            rgba=list(self.wire_box_geom_rgba),
+            material=box_geom_material,
+            joints=None,
+        )
+        chain_ind += 1
+        parent = objects[-1].root_body
+        objects.append(vert_chain_obj)
+        object_locations.append([vert_chain_size[1], -left_chain_size[1], 0.0])
+        object_quats.append(rot_quat)
+        object_parents.append(parent)
+
+        ball_joint = dict(ball_joint_spec)
+        ball_joint["name"] = "ball_joint_{}".format(chain_ind)
+        ball_joint["pos"] = "0 {} 0".format(vert_chain_size[1])
+        object_joints[vert_chain_obj.root_body] = [ball_joint]
+
+        for i in range(num_geoms_iter):
+            vert_chain_obj = BoxObject(
+                name="chain_{}".format(chain_ind),
+                size=list(vert_chain_size),
+                rgba=list(self.wire_box_geom_rgba),
+                material=box_geom_material,
+                joints=None,
+            )
+            chain_ind += 1
+            parent = objects[-1].root_body
+            objects.append(vert_chain_obj)
+            object_locations.append([0.0, -2.0 * vert_chain_size[1], 0.0])
+            object_quats.append([1.0, 0.0, 0.0, 0.0])
+            object_parents.append(parent)
+
+            ball_joint = dict(ball_joint_spec)
+            ball_joint["name"] = "ball_joint_{}".format(chain_ind)
+            ball_joint["pos"] = "0 {} 0".format(vert_chain_size[1])
+            object_joints[vert_chain_obj.root_body] = [ball_joint]
+
+        # add chain in rightward direction
+        rot_quat = T.convert_quat(
+            T.axisangle2quat(np.array([0.0, 0.0, 1.0]) * (np.pi / 2.0)), to="wxyz"
+        )
+        right_chain_size = list(self.wire_box_geom_size)
+        num_geoms_iter = self.num_box_geoms_right - 1
+        if self.merge_box_geoms:
+            # only one big geom instead of chain of geoms
+            right_chain_size[1] *= self.num_box_geoms_right
+            right_chain_size[1] /= self.merge_size
+            # number of additional geoms to add
+            num_geoms_iter = self.merge_size - 1
+        right_chain_obj = BoxObject(
+            name="chain_{}".format(chain_ind),
+            size=list(right_chain_size),
+            rgba=list(self.wire_box_geom_rgba),
+            material=box_geom_material,
+            joints=None,
+        )
+        chain_ind += 1
+        parent = objects[-1].root_body
+        objects.append(right_chain_obj)
+        object_locations.append([right_chain_size[1], -vert_chain_size[1], 0.0])
+        object_quats.append(rot_quat)
+        object_parents.append(parent)
+
+        ball_joint = dict(ball_joint_spec)
+        ball_joint["name"] = "ball_joint_{}".format(chain_ind)
+        ball_joint["pos"] = "0 {} 0".format(right_chain_size[1])
+        object_joints[right_chain_obj.root_body] = [ball_joint]
+
+        for i in range(num_geoms_iter):
+            right_chain_obj = BoxObject(
+                name="chain_{}".format(chain_ind),
+                size=list(right_chain_size),
+                rgba=list(self.wire_box_geom_rgba),
+                material=box_geom_material,
+                joints=None,
+            )
+            chain_ind += 1
+            parent = objects[-1].root_body
+            objects.append(right_chain_obj)
+            object_locations.append([0.0, -2.0 * right_chain_size[1], 0.0])
+            object_quats.append([1.0, 0.0, 0.0, 0.0])
+            object_parents.append(parent)
+
+            ball_joint = dict(ball_joint_spec)
+            ball_joint["name"] = "ball_joint_{}".format(chain_ind)
+            ball_joint["pos"] = "0 {} 0".format(right_chain_size[1])
+            object_joints[right_chain_obj.root_body] = [ball_joint]
+
+        # add cylinder object
+        self.cylinder_obj = InverseStackedCylinderObject(
+            name="cylinder_obj",
+            **self.cylinder_args,
+        )
+
+        x_off = 0.0
+        y_off = -(right_chain_size[1] + self.cylinder_obj.square_base_width)
+        z_off = ((self.cylinder_obj.h1 + self.cylinder_obj.h2) / 2.0) - right_chain_size[2]
+        parent = objects[-1].root_body
+        objects.append(self.cylinder_obj)
+        object_locations.append([x_off, y_off, z_off])
+        object_quats.append([1.0, 0.0, 0.0, 0.0])
+        object_parents.append(parent)
+
+        # # add debug site to see object center
+        # sites = [
+        #     dict(
+        #         name="TMP",
+        #         pos=array_to_string([0., 0., 0.]),
+        #         size="{}".format(0.1),
+        #         rgba=array_to_string([1., 0., 0., 1.]),
+        #     )
+        # ]
+
+        # Run super init
+        super().__init__(
+            name=name,
+            objects=objects,
+            object_locations=object_locations,
+            object_quats=object_quats,
+            object_parents=object_parents,
+            joints=joints,
+            body_joints=object_joints,
+            # sites=sites,
+        )

GR00T-WholeBodyControl/decoupled_wbc/docker/entrypoint/bash.sh

@@ -0,0 +1,5 @@
+#!/bin/bash
+set -e  # Exit on error
+
+echo "Dependencies installed successfully. Starting interactive bash shell..."
+exec /bin/bash

GR00T-WholeBodyControl/decoupled_wbc/docker/entrypoint/deploy.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+set -e # Exit on error
+
+# Run the deployment script
+# Check for script existence before running
+DEPLOY_SCRIPT="decoupled_wbc/scripts/deploy_g1.py"
+if [ -f "$DEPLOY_SCRIPT" ]; then
+    echo "Running deployment script at $DEPLOY_SCRIPT"
+    echo "Using python from $(which python)"
+    echo "Deploy args: $@"
+    exec python "$DEPLOY_SCRIPT" "$@"
+else
+    echo "ERROR: Deployment script not found at $DEPLOY_SCRIPT"
+    echo "Current directory structure:"
+    find . -type f -name "*.py" | grep -i deploy
+    echo "Available script options:"
+    find . -type f -name "*.py" | sort
+    echo "Starting a bash shell for troubleshooting..."
+    exec /bin/bash
+fi

GR00T-WholeBodyControl/decoupled_wbc/docker/entrypoint/install_deps.sh

@@ -0,0 +1,23 @@
+#!/bin/bash
+set -e
+
+# Source virtual environment and ROS2
+source ${HOME}/venv/bin/activate
+source /opt/ros/humble/setup.bash
+export ROS_LOCALHOST_ONLY=1
+
+# Install external dependencies
+echo "Current directory: $(pwd)"
+echo "Installing dependencies..."
+
+# Install Unitree SDK and LeRobot
+if [ -d "external_dependencies/unitree_sdk2_python" ]; then
+    cd external_dependencies/unitree_sdk2_python/
+    uv pip install -e . --no-deps
+    cd ../..
+fi
+
+# Install project packages
+if [ -f "decoupled_wbc/pyproject.toml" ]; then
+    UV_GIT_LFS=1 uv pip install -e "decoupled_wbc[full,dev]" -e "gear_sonic[sim]"
+fi

GR00T-WholeBodyControl/decoupled_wbc/tests/control/main/teleop/test_g1_control_loop.py

@@ -0,0 +1,469 @@
+import argparse
+import os
+from pathlib import Path
+import signal
+import subprocess
+import threading
+import time
+
+import numpy as np
+import pytest
+import rclpy
+from scipy.spatial.transform import Rotation as R
+from std_msgs.msg import String as RosStringMsg
+
+from decoupled_wbc.control.main.constants import (
+    CONTROL_GOAL_TOPIC,
+    KEYBOARD_INPUT_TOPIC,
+    STATE_TOPIC_NAME,
+)
+from decoupled_wbc.control.utils.ros_utils import ROSMsgPublisher, ROSMsgSubscriber
+from decoupled_wbc.control.utils.term_color_constants import GREEN_BOLD, RESET, YELLOW_BOLD
+from decoupled_wbc.data.viz.rerun_viz import RerunViz
+
+
+class KeyboardPublisher:
+    def __init__(self, topic_name: str = KEYBOARD_INPUT_TOPIC):
+        assert rclpy.ok(), "Expected ROS2 to be initialized in this process..."
+        executor = rclpy.get_global_executor()
+        self.node = executor.get_nodes()[0]
+        self.publisher = self.node.create_publisher(RosStringMsg, topic_name, 1)
+
+    def publish(self, key: str):
+        msg = RosStringMsg()
+        msg.data = key
+        self.publisher.publish(msg)
+
+
+def is_robot_fallen_from_quat(mujoco_quat):
+    # Convert MuJoCo [w, x, y, z] → SciPy [x, y, z, w]
+    w, x, y, z = mujoco_quat
+    scipy_quat = [x, y, z, w]
+
+    r = R.from_quat(scipy_quat)
+    roll, pitch, _ = r.as_euler("xyz", degrees=False)
+
+    MAX_ROLL_PITCH = np.radians(60)
+    print(f"[Fall Check] roll={roll:.3f} rad, pitch={pitch:.3f} rad")
+    return abs(roll) > MAX_ROLL_PITCH or abs(pitch) > MAX_ROLL_PITCH
+
+
+class LocomotionRunner:
+    def __init__(self, test_mode: str = "squat"):
+        self.test_mode = test_mode
+        if not rclpy.ok():
+            rclpy.init(args=None)
+        self.node = rclpy.create_node(f"EvalDriver_{test_mode}_{int(time.time())}")
+
+        # gracefully shutdown the spin thread when the test is done
+        self._stop_event = threading.Event()
+
+        self.spin_thread = threading.Thread(target=self._spin_loop, daemon=False)
+        self.spin_thread.start()
+
+        self.keyboard_event_publisher = KeyboardPublisher(KEYBOARD_INPUT_TOPIC)
+        self.control_publisher = ROSMsgPublisher(CONTROL_GOAL_TOPIC)
+        self.state_subscriber = ROSMsgSubscriber(STATE_TOPIC_NAME)
+        print(f"{test_mode} test initialized...")
+
+    def _spin_loop(self):
+        try:
+            while rclpy.ok() and not self._stop_event.is_set():
+                rclpy.spin_once(self.node)
+        except rclpy.executors.ExternalShutdownException:
+            print("[INFO] Spin thread exiting due to shutdown.")
+        finally:
+            print("spin loop stopped...")
+
+    def warm_up(self):
+        """Stabilize and release the robot."""
+        print("waiting for 2 seconds...")
+        time.sleep(2)
+        print(f"running {self.test_mode} test...")
+        self.activate()
+        print("activated...")
+        time.sleep(1)
+        self.release()
+        print("released...")
+        time.sleep(5)
+
+    def _run_walk_test(self):
+        self.walk_forward()  # speed up to 0.2 m/s
+        time.sleep(1)
+        self.walk_forward()  # speed up to 0.4 m/s
+
+        rate = self.node.create_rate(0.5)
+        start_time = time.time()
+        while rclpy.ok() and (time.time() - start_time) < 10.0:
+            obs = self.state_subscriber.get_msg()
+
+            if is_robot_fallen_from_quat(obs["torso_quat"]):
+                print("robot fallen...")
+                return 0
+            elif self._check_success_condition(obs):
+                print(f"robot reaching target ({self.test_mode})...")
+                return 1, {}
+            else:
+                rate.sleep()
+
+        print("test timed out after 10 seconds...")
+        return 0, {}
+
+    def _run_squat_test(self):
+        rate = self.node.create_rate(0.5)
+        start_time = time.time()
+        while rclpy.ok() and (time.time() - start_time) < 10.0:
+            obs = self.state_subscriber.get_msg()
+
+            if is_robot_fallen_from_quat(obs["torso_quat"]):
+                print("robot fallen...")
+                return 0, {}
+            elif self._check_success_condition(obs):
+                print(f"robot reaching target ({self.test_mode})...")
+                return 1, {}
+            else:
+                self.go_down()
+                rate.sleep()
+
+        print("test timed out after 10 seconds...")
+        return 0, {}
+
+    def cmd_to_velocity(self, cmd_list):
+        cmd_to_velocity = {
+            "w": np.array([0.2, 0.0, 0.0]),
+            "s": np.array([-0.2, 0.0, 0.0]),
+            "q": np.array([0.0, 0.2, 0.0]),
+            "e": np.array([0.0, -0.2, 0.0]),
+            "z": np.array([0.0, 0.0, 0.0]),
+        }
+
+        accumulated_velocity = np.array([0.0, 0.0, 0.0])
+        velocity_list = []
+        for cmd in cmd_list:
+            if cmd == "z":
+                accumulated_velocity = [0.0, 0.0, 0.0]
+            elif cmd in ["CHECK", "SKIP"]:
+                accumulated_velocity = velocity_list[-1]
+            else:
+                accumulated_velocity += cmd_to_velocity[cmd]
+            velocity_list.append(accumulated_velocity.copy())
+
+        return velocity_list
+
+    def _run_stop_test(self):
+        base_vel_thres = 0.25
+
+        cmd_list = (
+            ["w", "w", "w", "w", "s", "s", "s", "z", "SKIP", "CHECK"]
+            + ["s", "s", "q", "w", "w", "w", "e", "s", "s", "z", "SKIP", "CHECK"]
+            + ["q", "q", "w", "q", "e", "s", "s", "e", "w", "z", "SKIP", "CHECK"]
+            + ["w", "w", "w", "w", "w", "s", "s", "s", "s", "z", "SKIP", "CHECK"]
+        )
+
+        success_flag = 1
+
+        statistics = {
+            "floating_base_pose": {"state": []},
+            "floating_base_vel": {"state": [], "cmd": []},
+            "timestamp": [],
+        }
+        for cmd in cmd_list:
+            self.keyboard_event_publisher.publish(cmd)
+            time.sleep(0.5)
+            obs = self.state_subscriber.get_msg()
+            statistics["floating_base_pose"]["state"].append(
+                np.linalg.norm(obs["floating_base_pose"])
+            )
+            statistics["floating_base_vel"]["state"].append(
+                np.linalg.norm(obs["floating_base_vel"])
+            )
+            statistics["timestamp"].append(time.time())
+
+            if cmd == "CHECK" and np.linalg.norm(obs["floating_base_vel"]) > base_vel_thres:
+                print(
+                    f" [{YELLOW_BOLD}WARNING{RESET}] robot is not stopped fully. "
+                    f"Current base velocity: {np.linalg.norm(obs['floating_base_vel']):.3f} > {base_vel_thres:.3f}"
+                )
+                # success_flag = 0  # robot is not stopped
+
+            time.sleep(0.5)
+
+        vel_cmd = self.cmd_to_velocity(cmd_list)
+        vel_cmd = [np.linalg.norm(v) for v in vel_cmd]
+        statistics["floating_base_vel"]["cmd"] = vel_cmd
+        return success_flag, statistics
+
+    def _run_eef_track_test(self):
+        from decoupled_wbc.control.policy.lerobot_replay_policy import LerobotReplayPolicy
+
+        parquet_path = (
+            Path(__file__).parent.parent.parent.parent / "replay_data" / "g1_pnpbottle.parquet"
+        )
+        replay_policy = LerobotReplayPolicy(parquet_path=str(parquet_path))
+
+        freq = 50
+        rate = self.node.create_rate(freq)
+
+        statistics = {
+            # "floating_base_pose": {"state": [], "cmd": []},
+            "eef_base_pose": {"state": [], "cmd": []},
+            "timestamp": [],
+        }
+
+        for ii in range(500):
+            action = replay_policy.get_action()
+            action = replay_policy.action_to_cmd(action)
+            action["timestamp"] = time.monotonic()
+            action["target_time"] = time.monotonic() + ii / freq
+            self.control_publisher.publish(action)
+            obs = self.state_subscriber.get_msg()
+            if obs is None:
+                print("no obs...")
+                continue
+            gt_obs = replay_policy.get_observation()
+
+            # statistics["floating_base_pose"]["state"].append(obs["floating_base_pose"])
+            # statistics["floating_base_pose"]["cmd"].append(np.zeros_like(obs["floating_base_pose"]))
+            statistics["eef_base_pose"]["state"].append(obs["wrist_pose"])
+            statistics["eef_base_pose"]["cmd"].append(gt_obs["wrist_pose"])
+            statistics["timestamp"].append(time.time())
+
+            pos_err = np.linalg.norm(obs["wrist_pose"][:3] - gt_obs["wrist_pose"][:3])
+            if pos_err > 1e-1:
+                print(
+                    f" [{YELLOW_BOLD}WARNING{RESET}] robot failed to track the eef, "
+                    f"error: {pos_err:.3f} ({self.test_mode})..."
+                )
+                return 0, statistics
+
+            if is_robot_fallen_from_quat(obs["torso_quat"]):
+                print("robot fallen...")
+                return 0, statistics
+            else:
+                rate.sleep()
+
+        return 1, statistics
+
+    def run(self):
+        self.warm_up()
+
+        test_mode_to_func = {
+            "squat": self._run_squat_test,
+            "walk": self._run_walk_test,
+            "stop": self._run_stop_test,
+            "eef_track": self._run_eef_track_test,
+        }
+
+        result, statistics = test_mode_to_func[self.test_mode]()
+
+        self.post_process(statistics)
+        return result
+
+    def _check_success_condition(self, obs):
+        if self.test_mode == "squat":
+            return obs["floating_base_pose"][2] < 0.4
+        elif self.test_mode == "walk":
+            return np.linalg.norm(obs["floating_base_pose"][0:2]) > 1.0
+        return False
+
+    def activate(self):
+        self.keyboard_event_publisher.publish("]")
+
+    def release(self):
+        self.keyboard_event_publisher.publish("9")
+
+    def go_down(self):
+        self.keyboard_event_publisher.publish("2")
+
+    def walk_forward(self):
+        self.keyboard_event_publisher.publish("w")
+
+    def walk_stop(self):
+        self.keyboard_event_publisher.publish("z")
+
+    def post_process(self, statistics):
+        if len(statistics) == 0:
+            return
+
+        # plot the statistics
+        plot_keys = [key for key in statistics.keys() if key != "timestamp"]
+        viz = RerunViz(
+            image_keys=[],
+            tensor_keys=plot_keys,
+            window_size=10.0,
+            app_name=f"{self.test_mode}_test",
+        )
+
+        for ii in range(len(statistics[plot_keys[0]]["state"])):
+            tensor_data = {}
+            for k in plot_keys:
+                if "state" in statistics[k] and "cmd" in statistics[k]:
+                    tensor_data[k] = np.array(
+                        (statistics[k]["state"][ii], statistics[k]["cmd"][ii])
+                    ).reshape(2, -1)
+                else:
+                    tensor_data[k] = np.asarray(statistics[k]["state"][ii]).reshape(1, -1)
+            viz.plot_tensors(
+                tensor_data,
+                statistics["timestamp"][ii],
+            )
+
+        if self.test_mode == "stop":
+            base_velocity = statistics["floating_base_vel"]["state"]
+            base_velocity_cmd = statistics["floating_base_vel"]["cmd"]
+
+            base_velocity_tracking_err = []
+            for v_cmd, v in zip(base_velocity_cmd, base_velocity):  # TODO: check if this is correct
+                if v_cmd.max() < 1e-4:
+                    base_velocity_tracking_err.append(v)
+            print(
+                f" [{GREEN_BOLD}INFO{RESET}] Base velocity tracking when stopped: "
+                f"{np.mean(base_velocity_tracking_err):.3f}"
+            )
+
+        if self.test_mode == "eef_track":
+            eef_pose = statistics["eef_base_pose"]["state"]
+            eef_pose_cmd = statistics["eef_base_pose"]["cmd"]
+            eef_pose_tracking_err = []
+            for p_cmd, p in zip(eef_pose_cmd, eef_pose):
+                eef_pose_tracking_err.append(np.linalg.norm(p - p_cmd))
+            print(
+                f" [{GREEN_BOLD}INFO{RESET}] Eef pose tracking error: {np.mean(eef_pose_tracking_err):.3f}"
+            )
+
+    def shutdown(self):
+        self._stop_event.set()
+        self.spin_thread.join()
+        del self.state_subscriber
+        del self.keyboard_event_publisher
+        # Don't shutdown ROS between tests - let pytest handle it
+
+
+def start_g1_control_loop():
+    proc = subprocess.Popen(
+        [
+            "python3",
+            "decoupled_wbc/control/main/teleop/run_g1_control_loop.py",
+            "--keyboard_dispatcher_type",
+            "ros",
+            "--enable-offscreen",
+        ],
+        preexec_fn=os.setsid,
+    )
+    time.sleep(10)
+    return proc
+
+
+def run_test(test_mode: str):
+    """Run a single test with the specified mode."""
+    proc = start_g1_control_loop()
+    print(f"G1 control loop started for {test_mode} test...")
+
+    test = LocomotionRunner(test_mode)
+    result = test.run()
+
+    print("Shutting down...")
+    test.shutdown()
+    proc.send_signal(signal.SIGKILL)
+    proc.wait()
+
+    return result
+
+
+def test_squat():
+    """Pytest function for squat test."""
+    result = run_test("squat")
+    assert result == 1, "Squat test failed - robot either fell or didn't reach target height"
+
+
+def test_walk():
+    """Pytest function for walk test."""
+    result = run_test("walk")
+    assert result == 1, "Walk test failed - robot either fell or didn't reach target distance"
+
+
+@pytest.mark.skip(reason="skipping test for now, cicd test always gets killed")
+def test_stop():
+    """Pytest function for walking to a nearby position and stop test."""
+    result = run_test("stop")
+    assert result == 1, "Stop test failed - robot either fell or didn't reach target distance"
+
+
+@pytest.mark.skip(reason="skipping test for now, cicd test always gets killed")
+def test_eef_track():
+    """Pytest function for eef track test."""
+    result = run_test("eef_track")
+    assert result == 1, "Eef track test failed - robot either fell or didn't reach target distance"
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Run locomotion tests")
+    parser.add_argument("--squat", action="store_true", help="Run squat test only")
+    parser.add_argument("--walk", action="store_true", help="Run walk test only")
+    parser.add_argument("--stop", action="store_true", help="Run stop test only")
+    parser.add_argument("--eef_track", action="store_true", help="Run eef track test only")
+
+    args = parser.parse_args()
+
+    if args.squat and args.walk:
+        print("Error: Cannot specify both --squat and --walk")
+        return 1
+
+    if args.squat:
+        print("Running squat test only...")
+        result = run_test("squat")
+        if result == 1:
+            print("✓ Squat test PASSED")
+            return 0
+        else:
+            print("✗ Squat test FAILED")
+            return 1
+
+    elif args.walk:
+        print("Running walk test only...")
+        result = run_test("walk")
+        if result == 1:
+            print("✓ Walk test PASSED")
+            return 0
+        else:
+            print("✗ Walk test FAILED")
+            return 1
+
+    elif args.stop:
+        print("Running stop test only...")
+        result = run_test("stop")
+        if result == 1:
+            print("✓ Stop test PASSED")
+            return 0
+        else:
+            print("✗ Stop test FAILED")
+            return 1
+
+    elif args.eef_track:
+        print("Running eef track test only...")
+        result = run_test("eef_track")
+        if result == 1:
+            print("✓ Eef track test PASSED")
+            return 0
+        else:
+            print("✗ Eef track test FAILED")
+            return 1
+
+    else:
+        print("Running both tests...")
+        squat_result = run_test("squat")
+        walk_result = run_test("walk")
+
+        if squat_result == 1 and walk_result == 1:
+            print("✓ All tests PASSED")
+            return 0
+        else:
+            print(
+                f"✗ Test results: squat={'PASSED' if squat_result == 1 else 'FAILED'}, "
+                f"walk={'PASSED' if walk_result == 1 else 'FAILED'}"
+            )
+            return 1
+
+
+if __name__ == "__main__":
+    exit(main())

GR00T-WholeBodyControl/decoupled_wbc/tests/control/main/test_data_exporter_loop.py

@@ -0,0 +1,403 @@
+import glob
+import os
+import tempfile
+from unittest.mock import MagicMock, patch
+
+import numpy as np
+import pytest
+
+try:
+    from decoupled_wbc.control.main.teleop.run_g1_data_exporter import Gr00tDataCollector
+    from decoupled_wbc.control.robot_model.instantiation.g1 import instantiate_g1_robot_model
+    from decoupled_wbc.data.constants import RS_VIEW_CAMERA_HEIGHT, RS_VIEW_CAMERA_WIDTH
+    from decoupled_wbc.data.exporter import Gr00tDataExporter
+    from decoupled_wbc.data.utils import get_dataset_features
+except ModuleNotFoundError as e:
+    if "No module named 'rclpy'" in str(e):
+        pytestmark = pytest.mark.skip(reason="ROS (rclpy) is not installed")
+    else:
+        raise e
+
+
+import json
+
+# How does mocking ROS work?
+#
+# This test file uses mocking to simulate a ROS environment without requiring actual ROS hardware:
+#
+# 1. ros_ok_side_effect: Controls how long the ROS loop runs by returning a sequence of
+#    True/False values. [True, True, False] means "run for 2 iterations then stop"
+#
+# 2. MockROSMsgSubscriber: Simulates sensors (camera/state) by returning pre-defined data:
+#
+# 3. MockKeyboardListenerSubscriber: Simulates user input:
+#    - 'c' = start/stop recording
+#    - 'd' = discard episode
+#    - KeyboardInterrupt = simulate Ctrl+C
+#    - None = no input
+#
+# 4. MockROSEnvironment: A context manager that patches all ROS dependencies to use our mocks,
+#    allowing us to test ROS-dependent code without actual ROS running.
+
+
+class MockROSMsgSubscriber:
+    def __init__(self, return_value: list[dict]):
+        self.return_value = return_value
+        self.counter = 0
+
+    def get_image(self):
+        if self.counter < len(self.return_value):
+            self.counter += 1
+            return self.return_value[self.counter - 1]
+        else:
+            return None
+
+    def get_msg(self):
+        if self.counter < len(self.return_value):
+            self.counter += 1
+            return self.return_value[self.counter - 1]
+        else:
+            return None
+
+
+class MockKeyboardListenerSubscriber:
+    def __init__(self, return_value: list[str]):
+        self.return_value = return_value
+        self.counter = 0
+
+    def get_keyboard_input(self):
+        return self.return_value[self.counter]
+
+    def read_msg(self):
+        if self.counter < len(self.return_value):
+            result = self.return_value[self.counter]
+            if isinstance(result, KeyboardInterrupt):
+                raise result
+            self.counter += 1
+            return result
+        return None
+
+
+class MockROSEnvironment:
+    """Context manager for mocking ROS environment and subscribers."""
+
+    def __init__(self, ok_side_effect, keyboard_listener, img_subscriber, state_subscriber):
+        self.ok_side_effect = ok_side_effect
+        self.keyboard_listener = keyboard_listener
+        self.img_subscriber = img_subscriber
+        self.state_subscriber = state_subscriber
+        self.patches = []
+
+    def __enter__(self):
+        self.patches = [
+            patch("rclpy.init"),
+            patch("rclpy.create_node"),
+            patch("rclpy.spin"),
+            patch("rclpy.ok", side_effect=self.ok_side_effect),
+            patch("rclpy.shutdown"),
+            patch(
+                "decoupled_wbc.control.main.teleop.run_g1_data_exporter.KeyboardListenerSubscriber",
+                return_value=self.keyboard_listener,
+            ),
+            patch(
+                "decoupled_wbc.control.main.teleop.run_g1_data_exporter.ROSImgMsgSubscriber",
+                return_value=self.img_subscriber,
+            ),
+            patch(
+                "decoupled_wbc.control.main.teleop.run_g1_data_exporter.ROSMsgSubscriber",
+                return_value=self.state_subscriber,
+            ),
+        ]
+
+        for p in self.patches:
+            p.start()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        for p in reversed(self.patches):
+            p.stop()
+        return False
+
+
+def verify_parquet_files_exist(file_path: str, num_episodes: int):
+    parquet_files = glob.glob(os.path.join(file_path, "data/chunk-*/episode_*.parquet"))
+    assert (
+        len(parquet_files) == num_episodes
+    ), f"Expected {num_episodes} parquet files, but found {len(parquet_files)}"
+
+
+def verify_video_files_exist(file_path: str, observation_keys: list[str], num_episodes: int):
+    for observation_key in observation_keys:
+        video_files = glob.glob(
+            os.path.join(file_path, f"videos/chunk-*/{observation_key}/episode_*.mp4")
+        )
+        assert (
+            len(video_files) == num_episodes
+        ), f"Expected {num_episodes} video files, but found {len(video_files)}"
+
+
+def verify_metadata_files(file_path: str):
+    files_to_check = ["episodes.jsonl", "info.json", "tasks.jsonl", "modality.json"]
+    for file in files_to_check:
+        assert os.path.exists(os.path.join(file_path, "meta", file)), f"meta/{file} not created"
+
+
+@pytest.fixture
+def lerobot_features():
+    robot_model = instantiate_g1_robot_model()
+    return get_dataset_features(robot_model)
+
+
+@pytest.fixture
+def modality_config():
+    return {
+        "state": {"feature1": {"start": 0, "end": 4}, "feature2": {"start": 4, "end": 9}},
+        "action": {"feature1": {"start": 0, "end": 4}, "feature2": {"start": 4, "end": 9}},
+        "video": {"rs_view": {"original_key": "observation.images.ego_view"}},
+        "annotation": {"human.task_description": {"original_key": "task_index"}},
+    }
+
+
+def _get_image_stream_data(episode_length: int, frame_rate: int, img_height: int, img_width: int):
+    return [
+        {
+            "image": np.zeros((img_height, img_width, 3), dtype=np.uint8),
+            "timestamp": (i * 1 / frame_rate),
+        }
+        for i in range(episode_length)
+    ]
+
+
+def _get_state_act_stream_data(
+    episode_length: int, frame_rate: int, state_dim: int, action_dim: int
+):
+    return [
+        {
+            "q": np.zeros(state_dim),
+            "action": np.zeros(action_dim),
+            "timestamp": (i * 1 / frame_rate),
+            "navigate_command": np.zeros(3, dtype=np.float64),
+            "base_height_command": 0.0,
+            "wrist_pose": np.zeros(14, dtype=np.float64),
+            "action.eef": np.zeros(14, dtype=np.float64),
+        }
+        for i in range(episode_length)
+    ]
+
+
+def test_control_loop_happy_path_workflow(lerobot_features, modality_config):
+    """
+    This test records a single episode and saves it to disk.
+    """
+    episode_length = 10
+    frame_rate = 20
+    img_stream_data = _get_image_stream_data(
+        episode_length, frame_rate, RS_VIEW_CAMERA_HEIGHT, RS_VIEW_CAMERA_WIDTH
+    )
+    robot_model = instantiate_g1_robot_model()
+    state_act_stream_data = _get_state_act_stream_data(
+        episode_length, frame_rate, robot_model.num_joints, robot_model.num_joints
+    )
+
+    keyboard_sub_output = [None for _ in range(episode_length)]
+    keyboard_sub_output[0] = "c"  # Start recording
+    keyboard_sub_output[-1] = "c"  # Stop recording and save
+
+    # --------- Save the first episode ---------
+    mock_img_sub = MockROSMsgSubscriber(img_stream_data)
+    mock_state_sub = MockROSMsgSubscriber(state_act_stream_data)
+    mock_keyboard_listner = MockKeyboardListenerSubscriber(keyboard_sub_output)
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        dataset_dir = os.path.join(temp_dir, "dataset")
+
+        data_exporter = Gr00tDataExporter.create(
+            save_root=dataset_dir,
+            fps=frame_rate,
+            features=lerobot_features,
+            modality_config=modality_config,
+            task="test",
+        )
+
+        ros_ok_side_effect = [True] * (episode_length + 1) + [False]
+        with MockROSEnvironment(
+            ros_ok_side_effect, mock_keyboard_listner, mock_img_sub, mock_state_sub
+        ):
+            data_collector = Gr00tDataCollector(
+                camera_topic_name="mock_camera_topic",
+                state_topic_name="mock_state_topic",
+                data_exporter=data_exporter,
+                frequency=frame_rate,
+            )
+
+            # mocking to avoid actual sleeping
+            data_collector.rate = MagicMock()
+
+            data_collector.run()
+
+        verify_parquet_files_exist(dataset_dir, 1)
+        verify_video_files_exist(dataset_dir, data_exporter.meta.video_keys, 1)
+        verify_metadata_files(dataset_dir)
+
+        # --------- Save the second episode ---------
+        # we reset the mock subscribers and re-run the control loop
+        # This immitates the case where the user starts recording a new episode on an existing dataset
+        mock_img_sub = MockROSMsgSubscriber(img_stream_data)
+        mock_state_sub = MockROSMsgSubscriber(state_act_stream_data)
+        ros_ok_side_effect = [True] * (episode_length + 1) + [False]
+        mock_keyboard_listner = MockKeyboardListenerSubscriber(keyboard_sub_output)
+        with MockROSEnvironment(
+            ros_ok_side_effect, mock_keyboard_listner, mock_img_sub, mock_state_sub
+        ):
+            data_collector = Gr00tDataCollector(
+                camera_topic_name="mock_camera_topic",
+                state_topic_name="mock_state_topic",
+                data_exporter=data_exporter,
+                frequency=frame_rate,
+            )
+
+            # mocking to avoid actual sleeping
+            data_collector.rate = MagicMock()
+
+            data_collector.run()
+
+        # now there should be 2 episodes in the dataset
+        verify_parquet_files_exist(dataset_dir, 2)
+        verify_video_files_exist(dataset_dir, data_exporter.meta.video_keys, 2)
+        verify_metadata_files(dataset_dir)
+
+
+def test_control_loop_keyboard_interrupt_workflow(lerobot_features, modality_config):
+    """
+    This test simulates a keyboard interruption in the middle of recording.
+    Expected behavior:
+    - The episode is saved to disk
+    - The episode is marked as discarded
+    """
+    episode_length = 15
+    frame_rate = 20
+    img_stream_data = _get_image_stream_data(
+        episode_length, frame_rate, RS_VIEW_CAMERA_HEIGHT, RS_VIEW_CAMERA_WIDTH
+    )
+    robot_model = instantiate_g1_robot_model()
+    state_act_stream_data = _get_state_act_stream_data(
+        episode_length, frame_rate, robot_model.num_joints, robot_model.num_joints
+    )
+
+    keyboard_sub_output = [None for _ in range(episode_length)]
+    keyboard_sub_output[0] = "c"  # Start recording
+    keyboard_sub_output[5] = KeyboardInterrupt()  # keyboard interruption in the middle of recording
+
+    mock_img_sub = MockROSMsgSubscriber(img_stream_data)
+    mock_state_sub = MockROSMsgSubscriber(state_act_stream_data)
+    mock_keyboard_listener = MockKeyboardListenerSubscriber(keyboard_sub_output)
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        dataset_dir = os.path.join(temp_dir, "dataset")
+
+        data_exporter = Gr00tDataExporter.create(
+            save_root=dataset_dir,
+            fps=frame_rate,
+            features=lerobot_features,
+            modality_config=modality_config,
+            task="test",
+        )
+
+        ros_ok_side_effect = [True] * episode_length + [False]
+        with MockROSEnvironment(
+            ros_ok_side_effect, mock_keyboard_listener, mock_img_sub, mock_state_sub
+        ):
+            data_collector = Gr00tDataCollector(
+                camera_topic_name="mock_camera_topic",
+                state_topic_name="mock_state_topic",
+                data_exporter=data_exporter,
+                frequency=frame_rate,
+            )
+
+            data_collector.rate = MagicMock()
+            # try:
+            data_collector.run()
+            # except KeyboardInterrupt:
+            #     pass
+
+        verify_parquet_files_exist(dataset_dir, 1)
+        verify_video_files_exist(dataset_dir, data_exporter.meta.video_keys, 1)
+        verify_metadata_files(dataset_dir)
+
+        # verify that the episode is marked as discarded
+        ep_info = json.load(open(os.path.join(dataset_dir, "meta", "info.json")))
+        assert ep_info["discarded_episode_indices"][0] == 0
+        assert ep_info["total_frames"] == 5
+        assert ep_info["total_episodes"] == 1
+
+
+def test_discarded_episode_workflow(lerobot_features, modality_config):
+    """
+    This test simulates a case where the user discards an episode in the middle of recording.
+    Expected behavior:
+    - Record 3 episodes, discard episode 0 and 2
+    - There should be 3 episodes saved to disk
+    - Episode 0 and 2 should be flagged as discarded
+    """
+    episode_length = 17
+    frame_rate = 20
+    robot_model = instantiate_g1_robot_model()
+    state_dim = robot_model.num_joints
+    action_dim = robot_model.num_joints
+    img_stream_data = _get_image_stream_data(
+        episode_length, frame_rate, RS_VIEW_CAMERA_HEIGHT, RS_VIEW_CAMERA_WIDTH
+    )
+    state_act_stream_data = _get_state_act_stream_data(
+        episode_length, frame_rate, state_dim, action_dim
+    )
+
+    keyboard_sub_output = [None for _ in range(episode_length)]
+    keyboard_sub_output[0] = "c"  # Start recording episode index 0
+    keyboard_sub_output[5] = "x"  # Discard episode index 0
+    keyboard_sub_output[7] = "c"  # Start recording episode index 1
+    keyboard_sub_output[10] = "c"  # stop recording and save episode index 1
+    keyboard_sub_output[12] = "c"  # start recording episode index 2
+    keyboard_sub_output[15] = "x"  # discard episode index 2
+
+    mock_img_sub = MockROSMsgSubscriber(img_stream_data)
+    mock_state_sub = MockROSMsgSubscriber(state_act_stream_data)
+    mock_keyboard_listener = MockKeyboardListenerSubscriber(keyboard_sub_output)
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        dataset_dir = os.path.join(temp_dir, "dataset")
+
+        data_exporter = Gr00tDataExporter.create(
+            save_root=dataset_dir,
+            fps=frame_rate,
+            features=lerobot_features,
+            modality_config=modality_config,
+            task="test",
+        )
+
+        ros_ok_side_effect = [True] * episode_length + [False]
+        with MockROSEnvironment(
+            ros_ok_side_effect, mock_keyboard_listener, mock_img_sub, mock_state_sub
+        ):
+            data_collector = Gr00tDataCollector(
+                camera_topic_name="mock_camera_topic",
+                state_topic_name="mock_state_topic",
+                data_exporter=data_exporter,
+                frequency=frame_rate,
+            )
+
+            data_collector.rate = MagicMock()
+            try:
+                data_collector.run()
+            except Exception:
+                pass
+
+            # vrify if the episode is marked as discarded
+            ep_info = json.load(open(os.path.join(dataset_dir, "meta", "info.json")))
+            assert len(ep_info["discarded_episode_indices"]) == 2
+            assert ep_info["discarded_episode_indices"][0] == 0
+            assert ep_info["discarded_episode_indices"][1] == 2
+
+            # verify that all episodes are saved regardless of being discarded
+            verify_parquet_files_exist(dataset_dir, 3)
+            verify_video_files_exist(dataset_dir, data_exporter.meta.video_keys, 3)
+            verify_metadata_files(dataset_dir)

GR00T-WholeBodyControl/decoupled_wbc/tests/control/policy/interpolation_policy/test_interpolation_policy.py

@@ -0,0 +1,47 @@
+from pathlib import Path
+import pickle
+
+import numpy as np
+import pytest
+
+from decoupled_wbc.control.policy.interpolation_policy import (
+    InterpolationPolicy,
+)
+
+
+def get_test_data_path(filename: str) -> str:
+    """Get the absolute path to a test data file."""
+    test_dir = Path(__file__).parent
+    return str(test_dir / ".." / ".." / ".." / "replay_data" / filename)
+
+
+@pytest.fixture
+def logged_data():
+    """Load the logged data from file."""
+    data_path = get_test_data_path("interpolation_data.pkl")
+    with open(data_path, "rb") as f:
+        return pickle.load(f)
+
+
+def test_replay_logged_data(logged_data):
+    """Test that the wrapper produces the same pose commands as logged data."""
+    init_args = logged_data["init_args"]
+    interp = InterpolationPolicy(
+        init_time=init_args["curr_t"],
+        init_values={"target_pose": init_args["curr_pose"]},
+        max_change_rate=np.inf,
+    )
+
+    # Test all data points including the first one
+    for c in logged_data["calls"]:
+        # Get the action from wrapper
+        if c["type"] == "get_action":
+            action = interp.get_action(**c["args"])
+            expected_action = c["result"]
+            np.testing.assert_allclose(
+                action["target_pose"], expected_action["q"], rtol=1e-9, atol=1e-9
+            )
+            # print(action, expected_action)
+
+        else:
+            interp.set_goal(**c["args"])

GR00T-WholeBodyControl/decoupled_wbc/tests/control/policy/interpolation_policy/test_interpolation_ramp_up.py

@@ -0,0 +1,78 @@
+import numpy as np
+import pytest
+
+from decoupled_wbc.control.policy.interpolation_policy import (
+    InterpolationPolicy,
+)
+
+
+def test_trajectory_interpolation():
+    """
+    Test that the InterpolationPolicy correctly interpolates between waypoints.
+
+    Initial pose is at all zeros.
+    At t=4sec, the index 27 position (right_shoulder_yaw_joint) should be -1.5.
+    We run at 100Hz to see all intermediate waypoints.
+
+    Notes:
+        - The trajectory data is at 'trajectory_data.npy' in the current directory
+        - The visualization is at 'trajectory.png' in the current directory
+    """
+    # Create a pose with 32 joints (all zeros initially)
+    num_joints = 32
+    initial_pose = np.zeros(num_joints)
+
+    # Initial time (use a fixed value for reproducibility)
+    initial_time = 0.0
+
+    # Create the wrapper with initial pose
+    interpolator = InterpolationPolicy(
+        init_time=initial_time,
+        init_values={"target_pose": initial_pose},
+        max_change_rate=np.inf,
+    )
+
+    # Target pose: all zeros except index 27 which should be -1.5
+    target_pose = np.zeros(num_joints)
+    target_pose[27] = -1.5  # right_shoulder_yaw_joint
+    target_time = 4.0  # 4 seconds from now
+
+    # Set the planner command to schedule the waypoint
+    interpolator.set_goal(
+        {
+            "target_pose": target_pose,
+            "target_time": target_time,
+            "interpolation_garbage_collection_time": initial_time,
+        }
+    )
+
+    # Sample the trajectory at 100Hz
+    frequency = 100
+    dt = 1.0 / frequency
+    sample_times = np.arange(initial_time, target_time + dt, dt)
+
+    # Collect the interpolated poses
+    sampled_poses = []
+    for t in sample_times:
+        action = interpolator.get_action(t)
+        sampled_poses.append(action["target_pose"])
+
+    # Convert to numpy array for easier analysis
+    sampled_poses = np.array(sampled_poses)
+
+    # Check specific requirements
+    # Verify we actually moved from 0 to -1.5
+    joint_27_positions = sampled_poses[:, 27]
+    assert joint_27_positions[0] == pytest.approx(0.0)
+    assert joint_27_positions[-1] == pytest.approx(-1.5)
+
+    # Calculate the absolute changes between each step
+    changes = np.abs(np.diff(joint_27_positions))
+    assert np.all(changes < 0.004), "Joint 27 position should change by less than 0.004"
+
+    # Print some statistics about the trajectory
+    print(f"Total time steps: {len(sample_times)}")
+    print(f"Joint 27 trajectory start: {joint_27_positions[0]}")
+    print(f"Joint 27 trajectory end: {joint_27_positions[-1]}")
+    print(f"Joint 27 max velocity: {np.max(np.abs(np.diff(joint_27_positions) / dt))}")
+    print(f"Max velocity timestep: {np.argmax(np.abs(np.diff(joint_27_positions) / dt))}")

GR00T-WholeBodyControl/decoupled_wbc/tests/control/robot_model/robot_model_test.py

@@ -0,0 +1,911 @@
+# test_robot_model.py
+
+import numpy as np
+import pinocchio as pin
+import pytest
+
+from decoupled_wbc.control.robot_model import ReducedRobotModel
+from decoupled_wbc.control.robot_model.instantiation.g1 import instantiate_g1_robot_model
+
+
+@pytest.fixture
+def g1_robot_model():
+    """
+    Fixture that creates and returns a G1 RobotModel instance.
+    """
+    return instantiate_g1_robot_model()
+
+
+def test_robot_model_initialization(g1_robot_model):
+    """
+    Test initialization of the RobotModel and its main attributes.
+    """
+    for robot_model in [g1_robot_model]:
+        # Check that the Pinocchio wrapper exists
+        assert robot_model.pinocchio_wrapper is not None
+
+        # Check number of degrees of freedom (nq)
+        assert robot_model.num_dofs > 0
+
+        # Check we have the expected number of joints beyond the floating base
+        assert len(robot_model.joint_names) > 0
+
+        # Check that supplemental info is present
+        assert robot_model.supplemental_info is not None
+
+
+def test_robot_model_joint_names(g1_robot_model):
+    """
+    Test that joint_names is populated correctly
+    and that dof_index works.
+    """
+    for robot_model in [g1_robot_model]:
+        # Extract joint names
+        joint_names = robot_model.joint_names
+
+        # Pick the first joint name and get its index
+        first_joint_name = joint_names[0]
+        idx = robot_model.dof_index(first_joint_name)
+        assert idx >= 0
+
+        # Test that an unknown joint name raises an error
+        with pytest.raises(ValueError, match="Unknown joint name"):
+            _ = robot_model.dof_index("non_existent_joint")
+
+
+def test_robot_model_forward_kinematics_valid_q(g1_robot_model):
+    """
+    Test that cache_forward_kinematics works with a valid q.
+    """
+    for robot_model in [g1_robot_model]:
+        nq = robot_model.num_dofs
+
+        # Construct a valid configuration (e.g., zero vector)
+        q_valid = np.zeros(nq)
+
+        # Should not raise any exception
+        robot_model.cache_forward_kinematics(q_valid)
+
+
+def test_robot_model_forward_kinematics_invalid_q(g1_robot_model):
+    """
+    Test that cache_forward_kinematics raises an error with an invalid q.
+    """
+    for robot_model in [g1_robot_model]:
+        nq = robot_model.num_dofs
+
+        # Construct an invalid configuration (wrong size)
+        q_invalid = np.zeros(nq + 1)
+
+        with pytest.raises(ValueError, match="Expected q of length"):
+            robot_model.cache_forward_kinematics(q_invalid)
+
+
+def test_robot_model_frame_placement(g1_robot_model):
+    """
+    Test the frame_placement method with a valid and invalid frame name.
+    Also test that frame placements change with different configurations.
+    """
+    for robot_model in [g1_robot_model]:
+        # Skip if no supplemental info
+        if robot_model.supplemental_info is None:
+            pytest.skip("No supplemental info available for testing")
+
+        # Use the hand frame from supplemental info
+        test_frame = robot_model.supplemental_info.hand_frame_names["left"]
+
+        # Test with zero configuration
+        q_zero = np.zeros(robot_model.num_dofs)
+        robot_model.cache_forward_kinematics(q_zero)
+        placement_zero = robot_model.frame_placement(test_frame)
+        assert isinstance(placement_zero, pin.SE3)
+
+        # Test with non-zero configuration
+        q_non_zero = np.zeros(robot_model.num_dofs)
+        root_nq = 7 if robot_model.is_floating_base_model else 0
+
+        # Set a more significant configuration change
+        # Use π/2 for all joints to create a more noticeable difference
+        q_non_zero[root_nq:] = np.pi / 2  # 90 degrees for all joints
+
+        robot_model.cache_forward_kinematics(q_non_zero)
+        placement_non_zero = robot_model.frame_placement(test_frame)
+
+        # Verify that frame placements are different with different configurations
+        assert not np.allclose(
+            placement_zero.translation, placement_non_zero.translation
+        ) or not np.allclose(placement_zero.rotation, placement_non_zero.rotation)
+
+        # Should raise an error for an invalid frame
+        with pytest.raises(ValueError, match="Unknown frame"):
+            robot_model.frame_placement("non_existent_frame")
+
+
+# Tests for ReducedRobotModel
+def test_reduced_robot_model_initialization(g1_robot_model):
+    """
+    Test initialization of the ReducedRobotModel.
+    """
+    for robot_model in [g1_robot_model]:
+        # Create a reduced model by fixing some actual joints from the robot
+        fixed_joints = robot_model.joint_names[:2]  # Use first two joints from the robot
+        reduced_robot = ReducedRobotModel(robot_model, fixed_joints)
+
+        # Check that the full robot is stored
+        assert reduced_robot.full_robot is robot_model
+
+        # Check that fixed joints are stored correctly
+        assert reduced_robot.fixed_joints == fixed_joints
+        assert len(reduced_robot.fixed_values) == len(fixed_joints)
+
+        # Check that the number of dofs is reduced
+        assert reduced_robot.num_dofs == robot_model.num_dofs - len(fixed_joints)
+
+
+def test_reduced_robot_model_joint_names(g1_robot_model):
+    """
+    Test that joint_names in ReducedRobotModel excludes fixed joints.
+    """
+    for robot_model in [g1_robot_model]:
+        # Use actual joints from the robot
+        fixed_joints = robot_model.joint_names[:2]  # Use first two joints from the robot
+        reduced_robot = ReducedRobotModel(robot_model, fixed_joints)
+
+        # Check that fixed joints are not in the reduced model's joint names
+        for joint in fixed_joints:
+            assert joint not in reduced_robot.joint_names
+
+        # Check that other joints are still present
+        for joint in robot_model.joint_names:
+            if joint not in fixed_joints:
+                assert joint in reduced_robot.joint_names
+
+
+def test_reduced_robot_model_configuration_conversion(g1_robot_model):
+    """
+    Test conversion between reduced and full configurations.
+    """
+    for robot_model in [g1_robot_model]:
+        # Use actual joints from the robot
+        fixed_joints = robot_model.joint_names[:2]  # Use first two joints from the robot
+        fixed_values = [0.5, 1.0]
+        reduced_robot = ReducedRobotModel(robot_model, fixed_joints, fixed_values)
+
+        # Create a reduced configuration
+        q_reduced = np.zeros(reduced_robot.num_dofs)
+        q_reduced[0] = 0.3  # Set some value for testing
+
+        # Convert to full configuration
+        q_full = reduced_robot.reduced_to_full_configuration(q_reduced)
+
+        # Check that fixed joints have the correct values
+        for joint_name, value in zip(fixed_joints, fixed_values):
+            full_idx = robot_model.dof_index(joint_name)
+            assert q_full[full_idx] == value
+
+        # Convert back to reduced configuration
+        q_reduced_back = reduced_robot.full_to_reduced_configuration(q_full)
+
+        # Check that the conversion is reversible
+        np.testing.assert_array_almost_equal(q_reduced, q_reduced_back)
+
+
+def test_reduced_robot_model_forward_kinematics(g1_robot_model):
+    """
+    Test forward kinematics with the reduced model.
+    """
+    for robot_model in [g1_robot_model]:
+        # Use actual joints from the robot
+        fixed_joints = robot_model.joint_names[:2]  # Use first two joints from the robot
+        reduced_robot = ReducedRobotModel(robot_model, fixed_joints)
+
+        # Create a reduced configuration
+        q_reduced = np.zeros(reduced_robot.num_dofs)
+
+        # Should not raise any exception
+        reduced_robot.cache_forward_kinematics(q_reduced)
+
+        # Check that frame placement works
+        model = robot_model.pinocchio_wrapper.model
+        if len(model.frames) > 1:
+            valid_frame = model.frames[1].name
+            placement = reduced_robot.frame_placement(valid_frame)
+            assert isinstance(placement, pin.SE3)
+
+
+def test_robot_model_clip_configuration(g1_robot_model):
+    """
+    Test that clip_configuration properly clips values to joint limits.
+    """
+    for robot_model in [g1_robot_model]:
+        # Create a configuration with some values outside limits
+        q = np.zeros(robot_model.num_dofs)
+        root_nq = 7 if robot_model.is_floating_base_model else 0
+        # Create extreme values for all joints
+        q[root_nq:] = np.array([100.0, -100.0, 50.0, -50.0] * (robot_model.num_joints // 4 + 1))[
+            : robot_model.num_joints
+        ]
+
+        # Clip the configuration
+        q_clipped = robot_model.clip_configuration(q)
+
+        # Check that values are within limits
+        assert np.all(q_clipped[root_nq:] <= robot_model.upper_joint_limits)
+        assert np.all(q_clipped[root_nq:] >= robot_model.lower_joint_limits)
+
+
+def test_robot_model_get_actuated_joints(g1_robot_model):
+    """
+    Test getting body and hand actuated joints from configuration.
+    """
+    for robot_model in [g1_robot_model]:
+        # Skip if no supplemental info
+        if robot_model.supplemental_info is None:
+            pytest.skip("No supplemental info available for testing actuated joints")
+
+        # Create a test configuration
+        q = np.zeros(robot_model.num_dofs)
+        root_nq = 7 if robot_model.is_floating_base_model else 0
+        q[root_nq:] = np.arange(robot_model.num_joints)  # Set some values for joints
+
+        # Test body actuated joints
+        body_joints = robot_model.get_body_actuated_joints(q)
+        assert len(body_joints) == len(robot_model.get_body_actuated_joint_indices())
+
+        # Test hand actuated joints
+        hand_joints = robot_model.get_hand_actuated_joints(q)
+        assert len(hand_joints) == len(robot_model.get_hand_actuated_joint_indices())
+
+        # Test left hand joints
+        left_hand_joints = robot_model.get_hand_actuated_joints(q, side="left")
+        assert len(left_hand_joints) == len(robot_model.get_hand_actuated_joint_indices("left"))
+
+        # Test right hand joints
+        right_hand_joints = robot_model.get_hand_actuated_joints(q, side="right")
+        assert len(right_hand_joints) == len(robot_model.get_hand_actuated_joint_indices("right"))
+
+
+def test_robot_model_get_configuration_from_actuated_joints(g1_robot_model):
+    """
+    Test creating full configuration from actuated joint values.
+    """
+    for robot_model in [g1_robot_model]:
+        # Skip if no supplemental info
+        if robot_model.supplemental_info is None:
+            pytest.skip("No supplemental info available for testing actuated joints")
+
+        # Create test values for body and hands
+        body_values = np.ones(len(robot_model.get_body_actuated_joint_indices()))
+        hand_values = np.ones(len(robot_model.get_hand_actuated_joint_indices()))
+        left_hand_values = np.ones(len(robot_model.get_hand_actuated_joint_indices("left")))
+        right_hand_values = np.ones(len(robot_model.get_hand_actuated_joint_indices("right")))
+
+        # Test with combined hand values
+        q = robot_model.get_configuration_from_actuated_joints(
+            body_actuated_joint_values=body_values, hand_actuated_joint_values=hand_values
+        )
+        assert q.shape == (robot_model.num_dofs,)
+
+        # Test with separate hand values
+        q = robot_model.get_configuration_from_actuated_joints(
+            body_actuated_joint_values=body_values,
+            left_hand_actuated_joint_values=left_hand_values,
+            right_hand_actuated_joint_values=right_hand_values,
+        )
+        assert q.shape == (robot_model.num_dofs,)
+
+
+def test_robot_model_reset_forward_kinematics(g1_robot_model):
+    """
+    Test resetting forward kinematics to default configuration.
+    """
+    for robot_model in [g1_robot_model]:
+        # Skip if no supplemental info
+        if robot_model.supplemental_info is None:
+            pytest.skip("No supplemental info available for testing")
+
+        # Create a more significant configuration change
+        q = np.zeros(robot_model.num_dofs)
+        root_nq = 7 if robot_model.is_floating_base_model else 0
+        # Set some extreme joint angles
+        q[root_nq:] = np.pi / 2  # 90 degrees for all joints
+        robot_model.cache_forward_kinematics(q)
+
+        # Use a hand frame from supplemental info
+        test_frame = robot_model.supplemental_info.hand_frame_names["left"]
+
+        # Reset to default
+        robot_model.reset_forward_kinematics()
+        # Get frame placement after reset
+        placement_default = robot_model.frame_placement(test_frame)
+
+        # Check that frame placement matches what we get with q_zero
+        robot_model.cache_forward_kinematics(robot_model.q_zero)
+        placement_q_zero = robot_model.frame_placement(test_frame)
+        np.testing.assert_array_almost_equal(
+            placement_default.translation, placement_q_zero.translation
+        )
+        np.testing.assert_array_almost_equal(placement_default.rotation, placement_q_zero.rotation)
+
+
+# Additional tests for ReducedRobotModel
+def test_reduced_robot_model_clip_configuration(g1_robot_model):
+    """
+    Test that clip_configuration works in reduced space.
+    """
+    for robot_model in [g1_robot_model]:
+        fixed_joints = robot_model.joint_names[:2]
+        reduced_robot = ReducedRobotModel(robot_model, fixed_joints)
+
+        # Create a configuration with some values outside limits
+        q_reduced = np.zeros(reduced_robot.num_dofs)
+        root_nq = 7 if reduced_robot.full_robot.is_floating_base_model else 0
+        # Create extreme values for all joints
+        q_reduced[root_nq:] = np.array(
+            [100.0, -100.0, 50.0, -50.0] * (reduced_robot.num_joints // 4 + 1)
+        )[: reduced_robot.num_joints]
+
+        # Clip the configuration
+        q_clipped = reduced_robot.clip_configuration(q_reduced)
+
+        # Check that values are within limits
+        assert np.all(q_clipped[root_nq:] <= reduced_robot.upper_joint_limits)
+        assert np.all(q_clipped[root_nq:] >= reduced_robot.lower_joint_limits)
+
+
+def test_reduced_robot_model_get_actuated_joints(g1_robot_model):
+    """
+    Test getting body and hand actuated joints from reduced configuration.
+    """
+    for robot_model in [g1_robot_model]:
+        # Skip if no supplemental info
+        if robot_model.supplemental_info is None:
+            pytest.skip("No supplemental info available for testing actuated joints")
+
+        fixed_joints = robot_model.joint_names[:2]
+        reduced_robot = ReducedRobotModel(robot_model, fixed_joints)
+
+        # Create a test configuration
+        q_reduced = np.zeros(reduced_robot.num_dofs)
+        root_nq = 7 if reduced_robot.full_robot.is_floating_base_model else 0
+        q_reduced[root_nq:] = np.arange(reduced_robot.num_joints)
+
+        # Test body actuated joints
+        body_joints = reduced_robot.get_body_actuated_joints(q_reduced)
+        assert len(body_joints) == len(reduced_robot.get_body_actuated_joint_indices())
+
+        # Test hand actuated joints
+        hand_joints = reduced_robot.get_hand_actuated_joints(q_reduced)
+        assert len(hand_joints) == len(reduced_robot.get_hand_actuated_joint_indices())
+
+
+def test_reduced_robot_model_get_configuration_from_actuated_joints(g1_robot_model):
+    """
+    Test creating reduced configuration from actuated joint values.
+    """
+    for robot_model in [g1_robot_model]:
+        # Skip if no supplemental info
+        if robot_model.supplemental_info is None:
+            pytest.skip("No supplemental info available for testing actuated joints")
+
+        fixed_joints = robot_model.joint_names[:2]
+        reduced_robot = ReducedRobotModel(robot_model, fixed_joints)
+
+        # Create test values for body and hands
+        body_values = np.ones(len(reduced_robot.get_body_actuated_joint_indices()))
+        hand_values = np.ones(len(reduced_robot.get_hand_actuated_joint_indices()))
+        left_hand_values = np.ones(len(reduced_robot.get_hand_actuated_joint_indices("left")))
+        right_hand_values = np.ones(len(reduced_robot.get_hand_actuated_joint_indices("right")))
+
+        # Test with combined hand values
+        q_reduced = reduced_robot.get_configuration_from_actuated_joints(
+            body_actuated_joint_values=body_values, hand_actuated_joint_values=hand_values
+        )
+        assert q_reduced.shape == (reduced_robot.num_dofs,)
+
+        # Test with separate hand values
+        q_reduced = reduced_robot.get_configuration_from_actuated_joints(
+            body_actuated_joint_values=body_values,
+            left_hand_actuated_joint_values=left_hand_values,
+            right_hand_actuated_joint_values=right_hand_values,
+        )
+        assert q_reduced.shape == (reduced_robot.num_dofs,)
+
+        # Verify that the values were set correctly in the reduced configuration
+        # Check body actuated joints
+        body_indices = reduced_robot.get_body_actuated_joint_indices()
+        np.testing.assert_array_almost_equal(q_reduced[body_indices], body_values)
+
+        # Check left hand actuated joints
+        left_hand_indices = reduced_robot.get_hand_actuated_joint_indices("left")
+        np.testing.assert_array_almost_equal(q_reduced[left_hand_indices], left_hand_values)
+
+        # Check right hand actuated joints
+        right_hand_indices = reduced_robot.get_hand_actuated_joint_indices("right")
+        np.testing.assert_array_almost_equal(q_reduced[right_hand_indices], right_hand_values)
+
+
+def test_reduced_robot_model_reset_forward_kinematics(g1_robot_model):
+    """
+    Test resetting forward kinematics in reduced model.
+    """
+    for robot_model in [g1_robot_model]:
+        # Skip if no supplemental info
+        if robot_model.supplemental_info is None:
+            pytest.skip("No supplemental info available for testing")
+
+        fixed_joints = robot_model.joint_names[:2]
+        reduced_robot = ReducedRobotModel(robot_model, fixed_joints)
+
+        # Create a more significant configuration change
+        q_reduced = np.zeros(reduced_robot.num_dofs)
+        root_nq = 7 if reduced_robot.full_robot.is_floating_base_model else 0
+        # Set some extreme joint angles
+        q_reduced[root_nq:] = np.pi / 2  # 90 degrees for all joints
+        reduced_robot.cache_forward_kinematics(q_reduced)
+
+        # Reset to default
+        reduced_robot.reset_forward_kinematics()
+
+        # Check that frame placement matches what we get with q_zero
+        reduced_robot.cache_forward_kinematics(reduced_robot.q_zero)
+        placement_q_zero = reduced_robot.frame_placement(
+            reduced_robot.supplemental_info.hand_frame_names["left"]
+        )
+        placement_reset = reduced_robot.frame_placement(
+            reduced_robot.supplemental_info.hand_frame_names["left"]
+        )
+        np.testing.assert_array_almost_equal(
+            placement_reset.translation, placement_q_zero.translation
+        )
+        np.testing.assert_array_almost_equal(placement_reset.rotation, placement_q_zero.rotation)
+
+
+def test_reduced_robot_model_from_fixed_groups(g1_robot_model):
+    """
+    Test creating reduced model from fixed joint groups.
+    """
+    for robot_model in [g1_robot_model]:
+        # Skip if no supplemental info
+        if robot_model.supplemental_info is None:
+            pytest.skip("No supplemental info available for testing joint groups")
+
+        # Get a group name from the supplemental info
+        group_name = next(iter(robot_model.supplemental_info.joint_groups.keys()))
+        group_info = robot_model.supplemental_info.joint_groups[group_name]
+
+        # Get all joints that should be fixed (including those from subgroups)
+        expected_fixed_joints = set()
+        # Add direct joints
+        expected_fixed_joints.update(group_info["joints"])
+        # Add joints from subgroups
+        for subgroup_name in group_info["groups"]:
+            subgroup_joints = robot_model.get_joint_group_indices(subgroup_name)
+            expected_fixed_joints.update([robot_model.joint_names[idx] for idx in subgroup_joints])
+
+        # Test from_fixed_groups
+        reduced_robot = ReducedRobotModel.from_fixed_groups(robot_model, [group_name])
+        assert reduced_robot.full_robot is robot_model
+
+        # Verify that fixed joints are not in reduced model's joint names
+        for joint in expected_fixed_joints:
+            assert joint not in reduced_robot.joint_names
+
+        # Verify that fixed joints maintain their values in configuration
+        q_reduced = np.ones(reduced_robot.num_dofs)  # Set some non-zero values
+        q_full = reduced_robot.reduced_to_full_configuration(q_reduced)
+
+        # Get the fixed values from the reduced model
+        fixed_values = dict(zip(reduced_robot.fixed_joints, reduced_robot.fixed_values))
+
+        # Check that all expected fixed joints have their values preserved
+        for joint in expected_fixed_joints:
+            full_idx = robot_model.dof_index(joint)
+            assert q_full[full_idx] == fixed_values[joint]
+
+        # Test from_fixed_group (convenience method)
+        reduced_robot = ReducedRobotModel.from_fixed_group(robot_model, group_name)
+        assert reduced_robot.full_robot is robot_model
+
+        # Verify that fixed joints are not in reduced model's joint names
+        for joint in expected_fixed_joints:
+            assert joint not in reduced_robot.joint_names
+
+        # Verify that fixed joints maintain their values in configuration
+        q_reduced = np.ones(reduced_robot.num_dofs)  # Set some non-zero values
+        q_full = reduced_robot.reduced_to_full_configuration(q_reduced)
+
+        # Get the fixed values from the reduced model
+        fixed_values = dict(zip(reduced_robot.fixed_joints, reduced_robot.fixed_values))
+
+        # Check that all expected fixed joints have their values preserved
+        for joint in expected_fixed_joints:
+            full_idx = robot_model.dof_index(joint)
+            assert q_full[full_idx] == fixed_values[joint]
+
+
+def test_reduced_robot_model_from_active_groups(g1_robot_model):
+    """
+    Test creating reduced model from active joint groups.
+    """
+    for robot_model in [g1_robot_model]:
+        # Skip if no supplemental info
+        if robot_model.supplemental_info is None:
+            pytest.skip("No supplemental info available for testing joint groups")
+
+        # Get a group name from the supplemental info
+        group_name = next(iter(robot_model.supplemental_info.joint_groups.keys()))
+        group_info = robot_model.supplemental_info.joint_groups[group_name]
+
+        # Get all joints that should be active (including those from subgroups)
+        expected_active_joints = set()
+        # Add direct joints
+        expected_active_joints.update(group_info["joints"])
+        # Add joints from subgroups
+        for subgroup_name in group_info["groups"]:
+            subgroup_joints = robot_model.get_joint_group_indices(subgroup_name)
+            expected_active_joints.update([robot_model.joint_names[idx] for idx in subgroup_joints])
+
+        # Get all joints from the model
+        all_joints = set(robot_model.joint_names)
+        # The fixed joints should be all joints minus the active joints
+        expected_fixed_joints = all_joints - expected_active_joints
+
+        # Test from_active_groups
+        reduced_robot = ReducedRobotModel.from_active_groups(robot_model, [group_name])
+        assert reduced_robot.full_robot is robot_model
+
+        # Verify that active joints are in reduced model's joint names
+        for joint in expected_active_joints:
+            assert joint in reduced_robot.joint_names
+
+        # Verify that fixed joints are not in reduced model's joint names
+        for joint in expected_fixed_joints:
+            assert joint not in reduced_robot.joint_names
+
+        # Verify that fixed joints maintain their values in configuration
+        q_reduced = np.ones(reduced_robot.num_dofs)  # Set some non-zero values
+        q_full = reduced_robot.reduced_to_full_configuration(q_reduced)
+
+        # Get the fixed values from the reduced model
+        fixed_values = dict(zip(reduced_robot.fixed_joints, reduced_robot.fixed_values))
+
+        # Check that all expected fixed joints have their values preserved
+        for joint in expected_fixed_joints:
+            full_idx = robot_model.dof_index(joint)
+            assert q_full[full_idx] == fixed_values[joint]
+
+        # Test from_active_group (convenience method)
+        reduced_robot = ReducedRobotModel.from_active_group(robot_model, group_name)
+        assert reduced_robot.full_robot is robot_model
+
+        # Verify that active joints are in reduced model's joint names
+        for joint in expected_active_joints:
+            assert joint in reduced_robot.joint_names
+
+        # Verify that fixed joints are not in reduced model's joint names
+        for joint in expected_fixed_joints:
+            assert joint not in reduced_robot.joint_names
+
+        # Verify that fixed joints maintain their values in configuration
+        q_reduced = np.ones(reduced_robot.num_dofs)  # Set some non-zero values
+        q_full = reduced_robot.reduced_to_full_configuration(q_reduced)
+
+        # Get the fixed values from the reduced model
+        fixed_values = dict(zip(reduced_robot.fixed_joints, reduced_robot.fixed_values))
+
+        # Check that all expected fixed joints have their values preserved
+        for joint in expected_fixed_joints:
+            full_idx = robot_model.dof_index(joint)
+            assert q_full[full_idx] == fixed_values[joint]
+
+
+def test_reduced_robot_model_frame_placement(g1_robot_model):
+    """
+    Test the frame_placement method in reduced model with a valid and invalid frame name.
+    Also test that frame placements change with different configurations.
+    """
+    for robot_model in [g1_robot_model]:
+        # Skip if no supplemental info
+        if robot_model.supplemental_info is None:
+            pytest.skip("No supplemental info available for testing")
+
+        # Create a reduced model by fixing some joints
+        fixed_joints = robot_model.joint_names[:2]
+        reduced_robot = ReducedRobotModel(robot_model, fixed_joints)
+
+        # Use the hand frame from supplemental info
+        test_frame = reduced_robot.supplemental_info.hand_frame_names["left"]
+
+        # Test with zero configuration
+        q_reduced_zero = np.zeros(reduced_robot.num_dofs)
+        reduced_robot.cache_forward_kinematics(q_reduced_zero)
+        placement_zero = reduced_robot.frame_placement(test_frame)
+        assert isinstance(placement_zero, pin.SE3)
+
+        # Test with non-zero configuration
+        q_reduced_non_zero = np.zeros(reduced_robot.num_dofs)
+        root_nq = 7 if reduced_robot.full_robot.is_floating_base_model else 0
+
+        # Set a valid non-zero value for each joint
+        for i in range(root_nq, reduced_robot.num_dofs):
+            # Use a value that's within the joint limits
+            q_reduced_non_zero[i] = 0.5  # 0.5 radians is within most joint limits
+
+        reduced_robot.cache_forward_kinematics(q_reduced_non_zero)
+        placement_non_zero = reduced_robot.frame_placement(test_frame)
+
+        # Verify that frame placements are different with different configurations
+        assert not np.allclose(
+            placement_zero.translation, placement_non_zero.translation
+        ) or not np.allclose(placement_zero.rotation, placement_non_zero.rotation)
+
+        # Should raise an error for an invalid frame
+        with pytest.raises(ValueError, match="Unknown frame"):
+            reduced_robot.frame_placement("non_existent_frame")
+
+
+def test_robot_model_gravity_compensation_basic(g1_robot_model):
+    """
+    Test basic gravity compensation functionality.
+    """
+    for robot_model in [g1_robot_model]:
+        # Skip if no supplemental info
+        if robot_model.supplemental_info is None:
+            pytest.skip("No supplemental info available for testing gravity compensation")
+
+        # Create a valid configuration
+        q = np.zeros(robot_model.num_dofs)
+        if robot_model.is_floating_base_model:
+            # Set floating base to upright position
+            q[:7] = [0, 0, 1.0, 0, 0, 0, 1]  # [x, y, z, qx, qy, qz, qw]
+
+        # Test gravity compensation for all joints
+        gravity_torques = robot_model.compute_gravity_compensation_torques(q)
+
+        # Check output shape
+        assert gravity_torques.shape == (robot_model.num_dofs,)
+
+        # For a humanoid robot with arms, there should be some non-zero gravity torques
+        assert np.any(np.abs(gravity_torques) > 1e-6), "Expected some non-zero gravity torques"
+
+
+def test_robot_model_gravity_compensation_joint_groups(g1_robot_model):
+    """
+    Test gravity compensation with different joint group specifications.
+    """
+    for robot_model in [g1_robot_model]:
+        # Skip if no supplemental info
+        if robot_model.supplemental_info is None:
+            pytest.skip("No supplemental info available for testing gravity compensation")
+
+        # Create a valid configuration
+        q = np.zeros(robot_model.num_dofs)
+        if robot_model.is_floating_base_model:
+            q[:7] = [0, 0, 1.0, 0, 0, 0, 1]
+
+        # Get available joint groups
+        available_groups = list(robot_model.supplemental_info.joint_groups.keys())
+        if not available_groups:
+            pytest.skip("No joint groups available for testing")
+
+        test_group = available_groups[0]  # Use first available group
+
+        # Test with string input
+        gravity_str = robot_model.compute_gravity_compensation_torques(q, test_group)
+        assert gravity_str.shape == (robot_model.num_dofs,)
+
+        # Test with list input
+        gravity_list = robot_model.compute_gravity_compensation_torques(q, [test_group])
+        np.testing.assert_array_equal(gravity_str, gravity_list)
+
+        # Test with set input
+        gravity_set = robot_model.compute_gravity_compensation_torques(q, {test_group})
+        np.testing.assert_array_equal(gravity_str, gravity_set)
+
+        # Test that compensation is selective (some joints should be zero)
+        group_indices = robot_model.get_joint_group_indices(test_group)
+        if len(group_indices) < robot_model.num_dofs:
+            # Check that only specified joints have compensation
+            non_zero_mask = np.abs(gravity_str) > 1e-6
+            compensated_indices = np.where(non_zero_mask)[0]
+            # The compensated indices should be a subset of the group indices
+            assert len(compensated_indices) <= len(group_indices)
+
+
+def test_robot_model_gravity_compensation_multiple_groups(g1_robot_model):
+    """
+    Test gravity compensation with multiple joint groups.
+    """
+    for robot_model in [g1_robot_model]:
+        # Skip if no supplemental info
+        if robot_model.supplemental_info is None:
+            pytest.skip("No supplemental info available for testing gravity compensation")
+
+        # Create a valid configuration
+        q = np.zeros(robot_model.num_dofs)
+        if robot_model.is_floating_base_model:
+            q[:7] = [0, 0, 1.0, 0, 0, 0, 1]
+
+        # Get available joint groups
+        available_groups = list(robot_model.supplemental_info.joint_groups.keys())
+        if len(available_groups) < 2:
+            pytest.skip("Need at least 2 joint groups for testing")
+
+        # Test with multiple groups
+        test_groups = available_groups[:2]
+        gravity_multiple = robot_model.compute_gravity_compensation_torques(q, test_groups)
+        assert gravity_multiple.shape == (robot_model.num_dofs,)
+
+        # Test individual groups
+        gravity_1 = robot_model.compute_gravity_compensation_torques(q, test_groups[0])
+        gravity_2 = robot_model.compute_gravity_compensation_torques(q, test_groups[1])
+
+        # The multiple group result should have at least as many non-zero elements
+        # as either individual group (could be more due to overlaps)
+        nonzero_multiple = np.count_nonzero(np.abs(gravity_multiple) > 1e-6)
+        nonzero_1 = np.count_nonzero(np.abs(gravity_1) > 1e-6)
+        nonzero_2 = np.count_nonzero(np.abs(gravity_2) > 1e-6)
+        assert nonzero_multiple >= max(nonzero_1, nonzero_2)
+
+
+def test_robot_model_gravity_compensation_configuration_dependency(g1_robot_model):
+    """
+    Test that gravity compensation changes with robot configuration.
+    """
+    for robot_model in [g1_robot_model]:
+        # Skip if no supplemental info
+        if robot_model.supplemental_info is None:
+            pytest.skip("No supplemental info available for testing gravity compensation")
+
+        # Get available joint groups - prefer arms if available
+        available_groups = list(robot_model.supplemental_info.joint_groups.keys())
+        test_group = None
+        for group in ["arms", "left_arm", "right_arm"]:
+            if group in available_groups:
+                test_group = group
+                break
+        if test_group is None and available_groups:
+            test_group = available_groups[0]
+        if test_group is None:
+            pytest.skip("No joint groups available for testing")
+
+        # Test with different configurations
+        q1 = np.zeros(robot_model.num_dofs)
+        q2 = np.zeros(robot_model.num_dofs)
+
+        if robot_model.is_floating_base_model:
+            # Both configurations upright but different joint positions
+            q1[:7] = [0, 0, 1.0, 0, 0, 0, 1]
+            q2[:7] = [0, 0, 1.0, 0, 0, 0, 1]
+
+        # Change arm joint positions specifically (not random joints)
+        # This ensures we actually change joints that affect the gravity compensation
+        try:
+            arm_indices = robot_model.get_joint_group_indices(test_group)
+            if len(arm_indices) >= 2:
+                # Change first two arm joints significantly
+                q2[arm_indices[0]] = np.pi / 4  # 45 degrees
+                q2[arm_indices[1]] = np.pi / 6  # 30 degrees
+            elif len(arm_indices) >= 1:
+                # Change first arm joint if only one available
+                q2[arm_indices[0]] = np.pi / 3  # 60 degrees
+        except Exception:
+            # Fallback to changing some joints if arm indices not available
+            if robot_model.is_floating_base_model and robot_model.num_dofs > 9:
+                q2[7] = np.pi / 4
+                q2[8] = np.pi / 6
+            elif not robot_model.is_floating_base_model and robot_model.num_dofs > 2:
+                q2[0] = np.pi / 4
+                q2[1] = np.pi / 6
+
+        # Compute gravity compensation for both configurations
+        gravity_1 = robot_model.compute_gravity_compensation_torques(q1, test_group)
+        gravity_2 = robot_model.compute_gravity_compensation_torques(q2, test_group)
+
+        # They should be different (unless all compensated joints didn't change)
+        # Allow for small numerical differences
+        assert not np.allclose(
+            gravity_1, gravity_2, atol=1e-10
+        ), "Gravity compensation should change with configuration"
+
+
+def test_robot_model_gravity_compensation_error_handling(g1_robot_model):
+    """
+    Test error handling in gravity compensation.
+    """
+    for robot_model in [g1_robot_model]:
+        # Test with wrong configuration size
+        q_wrong = np.zeros(robot_model.num_dofs + 1)
+        with pytest.raises(ValueError, match="Expected q of length"):
+            robot_model.compute_gravity_compensation_torques(q_wrong)
+
+        # Test with invalid joint group
+        q_valid = np.zeros(robot_model.num_dofs)
+        if robot_model.is_floating_base_model:
+            q_valid[:7] = [0, 0, 1.0, 0, 0, 0, 1]
+
+        with pytest.raises(RuntimeError, match="Error computing gravity compensation"):
+            robot_model.compute_gravity_compensation_torques(q_valid, "non_existent_group")
+
+        # Test with mixed valid/invalid groups
+        if robot_model.supplemental_info is not None:
+            available_groups = list(robot_model.supplemental_info.joint_groups.keys())
+            if available_groups:
+                valid_group = available_groups[0]
+                with pytest.raises(RuntimeError, match="Error computing gravity compensation"):
+                    robot_model.compute_gravity_compensation_torques(
+                        q_valid, [valid_group, "non_existent_group"]
+                    )
+
+
+def test_robot_model_gravity_compensation_auto_clip(g1_robot_model):
+    """
+    Test auto-clipping functionality in gravity compensation.
+    """
+    for robot_model in [g1_robot_model]:
+        # Skip if no supplemental info
+        if robot_model.supplemental_info is None:
+            pytest.skip("No supplemental info available for testing gravity compensation")
+
+        # Create configuration with values outside joint limits
+        q = np.zeros(robot_model.num_dofs)
+        root_nq = 7 if robot_model.is_floating_base_model else 0
+
+        if robot_model.is_floating_base_model:
+            q[:7] = [0, 0, 1.0, 0, 0, 0, 1]  # Valid floating base
+
+        # Set extreme joint values (outside limits)
+        if robot_model.num_dofs > root_nq:
+            q[root_nq:] = 100.0  # Very large values
+
+        # Should work with auto_clip=True (default)
+        try:
+            gravity_clipped = robot_model.compute_gravity_compensation_torques(q, auto_clip=True)
+            assert gravity_clipped.shape == (robot_model.num_dofs,)
+        except Exception as e:
+            pytest.skip(f"Auto-clip test skipped due to: {e}")
+
+        # Test with auto_clip=False - might work or might not depending on limits
+        try:
+            gravity_no_clip = robot_model.compute_gravity_compensation_torques(q, auto_clip=False)
+            assert gravity_no_clip.shape == (robot_model.num_dofs,)
+        except Exception:
+            # This is expected if the configuration is invalid
+            pass
+
+
+def test_robot_model_gravity_compensation_arms_specific(g1_robot_model):
+    """
+    Test gravity compensation specifically for arm joints (if available).
+    """
+    for robot_model in [g1_robot_model]:
+        # Skip if no supplemental info
+        if robot_model.supplemental_info is None:
+            pytest.skip("No supplemental info available for testing gravity compensation")
+
+        available_groups = list(robot_model.supplemental_info.joint_groups.keys())
+
+        # Test arms specifically if available
+        if "arms" in available_groups:
+            q = np.zeros(robot_model.num_dofs)
+            if robot_model.is_floating_base_model:
+                q[:7] = [0, 0, 1.0, 0, 0, 0, 1]
+
+            # Test arms gravity compensation
+            gravity_arms = robot_model.compute_gravity_compensation_torques(q, "arms")
+            assert gravity_arms.shape == (robot_model.num_dofs,)
+
+            # Test left and right arms separately if available
+            if "left_arm" in available_groups and "right_arm" in available_groups:
+                gravity_left = robot_model.compute_gravity_compensation_torques(q, "left_arm")
+                gravity_right = robot_model.compute_gravity_compensation_torques(q, "right_arm")
+
+                # Both arms should have non-zero compensation (for typical configurations)
+                if np.any(np.abs(gravity_arms) > 1e-6):
+                    # If arms have compensation, at least one of left/right should too
+                    assert np.any(np.abs(gravity_left) > 1e-6) or np.any(
+                        np.abs(gravity_right) > 1e-6
+                    )
+        else:
+            pytest.skip("No arm joint groups available for testing")

GR00T-WholeBodyControl/decoupled_wbc/tests/control/teleop/test_teleop_retargeting_ik.py

@@ -0,0 +1,196 @@
+import time
+
+import numpy as np
+import pytest
+
+from decoupled_wbc.control.robot_model.instantiation.g1 import instantiate_g1_robot_model
+from decoupled_wbc.control.robot_model.robot_model import RobotModel
+from decoupled_wbc.control.teleop.solver.hand.instantiation.g1_hand_ik_instantiation import (
+    instantiate_g1_hand_ik_solver,
+)
+from decoupled_wbc.control.teleop.teleop_retargeting_ik import TeleopRetargetingIK
+
+
+@pytest.fixture(params=["lower_body", "lower_and_upper_body"])
+def retargeting_ik(request):
+    waist_location = request.param
+    robot_model = instantiate_g1_robot_model(waist_location=waist_location)
+    left_hand_ik_solver, right_hand_ik_solver = instantiate_g1_hand_ik_solver()
+    return TeleopRetargetingIK(
+        robot_model=robot_model,
+        left_hand_ik_solver=left_hand_ik_solver,
+        right_hand_ik_solver=right_hand_ik_solver,
+        enable_visualization=False,  # Change to true to visualize movements
+        body_active_joint_groups=["upper_body"],
+    )
+
+
+def generate_target_wrist_poses(mode: str, side: str, full_robot: RobotModel) -> dict:
+    """
+    Args:
+        mode: One of "rotation" or "translation"
+        side: One of "left" or "right" - specifies which side to animate
+    Returns:
+        Dictionary mapping link names to target poses for both wrists
+    """
+
+    assert mode in ["rotation", "translation", "both"]
+    assert side in ["left", "right", "both"]
+
+    # Set up initial state
+    full_robot.cache_forward_kinematics(full_robot.q_zero)
+
+    # Get both wrist link names
+    left_wrist_link = full_robot.supplemental_info.hand_frame_names["left"]
+    right_wrist_link = full_robot.supplemental_info.hand_frame_names["right"]
+
+    # Initialize default poses for both sides
+    left_default_pose = full_robot.frame_placement(left_wrist_link).np
+    right_default_pose = full_robot.frame_placement(right_wrist_link).np
+
+    left_initial_pose_matrix = full_robot.frame_placement(left_wrist_link).np
+    right_initial_pose_matrix = full_robot.frame_placement(right_wrist_link).np
+
+    # Constants
+    translation_cycle_duration = 4.0
+    rotation_cycle_duration = 4.0
+    total_duration = 12.0
+    translation_amplitude = 0.1
+    rotation_amplitude = np.deg2rad(60)  # 30 degrees
+
+    body_data_list = []
+    for t in np.linspace(0, total_duration, 100):
+        rotation_matrix = np.eye(3)
+        current_left_translation_vector = left_initial_pose_matrix[:3, 3].copy()
+        current_right_translation_vector = right_initial_pose_matrix[:3, 3].copy()
+
+        if mode == "rotation" or mode == "both":
+            # For rotation-only mode, start rotating immediately
+            rotation_axis_index = int(t // rotation_cycle_duration) % 3
+            time_within_cycle = t % rotation_cycle_duration
+            angle = rotation_amplitude * np.sin(
+                (2 * np.pi / rotation_cycle_duration) * time_within_cycle
+            )
+
+            if rotation_axis_index == 0:  # Roll
+                rotation_matrix = np.array(
+                    [
+                        [1, 0, 0],
+                        [0, np.cos(angle), -np.sin(angle)],
+                        [0, np.sin(angle), np.cos(angle)],
+                    ]
+                )
+            elif rotation_axis_index == 2:  # Pitch
+                rotation_matrix = np.array(
+                    [
+                        [np.cos(angle), 0, np.sin(angle)],
+                        [0, 1, 0],
+                        [-np.sin(angle), 0, np.cos(angle)],
+                    ]
+                )
+            else:  # Yaw
+                rotation_matrix = np.array(
+                    [
+                        [np.cos(angle), -np.sin(angle), 0],
+                        [np.sin(angle), np.cos(angle), 0],
+                        [0, 0, 1],
+                    ]
+                )
+
+        if mode == "translation" or mode == "both":
+            translation_axis_index = int(t // translation_cycle_duration) % 3
+            time_within_cycle = t % translation_cycle_duration
+            offset = translation_amplitude * np.sin(
+                (2 * np.pi / translation_cycle_duration) * time_within_cycle
+            )
+            current_left_translation_vector[translation_axis_index] += offset
+            current_right_translation_vector[translation_axis_index] += offset
+
+        # Construct the 4x4 pose matrix for the animated side
+        left_animated_pose = np.eye(4)
+        left_animated_pose[:3, :3] = rotation_matrix
+        left_animated_pose[:3, 3] = current_left_translation_vector
+
+        right_animated_pose = np.eye(4)
+        right_animated_pose[:3, :3] = rotation_matrix
+        right_animated_pose[:3, 3] = current_right_translation_vector
+
+        # Create body_data dictionary with both wrists
+        body_data = {}
+        if side == "left":
+            body_data[left_wrist_link] = left_animated_pose
+            body_data[right_wrist_link] = right_default_pose
+        elif side == "right":
+            body_data[left_wrist_link] = left_default_pose
+            body_data[right_wrist_link] = right_animated_pose
+        elif side == "both":
+            body_data[left_wrist_link] = left_animated_pose
+            body_data[right_wrist_link] = right_animated_pose
+
+        body_data_list.append(body_data)
+
+    return body_data_list
+
+
+@pytest.mark.parametrize("mode", ["translation", "rotation"])
+@pytest.mark.parametrize("side", ["both", "left", "right"])
+def test_ik_matches_fk(retargeting_ik, mode, side):
+    full_robot = retargeting_ik.full_robot
+
+    # Generate target wrist poses
+    body_data_list = generate_target_wrist_poses(mode, side, full_robot)
+
+    max_pos_error = 0
+    max_rot_error = 0
+
+    for body_data in body_data_list:
+
+        time_start = time.time()
+
+        # Run IK to get joint angles
+        q = retargeting_ik.compute_joint_positions(
+            body_data,
+            left_hand_data=None,  # Hand IK not tested
+            right_hand_data=None,  # Hand IK not tested
+        )
+
+        time_end = time.time()
+        ik_time = time_end - time_start
+        print(f"IK time: {ik_time} s")
+        # Test commented out because of inconsistency in CI/CD computation time
+        # assert ik_time < 0.05, f"IK time too high for 20Hz loop: {ik_time} s"
+
+        # Run FK to compute where the wrists actually ended up
+        full_robot.cache_forward_kinematics(q, auto_clip=False)
+        left_wrist_link = full_robot.supplemental_info.hand_frame_names["left"]
+        right_wrist_link = full_robot.supplemental_info.hand_frame_names["right"]
+        T_fk_left = full_robot.frame_placement(left_wrist_link).np
+        T_fk_right = full_robot.frame_placement(right_wrist_link).np
+        T_target_left = body_data[left_wrist_link]
+        T_target_right = body_data[right_wrist_link]
+
+        # Check that FK translation matches target translation
+        pos_fk_left = T_fk_left[:3, 3]
+        pos_target_left = T_target_left[:3, 3]
+        pos_fk_right = T_fk_right[:3, 3]
+        pos_target_right = T_target_right[:3, 3]
+
+        max_pos_error = max(max_pos_error, np.linalg.norm(pos_fk_left - pos_target_left))
+        max_pos_error = max(max_pos_error, np.linalg.norm(pos_fk_right - pos_target_right))
+
+        # Check that FK rotation matches target rotation
+        rot_fk_left = T_fk_left[:3, :3]
+        rot_target_left = T_target_left[:3, :3]
+        rot_diff_left = rot_fk_left @ rot_target_left.T
+        rot_error_left = np.arccos(np.clip((np.trace(rot_diff_left) - 1) / 2, -1, 1))
+        rot_fk_right = T_fk_right[:3, :3]
+        rot_target_right = T_target_right[:3, :3]
+        rot_diff_right = rot_fk_right @ rot_target_right.T
+        rot_error_right = np.arccos(np.clip((np.trace(rot_diff_right) - 1) / 2, -1, 1))
+
+        max_rot_error = max(max_rot_error, rot_error_left)
+        max_rot_error = max(max_rot_error, rot_error_right)
+
+    assert max_pos_error < 0.01 and max_rot_error < np.deg2rad(
+        1
+    ), f"Max position error: {max_pos_error}, Max rotation error: {np.rad2deg(max_rot_error)} deg"

GR00T-WholeBodyControl/decoupled_wbc/tests/control/visualization/test_meshcat_visualizer_env.py

@@ -0,0 +1,88 @@
+import pathlib
+import time
+
+import numpy as np
+import pytest
+
+from decoupled_wbc.control.robot_model import RobotModel
+from decoupled_wbc.control.robot_model.supplemental_info.g1.g1_supplemental_info import (
+    G1SupplementalInfo,
+)
+
+
+@pytest.fixture
+def env_fixture():
+    """
+    Pytest fixture that creates and yields the MeshcatVisualizerEnv.
+    After the test, it closes the environment to clean up.
+    """
+    from decoupled_wbc.control.visualization.meshcat_visualizer_env import MeshcatVisualizerEnv
+
+    root_dir = pathlib.Path(__file__).parent.parent.parent.parent
+    urdf_path = str(
+        root_dir / "decoupled_wbc/control/robot_model/model_data/g1/g1_29dof_with_hand.urdf"
+    )
+    asset_path = str(root_dir / "decoupled_wbc/control/robot_model/model_data/g1")
+    robot_config = {
+        "asset_path": asset_path,
+        "urdf_path": urdf_path,
+    }
+    robot_model = RobotModel(
+        robot_config["urdf_path"],
+        robot_config["asset_path"],
+        supplemental_info=G1SupplementalInfo(),
+    )
+    env = MeshcatVisualizerEnv(robot_model)
+    time.sleep(0.5)
+    yield env
+    env.close()
+
+
+def test_meshcat_env_init(env_fixture):
+    """
+    Test that the environment initializes without errors
+    and that reset() returns the proper data structure.
+    """
+    env = env_fixture
+    initial_obs = env.reset()
+    assert isinstance(initial_obs, dict), "reset() should return a dictionary."
+    assert "q" in initial_obs, "The returned dictionary should contain key 'q'."
+    assert (
+        len(initial_obs["q"]) == env.robot_model.num_dofs
+    ), "Length of 'q' should match the robot's DOF."
+
+
+def test_meshcat_env_observation(env_fixture):
+    """
+    Test that the observe() method returns a valid observation
+    conforming to the environment's observation space.
+    """
+    env = env_fixture
+    observation = env.observe()
+    assert isinstance(observation, dict), "observe() should return a dictionary."
+    assert "q" in observation, "The returned dictionary should contain key 'q'."
+    assert (
+        len(observation["q"]) == env.robot_model.num_dofs
+    ), "Length of 'q' should match the robot's DOF."
+
+
+def test_meshcat_env_action(env_fixture):
+    """
+    Test that we can queue an action and visualize it without error.
+    """
+    env = env_fixture
+    # Build a dummy action within the action space
+    test_action = {"q": 0.2 * np.ones(env.robot_model.num_dofs)}
+
+    # This should not raise an exception and should visualize the correct configuration
+    env.queue_action(test_action)
+
+
+def test_meshcat_env_close(env_fixture):
+    """
+    Test closing the environment. (Though the fixture calls env.close()
+    automatically, we can invoke it here to ensure it's safe to do so.)
+    """
+    env = env_fixture
+    env.close()
+    # If close() triggers no exceptions, we're good.

GR00T-WholeBodyControl/decoupled_wbc/tests/data/test_exporter.py

@@ -0,0 +1,522 @@
+import json
+from pathlib import Path
+import shutil
+import tempfile
+import time
+
+from lerobot.common.datasets.lerobot_dataset import LeRobotDataset
+import numpy as np
+import pytest
+
+from decoupled_wbc.data.exporter import DataCollectionInfo, Gr00tDataExporter
+
+
+@pytest.fixture
+def test_features():
+    """Fixture providing test features dict."""
+    return {
+        "observation.images.ego_view": {
+            "dtype": "video",
+            "shape": [64, 64, 3],  # Small images for faster tests
+            "names": ["height", "width", "channel"],
+        },
+        "observation.state": {
+            "dtype": "float32",
+            "shape": (8,),
+            "names": ["x1", "x2", "x3", "x4", "x5", "x6", "x7", "x8"],
+        },
+        "action": {
+            "dtype": "float32",
+            "shape": (8,),
+            "names": ["a1", "a2", "a3", "a4", "a5", "a6", "a7", "a8"],
+        },
+    }
+
+
+@pytest.fixture
+def test_modality_config():
+    return {
+        "state": {"feature1": {"start": 0, "end": 4}, "feature2": {"start": 4, "end": 9}},
+        "action": {"feature1": {"start": 0, "end": 4}, "feature2": {"start": 4, "end": 9}},
+        "video": {"rs_view": {"original_key": "observation.images.ego_view"}},
+        "annotation": {"human.task_description": {"original_key": "task_index"}},
+    }
+
+
+@pytest.fixture
+def test_data_collection_info():
+    return DataCollectionInfo(
+        teleoperator_username="test_user",
+        support_operator_username="test_user",
+        robot_type="test_robot",
+        lower_body_policy="test_policy",
+        wbc_model_path="test_path",
+    )
+
+
+def get_test_frame(step: int):
+    """Generate a test frame with data that varies by step."""
+    # Create a simple, small image that will encode quickly
+    img = np.ones((64, 64, 3), dtype=np.uint8) * (step % 255)
+    # Add a pattern to make each frame unique and verifiable
+    img[step % 64, :, :] = 255 - (step % 255)
+
+    return {
+        "observation.images.ego_view": img,
+        "observation.state": np.ones(8, dtype=np.float32) * step,
+        "action": np.ones(8, dtype=np.float32) * step,
+    }
+
+
+@pytest.fixture
+def temp_dir():
+    """Create a temporary directory for test data that's cleaned up after tests."""
+    temp_dir = tempfile.mkdtemp()
+    yield Path(temp_dir) / "dataset"
+    shutil.rmtree(temp_dir)
+
+
+class TestInterruptAndResume:
+    """Test class for simulating interruption and resumption of recording."""
+
+    # Skip this test if ffmpeg is not installed
+    @pytest.mark.skipif(
+        shutil.which("ffmpeg") is None, reason="ffmpeg not installed, skipping test"
+    )
+    def test_interrupted_mid_episode(
+        self, temp_dir, test_features, test_modality_config, test_data_collection_info
+    ):
+        """
+        Test that simulates a recording session that gets interrupted and then resumes.
+
+        This test uses the actual Gr00tDataExporter implementation with no mocks.
+        """
+        # Constants for the test
+        NUM_EPISODES = 2
+        FRAMES_PER_EPISODE = 5
+
+        # Pick a random episode and frame to interrupt at
+        interrupt_episode = 1
+        interrupt_frame = 3
+
+        print(f"Will interrupt at episode {interrupt_episode}, frame {interrupt_frame}")
+
+        # Track what we've added to verify later
+        completed_episodes = []
+        frames_added_first_session = 0
+
+        # Initial recording session
+        try:
+            # Start recording with real Gr00tDataExporter
+            exporter1 = Gr00tDataExporter.create(
+                save_root=temp_dir,
+                fps=30,
+                features=test_features,
+                modality_config=test_modality_config,
+                task="test_task",
+                robot_type="test_robot",
+                vcodec="libx264",  # Use a common codec that should be available
+                data_collection_info=test_data_collection_info,
+            )
+
+            # Record episodes until interruption
+            for episode in range(NUM_EPISODES):
+                for frame in range(FRAMES_PER_EPISODE):
+                    # Simulate interruption
+                    if episode == interrupt_episode and frame == interrupt_frame:
+                        print(f"Simulating interruption at episode {episode}, frame {frame}")
+                        raise KeyboardInterrupt("Simulated interruption")
+
+                    # Add frame
+                    exporter1.add_frame(get_test_frame(frame))
+                    frames_added_first_session += 1
+
+                # Save episode
+                exporter1.save_episode()
+                completed_episodes.append(episode)
+
+        except KeyboardInterrupt:
+            print(f"Recording interrupted at episode {interrupt_episode}, frame {interrupt_frame}")
+            print(f"Completed episodes: {completed_episodes}")
+            # Don't consolidate since we're interrupted
+            pass
+
+        # Verify what was recorded before interruption
+        assert len(completed_episodes) == interrupt_episode
+        assert (
+            frames_added_first_session == interrupt_episode * FRAMES_PER_EPISODE + interrupt_frame
+        )
+
+        # Let file system operations complete
+        time.sleep(0.5)
+
+        # Resume recording - create a new exporter pointing to the same directory
+        exporter2 = Gr00tDataExporter.create(
+            save_root=temp_dir,
+            fps=30,
+            features=test_features,
+            modality_config=test_modality_config,
+            task="test_task",
+            robot_type="test_robot",
+            vcodec="libx264",
+        )
+
+        # The interrupted episode had frames added but wasn't saved
+        # In a real scenario with the current implementation, we need to restart that episode
+
+        # Record all episodes from the beginning
+        frames_added_second_session = 0
+        episodes_saved_second_session = 0
+
+        for episode in range(NUM_EPISODES):
+            for frame in range(FRAMES_PER_EPISODE):
+                exporter2.add_frame(get_test_frame(frame))
+                frames_added_second_session += 1
+
+            # Save episode
+            exporter2.save_episode()
+            episodes_saved_second_session += 1
+
+        # Verify the result
+        assert frames_added_second_session == NUM_EPISODES * FRAMES_PER_EPISODE
+        assert episodes_saved_second_session == NUM_EPISODES
+
+        # Verify actual files were created
+        for episode_idx in range(NUM_EPISODES):
+            video_path = exporter2.root / exporter2.meta.get_video_file_path(
+                episode_idx, "observation.images.ego_view"
+            )
+            assert video_path.exists(), f"Video file not found: {video_path}"
+
+    @pytest.mark.skipif(
+        shutil.which("ffmpeg") is None, reason="ffmpeg not installed, skipping test"
+    )
+    def test_interrupted_after_episode_completion(
+        self, temp_dir, test_features, test_modality_config, test_data_collection_info
+    ):
+        """
+        Test specifically for the case when interruption happens after an episode is completed.
+        Uses the real Gr00tDataExporter implementation.
+        """
+        # First session - record 1 complete episode and then interrupt
+        exporter1 = Gr00tDataExporter.create(
+            save_root=temp_dir,
+            fps=30,
+            features=test_features,
+            modality_config=test_modality_config,
+            task="test_task",
+            data_collection_info=test_data_collection_info,
+            vcodec="libx264",
+        )
+
+        # Record 1 complete episode
+        for frame in range(5):
+            exporter1.add_frame(get_test_frame(frame))
+        exporter1.save_episode()
+
+        # Let file system operations complete
+        time.sleep(0.5)
+
+        # Verify the first episode was saved
+        video_path = exporter1.root / exporter1.meta.get_video_file_path(
+            0, "observation.images.ego_view"
+        )
+        assert video_path.exists(), f"First episode video file not found: {video_path}"
+
+        # Second session - resume and record another episode
+        exporter2 = Gr00tDataExporter.create(
+            save_root=temp_dir,
+            fps=30,
+            features=test_features,
+            modality_config=test_modality_config,
+            task="test_task",
+            vcodec="libx264",
+        )
+
+        # Record the second episode
+        for frame in range(5):
+            exporter2.add_frame(get_test_frame(frame))
+        exporter2.save_episode()
+
+        # Verify the second episode was saved
+        video_path = exporter2.root / exporter2.meta.get_video_file_path(
+            1, "observation.images.ego_view"
+        )
+        assert video_path.exists(), f"Second episode video file not found: {video_path}"
+
+    @pytest.mark.skipif(
+        shutil.which("ffmpeg") is None, reason="ffmpeg not installed, skipping test"
+    )
+    def test_interrupted_no_episode_completion(
+        self, temp_dir, test_features, test_modality_config, test_data_collection_info
+    ):
+        """
+        Test specifically for the case when interruption happens in the middle of recording an episode.
+        Uses the real Gr00tDataExporter implementation.
+        """
+        # First session - add some frames and interrupt before saving
+        exporter1 = Gr00tDataExporter.create(
+            save_root=temp_dir,
+            fps=30,
+            features=test_features,
+            modality_config=test_modality_config,
+            task="test_task",
+            data_collection_info=test_data_collection_info,
+            vcodec="libx264",
+        )
+
+        # Add 3 frames but don't save
+        for frame in range(3):
+            exporter1.add_frame(get_test_frame(frame))
+        # Don't save episode or consolidate to simulate interruption
+        # The episode buffer is only in memory and will be lost on interruption
+
+        # Let file system operations complete
+        time.sleep(0.5)
+
+        # Verify no episode was saved
+        video_path = exporter1.root / exporter1.meta.get_video_file_path(
+            0, "observation.images.ego_view"
+        )
+        assert not video_path.exists(), f"Episode should not have been saved: {video_path}"
+
+        # Second session - will raise an error because no meta file exist, so we can't resume
+        with pytest.raises(ValueError):
+            _ = Gr00tDataExporter.create(
+                save_root=temp_dir,
+                fps=30,
+                features=test_features,
+                modality_config=test_modality_config,
+                task="test_task",
+                vcodec="libx264",
+            )
+
+
+@pytest.mark.skipif(shutil.which("ffmpeg") is None, reason="ffmpeg not installed, skipping test")
+def test_full_workflow(temp_dir, test_features, test_modality_config, test_data_collection_info):
+    """
+    Test that simulates the complete workflow from the record_session.py example.
+    """
+    NUM_EPISODES = 2
+    FRAMES_PER_EPISODE = 3
+
+    # Create the exporter
+    exporter = Gr00tDataExporter.create(
+        save_root=temp_dir,
+        fps=20,
+        features=test_features,
+        modality_config=test_modality_config,
+        task="test_task",
+        data_collection_info=test_data_collection_info,
+        robot_type="dummy",
+    )
+
+    # Create a small dataset
+    for episode_index in range(NUM_EPISODES):
+        for frame_index in range(FRAMES_PER_EPISODE):
+            exporter.add_frame(get_test_frame(frame_index))
+        exporter.save_episode()
+
+    # check modality config
+    modality_config_path = exporter.root / "meta" / "modality.json"
+    assert modality_config_path.exists(), f"{modality_config_path} does not exists."
+    with open(modality_config_path, "rb") as f:
+        actual_modality_config = json.load(f)
+
+    assert (
+        actual_modality_config == test_modality_config
+    ), f"Modality configs don't match.\nActual: {actual_modality_config}\nExpected: {test_modality_config}"
+
+    # Verify results
+    for episode_idx in range(NUM_EPISODES):
+        video_path = exporter.root / exporter.meta.get_video_file_path(
+            episode_idx, "observation.images.ego_view"
+        )
+        assert video_path.exists(), f"Video file not found: {video_path}"
+
+    # Check that the expected number of episodes exists
+    episode_count = 0
+    for path in exporter.root.glob("**/*.mp4"):
+        episode_count += 1
+    assert episode_count == NUM_EPISODES, f"Expected {NUM_EPISODES} episodes, found {episode_count}"
+
+    # Check the values of the dataset
+    dataset = LeRobotDataset(
+        repo_id="dataset",
+        root=temp_dir,
+    )
+    for episode_idx in range(NUM_EPISODES):
+        for frame_idx in range(FRAMES_PER_EPISODE):
+            expected_frame = get_test_frame(frame_idx)
+            actual_frame = dataset[episode_idx * FRAMES_PER_EPISODE + frame_idx]
+            print(actual_frame["observation.images.ego_view"])
+            actual_image_frame = actual_frame["observation.images.ego_view"].permute(1, 2, 0) * 255
+            assert np.allclose(
+                actual_image_frame.numpy(), expected_frame["observation.images.ego_view"], atol=10
+            )  # Allow some tolerance for video compression
+            assert np.allclose(
+                actual_frame["observation.state"], expected_frame["observation.state"]
+            )
+            assert np.allclose(actual_frame["action"], expected_frame["action"])
+
+        # validate data_collection_info
+        assert dataset.meta.info["data_collection_info"] == test_data_collection_info.to_dict()
+
+
+@pytest.mark.skipif(shutil.which("ffmpeg") is None, reason="ffmpeg not installed, skipping test")
+def test_overwrite_existing_dataset_false(
+    temp_dir, test_features, test_modality_config, test_data_collection_info
+):
+    """
+    Test that appends to the existing dataset when overwrite_existing is set to false.
+    """
+    # first dataset
+    FIRST_NUM_EPISODES = 2
+    FIRST_FRAMES_PER_EPISODE = 3
+
+    exporter = Gr00tDataExporter.create(
+        save_root=temp_dir,
+        fps=20,
+        features=test_features,
+        modality_config=test_modality_config,
+        task="test_task",
+        data_collection_info=test_data_collection_info,
+        robot_type="dummy",
+    )
+    # !! `overwrite_existing` should always be set to false by default
+    # So we're deliberately not setting the overwrite_existing argument here.
+    # This test ensures that
+    # i. the default behavior is overwrite_existing=False
+    # ii. the dataset appends to the existing dataset (instead of overwriting)
+
+    # Create a first dataset
+    for episode_index in range(FIRST_NUM_EPISODES):
+        for frame_index in range(FIRST_FRAMES_PER_EPISODE):
+            exporter.add_frame(get_test_frame(frame_index))
+        exporter.save_episode()
+
+    # second dataset
+    del exporter
+    SECOND_NUM_EPISODES = 3
+    SECOND_FRAMES_PER_EPISODE = 2
+
+    exporter = Gr00tDataExporter.create(
+        save_root=temp_dir,
+        fps=20,
+        features=test_features,
+        modality_config=test_modality_config,
+        task="test_task",
+        robot_type="dummy",
+    )
+    for episode_index in range(SECOND_NUM_EPISODES):
+        for frame_index in range(SECOND_FRAMES_PER_EPISODE):
+            exporter.add_frame(get_test_frame(frame_index))
+        exporter.save_episode()
+
+    # verify that there are
+    EXPECTED_NUM_EPISODES = FIRST_NUM_EPISODES + SECOND_NUM_EPISODES
+    assert len(list(exporter.root.glob("**/*.mp4"))) == EXPECTED_NUM_EPISODES
+    assert len(list(exporter.root.glob("**/*.parquet"))) == EXPECTED_NUM_EPISODES
+
+
+def test_overwrite_existing_dataset_true(
+    temp_dir, test_features, test_modality_config, test_data_collection_info
+):
+    """
+    Test that overwrites to an existing dataset when overwrite_existing=True.
+    """
+    # first dataset
+    FIRST_NUM_EPISODES = 2
+    FIRST_FRAMES_PER_EPISODE = 3
+
+    exporter = Gr00tDataExporter.create(
+        save_root=temp_dir,
+        fps=20,
+        features=test_features,
+        modality_config=test_modality_config,
+        task="test_task",
+        data_collection_info=test_data_collection_info,
+        robot_type="dummy",
+    )
+
+    # Create a first dataset
+    for episode_index in range(FIRST_NUM_EPISODES):
+        for frame_index in range(FIRST_FRAMES_PER_EPISODE):
+            exporter.add_frame(get_test_frame(frame_index))
+        exporter.save_episode()
+
+    # verify that the dataset is written to the disk
+    assert len(list(exporter.root.glob("**/*.mp4"))) == FIRST_NUM_EPISODES
+    assert len(list(exporter.root.glob("**/*.parquet"))) == FIRST_NUM_EPISODES
+
+    # second dataset
+    SECOND_NUM_EPISODES = 3
+    SECOND_FRAMES_PER_EPISODE = 2
+
+    # re-initialize the exporter
+    del exporter
+    exporter = Gr00tDataExporter.create(
+        save_root=temp_dir,
+        fps=20,
+        features=test_features,
+        modality_config=test_modality_config,
+        task="test_task",
+        data_collection_info=test_data_collection_info,
+        robot_type="dummy",
+        overwrite_existing=True,
+    )
+    for episode_index in range(SECOND_NUM_EPISODES):
+        for frame_index in range(SECOND_FRAMES_PER_EPISODE):
+            exporter.add_frame(get_test_frame(frame_index))
+        exporter.save_episode()
+
+    # verify that the dataset is overwritten
+    assert len(list(exporter.root.glob("**/*.mp4"))) == SECOND_NUM_EPISODES
+    assert len(list(exporter.root.glob("**/*.parquet"))) == SECOND_NUM_EPISODES
+
+
+def test_save_episode_as_discarded_and_skip(
+    temp_dir, test_features, test_modality_config, test_data_collection_info
+):
+    """
+    Test that verifies the functionality of saving an episode as discarded and skipping an episode.
+    """
+    FIRST_NUM_EPISODES = 10
+    FIRST_FRAMES_PER_EPISODE = 3
+
+    exporter = Gr00tDataExporter.create(
+        save_root=temp_dir,
+        fps=20,
+        features=test_features,
+        modality_config=test_modality_config,
+        task="test_task",
+        data_collection_info=test_data_collection_info,
+        robot_type="dummy",
+    )
+
+    # Create a first dataset
+    saved_episodes = 0
+    discarded_episode_indices = []
+    for episode_index in range(FIRST_NUM_EPISODES):
+        for frame_index in range(FIRST_FRAMES_PER_EPISODE):
+            exporter.add_frame(get_test_frame(frame_index))
+        if episode_index % 3 == 0:
+            exporter.save_episode_as_discarded()
+            discarded_episode_indices.append(saved_episodes)
+            saved_episodes += 1
+        elif episode_index % 3 == 1:
+            exporter.skip_and_start_new_episode()
+        else:
+            exporter.save_episode()
+            saved_episodes += 1
+
+    # verify that the dataset is written to the disk
+    assert len(list(exporter.root.glob("**/*.mp4"))) == saved_episodes
+    assert len(list(exporter.root.glob("**/*.parquet"))) == saved_episodes
+
+    dataset = LeRobotDataset(
+        repo_id="dataset",
+        root=temp_dir,
+    )
+
+    assert dataset.meta.info["discarded_episode_indices"] == discarded_episode_indices

GR00T-WholeBodyControl/decoupled_wbc/tests/sim/test_sim_data_collection.py

@@ -0,0 +1,64 @@
+from decoupled_wbc.control.main.teleop.configs.configs import SyncSimDataCollectionConfig
+from decoupled_wbc.control.main.teleop.run_sync_sim_data_collection import (
+    main as data_collection_main,
+)
+
+
+def test_sim_data_collection_unit(robot_name="G1", task_name="GroundOnly"):
+    """
+    Fast CI unit test for simulation data collection (50 steps, no tracking checks).
+
+    This test validates that:
+    1. Data collection completes successfully
+    2. Upper body joints are moving (velocity check)
+
+    Note: This test runs for only 50 steps and does not perform end effector tracking validation
+    for faster CI execution.
+    """
+    config = SyncSimDataCollectionConfig()
+    config.robot = robot_name
+    config.task_name = task_name
+    config.enable_visualization = False
+    config.enable_real_device = False
+    config.enable_onscreen = False
+    config.save_img_obs = True
+    config.ci_test = True
+    config.ci_test_mode = "unit"
+    config.replay_data_path = "decoupled_wbc/tests/replay_data/all_joints_raw_data_replay.pkl"
+    config.remove_existing_dir = True
+    config.enable_gravity_compensation = True
+    res = data_collection_main(config)
+    assert res, "Data collection did not pass for all datasets"
+
+
+def test_sim_data_collection_pre_merge(robot_name="G1", task_name="GroundOnly"):
+    """
+    Pre-merge test for simulation data collection with end effector tracking validation (500 steps).
+
+    This test validates that:
+    1. Data collection completes successfully
+    2. Upper body joints are moving (velocity check)
+    3. End effector tracking error is within thresholds:
+    - G1 robots:
+        Max position error < 7cm (0.07m), Max rotation error < 17°,
+        Average position error < 5cm (0.05m), Average rotation error < 12°
+    """
+    config = SyncSimDataCollectionConfig()
+    config.robot = robot_name
+    config.task_name = task_name
+    config.enable_visualization = False
+    config.enable_real_device = False
+    config.enable_onscreen = False
+    config.save_img_obs = True
+    config.ci_test = True
+    config.ci_test_mode = "pre_merge"
+    config.replay_data_path = "decoupled_wbc/tests/replay_data/all_joints_raw_data_replay.pkl"
+    config.remove_existing_dir = True
+    config.enable_gravity_compensation = True
+    res = data_collection_main(config)
+    assert res, "Data collection did not pass for all datasets"
+
+
+if __name__ == "__main__":
+    # Run unit tests for fast CI
+    test_sim_data_collection_unit("G1", "GroundOnly")

GR00T-WholeBodyControl/external_dependencies/XRoboToolkit-PC-Service-Pybind_X86_and_ARM64/CLAUDE.md

@@ -0,0 +1,108 @@
+# CLAUDE.md
+
+This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
+
+## Project Overview
+
+This project provides Python bindings for the XRoboToolkit PC Service SDK, enabling Python applications to extract XR state data including controller poses, hand tracking, and body motion capture from XR devices (primarily PICO headsets).
+
+## Architecture
+
+The project consists of:
+
+- **Core C++ Bindings** (`bindings/py_bindings.cpp`): Pybind11-based C++ module that wraps the PXREARobotSDK
+- **SDK Integration**: Uses the XRoboToolkit-PC-Service SDK (cloned from external repository)
+- **Build System**: CMake-based build with Python setuptools integration
+- **Multi-platform Support**: Linux (x86_64/aarch64) and Windows
+
+Key components:
+- `PXREARobotSDK.h`: Main SDK header providing device connectivity and data parsing
+- `py_bindings.cpp`: Thread-safe C++ wrapper with mutex-protected global state variables
+- JSON parsing using nlohmann/json for device state data
+- Callback-based data updates from the SDK
+
+## Build Commands
+
+### Ubuntu/Linux Setup and Build
+```bash
+# Full setup (downloads dependencies and builds)
+bash setup_ubuntu.sh
+
+# Manual build after setup
+python setup.py install
+
+# Clean build artifacts
+python setup.py clean
+```
+
+### Windows Setup and Build
+```batch
+# Full setup (downloads dependencies and builds)
+setup_windows.bat
+
+# Manual build after setup
+python setup.py install
+```
+
+### Development Commands
+```bash
+# Uninstall existing package
+pip uninstall -y xrobotoolkit_sdk
+
+# Install pybind11 dependency
+conda install -c conda-forge pybind11
+# or
+pip install pybind11
+
+# Build and install
+python setup.py install
+```
+
+## Data Flow and Threading
+
+The SDK uses a callback-based architecture:
+- `OnPXREAClientCallback`: Main callback function that receives JSON data from connected devices
+- Global state variables (poses, button states, etc.) are updated in real-time
+- Thread-safe access via mutex locks for each data category
+- Data parsing from comma-separated pose strings to arrays
+
+## Key Functions and Data Types
+
+### Controller Data
+- Poses: `std::array<double, 7>` (x,y,z,qx,qy,qz,qw)
+- Buttons: Menu, Primary, Secondary, Axis Click
+- Analog: Trigger, Grip, Axis (x,y)
+
+### Hand Tracking
+- 26 joints per hand with 7 values each (position + quaternion)
+- Hand scale factor
+
+### Body Tracking
+- 24 body joints with pose, velocity, acceleration data
+- IMU timestamps for each joint
+- Availability flag for body tracking system
+
+## Dependencies
+
+### Required
+- pybind11 (Python binding framework)
+- CMake (build system)
+- XRoboToolkit-PC-Service SDK (automatically downloaded during setup)
+
+### Platform-specific Libraries
+- Linux: `libPXREARobotSDK.so`
+- Windows: `PXREARobotSDK.dll` and `PXREARobotSDK.lib`
+
+## Testing
+
+No formal test suite is included. Test functionality using the example scripts in `examples/`:
+- `example.py`: Basic controller and headset pose testing
+- `example_body_tracking.py`: Body tracking functionality
+- `run_binding_continuous.py`: Continuous data capture
+
+## Important Notes
+
+- The SDK requires active XR device connection (PICO headset)
+- Body tracking requires at least two Pico Swift devices
+- All data access is thread-safe but real-time dependent on device connectivity
+- The project builds a Python extension module that must be installed to site-packages

GR00T-WholeBodyControl/external_dependencies/XRoboToolkit-PC-Service-Pybind_X86_and_ARM64/CMakeLists.txt

@@ -0,0 +1,54 @@
+cmake_minimum_required(VERSION 3.14)
+
+project(MyPybind11Project LANGUAGES CXX)
+
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+
+cmake_host_system_information(RESULT ISA_NAME QUERY OS_PLATFORM) # Added: Important for UNIX specific logic
+message(STATUS "OS_PLATFORM (ISA_NAME): ${ISA_NAME}")
+
+include(GNUInstallDirs) # Add this line
+find_package(pybind11 REQUIRED)
+
+# Python Bindings for py_bindings.cpp
+pybind11_add_module(xrobotoolkit_sdk MODULE bindings/py_bindings.cpp)
+
+# Link xrobotoolkit_sdk module against pybind11
+target_link_libraries(xrobotoolkit_sdk PRIVATE pybind11::module)
+
+# Add include directories and link libraries for PXREARobotSDK to xrobotoolkit_sdk target
+if(WIN32)
+  target_include_directories(xrobotoolkit_sdk PUBLIC
+      ${PROJECT_SOURCE_DIR}/include
+  )
+  target_link_directories(xrobotoolkit_sdk PUBLIC ${PROJECT_SOURCE_DIR}/lib)
+  target_link_libraries(xrobotoolkit_sdk PUBLIC
+      PXREARobotSDK.dll # Assuming this is how PXREARobotSDK is linked, mirroring ConsoleDemo
+  )
+endif()
+
+if(UNIX)
+  # ISA_NAME is set by cmake_host_system_information above
+  if(ISA_NAME STREQUAL "aarch64")
+    target_include_directories(xrobotoolkit_sdk PUBLIC
+        ${PROJECT_SOURCE_DIR}/include/aarch64
+    )
+    target_link_directories(xrobotoolkit_sdk PUBLIC ${PROJECT_SOURCE_DIR}/lib/aarch64)
+  else()
+    target_include_directories(xrobotoolkit_sdk PUBLIC
+        ${PROJECT_SOURCE_DIR}/include
+    )
+    target_link_directories(xrobotoolkit_sdk PUBLIC ${PROJECT_SOURCE_DIR}/lib)
+  endif()
+  target_link_libraries(xrobotoolkit_sdk PUBLIC
+      PXREARobotSDK
+  )
+endif()
+
+# Install the Python module
+install(TARGETS xrobotoolkit_sdk
+    LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR} # Installs to <prefix>/lib
+                                                # You might want a Python-specific path like:
+                                                # DESTINATION lib/python${Python_VERSION_MAJOR}.${Python_VERSION_MINOR}/site-packages
+)

GR00T-WholeBodyControl/external_dependencies/XRoboToolkit-PC-Service-Pybind_X86_and_ARM64/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 XR Robotics
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

GR00T-WholeBodyControl/external_dependencies/XRoboToolkit-PC-Service-Pybind_X86_and_ARM64/README.md

@@ -0,0 +1,169 @@
+# XRoboToolkit-PC-Service-Pybind
+
+This project provides a python interface to extract XR state using XRoboToolkit-PC-Service sdk.
+
+## Requirements
+
+- [`pybind11`](https://github.com/pybind/pybind11)
+- [`XRoboRoolkit PC Service`](https://github.com/XR-Robotics/XRoboToolkit-PC-Service#)
+
+## Building the Project
+### Ubuntu 22.04
+
+```
+conda remove --name xr --all
+conda create -n xr python=3.10
+conda activate xr
+
+mkdir -p tmp
+cd tmp
+git clone https://github.com/XR-Robotics/XRoboToolkit-PC-Service.git
+cd XRoboToolkit-PC-Service/RoboticsService/PXREARobotSDK 
+bash build.sh
+cd ../../../..
+
+mkdir -p lib
+mkdir -p include
+cp tmp/XRoboToolkit-PC-Service/RoboticsService/PXREARobotSDK/PXREARobotSDK.h include/
+cp -r tmp/XRoboToolkit-PC-Service/RoboticsService/PXREARobotSDK/nlohmann include/nlohmann/
+cp tmp/XRoboToolkit-PC-Service/RoboticsService/PXREARobotSDK/build/libPXREARobotSDK.so lib/
+# rm -rf tmp
+
+# Build the project
+conda install -c conda-forge pybind11
+
+pip uninstall -y xrobotoolkit_sdk
+python setup.py install
+```
+### Linux Ubuntu 22.04 arm64 version (Nvidia orin supported)
+```
+bash setup_orin.sh
+```
+### Windows
+
+**Ensure pybind11 is installed before running the following command.**
+
+```
+setup_windows.bat
+```
+
+## Using the Python Bindings
+
+**1. Get Controller and Headset Poses**
+
+```python
+import xrobotoolkit_sdk as xrt
+
+xrt.init()
+
+left_pose = xrt.get_left_controller_pose()
+right_pose = xrt.get_right_controller_pose()
+headset_pose = xrt.get_headset_pose()
+
+print(f"Left Controller Pose: {left_pose}")
+print(f"Right Controller Pose: {right_pose}")
+print(f"Headset Pose: {headset_pose}")
+
+xrt.close()
+```
+
+**2. Get Controller Inputs (Triggers, Grips, Buttons, Axes)**
+
+```python
+import xrobotoolkit_sdk as xrt
+
+xrt.init()
+
+# Triggers and Grips
+left_trigger = xrt.get_left_trigger()
+right_grip = xrt.get_right_grip()
+print(f"Left Trigger: {left_trigger}, Right Grip: {right_grip}")
+
+# Buttons
+a_button_pressed = xrt.get_A_button()
+x_button_pressed = xrt.get_X_button()
+print(f"A Button Pressed: {a_button_pressed}, X Button Pressed: {x_button_pressed}")
+
+# Axes
+left_axis = xrt.get_left_axis()
+right_axis_click = xrt.get_right_axis_click()
+print(f"Left Axis: {left_axis}, Right Axis Clicked: {right_axis_click}")
+
+# Timestamp
+timestamp = xrt.get_time_stamp_ns()
+print(f"Current Timestamp (ns): {timestamp}")
+
+xrt.close()
+```
+
+**3. Get hand tracking state**
+```python
+import xrobotoolkit_sdk as xrt
+
+xrt.init()
+
+# Left Hand State
+left_hand_tracking_state = xrt.get_left_hand_tracking_state()
+print(f"Left Hand State: {left_hand_tracking_state}")
+
+# Left Hand isActive
+left_hand_is_active = xrt.get_left_hand_is_active()
+print(f"Left Hand isActive: {left_hand_is_active}")
+
+# Right Hand State
+right_hand_tracking_state = xrt.get_right_hand_tracking_state()
+print(f"Right Hand State: {right_hand_tracking_state}")
+
+# Right Hand isActive
+right_hand_is_active = xrt.get_right_hand_is_active()
+print(f"Right Hand isActive: {right_hand_is_active}")
+
+xrt.close()
+```
+
+**4. Get whole body motion tracking （please refer to this example when check Full Body tracking mode in UNITY app）**
+```python
+import xrobotoolkit_sdk as xrt
+
+xrt.init()
+
+# Check if body tracking data is available
+if xrt.is_body_data_available():
+    # Get body joint poses (24 joints, 7 values each: x,y,z,qx,qy,qz,qw)
+    body_poses = xrt.get_body_joints_pose()
+    print(f"Body joints pose data: {body_poses}")
+    
+    # Get body joint velocities (24 joints, 6 values each: vx,vy,vz,wx,wy,wz)
+    body_velocities = xrt.get_body_joints_velocity()
+    print(f"Body joints velocity data: {body_velocities}")
+    
+    # Get body joint accelerations (24 joints, 6 values each: ax,ay,az,wax,way,waz)
+    body_accelerations = xrt.get_body_joints_acceleration()
+    print(f"Body joints acceleration data: {body_accelerations}")
+    
+    # Get IMU timestamps for each joint
+    imu_timestamps = xrt.get_body_joints_timestamp()
+    print(f"IMU timestamps: {imu_timestamps}")
+    
+    # Get body data timestamp
+    body_timestamp = xrt.get_body_timestamp_ns()
+    print(f"Body data timestamp: {body_timestamp}")
+    
+    # Example: Get specific joint data (Head joint is index 15)
+    head_pose = body_poses[15]  # Head joint
+    x, y, z, qx, qy, qz, qw = head_pose
+    print(f"Head pose: Position({x:.3f}, {y:.3f}, {z:.3f}) Rotation({qx:.3f}, {qy:.3f}, {qz:.3f}, {qw:.3f})")
+else:
+    print("Body tracking data not available. Make sure:")
+    print("1. PICO headset is connected")
+    print("2. Body tracking is enabled in the control panel")
+    print("3. At least two Pico Swift devices are connected and calibrated")
+
+xrt.close()
+```
+
+**Body Joint Indices (similar to SMPL, 24 joints total):**
+- 0: Pelvis, 1: Left Hip, 2: Right Hip, 3: Spine1, 4: Left Knee, 5: Right Knee
+- 6: Spine2, 7: Left Ankle, 8: Right Ankle, 9: Spine3, 10: Left Foot, 11: Right Foot
+- 12: Neck, 13: Left Collar, 14: Right Collar, 15: Head, 16: Left Shoulder, 17: Right Shoulder
+- 18: Left Elbow, 19: Right Elbow, 20: Left Wrist, 21: Right Wrist, 22: Left Hand, 23: Right Hand

GR00T-WholeBodyControl/external_dependencies/XRoboToolkit-PC-Service-Pybind_X86_and_ARM64/setup.py

@@ -0,0 +1,148 @@
+import os
+import platform
+import re
+import shutil  # Added for shutil.rmtree
+import subprocess
+import sys
+from distutils.version import LooseVersion
+
+from setuptools import Command, Extension, find_packages, setup  # Added Command
+from setuptools.command.build_ext import build_ext
+
+
+class CMakeExtension(Extension):
+    def __init__(self, name, sourcedir=""):
+        Extension.__init__(self, name, sources=[])
+        self.sourcedir = os.path.abspath(sourcedir)
+
+
+class CMakeBuild(build_ext):
+    def run(self):
+        try:
+            out = subprocess.check_output(["cmake", "--version"])
+        except OSError:
+            raise RuntimeError(
+                "CMake must be installed to build the following extensions: "
+                + ", ".join(e.name for e in self.extensions)
+            )
+
+        if platform.system() == "Windows":
+            cmake_version = LooseVersion(re.search(r"version\s*([\d.]+)", out.decode()).group(1))
+            if cmake_version < "3.1.0":
+                raise RuntimeError("CMake >= 3.1.0 is required on Windows")
+
+        for ext in self.extensions:
+            self.build_extension(ext)
+
+    def build_extension(self, ext):
+        extdir = os.path.abspath(os.path.dirname(self.get_ext_fullpath(ext.name)))
+        # required for auto-detection of auxiliary "native" libs
+        if not extdir.endswith(os.path.sep):
+            extdir += os.path.sep
+
+        # Get pybind11 include paths
+        cmake_args = [
+            "-DCMAKE_LIBRARY_OUTPUT_DIRECTORY=" + extdir,
+            "-DPYTHON_EXECUTABLE=" + sys.executable,
+            "-DCMAKE_BUILD_TYPE=Release",
+        ]
+
+        cfg = "Debug" if self.debug else "Release"
+        build_args = ["--config", cfg]
+
+        if platform.system() == "Windows":
+            cmake_args += ["-DCMAKE_LIBRARY_OUTPUT_DIRECTORY_{}={}".format(cfg.upper(), extdir)]
+            if sys.maxsize > 2**32:
+                cmake_args += ["-A", "x64"]
+            build_args += ["--", "/m"]
+        else:
+            cmake_args += ["-DCMAKE_BUILD_TYPE=" + cfg]
+            build_args += ["--", "-j2"]  # Adjust core count as needed
+
+        env = os.environ.copy()
+        env["CXXFLAGS"] = '{} -DVERSION_INFO=\\"{}\\"'.format(
+            env.get("CXXFLAGS", ""), self.distribution.get_version()
+        )
+        if not os.path.exists(self.build_temp):
+            os.makedirs(self.build_temp)
+
+        subprocess.check_call(["cmake", ext.sourcedir] + cmake_args, cwd=self.build_temp, env=env)
+        subprocess.check_call(["cmake", "--build", "."] + build_args, cwd=self.build_temp)
+
+
+# New Clean Command
+class CleanCommand(Command):
+    """Custom clean command to tidy up the project root."""
+
+    user_options = []
+
+    def initialize_options(self):
+        pass
+
+    def finalize_options(self):
+        pass
+
+    def run(self):
+        # Remove build directory
+        if os.path.exists("build"):
+            print("Removing 'build/' directory")
+            shutil.rmtree("build")
+        # Remove .egg-info directory
+        for item in os.listdir("."):
+            if item.endswith(".egg-info"):
+                print(f"Removing '{item}' directory")
+                shutil.rmtree(item)
+        for item in os.listdir("."):
+            if item.endswith(".eggs"):
+                print(f"Removing '{item}' directory")
+                shutil.rmtree(item)
+        # Optionally, remove dist directory if you generate distributions
+        if os.path.exists("dist"):
+            print("Removing 'dist/' directory")
+            shutil.rmtree("dist")
+
+
+# New Uninstall Command
+class UninstallCommand(Command):
+    """Custom command to uninstall the package."""
+
+    user_options = []
+
+    def initialize_options(self):
+        pass
+
+    def finalize_options(self):
+        pass
+
+    def run(self):
+        package_name = self.distribution.get_name()
+        print(f"Attempting to uninstall {package_name}...")
+        try:
+            subprocess.check_call([sys.executable, "-m", "pip", "uninstall", "-y", package_name])
+            print(f"{package_name} uninstalled successfully.")
+        except subprocess.CalledProcessError as e:
+            print(
+                f"Failed to uninstall {package_name}. It may not be installed or pip uninstall failed."
+            )
+            print(f"Error: {e}")
+        except FileNotFoundError:
+            print("pip command not found. Please ensure pip is installed and in your PATH.")
+
+
+setup(
+    name="xrobotoolkit_sdk",
+    version="1.0.2",
+    author="Zhigen Zhao",
+    author_email="zhigen.zhao@bytedance.com",
+    description="A Python binding for XRobotoolkit PC Service SDK using pybind11 and CMake",
+    long_description="",  # Optionally, load from a README.md file
+    ext_modules=[CMakeExtension("xrobotoolkit_sdk")],
+    cmdclass=dict(
+        build_ext=CMakeBuild,
+        clean=CleanCommand,  # Add clean command
+        uninstall=UninstallCommand,  # Add uninstall command
+    ),
+    zip_safe=False,
+    python_requires=">=3.10",  # Specify your Python version requirement
+    packages=find_packages(),  # If you have other Python packages in your project
+)