From a9c9f907686dbb0dac4e38e29af6c472d8144f2e Mon Sep 17 00:00:00 2001 From: Abhishek Parmar Date: Sun, 10 May 2026 13:02:35 +0800 Subject: [PATCH 01/13] docs: clarify mujoco viewer setup --- .gitignore | 3 ++ README.md | 28 +++++++++- scripts/kinematics/check_mujoco_stack.py | 65 ++++++++++++++++++++++++ scripts/kinematics/parallel_ik_demo.py | 4 +- scripts/kinematics/test_robot_ik.py | 11 +++- 5 files changed, 107 insertions(+), 4 deletions(-) create mode 100644 scripts/kinematics/check_mujoco_stack.py diff --git a/.gitignore b/.gitignore index 0b71289a..94b67289 100644 --- a/.gitignore +++ b/.gitignore @@ -53,3 +53,6 @@ mlspaces_tests/**/test_data # pytest .pytest_cache/ __pycache__/ + +# Local agent notes +agents/documentation.md diff --git a/README.md b/README.md index d321b1b3..7ec69bfb 100644 --- a/README.md +++ b/README.md @@ -48,7 +48,7 @@ cd molmospaces Then, set up the virtual environment and install. -> Note: If you want to use the debug viewer on macOS you need to use conda or a Homebrew Python. This is because `mjpython` used by the debug viewer requires a shared `libpython3.11.dylib`, which `uv`'s standalone CPython does not ship. +> Note: If you want to use the debug viewer on macOS you need a Python environment that exposes a shared `libpython3.11.dylib` where `mjpython` can find it. Conda or Homebrew Python are the most reliable options; `uv` environments can run headless checks but may not expose the shared library at the path `mjpython` probes. With conda: @@ -67,6 +67,16 @@ source .venv/bin/activate uv pip install -e ".[mujoco]" ``` +To check a headless MuJoCo + MuJoCo Warp install: + +```bash +.venv/bin/python scripts/kinematics/check_mujoco_stack.py +``` + +On macOS, use conda or Homebrew Python for commands that open a MuJoCo viewer with +`mjpython` if your `uv` environment does not expose `libpython3.11.dylib`. The +`uv` environment is fine for non-viewer simulation, imports, and data generation. + One of the following options must be provided: - `mujoco` to use the classic MuJoCo renderer - `mujoco-filament` to use the improved Filament renderer for MuJoCo @@ -142,6 +152,22 @@ python scripts/datagen/run_pipeline.py --viewer --seed 3 mjpython scripts/datagen/run_pipeline.py --viewer --seed 3 ``` +For IK-specific viewer checks, use the kinematics test script rather than running +the kinematics module directly: + +```bash +# Standard MuJoCo IK +mjpython scripts/kinematics/test_robot_ik.py FrankaRobotConfig + +# Parallel/Warp IK +mjpython scripts/kinematics/test_robot_ik.py FrankaRobotConfig --parallel +``` + +If `mjpython` fails on macOS with `Library not loaded: @rpath/libpython3.11.dylib`, +recreate the environment with conda or Homebrew Python and reinstall +`pip install -e ".[mujoco]"`. That error is a Python runtime issue, not a MuJoCo +or MuJoCo Warp version mismatch. + The MolmoSpaces codebase has three entry points for data generation, evaluation, and debugging. The two initial entry points make use of experiment configs to configure runs. The third is more easily modifiable, with some logic for constructing runs on the fly, however constructing experiments is complicated and not all permutations have been tested fully. ```bash diff --git a/scripts/kinematics/check_mujoco_stack.py b/scripts/kinematics/check_mujoco_stack.py new file mode 100644 index 00000000..12fd805f --- /dev/null +++ b/scripts/kinematics/check_mujoco_stack.py @@ -0,0 +1,65 @@ +"""Check the MuJoCo, MuJoCo Warp, and mjpython runtime setup. + +This script does not open a viewer, so it is safe to run from a uv virtualenv: + + .venv/bin/python scripts/kinematics/check_mujoco_stack.py + +On macOS, mjpython requires a Python environment that exposes a shared +libpython3.11.dylib where mjpython can find it. Conda and Homebrew Python are the +most reliable viewer environments. +""" + +from __future__ import annotations + +import importlib.metadata as metadata +import os +import platform +import sys +import sysconfig + + +def _version(package_name: str) -> str: + try: + return metadata.version(package_name) + except metadata.PackageNotFoundError: + return "not installed" + + +def main() -> None: + import mujoco + + print(f"python executable: {sys.executable}") + print(f"python version: {platform.python_version()}") + print(f"platform: {platform.platform()}") + print() + print(f"mujoco import: {mujoco.__version__}") + print(f"mujoco dist: {_version('mujoco')}") + print(f"mujoco-warp dist: {_version('mujoco-warp')}") + print(f"mujoco-mjx dist: {_version('mujoco-mjx')}") + print(f"warp-lang dist: {_version('warp-lang')}") + print() + + model = mujoco.MjModel.from_xml_string( + '' + ) + data = mujoco.MjData(model) + mujoco.mj_step(model, data) + print(f"mujoco step: ok, time={data.time}") + + shared_library = sysconfig.get_config_var("LDLIBRARY") or "libpython3.11.dylib" + shared_library_path = os.path.join(sys.prefix, shared_library) + has_shared_python = os.path.exists(shared_library_path) + print() + print(f"shared python lib: {shared_library_path}") + print(f"exists: {has_shared_python}") + if sys.platform == "darwin" and not has_shared_python: + print() + print("mjpython viewer note:") + print(" This Python can run headless MuJoCo checks, but mjpython may fail") + print(" because the shared Python library is not exposed where mjpython") + print(" probes for it. Use conda or Homebrew Python for viewer commands") + print(" on macOS, or expose the dylib in this virtualenv.") + + +if __name__ == "__main__": + main() diff --git a/scripts/kinematics/parallel_ik_demo.py b/scripts/kinematics/parallel_ik_demo.py index 65cf6f9a..c53560cd 100644 --- a/scripts/kinematics/parallel_ik_demo.py +++ b/scripts/kinematics/parallel_ik_demo.py @@ -1,7 +1,9 @@ """ Demo script that demonstrates the parallel IK solver for multiple robots simultaneously. -On mac, run with mjpython. +On macOS, run viewer commands with mjpython from a conda or Homebrew Python +environment if your current environment does not expose libpython3.11.dylib. +uv environments can run headless checks, but mjpython needs that shared library. """ from collections import defaultdict diff --git a/scripts/kinematics/test_robot_ik.py b/scripts/kinematics/test_robot_ik.py index ead22b34..c9d42c7d 100644 --- a/scripts/kinematics/test_robot_ik.py +++ b/scripts/kinematics/test_robot_ik.py @@ -1,7 +1,14 @@ """ Interactive script to test a robot with the parallel and non-parallel IK solvers. -On mac, run with mjpython. +On macOS, run viewer commands with mjpython from a conda or Homebrew Python +environment if your current environment does not expose libpython3.11.dylib. +uv environments can run headless checks, but mjpython needs that shared library. + +Examples: + + mjpython scripts/kinematics/test_robot_ik.py FrankaRobotConfig + mjpython scripts/kinematics/test_robot_ik.py FrankaRobotConfig --parallel """ import argparse @@ -92,7 +99,7 @@ def main() -> None: move_group_id, target_pose, args.unlocked_move_groups, - robot_config.init_qpos, + robot_view.get_qpos_dict(), robot_view.base.pose, ) print(f"IK iteration {i}: {'Success' if ret is not None else 'Failed'}") From b533e107a9ab0a03232a83f0e9ea0ba9241cd1c4 Mon Sep 17 00:00:00 2001 From: Abhishek Parmar Date: Sun, 10 May 2026 18:08:36 +0800 Subject: [PATCH 02/13] feat: add Unitree G1 Dex1 embodiment --- .../component_tests/test_unitree_g1.py | 76 +++++++ molmo_spaces/configs/robot_configs.py | 36 +++ .../robots/robot_views/unitree_g1_view.py | 197 ++++++++++++++++ molmo_spaces/robots/unitree_g1.py | 103 +++++++++ scripts/assets/prepare_unitree_g1.py | 215 ++++++++++++++++++ 5 files changed, 627 insertions(+) create mode 100644 mlspaces_tests/component_tests/test_unitree_g1.py create mode 100644 molmo_spaces/robots/robot_views/unitree_g1_view.py create mode 100644 molmo_spaces/robots/unitree_g1.py create mode 100644 scripts/assets/prepare_unitree_g1.py diff --git a/mlspaces_tests/component_tests/test_unitree_g1.py b/mlspaces_tests/component_tests/test_unitree_g1.py new file mode 100644 index 00000000..fdacb7e5 --- /dev/null +++ b/mlspaces_tests/component_tests/test_unitree_g1.py @@ -0,0 +1,76 @@ +import os +from types import SimpleNamespace + +import mujoco +import numpy as np +import pytest + +import molmo_spaces.molmo_spaces_constants as constants +from molmo_spaces.configs.robot_configs import UnitreeG1Dex1RobotConfig +from molmo_spaces.robots.unitree_g1 import UnitreeG1Robot +from scripts.assets.prepare_unitree_g1 import ROBOT_ASSET_NAME, prepare_unitree_g1 + + +def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): + unitree_root = os.environ.get("UNITREE_URDF_ROOT") + if unitree_root is None: + pytest.skip("UNITREE_URDF_ROOT is not set") + + robot_root = tmp_path / "robots" + monkeypatch.setattr(constants, "ROBOTS_DIR", robot_root) + output_dir = robot_root / ROBOT_ASSET_NAME + xml_path = prepare_unitree_g1(unitree_root, output_dir) + + model = mujoco.MjModel.from_xml_path(str(xml_path)) + assert (model.nq, model.nv, model.nu) == (40, 39, 33) + + config = UnitreeG1Dex1RobotConfig() + scene_spec = mujoco.MjSpec() + robot_spec = mujoco.MjSpec.from_file(str(xml_path)) + config.robot_cls.add_robot_to_scene( + config, + scene_spec, + robot_spec, + config.robot_namespace, + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0, 0.0], + ) + scene_model = scene_spec.compile() + scene_data = mujoco.MjData(scene_model) + + view = config.robot_view_factory(scene_data, config.robot_namespace) + expected_counts = { + "base": (1, 0), + "left_leg": (6, 6), + "right_leg": (6, 6), + "waist": (3, 3), + "left_arm": (7, 7), + "right_arm": (7, 7), + "left_hand": (2, 2), + "right_hand": (2, 2), + } + for move_group_id, (n_joints, n_actuators) in expected_counts.items(): + move_group = view.get_move_group(move_group_id) + assert move_group.n_joints == n_joints + assert move_group.n_actuators == n_actuators + + robot = UnitreeG1Robot(scene_data, SimpleNamespace(robot_config=config)) + robot.reset() + mujoco.mj_forward(scene_model, scene_data) + + action = { + move_group_id: robot.robot_view.get_move_group(move_group_id).joint_pos.copy() + for move_group_id in robot.controllers + } + action["left_arm"] = action["left_arm"].copy() + action["left_arm"][3] = 0.1 + robot.update_control(action) + robot.compute_control() + + for _ in range(10): + mujoco.mj_step(scene_model, scene_data) + + assert np.isfinite(scene_data.qpos).all() + assert np.isfinite(scene_data.ctrl).all() + assert robot.state_dim == 40 + assert robot.action_dim(list(robot.controllers)) == 33 diff --git a/molmo_spaces/configs/robot_configs.py b/molmo_spaces/configs/robot_configs.py index a4a26196..8547726c 100644 --- a/molmo_spaces/configs/robot_configs.py +++ b/molmo_spaces/configs/robot_configs.py @@ -36,6 +36,8 @@ from molmo_spaces.robots.robot_views.mobile_franka_droid_view import MobileFrankaDroidRobotView from molmo_spaces.robots.robot_views.rby1_view import RBY1RobotView from molmo_spaces.robots.robot_views.rum_gripper_view import FloatingRUMRobotView +from molmo_spaces.robots.robot_views.unitree_g1_view import UnitreeG1RobotView +from molmo_spaces.robots.unitree_g1 import UnitreeG1Robot class ActionNoiseConfig(Config): @@ -419,3 +421,37 @@ def model_post_init(self, __context): assert self.command_mode["gripper"] == "joint_position" if "arm" in self.command_mode: assert self.command_mode["arm"] in ["joint_position", "joint_rel_position"] + + +class UnitreeG1Dex1RobotConfig(BaseRobotConfig): + """Configuration for Unitree G1 29-DoF humanoid with Dex1.1 hands.""" + + robot_cls: type[UnitreeG1Robot] | None = UnitreeG1Robot + robot_factory: Callable[[MjData, Any], Robot] | None = UnitreeG1Robot + robot_view_factory: RobotViewFactory | None = UnitreeG1RobotView + robot_namespace: str = "robot_0/" + default_world_pose: list[float] = [0.0, 0.0, 0.793, 1.0, 0.0, 0.0, 0.0] + name: str = "unitree_g1_29dof_dex1_1" + robot_xml_path: Path = Path("model.xml") + init_qpos: dict[str, list[float]] = { + "base": [0.0, 0.0, 0.793, 1.0, 0.0, 0.0, 0.0], + "left_leg": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + "right_leg": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + "waist": [0.0, 0.0, 0.0], + "left_arm": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + "right_arm": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + "left_hand": [0.0245, 0.0245], + "right_hand": [0.0245, 0.0245], + } + init_qpos_noise_range: dict[str, list[float]] | None = None + command_mode: dict[str, str | None] = { + "base": None, + "left_leg": "joint_position", + "right_leg": "joint_position", + "waist": "joint_position", + "left_arm": "joint_position", + "right_arm": "joint_position", + "left_hand": "joint_position", + "right_hand": "joint_position", + } + gravcomp: bool = True diff --git a/molmo_spaces/robots/robot_views/unitree_g1_view.py b/molmo_spaces/robots/robot_views/unitree_g1_view.py new file mode 100644 index 00000000..516773e9 --- /dev/null +++ b/molmo_spaces/robots/robot_views/unitree_g1_view.py @@ -0,0 +1,197 @@ +"""Robot view for the Unitree G1 Dex1.1 humanoid.""" + +from typing import Literal + +import numpy as np +from mujoco import MjData + +from molmo_spaces.robots.robot_views.abstract import ( + FreeJointRobotBaseGroup, + MJCFFrameMixin, + RobotBaseGroup, + RobotView, + SimplyActuatedMoveGroup, +) +from molmo_spaces.utils.mj_model_and_data_utils import body_pose + +UNITREE_G1_LEG_JOINTS = [ + "hip_pitch_joint", + "hip_roll_joint", + "hip_yaw_joint", + "knee_joint", + "ankle_pitch_joint", + "ankle_roll_joint", +] + +UNITREE_G1_WAIST_JOINTS = [ + "waist_yaw_joint", + "waist_roll_joint", + "waist_pitch_joint", +] + +UNITREE_G1_ARM_JOINTS = [ + "shoulder_pitch_joint", + "shoulder_roll_joint", + "shoulder_yaw_joint", + "elbow_joint", + "wrist_roll_joint", + "wrist_pitch_joint", + "wrist_yaw_joint", +] + +UNITREE_G1_HAND_JOINTS = [ + "dex1_finger_joint_1", + "dex1_finger_joint_2", +] + + +class UnitreeG1BaseGroup(FreeJointRobotBaseGroup): + """Floating pelvis base for Unitree G1.""" + + def __init__(self, mj_data: MjData, namespace: str = "") -> None: + model = mj_data.model + base_joint_id = model.joint(f"{namespace}floating_base_joint").id + super().__init__(mj_data, base_joint_id, [], [], floating=True) + + @property + def noop_ctrl(self) -> np.ndarray: + return np.array([]) + + +class UnitreeG1BodyMoveGroup(MJCFFrameMixin, SimplyActuatedMoveGroup): + """A simple joint-position-controlled G1 move group.""" + + def __init__( + self, + mj_data: MjData, + joint_names: list[str], + root_body_name: str, + leaf_body_name: str, + base: RobotBaseGroup, + namespace: str = "", + ) -> None: + model = mj_data.model + joint_ids = [model.joint(f"{namespace}{joint_name}").id for joint_name in joint_names] + act_ids = [model.actuator(f"{namespace}{joint_name}").id for joint_name in joint_names] + self._root_body_id = model.body(f"{namespace}{root_body_name}").id + self._leaf_body_id = model.body(f"{namespace}{leaf_body_name}").id + super().__init__(mj_data, joint_ids, act_ids, self._root_body_id, base) + + @property + def leaf_frame_id(self) -> int: + return self._leaf_body_id + + @property + def leaf_frame_type(self): + return "body" + + @property + def noop_ctrl(self) -> np.ndarray: + return self.joint_pos.copy() + + @property + def leaf_frame_to_world(self) -> np.ndarray: + return body_pose(self.mj_data, self._leaf_body_id) + + @property + def root_frame_to_world(self) -> np.ndarray: + return body_pose(self.mj_data, self._root_body_id) + + +class UnitreeG1SideMoveGroup(UnitreeG1BodyMoveGroup): + """A side-specific G1 move group with conventional joint-name prefixes.""" + + def __init__( + self, + mj_data: MjData, + side: Literal["left", "right"], + joint_suffixes: list[str], + root_body_name: str, + leaf_body_name: str, + base: RobotBaseGroup, + namespace: str = "", + ) -> None: + joint_names = [f"{side}_{joint_suffix}" for joint_suffix in joint_suffixes] + super().__init__(mj_data, joint_names, root_body_name, leaf_body_name, base, namespace) + + +class UnitreeG1RobotView(RobotView): + """Move-group view for the Unitree G1 Dex1.1 model.""" + + def __init__(self, mj_data: MjData, namespace: str = "") -> None: + self._namespace = namespace + base = UnitreeG1BaseGroup(mj_data, namespace=namespace) + move_groups = { + "base": base, + "left_leg": UnitreeG1SideMoveGroup( + mj_data, + "left", + UNITREE_G1_LEG_JOINTS, + "pelvis", + "left_ankle_roll_link", + base, + namespace, + ), + "right_leg": UnitreeG1SideMoveGroup( + mj_data, + "right", + UNITREE_G1_LEG_JOINTS, + "pelvis", + "right_ankle_roll_link", + base, + namespace, + ), + "waist": UnitreeG1BodyMoveGroup( + mj_data, + UNITREE_G1_WAIST_JOINTS, + "pelvis", + "torso_link", + base, + namespace, + ), + "left_arm": UnitreeG1SideMoveGroup( + mj_data, + "left", + UNITREE_G1_ARM_JOINTS, + "torso_link", + "left_wrist_yaw_link", + base, + namespace, + ), + "right_arm": UnitreeG1SideMoveGroup( + mj_data, + "right", + UNITREE_G1_ARM_JOINTS, + "torso_link", + "right_wrist_yaw_link", + base, + namespace, + ), + "left_hand": UnitreeG1SideMoveGroup( + mj_data, + "left", + UNITREE_G1_HAND_JOINTS, + "left_wrist_yaw_link", + "left_dex1_finger_link_1", + base, + namespace, + ), + "right_hand": UnitreeG1SideMoveGroup( + mj_data, + "right", + UNITREE_G1_HAND_JOINTS, + "right_wrist_yaw_link", + "right_dex1_finger_link_1", + base, + namespace, + ), + } + super().__init__(mj_data, move_groups) + + @property + def name(self) -> str: + return f"{self._namespace}unitree_g1_29dof_dex1_1" + + @property + def base(self) -> UnitreeG1BaseGroup: + return self._move_groups["base"] diff --git a/molmo_spaces/robots/unitree_g1.py b/molmo_spaces/robots/unitree_g1.py new file mode 100644 index 00000000..11961782 --- /dev/null +++ b/molmo_spaces/robots/unitree_g1.py @@ -0,0 +1,103 @@ +"""Unitree G1 Dex1.1 robot implementation.""" + +from typing import TYPE_CHECKING, Any + +import numpy as np +from mujoco import MjData + +from molmo_spaces.controllers.abstract import Controller +from molmo_spaces.controllers.joint_pos import JointPosController +from molmo_spaces.kinematics.mujoco_kinematics import MlSpacesKinematics +from molmo_spaces.kinematics.parallel.dummy_parallel_kinematics import DummyParallelKinematics +from molmo_spaces.robots.abstract import Robot + +if TYPE_CHECKING: + from molmo_spaces.configs.abstract_exp_config import MlSpacesExpConfig + + +class UnitreeG1Robot(Robot): + """Unitree G1 Dex1.1 with joint-position control for all non-base move groups.""" + + def __init__( + self, + mj_data: MjData, + config: "MlSpacesExpConfig", + ) -> None: + super().__init__(mj_data, config) + self._robot_view = config.robot_config.robot_view_factory( + mj_data, config.robot_config.robot_namespace + ) + self._kinematics = MlSpacesKinematics(config.robot_config) + self._parallel_kinematics = DummyParallelKinematics( + config.robot_config, + self._kinematics, + ) + self._controllers = { + move_group_id: JointPosController(self._robot_view.get_move_group(move_group_id)) + for move_group_id in self._robot_view.move_group_ids() + if move_group_id != "base" + and config.robot_config.command_mode.get(move_group_id) == "joint_position" + } + + @property + def namespace(self): + return self.exp_config.robot_config.robot_namespace + + @property + def robot_view(self): + return self._robot_view + + @property + def kinematics(self): + return self._kinematics + + @property + def parallel_kinematics(self): + return self._parallel_kinematics + + @property + def controllers(self) -> dict[str, Controller]: + return self._controllers + + @property + def state_dim(self) -> int: + return sum( + self._robot_view.get_move_group(move_group_id).pos_dim + for move_group_id in self._robot_view.move_group_ids() + ) + + def action_dim(self, move_group_ids: list[str]) -> int: + return sum(self._robot_view.get_move_group(mg_id).n_actuators for mg_id in move_group_ids) + + def get_arm_move_group_ids(self) -> list[str]: + return ["left_arm", "right_arm"] + + def update_control(self, action_command_dict: dict[str, Any]) -> None: + action_command_dict = self._apply_action_noise_and_save_unnoised_cmd_jp(action_command_dict) + + for mg_id, controller in self.controllers.items(): + if mg_id in action_command_dict and action_command_dict[mg_id] is not None: + controller.set_target(action_command_dict[mg_id]) + elif not controller.stationary: + controller.set_to_stationary() + + def compute_control(self) -> None: + for controller in self.controllers.values(): + ctrl_inputs = controller.compute_ctrl_inputs() + controller.robot_move_group.ctrl = ctrl_inputs + + def set_joint_pos(self, robot_joint_pos_dict) -> None: + for mg_id, joint_pos in robot_joint_pos_dict.items(): + self._robot_view.get_move_group(mg_id).joint_pos = joint_pos + + def set_world_pose(self, robot_world_pose) -> None: + self._robot_view.base.pose = robot_world_pose + + def reset(self) -> None: + for mg_id, default_pos in self.exp_config.robot_config.init_qpos.items(): + if mg_id in self._robot_view.move_group_ids(): + self._robot_view.get_move_group(mg_id).joint_pos = np.asarray(default_pos) + + @staticmethod + def robot_model_root_name() -> str: + return "pelvis" diff --git a/scripts/assets/prepare_unitree_g1.py b/scripts/assets/prepare_unitree_g1.py new file mode 100644 index 00000000..f84ba028 --- /dev/null +++ b/scripts/assets/prepare_unitree_g1.py @@ -0,0 +1,215 @@ +"""Prepare the Unitree G1 Dex1.1 model for MolmoSpaces. + +The upstream Unitree Dex1.1 URDF has the geometry we want, but it does not +include a free base or MuJoCo actuators. This script creates a deterministic +MuJoCo XML asset with the MolmoSpaces control contract. +""" + +from __future__ import annotations + +import argparse +import os +import shutil +import xml.etree.ElementTree as ET +from dataclasses import dataclass +from pathlib import Path + +import mujoco + +from molmo_spaces.molmo_spaces_constants import ROBOTS_DIR + +ROBOT_ASSET_NAME = "unitree_g1_29dof_dex1_1" +SOURCE_URDF_REL_PATH = Path("g1_description/g1_29dof_mode_15_with_dex1_1.urdf") +SOURCE_MESH_DIR_REL_PATH = Path("g1_description/meshes") +OUTPUT_XML_NAME = "model.xml" + + +GAIN_BY_JOINT_GROUP = { + "leg": (120.0, 12.0), + "waist": (80.0, 8.0), + "arm": (60.0, 6.0), + "hand": (25.0, 2.5), +} + + +@dataclass(frozen=True) +class JointSpec: + name: str + lower: float + upper: float + effort: float | None + group: str + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--unitree-urdf-root", + type=Path, + default=Path(os.environ["UNITREE_URDF_ROOT"]) if "UNITREE_URDF_ROOT" in os.environ else None, + help="Path to the Unitree `unitree_ros/robots` directory. Defaults to UNITREE_URDF_ROOT.", + ) + parser.add_argument( + "--output-dir", + type=Path, + default=ROBOTS_DIR / ROBOT_ASSET_NAME, + help=f"Output robot asset directory. Defaults to $MLSPACES_ASSETS_DIR/robots/{ROBOT_ASSET_NAME}.", + ) + parser.add_argument( + "--force", + action="store_true", + help="Overwrite an existing output directory.", + ) + return parser.parse_args() + + +def _joint_group(joint_name: str) -> str: + if "_hip_" in joint_name or "_knee_" in joint_name or "_ankle_" in joint_name: + return "leg" + if joint_name.startswith("waist_"): + return "waist" + if "_shoulder_" in joint_name or "_elbow_" in joint_name or "_wrist_" in joint_name: + return "arm" + if "_dex1_" in joint_name: + return "hand" + raise ValueError(f"Unexpected controllable G1 joint: {joint_name}") + + +def _parse_controllable_joints(urdf_path: Path) -> list[JointSpec]: + root = ET.parse(urdf_path).getroot() + joints = [] + for joint in root.findall("joint"): + joint_name = joint.get("name") + joint_type = joint.get("type") + if joint_name is None or joint_type in {None, "fixed", "floating"}: + continue + + limit = joint.find("limit") + if limit is None: + raise ValueError(f"Joint {joint_name} is missing a element") + + lower = limit.get("lower") + upper = limit.get("upper") + if lower is None or upper is None: + raise ValueError(f"Joint {joint_name} is missing lower/upper limits") + + effort = limit.get("effort") + joints.append( + JointSpec( + name=joint_name, + lower=float(lower), + upper=float(upper), + effort=float(effort) if effort is not None else None, + group=_joint_group(joint_name), + ) + ) + return joints + + +def _mesh_files_referenced_by_urdf(urdf_path: Path) -> set[Path]: + root = ET.parse(urdf_path).getroot() + mesh_files = set() + for mesh in root.findall(".//mesh"): + filename = mesh.get("filename") + if filename is None: + continue + mesh_files.add(Path(filename)) + return mesh_files + + +def _copy_referenced_meshes(source_root: Path, urdf_path: Path, output_dir: Path) -> None: + source_mesh_root = source_root / SOURCE_MESH_DIR_REL_PATH + output_mesh_root = output_dir / "meshes" + output_mesh_root.mkdir(parents=True, exist_ok=True) + + for mesh_file in sorted(_mesh_files_referenced_by_urdf(urdf_path)): + if mesh_file.is_absolute(): + source_mesh_path = mesh_file + relative_mesh_path = Path(mesh_file.name) + else: + relative_mesh_path = Path(mesh_file.name) + source_mesh_path = source_mesh_root / relative_mesh_path + + if not source_mesh_path.is_file(): + raise FileNotFoundError(f"Referenced mesh does not exist: {source_mesh_path}") + + destination = output_mesh_root / relative_mesh_path + destination.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(source_mesh_path, destination) + + +def _add_position_actuators(spec: mujoco.MjSpec, joints: list[JointSpec]) -> None: + for joint in joints: + stiffness, damping = GAIN_BY_JOINT_GROUP[joint.group] + kwargs = {} + if joint.effort is not None: + kwargs["forcelimited"] = 1 + kwargs["forcerange"] = [-joint.effort, joint.effort] + + spec.add_actuator( + name=joint.name, + trntype=mujoco.mjtTrn.mjTRN_JOINT, + target=joint.name, + gaintype=mujoco.mjtGain.mjGAIN_FIXED, + gainprm=[stiffness, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + biastype=mujoco.mjtBias.mjBIAS_AFFINE, + biasprm=[0.0, -stiffness, -damping, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ctrllimited=1, + ctrlrange=[joint.lower, joint.upper], + **kwargs, + ) + + +def prepare_unitree_g1( + unitree_urdf_root: str | os.PathLike[str], + output_dir: str | os.PathLike[str], + force: bool = False, +) -> Path: + source_root = Path(unitree_urdf_root).expanduser().resolve() + urdf_path = source_root / SOURCE_URDF_REL_PATH + if not urdf_path.is_file(): + raise FileNotFoundError(f"Could not find Unitree G1 Dex1.1 URDF: {urdf_path}") + + output_dir = Path(output_dir).expanduser().resolve() + if output_dir.exists(): + if not force: + raise FileExistsError(f"Output directory already exists: {output_dir}") + shutil.rmtree(output_dir) + output_dir.mkdir(parents=True) + + joints = _parse_controllable_joints(urdf_path) + if len(joints) != 33: + raise ValueError(f"Expected 33 controllable joints, found {len(joints)}") + + _copy_referenced_meshes(source_root, urdf_path, output_dir) + + spec = mujoco.MjSpec.from_file(str(urdf_path)) + pelvis = spec.body("pelvis") + if pelvis is None: + raise ValueError("Expected root body `pelvis` in Unitree G1 URDF") + pelvis.add_freejoint(name="floating_base_joint") + _add_position_actuators(spec, joints) + + model = spec.compile() + expected_shape = (40, 39, 33) + actual_shape = (model.nq, model.nv, model.nu) + if actual_shape != expected_shape: + raise ValueError(f"Expected model shape nq/nv/nu={expected_shape}, got {actual_shape}") + + output_xml_path = output_dir / OUTPUT_XML_NAME + spec.to_file(str(output_xml_path)) + mujoco.MjModel.from_xml_path(str(output_xml_path)) + return output_xml_path + + +def main() -> None: + args = parse_args() + if args.unitree_urdf_root is None: + raise SystemExit("Set UNITREE_URDF_ROOT or pass --unitree-urdf-root") + + output_xml_path = prepare_unitree_g1(args.unitree_urdf_root, args.output_dir, args.force) + print(f"Wrote Unitree G1 Dex1.1 model: {output_xml_path}") + + +if __name__ == "__main__": + main() From 11f47c12fc350c60f0ca1c7eb46e4b504f28fe00 Mon Sep 17 00:00:00 2001 From: Abhishek Parmar Date: Sun, 10 May 2026 18:31:54 +0800 Subject: [PATCH 03/13] feat: add Unitree G1 FK IK smoke path --- .../component_tests/test_unitree_g1.py | 52 +++++++++++++++++++ .../object_manipulation_datagen_configs.py | 17 ++++++ .../robots/robot_views/unitree_g1_view.py | 40 ++++++++++---- scripts/assets/prepare_unitree_g1.py | 24 ++++++++- scripts/kinematics/test_robot_ik.py | 43 +++++++++++++-- 5 files changed, 160 insertions(+), 16 deletions(-) diff --git a/mlspaces_tests/component_tests/test_unitree_g1.py b/mlspaces_tests/component_tests/test_unitree_g1.py index fdacb7e5..d8fe5dce 100644 --- a/mlspaces_tests/component_tests/test_unitree_g1.py +++ b/mlspaces_tests/component_tests/test_unitree_g1.py @@ -1,3 +1,4 @@ +import importlib import os from types import SimpleNamespace @@ -7,9 +8,18 @@ import molmo_spaces.molmo_spaces_constants as constants from molmo_spaces.configs.robot_configs import UnitreeG1Dex1RobotConfig +from molmo_spaces.data_generation.config_registry import get_config_class +from molmo_spaces.kinematics.mujoco_kinematics import MlSpacesKinematics from molmo_spaces.robots.unitree_g1 import UnitreeG1Robot from scripts.assets.prepare_unitree_g1 import ROBOT_ASSET_NAME, prepare_unitree_g1 +KINEMATICS_SITE_NAMES = [ + "left_wrist_site", + "right_wrist_site", + "left_grasp_site", + "right_grasp_site", +] + def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): unitree_root = os.environ.get("UNITREE_URDF_ROOT") @@ -23,6 +33,8 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): model = mujoco.MjModel.from_xml_path(str(xml_path)) assert (model.nq, model.nv, model.nu) == (40, 39, 33) + for site_name in KINEMATICS_SITE_NAMES: + assert model.site(site_name).id >= 0 config = UnitreeG1Dex1RobotConfig() scene_spec = mujoco.MjSpec() @@ -53,6 +65,46 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): move_group = view.get_move_group(move_group_id) assert move_group.n_joints == n_joints assert move_group.n_actuators == n_actuators + assert view.get_move_group("left_arm").leaf_frame_type == "site" + assert ( + scene_model.site("robot_0/left_wrist_site").id + == view.get_move_group("left_arm").leaf_frame_id + ) + assert view.get_move_group("right_arm").leaf_frame_type == "site" + assert ( + scene_model.site("robot_0/right_wrist_site").id + == view.get_move_group("right_arm").leaf_frame_id + ) + + kinematics = MlSpacesKinematics(config) + fk = kinematics.fk(config.init_qpos, np.eye(4)) + for move_group_id in ("left_arm", "right_arm"): + assert fk[move_group_id].shape == (4, 4) + assert np.isfinite(fk[move_group_id]).all() + assert np.allclose(fk[move_group_id][3], [0.0, 0.0, 0.0, 1.0]) + + importlib.import_module( + "molmo_spaces.data_generation.config.object_manipulation_datagen_configs" + ) + datagen_config_cls = get_config_class("UnitreeG1SceneSmokeDataGenConfig") + datagen_config = datagen_config_cls() + assert isinstance(datagen_config.robot_config, UnitreeG1Dex1RobotConfig) + + base_scene_path = constants.ABS_PATH_OF_TOP_LEVEL_MOLMO_SPACES_DIR / ( + "molmo_spaces/resources/base_scene.xml" + ) + datagen_scene_spec = mujoco.MjSpec.from_file(str(base_scene_path)) + datagen_robot_spec = mujoco.MjSpec.from_file(str(xml_path)) + datagen_config.robot_config.robot_cls.add_robot_to_scene( + datagen_config.robot_config, + datagen_scene_spec, + datagen_robot_spec, + datagen_config.robot_config.robot_namespace, + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0, 0.0], + ) + datagen_scene_model = datagen_scene_spec.compile() + assert datagen_scene_model.body("robot_0/pelvis").id >= 0 robot = UnitreeG1Robot(scene_data, SimpleNamespace(robot_config=config)) robot.reset() diff --git a/molmo_spaces/data_generation/config/object_manipulation_datagen_configs.py b/molmo_spaces/data_generation/config/object_manipulation_datagen_configs.py index 7b18789d..4845cf08 100644 --- a/molmo_spaces/data_generation/config/object_manipulation_datagen_configs.py +++ b/molmo_spaces/data_generation/config/object_manipulation_datagen_configs.py @@ -19,6 +19,7 @@ from molmo_spaces.configs.base_pick_and_place_next_to_configs import PickAndPlaceNextToDataGenConfig from molmo_spaces.configs.base_pick_config import PickBaseConfig from molmo_spaces.configs.camera_configs import ( + CameraSystemConfig, FrankaDroidCameraSystem, FrankaEasyRandomizedDroidCameraSystem, FrankaGoProD405D455CameraSystem, @@ -38,6 +39,7 @@ FrankaRobotConfig, RBY1MConfig, RBY1MOpenCloseConfig, + UnitreeG1Dex1RobotConfig, ) from molmo_spaces.configs.task_sampler_configs import ( OpenTaskSamplerConfig, @@ -123,6 +125,21 @@ def tag(self) -> str: return "rum_pick_datagen" +@register_config("UnitreeG1SceneSmokeDataGenConfig") +class UnitreeG1SceneSmokeDataGenConfig(PickBaseConfig): + """Minimal config-path smoke test for loading Unitree G1 into MolmoSpaces scenes.""" + + robot_config: UnitreeG1Dex1RobotConfig = UnitreeG1Dex1RobotConfig() + camera_config: CameraSystemConfig | None = None + task_horizon: int = 1 + profile: bool = False + output_dir: Path = ASSETS_DIR / "experiment_output" / "datagen" / "unitree_g1_scene_smoke_v1" + + @property + def tag(self) -> str: + return "unitree_g1_scene_smoke_datagen" + + @register_config("FrankaPickAndPlaceDataGenConfig") class FrankaPickAndPlaceDataGenConfig(PickAndPlaceDataGenConfig): robot_config: BaseRobotConfig = FrankaRobotConfig() diff --git a/molmo_spaces/robots/robot_views/unitree_g1_view.py b/molmo_spaces/robots/robot_views/unitree_g1_view.py index 516773e9..6aa57fe9 100644 --- a/molmo_spaces/robots/robot_views/unitree_g1_view.py +++ b/molmo_spaces/robots/robot_views/unitree_g1_view.py @@ -12,7 +12,7 @@ RobotView, SimplyActuatedMoveGroup, ) -from molmo_spaces.utils.mj_model_and_data_utils import body_pose +from molmo_spaces.utils.mj_model_and_data_utils import body_pose, site_pose UNITREE_G1_LEG_JOINTS = [ "hip_pitch_joint", @@ -66,24 +66,31 @@ def __init__( mj_data: MjData, joint_names: list[str], root_body_name: str, - leaf_body_name: str, + leaf_frame_name: str, base: RobotBaseGroup, namespace: str = "", + leaf_frame_type: Literal["body", "site"] = "body", ) -> None: model = mj_data.model joint_ids = [model.joint(f"{namespace}{joint_name}").id for joint_name in joint_names] act_ids = [model.actuator(f"{namespace}{joint_name}").id for joint_name in joint_names] self._root_body_id = model.body(f"{namespace}{root_body_name}").id - self._leaf_body_id = model.body(f"{namespace}{leaf_body_name}").id + self._leaf_frame_type = leaf_frame_type + if leaf_frame_type == "body": + self._leaf_frame_id = model.body(f"{namespace}{leaf_frame_name}").id + elif leaf_frame_type == "site": + self._leaf_frame_id = model.site(f"{namespace}{leaf_frame_name}").id + else: + raise ValueError(f"Unsupported leaf_frame_type: {leaf_frame_type}") super().__init__(mj_data, joint_ids, act_ids, self._root_body_id, base) @property def leaf_frame_id(self) -> int: - return self._leaf_body_id + return self._leaf_frame_id @property def leaf_frame_type(self): - return "body" + return self._leaf_frame_type @property def noop_ctrl(self) -> np.ndarray: @@ -91,7 +98,9 @@ def noop_ctrl(self) -> np.ndarray: @property def leaf_frame_to_world(self) -> np.ndarray: - return body_pose(self.mj_data, self._leaf_body_id) + if self.leaf_frame_type == "body": + return body_pose(self.mj_data, self._leaf_frame_id) + return site_pose(self.mj_data, self._leaf_frame_id) @property def root_frame_to_world(self) -> np.ndarray: @@ -107,12 +116,21 @@ def __init__( side: Literal["left", "right"], joint_suffixes: list[str], root_body_name: str, - leaf_body_name: str, + leaf_frame_name: str, base: RobotBaseGroup, namespace: str = "", + leaf_frame_type: Literal["body", "site"] = "body", ) -> None: joint_names = [f"{side}_{joint_suffix}" for joint_suffix in joint_suffixes] - super().__init__(mj_data, joint_names, root_body_name, leaf_body_name, base, namespace) + super().__init__( + mj_data, + joint_names, + root_body_name, + leaf_frame_name, + base, + namespace, + leaf_frame_type, + ) class UnitreeG1RobotView(RobotView): @@ -154,18 +172,20 @@ def __init__(self, mj_data: MjData, namespace: str = "") -> None: "left", UNITREE_G1_ARM_JOINTS, "torso_link", - "left_wrist_yaw_link", + "left_wrist_site", base, namespace, + "site", ), "right_arm": UnitreeG1SideMoveGroup( mj_data, "right", UNITREE_G1_ARM_JOINTS, "torso_link", - "right_wrist_yaw_link", + "right_wrist_site", base, namespace, + "site", ), "left_hand": UnitreeG1SideMoveGroup( mj_data, diff --git a/scripts/assets/prepare_unitree_g1.py b/scripts/assets/prepare_unitree_g1.py index f84ba028..e64b612c 100644 --- a/scripts/assets/prepare_unitree_g1.py +++ b/scripts/assets/prepare_unitree_g1.py @@ -46,7 +46,9 @@ def parse_args() -> argparse.Namespace: parser.add_argument( "--unitree-urdf-root", type=Path, - default=Path(os.environ["UNITREE_URDF_ROOT"]) if "UNITREE_URDF_ROOT" in os.environ else None, + default=Path(os.environ["UNITREE_URDF_ROOT"]) + if "UNITREE_URDF_ROOT" in os.environ + else None, help="Path to the Unitree `unitree_ros/robots` directory. Defaults to UNITREE_URDF_ROOT.", ) parser.add_argument( @@ -160,6 +162,25 @@ def _add_position_actuators(spec: mujoco.MjSpec, joints: list[JointSpec]) -> Non ) +def _add_kinematics_sites(spec: mujoco.MjSpec) -> None: + for side in ("left", "right"): + wrist_body = spec.body(f"{side}_wrist_yaw_link") + if wrist_body is None: + raise ValueError(f"Expected body `{side}_wrist_yaw_link` in Unitree G1 model") + wrist_body.add_site( + name=f"{side}_wrist_site", + pos=[0.0, 0.0, 0.0], + size=[0.015], + rgba=[0.1, 0.4, 1.0, 1.0], + ) + wrist_body.add_site( + name=f"{side}_grasp_site", + pos=[0.095, 0.0, 0.0], + size=[0.015], + rgba=[0.1, 1.0, 0.4, 1.0], + ) + + def prepare_unitree_g1( unitree_urdf_root: str | os.PathLike[str], output_dir: str | os.PathLike[str], @@ -188,6 +209,7 @@ def prepare_unitree_g1( if pelvis is None: raise ValueError("Expected root body `pelvis` in Unitree G1 URDF") pelvis.add_freejoint(name="floating_base_joint") + _add_kinematics_sites(spec) _add_position_actuators(spec, joints) model = spec.compile() diff --git a/scripts/kinematics/test_robot_ik.py b/scripts/kinematics/test_robot_ik.py index c9d42c7d..732f6e1f 100644 --- a/scripts/kinematics/test_robot_ik.py +++ b/scripts/kinematics/test_robot_ik.py @@ -18,10 +18,10 @@ import mujoco from mujoco.viewer import launch_passive -from molmo_spaces.kinematics.parallel.warp_kinematics import SimpleWarpKinematics +from molmo_spaces.configs.robot_configs import BaseRobotConfig from molmo_spaces.kinematics.mujoco_kinematics import MlSpacesKinematics +from molmo_spaces.kinematics.parallel.warp_kinematics import SimpleWarpKinematics from molmo_spaces.molmo_spaces_constants import get_robot_path -from molmo_spaces.configs.robot_configs import BaseRobotConfig def main() -> None: @@ -48,11 +48,16 @@ def main() -> None: action="store_true", help="Use the parallel IK solver instead of the non-parallel one", ) + parser.add_argument( + "--headless-once", + action="store_true", + help="Run one IK solve without opening the viewer", + ) args = parser.parse_args() config_module = importlib.import_module(args.config_module) config_class = getattr(config_module, args.config_class) - robot_config: "BaseRobotConfig" = config_class() + robot_config: BaseRobotConfig = config_class() if args.parallel: kinematics = SimpleWarpKinematics(robot_config) @@ -77,18 +82,46 @@ def main() -> None: robot_view.set_qpos_dict(robot_config.init_qpos) mujoco.mj_forward(model, data) - if args.move_group is None: + if args.move_group is None and robot_view.get_gripper_movegroup_ids(): move_group_id = robot_view.get_gripper_movegroup_ids()[0] move_group = robot_view.get_move_group(move_group_id) + elif args.move_group is None: + arm_move_group_ids = [mg_id for mg_id in robot_view.move_group_ids() if "arm" in mg_id] + if not arm_move_group_ids: + raise ValueError( + "No move group provided, and robot has no gripper or arm move groups to test" + ) + move_group_id = arm_move_group_ids[0] + move_group = robot_view.get_move_group(move_group_id) else: move_group_id = args.move_group move_group = robot_view.get_move_group(move_group_id) + unlocked_move_groups = args.unlocked_move_groups + if ( + unlocked_move_groups is None + and "arm" in move_group_id + and "waist" in robot_view.move_group_ids() + ): + unlocked_move_groups = ["waist", move_group_id] + pose0 = move_group.leaf_frame_to_world.copy() pose1 = pose0.copy() pose0[2, 3] += 0.05 # Move up 5cm pose1[2, 3] -= 0.05 # Move down 5cm + if args.headless_once: + ret = kinematics.ik( + move_group_id, + pose0, + unlocked_move_groups, + robot_view.get_qpos_dict(), + robot_view.base.pose, + ) + if ret is None: + raise SystemExit(f"IK failed for {move_group_id}") + return + with launch_passive(model, data) as viewer: viewer.sync() i = 0 @@ -98,7 +131,7 @@ def main() -> None: ret = kinematics.ik( move_group_id, target_pose, - args.unlocked_move_groups, + unlocked_move_groups, robot_view.get_qpos_dict(), robot_view.base.pose, ) From f64b9a3a165c0d5f8cecc261866cf850da60b6a8 Mon Sep 17 00:00:00 2001 From: Abhishek Parmar Date: Sun, 10 May 2026 18:33:55 +0800 Subject: [PATCH 04/13] feat: use G1 grasp sites for arm IK --- mlspaces_tests/component_tests/test_unitree_g1.py | 4 ++-- molmo_spaces/robots/robot_views/unitree_g1_view.py | 4 ++-- scripts/kinematics/test_robot_ik.py | 8 ++------ 3 files changed, 6 insertions(+), 10 deletions(-) diff --git a/mlspaces_tests/component_tests/test_unitree_g1.py b/mlspaces_tests/component_tests/test_unitree_g1.py index d8fe5dce..980abd6a 100644 --- a/mlspaces_tests/component_tests/test_unitree_g1.py +++ b/mlspaces_tests/component_tests/test_unitree_g1.py @@ -67,12 +67,12 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): assert move_group.n_actuators == n_actuators assert view.get_move_group("left_arm").leaf_frame_type == "site" assert ( - scene_model.site("robot_0/left_wrist_site").id + scene_model.site("robot_0/left_grasp_site").id == view.get_move_group("left_arm").leaf_frame_id ) assert view.get_move_group("right_arm").leaf_frame_type == "site" assert ( - scene_model.site("robot_0/right_wrist_site").id + scene_model.site("robot_0/right_grasp_site").id == view.get_move_group("right_arm").leaf_frame_id ) diff --git a/molmo_spaces/robots/robot_views/unitree_g1_view.py b/molmo_spaces/robots/robot_views/unitree_g1_view.py index 6aa57fe9..57933be7 100644 --- a/molmo_spaces/robots/robot_views/unitree_g1_view.py +++ b/molmo_spaces/robots/robot_views/unitree_g1_view.py @@ -172,7 +172,7 @@ def __init__(self, mj_data: MjData, namespace: str = "") -> None: "left", UNITREE_G1_ARM_JOINTS, "torso_link", - "left_wrist_site", + "left_grasp_site", base, namespace, "site", @@ -182,7 +182,7 @@ def __init__(self, mj_data: MjData, namespace: str = "") -> None: "right", UNITREE_G1_ARM_JOINTS, "torso_link", - "right_wrist_site", + "right_grasp_site", base, namespace, "site", diff --git a/scripts/kinematics/test_robot_ik.py b/scripts/kinematics/test_robot_ik.py index 732f6e1f..89257b10 100644 --- a/scripts/kinematics/test_robot_ik.py +++ b/scripts/kinematics/test_robot_ik.py @@ -98,12 +98,8 @@ def main() -> None: move_group = robot_view.get_move_group(move_group_id) unlocked_move_groups = args.unlocked_move_groups - if ( - unlocked_move_groups is None - and "arm" in move_group_id - and "waist" in robot_view.move_group_ids() - ): - unlocked_move_groups = ["waist", move_group_id] + if unlocked_move_groups is None and "arm" in move_group_id: + unlocked_move_groups = [move_group_id] pose0 = move_group.leaf_frame_to_world.copy() pose1 = pose0.copy() From f0fa9c6f9f285ad4d93e80e40457ccb7d5ca49db Mon Sep 17 00:00:00 2001 From: Abhishek Parmar Date: Sun, 10 May 2026 18:56:08 +0800 Subject: [PATCH 05/13] feat: add G1 right-arm pick datagen --- .../component_tests/test_unitree_g1.py | 89 ++++++++++++++++++- molmo_spaces/configs/camera_configs.py | 35 ++++++++ molmo_spaces/configs/policy_configs.py | 24 +++++ molmo_spaces/configs/robot_configs.py | 21 ++++- .../object_manipulation_datagen_configs.py | 33 +++++++ .../pick_planner_policy.py | 66 ++++++++++++++ .../robots/robot_views/unitree_g1_view.py | 81 +++++++++++++++++ 7 files changed, 347 insertions(+), 2 deletions(-) diff --git a/mlspaces_tests/component_tests/test_unitree_g1.py b/mlspaces_tests/component_tests/test_unitree_g1.py index 980abd6a..d44131cd 100644 --- a/mlspaces_tests/component_tests/test_unitree_g1.py +++ b/mlspaces_tests/component_tests/test_unitree_g1.py @@ -7,7 +7,10 @@ import pytest import molmo_spaces.molmo_spaces_constants as constants -from molmo_spaces.configs.robot_configs import UnitreeG1Dex1RobotConfig +from molmo_spaces.configs.robot_configs import ( + UnitreeG1Dex1RobotConfig, + UnitreeG1RightArmPickRobotConfig, +) from molmo_spaces.data_generation.config_registry import get_config_class from molmo_spaces.kinematics.mujoco_kinematics import MlSpacesKinematics from molmo_spaces.robots.unitree_g1 import UnitreeG1Robot @@ -89,6 +92,15 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): datagen_config_cls = get_config_class("UnitreeG1SceneSmokeDataGenConfig") datagen_config = datagen_config_cls() assert isinstance(datagen_config.robot_config, UnitreeG1Dex1RobotConfig) + pick_datagen_config_cls = get_config_class("UnitreeG1RightArmPickDataGenConfig") + pick_datagen_config = pick_datagen_config_cls() + assert isinstance(pick_datagen_config.robot_config, UnitreeG1RightArmPickRobotConfig) + assert pick_datagen_config.policy_config.policy_cls.__name__ == ( + "UnitreeG1RightArmPickPlannerPolicy" + ) + assert not pick_datagen_config.policy_config.filter_colliding_grasps + assert not pick_datagen_config.policy_config.filter_feasible_grasps + assert len(pick_datagen_config.camera_config.cameras) == 2 base_scene_path = constants.ABS_PATH_OF_TOP_LEVEL_MOLMO_SPACES_DIR / ( "molmo_spaces/resources/base_scene.xml" @@ -105,6 +117,81 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): ) datagen_scene_model = datagen_scene_spec.compile() assert datagen_scene_model.body("robot_0/pelvis").id >= 0 + pick_scene_spec = mujoco.MjSpec.from_file(str(base_scene_path)) + pick_robot_spec = mujoco.MjSpec.from_file(str(xml_path)) + pick_datagen_config.robot_config.robot_cls.add_robot_to_scene( + pick_datagen_config.robot_config, + pick_scene_spec, + pick_robot_spec, + pick_datagen_config.robot_config.robot_namespace, + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0, 0.0], + ) + pick_scene_model = pick_scene_spec.compile() + pick_scene_data = mujoco.MjData(pick_scene_model) + pick_view = pick_datagen_config.robot_config.robot_view_factory( + pick_scene_data, pick_datagen_config.robot_config.robot_namespace + ) + assert pick_view.move_group_ids() == ["base", "right_arm", "gripper"] + assert pick_view.get_gripper_movegroup_ids() == ["gripper"] + assert pick_view.get_move_group("right_arm").leaf_frame_type == "site" + assert pick_view.get_move_group("gripper").leaf_frame_type == "site" + assert ( + pick_scene_model.site("robot_0/right_grasp_site").id + == pick_view.get_move_group("right_arm").leaf_frame_id + ) + assert ( + pick_scene_model.site("robot_0/right_grasp_site").id + == pick_view.get_move_group("gripper").leaf_frame_id + ) + + gripper = pick_view.get_gripper("gripper") + gripper.set_gripper_ctrl_open(True) + assert np.all(gripper.ctrl <= gripper.ctrl_limits[:, 1]) + assert np.all(gripper.ctrl >= gripper.ctrl_limits[:, 0]) + gripper.set_gripper_ctrl_open(False) + assert np.all(gripper.ctrl <= gripper.ctrl_limits[:, 1]) + assert np.all(gripper.ctrl >= gripper.ctrl_limits[:, 0]) + + pick_kinematics = MlSpacesKinematics(pick_datagen_config.robot_config) + pick_fk = pick_kinematics.fk(pick_datagen_config.robot_config.init_qpos, np.eye(4)) + for move_group_id in ("right_arm", "gripper"): + assert pick_fk[move_group_id].shape == (4, 4) + assert np.isfinite(pick_fk[move_group_id]).all() + + target_pose = pick_fk["gripper"].copy() + target_pose[2, 3] += 0.02 + pick_ik = pick_kinematics.ik( + "gripper", + target_pose, + ["right_arm"], + pick_datagen_config.robot_config.init_qpos, + np.eye(4), + max_iter=250, + ) + assert pick_ik is not None + assert set(pick_ik) == {"base", "right_arm", "gripper"} + + pick_robot = UnitreeG1Robot( + pick_scene_data, SimpleNamespace(robot_config=pick_datagen_config.robot_config) + ) + pick_robot.reset() + mujoco.mj_forward(pick_scene_model, pick_scene_data) + pick_action = { + move_group_id: pick_robot.robot_view.get_move_group(move_group_id).joint_pos.copy() + for move_group_id in pick_robot.controllers + } + pick_action["right_arm"] = pick_action["right_arm"].copy() + pick_action["right_arm"][3] = 0.1 + pick_action["gripper"] = np.array([-0.02, -0.02]) + pick_robot.update_control(pick_action) + pick_robot.compute_control() + for _ in range(10): + mujoco.mj_step(pick_scene_model, pick_scene_data) + assert np.isfinite(pick_scene_data.qpos).all() + assert np.isfinite(pick_scene_data.ctrl).all() + assert pick_robot.state_dim == 16 + assert pick_robot.action_dim(list(pick_robot.controllers)) == 9 robot = UnitreeG1Robot(scene_data, SimpleNamespace(robot_config=config)) robot.reset() diff --git a/molmo_spaces/configs/camera_configs.py b/molmo_spaces/configs/camera_configs.py index 62d25224..59c33f1d 100644 --- a/molmo_spaces/configs/camera_configs.py +++ b/molmo_spaces/configs/camera_configs.py @@ -349,6 +349,40 @@ class FrankaRandomizedD405D455CameraSystem(CameraSystemConfig): ] +class UnitreeG1RightArmPickCameraSystem(CameraSystemConfig): + """Exocentric-only cameras for Unitree G1 right-arm pick smoke datagen.""" + + img_resolution: tuple[int, int] = (624, 352) + cameras: list[AllCameraTypes] = [ + RandomizedExocentricCameraConfig( + name="exo_camera_1", + distance_range=(0.6, 1.4), + height_range=(0.5, 1.1), + azimuth_range=(-np.pi, np.pi), + fov_range=(55, 80), + lookat_noise_range=(-0.08, 0.08), + visibility_constraints={ + "__task_objects__": 0.0001, + "__gripper__": 0.0001, + }, + allow_relaxed_constraints=True, + ), + RandomizedExocentricCameraConfig( + name="exo_camera_2", + distance_range=(0.7, 1.6), + height_range=(0.7, 1.3), + azimuth_range=(-np.pi, np.pi), + fov_range=(55, 85), + lookat_noise_range=(-0.08, 0.08), + visibility_constraints={ + "__task_objects__": 0.0001, + "__gripper__": 0.0001, + }, + allow_relaxed_constraints=True, + ), + ] + + class FrankaDroidCameraSystem(CameraSystemConfig): """Camera system for Franka with DROID-style fixed cameras. @@ -953,6 +987,7 @@ class FrankaEvalCameraSystem(CameraSystemConfig): | FrankaGoProD405RandomizedCameraSystem | FrankaRobotiq2f85CameraSystem | FrankaEvalCameraSystem + | UnitreeG1RightArmPickCameraSystem | I2rtYamCameraSystem | BimanualYamCameraSystem | FrankaEvalCameraSystem diff --git a/molmo_spaces/configs/policy_configs.py b/molmo_spaces/configs/policy_configs.py index 2d81b4d3..fae59138 100644 --- a/molmo_spaces/configs/policy_configs.py +++ b/molmo_spaces/configs/policy_configs.py @@ -139,6 +139,30 @@ def model_post_init(self, __context) -> None: self.policy_cls = PickPlannerPolicy +class UnitreeG1RightArmPickPlannerPolicyConfig(PickPlannerPolicyConfig): + """Pick planner config for Unitree G1 right-arm-only smoke datagen.""" + + policy_cls: type = None + filter_colliding_grasps: bool = False + filter_feasible_grasps: bool = False + grasp_feasibility_batch_size: int = 32 + grasp_feasibility_max_grasps: int = 32 + grasp_collision_batch_size: int = 32 + grasp_collision_max_grasps: int = 64 + max_retries: int = 0 + max_sequential_ik_failures: int = 4 + phase_timeout: float = 20.0 + + def model_post_init(self, __context) -> None: + """Set policy_cls after initialization to avoid circular imports.""" + super().model_post_init(__context) + from molmo_spaces.policy.solvers.object_manipulation.pick_planner_policy import ( + UnitreeG1RightArmPickPlannerPolicy, + ) + + self.policy_cls = UnitreeG1RightArmPickPlannerPolicy + + class PickAndPlacePlannerPolicyConfig(ObjectManipulationPlannerPolicyConfig): policy_cls: type = None # Will be set in model_post_init to avoid circular imports move_settle_time: float = 0.5 diff --git a/molmo_spaces/configs/robot_configs.py b/molmo_spaces/configs/robot_configs.py index 8547726c..5f7b8b8f 100644 --- a/molmo_spaces/configs/robot_configs.py +++ b/molmo_spaces/configs/robot_configs.py @@ -36,7 +36,10 @@ from molmo_spaces.robots.robot_views.mobile_franka_droid_view import MobileFrankaDroidRobotView from molmo_spaces.robots.robot_views.rby1_view import RBY1RobotView from molmo_spaces.robots.robot_views.rum_gripper_view import FloatingRUMRobotView -from molmo_spaces.robots.robot_views.unitree_g1_view import UnitreeG1RobotView +from molmo_spaces.robots.robot_views.unitree_g1_view import ( + UnitreeG1RightArmPickRobotView, + UnitreeG1RobotView, +) from molmo_spaces.robots.unitree_g1 import UnitreeG1Robot @@ -455,3 +458,19 @@ class UnitreeG1Dex1RobotConfig(BaseRobotConfig): "right_hand": "joint_position", } gravcomp: bool = True + + +class UnitreeG1RightArmPickRobotConfig(UnitreeG1Dex1RobotConfig): + """Unitree G1 config scoped to right-arm pick datagen.""" + + robot_view_factory: RobotViewFactory | None = UnitreeG1RightArmPickRobotView + init_qpos: dict[str, list[float]] = { + "base": [0.0, 0.0, 0.793, 1.0, 0.0, 0.0, 0.0], + "right_arm": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + "gripper": [0.0245, 0.0245], + } + command_mode: dict[str, str | None] = { + "base": None, + "right_arm": "joint_position", + "gripper": "joint_position", + } diff --git a/molmo_spaces/data_generation/config/object_manipulation_datagen_configs.py b/molmo_spaces/data_generation/config/object_manipulation_datagen_configs.py index 4845cf08..f516642f 100644 --- a/molmo_spaces/data_generation/config/object_manipulation_datagen_configs.py +++ b/molmo_spaces/data_generation/config/object_manipulation_datagen_configs.py @@ -27,12 +27,14 @@ FrankaRandomizedD405D455CameraSystem, FrankaRandomizedDroidCameraSystem, RBY1GoProD455CameraSystem, + UnitreeG1RightArmPickCameraSystem, ) from molmo_spaces.configs.policy_configs import ( CuroboOpenClosePlannerPolicyConfig, CuroboPickAndPlacePlannerPolicyConfig, OpenClosePlannerPolicyConfig, PickPlannerPolicyConfig, + UnitreeG1RightArmPickPlannerPolicyConfig, ) from molmo_spaces.configs.robot_configs import ( FloatingRUMRobotConfig, @@ -40,6 +42,7 @@ RBY1MConfig, RBY1MOpenCloseConfig, UnitreeG1Dex1RobotConfig, + UnitreeG1RightArmPickRobotConfig, ) from molmo_spaces.configs.task_sampler_configs import ( OpenTaskSamplerConfig, @@ -140,6 +143,36 @@ def tag(self) -> str: return "unitree_g1_scene_smoke_datagen" +@register_config("UnitreeG1RightArmPickDataGenConfig") +class UnitreeG1RightArmPickDataGenConfig(PickBaseConfig): + """Right-arm-only Unitree G1 pick smoke datagen.""" + + robot_config: UnitreeG1RightArmPickRobotConfig = UnitreeG1RightArmPickRobotConfig() + camera_config: UnitreeG1RightArmPickCameraSystem = UnitreeG1RightArmPickCameraSystem() + task_sampler_config: PickTaskSamplerConfig = PickTaskSamplerConfig( + task_sampler_class=PickTaskSampler, + house_inds=[1], + samples_per_house=5, + robot_object_z_offset=0.0, + robot_object_z_offset_random_min=0.0, + robot_object_z_offset_random_max=0.0, + base_pose_sampling_radius_range=(0.2, 0.8), + robot_safety_radius=0.25, + check_robot_placement_visibility=False, + ) + policy_config: UnitreeG1RightArmPickPlannerPolicyConfig = ( + UnitreeG1RightArmPickPlannerPolicyConfig() + ) + task_horizon: int = 120 + filter_for_successful_trajectories: bool = False + profile: bool = False + output_dir: Path = ASSETS_DIR / "experiment_output" / "datagen" / "unitree_g1_right_arm_pick_v1" + + @property + def tag(self) -> str: + return "unitree_g1_right_arm_pick_datagen" + + @register_config("FrankaPickAndPlaceDataGenConfig") class FrankaPickAndPlaceDataGenConfig(PickAndPlaceDataGenConfig): robot_config: BaseRobotConfig = FrankaRobotConfig() diff --git a/molmo_spaces/policy/solvers/object_manipulation/pick_planner_policy.py b/molmo_spaces/policy/solvers/object_manipulation/pick_planner_policy.py index 6355618a..9c290680 100644 --- a/molmo_spaces/policy/solvers/object_manipulation/pick_planner_policy.py +++ b/molmo_spaces/policy/solvers/object_manipulation/pick_planner_policy.py @@ -169,3 +169,69 @@ def _compute_target_poses(self) -> dict[str, np.ndarray]: log.info(f"Planning completed. w/ {len(target_poses)} steps\n") return target_poses + + +class UnitreeG1RightArmPickPlannerPolicy(PickPlannerPolicy): + """Pick planner variant that solves IK through the G1 right arm only.""" + + def _tcp_to_jp_fn(self, mg_id: str, target_pose: np.ndarray) -> dict[str, np.ndarray]: + kinematics = self.task.env.current_robot.kinematics + + jp = kinematics.ik( + mg_id, + target_pose, + ["right_arm"], + self.robot_view.get_qpos_dict(), + self.robot_view.base.pose, + ) + + action = self.robot_view.get_ctrl_dict() + if jp is not None: + self.sequential_ik_failures = 0 + action["right_arm"] = jp["right_arm"] + else: + self.sequential_ik_failures += 1 + log.info(f"⚠️ IK failed, holding current position, fails:{self.sequential_ik_failures}") + if self.sequential_ik_failures >= self.policy_config.max_sequential_ik_failures: + log.info("❌ Too many sequential IK failures, triggering retry.") + return self._handle_failure() + + return action + + def check_feasible_ik(self, pose: np.ndarray) -> bool: + if not self.policy_config.filter_feasible_grasps: + if pose.ndim > 2: + return np.ones(pose.shape[0], dtype=bool) + return True + + robot_view = self.task.env.current_robot.robot_view + kinematics = self.task.env.current_robot.kinematics + gripper_mg_id = robot_view.get_gripper_movegroup_ids()[0] + + if pose.ndim > 2: + return np.array( + [ + kinematics.ik( + gripper_mg_id, + single_pose, + ["right_arm"], + robot_view.get_qpos_dict(), + robot_view.base.pose, + max_iter=100, + ) + is not None + for single_pose in pose + ], + dtype=bool, + ) + + assert pose.shape == (4, 4) + jp_dict = kinematics.ik( + gripper_mg_id, + pose, + ["right_arm"], + robot_view.get_qpos_dict(), + robot_view.base.pose, + max_iter=100, + ) + return jp_dict is not None diff --git a/molmo_spaces/robots/robot_views/unitree_g1_view.py b/molmo_spaces/robots/robot_views/unitree_g1_view.py index 57933be7..525824de 100644 --- a/molmo_spaces/robots/robot_views/unitree_g1_view.py +++ b/molmo_spaces/robots/robot_views/unitree_g1_view.py @@ -7,6 +7,7 @@ from molmo_spaces.robots.robot_views.abstract import ( FreeJointRobotBaseGroup, + GripperGroup, MJCFFrameMixin, RobotBaseGroup, RobotView, @@ -133,6 +134,55 @@ def __init__( ) +class UnitreeG1DexGripperGroup(MJCFFrameMixin, GripperGroup): + """Dex1.1 two-finger hand exposed through the MolmoSpaces gripper API.""" + + OPEN_CTRL = np.array([0.0245, 0.0245]) + CLOSED_CTRL = np.array([-0.02, -0.02]) + + def __init__( + self, + mj_data: MjData, + side: Literal["left", "right"], + base: RobotBaseGroup, + namespace: str = "", + ) -> None: + model = mj_data.model + joint_names = [f"{side}_{joint_suffix}" for joint_suffix in UNITREE_G1_HAND_JOINTS] + joint_ids = [model.joint(f"{namespace}{joint_name}").id for joint_name in joint_names] + act_ids = [model.actuator(f"{namespace}{joint_name}").id for joint_name in joint_names] + root_body_id = model.body(f"{namespace}{side}_wrist_yaw_link").id + self._ee_site_id = model.site(f"{namespace}{side}_grasp_site").id + super().__init__(mj_data, joint_ids, act_ids, root_body_id, base) + + @property + def leaf_frame_id(self) -> int: + return self._ee_site_id + + @property + def leaf_frame_type(self): + return "site" + + def set_gripper_ctrl_open(self, open: bool) -> None: + self.ctrl = self.OPEN_CTRL if open else self.CLOSED_CTRL + + @property + def inter_finger_dist_range(self) -> tuple[float, float]: + return 0.0, float(np.sum(self.OPEN_CTRL - self.CLOSED_CTRL)) + + @property + def inter_finger_dist(self) -> float: + return float(np.sum(np.clip(self.joint_pos - self.CLOSED_CTRL, 0.0, None))) + + @property + def leaf_frame_to_world(self) -> np.ndarray: + return site_pose(self.mj_data, self._ee_site_id) + + @property + def root_frame_to_world(self) -> np.ndarray: + return self.leaf_frame_to_world + + class UnitreeG1RobotView(RobotView): """Move-group view for the Unitree G1 Dex1.1 model.""" @@ -215,3 +265,34 @@ def name(self) -> str: @property def base(self) -> UnitreeG1BaseGroup: return self._move_groups["base"] + + +class UnitreeG1RightArmPickRobotView(RobotView): + """Policy-facing G1 view for right-arm-only pick attempts.""" + + def __init__(self, mj_data: MjData, namespace: str = "") -> None: + self._namespace = namespace + base = UnitreeG1BaseGroup(mj_data, namespace=namespace) + move_groups = { + "base": base, + "right_arm": UnitreeG1SideMoveGroup( + mj_data, + "right", + UNITREE_G1_ARM_JOINTS, + "torso_link", + "right_grasp_site", + base, + namespace, + "site", + ), + "gripper": UnitreeG1DexGripperGroup(mj_data, "right", base, namespace), + } + super().__init__(mj_data, move_groups) + + @property + def name(self) -> str: + return f"{self._namespace}unitree_g1_29dof_dex1_1_right_arm_pick" + + @property + def base(self) -> UnitreeG1BaseGroup: + return self._move_groups["base"] From c5a37ec5eaa6cd7544006be0584d404153368ad8 Mon Sep 17 00:00:00 2001 From: Abhishek Parmar Date: Mon, 11 May 2026 02:12:18 +0800 Subject: [PATCH 06/13] fix: stabilize G1 right-arm pick placement --- .../component_tests/test_unitree_g1.py | 18 +++++++++- molmo_spaces/configs/robot_configs.py | 2 ++ .../object_manipulation_datagen_configs.py | 4 +-- molmo_spaces/robots/unitree_g1.py | 33 ++++++++++++++++++- molmo_spaces/tasks/pick_task_sampler.py | 32 +++++++++++++----- 5 files changed, 77 insertions(+), 12 deletions(-) diff --git a/mlspaces_tests/component_tests/test_unitree_g1.py b/mlspaces_tests/component_tests/test_unitree_g1.py index d44131cd..247668da 100644 --- a/mlspaces_tests/component_tests/test_unitree_g1.py +++ b/mlspaces_tests/component_tests/test_unitree_g1.py @@ -14,6 +14,7 @@ from molmo_spaces.data_generation.config_registry import get_config_class from molmo_spaces.kinematics.mujoco_kinematics import MlSpacesKinematics from molmo_spaces.robots.unitree_g1 import UnitreeG1Robot +from molmo_spaces.tasks.pick_task_sampler import UnitreeG1RightArmPickTaskSampler from scripts.assets.prepare_unitree_g1 import ROBOT_ASSET_NAME, prepare_unitree_g1 KINEMATICS_SITE_NAMES = [ @@ -95,6 +96,11 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): pick_datagen_config_cls = get_config_class("UnitreeG1RightArmPickDataGenConfig") pick_datagen_config = pick_datagen_config_cls() assert isinstance(pick_datagen_config.robot_config, UnitreeG1RightArmPickRobotConfig) + assert pick_datagen_config.robot_config.pin_base_in_place + assert ( + pick_datagen_config.task_sampler_config.task_sampler_class + is UnitreeG1RightArmPickTaskSampler + ) assert pick_datagen_config.policy_config.policy_cls.__name__ == ( "UnitreeG1RightArmPickPlannerPolicy" ) @@ -177,6 +183,11 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): ) pick_robot.reset() mujoco.mj_forward(pick_scene_model, pick_scene_data) + initial_pick_base_pose = pick_robot.robot_view.base.pose.copy() + initial_pick_base_pose[:3, 3] = [0.2, -0.3, initial_pick_base_pose[2, 3]] + pick_robot.robot_view.base.pose = initial_pick_base_pose + pick_robot.robot_view.base.joint_vel = np.zeros(pick_robot.robot_view.base.vel_dim) + mujoco.mj_forward(pick_scene_model, pick_scene_data) pick_action = { move_group_id: pick_robot.robot_view.get_move_group(move_group_id).joint_pos.copy() for move_group_id in pick_robot.controllers @@ -185,11 +196,16 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): pick_action["right_arm"][3] = 0.1 pick_action["gripper"] = np.array([-0.02, -0.02]) pick_robot.update_control(pick_action) - pick_robot.compute_control() for _ in range(10): + pick_robot.compute_control() mujoco.mj_step(pick_scene_model, pick_scene_data) assert np.isfinite(pick_scene_data.qpos).all() assert np.isfinite(pick_scene_data.ctrl).all() + assert np.allclose( + pick_robot.robot_view.base.pose[:3, 3], + initial_pick_base_pose[:3, 3], + atol=1e-3, + ) assert pick_robot.state_dim == 16 assert pick_robot.action_dim(list(pick_robot.controllers)) == 9 diff --git a/molmo_spaces/configs/robot_configs.py b/molmo_spaces/configs/robot_configs.py index 5f7b8b8f..8612b104 100644 --- a/molmo_spaces/configs/robot_configs.py +++ b/molmo_spaces/configs/robot_configs.py @@ -458,12 +458,14 @@ class UnitreeG1Dex1RobotConfig(BaseRobotConfig): "right_hand": "joint_position", } gravcomp: bool = True + pin_base_in_place: bool = False class UnitreeG1RightArmPickRobotConfig(UnitreeG1Dex1RobotConfig): """Unitree G1 config scoped to right-arm pick datagen.""" robot_view_factory: RobotViewFactory | None = UnitreeG1RightArmPickRobotView + pin_base_in_place: bool = True init_qpos: dict[str, list[float]] = { "base": [0.0, 0.0, 0.793, 1.0, 0.0, 0.0, 0.0], "right_arm": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], diff --git a/molmo_spaces/data_generation/config/object_manipulation_datagen_configs.py b/molmo_spaces/data_generation/config/object_manipulation_datagen_configs.py index f516642f..54035da7 100644 --- a/molmo_spaces/data_generation/config/object_manipulation_datagen_configs.py +++ b/molmo_spaces/data_generation/config/object_manipulation_datagen_configs.py @@ -64,7 +64,7 @@ PickAndPlaceMultiTaskSampler, PickAndPlaceTaskSampler, ) -from molmo_spaces.tasks.pick_task_sampler import PickTaskSampler +from molmo_spaces.tasks.pick_task_sampler import PickTaskSampler, UnitreeG1RightArmPickTaskSampler from molmo_spaces.utils.constants.object_constants import PICK_AND_PLACE_OBJECTS from molmo_spaces.utils.synset_utils import get_valid_pickupable_obja_uids @@ -150,7 +150,7 @@ class UnitreeG1RightArmPickDataGenConfig(PickBaseConfig): robot_config: UnitreeG1RightArmPickRobotConfig = UnitreeG1RightArmPickRobotConfig() camera_config: UnitreeG1RightArmPickCameraSystem = UnitreeG1RightArmPickCameraSystem() task_sampler_config: PickTaskSamplerConfig = PickTaskSamplerConfig( - task_sampler_class=PickTaskSampler, + task_sampler_class=UnitreeG1RightArmPickTaskSampler, house_inds=[1], samples_per_house=5, robot_object_z_offset=0.0, diff --git a/molmo_spaces/robots/unitree_g1.py b/molmo_spaces/robots/unitree_g1.py index 11961782..c09698c1 100644 --- a/molmo_spaces/robots/unitree_g1.py +++ b/molmo_spaces/robots/unitree_g1.py @@ -32,6 +32,7 @@ def __init__( config.robot_config, self._kinematics, ) + self._pinned_base_pose: np.ndarray | None = None self._controllers = { move_group_id: JointPosController(self._robot_view.get_move_group(move_group_id)) for move_group_id in self._robot_view.move_group_ids() @@ -82,21 +83,51 @@ def update_control(self, action_command_dict: dict[str, Any]) -> None: controller.set_to_stationary() def compute_control(self) -> None: + self._pin_base_if_configured() for controller in self.controllers.values(): ctrl_inputs = controller.compute_ctrl_inputs() controller.robot_move_group.ctrl = ctrl_inputs + self._pin_base_if_configured() def set_joint_pos(self, robot_joint_pos_dict) -> None: for mg_id, joint_pos in robot_joint_pos_dict.items(): self._robot_view.get_move_group(mg_id).joint_pos = joint_pos + if mg_id == "base": + self._set_pinned_base_pose_if_configured() def set_world_pose(self, robot_world_pose) -> None: self._robot_view.base.pose = robot_world_pose + self._set_pinned_base_pose_if_configured() + self._zero_base_velocity_if_configured() def reset(self) -> None: for mg_id, default_pos in self.exp_config.robot_config.init_qpos.items(): if mg_id in self._robot_view.move_group_ids(): - self._robot_view.get_move_group(mg_id).joint_pos = np.asarray(default_pos) + move_group = self._robot_view.get_move_group(mg_id) + move_group.joint_pos = np.asarray(default_pos) + move_group.joint_vel = np.zeros(move_group.vel_dim) + if self._should_pin_base(): + self._pinned_base_pose = None + self._zero_base_velocity_if_configured() + + def _should_pin_base(self) -> bool: + return bool(getattr(self.exp_config.robot_config, "pin_base_in_place", False)) + + def _set_pinned_base_pose_if_configured(self) -> None: + if self._should_pin_base(): + self._pinned_base_pose = self._robot_view.base.pose.copy() + + def _zero_base_velocity_if_configured(self) -> None: + if self._should_pin_base(): + self._robot_view.base.joint_vel = np.zeros(self._robot_view.base.vel_dim) + + def _pin_base_if_configured(self) -> None: + if not self._should_pin_base(): + return + if self._pinned_base_pose is None: + self._pinned_base_pose = self._robot_view.base.pose.copy() + self._robot_view.base.pose = self._pinned_base_pose + self._robot_view.base.joint_vel = np.zeros(self._robot_view.base.vel_dim) @staticmethod def robot_model_root_name() -> str: diff --git a/molmo_spaces/tasks/pick_task_sampler.py b/molmo_spaces/tasks/pick_task_sampler.py index fe949eb9..64b38f29 100644 --- a/molmo_spaces/tasks/pick_task_sampler.py +++ b/molmo_spaces/tasks/pick_task_sampler.py @@ -962,14 +962,7 @@ def _sample_and_place_robot(self, env: CPUMujocoEnv) -> None: else: raise ValueError(f"Invalid pickup object type: {type(pickup_obj)}") - initial_robot_z = ( - target_pos[2] - + self.config.task_sampler_config.robot_object_z_offset - + np.random.uniform( - self.config.task_sampler_config.robot_object_z_offset_random_min, - self.config.task_sampler_config.robot_object_z_offset_random_max, - ) - ) + initial_robot_z = self._initial_robot_z_for_pickup(target_pos, robot_view) # place robot near receptacle - this is the expensive call with collision/visibility checks if self._datagen_profiler is not None: @@ -1006,6 +999,18 @@ def _sample_and_place_robot(self, env: CPUMujocoEnv) -> None: log.info(f"Supporting receptacle: {self.config.task_config.receptacle_name}") + def _initial_robot_z_for_pickup(self, target_pos: np.ndarray, robot_view) -> float: + """Return the robot base Z used by placement sampling for a pickup target.""" + del robot_view + return float( + target_pos[2] + + self.config.task_sampler_config.robot_object_z_offset + + np.random.uniform( + self.config.task_sampler_config.robot_object_z_offset_random_min, + self.config.task_sampler_config.robot_object_z_offset_random_max, + ) + ) + def _place_target_near_object( self, env: CPUMujocoEnv, object_pos: np.ndarray, placement_region=None ) -> None: @@ -1179,3 +1184,14 @@ def add_pickup_target( ) return spec + + +class UnitreeG1RightArmPickTaskSampler(PickTaskSampler): + """Pick sampler variant that keeps the Unitree G1 pelvis at its standing height.""" + + def _initial_robot_z_for_pickup(self, target_pos: np.ndarray, robot_view) -> float: + del target_pos, robot_view + base_qpos = self.config.robot_config.init_qpos.get("base") + if base_qpos is None or len(base_qpos) < 3: + raise ValueError("Unitree G1 pick sampling requires a base init_qpos with xyz.") + return float(base_qpos[2]) From 19e922540e51f7c5950e3e9f883408fdd16c233f Mon Sep 17 00:00:00 2001 From: Abhishek Parmar Date: Mon, 11 May 2026 02:23:48 +0800 Subject: [PATCH 07/13] fix: prevent G1 default pose floor collision --- .../component_tests/test_unitree_g1.py | 36 +++++++++++++++++ molmo_spaces/robots/unitree_g1.py | 40 ++++++++++++++++++- scripts/assets/prepare_unitree_g1.py | 2 + 3 files changed, 77 insertions(+), 1 deletion(-) diff --git a/mlspaces_tests/component_tests/test_unitree_g1.py b/mlspaces_tests/component_tests/test_unitree_g1.py index 247668da..af252018 100644 --- a/mlspaces_tests/component_tests/test_unitree_g1.py +++ b/mlspaces_tests/component_tests/test_unitree_g1.py @@ -25,6 +25,25 @@ ] +def _has_deep_robot_floor_penetration( + model: mujoco.MjModel, + data: mujoco.MjData, + threshold_m: float = 0.01, +) -> bool: + for i in range(data.ncon): + contact = data.contact[i] + if contact.dist > -threshold_m: + continue + geom1 = model.geom(contact.geom1) + geom2 = model.geom(contact.geom2) + body1 = model.body(model.geom_bodyid[contact.geom1]).name + body2 = model.body(model.geom_bodyid[contact.geom2]).name + names = " ".join([geom1.name, geom2.name, body1, body2]).lower() + if "robot_0/" in names and "floor" in names: + return True + return False + + def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): unitree_root = os.environ.get("UNITREE_URDF_ROOT") if unitree_root is None: @@ -37,6 +56,9 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): model = mujoco.MjModel.from_xml_path(str(xml_path)) assert (model.nq, model.nv, model.nu) == (40, 39, 33) + base_joint_id = model.joint("floating_base_joint").id + base_qposadr = model.jnt_qposadr[base_joint_id] + assert model.qpos0[base_qposadr + 2] == pytest.approx(0.793) for site_name in KINEMATICS_SITE_NAMES: assert model.site(site_name).id >= 0 @@ -51,8 +73,14 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): [0.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0], ) + config.robot_cls.apply_control_overrides(scene_spec, config) scene_model = scene_spec.compile() scene_data = mujoco.MjData(scene_model) + scene_base_joint_id = scene_model.joint("robot_0/floating_base_joint").id + scene_base_qposadr = scene_model.jnt_qposadr[scene_base_joint_id] + assert scene_model.qpos0[scene_base_qposadr + 2] == pytest.approx(0.793) + mujoco.mj_forward(scene_model, scene_data) + assert not _has_deep_robot_floor_penetration(scene_model, scene_data) view = config.robot_view_factory(scene_data, config.robot_namespace) expected_counts = { @@ -121,6 +149,10 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): [0.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0], ) + datagen_config.robot_config.robot_cls.apply_control_overrides( + datagen_scene_spec, + datagen_config.robot_config, + ) datagen_scene_model = datagen_scene_spec.compile() assert datagen_scene_model.body("robot_0/pelvis").id >= 0 pick_scene_spec = mujoco.MjSpec.from_file(str(base_scene_path)) @@ -133,6 +165,10 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): [0.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0], ) + pick_datagen_config.robot_config.robot_cls.apply_control_overrides( + pick_scene_spec, + pick_datagen_config.robot_config, + ) pick_scene_model = pick_scene_spec.compile() pick_scene_data = mujoco.MjData(pick_scene_model) pick_view = pick_datagen_config.robot_config.robot_view_factory( diff --git a/molmo_spaces/robots/unitree_g1.py b/molmo_spaces/robots/unitree_g1.py index c09698c1..6f25c1f3 100644 --- a/molmo_spaces/robots/unitree_g1.py +++ b/molmo_spaces/robots/unitree_g1.py @@ -3,7 +3,7 @@ from typing import TYPE_CHECKING, Any import numpy as np -from mujoco import MjData +from mujoco import MjData, MjSpec from molmo_spaces.controllers.abstract import Controller from molmo_spaces.controllers.joint_pos import JointPosController @@ -39,6 +39,7 @@ def __init__( if move_group_id != "base" and config.robot_config.command_mode.get(move_group_id) == "joint_position" } + self.reset() @property def namespace(self): @@ -60,6 +61,40 @@ def parallel_kinematics(self): def controllers(self) -> dict[str, Controller]: return self._controllers + @classmethod + def add_robot_to_scene( + cls, + robot_config, + spec: MjSpec, + robot_spec: MjSpec, + prefix: str, + pos: list[float], + quat: list[float], + randomize_textures: bool = False, + ) -> None: + super().add_robot_to_scene( + robot_config, + spec, + robot_spec, + prefix, + pos, + quat, + randomize_textures, + ) + root_body = spec.body(prefix + cls.robot_model_root_name()) + if root_body is None: + raise ValueError(f"Robot root body {prefix}{cls.robot_model_root_name()} not found") + root_body.pos = robot_config.default_world_pose[:3] + root_body.quat = robot_config.default_world_pose[3:7] + + @classmethod + def apply_control_overrides(cls, spec: MjSpec, robot_config) -> None: + super().apply_control_overrides(spec, robot_config) + if robot_config.gravcomp: + root_body = spec.body(robot_config.robot_namespace + cls.robot_model_root_name()) + if root_body is not None: + root_body.gravcomp = 1.0 + @property def state_dim(self) -> int: return sum( @@ -101,6 +136,9 @@ def set_world_pose(self, robot_world_pose) -> None: self._zero_base_velocity_if_configured() def reset(self) -> None: + self.mj_data.qpos[:] = self.mj_model.qpos0 + self.mj_data.qvel[:] = 0.0 + self.mj_data.ctrl[:] = 0.0 for mg_id, default_pos in self.exp_config.robot_config.init_qpos.items(): if mg_id in self._robot_view.move_group_ids(): move_group = self._robot_view.get_move_group(mg_id) diff --git a/scripts/assets/prepare_unitree_g1.py b/scripts/assets/prepare_unitree_g1.py index e64b612c..6e914fbc 100644 --- a/scripts/assets/prepare_unitree_g1.py +++ b/scripts/assets/prepare_unitree_g1.py @@ -22,6 +22,7 @@ SOURCE_URDF_REL_PATH = Path("g1_description/g1_29dof_mode_15_with_dex1_1.urdf") SOURCE_MESH_DIR_REL_PATH = Path("g1_description/meshes") OUTPUT_XML_NAME = "model.xml" +DEFAULT_PELVIS_HEIGHT_M = 0.793 GAIN_BY_JOINT_GROUP = { @@ -208,6 +209,7 @@ def prepare_unitree_g1( pelvis = spec.body("pelvis") if pelvis is None: raise ValueError("Expected root body `pelvis` in Unitree G1 URDF") + pelvis.pos = [0.0, 0.0, DEFAULT_PELVIS_HEIGHT_M] pelvis.add_freejoint(name="floating_base_joint") _add_kinematics_sites(spec) _add_position_actuators(spec, joints) From e196d1076ba37dec3e3539a6a048a9d4a186a01a Mon Sep 17 00:00:00 2001 From: Abhishek Parmar Date: Mon, 11 May 2026 02:30:30 +0800 Subject: [PATCH 08/13] fix: sync G1 base pin after placement --- mlspaces_tests/component_tests/test_unitree_g1.py | 12 ++++++++++++ molmo_spaces/robots/unitree_g1.py | 4 ++++ molmo_spaces/tasks/pick_task_sampler.py | 4 ++++ 3 files changed, 20 insertions(+) diff --git a/mlspaces_tests/component_tests/test_unitree_g1.py b/mlspaces_tests/component_tests/test_unitree_g1.py index af252018..35f465cc 100644 --- a/mlspaces_tests/component_tests/test_unitree_g1.py +++ b/mlspaces_tests/component_tests/test_unitree_g1.py @@ -245,6 +245,18 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): assert pick_robot.state_dim == 16 assert pick_robot.action_dim(list(pick_robot.controllers)) == 9 + next_pick_base_pose = initial_pick_base_pose.copy() + next_pick_base_pose[:3, 3] = [0.6, -0.1, initial_pick_base_pose[2, 3]] + pick_robot.robot_view.base.pose = next_pick_base_pose + pick_robot.sync_pinned_base_pose() + pick_robot.compute_control() + mujoco.mj_step(pick_scene_model, pick_scene_data) + assert np.allclose( + pick_robot.robot_view.base.pose[:3, 3], + next_pick_base_pose[:3, 3], + atol=1e-3, + ) + robot = UnitreeG1Robot(scene_data, SimpleNamespace(robot_config=config)) robot.reset() mujoco.mj_forward(scene_model, scene_data) diff --git a/molmo_spaces/robots/unitree_g1.py b/molmo_spaces/robots/unitree_g1.py index 6f25c1f3..d8f0ce14 100644 --- a/molmo_spaces/robots/unitree_g1.py +++ b/molmo_spaces/robots/unitree_g1.py @@ -132,6 +132,10 @@ def set_joint_pos(self, robot_joint_pos_dict) -> None: def set_world_pose(self, robot_world_pose) -> None: self._robot_view.base.pose = robot_world_pose + self.sync_pinned_base_pose() + + def sync_pinned_base_pose(self) -> None: + """Latch the current base pose as the fixed pose for pinned-base workflows.""" self._set_pinned_base_pose_if_configured() self._zero_base_velocity_if_configured() diff --git a/molmo_spaces/tasks/pick_task_sampler.py b/molmo_spaces/tasks/pick_task_sampler.py index 64b38f29..5c6f3d31 100644 --- a/molmo_spaces/tasks/pick_task_sampler.py +++ b/molmo_spaces/tasks/pick_task_sampler.py @@ -1189,6 +1189,10 @@ def add_pickup_target( class UnitreeG1RightArmPickTaskSampler(PickTaskSampler): """Pick sampler variant that keeps the Unitree G1 pelvis at its standing height.""" + def _sample_and_place_robot(self, env: CPUMujocoEnv) -> None: + super()._sample_and_place_robot(env) + env.current_robot.sync_pinned_base_pose() + def _initial_robot_z_for_pickup(self, target_pos: np.ndarray, robot_view) -> float: del target_pos, robot_view base_qpos = self.config.robot_config.init_qpos.get("base") From 4cfa48c344e631066fba2afa6695914c171406fa Mon Sep 17 00:00:00 2001 From: Abhishek Parmar Date: Mon, 11 May 2026 02:49:54 +0800 Subject: [PATCH 09/13] fixup! fix: stabilize G1 right-arm pick placement --- mlspaces_tests/component_tests/test_unitree_g1.py | 5 +++++ .../config/object_manipulation_datagen_configs.py | 2 +- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/mlspaces_tests/component_tests/test_unitree_g1.py b/mlspaces_tests/component_tests/test_unitree_g1.py index 35f465cc..807ca80c 100644 --- a/mlspaces_tests/component_tests/test_unitree_g1.py +++ b/mlspaces_tests/component_tests/test_unitree_g1.py @@ -129,6 +129,11 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): pick_datagen_config.task_sampler_config.task_sampler_class is UnitreeG1RightArmPickTaskSampler ) + assert pick_datagen_config.task_sampler_config.base_pose_sampling_radius_range == ( + 0.15, + 0.4, + ) + assert pick_datagen_config.task_sampler_config.robot_safety_radius == 0.25 assert pick_datagen_config.policy_config.policy_cls.__name__ == ( "UnitreeG1RightArmPickPlannerPolicy" ) diff --git a/molmo_spaces/data_generation/config/object_manipulation_datagen_configs.py b/molmo_spaces/data_generation/config/object_manipulation_datagen_configs.py index 54035da7..5e483bea 100644 --- a/molmo_spaces/data_generation/config/object_manipulation_datagen_configs.py +++ b/molmo_spaces/data_generation/config/object_manipulation_datagen_configs.py @@ -156,7 +156,7 @@ class UnitreeG1RightArmPickDataGenConfig(PickBaseConfig): robot_object_z_offset=0.0, robot_object_z_offset_random_min=0.0, robot_object_z_offset_random_max=0.0, - base_pose_sampling_radius_range=(0.2, 0.8), + base_pose_sampling_radius_range=(0.15, 0.4), robot_safety_radius=0.25, check_robot_placement_visibility=False, ) From afca46b5a769d23f5ddfd29b278261d3c47b0167 Mon Sep 17 00:00:00 2001 From: Abhishek Parmar Date: Mon, 11 May 2026 17:06:22 +0800 Subject: [PATCH 10/13] feat: add G1 tabletop pick-and-place datagen --- .../component_tests/test_unitree_g1.py | 110 ++++++- molmo_spaces/configs/policy_configs.py | 24 ++ molmo_spaces/configs/robot_configs.py | 12 + molmo_spaces/configs/task_sampler_configs.py | 14 + .../object_manipulation_datagen_configs.py | 53 ++++ .../pick_and_place_planner_policy.py | 66 ++++ molmo_spaces/robots/unitree_g1.py | 21 ++ molmo_spaces/tasks/pick_task_sampler.py | 3 + molmo_spaces/tasks/task_sampler.py | 7 +- .../tasks/unitree_g1_tabletop_task_sampler.py | 231 ++++++++++++++ scripts/assets/prepare_unitree_g1.py | 13 + scripts/assets/prepare_unitree_g1_tabletop.py | 296 ++++++++++++++++++ 12 files changed, 847 insertions(+), 3 deletions(-) create mode 100644 molmo_spaces/tasks/unitree_g1_tabletop_task_sampler.py create mode 100644 scripts/assets/prepare_unitree_g1_tabletop.py diff --git a/mlspaces_tests/component_tests/test_unitree_g1.py b/mlspaces_tests/component_tests/test_unitree_g1.py index 807ca80c..3bda6bc5 100644 --- a/mlspaces_tests/component_tests/test_unitree_g1.py +++ b/mlspaces_tests/component_tests/test_unitree_g1.py @@ -15,7 +15,22 @@ from molmo_spaces.kinematics.mujoco_kinematics import MlSpacesKinematics from molmo_spaces.robots.unitree_g1 import UnitreeG1Robot from molmo_spaces.tasks.pick_task_sampler import UnitreeG1RightArmPickTaskSampler -from scripts.assets.prepare_unitree_g1 import ROBOT_ASSET_NAME, prepare_unitree_g1 +from scripts.assets.prepare_unitree_g1 import ( + LEFT_ARM_STOW_QPOS, + LEFT_HAND_OPEN_QPOS, + ROBOT_ASSET_NAME, + prepare_unitree_g1, +) +from scripts.assets.prepare_unitree_g1_tabletop import ( + PLACE_RECEPTACLE_BODY_NAME, + PLACE_RECEPTACLE_SITE_NAME, + TABLE_BODY_NAME, + TABLE_MATERIAL_NAME, + TABLE_RGBA, + TABLETOP_GEOM_NAME, + TABLETOP_VISUAL_GEOM_NAME, + prepare_unitree_g1_tabletop, +) KINEMATICS_SITE_NAMES = [ "left_wrist_site", @@ -25,6 +40,36 @@ ] +def test_prepare_unitree_g1_tabletop_scenes_smoke(tmp_path): + scene_dir = tmp_path / "scenes" / "unitree_g1_tabletop_v1" + scene_paths = prepare_unitree_g1_tabletop(scene_dir) + assert [path.name for path in scene_paths] == [ + "unitree_g1_tabletop_pelvis_minus_10cm_v1.xml", + "unitree_g1_tabletop_pelvis_height_v1.xml", + ] + + expected_heights = [0.693, 0.793] + for scene_path, expected_height in zip(scene_paths, expected_heights, strict=True): + metadata_path = scene_path.with_name(scene_path.stem + "_metadata.json") + assert metadata_path.exists() + model = mujoco.MjModel.from_xml_path(str(scene_path)) + data = mujoco.MjData(model) + mujoco.mj_forward(model, data) + assert model.body(TABLE_BODY_NAME).id >= 0 + assert model.body(PLACE_RECEPTACLE_BODY_NAME).id >= 0 + assert model.site(PLACE_RECEPTACLE_SITE_NAME).id >= 0 + table_material_id = model.material(TABLE_MATERIAL_NAME).id + np.testing.assert_allclose(model.mat_rgba[table_material_id], TABLE_RGBA) + tabletop_geom = model.geom(TABLETOP_GEOM_NAME) + tabletop_visual_geom = model.geom(TABLETOP_VISUAL_GEOM_NAME) + assert model.geom_group[tabletop_geom.id] == 4 + assert model.geom_group[tabletop_visual_geom.id] == 0 + assert model.geom_contype[tabletop_visual_geom.id] == 0 + assert model.geom_conaffinity[tabletop_visual_geom.id] == 0 + tabletop_top_z = float(tabletop_geom.pos[2] + tabletop_geom.size[2]) + assert tabletop_top_z == pytest.approx(expected_height) + + def _has_deep_robot_floor_penetration( model: mujoco.MjModel, data: mujoco.MjData, @@ -44,6 +89,21 @@ def _has_deep_robot_floor_penetration( return False +def _assert_locked_joints( + model: mujoco.MjModel, + data: mujoco.MjData, + namespace: str, + locked_joint_qpos: dict[str, float], +) -> None: + for joint_name, expected_qpos in locked_joint_qpos.items(): + namespaced_joint_name = f"{namespace}{joint_name}" + joint_id = model.joint(namespaced_joint_name).id + qposadr = model.jnt_qposadr[joint_id] + actuator_id = model.actuator(namespaced_joint_name).id + assert data.qpos[qposadr] == pytest.approx(expected_qpos) + assert data.ctrl[actuator_id] == pytest.approx(expected_qpos) + + def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): unitree_root = os.environ.get("UNITREE_URDF_ROOT") if unitree_root is None: @@ -124,6 +184,9 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): pick_datagen_config_cls = get_config_class("UnitreeG1RightArmPickDataGenConfig") pick_datagen_config = pick_datagen_config_cls() assert isinstance(pick_datagen_config.robot_config, UnitreeG1RightArmPickRobotConfig) + assert pick_datagen_config.robot_config.locked_joint_qpos == ( + LEFT_ARM_STOW_QPOS | LEFT_HAND_OPEN_QPOS + ) assert pick_datagen_config.robot_config.pin_base_in_place assert ( pick_datagen_config.task_sampler_config.task_sampler_class @@ -140,6 +203,24 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): assert not pick_datagen_config.policy_config.filter_colliding_grasps assert not pick_datagen_config.policy_config.filter_feasible_grasps assert len(pick_datagen_config.camera_config.cameras) == 2 + tabletop_datagen_config_cls = get_config_class( + "UnitreeG1RightArmTabletopPickAndPlaceDataGenConfig" + ) + tabletop_datagen_config = tabletop_datagen_config_cls() + assert tabletop_datagen_config.scene_dataset == "user" + assert isinstance(tabletop_datagen_config.robot_config, UnitreeG1RightArmPickRobotConfig) + assert ( + tabletop_datagen_config.task_sampler_config.task_sampler_class.__name__ + == "UnitreeG1RightArmTabletopPickAndPlaceTaskSampler" + ) + assert tabletop_datagen_config.task_sampler_config.house_inds == [0, 1] + assert tabletop_datagen_config.task_sampler_config.samples_per_house == 5 + assert len(tabletop_datagen_config.task_sampler_config.scene_xml_paths) == 2 + assert tabletop_datagen_config.policy_config.policy_cls.__name__ == ( + "UnitreeG1RightArmPickAndPlacePlannerPolicy" + ) + assert not tabletop_datagen_config.policy_config.filter_colliding_grasps + assert not tabletop_datagen_config.policy_config.filter_feasible_grasps base_scene_path = constants.ABS_PATH_OF_TOP_LEVEL_MOLMO_SPACES_DIR / ( "molmo_spaces/resources/base_scene.xml" @@ -222,8 +303,28 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): pick_robot = UnitreeG1Robot( pick_scene_data, SimpleNamespace(robot_config=pick_datagen_config.robot_config) ) + mujoco.mj_resetData(pick_scene_model, pick_scene_data) + for group_name, qpos in pick_datagen_config.robot_config.init_qpos.items(): + pick_robot.robot_view.get_move_group(group_name).joint_pos = np.asarray(qpos) + for controller in pick_robot.controllers.values(): + controller.reset() + pick_robot.apply_initial_state_overrides() + mujoco.mj_forward(pick_scene_model, pick_scene_data) + _assert_locked_joints( + pick_scene_model, + pick_scene_data, + pick_datagen_config.robot_config.robot_namespace, + pick_datagen_config.robot_config.locked_joint_qpos, + ) + pick_robot.reset() mujoco.mj_forward(pick_scene_model, pick_scene_data) + _assert_locked_joints( + pick_scene_model, + pick_scene_data, + pick_datagen_config.robot_config.robot_namespace, + pick_datagen_config.robot_config.locked_joint_qpos, + ) initial_pick_base_pose = pick_robot.robot_view.base.pose.copy() initial_pick_base_pose[:3, 3] = [0.2, -0.3, initial_pick_base_pose[2, 3]] pick_robot.robot_view.base.pose = initial_pick_base_pose @@ -240,8 +341,15 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): for _ in range(10): pick_robot.compute_control() mujoco.mj_step(pick_scene_model, pick_scene_data) + pick_robot.compute_control() assert np.isfinite(pick_scene_data.qpos).all() assert np.isfinite(pick_scene_data.ctrl).all() + _assert_locked_joints( + pick_scene_model, + pick_scene_data, + pick_datagen_config.robot_config.robot_namespace, + pick_datagen_config.robot_config.locked_joint_qpos, + ) assert np.allclose( pick_robot.robot_view.base.pose[:3, 3], initial_pick_base_pose[:3, 3], diff --git a/molmo_spaces/configs/policy_configs.py b/molmo_spaces/configs/policy_configs.py index fae59138..464032c9 100644 --- a/molmo_spaces/configs/policy_configs.py +++ b/molmo_spaces/configs/policy_configs.py @@ -178,6 +178,30 @@ def model_post_init(self, __context) -> None: self.policy_cls = PickAndPlacePlannerPolicy +class UnitreeG1RightArmPickAndPlacePlannerPolicyConfig(PickAndPlacePlannerPolicyConfig): + """Pick-and-place planner config for fixed-base G1 right-arm tabletop datagen.""" + + policy_cls: type = None + filter_colliding_grasps: bool = False + filter_feasible_grasps: bool = False + grasp_feasibility_batch_size: int = 32 + grasp_feasibility_max_grasps: int = 32 + grasp_collision_batch_size: int = 32 + grasp_collision_max_grasps: int = 64 + max_retries: int = 0 + max_sequential_ik_failures: int = 4 + phase_timeout: float = 24.0 + + def model_post_init(self, __context) -> None: + """Set policy_cls after initialization to avoid circular imports.""" + super().model_post_init(__context) + from molmo_spaces.policy.solvers.object_manipulation.pick_and_place_planner_policy import ( + UnitreeG1RightArmPickAndPlacePlannerPolicy, + ) + + self.policy_cls = UnitreeG1RightArmPickAndPlacePlannerPolicy + + class CuroboOpenClosePlannerPolicyConfig(OpenClosePlannerPolicyConfig): policy_cls: type = None # Will be set in model_post_init to avoid circular imports left_curobo_planner_config: CuroboPlannerConfig | None = None # will be set in model_post_init diff --git a/molmo_spaces/configs/robot_configs.py b/molmo_spaces/configs/robot_configs.py index 8612b104..c97255f6 100644 --- a/molmo_spaces/configs/robot_configs.py +++ b/molmo_spaces/configs/robot_configs.py @@ -101,6 +101,7 @@ class BaseRobotConfig(Config): ] # move_group to command_mode e.g., "joint", "cartesian", "velocity" init_qpos: dict[str, list[float]] init_qpos_noise_range: dict[str, list[float]] | None + locked_joint_qpos: dict[str, float] | None = None name: str | None robot_xml_path: ( Path | None @@ -471,6 +472,17 @@ class UnitreeG1RightArmPickRobotConfig(UnitreeG1Dex1RobotConfig): "right_arm": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], "gripper": [0.0245, 0.0245], } + locked_joint_qpos: dict[str, float] | None = { + "left_shoulder_pitch_joint": 0.25, + "left_shoulder_roll_joint": 0.35, + "left_shoulder_yaw_joint": -0.1, + "left_elbow_joint": 1.0, + "left_wrist_roll_joint": 0.0, + "left_wrist_pitch_joint": 0.0, + "left_wrist_yaw_joint": 0.0, + "left_dex1_finger_joint_1": 0.0245, + "left_dex1_finger_joint_2": 0.0245, + } command_mode: dict[str, str | None] = { "base": None, "right_arm": "joint_position", diff --git a/molmo_spaces/configs/task_sampler_configs.py b/molmo_spaces/configs/task_sampler_configs.py index 070548b6..6a2bc551 100644 --- a/molmo_spaces/configs/task_sampler_configs.py +++ b/molmo_spaces/configs/task_sampler_configs.py @@ -210,6 +210,20 @@ class PickAndPlaceTaskSamplerConfig(PickTaskSamplerConfig): episodes_per_receptacle: int = 2 +class UnitreeG1TabletopPickAndPlaceTaskSamplerConfig(PickAndPlaceTaskSamplerConfig): + """Sampler settings for the fixed Unitree G1 tabletop user scene.""" + + table_body_name: str = "g1_table" + table_geom_name: str = "g1_tabletop_geom" + place_receptacle_name: str = "g1_place_bin" + pickup_workspace_center_xy: tuple[float, float] = (0.42, -0.16) + pickup_workspace_size_xy: tuple[float, float] = (0.18, 0.22) + place_workspace_center_xy: tuple[float, float] = (0.42, 0.16) + place_workspace_size_xy: tuple[float, float] = (0.16, 0.20) + object_table_clearance: float = 0.006 + receptacle_table_clearance: float = 0.001 + + class PickAndPlaceNextToTaskSamplerConfig(PickAndPlaceTaskSamplerConfig): place_receptacle_types: list[str] = [] # Empty = any object on bench min_object_to_receptacle_dist: float = 0.3 # avoid insta-success by keeping this large diff --git a/molmo_spaces/data_generation/config/object_manipulation_datagen_configs.py b/molmo_spaces/data_generation/config/object_manipulation_datagen_configs.py index 5e483bea..85ea88f4 100644 --- a/molmo_spaces/data_generation/config/object_manipulation_datagen_configs.py +++ b/molmo_spaces/data_generation/config/object_manipulation_datagen_configs.py @@ -34,6 +34,7 @@ CuroboPickAndPlacePlannerPolicyConfig, OpenClosePlannerPolicyConfig, PickPlannerPolicyConfig, + UnitreeG1RightArmPickAndPlacePlannerPolicyConfig, UnitreeG1RightArmPickPlannerPolicyConfig, ) from molmo_spaces.configs.robot_configs import ( @@ -51,6 +52,7 @@ PickAndPlaceTaskSamplerConfig, PickTaskSamplerConfig, RUMPickTaskSamplerConfig, + UnitreeG1TabletopPickAndPlaceTaskSamplerConfig, ) # Oder of configs should be order the code is executed in @@ -65,6 +67,9 @@ PickAndPlaceTaskSampler, ) from molmo_spaces.tasks.pick_task_sampler import PickTaskSampler, UnitreeG1RightArmPickTaskSampler +from molmo_spaces.tasks.unitree_g1_tabletop_task_sampler import ( + UnitreeG1RightArmTabletopPickAndPlaceTaskSampler, +) from molmo_spaces.utils.constants.object_constants import PICK_AND_PLACE_OBJECTS from molmo_spaces.utils.synset_utils import get_valid_pickupable_obja_uids @@ -173,6 +178,54 @@ def tag(self) -> str: return "unitree_g1_right_arm_pick_datagen" +@register_config("UnitreeG1RightArmTabletopPickAndPlaceDataGenConfig") +class UnitreeG1RightArmTabletopPickAndPlaceDataGenConfig(PickAndPlaceDataGenConfig): + """Right-arm-only Unitree G1 pick-and-place datagen for a fixed tabletop user scene.""" + + scene_dataset: str = "user" + robot_config: UnitreeG1RightArmPickRobotConfig = UnitreeG1RightArmPickRobotConfig() + camera_config: UnitreeG1RightArmPickCameraSystem = UnitreeG1RightArmPickCameraSystem() + task_sampler_config: UnitreeG1TabletopPickAndPlaceTaskSamplerConfig = ( + UnitreeG1TabletopPickAndPlaceTaskSamplerConfig( + task_sampler_class=UnitreeG1RightArmTabletopPickAndPlaceTaskSampler, + house_inds=[0, 1], + scene_xml_paths=[ + str( + ASSETS_DIR + / "scenes" + / "unitree_g1_tabletop_v1" + / "unitree_g1_tabletop_pelvis_minus_10cm_v1.xml" + ), + str( + ASSETS_DIR + / "scenes" + / "unitree_g1_tabletop_v1" + / "unitree_g1_tabletop_pelvis_height_v1.xml" + ), + ], + samples_per_house=5, + added_pickup_objects=get_valid_pickupable_obja_uids(), + num_added_pickups=12, + episodes_per_added_pickup=1, + check_robot_placement_visibility=False, + robot_safety_radius=0.25, + ) + ) + policy_config: UnitreeG1RightArmPickAndPlacePlannerPolicyConfig = ( + UnitreeG1RightArmPickAndPlacePlannerPolicyConfig() + ) + task_horizon: int = 180 + filter_for_successful_trajectories: bool = False + profile: bool = False + output_dir: Path = ( + ASSETS_DIR / "experiment_output" / "datagen" / "unitree_g1_tabletop_pnp_v1" + ) + + @property + def tag(self) -> str: + return "unitree_g1_right_arm_tabletop_pick_and_place_datagen" + + @register_config("FrankaPickAndPlaceDataGenConfig") class FrankaPickAndPlaceDataGenConfig(PickAndPlaceDataGenConfig): robot_config: BaseRobotConfig = FrankaRobotConfig() diff --git a/molmo_spaces/policy/solvers/object_manipulation/pick_and_place_planner_policy.py b/molmo_spaces/policy/solvers/object_manipulation/pick_and_place_planner_policy.py index 42becb26..355ebcf0 100644 --- a/molmo_spaces/policy/solvers/object_manipulation/pick_and_place_planner_policy.py +++ b/molmo_spaces/policy/solvers/object_manipulation/pick_and_place_planner_policy.py @@ -287,3 +287,69 @@ def _compute_target_poses(self) -> dict[str, np.ndarray]: self.task.viewer.sync() return target_poses + + +class UnitreeG1RightArmPickAndPlacePlannerPolicy(PickAndPlacePlannerPolicy): + """Pick-and-place planner variant that solves IK through the G1 right arm only.""" + + def _tcp_to_jp_fn(self, mg_id: str, target_pose: np.ndarray) -> dict[str, np.ndarray]: + kinematics = self.task.env.current_robot.kinematics + + jp = kinematics.ik( + mg_id, + target_pose, + ["right_arm"], + self.robot_view.get_qpos_dict(), + self.robot_view.base.pose, + ) + + action = self.robot_view.get_ctrl_dict() + if jp is not None: + self.sequential_ik_failures = 0 + action["right_arm"] = jp["right_arm"] + else: + self.sequential_ik_failures += 1 + log.info(f"IK failed, holding current position, fails:{self.sequential_ik_failures}") + if self.sequential_ik_failures >= self.policy_config.max_sequential_ik_failures: + log.info("Too many sequential IK failures, triggering retry.") + return self._handle_failure() + + return action + + def check_feasible_ik(self, pose: np.ndarray) -> bool: + if not self.policy_config.filter_feasible_grasps: + if pose.ndim > 2: + return np.ones(pose.shape[0], dtype=bool) + return True + + robot_view = self.task.env.current_robot.robot_view + kinematics = self.task.env.current_robot.kinematics + gripper_mg_id = robot_view.get_gripper_movegroup_ids()[0] + + if pose.ndim > 2: + return np.array( + [ + kinematics.ik( + gripper_mg_id, + single_pose, + ["right_arm"], + robot_view.get_qpos_dict(), + robot_view.base.pose, + max_iter=100, + ) + is not None + for single_pose in pose + ], + dtype=bool, + ) + + assert pose.shape == (4, 4) + jp_dict = kinematics.ik( + gripper_mg_id, + pose, + ["right_arm"], + robot_view.get_qpos_dict(), + robot_view.base.pose, + max_iter=100, + ) + return jp_dict is not None diff --git a/molmo_spaces/robots/unitree_g1.py b/molmo_spaces/robots/unitree_g1.py index d8f0ce14..d5c09c9e 100644 --- a/molmo_spaces/robots/unitree_g1.py +++ b/molmo_spaces/robots/unitree_g1.py @@ -119,9 +119,11 @@ def update_control(self, action_command_dict: dict[str, Any]) -> None: def compute_control(self) -> None: self._pin_base_if_configured() + self._apply_locked_joint_targets() for controller in self.controllers.values(): ctrl_inputs = controller.compute_ctrl_inputs() controller.robot_move_group.ctrl = ctrl_inputs + self._apply_locked_joint_targets() self._pin_base_if_configured() def set_joint_pos(self, robot_joint_pos_dict) -> None: @@ -148,10 +150,29 @@ def reset(self) -> None: move_group = self._robot_view.get_move_group(mg_id) move_group.joint_pos = np.asarray(default_pos) move_group.joint_vel = np.zeros(move_group.vel_dim) + self._apply_locked_joint_targets() if self._should_pin_base(): self._pinned_base_pose = None self._zero_base_velocity_if_configured() + def apply_initial_state_overrides(self) -> None: + """Apply non-policy initial state constraints after sampler-level resets.""" + self._apply_locked_joint_targets() + self._pin_base_if_configured() + + def _apply_locked_joint_targets(self) -> None: + """Hold joints that are intentionally outside the policy-facing move groups.""" + locked_joint_qpos = getattr(self.exp_config.robot_config, "locked_joint_qpos", None) or {} + for joint_name, joint_qpos in locked_joint_qpos.items(): + namespaced_joint_name = f"{self.namespace}{joint_name}" + joint_id = self.mj_model.joint(namespaced_joint_name).id + qpos_addr = self.mj_model.jnt_qposadr[joint_id] + qvel_addr = self.mj_model.jnt_dofadr[joint_id] + self.mj_data.qpos[qpos_addr] = joint_qpos + self.mj_data.qvel[qvel_addr] = 0.0 + actuator_id = self.mj_model.actuator(namespaced_joint_name).id + self.mj_data.ctrl[actuator_id] = joint_qpos + def _should_pin_base(self) -> bool: return bool(getattr(self.exp_config.robot_config, "pin_base_in_place", False)) diff --git a/molmo_spaces/tasks/pick_task_sampler.py b/molmo_spaces/tasks/pick_task_sampler.py index 5c6f3d31..ad7090dc 100644 --- a/molmo_spaces/tasks/pick_task_sampler.py +++ b/molmo_spaces/tasks/pick_task_sampler.py @@ -245,6 +245,9 @@ def randomize_scene(self, env: CPUMujocoEnv, robot_view) -> None: if "head" in robot.robot_view.move_group_ids(): head_mg = robot.robot_view.get_move_group("head") head_mg.ctrl = head_mg.noop_ctrl + apply_initial_state_overrides = getattr(robot, "apply_initial_state_overrides", None) + if apply_initial_state_overrides is not None: + apply_initial_state_overrides() # robot_color = None # robot_color = [.941, .322, .612,1.] # example: red diff --git a/molmo_spaces/tasks/task_sampler.py b/molmo_spaces/tasks/task_sampler.py index 3eeca723..a6691837 100644 --- a/molmo_spaces/tasks/task_sampler.py +++ b/molmo_spaces/tasks/task_sampler.py @@ -719,7 +719,10 @@ def update_scene(self, scene_path: str | None = None, variant: str = "base") -> scene_path = self._current_house_scene_path(variant=variant) # If using a MolmoSpaces scene, install it - if get_scenes_root().resolve() in Path(scene_path).resolve().parents: + if ( + self.config.scene_dataset != "user" + and get_scenes_root().resolve() in Path(scene_path).resolve().parents + ): # Track asset installation time (fetching/extracting scene, objects, grasps) # Use detailed profiling to identify which asset type is slow if self._datagen_profiler is not None: @@ -897,7 +900,7 @@ def _get_dataset_index_map(self) -> dict: else: mapping = { self.config.data_split: { - i: {"base": scene_xml_path} + i: {"base": scene_xml_path, "ceiling": scene_xml_path} for i, scene_xml_path in enumerate( self.config.task_sampler_config.scene_xml_paths ) diff --git a/molmo_spaces/tasks/unitree_g1_tabletop_task_sampler.py b/molmo_spaces/tasks/unitree_g1_tabletop_task_sampler.py new file mode 100644 index 00000000..38c9fc9f --- /dev/null +++ b/molmo_spaces/tasks/unitree_g1_tabletop_task_sampler.py @@ -0,0 +1,231 @@ +"""Unitree G1 tabletop task samplers.""" + +from __future__ import annotations + +import logging + +import mujoco +import numpy as np + +from molmo_spaces.env.data_views import create_mlspaces_body +from molmo_spaces.env.env import CPUMujocoEnv +from molmo_spaces.tasks.pick_and_place_task import PickAndPlaceTask +from molmo_spaces.tasks.pick_and_place_task_sampler import PickAndPlaceTaskSampler +from molmo_spaces.tasks.task_sampler import BaseMujocoTaskSampler +from molmo_spaces.tasks.task_sampler_errors import HouseInvalidForTask +from molmo_spaces.utils.grasp_sample import has_valid_grasp_file +from molmo_spaces.utils.mj_model_and_data_utils import body_aabb +from molmo_spaces.utils.mujoco_scene_utils import get_supporting_geom +from molmo_spaces.utils.pose import pos_quat_to_pose_mat, pose_mat_to_7d + +log = logging.getLogger(__name__) + + +class UnitreeG1RightArmTabletopPickAndPlaceTaskSampler(PickAndPlaceTaskSampler): + """Pick-and-place sampler for the fixed G1 real-setup tabletop user scene.""" + + def add_auxiliary_objects(self, spec: mujoco.MjSpec) -> None: + from molmo_spaces.tasks.pick_task_sampler import PickTaskSampler + + PickTaskSampler.add_auxiliary_objects(self, spec) + + def init_scene(self, env) -> None: + BaseMujocoTaskSampler.init_scene(self, env) + self._task_counter = 0 + self._grasp_failure_counts = {} + self._metadata_adder.add_meta(env.current_scene_metadata) + self.candidate_objects = self._get_added_pickup_candidates(env) + if not self.candidate_objects: + raise HouseInvalidForTask("No dynamically added tabletop pickup objects available") + + def _get_added_pickup_candidates(self, env: CPUMujocoEnv): + om = env.object_managers[env.current_batch_index] + candidates = [] + for name in self._added_pickup_names: + try: + candidates.append(om.get_object_by_name(name)) + except KeyError: + log.debug("Skipping unavailable tabletop pickup object %s", name) + return candidates + + def _restore_pickupables_to_staging(self, env: CPUMujocoEnv) -> None: + for name, pose7 in self._added_pickup_staging_poses.items(): + body = create_mlspaces_body(env.current_data, name) + body.pose = pos_quat_to_pose_mat(pose7[:3], pose7[3:7]) + self._zero_free_body_velocity(env, name) + + def _zero_free_body_velocity(self, env: CPUMujocoEnv, body_name: str) -> None: + model = env.current_model + data = env.current_data + body_id = model.body(body_name).id + joint_id = model.body_jntadr[body_id] + if joint_id < 0: + return + dof_adr = model.jnt_dofadr[joint_id] + dof_num = 6 if model.jnt_type[joint_id] == mujoco.mjtJoint.mjJNT_FREE else 1 + data.qvel[dof_adr : dof_adr + dof_num] = 0.0 + + def _table_top_z(self, env: CPUMujocoEnv) -> float: + table_body = env.current_data.body(self.config.task_sampler_config.table_body_name).id + center, size = body_aabb(env.current_model, env.current_data, table_body, visual_only=False) + return float(center[2] + size[2] / 2.0) + + def _sample_xy(self, center_xy: tuple[float, float], size_xy: tuple[float, float]) -> np.ndarray: + center = np.asarray(center_xy, dtype=float) + half_size = np.asarray(size_xy, dtype=float) / 2.0 + return center + np.random.uniform(-half_size, half_size) + + def _place_pickup_on_table(self, env: CPUMujocoEnv, pickup_obj_name: str) -> None: + body = create_mlspaces_body(env.current_data, pickup_obj_name) + model = env.current_model + data = env.current_data + mujoco.mj_forward(model, data) + center, size = body_aabb(model, data, body.body_id, visual_only=False) + bottom_z = center[2] - size[2] / 2.0 + body_to_bottom = body.position[2] - bottom_z + + xy = self._sample_xy( + self.config.task_sampler_config.pickup_workspace_center_xy, + self.config.task_sampler_config.pickup_workspace_size_xy, + ) + body.position = np.array( + [ + xy[0], + xy[1], + self._table_top_z(env) + + body_to_bottom + + self.config.task_sampler_config.object_table_clearance, + ] + ) + self._zero_free_body_velocity(env, pickup_obj_name) + mujoco.mj_forward(model, data) + + def _place_receptacle_on_table(self, env: CPUMujocoEnv) -> None: + receptacle_name = self.config.task_sampler_config.place_receptacle_name + body = create_mlspaces_body(env.current_data, receptacle_name) + xy = self._sample_xy( + self.config.task_sampler_config.place_workspace_center_xy, + self.config.task_sampler_config.place_workspace_size_xy, + ) + body.position = np.array( + [ + xy[0], + xy[1], + self._table_top_z(env) + self.config.task_sampler_config.receptacle_table_clearance, + ] + ) + body.quat = np.array([1.0, 0.0, 0.0, 0.0]) + mujoco.mj_forward(env.current_model, env.current_data) + + def _sample_and_place_robot(self, env: CPUMujocoEnv) -> None: + robot_view = env.current_robot.robot_view + base_qpos = self.config.robot_config.init_qpos["base"] + robot_view.base.pose = pos_quat_to_pose_mat(np.asarray(base_qpos[:3]), base_qpos[3:7]) + robot_view.base.joint_vel = np.zeros(robot_view.base.vel_dim) + env.current_robot.sync_pinned_base_pose() + self.config.task_config.robot_base_pose = pose_mat_to_7d(robot_view.base.pose).tolist() + + def _select_pickup_object(self, env: CPUMujocoEnv) -> int: + if not self._added_pickup_names: + raise HouseInvalidForTask("No added pickup objects were loaded") + + max_attempts = len(self._added_pickup_names) + for _ in range(max_attempts): + object_index = self._task_counter % len(self._added_pickup_names) + self._task_counter += 1 + pickup_obj_name = self._added_pickup_names[object_index] + self._added_pickup_obj_name = pickup_obj_name + self.config.task_config.pickup_obj_name = pickup_obj_name + + self._restore_pickupables_to_staging(env) + self._place_pickup_on_table(env, pickup_obj_name) + self._place_receptacle_on_table(env) + self._sample_and_place_robot(env) + + om = env.object_managers[env.current_batch_index] + pickup_obj = om.get_object_by_name(pickup_obj_name) + self.config.task_config.pickup_obj_start_pose = pose_mat_to_7d( + pickup_obj.pose + ).tolist() + pickup_obj_goal_pose = pose_mat_to_7d(pickup_obj.pose) + pickup_obj_goal_pose[2] += 0.05 + self.config.task_config.pickup_obj_goal_pose = pickup_obj_goal_pose.tolist() + self.config.task_config.object_poses = { + pickup_obj_name: pose_mat_to_7d(pickup_obj.pose).tolist() + } + + self.place_receptacle_name = self.config.task_sampler_config.place_receptacle_name + receptacle = om.get_object_by_name(self.place_receptacle_name) + self.config.task_config.place_receptacle_name = self.place_receptacle_name + self.config.task_config.place_target_name = self.place_receptacle_name + self.config.task_config.place_receptacle_start_pose = pose_mat_to_7d( + receptacle.pose + ).tolist() + + asset_uid = self.get_asset_uid_from_object(env, pickup_obj_name) + if asset_uid and not has_valid_grasp_file(asset_uid): + log.info("Skipping %s because no valid grasp file exists", pickup_obj_name) + continue + + supporting_geom_id = get_supporting_geom(env.current_data, pickup_obj.body_id) + if supporting_geom_id is None: + supporting_geom_id = env.current_model.geom( + self.config.task_sampler_config.table_geom_name + ).id + + self.setup_cameras(env) + return supporting_geom_id + + raise HouseInvalidForTask("Unable to sample a valid G1 tabletop pickup object") + + def _build_context_objects(self, env, om, pickup_obj_name, supporting_geom_id): + del supporting_geom_id + return [ + om.get_object_by_name(pickup_obj_name), + om.get_object_by_name(self.config.task_sampler_config.place_receptacle_name), + ] + + def _generate_referral_expressions(self, env, pickup_obj_name: str, context_objects: list): + if pickup_obj_name == self.config.task_sampler_config.place_receptacle_name: + priority = [(1.0, 1.0, "box")] + return priority, priority + try: + return super()._generate_referral_expressions(env, pickup_obj_name, context_objects) + except (AssertionError, KeyError, ValueError): + om = env.object_managers[env.current_batch_index] + try: + expression = om.fallback_expression(pickup_obj_name) + except (KeyError, ValueError): + expression = pickup_obj_name.split("/")[-1].replace("_", " ") + priority = [(1.0, 1.0, expression)] + return priority, priority + + def _get_place_target_candidates(self, env, pickup_obj_name, supporting_geom_id) -> list[str]: + del env, pickup_obj_name, supporting_geom_id + return [self.config.task_sampler_config.place_receptacle_name] + + def _prepare_place_target( + self, + env, + place_target_name, + pickup_obj_name, + pickup_obj_pos, + supporting_geom_id, + ) -> bool: + del env, place_target_name, pickup_obj_name, pickup_obj_pos, supporting_geom_id + return True + + def _filter_place_target(self, env, pickup_obj_name, place_target_name) -> bool: + del env, pickup_obj_name, place_target_name + return True + + def _finalize_task_config(self, env: CPUMujocoEnv): + del env + added = {} + if self._added_pickup_obj_name in self.added_objects: + added[self._added_pickup_obj_name] = self.added_objects[self._added_pickup_obj_name] + self.config.task_config.added_objects = added + + def _sample_task(self, env: CPUMujocoEnv): + self._configure_pick_and_place(env) + return PickAndPlaceTask(env, self.config) diff --git a/scripts/assets/prepare_unitree_g1.py b/scripts/assets/prepare_unitree_g1.py index 6e914fbc..6fb82398 100644 --- a/scripts/assets/prepare_unitree_g1.py +++ b/scripts/assets/prepare_unitree_g1.py @@ -23,6 +23,19 @@ SOURCE_MESH_DIR_REL_PATH = Path("g1_description/meshes") OUTPUT_XML_NAME = "model.xml" DEFAULT_PELVIS_HEIGHT_M = 0.793 +LEFT_ARM_STOW_QPOS = { + "left_shoulder_pitch_joint": 0.25, + "left_shoulder_roll_joint": 0.35, + "left_shoulder_yaw_joint": -0.1, + "left_elbow_joint": 1.0, + "left_wrist_roll_joint": 0.0, + "left_wrist_pitch_joint": 0.0, + "left_wrist_yaw_joint": 0.0, +} +LEFT_HAND_OPEN_QPOS = { + "left_dex1_finger_joint_1": 0.0245, + "left_dex1_finger_joint_2": 0.0245, +} GAIN_BY_JOINT_GROUP = { diff --git a/scripts/assets/prepare_unitree_g1_tabletop.py b/scripts/assets/prepare_unitree_g1_tabletop.py new file mode 100644 index 00000000..36ed8593 --- /dev/null +++ b/scripts/assets/prepare_unitree_g1_tabletop.py @@ -0,0 +1,296 @@ +"""Prepare simple Unitree G1 tabletop user scenes for MolmoSpaces datagen. + +The generated scenes are intentionally small: a floor, a lowered/pelvis-height +table, and a mocap bin that the task sampler can reposition on the tabletop. +Pickup objects are still injected by the task sampler from existing MolmoSpaces +assets so grasp files and episode serialization follow the normal pipeline. +""" + +from __future__ import annotations + +import argparse +import json +import shutil +import xml.etree.ElementTree as ET +from dataclasses import dataclass +from pathlib import Path + +import mujoco + +from molmo_spaces.molmo_spaces_constants import ( + ABS_PATH_OF_TOP_LEVEL_MOLMO_SPACES_DIR, + ASSETS_DIR, +) + +SCENE_ASSET_NAME = "unitree_g1_tabletop_v1" +OUTPUT_SCENE_DIR = ASSETS_DIR / "scenes" / SCENE_ASSET_NAME + +TABLE_BODY_NAME = "g1_table" +TABLETOP_GEOM_NAME = "g1_tabletop_geom" +TABLETOP_VISUAL_GEOM_NAME = "g1_tabletop_visual" +TABLE_MATERIAL_NAME = "g1_table_mat" +TABLE_RGBA = (0.74, 0.56, 0.36, 1.0) +PLACE_RECEPTACLE_BODY_NAME = "g1_place_bin" +PLACE_RECEPTACLE_SITE_NAME = "g1_place_bin_site" + +PELVIS_MINUS_10CM_TABLETOP_HEIGHT_M = 0.693 +PELVIS_TABLETOP_HEIGHT_M = 0.793 + + +@dataclass(frozen=True) +class TabletopSceneSpec: + name: str + tabletop_height_m: float + + +SCENE_SPECS = ( + TabletopSceneSpec( + "unitree_g1_tabletop_pelvis_minus_10cm_v1", + PELVIS_MINUS_10CM_TABLETOP_HEIGHT_M, + ), + TabletopSceneSpec("unitree_g1_tabletop_pelvis_height_v1", PELVIS_TABLETOP_HEIGHT_M), +) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--output-dir", + type=Path, + default=OUTPUT_SCENE_DIR, + help="Output scene directory. Defaults to $MLSPACES_ASSETS_DIR/scenes/unitree_g1_tabletop_v1.", + ) + parser.add_argument( + "--force", + action="store_true", + help="Overwrite an existing output directory.", + ) + return parser.parse_args() + + +def _set_attrs(elem: ET.Element, **attrs: object) -> ET.Element: + for key, value in attrs.items(): + if isinstance(value, (list, tuple)): + elem.set(key, " ".join(str(x) for x in value)) + else: + elem.set(key, str(value)) + return elem + + +def _add_materials(root: ET.Element) -> None: + asset = root.find("asset") + if asset is None: + asset = ET.SubElement(root, "asset") + _set_attrs( + ET.SubElement(asset, "material"), + name=TABLE_MATERIAL_NAME, + rgba=TABLE_RGBA, + ) + _set_attrs( + ET.SubElement(asset, "material"), + name="g1_bin_mat", + rgba=[0.12, 0.33, 0.62, 1.0], + ) + _set_attrs( + ET.SubElement(asset, "material"), + name="g1_floor_mat", + rgba=[0.72, 0.74, 0.70, 1.0], + ) + + +def _add_visible_collision_geom( + parent: ET.Element, + *, + name: str, + geom_type: str, + pos: list[float] | tuple[float, ...], + size: list[float] | tuple[float, ...], + material: str, + visual_name: str | None = None, +) -> None: + _set_attrs( + ET.SubElement(parent, "geom"), + name=name, + type=geom_type, + pos=pos, + size=size, + material=material, + **{"class": "__STRUCTURAL_MJT__"}, + ) + _set_attrs( + ET.SubElement(parent, "geom"), + name=visual_name or f"{name}_visual", + type=geom_type, + pos=pos, + size=size, + material=material, + **{"class": "__VISUAL_MJT__"}, + ) + + +def _add_floor(worldbody: ET.Element) -> None: + _add_visible_collision_geom( + worldbody, + name="g1_tabletop_floor", + geom_type="plane", + pos=[0, 0, 0], + size=[3.0, 3.0, 0.02], + material="g1_floor_mat", + ) + + +def _add_table(worldbody: ET.Element, tabletop_height_m: float) -> None: + table_body = _set_attrs(ET.SubElement(worldbody, "body"), name=TABLE_BODY_NAME) + tabletop_half_size = (0.35, 0.45, 0.025) + tabletop_center = (0.55, 0.0, tabletop_height_m - tabletop_half_size[2]) + _add_visible_collision_geom( + table_body, + name=TABLETOP_GEOM_NAME, + geom_type="box", + pos=tabletop_center, + size=tabletop_half_size, + material="g1_table_mat", + visual_name=TABLETOP_VISUAL_GEOM_NAME, + ) + + leg_half_size = (0.025, 0.025, tabletop_height_m / 2.0) + leg_z = tabletop_height_m / 2.0 + for i, x in enumerate((0.25, 0.85)): + for j, y in enumerate((-0.35, 0.35)): + _add_visible_collision_geom( + table_body, + name=f"g1_table_leg_{i}_{j}", + geom_type="box", + pos=[x, y, leg_z], + size=leg_half_size, + material="g1_table_mat", + ) + + +def _add_bin(worldbody: ET.Element, tabletop_height_m: float) -> None: + bin_body = _set_attrs( + ET.SubElement(worldbody, "body"), + name=PLACE_RECEPTACLE_BODY_NAME, + mocap="true", + pos=[0.43, 0.18, tabletop_height_m + 0.001], + quat=[1, 0, 0, 0], + ) + + outer_x = 0.14 + outer_y = 0.11 + wall = 0.012 + bottom_h = 0.012 + wall_h = 0.095 + + _add_visible_collision_geom( + bin_body, + name="g1_place_bin_bottom", + geom_type="box", + pos=[0, 0, bottom_h / 2.0], + size=[outer_x, outer_y, bottom_h / 2.0], + material="g1_bin_mat", + ) + for name, pos, size in ( + ( + "g1_place_bin_wall_pos_x", + [outer_x - wall / 2.0, 0, bottom_h + wall_h / 2.0], + [wall / 2.0, outer_y, wall_h / 2.0], + ), + ( + "g1_place_bin_wall_neg_x", + [-outer_x + wall / 2.0, 0, bottom_h + wall_h / 2.0], + [wall / 2.0, outer_y, wall_h / 2.0], + ), + ( + "g1_place_bin_wall_pos_y", + [0, outer_y - wall / 2.0, bottom_h + wall_h / 2.0], + [outer_x, wall / 2.0, wall_h / 2.0], + ), + ( + "g1_place_bin_wall_neg_y", + [0, -outer_y + wall / 2.0, bottom_h + wall_h / 2.0], + [outer_x, wall / 2.0, wall_h / 2.0], + ), + ): + _add_visible_collision_geom( + bin_body, + name=name, + geom_type="box", + pos=pos, + size=size, + material="g1_bin_mat", + ) + + _set_attrs( + ET.SubElement(bin_body, "site"), + name=PLACE_RECEPTACLE_SITE_NAME, + pos=[0, 0, bottom_h + wall_h], + size=[0.015], + rgba=[0.0, 0.8, 1.0, 0.6], + ) + + +def _metadata(scene_name: str, tabletop_height_m: float) -> dict: + return { + "scene_name": scene_name, + "objects": { + TABLE_BODY_NAME: { + "category": "Table", + "object_enum": "Table", + "is_static": True, + "tabletop_height_m": tabletop_height_m, + }, + PLACE_RECEPTACLE_BODY_NAME: { + "category": "Box", + "object_enum": "Box", + "is_static": True, + "boundingBox": {"x": 0.28, "y": 0.22, "z": 0.107}, + "name_map": {"sites": {PLACE_RECEPTACLE_SITE_NAME: PLACE_RECEPTACLE_SITE_NAME}}, + }, + }, + } + + +def _write_scene(output_dir: Path, scene_spec: TabletopSceneSpec) -> Path: + base_scene_path = ( + ABS_PATH_OF_TOP_LEVEL_MOLMO_SPACES_DIR / "molmo_spaces" / "resources" / "base_scene.xml" + ) + root = ET.parse(base_scene_path).getroot() + root.set("model", scene_spec.name) + _add_materials(root) + worldbody = root.find("worldbody") + if worldbody is None: + worldbody = ET.SubElement(root, "worldbody") + _set_attrs(ET.SubElement(worldbody, "light"), name="key", pos=[0, -1.5, 2.5], dir=[0, 0.5, -1]) + _add_floor(worldbody) + _add_table(worldbody, scene_spec.tabletop_height_m) + _add_bin(worldbody, scene_spec.tabletop_height_m) + + xml_path = output_dir / f"{scene_spec.name}.xml" + metadata_path = output_dir / f"{scene_spec.name}_metadata.json" + ET.indent(root, space=" ") + ET.ElementTree(root).write(xml_path, encoding="utf-8", xml_declaration=True) + with open(metadata_path, "w") as f: + json.dump(_metadata(scene_spec.name, scene_spec.tabletop_height_m), f, indent=2) + f.write("\n") + mujoco.MjModel.from_xml_path(str(xml_path)) + return xml_path + + +def prepare_unitree_g1_tabletop(output_dir: Path = OUTPUT_SCENE_DIR, force: bool = False) -> list[Path]: + output_dir = Path(output_dir) + if output_dir.exists() and force: + shutil.rmtree(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + return [_write_scene(output_dir, scene_spec) for scene_spec in SCENE_SPECS] + + +def main() -> None: + args = parse_args() + paths = prepare_unitree_g1_tabletop(args.output_dir, args.force) + for path in paths: + print(path) + + +if __name__ == "__main__": + main() From 9fb3afc3b5b0c66659542a8734d5b079860ebd94 Mon Sep 17 00:00:00 2001 From: Abhishek Parmar Date: Tue, 12 May 2026 02:40:15 +0800 Subject: [PATCH 11/13] Integrate G1 tabletop pick-and-place --- .../component_tests/test_unitree_g1.py | 405 +++- molmo_spaces/configs/policy_configs.py | 51 +- molmo_spaces/configs/robot_configs.py | 10 + molmo_spaces/configs/task_sampler_configs.py | 12 +- .../object_manipulation_datagen_configs.py | 142 +- molmo_spaces/data_generation/main.py | 36 + molmo_spaces/kinematics/mujoco_kinematics.py | 132 +- .../pick_and_place_planner_policy.py | 2028 ++++++++++++++++- .../tasks/unitree_g1_tabletop_task_sampler.py | 45 +- molmo_spaces/utils/grasp_sample.py | 20 + scripts/assets/prepare_unitree_g1.py | 47 +- scripts/assets/prepare_unitree_g1_tabletop.py | 6 +- 12 files changed, 2895 insertions(+), 39 deletions(-) diff --git a/mlspaces_tests/component_tests/test_unitree_g1.py b/mlspaces_tests/component_tests/test_unitree_g1.py index 3bda6bc5..c1ae3747 100644 --- a/mlspaces_tests/component_tests/test_unitree_g1.py +++ b/mlspaces_tests/component_tests/test_unitree_g1.py @@ -7,15 +7,27 @@ import pytest import molmo_spaces.molmo_spaces_constants as constants +from molmo_spaces.configs.policy_configs import PickAndPlacePlannerPolicyConfig from molmo_spaces.configs.robot_configs import ( UnitreeG1Dex1RobotConfig, UnitreeG1RightArmPickRobotConfig, + UnitreeG1RightArmTabletopPickRobotConfig, ) from molmo_spaces.data_generation.config_registry import get_config_class from molmo_spaces.kinematics.mujoco_kinematics import MlSpacesKinematics +from molmo_spaces.policy.solvers.object_manipulation.pick_and_place_planner_policy import ( + UnitreeG1RightArmPickAndPlacePlannerPolicy, +) from molmo_spaces.robots.unitree_g1 import UnitreeG1Robot from molmo_spaces.tasks.pick_task_sampler import UnitreeG1RightArmPickTaskSampler from scripts.assets.prepare_unitree_g1 import ( + DEX1_FINGERTIP_PAD_CONDIM, + DEX1_FINGERTIP_PAD_FRICTION, + DEX1_FINGERTIP_PAD_GROUP, + DEX1_FINGERTIP_PAD_POSITIONS, + DEX1_FINGERTIP_PAD_SIZE, + DEX1_GRASP_SITE_POS, + DEX1_HAND_FORCE_LIMIT_MULTIPLIER, LEFT_ARM_STOW_QPOS, LEFT_HAND_OPEN_QPOS, ROBOT_ASSET_NAME, @@ -45,10 +57,9 @@ def test_prepare_unitree_g1_tabletop_scenes_smoke(tmp_path): scene_paths = prepare_unitree_g1_tabletop(scene_dir) assert [path.name for path in scene_paths] == [ "unitree_g1_tabletop_pelvis_minus_10cm_v1.xml", - "unitree_g1_tabletop_pelvis_height_v1.xml", ] - expected_heights = [0.693, 0.793] + expected_heights = [0.693] for scene_path, expected_height in zip(scene_paths, expected_heights, strict=True): metadata_path = scene_path.with_name(scene_path.stem + "_metadata.json") assert metadata_path.exists() @@ -104,6 +115,85 @@ def _assert_locked_joints( assert data.ctrl[actuator_id] == pytest.approx(expected_qpos) +def test_unitree_g1_grasp_contact_filter_requires_right_fingertip_pad(): + policy = object.__new__(UnitreeG1RightArmPickAndPlacePlannerPolicy) + policy.config = SimpleNamespace( + robot_config=SimpleNamespace(robot_namespace="robot_0/") + ) + policy.policy_config = SimpleNamespace( + g1_reject_grasp_table_contact=True, + g1_reject_open_grasp_object_contact=True, + g1_closed_grasp_min_pad_geom_count=2, + g1_closed_grasp_max_object_shift_m=0.05, + ) + base_contact = { + "geom2_name": "Salt_Shaker_1/contact", + "body2_name": "Salt_Shaker_1", + "root2_name": "Salt_Shaker_1", + "class": "robot_pickup_object", + } + + pad_contact = { + **base_contact, + "geom1_name": "robot_0/right_dex1_fingertip_pad_1", + "body1_name": "robot_0/right_dex1_finger_link_1", + "root1_name": "robot_0/pelvis", + } + assert policy._g1_is_allowed_grasp_object_contact(pad_contact) + second_pad_contact = { + **base_contact, + "geom1_name": "robot_0/right_dex1_fingertip_pad_2", + "body1_name": "robot_0/right_dex1_finger_link_2", + "root1_name": "robot_0/pelvis", + } + + finger_link_contact = { + **base_contact, + "geom1_name": "robot_0/right_dex1_finger_collision", + "body1_name": "robot_0/right_dex1_finger_link_1", + "root1_name": "robot_0/pelvis", + } + assert not policy._g1_is_allowed_grasp_object_contact(finger_link_contact) + + wrist_contact = { + **base_contact, + "geom1_name": "robot_0/right_wrist_yaw_link_collision", + "body1_name": "robot_0/right_wrist_yaw_link", + "root1_name": "robot_0/pelvis", + } + assert not policy._g1_is_allowed_grasp_object_contact(wrist_contact) + assert ( + policy._g1_grasp_contact_quality_failure( + policy._g1_contact_quality_from_contacts([pad_contact]) + ) + == "open_grasp_object_contact" + ) + + quality = policy._g1_contact_quality_from_contacts( + [pad_contact, second_pad_contact, finger_link_contact] + ) + assert quality["pad_geom_count"] == 2 + assert not policy._g1_closed_grasp_contact_quality_failure(quality, 0.0) + + wrist_quality = policy._g1_contact_quality_from_contacts( + [pad_contact, second_pad_contact, wrist_contact] + ) + assert ( + policy._g1_closed_grasp_contact_quality_failure(wrist_quality, 0.0) + == "non_fingertip_object_contact" + ) + + single_pad_quality = policy._g1_contact_quality_from_contacts([pad_contact]) + assert ( + policy._g1_closed_grasp_contact_quality_failure(single_pad_quality, 0.0) + == "missing_closed_fingertip_pad_contact" + ) + assert ( + policy._g1_closed_grasp_contact_quality_failure(quality, 0.051) + == "closed_grasp_object_shift" + ) + + def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): unitree_root = os.environ.get("UNITREE_URDF_ROOT") if unitree_root is None: @@ -121,6 +211,27 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): assert model.qpos0[base_qposadr + 2] == pytest.approx(0.793) for site_name in KINEMATICS_SITE_NAMES: assert model.site(site_name).id >= 0 + for site_name in ("left_grasp_site", "right_grasp_site"): + np.testing.assert_allclose(model.site_pos[model.site(site_name).id], DEX1_GRASP_SITE_POS) + for side in ("left", "right"): + for finger_id, finger_suffix in (("1", "dex1_finger_link_1"), ("2", "dex1_finger_link_2")): + pad = model.geom(f"{side}_dex1_fingertip_pad_{finger_id}") + assert model.geom_type[pad.id] == mujoco.mjtGeom.mjGEOM_BOX + assert model.geom_contype[pad.id] == 1 + assert model.geom_conaffinity[pad.id] == 1 + assert model.geom_condim[pad.id] == DEX1_FINGERTIP_PAD_CONDIM + assert model.geom_group[pad.id] == DEX1_FINGERTIP_PAD_GROUP + np.testing.assert_allclose(model.geom_size[pad.id], DEX1_FINGERTIP_PAD_SIZE) + np.testing.assert_allclose( + model.geom_pos[pad.id], + DEX1_FINGERTIP_PAD_POSITIONS[finger_suffix], + ) + np.testing.assert_allclose(model.geom_friction[pad.id], DEX1_FINGERTIP_PAD_FRICTION) + actuator = model.actuator(f"{side}_dex1_finger_joint_{finger_id}") + np.testing.assert_allclose( + model.actuator_forcerange[actuator.id], + [-20.0 * DEX1_HAND_FORCE_LIMIT_MULTIPLIER, 20.0 * DEX1_HAND_FORCE_LIMIT_MULTIPLIER], + ) config = UnitreeG1Dex1RobotConfig() scene_spec = mujoco.MjSpec() @@ -208,19 +319,147 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): ) tabletop_datagen_config = tabletop_datagen_config_cls() assert tabletop_datagen_config.scene_dataset == "user" - assert isinstance(tabletop_datagen_config.robot_config, UnitreeG1RightArmPickRobotConfig) + assert isinstance( + tabletop_datagen_config.robot_config, + UnitreeG1RightArmTabletopPickRobotConfig, + ) + assert tabletop_datagen_config.robot_config.init_qpos["base"][:2] == [0.08, 0.0] assert ( tabletop_datagen_config.task_sampler_config.task_sampler_class.__name__ == "UnitreeG1RightArmTabletopPickAndPlaceTaskSampler" ) - assert tabletop_datagen_config.task_sampler_config.house_inds == [0, 1] + assert tabletop_datagen_config.task_sampler_config.house_inds == [0] assert tabletop_datagen_config.task_sampler_config.samples_per_house == 5 - assert len(tabletop_datagen_config.task_sampler_config.scene_xml_paths) == 2 + assert len(tabletop_datagen_config.task_sampler_config.scene_xml_paths) == 1 + assert tabletop_datagen_config.task_horizon == 360 + assert tabletop_datagen_config.end_on_success + assert tabletop_datagen_config.task_sampler_config.robot_base_xy == (0.08, 0.0) + assert tabletop_datagen_config.task_sampler_config.pickup_workspace_center_xy == (0.42, -0.20) + assert tabletop_datagen_config.task_sampler_config.pickup_workspace_size_xy == (0.16, 0.10) + assert tabletop_datagen_config.task_sampler_config.place_workspace_center_xy == (0.42, 0.08) + assert tabletop_datagen_config.task_sampler_config.place_workspace_size_xy == (0.10, 0.08) + assert tabletop_datagen_config.task_sampler_config.place_receptacle_half_size_xy == ( + 0.14, + 0.11, + ) + assert tabletop_datagen_config.task_sampler_config.pickup_receptacle_min_clearance == 0.015 + assert tabletop_datagen_config.task_sampler_config.max_place_receptacle_sampling_attempts == 100 + pickup_y_max = ( + tabletop_datagen_config.task_sampler_config.pickup_workspace_center_xy[1] + + tabletop_datagen_config.task_sampler_config.pickup_workspace_size_xy[1] / 2.0 + ) + place_y_min = ( + tabletop_datagen_config.task_sampler_config.place_workspace_center_xy[1] + - tabletop_datagen_config.task_sampler_config.place_workspace_size_xy[1] / 2.0 + ) + assert place_y_min - pickup_y_max > ( + tabletop_datagen_config.task_sampler_config.place_receptacle_half_size_xy[1] + + tabletop_datagen_config.task_sampler_config.pickup_receptacle_min_clearance + ) assert tabletop_datagen_config.policy_config.policy_cls.__name__ == ( "UnitreeG1RightArmPickAndPlacePlannerPolicy" ) assert not tabletop_datagen_config.policy_config.filter_colliding_grasps assert not tabletop_datagen_config.policy_config.filter_feasible_grasps + assert PickAndPlacePlannerPolicyConfig().tcp_rot_err_threshold == np.radians(30.0) + assert PickAndPlacePlannerPolicyConfig().gripper_empty_threshold == 0.002 + assert tabletop_datagen_config.policy_config.pregrasp_z_offset == 0.12 + assert tabletop_datagen_config.policy_config.gripper_close_duration == 1.0 + assert tabletop_datagen_config.policy_config.gripper_empty_threshold == 0.0005 + assert tabletop_datagen_config.policy_config.grasp_vertical_cost_weight == 0.5 + assert tabletop_datagen_config.policy_config.grasp_rot_cost_weight == 0.02 + assert tabletop_datagen_config.policy_config.pregrasp_tcp_pos_err_threshold == 0.1 + assert np.isinf(tabletop_datagen_config.policy_config.pregrasp_tcp_rot_err_threshold) + assert tabletop_datagen_config.policy_config.grasp_feasibility_max_grasps == 128 + assert tabletop_datagen_config.policy_config.g1_online_grasp_selector + assert tabletop_datagen_config.policy_config.g1_grasp_candidate_limit == 256 + assert tabletop_datagen_config.policy_config.g1_grasp_ik_eval_limit == 256 + assert tabletop_datagen_config.policy_config.g1_grasp_require_all_pick_place_phases + assert tabletop_datagen_config.policy_config.g1_grasp_joint_margin_weight == 0.5 + assert tabletop_datagen_config.policy_config.g1_grasp_joint_motion_weight == 0.25 + assert tabletop_datagen_config.policy_config.g1_grasp_topdown_weight == 1.0 + assert tabletop_datagen_config.policy_config.g1_grasp_min_vertical_axis_z == 0.75 + assert tabletop_datagen_config.policy_config.g1_grasp_max_tcp_rot_deg == 120.0 + assert tabletop_datagen_config.policy_config.g1_ignore_flipped_grasps + assert tabletop_datagen_config.policy_config.g1_grasp_inward_xy_offset == 0.006 + assert tabletop_datagen_config.policy_config.g1_grasp_table_clearance == 0.065 + assert tabletop_datagen_config.policy_config.g1_center_grasp_lateral + assert tabletop_datagen_config.policy_config.g1_grasp_lateral_centering_scale == 1.0 + assert tabletop_datagen_config.policy_config.g1_grasp_lateral_centering_max_offset == 0.02 + assert not tabletop_datagen_config.policy_config.g1_level_grasp_orientation + assert tabletop_datagen_config.policy_config.g1_grasp_level_max_tilt_deg == 35.0 + assert not tabletop_datagen_config.policy_config.g1_require_fingertip_pad_grasp_contact + assert tabletop_datagen_config.policy_config.g1_reject_non_fingertip_grasp_object_contact + assert tabletop_datagen_config.policy_config.g1_reject_grasp_table_contact + assert tabletop_datagen_config.policy_config.g1_reject_open_grasp_object_contact + assert tabletop_datagen_config.policy_config.g1_grasp_single_pad_contact_penalty == 0.5 + assert tabletop_datagen_config.policy_config.g1_closed_grasp_quality_enabled + assert tabletop_datagen_config.policy_config.g1_closed_grasp_settle_steps == 25 + assert tabletop_datagen_config.policy_config.g1_closed_grasp_min_pad_geom_count == 2 + assert tabletop_datagen_config.policy_config.g1_closed_grasp_max_object_shift_m == 0.05 + assert tabletop_datagen_config.policy_config.g1_closed_grasp_penalty_per_missing_pad == 1.0 + assert tabletop_datagen_config.policy_config.g1_pregrasp_min_vertical_lift == 0.10 + assert tabletop_datagen_config.policy_config.g1_pregrasp_object_clearance == 0.12 + assert tabletop_datagen_config.policy_config.g1_place_travel_object_clearance == 0.13 + assert tabletop_datagen_config.policy_config.g1_held_object_speed == 0.04 + assert tabletop_datagen_config.policy_config.g1_postgrasp_hold_duration == 0.4 + assert tabletop_datagen_config.policy_config.g1_record_partial_attempt_on_no_full_grasp_candidate + assert not tabletop_datagen_config.policy_config.g1_pick_lift_only + assert tabletop_datagen_config.policy_config.enable_failure_diagnostics + assert tabletop_datagen_config.policy_config.diagnostic_large_object_max_extent_m == 0.18 + assert tabletop_datagen_config.policy_config.diagnostic_failure_hold_duration_s == 2.0 + assert tabletop_datagen_config.policy_config.record_unfiltered_attempt_on_no_feasible_grasp + assert not tabletop_datagen_config.policy_config.g1_ik_debug + assert tabletop_datagen_config.policy_config.g1_ik_debug_higher_z_offset == 0.05 + assert tabletop_datagen_config.policy_config.g1_ik_debug_top_k_grasps == 5 + diagnostic_datagen_config_cls = get_config_class( + "UnitreeG1RightArmTabletopPickAndPlaceDiagnosticDataGenConfig" + ) + diagnostic_datagen_config = diagnostic_datagen_config_cls() + assert diagnostic_datagen_config.task_sampler_config.added_pickup_objects == [ + "Apple_22", + "Egg_9", + "Potato_12", + "Salt_Shaker_1", + "Cup_5", + "Apple_18", + "Tomato_26", + "Apple_4", + ] + assert diagnostic_datagen_config.task_sampler_config.num_added_pickups == 8 + assert diagnostic_datagen_config.task_sampler_config.samples_per_house == 8 + assert diagnostic_datagen_config.policy_config.g1_ik_debug + assert diagnostic_datagen_config.policy_config.g1_ik_debug_top_k_grasps == 8 + assert not diagnostic_datagen_config.policy_config.filter_feasible_grasps + assert diagnostic_datagen_config.policy_config.pregrasp_tcp_pos_err_threshold == 0.1 + assert np.isinf(diagnostic_datagen_config.policy_config.pregrasp_tcp_rot_err_threshold) + assert diagnostic_datagen_config.policy_config.g1_online_grasp_selector + assert diagnostic_datagen_config.policy_config.g1_record_partial_attempt_on_no_full_grasp_candidate + viewer_debug_datagen_config_cls = get_config_class( + "UnitreeG1RightArmTabletopPickAndPlaceViewerDebugConfig" + ) + viewer_debug_datagen_config = viewer_debug_datagen_config_cls() + assert viewer_debug_datagen_config.use_passive_viewer + assert viewer_debug_datagen_config.task_sampler_config.added_pickup_objects == [ + "Salt_Shaker_1" + ] + assert viewer_debug_datagen_config.task_sampler_config.samples_per_house == 1 + assert viewer_debug_datagen_config.policy_config.debug_poses + assert viewer_debug_datagen_config.policy_config.g1_ik_debug + pick_lift_debug_datagen_config_cls = get_config_class( + "UnitreeG1RightArmTabletopPickLiftViewerDebugConfig" + ) + pick_lift_debug_datagen_config = pick_lift_debug_datagen_config_cls() + assert pick_lift_debug_datagen_config.use_passive_viewer + assert pick_lift_debug_datagen_config.task_sampler_config.added_pickup_objects == [ + "Salt_Shaker_1" + ] + assert pick_lift_debug_datagen_config.task_sampler_config.samples_per_house == 1 + assert pick_lift_debug_datagen_config.task_horizon == 120 + assert pick_lift_debug_datagen_config.policy_config.debug_poses + assert pick_lift_debug_datagen_config.policy_config.g1_ik_debug + assert not pick_lift_debug_datagen_config.policy_config.g1_grasp_require_all_pick_place_phases + assert pick_lift_debug_datagen_config.policy_config.g1_pick_lift_only base_scene_path = constants.ABS_PATH_OF_TOP_LEVEL_MOLMO_SPACES_DIR / ( "molmo_spaces/resources/base_scene.xml" @@ -299,6 +538,162 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): ) assert pick_ik is not None assert set(pick_ik) == {"base", "right_arm", "gripper"} + pick_ik_diag = pick_kinematics.diagnose_ik( + "gripper", + target_pose, + ["right_arm"], + pick_datagen_config.robot_config.init_qpos, + np.eye(4), + max_iter=250, + ) + assert pick_ik_diag["success"] + assert pick_ik_diag["final_pos_error_norm"] < 1e-3 + assert pick_ik_diag["clamp_event_count"] >= 0 + unreachable_target_pose = target_pose.copy() + unreachable_target_pose[0, 3] += 10.0 + unreachable_ik_diag = pick_kinematics.diagnose_ik( + "gripper", + unreachable_target_pose, + ["right_arm"], + pick_datagen_config.robot_config.init_qpos, + np.eye(4), + max_iter=25, + ) + assert not unreachable_ik_diag["success"] + assert unreachable_ik_diag["final_pos_error_norm"] > 1.0 + fake_tabletop_policy = object.__new__(UnitreeG1RightArmPickAndPlacePlannerPolicy) + fake_tabletop_policy.robot_view = pick_view + fake_tabletop_policy.policy_config = tabletop_datagen_config.policy_config + fake_tabletop_policy.task = SimpleNamespace( + env=SimpleNamespace( + current_robot=SimpleNamespace( + robot_view=pick_view, + kinematics=pick_kinematics, + ) + ) + ) + for group_name, qpos in pick_datagen_config.robot_config.init_qpos.items(): + pick_view.get_move_group(group_name).joint_pos = np.asarray(qpos) + reachable_targets = { + phase: target_pose.copy() + for phase in ("pregrasp", "grasp", "lift", "preplace", "place") + } + reachable_eval = fake_tabletop_policy._g1_eval_candidate_target_poses( + reachable_targets, + pick_datagen_config.robot_config.init_qpos, + np.eye(4), + ) + assert reachable_eval["success"] + unreachable_targets = { + phase: unreachable_target_pose.copy() + for phase in ("pregrasp", "grasp", "lift", "preplace", "place") + } + unreachable_eval = fake_tabletop_policy._g1_eval_candidate_target_poses( + unreachable_targets, + pick_datagen_config.robot_config.init_qpos, + np.eye(4), + ) + assert not unreachable_eval["success"] + assert unreachable_eval["failed_phase"] == "pregrasp" + + clearance_model = mujoco.MjModel.from_xml_string( + """ + + + + + + + + + + + + + + """ + ) + clearance_data = mujoco.MjData(clearance_model) + mujoco.mj_forward(clearance_model, clearance_data) + fake_clearance_policy = object.__new__(UnitreeG1RightArmPickAndPlacePlannerPolicy) + fake_clearance_policy.policy_config = tabletop_datagen_config.policy_config + fake_clearance_policy.config = SimpleNamespace( + task_sampler_config=SimpleNamespace(table_body_name="g1_table") + ) + fake_clearance_policy.task = SimpleNamespace( + env=SimpleNamespace(current_data=clearance_data) + ) + pickup_obj = SimpleNamespace( + object_id=clearance_model.body("pickup").id, + position=np.array([0.0, 0.0, 0.75]), + ) + place_receptacle = SimpleNamespace( + object_id=clearance_model.body("bin").id, + position=np.array([0.2, 0.0, 0.70]), + ) + grasp_pose = np.eye(4) + tilt_angle = np.radians(24.0) + grasp_pose[:3, :3] = np.array( + [ + [1.0, 0.0, 0.0], + [0.0, np.cos(tilt_angle), -np.sin(tilt_angle)], + [0.0, np.sin(tilt_angle), np.cos(tilt_angle)], + ] + ) + grasp_pose[:3, 3] = [0.0, -0.015, 0.75] + g1_target_poses = fake_clearance_policy._build_g1_candidate_target_poses( + grasp_pose, + pickup_obj, + place_receptacle, + ) + pickup_top_z = 0.80 + table_top_z = 0.70 + assert g1_target_poses["grasp"][2, 3] >= ( + table_top_z + tabletop_datagen_config.policy_config.g1_grasp_table_clearance + ) + object_in_grasp = g1_target_poses["grasp"][:3, :3].T @ ( + pickup_obj.position - g1_target_poses["grasp"][:3, 3] + ) + assert abs(object_in_grasp[1]) < 1e-8 + np.testing.assert_allclose(g1_target_poses["grasp"][:3, 2], grasp_pose[:3, 2]) + + fake_clearance_policy.policy_config = tabletop_datagen_config.policy_config.model_copy( + update={"g1_level_grasp_orientation": True} + ) + leveled_target_poses = fake_clearance_policy._build_g1_candidate_target_poses( + grasp_pose, + pickup_obj, + place_receptacle, + ) + np.testing.assert_allclose(leveled_target_poses["grasp"][:3, 2], [0.0, 0.0, 1.0]) + np.testing.assert_allclose( + leveled_target_poses["grasp"][:3, :3].T @ leveled_target_poses["grasp"][:3, :3], + np.eye(3), + atol=1e-7, + ) + assert g1_target_poses["pregrasp"][2, 3] >= grasp_pose[2, 3] + 0.10 + assert g1_target_poses["pregrasp"][2, 3] >= pickup_top_z + 0.12 + pickup_bottom_z = 0.70 + receptacle_top_z = 0.75 + min_travel_bottom_z = ( + receptacle_top_z + + tabletop_datagen_config.policy_config.g1_place_travel_object_clearance + ) + for phase in ("lift", "preplace"): + carried_bottom_z = ( + g1_target_poses[phase][2, 3] + + pickup_bottom_z + - g1_target_poses["grasp"][2, 3] + ) + assert carried_bottom_z >= min_travel_bottom_z + place_bottom_z = ( + g1_target_poses["place"][2, 3] + + pickup_bottom_z + - g1_target_poses["grasp"][2, 3] + ) + assert place_bottom_z < min_travel_bottom_z pick_robot = UnitreeG1Robot( pick_scene_data, SimpleNamespace(robot_config=pick_datagen_config.robot_config) diff --git a/molmo_spaces/configs/policy_configs.py b/molmo_spaces/configs/policy_configs.py index 464032c9..dc41f45e 100644 --- a/molmo_spaces/configs/policy_configs.py +++ b/molmo_spaces/configs/policy_configs.py @@ -182,15 +182,64 @@ class UnitreeG1RightArmPickAndPlacePlannerPolicyConfig(PickAndPlacePlannerPolicy """Pick-and-place planner config for fixed-base G1 right-arm tabletop datagen.""" policy_cls: type = None + pregrasp_z_offset: float = 0.12 + gripper_close_duration: float = 1.0 + gripper_empty_threshold: float = 0.0005 + grasp_rot_cost_weight: float = 0.02 + grasp_vertical_cost_weight: float = 0.5 filter_colliding_grasps: bool = False filter_feasible_grasps: bool = False grasp_feasibility_batch_size: int = 32 - grasp_feasibility_max_grasps: int = 32 + grasp_feasibility_max_grasps: int = 128 grasp_collision_batch_size: int = 32 grasp_collision_max_grasps: int = 64 max_retries: int = 0 max_sequential_ik_failures: int = 4 phase_timeout: float = 24.0 + pregrasp_tcp_pos_err_threshold: float = 0.1 + pregrasp_tcp_rot_err_threshold: float = float("inf") + enable_failure_diagnostics: bool = True + diagnostic_large_object_max_extent_m: float = 0.18 + diagnostic_large_object_volume_m3: float = 0.004 + diagnostic_failure_hold_duration_s: float = 2.0 + record_unfiltered_attempt_on_no_feasible_grasp: bool = True + g1_ik_debug: bool = False + g1_ik_debug_higher_z_offset: float = 0.05 + g1_ik_debug_top_k_grasps: int = 5 + g1_online_grasp_selector: bool = True + g1_grasp_candidate_limit: int = 256 + g1_grasp_ik_eval_limit: int = 256 + g1_grasp_require_all_pick_place_phases: bool = True + g1_grasp_joint_margin_weight: float = 0.5 + g1_grasp_joint_motion_weight: float = 0.25 + g1_grasp_topdown_weight: float = 1.0 + g1_grasp_min_vertical_axis_z: float = 0.75 + g1_grasp_max_tcp_rot_deg: float = 120.0 + g1_ignore_flipped_grasps: bool = True + g1_grasp_inward_xy_offset: float = 0.006 + g1_grasp_table_clearance: float = 0.065 + g1_center_grasp_lateral: bool = True + g1_grasp_lateral_centering_scale: float = 1.0 + g1_grasp_lateral_centering_max_offset: float = 0.02 + g1_level_grasp_orientation: bool = False + g1_grasp_level_max_tilt_deg: float = 35.0 + g1_require_fingertip_pad_grasp_contact: bool = False + g1_reject_non_fingertip_grasp_object_contact: bool = True + g1_reject_grasp_table_contact: bool = True + g1_reject_open_grasp_object_contact: bool = True + g1_grasp_single_pad_contact_penalty: float = 0.5 + g1_closed_grasp_quality_enabled: bool = True + g1_closed_grasp_settle_steps: int = 25 + g1_closed_grasp_min_pad_geom_count: int = 2 + g1_closed_grasp_max_object_shift_m: float = 0.05 + g1_closed_grasp_penalty_per_missing_pad: float = 1.0 + g1_pregrasp_min_vertical_lift: float = 0.10 + g1_pregrasp_object_clearance: float = 0.12 + g1_place_travel_object_clearance: float = 0.13 + g1_held_object_speed: float = 0.04 + g1_postgrasp_hold_duration: float = 0.4 + g1_record_partial_attempt_on_no_full_grasp_candidate: bool = True + g1_pick_lift_only: bool = False def model_post_init(self, __context) -> None: """Set policy_cls after initialization to avoid circular imports.""" diff --git a/molmo_spaces/configs/robot_configs.py b/molmo_spaces/configs/robot_configs.py index c97255f6..4a7bf2de 100644 --- a/molmo_spaces/configs/robot_configs.py +++ b/molmo_spaces/configs/robot_configs.py @@ -488,3 +488,13 @@ class UnitreeG1RightArmPickRobotConfig(UnitreeG1Dex1RobotConfig): "right_arm": "joint_position", "gripper": "joint_position", } + + +class UnitreeG1RightArmTabletopPickRobotConfig(UnitreeG1RightArmPickRobotConfig): + """Right-arm G1 config that starts at the fixed tabletop base pose.""" + + init_qpos: dict[str, list[float]] = { + "base": [0.08, 0.0, 0.793, 1.0, 0.0, 0.0, 0.0], + "right_arm": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + "gripper": [0.0245, 0.0245], + } diff --git a/molmo_spaces/configs/task_sampler_configs.py b/molmo_spaces/configs/task_sampler_configs.py index 6a2bc551..559cd2ae 100644 --- a/molmo_spaces/configs/task_sampler_configs.py +++ b/molmo_spaces/configs/task_sampler_configs.py @@ -216,10 +216,14 @@ class UnitreeG1TabletopPickAndPlaceTaskSamplerConfig(PickAndPlaceTaskSamplerConf table_body_name: str = "g1_table" table_geom_name: str = "g1_tabletop_geom" place_receptacle_name: str = "g1_place_bin" - pickup_workspace_center_xy: tuple[float, float] = (0.42, -0.16) - pickup_workspace_size_xy: tuple[float, float] = (0.18, 0.22) - place_workspace_center_xy: tuple[float, float] = (0.42, 0.16) - place_workspace_size_xy: tuple[float, float] = (0.16, 0.20) + robot_base_xy: tuple[float, float] = (0.08, 0.0) + pickup_workspace_center_xy: tuple[float, float] = (0.42, -0.20) + pickup_workspace_size_xy: tuple[float, float] = (0.16, 0.10) + place_workspace_center_xy: tuple[float, float] = (0.42, 0.08) + place_workspace_size_xy: tuple[float, float] = (0.10, 0.08) + place_receptacle_half_size_xy: tuple[float, float] = (0.14, 0.11) + pickup_receptacle_min_clearance: float = 0.015 + max_place_receptacle_sampling_attempts: int = 100 object_table_clearance: float = 0.006 receptacle_table_clearance: float = 0.001 diff --git a/molmo_spaces/data_generation/config/object_manipulation_datagen_configs.py b/molmo_spaces/data_generation/config/object_manipulation_datagen_configs.py index 85ea88f4..c43a1a18 100644 --- a/molmo_spaces/data_generation/config/object_manipulation_datagen_configs.py +++ b/molmo_spaces/data_generation/config/object_manipulation_datagen_configs.py @@ -44,6 +44,7 @@ RBY1MOpenCloseConfig, UnitreeG1Dex1RobotConfig, UnitreeG1RightArmPickRobotConfig, + UnitreeG1RightArmTabletopPickRobotConfig, ) from molmo_spaces.configs.task_sampler_configs import ( OpenTaskSamplerConfig, @@ -73,6 +74,17 @@ from molmo_spaces.utils.constants.object_constants import PICK_AND_PLACE_OBJECTS from molmo_spaces.utils.synset_utils import get_valid_pickupable_obja_uids +G1_TABLETOP_DIAGNOSTIC_PICKUP_OBJECTS = [ + "Apple_22", + "Egg_9", + "Potato_12", + "Salt_Shaker_1", + "Cup_5", + "Apple_18", + "Tomato_26", + "Apple_4", +] + @register_config("FrankaPickDroidDataGenConfig") class FrankaPickDroidDataGenConfig(PickBaseConfig): @@ -183,12 +195,14 @@ class UnitreeG1RightArmTabletopPickAndPlaceDataGenConfig(PickAndPlaceDataGenConf """Right-arm-only Unitree G1 pick-and-place datagen for a fixed tabletop user scene.""" scene_dataset: str = "user" - robot_config: UnitreeG1RightArmPickRobotConfig = UnitreeG1RightArmPickRobotConfig() + robot_config: UnitreeG1RightArmTabletopPickRobotConfig = ( + UnitreeG1RightArmTabletopPickRobotConfig() + ) camera_config: UnitreeG1RightArmPickCameraSystem = UnitreeG1RightArmPickCameraSystem() task_sampler_config: UnitreeG1TabletopPickAndPlaceTaskSamplerConfig = ( UnitreeG1TabletopPickAndPlaceTaskSamplerConfig( task_sampler_class=UnitreeG1RightArmTabletopPickAndPlaceTaskSampler, - house_inds=[0, 1], + house_inds=[0], scene_xml_paths=[ str( ASSETS_DIR @@ -196,12 +210,6 @@ class UnitreeG1RightArmTabletopPickAndPlaceDataGenConfig(PickAndPlaceDataGenConf / "unitree_g1_tabletop_v1" / "unitree_g1_tabletop_pelvis_minus_10cm_v1.xml" ), - str( - ASSETS_DIR - / "scenes" - / "unitree_g1_tabletop_v1" - / "unitree_g1_tabletop_pelvis_height_v1.xml" - ), ], samples_per_house=5, added_pickup_objects=get_valid_pickupable_obja_uids(), @@ -214,7 +222,8 @@ class UnitreeG1RightArmTabletopPickAndPlaceDataGenConfig(PickAndPlaceDataGenConf policy_config: UnitreeG1RightArmPickAndPlacePlannerPolicyConfig = ( UnitreeG1RightArmPickAndPlacePlannerPolicyConfig() ) - task_horizon: int = 180 + task_horizon: int = 360 + end_on_success: bool = True filter_for_successful_trajectories: bool = False profile: bool = False output_dir: Path = ( @@ -226,6 +235,121 @@ def tag(self) -> str: return "unitree_g1_right_arm_tabletop_pick_and_place_datagen" +@register_config("UnitreeG1RightArmTabletopPickAndPlaceDiagnosticDataGenConfig") +class UnitreeG1RightArmTabletopPickAndPlaceDiagnosticDataGenConfig( + UnitreeG1RightArmTabletopPickAndPlaceDataGenConfig +): + """Diagnostic G1 tabletop run with small fixed objects and detailed IK logs.""" + + task_sampler_config: UnitreeG1TabletopPickAndPlaceTaskSamplerConfig = ( + UnitreeG1TabletopPickAndPlaceTaskSamplerConfig( + task_sampler_class=UnitreeG1RightArmTabletopPickAndPlaceTaskSampler, + house_inds=[0], + scene_xml_paths=[ + str( + ASSETS_DIR + / "scenes" + / "unitree_g1_tabletop_v1" + / "unitree_g1_tabletop_pelvis_minus_10cm_v1.xml" + ), + ], + samples_per_house=len(G1_TABLETOP_DIAGNOSTIC_PICKUP_OBJECTS), + added_pickup_objects=G1_TABLETOP_DIAGNOSTIC_PICKUP_OBJECTS, + num_added_pickups=len(G1_TABLETOP_DIAGNOSTIC_PICKUP_OBJECTS), + episodes_per_added_pickup=1, + check_robot_placement_visibility=False, + robot_safety_radius=0.25, + ) + ) + policy_config: UnitreeG1RightArmPickAndPlacePlannerPolicyConfig = ( + UnitreeG1RightArmPickAndPlacePlannerPolicyConfig( + g1_ik_debug=True, + g1_ik_debug_top_k_grasps=8, + ) + ) + output_dir: Path = ( + ASSETS_DIR / "experiment_output" / "datagen" / "unitree_g1_tabletop_pnp_diagnostic_v1" + ) + + @property + def tag(self) -> str: + return "unitree_g1_right_arm_tabletop_pick_and_place_diagnostic_datagen" + + +@register_config("UnitreeG1RightArmTabletopPickAndPlaceViewerDebugConfig") +class UnitreeG1RightArmTabletopPickAndPlaceViewerDebugConfig( + UnitreeG1RightArmTabletopPickAndPlaceDiagnosticDataGenConfig +): + """Single-object G1 tabletop run for inspecting planner waypoints in the MuJoCo viewer.""" + + use_passive_viewer: bool = True + task_sampler_config: UnitreeG1TabletopPickAndPlaceTaskSamplerConfig = ( + UnitreeG1TabletopPickAndPlaceTaskSamplerConfig( + task_sampler_class=UnitreeG1RightArmTabletopPickAndPlaceTaskSampler, + house_inds=[0], + scene_xml_paths=[ + str( + ASSETS_DIR + / "scenes" + / "unitree_g1_tabletop_v1" + / "unitree_g1_tabletop_pelvis_minus_10cm_v1.xml" + ), + ], + samples_per_house=1, + added_pickup_objects=["Salt_Shaker_1"], + num_added_pickups=1, + episodes_per_added_pickup=1, + check_robot_placement_visibility=False, + robot_safety_radius=0.25, + ) + ) + policy_config: UnitreeG1RightArmPickAndPlacePlannerPolicyConfig = ( + UnitreeG1RightArmPickAndPlacePlannerPolicyConfig( + debug_poses=True, + g1_ik_debug=True, + g1_ik_debug_top_k_grasps=8, + ) + ) + output_dir: Path = ( + ASSETS_DIR + / "experiment_output" + / "datagen" + / "unitree_g1_tabletop_pnp_viewer_debug_v1" + ) + + @property + def tag(self) -> str: + return "unitree_g1_right_arm_tabletop_pick_and_place_viewer_debug" + + +@register_config("UnitreeG1RightArmTabletopPickLiftViewerDebugConfig") +class UnitreeG1RightArmTabletopPickLiftViewerDebugConfig( + UnitreeG1RightArmTabletopPickAndPlaceViewerDebugConfig +): + """Single-object G1 tabletop run that stops after lift for grasp/contact inspection.""" + + policy_config: UnitreeG1RightArmPickAndPlacePlannerPolicyConfig = ( + UnitreeG1RightArmPickAndPlacePlannerPolicyConfig( + debug_poses=True, + g1_ik_debug=True, + g1_ik_debug_top_k_grasps=8, + g1_grasp_require_all_pick_place_phases=False, + g1_pick_lift_only=True, + ) + ) + task_horizon: int = 120 + output_dir: Path = ( + ASSETS_DIR + / "experiment_output" + / "datagen" + / "unitree_g1_tabletop_pick_lift_viewer_debug_v1" + ) + + @property + def tag(self) -> str: + return "unitree_g1_right_arm_tabletop_pick_lift_viewer_debug" + + @register_config("FrankaPickAndPlaceDataGenConfig") class FrankaPickAndPlaceDataGenConfig(PickAndPlaceDataGenConfig): robot_config: BaseRobotConfig = FrankaRobotConfig() diff --git a/molmo_spaces/data_generation/main.py b/molmo_spaces/data_generation/main.py index 03a585ec..9b88dc6a 100644 --- a/molmo_spaces/data_generation/main.py +++ b/molmo_spaces/data_generation/main.py @@ -40,6 +40,24 @@ def get_args(): "optionally with the module name prepended with a colon (e.g. molmo_spaces.data_generation.config.object_manipulation_datagen_configs:FrankaPickDroidDataGenConfig). " "If the module is specified, only that module will be imported to populate the registry. Otherwise, all config files will be imported.", ) + parser.add_argument( + "--enable-g1-ik-grasp-filtering", + action="store_true", + help=( + "Enable IK-feasible grasp filtering for G1 tabletop debugging. " + "By default, G1 uses the unfiltered grasp ranking path so each episode records an attempt." + ), + ) + parser.add_argument( + "--g1-ik-debug", + action="store_true", + help="Write detailed G1 tabletop IK diagnostics next to the datagen run log.", + ) + parser.add_argument( + "--viewer", + action="store_true", + help="Launch the configured experiment with the existing MuJoCo passive viewer.", + ) return parser.parse_args() @@ -94,8 +112,26 @@ def main() -> None: exp_config_args = vars(args) # pop exp_config_cls from args exp_config_args.pop("exp_config_cls") + enable_g1_ik_grasp_filtering = exp_config_args.pop("enable_g1_ik_grasp_filtering") + g1_ik_debug = exp_config_args.pop("g1_ik_debug") + viewer = exp_config_args.pop("viewer") # Create exp_config instance with args exp_config = ExpConfigClass(**exp_config_args) + if viewer: + exp_config.use_passive_viewer = True + if enable_g1_ik_grasp_filtering: + policy_config = getattr(exp_config, "policy_config", None) + if policy_config is None or not hasattr(policy_config, "filter_feasible_grasps"): + raise ValueError( + "--enable-g1-ik-grasp-filtering requires a policy config with " + "filter_feasible_grasps" + ) + policy_config.filter_feasible_grasps = True + if g1_ik_debug: + policy_config = getattr(exp_config, "policy_config", None) + if policy_config is None or not hasattr(policy_config, "g1_ik_debug"): + raise ValueError("--g1-ik-debug requires a G1 policy config with g1_ik_debug") + policy_config.g1_ik_debug = True # Optional: Modify the config parameters here if needed # Eg. for hyperparamter sweeps etc. diff --git a/molmo_spaces/kinematics/mujoco_kinematics.py b/molmo_spaces/kinematics/mujoco_kinematics.py index 087d23f0..445ce894 100644 --- a/molmo_spaces/kinematics/mujoco_kinematics.py +++ b/molmo_spaces/kinematics/mujoco_kinematics.py @@ -4,7 +4,7 @@ this is not suitable for large batches. """ -from typing import Literal, TYPE_CHECKING +from typing import TYPE_CHECKING, Literal import mujoco import numpy as np @@ -252,3 +252,133 @@ def ik( if succ: return self._robot_view.get_qpos_dict() return None + + def diagnose_ik( + self, + move_group_id: str, + pose: np.ndarray, + unlocked_move_group_ids: list[str] | None, + q0: dict[str, np.ndarray], + base_pose: np.ndarray, + rel_to_base: bool = False, + eps: float = 1e-4, + max_iter: int = 1000, + damping: float = 1e-12, + dt: float = 1.0, + ) -> dict: + """Solve IK and return structured diagnostics instead of only success/failure.""" + if set(q0.keys()) != set(self._robot_view.move_group_ids()): + raise ValueError( + f"q0 keys must match move group ids: {set(q0.keys())} != {set(self._robot_view.move_group_ids())}" + ) + + if unlocked_move_group_ids is None: + unlocked_move_group_ids = self._robot_view.move_group_ids() + + requested_unlocked_groups = list(unlocked_move_group_ids) + unavailable_groups = [ + mg_id for mg_id in unlocked_move_group_ids if mg_id not in self._robot_view.move_group_ids() + ] + unlocked_move_group_ids = [ + mg_id for mg_id in unlocked_move_group_ids if mg_id in self._robot_view.move_group_ids() + ] + if not unlocked_move_group_ids: + return { + "success": False, + "reason": "no_available_unlocked_groups", + "requested_unlocked_groups": requested_unlocked_groups, + "unavailable_groups": unavailable_groups, + "qpos": None, + } + + self._robot_view.base.pose = base_pose + self._robot_view.set_qpos_dict(q0) + if rel_to_base: + pose = self._robot_view.base.pose @ pose + + move_group = self._robot_view.get_move_group(move_group_id) + singular_iterations = 0 + min_singular_value = float("inf") + min_jjt_det = float("inf") + clamp_events: list[dict] = [] + err_norm = float("inf") + pos_err_norm = float("inf") + rot_err_norm = float("inf") + iterations = 0 + succ = False + + for i in range(max_iter): + iterations = i + 1 + mujoco.mj_fwdPosition(self._mj_model, self._mj_data) + mujoco.mj_sensorPos(self._mj_model, self._mj_data) + + ee_pose = relative_to_global_transform( + move_group.leaf_frame_to_robot, self._robot_view.base.pose + ) + err_trf = inverse_homogeneous_matrix(ee_pose) @ pose + twist_lin, twist_ang = transform_to_twist(err_trf) + err = np.concatenate([ee_pose[:3, :3] @ twist_lin, ee_pose[:3, :3] @ twist_ang]) + pos_err_norm = float(np.linalg.norm(err[:3])) + rot_err_norm = float(np.linalg.norm(err[3:])) + err_norm = float(np.linalg.norm(err)) + if err_norm < eps: + succ = True + break + if i == max_iter - 1: + break + + J: np.ndarray = self._robot_view.get_jacobian(move_group_id, unlocked_move_group_ids) + singular_values = np.linalg.svd(J, compute_uv=False) + if singular_values.size: + min_singular_value = min(min_singular_value, float(singular_values[-1])) + jjt_det = float(np.linalg.det(J @ J.T)) + min_jjt_det = min(min_jjt_det, jjt_det) + if jjt_det < 1e-20: + singular_iterations += 1 + + H = J @ J.T + damping * np.eye(J.shape[0]) + q_dot = J.T @ np.linalg.solve(H, err) + dq = q_dot * dt + + j = 0 + for mg_id in unlocked_move_group_ids: + mg = self._robot_view.get_move_group(mg_id) + before = mg.joint_pos.copy() + proposed = mg.integrate_joint_vel(mg.joint_pos, dq[j : j + mg.vel_dim]) + clipped = np.clip(proposed, mg.joint_pos_limits[:, 0], mg.joint_pos_limits[:, 1]) + clamped = np.abs(proposed - clipped) > 1e-9 + if np.any(clamped): + clamp_events.append( + { + "iteration": i, + "move_group": mg_id, + "joint_indices": np.flatnonzero(clamped).astype(int).tolist(), + "before": np.round(before[clamped], 6).tolist(), + "proposed": np.round(proposed[clamped], 6).tolist(), + "clipped": np.round(clipped[clamped], 6).tolist(), + } + ) + mg.joint_pos = clipped + j += mg.vel_dim + + qpos = self._robot_view.get_qpos_dict() + final_pose = relative_to_global_transform(move_group.leaf_frame_to_robot, self._robot_view.base.pose) + return { + "success": bool(succ), + "reason": "converged" if succ else "max_iter", + "requested_unlocked_groups": requested_unlocked_groups, + "unlocked_groups": list(unlocked_move_group_ids), + "unavailable_groups": unavailable_groups, + "iterations": iterations, + "final_error_norm": err_norm, + "final_pos_error_norm": pos_err_norm, + "final_rot_error_norm": rot_err_norm, + "singular_iterations": singular_iterations, + "min_singular_value": min_singular_value if np.isfinite(min_singular_value) else None, + "min_jjt_det": min_jjt_det if np.isfinite(min_jjt_det) else None, + "clamp_event_count": len(clamp_events), + "clamp_events": clamp_events[:20], + "qpos": qpos if succ else None, + "final_qpos": qpos, + "final_pose": final_pose, + } diff --git a/molmo_spaces/policy/solvers/object_manipulation/pick_and_place_planner_policy.py b/molmo_spaces/policy/solvers/object_manipulation/pick_and_place_planner_policy.py index 355ebcf0..bf088960 100644 --- a/molmo_spaces/policy/solvers/object_manipulation/pick_and_place_planner_policy.py +++ b/molmo_spaces/policy/solvers/object_manipulation/pick_and_place_planner_policy.py @@ -1,6 +1,9 @@ import logging +from pathlib import Path +import mujoco import numpy as np +from scipy.spatial.transform import Rotation as R from molmo_spaces.configs.task_configs import PickAndPlaceTaskConfig from molmo_spaces.env.data_views import MlSpacesObject @@ -14,12 +17,35 @@ TCPMoveSegment, TCPMoveSequence, ) -from molmo_spaces.utils.grasp_sample import compute_grasp_pose +from molmo_spaces.utils.grasp_sample import ( + compute_grasp_pose, + get_all_grasp_poses, + get_noncolliding_grasp_mask, +) from molmo_spaces.utils.mj_model_and_data_utils import body_aabb +from molmo_spaces.utils.pose import pose_mat_to_7d log = logging.getLogger(__name__) +class HoldCurrentControlAction(ActionPrimitive): + """Hold the current commanded controls for a fixed duration.""" + + def __init__(self, robot_view, duration: float) -> None: + super().__init__(robot_view, duration) + self._action = None + + def execute(self) -> bool: + if self.start_time is None: + self.start_time = self.robot_view.mj_data.time + return self.elapsed_time() >= self.duration + + def get_current_action(self) -> dict: + if self._action is None: + self._action = self.robot_view.get_ctrl_dict() + return self._action + + class PickAndPlacePlannerPolicy(BaseObjectManipulationPlannerPolicy): def _compute_trajectory(self) -> list[ActionPrimitive]: robot_view = self.task.env.current_robot.robot_view @@ -112,7 +138,7 @@ def _compute_trajectory(self) -> list[ActionPrimitive]: ) ], ), - NoopAction(robot_view, 2.0), + HoldCurrentControlAction(robot_view, 2.0), ] def _get_grasp_poses( @@ -279,9 +305,7 @@ def _compute_target_poses(self) -> dict[str, np.ndarray]: target_poses["place"] = place_pose target_poses["postplace"] = postplace_pose - # debug - visualize_poses = True - if visualize_poses and self.task.viewer is not None: + if self.policy_config.debug_poses and self.task.viewer is not None: self._show_poses(np.stack(list(target_poses.values()), axis=0), style="tcp") if self.task.viewer: self.task.viewer.sync() @@ -292,6 +316,1760 @@ def _compute_target_poses(self) -> dict[str, np.ndarray]: class UnitreeG1RightArmPickAndPlacePlannerPolicy(PickAndPlacePlannerPolicy): """Pick-and-place planner variant that solves IK through the G1 right arm only.""" + def reset(self, reset_retries: bool = True): + self._g1_failure_diag_steps: list[dict] = [] + self._g1_last_target_pose: np.ndarray | None = None + self._g1_initial_tcp_pose = self._current_g1_tcp_pose() + self._g1_failure_diagnostics_logged = False + self._g1_failure_reason_override: str | None = None + self._g1_failure_phase_override: str | None = None + self._g1_bypass_feasible_ik = False + self._g1_unfiltered_grasp_fallback_used = False + self._g1_unfiltered_grasp_fallback_reason: str | None = None + self._g1_selected_grasp_debug: dict | None = None + self._g1_grasp_selector_debug: dict | None = None + self._g1_last_ik_debug: dict | None = None + self._g1_failure_hold_action: NoopAction | None = None + self._g1_terminal_failure_pending = False + self._g1_grip_diag_samples: list[dict] = [] + self._g1_grip_diag_last_log_time = -np.inf + self._g1_grip_diag_last_key: tuple | None = None + super().reset(reset_retries=reset_retries) + if self._g1_failure_reason_override: + self._populate_failed_trajectory_target_poses() + + def _record_g1_selected_grasp_debug(self, grasp_debug: dict) -> None: + self._g1_selected_grasp_debug = grasp_debug + + def _populate_failed_trajectory_target_poses(self) -> None: + tcp_pose = self._current_g1_tcp_pose() + self.target_poses.update( + { + "pregrasp": tcp_pose, + "grasp": tcp_pose, + "lift": tcp_pose, + "preplace": tcp_pose, + "place": tcp_pose, + "postplace": tcp_pose, + } + ) + + def _compute_trajectory(self) -> list[ActionPrimitive]: + try: + if getattr(self.policy_config, "g1_pick_lift_only", False): + return self._compute_g1_pick_lift_only_trajectory() + return self._compute_g1_pick_and_place_trajectory() + except ValueError as exc: + reason = self._classify_trajectory_build_failure(exc) + if reason is None: + raise + self._g1_failure_reason_override = reason + self._g1_failure_phase_override = "grasp_selection" + log.info("[G1_DIAG] trajectory_build_failure=%s error=%r", reason, str(exc)) + return [ + NoopAction( + self.robot_view, + getattr(self.policy_config, "diagnostic_failure_hold_duration_s", 2.0), + ) + ] + + def _g1_postgrasp_hold_action(self, robot_view) -> HoldCurrentControlAction | None: + hold_duration = float(getattr(self.policy_config, "g1_postgrasp_hold_duration", 0.0)) + if hold_duration <= 0.0: + return None + return HoldCurrentControlAction(robot_view, hold_duration) + + def _compute_g1_pick_and_place_trajectory(self) -> list[ActionPrimitive]: + robot_view = self.task.env.current_robot.robot_view + target_poses = self._compute_target_poses() + held_object_speed = float( + getattr(self.policy_config, "g1_held_object_speed", self.policy_config.speed_slow) + ) + + gripper_mg_id = robot_view.get_gripper_movegroup_ids()[0] + start_ee_pose = robot_view.get_move_group(gripper_mg_id).leaf_frame_to_world + + postgrasp_hold = self._g1_postgrasp_hold_action(robot_view) + trajectory: list[ActionPrimitive] = [ + GripperAction(robot_view, True, 0.0), + TCPMoveSequence( + robot_view, + self._tcp_to_jp_fn, + self.policy_config.move_settle_time, + gripper_empty_threshold=self.policy_config.gripper_empty_threshold, + tcp_pos_err_threshold=self.policy_config.tcp_pos_err_threshold, + tcp_rot_err_threshold=self.policy_config.tcp_rot_err_threshold, + move_segments=[ + TCPMoveSegment( + name="pregrasp", + start_pose=start_ee_pose, + end_pose=target_poses["pregrasp"], + speed=self.policy_config.speed_fast, + ), + TCPMoveSegment( + name="grasp", + start_pose=target_poses["pregrasp"], + end_pose=target_poses["grasp"], + speed=self.policy_config.speed_slow, + ), + ], + ), + GripperAction(robot_view, False, self.policy_config.gripper_close_duration), + ] + if postgrasp_hold is not None: + trajectory.append(postgrasp_hold) + trajectory.extend( + [ + TCPMoveSequence( + robot_view, + self._tcp_to_jp_fn, + self.policy_config.move_settle_time, + is_holding_object=True, + gripper_empty_threshold=self.policy_config.gripper_empty_threshold, + tcp_pos_err_threshold=self.policy_config.tcp_pos_err_threshold, + tcp_rot_err_threshold=self.policy_config.tcp_rot_err_threshold, + move_segments=[ + TCPMoveSegment( + name="lift", + start_pose=target_poses["grasp"], + end_pose=target_poses["lift"], + speed=held_object_speed, + ), + TCPMoveSegment( + name="preplace", + start_pose=target_poses["lift"], + end_pose=target_poses["preplace"], + speed=held_object_speed, + ), + TCPMoveSegment( + name="place", + start_pose=target_poses["preplace"], + end_pose=target_poses["place"], + speed=held_object_speed, + ), + ], + ), + GripperAction(robot_view, True, self.policy_config.gripper_open_duration), + TCPMoveSequence( + robot_view, + self._tcp_to_jp_fn, + self.policy_config.move_settle_time, + gripper_empty_threshold=self.policy_config.gripper_empty_threshold, + tcp_pos_err_threshold=self.policy_config.tcp_pos_err_threshold, + tcp_rot_err_threshold=self.policy_config.tcp_rot_err_threshold, + move_segments=[ + TCPMoveSegment( + name="retreat", + start_pose=target_poses["place"], + end_pose=target_poses["postplace"], + speed=self.policy_config.speed_fast, + ) + ], + ), + JointMoveSequence( + robot_view, + self.policy_config.move_settle_time, + gripper_empty_threshold=self.policy_config.gripper_empty_threshold, + move_segments=[ + JointMoveSegment( + name="go_home", + start_qpos=None, + end_qpos=self.config.robot_config.init_qpos, + duration_s=4.0, + ) + ], + ), + HoldCurrentControlAction(robot_view, 2.0), + ] + ) + return trajectory + + def _compute_g1_pick_lift_only_trajectory(self) -> list[ActionPrimitive]: + robot_view = self.task.env.current_robot.robot_view + target_poses = self._compute_target_poses() + held_object_speed = float( + getattr(self.policy_config, "g1_held_object_speed", self.policy_config.speed_slow) + ) + + gripper_mg_id = robot_view.get_gripper_movegroup_ids()[0] + start_ee_pose = robot_view.get_move_group(gripper_mg_id).leaf_frame_to_world + + postgrasp_hold = self._g1_postgrasp_hold_action(robot_view) + trajectory: list[ActionPrimitive] = [ + GripperAction(robot_view, True, 0.0), + TCPMoveSequence( + robot_view, + self._tcp_to_jp_fn, + self.policy_config.move_settle_time, + gripper_empty_threshold=self.policy_config.gripper_empty_threshold, + tcp_pos_err_threshold=self.policy_config.tcp_pos_err_threshold, + tcp_rot_err_threshold=self.policy_config.tcp_rot_err_threshold, + move_segments=[ + TCPMoveSegment( + name="pregrasp", + start_pose=start_ee_pose, + end_pose=target_poses["pregrasp"], + speed=self.policy_config.speed_fast, + ), + TCPMoveSegment( + name="grasp", + start_pose=target_poses["pregrasp"], + end_pose=target_poses["grasp"], + speed=self.policy_config.speed_slow, + ), + ], + ), + GripperAction(robot_view, False, self.policy_config.gripper_close_duration), + ] + if postgrasp_hold is not None: + trajectory.append(postgrasp_hold) + trajectory.extend( + [ + TCPMoveSequence( + robot_view, + self._tcp_to_jp_fn, + self.policy_config.move_settle_time, + is_holding_object=True, + gripper_empty_threshold=self.policy_config.gripper_empty_threshold, + tcp_pos_err_threshold=self.policy_config.tcp_pos_err_threshold, + tcp_rot_err_threshold=self.policy_config.tcp_rot_err_threshold, + move_segments=[ + TCPMoveSegment( + name="lift", + start_pose=target_poses["grasp"], + end_pose=target_poses["lift"], + speed=held_object_speed, + ) + ], + ), + HoldCurrentControlAction(robot_view, 2.0), + ] + ) + return trajectory + + def _compute_target_poses(self) -> dict[str, np.ndarray]: + if getattr(self.policy_config, "g1_online_grasp_selector", False): + return self._compute_g1_online_grasp_target_poses() + + try: + return super()._compute_target_poses() + except ValueError as exc: + if ( + str(exc) != "No feasible grasp found" + or not self.policy_config.record_unfiltered_attempt_on_no_feasible_grasp + ): + raise + + self._g1_unfiltered_grasp_fallback_used = True + self._g1_unfiltered_grasp_fallback_reason = "no_feasible_g1_grasp" + log.info( + "[G1_DIAG] no_feasible_g1_grasp; recording one unfiltered grasp attempt" + ) + self._g1_bypass_feasible_ik = True + try: + return super()._compute_target_poses() + finally: + self._g1_bypass_feasible_ik = False + + def _classify_trajectory_build_failure(self, exc: ValueError) -> str | None: + message = str(exc) + if message == "No feasible grasp found": + return "no_feasible_g1_grasp" + if message == "No G1 tabletop grasp candidate found": + return "no_g1_tabletop_grasp_candidate" + if message == "IK failed for preplace pose": + return "no_feasible_g1_preplace" + if message == "IK failed for place pose": + return "no_feasible_g1_place" + return None + + def _compute_g1_online_grasp_target_poses(self) -> dict[str, np.ndarray]: + task_config = self.config.task_config + assert isinstance(task_config, PickAndPlaceTaskConfig) + + robot_view = self.task.env.current_robot.robot_view + om = self.task.env.object_managers[self.task.env.current_batch_index] + pickup_obj: MlSpacesObject = om.get_object_by_name(task_config.pickup_obj_name) + place_receptacle: MlSpacesObject = om.get_object_by_name(task_config.place_receptacle_name) + + grasp_pose_world = self._select_g1_tabletop_grasp_pose( + pickup_obj, + place_receptacle, + robot_view, + ) + target_poses = self._build_g1_candidate_target_poses( + grasp_pose_world, + pickup_obj, + place_receptacle, + ) + + if self.policy_config.debug_poses and self.task.viewer is not None: + self._show_poses(np.stack(list(target_poses.values()), axis=0), style="tcp") + if self.task.viewer: + self.task.viewer.sync() + + return target_poses + + def _build_g1_candidate_target_poses( + self, + grasp_pose_world: np.ndarray, + pickup_obj: MlSpacesObject, + place_receptacle: MlSpacesObject, + ) -> dict[str, np.ndarray]: + grasp_pose_world = grasp_pose_world.copy() + grasp_pose_world = self._g1_apply_grasp_inward_offset( + grasp_pose_world, + pickup_obj, + ) + grasp_pose_world = self._g1_level_grasp_orientation(grasp_pose_world) + table_top_z = self._g1_tabletop_top_z() + min_grasp_z = table_top_z + float( + getattr(self.policy_config, "g1_grasp_table_clearance", 0.0) + ) + if np.isfinite(min_grasp_z): + grasp_pose_world[2, 3] = max(grasp_pose_world[2, 3], min_grasp_z) + grasp_pose_world = self._g1_center_grasp_lateral( + grasp_pose_world, + pickup_obj, + ) + + pregrasp_pose = grasp_pose_world.copy() + pregrasp_pose[:3, 3] -= self.policy_config.pregrasp_z_offset * pregrasp_pose[:3, 2] + + model = self.task.env.current_data.model + data = self.task.env.current_data + place_receptacle_aabb_center, place_receptacle_aabb_size = body_aabb( + model, + data, + place_receptacle.object_id, + ) + receptacle_top_z = place_receptacle_aabb_center[2] + place_receptacle_aabb_size[2] / 2 + pickup_obj_aabb_center, pickup_obj_aabb_size = body_aabb( + model, + data, + pickup_obj.object_id, + ) + pickup_obj_top_z = pickup_obj_aabb_center[2] + pickup_obj_aabb_size[2] / 2 + pickup_obj_bottom_z = pickup_obj_aabb_center[2] - pickup_obj_aabb_size[2] / 2 + pickup_obj_clearance_offset = max(grasp_pose_world[2, 3] - pickup_obj_bottom_z, 0.0) + min_pregrasp_z = max( + pregrasp_pose[2, 3], + grasp_pose_world[2, 3] + + float(getattr(self.policy_config, "g1_pregrasp_min_vertical_lift", 0.0)), + pickup_obj_top_z + + float(getattr(self.policy_config, "g1_pregrasp_object_clearance", 0.0)), + ) + pregrasp_pose[2, 3] = min_pregrasp_z + + lift_pose = grasp_pose_world.copy() + lift_pose[2, 3] = ( + receptacle_top_z + + pickup_obj_clearance_offset + + self.policy_config.place_z_offset + ) + lift_pose = self._g1_raise_pose_for_carried_object_bottom_clearance( + lift_pose, + grasp_pose_world, + pickup_obj_bottom_z, + receptacle_top_z, + ) + + preplace_pose = grasp_pose_world.copy() + preplace_pose[:2, 3] = place_receptacle.position[:2] + preplace_pose[2, 3] = ( + receptacle_top_z + + pickup_obj_clearance_offset + + self.policy_config.place_z_offset + ) + preplace_pose[:3, 3] += grasp_pose_world[:3, 3] - pickup_obj.position + preplace_pose = self._g1_raise_pose_for_carried_object_bottom_clearance( + preplace_pose, + grasp_pose_world, + pickup_obj_bottom_z, + receptacle_top_z, + ) + + place_pose = preplace_pose.copy() + place_pose[2, 3] = receptacle_top_z + pickup_obj_clearance_offset + + postplace_pose = place_pose.copy() + postplace_pose[:3, 3] -= self.policy_config.end_z_offset * postplace_pose[:3, 2] + + return { + "pregrasp": pregrasp_pose, + "grasp": grasp_pose_world, + "lift": lift_pose, + "preplace": preplace_pose, + "place": place_pose, + "postplace": postplace_pose, + } + + def _g1_raise_pose_for_carried_object_bottom_clearance( + self, + pose: np.ndarray, + grasp_pose_world: np.ndarray, + pickup_obj_bottom_z: float, + receptacle_top_z: float, + ) -> np.ndarray: + clearance = float(getattr(self.policy_config, "g1_place_travel_object_clearance", 0.0)) + if clearance <= 0.0: + return pose + + raised_pose = pose.copy() + carried_bottom_z = raised_pose[2, 3] + pickup_obj_bottom_z - grasp_pose_world[2, 3] + min_bottom_z = receptacle_top_z + clearance + if carried_bottom_z < min_bottom_z: + raised_pose[2, 3] += min_bottom_z - carried_bottom_z + return raised_pose + + def _g1_center_grasp_lateral( + self, + grasp_pose_world: np.ndarray, + pickup_obj: MlSpacesObject, + ) -> np.ndarray: + if not getattr(self.policy_config, "g1_center_grasp_lateral", False): + return grasp_pose_world + + pose = grasp_pose_world.copy() + object_pos = np.asarray(pickup_obj.position[:3], dtype=float) + object_in_grasp = pose[:3, :3].T @ (object_pos - pose[:3, 3]) + lateral_offset = float(object_in_grasp[1]) + lateral_offset *= float( + getattr(self.policy_config, "g1_grasp_lateral_centering_scale", 1.0) + ) + max_offset = float( + getattr(self.policy_config, "g1_grasp_lateral_centering_max_offset", np.inf) + ) + if np.isfinite(max_offset): + lateral_offset = float(np.clip(lateral_offset, -max_offset, max_offset)) + + pose[:3, 3] += lateral_offset * pose[:3, 1] + return pose + + def _g1_level_grasp_orientation(self, grasp_pose_world: np.ndarray) -> np.ndarray: + if not getattr(self.policy_config, "g1_level_grasp_orientation", False): + return grasp_pose_world + + pose = grasp_pose_world.copy() + z_axis = pose[:3, 2] + z_norm = float(np.linalg.norm(z_axis)) + if z_norm <= 1e-6: + return pose + + z_axis = z_axis / z_norm + vertical_tilt_deg = float(np.degrees(np.arccos(np.clip(abs(z_axis[2]), -1.0, 1.0)))) + max_tilt_deg = float( + getattr(self.policy_config, "g1_grasp_level_max_tilt_deg", np.inf) + ) + if np.isfinite(max_tilt_deg) and vertical_tilt_deg > max_tilt_deg: + return pose + + target_z = np.array([0.0, 0.0, 1.0 if z_axis[2] >= 0.0 else -1.0]) + x_axis = pose[:3, 0] + x_axis = x_axis - target_z * float(np.dot(x_axis, target_z)) + if np.linalg.norm(x_axis) <= 1e-6: + x_axis = pose[:3, 1] - target_z * float(np.dot(pose[:3, 1], target_z)) + if np.linalg.norm(x_axis) <= 1e-6: + x_axis = np.array([1.0, 0.0, 0.0]) + + x_axis = x_axis / np.linalg.norm(x_axis) + y_axis = np.cross(target_z, x_axis) + y_axis = y_axis / np.linalg.norm(y_axis) + pose[:3, :3] = np.column_stack([x_axis, y_axis, target_z]) + return pose + + def _g1_apply_grasp_inward_offset( + self, + grasp_pose_world: np.ndarray, + pickup_obj: MlSpacesObject, + ) -> np.ndarray: + inward_offset = float( + getattr(self.policy_config, "g1_grasp_inward_xy_offset", 0.0) + ) + if inward_offset <= 0.0: + return grasp_pose_world + + offset_pose = grasp_pose_world.copy() + object_xy = np.asarray(pickup_obj.position[:2], dtype=float) + grasp_xy = offset_pose[:2, 3] + object_delta_xy = object_xy - grasp_xy + object_dist_xy = float(np.linalg.norm(object_delta_xy)) + if object_dist_xy <= 1e-6: + return offset_pose + + offset_pose[:2, 3] += object_delta_xy / object_dist_xy * min( + inward_offset, + object_dist_xy, + ) + return offset_pose + + def _g1_tabletop_top_z(self) -> float: + task_sampler_config = getattr(getattr(self, "config", None), "task_sampler_config", None) + table_body_name = getattr(task_sampler_config, "table_body_name", None) + if not table_body_name: + return -np.inf + + model = self.task.env.current_data.model + data = self.task.env.current_data + try: + table_body_id = model.body(table_body_name).id + except KeyError: + return -np.inf + + table_center, table_size = body_aabb(model, data, table_body_id, visual_only=False) + return float(table_center[2] + table_size[2] / 2) + + def _g1_pregrasp_debug( + self, + pregrasp_pose: np.ndarray, + grasp_pose_world: np.ndarray, + pickup_obj: MlSpacesObject, + ) -> dict[str, float]: + model = self.task.env.current_data.model + data = self.task.env.current_data + pickup_obj_aabb_center, pickup_obj_aabb_size = body_aabb( + model, + data, + pickup_obj.object_id, + ) + pickup_obj_top_z = pickup_obj_aabb_center[2] + pickup_obj_aabb_size[2] / 2 + object_in_grasp = grasp_pose_world[:3, :3].T @ ( + np.asarray(pickup_obj.position[:3], dtype=float) - grasp_pose_world[:3, 3] + ) + raw_grasp_pose = getattr(self, "_g1_selected_raw_grasp_pose", None) + table_top_z = self._g1_tabletop_top_z() + return { + "object_top_z": float(pickup_obj_top_z), + "table_top_z": float(table_top_z), + "raw_grasp_z": float(raw_grasp_pose[2, 3]) if raw_grasp_pose is not None else None, + "grasp_z": float(grasp_pose_world[2, 3]), + "grasp_vertical_axis_z": float(grasp_pose_world[2, 2]), + "grasp_vertical_tilt_deg": float( + np.degrees(np.arccos(np.clip(abs(grasp_pose_world[2, 2]), -1.0, 1.0))) + ), + "object_local_y_in_grasp_m": float(object_in_grasp[1]), + "grasp_table_clearance_m": float(grasp_pose_world[2, 3] - table_top_z) + if np.isfinite(table_top_z) + else None, + "pregrasp_z": float(pregrasp_pose[2, 3]), + "pregrasp_above_grasp_m": float(pregrasp_pose[2, 3] - grasp_pose_world[2, 3]), + "pregrasp_object_clearance_m": float(pregrasp_pose[2, 3] - pickup_obj_top_z), + } + + def _copy_qpos_dict(self, qpos: dict[str, np.ndarray]) -> dict[str, np.ndarray]: + return {move_group_id: np.asarray(values).copy() for move_group_id, values in qpos.items()} + + def _g1_eval_candidate_target_poses( + self, + target_poses: dict[str, np.ndarray], + initial_qpos: dict[str, np.ndarray] | None = None, + base_pose: np.ndarray | None = None, + ) -> dict: + robot_view = self.task.env.current_robot.robot_view + kinematics = self.task.env.current_robot.kinematics + gripper_mg_id = robot_view.get_gripper_movegroup_ids()[0] + phase_order = ( + ["pregrasp", "grasp", "lift", "preplace", "place"] + if getattr(self.policy_config, "g1_grasp_require_all_pick_place_phases", True) + else ["pregrasp", "grasp", "lift"] + ) + + live_qpos = self._copy_qpos_dict(robot_view.get_qpos_dict()) + live_base_pose = robot_view.base.pose.copy() + qpos = self._copy_qpos_dict(initial_qpos or live_qpos) + if base_pose is None: + base_pose = live_base_pose + + phase_results = {} + try: + for phase in phase_order: + result = kinematics.diagnose_ik( + gripper_mg_id, + target_poses[phase], + ["right_arm"], + qpos, + base_pose, + max_iter=250, + ) + phase_results[phase] = result + if not result["success"]: + return { + "success": False, + "failed_phase": phase, + "phase_results": phase_results, + } + qpos = self._copy_qpos_dict(result["qpos"]) + + return { + "success": True, + "failed_phase": None, + "phase_results": phase_results, + "final_qpos": qpos, + } + finally: + robot_view.base.pose = live_base_pose + robot_view.set_qpos_dict(live_qpos) + + def _g1_contact_quality_from_contacts(self, contacts: list[dict]) -> dict: + robot_namespace = self.config.robot_config.robot_namespace + object_contacts = [ + contact + for contact in contacts + if contact["class"] == "robot_pickup_object" + ] + pad_contacts = [ + contact + for contact in object_contacts + if self._g1_has_right_fingertip_pad_contact(contact) + ] + non_pad_object_contacts = [ + contact + for contact in object_contacts + if not self._g1_has_right_fingertip_pad_contact(contact) + ] + finger_link_object_contacts = [ + contact + for contact in non_pad_object_contacts + if any( + body_name.startswith(f"{robot_namespace}right_dex1_finger_link") + for _geom_name, body_name in self._g1_robot_contact_endpoints(contact) + ) + ] + other_robot_object_contacts = [ + contact + for contact in non_pad_object_contacts + if contact not in finger_link_object_contacts + ] + table_contacts = [ + contact + for contact in contacts + if contact["class"] == "robot_table" + ] + bin_contacts = [ + contact + for contact in contacts + if contact["class"] == "robot_bin" + ] + pad_geom_names = sorted( + { + geom_name + for contact in pad_contacts + for geom_name, _body_name in self._g1_robot_contact_endpoints(contact) + if "right_dex1_fingertip_pad" in geom_name + } + ) + return { + "object_contact_count": len(object_contacts), + "pad_contact_count": len(pad_contacts), + "pad_geom_count": len(pad_geom_names), + "pad_geom_names": pad_geom_names, + "non_pad_object_contacts": non_pad_object_contacts, + "finger_link_object_contacts": finger_link_object_contacts, + "other_robot_object_contacts": other_robot_object_contacts, + "table_contacts": table_contacts, + "bin_contacts": bin_contacts, + } + + def _g1_grip_diagnostics(self, contacts: list[dict] | None = None) -> dict: + robot_view = self.task.env.current_robot.robot_view + gripper_mg_id = robot_view.get_gripper_movegroup_ids()[0] + gripper = robot_view.get_gripper(gripper_mg_id) + data = self.task.env.current_data + model = self.task.env.current_model + actuator_ids = np.asarray(getattr(gripper, "_actuator_ids", []), dtype=int) + actuator_force = ( + data.actuator_force[actuator_ids].copy() + if actuator_ids.size and hasattr(data, "actuator_force") + else np.array([]) + ) + actuator_range = ( + model.actuator_forcerange[actuator_ids].copy() + if actuator_ids.size + else np.empty((0, 2)) + ) + force_limit = ( + np.maximum(np.abs(actuator_range[:, 0]), np.abs(actuator_range[:, 1])) + if actuator_range.size + else np.array([]) + ) + force_fraction = ( + np.divide( + np.abs(actuator_force), + np.maximum(force_limit, 1e-9), + out=np.zeros_like(actuator_force, dtype=float), + where=force_limit > 0.0, + ) + if actuator_force.size + else np.array([]) + ) + contacts = self._collect_contact_diagnostics() if contacts is None else contacts + contact_quality = self._g1_contact_quality_from_contacts(contacts) + ctrl = gripper.ctrl.copy() + joint_pos = gripper.joint_pos.copy() + return { + "sim_time": float(data.time), + "phase": self.get_phase(), + "ctrl": np.round(ctrl, 5).tolist(), + "joint_pos": np.round(joint_pos, 5).tolist(), + "ctrl_minus_joint_pos": np.round(ctrl - joint_pos, 5).tolist(), + "inter_finger_dist": float(gripper.inter_finger_dist), + "actuator_force": np.round(actuator_force, 3).tolist(), + "actuator_force_fraction": np.round(force_fraction, 3).tolist(), + "pad_geom_count": int(contact_quality["pad_geom_count"]), + "pad_geom_names": contact_quality["pad_geom_names"], + "pad_contact_count": int(contact_quality["pad_contact_count"]), + "finger_link_contact_count": len(contact_quality["finger_link_object_contacts"]), + "other_robot_object_contact_count": len(contact_quality["other_robot_object_contacts"]), + "table_contact_count": len(contact_quality["table_contacts"]), + "bin_contact_count": len(contact_quality["bin_contacts"]), + } + + def _record_g1_grip_diagnostic_sample(self) -> None: + if not self._failure_diagnostics_enabled(): + return + phase = self.get_phase() + if phase not in {"gripper-close", "lift", "preplace", "place"}: + return + + sample = self._g1_grip_diagnostics() + self._g1_grip_diag_samples.append(sample) + self._g1_grip_diag_samples = self._g1_grip_diag_samples[-240:] + + log_key = ( + sample["phase"], + sample["pad_geom_count"], + sample["finger_link_contact_count"], + sample["other_robot_object_contact_count"], + sample["bin_contact_count"], + ) + log_interval_s = 0.5 + should_log = ( + log_key != self._g1_grip_diag_last_key + or sample["sim_time"] - self._g1_grip_diag_last_log_time >= log_interval_s + ) + if not should_log: + return + + self._g1_grip_diag_last_key = log_key + self._g1_grip_diag_last_log_time = sample["sim_time"] + log.info( + "[G1_GRIP] phase=%s ctrl=%s q=%s ctrl_minus_q=%s dist=%.5f " + "force=%s force_frac=%s pads=%s finger_links=%s other=%s bin=%s", + sample["phase"], + sample["ctrl"], + sample["joint_pos"], + sample["ctrl_minus_joint_pos"], + sample["inter_finger_dist"], + sample["actuator_force"], + sample["actuator_force_fraction"], + sample["pad_geom_names"], + sample["finger_link_contact_count"], + sample["other_robot_object_contact_count"], + sample["bin_contact_count"], + ) + + def _g1_open_grasp_contact_quality( + self, + qpos: dict[str, np.ndarray], + base_pose: np.ndarray, + ) -> dict: + robot_view = self.task.env.current_robot.robot_view + live_qpos = self._copy_qpos_dict(robot_view.get_qpos_dict()) + live_base_pose = robot_view.base.pose.copy() + model = self.task.env.current_model + data = self.task.env.current_data + + try: + robot_view.base.pose = base_pose + robot_view.set_qpos_dict(qpos) + mujoco.mj_forward(model, data) + contacts = self._collect_contact_diagnostics() + finally: + robot_view.base.pose = live_base_pose + robot_view.set_qpos_dict(live_qpos) + mujoco.mj_forward(model, data) + + quality = self._g1_contact_quality_from_contacts(contacts) + quality.update( + { + "closed_probe": False, + "object_shift_m": 0.0, + "failure_reason": self._g1_grasp_contact_quality_failure(quality), + } + ) + return quality + + def _g1_mujoco_state_snapshot(self) -> dict: + data = self.task.env.current_data + return { + "time": float(data.time), + "qpos": data.qpos.copy(), + "qvel": data.qvel.copy(), + "qacc": data.qacc.copy(), + "qacc_warmstart": data.qacc_warmstart.copy(), + "ctrl": data.ctrl.copy(), + "act": data.act.copy(), + "mocap_pos": data.mocap_pos.copy(), + "mocap_quat": data.mocap_quat.copy(), + } + + def _g1_restore_mujoco_state(self, snapshot: dict) -> None: + model = self.task.env.current_model + data = self.task.env.current_data + data.time = snapshot["time"] + data.qpos[:] = snapshot["qpos"] + data.qvel[:] = snapshot["qvel"] + data.qacc[:] = snapshot["qacc"] + data.qacc_warmstart[:] = snapshot["qacc_warmstart"] + data.ctrl[:] = snapshot["ctrl"] + data.act[:] = snapshot["act"] + data.mocap_pos[:] = snapshot["mocap_pos"] + data.mocap_quat[:] = snapshot["mocap_quat"] + mujoco.mj_forward(model, data) + + def _g1_closed_grasp_contact_quality( + self, + qpos: dict[str, np.ndarray], + base_pose: np.ndarray, + pickup_obj: MlSpacesObject, + ) -> dict: + robot_view = self.task.env.current_robot.robot_view + model = self.task.env.current_model + data = self.task.env.current_data + snapshot = self._g1_mujoco_state_snapshot() + settle_steps = max( + int(getattr(self.policy_config, "g1_closed_grasp_settle_steps", 25)), + 0, + ) + + try: + robot_view.base.pose = base_pose + robot_view.set_qpos_dict(qpos) + gripper_mg_id = robot_view.get_gripper_movegroup_ids()[0] + gripper = robot_view.get_gripper(gripper_mg_id) + gripper.set_gripper_ctrl_open(False) + gripper.joint_pos = gripper.ctrl.copy() + data.qvel[:] = 0.0 + data.qacc[:] = 0.0 + data.qacc_warmstart[:] = 0.0 + mujoco.mj_forward(model, data) + object_start_pos = pickup_obj.position.copy() + for _ in range(settle_steps): + mujoco.mj_step(model, data) + mujoco.mj_forward(model, data) + object_shift_m = float(np.linalg.norm(pickup_obj.position - object_start_pos)) + contacts = self._collect_contact_diagnostics() + quality = self._g1_contact_quality_from_contacts(contacts) + quality.update( + { + "closed_probe": True, + "object_shift_m": object_shift_m, + "settle_steps": settle_steps, + "failure_reason": self._g1_closed_grasp_contact_quality_failure( + quality, + object_shift_m, + ), + } + ) + return quality + finally: + self._g1_restore_mujoco_state(snapshot) + + def _g1_candidate_grasp_contact_quality( + self, + qpos: dict[str, np.ndarray], + base_pose: np.ndarray, + pickup_obj: MlSpacesObject, + ) -> dict: + open_quality = self._g1_open_grasp_contact_quality(qpos, base_pose) + if open_quality["failure_reason"]: + return open_quality + if getattr(self.policy_config, "g1_closed_grasp_quality_enabled", True): + return self._g1_closed_grasp_contact_quality(qpos, base_pose, pickup_obj) + return open_quality + + def _g1_grasp_contact_quality_failure(self, quality: dict) -> str | None: + if ( + getattr(self.policy_config, "g1_reject_open_grasp_object_contact", True) + and quality["object_contact_count"] > 0 + ): + return "open_grasp_object_contact" + if ( + getattr(self.policy_config, "g1_reject_grasp_table_contact", True) + and quality["table_contacts"] + ): + return "table_contact" + if ( + getattr(self.policy_config, "g1_reject_non_fingertip_grasp_object_contact", True) + and quality["non_pad_object_contacts"] + ): + return "non_fingertip_object_contact" + if ( + getattr(self.policy_config, "g1_require_fingertip_pad_grasp_contact", True) + and quality["pad_contact_count"] == 0 + ): + return "missing_fingertip_pad_contact" + return None + + def _g1_grasp_quality_rejection_key(self, quality: dict) -> str: + if quality["closed_probe"]: + return "closed_grasp_quality" + if quality["failure_reason"] == "open_grasp_object_contact": + return "open_grasp_contact" + return "grasp_contact" + + def _g1_closed_grasp_contact_quality_failure( + self, + quality: dict, + object_shift_m: float, + ) -> str | None: + if ( + getattr(self.policy_config, "g1_reject_grasp_table_contact", True) + and quality["table_contacts"] + ): + return "table_contact" + if quality["bin_contacts"]: + return "bin_contact" + if quality["other_robot_object_contacts"]: + return "non_fingertip_object_contact" + min_pad_geom_count = int( + getattr(self.policy_config, "g1_closed_grasp_min_pad_geom_count", 2) + ) + if quality["pad_geom_count"] < min_pad_geom_count: + return "missing_closed_fingertip_pad_contact" + max_object_shift_m = float( + getattr(self.policy_config, "g1_closed_grasp_max_object_shift_m", 0.02) + ) + if object_shift_m > max_object_shift_m: + return "closed_grasp_object_shift" + return None + + def _g1_grasp_contact_score_penalty(self, quality: dict | None) -> float: + if not quality: + return 0.0 + if quality.get("closed_probe"): + min_pad_geom_count = int( + getattr(self.policy_config, "g1_closed_grasp_min_pad_geom_count", 2) + ) + missing_pad_count = max(min_pad_geom_count - quality["pad_geom_count"], 0) + return missing_pad_count * float( + getattr(self.policy_config, "g1_closed_grasp_penalty_per_missing_pad", 1.0) + ) + if quality["pad_geom_count"] >= 2: + return 0.0 + return float(getattr(self.policy_config, "g1_grasp_single_pad_contact_penalty", 0.0)) + + def _g1_closed_grasp_quality_debug(self, grasp_idx: int, quality: dict) -> dict: + return { + "grasp_idx": int(grasp_idx), + "failure_reason": quality["failure_reason"], + "pad_geom_count": int(quality["pad_geom_count"]), + "pad_geom_names": quality["pad_geom_names"], + "object_shift_m": float(quality["object_shift_m"]), + "finger_link_contact_count": len(quality["finger_link_object_contacts"]), + "other_robot_object_contact_count": len(quality["other_robot_object_contacts"]), + "table_contact_count": len(quality["table_contacts"]), + "bin_contact_count": len(quality["bin_contacts"]), + } + + def _g1_robot_contact_endpoints(self, contact_info: dict) -> list[tuple[str, str]]: + robot_namespace = self.config.robot_config.robot_namespace + robot_endpoints = [] + if str(contact_info["root1_name"]).startswith(robot_namespace): + robot_endpoints.append( + (str(contact_info["geom1_name"]), str(contact_info["body1_name"])) + ) + if str(contact_info["root2_name"]).startswith(robot_namespace): + robot_endpoints.append( + (str(contact_info["geom2_name"]), str(contact_info["body2_name"])) + ) + return robot_endpoints + + def _g1_is_allowed_grasp_object_contact(self, contact_info: dict) -> bool: + return self._g1_has_right_fingertip_pad_contact(contact_info) + + def _g1_has_right_fingertip_pad_contact(self, contact_info: dict) -> bool: + return any( + "right_dex1_fingertip_pad" in geom_name + for geom_name, _body_name in self._g1_robot_contact_endpoints(contact_info) + ) + + def _g1_joint_score_terms( + self, + grasp_pose_world: np.ndarray, + phase_results: dict[str, dict], + initial_qpos: dict[str, np.ndarray], + ) -> dict[str, float]: + right_arm_group = self.robot_view.get_move_group("right_arm") + limits = right_arm_group.joint_pos_limits + ranges = np.maximum(limits[:, 1] - limits[:, 0], 1e-6) + initial_right_arm = np.asarray(initial_qpos["right_arm"]) + + min_normalized_margin = 0.5 + joint_motion = 0.0 + n_solutions = 0 + for result in phase_results.values(): + qpos = result.get("qpos") + if not qpos: + continue + right_arm_qpos = np.asarray(qpos["right_arm"]) + normalized_margin = np.minimum( + right_arm_qpos - limits[:, 0], + limits[:, 1] - right_arm_qpos, + ) / ranges + min_normalized_margin = min( + min_normalized_margin, + float(np.min(normalized_margin)), + ) + joint_motion += float( + np.linalg.norm(right_arm_qpos - initial_right_arm) + / np.sqrt(max(len(initial_right_arm), 1)) + ) + n_solutions += 1 + + if n_solutions: + joint_motion /= n_solutions + joint_margin_cost = 1.0 - np.clip(min_normalized_margin / 0.5, 0.0, 1.0) + topdown_cost = float(np.clip(grasp_pose_world[2, 2] + 1.0, 0.0, 2.0)) + return { + "joint_margin_cost": float(joint_margin_cost), + "joint_motion_cost": float(joint_motion), + "topdown_cost": topdown_cost, + "min_normalized_joint_margin": float(min_normalized_margin), + } + + def _select_g1_tabletop_grasp_pose( + self, + pickup_obj: MlSpacesObject, + place_receptacle: MlSpacesObject, + robot_view, + ) -> np.ndarray: + model = self.task.env.current_model + data = self.task.env.current_data + grasp_poses_world, _, object_pose = get_all_grasp_poses(self, pickup_obj) + original_grasp_count = len(grasp_poses_world) // 2 + + tcp_pose_world = self._current_g1_tcp_pose() + dist_tcp = np.linalg.inv(tcp_pose_world) @ grasp_poses_world + dists_tcp_p = np.linalg.norm(dist_tcp[:, :3, 3], axis=1) + dists_tcp_o = R.from_matrix(dist_tcp[:, :3, :3]).magnitude() * 180 / np.pi + dists_up = grasp_poses_world[:, 2, 2] + dists_com = np.linalg.norm( + (np.linalg.inv(object_pose) @ grasp_poses_world)[:, :3, 3], + axis=1, + ) + base_cost = ( + self.policy_config.grasp_pos_cost_weight * dists_tcp_p + + self.policy_config.grasp_rot_cost_weight * dists_tcp_o + + self.policy_config.grasp_vertical_cost_weight * dists_up + + self.policy_config.grasp_com_dist_cost_weight * dists_com + ) + + eligible_mask = np.ones(len(grasp_poses_world), dtype=bool) + if getattr(self.policy_config, "g1_ignore_flipped_grasps", True): + eligible_mask &= np.arange(len(grasp_poses_world)) < original_grasp_count + + max_tcp_rot_deg = float( + getattr(self.policy_config, "g1_grasp_max_tcp_rot_deg", np.inf) + ) + tcp_rotation_mask = np.zeros(len(grasp_poses_world), dtype=bool) + if np.isfinite(max_tcp_rot_deg): + tcp_eligible_mask = eligible_mask.copy() + eligible_mask &= dists_tcp_o <= max_tcp_rot_deg + tcp_rotation_mask = tcp_eligible_mask & ~eligible_mask + min_vertical_axis_z = float( + getattr(self.policy_config, "g1_grasp_min_vertical_axis_z", -np.inf) + ) + vertical_eligible_mask = eligible_mask.copy() + eligible_mask &= dists_up >= min_vertical_axis_z + vertical_axis_mask = vertical_eligible_mask & ~eligible_mask + eligible_ids = np.where(eligible_mask)[0] + + candidate_limit = min( + int(getattr(self.policy_config, "g1_grasp_candidate_limit", 256)), + len(eligible_ids), + ) + candidate_ids = eligible_ids[np.argsort(base_cost[eligible_ids], kind="stable")][ + :candidate_limit + ] + + if self.policy_config.filter_colliding_grasps: + noncolliding_mask = get_noncolliding_grasp_mask( + model, + data, + grasp_poses_world[candidate_ids], + self.policy_config.grasp_collision_batch_size, + ) + else: + noncolliding_mask = np.ones(len(candidate_ids), dtype=bool) + + collision_free_ids = candidate_ids[noncolliding_mask] + ik_eval_limit = min( + int(getattr(self.policy_config, "g1_grasp_ik_eval_limit", 128)), + len(collision_free_ids), + ) + eval_ids = collision_free_ids[:ik_eval_limit] + rejection_counts = { + "flipped": ( + int(len(grasp_poses_world) - original_grasp_count) + if getattr(self.policy_config, "g1_ignore_flipped_grasps", True) + else 0 + ), + "tcp_rotation": int(np.count_nonzero(tcp_rotation_mask)), + "vertical_axis": int(np.count_nonzero(vertical_axis_mask)), + "collision": int(len(candidate_ids) - len(collision_free_ids)), + "pregrasp": 0, + "grasp": 0, + "lift": 0, + "preplace": 0, + "place": 0, + "grasp_contact": 0, + "open_grasp_contact": 0, + "closed_grasp_quality": 0, + } + initial_qpos = self._copy_qpos_dict(robot_view.get_qpos_dict()) + base_pose = robot_view.base.pose.copy() + feasible_candidates = [] + partial_candidates = [] + phase_progress_rank = { + "pregrasp": 0, + "grasp": 1, + "lift": 2, + "preplace": 3, + "place": 4, + } + closed_grasp_quality_failures = [] + + for grasp_idx in eval_ids: + grasp_pose_world = grasp_poses_world[grasp_idx] + target_poses = self._build_g1_candidate_target_poses( + grasp_pose_world, + pickup_obj, + place_receptacle, + ) + eval_result = self._g1_eval_candidate_target_poses( + target_poses, + initial_qpos, + base_pose, + ) + if not eval_result["success"]: + failed_phase = eval_result["failed_phase"] + rejection_counts[failed_phase] = rejection_counts.get(failed_phase, 0) + 1 + if phase_progress_rank.get(failed_phase, 0) > phase_progress_rank["grasp"]: + grasp_quality = None + grasp_result = eval_result["phase_results"].get("grasp") + if ( + grasp_result + and grasp_result.get("success") + and grasp_result.get("qpos") + ): + grasp_quality = self._g1_candidate_grasp_contact_quality( + grasp_result["qpos"], + base_pose, + pickup_obj, + ) + if grasp_quality["failure_reason"]: + rejection_key = self._g1_grasp_quality_rejection_key(grasp_quality) + rejection_counts[rejection_key] += 1 + reason_key = f"{rejection_key}:{grasp_quality['failure_reason']}" + rejection_counts[reason_key] = rejection_counts.get(reason_key, 0) + 1 + if grasp_quality["closed_probe"]: + closed_grasp_quality_failures.append( + self._g1_closed_grasp_quality_debug( + grasp_idx, + grasp_quality, + ) + ) + continue + g1_terms = self._g1_joint_score_terms( + grasp_pose_world, + eval_result["phase_results"], + initial_qpos, + ) + progress_penalty = 10.0 - phase_progress_rank.get(failed_phase, 0) + total_score = ( + progress_penalty + + float(base_cost[grasp_idx]) + + self.policy_config.g1_grasp_joint_margin_weight + * g1_terms["joint_margin_cost"] + + self.policy_config.g1_grasp_joint_motion_weight + * g1_terms["joint_motion_cost"] + + self.policy_config.g1_grasp_topdown_weight * g1_terms["topdown_cost"] + + self._g1_grasp_contact_score_penalty(grasp_quality) + ) + partial_candidates.append( + { + "grasp_idx": int(grasp_idx), + "grasp_pose_world": grasp_pose_world.copy(), + "target_poses": target_poses, + "score": total_score, + "failed_phase": failed_phase, + "score_terms": { + "base_cost": float(base_cost[grasp_idx]), + "tcp_pos_cost_m": float(dists_tcp_p[grasp_idx]), + "tcp_rot_cost_deg": float(dists_tcp_o[grasp_idx]), + "vertical_axis_z": float(dists_up[grasp_idx]), + "com_dist_cost_m": float(dists_com[grasp_idx]), + "progress_penalty": progress_penalty, + "is_flipped": bool(grasp_idx >= original_grasp_count), + "grasp_pad_contact_count": ( + int(grasp_quality["pad_contact_count"]) + if grasp_quality + else 0 + ), + "grasp_pad_geom_count": ( + int(grasp_quality["pad_geom_count"]) + if grasp_quality + else 0 + ), + "grasp_contact_penalty": self._g1_grasp_contact_score_penalty( + grasp_quality + ), + "closed_grasp_object_shift_m": ( + float(grasp_quality["object_shift_m"]) + if grasp_quality + else 0.0 + ), + "closed_grasp_failure_reason": ( + grasp_quality["failure_reason"] + if grasp_quality + else None + ), + **g1_terms, + }, + "pregrasp_debug": self._g1_pregrasp_debug( + target_poses["pregrasp"], + target_poses["grasp"], + pickup_obj, + ), + } + ) + continue + + grasp_quality = self._g1_candidate_grasp_contact_quality( + eval_result["phase_results"]["grasp"]["qpos"], + base_pose, + pickup_obj, + ) + if grasp_quality["failure_reason"]: + rejection_key = self._g1_grasp_quality_rejection_key(grasp_quality) + rejection_counts[rejection_key] += 1 + reason_key = f"{rejection_key}:{grasp_quality['failure_reason']}" + rejection_counts[reason_key] = rejection_counts.get(reason_key, 0) + 1 + if grasp_quality["closed_probe"]: + closed_grasp_quality_failures.append( + self._g1_closed_grasp_quality_debug( + grasp_idx, + grasp_quality, + ) + ) + continue + + g1_terms = self._g1_joint_score_terms( + grasp_pose_world, + eval_result["phase_results"], + initial_qpos, + ) + total_score = ( + float(base_cost[grasp_idx]) + + self.policy_config.g1_grasp_joint_margin_weight + * g1_terms["joint_margin_cost"] + + self.policy_config.g1_grasp_joint_motion_weight + * g1_terms["joint_motion_cost"] + + self.policy_config.g1_grasp_topdown_weight * g1_terms["topdown_cost"] + + self._g1_grasp_contact_score_penalty(grasp_quality) + ) + feasible_candidates.append( + { + "grasp_idx": int(grasp_idx), + "grasp_pose_world": grasp_pose_world.copy(), + "target_poses": target_poses, + "score": total_score, + "score_terms": { + "base_cost": float(base_cost[grasp_idx]), + "tcp_pos_cost_m": float(dists_tcp_p[grasp_idx]), + "tcp_rot_cost_deg": float(dists_tcp_o[grasp_idx]), + "vertical_axis_z": float(dists_up[grasp_idx]), + "com_dist_cost_m": float(dists_com[grasp_idx]), + "is_flipped": bool(grasp_idx >= original_grasp_count), + "grasp_pad_contact_count": int(grasp_quality["pad_contact_count"]), + "grasp_pad_geom_count": int(grasp_quality["pad_geom_count"]), + "grasp_contact_penalty": self._g1_grasp_contact_score_penalty( + grasp_quality + ), + "closed_grasp_object_shift_m": float( + grasp_quality["object_shift_m"] + ), + "closed_grasp_failure_reason": grasp_quality["failure_reason"], + **g1_terms, + }, + "pregrasp_debug": self._g1_pregrasp_debug( + target_poses["pregrasp"], + target_poses["grasp"], + pickup_obj, + ), + } + ) + + selector_summary = { + "grasp_count": int(len(grasp_poses_world)), + "eligible_count": int(len(eligible_ids)), + "candidate_count": int(candidate_limit), + "collision_free_count": int(len(collision_free_ids)), + "ik_tested_count": int(len(eval_ids)), + "ik_feasible_count": int(len(feasible_candidates)), + "partial_candidate_count": int(len(partial_candidates)), + "rejection_counts": rejection_counts, + "closed_grasp_quality_failures": closed_grasp_quality_failures[:8], + } + self._g1_grasp_selector_debug = selector_summary + + if not feasible_candidates: + if ( + partial_candidates + and getattr( + self.policy_config, + "g1_record_partial_attempt_on_no_full_grasp_candidate", + True, + ) + ): + selected = min(partial_candidates, key=lambda candidate: candidate["score"]) + log.info( + "[G1_DIAG] reason=no_full_g1_tabletop_grasp_candidate_using_partial " + "failed_phase=%s summary=%s", + selected["failed_phase"], + selector_summary, + ) + return self._finalize_g1_selected_grasp( + selected, + selector_summary, + grasp_poses_world, + candidate_ids, + dists_tcp_p, + dists_tcp_o, + partial=True, + ) + log.info("[G1_DIAG] reason=no_g1_tabletop_grasp_candidate summary=%s", selector_summary) + raise ValueError("No G1 tabletop grasp candidate found") + + selected = min(feasible_candidates, key=lambda candidate: candidate["score"]) + return self._finalize_g1_selected_grasp( + selected, + selector_summary, + grasp_poses_world, + candidate_ids, + dists_tcp_p, + dists_tcp_o, + partial=False, + ) + + def _finalize_g1_selected_grasp( + self, + selected: dict, + selector_summary: dict, + grasp_poses_world: np.ndarray, + candidate_ids: np.ndarray, + dists_tcp_p: np.ndarray, + dists_tcp_o: np.ndarray, + partial: bool, + ) -> np.ndarray: + original_grasp_count = len(grasp_poses_world) // 2 + original_grasp_idx = selected["grasp_idx"] % original_grasp_count + selected_is_flipped = selected["grasp_idx"] >= original_grasp_count + selected_summary = { + **selector_summary, + "selected_grasp_idx": selected["grasp_idx"], + "selected_original_grasp_idx": int(original_grasp_idx), + "selected_is_flipped": bool(selected_is_flipped), + "selected_score": float(selected["score"]), + "selected_score_terms": selected["score_terms"], + "selected_pregrasp": selected.get("pregrasp_debug"), + "selected_partial_candidate": partial, + "selected_partial_failed_phase": selected.get("failed_phase"), + } + self._g1_grasp_selector_debug = selected_summary + self._record_g1_selected_grasp_debug( + { + "grasp_idx": selected["grasp_idx"], + "original_grasp_idx": int(original_grasp_idx), + "is_flipped": bool(selected_is_flipped), + "grasp_pose_world": selected["grasp_pose_world"].copy(), + "top_candidate_ids": candidate_ids[ + : getattr(self.policy_config, "g1_ik_debug_top_k_grasps", 5) + ] + .astype(int) + .tolist(), + "top_candidate_poses": grasp_poses_world[ + candidate_ids[ + : getattr(self.policy_config, "g1_ik_debug_top_k_grasps", 5) + ] + ].copy(), + "selected_pos_cost_m": float(dists_tcp_p[selected["grasp_idx"]]), + "selected_rot_cost_deg": float(dists_tcp_o[selected["grasp_idx"]]), + "selected_pregrasp": selected.get("pregrasp_debug"), + "selector_summary": selected_summary, + } + ) + log.info("[G1_GRASP] selector_summary=%s", selected_summary) + return selected["grasp_pose_world"] + + def _current_g1_tcp_pose(self) -> np.ndarray: + return ( + self.robot_view.get_move_group(self.robot_view.get_gripper_movegroup_ids()[0]) + .leaf_frame_to_world.copy() + ) + + def _failure_diagnostics_enabled(self) -> bool: + return bool(getattr(self.policy_config, "enable_failure_diagnostics", False)) + + def _g1_ik_debug_enabled(self) -> bool: + return bool(getattr(self.policy_config, "g1_ik_debug", False)) + + def _names_for_contact(self, geom_id: int) -> dict[str, str | int]: + model = self.task.env.current_model + body_id = model.geom_bodyid[geom_id] + root_body_id = model.body_rootid[body_id] + return { + "geom_id": int(geom_id), + "geom_name": model.geom(geom_id).name, + "body_id": int(body_id), + "body_name": model.body(body_id).name, + "root_body_id": int(root_body_id), + "root_body_name": model.body(root_body_id).name, + } + + def _contact_has_name(self, contact_info: dict, needle: str) -> bool: + if not needle: + return False + needle = needle.lower() + return any( + needle in str(contact_info[key]).lower() + for key in ( + "geom1_name", + "geom2_name", + "body1_name", + "body2_name", + "root1_name", + "root2_name", + ) + ) + + def _contact_has_exact_name(self, contact_info: dict, needle: str) -> bool: + if not needle: + return False + return any( + needle == str(contact_info[key]) + for key in ( + "geom1_name", + "geom2_name", + "body1_name", + "body2_name", + "root1_name", + "root2_name", + ) + ) + + def _classify_contact(self, contact_info: dict) -> str: + robot_namespace = self.config.robot_config.robot_namespace + robot1 = str(contact_info["root1_name"]).startswith(robot_namespace) + robot2 = str(contact_info["root2_name"]).startswith(robot_namespace) + if not (robot1 or robot2): + return "other" + if robot1 and robot2: + return "robot_self" + + task_sampler_config = self.config.task_sampler_config + task_config = self.config.task_config + if self._contact_has_exact_name(contact_info, getattr(task_sampler_config, "table_body_name", "")): + return "robot_table" + if self._contact_has_exact_name(contact_info, getattr(task_sampler_config, "table_geom_name", "")): + return "robot_table" + if self._contact_has_name(contact_info, getattr(task_sampler_config, "place_receptacle_name", "")): + return "robot_bin" + if self._contact_has_name(contact_info, getattr(task_config, "pickup_obj_name", "")): + return "robot_pickup_object" + return "other" + + def _collect_contact_diagnostics(self) -> list[dict]: + data = self.task.env.current_data + contact_infos = [] + for contact_id in range(data.ncon): + contact = data.contact[contact_id] + if contact.dist > 0.0: + continue + geom1 = self._names_for_contact(contact.geom1) + geom2 = self._names_for_contact(contact.geom2) + contact_info = { + "contact_id": int(contact_id), + "dist": float(contact.dist), + "geom1_name": geom1["geom_name"], + "geom2_name": geom2["geom_name"], + "body1_name": geom1["body_name"], + "body2_name": geom2["body_name"], + "root1_name": geom1["root_body_name"], + "root2_name": geom2["root_body_name"], + } + contact_info["class"] = self._classify_contact(contact_info) + if contact_info["class"] != "other": + contact_infos.append(contact_info) + return contact_infos + + def _record_failure_diagnostic_step( + self, + target_pose: np.ndarray, + ik_succeeded: bool, + ) -> None: + if not self._failure_diagnostics_enabled(): + return + + tcp_pose = self._current_g1_tcp_pose() + target_pos_err = float(np.linalg.norm(tcp_pose[:3, 3] - target_pose[:3, 3])) + self._g1_failure_diag_steps.append( + { + "step_index": len(self._g1_failure_diag_steps), + "sim_time": float(self.task.env.current_data.time), + "phase": self.get_phase(), + "ik_succeeded": bool(ik_succeeded), + "sequential_ik_failures": int(self.sequential_ik_failures), + "tcp_pos": np.round(tcp_pose[:3, 3], 4).tolist(), + "target_pos": np.round(target_pose[:3, 3], 4).tolist(), + "target_pos_err": target_pos_err, + "contacts": self._collect_contact_diagnostics(), + } + ) + + def _pickup_object_diagnostics(self) -> dict: + task_config = self.config.task_config + pickup_obj_name = getattr(task_config, "pickup_obj_name", None) + if not pickup_obj_name: + return {"pickup_obj_name": None} + + diagnostics = { + "pickup_obj_name": pickup_obj_name, + "pickup_obj_start_pose": getattr(task_config, "pickup_obj_start_pose", None), + "place_receptacle_name": getattr(task_config, "place_receptacle_name", None), + } + try: + diagnostics["task_description"] = self.task.get_obs_scene().get("task_description") + except Exception: + diagnostics["task_description"] = None + try: + om = self.task.env.object_managers[self.task.env.current_batch_index] + pickup_obj = om.get_object_by_name(pickup_obj_name) + center, size = body_aabb( + self.task.env.current_model, + self.task.env.current_data, + pickup_obj.body_id, + visual_only=False, + ) + max_extent = float(np.max(size)) + volume = float(np.prod(size)) + diagnostics.update( + { + "pickup_obj_pos": np.round(pickup_obj.position, 4).tolist(), + "pickup_obj_aabb_center": np.round(center, 4).tolist(), + "pickup_obj_aabb_size": np.round(size, 4).tolist(), + "pickup_obj_max_extent_m": max_extent, + "pickup_obj_volume_m3": volume, + "object_size_outlier": bool( + max_extent + > getattr( + self.policy_config, + "diagnostic_large_object_max_extent_m", + 0.18, + ) + or volume + > getattr(self.policy_config, "diagnostic_large_object_volume_m3", 0.004) + ), + } + ) + except Exception as exc: + diagnostics["pickup_obj_error"] = repr(exc) + return diagnostics + + def _tcp_tracking_diagnostics(self) -> dict | None: + if self.action_idx >= len(self.action_primitives): + return None + action_primitive = self.action_primitives[self.action_idx] + if not isinstance(action_primitive, TCPMoveSequence) or action_primitive.move_seg_idx is None: + return None + + phase = action_primitive.get_current_phase() + target_pose = action_primitive.get_current_target_pose() + gripper_mg_id = self.robot_view.get_gripper_movegroup_ids()[0] + gripper = self.robot_view.get_gripper(gripper_mg_id) + target_from_tcp = np.linalg.inv(gripper.leaf_frame_to_world) @ target_pose + pos_err = float(np.linalg.norm(target_from_tcp[:3, 3])) + rot_err = float(R.from_matrix(target_from_tcp[:3, :3]).magnitude()) + pos_threshold, rot_threshold = self._tcp_tracking_thresholds_for_phase(phase) + return { + "phase": phase, + "tcp_tracking_pos_error": pos_err, + "tcp_tracking_rot_error": rot_err, + "tcp_tracking_pos_threshold": pos_threshold, + "tcp_tracking_rot_threshold": rot_threshold, + "tcp_tracking_position_failed": bool(pos_err > pos_threshold), + "tcp_tracking_rotation_failed": bool(rot_err > rot_threshold), + } + + def _tcp_tracking_thresholds_for_phase(self, phase: str) -> tuple[float, float]: + if phase == "pregrasp": + return ( + float(getattr(self.policy_config, "pregrasp_tcp_pos_err_threshold", 0.1)), + float(getattr(self.policy_config, "pregrasp_tcp_rot_err_threshold", np.inf)), + ) + return ( + float(self.policy_config.tcp_pos_err_threshold), + float(self.policy_config.tcp_rot_err_threshold), + ) + + def _failure_reason_guess( + self, + had_ik_failure: bool, + had_robot_collision: bool, + tracking_diagnostics: dict | None, + ) -> str: + if self._g1_last_ik_debug and self._g1_last_ik_debug.get("root_cause"): + return self._g1_last_ik_debug["root_cause"] + if self._g1_failure_reason_override: + return self._g1_failure_reason_override + if had_robot_collision and had_ik_failure: + return "ik_failure_with_robot_contact" + if had_robot_collision: + return "robot_collision" + if had_ik_failure: + return "unreachable_or_poor_grasp_pose" + if tracking_diagnostics: + position_failed = tracking_diagnostics["tcp_tracking_position_failed"] + rotation_failed = tracking_diagnostics["tcp_tracking_rotation_failed"] + if position_failed and rotation_failed: + return "tcp_pose_tracking_error" + if position_failed: + return "tcp_position_tracking_error" + if rotation_failed: + return "tcp_orientation_tracking_error" + return "planner_failure" + + def _log_failure_diagnostics(self) -> None: + if ( + not self._failure_diagnostics_enabled() + or getattr(self, "_g1_failure_diagnostics_logged", False) + ): + return + self._g1_failure_diagnostics_logged = True + + steps = getattr(self, "_g1_failure_diag_steps", []) + contacts_by_class: dict[str, int] = {} + for step in steps: + for contact in step["contacts"]: + contacts_by_class[contact["class"]] = contacts_by_class.get(contact["class"], 0) + 1 + + robot_contact_classes = { + "robot_table", + "robot_bin", + "robot_pickup_object", + "robot_self", + } + had_robot_collision = any(contacts_by_class.get(k, 0) > 0 for k in robot_contact_classes) + had_ik_failure = any(not step["ik_succeeded"] for step in steps) + + current_tcp_pose = self._current_g1_tcp_pose() + initial_tcp_pose = getattr(self, "_g1_initial_tcp_pose", current_tcp_pose) + tcp_displacement = float( + np.linalg.norm(current_tcp_pose[:3, 3] - initial_tcp_pose[:3, 3]) + ) + failure_before_motion = bool(tcp_displacement < 0.02 and had_ik_failure) + tracking_diagnostics = self._tcp_tracking_diagnostics() + + summary = { + **self._pickup_object_diagnostics(), + "phase_at_failure": self._g1_failure_phase_override or self.get_phase(), + "failure_reason_guess": self._failure_reason_guess( + had_ik_failure, + had_robot_collision, + tracking_diagnostics, + ), + "unfiltered_grasp_fallback_used": self._g1_unfiltered_grasp_fallback_used, + "unfiltered_grasp_fallback_reason": self._g1_unfiltered_grasp_fallback_reason, + "grasp_selector": self._g1_grasp_selector_debug, + "failure_before_motion": failure_before_motion, + "tcp_displacement_m": tcp_displacement, + "tcp_tracking": tracking_diagnostics, + "sequential_ik_failures": int(self.sequential_ik_failures), + "contact_counts_by_class": contacts_by_class, + "last_target_pose": ( + np.round(pose_mat_to_7d(self._g1_last_target_pose), 4).tolist() + if self._g1_last_target_pose is not None + else None + ), + "current_tcp_pose": np.round(pose_mat_to_7d(current_tcp_pose), 4).tolist(), + "num_diagnostic_steps": len(steps), + "held_grip_diagnostics": getattr(self, "_g1_grip_diag_samples", [])[-12:], + } + log.info("[G1_DIAG] failure_summary=%s", summary) + self._write_g1_failure_summary_debug_row(summary, contacts_by_class) + + for step in steps: + if step["contacts"] or not step["ik_succeeded"]: + log.info( + "[G1_DIAG] step=%s phase=%s ik=%s ik_fails=%s target_err=%.4f " + "tcp=%s target=%s contacts=%s", + step["step_index"], + step["phase"], + step["ik_succeeded"], + step["sequential_ik_failures"], + step["target_pos_err"], + step["tcp_pos"], + step["target_pos"], + step["contacts"], + ) + + def _handle_failure(self) -> dict[str, np.ndarray | bool]: + self._log_failure_diagnostics() + hold_duration = float( + getattr(self.policy_config, "diagnostic_failure_hold_duration_s", 0.0) + ) + if ( + self._failure_diagnostics_enabled() + and self.retry_count >= self.policy_config.max_retries + and hold_duration > 0.0 + ): + self._g1_terminal_failure_pending = True + if self._g1_failure_hold_action is None: + log.info("[G1_DIAG] holding failed attempt for %.2fs before done", hold_duration) + self._g1_failure_hold_action = NoopAction(self.robot_view, hold_duration) + if not self._g1_failure_hold_action.execute(): + return self._g1_failure_hold_action.get_current_action() + + return super()._handle_failure() + + def _check_for_failures(self) -> bool: + self._record_g1_grip_diagnostic_sample() + if self._g1_terminal_failure_pending: + return True + if self._g1_failure_reason_override: + if self.action_idx >= len(self.action_primitives): + return True + action_primitive = self.action_primitives[self.action_idx] + if action_primitive.start_time is None: + return False + return action_primitive.elapsed_time() >= action_primitive.duration + if self.action_idx < len(self.action_primitives): + action_primitive = self.action_primitives[self.action_idx] + if ( + isinstance(action_primitive, TCPMoveSequence) + and action_primitive.start_time is not None + and action_primitive.get_current_phase() == "pregrasp" + ): + tracking_diagnostics = self._tcp_tracking_diagnostics() + if tracking_diagnostics is None: + return False + return bool( + tracking_diagnostics["tcp_tracking_position_failed"] + or tracking_diagnostics["tcp_tracking_rotation_failed"] + ) + return super()._check_for_failures() + def _tcp_to_jp_fn(self, mg_id: str, target_pose: np.ndarray) -> dict[str, np.ndarray]: kinematics = self.task.env.current_robot.kinematics @@ -302,21 +2080,261 @@ def _tcp_to_jp_fn(self, mg_id: str, target_pose: np.ndarray) -> dict[str, np.nda self.robot_view.get_qpos_dict(), self.robot_view.base.pose, ) + self._g1_last_target_pose = target_pose.copy() action = self.robot_view.get_ctrl_dict() if jp is not None: self.sequential_ik_failures = 0 action["right_arm"] = jp["right_arm"] + self._record_failure_diagnostic_step(target_pose, True) else: self.sequential_ik_failures += 1 log.info(f"IK failed, holding current position, fails:{self.sequential_ik_failures}") + self._record_failure_diagnostic_step(target_pose, False) + self._run_g1_ik_debug(mg_id, target_pose) if self.sequential_ik_failures >= self.policy_config.max_sequential_ik_failures: log.info("Too many sequential IK failures, triggering retry.") return self._handle_failure() return action + def _run_diagnostic_ik( + self, + mg_id: str, + target_pose: np.ndarray, + unlocked_groups: list[str], + max_iter: int = 250, + ) -> dict: + return self.task.env.current_robot.kinematics.diagnose_ik( + mg_id, + target_pose, + unlocked_groups, + self.robot_view.get_qpos_dict(), + self.robot_view.base.pose, + max_iter=max_iter, + ) + + def _candidate_pose_for_phase(self, candidate_grasp_pose: np.ndarray) -> np.ndarray: + pose = candidate_grasp_pose.copy() + if self.get_phase() == "pregrasp": + pose[:3, 3] -= self.policy_config.pregrasp_z_offset * pose[:3, 2] + return pose + + def _summarize_ik_result(self, result: dict) -> str: + status = "ok" if result.get("success") else "fail" + pos_err = result.get("final_pos_error_norm") + rot_err = result.get("final_rot_error_norm") + clamps = result.get("clamp_event_count", 0) + singular = result.get("singular_iterations", 0) + unavailable = result.get("unavailable_groups") or [] + parts = [status] + if pos_err is not None: + parts.append(f"p={pos_err:.3f}") + if rot_err is not None: + parts.append(f"r={rot_err:.3f}") + if clamps: + parts.append(f"clamp={clamps}") + if singular: + parts.append(f"sing={singular}") + if unavailable: + parts.append(f"missing={','.join(unavailable)}") + return " ".join(parts) + + def _target_distance_diagnostics(self, target_pose: np.ndarray) -> dict: + base_pos = self.robot_view.base.pose[:3, 3] + target_pos = target_pose[:3, 3] + diagnostics = { + "target_z": float(target_pos[2]), + "target_base_distance": float(np.linalg.norm(target_pos - base_pos)), + } + try: + shoulder_pos = self.task.env.current_data.body( + f"{self.config.robot_config.robot_namespace}right_shoulder_pitch_link" + ).xpos + diagnostics["target_shoulder_distance"] = float(np.linalg.norm(target_pos - shoulder_pos)) + except KeyError: + diagnostics["target_shoulder_distance"] = None + return diagnostics + + def _ik_debug_root_cause(self, debug: dict, contacts: list[dict]) -> str: + contact_classes = {contact["class"] for contact in contacts} + if contact_classes & {"robot_table", "robot_bin", "robot_pickup_object"}: + return "table/object_collision" + if "robot_self" in contact_classes: + return "robot_self_collision" + + right_arm = debug["variants"]["right_arm"] + waist_right_arm = debug["variants"]["waist_right_arm"] + base_waist_right_arm = debug["variants"]["base_waist_right_arm"] + relaxed_orientation = debug["variants"]["relaxed_orientation"] + higher_z = debug["variants"]["higher_z"] + + if not right_arm.get("success") and waist_right_arm.get("success"): + return "needs_waist" + if not right_arm.get("success") and relaxed_orientation.get("success"): + return "orientation_constrained" + if not right_arm.get("success") and higher_z.get("success"): + return "orientation_constrained" + if not right_arm.get("success") and any( + candidate.get("success") for candidate in debug.get("top_k_grasp_results", []) + ): + return "orientation_constrained" + if right_arm.get("clamp_event_count", 0) > 0: + return "joint_limit" + if ( + right_arm.get("singular_iterations", 0) > 0 + or (right_arm.get("min_singular_value") or 1.0) < 1e-4 + ): + return "singularity" + if not right_arm.get("success") and base_waist_right_arm.get("success"): + return "object_too_far" + return "object_too_far" + + def _g1_ik_debug_table_path(self) -> Path: + output_dir = Path(self.config.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + return output_dir / "g1_ik_diagnostics.md" + + def _ensure_g1_ik_debug_table(self) -> Path: + table_path = self._g1_ik_debug_table_path() + if not table_path.exists(): + table_path.write_text( + "| episode_object | phase | grasp_id | target_z | base_dist | shoulder_dist | " + "right_arm | waist_right_arm | base_waist_right_arm | relaxed_orientation | " + "higher_z | top_k_success | contacts | root_cause |\n" + "| --- | --- | --- | ---: | ---: | ---: | --- | --- | --- | --- | --- | ---: | --- | --- |\n" + ) + return table_path + + def _write_g1_ik_debug_table(self, debug: dict) -> None: + table_path = self._ensure_g1_ik_debug_table() + selected_grasp = debug.get("selected_grasp") or {} + distances = debug["distances"] + variants = debug["variants"] + top_k_success = sum(1 for item in debug.get("top_k_grasp_results", []) if item.get("success")) + contacts = ",".join(sorted({contact["class"] for contact in debug["contacts"]})) or "none" + row = ( + f"| {debug['pickup_obj_name']} | {debug['phase']} | " + f"{selected_grasp.get('grasp_idx', 'n/a')} | " + f"{distances['target_z']:.3f} | {distances['target_base_distance']:.3f} | " + f"{(distances['target_shoulder_distance'] or float('nan')):.3f} | " + f"{self._summarize_ik_result(variants['right_arm'])} | " + f"{self._summarize_ik_result(variants['waist_right_arm'])} | " + f"{self._summarize_ik_result(variants['base_waist_right_arm'])} | " + f"{self._summarize_ik_result(variants['relaxed_orientation'])} | " + f"{self._summarize_ik_result(variants['higher_z'])} | " + f"{top_k_success} | {contacts} | {debug['root_cause']} |\n" + ) + with table_path.open("a") as f: + f.write(row) + + def _write_g1_failure_summary_debug_row( + self, + summary: dict, + contacts_by_class: dict[str, int], + ) -> None: + if not self._g1_ik_debug_enabled() or self._g1_last_ik_debug is not None: + return + + table_path = self._ensure_g1_ik_debug_table() + target_pose = self._g1_last_target_pose if self._g1_last_target_pose is not None else self._current_g1_tcp_pose() + distances = self._target_distance_diagnostics(target_pose) + selected_grasp = self._g1_selected_grasp_debug or {} + contacts = ",".join(sorted(contacts_by_class)) or "none" + row = ( + f"| {summary.get('pickup_obj_name')} | {summary.get('phase_at_failure')} | " + f"{selected_grasp.get('grasp_idx', 'n/a')} | " + f"{distances['target_z']:.3f} | {distances['target_base_distance']:.3f} | " + f"{(distances['target_shoulder_distance'] or float('nan')):.3f} | " + "n/a | n/a | n/a | n/a | n/a | 0 | " + f"{contacts} | {summary.get('failure_reason_guess')} |\n" + ) + with table_path.open("a") as f: + f.write(row) + + def _run_g1_ik_debug(self, mg_id: str, target_pose: np.ndarray) -> None: + if not self._g1_ik_debug_enabled(): + return + + current_tcp_pose = self._current_g1_tcp_pose() + relaxed_orientation_pose = target_pose.copy() + relaxed_orientation_pose[:3, :3] = current_tcp_pose[:3, :3] + higher_z_pose = target_pose.copy() + higher_z_pose[2, 3] += getattr(self.policy_config, "g1_ik_debug_higher_z_offset", 0.05) + + variants = { + "right_arm": self._run_diagnostic_ik(mg_id, target_pose, ["right_arm"]), + "waist_right_arm": self._run_diagnostic_ik( + mg_id, + target_pose, + ["waist", "right_arm"], + ), + "base_waist_right_arm": self._run_diagnostic_ik( + mg_id, + target_pose, + ["base", "waist", "right_arm"], + ), + "relaxed_orientation": self._run_diagnostic_ik( + mg_id, + relaxed_orientation_pose, + ["right_arm"], + ), + "higher_z": self._run_diagnostic_ik(mg_id, higher_z_pose, ["right_arm"]), + } + + top_k_grasp_results = [] + selected_grasp = self._g1_selected_grasp_debug or {} + candidate_poses = selected_grasp.get("top_candidate_poses") + candidate_ids = selected_grasp.get("top_candidate_ids") or [] + if candidate_poses is not None: + for candidate_id, candidate_pose in zip(candidate_ids, candidate_poses, strict=False): + candidate_target = self._candidate_pose_for_phase(candidate_pose) + result = self._run_diagnostic_ik(mg_id, candidate_target, ["right_arm"], max_iter=150) + result["candidate_id"] = int(candidate_id) + top_k_grasp_results.append(result) + + contacts = self._collect_contact_diagnostics() + debug = { + "pickup_obj_name": getattr(self.config.task_config, "pickup_obj_name", None), + "phase": self.get_phase(), + "selected_grasp": { + key: value + for key, value in selected_grasp.items() + if key not in {"grasp_pose_world", "top_candidate_poses"} + }, + "target_pose": np.round(pose_mat_to_7d(target_pose), 6).tolist(), + "current_tcp_pose": np.round(pose_mat_to_7d(current_tcp_pose), 6).tolist(), + "right_arm_qpos": np.round(self.robot_view.get_move_group("right_arm").joint_pos, 6).tolist(), + "distances": self._target_distance_diagnostics(target_pose), + "variants": variants, + "top_k_grasp_results": top_k_grasp_results, + "contacts": contacts, + } + debug["root_cause"] = self._ik_debug_root_cause(debug, contacts) + self._g1_last_ik_debug = debug + + log.info( + "[G1_IK_DEBUG] object=%s phase=%s grasp=%s root_cause=%s right_arm=%s " + "waist_right_arm=%s base_waist_right_arm=%s relaxed=%s higher_z=%s contacts=%s", + debug["pickup_obj_name"], + debug["phase"], + debug["selected_grasp"].get("grasp_idx"), + debug["root_cause"], + self._summarize_ik_result(variants["right_arm"]), + self._summarize_ik_result(variants["waist_right_arm"]), + self._summarize_ik_result(variants["base_waist_right_arm"]), + self._summarize_ik_result(variants["relaxed_orientation"]), + self._summarize_ik_result(variants["higher_z"]), + contacts, + ) + self._write_g1_ik_debug_table(debug) + def check_feasible_ik(self, pose: np.ndarray) -> bool: + if self._g1_bypass_feasible_ik: + if pose.ndim > 2: + return np.ones(pose.shape[0], dtype=bool) + return True + if not self.policy_config.filter_feasible_grasps: if pose.ndim > 2: return np.ones(pose.shape[0], dtype=bool) diff --git a/molmo_spaces/tasks/unitree_g1_tabletop_task_sampler.py b/molmo_spaces/tasks/unitree_g1_tabletop_task_sampler.py index 38c9fc9f..7d4faeb9 100644 --- a/molmo_spaces/tasks/unitree_g1_tabletop_task_sampler.py +++ b/molmo_spaces/tasks/unitree_g1_tabletop_task_sampler.py @@ -75,6 +75,29 @@ def _sample_xy(self, center_xy: tuple[float, float], size_xy: tuple[float, float half_size = np.asarray(size_xy, dtype=float) / 2.0 return center + np.random.uniform(-half_size, half_size) + def _pickup_clear_of_receptacle( + self, + env: CPUMujocoEnv, + pickup_obj_name: str, + receptacle_xy: np.ndarray, + ) -> bool: + pickup_body = create_mlspaces_body(env.current_data, pickup_obj_name) + center, size = body_aabb( + env.current_model, + env.current_data, + pickup_body.body_id, + visual_only=False, + ) + pickup_half_size_xy = np.asarray(size[:2], dtype=float) / 2.0 + receptacle_half_size_xy = np.asarray( + self.config.task_sampler_config.place_receptacle_half_size_xy, + dtype=float, + ) + clearance = float(self.config.task_sampler_config.pickup_receptacle_min_clearance) + required_separation = pickup_half_size_xy + receptacle_half_size_xy + clearance + xy_delta = np.abs(np.asarray(center[:2], dtype=float) - receptacle_xy) + return bool(np.any(xy_delta > required_separation)) + def _place_pickup_on_table(self, env: CPUMujocoEnv, pickup_obj_name: str) -> None: body = create_mlspaces_body(env.current_data, pickup_obj_name) model = env.current_model @@ -100,13 +123,20 @@ def _place_pickup_on_table(self, env: CPUMujocoEnv, pickup_obj_name: str) -> Non self._zero_free_body_velocity(env, pickup_obj_name) mujoco.mj_forward(model, data) - def _place_receptacle_on_table(self, env: CPUMujocoEnv) -> None: + def _place_receptacle_on_table(self, env: CPUMujocoEnv, pickup_obj_name: str) -> None: receptacle_name = self.config.task_sampler_config.place_receptacle_name body = create_mlspaces_body(env.current_data, receptacle_name) - xy = self._sample_xy( - self.config.task_sampler_config.place_workspace_center_xy, - self.config.task_sampler_config.place_workspace_size_xy, - ) + for _ in range(self.config.task_sampler_config.max_place_receptacle_sampling_attempts): + xy = self._sample_xy( + self.config.task_sampler_config.place_workspace_center_xy, + self.config.task_sampler_config.place_workspace_size_xy, + ) + if self._pickup_clear_of_receptacle(env, pickup_obj_name, xy): + break + else: + raise HouseInvalidForTask( + f"Unable to place {receptacle_name} clear of {pickup_obj_name}" + ) body.position = np.array( [ xy[0], @@ -119,7 +149,8 @@ def _place_receptacle_on_table(self, env: CPUMujocoEnv) -> None: def _sample_and_place_robot(self, env: CPUMujocoEnv) -> None: robot_view = env.current_robot.robot_view - base_qpos = self.config.robot_config.init_qpos["base"] + base_qpos = np.asarray(self.config.robot_config.init_qpos["base"], dtype=float).copy() + base_qpos[:2] = np.asarray(self.config.task_sampler_config.robot_base_xy, dtype=float) robot_view.base.pose = pos_quat_to_pose_mat(np.asarray(base_qpos[:3]), base_qpos[3:7]) robot_view.base.joint_vel = np.zeros(robot_view.base.vel_dim) env.current_robot.sync_pinned_base_pose() @@ -139,7 +170,7 @@ def _select_pickup_object(self, env: CPUMujocoEnv) -> int: self._restore_pickupables_to_staging(env) self._place_pickup_on_table(env, pickup_obj_name) - self._place_receptacle_on_table(env) + self._place_receptacle_on_table(env, pickup_obj_name) self._sample_and_place_robot(env) om = env.object_managers[env.current_batch_index] diff --git a/molmo_spaces/utils/grasp_sample.py b/molmo_spaces/utils/grasp_sample.py index 789d9df3..00c22255 100644 --- a/molmo_spaces/utils/grasp_sample.py +++ b/molmo_spaces/utils/grasp_sample.py @@ -519,6 +519,26 @@ def compute_grasp_pose( raise ValueError("No feasible grasp found") original_grasp_idx = grasp_idx % (len(grasp_poses_world) // 2) + if hasattr(policy, "_record_g1_selected_grasp_debug"): + policy._record_g1_selected_grasp_debug( + { + "grasp_idx": int(grasp_idx), + "original_grasp_idx": int(original_grasp_idx), + "grasp_pose_world": grasp_pose_world.copy(), + "top_candidate_ids": close_grasp_ids[ + : getattr(policy.policy_config, "g1_ik_debug_top_k_grasps", 5) + ] + .astype(int) + .tolist(), + "top_candidate_poses": grasp_poses_world[ + close_grasp_ids[ + : getattr(policy.policy_config, "g1_ik_debug_top_k_grasps", 5) + ] + ].copy(), + "selected_pos_cost_m": float(dists_tcp_p[grasp_idx]), + "selected_rot_cost_deg": float(dist_tcp_o[grasp_idx]), + } + ) log.info( f"Feasible grasp found {grasp_idx} (originally {original_grasp_idx}): w/ {dists_tcp_p[grasp_idx]:.3f}[m] {dist_tcp_o[grasp_idx]:.3f}[deg]" ) diff --git a/scripts/assets/prepare_unitree_g1.py b/scripts/assets/prepare_unitree_g1.py index 6fb82398..a633eb0b 100644 --- a/scripts/assets/prepare_unitree_g1.py +++ b/scripts/assets/prepare_unitree_g1.py @@ -23,6 +23,16 @@ SOURCE_MESH_DIR_REL_PATH = Path("g1_description/meshes") OUTPUT_XML_NAME = "model.xml" DEFAULT_PELVIS_HEIGHT_M = 0.793 +DEX1_GRASP_SITE_POS = [0.148, 0.0, 0.0] +DEX1_FINGERTIP_PAD_SIZE = [0.02, 0.01, 0.045] +DEX1_FINGERTIP_PAD_POSITIONS = { + "dex1_finger_link_1": [0.1065, -0.0285, 0.0], + "dex1_finger_link_2": [0.1065, 0.0285, 0.0], +} +DEX1_FINGERTIP_PAD_FRICTION = [12.0, 0.2, 0.02] +DEX1_FINGERTIP_PAD_CONDIM = 6 +DEX1_FINGERTIP_PAD_GROUP = 4 +DEX1_HAND_FORCE_LIMIT_MULTIPLIER = 3.0 LEFT_ARM_STOW_QPOS = { "left_shoulder_pitch_joint": 0.25, "left_shoulder_roll_joint": 0.35, @@ -42,7 +52,7 @@ "leg": (120.0, 12.0), "waist": (80.0, 8.0), "arm": (60.0, 6.0), - "hand": (25.0, 2.5), + "hand": (80.0, 8.0), } @@ -159,8 +169,13 @@ def _add_position_actuators(spec: mujoco.MjSpec, joints: list[JointSpec]) -> Non stiffness, damping = GAIN_BY_JOINT_GROUP[joint.group] kwargs = {} if joint.effort is not None: + effort = ( + joint.effort * DEX1_HAND_FORCE_LIMIT_MULTIPLIER + if joint.group == "hand" + else joint.effort + ) kwargs["forcelimited"] = 1 - kwargs["forcerange"] = [-joint.effort, joint.effort] + kwargs["forcerange"] = [-effort, effort] spec.add_actuator( name=joint.name, @@ -189,12 +204,37 @@ def _add_kinematics_sites(spec: mujoco.MjSpec) -> None: ) wrist_body.add_site( name=f"{side}_grasp_site", - pos=[0.095, 0.0, 0.0], + pos=DEX1_GRASP_SITE_POS, size=[0.015], rgba=[0.1, 1.0, 0.4, 1.0], ) +def _add_dex1_fingertip_contact_pads(spec: mujoco.MjSpec) -> None: + for side in ("left", "right"): + for finger_suffix, pad_pos in DEX1_FINGERTIP_PAD_POSITIONS.items(): + finger_body_name = f"{side}_{finger_suffix}" + finger_body = spec.body(finger_body_name) + if finger_body is None: + raise ValueError(f"Expected body `{finger_body_name}` in Unitree G1 model") + + finger_id = finger_suffix.rsplit("_", maxsplit=1)[-1] + finger_body.add_geom( + name=f"{side}_dex1_fingertip_pad_{finger_id}", + type=mujoco.mjtGeom.mjGEOM_BOX, + pos=pad_pos, + size=DEX1_FINGERTIP_PAD_SIZE, + contype=1, + conaffinity=1, + condim=DEX1_FINGERTIP_PAD_CONDIM, + priority=2, + friction=DEX1_FINGERTIP_PAD_FRICTION, + density=0.0, + group=DEX1_FINGERTIP_PAD_GROUP, + rgba=[0.05, 0.8, 0.2, 0.35], + ) + + def prepare_unitree_g1( unitree_urdf_root: str | os.PathLike[str], output_dir: str | os.PathLike[str], @@ -225,6 +265,7 @@ def prepare_unitree_g1( pelvis.pos = [0.0, 0.0, DEFAULT_PELVIS_HEIGHT_M] pelvis.add_freejoint(name="floating_base_joint") _add_kinematics_sites(spec) + _add_dex1_fingertip_contact_pads(spec) _add_position_actuators(spec, joints) model = spec.compile() diff --git a/scripts/assets/prepare_unitree_g1_tabletop.py b/scripts/assets/prepare_unitree_g1_tabletop.py index 36ed8593..ae7d8b40 100644 --- a/scripts/assets/prepare_unitree_g1_tabletop.py +++ b/scripts/assets/prepare_unitree_g1_tabletop.py @@ -33,8 +33,7 @@ PLACE_RECEPTACLE_BODY_NAME = "g1_place_bin" PLACE_RECEPTACLE_SITE_NAME = "g1_place_bin_site" -PELVIS_MINUS_10CM_TABLETOP_HEIGHT_M = 0.693 -PELVIS_TABLETOP_HEIGHT_M = 0.793 +TABLETOP_HEIGHT_M = 0.693 @dataclass(frozen=True) @@ -46,9 +45,8 @@ class TabletopSceneSpec: SCENE_SPECS = ( TabletopSceneSpec( "unitree_g1_tabletop_pelvis_minus_10cm_v1", - PELVIS_MINUS_10CM_TABLETOP_HEIGHT_M, + TABLETOP_HEIGHT_M, ), - TabletopSceneSpec("unitree_g1_tabletop_pelvis_height_v1", PELVIS_TABLETOP_HEIGHT_M), ) From 2b7a4cf3e6de196fd72cdbb3e3747f605e70123d Mon Sep 17 00:00:00 2001 From: Abhishek Parmar Date: Tue, 12 May 2026 20:40:50 +0800 Subject: [PATCH 12/13] fix: stabilize G1 tabletop pick-and-place carry MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Stack of debugging-driven changes to eliminate grip/slip/wedge/hidden- collision failure modes during the right-arm tabletop pick-and-place trajectory. The remaining residual is the runtime DLS IK getting stuck on cross-body preplace approaches; addressed via planning-time mitigations only — no runtime planner replacement yet. Asset side (scripts/assets/prepare_unitree_g1.py): - Fingertip pad reoriented from (0.02, 0.01, 0.045) (vertical fin) to (0.045, 0.01, 0.02) (flat plate along fingertip face). - Pad sliding/torsional/rolling friction (12, 0.2, 0.02) -> (20, 0.4, 0.04) for slip resistance. - Arm actuator damping doubled to (60, 12); stiffness unchanged to avoid hard-approach push-aside of pickup objects. Workspace + clearance (molmo_spaces/configs/task_sampler_configs.py and the policy config): - pickup_workspace_size_xy (0.16, 0.10) -> (0.16, 0.06). - place_workspace_center_xy (0.42, 0.08) -> (0.42, 0.02), place_workspace_size_xy (0.10, 0.08) -> (0.04, 0.02). - g1_place_travel_object_clearance 0.13 -> 0.07; previous 0.12 pushed the preplace target out of the right-arm reach envelope. Selector + runtime IK (molmo_spaces/policy/solvers/...): - max_sequential_ik_failures 4 -> 20. - g1_grasp_joint_margin_weight 0.5 -> 1.5; g1_grasp_joint_motion_weight 0.25 -> 1.0 so the selector prefers grasps whose IK chain stays away from limits and has small phase-to-phase joint excursions. - g1_runtime_ik_damping 1e-3 (was using shared default 1e-12), g1_selector_ik_damping 1e-4, g1_runtime_ik_dt 0.5. - Null-space posture nudge in _tcp_to_jp_fn biases the redundant 1 DoF of the 7-DoF right arm toward init_qpos without changing TCP pose. - Smoothed ctrl EMA scaffolding present but disabled by default (g1_runtime_ctrl_smoothing 0.0) — averaging across IK basins worsened rotation tracking. - HoldCurrentControlAction (not NoopAction) for the runtime failure hold so a still-gripped held object isn't released during the 2 s terminal hold. - Closed-grasp probe now sets arm ctrl from the candidate qpos and re-pins the base / locked left arm between mj_step calls so the wrist actually holds the candidate during the settle window. Diagnostics (G1-only, additive): - [G1_CONV] convention check at reset (pelvis + right_grasp_site axes in world). - [G1_BASIN_FLIP] warning when right-arm IK output ctrl jumps more than g1_runtime_basin_flip_threshold_rad (0.30) between frames. - [G1_CONTACTS] full per-frame contact dump with body#geomNNN fallback for unnamed structural geoms. Test updates mirror the new defaults. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../component_tests/test_unitree_g1.py | 133 +++++- molmo_spaces/configs/policy_configs.py | 96 +++- molmo_spaces/configs/task_sampler_configs.py | 17 +- .../pick_and_place_planner_policy.py | 440 +++++++++++++++++- scripts/assets/prepare_unitree_g1.py | 95 +++- 5 files changed, 727 insertions(+), 54 deletions(-) diff --git a/mlspaces_tests/component_tests/test_unitree_g1.py b/mlspaces_tests/component_tests/test_unitree_g1.py index c1ae3747..6aea149b 100644 --- a/mlspaces_tests/component_tests/test_unitree_g1.py +++ b/mlspaces_tests/component_tests/test_unitree_g1.py @@ -27,6 +27,7 @@ DEX1_FINGERTIP_PAD_POSITIONS, DEX1_FINGERTIP_PAD_SIZE, DEX1_GRASP_SITE_POS, + DEX1_GRASP_SITE_QUAT, DEX1_HAND_FORCE_LIMIT_MULTIPLIER, LEFT_ARM_STOW_QPOS, LEFT_HAND_OPEN_QPOS, @@ -123,6 +124,7 @@ def test_unitree_g1_grasp_contact_filter_requires_right_fingertip_pad(): policy.policy_config = SimpleNamespace( g1_reject_grasp_table_contact=True, g1_reject_open_grasp_object_contact=True, + g1_allow_open_fingertip_pad_contact=True, g1_closed_grasp_min_pad_geom_count=2, g1_closed_grasp_max_object_shift_m=0.05, ) @@ -162,9 +164,12 @@ def test_unitree_g1_grasp_contact_filter_requires_right_fingertip_pad(): "root1_name": "robot_0/pelvis", } assert not policy._g1_is_allowed_grasp_object_contact(wrist_contact) + assert not policy._g1_grasp_contact_quality_failure( + policy._g1_contact_quality_from_contacts([pad_contact]) + ) assert ( policy._g1_grasp_contact_quality_failure( - policy._g1_contact_quality_from_contacts([pad_contact]) + policy._g1_contact_quality_from_contacts([finger_link_contact]) ) == "open_grasp_object_contact" ) @@ -212,7 +217,17 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): for site_name in KINEMATICS_SITE_NAMES: assert model.site(site_name).id >= 0 for site_name in ("left_grasp_site", "right_grasp_site"): - np.testing.assert_allclose(model.site_pos[model.site(site_name).id], DEX1_GRASP_SITE_POS) + site_id = model.site(site_name).id + np.testing.assert_allclose(model.site_pos[site_id], DEX1_GRASP_SITE_POS) + np.testing.assert_allclose(model.site_quat[site_id], DEX1_GRASP_SITE_QUAT) + assert ( + mujoco.mj_name2id( + model, + mujoco.mjtObj.mjOBJ_GEOM, + "right_grasp_site_debug_axis_x", + ) + == -1 + ) for side in ("left", "right"): for finger_id, finger_suffix in (("1", "dex1_finger_link_1"), ("2", "dex1_finger_link_2")): pad = model.geom(f"{side}_dex1_fingertip_pad_{finger_id}") @@ -252,6 +267,12 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): assert scene_model.qpos0[scene_base_qposadr + 2] == pytest.approx(0.793) mujoco.mj_forward(scene_model, scene_data) assert not _has_deep_robot_floor_penetration(scene_model, scene_data) + wrist_site_id = scene_model.site("robot_0/right_wrist_site").id + grasp_site_id = scene_model.site("robot_0/right_grasp_site").id + wrist_to_grasp = scene_data.site_xpos[grasp_site_id] - scene_data.site_xpos[wrist_site_id] + wrist_to_grasp /= np.linalg.norm(wrist_to_grasp) + grasp_xmat = scene_data.site_xmat[grasp_site_id].reshape(3, 3) + np.testing.assert_allclose(grasp_xmat[:, 2], wrist_to_grasp, atol=1e-6) view = config.robot_view_factory(scene_data, config.robot_namespace) expected_counts = { @@ -335,9 +356,9 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): assert tabletop_datagen_config.end_on_success assert tabletop_datagen_config.task_sampler_config.robot_base_xy == (0.08, 0.0) assert tabletop_datagen_config.task_sampler_config.pickup_workspace_center_xy == (0.42, -0.20) - assert tabletop_datagen_config.task_sampler_config.pickup_workspace_size_xy == (0.16, 0.10) - assert tabletop_datagen_config.task_sampler_config.place_workspace_center_xy == (0.42, 0.08) - assert tabletop_datagen_config.task_sampler_config.place_workspace_size_xy == (0.10, 0.08) + assert tabletop_datagen_config.task_sampler_config.pickup_workspace_size_xy == (0.16, 0.06) + assert tabletop_datagen_config.task_sampler_config.place_workspace_center_xy == (0.42, 0.02) + assert tabletop_datagen_config.task_sampler_config.place_workspace_size_xy == (0.04, 0.02) assert tabletop_datagen_config.task_sampler_config.place_receptacle_half_size_xy == ( 0.14, 0.11, @@ -363,7 +384,7 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): assert not tabletop_datagen_config.policy_config.filter_feasible_grasps assert PickAndPlacePlannerPolicyConfig().tcp_rot_err_threshold == np.radians(30.0) assert PickAndPlacePlannerPolicyConfig().gripper_empty_threshold == 0.002 - assert tabletop_datagen_config.policy_config.pregrasp_z_offset == 0.12 + assert tabletop_datagen_config.policy_config.pregrasp_z_offset == 0.06 assert tabletop_datagen_config.policy_config.gripper_close_duration == 1.0 assert tabletop_datagen_config.policy_config.gripper_empty_threshold == 0.0005 assert tabletop_datagen_config.policy_config.grasp_vertical_cost_weight == 0.5 @@ -375,32 +396,44 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): assert tabletop_datagen_config.policy_config.g1_grasp_candidate_limit == 256 assert tabletop_datagen_config.policy_config.g1_grasp_ik_eval_limit == 256 assert tabletop_datagen_config.policy_config.g1_grasp_require_all_pick_place_phases - assert tabletop_datagen_config.policy_config.g1_grasp_joint_margin_weight == 0.5 - assert tabletop_datagen_config.policy_config.g1_grasp_joint_motion_weight == 0.25 + assert tabletop_datagen_config.policy_config.g1_grasp_joint_margin_weight == 1.5 + assert tabletop_datagen_config.policy_config.g1_grasp_joint_motion_weight == 1.0 assert tabletop_datagen_config.policy_config.g1_grasp_topdown_weight == 1.0 + assert tabletop_datagen_config.policy_config.g1_runtime_ik_damping == 5e-3 + assert tabletop_datagen_config.policy_config.g1_selector_ik_damping == 1e-4 + assert tabletop_datagen_config.policy_config.g1_runtime_ik_dt == 0.5 + assert tabletop_datagen_config.policy_config.g1_runtime_ctrl_smoothing == 0.0 assert tabletop_datagen_config.policy_config.g1_grasp_min_vertical_axis_z == 0.75 assert tabletop_datagen_config.policy_config.g1_grasp_max_tcp_rot_deg == 120.0 - assert tabletop_datagen_config.policy_config.g1_ignore_flipped_grasps - assert tabletop_datagen_config.policy_config.g1_grasp_inward_xy_offset == 0.006 + assert not tabletop_datagen_config.policy_config.g1_ignore_flipped_grasps + assert tabletop_datagen_config.policy_config.g1_grasp_inward_xy_offset == 0.0 assert tabletop_datagen_config.policy_config.g1_grasp_table_clearance == 0.065 - assert tabletop_datagen_config.policy_config.g1_center_grasp_lateral + assert not tabletop_datagen_config.policy_config.g1_center_grasp_lateral assert tabletop_datagen_config.policy_config.g1_grasp_lateral_centering_scale == 1.0 assert tabletop_datagen_config.policy_config.g1_grasp_lateral_centering_max_offset == 0.02 + assert not tabletop_datagen_config.policy_config.g1_center_grasp_forward + assert tabletop_datagen_config.policy_config.g1_grasp_forward_centering_scale == 1.0 + assert tabletop_datagen_config.policy_config.g1_grasp_forward_centering_max_offset == 0.03 + assert tabletop_datagen_config.policy_config.g1_grasp_forward_centering_target_m == 0.0 assert not tabletop_datagen_config.policy_config.g1_level_grasp_orientation assert tabletop_datagen_config.policy_config.g1_grasp_level_max_tilt_deg == 35.0 assert not tabletop_datagen_config.policy_config.g1_require_fingertip_pad_grasp_contact assert tabletop_datagen_config.policy_config.g1_reject_non_fingertip_grasp_object_contact assert tabletop_datagen_config.policy_config.g1_reject_grasp_table_contact assert tabletop_datagen_config.policy_config.g1_reject_open_grasp_object_contact + assert tabletop_datagen_config.policy_config.g1_allow_open_fingertip_pad_contact assert tabletop_datagen_config.policy_config.g1_grasp_single_pad_contact_penalty == 0.5 assert tabletop_datagen_config.policy_config.g1_closed_grasp_quality_enabled - assert tabletop_datagen_config.policy_config.g1_closed_grasp_settle_steps == 25 + assert tabletop_datagen_config.policy_config.g1_closed_grasp_settle_steps == 120 assert tabletop_datagen_config.policy_config.g1_closed_grasp_min_pad_geom_count == 2 assert tabletop_datagen_config.policy_config.g1_closed_grasp_max_object_shift_m == 0.05 assert tabletop_datagen_config.policy_config.g1_closed_grasp_penalty_per_missing_pad == 1.0 - assert tabletop_datagen_config.policy_config.g1_pregrasp_min_vertical_lift == 0.10 - assert tabletop_datagen_config.policy_config.g1_pregrasp_object_clearance == 0.12 - assert tabletop_datagen_config.policy_config.g1_place_travel_object_clearance == 0.13 + assert tabletop_datagen_config.policy_config.g1_pregrasp_min_vertical_lift == 0.0 + assert tabletop_datagen_config.policy_config.g1_pregrasp_object_clearance == 0.04 + assert tabletop_datagen_config.policy_config.g1_place_travel_object_clearance == 0.07 + assert tabletop_datagen_config.policy_config.g1_runtime_null_space_weight == 0.3 + assert tabletop_datagen_config.policy_config.g1_runtime_null_space_damping == 1e-4 + assert tabletop_datagen_config.policy_config.g1_runtime_null_space_max_step_rad == 0.05 assert tabletop_datagen_config.policy_config.g1_held_object_speed == 0.04 assert tabletop_datagen_config.policy_config.g1_postgrasp_hold_duration == 0.4 assert tabletop_datagen_config.policy_config.g1_record_partial_attempt_on_no_full_grasp_candidate @@ -595,6 +628,40 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): ) assert not unreachable_eval["success"] assert unreachable_eval["failed_phase"] == "pregrasp" + relaxed_pregrasp_result = fake_tabletop_policy._g1_accept_relaxed_candidate_ik_result( + "pregrasp", + { + "success": False, + "final_pos_error_norm": 0.05, + "final_rot_error_norm": 10.0, + "final_qpos": pick_datagen_config.robot_config.init_qpos, + "qpos": None, + }, + ) + assert relaxed_pregrasp_result["success"] + assert relaxed_pregrasp_result["accepted_with_g1_candidate_thresholds"] + strict_grasp_result = fake_tabletop_policy._g1_accept_relaxed_candidate_ik_result( + "grasp", + { + "success": False, + "final_pos_error_norm": 0.05, + "final_rot_error_norm": 0.4, + "final_qpos": pick_datagen_config.robot_config.init_qpos, + "qpos": None, + }, + ) + assert strict_grasp_result["success"] + rejected_grasp_result = fake_tabletop_policy._g1_accept_relaxed_candidate_ik_result( + "grasp", + { + "success": False, + "final_pos_error_norm": 0.05, + "final_rot_error_norm": 1.0, + "final_qpos": pick_datagen_config.robot_config.init_qpos, + "qpos": None, + }, + ) + assert not rejected_grasp_result["success"] clearance_model = mujoco.MjModel.from_xml_string( """ @@ -653,12 +720,29 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): assert g1_target_poses["grasp"][2, 3] >= ( table_top_z + tabletop_datagen_config.policy_config.g1_grasp_table_clearance ) - object_in_grasp = g1_target_poses["grasp"][:3, :3].T @ ( - pickup_obj.position - g1_target_poses["grasp"][:3, 3] - ) - assert abs(object_in_grasp[1]) < 1e-8 + # Default G1 candidate construction must not shift the grasp pose in XY: + # the DROID grasp bank already encodes pinch placement in the corrected + # grasp-site frame, and lateral/forward centering and inward XY offsets + # default off. + np.testing.assert_allclose(g1_target_poses["grasp"][:2, 3], grasp_pose[:2, 3]) np.testing.assert_allclose(g1_target_poses["grasp"][:3, 2], grasp_pose[:3, 2]) + # Opt-in centering still applies its math when enabled. + centered_policy = object.__new__(UnitreeG1RightArmPickAndPlacePlannerPolicy) + centered_policy.policy_config = tabletop_datagen_config.policy_config.model_copy( + update={ + "g1_center_grasp_forward": True, + "g1_grasp_forward_centering_target_m": 0.02, + } + ) + forward_offset_pose = np.eye(4) + forward_offset_pose[:3, 3] = [0.0, 0.0, -0.05] + forward_centered_pose = centered_policy._g1_center_grasp_forward( + forward_offset_pose, + SimpleNamespace(position=np.array([0.0, 0.0, 0.0])), + ) + np.testing.assert_allclose(forward_centered_pose[:3, 3], [0.0, 0.0, -0.02]) + fake_clearance_policy.policy_config = tabletop_datagen_config.policy_config.model_copy( update={"g1_level_grasp_orientation": True} ) @@ -673,8 +757,13 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): np.eye(3), atol=1e-7, ) - assert g1_target_poses["pregrasp"][2, 3] >= grasp_pose[2, 3] + 0.10 - assert g1_target_poses["pregrasp"][2, 3] >= pickup_top_z + 0.12 + assert g1_target_poses["pregrasp"][2, 3] >= ( + pickup_top_z + + tabletop_datagen_config.policy_config.g1_pregrasp_object_clearance + ) + # Pregrasp also lifts strictly above the (clamped) grasp pose along the + # gripper-forward axis, so the wrist always sits above the grasp. + assert g1_target_poses["pregrasp"][2, 3] > g1_target_poses["grasp"][2, 3] pickup_bottom_z = 0.70 receptacle_top_z = 0.75 min_travel_bottom_z = ( @@ -687,7 +776,7 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): + pickup_bottom_z - g1_target_poses["grasp"][2, 3] ) - assert carried_bottom_z >= min_travel_bottom_z + assert carried_bottom_z + 1e-9 >= min_travel_bottom_z place_bottom_z = ( g1_target_poses["place"][2, 3] + pickup_bottom_z diff --git a/molmo_spaces/configs/policy_configs.py b/molmo_spaces/configs/policy_configs.py index dc41f45e..39d4ec09 100644 --- a/molmo_spaces/configs/policy_configs.py +++ b/molmo_spaces/configs/policy_configs.py @@ -182,7 +182,7 @@ class UnitreeG1RightArmPickAndPlacePlannerPolicyConfig(PickAndPlacePlannerPolicy """Pick-and-place planner config for fixed-base G1 right-arm tabletop datagen.""" policy_cls: type = None - pregrasp_z_offset: float = 0.12 + pregrasp_z_offset: float = 0.06 gripper_close_duration: float = 1.0 gripper_empty_threshold: float = 0.0005 grasp_rot_cost_weight: float = 0.02 @@ -194,7 +194,12 @@ class UnitreeG1RightArmPickAndPlacePlannerPolicyConfig(PickAndPlacePlannerPolicy grasp_collision_batch_size: int = 32 grasp_collision_max_grasps: int = 64 max_retries: int = 0 - max_sequential_ik_failures: int = 4 + # Bumped from 4 → 20. The standalone IK teleop shows the right arm can + # reach the preplace region; the runtime was just aborting too early on + # transient near-singular sub-targets during the continuous SE(3) sweep. + # Higher threshold gives the seed chain a chance to recover before the + # planner declares the carry unreachable. + max_sequential_ik_failures: int = 20 phase_timeout: float = 24.0 pregrasp_tcp_pos_err_threshold: float = 0.1 pregrasp_tcp_rot_err_threshold: float = float("inf") @@ -210,32 +215,101 @@ class UnitreeG1RightArmPickAndPlacePlannerPolicyConfig(PickAndPlacePlannerPolicy g1_grasp_candidate_limit: int = 256 g1_grasp_ik_eval_limit: int = 256 g1_grasp_require_all_pick_place_phases: bool = True - g1_grasp_joint_margin_weight: float = 0.5 - g1_grasp_joint_motion_weight: float = 0.25 + # Bumped from the original 0.5 / 0.25 / 1.0 defaults. The joint-margin + # and joint-motion weights are the selector's lever for biasing toward + # grasps whose IK chain stays away from joint limits and whose + # phase-to-phase joint excursion is small — directly addressing the + # runtime "elbow flips mid-carry" failure mode by pre-filtering for + # grasps with a smoother kinematic path. + g1_grasp_joint_margin_weight: float = 1.5 + g1_grasp_joint_motion_weight: float = 1.0 g1_grasp_topdown_weight: float = 1.0 + # Damped-least-squares damping passed to the shared IK solver at G1 + # call sites only. Higher damping = solver stays closer to seed near + # singularities, fewer elbow-branch flips, at the cost of slightly + # slower / less precise convergence. Defaults unchanged for every + # other robot. + g1_runtime_ik_damping: float = 5e-3 + g1_selector_ik_damping: float = 1e-4 + # Runtime IK integration step size. Default 1.0 = full Newton step; + # smaller values produce smoother per-iteration joint changes near + # singular configurations at the cost of slightly more iterations to + # converge. 0.5 is a good balance. + g1_runtime_ik_dt: float = 0.5 + # Exponential-moving-average smoothing factor for the right-arm ctrl + # output. `ctrl_new = α·ctrl_prev + (1-α)·ik_output`. 0.0 = no + # smoothing (raw IK output); 1.0 = freeze (never update). Disabled + # for now: α=0.4 introduced bimodal-basin averaging (mixing + # elbow-up/down IK solutions into a between-state) that produced a + # 43° wrist tracking error in run 23. The arm-gain + IK-dt changes + # alone provided the IK-stability improvement. + g1_runtime_ctrl_smoothing: float = 0.0 + # If any right-arm joint changes by more than this between consecutive + # IK calls, log a [G1_BASIN_FLIP] warning. Indicates the DLS solver + # jumped kinematic branches (elbow swivel or wrist flip) between + # adjacent sub-targets along the Cartesian sweep. + g1_runtime_basin_flip_threshold_rad: float = 0.30 + # If True, log every active MuJoCo contact at runtime (throttled to + # 100 ms intervals) with both body names. Surfaces hidden collisions + # the classifier filters out as `other`. + g1_runtime_full_contact_diagnostic: bool = True + # Null-space posture nudge weight. After each IK call, project + # `(nominal_qpos - ik_qpos)` onto the null space of the gripper-site + # Jacobian and add `weight * v_null` to the IK output. This biases + # the 1-DoF redundancy of the 7-DoF right arm toward a fixed nominal + # posture (`init_qpos["right_arm"]`) without changing the commanded + # TCP pose — preventing the wrist roll/yaw basin flips that + # appeared in run 27 as 19° one-frame joint jumps. 0.0 = disabled. + g1_runtime_null_space_weight: float = 0.3 + g1_runtime_null_space_damping: float = 1e-4 + g1_runtime_null_space_max_step_rad: float = 0.05 g1_grasp_min_vertical_axis_z: float = 0.75 g1_grasp_max_tcp_rot_deg: float = 120.0 - g1_ignore_flipped_grasps: bool = True - g1_grasp_inward_xy_offset: float = 0.006 + g1_ignore_flipped_grasps: bool = False + # The G1 right_grasp_site is at the physical pinch center and uses the + # MolmoSpaces parallel-jaw frame (+z forward, +y finger-opening). DROID + # grasp poses are stored in the Robotiq grasp_site frame which is also at + # the pinch center with the same axis convention, so DROID grasp poses + # apply to the G1 right_grasp_site directly. The inward / lateral / + # forward centering knobs below were empirically tuned against the prior + # rotated G1 grasp site and now distort otherwise-good poses; default + # them off and keep the levers available for follow-up tuning. + g1_grasp_inward_xy_offset: float = 0.0 g1_grasp_table_clearance: float = 0.065 - g1_center_grasp_lateral: bool = True + g1_center_grasp_lateral: bool = False g1_grasp_lateral_centering_scale: float = 1.0 g1_grasp_lateral_centering_max_offset: float = 0.02 + g1_center_grasp_forward: bool = False + g1_grasp_forward_centering_scale: float = 1.0 + g1_grasp_forward_centering_max_offset: float = 0.03 + g1_grasp_forward_centering_target_m: float = 0.0 g1_level_grasp_orientation: bool = False g1_grasp_level_max_tilt_deg: float = 35.0 g1_require_fingertip_pad_grasp_contact: bool = False g1_reject_non_fingertip_grasp_object_contact: bool = True g1_reject_grasp_table_contact: bool = True g1_reject_open_grasp_object_contact: bool = True + g1_allow_open_fingertip_pad_contact: bool = True g1_grasp_single_pad_contact_penalty: float = 0.5 g1_closed_grasp_quality_enabled: bool = True - g1_closed_grasp_settle_steps: int = 25 + g1_closed_grasp_settle_steps: int = 120 g1_closed_grasp_min_pad_geom_count: int = 2 g1_closed_grasp_max_object_shift_m: float = 0.05 g1_closed_grasp_penalty_per_missing_pad: float = 1.0 - g1_pregrasp_min_vertical_lift: float = 0.10 - g1_pregrasp_object_clearance: float = 0.12 - g1_place_travel_object_clearance: float = 0.13 + # With the corrected grasp-site frame, `pregrasp_z_offset` along the + # gripper's true forward axis already lifts pregrasp above the grasp by + # the configured offset for top-down approaches; the additional + # "minimum vertical lift" and "object clearance" margins below were + # over-padding under the prior misaligned site frame. + g1_pregrasp_min_vertical_lift: float = 0.0 + g1_pregrasp_object_clearance: float = 0.04 + # Carried bottle's bottom must stay this far above the bin's AABB top + # during the lift→preplace travel. Reduced from 0.13 — the original + # value was over-cautious and forced the right arm into a high-reach + # envelope (~24 cm above table) where the runtime IK frequently + # singular-fails. The bin walls are 9.5 cm tall, so 5 cm of clearance + # is plenty of margin during the slow held-object carry (4 cm/s). + g1_place_travel_object_clearance: float = 0.07 g1_held_object_speed: float = 0.04 g1_postgrasp_hold_duration: float = 0.4 g1_record_partial_attempt_on_no_full_grasp_candidate: bool = True diff --git a/molmo_spaces/configs/task_sampler_configs.py b/molmo_spaces/configs/task_sampler_configs.py index 559cd2ae..e9851d95 100644 --- a/molmo_spaces/configs/task_sampler_configs.py +++ b/molmo_spaces/configs/task_sampler_configs.py @@ -218,9 +218,20 @@ class UnitreeG1TabletopPickAndPlaceTaskSamplerConfig(PickAndPlaceTaskSamplerConf place_receptacle_name: str = "g1_place_bin" robot_base_xy: tuple[float, float] = (0.08, 0.0) pickup_workspace_center_xy: tuple[float, float] = (0.42, -0.20) - pickup_workspace_size_xy: tuple[float, float] = (0.16, 0.10) - place_workspace_center_xy: tuple[float, float] = (0.42, 0.08) - place_workspace_size_xy: tuple[float, float] = (0.10, 0.08) + # y range tightened from 0.10 to 0.06 so the bottle's far edge cannot + # crowd the bin under the constrained place workspace below. The + # earlier wider y range left less than the 15 mm + # pickup_receptacle_min_clearance at the unlucky end of the sample + # range, which silently dropped episodes. + pickup_workspace_size_xy: tuple[float, float] = (0.16, 0.06) + # The bin is 22 cm wide in y, so the bin center can sit only slightly + # to the right of the pelvis centerline before its physical extent + # crowds the pickup workspace. Constrain the bin to a tight band just + # past y=0 — close enough to the right arm to keep the carry path + # short, far enough from the pickup to leave a visible breathing-room + # gap (~5 cm worst case) between the bottle and the bin wall. + place_workspace_center_xy: tuple[float, float] = (0.42, 0.02) + place_workspace_size_xy: tuple[float, float] = (0.04, 0.02) place_receptacle_half_size_xy: tuple[float, float] = (0.14, 0.11) pickup_receptacle_min_clearance: float = 0.015 max_place_receptacle_sampling_attempts: int = 100 diff --git a/molmo_spaces/policy/solvers/object_manipulation/pick_and_place_planner_policy.py b/molmo_spaces/policy/solvers/object_manipulation/pick_and_place_planner_policy.py index bf088960..f55a9638 100644 --- a/molmo_spaces/policy/solvers/object_manipulation/pick_and_place_planner_policy.py +++ b/molmo_spaces/policy/solvers/object_manipulation/pick_and_place_planner_policy.py @@ -319,7 +319,10 @@ class UnitreeG1RightArmPickAndPlacePlannerPolicy(PickAndPlacePlannerPolicy): def reset(self, reset_retries: bool = True): self._g1_failure_diag_steps: list[dict] = [] self._g1_last_target_pose: np.ndarray | None = None + self._g1_site_frame_debug_logged = False self._g1_initial_tcp_pose = self._current_g1_tcp_pose() + self._log_g1_site_frame_debug_once("reset") + self._log_g1_frame_conventions("reset") self._g1_failure_diagnostics_logged = False self._g1_failure_reason_override: str | None = None self._g1_failure_phase_override: str | None = None @@ -329,8 +332,16 @@ def reset(self, reset_retries: bool = True): self._g1_selected_grasp_debug: dict | None = None self._g1_grasp_selector_debug: dict | None = None self._g1_last_ik_debug: dict | None = None - self._g1_failure_hold_action: NoopAction | None = None + self._g1_failure_hold_action: HoldCurrentControlAction | None = None self._g1_terminal_failure_pending = False + # Smoothed right-arm ctrl carried across runtime IK calls so the + # EMA filter in `_tcp_to_jp_fn` has a previous-value reference. + self._g1_right_arm_ctrl_smoothed: np.ndarray | None = None + # State for the per-frame runtime IK diagnostic (basin-flip + # detection + full-contact dump). Reset per episode. + self._g1_runtime_prev_arm_qpos: np.ndarray | None = None + self._g1_runtime_diag_step = 0 + self._g1_runtime_contact_last_log_t = -1e9 self._g1_grip_diag_samples: list[dict] = [] self._g1_grip_diag_last_log_time = -np.inf self._g1_grip_diag_last_key: tuple | None = None @@ -632,6 +643,10 @@ def _build_g1_candidate_target_poses( grasp_pose_world, pickup_obj, ) + grasp_pose_world = self._g1_center_grasp_forward( + grasp_pose_world, + pickup_obj, + ) pregrasp_pose = grasp_pose_world.copy() pregrasp_pose[:3, 3] -= self.policy_config.pregrasp_z_offset * pregrasp_pose[:3, 2] @@ -746,6 +761,33 @@ def _g1_center_grasp_lateral( pose[:3, 3] += lateral_offset * pose[:3, 1] return pose + def _g1_center_grasp_forward( + self, + grasp_pose_world: np.ndarray, + pickup_obj: MlSpacesObject, + ) -> np.ndarray: + if not getattr(self.policy_config, "g1_center_grasp_forward", False): + return grasp_pose_world + + pose = grasp_pose_world.copy() + object_pos = np.asarray(pickup_obj.position[:3], dtype=float) + object_in_grasp = pose[:3, :3].T @ (object_pos - pose[:3, 3]) + target_forward_offset = float( + getattr(self.policy_config, "g1_grasp_forward_centering_target_m", 0.0) + ) + forward_offset = float(object_in_grasp[2] - target_forward_offset) + forward_offset *= float( + getattr(self.policy_config, "g1_grasp_forward_centering_scale", 1.0) + ) + max_offset = float( + getattr(self.policy_config, "g1_grasp_forward_centering_max_offset", np.inf) + ) + if np.isfinite(max_offset): + forward_offset = float(np.clip(forward_offset, -max_offset, max_offset)) + + pose[:3, 3] += forward_offset * pose[:3, 2] + return pose + def _g1_level_grasp_orientation(self, grasp_pose_world: np.ndarray) -> np.ndarray: if not getattr(self.policy_config, "g1_level_grasp_orientation", False): return grasp_pose_world @@ -848,6 +890,7 @@ def _g1_pregrasp_debug( np.degrees(np.arccos(np.clip(abs(grasp_pose_world[2, 2]), -1.0, 1.0))) ), "object_local_y_in_grasp_m": float(object_in_grasp[1]), + "object_local_z_in_grasp_m": float(object_in_grasp[2]), "grasp_table_clearance_m": float(grasp_pose_world[2, 3] - table_top_z) if np.isfinite(table_top_z) else None, @@ -890,7 +933,11 @@ def _g1_eval_candidate_target_poses( qpos, base_pose, max_iter=250, + damping=float( + getattr(self.policy_config, "g1_selector_ik_damping", 1e-12) + ), ) + result = self._g1_accept_relaxed_candidate_ik_result(phase, result) phase_results[phase] = result if not result["success"]: return { @@ -910,6 +957,40 @@ def _g1_eval_candidate_target_poses( robot_view.base.pose = live_base_pose robot_view.set_qpos_dict(live_qpos) + def _g1_candidate_ik_thresholds_for_phase(self, phase: str) -> tuple[float, float]: + if phase == "pregrasp": + return ( + float(getattr(self.policy_config, "pregrasp_tcp_pos_err_threshold", 0.1)), + float(getattr(self.policy_config, "pregrasp_tcp_rot_err_threshold", np.inf)), + ) + return ( + float(self.policy_config.tcp_pos_err_threshold), + float(self.policy_config.tcp_rot_err_threshold), + ) + + def _g1_accept_relaxed_candidate_ik_result(self, phase: str, result: dict) -> dict: + if result["success"]: + return result + + final_qpos = result.get("final_qpos") + if final_qpos is None: + return result + + pos_threshold, rot_threshold = self._g1_candidate_ik_thresholds_for_phase(phase) + pos_err = float(result.get("final_pos_error_norm", np.inf)) + rot_err = float(result.get("final_rot_error_norm", np.inf)) + if pos_err > pos_threshold or rot_err > rot_threshold: + return result + + accepted = dict(result) + accepted["success"] = True + accepted["strict_success"] = False + accepted["accepted_with_g1_candidate_thresholds"] = True + accepted["qpos"] = self._copy_qpos_dict(final_qpos) + accepted["g1_candidate_pos_err_threshold"] = pos_threshold + accepted["g1_candidate_rot_err_threshold"] = rot_threshold + return accepted + def _g1_contact_quality_from_contacts(self, contacts: list[dict]) -> dict: robot_namespace = self.config.robot_config.robot_namespace object_contacts = [ @@ -1133,22 +1214,42 @@ def _g1_closed_grasp_contact_quality( base_pose: np.ndarray, pickup_obj: MlSpacesObject, ) -> dict: - robot_view = self.task.env.current_robot.robot_view + robot = self.task.env.current_robot + robot_view = robot.robot_view model = self.task.env.current_model data = self.task.env.current_data snapshot = self._g1_mujoco_state_snapshot() + # Without restoring this, the runtime would keep pinning the base to + # whatever pose we latched during the probe. + prior_pinned_base_pose = getattr(robot, "_pinned_base_pose", None) + if prior_pinned_base_pose is not None: + prior_pinned_base_pose = prior_pinned_base_pose.copy() settle_steps = max( int(getattr(self.policy_config, "g1_closed_grasp_settle_steps", 25)), 0, ) + gripper_mg_ids = set(robot_view.get_gripper_movegroup_ids()) try: robot_view.base.pose = base_pose + # Latch the probe's base pose so per-step re-pinning doesn't snap + # back to the live runtime base pose. + if hasattr(robot, "sync_pinned_base_pose"): + robot.sync_pinned_base_pose() robot_view.set_qpos_dict(qpos) + # Hold candidate qpos via position-servo ctrl on non-gripper move + # groups so the wrist doesn't drift back to the prior live ctrl + # target during the settle. + for mg_id in robot_view.move_group_ids(): + if mg_id in gripper_mg_ids or mg_id not in qpos: + continue + mg = robot_view.get_move_group(mg_id) + if mg.n_actuators == 0 or mg.n_actuators != mg.pos_dim: + continue + mg.ctrl = np.asarray(qpos[mg_id]) gripper_mg_id = robot_view.get_gripper_movegroup_ids()[0] gripper = robot_view.get_gripper(gripper_mg_id) gripper.set_gripper_ctrl_open(False) - gripper.joint_pos = gripper.ctrl.copy() data.qvel[:] = 0.0 data.qacc[:] = 0.0 data.qacc_warmstart[:] = 0.0 @@ -1156,6 +1257,10 @@ def _g1_closed_grasp_contact_quality( object_start_pos = pickup_obj.position.copy() for _ in range(settle_steps): mujoco.mj_step(model, data) + # Re-enforce pinned base + locked left arm/hand between + # steps; right arm is held via the ctrl assignment above. + if hasattr(robot, "apply_initial_state_overrides"): + robot.apply_initial_state_overrides() mujoco.mj_forward(model, data) object_shift_m = float(np.linalg.norm(pickup_obj.position - object_start_pos)) contacts = self._collect_contact_diagnostics() @@ -1174,6 +1279,8 @@ def _g1_closed_grasp_contact_quality( return quality finally: self._g1_restore_mujoco_state(snapshot) + if hasattr(robot, "_pinned_base_pose"): + robot._pinned_base_pose = prior_pinned_base_pose def _g1_candidate_grasp_contact_quality( self, @@ -1192,6 +1299,7 @@ def _g1_grasp_contact_quality_failure(self, quality: dict) -> str | None: if ( getattr(self.policy_config, "g1_reject_open_grasp_object_contact", True) and quality["object_contact_count"] > 0 + and not self._g1_is_allowed_open_grasp_object_contact(quality) ): return "open_grasp_object_contact" if ( @@ -1211,6 +1319,21 @@ def _g1_grasp_contact_quality_failure(self, quality: dict) -> str | None: return "missing_fingertip_pad_contact" return None + def _g1_is_allowed_open_grasp_object_contact(self, quality: dict) -> bool: + if not getattr(self.policy_config, "g1_allow_open_fingertip_pad_contact", False): + return False + if not quality["non_pad_object_contacts"]: + return quality["pad_geom_count"] > 0 + + min_pad_geom_count = int( + getattr(self.policy_config, "g1_closed_grasp_min_pad_geom_count", 2) + ) + return ( + quality["pad_geom_count"] >= min_pad_geom_count + and bool(quality["finger_link_object_contacts"]) + and not quality["other_robot_object_contacts"] + ) + def _g1_grasp_quality_rejection_key(self, quality: dict) -> str: if quality["closed_probe"]: return "closed_grasp_quality" @@ -1353,7 +1476,8 @@ def _select_g1_tabletop_grasp_pose( dist_tcp = np.linalg.inv(tcp_pose_world) @ grasp_poses_world dists_tcp_p = np.linalg.norm(dist_tcp[:, :3, 3], axis=1) dists_tcp_o = R.from_matrix(dist_tcp[:, :3, :3]).magnitude() * 180 / np.pi - dists_up = grasp_poses_world[:, 2, 2] + grasp_forward_axis_z = grasp_poses_world[:, 2, 2] + grasp_down_alignment = -grasp_forward_axis_z dists_com = np.linalg.norm( (np.linalg.inv(object_pose) @ grasp_poses_world)[:, :3, 3], axis=1, @@ -1361,7 +1485,7 @@ def _select_g1_tabletop_grasp_pose( base_cost = ( self.policy_config.grasp_pos_cost_weight * dists_tcp_p + self.policy_config.grasp_rot_cost_weight * dists_tcp_o - + self.policy_config.grasp_vertical_cost_weight * dists_up + + self.policy_config.grasp_vertical_cost_weight * grasp_forward_axis_z + self.policy_config.grasp_com_dist_cost_weight * dists_com ) @@ -1381,7 +1505,7 @@ def _select_g1_tabletop_grasp_pose( getattr(self.policy_config, "g1_grasp_min_vertical_axis_z", -np.inf) ) vertical_eligible_mask = eligible_mask.copy() - eligible_mask &= dists_up >= min_vertical_axis_z + eligible_mask &= grasp_down_alignment >= min_vertical_axis_z vertical_axis_mask = vertical_eligible_mask & ~eligible_mask eligible_ids = np.where(eligible_mask)[0] @@ -1508,7 +1632,10 @@ def _select_g1_tabletop_grasp_pose( "base_cost": float(base_cost[grasp_idx]), "tcp_pos_cost_m": float(dists_tcp_p[grasp_idx]), "tcp_rot_cost_deg": float(dists_tcp_o[grasp_idx]), - "vertical_axis_z": float(dists_up[grasp_idx]), + "vertical_axis_z": float(grasp_forward_axis_z[grasp_idx]), + "grasp_down_alignment": float( + grasp_down_alignment[grasp_idx] + ), "com_dist_cost_m": float(dists_com[grasp_idx]), "progress_penalty": progress_penalty, "is_flipped": bool(grasp_idx >= original_grasp_count), @@ -1589,7 +1716,8 @@ def _select_g1_tabletop_grasp_pose( "base_cost": float(base_cost[grasp_idx]), "tcp_pos_cost_m": float(dists_tcp_p[grasp_idx]), "tcp_rot_cost_deg": float(dists_tcp_o[grasp_idx]), - "vertical_axis_z": float(dists_up[grasp_idx]), + "vertical_axis_z": float(grasp_forward_axis_z[grasp_idx]), + "grasp_down_alignment": float(grasp_down_alignment[grasp_idx]), "com_dist_cost_m": float(dists_com[grasp_idx]), "is_flipped": bool(grasp_idx >= original_grasp_count), "grasp_pad_contact_count": int(grasp_quality["pad_contact_count"]), @@ -1725,6 +1853,257 @@ def _failure_diagnostics_enabled(self) -> bool: def _g1_ik_debug_enabled(self) -> bool: return bool(getattr(self.policy_config, "g1_ik_debug", False)) + @staticmethod + def _unit_vector(vec: np.ndarray) -> list[float] | None: + norm = float(np.linalg.norm(vec)) + if norm < 1e-8: + return None + return np.round(vec / norm, 4).tolist() + + def _g1_named_site_id(self, site_name: str) -> int: + model = self.task.env.current_model + namespace = self.task.env.current_robot.namespace + for candidate in (f"{namespace}{site_name}", site_name): + site_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE, candidate) + if site_id >= 0: + return int(site_id) + raise ValueError(f"Expected MuJoCo site `{site_name}`") + + def _g1_named_geom_id(self, geom_name: str) -> int: + model = self.task.env.current_model + namespace = self.task.env.current_robot.namespace + for candidate in (f"{namespace}{geom_name}", geom_name): + geom_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, candidate) + if geom_id >= 0: + return int(geom_id) + raise ValueError(f"Expected MuJoCo geom `{geom_name}`") + + def _g1_site_frame_debug(self) -> dict: + data = self.task.env.current_data + wrist_site_id = self._g1_named_site_id("right_wrist_site") + grasp_site_id = self._g1_named_site_id("right_grasp_site") + pad_1_geom_id = self._g1_named_geom_id("right_dex1_fingertip_pad_1") + pad_2_geom_id = self._g1_named_geom_id("right_dex1_fingertip_pad_2") + + wrist_pos = data.site_xpos[wrist_site_id].copy() + grasp_pos = data.site_xpos[grasp_site_id].copy() + grasp_xmat = data.site_xmat[grasp_site_id].reshape(3, 3).copy() + grasp_axes = { + "+x": np.round(grasp_xmat[:, 0], 4).tolist(), + "+y": np.round(grasp_xmat[:, 1], 4).tolist(), + "+z": np.round(grasp_xmat[:, 2], 4).tolist(), + } + wrist_to_grasp = grasp_pos - wrist_pos + pad_axis = data.geom_xpos[pad_2_geom_id] - data.geom_xpos[pad_1_geom_id] + wrist_to_grasp_unit = self._unit_vector(wrist_to_grasp) + pad_axis_unit = self._unit_vector(pad_axis) + + return { + "right_grasp_site_pos": np.round(grasp_pos, 4).tolist(), + "right_grasp_site_axes_world": grasp_axes, + "wrist_to_grasp_unit": wrist_to_grasp_unit, + "finger_pad_1_to_2_unit": pad_axis_unit, + "dot_site_x_with_wrist_to_grasp": float( + np.dot(grasp_xmat[:, 0], wrist_to_grasp / np.linalg.norm(wrist_to_grasp)) + ), + "dot_site_y_with_finger_pad_axis": float( + np.dot(grasp_xmat[:, 1], pad_axis / np.linalg.norm(pad_axis)) + ), + "dot_site_z_with_wrist_to_grasp": float( + np.dot(grasp_xmat[:, 2], wrist_to_grasp / np.linalg.norm(wrist_to_grasp)) + ), + } + + def _log_g1_site_frame_debug_once(self, label: str) -> None: + if not self._g1_ik_debug_enabled() or self._g1_site_frame_debug_logged: + return + self._g1_site_frame_debug_logged = True + log.info("[G1_SITE] label=%s debug=%s", label, self._g1_site_frame_debug()) + + def _g1_apply_null_space_posture_nudge( + self, gripper_mg_id: str, arm_ctrl: np.ndarray + ) -> np.ndarray: + """Bias `arm_ctrl` toward the nominal right-arm posture along the + null space of the gripper-site Jacobian. Returns a modified ctrl + whose forward-kinematic TCP pose is (to first order) identical to + the input. The redundant 1 DoF of the 7-DoF arm is steered toward + the configured nominal, preventing the wrist roll/yaw basin flips + seen in the runtime IK output across consecutive frames. + """ + weight = float( + getattr(self.policy_config, "g1_runtime_null_space_weight", 0.0) + ) + if weight <= 0.0: + return arm_ctrl + + try: + nominal = self.config.robot_config.init_qpos.get("right_arm") + except AttributeError: + return arm_ctrl + if nominal is None: + return arm_ctrl + nominal_arr = np.asarray(nominal, dtype=float) + if nominal_arr.shape != arm_ctrl.shape: + return arm_ctrl + + # Jacobian of the gripper site w.r.t. right_arm joints, at the + # current runtime qpos (close enough to the IK-solved qpos for a + # first-order null-space projection). + J = self.robot_view.get_jacobian(gripper_mg_id, ["right_arm"]) + if J.shape[0] != 6 or J.shape[1] != arm_ctrl.shape[0]: + return arm_ctrl + + # Damped null-space projector N = I - J⁺ J where J⁺ uses damping + # for numerical stability near singular configurations. + damping = float( + getattr(self.policy_config, "g1_runtime_null_space_damping", 1e-4) + ) + try: + J_pinv = J.T @ np.linalg.solve( + J @ J.T + damping * np.eye(6), np.eye(6) + ) + except np.linalg.LinAlgError: + return arm_ctrl + N = np.eye(arm_ctrl.shape[0]) - J_pinv @ J + + delta = N @ (nominal_arr - arm_ctrl) + max_step = float( + getattr(self.policy_config, "g1_runtime_null_space_max_step_rad", 0.05) + ) + scaled = weight * delta + per_joint_cap = float(max_step) if np.isfinite(max_step) else np.inf + scaled = np.clip(scaled, -per_joint_cap, per_joint_cap) + + adjusted = arm_ctrl + scaled + # Clamp to actuator ctrlrange (joint limits) so the nudge cannot + # produce an out-of-range command. + arm_mg = self.robot_view.get_move_group("right_arm") + if hasattr(arm_mg, "ctrl_limits"): + limits = arm_mg.ctrl_limits + adjusted = np.clip(adjusted, limits[:, 0], limits[:, 1]) + return adjusted + + def _g1_runtime_diagnostic( + self, + target_pose: np.ndarray, + ik_succeeded: bool, + new_arm_ctrl: np.ndarray | None = None, + ) -> None: + """Per-frame runtime diagnostic: basin-flip detection on the + right-arm **IK output ctrl** and a full contact dump. + + Compare the *commanded* qpos (= IK output → ctrl) across + consecutive successful IK calls. The actual joint pos lags via + servo dynamics and would not reveal an IK-output jump. + """ + self._g1_runtime_diag_step += 1 + step = self._g1_runtime_diag_step + model = self.task.env.current_model + data = self.task.env.current_data + phase = self.get_phase() + threshold = float( + getattr(self.policy_config, "g1_runtime_basin_flip_threshold_rad", 0.30) + ) + force_dump = False + + # 1) Basin-flip detection on the IK *output* ctrl + if ik_succeeded and new_arm_ctrl is not None: + if self._g1_runtime_prev_arm_qpos is not None: + delta = new_arm_ctrl - self._g1_runtime_prev_arm_qpos + abs_delta = np.abs(delta) + max_delta = float(abs_delta.max()) + if max_delta > threshold: + worst = int(abs_delta.argmax()) + force_dump = True + log.warning( + "[G1_BASIN_FLIP] step=%d phase=%s max_delta=%.4frad " + "on_joint=%d prev_ctrl=%s new_ctrl=%s per_joint_delta=%s " + "target_pos=%s", + step, + phase, + max_delta, + worst, + np.round(self._g1_runtime_prev_arm_qpos, 4).tolist(), + np.round(new_arm_ctrl, 4).tolist(), + np.round(delta, 4).tolist(), + np.round(target_pose[:3, 3], 4).tolist(), + ) + self._g1_runtime_prev_arm_qpos = new_arm_ctrl.copy() + + # 2) Full contact dump (throttled to ~10 Hz; also fires immediately + # on basin-flip) + if not getattr( + self.policy_config, "g1_runtime_full_contact_diagnostic", False + ): + return + sim_t = float(data.time) + if (sim_t - self._g1_runtime_contact_last_log_t < 0.1) and not force_dump: + return + self._g1_runtime_contact_last_log_t = sim_t + + def geom_label(geom_id: int) -> str: + gn = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_GEOM, geom_id) + if gn: + return gn + # Fall back to body name + geom id so unnamed structural + # collision geoms (e.g., the bin's hidden walls) are + # identifiable in the log. + body_id = int(model.geom_bodyid[geom_id]) + bn = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_BODY, body_id) or "?" + return f"{bn}#geom{geom_id}" + + contacts = [] + for i in range(data.ncon): + con = data.contact[i] + g1n = geom_label(con.geom1) + g2n = geom_label(con.geom2) + contacts.append((float(con.dist), g1n, g2n)) + contacts.sort(key=lambda c: c[0]) # most-interpenetrating first + formatted = ";".join( + f"{g1}↔{g2}@{d:+.4f}" for d, g1, g2 in contacts[:12] + ) + log.info( + "[G1_CONTACTS] step=%d phase=%s t=%.3f n=%d %s", + step, + phase, + sim_t, + data.ncon, + formatted if formatted else "(none)", + ) + + def _log_g1_frame_conventions(self, label: str) -> None: + """Unconditional convention check at reset. Confirms: + - Pelvis frame: +x forward, +y left, +z up (MolmoSpaces robot base) + - Right grasp_site: +z forward, +y finger-opening (parallel-jaw) + Logs the world-frame axes of both so any deviation from convention + is visible at a glance. + """ + try: + model = self.task.env.current_model + data = self.task.env.current_data + pelvis_id = mujoco.mj_name2id( + model, + mujoco.mjtObj.mjOBJ_BODY, + f"{self.robot_view._namespace}pelvis", + ) + if pelvis_id < 0: + return + pelvis_xmat = data.xmat[pelvis_id].reshape(3, 3) + site_debug = self._g1_site_frame_debug() + log.info( + "[G1_CONV] label=%s pelvis_world_axes={+x:%s, +y:%s, +z:%s} " + "grasp_axes=%s dot_z_forward=%.3f dot_y_finger=%.3f", + label, + np.round(pelvis_xmat[:, 0], 3).tolist(), + np.round(pelvis_xmat[:, 1], 3).tolist(), + np.round(pelvis_xmat[:, 2], 3).tolist(), + site_debug["right_grasp_site_axes_world"], + site_debug["dot_site_z_with_wrist_to_grasp"], + site_debug["dot_site_y_with_finger_pad_axis"], + ) + except Exception as exc: # diagnostic only — never block the run + log.debug("[G1_CONV] frame-convention diagnostic failed: %r", exc) + def _names_for_contact(self, geom_id: int) -> dict[str, str | int]: model = self.task.env.current_model body_id = model.geom_bodyid[geom_id] @@ -2037,7 +2416,14 @@ def _handle_failure(self) -> dict[str, np.ndarray | bool]: self._g1_terminal_failure_pending = True if self._g1_failure_hold_action is None: log.info("[G1_DIAG] holding failed attempt for %.2fs before done", hold_duration) - self._g1_failure_hold_action = NoopAction(self.robot_view, hold_duration) + # Hold the current commanded controls (including a closed + # gripper carrying a held object) instead of NoopAction, which + # can infer an open gripper when the held object keeps the + # fingers apart and cause a visible drop at the end of a + # failed attempt. + self._g1_failure_hold_action = HoldCurrentControlAction( + self.robot_view, hold_duration + ) if not self._g1_failure_hold_action.execute(): return self._g1_failure_hold_action.get_current_action() @@ -2079,19 +2465,47 @@ def _tcp_to_jp_fn(self, mg_id: str, target_pose: np.ndarray) -> dict[str, np.nda ["right_arm"], self.robot_view.get_qpos_dict(), self.robot_view.base.pose, + damping=float(getattr(self.policy_config, "g1_runtime_ik_damping", 1e-12)), + dt=float(getattr(self.policy_config, "g1_runtime_ik_dt", 1.0)), ) self._g1_last_target_pose = target_pose.copy() action = self.robot_view.get_ctrl_dict() if jp is not None: self.sequential_ik_failures = 0 - action["right_arm"] = jp["right_arm"] + new_arm_ctrl = np.asarray(jp["right_arm"], dtype=float) + # Null-space posture nudge: bias the redundant 1 DoF of the 7-DoF + # right arm toward a fixed nominal posture without changing the + # commanded TCP pose. Prevents the wrist-orientation basin flips + # caught by the [G1_BASIN_FLIP] diagnostic. + new_arm_ctrl = self._g1_apply_null_space_posture_nudge( + mg_id, new_arm_ctrl + ) + # Exponential-moving-average smoothing across frames to kill + # high-frequency IK output jitter that an underdamped servo + # would otherwise overshoot. + alpha = float( + getattr(self.policy_config, "g1_runtime_ctrl_smoothing", 0.0) + ) + if 0.0 < alpha < 1.0 and self._g1_right_arm_ctrl_smoothed is not None: + new_arm_ctrl = ( + alpha * self._g1_right_arm_ctrl_smoothed + + (1.0 - alpha) * new_arm_ctrl + ) + self._g1_right_arm_ctrl_smoothed = new_arm_ctrl.copy() + action["right_arm"] = new_arm_ctrl self._record_failure_diagnostic_step(target_pose, True) + self._g1_runtime_diagnostic( + target_pose, ik_succeeded=True, new_arm_ctrl=new_arm_ctrl + ) else: self.sequential_ik_failures += 1 log.info(f"IK failed, holding current position, fails:{self.sequential_ik_failures}") self._record_failure_diagnostic_step(target_pose, False) self._run_g1_ik_debug(mg_id, target_pose) + self._g1_runtime_diagnostic( + target_pose, ik_succeeded=False, new_arm_ctrl=None + ) if self.sequential_ik_failures >= self.policy_config.max_sequential_ik_failures: log.info("Too many sequential IK failures, triggering retry.") return self._handle_failure() @@ -2112,6 +2526,7 @@ def _run_diagnostic_ik( self.robot_view.get_qpos_dict(), self.robot_view.base.pose, max_iter=max_iter, + damping=float(getattr(self.policy_config, "g1_selector_ik_damping", 1e-12)), ) def _candidate_pose_for_phase(self, candidate_grasp_pose: np.ndarray) -> np.ndarray: @@ -2344,6 +2759,9 @@ def check_feasible_ik(self, pose: np.ndarray) -> bool: kinematics = self.task.env.current_robot.kinematics gripper_mg_id = robot_view.get_gripper_movegroup_ids()[0] + selector_damping = float( + getattr(self.policy_config, "g1_selector_ik_damping", 1e-12) + ) if pose.ndim > 2: return np.array( [ @@ -2354,6 +2772,7 @@ def check_feasible_ik(self, pose: np.ndarray) -> bool: robot_view.get_qpos_dict(), robot_view.base.pose, max_iter=100, + damping=selector_damping, ) is not None for single_pose in pose @@ -2369,5 +2788,6 @@ def check_feasible_ik(self, pose: np.ndarray) -> bool: robot_view.get_qpos_dict(), robot_view.base.pose, max_iter=100, + damping=selector_damping, ) return jp_dict is not None diff --git a/scripts/assets/prepare_unitree_g1.py b/scripts/assets/prepare_unitree_g1.py index a633eb0b..76f9214f 100644 --- a/scripts/assets/prepare_unitree_g1.py +++ b/scripts/assets/prepare_unitree_g1.py @@ -24,12 +24,18 @@ OUTPUT_XML_NAME = "model.xml" DEFAULT_PELVIS_HEIGHT_M = 0.793 DEX1_GRASP_SITE_POS = [0.148, 0.0, 0.0] -DEX1_FINGERTIP_PAD_SIZE = [0.02, 0.01, 0.045] +# Rotate Unitree's Dex1.1 wrist frame into MolmoSpaces' parallel-jaw gripper frame: +# local +z points forward through the gripper, local +y spans the finger opening axis. +DEX1_GRASP_SITE_QUAT = [0.7071067811865476, 0.0, 0.7071067811865475, 0.0] +DEX1_KINEMATICS_SITE_SIZE = [0.04] +DEX1_DEBUG_AXIS_LENGTH = 0.08 +DEX1_DEBUG_AXIS_RADIUS = 0.004 +DEX1_FINGERTIP_PAD_SIZE = [0.045, 0.01, 0.02] DEX1_FINGERTIP_PAD_POSITIONS = { "dex1_finger_link_1": [0.1065, -0.0285, 0.0], "dex1_finger_link_2": [0.1065, 0.0285, 0.0], } -DEX1_FINGERTIP_PAD_FRICTION = [12.0, 0.2, 0.02] +DEX1_FINGERTIP_PAD_FRICTION = [20.0, 0.4, 0.04] DEX1_FINGERTIP_PAD_CONDIM = 6 DEX1_FINGERTIP_PAD_GROUP = 4 DEX1_HAND_FORCE_LIMIT_MULTIPLIER = 3.0 @@ -51,7 +57,14 @@ GAIN_BY_JOINT_GROUP = { "leg": (120.0, 12.0), "waist": (80.0, 8.0), - "arm": (60.0, 6.0), + # Arm damping doubled from the original (60, 6) to (60, 12). Keeping + # stiffness at 60 preserves the "soft" approach behavior — at 100, + # the arm was tracking pregrasp/grasp targets stiffly enough that + # the invisible Dex1.1 fingertip pad collision geoms shoved the + # pickup object aside during descent (run 24 lost the bottle 11 cm + # to the right of its sampled pose). Higher damping reduces servo + # ringing without making the arm push harder on tracking errors. + "arm": (60.0, 12.0), "hand": (80.0, 8.0), } @@ -86,6 +99,11 @@ def parse_args() -> argparse.Namespace: action="store_true", help="Overwrite an existing output directory.", ) + parser.add_argument( + "--add-debug-grasp-site-axes", + action="store_true", + help="Add non-collision RGB visual axes to the right grasp site for temporary debugging.", + ) return parser.parse_args() @@ -191,7 +209,10 @@ def _add_position_actuators(spec: mujoco.MjSpec, joints: list[JointSpec]) -> Non ) -def _add_kinematics_sites(spec: mujoco.MjSpec) -> None: +def _add_kinematics_sites( + spec: mujoco.MjSpec, + add_debug_grasp_site_axes: bool = False, +) -> None: for side in ("left", "right"): wrist_body = spec.body(f"{side}_wrist_yaw_link") if wrist_body is None: @@ -199,15 +220,67 @@ def _add_kinematics_sites(spec: mujoco.MjSpec) -> None: wrist_body.add_site( name=f"{side}_wrist_site", pos=[0.0, 0.0, 0.0], - size=[0.015], + size=DEX1_KINEMATICS_SITE_SIZE, rgba=[0.1, 0.4, 1.0, 1.0], ) wrist_body.add_site( name=f"{side}_grasp_site", pos=DEX1_GRASP_SITE_POS, - size=[0.015], + quat=DEX1_GRASP_SITE_QUAT, + size=DEX1_KINEMATICS_SITE_SIZE, rgba=[0.1, 1.0, 0.4, 1.0], ) + if side == "right" and add_debug_grasp_site_axes: + axis_specs = ( + ( + "x", + [0.0, 0.0, -DEX1_DEBUG_AXIS_LENGTH / 2], + [ + DEX1_DEBUG_AXIS_LENGTH / 2, + DEX1_DEBUG_AXIS_RADIUS, + DEX1_DEBUG_AXIS_RADIUS, + ], + [1.0, 0.0, 0.0, 1.0], + ), + ( + "y", + [0.0, DEX1_DEBUG_AXIS_LENGTH / 2, 0.0], + [ + DEX1_DEBUG_AXIS_RADIUS, + DEX1_DEBUG_AXIS_LENGTH / 2, + DEX1_DEBUG_AXIS_RADIUS, + ], + [0.0, 1.0, 0.0, 1.0], + ), + ( + "z", + [DEX1_DEBUG_AXIS_LENGTH / 2, 0.0, 0.0], + [ + DEX1_DEBUG_AXIS_RADIUS, + DEX1_DEBUG_AXIS_RADIUS, + DEX1_DEBUG_AXIS_LENGTH / 2, + ], + [0.1, 0.3, 1.0, 1.0], + ), + ) + for axis_name, axis_offset, axis_size, rgba in axis_specs: + axis_pos = [ + DEX1_GRASP_SITE_POS[0] + axis_offset[0], + DEX1_GRASP_SITE_POS[1] + axis_offset[1], + DEX1_GRASP_SITE_POS[2] + axis_offset[2], + ] + wrist_body.add_geom( + name=f"{side}_grasp_site_debug_axis_{axis_name}", + type=mujoco.mjtGeom.mjGEOM_BOX, + pos=axis_pos, + quat=DEX1_GRASP_SITE_QUAT, + size=axis_size, + contype=0, + conaffinity=0, + density=0.0, + group=5, + rgba=rgba, + ) def _add_dex1_fingertip_contact_pads(spec: mujoco.MjSpec) -> None: @@ -239,6 +312,7 @@ def prepare_unitree_g1( unitree_urdf_root: str | os.PathLike[str], output_dir: str | os.PathLike[str], force: bool = False, + add_debug_grasp_site_axes: bool = False, ) -> Path: source_root = Path(unitree_urdf_root).expanduser().resolve() urdf_path = source_root / SOURCE_URDF_REL_PATH @@ -264,7 +338,7 @@ def prepare_unitree_g1( raise ValueError("Expected root body `pelvis` in Unitree G1 URDF") pelvis.pos = [0.0, 0.0, DEFAULT_PELVIS_HEIGHT_M] pelvis.add_freejoint(name="floating_base_joint") - _add_kinematics_sites(spec) + _add_kinematics_sites(spec, add_debug_grasp_site_axes) _add_dex1_fingertip_contact_pads(spec) _add_position_actuators(spec, joints) @@ -285,7 +359,12 @@ def main() -> None: if args.unitree_urdf_root is None: raise SystemExit("Set UNITREE_URDF_ROOT or pass --unitree-urdf-root") - output_xml_path = prepare_unitree_g1(args.unitree_urdf_root, args.output_dir, args.force) + output_xml_path = prepare_unitree_g1( + args.unitree_urdf_root, + args.output_dir, + args.force, + args.add_debug_grasp_site_axes, + ) print(f"Wrote Unitree G1 Dex1.1 model: {output_xml_path}") From 43e38891f156c48f29de26dfef769865778a25df Mon Sep 17 00:00:00 2001 From: Abhishek Parmar Date: Wed, 13 May 2026 01:49:17 +0800 Subject: [PATCH 13/13] fix: remove inverted G1 vertical-axis filter; shorten pregrasp setback MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The G1 tabletop grasp selector had an inverted vertical-axis filter: `dists_up >= 0.75` required the gripper +z axis to point UP in world, which rejected top-down grasps under the +z=forward convention. The score weight already biased toward top-down grasps (low dists_up), so the selector and score disagreed — the filter rejected what the score preferred. The remaining "eligible" candidates were upside-down grasps that put the wrist at the right arm's kinematic edge, so runtime IK failed. Drop the filter entirely. The shared compute_grasp_pose pipeline used by Franka / RBY / YAM / RUM has no hard vertical-axis rejection — only a soft vertical_cost_weight in the score. G1 now matches that pattern. Also reduce pregrasp_z_offset from 0.12 m to 0.06 m on the G1 right-arm pick-and-place config: the 12 cm setback along the approach axis was pushing the pregrasp pose past the right arm's reach envelope. The 2b7a4cf runtime additions (null-space posture nudge, basin-flip detector, EMA ctrl smoothing scaffold, G1-specific IK damping overrides, grasp-centering knobs, full-contact diagnostics, [G1_*] logging, closed-grasp probe sync) were all bandaids compensating for the filter bug. With the root cause fixed they're unnecessary; revert to the 9fb3afc baseline policy + shared TCPMove + shared DLS IK. The asset-side fixes from 2b7a4cf are kept: grasp-site quat for the +z-forward README convention, fingertip pad orientation/friction, arm gains, hand force-limit multiplier. Task-sampler workspace tightening is also kept. Component tests pass; integration smoke test (skipped in CI without UNITREE_URDF_ROOT) updated to assert the new pickup/place workspace sizes and pregrasp offset. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../component_tests/test_unitree_g1.py | 126 +---- molmo_spaces/configs/policy_configs.py | 95 +--- .../pick_and_place_planner_policy.py | 445 +----------------- 3 files changed, 37 insertions(+), 629 deletions(-) diff --git a/mlspaces_tests/component_tests/test_unitree_g1.py b/mlspaces_tests/component_tests/test_unitree_g1.py index 6aea149b..397b9804 100644 --- a/mlspaces_tests/component_tests/test_unitree_g1.py +++ b/mlspaces_tests/component_tests/test_unitree_g1.py @@ -27,7 +27,6 @@ DEX1_FINGERTIP_PAD_POSITIONS, DEX1_FINGERTIP_PAD_SIZE, DEX1_GRASP_SITE_POS, - DEX1_GRASP_SITE_QUAT, DEX1_HAND_FORCE_LIMIT_MULTIPLIER, LEFT_ARM_STOW_QPOS, LEFT_HAND_OPEN_QPOS, @@ -124,7 +123,6 @@ def test_unitree_g1_grasp_contact_filter_requires_right_fingertip_pad(): policy.policy_config = SimpleNamespace( g1_reject_grasp_table_contact=True, g1_reject_open_grasp_object_contact=True, - g1_allow_open_fingertip_pad_contact=True, g1_closed_grasp_min_pad_geom_count=2, g1_closed_grasp_max_object_shift_m=0.05, ) @@ -164,12 +162,9 @@ def test_unitree_g1_grasp_contact_filter_requires_right_fingertip_pad(): "root1_name": "robot_0/pelvis", } assert not policy._g1_is_allowed_grasp_object_contact(wrist_contact) - assert not policy._g1_grasp_contact_quality_failure( - policy._g1_contact_quality_from_contacts([pad_contact]) - ) assert ( policy._g1_grasp_contact_quality_failure( - policy._g1_contact_quality_from_contacts([finger_link_contact]) + policy._g1_contact_quality_from_contacts([pad_contact]) ) == "open_grasp_object_contact" ) @@ -217,17 +212,7 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): for site_name in KINEMATICS_SITE_NAMES: assert model.site(site_name).id >= 0 for site_name in ("left_grasp_site", "right_grasp_site"): - site_id = model.site(site_name).id - np.testing.assert_allclose(model.site_pos[site_id], DEX1_GRASP_SITE_POS) - np.testing.assert_allclose(model.site_quat[site_id], DEX1_GRASP_SITE_QUAT) - assert ( - mujoco.mj_name2id( - model, - mujoco.mjtObj.mjOBJ_GEOM, - "right_grasp_site_debug_axis_x", - ) - == -1 - ) + np.testing.assert_allclose(model.site_pos[model.site(site_name).id], DEX1_GRASP_SITE_POS) for side in ("left", "right"): for finger_id, finger_suffix in (("1", "dex1_finger_link_1"), ("2", "dex1_finger_link_2")): pad = model.geom(f"{side}_dex1_fingertip_pad_{finger_id}") @@ -267,12 +252,6 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): assert scene_model.qpos0[scene_base_qposadr + 2] == pytest.approx(0.793) mujoco.mj_forward(scene_model, scene_data) assert not _has_deep_robot_floor_penetration(scene_model, scene_data) - wrist_site_id = scene_model.site("robot_0/right_wrist_site").id - grasp_site_id = scene_model.site("robot_0/right_grasp_site").id - wrist_to_grasp = scene_data.site_xpos[grasp_site_id] - scene_data.site_xpos[wrist_site_id] - wrist_to_grasp /= np.linalg.norm(wrist_to_grasp) - grasp_xmat = scene_data.site_xmat[grasp_site_id].reshape(3, 3) - np.testing.assert_allclose(grasp_xmat[:, 2], wrist_to_grasp, atol=1e-6) view = config.robot_view_factory(scene_data, config.robot_namespace) expected_counts = { @@ -396,44 +375,31 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): assert tabletop_datagen_config.policy_config.g1_grasp_candidate_limit == 256 assert tabletop_datagen_config.policy_config.g1_grasp_ik_eval_limit == 256 assert tabletop_datagen_config.policy_config.g1_grasp_require_all_pick_place_phases - assert tabletop_datagen_config.policy_config.g1_grasp_joint_margin_weight == 1.5 - assert tabletop_datagen_config.policy_config.g1_grasp_joint_motion_weight == 1.0 + assert tabletop_datagen_config.policy_config.g1_grasp_joint_margin_weight == 0.5 + assert tabletop_datagen_config.policy_config.g1_grasp_joint_motion_weight == 0.25 assert tabletop_datagen_config.policy_config.g1_grasp_topdown_weight == 1.0 - assert tabletop_datagen_config.policy_config.g1_runtime_ik_damping == 5e-3 - assert tabletop_datagen_config.policy_config.g1_selector_ik_damping == 1e-4 - assert tabletop_datagen_config.policy_config.g1_runtime_ik_dt == 0.5 - assert tabletop_datagen_config.policy_config.g1_runtime_ctrl_smoothing == 0.0 - assert tabletop_datagen_config.policy_config.g1_grasp_min_vertical_axis_z == 0.75 assert tabletop_datagen_config.policy_config.g1_grasp_max_tcp_rot_deg == 120.0 - assert not tabletop_datagen_config.policy_config.g1_ignore_flipped_grasps - assert tabletop_datagen_config.policy_config.g1_grasp_inward_xy_offset == 0.0 + assert tabletop_datagen_config.policy_config.g1_ignore_flipped_grasps + assert tabletop_datagen_config.policy_config.g1_grasp_inward_xy_offset == 0.006 assert tabletop_datagen_config.policy_config.g1_grasp_table_clearance == 0.065 - assert not tabletop_datagen_config.policy_config.g1_center_grasp_lateral + assert tabletop_datagen_config.policy_config.g1_center_grasp_lateral assert tabletop_datagen_config.policy_config.g1_grasp_lateral_centering_scale == 1.0 assert tabletop_datagen_config.policy_config.g1_grasp_lateral_centering_max_offset == 0.02 - assert not tabletop_datagen_config.policy_config.g1_center_grasp_forward - assert tabletop_datagen_config.policy_config.g1_grasp_forward_centering_scale == 1.0 - assert tabletop_datagen_config.policy_config.g1_grasp_forward_centering_max_offset == 0.03 - assert tabletop_datagen_config.policy_config.g1_grasp_forward_centering_target_m == 0.0 assert not tabletop_datagen_config.policy_config.g1_level_grasp_orientation assert tabletop_datagen_config.policy_config.g1_grasp_level_max_tilt_deg == 35.0 assert not tabletop_datagen_config.policy_config.g1_require_fingertip_pad_grasp_contact assert tabletop_datagen_config.policy_config.g1_reject_non_fingertip_grasp_object_contact assert tabletop_datagen_config.policy_config.g1_reject_grasp_table_contact assert tabletop_datagen_config.policy_config.g1_reject_open_grasp_object_contact - assert tabletop_datagen_config.policy_config.g1_allow_open_fingertip_pad_contact assert tabletop_datagen_config.policy_config.g1_grasp_single_pad_contact_penalty == 0.5 assert tabletop_datagen_config.policy_config.g1_closed_grasp_quality_enabled - assert tabletop_datagen_config.policy_config.g1_closed_grasp_settle_steps == 120 + assert tabletop_datagen_config.policy_config.g1_closed_grasp_settle_steps == 25 assert tabletop_datagen_config.policy_config.g1_closed_grasp_min_pad_geom_count == 2 assert tabletop_datagen_config.policy_config.g1_closed_grasp_max_object_shift_m == 0.05 assert tabletop_datagen_config.policy_config.g1_closed_grasp_penalty_per_missing_pad == 1.0 - assert tabletop_datagen_config.policy_config.g1_pregrasp_min_vertical_lift == 0.0 - assert tabletop_datagen_config.policy_config.g1_pregrasp_object_clearance == 0.04 - assert tabletop_datagen_config.policy_config.g1_place_travel_object_clearance == 0.07 - assert tabletop_datagen_config.policy_config.g1_runtime_null_space_weight == 0.3 - assert tabletop_datagen_config.policy_config.g1_runtime_null_space_damping == 1e-4 - assert tabletop_datagen_config.policy_config.g1_runtime_null_space_max_step_rad == 0.05 + assert tabletop_datagen_config.policy_config.g1_pregrasp_min_vertical_lift == 0.10 + assert tabletop_datagen_config.policy_config.g1_pregrasp_object_clearance == 0.12 + assert tabletop_datagen_config.policy_config.g1_place_travel_object_clearance == 0.13 assert tabletop_datagen_config.policy_config.g1_held_object_speed == 0.04 assert tabletop_datagen_config.policy_config.g1_postgrasp_hold_duration == 0.4 assert tabletop_datagen_config.policy_config.g1_record_partial_attempt_on_no_full_grasp_candidate @@ -628,40 +594,6 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): ) assert not unreachable_eval["success"] assert unreachable_eval["failed_phase"] == "pregrasp" - relaxed_pregrasp_result = fake_tabletop_policy._g1_accept_relaxed_candidate_ik_result( - "pregrasp", - { - "success": False, - "final_pos_error_norm": 0.05, - "final_rot_error_norm": 10.0, - "final_qpos": pick_datagen_config.robot_config.init_qpos, - "qpos": None, - }, - ) - assert relaxed_pregrasp_result["success"] - assert relaxed_pregrasp_result["accepted_with_g1_candidate_thresholds"] - strict_grasp_result = fake_tabletop_policy._g1_accept_relaxed_candidate_ik_result( - "grasp", - { - "success": False, - "final_pos_error_norm": 0.05, - "final_rot_error_norm": 0.4, - "final_qpos": pick_datagen_config.robot_config.init_qpos, - "qpos": None, - }, - ) - assert strict_grasp_result["success"] - rejected_grasp_result = fake_tabletop_policy._g1_accept_relaxed_candidate_ik_result( - "grasp", - { - "success": False, - "final_pos_error_norm": 0.05, - "final_rot_error_norm": 1.0, - "final_qpos": pick_datagen_config.robot_config.init_qpos, - "qpos": None, - }, - ) - assert not rejected_grasp_result["success"] clearance_model = mujoco.MjModel.from_xml_string( """ @@ -720,28 +652,11 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): assert g1_target_poses["grasp"][2, 3] >= ( table_top_z + tabletop_datagen_config.policy_config.g1_grasp_table_clearance ) - # Default G1 candidate construction must not shift the grasp pose in XY: - # the DROID grasp bank already encodes pinch placement in the corrected - # grasp-site frame, and lateral/forward centering and inward XY offsets - # default off. - np.testing.assert_allclose(g1_target_poses["grasp"][:2, 3], grasp_pose[:2, 3]) - np.testing.assert_allclose(g1_target_poses["grasp"][:3, 2], grasp_pose[:3, 2]) - - # Opt-in centering still applies its math when enabled. - centered_policy = object.__new__(UnitreeG1RightArmPickAndPlacePlannerPolicy) - centered_policy.policy_config = tabletop_datagen_config.policy_config.model_copy( - update={ - "g1_center_grasp_forward": True, - "g1_grasp_forward_centering_target_m": 0.02, - } - ) - forward_offset_pose = np.eye(4) - forward_offset_pose[:3, 3] = [0.0, 0.0, -0.05] - forward_centered_pose = centered_policy._g1_center_grasp_forward( - forward_offset_pose, - SimpleNamespace(position=np.array([0.0, 0.0, 0.0])), + object_in_grasp = g1_target_poses["grasp"][:3, :3].T @ ( + pickup_obj.position - g1_target_poses["grasp"][:3, 3] ) - np.testing.assert_allclose(forward_centered_pose[:3, 3], [0.0, 0.0, -0.02]) + assert abs(object_in_grasp[1]) < 1e-8 + np.testing.assert_allclose(g1_target_poses["grasp"][:3, 2], grasp_pose[:3, 2]) fake_clearance_policy.policy_config = tabletop_datagen_config.policy_config.model_copy( update={"g1_level_grasp_orientation": True} @@ -757,13 +672,8 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): np.eye(3), atol=1e-7, ) - assert g1_target_poses["pregrasp"][2, 3] >= ( - pickup_top_z - + tabletop_datagen_config.policy_config.g1_pregrasp_object_clearance - ) - # Pregrasp also lifts strictly above the (clamped) grasp pose along the - # gripper-forward axis, so the wrist always sits above the grasp. - assert g1_target_poses["pregrasp"][2, 3] > g1_target_poses["grasp"][2, 3] + assert g1_target_poses["pregrasp"][2, 3] >= grasp_pose[2, 3] + 0.10 + assert g1_target_poses["pregrasp"][2, 3] >= pickup_top_z + 0.12 pickup_bottom_z = 0.70 receptacle_top_z = 0.75 min_travel_bottom_z = ( @@ -776,7 +686,7 @@ def test_prepare_unitree_g1_dex1_smoke(tmp_path, monkeypatch): + pickup_bottom_z - g1_target_poses["grasp"][2, 3] ) - assert carried_bottom_z + 1e-9 >= min_travel_bottom_z + assert carried_bottom_z >= min_travel_bottom_z place_bottom_z = ( g1_target_poses["place"][2, 3] + pickup_bottom_z diff --git a/molmo_spaces/configs/policy_configs.py b/molmo_spaces/configs/policy_configs.py index 39d4ec09..50f06da3 100644 --- a/molmo_spaces/configs/policy_configs.py +++ b/molmo_spaces/configs/policy_configs.py @@ -194,12 +194,7 @@ class UnitreeG1RightArmPickAndPlacePlannerPolicyConfig(PickAndPlacePlannerPolicy grasp_collision_batch_size: int = 32 grasp_collision_max_grasps: int = 64 max_retries: int = 0 - # Bumped from 4 → 20. The standalone IK teleop shows the right arm can - # reach the preplace region; the runtime was just aborting too early on - # transient near-singular sub-targets during the continuous SE(3) sweep. - # Higher threshold gives the seed chain a chance to recover before the - # planner declares the carry unreachable. - max_sequential_ik_failures: int = 20 + max_sequential_ik_failures: int = 4 phase_timeout: float = 24.0 pregrasp_tcp_pos_err_threshold: float = 0.1 pregrasp_tcp_rot_err_threshold: float = float("inf") @@ -215,101 +210,31 @@ class UnitreeG1RightArmPickAndPlacePlannerPolicyConfig(PickAndPlacePlannerPolicy g1_grasp_candidate_limit: int = 256 g1_grasp_ik_eval_limit: int = 256 g1_grasp_require_all_pick_place_phases: bool = True - # Bumped from the original 0.5 / 0.25 / 1.0 defaults. The joint-margin - # and joint-motion weights are the selector's lever for biasing toward - # grasps whose IK chain stays away from joint limits and whose - # phase-to-phase joint excursion is small — directly addressing the - # runtime "elbow flips mid-carry" failure mode by pre-filtering for - # grasps with a smoother kinematic path. - g1_grasp_joint_margin_weight: float = 1.5 - g1_grasp_joint_motion_weight: float = 1.0 + g1_grasp_joint_margin_weight: float = 0.5 + g1_grasp_joint_motion_weight: float = 0.25 g1_grasp_topdown_weight: float = 1.0 - # Damped-least-squares damping passed to the shared IK solver at G1 - # call sites only. Higher damping = solver stays closer to seed near - # singularities, fewer elbow-branch flips, at the cost of slightly - # slower / less precise convergence. Defaults unchanged for every - # other robot. - g1_runtime_ik_damping: float = 5e-3 - g1_selector_ik_damping: float = 1e-4 - # Runtime IK integration step size. Default 1.0 = full Newton step; - # smaller values produce smoother per-iteration joint changes near - # singular configurations at the cost of slightly more iterations to - # converge. 0.5 is a good balance. - g1_runtime_ik_dt: float = 0.5 - # Exponential-moving-average smoothing factor for the right-arm ctrl - # output. `ctrl_new = α·ctrl_prev + (1-α)·ik_output`. 0.0 = no - # smoothing (raw IK output); 1.0 = freeze (never update). Disabled - # for now: α=0.4 introduced bimodal-basin averaging (mixing - # elbow-up/down IK solutions into a between-state) that produced a - # 43° wrist tracking error in run 23. The arm-gain + IK-dt changes - # alone provided the IK-stability improvement. - g1_runtime_ctrl_smoothing: float = 0.0 - # If any right-arm joint changes by more than this between consecutive - # IK calls, log a [G1_BASIN_FLIP] warning. Indicates the DLS solver - # jumped kinematic branches (elbow swivel or wrist flip) between - # adjacent sub-targets along the Cartesian sweep. - g1_runtime_basin_flip_threshold_rad: float = 0.30 - # If True, log every active MuJoCo contact at runtime (throttled to - # 100 ms intervals) with both body names. Surfaces hidden collisions - # the classifier filters out as `other`. - g1_runtime_full_contact_diagnostic: bool = True - # Null-space posture nudge weight. After each IK call, project - # `(nominal_qpos - ik_qpos)` onto the null space of the gripper-site - # Jacobian and add `weight * v_null` to the IK output. This biases - # the 1-DoF redundancy of the 7-DoF right arm toward a fixed nominal - # posture (`init_qpos["right_arm"]`) without changing the commanded - # TCP pose — preventing the wrist roll/yaw basin flips that - # appeared in run 27 as 19° one-frame joint jumps. 0.0 = disabled. - g1_runtime_null_space_weight: float = 0.3 - g1_runtime_null_space_damping: float = 1e-4 - g1_runtime_null_space_max_step_rad: float = 0.05 - g1_grasp_min_vertical_axis_z: float = 0.75 g1_grasp_max_tcp_rot_deg: float = 120.0 - g1_ignore_flipped_grasps: bool = False - # The G1 right_grasp_site is at the physical pinch center and uses the - # MolmoSpaces parallel-jaw frame (+z forward, +y finger-opening). DROID - # grasp poses are stored in the Robotiq grasp_site frame which is also at - # the pinch center with the same axis convention, so DROID grasp poses - # apply to the G1 right_grasp_site directly. The inward / lateral / - # forward centering knobs below were empirically tuned against the prior - # rotated G1 grasp site and now distort otherwise-good poses; default - # them off and keep the levers available for follow-up tuning. - g1_grasp_inward_xy_offset: float = 0.0 + g1_ignore_flipped_grasps: bool = True + g1_grasp_inward_xy_offset: float = 0.006 g1_grasp_table_clearance: float = 0.065 - g1_center_grasp_lateral: bool = False + g1_center_grasp_lateral: bool = True g1_grasp_lateral_centering_scale: float = 1.0 g1_grasp_lateral_centering_max_offset: float = 0.02 - g1_center_grasp_forward: bool = False - g1_grasp_forward_centering_scale: float = 1.0 - g1_grasp_forward_centering_max_offset: float = 0.03 - g1_grasp_forward_centering_target_m: float = 0.0 g1_level_grasp_orientation: bool = False g1_grasp_level_max_tilt_deg: float = 35.0 g1_require_fingertip_pad_grasp_contact: bool = False g1_reject_non_fingertip_grasp_object_contact: bool = True g1_reject_grasp_table_contact: bool = True g1_reject_open_grasp_object_contact: bool = True - g1_allow_open_fingertip_pad_contact: bool = True g1_grasp_single_pad_contact_penalty: float = 0.5 g1_closed_grasp_quality_enabled: bool = True - g1_closed_grasp_settle_steps: int = 120 + g1_closed_grasp_settle_steps: int = 25 g1_closed_grasp_min_pad_geom_count: int = 2 g1_closed_grasp_max_object_shift_m: float = 0.05 g1_closed_grasp_penalty_per_missing_pad: float = 1.0 - # With the corrected grasp-site frame, `pregrasp_z_offset` along the - # gripper's true forward axis already lifts pregrasp above the grasp by - # the configured offset for top-down approaches; the additional - # "minimum vertical lift" and "object clearance" margins below were - # over-padding under the prior misaligned site frame. - g1_pregrasp_min_vertical_lift: float = 0.0 - g1_pregrasp_object_clearance: float = 0.04 - # Carried bottle's bottom must stay this far above the bin's AABB top - # during the lift→preplace travel. Reduced from 0.13 — the original - # value was over-cautious and forced the right arm into a high-reach - # envelope (~24 cm above table) where the runtime IK frequently - # singular-fails. The bin walls are 9.5 cm tall, so 5 cm of clearance - # is plenty of margin during the slow held-object carry (4 cm/s). - g1_place_travel_object_clearance: float = 0.07 + g1_pregrasp_min_vertical_lift: float = 0.10 + g1_pregrasp_object_clearance: float = 0.12 + g1_place_travel_object_clearance: float = 0.13 g1_held_object_speed: float = 0.04 g1_postgrasp_hold_duration: float = 0.4 g1_record_partial_attempt_on_no_full_grasp_candidate: bool = True diff --git a/molmo_spaces/policy/solvers/object_manipulation/pick_and_place_planner_policy.py b/molmo_spaces/policy/solvers/object_manipulation/pick_and_place_planner_policy.py index f55a9638..075d807a 100644 --- a/molmo_spaces/policy/solvers/object_manipulation/pick_and_place_planner_policy.py +++ b/molmo_spaces/policy/solvers/object_manipulation/pick_and_place_planner_policy.py @@ -319,10 +319,7 @@ class UnitreeG1RightArmPickAndPlacePlannerPolicy(PickAndPlacePlannerPolicy): def reset(self, reset_retries: bool = True): self._g1_failure_diag_steps: list[dict] = [] self._g1_last_target_pose: np.ndarray | None = None - self._g1_site_frame_debug_logged = False self._g1_initial_tcp_pose = self._current_g1_tcp_pose() - self._log_g1_site_frame_debug_once("reset") - self._log_g1_frame_conventions("reset") self._g1_failure_diagnostics_logged = False self._g1_failure_reason_override: str | None = None self._g1_failure_phase_override: str | None = None @@ -332,16 +329,8 @@ def reset(self, reset_retries: bool = True): self._g1_selected_grasp_debug: dict | None = None self._g1_grasp_selector_debug: dict | None = None self._g1_last_ik_debug: dict | None = None - self._g1_failure_hold_action: HoldCurrentControlAction | None = None + self._g1_failure_hold_action: NoopAction | None = None self._g1_terminal_failure_pending = False - # Smoothed right-arm ctrl carried across runtime IK calls so the - # EMA filter in `_tcp_to_jp_fn` has a previous-value reference. - self._g1_right_arm_ctrl_smoothed: np.ndarray | None = None - # State for the per-frame runtime IK diagnostic (basin-flip - # detection + full-contact dump). Reset per episode. - self._g1_runtime_prev_arm_qpos: np.ndarray | None = None - self._g1_runtime_diag_step = 0 - self._g1_runtime_contact_last_log_t = -1e9 self._g1_grip_diag_samples: list[dict] = [] self._g1_grip_diag_last_log_time = -np.inf self._g1_grip_diag_last_key: tuple | None = None @@ -643,10 +632,6 @@ def _build_g1_candidate_target_poses( grasp_pose_world, pickup_obj, ) - grasp_pose_world = self._g1_center_grasp_forward( - grasp_pose_world, - pickup_obj, - ) pregrasp_pose = grasp_pose_world.copy() pregrasp_pose[:3, 3] -= self.policy_config.pregrasp_z_offset * pregrasp_pose[:3, 2] @@ -761,33 +746,6 @@ def _g1_center_grasp_lateral( pose[:3, 3] += lateral_offset * pose[:3, 1] return pose - def _g1_center_grasp_forward( - self, - grasp_pose_world: np.ndarray, - pickup_obj: MlSpacesObject, - ) -> np.ndarray: - if not getattr(self.policy_config, "g1_center_grasp_forward", False): - return grasp_pose_world - - pose = grasp_pose_world.copy() - object_pos = np.asarray(pickup_obj.position[:3], dtype=float) - object_in_grasp = pose[:3, :3].T @ (object_pos - pose[:3, 3]) - target_forward_offset = float( - getattr(self.policy_config, "g1_grasp_forward_centering_target_m", 0.0) - ) - forward_offset = float(object_in_grasp[2] - target_forward_offset) - forward_offset *= float( - getattr(self.policy_config, "g1_grasp_forward_centering_scale", 1.0) - ) - max_offset = float( - getattr(self.policy_config, "g1_grasp_forward_centering_max_offset", np.inf) - ) - if np.isfinite(max_offset): - forward_offset = float(np.clip(forward_offset, -max_offset, max_offset)) - - pose[:3, 3] += forward_offset * pose[:3, 2] - return pose - def _g1_level_grasp_orientation(self, grasp_pose_world: np.ndarray) -> np.ndarray: if not getattr(self.policy_config, "g1_level_grasp_orientation", False): return grasp_pose_world @@ -890,7 +848,6 @@ def _g1_pregrasp_debug( np.degrees(np.arccos(np.clip(abs(grasp_pose_world[2, 2]), -1.0, 1.0))) ), "object_local_y_in_grasp_m": float(object_in_grasp[1]), - "object_local_z_in_grasp_m": float(object_in_grasp[2]), "grasp_table_clearance_m": float(grasp_pose_world[2, 3] - table_top_z) if np.isfinite(table_top_z) else None, @@ -933,11 +890,7 @@ def _g1_eval_candidate_target_poses( qpos, base_pose, max_iter=250, - damping=float( - getattr(self.policy_config, "g1_selector_ik_damping", 1e-12) - ), ) - result = self._g1_accept_relaxed_candidate_ik_result(phase, result) phase_results[phase] = result if not result["success"]: return { @@ -957,40 +910,6 @@ def _g1_eval_candidate_target_poses( robot_view.base.pose = live_base_pose robot_view.set_qpos_dict(live_qpos) - def _g1_candidate_ik_thresholds_for_phase(self, phase: str) -> tuple[float, float]: - if phase == "pregrasp": - return ( - float(getattr(self.policy_config, "pregrasp_tcp_pos_err_threshold", 0.1)), - float(getattr(self.policy_config, "pregrasp_tcp_rot_err_threshold", np.inf)), - ) - return ( - float(self.policy_config.tcp_pos_err_threshold), - float(self.policy_config.tcp_rot_err_threshold), - ) - - def _g1_accept_relaxed_candidate_ik_result(self, phase: str, result: dict) -> dict: - if result["success"]: - return result - - final_qpos = result.get("final_qpos") - if final_qpos is None: - return result - - pos_threshold, rot_threshold = self._g1_candidate_ik_thresholds_for_phase(phase) - pos_err = float(result.get("final_pos_error_norm", np.inf)) - rot_err = float(result.get("final_rot_error_norm", np.inf)) - if pos_err > pos_threshold or rot_err > rot_threshold: - return result - - accepted = dict(result) - accepted["success"] = True - accepted["strict_success"] = False - accepted["accepted_with_g1_candidate_thresholds"] = True - accepted["qpos"] = self._copy_qpos_dict(final_qpos) - accepted["g1_candidate_pos_err_threshold"] = pos_threshold - accepted["g1_candidate_rot_err_threshold"] = rot_threshold - return accepted - def _g1_contact_quality_from_contacts(self, contacts: list[dict]) -> dict: robot_namespace = self.config.robot_config.robot_namespace object_contacts = [ @@ -1214,42 +1133,22 @@ def _g1_closed_grasp_contact_quality( base_pose: np.ndarray, pickup_obj: MlSpacesObject, ) -> dict: - robot = self.task.env.current_robot - robot_view = robot.robot_view + robot_view = self.task.env.current_robot.robot_view model = self.task.env.current_model data = self.task.env.current_data snapshot = self._g1_mujoco_state_snapshot() - # Without restoring this, the runtime would keep pinning the base to - # whatever pose we latched during the probe. - prior_pinned_base_pose = getattr(robot, "_pinned_base_pose", None) - if prior_pinned_base_pose is not None: - prior_pinned_base_pose = prior_pinned_base_pose.copy() settle_steps = max( int(getattr(self.policy_config, "g1_closed_grasp_settle_steps", 25)), 0, ) - gripper_mg_ids = set(robot_view.get_gripper_movegroup_ids()) try: robot_view.base.pose = base_pose - # Latch the probe's base pose so per-step re-pinning doesn't snap - # back to the live runtime base pose. - if hasattr(robot, "sync_pinned_base_pose"): - robot.sync_pinned_base_pose() robot_view.set_qpos_dict(qpos) - # Hold candidate qpos via position-servo ctrl on non-gripper move - # groups so the wrist doesn't drift back to the prior live ctrl - # target during the settle. - for mg_id in robot_view.move_group_ids(): - if mg_id in gripper_mg_ids or mg_id not in qpos: - continue - mg = robot_view.get_move_group(mg_id) - if mg.n_actuators == 0 or mg.n_actuators != mg.pos_dim: - continue - mg.ctrl = np.asarray(qpos[mg_id]) gripper_mg_id = robot_view.get_gripper_movegroup_ids()[0] gripper = robot_view.get_gripper(gripper_mg_id) gripper.set_gripper_ctrl_open(False) + gripper.joint_pos = gripper.ctrl.copy() data.qvel[:] = 0.0 data.qacc[:] = 0.0 data.qacc_warmstart[:] = 0.0 @@ -1257,10 +1156,6 @@ def _g1_closed_grasp_contact_quality( object_start_pos = pickup_obj.position.copy() for _ in range(settle_steps): mujoco.mj_step(model, data) - # Re-enforce pinned base + locked left arm/hand between - # steps; right arm is held via the ctrl assignment above. - if hasattr(robot, "apply_initial_state_overrides"): - robot.apply_initial_state_overrides() mujoco.mj_forward(model, data) object_shift_m = float(np.linalg.norm(pickup_obj.position - object_start_pos)) contacts = self._collect_contact_diagnostics() @@ -1279,8 +1174,6 @@ def _g1_closed_grasp_contact_quality( return quality finally: self._g1_restore_mujoco_state(snapshot) - if hasattr(robot, "_pinned_base_pose"): - robot._pinned_base_pose = prior_pinned_base_pose def _g1_candidate_grasp_contact_quality( self, @@ -1299,7 +1192,6 @@ def _g1_grasp_contact_quality_failure(self, quality: dict) -> str | None: if ( getattr(self.policy_config, "g1_reject_open_grasp_object_contact", True) and quality["object_contact_count"] > 0 - and not self._g1_is_allowed_open_grasp_object_contact(quality) ): return "open_grasp_object_contact" if ( @@ -1319,21 +1211,6 @@ def _g1_grasp_contact_quality_failure(self, quality: dict) -> str | None: return "missing_fingertip_pad_contact" return None - def _g1_is_allowed_open_grasp_object_contact(self, quality: dict) -> bool: - if not getattr(self.policy_config, "g1_allow_open_fingertip_pad_contact", False): - return False - if not quality["non_pad_object_contacts"]: - return quality["pad_geom_count"] > 0 - - min_pad_geom_count = int( - getattr(self.policy_config, "g1_closed_grasp_min_pad_geom_count", 2) - ) - return ( - quality["pad_geom_count"] >= min_pad_geom_count - and bool(quality["finger_link_object_contacts"]) - and not quality["other_robot_object_contacts"] - ) - def _g1_grasp_quality_rejection_key(self, quality: dict) -> str: if quality["closed_probe"]: return "closed_grasp_quality" @@ -1476,8 +1353,7 @@ def _select_g1_tabletop_grasp_pose( dist_tcp = np.linalg.inv(tcp_pose_world) @ grasp_poses_world dists_tcp_p = np.linalg.norm(dist_tcp[:, :3, 3], axis=1) dists_tcp_o = R.from_matrix(dist_tcp[:, :3, :3]).magnitude() * 180 / np.pi - grasp_forward_axis_z = grasp_poses_world[:, 2, 2] - grasp_down_alignment = -grasp_forward_axis_z + dists_up = grasp_poses_world[:, 2, 2] dists_com = np.linalg.norm( (np.linalg.inv(object_pose) @ grasp_poses_world)[:, :3, 3], axis=1, @@ -1485,7 +1361,7 @@ def _select_g1_tabletop_grasp_pose( base_cost = ( self.policy_config.grasp_pos_cost_weight * dists_tcp_p + self.policy_config.grasp_rot_cost_weight * dists_tcp_o - + self.policy_config.grasp_vertical_cost_weight * grasp_forward_axis_z + + self.policy_config.grasp_vertical_cost_weight * dists_up + self.policy_config.grasp_com_dist_cost_weight * dists_com ) @@ -1501,12 +1377,6 @@ def _select_g1_tabletop_grasp_pose( tcp_eligible_mask = eligible_mask.copy() eligible_mask &= dists_tcp_o <= max_tcp_rot_deg tcp_rotation_mask = tcp_eligible_mask & ~eligible_mask - min_vertical_axis_z = float( - getattr(self.policy_config, "g1_grasp_min_vertical_axis_z", -np.inf) - ) - vertical_eligible_mask = eligible_mask.copy() - eligible_mask &= grasp_down_alignment >= min_vertical_axis_z - vertical_axis_mask = vertical_eligible_mask & ~eligible_mask eligible_ids = np.where(eligible_mask)[0] candidate_limit = min( @@ -1540,7 +1410,6 @@ def _select_g1_tabletop_grasp_pose( else 0 ), "tcp_rotation": int(np.count_nonzero(tcp_rotation_mask)), - "vertical_axis": int(np.count_nonzero(vertical_axis_mask)), "collision": int(len(candidate_ids) - len(collision_free_ids)), "pregrasp": 0, "grasp": 0, @@ -1632,10 +1501,7 @@ def _select_g1_tabletop_grasp_pose( "base_cost": float(base_cost[grasp_idx]), "tcp_pos_cost_m": float(dists_tcp_p[grasp_idx]), "tcp_rot_cost_deg": float(dists_tcp_o[grasp_idx]), - "vertical_axis_z": float(grasp_forward_axis_z[grasp_idx]), - "grasp_down_alignment": float( - grasp_down_alignment[grasp_idx] - ), + "vertical_axis_z": float(dists_up[grasp_idx]), "com_dist_cost_m": float(dists_com[grasp_idx]), "progress_penalty": progress_penalty, "is_flipped": bool(grasp_idx >= original_grasp_count), @@ -1716,8 +1582,7 @@ def _select_g1_tabletop_grasp_pose( "base_cost": float(base_cost[grasp_idx]), "tcp_pos_cost_m": float(dists_tcp_p[grasp_idx]), "tcp_rot_cost_deg": float(dists_tcp_o[grasp_idx]), - "vertical_axis_z": float(grasp_forward_axis_z[grasp_idx]), - "grasp_down_alignment": float(grasp_down_alignment[grasp_idx]), + "vertical_axis_z": float(dists_up[grasp_idx]), "com_dist_cost_m": float(dists_com[grasp_idx]), "is_flipped": bool(grasp_idx >= original_grasp_count), "grasp_pad_contact_count": int(grasp_quality["pad_contact_count"]), @@ -1853,257 +1718,6 @@ def _failure_diagnostics_enabled(self) -> bool: def _g1_ik_debug_enabled(self) -> bool: return bool(getattr(self.policy_config, "g1_ik_debug", False)) - @staticmethod - def _unit_vector(vec: np.ndarray) -> list[float] | None: - norm = float(np.linalg.norm(vec)) - if norm < 1e-8: - return None - return np.round(vec / norm, 4).tolist() - - def _g1_named_site_id(self, site_name: str) -> int: - model = self.task.env.current_model - namespace = self.task.env.current_robot.namespace - for candidate in (f"{namespace}{site_name}", site_name): - site_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE, candidate) - if site_id >= 0: - return int(site_id) - raise ValueError(f"Expected MuJoCo site `{site_name}`") - - def _g1_named_geom_id(self, geom_name: str) -> int: - model = self.task.env.current_model - namespace = self.task.env.current_robot.namespace - for candidate in (f"{namespace}{geom_name}", geom_name): - geom_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, candidate) - if geom_id >= 0: - return int(geom_id) - raise ValueError(f"Expected MuJoCo geom `{geom_name}`") - - def _g1_site_frame_debug(self) -> dict: - data = self.task.env.current_data - wrist_site_id = self._g1_named_site_id("right_wrist_site") - grasp_site_id = self._g1_named_site_id("right_grasp_site") - pad_1_geom_id = self._g1_named_geom_id("right_dex1_fingertip_pad_1") - pad_2_geom_id = self._g1_named_geom_id("right_dex1_fingertip_pad_2") - - wrist_pos = data.site_xpos[wrist_site_id].copy() - grasp_pos = data.site_xpos[grasp_site_id].copy() - grasp_xmat = data.site_xmat[grasp_site_id].reshape(3, 3).copy() - grasp_axes = { - "+x": np.round(grasp_xmat[:, 0], 4).tolist(), - "+y": np.round(grasp_xmat[:, 1], 4).tolist(), - "+z": np.round(grasp_xmat[:, 2], 4).tolist(), - } - wrist_to_grasp = grasp_pos - wrist_pos - pad_axis = data.geom_xpos[pad_2_geom_id] - data.geom_xpos[pad_1_geom_id] - wrist_to_grasp_unit = self._unit_vector(wrist_to_grasp) - pad_axis_unit = self._unit_vector(pad_axis) - - return { - "right_grasp_site_pos": np.round(grasp_pos, 4).tolist(), - "right_grasp_site_axes_world": grasp_axes, - "wrist_to_grasp_unit": wrist_to_grasp_unit, - "finger_pad_1_to_2_unit": pad_axis_unit, - "dot_site_x_with_wrist_to_grasp": float( - np.dot(grasp_xmat[:, 0], wrist_to_grasp / np.linalg.norm(wrist_to_grasp)) - ), - "dot_site_y_with_finger_pad_axis": float( - np.dot(grasp_xmat[:, 1], pad_axis / np.linalg.norm(pad_axis)) - ), - "dot_site_z_with_wrist_to_grasp": float( - np.dot(grasp_xmat[:, 2], wrist_to_grasp / np.linalg.norm(wrist_to_grasp)) - ), - } - - def _log_g1_site_frame_debug_once(self, label: str) -> None: - if not self._g1_ik_debug_enabled() or self._g1_site_frame_debug_logged: - return - self._g1_site_frame_debug_logged = True - log.info("[G1_SITE] label=%s debug=%s", label, self._g1_site_frame_debug()) - - def _g1_apply_null_space_posture_nudge( - self, gripper_mg_id: str, arm_ctrl: np.ndarray - ) -> np.ndarray: - """Bias `arm_ctrl` toward the nominal right-arm posture along the - null space of the gripper-site Jacobian. Returns a modified ctrl - whose forward-kinematic TCP pose is (to first order) identical to - the input. The redundant 1 DoF of the 7-DoF arm is steered toward - the configured nominal, preventing the wrist roll/yaw basin flips - seen in the runtime IK output across consecutive frames. - """ - weight = float( - getattr(self.policy_config, "g1_runtime_null_space_weight", 0.0) - ) - if weight <= 0.0: - return arm_ctrl - - try: - nominal = self.config.robot_config.init_qpos.get("right_arm") - except AttributeError: - return arm_ctrl - if nominal is None: - return arm_ctrl - nominal_arr = np.asarray(nominal, dtype=float) - if nominal_arr.shape != arm_ctrl.shape: - return arm_ctrl - - # Jacobian of the gripper site w.r.t. right_arm joints, at the - # current runtime qpos (close enough to the IK-solved qpos for a - # first-order null-space projection). - J = self.robot_view.get_jacobian(gripper_mg_id, ["right_arm"]) - if J.shape[0] != 6 or J.shape[1] != arm_ctrl.shape[0]: - return arm_ctrl - - # Damped null-space projector N = I - J⁺ J where J⁺ uses damping - # for numerical stability near singular configurations. - damping = float( - getattr(self.policy_config, "g1_runtime_null_space_damping", 1e-4) - ) - try: - J_pinv = J.T @ np.linalg.solve( - J @ J.T + damping * np.eye(6), np.eye(6) - ) - except np.linalg.LinAlgError: - return arm_ctrl - N = np.eye(arm_ctrl.shape[0]) - J_pinv @ J - - delta = N @ (nominal_arr - arm_ctrl) - max_step = float( - getattr(self.policy_config, "g1_runtime_null_space_max_step_rad", 0.05) - ) - scaled = weight * delta - per_joint_cap = float(max_step) if np.isfinite(max_step) else np.inf - scaled = np.clip(scaled, -per_joint_cap, per_joint_cap) - - adjusted = arm_ctrl + scaled - # Clamp to actuator ctrlrange (joint limits) so the nudge cannot - # produce an out-of-range command. - arm_mg = self.robot_view.get_move_group("right_arm") - if hasattr(arm_mg, "ctrl_limits"): - limits = arm_mg.ctrl_limits - adjusted = np.clip(adjusted, limits[:, 0], limits[:, 1]) - return adjusted - - def _g1_runtime_diagnostic( - self, - target_pose: np.ndarray, - ik_succeeded: bool, - new_arm_ctrl: np.ndarray | None = None, - ) -> None: - """Per-frame runtime diagnostic: basin-flip detection on the - right-arm **IK output ctrl** and a full contact dump. - - Compare the *commanded* qpos (= IK output → ctrl) across - consecutive successful IK calls. The actual joint pos lags via - servo dynamics and would not reveal an IK-output jump. - """ - self._g1_runtime_diag_step += 1 - step = self._g1_runtime_diag_step - model = self.task.env.current_model - data = self.task.env.current_data - phase = self.get_phase() - threshold = float( - getattr(self.policy_config, "g1_runtime_basin_flip_threshold_rad", 0.30) - ) - force_dump = False - - # 1) Basin-flip detection on the IK *output* ctrl - if ik_succeeded and new_arm_ctrl is not None: - if self._g1_runtime_prev_arm_qpos is not None: - delta = new_arm_ctrl - self._g1_runtime_prev_arm_qpos - abs_delta = np.abs(delta) - max_delta = float(abs_delta.max()) - if max_delta > threshold: - worst = int(abs_delta.argmax()) - force_dump = True - log.warning( - "[G1_BASIN_FLIP] step=%d phase=%s max_delta=%.4frad " - "on_joint=%d prev_ctrl=%s new_ctrl=%s per_joint_delta=%s " - "target_pos=%s", - step, - phase, - max_delta, - worst, - np.round(self._g1_runtime_prev_arm_qpos, 4).tolist(), - np.round(new_arm_ctrl, 4).tolist(), - np.round(delta, 4).tolist(), - np.round(target_pose[:3, 3], 4).tolist(), - ) - self._g1_runtime_prev_arm_qpos = new_arm_ctrl.copy() - - # 2) Full contact dump (throttled to ~10 Hz; also fires immediately - # on basin-flip) - if not getattr( - self.policy_config, "g1_runtime_full_contact_diagnostic", False - ): - return - sim_t = float(data.time) - if (sim_t - self._g1_runtime_contact_last_log_t < 0.1) and not force_dump: - return - self._g1_runtime_contact_last_log_t = sim_t - - def geom_label(geom_id: int) -> str: - gn = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_GEOM, geom_id) - if gn: - return gn - # Fall back to body name + geom id so unnamed structural - # collision geoms (e.g., the bin's hidden walls) are - # identifiable in the log. - body_id = int(model.geom_bodyid[geom_id]) - bn = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_BODY, body_id) or "?" - return f"{bn}#geom{geom_id}" - - contacts = [] - for i in range(data.ncon): - con = data.contact[i] - g1n = geom_label(con.geom1) - g2n = geom_label(con.geom2) - contacts.append((float(con.dist), g1n, g2n)) - contacts.sort(key=lambda c: c[0]) # most-interpenetrating first - formatted = ";".join( - f"{g1}↔{g2}@{d:+.4f}" for d, g1, g2 in contacts[:12] - ) - log.info( - "[G1_CONTACTS] step=%d phase=%s t=%.3f n=%d %s", - step, - phase, - sim_t, - data.ncon, - formatted if formatted else "(none)", - ) - - def _log_g1_frame_conventions(self, label: str) -> None: - """Unconditional convention check at reset. Confirms: - - Pelvis frame: +x forward, +y left, +z up (MolmoSpaces robot base) - - Right grasp_site: +z forward, +y finger-opening (parallel-jaw) - Logs the world-frame axes of both so any deviation from convention - is visible at a glance. - """ - try: - model = self.task.env.current_model - data = self.task.env.current_data - pelvis_id = mujoco.mj_name2id( - model, - mujoco.mjtObj.mjOBJ_BODY, - f"{self.robot_view._namespace}pelvis", - ) - if pelvis_id < 0: - return - pelvis_xmat = data.xmat[pelvis_id].reshape(3, 3) - site_debug = self._g1_site_frame_debug() - log.info( - "[G1_CONV] label=%s pelvis_world_axes={+x:%s, +y:%s, +z:%s} " - "grasp_axes=%s dot_z_forward=%.3f dot_y_finger=%.3f", - label, - np.round(pelvis_xmat[:, 0], 3).tolist(), - np.round(pelvis_xmat[:, 1], 3).tolist(), - np.round(pelvis_xmat[:, 2], 3).tolist(), - site_debug["right_grasp_site_axes_world"], - site_debug["dot_site_z_with_wrist_to_grasp"], - site_debug["dot_site_y_with_finger_pad_axis"], - ) - except Exception as exc: # diagnostic only — never block the run - log.debug("[G1_CONV] frame-convention diagnostic failed: %r", exc) - def _names_for_contact(self, geom_id: int) -> dict[str, str | int]: model = self.task.env.current_model body_id = model.geom_bodyid[geom_id] @@ -2416,14 +2030,7 @@ def _handle_failure(self) -> dict[str, np.ndarray | bool]: self._g1_terminal_failure_pending = True if self._g1_failure_hold_action is None: log.info("[G1_DIAG] holding failed attempt for %.2fs before done", hold_duration) - # Hold the current commanded controls (including a closed - # gripper carrying a held object) instead of NoopAction, which - # can infer an open gripper when the held object keeps the - # fingers apart and cause a visible drop at the end of a - # failed attempt. - self._g1_failure_hold_action = HoldCurrentControlAction( - self.robot_view, hold_duration - ) + self._g1_failure_hold_action = NoopAction(self.robot_view, hold_duration) if not self._g1_failure_hold_action.execute(): return self._g1_failure_hold_action.get_current_action() @@ -2465,47 +2072,19 @@ def _tcp_to_jp_fn(self, mg_id: str, target_pose: np.ndarray) -> dict[str, np.nda ["right_arm"], self.robot_view.get_qpos_dict(), self.robot_view.base.pose, - damping=float(getattr(self.policy_config, "g1_runtime_ik_damping", 1e-12)), - dt=float(getattr(self.policy_config, "g1_runtime_ik_dt", 1.0)), ) self._g1_last_target_pose = target_pose.copy() action = self.robot_view.get_ctrl_dict() if jp is not None: self.sequential_ik_failures = 0 - new_arm_ctrl = np.asarray(jp["right_arm"], dtype=float) - # Null-space posture nudge: bias the redundant 1 DoF of the 7-DoF - # right arm toward a fixed nominal posture without changing the - # commanded TCP pose. Prevents the wrist-orientation basin flips - # caught by the [G1_BASIN_FLIP] diagnostic. - new_arm_ctrl = self._g1_apply_null_space_posture_nudge( - mg_id, new_arm_ctrl - ) - # Exponential-moving-average smoothing across frames to kill - # high-frequency IK output jitter that an underdamped servo - # would otherwise overshoot. - alpha = float( - getattr(self.policy_config, "g1_runtime_ctrl_smoothing", 0.0) - ) - if 0.0 < alpha < 1.0 and self._g1_right_arm_ctrl_smoothed is not None: - new_arm_ctrl = ( - alpha * self._g1_right_arm_ctrl_smoothed - + (1.0 - alpha) * new_arm_ctrl - ) - self._g1_right_arm_ctrl_smoothed = new_arm_ctrl.copy() - action["right_arm"] = new_arm_ctrl + action["right_arm"] = jp["right_arm"] self._record_failure_diagnostic_step(target_pose, True) - self._g1_runtime_diagnostic( - target_pose, ik_succeeded=True, new_arm_ctrl=new_arm_ctrl - ) else: self.sequential_ik_failures += 1 log.info(f"IK failed, holding current position, fails:{self.sequential_ik_failures}") self._record_failure_diagnostic_step(target_pose, False) self._run_g1_ik_debug(mg_id, target_pose) - self._g1_runtime_diagnostic( - target_pose, ik_succeeded=False, new_arm_ctrl=None - ) if self.sequential_ik_failures >= self.policy_config.max_sequential_ik_failures: log.info("Too many sequential IK failures, triggering retry.") return self._handle_failure() @@ -2526,7 +2105,6 @@ def _run_diagnostic_ik( self.robot_view.get_qpos_dict(), self.robot_view.base.pose, max_iter=max_iter, - damping=float(getattr(self.policy_config, "g1_selector_ik_damping", 1e-12)), ) def _candidate_pose_for_phase(self, candidate_grasp_pose: np.ndarray) -> np.ndarray: @@ -2759,9 +2337,6 @@ def check_feasible_ik(self, pose: np.ndarray) -> bool: kinematics = self.task.env.current_robot.kinematics gripper_mg_id = robot_view.get_gripper_movegroup_ids()[0] - selector_damping = float( - getattr(self.policy_config, "g1_selector_ik_damping", 1e-12) - ) if pose.ndim > 2: return np.array( [ @@ -2772,7 +2347,6 @@ def check_feasible_ik(self, pose: np.ndarray) -> bool: robot_view.get_qpos_dict(), robot_view.base.pose, max_iter=100, - damping=selector_damping, ) is not None for single_pose in pose @@ -2788,6 +2362,5 @@ def check_feasible_ik(self, pose: np.ndarray) -> bool: robot_view.get_qpos_dict(), robot_view.base.pose, max_iter=100, - damping=selector_damping, ) return jp_dict is not None