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/mlspaces_tests/component_tests/test_unitree_g1.py b/mlspaces_tests/component_tests/test_unitree_g1.py
new file mode 100644
index 00000000..397b9804
--- /dev/null
+++ b/mlspaces_tests/component_tests/test_unitree_g1.py
@@ -0,0 +1,786 @@
+import importlib
+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.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,
+ 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",
+ "right_wrist_site",
+ "left_grasp_site",
+ "right_grasp_site",
+]
+
+
+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",
+ ]
+
+ 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()
+ 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,
+ 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 _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_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:
+ 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)
+ 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
+ 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()
+ 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],
+ )
+ 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 = {
+ "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
+ assert view.get_move_group("left_arm").leaf_frame_type == "site"
+ assert (
+ 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_grasp_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)
+ 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
+ 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"
+ )
+ 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,
+ 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]
+ assert tabletop_datagen_config.task_sampler_config.samples_per_house == 5
+ 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.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,
+ )
+ 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.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
+ 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_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"
+ )
+ 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_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))
+ 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_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(
+ 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_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)
+ )
+ 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
+ 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
+ }
+ 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)
+ 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],
+ atol=1e-3,
+ )
+ 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)
+
+ 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/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..50f06da3 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
@@ -154,6 +178,78 @@ 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
+ 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
+ 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 = 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_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."""
+ 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 a4a26196..4a7bf2de 100644
--- a/molmo_spaces/configs/robot_configs.py
+++ b/molmo_spaces/configs/robot_configs.py
@@ -36,6 +36,11 @@
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 (
+ UnitreeG1RightArmPickRobotView,
+ UnitreeG1RobotView,
+)
+from molmo_spaces.robots.unitree_g1 import UnitreeG1Robot
class ActionNoiseConfig(Config):
@@ -96,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
@@ -419,3 +425,76 @@ 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
+ 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],
+ "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",
+ "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 070548b6..e9851d95 100644
--- a/molmo_spaces/configs/task_sampler_configs.py
+++ b/molmo_spaces/configs/task_sampler_configs.py
@@ -210,6 +210,35 @@ 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"
+ robot_base_xy: tuple[float, float] = (0.08, 0.0)
+ pickup_workspace_center_xy: tuple[float, float] = (0.42, -0.20)
+ # 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
+ 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 7b18789d..c43a1a18 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,
@@ -26,18 +27,24 @@
FrankaRandomizedD405D455CameraSystem,
FrankaRandomizedDroidCameraSystem,
RBY1GoProD455CameraSystem,
+ UnitreeG1RightArmPickCameraSystem,
)
from molmo_spaces.configs.policy_configs import (
CuroboOpenClosePlannerPolicyConfig,
CuroboPickAndPlacePlannerPolicyConfig,
OpenClosePlannerPolicyConfig,
PickPlannerPolicyConfig,
+ UnitreeG1RightArmPickAndPlacePlannerPolicyConfig,
+ UnitreeG1RightArmPickPlannerPolicyConfig,
)
from molmo_spaces.configs.robot_configs import (
FloatingRUMRobotConfig,
FrankaRobotConfig,
RBY1MConfig,
RBY1MOpenCloseConfig,
+ UnitreeG1Dex1RobotConfig,
+ UnitreeG1RightArmPickRobotConfig,
+ UnitreeG1RightArmTabletopPickRobotConfig,
)
from molmo_spaces.configs.task_sampler_configs import (
OpenTaskSamplerConfig,
@@ -46,6 +53,7 @@
PickAndPlaceTaskSamplerConfig,
PickTaskSamplerConfig,
RUMPickTaskSamplerConfig,
+ UnitreeG1TabletopPickAndPlaceTaskSamplerConfig,
)
# Oder of configs should be order the code is executed in
@@ -59,10 +67,24 @@
PickAndPlaceMultiTaskSampler,
PickAndPlaceTaskSampler,
)
-from molmo_spaces.tasks.pick_task_sampler import PickTaskSampler
+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
+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):
@@ -123,6 +145,211 @@ 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("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=UnitreeG1RightArmPickTaskSampler,
+ 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.15, 0.4),
+ 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("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: UnitreeG1RightArmTabletopPickRobotConfig = (
+ UnitreeG1RightArmTabletopPickRobotConfig()
+ )
+ camera_config: UnitreeG1RightArmPickCameraSystem = UnitreeG1RightArmPickCameraSystem()
+ 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=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 = 360
+ end_on_success: bool = True
+ 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("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 42becb26..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
@@ -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,11 +305,2062 @@ 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()
+
+ return target_poses
+
+
+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
+ 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)),
+ "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
+
+ jp = kinematics.ik(
+ mg_id,
+ target_pose,
+ ["right_arm"],
+ 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)
+ 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/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
new file mode 100644
index 00000000..525824de
--- /dev/null
+++ b/molmo_spaces/robots/robot_views/unitree_g1_view.py
@@ -0,0 +1,298 @@
+"""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,
+ GripperGroup,
+ MJCFFrameMixin,
+ RobotBaseGroup,
+ RobotView,
+ SimplyActuatedMoveGroup,
+)
+from molmo_spaces.utils.mj_model_and_data_utils import body_pose, site_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_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_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_frame_id
+
+ @property
+ def leaf_frame_type(self):
+ return self._leaf_frame_type
+
+ @property
+ def noop_ctrl(self) -> np.ndarray:
+ return self.joint_pos.copy()
+
+ @property
+ def leaf_frame_to_world(self) -> np.ndarray:
+ 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:
+ 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_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_frame_name,
+ base,
+ namespace,
+ leaf_frame_type,
+ )
+
+
+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."""
+
+ 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_grasp_site",
+ base,
+ namespace,
+ "site",
+ ),
+ "right_arm": UnitreeG1SideMoveGroup(
+ mj_data,
+ "right",
+ UNITREE_G1_ARM_JOINTS,
+ "torso_link",
+ "right_grasp_site",
+ base,
+ namespace,
+ "site",
+ ),
+ "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"]
+
+
+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"]
diff --git a/molmo_spaces/robots/unitree_g1.py b/molmo_spaces/robots/unitree_g1.py
new file mode 100644
index 00000000..d5c09c9e
--- /dev/null
+++ b/molmo_spaces/robots/unitree_g1.py
@@ -0,0 +1,197 @@
+"""Unitree G1 Dex1.1 robot implementation."""
+
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
+from mujoco import MjData, MjSpec
+
+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._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()
+ if move_group_id != "base"
+ and config.robot_config.command_mode.get(move_group_id) == "joint_position"
+ }
+ self.reset()
+
+ @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
+
+ @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(
+ 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:
+ 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:
+ 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.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()
+
+ 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)
+ 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))
+
+ 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:
+ return "pelvis"
diff --git a/molmo_spaces/tasks/pick_task_sampler.py b/molmo_spaces/tasks/pick_task_sampler.py
index fe949eb9..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
@@ -962,14 +965,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 +1002,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 +1187,18 @@ def add_pickup_target(
)
return spec
+
+
+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")
+ 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])
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..7d4faeb9
--- /dev/null
+++ b/molmo_spaces/tasks/unitree_g1_tabletop_task_sampler.py
@@ -0,0 +1,262 @@
+"""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 _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
+ 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, pickup_obj_name: str) -> None:
+ receptacle_name = self.config.task_sampler_config.place_receptacle_name
+ body = create_mlspaces_body(env.current_data, receptacle_name)
+ 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],
+ 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 = 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()
+ 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, pickup_obj_name)
+ 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/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
new file mode 100644
index 00000000..76f9214f
--- /dev/null
+++ b/scripts/assets/prepare_unitree_g1.py
@@ -0,0 +1,372 @@
+"""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"
+DEFAULT_PELVIS_HEIGHT_M = 0.793
+DEX1_GRASP_SITE_POS = [0.148, 0.0, 0.0]
+# 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 = [20.0, 0.4, 0.04]
+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,
+ "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 = {
+ "leg": (120.0, 12.0),
+ "waist": (80.0, 8.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),
+}
+
+
+@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.",
+ )
+ 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()
+
+
+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:
+ effort = (
+ joint.effort * DEX1_HAND_FORCE_LIMIT_MULTIPLIER
+ if joint.group == "hand"
+ else joint.effort
+ )
+ kwargs["forcelimited"] = 1
+ kwargs["forcerange"] = [-effort, 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 _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:
+ 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=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,
+ 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:
+ 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],
+ 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
+ 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.pos = [0.0, 0.0, DEFAULT_PELVIS_HEIGHT_M]
+ pelvis.add_freejoint(name="floating_base_joint")
+ _add_kinematics_sites(spec, add_debug_grasp_site_axes)
+ _add_dex1_fingertip_contact_pads(spec)
+ _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,
+ args.add_debug_grasp_site_axes,
+ )
+ print(f"Wrote Unitree G1 Dex1.1 model: {output_xml_path}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/scripts/assets/prepare_unitree_g1_tabletop.py b/scripts/assets/prepare_unitree_g1_tabletop.py
new file mode 100644
index 00000000..ae7d8b40
--- /dev/null
+++ b/scripts/assets/prepare_unitree_g1_tabletop.py
@@ -0,0 +1,294 @@
+"""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"
+
+TABLETOP_HEIGHT_M = 0.693
+
+
+@dataclass(frozen=True)
+class TabletopSceneSpec:
+ name: str
+ tabletop_height_m: float
+
+
+SCENE_SPECS = (
+ TabletopSceneSpec(
+ "unitree_g1_tabletop_pelvis_minus_10cm_v1",
+ 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()
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..89257b10 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
@@ -11,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:
@@ -41,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)
@@ -70,18 +82,42 @@ 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:
+ unlocked_move_groups = [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
@@ -91,8 +127,8 @@ def main() -> None:
ret = kinematics.ik(
move_group_id,
target_pose,
- args.unlocked_move_groups,
- robot_config.init_qpos,
+ unlocked_move_groups,
+ robot_view.get_qpos_dict(),
robot_view.base.pose,
)
print(f"IK iteration {i}: {'Success' if ret is not None else 'Failed'}")