32a241c28f
Some checks failed
nightly / Test against latest dependencies (py3.10) (push) Has been cancelled
nightly / Test against latest dependencies (py3.13) (push) Has been cancelled
tests / tests (3.13, locked) (push) Has been cancelled
tests / tests (3.13, unlocked) (push) Has been cancelled
tests / pyright (3.10) (push) Has been cancelled
tests / lint-format (push) Has been cancelled
tests / tests (3.10, locked) (push) Has been cancelled
tests / tests (3.11, locked) (push) Has been cancelled
tests / tests (3.12, locked) (push) Has been cancelled
tests / pyright (3.11) (push) Has been cancelled
tests / pyright (3.12) (push) Has been cancelled
tests / pyright (3.13) (push) Has been cancelled
tests / ty-check (3.10) (push) Has been cancelled
tests / ty-check (3.11) (push) Has been cancelled
tests / ty-check (3.12) (push) Has been cancelled
tests / ty-check (3.13) (push) Has been cancelled
tests / stubs (push) Has been cancelled
tests / smoke-test (push) Has been cancelled
Docker / check_paths (push) Has been cancelled
docs / build (push) Has been cancelled
Docker / build (push) Has been cancelled
Upstream: https://github.com/michaelgillett/mjlab Upstream-Commit: c19f713c415a699a79d71cd96aa13c3104a05047 Upstream-Branch: main
1773 lines
54 KiB
Python
1773 lines
54 KiB
Python
"""Tests for per-world mesh variant support."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any, cast
|
|
|
|
import mujoco
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
|
|
from mjlab.entity import (
|
|
EntityCfg,
|
|
VariantEntityCfg,
|
|
)
|
|
from mjlab.entity.variants import (
|
|
SlotKey,
|
|
VariantGeomSpec,
|
|
VariantSlot,
|
|
allocate_worlds,
|
|
build_variant_model,
|
|
)
|
|
from mjlab.viewer.model_sync import (
|
|
disable_model_sameframe_shortcuts,
|
|
sync_model_fields,
|
|
)
|
|
|
|
# Helpers: variant specs with visual + collision mesh geoms.
|
|
|
|
|
|
def _sphere_2col_spec() -> mujoco.MjSpec:
|
|
"""Sphere: 1 visual + 2 collision geoms."""
|
|
spec = mujoco.MjSpec()
|
|
mv = spec.add_mesh()
|
|
mv.name = "visual"
|
|
mv.make_sphere(subdivision=3)
|
|
for i in range(2):
|
|
mc = spec.add_mesh()
|
|
mc.name = f"col_{i}"
|
|
mc.make_sphere(subdivision=1)
|
|
body = spec.worldbody.add_body()
|
|
body.name = "prop"
|
|
body.add_freejoint()
|
|
gv = body.add_geom()
|
|
gv.name = "visual"
|
|
gv.type = mujoco.mjtGeom.mjGEOM_MESH
|
|
gv.meshname = "visual"
|
|
gv.contype = 0
|
|
gv.conaffinity = 0
|
|
for i in range(2):
|
|
gc = body.add_geom()
|
|
gc.name = f"col_{i}"
|
|
gc.type = mujoco.mjtGeom.mjGEOM_MESH
|
|
gc.meshname = f"col_{i}"
|
|
return spec
|
|
|
|
|
|
def _cone_4col_spec() -> mujoco.MjSpec:
|
|
"""Cone: 1 visual + 4 collision geoms (more than sphere)."""
|
|
spec = mujoco.MjSpec()
|
|
mv = spec.add_mesh()
|
|
mv.name = "visual"
|
|
mv.make_cone(nedge=8, radius=0.05)
|
|
for i in range(4):
|
|
mc = spec.add_mesh()
|
|
mc.name = f"col_{i}"
|
|
mc.make_sphere(subdivision=1)
|
|
body = spec.worldbody.add_body()
|
|
body.name = "prop"
|
|
body.add_freejoint()
|
|
gv = body.add_geom()
|
|
gv.name = "visual"
|
|
gv.type = mujoco.mjtGeom.mjGEOM_MESH
|
|
gv.meshname = "visual"
|
|
gv.contype = 0
|
|
gv.conaffinity = 0
|
|
for i in range(4):
|
|
gc = body.add_geom()
|
|
gc.name = f"col_{i}"
|
|
gc.type = mujoco.mjtGeom.mjGEOM_MESH
|
|
gc.meshname = f"col_{i}"
|
|
return spec
|
|
|
|
|
|
def _simple_sphere_spec() -> mujoco.MjSpec:
|
|
"""Single-geom sphere for simple tests."""
|
|
spec = mujoco.MjSpec()
|
|
m = spec.add_mesh()
|
|
m.name = "sphere"
|
|
m.make_sphere(subdivision=2)
|
|
body = spec.worldbody.add_body()
|
|
body.name = "prop"
|
|
body.add_freejoint()
|
|
g = body.add_geom()
|
|
g.name = "visual"
|
|
g.type = mujoco.mjtGeom.mjGEOM_MESH
|
|
g.meshname = "sphere"
|
|
return spec
|
|
|
|
|
|
def _simple_cone_spec() -> mujoco.MjSpec:
|
|
"""Single-geom cone for simple tests."""
|
|
spec = mujoco.MjSpec()
|
|
m = spec.add_mesh()
|
|
m.name = "cone"
|
|
m.make_cone(nedge=8, radius=0.05)
|
|
body = spec.worldbody.add_body()
|
|
body.name = "prop"
|
|
body.add_freejoint()
|
|
g = body.add_geom()
|
|
g.name = "visual"
|
|
g.type = mujoco.mjtGeom.mjGEOM_MESH
|
|
g.meshname = "cone"
|
|
return spec
|
|
|
|
|
|
def _hinge_spec() -> mujoco.MjSpec:
|
|
"""Object with a hinge joint (incompatible with freejoint variants)."""
|
|
spec = mujoco.MjSpec()
|
|
m = spec.add_mesh()
|
|
m.name = "box"
|
|
m.make_sphere(subdivision=1)
|
|
body = spec.worldbody.add_body()
|
|
body.name = "prop"
|
|
j = body.add_joint()
|
|
j.name = "hinge"
|
|
j.type = mujoco.mjtJoint.mjJNT_HINGE
|
|
g = body.add_geom()
|
|
g.name = "visual"
|
|
g.type = mujoco.mjtGeom.mjGEOM_MESH
|
|
g.meshname = "box"
|
|
return spec
|
|
|
|
|
|
def _build_scene_with_variants(
|
|
variant_a_fn, variant_b_fn, *, weight_a=0.5, weight_b=0.5
|
|
):
|
|
"""Build a scene spec + variant_info from two variant spec_fns."""
|
|
cfg = VariantEntityCfg(
|
|
variants={"a": variant_a_fn, "b": variant_b_fn},
|
|
assignment={"a": weight_a, "b": weight_b},
|
|
)
|
|
entity = cfg.build()
|
|
assert entity.variant_metadata is not None
|
|
scene_spec = mujoco.MjSpec()
|
|
frame = scene_spec.worldbody.add_frame()
|
|
scene_spec.attach(entity.spec, prefix="object/", frame=frame)
|
|
return scene_spec, [("object/", entity.variant_metadata)]
|
|
|
|
|
|
# allocate_worlds.
|
|
|
|
|
|
def test_allocate_worlds_proportional():
|
|
result = allocate_worlds((0.6, 0.4), 10)
|
|
assert len(result) == 10
|
|
assert result.count(0) == 6
|
|
assert result.count(1) == 4
|
|
|
|
|
|
def test_allocate_worlds_uniform():
|
|
result = allocate_worlds((1.0, 1.0), 8)
|
|
assert result.count(0) == 4
|
|
assert result.count(1) == 4
|
|
|
|
|
|
def test_allocate_worlds_single_variant():
|
|
result = allocate_worlds((1.0,), 5)
|
|
assert result == [0, 0, 0, 0, 0]
|
|
|
|
|
|
def test_allocate_worlds_zero_weight_skips_variant():
|
|
"""A zero-weight variant gets zero worlds; the rest split nworld."""
|
|
result = allocate_worlds((1.0, 0.0, 1.0), 10)
|
|
assert len(result) == 10
|
|
assert result.count(1) == 0
|
|
assert result.count(0) == 5
|
|
assert result.count(2) == 5
|
|
|
|
|
|
def test_allocate_worlds_rejects_negative_weight():
|
|
with pytest.raises(ValueError, match="non-negative"):
|
|
allocate_worlds((1.0, -0.1), 10)
|
|
|
|
|
|
def test_allocate_worlds_rejects_all_zero():
|
|
with pytest.raises(ValueError, match="positive sum"):
|
|
allocate_worlds((0.0, 0.0), 10)
|
|
|
|
|
|
def test_allocate_worlds_largest_remainder_sums_to_nworld():
|
|
"""Largest-remainder rounding must always allocate exactly nworld worlds."""
|
|
for nworld in (3, 7, 100, 1000):
|
|
result = allocate_worlds((1.0, 1.0, 1.0), nworld)
|
|
assert len(result) == nworld
|
|
# Difference between any two variant counts is at most 1 (uniform).
|
|
counts = [result.count(i) for i in range(3)]
|
|
assert max(counts) - min(counts) <= 1
|
|
|
|
|
|
# assignment_fn override.
|
|
|
|
|
|
def test_assignment_fn_overrides_weights():
|
|
"""When assignment_fn is set, it dictates the per-world variant indices."""
|
|
cfg = VariantEntityCfg(
|
|
variants={
|
|
"sphere": _simple_sphere_spec,
|
|
"cone": _simple_cone_spec, # ignored
|
|
},
|
|
assignment=lambda nworld: [0, 1, 0, 1] * (nworld // 4),
|
|
)
|
|
entity = cfg.build()
|
|
assert entity.variant_metadata is not None
|
|
scene_spec = mujoco.MjSpec()
|
|
frame = scene_spec.worldbody.add_frame()
|
|
scene_spec.attach(entity.spec, prefix="object/", frame=frame)
|
|
result = build_variant_model(scene_spec, 8, [("object/", entity.variant_metadata)])
|
|
w2v = result.world_to_variant["object/"]
|
|
assert list(w2v) == [0, 1, 0, 1, 0, 1, 0, 1]
|
|
|
|
|
|
def test_assignment_fn_seeded_is_nworld_invariant():
|
|
"""Per-world independent RNG draws make world W's variant a function of W
|
|
alone, independent of nworld."""
|
|
weights = (1.0, 2.0, 1.0)
|
|
cum = np.cumsum(np.asarray(weights) / sum(weights))
|
|
|
|
def seeded_assignment(seed: int):
|
|
def fn(nworld: int) -> list[int]:
|
|
return [
|
|
int(np.searchsorted(cum, np.random.default_rng((seed, w)).random()))
|
|
for w in range(nworld)
|
|
]
|
|
|
|
return fn
|
|
|
|
fn_64 = seeded_assignment(seed=42)(64)
|
|
fn_256 = seeded_assignment(seed=42)(256)
|
|
# World 0..63 must agree across both batch sizes.
|
|
assert fn_64 == fn_256[:64]
|
|
|
|
|
|
def test_assignment_fn_rejects_wrong_length():
|
|
cfg = VariantEntityCfg(
|
|
variants={
|
|
"sphere": _simple_sphere_spec,
|
|
"cone": _simple_cone_spec,
|
|
},
|
|
assignment=lambda nworld: [0] * (nworld - 1), # one short
|
|
)
|
|
entity = cfg.build()
|
|
assert entity.variant_metadata is not None
|
|
scene_spec = mujoco.MjSpec()
|
|
frame = scene_spec.worldbody.add_frame()
|
|
scene_spec.attach(entity.spec, prefix="object/", frame=frame)
|
|
with pytest.raises(ValueError, match="returned .* indices but nworld="):
|
|
build_variant_model(scene_spec, 4, [("object/", entity.variant_metadata)])
|
|
|
|
|
|
def test_assignment_fn_rejects_out_of_range_index():
|
|
cfg = VariantEntityCfg(
|
|
variants={
|
|
"sphere": _simple_sphere_spec,
|
|
"cone": _simple_cone_spec,
|
|
},
|
|
assignment=lambda nworld: [0, 1, 0, 99], # 99 is out of range
|
|
)
|
|
entity = cfg.build()
|
|
assert entity.variant_metadata is not None
|
|
scene_spec = mujoco.MjSpec()
|
|
frame = scene_spec.worldbody.add_frame()
|
|
scene_spec.attach(entity.spec, prefix="object/", frame=frame)
|
|
with pytest.raises(ValueError, match="returned variant index 99"):
|
|
build_variant_model(scene_spec, 4, [("object/", entity.variant_metadata)])
|
|
|
|
|
|
# Entity merging.
|
|
|
|
|
|
def test_entity_builds_with_variants():
|
|
cfg = VariantEntityCfg(
|
|
variants={
|
|
"sphere": _simple_sphere_spec,
|
|
"cone": _simple_cone_spec,
|
|
},
|
|
)
|
|
entity = cfg.build()
|
|
meta = entity.variant_metadata
|
|
assert meta is not None
|
|
assert meta.variant_names == ("sphere", "cone")
|
|
assert meta.num_mesh_geoms == 1
|
|
mesh_names = [m.name for m in entity.spec.meshes]
|
|
assert any("sphere" in n for n in mesh_names)
|
|
assert any("cone" in n for n in mesh_names)
|
|
|
|
|
|
def test_multi_geom_body_padding():
|
|
"""Sphere (3 geoms) + cone (5 geoms) -> body padded to 5 mesh geoms."""
|
|
cfg = VariantEntityCfg(
|
|
variants={
|
|
"sphere": _sphere_2col_spec,
|
|
"cone": _cone_4col_spec,
|
|
},
|
|
)
|
|
entity = cfg.build()
|
|
meta = entity.variant_metadata
|
|
assert meta is not None
|
|
assert meta.num_mesh_geoms == 5 # max(3, 5)
|
|
# Sphere: 3 real + 2 padding (None).
|
|
assert sum(1 for n in meta.variant_mesh_names[0] if n is None) == 2
|
|
# Cone: 5 real, no padding.
|
|
assert all(n is not None for n in meta.variant_mesh_names[1])
|
|
|
|
|
|
# Validation.
|
|
|
|
|
|
def test_mismatched_joint_structure_raises():
|
|
cfg = VariantEntityCfg(
|
|
variants={
|
|
"sphere": _simple_sphere_spec,
|
|
"hinge": _hinge_spec,
|
|
},
|
|
)
|
|
with pytest.raises(ValueError, match="joint"):
|
|
cfg.build()
|
|
|
|
|
|
def test_single_variant_builds():
|
|
"""A single variant degenerates cleanly; useful for templated variant sets."""
|
|
cfg = VariantEntityCfg(
|
|
variants={"only": _simple_sphere_spec},
|
|
)
|
|
entity = cfg.build()
|
|
assert entity.variant_metadata is not None
|
|
assert entity.variant_metadata.variant_names == ("only",)
|
|
|
|
|
|
def test_empty_variants_raises():
|
|
cfg = VariantEntityCfg(variants={})
|
|
with pytest.raises(ValueError, match="at least one"):
|
|
cfg.build()
|
|
|
|
|
|
def _fixed_base_sphere_spec() -> mujoco.MjSpec:
|
|
"""Fixed-base sphere variant (no free joint): currently unsupported."""
|
|
spec = mujoco.MjSpec()
|
|
m = spec.add_mesh(name="sphere")
|
|
m.make_sphere(subdivision=2)
|
|
body = spec.worldbody.add_body(name="prop")
|
|
body.add_geom(type=mujoco.mjtGeom.mjGEOM_MESH, meshname="sphere")
|
|
return spec
|
|
|
|
|
|
def test_fixed_base_variants_rejected():
|
|
"""Variants must be floating-base; fixed-base raises with a clear message."""
|
|
cfg = VariantEntityCfg(
|
|
variants={
|
|
"a": _fixed_base_sphere_spec,
|
|
"b": _fixed_base_sphere_spec,
|
|
},
|
|
)
|
|
with pytest.raises(ValueError, match="floating-base"):
|
|
cfg.build()
|
|
|
|
|
|
def test_setting_spec_fn_on_variant_cfg_raises():
|
|
"""VariantEntityCfg.spec_fn is unused; setting it should fail loudly."""
|
|
with pytest.raises(ValueError, match="spec_fn cannot be set"):
|
|
VariantEntityCfg(
|
|
variants={"only": _simple_sphere_spec},
|
|
spec_fn=_simple_sphere_spec,
|
|
)
|
|
|
|
|
|
# Recursive validation: helpers and tests.
|
|
|
|
|
|
def _articulated_spec(
|
|
*,
|
|
root_mesh: str = "root_mesh",
|
|
child_mesh: str = "child_mesh",
|
|
with_grandchild: bool = False,
|
|
) -> mujoco.MjSpec:
|
|
"""Root + child body (hinge joint). Optional grandchild for arity tests."""
|
|
spec = mujoco.MjSpec()
|
|
rm = spec.add_mesh(name=root_mesh)
|
|
rm.make_sphere(subdivision=2)
|
|
cm = spec.add_mesh(name=child_mesh)
|
|
cm.make_sphere(subdivision=2)
|
|
root = spec.worldbody.add_body(name="prop")
|
|
root.add_freejoint()
|
|
rg = root.add_geom()
|
|
rg.name = "root_geom"
|
|
rg.type = mujoco.mjtGeom.mjGEOM_MESH
|
|
rg.meshname = root_mesh
|
|
child = root.add_body(name="lid")
|
|
cj = child.add_joint()
|
|
cj.name = "hinge"
|
|
cj.type = mujoco.mjtJoint.mjJNT_HINGE
|
|
cg = child.add_geom()
|
|
cg.name = "child_geom"
|
|
cg.type = mujoco.mjtGeom.mjGEOM_MESH
|
|
cg.meshname = child_mesh
|
|
if with_grandchild:
|
|
gm = spec.add_mesh(name="grand_mesh")
|
|
gm.make_sphere(subdivision=2)
|
|
grand = child.add_body(name="grand")
|
|
grand.add_geom(
|
|
name="grand_geom", type=mujoco.mjtGeom.mjGEOM_MESH, meshname="grand_mesh"
|
|
)
|
|
return spec
|
|
|
|
|
|
def _spec_with_actuator(actuator_name: str = "act") -> mujoco.MjSpec:
|
|
"""Single-body sphere with a hinge child + a position actuator."""
|
|
spec = mujoco.MjSpec()
|
|
m = spec.add_mesh(name="sphere")
|
|
m.make_sphere(subdivision=2)
|
|
root = spec.worldbody.add_body(name="prop")
|
|
root.add_freejoint()
|
|
root.add_geom(name="visual", type=mujoco.mjtGeom.mjGEOM_MESH, meshname="sphere")
|
|
child = root.add_body(name="lid")
|
|
cj = child.add_joint()
|
|
cj.name = "hinge"
|
|
cj.type = mujoco.mjtJoint.mjJNT_HINGE
|
|
child.add_geom(name="lid_geom", type=mujoco.mjtGeom.mjGEOM_MESH, meshname="sphere")
|
|
act = spec.add_actuator()
|
|
act.name = actuator_name
|
|
act.set_to_motor()
|
|
act.target = "hinge"
|
|
return spec
|
|
|
|
|
|
def _spec_with_primitive(primitive_role: str = "collision") -> mujoco.MjSpec:
|
|
"""Sphere variant with an additional primitive box (visual or collision)."""
|
|
spec = mujoco.MjSpec()
|
|
m = spec.add_mesh(name="sphere")
|
|
m.make_sphere(subdivision=2)
|
|
body = spec.worldbody.add_body(name="prop")
|
|
body.add_freejoint()
|
|
box = body.add_geom()
|
|
box.name = "primitive"
|
|
box.type = mujoco.mjtGeom.mjGEOM_BOX
|
|
box.size = np.array([0.05, 0.05, 0.05])
|
|
if primitive_role == "visual":
|
|
box.contype = 0
|
|
box.conaffinity = 0
|
|
body.add_geom(name="mesh_geom", type=mujoco.mjtGeom.mjGEOM_MESH, meshname="sphere")
|
|
return spec
|
|
|
|
|
|
def _spec_with_diagonal_inertia() -> mujoco.MjSpec:
|
|
spec = _simple_sphere_spec()
|
|
body = list(spec.worldbody.bodies)[0]
|
|
body.explicitinertial = True
|
|
body.mass = 1.0
|
|
body.ipos = np.array([0.0, 0.0, 0.0])
|
|
body.inertia = np.array([0.001, 0.001, 0.001])
|
|
body.iquat = np.array([1.0, 0.0, 0.0, 0.0])
|
|
return spec
|
|
|
|
|
|
def _spec_with_fullinertia() -> mujoco.MjSpec:
|
|
spec = _simple_sphere_spec()
|
|
body = list(spec.worldbody.bodies)[0]
|
|
body.explicitinertial = True
|
|
body.mass = 1.0
|
|
body.ipos = np.array([0.0, 0.0, 0.0])
|
|
body.fullinertia = np.array([0.001, 0.001, 0.001, 0.0, 0.0, 0.0])
|
|
return spec
|
|
|
|
|
|
def _spec_with_reserved_mesh_name() -> mujoco.MjSpec:
|
|
spec = mujoco.MjSpec()
|
|
m = spec.add_mesh(name="mjlab/pad/sneaky")
|
|
m.make_sphere(subdivision=2)
|
|
body = spec.worldbody.add_body(name="prop")
|
|
body.add_freejoint()
|
|
body.add_geom(
|
|
name="visual", type=mujoco.mjtGeom.mjGEOM_MESH, meshname="mjlab/pad/sneaky"
|
|
)
|
|
return spec
|
|
|
|
|
|
def test_articulated_same_topology_validates():
|
|
"""Articulated variants with matching topology pass validation
|
|
(build still rejects via floating-base check; this verifies validation
|
|
itself does not complain)."""
|
|
cfg = VariantEntityCfg(
|
|
variants={
|
|
"a": lambda: _articulated_spec(root_mesh="r_a", child_mesh="c_a"),
|
|
"b": lambda: _articulated_spec(root_mesh="r_b", child_mesh="c_b"),
|
|
},
|
|
)
|
|
entity = cfg.build()
|
|
assert entity.variant_metadata is not None
|
|
|
|
|
|
def test_recursive_child_body_count_mismatch_rejected():
|
|
"""Variants with different grandchild counts fail recursive validation."""
|
|
cfg = VariantEntityCfg(
|
|
variants={
|
|
"shallow": lambda: _articulated_spec(with_grandchild=False),
|
|
"deep": lambda: _articulated_spec(with_grandchild=True),
|
|
},
|
|
)
|
|
with pytest.raises(ValueError, match="child bodies"):
|
|
cfg.build()
|
|
|
|
|
|
def test_recursive_child_body_name_mismatch_rejected():
|
|
def variant_lid():
|
|
return _articulated_spec()
|
|
|
|
def variant_renamed_child():
|
|
spec = _articulated_spec()
|
|
list(spec.worldbody.bodies)[0].bodies[0].name = "drawer_top"
|
|
return spec
|
|
|
|
cfg = VariantEntityCfg(
|
|
variants={
|
|
"lid": variant_lid,
|
|
"drawer": variant_renamed_child,
|
|
},
|
|
)
|
|
with pytest.raises(ValueError, match="body path"):
|
|
cfg.build()
|
|
|
|
|
|
def test_recursive_joint_mismatch_in_child_body_rejected():
|
|
def variant_a():
|
|
return _articulated_spec()
|
|
|
|
def variant_b_slide():
|
|
spec = _articulated_spec()
|
|
child = list(spec.worldbody.bodies)[0].bodies[0]
|
|
list(child.joints)[0].type = mujoco.mjtJoint.mjJNT_SLIDE
|
|
return spec
|
|
|
|
cfg = VariantEntityCfg(
|
|
variants={
|
|
"hinge": variant_a,
|
|
"slide": variant_b_slide,
|
|
},
|
|
)
|
|
with pytest.raises(ValueError, match="joint"):
|
|
cfg.build()
|
|
|
|
|
|
def test_primitive_geom_count_mismatch_rejected():
|
|
cfg = VariantEntityCfg(
|
|
variants={
|
|
"with_box": _spec_with_primitive,
|
|
"without_box": _simple_sphere_spec,
|
|
},
|
|
)
|
|
with pytest.raises(ValueError, match="non-mesh geoms"):
|
|
cfg.build()
|
|
|
|
|
|
def test_primitive_geom_role_mismatch_rejected():
|
|
cfg = VariantEntityCfg(
|
|
variants={
|
|
"col": lambda: _spec_with_primitive("collision"),
|
|
"vis": lambda: _spec_with_primitive("visual"),
|
|
},
|
|
)
|
|
with pytest.raises(ValueError, match="primitive geom"):
|
|
cfg.build()
|
|
|
|
|
|
def test_actuator_count_mismatch_rejected():
|
|
cfg = VariantEntityCfg(
|
|
variants={
|
|
"no_act": _articulated_spec,
|
|
"with_act": _spec_with_actuator,
|
|
},
|
|
)
|
|
with pytest.raises(ValueError, match="actuator count"):
|
|
cfg.build()
|
|
|
|
|
|
def test_actuator_name_mismatch_rejected():
|
|
cfg = VariantEntityCfg(
|
|
variants={
|
|
"act_a": lambda: _spec_with_actuator("motor_a"),
|
|
"act_b": lambda: _spec_with_actuator("motor_b"),
|
|
},
|
|
)
|
|
with pytest.raises(ValueError, match="actuator #0"):
|
|
cfg.build()
|
|
|
|
|
|
def test_fullinertia_diagonal_mixing_rejected():
|
|
cfg = VariantEntityCfg(
|
|
variants={
|
|
"diag": _spec_with_diagonal_inertia,
|
|
"full": _spec_with_fullinertia,
|
|
},
|
|
)
|
|
with pytest.raises(ValueError, match="inertial representation"):
|
|
cfg.build()
|
|
|
|
|
|
def test_diagonal_inertia_consistent_accepted():
|
|
cfg = VariantEntityCfg(
|
|
variants={
|
|
"a": _spec_with_diagonal_inertia,
|
|
"b": _spec_with_diagonal_inertia,
|
|
},
|
|
)
|
|
entity = cfg.build()
|
|
assert entity.variant_metadata is not None
|
|
|
|
|
|
def test_reserved_prefix_in_mesh_rejected():
|
|
cfg = VariantEntityCfg(
|
|
variants={
|
|
"good": _simple_sphere_spec,
|
|
"bad": _spec_with_reserved_mesh_name,
|
|
},
|
|
)
|
|
with pytest.raises(ValueError, match="reserved name prefix"):
|
|
cfg.build()
|
|
|
|
|
|
def test_validation_error_format():
|
|
"""Error messages use the standardized mjlab.entity prefix and Hint suffix."""
|
|
cfg = VariantEntityCfg(
|
|
variants={
|
|
"ok": _simple_sphere_spec,
|
|
"bad": _hinge_spec,
|
|
},
|
|
)
|
|
with pytest.raises(ValueError) as excinfo:
|
|
cfg.build()
|
|
msg = str(excinfo.value)
|
|
assert msg.startswith("mjlab.entity: VariantEntityCfg 'bad': ")
|
|
assert "Hint:" in msg
|
|
|
|
|
|
def _spec_with_sensor() -> mujoco.MjSpec:
|
|
"""Sphere with a free joint and a velocity sensor on it."""
|
|
spec = _simple_sphere_spec()
|
|
s = spec.add_sensor()
|
|
s.name = "vel"
|
|
s.type = mujoco.mjtSensor.mjSENS_VELOCIMETER
|
|
s.objtype = mujoco.mjtObj.mjOBJ_SITE
|
|
s.objname = "site_a"
|
|
# Sensor needs a site target; add one.
|
|
body = list(spec.worldbody.bodies)[0]
|
|
site = body.add_site()
|
|
site.name = "site_a"
|
|
return spec
|
|
|
|
|
|
def test_sensor_count_mismatch_rejected():
|
|
cfg = VariantEntityCfg(
|
|
variants={
|
|
"no_sens": _simple_sphere_spec,
|
|
"with_sens": _spec_with_sensor,
|
|
},
|
|
)
|
|
with pytest.raises(ValueError, match="sensor count"):
|
|
cfg.build()
|
|
|
|
|
|
def test_validate_specs_directly_rejects_zero_root_bodies():
|
|
"""Empty worldbody fails with a clear root-body message."""
|
|
from mjlab.entity.variants import validate_variant_specs
|
|
|
|
empty = mujoco.MjSpec()
|
|
ok = _simple_sphere_spec()
|
|
with pytest.raises(ValueError, match="exactly one root body"):
|
|
validate_variant_specs(["empty", "ok"], [empty, ok])
|
|
|
|
|
|
# Slot metadata (Workstream 3).
|
|
|
|
|
|
def _slot_is_padding(meta, slot_index: int) -> bool:
|
|
"""True if any variant leaves the slot at ``slot_index`` unfilled."""
|
|
return any(specs[slot_index] is None for specs in meta.variant_slot_specs)
|
|
|
|
|
|
def test_slot_metadata_single_variant_single_geom():
|
|
cfg = VariantEntityCfg(variants={"only": _simple_sphere_spec})
|
|
meta = cfg.build().variant_metadata
|
|
assert meta is not None
|
|
assert len(meta.slots) == 1
|
|
slot = meta.slots[0]
|
|
assert isinstance(slot, VariantSlot)
|
|
assert slot.key == SlotKey(body_path="/prop", role="collision", ordinal=0)
|
|
assert _slot_is_padding(meta, 0) is False
|
|
assert slot.template_geom_name == "mjlab/pad/prop/collision/0"
|
|
# source_geom_names has one entry per variant.
|
|
assert slot.source_geom_names == ("visual",)
|
|
# variant_slot_specs aligns with slots.
|
|
assert len(meta.variant_slot_specs) == 1
|
|
assert len(meta.variant_slot_specs[0]) == 1
|
|
vgs = meta.variant_slot_specs[0][0]
|
|
assert isinstance(vgs, VariantGeomSpec)
|
|
assert vgs.geom_name == "visual"
|
|
assert vgs.mesh_name == "sphere"
|
|
|
|
|
|
def test_slot_metadata_visual_collision_split():
|
|
"""Visual and collision geoms on the same body get distinct slots."""
|
|
cfg = VariantEntityCfg(
|
|
variants={
|
|
"sphere": _sphere_2col_spec,
|
|
"cone": _cone_4col_spec,
|
|
}
|
|
)
|
|
meta = cfg.build().variant_metadata
|
|
assert meta is not None
|
|
# Body /prop has 1 visual + max(2, 4) = 4 collision slots.
|
|
visual_slots = [s for s in meta.slots if s.key.role == "visual"]
|
|
collision_slots = [s for s in meta.slots if s.key.role == "collision"]
|
|
assert len(visual_slots) == 1
|
|
assert len(collision_slots) == 4
|
|
assert all(s.key.body_path == "/prop" for s in meta.slots)
|
|
|
|
|
|
def test_slot_metadata_visual_before_collision():
|
|
"""Slot ordering: visual slots come before collision slots per body."""
|
|
cfg = VariantEntityCfg(
|
|
variants={
|
|
"sphere": _sphere_2col_spec,
|
|
"cone": _cone_4col_spec,
|
|
}
|
|
)
|
|
meta = cfg.build().variant_metadata
|
|
assert meta is not None
|
|
roles = [s.key.role for s in meta.slots]
|
|
# Find first collision; all visual should come before it.
|
|
first_col = roles.index("collision")
|
|
assert all(r == "visual" for r in roles[:first_col])
|
|
assert all(r == "collision" for r in roles[first_col:])
|
|
|
|
|
|
def test_slot_metadata_padding_derivable_from_specs():
|
|
"""A slot held by some variants but not others is unfilled (None) for those."""
|
|
cfg = VariantEntityCfg(
|
|
variants={
|
|
"sphere": _sphere_2col_spec,
|
|
"cone": _cone_4col_spec,
|
|
}
|
|
)
|
|
meta = cfg.build().variant_metadata
|
|
assert meta is not None
|
|
col_positions = [i for i, s in enumerate(meta.slots) if s.key.role == "collision"]
|
|
# Sphere has 2 collision; cone has 4. Slots 0, 1 fully populated; 2, 3 are padding.
|
|
assert _slot_is_padding(meta, col_positions[0]) is False
|
|
assert _slot_is_padding(meta, col_positions[1]) is False
|
|
assert _slot_is_padding(meta, col_positions[2]) is True
|
|
assert _slot_is_padding(meta, col_positions[3]) is True
|
|
# Sphere has None at the padding slots.
|
|
variant_idx_sphere = meta.variant_names.index("sphere")
|
|
specs_sphere = meta.variant_slot_specs[variant_idx_sphere]
|
|
assert specs_sphere[col_positions[2]] is None
|
|
assert specs_sphere[col_positions[3]] is None
|
|
# Cone has VariantGeomSpec for all collision slots.
|
|
variant_idx_cone = meta.variant_names.index("cone")
|
|
specs_cone = meta.variant_slot_specs[variant_idx_cone]
|
|
for cp in col_positions:
|
|
assert specs_cone[cp] is not None
|
|
|
|
|
|
def test_slot_metadata_articulated_per_body_slots():
|
|
"""Articulated variants produce slots per (body, role)."""
|
|
cfg = VariantEntityCfg(
|
|
variants={
|
|
"a": lambda: _articulated_spec(root_mesh="ar", child_mesh="ac"),
|
|
"b": lambda: _articulated_spec(root_mesh="br", child_mesh="bc"),
|
|
}
|
|
)
|
|
meta = cfg.build().variant_metadata
|
|
assert meta is not None
|
|
# Two bodies (/prop, /prop/lid), each with 1 collision mesh -> 2 slots.
|
|
paths = sorted({s.key.body_path for s in meta.slots})
|
|
assert paths == ["/prop", "/prop/lid"]
|
|
# Each body has exactly one collision slot, no visuals.
|
|
for path in paths:
|
|
body_slots = [s for s in meta.slots if s.key.body_path == path]
|
|
assert len(body_slots) == 1
|
|
assert body_slots[0].key.role == "collision"
|
|
slot_idx = meta.slots.index(body_slots[0])
|
|
assert _slot_is_padding(meta, slot_idx) is False
|
|
# Mesh names captured per variant per slot.
|
|
variant_a_specs = meta.variant_slot_specs[meta.variant_names.index("a")]
|
|
variant_b_specs = meta.variant_slot_specs[meta.variant_names.index("b")]
|
|
a_meshes = sorted(s.mesh_name for s in variant_a_specs if s is not None)
|
|
b_meshes = sorted(s.mesh_name for s in variant_b_specs if s is not None)
|
|
assert a_meshes == ["ac", "ar"]
|
|
assert b_meshes == ["bc", "br"]
|
|
|
|
|
|
def test_slot_metadata_template_name_under_reserved_prefix():
|
|
"""Every template slot name starts with the reserved mjlab/pad/ prefix."""
|
|
cfg = VariantEntityCfg(
|
|
variants={
|
|
"a": _simple_sphere_spec,
|
|
"b": _simple_cone_spec,
|
|
}
|
|
)
|
|
meta = cfg.build().variant_metadata
|
|
assert meta is not None
|
|
for slot in meta.slots:
|
|
assert slot.template_geom_name.startswith("mjlab/pad/")
|
|
|
|
|
|
def test_slot_metadata_captures_visual_role_from_zero_contact_bits():
|
|
"""A geom with contype=0 and conaffinity=0 is classified as visual."""
|
|
cfg = VariantEntityCfg(
|
|
variants={
|
|
"sphere": _sphere_2col_spec,
|
|
"cone": _cone_4col_spec,
|
|
}
|
|
)
|
|
meta = cfg.build().variant_metadata
|
|
assert meta is not None
|
|
visual_slots = [s for s in meta.slots if s.key.role == "visual"]
|
|
assert len(visual_slots) == 1
|
|
visual_slot = visual_slots[0]
|
|
for variant_specs in meta.variant_slot_specs:
|
|
spec_at_visual = variant_specs[meta.slots.index(visual_slot)]
|
|
assert spec_at_visual is not None
|
|
assert spec_at_visual.contype == 0
|
|
assert spec_at_visual.conaffinity == 0
|
|
|
|
|
|
def test_slot_metadata_ordinals_are_zero_based_per_body_and_role():
|
|
"""Within a (body, role), slot ordinals are 0..max-1."""
|
|
cfg = VariantEntityCfg(
|
|
variants={
|
|
"sphere": _sphere_2col_spec,
|
|
"cone": _cone_4col_spec,
|
|
}
|
|
)
|
|
meta = cfg.build().variant_metadata
|
|
assert meta is not None
|
|
collision_slots = [s for s in meta.slots if s.key.role == "collision"]
|
|
assert [s.key.ordinal for s in collision_slots] == [0, 1, 2, 3]
|
|
visual_slots = [s for s in meta.slots if s.key.role == "visual"]
|
|
assert [s.key.ordinal for s in visual_slots] == [0]
|
|
|
|
|
|
def test_slot_metadata_alignment_invariant():
|
|
"""variant_slot_specs[v] aligns with slots positionally for every variant."""
|
|
cfg = VariantEntityCfg(
|
|
variants={
|
|
"sphere": _sphere_2col_spec,
|
|
"cone": _cone_4col_spec,
|
|
}
|
|
)
|
|
meta = cfg.build().variant_metadata
|
|
assert meta is not None
|
|
for variant_specs in meta.variant_slot_specs:
|
|
assert len(variant_specs) == len(meta.slots)
|
|
|
|
|
|
def test_template_geom_contype_matches_slot_role():
|
|
"""Template contype/conaffinity are derived from slot role (union of variants)."""
|
|
scene_spec, vi = _build_scene_with_variants(_sphere_2col_spec, _cone_4col_spec)
|
|
result = build_variant_model(scene_spec, 4, vi)
|
|
metadata = vi[0][1]
|
|
for slot in metadata.slots:
|
|
full_name = f"object/{slot.template_geom_name}"
|
|
gid = mujoco.mj_name2id(result.mj_model, mujoco.mjtObj.mjOBJ_GEOM, full_name)
|
|
assert gid >= 0, f"slot geom '{full_name}' missing from compiled model"
|
|
contype = int(result.mj_model.geom_contype[gid])
|
|
conaffinity = int(result.mj_model.geom_conaffinity[gid])
|
|
if slot.key.role == "visual":
|
|
assert contype == 0, f"visual slot {slot.key} has contype={contype}"
|
|
assert conaffinity == 0, f"visual slot {slot.key} has conaffinity={conaffinity}"
|
|
else:
|
|
assert contype == 1, f"collision slot {slot.key} has contype={contype}"
|
|
assert conaffinity == 1, (
|
|
f"collision slot {slot.key} has conaffinity={conaffinity}"
|
|
)
|
|
|
|
|
|
def test_template_geoms_use_mjlab_pad_prefix():
|
|
"""All entity mesh-geom names in the compiled template use mjlab/pad/ prefix."""
|
|
scene_spec, vi = _build_scene_with_variants(_sphere_2col_spec, _cone_4col_spec)
|
|
result = build_variant_model(scene_spec, 4, vi)
|
|
for gid in range(result.mj_model.ngeom):
|
|
if result.mj_model.geom_type[gid] != mujoco.mjtGeom.mjGEOM_MESH:
|
|
continue
|
|
name = mujoco.mj_id2name(result.mj_model, mujoco.mjtObj.mjOBJ_GEOM, gid) or ""
|
|
if not name.startswith("object/"):
|
|
continue
|
|
suffix = name[len("object/") :]
|
|
assert suffix.startswith("mjlab/pad/"), (
|
|
f"entity mesh geom '{name}' does not use mjlab/pad/ prefix"
|
|
)
|
|
|
|
|
|
def test_visual_collision_split_inertia_matches_independent_compile():
|
|
"""Per-world body_mass matches independent compile when variants have visual+collision split.
|
|
|
|
This verifies the slot-driven reference compile preserves the visual
|
|
role (contype=0/conaffinity=0) on the visual mesh; if the old
|
|
contype=1/conaffinity=1 reset still ran, the visual mesh's
|
|
inertia-inference behavior would not change for default groups, but
|
|
this exercise pins the contract end-to-end.
|
|
"""
|
|
scene_spec, vi = _build_scene_with_variants(_sphere_2col_spec, _cone_4col_spec)
|
|
result = build_variant_model(scene_spec, 4, vi)
|
|
|
|
sphere_model = _sphere_2col_spec().compile()
|
|
cone_model = _cone_4col_spec().compile()
|
|
|
|
body_mass = result.wp_model.body_mass.numpy()
|
|
w2v = result.world_to_variant["object/"]
|
|
obj_body = result.mj_model.nbody - 1
|
|
|
|
sphere_w = int(np.where(w2v == 0)[0][0])
|
|
cone_w = int(np.where(w2v == 1)[0][0])
|
|
|
|
np.testing.assert_allclose(
|
|
body_mass[sphere_w, obj_body],
|
|
sphere_model.body_mass[-1],
|
|
atol=1e-4,
|
|
)
|
|
np.testing.assert_allclose(
|
|
body_mass[cone_w, obj_body],
|
|
cone_model.body_mass[-1],
|
|
atol=1e-4,
|
|
)
|
|
|
|
|
|
def test_variant_order_irrelevant_per_variant_compile():
|
|
"""Reordering the variant dict does not change per-variant per-world fields."""
|
|
cfg_ab = VariantEntityCfg(
|
|
variants={
|
|
"a": _simple_sphere_spec,
|
|
"b": _simple_cone_spec,
|
|
}
|
|
)
|
|
cfg_ba = VariantEntityCfg(
|
|
variants={
|
|
"b": _simple_cone_spec,
|
|
"a": _simple_sphere_spec,
|
|
}
|
|
)
|
|
|
|
def _build_scene(cfg: VariantEntityCfg):
|
|
entity = cfg.build()
|
|
assert entity.variant_metadata is not None
|
|
scene_spec = mujoco.MjSpec()
|
|
frame = scene_spec.worldbody.add_frame()
|
|
scene_spec.attach(entity.spec, prefix="object/", frame=frame)
|
|
return scene_spec, [("object/", entity.variant_metadata)]
|
|
|
|
scene_ab, vi_ab = _build_scene(cfg_ab)
|
|
scene_ba, vi_ba = _build_scene(cfg_ba)
|
|
res_ab = build_variant_model(scene_ab, 4, vi_ab)
|
|
res_ba = build_variant_model(scene_ba, 4, vi_ba)
|
|
|
|
obj_body_ab = res_ab.mj_model.nbody - 1
|
|
obj_body_ba = res_ba.mj_model.nbody - 1
|
|
|
|
body_mass_ab = res_ab.wp_model.body_mass.numpy()
|
|
body_mass_ba = res_ba.wp_model.body_mass.numpy()
|
|
|
|
# In cfg_ab, "a" is variant index 0; in cfg_ba, "a" is variant index 1.
|
|
w2v_ab = res_ab.world_to_variant["object/"]
|
|
w2v_ba = res_ba.world_to_variant["object/"]
|
|
a_world_ab = int(np.where(w2v_ab == 0)[0][0])
|
|
a_world_ba = int(np.where(w2v_ba == 1)[0][0])
|
|
b_world_ab = int(np.where(w2v_ab == 1)[0][0])
|
|
b_world_ba = int(np.where(w2v_ba == 0)[0][0])
|
|
|
|
np.testing.assert_allclose(
|
|
body_mass_ab[a_world_ab, obj_body_ab],
|
|
body_mass_ba[a_world_ba, obj_body_ba],
|
|
atol=1e-5,
|
|
err_msg="variant 'a' body_mass differs across orderings",
|
|
)
|
|
np.testing.assert_allclose(
|
|
body_mass_ab[b_world_ab, obj_body_ab],
|
|
body_mass_ba[b_world_ba, obj_body_ba],
|
|
atol=1e-5,
|
|
err_msg="variant 'b' body_mass differs across orderings",
|
|
)
|
|
|
|
|
|
def test_slot_metadata_source_geom_names_record_padding_as_none():
|
|
"""source_geom_names has None where a variant doesn't fill a slot."""
|
|
cfg = VariantEntityCfg(
|
|
variants={
|
|
"sphere": _sphere_2col_spec,
|
|
"cone": _cone_4col_spec,
|
|
}
|
|
)
|
|
meta = cfg.build().variant_metadata
|
|
assert meta is not None
|
|
sphere_idx = meta.variant_names.index("sphere")
|
|
collision_slots = [s for s in meta.slots if s.key.role == "collision"]
|
|
# Sphere has 2 collision -> ordinals 2, 3 are None for sphere.
|
|
assert collision_slots[2].source_geom_names[sphere_idx] is None
|
|
assert collision_slots[3].source_geom_names[sphere_idx] is None
|
|
# Cone fills all 4.
|
|
cone_idx = meta.variant_names.index("cone")
|
|
for cs in collision_slots:
|
|
assert cs.source_geom_names[cone_idx] is not None
|
|
|
|
|
|
def test_no_variants_unchanged():
|
|
cfg = EntityCfg(spec_fn=_simple_sphere_spec)
|
|
entity = cfg.build()
|
|
assert entity.variant_metadata is None
|
|
|
|
|
|
# build_variant_model: dataid and dependent fields.
|
|
|
|
|
|
def test_dataid_assigned_per_world():
|
|
"""Each world's geom_dataid points to its variant's meshes."""
|
|
scene_spec, vi = _build_scene_with_variants(_simple_sphere_spec, _simple_cone_spec)
|
|
result = build_variant_model(scene_spec, 4, vi)
|
|
|
|
dataid = result.wp_model.geom_dataid.numpy()
|
|
assert dataid.shape == (4, result.mj_model.ngeom)
|
|
assert dataid.ndim == 2
|
|
|
|
w2v = result.world_to_variant["object/"]
|
|
assert w2v[0] == 0 # variant a (sphere)
|
|
assert w2v[2] == 1 # variant b (cone)
|
|
|
|
# Sphere and cone worlds must have different dataid values.
|
|
assert not np.array_equal(dataid[0], dataid[2])
|
|
|
|
|
|
def _sphere_with_material_spec() -> mujoco.MjSpec:
|
|
"""Single-geom sphere whose visual references a named material."""
|
|
spec = mujoco.MjSpec()
|
|
m = spec.add_mesh()
|
|
m.name = "sphere"
|
|
m.make_sphere(subdivision=2)
|
|
mat = spec.add_material()
|
|
mat.name = "red_mat"
|
|
mat.rgba[:] = (1.0, 0.0, 0.0, 1.0)
|
|
body = spec.worldbody.add_body()
|
|
body.name = "prop"
|
|
body.add_freejoint()
|
|
g = body.add_geom()
|
|
g.name = "visual"
|
|
g.type = mujoco.mjtGeom.mjGEOM_MESH
|
|
g.meshname = "sphere"
|
|
g.material = "red_mat"
|
|
return spec
|
|
|
|
|
|
def _cone_with_material_spec() -> mujoco.MjSpec:
|
|
"""Single-geom cone whose visual references a different named material."""
|
|
spec = mujoco.MjSpec()
|
|
m = spec.add_mesh()
|
|
m.name = "cone"
|
|
m.make_cone(nedge=8, radius=0.05)
|
|
mat = spec.add_material()
|
|
mat.name = "blue_mat"
|
|
mat.rgba[:] = (0.0, 0.0, 1.0, 1.0)
|
|
body = spec.worldbody.add_body()
|
|
body.name = "prop"
|
|
body.add_freejoint()
|
|
g = body.add_geom()
|
|
g.name = "visual"
|
|
g.type = mujoco.mjtGeom.mjGEOM_MESH
|
|
g.meshname = "cone"
|
|
g.material = "blue_mat"
|
|
return spec
|
|
|
|
|
|
def test_materials_merged_under_variant_prefix():
|
|
"""Both variants' materials end up in the merged spec, name-prefixed."""
|
|
scene_spec, vi = _build_scene_with_variants(
|
|
_sphere_with_material_spec, _cone_with_material_spec
|
|
)
|
|
model = scene_spec.compile()
|
|
mat_names = {model.material(i).name for i in range(model.nmat)}
|
|
assert "object/a/red_mat" in mat_names
|
|
assert "object/b/blue_mat" in mat_names
|
|
|
|
|
|
def test_matid_assigned_per_world():
|
|
"""Each world's geom_matid points to its variant's material."""
|
|
scene_spec, vi = _build_scene_with_variants(
|
|
_sphere_with_material_spec, _cone_with_material_spec
|
|
)
|
|
result = build_variant_model(scene_spec, 4, vi)
|
|
|
|
matid = result.wp_model.geom_matid.numpy()
|
|
assert matid.shape == (4, result.mj_model.ngeom)
|
|
|
|
w2v = result.world_to_variant["object/"]
|
|
red_id = mujoco.mj_name2id(
|
|
result.mj_model, mujoco.mjtObj.mjOBJ_MATERIAL, "object/a/red_mat"
|
|
)
|
|
blue_id = mujoco.mj_name2id(
|
|
result.mj_model, mujoco.mjtObj.mjOBJ_MATERIAL, "object/b/blue_mat"
|
|
)
|
|
assert red_id >= 0 and blue_id >= 0 and red_id != blue_id
|
|
|
|
# Slot geom is the last mesh geom (single-geom variants -> ordinal 0).
|
|
slot_gid = next(
|
|
gid
|
|
for gid in range(result.mj_model.ngeom - 1, -1, -1)
|
|
if result.mj_model.geom_type[gid] == mujoco.mjtGeom.mjGEOM_MESH
|
|
)
|
|
|
|
for w in range(4):
|
|
expected = red_id if w2v[w] == 0 else blue_id
|
|
assert int(matid[w, slot_gid]) == expected
|
|
|
|
|
|
def test_matid_minus_one_when_variant_has_no_material():
|
|
"""A variant slot without a material yields geom_matid == -1 in its worlds."""
|
|
scene_spec, vi = _build_scene_with_variants(
|
|
_sphere_with_material_spec,
|
|
_simple_cone_spec, # cone has no material
|
|
)
|
|
result = build_variant_model(scene_spec, 4, vi)
|
|
|
|
matid = result.wp_model.geom_matid.numpy()
|
|
w2v = result.world_to_variant["object/"]
|
|
slot_gid = next(
|
|
gid
|
|
for gid in range(result.mj_model.ngeom - 1, -1, -1)
|
|
if result.mj_model.geom_type[gid] == mujoco.mjtGeom.mjGEOM_MESH
|
|
)
|
|
cone_world = int(np.where(w2v == 1)[0][0])
|
|
assert int(matid[cone_world, slot_gid]) == -1
|
|
|
|
|
|
def test_padding_slots_get_disabled():
|
|
"""Shorter variant's padding geom slots have dataid == -1."""
|
|
scene_spec, vi = _build_scene_with_variants(_sphere_2col_spec, _cone_4col_spec)
|
|
result = build_variant_model(scene_spec, 4, vi)
|
|
|
|
dataid = result.wp_model.geom_dataid.numpy()
|
|
w2v = result.world_to_variant["object/"]
|
|
|
|
# Find a sphere world (variant 0, 3 mesh geoms -> 2 padding slots).
|
|
sphere_world = int(np.where(w2v == 0)[0][0])
|
|
# Find mesh geom columns (skip non-mesh geoms like worldbody).
|
|
mesh_geom_ids = [
|
|
gid
|
|
for gid in range(result.mj_model.ngeom)
|
|
if result.mj_model.geom_type[gid] == mujoco.mjtGeom.mjGEOM_MESH
|
|
]
|
|
sphere_dataid = dataid[sphere_world, mesh_geom_ids]
|
|
# Last 2 mesh geom slots should be -1 (disabled padding).
|
|
assert sphere_dataid[-1] == -1
|
|
assert sphere_dataid[-2] == -1
|
|
# Padding slots must still be collision-enabled in the template/warp model.
|
|
# Short variants are disabled by per-world dataid=-1; long variants need the
|
|
# same slots enabled so their extra hulls can collide.
|
|
assert np.all(result.mj_model.geom_contype[mesh_geom_ids[-2:]] == 1)
|
|
assert np.all(result.mj_model.geom_conaffinity[mesh_geom_ids[-2:]] == 1)
|
|
assert np.all(result.wp_model.geom_contype.numpy()[mesh_geom_ids[-2:]] == 1)
|
|
assert np.all(result.wp_model.geom_conaffinity.numpy()[mesh_geom_ids[-2:]] == 1)
|
|
# First 3 should be valid (>= 0).
|
|
assert all(d >= 0 for d in sphere_dataid[:3])
|
|
|
|
|
|
def test_dependent_fields_match_individual_compilation():
|
|
"""Per-world body_mass matches independently compiled variant models."""
|
|
scene_spec, vi = _build_scene_with_variants(_simple_sphere_spec, _simple_cone_spec)
|
|
result = build_variant_model(scene_spec, 4, vi)
|
|
|
|
# Compile each variant independently for reference values.
|
|
sphere_model = _simple_sphere_spec().compile()
|
|
cone_model = _simple_cone_spec().compile()
|
|
|
|
body_mass = result.wp_model.body_mass.numpy()
|
|
w2v = result.world_to_variant["object/"]
|
|
|
|
sphere_w = int(np.where(w2v == 0)[0][0])
|
|
cone_w = int(np.where(w2v == 1)[0][0])
|
|
|
|
# The object body is the last body in the scene.
|
|
obj_body = result.mj_model.nbody - 1
|
|
|
|
# Mass should match individually compiled models.
|
|
np.testing.assert_allclose(
|
|
body_mass[sphere_w, obj_body],
|
|
sphere_model.body_mass[-1],
|
|
atol=1e-4,
|
|
)
|
|
np.testing.assert_allclose(
|
|
body_mass[cone_w, obj_body],
|
|
cone_model.body_mass[-1],
|
|
atol=1e-4,
|
|
)
|
|
|
|
# Sphere and cone should have different masses.
|
|
assert not np.isclose(body_mass[sphere_w, obj_body], body_mass[cone_w, obj_body])
|
|
|
|
|
|
def test_select_default_values_uses_per_world_variant_defaults():
|
|
"""Per-world defaults are indexed by env first, then by entity."""
|
|
from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg
|
|
from mjlab.envs.mdp.dr._core import _select_default_values
|
|
from mjlab.scene import SceneCfg
|
|
from mjlab.terrains import TerrainEntityCfg
|
|
|
|
def _explicit_variant(
|
|
mesh_name: str,
|
|
mass: float,
|
|
inertia: tuple[float, float, float],
|
|
*,
|
|
cone: bool = False,
|
|
) -> mujoco.MjSpec:
|
|
spec = mujoco.MjSpec()
|
|
mesh = spec.add_mesh()
|
|
mesh.name = mesh_name
|
|
if cone:
|
|
mesh.make_cone(nedge=8, radius=0.05)
|
|
else:
|
|
mesh.make_sphere(subdivision=1)
|
|
body = spec.worldbody.add_body(name="prop")
|
|
body.add_freejoint()
|
|
body.explicitinertial = True
|
|
body.mass = mass
|
|
body.ipos[:] = (0.0, 0.0, 0.0)
|
|
body.inertia[:] = inertia
|
|
body.iquat[:] = (1.0, 0.0, 0.0, 0.0)
|
|
body.add_geom(
|
|
name="visual",
|
|
type=mujoco.mjtGeom.mjGEOM_MESH,
|
|
meshname=mesh_name,
|
|
contype=0,
|
|
conaffinity=0,
|
|
mass=0.0,
|
|
)
|
|
return spec
|
|
|
|
object_cfg = VariantEntityCfg(
|
|
variants={
|
|
"sphere": lambda: _explicit_variant("sphere", 0.2, (1e-4, 2e-4, 3e-4)),
|
|
"cone": lambda: _explicit_variant("cone", 0.7, (4e-4, 5e-4, 6e-4), cone=True),
|
|
},
|
|
init_state=EntityCfg.InitialStateCfg(pos=(0.0, 0.0, 0.2)),
|
|
)
|
|
env_cfg = ManagerBasedRlEnvCfg(
|
|
decimation=1,
|
|
scene=SceneCfg(
|
|
terrain=TerrainEntityCfg(terrain_type="plane"),
|
|
num_envs=4,
|
|
env_spacing=1.0,
|
|
entities={"object": object_cfg},
|
|
),
|
|
)
|
|
|
|
env = ManagerBasedRlEnv(cfg=env_cfg, device="cpu")
|
|
try:
|
|
obj_body = int(env.scene["object"].indexing.root_body_id)
|
|
env_ids = torch.arange(env.num_envs, device=env.device)
|
|
body_ids = torch.tensor([obj_body], device=env.device)
|
|
|
|
for field in ("body_mass", "body_ipos", "body_inertia", "body_iquat"):
|
|
selected = _select_default_values(env, field, env_ids, body_ids)
|
|
torch.testing.assert_close(
|
|
selected[:, 0],
|
|
getattr(env.sim.model, field)[:, obj_body],
|
|
)
|
|
finally:
|
|
env.close()
|
|
|
|
|
|
def test_viser_builds_per_world_mesh_handles_for_variants():
|
|
"""Viser dynamic meshes must not collapse all worlds onto env0's mesh."""
|
|
from contextlib import nullcontext
|
|
|
|
from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg
|
|
from mjlab.scene import SceneCfg
|
|
from mjlab.terrains import TerrainEntityCfg
|
|
from mjlab.viewer.viser.scene import MjlabViserScene, _PerWorldMeshGroup
|
|
|
|
class _Handle:
|
|
def __init__(self, **kwargs):
|
|
self.visible = kwargs.get("visible", True)
|
|
self.batched_positions = kwargs.get("batched_positions", np.zeros((0, 3)))
|
|
self.batched_wxyzs = kwargs.get("batched_wxyzs", np.zeros((0, 4)))
|
|
self.batched_scales = kwargs.get("batched_scales")
|
|
self.batched_colors = kwargs.get("batched_colors")
|
|
self.batched_opacities = kwargs.get("batched_opacities")
|
|
self.position = kwargs.get("position", np.zeros(3))
|
|
self.wxyz = kwargs.get("wxyz", np.array([1.0, 0.0, 0.0, 0.0]))
|
|
|
|
def remove(self) -> None:
|
|
pass
|
|
|
|
class _Scene:
|
|
def __init__(self):
|
|
self.batched: list[tuple[tuple, dict, _Handle]] = []
|
|
|
|
def configure_environment_map(self, **_kwargs) -> None:
|
|
pass
|
|
|
|
def add_frame(self, *_args, **kwargs) -> _Handle:
|
|
return _Handle(**kwargs)
|
|
|
|
def add_grid(self, *_args, **kwargs) -> _Handle:
|
|
return _Handle(**kwargs)
|
|
|
|
def add_mesh_trimesh(self, *_args, **kwargs) -> _Handle:
|
|
return _Handle(**kwargs)
|
|
|
|
def add_batched_meshes_trimesh(self, *args, **kwargs) -> _Handle:
|
|
handle = _Handle(**kwargs)
|
|
self.batched.append((args, kwargs, handle))
|
|
return handle
|
|
|
|
def add_batched_meshes_simple(self, *args, **kwargs) -> _Handle:
|
|
handle = _Handle(**kwargs)
|
|
self.batched.append((args, kwargs, handle))
|
|
return handle
|
|
|
|
class _Server:
|
|
def __init__(self):
|
|
self.scene = _Scene()
|
|
|
|
def atomic(self):
|
|
return nullcontext()
|
|
|
|
def flush(self) -> None:
|
|
pass
|
|
|
|
env_cfg = ManagerBasedRlEnvCfg(
|
|
decimation=1,
|
|
scene=SceneCfg(
|
|
terrain=TerrainEntityCfg(terrain_type="plane"),
|
|
num_envs=4,
|
|
env_spacing=1.0,
|
|
entities={
|
|
"object": VariantEntityCfg(
|
|
variants={
|
|
"sphere": _simple_sphere_spec,
|
|
"cone": _simple_cone_spec,
|
|
},
|
|
init_state=EntityCfg.InitialStateCfg(pos=(0.0, 0.0, 0.2)),
|
|
)
|
|
},
|
|
),
|
|
)
|
|
|
|
env = ManagerBasedRlEnv(cfg=env_cfg, device="cpu")
|
|
try:
|
|
env.sim.expand_model_fields(("geom_rgba",))
|
|
env.sim.model.geom_rgba[:, :, :3] = torch.linspace(
|
|
0.2,
|
|
0.9,
|
|
env.num_envs,
|
|
device=env.device,
|
|
)[:, None, None]
|
|
server = _Server()
|
|
scene = MjlabViserScene(
|
|
cast(Any, server),
|
|
env.sim.mj_model,
|
|
env.num_envs,
|
|
sim_model=env.sim.model,
|
|
expanded_fields=env.sim.expanded_fields,
|
|
)
|
|
groups = [mg for mg in scene._mesh_groups if isinstance(mg, _PerWorldMeshGroup)]
|
|
|
|
assert groups
|
|
assert sum(len(mg.env_ids) for mg in groups) >= env.num_envs
|
|
|
|
body_xpos = env.sim.data.xpos.cpu().numpy()
|
|
body_xmat = env.sim.data.xmat.cpu().numpy()
|
|
mocap_pos = (
|
|
env.sim.data.mocap_pos.cpu().numpy() if env.sim.mj_model.nmocap > 0 else None
|
|
)
|
|
mocap_quat = (
|
|
env.sim.data.mocap_quat.cpu().numpy() if env.sim.mj_model.nmocap > 0 else None
|
|
)
|
|
scene.show_only_selected = True
|
|
scene.update_from_arrays(body_xpos, body_xmat, mocap_pos, mocap_quat, env_idx=0)
|
|
scene.update_from_arrays(body_xpos, body_xmat, mocap_pos, mocap_quat, env_idx=1)
|
|
|
|
assert any(mg.handle.visible for mg in groups)
|
|
|
|
handle_count = len(server.scene.batched)
|
|
env.sim.model.geom_rgba[:, :, :3] = torch.linspace(
|
|
0.9,
|
|
0.2,
|
|
env.num_envs,
|
|
device=env.device,
|
|
)[:, None, None]
|
|
scene.update_from_arrays(body_xpos, body_xmat, mocap_pos, mocap_quat, env_idx=0)
|
|
assert len(server.scene.batched) > handle_count
|
|
finally:
|
|
env.close()
|
|
|
|
|
|
def test_viser_convex_hulls_are_per_variant():
|
|
"""Convex-hull handles must differ across variants, not all show env0's hull."""
|
|
from contextlib import nullcontext
|
|
|
|
from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg
|
|
from mjlab.scene import SceneCfg
|
|
from mjlab.terrains import TerrainEntityCfg
|
|
from mjlab.viewer.viser.scene import MjlabViserScene, _PerWorldHullGroup
|
|
|
|
class _Handle:
|
|
def __init__(self, **kwargs):
|
|
self.visible = kwargs.get("visible", True)
|
|
self.batched_positions = kwargs.get("batched_positions", np.zeros((0, 3)))
|
|
self.batched_wxyzs = kwargs.get("batched_wxyzs", np.zeros((0, 4)))
|
|
self.batched_scales = kwargs.get("batched_scales")
|
|
self.batched_colors = kwargs.get("batched_colors")
|
|
self.batched_opacities = kwargs.get("batched_opacities")
|
|
self.position = kwargs.get("position", np.zeros(3))
|
|
self.wxyz = kwargs.get("wxyz", np.array([1.0, 0.0, 0.0, 0.0]))
|
|
self.vertices = kwargs.get("vertices")
|
|
self.faces = kwargs.get("faces")
|
|
|
|
def remove(self) -> None:
|
|
pass
|
|
|
|
class _Scene:
|
|
def __init__(self):
|
|
self.batched: list[tuple[tuple, dict, _Handle]] = []
|
|
|
|
def configure_environment_map(self, **_kwargs) -> None:
|
|
pass
|
|
|
|
def add_frame(self, *_args, **kwargs) -> _Handle:
|
|
return _Handle(**kwargs)
|
|
|
|
def add_grid(self, *_args, **kwargs) -> _Handle:
|
|
return _Handle(**kwargs)
|
|
|
|
def add_mesh_trimesh(self, *_args, **kwargs) -> _Handle:
|
|
return _Handle(**kwargs)
|
|
|
|
def add_batched_meshes_trimesh(self, *args, **kwargs) -> _Handle:
|
|
handle = _Handle(**kwargs)
|
|
self.batched.append((args, kwargs, handle))
|
|
return handle
|
|
|
|
def add_batched_meshes_simple(self, path, vertices, faces, **kwargs) -> _Handle:
|
|
# Capture the mesh identity so the test can compare hull shapes.
|
|
kwargs = dict(kwargs)
|
|
kwargs["vertices"] = np.asarray(vertices)
|
|
kwargs["faces"] = np.asarray(faces)
|
|
handle = _Handle(**kwargs)
|
|
self.batched.append(((path,), kwargs, handle))
|
|
return handle
|
|
|
|
class _Server:
|
|
def __init__(self):
|
|
self.scene = _Scene()
|
|
|
|
def atomic(self):
|
|
return nullcontext()
|
|
|
|
def flush(self) -> None:
|
|
pass
|
|
|
|
# Sphere and cone produce visibly different convex hulls.
|
|
env_cfg = ManagerBasedRlEnvCfg(
|
|
decimation=1,
|
|
scene=SceneCfg(
|
|
terrain=TerrainEntityCfg(terrain_type="plane"),
|
|
num_envs=4,
|
|
env_spacing=1.0,
|
|
entities={
|
|
"object": VariantEntityCfg(
|
|
variants={
|
|
"sphere": _simple_sphere_spec,
|
|
"cone": _simple_cone_spec,
|
|
},
|
|
init_state=EntityCfg.InitialStateCfg(pos=(0.0, 0.0, 0.2)),
|
|
)
|
|
},
|
|
),
|
|
)
|
|
|
|
env = ManagerBasedRlEnv(cfg=env_cfg, device="cpu")
|
|
try:
|
|
server = _Server()
|
|
scene = MjlabViserScene(
|
|
cast(Any, server),
|
|
env.sim.mj_model,
|
|
env.num_envs,
|
|
sim_model=env.sim.model,
|
|
expanded_fields=env.sim.expanded_fields,
|
|
)
|
|
groups: list[_PerWorldHullGroup] = list(scene._hull_per_world_groups)
|
|
# Two distinct variants -> at least two hull handles on the same body.
|
|
assert len(groups) >= 2, f"expected >=2 hull variants, got {len(groups)}"
|
|
all_envs = np.concatenate([g.env_ids for g in groups])
|
|
assert sorted(all_envs.tolist()) == list(range(env.num_envs))
|
|
# Hulls must be shape-distinct, not all copies of env0's hull.
|
|
shapes = {(g.handle.vertices.shape, g.handle.faces.shape) for g in groups}
|
|
assert len(shapes) >= 2, (
|
|
f"hull variants collapsed to one shape: {shapes} "
|
|
"(all envs would share env0's hull)"
|
|
)
|
|
|
|
body_xpos = env.sim.data.xpos.cpu().numpy()
|
|
body_xmat = env.sim.data.xmat.cpu().numpy()
|
|
scene.show_convex_hull = True
|
|
scene.show_only_selected = True
|
|
for target_env in range(env.num_envs):
|
|
scene.update_from_arrays(body_xpos, body_xmat, env_idx=target_env)
|
|
visible_groups = [g for g in groups if g.handle.visible]
|
|
assert len(visible_groups) == 1
|
|
assert target_env in visible_groups[0].env_ids
|
|
assert visible_groups[0].handle.batched_positions.shape[0] == 1
|
|
|
|
scene.show_only_selected = False
|
|
scene.update_from_arrays(body_xpos, body_xmat, env_idx=0)
|
|
assert all(g.handle.visible for g in groups)
|
|
finally:
|
|
env.close()
|
|
|
|
|
|
# DR consistency on variant scenes.
|
|
|
|
|
|
def _explicit_mass_variant(
|
|
mesh_name: str,
|
|
mass: float,
|
|
*,
|
|
cone: bool = False,
|
|
) -> mujoco.MjSpec:
|
|
"""Build a single-geom freejoint variant with an explicit body mass."""
|
|
spec = mujoco.MjSpec()
|
|
mesh = spec.add_mesh()
|
|
mesh.name = mesh_name
|
|
if cone:
|
|
mesh.make_cone(nedge=8, radius=0.05)
|
|
else:
|
|
mesh.make_sphere(subdivision=1)
|
|
body = spec.worldbody.add_body(name="prop")
|
|
body.add_freejoint()
|
|
body.explicitinertial = True
|
|
body.mass = mass
|
|
body.ipos[:] = (0.0, 0.0, 0.0)
|
|
body.inertia[:] = (1e-4, 1e-4, 1e-4)
|
|
body.iquat[:] = (1.0, 0.0, 0.0, 0.0)
|
|
body.add_geom(
|
|
name="visual",
|
|
type=mujoco.mjtGeom.mjGEOM_MESH,
|
|
meshname=mesh_name,
|
|
contype=0,
|
|
conaffinity=0,
|
|
mass=0.0,
|
|
)
|
|
return spec
|
|
|
|
|
|
def test_dr_body_mass_scale_preserves_variant_baseline():
|
|
"""``dr.body_mass`` scale must use each variant's own baseline.
|
|
|
|
This is the load-bearing claim of ``_per_world_default_fields``: scaling
|
|
body_mass on a variant scene by a per-env factor must produce
|
|
``variant_default[env] * scale[env]``, not ``template_default * scale[env]``.
|
|
"""
|
|
from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg
|
|
from mjlab.envs.mdp import dr
|
|
from mjlab.managers.event_manager import EventTermCfg
|
|
from mjlab.managers.scene_entity_config import SceneEntityCfg
|
|
from mjlab.scene import SceneCfg
|
|
from mjlab.terrains import TerrainEntityCfg
|
|
|
|
light_mass = 0.1
|
|
heavy_mass = 1.0
|
|
scale = 2.0
|
|
|
|
object_cfg = VariantEntityCfg(
|
|
variants={
|
|
"light": lambda: _explicit_mass_variant("light", light_mass),
|
|
"heavy": lambda: _explicit_mass_variant("heavy", heavy_mass, cone=True),
|
|
},
|
|
init_state=EntityCfg.InitialStateCfg(pos=(0.0, 0.0, 0.2)),
|
|
)
|
|
env_cfg = ManagerBasedRlEnvCfg(
|
|
decimation=1,
|
|
scene=SceneCfg(
|
|
terrain=TerrainEntityCfg(terrain_type="plane"),
|
|
num_envs=4,
|
|
env_spacing=1.0,
|
|
entities={"object": object_cfg},
|
|
),
|
|
events={
|
|
"scale_mass": EventTermCfg(
|
|
func=dr.body_mass,
|
|
mode="startup",
|
|
params={
|
|
"asset_cfg": SceneEntityCfg("object", body_names=("prop",)),
|
|
"operation": "scale",
|
|
"ranges": (scale, scale), # deterministic factor
|
|
},
|
|
),
|
|
},
|
|
)
|
|
|
|
with pytest.warns(UserWarning, match="dr.body_mass only randomizes mass"):
|
|
env = ManagerBasedRlEnv(cfg=env_cfg, device="cpu")
|
|
try:
|
|
obj_body = int(env.scene["object"].indexing.root_body_id)
|
|
w2v = env.sim.world_to_variant["object"]
|
|
actual = env.sim.model.body_mass[:, obj_body].cpu()
|
|
|
|
variant_baseline = torch.tensor([light_mass, heavy_mass], dtype=actual.dtype)
|
|
expected = variant_baseline[w2v.cpu()] * scale
|
|
torch.testing.assert_close(actual, expected, atol=1e-5, rtol=1e-5)
|
|
|
|
# Sanity: at least one env per variant, otherwise the test is vacuous.
|
|
assert (w2v == 0).any() and (w2v == 1).any()
|
|
finally:
|
|
env.close()
|
|
|
|
|
|
# Full env lifecycle.
|
|
|
|
|
|
def test_env_step_with_variants():
|
|
"""Build a full ManagerBasedRlEnv with variants; step without crashing."""
|
|
from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg
|
|
from mjlab.envs.mdp.events import reset_root_state_uniform
|
|
from mjlab.managers.event_manager import EventTermCfg
|
|
from mjlab.managers.scene_entity_config import SceneEntityCfg
|
|
from mjlab.scene import SceneCfg
|
|
from mjlab.terrains import TerrainEntityCfg
|
|
|
|
object_cfg = VariantEntityCfg(
|
|
variants={
|
|
"sphere": _simple_sphere_spec,
|
|
"cone": _simple_cone_spec,
|
|
},
|
|
init_state=EntityCfg.InitialStateCfg(pos=(0.0, 0.0, 0.2)),
|
|
)
|
|
|
|
env_cfg = ManagerBasedRlEnvCfg(
|
|
decimation=2,
|
|
scene=SceneCfg(
|
|
terrain=TerrainEntityCfg(terrain_type="plane"),
|
|
num_envs=4,
|
|
env_spacing=1.0,
|
|
entities={"object": object_cfg},
|
|
),
|
|
events={
|
|
"reset": EventTermCfg(
|
|
func=reset_root_state_uniform,
|
|
mode="reset",
|
|
params={
|
|
"pose_range": {},
|
|
"velocity_range": {},
|
|
"asset_cfg": SceneEntityCfg("object"),
|
|
},
|
|
),
|
|
},
|
|
)
|
|
|
|
env = ManagerBasedRlEnv(cfg=env_cfg, device="cpu")
|
|
obs, _ = env.reset()
|
|
actions = torch.zeros(env.num_envs, 0)
|
|
for _ in range(10):
|
|
obs, rew, term, trunc, info = env.step(actions)
|
|
# No NaN in positions.
|
|
qpos = env.sim.data.qpos[:].cpu().numpy()
|
|
assert np.all(np.isfinite(qpos))
|
|
env.close()
|
|
|
|
|
|
# Viewer: sameframe shortcut fix.
|
|
|
|
|
|
def _viewer_regression_sphere_spec() -> mujoco.MjSpec:
|
|
spec = mujoco.MjSpec()
|
|
m = spec.add_mesh()
|
|
m.name = "sphere"
|
|
m.make_sphere(subdivision=3)
|
|
m.scale[:] = (0.05, 0.05, 0.05)
|
|
body = spec.worldbody.add_body()
|
|
body.name = "prop"
|
|
body.add_freejoint()
|
|
g = body.add_geom()
|
|
g.name = "visual"
|
|
g.type = mujoco.mjtGeom.mjGEOM_MESH
|
|
g.meshname = "sphere"
|
|
return spec
|
|
|
|
|
|
def _viewer_regression_cone_spec() -> mujoco.MjSpec:
|
|
spec = mujoco.MjSpec()
|
|
m = spec.add_mesh()
|
|
m.name = "cone"
|
|
m.make_cone(nedge=16, radius=0.04)
|
|
m.scale[:] = (0.05, 0.05, 0.05)
|
|
body = spec.worldbody.add_body()
|
|
body.name = "prop"
|
|
body.add_freejoint()
|
|
g = body.add_geom()
|
|
g.name = "visual"
|
|
g.type = mujoco.mjtGeom.mjGEOM_MESH
|
|
g.meshname = "cone"
|
|
return spec
|
|
|
|
|
|
def test_sameframe_fix_makes_host_forward_match_variant():
|
|
"""Clearing sameframe shortcuts aligns host mj_forward with variant."""
|
|
base_model = _viewer_regression_sphere_spec().compile()
|
|
cone_model = _viewer_regression_cone_spec().compile()
|
|
|
|
# Sync cone's kinematic fields onto sphere's model (like viewer does).
|
|
for field in (
|
|
"geom_size",
|
|
"geom_pos",
|
|
"geom_quat",
|
|
"body_mass",
|
|
"body_inertia",
|
|
"body_ipos",
|
|
"body_iquat",
|
|
):
|
|
getattr(base_model, field)[:] = getattr(cone_model, field)
|
|
|
|
base_data = mujoco.MjData(base_model)
|
|
base_data.qpos[:] = cone_model.qpos0
|
|
base_data.qpos[2] = 0.05
|
|
mujoco.mj_forward(base_model, base_data)
|
|
|
|
cone_data = mujoco.MjData(cone_model)
|
|
cone_data.qpos[:] = cone_model.qpos0
|
|
cone_data.qpos[2] = 0.05
|
|
mujoco.mj_forward(cone_model, cone_data)
|
|
|
|
# Before fix: positions differ due to stale sameframe flags.
|
|
assert not np.allclose(base_data.geom_xpos, cone_data.geom_xpos)
|
|
|
|
# After fix: clearing sameframe makes them match.
|
|
disable_model_sameframe_shortcuts(base_model)
|
|
mujoco.mj_forward(base_model, base_data)
|
|
np.testing.assert_allclose(base_data.geom_xpos, cone_data.geom_xpos, atol=1e-6)
|
|
|
|
|
|
def test_sync_model_fields_copies_only_requested_env_fields():
|
|
"""Viewer model sync copies explicit fields and leaves others unchanged."""
|
|
model = _simple_sphere_spec().compile()
|
|
|
|
class _SimModel:
|
|
geom_rgba = torch.tensor(
|
|
[
|
|
[[0.1, 0.2, 0.3, 0.4]],
|
|
[[0.5, 0.6, 0.7, 0.8]],
|
|
],
|
|
dtype=torch.float32,
|
|
)
|
|
geom_pos = torch.tensor(
|
|
[
|
|
[[1.0, 2.0, 3.0]],
|
|
[[4.0, 5.0, 6.0]],
|
|
],
|
|
dtype=torch.float32,
|
|
)
|
|
|
|
original_geom_pos = model.geom_pos.copy()
|
|
|
|
sync_model_fields(model, _SimModel(), {"geom_rgba"}, env_idx=1)
|
|
|
|
np.testing.assert_allclose(model.geom_rgba, [[0.5, 0.6, 0.7, 0.8]])
|
|
np.testing.assert_allclose(model.geom_pos, original_geom_pos)
|