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
871 lines
26 KiB
Python
871 lines
26 KiB
Python
"""Tests for entity module."""
|
|
|
|
from dataclasses import dataclass
|
|
|
|
import mujoco
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
from conftest import get_test_device, load_fixture_xml
|
|
|
|
from mjlab.actuator import BuiltinPositionActuatorCfg, XmlActuatorCfg
|
|
from mjlab.entity import Entity, EntityArticulationInfoCfg, EntityCfg
|
|
from mjlab.scene import Scene, SceneCfg
|
|
from mjlab.sim.sim import Simulation, SimulationCfg
|
|
from mjlab.utils.spec_config import CollisionCfg, GeomCfg
|
|
|
|
FIXED_BASE_XML = """
|
|
<mujoco>
|
|
<worldbody>
|
|
<body name="object" pos="0 0 0.5">
|
|
<geom name="object_geom" type="box" size="0.1 0.1 0.1" rgba="0.8 0.3 0.3 1"/>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
|
|
FLOATING_BASE_XML = """
|
|
<mujoco>
|
|
<worldbody>
|
|
<body name="object" pos="0 0 1">
|
|
<freejoint name="free_joint"/>
|
|
<geom name="object_geom" type="box" size="0.1 0.1 0.1" rgba="0.3 0.3 0.8 1" mass="0.1"/>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
|
|
FIXED_BASE_ARTICULATED_XML = """
|
|
<mujoco>
|
|
<worldbody>
|
|
<body name="base" pos="0 0 0.5">
|
|
<geom name="base_geom" type="cylinder" size="0.1 0.05" mass="5.0"/>
|
|
<body name="link1" pos="0 0 0.1">
|
|
<joint name="joint1" type="hinge" axis="0 0 1" range="-3.14 3.14"/>
|
|
<geom name="link1_geom" type="box" size="0.05 0.05 0.2" mass="1.0"/>
|
|
<body name="link2" pos="0 0 0.4">
|
|
<joint name="joint2" type="hinge" axis="0 1 0" range="-1.57 1.57"/>
|
|
<geom name="link2_geom" type="box" size="0.05 0.05 0.15" mass="0.5"/>
|
|
</body>
|
|
</body>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
|
|
FLOATING_BASE_ARTICULATED_XML = load_fixture_xml("floating_base_articulated")
|
|
|
|
ACTUATOR_ORDER_TEST_XML = """
|
|
<mujoco>
|
|
<worldbody>
|
|
<body name="base" pos="0 0 0">
|
|
<inertial pos="0 0 0" mass="1" diaginertia="0.01 0.01 0.01"/>
|
|
<geom type="box" size="0.1 0.1 0.1" rgba="0.5 0.5 0.5 1"/>
|
|
<joint name="joint_a" axis="1 0 0" range="-1 1"/>
|
|
<body name="link1" pos="0.2 0 0">
|
|
<inertial pos="0 0 0" mass="1" diaginertia="0.01 0.01 0.01"/>
|
|
<geom type="box" size="0.1 0.1 0.1" rgba="0.5 0.5 0.5 1"/>
|
|
<joint name="joint_b" axis="0 1 0" range="-1 1"/>
|
|
<body name="link2" pos="0.2 0 0">
|
|
<inertial pos="0 0 0" mass="1" diaginertia="0.01 0.01 0.01"/>
|
|
<geom type="box" size="0.1 0.1 0.1" rgba="0.5 0.5 0.5 1"/>
|
|
<joint name="joint_c" axis="0 0 1" range="-1 1"/>
|
|
</body>
|
|
</body>
|
|
</body>
|
|
</worldbody>
|
|
<actuator>
|
|
<position name="act_c" joint="joint_c" kp="10"/>
|
|
<position name="act_b" joint="joint_b" kp="10"/>
|
|
<position name="act_a" joint="joint_a" kp="10"/>
|
|
</actuator>
|
|
</mujoco>
|
|
"""
|
|
|
|
UNDERACTUATED_XML = """
|
|
<mujoco>
|
|
<worldbody>
|
|
<body name="base" pos="0 0 0">
|
|
<inertial pos="0 0 0" mass="1" diaginertia="0.01 0.01 0.01"/>
|
|
<geom type="box" size="0.1 0.1 0.1" rgba="0.5 0.5 0.5 1"/>
|
|
<joint name="joint_a" axis="1 0 0" range="-1 1"/>
|
|
<body name="link1" pos="0.2 0 0">
|
|
<inertial pos="0 0 0" mass="1" diaginertia="0.01 0.01 0.01"/>
|
|
<geom type="box" size="0.1 0.1 0.1" rgba="0.5 0.5 0.5 1"/>
|
|
<joint name="joint_b" axis="0 1 0" range="-1 1"/>
|
|
<body name="link2" pos="0.2 0 0">
|
|
<inertial pos="0 0 0" mass="1" diaginertia="0.01 0.01 0.01"/>
|
|
<geom type="box" size="0.1 0.1 0.1" rgba="0.5 0.5 0.5 1"/>
|
|
<joint name="joint_c" axis="0 0 1" range="-1 1"/>
|
|
</body>
|
|
</body>
|
|
</body>
|
|
</worldbody>
|
|
<actuator>
|
|
<position name="act_c" joint="joint_c" kp="10"/>
|
|
</actuator>
|
|
</mujoco>
|
|
"""
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def device():
|
|
"""Test device fixture."""
|
|
return get_test_device()
|
|
|
|
|
|
def create_fixed_base_entity():
|
|
"""Create a simple fixed-base entity."""
|
|
cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(FIXED_BASE_XML))
|
|
return Entity(cfg)
|
|
|
|
|
|
def create_floating_base_entity():
|
|
"""Create a floating-base entity."""
|
|
cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(FLOATING_BASE_XML))
|
|
return Entity(cfg)
|
|
|
|
|
|
def create_fixed_articulated_entity():
|
|
"""Create a fixed-base articulated entity (e.g., robot arm)."""
|
|
cfg = EntityCfg(
|
|
spec_fn=lambda: mujoco.MjSpec.from_string(FIXED_BASE_ARTICULATED_XML),
|
|
articulation=EntityArticulationInfoCfg(
|
|
actuators=(
|
|
BuiltinPositionActuatorCfg(
|
|
target_names_expr=("joint1", "joint2"),
|
|
effort_limit=1.0,
|
|
stiffness=1.0,
|
|
damping=1.0,
|
|
),
|
|
)
|
|
),
|
|
)
|
|
return Entity(cfg)
|
|
|
|
|
|
def create_floating_articulated_entity():
|
|
"""Create a floating-base articulated entity."""
|
|
cfg = EntityCfg(
|
|
spec_fn=lambda: mujoco.MjSpec.from_string(FLOATING_BASE_ARTICULATED_XML),
|
|
articulation=EntityArticulationInfoCfg(
|
|
actuators=(
|
|
BuiltinPositionActuatorCfg(
|
|
target_names_expr=("joint1", "joint2"),
|
|
effort_limit=1.0,
|
|
stiffness=1.0,
|
|
damping=1.0,
|
|
),
|
|
)
|
|
),
|
|
)
|
|
return Entity(cfg)
|
|
|
|
|
|
def initialize_entity_with_sim(entity, device, num_envs=1):
|
|
"""Initialize an entity with a simulation."""
|
|
model = entity.compile()
|
|
sim_cfg = SimulationCfg(njmax=75)
|
|
sim = Simulation(num_envs=num_envs, cfg=sim_cfg, model=model, device=device)
|
|
entity.initialize(model, sim.model, sim.data, device)
|
|
return entity, sim
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"entity_fn,expected",
|
|
[
|
|
(
|
|
create_fixed_base_entity,
|
|
{
|
|
"is_fixed_base": True,
|
|
"is_articulated": False,
|
|
"is_actuated": False,
|
|
"num_bodies": 2, # mocap_base wrapper + object
|
|
"num_joints": 0,
|
|
"num_actuators": 0,
|
|
},
|
|
),
|
|
(
|
|
create_floating_base_entity,
|
|
{
|
|
"is_fixed_base": False,
|
|
"is_articulated": False,
|
|
"is_actuated": False,
|
|
"num_bodies": 1,
|
|
"num_joints": 0,
|
|
"num_actuators": 0,
|
|
},
|
|
),
|
|
(
|
|
create_fixed_articulated_entity,
|
|
{
|
|
"is_fixed_base": True,
|
|
"is_articulated": True,
|
|
"is_actuated": True,
|
|
"num_bodies": 4, # mocap_base wrapper + base + link1 + link2
|
|
"num_joints": 2,
|
|
"num_actuators": 2,
|
|
},
|
|
),
|
|
(
|
|
create_floating_articulated_entity,
|
|
{
|
|
"is_fixed_base": False,
|
|
"is_articulated": True,
|
|
"is_actuated": True,
|
|
"num_bodies": 3,
|
|
"num_joints": 2,
|
|
"num_actuators": 2,
|
|
},
|
|
),
|
|
],
|
|
)
|
|
def test_entity_properties(entity_fn, expected):
|
|
"""Test entity type properties and element counts."""
|
|
entity = entity_fn()
|
|
for prop, value in expected.items():
|
|
assert getattr(entity, prop) == value
|
|
|
|
|
|
def test_unnamed_freejoint_gets_default_name():
|
|
"""Test that an unnamed freejoint is auto-named during entity init."""
|
|
xml = """
|
|
<mujoco>
|
|
<worldbody>
|
|
<body name="object" pos="0 0 1">
|
|
<freejoint/>
|
|
<geom name="object_geom" type="box" size="0.1 0.1 0.1" mass="0.1"/>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(xml))
|
|
entity = Entity(cfg)
|
|
assert "floating_base_joint" in entity.all_joint_names
|
|
|
|
|
|
def test_multiple_freejoints_raises():
|
|
"""An entity with more than one freejoint is rejected at construction."""
|
|
xml = """
|
|
<mujoco>
|
|
<worldbody>
|
|
<body name="object_a" pos="0 0 1">
|
|
<freejoint/>
|
|
<geom type="box" size="0.1 0.1 0.1" mass="0.1"/>
|
|
</body>
|
|
<body name="object_b" pos="1 0 1">
|
|
<freejoint/>
|
|
<geom type="box" size="0.1 0.1 0.1" mass="0.1"/>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(xml))
|
|
with pytest.raises(ValueError, match="2 freejoints"):
|
|
Entity(cfg)
|
|
|
|
|
|
def test_geom_editor_applied():
|
|
"""Test that geom editors are applied during entity init."""
|
|
cfg = EntityCfg(
|
|
spec_fn=lambda: mujoco.MjSpec.from_string(FIXED_BASE_ARTICULATED_XML),
|
|
geoms=(GeomCfg(geom_names_expr=("link.*_geom",), group=3),),
|
|
)
|
|
entity = Entity(cfg)
|
|
|
|
assert entity.spec.geom("link1_geom").group == 3
|
|
assert entity.spec.geom("link2_geom").group == 3
|
|
assert entity.spec.geom("base_geom").group == 0
|
|
|
|
|
|
def test_geom_collision_overlap_warns():
|
|
"""A GeomCfg collision patch clobbered by a CollisionCfg triggers a warning."""
|
|
cfg = EntityCfg(
|
|
spec_fn=lambda: mujoco.MjSpec.from_string(FIXED_BASE_ARTICULATED_XML),
|
|
geoms=(GeomCfg(geom_names_expr=("link1_geom",), condim=6),),
|
|
collisions=(
|
|
CollisionCfg(
|
|
geom_names_expr=("link.*_geom",), contype=1, conaffinity=1, condim=3, priority=0
|
|
),
|
|
),
|
|
)
|
|
with pytest.warns(UserWarning, match="link1_geom.condim"):
|
|
Entity(cfg)
|
|
|
|
|
|
def test_find_methods():
|
|
"""Test find methods with exact and regex matches."""
|
|
entity = create_floating_articulated_entity()
|
|
|
|
# Test exact matches.
|
|
assert entity.find_bodies("base")[1] == ["base"]
|
|
assert entity.find_joints("joint1")[1] == ["joint1"]
|
|
assert entity.find_sites("site1")[1] == ["site1"]
|
|
|
|
# Test regex matches.
|
|
assert entity.find_bodies("link.*")[1] == ["link1", "link2"]
|
|
assert entity.find_joints("joint.*")[1] == ["joint1", "joint2"]
|
|
|
|
|
|
def test_find_with_subset_filtering():
|
|
"""Test find methods with subset filtering."""
|
|
entity = create_floating_articulated_entity()
|
|
|
|
# Test subset filtering.
|
|
_, names = entity.find_joints("joint1", joint_subset=["joint1", "joint2"])
|
|
assert names == ["joint1"]
|
|
|
|
# Test error on invalid subset.
|
|
with pytest.raises(ValueError, match="Not all regular expressions are matched"):
|
|
entity.find_joints("joint1", joint_subset=["joint2"])
|
|
|
|
|
|
def test_root_state_read_write(device):
|
|
"""Test root state can be written and read from simulation."""
|
|
entity = create_floating_base_entity()
|
|
entity, sim = initialize_entity_with_sim(entity, device)
|
|
|
|
# fmt: off
|
|
root_state = torch.tensor([
|
|
1.0, 2.0, 3.0, # position
|
|
1.0, 0.0, 0.0, 0.0, # quaternion (identity)
|
|
0.5, 0.0, 0.0, # linear velocity in X
|
|
0.0, 0.0, 0.2 # angular velocity around Z
|
|
], device=device).unsqueeze(0)
|
|
# fmt: on
|
|
|
|
entity.write_root_state_to_sim(root_state)
|
|
|
|
# Verify the state was actually written.
|
|
q_slice = entity.data.indexing.free_joint_q_adr
|
|
v_slice = entity.data.indexing.free_joint_v_adr
|
|
assert torch.allclose(sim.data.qpos[:, q_slice], root_state[:, :7])
|
|
assert torch.allclose(sim.data.qvel[:, v_slice], root_state[:, 7:])
|
|
|
|
|
|
def test_external_force_and_torque(device):
|
|
"""Test forces translate, torques rotate, and forces can be cleared."""
|
|
entity = create_floating_base_entity()
|
|
entity, sim = initialize_entity_with_sim(entity, device)
|
|
|
|
# Apply force in X, torque around Z.
|
|
entity.write_external_wrench_to_sim(
|
|
forces=torch.tensor([[5.0, 0.0, 0.0]], device=sim.device),
|
|
torques=torch.tensor([[0.0, 0.0, 3.0]], device=sim.device),
|
|
)
|
|
|
|
initial_pos = sim.data.qpos[0, :3].clone()
|
|
initial_quat = sim.data.qpos[0, 3:7].clone()
|
|
|
|
for _ in range(10):
|
|
sim.step()
|
|
|
|
# Verify X translation and rotation occurred.
|
|
assert sim.data.qpos[0, 0] > initial_pos[0], "Force should cause X translation"
|
|
assert not torch.allclose(sim.data.qpos[0, 3:7], initial_quat), (
|
|
"Torque should cause rotation"
|
|
)
|
|
|
|
# Verify angular velocity is primarily around Z.
|
|
angular_vel = sim.data.qvel[0, 3:6]
|
|
z_rotation = abs(angular_vel[2])
|
|
xy_rotation = abs(angular_vel[0]) + abs(angular_vel[1])
|
|
assert z_rotation > xy_rotation * 5, "Rotation should be primarily around Z axis"
|
|
|
|
|
|
def test_external_force_clearing(device):
|
|
"""Test external forces can be cleared."""
|
|
entity = create_floating_base_entity()
|
|
entity, sim = initialize_entity_with_sim(entity, device)
|
|
|
|
# Apply force.
|
|
entity.write_external_wrench_to_sim(
|
|
forces=torch.tensor([[5.0, 0.0, 0.0]], device=sim.device),
|
|
torques=torch.tensor([[0.0, 0.0, 3.0]], device=sim.device),
|
|
)
|
|
|
|
# Clear forces.
|
|
entity.write_external_wrench_to_sim(
|
|
forces=torch.zeros((1, 3), device=sim.device),
|
|
torques=torch.zeros((1, 3), device=sim.device),
|
|
)
|
|
|
|
body_id = entity.indexing.body_ids[0]
|
|
assert torch.allclose(
|
|
sim.data.xfrc_applied[:, body_id, :], torch.zeros(6, device=sim.device)
|
|
)
|
|
|
|
|
|
def test_external_force_on_specific_body(device):
|
|
"""Test applying force to specific body in articulated system."""
|
|
entity = create_floating_articulated_entity()
|
|
entity, sim = initialize_entity_with_sim(entity, device)
|
|
|
|
# Apply force only to link1.
|
|
body_ids = entity.find_bodies("link1")[0]
|
|
entity.write_external_wrench_to_sim(
|
|
forces=torch.tensor([[3.0, 0.0, 0.0]], device=sim.device),
|
|
torques=torch.zeros((1, 3), device=sim.device),
|
|
body_ids=body_ids,
|
|
)
|
|
|
|
# Verify force applied only to link1.
|
|
link1_id = sim.mj_model.body("link1").id
|
|
base_id = sim.mj_model.body("base").id
|
|
assert torch.allclose(
|
|
sim.data.xfrc_applied[0, link1_id, :3],
|
|
torch.tensor([3.0, 0.0, 0.0], device=sim.device),
|
|
)
|
|
assert torch.allclose(
|
|
sim.data.xfrc_applied[0, base_id, :3], torch.zeros(3, device=sim.device)
|
|
)
|
|
|
|
# Verify motion occurs.
|
|
initial_pos = sim.data.xpos[0, link1_id, :].clone()
|
|
for _ in range(10):
|
|
sim.step()
|
|
assert not torch.allclose(sim.data.xpos[0, link1_id, :], initial_pos)
|
|
|
|
|
|
def test_fixed_base_initial_position():
|
|
"""Test fixed-base entity's initial pos/rot are applied to the mocap wrapper."""
|
|
cfg = EntityCfg(
|
|
spec_fn=lambda: mujoco.MjSpec.from_string(FIXED_BASE_XML),
|
|
init_state=EntityCfg.InitialStateCfg((1.0, 2.0, 3.0), (0.7071, 0.7071, 0.0, 0.0)),
|
|
)
|
|
entity = Entity(cfg)
|
|
model = entity.compile()
|
|
|
|
# init_state is applied to the auto-generated mocap_base wrapper body.
|
|
mocap_body = model.body("mocap_base")
|
|
np.testing.assert_allclose(mocap_body.pos, [1.0, 2.0, 3.0], rtol=1e-6)
|
|
np.testing.assert_allclose(mocap_body.quat, [0.7071, 0.7071, 0.0, 0.0], atol=1e-4)
|
|
|
|
|
|
def test_keyframe_ctrl_maps_joint_pos_to_actuators():
|
|
"""Test keyframe ctrl values match init_state joint positions."""
|
|
cfg = EntityCfg(
|
|
spec_fn=lambda: mujoco.MjSpec.from_string(FLOATING_BASE_ARTICULATED_XML),
|
|
articulation=EntityArticulationInfoCfg(
|
|
actuators=(
|
|
BuiltinPositionActuatorCfg(
|
|
target_names_expr=(
|
|
"joint1",
|
|
"joint2",
|
|
),
|
|
effort_limit=1.0,
|
|
stiffness=1.0,
|
|
damping=1.0,
|
|
),
|
|
)
|
|
),
|
|
init_state=EntityCfg.InitialStateCfg(joint_pos={"joint1": 0.5, "joint2": -0.25}),
|
|
)
|
|
model = Entity(cfg).compile()
|
|
|
|
assert model.nkey == 1
|
|
assert model.nu == 2
|
|
assert list(model.key("init_state").ctrl) == [0.5, -0.25]
|
|
|
|
|
|
def test_keyframe_ctrl_underactuated():
|
|
"""Test ctrl is correctly constructed for an underactuated system."""
|
|
cfg = EntityCfg(
|
|
spec_fn=lambda: mujoco.MjSpec.from_string(FLOATING_BASE_ARTICULATED_XML),
|
|
articulation=EntityArticulationInfoCfg(
|
|
actuators=(
|
|
BuiltinPositionActuatorCfg(
|
|
target_names_expr=("joint1",), # Only one actuator.
|
|
effort_limit=1.0,
|
|
stiffness=1.0,
|
|
damping=1.0,
|
|
),
|
|
)
|
|
),
|
|
init_state=EntityCfg.InitialStateCfg(joint_pos={"joint1": 0.42, "joint2": -0.99}),
|
|
)
|
|
model = Entity(cfg).compile()
|
|
|
|
assert model.nu == 1
|
|
assert model.key_ctrl[0, 0] == 0.42
|
|
|
|
|
|
def test_fixed_base_mocap_runtime_pose_change(device):
|
|
"""Test fixed-base mocap entity can have its pose changed at runtime."""
|
|
|
|
def spec_fn():
|
|
spec = mujoco.MjSpec.from_string(FIXED_BASE_ARTICULATED_XML)
|
|
spec.worldbody.first_body().mocap = True
|
|
return spec
|
|
|
|
cfg = EntityCfg(
|
|
spec_fn=spec_fn,
|
|
init_state=EntityCfg.InitialStateCfg((1.0, 2.0, 3.0), (1.0, 0.0, 0.0, 0.0)),
|
|
)
|
|
entity = Entity(cfg)
|
|
entity, sim = initialize_entity_with_sim(entity, device)
|
|
|
|
assert entity.indexing.mocap_id is not None
|
|
assert entity.is_mocap is True
|
|
|
|
# fmt: off
|
|
new_pose = torch.tensor([
|
|
5.0, 6.0, 7.0,
|
|
1.0, 0.0, 0.0, 0.0,
|
|
], device=device).unsqueeze(0)
|
|
# fmt: on
|
|
entity.write_mocap_pose_to_sim(new_pose)
|
|
|
|
sim.forward()
|
|
assert torch.allclose(entity.data.root_link_pose_w, new_pose, atol=1e-5)
|
|
|
|
|
|
def test_find_joints_by_actuator_names_preserves_natural_order(device):
|
|
"""Test that find_joints_by_actuator_names returns joints in natural joint order.
|
|
|
|
This is a regression test for a bug where joints were returned in actuator
|
|
definition order instead of natural joint order, breaking motion tracking tasks.
|
|
"""
|
|
robot_cfg = EntityCfg(
|
|
spec_fn=lambda: mujoco.MjSpec.from_string(ACTUATOR_ORDER_TEST_XML),
|
|
articulation=EntityArticulationInfoCfg(
|
|
actuators=(XmlActuatorCfg(target_names_expr=(".*",)),)
|
|
),
|
|
)
|
|
|
|
robot = Entity(robot_cfg)
|
|
robot.compile()
|
|
|
|
# Natural joint order should be: joint_a, joint_b, joint_c.
|
|
assert list(robot.joint_names) == ["joint_a", "joint_b", "joint_c"]
|
|
|
|
# Actuator order is: act_c, act_b, act_a (reverse).
|
|
# But find_joints_by_actuator_names should still return joints in natural order.
|
|
joint_ids, joint_names = robot.find_joints_by_actuator_names(".*")
|
|
|
|
# Critical: joints must be in natural order, not actuator order.
|
|
assert joint_names == ["joint_a", "joint_b", "joint_c"]
|
|
assert joint_ids == [0, 1, 2]
|
|
|
|
# Verify this differs from actuator order (which is reverse).
|
|
assert list(robot.actuator_names) == ["act_c", "act_b", "act_a"]
|
|
|
|
|
|
def test_ctrl_ids_follow_natural_joint_order(device):
|
|
"""Test that entity.indexing.ctrl_ids are in actuator definition order.
|
|
|
|
ctrl_ids follow actuator definition order for simplicity. ONNX export builds
|
|
the natural joint order mapping where needed.
|
|
"""
|
|
robot_cfg = EntityCfg(
|
|
spec_fn=lambda: mujoco.MjSpec.from_string(ACTUATOR_ORDER_TEST_XML),
|
|
articulation=EntityArticulationInfoCfg(
|
|
actuators=(XmlActuatorCfg(target_names_expr=(".*",)),)
|
|
),
|
|
)
|
|
|
|
robot = Entity(robot_cfg)
|
|
mj_model = robot.compile()
|
|
|
|
# Create simulation to initialize entity.
|
|
sim_cfg = SimulationCfg()
|
|
sim = Simulation(num_envs=1, cfg=sim_cfg, model=mj_model, device=device)
|
|
robot.initialize(sim.mj_model, sim.model, sim.data, device)
|
|
|
|
# Natural joint order: joint_a, joint_b, joint_c.
|
|
assert list(robot.joint_names) == ["joint_a", "joint_b", "joint_c"]
|
|
|
|
# Actuator definition order (from XML): act_c, act_b, act_a.
|
|
assert list(robot.actuator_names) == ["act_c", "act_b", "act_a"]
|
|
|
|
# ctrl_ids should be in actuator definition order (c, b, a).
|
|
ctrl_ids = robot.indexing.ctrl_ids.cpu().tolist()
|
|
|
|
# Map actuator names to their MuJoCo IDs in the compiled model.
|
|
actuator_name_to_id = {
|
|
mj_model.actuator(i).name.split("/")[-1]: i for i in range(mj_model.nu)
|
|
}
|
|
|
|
# ctrl_ids should be ordered as: act_c, act_b, act_a (actuator definition order).
|
|
expected_ctrl_ids = [
|
|
actuator_name_to_id["act_c"],
|
|
actuator_name_to_id["act_b"],
|
|
actuator_name_to_id["act_a"],
|
|
]
|
|
assert ctrl_ids == expected_ctrl_ids
|
|
|
|
|
|
def test_find_joints_by_actuator_names_returns_entity_local_indices():
|
|
"""Test that find_joints_by_actuator_names returns entity-local indices."""
|
|
robot_cfg = EntityCfg(
|
|
spec_fn=lambda: mujoco.MjSpec.from_string(UNDERACTUATED_XML),
|
|
articulation=EntityArticulationInfoCfg(
|
|
actuators=(XmlActuatorCfg(target_names_expr=(".*",)),)
|
|
),
|
|
)
|
|
|
|
robot = Entity(robot_cfg)
|
|
robot.compile()
|
|
|
|
# Natural joint order: joint_a (0), joint_b (1), joint_c (2).
|
|
assert list(robot.joint_names) == ["joint_a", "joint_b", "joint_c"]
|
|
|
|
# Only joint_c has an actuator.
|
|
joint_ids, joint_names = robot.find_joints_by_actuator_names(".*")
|
|
|
|
# Should return entity-local index [2], not subset-local [0].
|
|
assert joint_names == ["joint_c"]
|
|
assert joint_ids == [2] # Index of joint_c in self.joint_names.
|
|
|
|
|
|
@dataclass
|
|
class CustomEntityCfg(EntityCfg):
|
|
"""Custom entity config with additional fields."""
|
|
|
|
custom_threshold: float = 0.5
|
|
|
|
def build(self) -> "CustomEntity":
|
|
return CustomEntity(self)
|
|
|
|
|
|
class CustomEntity(Entity):
|
|
"""Custom entity with additional properties."""
|
|
|
|
cfg: CustomEntityCfg
|
|
|
|
@property
|
|
def custom_value(self) -> float:
|
|
"""Custom property that uses config field."""
|
|
return self.cfg.custom_threshold * 2
|
|
|
|
|
|
def test_custom_entity_subclass(device):
|
|
"""Test that custom Entity subclasses work through the scene."""
|
|
scene_cfg = SceneCfg(
|
|
num_envs=1,
|
|
entities={
|
|
"custom": CustomEntityCfg(
|
|
spec_fn=lambda: mujoco.MjSpec.from_string(FIXED_BASE_XML),
|
|
custom_threshold=0.9,
|
|
),
|
|
},
|
|
)
|
|
scene = Scene(scene_cfg, device)
|
|
|
|
# Scene should have instantiated our custom entity type.
|
|
custom_entity = scene.entities["custom"]
|
|
assert isinstance(custom_entity, CustomEntity)
|
|
assert custom_entity.cfg.custom_threshold == 0.9
|
|
assert custom_entity.custom_value == 1.8
|
|
|
|
|
|
# ============================================================================
|
|
# Keyframe Fallback Tests
|
|
# ============================================================================
|
|
|
|
XML_WITH_KEYFRAME = """
|
|
<mujoco>
|
|
<worldbody>
|
|
<body name="robot">
|
|
<freejoint name="root"/>
|
|
<geom type="box" size="0.1 0.1 0.1"/>
|
|
<body name="link" pos="0.2 0 0">
|
|
<joint name="joint1" type="hinge" axis="0 1 0"/>
|
|
<geom type="box" size="0.05 0.05 0.05"/>
|
|
</body>
|
|
</body>
|
|
</worldbody>
|
|
<keyframe>
|
|
<key name="home" qpos="0 0 1 1 0 0 0 0.5"/>
|
|
</keyframe>
|
|
</mujoco>
|
|
"""
|
|
|
|
XML_WITHOUT_KEYFRAME = """
|
|
<mujoco>
|
|
<worldbody>
|
|
<body name="robot">
|
|
<freejoint name="root"/>
|
|
<geom type="box" size="0.1 0.1 0.1"/>
|
|
<body name="link" pos="0.2 0 0">
|
|
<joint name="joint1" type="hinge" axis="0 1 0"/>
|
|
<geom type="box" size="0.05 0.05 0.05"/>
|
|
</body>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
|
|
|
|
def test_joint_pos_none_uses_model_keyframe():
|
|
"""Test that joint_pos=None uses the model's existing keyframe."""
|
|
cfg = EntityCfg(
|
|
init_state=EntityCfg.InitialStateCfg(joint_pos=None),
|
|
spec_fn=lambda: mujoco.MjSpec.from_string(XML_WITH_KEYFRAME),
|
|
)
|
|
entity = Entity(cfg)
|
|
model = entity.spec.compile()
|
|
|
|
assert model.nkey == 1
|
|
# Keyframe: qpos="0 0 1 1 0 0 0 0.5" (root pos/quat + joint1)
|
|
assert model.key(0).qpos[7] == 0.5 # joint1 position
|
|
|
|
|
|
def test_joint_pos_none_errors_without_keyframe():
|
|
"""Test that joint_pos=None raises error if model has no keyframe."""
|
|
cfg = EntityCfg(
|
|
init_state=EntityCfg.InitialStateCfg(joint_pos=None),
|
|
spec_fn=lambda: mujoco.MjSpec.from_string(XML_WITHOUT_KEYFRAME),
|
|
)
|
|
with pytest.raises(ValueError, match="requires the model to have a keyframe"):
|
|
Entity(cfg)
|
|
|
|
|
|
XML_FIXED_BASE_WITH_KEYFRAME = """
|
|
<mujoco>
|
|
<worldbody>
|
|
<body name="arm">
|
|
<joint name="joint1" type="hinge" axis="0 0 1"/>
|
|
<geom type="box" size="0.1 0.1 0.1"/>
|
|
</body>
|
|
</worldbody>
|
|
<keyframe>
|
|
<key name="home" qpos="0.5"/>
|
|
</keyframe>
|
|
</mujoco>
|
|
"""
|
|
|
|
|
|
def test_joint_pos_none_fixed_base_uses_keyframe():
|
|
"""Test that joint_pos=None works for fixed-base entities with keyframes."""
|
|
cfg = EntityCfg(
|
|
init_state=EntityCfg.InitialStateCfg(joint_pos=None),
|
|
spec_fn=lambda: mujoco.MjSpec.from_string(XML_FIXED_BASE_WITH_KEYFRAME),
|
|
)
|
|
entity = Entity(cfg)
|
|
model = entity.spec.compile()
|
|
|
|
assert model.nkey == 1
|
|
assert model.key(0).qpos[0] == 0.5
|
|
|
|
|
|
XML_WITH_SITES_AND_TENDONS = """
|
|
<mujoco>
|
|
<worldbody>
|
|
<body name="base">
|
|
<joint name="joint1" type="hinge" axis="0 0 1" range="-3.14 3.14"/>
|
|
<geom type="box" size="0.1 0.1 0.1" mass="1.0"/>
|
|
<site name="site1" pos="0.1 0 0"/>
|
|
<site name="site2" pos="0 0.1 0"/>
|
|
<body name="link1" pos="0 0 0.2">
|
|
<joint name="joint2" type="hinge" axis="0 1 0" range="-1.57 1.57"/>
|
|
<geom type="box" size="0.05 0.05 0.1" mass="0.5"/>
|
|
<site name="site3" pos="0 0 0.1"/>
|
|
</body>
|
|
</body>
|
|
</worldbody>
|
|
<tendon>
|
|
<fixed name="tendon1">
|
|
<joint joint="joint1" coef="1.0"/>
|
|
</fixed>
|
|
<fixed name="tendon2">
|
|
<joint joint="joint2" coef="1.0"/>
|
|
</fixed>
|
|
</tendon>
|
|
</mujoco>
|
|
"""
|
|
|
|
|
|
def test_tendon_and_site_targets_only_allocated_when_needed(device):
|
|
"""Test that tendon/site targets are only allocated when actuators use them."""
|
|
from mjlab.actuator import BuiltinMotorActuatorCfg
|
|
from mjlab.actuator.actuator import TransmissionType
|
|
|
|
# Entity with sites and tendons but NO site/tendon actuators.
|
|
# Should allocate empty tensors for site and tendon targets.
|
|
cfg = EntityCfg(
|
|
spec_fn=lambda: mujoco.MjSpec.from_string(XML_WITH_SITES_AND_TENDONS),
|
|
articulation=EntityArticulationInfoCfg(
|
|
actuators=(
|
|
# Only joint actuators, no tendon or site actuators.
|
|
BuiltinMotorActuatorCfg(
|
|
target_names_expr=("joint1",),
|
|
effort_limit=10.0,
|
|
transmission_type=TransmissionType.JOINT,
|
|
),
|
|
)
|
|
),
|
|
)
|
|
|
|
entity = Entity(cfg)
|
|
model = entity.compile()
|
|
sim = Simulation(num_envs=4, cfg=SimulationCfg(), model=model, device=device)
|
|
entity.initialize(model, sim.model, sim.data, device)
|
|
|
|
# Verify the entity has sites and tendons.
|
|
assert len(entity.site_names) == 3
|
|
assert len(entity.tendon_names) == 2
|
|
|
|
# Verify tendon and site targets are empty (not allocated).
|
|
assert entity.data.site_effort_target.shape == (4, 0)
|
|
assert entity.data.tendon_len_target.shape == (4, 0)
|
|
assert entity.data.tendon_vel_target.shape == (4, 0)
|
|
assert entity.data.tendon_effort_target.shape == (4, 0)
|
|
|
|
# Joint targets should still be allocated (2 joints).
|
|
assert entity.data.joint_pos_target.shape == (4, 2)
|
|
|
|
|
|
def test_add_actuators_wrong_namespace_error_includes_hint():
|
|
"""Error message includes namespace hint when name exists in another namespace."""
|
|
from mjlab.actuator.actuator import TransmissionType
|
|
|
|
cfg = EntityCfg(
|
|
spec_fn=lambda: mujoco.MjSpec.from_string(XML_WITH_SITES_AND_TENDONS),
|
|
articulation=EntityArticulationInfoCfg(
|
|
actuators=(
|
|
BuiltinPositionActuatorCfg(
|
|
target_names_expr=("tendon1",),
|
|
stiffness=10.0,
|
|
damping=1.0,
|
|
transmission_type=TransmissionType.JOINT,
|
|
),
|
|
)
|
|
),
|
|
)
|
|
with pytest.raises(ValueError, match="Matches were found in.*tendons.*tendon1"):
|
|
Entity(cfg)
|
|
|
|
|
|
@pytest.mark.filterwarnings("ignore:Actuator config matched")
|
|
def test_wildcard_warns_about_unactuated_namespaces():
|
|
"""Wildcard matching joints should warn about unactuated tendons."""
|
|
cfg = EntityCfg(
|
|
spec_fn=lambda: mujoco.MjSpec.from_string(XML_WITH_SITES_AND_TENDONS),
|
|
articulation=EntityArticulationInfoCfg(
|
|
actuators=(
|
|
BuiltinPositionActuatorCfg(
|
|
target_names_expr=(".*",),
|
|
stiffness=10.0,
|
|
damping=1.0,
|
|
),
|
|
)
|
|
),
|
|
)
|
|
with pytest.warns(match="also match.*tendon"):
|
|
Entity(cfg)
|
|
|
|
|
|
def test_set_joint_position_target_outer_product(device):
|
|
"""Tensor env_ids + tensor joint_ids select the outer product, not a diagonal."""
|
|
entity = create_fixed_articulated_entity()
|
|
entity, _ = initialize_entity_with_sim(entity, device, num_envs=2)
|
|
|
|
targets = torch.tensor([[1.0, 2.0], [3.0, 4.0]], device=device)
|
|
entity.set_joint_position_target(
|
|
targets,
|
|
joint_ids=torch.tensor([0, 1], device=device),
|
|
env_ids=torch.tensor([0, 1], device=device),
|
|
)
|
|
assert torch.equal(entity.data.joint_pos_target, targets)
|