Spaces:
Paused
Paused
| import mujoco | |
| import numpy as np | |
| import pytest | |
| from conftest import GOOD_SCENE | |
| pytest.importorskip("pxr") | |
| from pxr import Usd, UsdGeom # noqa: E402 | |
| from api.mjcf_check import prepare_scene_dir # noqa: E402 | |
| from worker.usd_scene import build_stage # noqa: E402 | |
| def test_replay_stage_matches_mujoco(tmp_path, menagerie): | |
| """Every drawn geom's world position in the USD equals MuJoCo's geom_xpos, frame by frame.""" | |
| scene_dir = prepare_scene_dir(GOOD_SCENE, "unitree_g1", menagerie) | |
| model = mujoco.MjModel.from_xml_path(str(scene_dir / "scene.xml")) | |
| data = mujoco.MjData(model) | |
| data.qvel[:] = np.random.default_rng(0).normal(0, 1, model.nv) # tumble things around | |
| pos, quat, geom_xpos = [], [], [] | |
| for _ in range(20): | |
| for _ in range(10): | |
| mujoco.mj_step(model, data) | |
| pos.append(data.xpos.copy()) | |
| quat.append(data.xquat.copy()) | |
| geom_xpos.append(data.geom_xpos.copy()) | |
| names = [mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_BODY, b) or f"body_{b}" for b in range(model.nbody)] | |
| np.savez(tmp_path / "traj.npz", body_pos=pos, body_quat=quat, body_names=names, fps=30) | |
| info = build_stage(scene_dir / "scene.xml", tmp_path / "traj.npz", tmp_path / "scene.usdc", "pelvis") | |
| assert info["frames"] == 20 and info["geoms_drawn"] >= 3 | |
| stage = Usd.Stage.Open(str(tmp_path / "scene.usdc")) | |
| checked = 0 | |
| for prim in stage.Traverse(): | |
| name = prim.GetName() | |
| if not (name.startswith("g") and name[1:].isdigit()): | |
| continue | |
| g = int(name[1:]) | |
| if model.geom_type[g] == mujoco.mjtGeom.mjGEOM_PLANE: | |
| continue | |
| for i in (0, 9, 19): | |
| m = UsdGeom.Xformable(prim).ComputeLocalToWorldTransform(Usd.TimeCode(i)) | |
| centre = np.array(m.Transform((0, 0, 0))) if model.geom_type[g] != mujoco.mjtGeom.mjGEOM_MESH else None | |
| if centre is not None: | |
| assert np.allclose(centre, geom_xpos[i][g], atol=1e-4), (g, i) | |
| checked += 1 | |
| assert checked >= 6 | |
| cam = UsdGeom.Xformable(stage.GetPrimAtPath("/World/Camera")) | |
| eye = np.array(cam.ComputeLocalToWorldTransform(Usd.TimeCode(0)).ExtractTranslation()) | |
| assert eye[2] > pos[0][1][2] # looking down at the robot | |