From 6551d1af2f68dbbb2d4503ea5d070e42d5d7fdf2 Mon Sep 17 00:00:00 2001 From: Tobias Juelg Date: Mon, 26 Oct 2026 10:47:13 -0700 Subject: [PATCH 1/2] fix(block_balance): sequential stages and success only for a free-standing balanced beam Stages advance one at a time, so a stage can no longer be reached without the ones before it. beam_on_cube is evaluated at the current step instead of being latched, and success requires the beam and both rectangles to touch only each other and the cube (not the table or a robot). Previously a beam that touched the cube once and then lay flat on the table with both rectangles on it counted as balanced. All 54 successful sim demos still reach the final stage. --- src/duobench/tasks/block_balance.py | 47 ++++++++++++++++++----------- 1 file changed, 30 insertions(+), 17 deletions(-) diff --git a/src/duobench/tasks/block_balance.py b/src/duobench/tasks/block_balance.py index 7ff341f..fcf6aa6 100644 --- a/src/duobench/tasks/block_balance.py +++ b/src/duobench/tasks/block_balance.py @@ -81,6 +81,18 @@ def _body_has_any_contact(self, sim: Sim, body_id: int) -> bool: return True return False + def _body_touches_only(self, sim: Sim, body_id: int, allowed_body_ids: set[int]) -> bool: + """True if every contact of the body is with one of the allowed bodies (e.g. not the table or a robot link).""" + for i in range(sim.data.ncon): + contact = sim.data.contact[i] + geom1_body_id = sim.model.geom_bodyid[contact.geom1] + geom2_body_id = sim.model.geom_bodyid[contact.geom2] + if body_id == geom1_body_id and geom2_body_id not in allowed_body_ids: + return False + if body_id == geom2_body_id and geom1_body_id not in allowed_body_ids: + return False + return True + def update_internal_state(self, sim: Sim): beam_body_id = mj.mj_name2id(sim.model, mj.mjtObj.mjOBJ_BODY, self.beam_body_name) blue_rectangle_body_id = mj.mj_name2id(sim.model, mj.mjtObj.mjOBJ_BODY, self.blue_rectangle_body_name) @@ -93,9 +105,8 @@ def update_internal_state(self, sim: Sim): beam_grasped = beam_body_id in left_arm_contacts or beam_body_id in right_arm_contacts - beam_on_cube = (self.red_cube_geom_name, self.beam_geom_name) in contact_pairs or self.internal_state[ - "beam_on_cube" - ] + # current step only: a beam that touched the cube once and then fell onto the table is not on the cube + beam_on_cube = (self.red_cube_geom_name, self.beam_geom_name) in contact_pairs cubes_contacted_by_left = { body_id for body_id in (blue_rectangle_body_id, green_rectangle_body_id) if body_id in left_arm_contacts @@ -120,6 +131,7 @@ def update_internal_state(self, sim: Sim): or right_arm_contacts & {beam_body_id, blue_rectangle_body_id, green_rectangle_body_id, red_cube_body_id} ) + task_body_ids = {beam_body_id, blue_rectangle_body_id, green_rectangle_body_id, red_cube_body_id} beam_qvel_adr = sim.model.jnt_dofadr[ mj.mj_name2id(sim.model, mj.mjtObj.mjOBJ_JOINT, self.prefix + "pink_beam_joint") ] @@ -140,6 +152,10 @@ def update_internal_state(self, sim: Sim): and self._body_has_any_contact(sim, beam_body_id) and self._body_has_any_contact(sim, blue_rectangle_body_id) and self._body_has_any_contact(sim, green_rectangle_body_id) + # the structure stands on its own: beam and rectangles touch only each other and the cube, not the table + and self._body_touches_only(sim, beam_body_id, task_body_ids) + and self._body_touches_only(sim, blue_rectangle_body_id, task_body_ids) + and self._body_touches_only(sim, green_rectangle_body_id, task_body_ids) ) if beam_balanced_now: @@ -157,20 +173,17 @@ def update_internal_state(self, sim: Sim): self.update_stage() def update_stage(self): - if self.internal_state["beam_balanced"]: - new_stage = 5 - elif self.internal_state["both_cubes_on_beam"]: - new_stage = 4 - elif self.internal_state["both_cubes_grasped"]: - new_stage = 3 - elif self.internal_state["beam_on_cube"]: - new_stage = 2 - elif self.internal_state["beam_grasped"]: - new_stage = 1 - else: - new_stage = 0 - - self.stage = max(self.stage, new_stage) + # stages are sequential: stage n + 1 can only be reached from stage n + if self.stage == 0 and self.internal_state["beam_grasped"]: + self.stage = 1 + elif self.stage == 1 and self.internal_state["beam_on_cube"]: + self.stage = 2 + elif self.stage == 2 and self.internal_state["both_cubes_grasped"]: + self.stage = 3 + elif self.stage == 3 and self.internal_state["both_cubes_on_beam"]: + self.stage = 4 + elif self.stage == 4 and self.internal_state["beam_balanced"]: + self.stage = 5 @dataclass(kw_only=True) From 1d7a33b705fe73c847d314187d67b7a6c090026f Mon Sep 17 00:00:00 2001 From: Tobias Juelg Date: Mon, 26 Oct 2026 10:47:13 -0700 Subject: [PATCH 2/2] test(block_balance): stages are sequential, a beam on the table is not balanced --- src/tests/test_block_balance_stage.py | 132 ++++++++++++++++++++++++++ 1 file changed, 132 insertions(+) create mode 100644 src/tests/test_block_balance_stage.py diff --git a/src/tests/test_block_balance_stage.py b/src/tests/test_block_balance_stage.py new file mode 100644 index 0000000..cd87bfd --- /dev/null +++ b/src/tests/test_block_balance_stage.py @@ -0,0 +1,132 @@ +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.block_balance # noqa: F401 registers duobench/block_balance +from duobench.tasks.block_balance import BlockBalanceStage, BlockBalanceTaskConfig + +warnings.filterwarnings("ignore", category=UserWarning, module=r"gymnasium\.core") + +PREFIX = BlockBalanceTaskConfig().prefix +# half heights of the box geoms +CUBE_HALF_Z, BEAM_HALF_Z, RECT_HALF_Z = 0.02, 0.01, 0.045 + + +def test_stages_are_sequential(): + stage = BlockBalanceStage(BlockBalanceTaskConfig(), prefix=PREFIX) + stage.internal_state["both_cubes_on_beam"] = True + stage.internal_state["beam_balanced"] = True + stage.update_stage() + assert stage.stage == 0 + + for key, expected in ( + ("beam_grasped", 1), + ("beam_on_cube", 2), + ("both_cubes_grasped", 3), + ("both_cubes_on_beam", 4), + ("beam_balanced", 5), + ): + stage.internal_state[key] = True + stage.update_stage() + assert stage.stage == expected + + +@pytest.fixture(scope="module") +def env(): + os.environ.setdefault("MUJOCO_GL", "egl") + env_creator = gym.envs.registry["duobench/block_balance"].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(sim, name: str, pos: np.ndarray): + model, data = sim.model, sim.data + joint = mj.mj_name2id(model, mj.mjtObj.mjOBJ_JOINT, PREFIX + name + "_joint") + qpos, qvel = model.jnt_qposadr[joint], model.jnt_dofadr[joint] + data.qpos[qpos : qpos + 3] = pos + data.qpos[qpos + 3 : qpos + 7] = [1.0, 0.0, 0.0, 0.0] + data.qvel[qvel : qvel + 6] = 0.0 + + +def _settle_and_track(sim, stage: BlockBalanceStage, start_stage: int, steps: int = 300, **state) -> int: + """Let the objects settle (robots untouched), then run the stage check for a few steps from start_stage + (with internal_state entries preset from earlier steps).""" + mj.mj_forward(sim.model, sim.data) + for _ in range(steps): + mj.mj_step(sim.model, sim.data) + stage.reset() + stage.stage = start_stage + stage.internal_state.update(state) + for _ in range(10): + mj.mj_step(sim.model, sim.data) + stage.update_internal_state(sim) + return stage.stage + + +def _table_and_center(sim) -> tuple[float, np.ndarray]: + cube = sim.data.body(PREFIX + "big_red_cube").xpos.copy() + return cube[2] - CUBE_HALF_Z, cube[:2] + + +def test_balanced_structure_succeeds(env): + sim = env.get_wrapper_attr("sim") + env.reset(seed=0) + table, (x, y) = _table_and_center(sim) + beam_z = table + 2 * CUBE_HALF_Z + BEAM_HALF_Z + rect_z = beam_z + BEAM_HALF_Z + RECT_HALF_Z + _place(sim, "big_red_cube", np.array([x, y, table + CUBE_HALF_Z])) + _place(sim, "pink_beam", np.array([x, y, beam_z])) + _place(sim, "blue_rectangle", np.array([x - 0.07, y, rect_z])) + _place(sim, "green_rectangle", np.array([x + 0.07, y, rect_z])) + + stage = BlockBalanceStage(BlockBalanceTaskConfig(), prefix=PREFIX) + assert _settle_and_track(sim, stage, start_stage=4) == 5 + + +def test_beam_on_table_is_not_balanced(env): + """The beam touched the cube once, then lies flat on the table with both rectangles on it: not a success.""" + sim = env.get_wrapper_attr("sim") + env.reset(seed=0) + table, (x, y) = _table_and_center(sim) + beam_z = table + BEAM_HALF_Z + rect_z = beam_z + BEAM_HALF_Z + RECT_HALF_Z + _place(sim, "big_red_cube", np.array([x, y + 0.0395, table + CUBE_HALF_Z])) # touching the beam's side + _place(sim, "pink_beam", np.array([x, y, beam_z])) + _place(sim, "blue_rectangle", np.array([x - 0.07, y, rect_z])) + _place(sim, "green_rectangle", np.array([x + 0.07, y, rect_z])) + + stage = BlockBalanceStage(BlockBalanceTaskConfig(), prefix=PREFIX) + final = _settle_and_track(sim, stage, start_stage=4) + assert stage.internal_state["both_cubes_on_beam"] # the rectangles do stand on the beam + assert final == 4 + + +def test_beam_on_table_without_cube_contact_is_not_on_cube(env): + """Old check: beam_on_cube stayed latched after the beam left the cube. New: current contact only.""" + sim = env.get_wrapper_attr("sim") + env.reset(seed=0) + table, (x, y) = _table_and_center(sim) + beam_z = table + BEAM_HALF_Z + rect_z = beam_z + BEAM_HALF_Z + RECT_HALF_Z + _place(sim, "big_red_cube", np.array([x, y + 0.15, table + CUBE_HALF_Z])) + _place(sim, "pink_beam", np.array([x, y, beam_z])) + _place(sim, "blue_rectangle", np.array([x - 0.07, y, rect_z])) + _place(sim, "green_rectangle", np.array([x + 0.07, y, rect_z])) + + stage = BlockBalanceStage(BlockBalanceTaskConfig(), prefix=PREFIX) + final = _settle_and_track(sim, stage, start_stage=4, beam_on_cube=True) # touched the cube earlier + assert not stage.internal_state["beam_on_cube"] + assert not stage.internal_state["both_cubes_on_beam"] + assert final == 4