From b27f27a42c5b432716d22696d21f4d76cdcee6d6 Mon Sep 17 00:00:00 2001 From: Tobias Juelg Date: Mon, 26 Oct 2026 00:00:44 -0700 Subject: [PATCH 1/2] fix(join_blocks): require the peg to be inside the socket for success The connection checks accepted any peg within 5 cm of the socket marker on every world axis (wall: 5 cm in x/y). The peg is only about 2 x 2.6 x 2.5 cm, so two blocks held side by side already counted as connected, and the latched blocks_connected let an episode succeed after the blocks came apart. block_dist / wall_dist were defined but unused. - Measure the peg offset in the socket marker frame and require it to be below block_dist / wall_dist on every axis (orientation independent). - Set both to 15 mm, calibrated on the 77 sim demos (replayed sim_state): the socket holds the peg sideways (clear of it only from ~22 mm) and is open along its z axis (insertion depth, 25 mm socket). With 15 mm, 48/50 successful demos reach stage 3 (10 mm: 40/50, 20 mm: 50/50). - Success (stage 2 -> 3) needs blocks_connected and wall_connected at the same step while both grippers hold the blocks; the latched blocks_connected only advances stage 1 -> 2. --- src/duobench/tasks/join_blocks.py | 44 +++++++++++++++++++++++++------ 1 file changed, 36 insertions(+), 8 deletions(-) diff --git a/src/duobench/tasks/join_blocks.py b/src/duobench/tasks/join_blocks.py index 36df93b..d31366b 100644 --- a/src/duobench/tasks/join_blocks.py +++ b/src/duobench/tasks/join_blocks.py @@ -55,8 +55,11 @@ class JoinBlocksTaskConfig(BaseTaskConfig): wall = "wall_wall" wall_peg_name = "wall_wall_peg" - block_dist = 0.01 - wall_dist = 0.01 + # max peg-center offset from the socket marker on every axis of the marker frame. The socket holds the peg + # sideways (it is clear of the socket only from ~22 mm on) and is open along its z axis (insertion depth, + # 25 mm socket height). 15 mm: 48/50 successful sim demos pass, a peg held next to the socket does not. + block_dist = 0.015 + wall_dist = 0.015 class JoinBlocksTask(Task[JoinBlocksTaskConfig]): @@ -119,6 +122,20 @@ def add_task_env(cfg: JoinBlocksTaskConfig, env: gym.Env, simulation: Sim, env_c return TaskStageWrapper(env, JoinBlocksStage(cfg)) +def peg_in_socket(peg_pos: np.ndarray, marker_pos: np.ndarray, marker_xmat: np.ndarray, max_dist: float) -> bool: + """Whether the peg center lies within max_dist of the socket marker on every axis of the marker frame.""" + offset = np.asarray(marker_xmat).reshape(3, 3).T @ (np.asarray(peg_pos) - np.asarray(marker_pos)) + return bool(np.all(np.abs(offset) < max_dist)) + + +def get_geom_xmat(model, data, name): + geom_id = mj.mj_name2id(model, mj.mjtObj.mjOBJ_GEOM, name) + if geom_id == -1: + msg = f"Geom '{name}' not found" + raise ValueError(msg) + return data.geom_xmat[geom_id] + + class JoinBlocksStage(TaskStage): def __init__(self, cfg: JoinBlocksTaskConfig): super().__init__( @@ -154,14 +171,25 @@ def update_internal_state(self, sim: Sim): approaching_done = bool(left_contacts & block_body_ids) and bool(right_contacts & block_body_ids) - diff_block = abs(block_peg_pos - block_marker_pos) - blocks_connected = diff_block[0] < 0.05 and diff_block[1] < 0.05 and diff_block[2] < 0.05 - - diff_wall = abs(wall_peg_pos - wall_marker_pos) - wall_connected = diff_wall[0] < 0.05 and diff_wall[1] < 0.05 and diff_wall[2] < 0.01 + blocks_connected = peg_in_socket( + block_peg_pos, + block_marker_pos, + get_geom_xmat(sim.model, sim.data, self.cfg.marker_block), + self.cfg.block_dist, + ) + wall_connected = peg_in_socket( + wall_peg_pos, + wall_marker_pos, + get_geom_xmat(sim.model, sim.data, self.cfg.marker_wall), + self.cfg.wall_dist, + ) + # success needs both joints at the same time: the latched blocks_connected below only advances the stage holding_on_wall = ( - wall_connected and bool(left_contacts & block_body_ids) and bool(right_contacts & block_body_ids) + blocks_connected + and wall_connected + and bool(left_contacts & block_body_ids) + and bool(right_contacts & block_body_ids) ) self.internal_state["approaching_done"] = self.internal_state["approaching_done"] or approaching_done From 281c72287fc74ecdc205211283f40eb8bec54a9f Mon Sep 17 00:00:00 2001 From: Tobias Juelg Date: Mon, 26 Oct 2026 00:00:44 -0700 Subject: [PATCH 2/2] test(join_blocks): blocks side by side are not connected, an inserted peg is --- src/tests/test_join_blocks_stage.py | 94 +++++++++++++++++++++++++++++ 1 file changed, 94 insertions(+) create mode 100644 src/tests/test_join_blocks_stage.py diff --git a/src/tests/test_join_blocks_stage.py b/src/tests/test_join_blocks_stage.py new file mode 100644 index 0000000..cd94bdd --- /dev/null +++ b/src/tests/test_join_blocks_stage.py @@ -0,0 +1,94 @@ +import os +import warnings + +import gymnasium as gym +import mujoco as mj +import numpy as np +import pytest +from rcs.envs.scenes import SimEnvCreator + +import duobench.tasks.join_blocks # noqa: F401 registers duobench/join_blocks +from duobench.tasks.join_blocks import ( + JoinBlocksStage, + JoinBlocksTaskConfig, + get_geom_xmat, + peg_in_socket, +) +from duobench.utils.helper_wrappers import get_geom_pos + +warnings.filterwarnings("ignore", category=UserWarning, module=r"gymnasium\.core") + +# p_block orientation relative to h_block with the peg in the socket (from the successful sim demos) +P_IN_H_QUAT = np.array([np.sqrt(0.5), 0.0, 0.0, -np.sqrt(0.5)]) + + +def test_peg_in_socket_uses_marker_frame(): + marker_xmat = np.array([[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]]) # 90 deg about z + marker = np.array([0.5, 0.2, 0.1]) + + assert peg_in_socket(marker + np.array([0.001, -0.002, 0.012]), marker, marker_xmat, 0.015) + assert not peg_in_socket(marker + np.array([0.03, 0.0, 0.0]), marker, marker_xmat, 0.015) + assert not peg_in_socket(marker + np.array([0.0, 0.0, 0.02]), marker, marker_xmat, 0.015) + + +@pytest.fixture(scope="module") +def env(): + os.environ.setdefault("MUJOCO_GL", "egl") + env_creator = gym.envs.registry["duobench/join_blocks"].entry_point + assert isinstance(env_creator, SimEnvCreator) + cfg = env_creator.config() + cfg.headless = True + cfg.camera_cfgs = None + cfg.camera_adds = None + env = env_creator.create_env(cfg) + env.reset(seed=0) + yield env + env.close() + + +def _place_p_block(sim, cfg: JoinBlocksTaskConfig, offset_in_marker_frame: np.ndarray) -> float: + """Put the p_block peg at the h_block socket marker (+ offset); returns the deepest block penetration.""" + model, data = sim.model, sim.data + h_block = mj.mj_name2id(model, mj.mjtObj.mjOBJ_BODY, cfg.h_block) + p_block = mj.mj_name2id(model, mj.mjtObj.mjOBJ_BODY, cfg.p_block) + qpos = model.jnt_qposadr[mj.mj_name2id(model, mj.mjtObj.mjOBJ_JOINT, cfg.p_block_joint)] + + mj.mju_mulQuat(data.qpos[qpos + 3 : qpos + 7], data.xquat[h_block], P_IN_H_QUAT) + for _ in range(2): # orientation first, then shift the peg center onto the target + mj.mj_forward(model, data) + marker_xmat = get_geom_xmat(model, data, cfg.marker_block).reshape(3, 3) + target = get_geom_pos(model, data, cfg.marker_block) + marker_xmat @ offset_in_marker_frame + data.qpos[qpos : qpos + 3] += target - get_geom_pos(model, data, cfg.block_peg_name) + mj.mj_forward(model, data) + + h_geoms = [g for g in range(model.ngeom) if model.geom_bodyid[g] == h_block and model.geom_contype[g]] + p_geoms = [g for g in range(model.ngeom) if model.geom_bodyid[g] == p_block and model.geom_contype[g]] + fromto = np.zeros(6) + return min(mj.mj_geomDistance(model, data, h, p, 0.05, fromto) for h in h_geoms for p in p_geoms) + + +@pytest.mark.parametrize( + ("offset", "connected"), + [ + ((0.0, 0.0, 0.0), True), # peg fully inserted + ((0.0, 0.0, 0.01), True), # peg 1 cm short of full insertion depth + ((0.035, 0.0, 0.0), False), # blocks side by side, peg next to the socket + ((0.0, -0.03, 0.0), False), # blocks side by side, peg in front of the socket opening + ((0.0, 0.0, 0.03), False), # peg on top of the socket + ], +) +def test_blocks_connected_requires_inserted_peg(env, offset, connected): + cfg = JoinBlocksTaskConfig() + sim = env.unwrapped.sim + penetration = _place_p_block(sim, cfg, np.array(offset)) + assert penetration > -0.003, "test pose is not physically reachable" + + stage = JoinBlocksStage(cfg) + stage.update_internal_state(sim) + assert stage.internal_state["blocks_connected"] is connected + + # the previous check (5 cm on every world axis) counted all of these as connected + diff = np.abs( + get_geom_pos(sim.model, sim.data, cfg.block_peg_name) - get_geom_pos(sim.model, sim.data, cfg.marker_block) + ) + assert np.all(diff < 0.05)