Keep free objects where the scene puts them when resetting to the robot keyframe

MuJoCo pads a robot keyframe with zeros for joints the scene adds, so every
free body (boxes, balls) started at the world origin inside the robot and
knocked the G1 over. The keyframe is now applied to the robot's own joints
only, in both the validator and the worker.

api/mjcf_check.py

@@ -18,7 +18,7 @@ import xml.etree.ElementTree as ET
 from dataclasses import asdict, dataclass, field
 from pathlib import Path
 
-from .robots import ROBOTS, robot_dir
+from .robots import ROBOTS, reset_to_robot_keyframe, robot_dir
 
 MAX_MJCF_BYTES = 300_000
 SCENE_FILE = "scene.xml"
@@ -174,8 +174,7 @@ def _load_and_step(directory: Path, sim_seconds: float, robot: str | None) -> Va
         return ValidationResult(False, "load", f"robot root body '{ROBOTS[robot].root_body}' missing", stats=stats)
 
     data = mujoco.MjData(model)
-    if model.nkey:
-        mujoco.mj_resetDataKeyframe(model, data, 0)
+    reset_to_robot_keyframe(model, data, ROBOTS[robot].root_body if robot else None)
     steps = max(1, int(math.ceil(sim_seconds / model.opt.timestep)))
     for _ in range(steps):
         mujoco.mj_step(model, data)

api/robots.py

@@ -72,3 +72,38 @@ def robot_facts(menagerie_dir: str, name: str) -> str:
         "unless the scenario says otherwise.",
     ]
     return "\n".join(lines)
+
+
+def reset_to_robot_keyframe(model: mujoco.MjModel, data: mujoco.MjData, root_body: str | None) -> None:
+    """Reset to the robot's first keyframe without disturbing the rest of the scene.
+
+    Robot keyframes (e.g. the G1's "stand") only list the robot's own joints. When a scene adds
+    free bodies, MuJoCo pads the keyframe with zeros for them, which would teleport every box
+    or ball to the world origin, inside the robot. So the keyframe is applied to the joints of
+    the robot's subtree only; every other joint keeps the scene's default pose (qpos0).
+    """
+    mujoco.mj_resetData(model, data)
+    if model.nkey == 0:
+        return
+    root = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, root_body) if root_body else -1
+    if root < 0:  # no robot: the scene's own keyframe describes the whole scene
+        mujoco.mj_resetDataKeyframe(model, data, 0)
+        return
+
+    def in_robot(body: int) -> bool:
+        while body > 0:
+            if body == root:
+                return True
+            body = model.body_parentid[body]
+        return False
+
+    key_qpos = model.key_qpos[0]
+    for j in range(model.njnt):
+        if not in_robot(model.jnt_bodyid[j]):
+            continue
+        start = model.jnt_qposadr[j]
+        end = model.jnt_qposadr[j + 1] if j + 1 < model.njnt else model.nq
+        data.qpos[start:end] = key_qpos[start:end]
+    if model.nu:
+        data.ctrl[:] = model.key_ctrl[0]
+    mujoco.mj_forward(model, data)

tests/test_runs.py

@@ -418,3 +418,22 @@ def test_importing_the_worker_does_not_change_the_environment():
 
     importlib.reload(worker.handler)
     assert dict(os.environ) == before
+
+
+def test_free_objects_keep_their_scene_pose(tmp_path, menagerie):
+    """The robot keyframe must not teleport the scene's free bodies to the origin (seen in CI: G1 knocked over)."""
+    import mujoco
+
+    from api.robots import reset_to_robot_keyframe
+
+    scene = tmp_path / "scene.xml"
+    scene.write_text(GOOD_SCENE)
+    (tmp_path / "g1.xml").write_text((menagerie / "unitree_g1" / "g1.xml").read_text())
+    model = mujoco.MjModel.from_xml_path(str(scene))
+    data = mujoco.MjData(model)
+    assert list(model.key_qpos[0][-7:]) == [0, 0, 0, 1, 0, 0, 0]  # MuJoCo pads the crate to the origin
+    reset_to_robot_keyframe(model, data, "pelvis")
+    crate = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "crate")
+    assert list(np.round(data.xpos[crate], 3)) == [1.0, 1.0, 0.3]  # where the scene put it
+    pelvis = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "pelvis")
+    assert data.xpos[pelvis][2] == pytest.approx(0.8)  # robot from its keyframe

worker/sim_runner.py

@@ -28,6 +28,8 @@ from pathlib import Path
 import mujoco
 import numpy as np
 
+from api.robots import reset_to_robot_keyframe
+
 NS = 1_000_000_000
 VISIBLE_GROUPS = (0, 1, 2)  # MuJoCo's viewer shows groups 0-2 by default; Menagerie puts collision meshes in 3
 
@@ -234,9 +236,7 @@ def simulate(scene: Path, out_dir: Path, config: RunConfig) -> dict:
     out_dir.mkdir(parents=True, exist_ok=True)
     model = mujoco.MjModel.from_xml_path(str(scene))
     data = mujoco.MjData(model)
-    if model.nkey:
-        mujoco.mj_resetDataKeyframe(model, data, 0)
-    mujoco.mj_forward(model, data)
+    reset_to_robot_keyframe(model, data, config.root_body)
     target = _hold_targets(model)
     root = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, config.root_body) if config.root_body else -1