From 39c0b125c56fa25b9d7c91e44f4b3d6450604367 Mon Sep 17 00:00:00 2001 From: Steven Zhang Date: Fri, 31 Jul 2026 14:34:03 -0400 Subject: [PATCH] feat: add six rigid labware checkpoint environments --- scripts/run_workflow.py | 9 +- .../matterix/envs/matterix_base_env_cfg.py | 2 +- source/matterix_assets/data | 2 +- .../rigid_labware_batch1_local_only.py | 78 +++++++++ .../matterix_tasks/test_dev_tasks/__init__.py | 67 +++++++- .../capped_labware_checkpoints.py | 137 ++++++++++++++++ .../test_franka_rigid_labware_duran_100.py | 20 +++ .../test_franka_rigid_labware_duran_500.py | 20 +++ .../test_franka_rigid_labware_falcon_15.py | 20 +++ .../test_franka_rigid_labware_falcon_50.py | 20 +++ .../test_franka_rigid_labware_flask_250.py | 105 +++++++++++++ .../test_franka_rigid_labware_flask_50.py | 105 +++++++++++++ .../test_franka_rigid_labware_flasks.py | 148 ++++++++++++++++++ 13 files changed, 729 insertions(+), 4 deletions(-) create mode 100644 source/matterix_assets/matterix_assets/labware/rigid_labware_batch1_local_only.py create mode 100644 source/matterix_tasks/matterix_tasks/test_dev_tasks/capped_labware_checkpoints.py create mode 100644 source/matterix_tasks/matterix_tasks/test_dev_tasks/test_franka_rigid_labware_duran_100.py create mode 100644 source/matterix_tasks/matterix_tasks/test_dev_tasks/test_franka_rigid_labware_duran_500.py create mode 100644 source/matterix_tasks/matterix_tasks/test_dev_tasks/test_franka_rigid_labware_falcon_15.py create mode 100644 source/matterix_tasks/matterix_tasks/test_dev_tasks/test_franka_rigid_labware_falcon_50.py create mode 100644 source/matterix_tasks/matterix_tasks/test_dev_tasks/test_franka_rigid_labware_flask_250.py create mode 100644 source/matterix_tasks/matterix_tasks/test_dev_tasks/test_franka_rigid_labware_flask_50.py create mode 100644 source/matterix_tasks/matterix_tasks/test_dev_tasks/test_franka_rigid_labware_flasks.py diff --git a/scripts/run_workflow.py b/scripts/run_workflow.py index 6ad9476..58bf78c 100644 --- a/scripts/run_workflow.py +++ b/scripts/run_workflow.py @@ -40,11 +40,14 @@ help="Environment/task name.", ) parser.add_argument("--workflow", type=str, default="pickup_beaker", help="Name of the workflow to run.") +parser.add_argument("--record_path", type=str, default=None, help="Optional unique HDF5 recorder path for this run.") +parser.add_argument("--episodes", type=int, default=0, help="Stop after this many episodes; 0 keeps the existing continuous behavior.") AppLauncher.add_app_launcher_args(parser) args_cli = parser.parse_args() # Launch omniverse app -app_launcher = AppLauncher(headless=args_cli.headless) +# Forward the complete parsed launcher configuration so --livestream reaches WebRTC. +app_launcher = AppLauncher(args_cli) simulation_app = app_launcher.app """Rest everything else.""" @@ -66,6 +69,8 @@ def main(): num_envs=args_cli.num_envs, use_fabric=not args_cli.disable_fabric, ) + if args_cli.record_path is not None: + env_cfg.record_path = args_cli.record_path # Validate workflow exists if not hasattr(env_cfg, "workflows") or not env_cfg.workflows: @@ -133,6 +138,8 @@ def main(): sm.print_status(step=step_count, episode=episode_count) sm.print_status(step=step_count, episode=episode_count) + if args_cli.episodes > 0 and episode_count >= args_cli.episodes: + break env.close() diff --git a/source/matterix/matterix/envs/matterix_base_env_cfg.py b/source/matterix/matterix/envs/matterix_base_env_cfg.py index 9977693..d9c6100 100644 --- a/source/matterix/matterix/envs/matterix_base_env_cfg.py +++ b/source/matterix/matterix/envs/matterix_base_env_cfg.py @@ -63,7 +63,7 @@ class MatterixBaseEnvCfg: sim: SimulationCfg = SimulationCfg( render=RenderCfg( - carb_settings={"rtx_translucency_enabled": True, "rtx_raytracing_fractionalCutoutOpacity": True} + carb_settings={"rtx_translucency_enabled": True} ) ) """Physics simulation configuration. Default is SimulationCfg().""" diff --git a/source/matterix_assets/data b/source/matterix_assets/data index 0d856a0..ebf0f71 160000 --- a/source/matterix_assets/data +++ b/source/matterix_assets/data @@ -1 +1 @@ -Subproject commit 0d856a0572d3e0823204264fd3d2700e15a43f4b +Subproject commit ebf0f7165ace0c3ae2abf56bc774f8eb925c3cc0 diff --git a/source/matterix_assets/matterix_assets/labware/rigid_labware_batch1_local_only.py b/source/matterix_assets/matterix_assets/labware/rigid_labware_batch1_local_only.py new file mode 100644 index 0000000..07f0eff --- /dev/null +++ b/source/matterix_assets/matterix_assets/labware/rigid_labware_batch1_local_only.py @@ -0,0 +1,78 @@ +# Copyright (c) 2022-2026, The Matterix Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""LOCAL-ONLY configs for the six Batch 1 labware checkpoint candidates. + +These configs are test-environment plumbing. The reusable USD payloads live in +the Matterix_assets data submodule; the capped payloads are spawned as generic +USD scene assets by their dedicated checkpoint environments because each has +two rigid bodies joined by a fixed joint. +""" + +import json +import os + +from matterix_assets import MATTERIX_ASSETS_DATA_DIR +from matterix.managers.semantics.primitive_semantics import IsInContactPhysicsCfg + +from isaaclab.utils import configclass + +from ..matterix_rigid_object import MatterixRigidObjectCfg + + +default_prim_path = "{ENV_REGEX_NS}/RigidObjects_Labware" + +CORNING_4980_50_MASS_KG = 0.0305 +CORNING_4980_250_MASS_KG = 0.1122 + + +def _load_frames(slug: str) -> dict[str, tuple[float, float, float]]: + """Load authored interface frames from the staged payload.""" + path = os.path.join(MATTERIX_ASSETS_DATA_DIR, "labware", slug, "frames.json") + with open(path) as handle: + contract = json.load(handle) + frames = contract.get("frames") + if not isinstance(frames, dict) or not frames: + raise ValueError(f"{path} carries no frames") + missing = {"grasp", "pre_grasp", "post_grasp"} - set(frames) + if missing: + raise ValueError(f"{path} is missing required frames: {sorted(missing)}") + return {name: tuple(offset) for name, offset in frames.items()} + + +@configclass +class CORNING_4980_50_LOCAL_ONLY_CFG(MatterixRigidObjectCfg): + """Corning PYREX 4980-50, 50 mL narrow-mouth Erlenmeyer flask.""" + + prim_path = default_prim_path + usd_path = f"{MATTERIX_ASSETS_DATA_DIR}/labware/corning-4980-50/corning-4980-50-inst.usda" + scale = (1.0, 1.0, 1.0) + mass = CORNING_4980_50_MASS_KG + activate_contact_sensors = True + frames = _load_frames("corning-4980-50") + semantic_tags = [("class", "flask")] + semantics = [ + IsInContactPhysicsCfg( + filter_prim_paths_expr=["robot/panda_leftfinger", "robot/panda_rightfinger"] + ) + ] + + +@configclass +class CORNING_4980_250_LOCAL_ONLY_CFG(MatterixRigidObjectCfg): + """Corning PYREX 4980-250, 250 mL narrow-mouth Erlenmeyer flask.""" + + prim_path = default_prim_path + usd_path = f"{MATTERIX_ASSETS_DATA_DIR}/labware/corning-4980-250/corning-4980-250-inst.usda" + scale = (1.0, 1.0, 1.0) + mass = CORNING_4980_250_MASS_KG + activate_contact_sensors = True + frames = _load_frames("corning-4980-250") + semantic_tags = [("class", "flask")] + semantics = [ + IsInContactPhysicsCfg( + filter_prim_paths_expr=["robot/panda_leftfinger", "robot/panda_rightfinger"] + ) + ] diff --git a/source/matterix_tasks/matterix_tasks/test_dev_tasks/__init__.py b/source/matterix_tasks/matterix_tasks/test_dev_tasks/__init__.py index e21147c..93078e7 100644 --- a/source/matterix_tasks/matterix_tasks/test_dev_tasks/__init__.py +++ b/source/matterix_tasks/matterix_tasks/test_dev_tasks/__init__.py @@ -6,7 +6,18 @@ import gymnasium as gym import os -from . import test_franka_beaker_lift, test_franka_beakers, test_particle_systems, test_semantics_heat_transfer +from . import ( + test_franka_beaker_lift, + test_franka_beakers, + test_franka_rigid_labware_duran_100, + test_franka_rigid_labware_duran_500, + test_franka_rigid_labware_falcon_15, + test_franka_rigid_labware_falcon_50, + test_franka_rigid_labware_flask_50, + test_franka_rigid_labware_flask_250, + test_particle_systems, + test_semantics_heat_transfer, +) ## # Register Gym environments. @@ -47,3 +58,57 @@ }, disable_env_checker=True, ) + +gym.register( + id="Matterix-Test-Rigid-Labware-Flask-50-Franka-v1", + entry_point="matterix.envs:MatterixBaseEnv", + kwargs={ + "env_cfg_entry_point": test_franka_rigid_labware_flask_50.FrankaRigidLabwareFlask50EnvTestCfg, + }, + disable_env_checker=True, +) + +gym.register( + id="Matterix-Test-Rigid-Labware-Flask-250-Franka-v1", + entry_point="matterix.envs:MatterixBaseEnv", + kwargs={ + "env_cfg_entry_point": test_franka_rigid_labware_flask_250.FrankaRigidLabwareFlask250EnvTestCfg, + }, + disable_env_checker=True, +) + +gym.register( + id="Matterix-Test-Rigid-Labware-Duran-100-Franka-v1", + entry_point="matterix.envs:MatterixBaseEnv", + kwargs={ + "env_cfg_entry_point": test_franka_rigid_labware_duran_100.FrankaRigidLabwareDuran100EnvTestCfg, + }, + disable_env_checker=True, +) + +gym.register( + id="Matterix-Test-Rigid-Labware-Duran-500-Franka-v1", + entry_point="matterix.envs:MatterixBaseEnv", + kwargs={ + "env_cfg_entry_point": test_franka_rigid_labware_duran_500.FrankaRigidLabwareDuran500EnvTestCfg, + }, + disable_env_checker=True, +) + +gym.register( + id="Matterix-Test-Rigid-Labware-Falcon-15-Franka-v1", + entry_point="matterix.envs:MatterixBaseEnv", + kwargs={ + "env_cfg_entry_point": test_franka_rigid_labware_falcon_15.FrankaRigidLabwareFalcon15EnvTestCfg, + }, + disable_env_checker=True, +) + +gym.register( + id="Matterix-Test-Rigid-Labware-Falcon-50-Franka-v1", + entry_point="matterix.envs:MatterixBaseEnv", + kwargs={ + "env_cfg_entry_point": test_franka_rigid_labware_falcon_50.FrankaRigidLabwareFalcon50EnvTestCfg, + }, + disable_env_checker=True, +) diff --git a/source/matterix_tasks/matterix_tasks/test_dev_tasks/capped_labware_checkpoints.py b/source/matterix_tasks/matterix_tasks/test_dev_tasks/capped_labware_checkpoints.py new file mode 100644 index 0000000..3afa42d --- /dev/null +++ b/source/matterix_tasks/matterix_tasks/test_dev_tasks/capped_labware_checkpoints.py @@ -0,0 +1,137 @@ +"""Shared helpers for capped DURAN and Falcon visual checkpoints. + +The capped payloads are intentionally spawned as generic USD assets. Each USD +contains two rigid bodies joined by a ``PhysicsFixedJoint`` and does not expose +an articulation root, so it must not be registered as a ``RigidObject``. +""" + +import torch + +from matterix.envs import MatterixBaseEnvCfg, mdp +from matterix_assets import MATTERIX_ASSETS_DATA_DIR, MatterixStaticObjectCfg +from matterix_assets.infrastructure.tables import TABLE_SEATTLE_INST_Cfg +from matterix_assets.robots import FRANKA_PANDA_HIGH_PD_IK_CFG + +from matterix.managers import EventManagerCfg +from matterix_sm import CloseGripperCfg, MoveRelativeCfg, OpenGripperCfg +from matterix_sm.primitive_actions.move_to_pose import MoveToPoseCfg +from matterix_sm.robot_action_spaces import FRANKA_IK_ACTION_SPACE + +import isaaclab.envs.mdp as isaaclab_mdp +from isaaclab.managers import EventTermCfg as EventTerm +from isaaclab.managers import ObservationGroupCfg as ObsGroup +from isaaclab.managers import ObservationTermCfg as ObsTerm +from isaaclab.utils import configclass + + +@configclass +class EventCfg(EventManagerCfg): + """Reset events for a capped labware checkpoint.""" + + reset_scene_to_default = EventTerm( + func=isaaclab_mdp.reset_scene_to_default, + mode="reset", + ) + + +@configclass +class ObservationManagerCfg: + """Robot observations used by the hard-coded visual manipulation sequence.""" + + @configclass + class ArticulationsGroup(ObsGroup): + robot__root_world_pos = ObsTerm(func=mdp.root_world_pos, params={"asset_name": "robot"}) + robot__root_world_quat = ObsTerm(func=mdp.root_world_quat, params={"asset_name": "robot"}) + robot__joint_pos = ObsTerm(func=mdp.joint_pos, params={"asset_name": "robot"}) + robot__joint_vel = ObsTerm(func=mdp.joint_vel, params={"asset_name": "robot"}) + robot__ee_world_pos = ObsTerm(func=mdp.ee_world_pos, params={"asset_name": "robot"}) + robot__ee_world_quat = ObsTerm(func=mdp.ee_world_quat, params={"asset_name": "robot"}) + robot__gripper_pos = ObsTerm(func=mdp.gripper_pos, params={"asset_name": "robot"}) + + def __post_init__(self): + self.enable_corruption = False + self.concatenate_terms = False + + articulations: ArticulationsGroup = ArticulationsGroup() + + +def capped_payload_cfg(slug: str, pos: tuple[float, float, float]) -> MatterixStaticObjectCfg: + """Build a generic scene asset for one two-body capped USD payload.""" + return MatterixStaticObjectCfg( + usd_path=f"{MATTERIX_ASSETS_DATA_DIR}/labware/{slug}/{slug}-inst.usda", + pos=pos, + scale=(1.0, 1.0, 1.0), + ) + + +def capped_pick_and_place( + vessel_pos: tuple[float, float, float], + pre_grasp_z: float, + grasp_z: float, + agent: str = "robot", +): + """Return the visual checkpoint sequence for one fixed-jointed payload.""" + x, y, z = vessel_pos + return [ + OpenGripperCfg(agent_assets=agent, action_space_info=FRANKA_IK_ACTION_SPACE), + MoveToPoseCfg( + agent_assets=agent, + target_positions=torch.tensor([[x, y, z + pre_grasp_z]]), + action_space_info=FRANKA_IK_ACTION_SPACE, + ), + MoveToPoseCfg( + agent_assets=agent, + target_positions=torch.tensor([[x, y, z + grasp_z]]), + action_space_info=FRANKA_IK_ACTION_SPACE, + ), + CloseGripperCfg(agent_assets=agent, action_space_info=FRANKA_IK_ACTION_SPACE), + MoveRelativeCfg( + agent_assets=agent, + position_offset=(0.0, 0.0, 0.1), + orientation_offset=None, + action_space_info=FRANKA_IK_ACTION_SPACE, + ), + MoveToPoseCfg( + agent_assets=agent, + target_positions=torch.tensor([[x, y, z + grasp_z + 0.005]]), + action_space_info=FRANKA_IK_ACTION_SPACE, + ), + OpenGripperCfg(agent_assets=agent, action_space_info=FRANKA_IK_ACTION_SPACE), + MoveRelativeCfg( + agent_assets=agent, + position_offset=(0.0, 0.0, 0.15), + orientation_offset=None, + action_space_info=FRANKA_IK_ACTION_SPACE, + ), + ] + + +def capped_env_fields( + slug: str, + vessel_pos: tuple[float, float, float], + pre_grasp_z: float, + grasp_z: float, + workflow_name: str, + description: str, +) -> dict: + """Return common config fields for a one-asset capped checkpoint.""" + return { + "env_spacing": 10.0, + "objects": { + "capped_labware": capped_payload_cfg(slug, vessel_pos), + "table": TABLE_SEATTLE_INST_Cfg(pos=(0.5, 0, 0)), + }, + "articulated_assets": { + "robot": FRANKA_PANDA_HIGH_PD_IK_CFG(pos=(0.0, 0, 0)), + }, + "gripper_joint_names": ["panda_finger_joint1", "panda_finger_joint2"], + "observations": ObservationManagerCfg(), + "events": EventCfg(), + "record_path": "datasets/dataset.hdf5", + "workflows": { + workflow_name: { + "description": description, + "actions": capped_pick_and_place(vessel_pos, pre_grasp_z, grasp_z), + } + }, + } diff --git a/source/matterix_tasks/matterix_tasks/test_dev_tasks/test_franka_rigid_labware_duran_100.py b/source/matterix_tasks/matterix_tasks/test_dev_tasks/test_franka_rigid_labware_duran_100.py new file mode 100644 index 0000000..8f35108 --- /dev/null +++ b/source/matterix_tasks/matterix_tasks/test_dev_tasks/test_franka_rigid_labware_duran_100.py @@ -0,0 +1,20 @@ +"""Dedicated physical checkpoint for the DURAN 100 mL capped bottle.""" + +from matterix.envs import MatterixBaseEnvCfg +from isaaclab.utils import configclass + +from .capped_labware_checkpoints import capped_env_fields + + +@configclass +class FrankaRigidLabwareDuran100EnvTestCfg(MatterixBaseEnvCfg): + """One-environment-per-asset DURAN 100 mL visual checkpoint.""" + + locals().update(capped_env_fields( + slug="dwk-218012458", + vessel_pos=(0.55, 0.0, 0.0), + pre_grasp_z=0.145, + grasp_z=0.075, + workflow_name="pick_and_place_duran_100", + description="Pick up and place the fixed-jointed DURAN 100 mL GL45 bottle", + )) diff --git a/source/matterix_tasks/matterix_tasks/test_dev_tasks/test_franka_rigid_labware_duran_500.py b/source/matterix_tasks/matterix_tasks/test_dev_tasks/test_franka_rigid_labware_duran_500.py new file mode 100644 index 0000000..4b4027a --- /dev/null +++ b/source/matterix_tasks/matterix_tasks/test_dev_tasks/test_franka_rigid_labware_duran_500.py @@ -0,0 +1,20 @@ +"""Dedicated physical checkpoint for the DURAN 500 mL capped bottle.""" + +from matterix.envs import MatterixBaseEnvCfg +from isaaclab.utils import configclass + +from .capped_labware_checkpoints import capped_env_fields + + +@configclass +class FrankaRigidLabwareDuran500EnvTestCfg(MatterixBaseEnvCfg): + """One-environment-per-asset DURAN 500 mL visual checkpoint.""" + + locals().update(capped_env_fields( + slug="dwk-218014459", + vessel_pos=(0.55, 0.0, 0.0), + pre_grasp_z=0.231, + grasp_z=0.145, + workflow_name="pick_and_place_duran_500", + description="Pick up and place the fixed-jointed DURAN 500 mL GL45 bottle", + )) diff --git a/source/matterix_tasks/matterix_tasks/test_dev_tasks/test_franka_rigid_labware_falcon_15.py b/source/matterix_tasks/matterix_tasks/test_dev_tasks/test_franka_rigid_labware_falcon_15.py new file mode 100644 index 0000000..6cb18ab --- /dev/null +++ b/source/matterix_tasks/matterix_tasks/test_dev_tasks/test_franka_rigid_labware_falcon_15.py @@ -0,0 +1,20 @@ +"""Dedicated physical checkpoint for the Falcon 15 mL capped tube.""" + +from matterix.envs import MatterixBaseEnvCfg +from isaaclab.utils import configclass + +from .capped_labware_checkpoints import capped_env_fields + + +@configclass +class FrankaRigidLabwareFalcon15EnvTestCfg(MatterixBaseEnvCfg): + """One-environment-per-asset Falcon 15 mL visual checkpoint.""" + + locals().update(capped_env_fields( + slug="falcon-352096", + vessel_pos=(0.55, 0.0, 0.0), + pre_grasp_z=0.1438, + grasp_z=0.104, + workflow_name="pick_and_place_falcon_15", + description="Pick up and place the fixed-jointed Falcon 15 mL tube", + )) diff --git a/source/matterix_tasks/matterix_tasks/test_dev_tasks/test_franka_rigid_labware_falcon_50.py b/source/matterix_tasks/matterix_tasks/test_dev_tasks/test_franka_rigid_labware_falcon_50.py new file mode 100644 index 0000000..a6a03c9 --- /dev/null +++ b/source/matterix_tasks/matterix_tasks/test_dev_tasks/test_franka_rigid_labware_falcon_50.py @@ -0,0 +1,20 @@ +"""Dedicated physical checkpoint for the Falcon 50 mL capped tube.""" + +from matterix.envs import MatterixBaseEnvCfg +from isaaclab.utils import configclass + +from .capped_labware_checkpoints import capped_env_fields + + +@configclass +class FrankaRigidLabwareFalcon50EnvTestCfg(MatterixBaseEnvCfg): + """One-environment-per-asset Falcon 50 mL visual checkpoint.""" + + locals().update(capped_env_fields( + slug="falcon-352070", + vessel_pos=(0.55, 0.0, 0.0), + pre_grasp_z=0.13955, + grasp_z=0.097, + workflow_name="pick_and_place_falcon_50", + description="Pick up and place the fixed-jointed Falcon 50 mL tube", + )) diff --git a/source/matterix_tasks/matterix_tasks/test_dev_tasks/test_franka_rigid_labware_flask_250.py b/source/matterix_tasks/matterix_tasks/test_dev_tasks/test_franka_rigid_labware_flask_250.py new file mode 100644 index 0000000..d506526 --- /dev/null +++ b/source/matterix_tasks/matterix_tasks/test_dev_tasks/test_franka_rigid_labware_flask_250.py @@ -0,0 +1,105 @@ +"""Dedicated physical checkpoint for the Corning 4980-250 flask.""" + +from matterix.envs import MatterixBaseEnvCfg, mdp +from matterix_assets.infrastructure.tables import TABLE_SEATTLE_INST_Cfg +from matterix_assets.labware.rigid_labware_batch1_local_only import CORNING_4980_250_LOCAL_ONLY_CFG +from matterix_assets.robots import FRANKA_PANDA_HIGH_PD_IK_CFG + +from matterix_sm import PickObjectCfg +from matterix_sm.robot_action_spaces import FRANKA_IK_ACTION_SPACE + +from isaaclab.managers import ObservationGroupCfg as ObsGroup +from isaaclab.managers import ObservationTermCfg as ObsTerm +from isaaclab.utils import configclass + +from .test_franka_rigid_labware_flasks import ( + FLASK_250_GRASP_OFFSET_M, + FLASK_250_POS, + EventCfg, + _put_back, +) + + +@configclass +class ObservationManagerCfg: + @configclass + class ArticulationsGroup(ObsGroup): + robot__root_world_pos = ObsTerm(func=mdp.root_world_pos, params={"asset_name": "robot"}) + robot__root_world_quat = ObsTerm(func=mdp.root_world_quat, params={"asset_name": "robot"}) + robot__joint_pos = ObsTerm(func=mdp.joint_pos, params={"asset_name": "robot"}) + robot__joint_vel = ObsTerm(func=mdp.joint_vel, params={"asset_name": "robot"}) + robot__ee_world_pos = ObsTerm(func=mdp.ee_world_pos, params={"asset_name": "robot"}) + robot__ee_world_quat = ObsTerm(func=mdp.ee_world_quat, params={"asset_name": "robot"}) + robot__gripper_pos = ObsTerm(func=mdp.gripper_pos, params={"asset_name": "robot"}) + + def __post_init__(self): + self.enable_corruption = False + self.concatenate_terms = False + + @configclass + class RigidObjectsGroup(ObsGroup): + flask_250__object_world_pos = ObsTerm( + func=mdp.object_world_pos, params={"asset_name": "flask_250"}) + flask_250__object_world_quat = ObsTerm( + func=mdp.object_world_quat, params={"asset_name": "flask_250"}) + flask_250__object_lin_vel = ObsTerm( + func=mdp.object_lin_vel, params={"asset_name": "flask_250"}) + flask_250__object_ang_vel = ObsTerm( + func=mdp.object_ang_vel, params={"asset_name": "flask_250"}) + flask_250__pre_grasp_frame = ObsTerm( + func=mdp.frame_world_pose, + params={"asset_name": "flask_250", "frame_name": "pre_grasp"}) + flask_250__grasp_frame = ObsTerm( + func=mdp.frame_world_pose, + params={"asset_name": "flask_250", "frame_name": "grasp"}) + flask_250__post_grasp_frame = ObsTerm( + func=mdp.frame_world_pose, + params={"asset_name": "flask_250", "frame_name": "post_grasp"}) + flask_250__opening_frame = ObsTerm( + func=mdp.frame_world_pose, + params={"asset_name": "flask_250", "frame_name": "opening"}) + flask_250__base_frame = ObsTerm( + func=mdp.frame_world_pose, + params={"asset_name": "flask_250", "frame_name": "base"}) + + def __post_init__(self): + self.enable_corruption = False + self.concatenate_terms = False + + articulations: ArticulationsGroup = ArticulationsGroup() + rigid_objects: RigidObjectsGroup = RigidObjectsGroup() + + +@configclass +class FrankaRigidLabwareFlask250EnvTestCfg(MatterixBaseEnvCfg): + """One-environment-per-asset flask-250 visual checkpoint.""" + + env_spacing = 10.0 + objects = { + "flask_250": CORNING_4980_250_LOCAL_ONLY_CFG(pos=FLASK_250_POS), + "table": TABLE_SEATTLE_INST_Cfg(pos=(0.5, 0, 0)), + } + articulated_assets = { + "robot": FRANKA_PANDA_HIGH_PD_IK_CFG(pos=(0.0, 0, 0)), + } + gripper_joint_names = ["panda_finger_joint1", "panda_finger_joint2"] + observations = ObservationManagerCfg() + events = EventCfg() + record_path = "datasets/dataset.hdf5" + workflows = { + "pickup_flask_250": PickObjectCfg( + description="Pick up the 250 mL Erlenmeyer flask", + agent_assets="robot", + object="flask_250", + action_space_info=FRANKA_IK_ACTION_SPACE, + ), + "pick_and_place_flask_250": [ + PickObjectCfg( + description="Pick up the 250 mL Erlenmeyer flask", + agent_assets="robot", + object="flask_250", + action_space_info=FRANKA_IK_ACTION_SPACE, + ), + *_put_back(FLASK_250_POS, FLASK_250_GRASP_OFFSET_M), + ], + } diff --git a/source/matterix_tasks/matterix_tasks/test_dev_tasks/test_franka_rigid_labware_flask_50.py b/source/matterix_tasks/matterix_tasks/test_dev_tasks/test_franka_rigid_labware_flask_50.py new file mode 100644 index 0000000..2e20cde --- /dev/null +++ b/source/matterix_tasks/matterix_tasks/test_dev_tasks/test_franka_rigid_labware_flask_50.py @@ -0,0 +1,105 @@ +"""Dedicated physical checkpoint for the Corning 4980-50 flask.""" + +from matterix.envs import MatterixBaseEnvCfg, mdp +from matterix_assets.infrastructure.tables import TABLE_SEATTLE_INST_Cfg +from matterix_assets.labware.rigid_labware_batch1_local_only import CORNING_4980_50_LOCAL_ONLY_CFG +from matterix_assets.robots import FRANKA_PANDA_HIGH_PD_IK_CFG + +from matterix_sm import PickObjectCfg +from matterix_sm.robot_action_spaces import FRANKA_IK_ACTION_SPACE + +from isaaclab.managers import ObservationGroupCfg as ObsGroup +from isaaclab.managers import ObservationTermCfg as ObsTerm +from isaaclab.utils import configclass + +from .test_franka_rigid_labware_flasks import ( + FLASK_50_GRASP_OFFSET_M, + FLASK_50_POS, + EventCfg, + _put_back, +) + + +@configclass +class ObservationManagerCfg: + @configclass + class ArticulationsGroup(ObsGroup): + robot__root_world_pos = ObsTerm(func=mdp.root_world_pos, params={"asset_name": "robot"}) + robot__root_world_quat = ObsTerm(func=mdp.root_world_quat, params={"asset_name": "robot"}) + robot__joint_pos = ObsTerm(func=mdp.joint_pos, params={"asset_name": "robot"}) + robot__joint_vel = ObsTerm(func=mdp.joint_vel, params={"asset_name": "robot"}) + robot__ee_world_pos = ObsTerm(func=mdp.ee_world_pos, params={"asset_name": "robot"}) + robot__ee_world_quat = ObsTerm(func=mdp.ee_world_quat, params={"asset_name": "robot"}) + robot__gripper_pos = ObsTerm(func=mdp.gripper_pos, params={"asset_name": "robot"}) + + def __post_init__(self): + self.enable_corruption = False + self.concatenate_terms = False + + @configclass + class RigidObjectsGroup(ObsGroup): + flask_50__object_world_pos = ObsTerm( + func=mdp.object_world_pos, params={"asset_name": "flask_50"}) + flask_50__object_world_quat = ObsTerm( + func=mdp.object_world_quat, params={"asset_name": "flask_50"}) + flask_50__object_lin_vel = ObsTerm( + func=mdp.object_lin_vel, params={"asset_name": "flask_50"}) + flask_50__object_ang_vel = ObsTerm( + func=mdp.object_ang_vel, params={"asset_name": "flask_50"}) + flask_50__pre_grasp_frame = ObsTerm( + func=mdp.frame_world_pose, + params={"asset_name": "flask_50", "frame_name": "pre_grasp"}) + flask_50__grasp_frame = ObsTerm( + func=mdp.frame_world_pose, + params={"asset_name": "flask_50", "frame_name": "grasp"}) + flask_50__post_grasp_frame = ObsTerm( + func=mdp.frame_world_pose, + params={"asset_name": "flask_50", "frame_name": "post_grasp"}) + flask_50__opening_frame = ObsTerm( + func=mdp.frame_world_pose, + params={"asset_name": "flask_50", "frame_name": "opening"}) + flask_50__base_frame = ObsTerm( + func=mdp.frame_world_pose, + params={"asset_name": "flask_50", "frame_name": "base"}) + + def __post_init__(self): + self.enable_corruption = False + self.concatenate_terms = False + + articulations: ArticulationsGroup = ArticulationsGroup() + rigid_objects: RigidObjectsGroup = RigidObjectsGroup() + + +@configclass +class FrankaRigidLabwareFlask50EnvTestCfg(MatterixBaseEnvCfg): + """One-environment-per-asset flask-50 visual checkpoint.""" + + env_spacing = 10.0 + objects = { + "flask_50": CORNING_4980_50_LOCAL_ONLY_CFG(pos=FLASK_50_POS), + "table": TABLE_SEATTLE_INST_Cfg(pos=(0.5, 0, 0)), + } + articulated_assets = { + "robot": FRANKA_PANDA_HIGH_PD_IK_CFG(pos=(0.0, 0, 0)), + } + gripper_joint_names = ["panda_finger_joint1", "panda_finger_joint2"] + observations = ObservationManagerCfg() + events = EventCfg() + record_path = "datasets/dataset.hdf5" + workflows = { + "pickup_flask_50": PickObjectCfg( + description="Pick up the 50 mL Erlenmeyer flask", + agent_assets="robot", + object="flask_50", + action_space_info=FRANKA_IK_ACTION_SPACE, + ), + "pick_and_place_flask_50": [ + PickObjectCfg( + description="Pick up the 50 mL Erlenmeyer flask", + agent_assets="robot", + object="flask_50", + action_space_info=FRANKA_IK_ACTION_SPACE, + ), + *_put_back(FLASK_50_POS, FLASK_50_GRASP_OFFSET_M), + ], + } diff --git a/source/matterix_tasks/matterix_tasks/test_dev_tasks/test_franka_rigid_labware_flasks.py b/source/matterix_tasks/matterix_tasks/test_dev_tasks/test_franka_rigid_labware_flasks.py new file mode 100644 index 0000000..c386764 --- /dev/null +++ b/source/matterix_tasks/matterix_tasks/test_dev_tasks/test_franka_rigid_labware_flasks.py @@ -0,0 +1,148 @@ +# Copyright (c) 2022-2026, The Matterix Project Developers. +# All rights reserved. +# +# SPDX-License-Identifier: BSD-3-Clause + +"""Shared helpers for the per-asset flask WebRTC checkpoint environments. + +Each registered task imports these frame observations and placement helpers while +keeping exactly one flask in its own environment. +""" + +from matterix.envs import MatterixBaseEnvCfg, mdp +from matterix.managers import EventManagerCfg +from matterix_assets.infrastructure.tables import TABLE_SEATTLE_INST_Cfg +from matterix_assets.labware.rigid_labware_batch1_local_only import ( + CORNING_4980_50_LOCAL_ONLY_CFG, + CORNING_4980_250_LOCAL_ONLY_CFG, +) +from matterix_assets.robots import FRANKA_PANDA_HIGH_PD_IK_CFG + +import torch + +from matterix_sm import MoveRelativeCfg, OpenGripperCfg, PickObjectCfg +from matterix_sm.primitive_actions.move_to_pose import MoveToPoseCfg +from matterix_sm.robot_action_spaces import FRANKA_IK_ACTION_SPACE + +import isaaclab.envs.mdp as isaaclab_mdp +from isaaclab.managers import EventTermCfg as EventTerm +from isaaclab.managers import ObservationGroupCfg as ObsGroup +from isaaclab.managers import ObservationTermCfg as ObsTerm +from isaaclab.utils import configclass + + +@configclass +class EventCfg(EventManagerCfg): + """Reset events for the dedicated flask checkpoint.""" + + reset_scene_to_default = EventTerm( + func=isaaclab_mdp.reset_scene_to_default, + mode="reset", + ) + + +APPROACH_CLEARANCE_M = 0.15 +RELEASE_CLEARANCE_M = 0.005 +FLASK_50_POS = (0.55, 0.0, 0.039) +FLASK_250_POS = (0.55, 0.24, 0.066) +FLASK_50_GRASP_OFFSET_M = 0.031 +FLASK_250_GRASP_OFFSET_M = 0.050 + + +def _put_back(vessel_pos, grasp_offset_z, agent="robot"): + """Return a flask to its authored world pose and release it.""" + x, y, z = vessel_pos + return [ + MoveToPoseCfg( + agent_assets=agent, + target_positions=torch.tensor([[x, y, z + grasp_offset_z + APPROACH_CLEARANCE_M]]), + action_space_info=FRANKA_IK_ACTION_SPACE, + ), + MoveToPoseCfg( + agent_assets=agent, + target_positions=torch.tensor([[x, y, z + grasp_offset_z + RELEASE_CLEARANCE_M]]), + action_space_info=FRANKA_IK_ACTION_SPACE, + ), + OpenGripperCfg(agent_assets=agent, action_space_info=FRANKA_IK_ACTION_SPACE), + MoveRelativeCfg( + agent_assets=agent, + position_offset=(0.0, 0.0, APPROACH_CLEARANCE_M), + orientation_offset=None, + action_space_info=FRANKA_IK_ACTION_SPACE, + ), + ] + + +@configclass +class ObservationManagerCfg: + """Robot and flask-frame observations consumed by the state machine.""" + + @configclass + class ArticulationsGroup(ObsGroup): + robot__root_world_pos = ObsTerm(func=mdp.root_world_pos, params={"asset_name": "robot"}) + robot__root_world_quat = ObsTerm(func=mdp.root_world_quat, params={"asset_name": "robot"}) + robot__joint_pos = ObsTerm(func=mdp.joint_pos, params={"asset_name": "robot"}) + robot__joint_vel = ObsTerm(func=mdp.joint_vel, params={"asset_name": "robot"}) + robot__ee_world_pos = ObsTerm(func=mdp.ee_world_pos, params={"asset_name": "robot"}) + robot__ee_world_quat = ObsTerm(func=mdp.ee_world_quat, params={"asset_name": "robot"}) + robot__gripper_pos = ObsTerm(func=mdp.gripper_pos, params={"asset_name": "robot"}) + + def __post_init__(self): + self.enable_corruption = False + self.concatenate_terms = False + + @configclass + class RigidObjectsGroup(ObsGroup): + flask_50__object_world_pos = ObsTerm( + func=mdp.object_world_pos, params={"asset_name": "flask_50"}) + flask_50__object_world_quat = ObsTerm( + func=mdp.object_world_quat, params={"asset_name": "flask_50"}) + flask_50__object_lin_vel = ObsTerm( + func=mdp.object_lin_vel, params={"asset_name": "flask_50"}) + flask_50__object_ang_vel = ObsTerm( + func=mdp.object_ang_vel, params={"asset_name": "flask_50"}) + flask_50__pre_grasp_frame = ObsTerm( + func=mdp.frame_world_pose, + params={"asset_name": "flask_50", "frame_name": "pre_grasp"}) + flask_50__grasp_frame = ObsTerm( + func=mdp.frame_world_pose, + params={"asset_name": "flask_50", "frame_name": "grasp"}) + flask_50__post_grasp_frame = ObsTerm( + func=mdp.frame_world_pose, + params={"asset_name": "flask_50", "frame_name": "post_grasp"}) + flask_50__opening_frame = ObsTerm( + func=mdp.frame_world_pose, + params={"asset_name": "flask_50", "frame_name": "opening"}) + flask_50__base_frame = ObsTerm( + func=mdp.frame_world_pose, + params={"asset_name": "flask_50", "frame_name": "base"}) + + flask_250__object_world_pos = ObsTerm( + func=mdp.object_world_pos, params={"asset_name": "flask_250"}) + flask_250__object_world_quat = ObsTerm( + func=mdp.object_world_quat, params={"asset_name": "flask_250"}) + flask_250__object_lin_vel = ObsTerm( + func=mdp.object_lin_vel, params={"asset_name": "flask_250"}) + flask_250__object_ang_vel = ObsTerm( + func=mdp.object_ang_vel, params={"asset_name": "flask_250"}) + flask_250__pre_grasp_frame = ObsTerm( + func=mdp.frame_world_pose, + params={"asset_name": "flask_250", "frame_name": "pre_grasp"}) + flask_250__grasp_frame = ObsTerm( + func=mdp.frame_world_pose, + params={"asset_name": "flask_250", "frame_name": "grasp"}) + flask_250__post_grasp_frame = ObsTerm( + func=mdp.frame_world_pose, + params={"asset_name": "flask_250", "frame_name": "post_grasp"}) + flask_250__opening_frame = ObsTerm( + func=mdp.frame_world_pose, + params={"asset_name": "flask_250", "frame_name": "opening"}) + flask_250__base_frame = ObsTerm( + func=mdp.frame_world_pose, params={"asset_name": "flask_250", "frame_name": "base"}) + + def __post_init__(self): + self.enable_corruption = False + self.concatenate_terms = False + + articulations: ArticulationsGroup = ArticulationsGroup() + rigid_objects: RigidObjectsGroup = RigidObjectsGroup()