Spaces:
Paused
Paused
File size: 2,260 Bytes
2fc178d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 | 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
|