Import upstream snapshot c19f713c415a699a79d71cd96aa13c3104a05047

Upstream: https://github.com/michaelgillett/mjlab
Upstream-Commit: c19f713c415a699a79d71cd96aa13c3104a05047
Upstream-Branch: main

tests/test_builtin_sensor.py

@@ -0,0 +1,378 @@
+"""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"