"""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 """ """ @pytest.fixture(scope="module") def robot_with_xml_sensors(): """XML for robot with sensors already defined in the XML.""" return """ """ 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"