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
379 lines
11 KiB
Python
379 lines
11 KiB
Python
"""Tests for builtin_sensor.py."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import copy
|
|
|
|
import mujoco
|
|
import pytest
|
|
import torch
|
|
from conftest import get_test_device
|
|
|
|
from mjlab.entity import EntityCfg
|
|
from mjlab.scene import Scene, SceneCfg
|
|
from mjlab.sensor.builtin_sensor import BuiltinSensorCfg, ObjRef
|
|
from mjlab.sim.sim import Simulation, SimulationCfg
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def device():
|
|
"""Test device fixture."""
|
|
return get_test_device()
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def articulated_robot_xml():
|
|
"""XML for a simple articulated robot with joints."""
|
|
return """
|
|
<mujoco>
|
|
<worldbody>
|
|
<geom name="floor" type="plane" size="5 5 0.1" pos="0 0 0"/>
|
|
<body name="base" pos="0 0 1">
|
|
<freejoint name="free_joint"/>
|
|
<geom name="base_geom" type="box" size="0.2 0.2 0.1" mass="5.0"/>
|
|
<site name="base_site" pos="0 0 0"/>
|
|
<body name="link1" pos="0.3 0 0">
|
|
<joint name="joint1" type="hinge" axis="0 0 1" range="-1.57 1.57"/>
|
|
<geom name="link1_geom" type="box" size="0.1 0.1 0.1" mass="1.0"/>
|
|
<site name="link1_site" pos="0 0 0"/>
|
|
</body>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def robot_with_xml_sensors():
|
|
"""XML for robot with sensors already defined in the XML."""
|
|
return """
|
|
<mujoco>
|
|
<worldbody>
|
|
<body name="base" pos="0 0 1">
|
|
<freejoint name="free_joint"/>
|
|
<geom name="base_geom" type="box" size="0.2 0.2 0.1" mass="5.0"/>
|
|
<site name="base_site" pos="0 0 0"/>
|
|
<body name="link1" pos="0.3 0 0">
|
|
<joint name="joint1" type="hinge" axis="0 0 1" range="-1.57 1.57"/>
|
|
<geom name="link1_geom" type="box" size="0.1 0.1 0.1" mass="1.0"/>
|
|
<site name="link1_site" pos="0 0 0"/>
|
|
</body>
|
|
</body>
|
|
</worldbody>
|
|
<sensor>
|
|
<jointpos name="xml_joint_sensor" joint="joint1"/>
|
|
<accelerometer name="xml_accel_sensor" site="base_site"/>
|
|
<gyro name="xml_gyro_sensor" site="link1_site"/>
|
|
</sensor>
|
|
</mujoco>
|
|
"""
|
|
|
|
|
|
def test_jointpos_sensor(articulated_robot_xml, device):
|
|
"""Verify joint pos sensor returns correctly shaped tensor for scalar joint values."""
|
|
entity_cfg = EntityCfg(
|
|
spec_fn=lambda: mujoco.MjSpec.from_string(articulated_robot_xml)
|
|
)
|
|
|
|
jointpos_sensor_cfg = BuiltinSensorCfg(
|
|
name="joint1_pos",
|
|
sensor_type="jointpos",
|
|
obj=ObjRef(type="joint", name="joint1", entity="robot"),
|
|
)
|
|
|
|
scene_cfg = SceneCfg(
|
|
num_envs=2,
|
|
env_spacing=3.0,
|
|
entities={"robot": entity_cfg},
|
|
sensors=(jointpos_sensor_cfg,),
|
|
)
|
|
|
|
scene = Scene(scene_cfg, device)
|
|
model = scene.compile()
|
|
sim_cfg = SimulationCfg(njmax=20)
|
|
sim = Simulation(num_envs=2, cfg=sim_cfg, model=model, device=device)
|
|
scene.initialize(sim.mj_model, sim.model, sim.data)
|
|
|
|
sensor = scene["robot/joint1_pos"]
|
|
sim.step()
|
|
data = sensor.data
|
|
|
|
assert isinstance(data, torch.Tensor)
|
|
assert data.shape == (2, 1)
|
|
|
|
|
|
def test_accelerometer_sensor(articulated_robot_xml, device):
|
|
"""Verify accelerometer reads non-zero acceleration when robot is on floor."""
|
|
entity_cfg = EntityCfg(
|
|
spec_fn=lambda: mujoco.MjSpec.from_string(articulated_robot_xml)
|
|
)
|
|
|
|
accel_sensor_cfg = BuiltinSensorCfg(
|
|
name="base_accel",
|
|
sensor_type="accelerometer",
|
|
obj=ObjRef(type="site", name="base_site", entity="robot"),
|
|
)
|
|
|
|
scene_cfg = SceneCfg(
|
|
num_envs=2,
|
|
env_spacing=3.0,
|
|
entities={"robot": entity_cfg},
|
|
sensors=(accel_sensor_cfg,),
|
|
)
|
|
|
|
scene = Scene(scene_cfg, device)
|
|
model = scene.compile()
|
|
# The robot falls and rests on the floor here, generating contact
|
|
# constraints; njmax must accommodate the resting nefc (~32), unlike the
|
|
# other tests which only step once before the robot lands.
|
|
sim_cfg = SimulationCfg(njmax=40)
|
|
sim = Simulation(num_envs=2, cfg=sim_cfg, model=model, device=device)
|
|
scene.initialize(sim.mj_model, sim.model, sim.data)
|
|
|
|
sensor = scene["robot/base_accel"]
|
|
|
|
# Step until the robot falls (from z=1) and rests on the floor. During free
|
|
# fall the accelerometer correctly reads ~0 (zero proper acceleration); it
|
|
# only registers the ~g ground reaction once the robot has landed (~280
|
|
# steps), so step well past that.
|
|
for _ in range(400):
|
|
sim.step()
|
|
|
|
data = sensor.data
|
|
|
|
assert isinstance(data, torch.Tensor)
|
|
assert data.shape == (2, 3)
|
|
assert torch.any(torch.abs(data) > 0)
|
|
|
|
|
|
def test_multiple_sensors(articulated_robot_xml, device):
|
|
"""Verify multiple sensors can be registered and return correctly shaped data."""
|
|
entity_cfg = EntityCfg(
|
|
spec_fn=lambda: mujoco.MjSpec.from_string(articulated_robot_xml)
|
|
)
|
|
|
|
jointpos_cfg = BuiltinSensorCfg(
|
|
name="joint1_pos",
|
|
sensor_type="jointpos",
|
|
obj=ObjRef(type="joint", name="joint1", entity="robot"),
|
|
)
|
|
jointvel_cfg = BuiltinSensorCfg(
|
|
name="joint1_vel",
|
|
sensor_type="jointvel",
|
|
obj=ObjRef(type="joint", name="joint1", entity="robot"),
|
|
)
|
|
gyro_cfg = BuiltinSensorCfg(
|
|
name="base_gyro",
|
|
sensor_type="gyro",
|
|
obj=ObjRef(type="site", name="base_site", entity="robot"),
|
|
)
|
|
|
|
scene_cfg = SceneCfg(
|
|
num_envs=1,
|
|
env_spacing=3.0,
|
|
entities={"robot": entity_cfg},
|
|
sensors=(jointpos_cfg, jointvel_cfg, gyro_cfg),
|
|
)
|
|
|
|
scene = Scene(scene_cfg, device)
|
|
model = scene.compile()
|
|
sim_cfg = SimulationCfg(njmax=20)
|
|
sim = Simulation(num_envs=1, cfg=sim_cfg, model=model, device=device)
|
|
scene.initialize(sim.mj_model, sim.model, sim.data)
|
|
|
|
jointpos_sensor = scene["robot/joint1_pos"]
|
|
jointvel_sensor = scene["robot/joint1_vel"]
|
|
gyro_sensor = scene["robot/base_gyro"]
|
|
|
|
sim.step()
|
|
|
|
pos_data = jointpos_sensor.data
|
|
vel_data = jointvel_sensor.data
|
|
gyro_data = gyro_sensor.data
|
|
|
|
assert pos_data.shape == (1, 1)
|
|
assert vel_data.shape == (1, 1)
|
|
assert gyro_data.shape == (1, 3)
|
|
|
|
|
|
def test_error_on_invalid_ref():
|
|
"""Verify ValueError is raised when ref provided for unsupported sensor type."""
|
|
with pytest.raises(ValueError, match="does not support ref specification"):
|
|
BuiltinSensorCfg(
|
|
name="invalid_sensor",
|
|
sensor_type="jointpos",
|
|
obj=ObjRef(type="joint", name="joint1", entity="robot"),
|
|
ref=ObjRef(type="body", name="base"),
|
|
)
|
|
|
|
|
|
def test_error_on_missing_obj():
|
|
"""Verify ValueError is raised when obj is not provided for required sensor type."""
|
|
with pytest.raises(ValueError, match="requires obj with type='joint'"):
|
|
BuiltinSensorCfg(
|
|
name="invalid_sensor",
|
|
sensor_type="jointpos",
|
|
)
|
|
|
|
|
|
def test_error_on_wrong_obj_type_for_site_sensor():
|
|
"""Verify ValueError is raised when wrong obj type is used for site sensor."""
|
|
with pytest.raises(ValueError, match="requires obj.type='site'"):
|
|
BuiltinSensorCfg(
|
|
name="invalid_sensor",
|
|
sensor_type="accelerometer",
|
|
obj=ObjRef(type="body", name="base"),
|
|
)
|
|
|
|
|
|
def test_error_on_wrong_obj_type_for_body_sensor():
|
|
"""Verify ValueError is raised when wrong obj type is used for body sensor."""
|
|
with pytest.raises(ValueError, match="requires obj.type='body'"):
|
|
BuiltinSensorCfg(
|
|
name="invalid_sensor",
|
|
sensor_type="subtreecom",
|
|
obj=ObjRef(type="site", name="base"),
|
|
)
|
|
|
|
|
|
def test_error_on_wrong_obj_type_for_joint_sensor():
|
|
"""Verify ValueError is raised when wrong obj type is used for joint sensor."""
|
|
with pytest.raises(ValueError, match="requires obj.type='joint'"):
|
|
BuiltinSensorCfg(
|
|
name="invalid_sensor",
|
|
sensor_type="jointvel",
|
|
obj=ObjRef(type="body", name="base"),
|
|
)
|
|
|
|
|
|
def test_spatial_frame_sensor_accepts_multiple_types():
|
|
"""Verify spatial frame sensors accept body, xbody, geom, site, camera."""
|
|
BuiltinSensorCfg(
|
|
name="frame_sensor_body",
|
|
sensor_type="framepos",
|
|
obj=ObjRef(type="body", name="test"),
|
|
)
|
|
BuiltinSensorCfg(
|
|
name="frame_sensor_xbody",
|
|
sensor_type="framepos",
|
|
obj=ObjRef(type="xbody", name="test"),
|
|
)
|
|
BuiltinSensorCfg(
|
|
name="frame_sensor_geom",
|
|
sensor_type="framepos",
|
|
obj=ObjRef(type="geom", name="test"),
|
|
)
|
|
BuiltinSensorCfg(
|
|
name="frame_sensor_site",
|
|
sensor_type="framepos",
|
|
obj=ObjRef(type="site", name="test"),
|
|
)
|
|
BuiltinSensorCfg(
|
|
name="frame_sensor_camera",
|
|
sensor_type="framepos",
|
|
obj=ObjRef(type="camera", name="test"),
|
|
)
|
|
|
|
|
|
def test_xml_sensors_auto_discovered(robot_with_xml_sensors, device):
|
|
"""Verify sensors defined in entity XML are automatically discovered and exposed."""
|
|
entity_cfg = EntityCfg(
|
|
spec_fn=lambda: mujoco.MjSpec.from_string(robot_with_xml_sensors)
|
|
)
|
|
|
|
scene_cfg = SceneCfg(
|
|
num_envs=2,
|
|
env_spacing=3.0,
|
|
entities={"robot": entity_cfg},
|
|
sensors=(),
|
|
)
|
|
|
|
scene = Scene(scene_cfg, device)
|
|
model = scene.compile()
|
|
sim_cfg = SimulationCfg(njmax=20)
|
|
sim = Simulation(num_envs=2, cfg=sim_cfg, model=model, device=device)
|
|
scene.initialize(sim.mj_model, sim.model, sim.data)
|
|
|
|
joint_sensor = scene["robot/xml_joint_sensor"]
|
|
accel_sensor = scene["robot/xml_accel_sensor"]
|
|
gyro_sensor = scene["robot/xml_gyro_sensor"]
|
|
|
|
sim.step()
|
|
|
|
joint_data = joint_sensor.data
|
|
accel_data = accel_sensor.data
|
|
gyro_data = gyro_sensor.data
|
|
|
|
assert isinstance(joint_data, torch.Tensor)
|
|
assert joint_data.shape == (2, 1)
|
|
assert isinstance(accel_data, torch.Tensor)
|
|
assert accel_data.shape == (2, 3)
|
|
assert isinstance(gyro_data, torch.Tensor)
|
|
assert gyro_data.shape == (2, 3)
|
|
|
|
|
|
def test_builtin_sensor_errors_on_duplicate_name(robot_with_xml_sensors, device):
|
|
"""Verify BuiltinSensorCfg throws error when name conflicts with XML sensor."""
|
|
entity_cfg = EntityCfg(
|
|
spec_fn=lambda: mujoco.MjSpec.from_string(robot_with_xml_sensors)
|
|
)
|
|
|
|
duplicate_sensor_cfg = BuiltinSensorCfg(
|
|
name="xml_joint_sensor",
|
|
sensor_type="jointpos",
|
|
obj=ObjRef(type="joint", name="joint1", entity="robot"),
|
|
)
|
|
|
|
scene_cfg = SceneCfg(
|
|
num_envs=2,
|
|
env_spacing=3.0,
|
|
entities={"robot": entity_cfg},
|
|
sensors=(duplicate_sensor_cfg,),
|
|
)
|
|
|
|
with pytest.raises(ValueError, match="defined in both entity XML and scene config"):
|
|
Scene(scene_cfg, device)
|
|
|
|
|
|
def test_cutoff_parameter(articulated_robot_xml, device):
|
|
"""Verify cutoff parameter propagates to MuJoCo model."""
|
|
entity_cfg = EntityCfg(
|
|
spec_fn=lambda: mujoco.MjSpec.from_string(articulated_robot_xml)
|
|
)
|
|
|
|
sensor_cfg = BuiltinSensorCfg(
|
|
name="joint1_pos",
|
|
sensor_type="jointpos",
|
|
obj=ObjRef(type="joint", name="joint1", entity="robot"),
|
|
cutoff=0.01,
|
|
)
|
|
|
|
scene_cfg = SceneCfg(
|
|
num_envs=1,
|
|
env_spacing=3.0,
|
|
entities={"robot": entity_cfg},
|
|
sensors=(sensor_cfg,),
|
|
)
|
|
|
|
scene = Scene(scene_cfg, device)
|
|
model = scene.compile()
|
|
sim_cfg = SimulationCfg(njmax=20)
|
|
sim = Simulation(num_envs=1, cfg=sim_cfg, model=model, device=device)
|
|
|
|
sensor = sim.mj_model.sensor("robot/joint1_pos")
|
|
assert sensor.cutoff[0] == 0.01
|
|
|
|
|
|
def test_deepcopy_no_double_prefix():
|
|
"""Regression test for #850: deepcopy should not double-prefix sensor name."""
|
|
cfg = BuiltinSensorCfg(
|
|
name="joint1_pos",
|
|
sensor_type="jointpos",
|
|
obj=ObjRef(type="joint", name="joint1", entity="robot"),
|
|
)
|
|
assert cfg.prefixed_name == "robot/joint1_pos"
|
|
|
|
cfg_copy = copy.deepcopy(cfg)
|
|
assert cfg_copy.prefixed_name == "robot/joint1_pos"
|
|
assert cfg_copy.name == "joint1_pos"
|