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) 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