duobench 0.1.0__py3-none-any.whl

duobench/__init__.py

@@ -0,0 +1,119 @@
+__version__ = "0.1.0"
+
+import os
+from pathlib import Path
+
+import numpy as np
+import rcs
+from rcs import get_prefix
+from rcs._core import common
+
+DUOBENCH_GITHUB_ASSET_ARCHIVE_URL = "https://github.com/RobotControlStack/duobench/archive/refs/tags/{tag}.zip"
+
+DUOBENCH_PREFIX = get_prefix(
+    os.environ.get("DUOBENCH_PREFIX"),
+    Path(__file__).resolve().parents[2],
+    Path.home() / ".duobench",
+    __version__,
+    DUOBENCH_GITHUB_ASSET_ARCHIVE_URL,
+)
+
+CAMERA_PATHS: dict[str, str] = {}
+GRIPPER_PATHS: dict[str, str] = {}
+OBJECT_PATHS: dict[str, str] = {
+    # transfer_reorient
+    "transfer_reorient_socket": "assets/objects/transfer_reorient/socket.xml",
+    "transfer_reorient_hex_peg": "assets/objects/transfer_reorient/hexpeg.xml",
+    "transfer_reorient_squere_peg": "assets/objects/transfer_reorient/squere_peg.xml",
+    "transfer_reorient_triangel_peg": "assets/objects/transfer_reorient/triangle_peg.xml",
+    # block_stacking
+    "block_stacking_big_red_cube": "assets/objects/block_balance/big_red_cube.xml",
+    "block_stacking_blue_rectangle": "assets/objects/block_balance/blue_rectangle.xml",
+    "block_stacking_green_rectangle": "assets/objects/block_balance/green_rectangle.xml",
+    "block_stacking_pink_beam": "assets/objects/block_balance/pink_beam.xml",
+    # bin_sort
+    "parallel_pick_white_bowl": "assets/objects/bin_sort/white_bowl.xml",
+    "parallel_pick_white_tri_cylinder": "assets/objects/bin_sort/white_tri_cylinder.xml",
+    "parallel_pick_white_box": "assets/objects/bin_sort/white_box.xml",
+    "parallel_pick_white_pent_cylinder": "assets/objects/bin_sort/white_pent_cylinder.xml",
+    "parallel_pick_white_hex_cylinder": "assets/objects/bin_sort/white_hex_cylinder.xml",
+    "parallel_pick_white_oct_cylinder": "assets/objects/bin_sort/white_oct_cylinder.xml",
+    "parallel_pick_white_dec_cylinder": "assets/objects/bin_sort/white_dec_cylinder.xml",
+    "parallel_pick_white_cylinder": "assets/objects/bin_sort/white_cylinder.xml",
+    "parallel_pick_black_bowl": "assets/objects/bin_sort/black_bowl.xml",
+    "parallel_pick_black_tri_cylinder": "assets/objects/bin_sort/black_tri_cylinder.xml",
+    "parallel_pick_black_box": "assets/objects/bin_sort/black_box.xml",
+    "parallel_pick_black_pent_cylinder": "assets/objects/bin_sort/black_pent_cylinder.xml",
+    "parallel_pick_black_hex_cylinder": "assets/objects/bin_sort/black_hex_cylinder.xml",
+    "parallel_pick_black_oct_cylinder": "assets/objects/bin_sort/black_oct_cylinder.xml",
+    "parallel_pick_black_dec_cylinder": "assets/objects/bin_sort/black_dec_cylinder.xml",
+    "parallel_pick_black_cylinder": "assets/objects/bin_sort/black_cylinder.xml",
+    "parallel_pick_blue_tri_cylinder": "assets/objects/bin_sort/blue_tri_cylinder.xml",
+    "parallel_pick_blue_box": "assets/objects/bin_sort/blue_box.xml",
+    "parallel_pick_blue_pent_cylinder": "assets/objects/bin_sort/blue_pent_cylinder.xml",
+    "parallel_pick_blue_hex_cylinder": "assets/objects/bin_sort/blue_hex_cylinder.xml",
+    "parallel_pick_blue_cylinder": "assets/objects/bin_sort/blue_cylinder.xml",
+    "parallel_pick_green_tri_cylinder": "assets/objects/bin_sort/green_tri_cylinder.xml",
+    "parallel_pick_green_box": "assets/objects/bin_sort/green_box.xml",
+    "parallel_pick_green_pent_cylinder": "assets/objects/bin_sort/green_pent_cylinder.xml",
+    "parallel_pick_green_hex_cylinder": "assets/objects/bin_sort/green_hex_cylinder.xml",
+    "parallel_pick_green_cylinder": "assets/objects/bin_sort/green_cylinder.xml",
+    # spring_door
+    "spring_door_gray_microwave": "assets/objects/spring_door/gray_microwave.xml",
+    "spring_door_mug": "assets/objects/spring_door/mug.xml",
+    "spring_door_simple_mug": "assets/objects/spring_door/simple_mug.xml",
+    # hinge chest
+    "hinge_chest_chest": "assets/objects/hinge_chest/chest.xml",
+    # transfer_gate
+    "handover_hole_stand": "assets/objects/transfer_gate/stand.xml",
+    "handover_hole_mat": "assets/objects/transfer_gate/mat.xml",
+    "handover_hole_white_long_box": "assets/objects/transfer_gate/white_long_box.xml",
+    # carry pot
+    "carry_pot_pot": "assets/objects/carry_pot/pot.xml",
+    "carry_pot_stove": "assets/objects/carry_pot/stove.xml",
+    # join_blocks
+    "h_block": "assets/objects/join_blocks/join_blocks_h_block.xml",
+    "p_block": "assets/objects/join_blocks/join_blocks_p_block.xml",
+    "wall": "assets/objects/join_blocks/join_blocks_wall.xml",
+    # pour_marbles
+    "marble": "assets/objects/pour_marbles/marble.xml",
+    "teacup": "assets/objects/pour_marbles/teacup.xml",
+    # ball_maze
+    "maze1": "assets/objects/ball_maze/maze_simple_1.xml",
+    "maze2": "assets/objects/ball_maze/maze_simple_2.xml",
+    "maze3": "assets/objects/ball_maze/maze_simple_3.xml",
+    "maze4": "assets/objects/ball_maze/maze_simple_4.xml",
+    "maze5": "assets/objects/ball_maze/maze_simple_5.xml",
+    "maze6": "assets/objects/ball_maze/maze_simple_6.xml",
+    "maze7": "assets/objects/ball_maze/maze_simple_7.xml",
+    "maze8": "assets/objects/ball_maze/maze_simple_8.xml",
+    "maze9": "assets/objects/ball_maze/maze_simple_9.xml",
+    "maze10": "assets/objects/ball_maze/maze_simple_10.xml",
+}
+
+
+SCENE_PATHS = {
+    "empty_world": "assets/scenes/base_world.xml",
+    "vention_world": "assets/scenes/vention_table/vention_world.xml",
+}
+
+# Append DUOBENCH package prefix to all asset paths
+for path_dict in (GRIPPER_PATHS, SCENE_PATHS, OBJECT_PATHS, CAMERA_PATHS):
+    for name, path in path_dict.items():
+        abs_path = os.path.join(DUOBENCH_PREFIX, path)
+        if not os.path.isfile(abs_path):
+            error_msg = f"Asset {name} not found at path: {abs_path}. Please make sure to download the assets."
+            raise FileNotFoundError(error_msg)
+        else:
+            path_dict[name] = abs_path
+
+# Update the global RCS registries with our assets and tasks
+rcs.SCENE_PATHS.update(SCENE_PATHS)
+rcs.OBJECT_PATHS.update(OBJECT_PATHS)
+rcs.GRIPPER_PATHS.update(GRIPPER_PATHS)
+rcs.CAMERA_PATHS.update(CAMERA_PATHS)
+
+
+rcs.DEFAULT_TRANSFORMS.update(
+    {"VENTION_HEIGHT_OFFSET": common.Pose(translation=np.array([0, 0, 0.8095]), quaternion=np.array([0, 0, 0, 1]))}
+)

duobench/__main__.py

@@ -0,0 +1,20 @@
+import pkgutil
+from importlib import import_module
+
+import duobench.tasks
+
+
+def _register_duobench_tasks() -> None:
+    for module in pkgutil.iter_modules(duobench.tasks.__path__):
+        import_module(f"{duobench.tasks.__name__}.{module.name}")
+
+
+def main() -> None:
+    # Import the base RCS CLI only after DuoBench tasks are registered so the
+    # existing commands can resolve `duobench/*` env ids.
+    _register_duobench_tasks()
+    import_module("rcs.__main__").app()
+
+
+if __name__ == "__main__":
+    main()

duobench/tasks/__init__.py

@@ -0,0 +1,81 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any
+
+import gymnasium as gym
+from rcs import sim as simulation
+
+
+@dataclass(kw_only=True)
+class TaskStage(ABC):
+    stage: int = 0
+    max_stage: int = 0
+    internal_state: dict[str, bool] = field(default_factory=dict)
+    stage_to_subinstructions: dict[int, str] = field(default_factory=dict)
+    instruction: str = ""
+
+    @abstractmethod
+    def update_stage(self):
+        msg = "TaskStage is an abstract class. Please implement the update_stage method in a subclass."
+        raise NotImplementedError(msg)
+
+    @abstractmethod
+    def update_internal_state(self, sim: simulation.Sim):
+        msg = "TaskStage is an abstract class. Please implement the update_internal_state method in a subclass."
+        raise NotImplementedError(msg)
+
+    def reset(self):
+        self.stage = 0
+        for key in self.internal_state:
+            self.internal_state[key] = False
+
+    @property
+    def current_subinstruction(self) -> str:
+        return self.stage_to_subinstructions[self.stage] if self.stage_to_subinstructions is not None else None
+
+    @property
+    def success(self) -> bool:
+        return self.stage == self.max_stage
+
+    @property
+    def normalized_stage(self) -> float:
+        return self.stage / self.max_stage
+
+    @property
+    def info(self) -> dict[str, Any]:
+        return {
+            "success": self.success,
+            "stage": self.stage,
+            "max_stage": self.max_stage,
+            "current_subinstruction": self.current_subinstruction,
+            "stage_to_subinstructions": {
+                str(stage): instruction for stage, instruction in self.stage_to_subinstructions.items()
+            },
+            "instruction": self.instruction,
+        }
+
+
+class TaskStageWrapper(gym.Wrapper):
+
+    def __init__(
+        self,
+        env: gym.Env,
+        stage_tracker: TaskStage,
+    ):
+        super().__init__(env)
+        self.stage_tracker = stage_tracker
+
+    def step(self, action: dict[str, Any]):
+        obs, _, _, truncated, info = super().step(action)
+        self.stage_tracker.update_internal_state(self.get_wrapper_attr("sim"))
+
+        info.update(self.stage_tracker.info)
+        reward = self.stage_tracker.normalized_stage
+
+        return obs, reward, self.stage_tracker.success, truncated, info
+
+    def reset(self, *, seed: int | None = None, options: dict[str, Any] | None = None):
+        self.stage_tracker.reset()
+        obs, info = super().reset(seed=seed, options=options)
+        info.update(self.stage_tracker.info)
+        return obs, info

duobench/tasks/ball_maze.py

@@ -0,0 +1,251 @@
+from dataclasses import dataclass, field
+from typing import Any
+
+import gymnasium as gym
+import mujoco as mj
+import numpy as np
+import rcs
+from rcs.envs.scenes import BaseTaskConfig, SimEnvCreatorConfig, Task
+from rcs.sim.composer import ModelComposer
+from rcs.sim.sim import Sim
+
+from duobench.tasks import TaskStage, TaskStageWrapper
+from duobench.utils.helper_wrappers import (
+    RandomSquareObjsPos,
+    get_bodies_in_contact_with_gripper_pad,
+)
+from duobench.utils.vention_config import VentionSceneFR3Duo
+
+
+@dataclass(kw_only=True)
+class BallMazeTaskConfig(BaseTaskConfig):
+
+    task_id: str = "ball_maze"
+
+    board_number = 2
+    object_body = "board"
+    include_rotation: bool = True
+    task_instructions = "pick up the board and tilt it so the ball roles onto the red square"
+    hard_reset = False
+    x_width = 0.3
+    y_width = 0.3
+    z_init = 0.02
+    obj_position_margin = 0.08
+    object_joint_names: list[str] = field(default_factory=lambda: ["board_board_joint"])
+    object2root_frame: rcs.common.Pose = field(
+        default_factory=lambda: rcs.common.Pose(translation=np.array([0.5, 0, 0.05]), quaternion=np.array([0, 0, 0, 1]))
+    )
+    object_center_to_root_frame: rcs.common.Pose = field(
+        default_factory=lambda: rcs.common.Pose(
+            # 0.02 = z_init
+            translation=np.array([0.5, 0.0, 0.02]),
+            quaternion=np.array([0, 0, 0, 1]),
+        )
+    )
+
+
+class BallMazeTask(Task[BallMazeTaskConfig]):
+
+    @staticmethod
+    def add_task_mujoco(cfg: BallMazeTaskConfig, composer: ModelComposer, env_cfg: SimEnvCreatorConfig):
+        """Add task-specific elements to the Mujoco scene."""
+        object2world = env_cfg.root_frame_to_world * cfg.object_center_to_root_frame
+        key = "maze{number}"
+        board_xml = rcs.OBJECT_PATHS[key.format(number=cfg.board_number)].format(number=cfg.board_number)
+
+        composer.add_object_world_frame(
+            board_xml,
+            object_prefix=cfg.object_body + "_",
+            pose=object2world,
+            register_root_relative_replay_free_joints=True,
+        )
+
+    @staticmethod
+    def add_task_env(cfg: BallMazeTaskConfig, env: gym.Env, simulation: Sim, env_cfg: SimEnvCreatorConfig) -> gym.Env:
+        """Add task-specific wrappers to the environment."""
+        _ = simulation
+        object2world = env_cfg.root_frame_to_world * cfg.object_center_to_root_frame
+
+        env = RandomSquareObjsPos(
+            env=env,
+            x_width=cfg.x_width,
+            y_width=cfg.y_width,
+            z_init=cfg.z_init,
+            center2world=object2world,
+            include_rotation=cfg.include_rotation,
+            obj_joint_names=cfg.object_joint_names,
+            obj_position_margin=cfg.obj_position_margin,
+        )
+
+        return BallMazeTaskWrapper(env, BallMazeStage(cfg))
+
+
+class BallMazeStage(TaskStage):
+    def __init__(self, cfg: BallMazeTaskConfig):
+        super().__init__(
+            stage=0,
+            max_stage=4,
+            internal_state={
+                "left_arm_contact": False,
+                "right_arm_contact": False,
+                "both_arms_contact": False,
+                "maze_lifted": False,
+                "ball_left_start": False,
+                "ball_at_goal": False,
+                "ball_near_goal": False,
+                "maze_on_table": True,
+                "hands_released": False,
+            },
+            stage_to_subinstructions={
+                0: "make contact with the maze",
+                1: "grasp the maze with both arms",
+                2: "lift the maze and move the ball out of the start area",
+                3: "guide the ball from the start area toward the goal",
+                4: "task completed; the ball has reached the goal",
+            },
+            instruction=cfg.task_instructions,
+        )
+        self.cfg = cfg
+        self.start_marker = "board_marker_start"
+        self.goal_marker = "board_marker_goal"
+        self.ball_joint = "board_ball_joint"
+        self.ball_name = "board_ball"
+        self.board_name = "board_board"
+        self.start_dist = 0.02
+        self.goal_dist = 0.015
+        self.done_goal_radius = 0.10
+        self.table_height: float | None = None
+
+    @staticmethod
+    def _planar_distance(a: np.ndarray, b: np.ndarray) -> float:
+        return float(np.linalg.norm(a[:2] - b[:2]))
+
+    def set_table_height(self, table_height: float):
+        self.table_height = table_height
+
+    def _board_body_ids(self, sim: Sim) -> set[int]:
+        body_ids: set[int] = set()
+        for body_id in range(sim.model.nbody):
+            body_name = sim.model.body(body_id).name
+            if body_name and body_name.startswith("board_") and body_name != self.ball_name:
+                body_ids.add(body_id)
+        return body_ids
+
+    def _maze_on_table(self, sim: Sim, board_body_ids: set[int], board_height: float) -> bool:
+        for i in range(sim.data.ncon):
+            contact = sim.data.contact[i]
+            geom1_name = sim.model.geom(contact.geom1).name
+            geom2_name = sim.model.geom(contact.geom2).name
+            body1_id = sim.model.geom_bodyid[contact.geom1]
+            body2_id = sim.model.geom_bodyid[contact.geom2]
+
+            if geom1_name == "floor" and body2_id in board_body_ids:
+                return True
+            if geom2_name == "floor" and body1_id in board_body_ids:
+                return True
+
+        if self.table_height is None:
+            return False
+
+        return board_height <= self.table_height + 0.02
+
+    def update_internal_state(self, sim: Sim):
+        board_body_id = mj.mj_name2id(sim.model, mj.mjtObj.mjOBJ_BODY, self.board_name)
+        board_body_ids = self._board_body_ids(sim)
+        left_arm_contacts = get_bodies_in_contact_with_gripper_pad(sim, "left")
+        right_arm_contacts = get_bodies_in_contact_with_gripper_pad(sim, "right")
+
+        left_arm_contact = any(body_id in board_body_ids for body_id in left_arm_contacts)
+        right_arm_contact = any(body_id in board_body_ids for body_id in right_arm_contacts)
+        both_arms_contact = left_arm_contact and right_arm_contact
+        any_arm_contact = left_arm_contact or right_arm_contact
+
+        board_height = sim.data.xpos[board_body_id][2]
+        maze_on_table = self._maze_on_table(sim, board_body_ids, board_height)
+
+        ball_pose = sim.data.body(self.ball_name).xpos
+        start_pose = sim.data.body(self.start_marker).xpos
+        goal_pose = sim.data.body(self.goal_marker).xpos
+        # The ball center is always above the board plane, so using full 3D distance
+        # can incorrectly mark "left start" immediately after reset.
+        start_dist = self._planar_distance(start_pose, ball_pose)
+        goal_dist = np.linalg.norm(goal_pose - ball_pose)
+
+        self.internal_state["left_arm_contact"] = left_arm_contact
+        self.internal_state["right_arm_contact"] = right_arm_contact
+        self.internal_state["both_arms_contact"] = both_arms_contact
+        self.internal_state["maze_lifted"] = both_arms_contact and not maze_on_table
+        self.internal_state["ball_left_start"] = bool(
+            self.internal_state["ball_left_start"] or start_dist >= self.start_dist
+        )
+        self.internal_state["ball_at_goal"] = bool(self.internal_state["ball_at_goal"] or goal_dist <= self.goal_dist)
+        self.internal_state["ball_near_goal"] = bool(goal_dist <= self.done_goal_radius)
+        self.internal_state["maze_on_table"] = maze_on_table
+        self.internal_state["hands_released"] = not any_arm_contact
+
+        self.update_stage()
+
+    def update_stage(self):
+        if self.internal_state["ball_at_goal"]:
+            self.stage = 4
+        elif self.internal_state["ball_left_start"]:
+            self.stage = 3
+        elif self.internal_state["maze_lifted"]:
+            self.stage = 2
+        elif self.internal_state["left_arm_contact"] or self.internal_state["right_arm_contact"]:
+            self.stage = 1
+        else:
+            self.stage = 0
+
+
+class BallMazeTaskWrapper(TaskStageWrapper):
+    """
+    Wrapper to conduct the task-specific soft reset for the ball maze while using the shared task-stage interface.
+    """
+
+    def __init__(self, env: gym.Env, stage_tracker: BallMazeStage):
+        super().__init__(env, stage_tracker)
+        self.stage_tracker: BallMazeStage = stage_tracker
+        self.sim = self.env.get_wrapper_attr("sim")
+        self.start_marker = stage_tracker.start_marker
+        self.ball_joint = stage_tracker.ball_joint
+
+    def reset(self, *, seed: int | None = None, options: dict[str, Any] | None = None):
+        obs, info = super().reset(seed=seed, options=options)
+
+        # The board pose is randomized by inner wrappers during reset, so place the
+        # ball only after the full reset chain has finished. Spawn the ball tangent
+        # to the maze floor and let the physics settle before the episode starts.
+        pos = self.sim.data.body(self.start_marker).xpos.copy()
+        ball_radius = 0.01
+        spawn_clearance = 0.001
+        pos[2] += ball_radius + spawn_clearance
+
+        joint_id = self.sim.model.joint(self.ball_joint).id
+        qpos_adr = self.sim.model.jnt_qposadr[joint_id]
+        qvel_adr = self.sim.model.jnt_dofadr[joint_id]
+        self.sim.data.qpos[qpos_adr : qpos_adr + 3] = pos
+        # Reset the free joint velocity so the ball does not inherit motion from the previous rollout.
+        self.sim.data.qvel[qvel_adr : qvel_adr + 6] = 0
+
+        mj.mj_forward(self.sim.model, self.sim.data)
+        for _ in range(10):
+            mj.mj_step(self.sim.model, self.sim.data)
+        self.stage_tracker.set_table_height(float(self.sim.data.body(self.stage_tracker.board_name).xpos[2]))
+        self.stage_tracker.update_internal_state(self.sim)
+        info.update(self.stage_tracker.info)
+        return obs, info
+
+
+rcs.TASKS.update({"ball_maze": BallMazeTask})
+
+
+class BallMazeEnvConfig(VentionSceneFR3Duo):
+
+    def config(self) -> SimEnvCreatorConfig:
+        cfg = super().config()
+        cfg.task_cfg = BallMazeTaskConfig()
+        return cfg
+
+
+gym.register(id="duobench/ball_maze", entry_point=BallMazeEnvConfig())