Import upstream snapshot c19f713c415a699a79d71cd96aa13c3104a05047

Upstream: https://github.com/michaelgillett/mjlab Upstream-Commit: c19f713c415a699a79d71cd96aa13c3104a05047 Upstream-Branch: main
2026-08-28 15:42:17 +08:00
commit 32a241c28f
513 changed files with 120771 additions and 0 deletions
@@ -0,0 +1,209 @@
+"""Shared test fixtures and utilities."""
+
+import os
+from pathlib import Path
+
+import mujoco
+import pytest
+import torch
+import warp as wp
+
+from mjlab.entity import Entity, EntityArticulationInfoCfg, EntityCfg
+from mjlab.scene import Scene, SceneCfg
+from mjlab.sim.sim import Simulation, SimulationCfg
+
+
+@pytest.fixture(scope="session", autouse=True)
+def configure_test_environment():
+  """Configure test environment settings automatically for all tests."""
+  wp.config.quiet = True
+
+
+def get_test_device() -> str:
+  """Get device for testing, preferring CUDA if available.
+
+  Can be overridden with FORCE_CPU=1 environment variable to test
+  CPU-only behavior on GPU machines.
+  """
+  if os.environ.get("FORCE_CPU") == "1":
+    return "cpu"
+  return "cuda" if torch.cuda.is_available() else "cpu"
+
+
+@pytest.fixture
+def fixtures_dir() -> Path:
+  """Path to test fixtures directory."""
+  return Path(__file__).parent / "fixtures"
+
+
+def load_fixture_xml(fixture_name: str) -> str:
+  """Load XML content from fixture file.
+
+  Args:
+    fixture_name: Name of the fixture file (without .xml extension) or full path.
+
+  Returns:
+    XML content as string.
+  """
+  fixtures_path = Path(__file__).parent / "fixtures"
+  if not fixture_name.endswith(".xml"):
+    fixture_name = f"{fixture_name}.xml"
+  fixture_file = fixtures_path / fixture_name
+  return fixture_file.read_text()
+
+
+def create_entity_with_actuator(xml_string: str, actuator_cfg):
+  """Create entity with actuator from XML string.
+
+  Args:
+    xml_string: MuJoCo XML model string.
+    actuator_cfg: Actuator configuration.
+
+  Returns:
+    Entity instance.
+  """
+  cfg = EntityCfg(
+    spec_fn=lambda: mujoco.MjSpec.from_string(xml_string),
+    articulation=EntityArticulationInfoCfg(actuators=(actuator_cfg,)),
+  )
+  return Entity(cfg)
+
+
+def create_entity_from_fixture(fixture_name: str, actuator_cfg=None):
+  """Create entity from fixture file.
+
+  Args:
+    fixture_name: Name of the fixture file (without .xml extension).
+    actuator_cfg: Optional actuator configuration.
+
+  Returns:
+    Entity instance.
+  """
+  xml_string = load_fixture_xml(fixture_name)
+  if actuator_cfg is not None:
+    return create_entity_with_actuator(xml_string, actuator_cfg)
+  cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(xml_string))
+  return Entity(cfg)
+
+
+def initialize_entity(entity: Entity, device: str, num_envs: int = 1):
+  """Initialize entity with simulation.
+
+  Args:
+    entity: Entity to initialize.
+    device: Device to use ("cpu" or "cuda").
+    num_envs: Number of environments.
+
+  Returns:
+    Tuple of (entity, simulation).
+  """
+  model = entity.compile()
+  sim_cfg = SimulationCfg()
+  sim = Simulation(num_envs=num_envs, cfg=sim_cfg, model=model, device=device)
+  entity.initialize(model, sim.model, sim.data, device)
+  return entity, sim
+
+
+def make_scene_and_sim(
+  device: str,
+  xml: str | dict[str, str],
+  sensors: tuple,
+  num_envs: int = 1,
+  sim_cfg: SimulationCfg | None = None,
+) -> tuple[Scene, Simulation]:
+  """Create a scene and simulation from inline XML with sensors wired up.
+
+  ``xml`` may be a single XML string (registered as the ``robot`` entity) or a
+  mapping of entity name to XML string for multi-entity scenes.
+  """
+  xml_by_entity = {"robot": xml} if isinstance(xml, str) else xml
+  entities = {
+    name: EntityCfg(spec_fn=lambda s=s: mujoco.MjSpec.from_string(s))
+    for name, s in xml_by_entity.items()
+  }
+  scene_cfg = SceneCfg(
+    num_envs=num_envs,
+    env_spacing=5.0,
+    entities=entities,
+    sensors=sensors,
+  )
+  scene = Scene(scene_cfg, device)
+  model = scene.compile()
+  if sim_cfg is None:
+    sim_cfg = SimulationCfg(njmax=20)
+  sim = Simulation(num_envs=num_envs, cfg=sim_cfg, model=model, device=device)
+  scene.initialize(sim.mj_model, sim.model, sim.data)
+  if scene.sensor_context is not None:
+    sim.set_sensor_context(scene.sensor_context)
+  return scene, sim
+
+
+# =============================================================================
+# XML Fixture Loaders
+# =============================================================================
+
+
+@pytest.fixture
+def fixed_base_box_xml() -> str:
+  """Load fixed base box XML fixture."""
+  return load_fixture_xml("fixed_base_box")
+
+
+@pytest.fixture
+def floating_base_box_xml() -> str:
+  """Load floating base box XML fixture."""
+  return load_fixture_xml("floating_base_box")
+
+
+@pytest.fixture
+def fixed_base_articulated_xml() -> str:
+  """Load fixed base articulated robot XML fixture."""
+  return load_fixture_xml("fixed_base_articulated")
+
+
+@pytest.fixture
+def floating_base_articulated_xml() -> str:
+  """Load floating base articulated robot XML fixture."""
+  return load_fixture_xml("floating_base_articulated")
+
+
+@pytest.fixture
+def biped_xml() -> str:
+  """Load biped robot XML fixture with ground plane."""
+  return load_fixture_xml("biped")
+
+
+@pytest.fixture
+def robot_with_floor_xml() -> str:
+  """XML for a floating body above a ground plane."""
+  return """
+    <mujoco>
+      <worldbody>
+        <geom name="floor" type="plane" size="10 10 0.1" pos="0 0 0"/>
+        <body name="base" pos="0 0 2">
+          <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.1"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  """
+
+
+@pytest.fixture
+def falling_box_xml() -> str:
+  """XML for a box that can fall onto a ground plane."""
+  return """
+    <mujoco>
+      <worldbody>
+        <body name="ground" pos="0 0 0">
+          <geom name="ground_geom" type="plane" size="5 5 0.1" rgba="0.5 0.5 0.5 1"/>
+        </body>
+        <body name="box" pos="0 0 0.5">
+          <freejoint name="box_joint"/>
+          <geom name="box_geom" type="box" size="0.1 0.1 0.1" rgba="0.8 0.3 0.3 1"
+            mass="1.0"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  """
@@ -0,0 +1,19 @@
+<mujoco>
+  <worldbody>
+    <body name="ground" pos="0 0 0">
+      <geom name="ground_geom" type="plane" size="5 5 0.1" rgba="0.5 0.5 0.5 1"/>
+    </body>
+    <body name="base" pos="0 0 0.5">
+      <freejoint name="base_joint"/>
+      <geom name="torso_geom" type="box" size="0.15 0.1 0.2" mass="5.0"/>
+      <body name="left_foot" pos="0.1 0 -0.25">
+        <joint name="left_ankle" type="hinge" axis="0 1 0" range="-0.5 0.5"/>
+        <geom name="left_foot_geom" type="box" size="0.05 0.08 0.02" mass="0.2"/>
+      </body>
+      <body name="right_foot" pos="-0.1 0 -0.25">
+        <joint name="right_ankle" type="hinge" axis="0 1 0" range="-0.5 0.5"/>
+        <geom name="right_foot_geom" type="box" size="0.05 0.08 0.02" mass="0.2"/>
+      </body>
+    </body>
+  </worldbody>
+</mujoco>
@@ -0,0 +1,15 @@
+<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>
@@ -0,0 +1,7 @@
+<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>
@@ -0,0 +1,20 @@
+<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="1.0"/>
+      <body name="link1" pos="0 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="0.1"/>
+        <site name="site1" pos="0 0 0"/>
+      </body>
+      <body name="link2" pos="0 0 0">
+        <joint name="joint2" type="hinge" axis="0 0 1" range="-1.57 1.57"/>
+        <geom name="link2_geom" type="box" size="0.1 0.1 0.1" mass="0.1"/>
+      </body>
+    </body>
+  </worldbody>
+  <sensor>
+    <jointpos name="joint1_pos" joint="joint1"/>
+  </sensor>
+</mujoco>
@@ -0,0 +1,8 @@
+<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>
@@ -0,0 +1,43 @@
+<mujoco model="quadcopter">
+  <compiler angle="radian" autolimits="true"/>
+
+  <option timestep="0.002" integrator="implicitfast" gravity="0 0 -9.81"/>
+
+  <default>
+    <geom rgba="0.5 0.5 0.5 1"/>
+    <site size="0.01" rgba="1 0 0 1"/>
+  </default>
+
+  <worldbody>
+    <light pos="0 0 3" dir="0 0 -1" diffuse="0.8 0.8 0.8"/>
+    <geom type="plane" size="5 5 0.1" rgba="0.9 0.9 0.9 1"/>
+
+    <body name="base" pos="0 0 0.5">
+      <freejoint name="root"/>
+      <inertial pos="0 0 0" mass="1.0" diaginertia="0.01 0.01 0.02"/>
+
+      <!-- Central body -->
+      <geom type="box" size="0.1 0.1 0.02" rgba="0.3 0.3 0.8 1"/>
+
+      <!-- Rotor sites at corners -->
+      <site name="rotor_fl" pos="0.1 0.1 0"/>
+      <site name="rotor_fr" pos="0.1 -0.1 0"/>
+      <site name="rotor_rl" pos="-0.1 0.1 0"/>
+      <site name="rotor_rr" pos="-0.1 -0.1 0"/>
+
+      <!-- Visual rotor geoms (non-collision) -->
+      <geom name="rotor_fl_vis" type="cylinder" size="0.05 0.005" pos="0.1 0.1 0" rgba="1 0 0 0.5" contype="0" conaffinity="0"/>
+      <geom name="rotor_fr_vis" type="cylinder" size="0.05 0.005" pos="0.1 -0.1 0" rgba="1 0 0 0.5" contype="0" conaffinity="0"/>
+      <geom name="rotor_rl_vis" type="cylinder" size="0.05 0.005" pos="-0.1 0.1 0" rgba="1 0 0 0.5" contype="0" conaffinity="0"/>
+      <geom name="rotor_rr_vis" type="cylinder" size="0.05 0.005" pos="-0.1 -0.1 0" rgba="1 0 0 0.5" contype="0" conaffinity="0"/>
+    </body>
+  </worldbody>
+
+  <actuator>
+    <!-- Site actuators applying vertical thrust forces -->
+    <motor name="thrust_fl" site="rotor_fl" gear="0 0 1 0 0 0" ctrlrange="0 20"/>
+    <motor name="thrust_fr" site="rotor_fr" gear="0 0 1 0 0 0" ctrlrange="0 20"/>
+    <motor name="thrust_rl" site="rotor_rl" gear="0 0 1 0 0 0" ctrlrange="0 20"/>
+    <motor name="thrust_rr" site="rotor_rr" gear="0 0 1 0 0 0" ctrlrange="0 20"/>
+  </actuator>
+</mujoco>
@@ -0,0 +1,43 @@
+<mujoco model="tendon_finger">
+  <compiler angle="radian" autolimits="true"/>
+
+  <default>
+    <geom type="capsule" size="0.01" rgba="0.8 0.6 0.4 1"/>
+    <joint type="hinge" axis="0 1 0" range="-1.57 1.57" damping="0.1"/>
+    <site size="0.005" rgba="1 0 0 1"/>
+  </default>
+
+  <worldbody>
+    <light pos="0 0 1" dir="0 0 -1" diffuse="0.8 0.8 0.8"/>
+    <geom type="plane" size="1 1 0.1" rgba="0.9 0.9 0.9 1"/>
+
+    <body name="proximal_link" pos="0 0 0.5">
+      <joint name="proximal_joint"/>
+      <geom fromto="0 0 0 0 0 0.15" size="0.015"/>
+      <site name="tendon_p1" pos="0 0.02 0.05"/>
+      <site name="tendon_p2" pos="0 0.02 0.15"/>
+
+      <body name="distal_link" pos="0 0 0.15">
+        <joint name="distal_joint"/>
+        <geom fromto="0 0 0 0 0 0.15" size="0.012"/>
+        <site name="tendon_p3" pos="0 0.02 0.05"/>
+        <site name="fingertip" pos="0 0 0.15"/>
+      </body>
+    </body>
+  </worldbody>
+
+  <tendon>
+    <spatial name="finger_tendon" width="0.003" rgba="0.2 0.2 0.8 1">
+      <site site="tendon_p1"/>
+      <site site="tendon_p2"/>
+      <site site="tendon_p3"/>
+    </spatial>
+  </tendon>
+
+  <actuator>
+    <!-- Tendon actuators for testing different control modes -->
+    <position name="tendon_position" tendon="finger_tendon" kp="100" ctrlrange="-0.1 0.1"/>
+    <velocity name="tendon_velocity" tendon="finger_tendon" kv="10" ctrlrange="-1 1"/>
+    <motor name="tendon_effort" tendon="finger_tendon" ctrlrange="-10 10"/>
+  </actuator>
+</mujoco>
@@ -0,0 +1,36 @@
+"""Smoke test for mjlab package."""
+
+import io
+import sys
+import warnings
+from contextlib import redirect_stderr, redirect_stdout
+
+try:
+  import pytest
+except ModuleNotFoundError:
+  pytest = None  # type: ignore[assignment]
+
+
+@pytest.mark.slow if pytest else lambda f: f
+def test_basic_functionality() -> None:
+  """Test that mjlab can create and close an environment."""
+  from mjlab.envs.manager_based_rl_env import ManagerBasedRlEnv
+  from mjlab.tasks.velocity.config.go1.env_cfgs import unitree_go1_flat_env_cfg
+
+  # Suppress env spam.
+  with warnings.catch_warnings():
+    warnings.simplefilter("ignore")
+    with redirect_stdout(io.StringIO()), redirect_stderr(io.StringIO()):
+      env = ManagerBasedRlEnv(unitree_go1_flat_env_cfg(), device="cpu")
+      assert env.sim.data.time == 0.0
+      env.close()
+
+
+if __name__ == "__main__":
+  try:
+    test_basic_functionality()
+    print("✓ Smoke test passed!")
+    sys.exit(0)
+  except Exception as e:
+    print(f"✗ Smoke test failed: {e}")
+    sys.exit(1)
@@ -0,0 +1,118 @@
+"""Tests for action manager functionality."""
+
+from unittest.mock import Mock
+
+import pytest
+import torch
+from conftest import get_test_device
+
+from mjlab.managers.action_manager import ActionManager
+
+
+@pytest.fixture
+def device():
+  """Test device fixture."""
+  return get_test_device()
+
+
+def _make_mock_action_term(action_dim: int):
+  """Create a mock action term factory."""
+
+  def factory(env):
+    term = Mock()
+    term.action_dim = action_dim
+    term.raw_action = torch.zeros(env.num_envs, action_dim, device=env.device)
+    term.process_actions = Mock()
+    term.apply_actions = Mock()
+    term.reset = Mock()
+    return term
+
+  return factory
+
+
+@pytest.fixture
+def mock_env(device):
+  """Create a mock environment for testing."""
+  env = Mock()
+  env.num_envs = 4
+  env.device = device
+  return env
+
+
+@pytest.fixture
+def action_term_cfg():
+  """Create a simple action term config."""
+  cfg = Mock()
+  cfg.build = _make_mock_action_term(action_dim=3)
+  cfg.entity_name = "robot"
+  return cfg
+
+
+def test_action_history_tracking(mock_env, action_term_cfg, device):
+  """Test that action, prev_action, and prev_prev_action track history correctly."""
+  manager = ActionManager({"action": action_term_cfg}, mock_env)
+
+  # Initial state: all zeros.
+  assert torch.all(manager.action == 0.0)
+  assert torch.all(manager.prev_action == 0.0)
+  assert torch.all(manager.prev_prev_action == 0.0)
+
+  # Process actions and verify history shifts correctly.
+  actions = [
+    torch.tensor([[float(i)] * 3] * mock_env.num_envs, device=device)
+    for i in range(1, 5)
+  ]
+
+  manager.process_action(actions[0])
+  assert torch.allclose(manager.action, actions[0])
+  assert torch.all(manager.prev_action == 0.0)
+  assert torch.all(manager.prev_prev_action == 0.0)
+
+  manager.process_action(actions[1])
+  assert torch.allclose(manager.action, actions[1])
+  assert torch.allclose(manager.prev_action, actions[0])
+  assert torch.all(manager.prev_prev_action == 0.0)
+
+  manager.process_action(actions[2])
+  assert torch.allclose(manager.action, actions[2])
+  assert torch.allclose(manager.prev_action, actions[1])
+  assert torch.allclose(manager.prev_prev_action, actions[0])
+
+  manager.process_action(actions[3])
+  assert torch.allclose(manager.action, actions[3])
+  assert torch.allclose(manager.prev_action, actions[2])
+  assert torch.allclose(manager.prev_prev_action, actions[1])
+
+
+def test_action_history_reset(mock_env, action_term_cfg, device):
+  """Test that reset clears action history for all or specific environments."""
+  manager = ActionManager({"action": action_term_cfg}, mock_env)
+
+  # Populate history.
+  actions = [
+    torch.tensor([[float(i)] * 3] * mock_env.num_envs, device=device)
+    for i in range(1, 4)
+  ]
+  for a in actions:
+    manager.process_action(a)
+
+  # Partial reset: only envs 0 and 2.
+  manager.reset(env_ids=torch.tensor([0, 2]))
+
+  # Reset envs should be zeros.
+  for env_id in [0, 2]:
+    assert torch.all(manager.action[env_id] == 0.0)
+    assert torch.all(manager.prev_action[env_id] == 0.0)
+    assert torch.all(manager.prev_prev_action[env_id] == 0.0)
+
+  # Non-reset envs should retain history.
+  for env_id in [1, 3]:
+    assert torch.allclose(manager.action[env_id], actions[2][env_id])
+    assert torch.allclose(manager.prev_action[env_id], actions[1][env_id])
+    assert torch.allclose(manager.prev_prev_action[env_id], actions[0][env_id])
+
+  # Full reset.
+  manager.reset()
+  assert torch.all(manager.action == 0.0)
+  assert torch.all(manager.prev_action == 0.0)
+  assert torch.all(manager.prev_prev_action == 0.0)
@@ -0,0 +1,318 @@
+"""Tests for actions."""
+
+from pathlib import Path
+from unittest.mock import Mock
+
+import mujoco
+import pytest
+import torch
+from conftest import get_test_device, load_fixture_xml
+
+from mjlab.actuator.actuator import TransmissionType
+from mjlab.actuator.builtin_actuator import BuiltinMotorActuatorCfg
+from mjlab.entity import Entity, EntityArticulationInfoCfg, EntityCfg
+from mjlab.envs import ManagerBasedRlEnv
+from mjlab.envs.mdp.actions import (
+  JointPositionActionCfg,
+  RelativeJointPositionActionCfg,
+  SiteEffortActionCfg,
+  TendonEffortActionCfg,
+  TendonLengthActionCfg,
+  TendonVelocityActionCfg,
+)
+from mjlab.sim.sim import Simulation, SimulationCfg
+
+
+@pytest.fixture(scope="module")
+def device():
+  return get_test_device()
+
+
+@pytest.fixture(scope="module")
+def fixtures_dir():
+  return Path(__file__).parent / "fixtures"
+
+
+@pytest.fixture(scope="module")
+def tendon_finger_entity(fixtures_dir, device):
+  return make_entity(
+    fixtures_dir / "tendon_finger.xml",
+    ("finger_tendon",),
+    TransmissionType.TENDON,
+    device,
+    from_file=True,
+  )
+
+
+@pytest.fixture(scope="module")
+def fixed_base_entity(fixtures_dir, device):
+  return make_entity(
+    fixtures_dir / "fixed_base_articulated.xml",
+    ("joint.*",),
+    TransmissionType.JOINT,
+    device,
+    from_file=True,
+  )
+
+
+@pytest.fixture(scope="module")
+def floating_base_entity(device):
+  return make_entity(
+    load_fixture_xml("floating_base_articulated"),
+    ("joint.*",),
+    TransmissionType.JOINT,
+    device,
+    from_file=False,
+  )
+
+
+def make_entity(xml_or_path, target_expr, transmission_type, device, from_file=False):
+  """Create and initialize entity."""
+
+  def spec_fn():
+    if from_file:
+      return mujoco.MjSpec.from_file(str(xml_or_path))
+    return mujoco.MjSpec.from_string(xml_or_path)
+
+  cfg = EntityCfg(
+    spec_fn=spec_fn,
+    articulation=EntityArticulationInfoCfg(
+      actuators=(
+        BuiltinMotorActuatorCfg(
+          target_names_expr=target_expr,
+          transmission_type=transmission_type,
+          effort_limit=10.0,
+        ),
+      )
+    ),
+  )
+  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)
+  return entity
+
+
+def make_env(entity, name, device):
+  """Create mock environment."""
+  env = Mock(spec=ManagerBasedRlEnv)
+  env.num_envs = 4
+  env.device = device
+  env.scene = {name: entity}
+  return env
+
+
+def test_base_action_applies_scale_and_offset(tendon_finger_entity, device):
+  """BaseAction: processed = raw * scale + offset."""
+  entity = tendon_finger_entity
+  env = make_env(entity, "finger", device)
+
+  cfg = TendonLengthActionCfg(
+    entity_name="finger",
+    actuator_names=("finger_tendon",),
+    scale=2.0,
+    offset=0.5,
+  )
+  action = cfg.build(env)
+
+  raw = torch.tensor([[1.0], [2.0], [3.0], [4.0]], device=device)
+  action.process_actions(raw)
+
+  assert torch.allclose(action._processed_actions, raw * 2.0 + 0.5)
+
+
+def test_base_action_reset_zeros_specific_envs(tendon_finger_entity, device):
+  """BaseAction.reset() zeros raw_action for specified env_ids only."""
+  entity = tendon_finger_entity
+  env = make_env(entity, "finger", device)
+
+  cfg = TendonLengthActionCfg(entity_name="finger", actuator_names=("finger_tendon",))
+  action = cfg.build(env)
+
+  action.process_actions(torch.ones(4, 1, device=device))
+  action.reset(env_ids=torch.tensor([0, 2], device=device))
+
+  assert torch.all(action.raw_action[[0, 2]] == 0.0)
+  assert torch.all(action.raw_action[[1, 3]] == 1.0)
+
+
+@pytest.mark.parametrize(
+  "cfg_cls,target_attr,fixture,target_expr,transmission,entity_name",
+  [
+    # Joints.
+    (
+      JointPositionActionCfg,
+      "joint_pos_target",
+      "floating_base_articulated",
+      ("joint.*",),
+      TransmissionType.JOINT,
+      "robot",
+    ),
+    # Tendons.
+    (
+      TendonLengthActionCfg,
+      "tendon_len_target",
+      "tendon_finger.xml",
+      ("finger_tendon",),
+      TransmissionType.TENDON,
+      "finger",
+    ),
+    (
+      TendonVelocityActionCfg,
+      "tendon_vel_target",
+      "tendon_finger.xml",
+      ("finger_tendon",),
+      TransmissionType.TENDON,
+      "finger",
+    ),
+    (
+      TendonEffortActionCfg,
+      "tendon_effort_target",
+      "tendon_finger.xml",
+      ("finger_tendon",),
+      TransmissionType.TENDON,
+      "finger",
+    ),
+    # Sites.
+    (
+      SiteEffortActionCfg,
+      "site_effort_target",
+      "quadcopter.xml",
+      ("rotor_.*",),
+      TransmissionType.SITE,
+      "drone",
+    ),
+  ],
+)
+def test_action_sets_entity_target(
+  fixtures_dir,
+  device,
+  cfg_cls,
+  target_attr,
+  fixture,
+  target_expr,
+  transmission,
+  entity_name,
+):
+  """Each action type writes to correct entity.data field."""
+  if fixture.endswith(".xml"):
+    entity = make_entity(
+      fixtures_dir / fixture, target_expr, transmission, device, from_file=True
+    )
+  else:
+    entity = make_entity(
+      load_fixture_xml(fixture), target_expr, transmission, device, from_file=False
+    )
+
+  env = make_env(entity, entity_name, device)
+  cfg = cfg_cls(entity_name=entity_name, actuator_names=target_expr)
+  action = cfg.build(env)
+
+  target = torch.arange(4 * action.action_dim, device=device, dtype=torch.float32)
+  target = target.reshape(4, action.action_dim) * 0.1
+
+  action.process_actions(target)
+  action.apply_actions()
+
+  entity_target = getattr(entity.data, target_attr)
+  assert torch.allclose(entity_target, target)
+
+
+def test_base_action_clip(fixed_base_entity, device):
+  """BaseAction: clip clamps only matched actuators; others stay unclipped."""
+  entity = fixed_base_entity
+  env = make_env(entity, "robot", device)
+
+  # Clip only joint1; joint2 should remain unclipped.
+  cfg = JointPositionActionCfg(
+    entity_name="robot",
+    actuator_names=("joint.*",),
+    scale=1.0,
+    use_default_offset=False,
+    clip={"joint1": (-0.5, 0.5)},
+  )
+  action = cfg.build(env)
+
+  # joint1=2.0 should be clipped to 0.5, joint2=2.0 should pass through.
+  raw = torch.tensor([[2.0, 2.0]], device=device).expand(4, -1)
+  action.process_actions(raw)
+
+  assert torch.allclose(
+    action._processed_actions[:, 0], torch.tensor(0.5, device=device)
+  )
+  assert torch.allclose(
+    action._processed_actions[:, 1], torch.tensor(2.0, device=device)
+  )
+
+
+def test_relative_joint_position_zero_action(floating_base_entity, device):
+  """With zero action, targets equal current_pos."""
+  entity = floating_base_entity
+  env = make_env(entity, "robot", device)
+
+  cfg = RelativeJointPositionActionCfg(
+    entity_name="robot", actuator_names=("joint.*",), scale=1.0
+  )
+  action = cfg.build(env)
+
+  current_pos = entity.data.joint_pos[:, action.target_ids].clone()
+
+  action.process_actions(torch.zeros(4, action.action_dim, device=device))
+  action.apply_actions()
+
+  assert torch.allclose(entity.data.joint_pos_target[:, action.target_ids], current_pos)
+
+
+def test_relative_joint_position_nonzero_action(floating_base_entity, device):
+  """With nonzero action, targets shift by action * scale from current."""
+  entity = floating_base_entity
+  env = make_env(entity, "robot", device)
+
+  scale = 0.1
+  cfg = RelativeJointPositionActionCfg(
+    entity_name="robot", actuator_names=("joint.*",), scale=scale
+  )
+  action = cfg.build(env)
+
+  current_pos = entity.data.joint_pos[:, action.target_ids].clone()
+
+  raw = torch.ones(4, action.action_dim, device=device)
+  action.process_actions(raw)
+  action.apply_actions()
+
+  expected = current_pos + raw * scale
+  assert torch.allclose(entity.data.joint_pos_target[:, action.target_ids], expected)
+
+
+def test_relative_joint_position_ignores_encoder_bias(floating_base_entity, device):
+  """Encoder bias must not affect the target: target = current_pos + delta."""
+  entity = floating_base_entity
+  env = make_env(entity, "robot", device)
+
+  cfg = RelativeJointPositionActionCfg(
+    entity_name="robot", actuator_names=("joint.*",), scale=1.0
+  )
+  action = cfg.build(env)
+
+  entity.data.encoder_bias[:, action.target_ids] = 0.05
+
+  current_pos = entity.data.joint_pos[:, action.target_ids].clone()
+  delta = 0.1
+  raw = torch.full((4, action.action_dim), delta, device=device)
+  action.process_actions(raw)
+  action.apply_actions()
+
+  assert torch.allclose(
+    entity.data.joint_pos_target[:, action.target_ids], current_pos + delta
+  )
+
+  # Reset bias so this shared fixture doesn't affect other tests.
+  entity.data.encoder_bias[:, action.target_ids] = 0.0
+
+
+def test_relative_joint_position_offset_raises(floating_base_entity, device):
+  """Setting offset on RelativeJointPositionActionCfg raises ValueError."""
+  with pytest.raises(ValueError, match="offset"):
+    RelativeJointPositionActionCfg(
+      entity_name="robot", actuator_names=("joint.*",), offset=0.5
+    )
@@ -0,0 +1,327 @@
+"""Tests for actuator module."""
+
+import mujoco
+import pytest
+import torch
+from conftest import (
+  create_entity_with_actuator,
+  get_test_device,
+  initialize_entity,
+  load_fixture_xml,
+)
+
+from mjlab.actuator import (
+  BuiltinMotorActuatorCfg,
+  BuiltinPositionActuatorCfg,
+  BuiltinVelocityActuatorCfg,
+  IdealPdActuatorCfg,
+  XmlActuatorCfg,
+)
+from mjlab.entity import Entity, EntityArticulationInfoCfg, EntityCfg
+
+
+@pytest.fixture(scope="module")
+def device():
+  return get_test_device()
+
+
+@pytest.fixture(scope="module")
+def robot_xml():
+  return load_fixture_xml("floating_base_articulated")
+
+
+def test_builtin_pd_actuator_compute(device, robot_xml):
+  """BuiltinPositionActuator writes position targets to ctrl."""
+  actuator_cfg = BuiltinPositionActuatorCfg(
+    target_names_expr=("joint.*",), stiffness=50.0, damping=5.0
+  )
+  entity = create_entity_with_actuator(robot_xml, actuator_cfg)
+  entity, sim = initialize_entity(entity, device)
+
+  entity.set_joint_position_target(torch.tensor([[0.5, -0.3]], device=device))
+  entity.write_data_to_sim()
+
+  ctrl = sim.data.ctrl[0]
+  assert torch.allclose(ctrl, torch.tensor([0.5, -0.3], device=device))
+
+
+def test_ideal_pd_actuator_compute(device, robot_xml):
+  """IdealPdActuator computes torques via explicit PD control."""
+  actuator_cfg = IdealPdActuatorCfg(
+    target_names_expr=("joint.*",), effort_limit=100.0, stiffness=50.0, damping=5.0
+  )
+  entity = create_entity_with_actuator(robot_xml, actuator_cfg)
+  entity, sim = initialize_entity(entity, device)
+
+  entity.write_joint_state_to_sim(
+    position=torch.tensor([[0.0, 0.0]], device=device),
+    velocity=torch.tensor([[0.0, 0.0]], device=device),
+  )
+
+  entity.set_joint_position_target(torch.tensor([[0.1, -0.1]], device=device))
+  entity.set_joint_velocity_target(torch.tensor([[0.0, 0.0]], device=device))
+  entity.set_joint_effort_target(torch.tensor([[0.0, 0.0]], device=device))
+  entity.write_data_to_sim()
+
+  ctrl = sim.data.ctrl[0]
+  assert torch.allclose(ctrl, torch.tensor([5.0, -5.0], device=device))
+
+
+def test_targets_cleared_on_reset(device, robot_xml):
+  """Entity.reset() zeros all targets."""
+  actuator_cfg = BuiltinPositionActuatorCfg(
+    target_names_expr=("joint.*",), stiffness=50.0, damping=5.0
+  )
+  entity = create_entity_with_actuator(robot_xml, actuator_cfg)
+  entity, sim = initialize_entity(entity, device)
+
+  entity.set_joint_position_target(torch.tensor([[0.5, -0.3]], device=device))
+  entity.write_data_to_sim()
+
+  assert not torch.allclose(
+    entity.data.joint_pos_target, torch.zeros(1, 2, device=device)
+  )
+
+  entity.reset()
+
+  assert torch.allclose(entity.data.joint_pos_target, torch.zeros(1, 2, device=device))
+  assert torch.allclose(entity.data.joint_vel_target, torch.zeros(1, 2, device=device))
+  assert torch.allclose(
+    entity.data.joint_effort_target, torch.zeros(1, 2, device=device)
+  )
+
+
+# ---------------------------------------------------------------------------
+# Internal attach prefix tests (issue #709)
+# ---------------------------------------------------------------------------
+
+
+def _make_arm_spec() -> mujoco.MjSpec:
+  """Helper: single-joint arm with a motor actuator."""
+  spec = mujoco.MjSpec()
+  body = spec.worldbody.add_body(name="link")
+  body.add_joint(name="elbow", type=mujoco.mjtJoint.mjJNT_HINGE)
+  body.add_geom(type=mujoco.mjtGeom.mjGEOM_SPHERE, size=[0.1, 0, 0])
+  act = spec.add_actuator(name="motor_elbow", target="elbow")
+  act.trntype = mujoco.mjtTrn.mjTRN_JOINT
+  act.gainprm[:] = [1] + [0] * 9
+  return spec
+
+
+def _prefixed_entity_spec() -> mujoco.MjSpec:
+  """Entity spec that composes a sub-model via internal attach."""
+  root = mujoco.MjSpec()
+  root.worldbody.add_geom(type=mujoco.mjtGeom.mjGEOM_PLANE, size=[1, 1, 0.01])
+  frame = root.worldbody.add_frame()
+  root.attach(_make_arm_spec(), prefix="arm/", frame=frame)
+  return root
+
+
+def test_builtin_actuator_with_internal_attach_prefix(device):
+  """Builtin actuator resolves joints through internal attach prefix."""
+  cfg = EntityCfg(
+    spec_fn=_prefixed_entity_spec,
+    articulation=EntityArticulationInfoCfg(
+      actuators=(
+        BuiltinMotorActuatorCfg(target_names_expr=("elbow",), effort_limit=100.0),
+      )
+    ),
+  )
+  entity = Entity(cfg)
+
+  # User-facing names should be stripped.
+  assert entity.joint_names == ("elbow",)
+
+  entity, _ = initialize_entity(entity, device)
+
+  assert len(entity.actuators) == 1
+  assert entity.actuators[0].target_names == ["elbow"]
+
+
+def test_xml_actuator_with_internal_attach_prefix(device):
+  """XML actuator matches targets through internal attach prefix."""
+  cfg = EntityCfg(
+    spec_fn=_prefixed_entity_spec,
+    articulation=EntityArticulationInfoCfg(
+      actuators=(XmlActuatorCfg(target_names_expr=("elbow",)),)
+    ),
+  )
+  entity = Entity(cfg)
+  entity, sim = initialize_entity(entity, device)
+
+  assert len(entity.actuators) == 1
+  assert entity.actuators[0].target_names == ["elbow"]
+
+
+# ---------------------------------------------------------------------------
+# Armature / frictionloss preservation tests
+# ---------------------------------------------------------------------------
+
+_XML_WITH_JOINT_PROPERTIES = """\
+<mujoco>
+  <worldbody>
+    <body name="base" pos="0 0 1">
+      <freejoint name="free_joint"/>
+      <geom type="box" size="0.2 0.2 0.1" mass="1.0"/>
+      <body name="link1">
+        <joint name="joint1" type="hinge" axis="0 0 1"
+               range="-1.57 1.57" armature="0.5" frictionloss="0.3"
+               damping="0.8"/>
+        <geom type="box" size="0.1 0.1 0.1" mass="0.1"/>
+      </body>
+    </body>
+  </worldbody>
+</mujoco>
+"""
+
+
+def test_xml_armature_frictionloss_preserved_by_default(device):
+  """Builtin actuator preserves XML armature/frictionloss when not set."""
+  cfg = BuiltinMotorActuatorCfg(target_names_expr=("joint1",), effort_limit=100.0)
+  entity = create_entity_with_actuator(_XML_WITH_JOINT_PROPERTIES, cfg)
+  entity, sim = initialize_entity(entity, device)
+
+  m = sim.mj_model
+  dof_id = m.jnt_dofadr[m.joint("joint1").id]
+  assert m.dof_armature[dof_id] == pytest.approx(0.5)
+  assert m.dof_frictionloss[dof_id] == pytest.approx(0.3)
+
+
+def test_explicit_armature_frictionloss_overrides_xml(device):
+  """Explicit values override XML armature/frictionloss."""
+  cfg = BuiltinMotorActuatorCfg(
+    target_names_expr=("joint1",),
+    effort_limit=100.0,
+    armature=1.0,
+    frictionloss=0.7,
+  )
+  entity = create_entity_with_actuator(_XML_WITH_JOINT_PROPERTIES, cfg)
+  entity, sim = initialize_entity(entity, device)
+
+  m = sim.mj_model
+  dof_id = m.jnt_dofadr[m.joint("joint1").id]
+  assert m.dof_armature[dof_id] == pytest.approx(1.0)
+  assert m.dof_frictionloss[dof_id] == pytest.approx(0.7)
+
+
+def test_explicit_zero_overrides_xml(device):
+  """Explicit 0.0 overrides XML values (distinguishes None from 0.0)."""
+  cfg = BuiltinMotorActuatorCfg(
+    target_names_expr=("joint1",),
+    effort_limit=100.0,
+    armature=0.0,
+    frictionloss=0.0,
+  )
+  entity = create_entity_with_actuator(_XML_WITH_JOINT_PROPERTIES, cfg)
+  entity, sim = initialize_entity(entity, device)
+
+  m = sim.mj_model
+  dof_id = m.jnt_dofadr[m.joint("joint1").id]
+  assert m.dof_armature[dof_id] == pytest.approx(0.0)
+  assert m.dof_frictionloss[dof_id] == pytest.approx(0.0)
+
+
+def test_partial_override_preserves_other_field(device):
+  """Setting only armature preserves XML frictionloss, and vice versa."""
+  cfg = BuiltinMotorActuatorCfg(
+    target_names_expr=("joint1",), effort_limit=100.0, armature=1.0
+  )
+  entity = create_entity_with_actuator(_XML_WITH_JOINT_PROPERTIES, cfg)
+  entity, sim = initialize_entity(entity, device)
+
+  m = sim.mj_model
+  dof_id = m.jnt_dofadr[m.joint("joint1").id]
+  assert m.dof_armature[dof_id] == pytest.approx(1.0)
+  assert m.dof_frictionloss[dof_id] == pytest.approx(0.3)
+
+
+@pytest.mark.parametrize(
+  "actuator_cfg",
+  [
+    BuiltinPositionActuatorCfg(
+      target_names_expr=("joint1",), stiffness=50.0, damping=5.0
+    ),
+    BuiltinVelocityActuatorCfg(target_names_expr=("joint1",), damping=5.0),
+    IdealPdActuatorCfg(
+      target_names_expr=("joint1",),
+      stiffness=50.0,
+      damping=5.0,
+      effort_limit=100.0,
+    ),
+  ],
+  ids=["position", "velocity", "ideal_pd"],
+)
+def test_all_actuator_types_preserve_xml_properties(device, actuator_cfg):
+  """All builtin actuator types preserve XML armature/frictionloss."""
+  entity = create_entity_with_actuator(_XML_WITH_JOINT_PROPERTIES, actuator_cfg)
+  entity, sim = initialize_entity(entity, device)
+
+  m = sim.mj_model
+  dof_id = m.jnt_dofadr[m.joint("joint1").id]
+  assert m.dof_armature[dof_id] == pytest.approx(0.5)
+  assert m.dof_frictionloss[dof_id] == pytest.approx(0.3)
+
+
+# ---------------------------------------------------------------------------
+# Viscous damping tests
+# ---------------------------------------------------------------------------
+
+
+def test_viscous_damping_default_preserves_xml(device):
+  """viscous_damping=None preserves XML joint damping (0.8)."""
+  cfg = BuiltinMotorActuatorCfg(target_names_expr=("joint1",), effort_limit=100.0)
+  entity = create_entity_with_actuator(_XML_WITH_JOINT_PROPERTIES, cfg)
+  entity, sim = initialize_entity(entity, device)
+
+  m = sim.mj_model
+  dof_id = m.jnt_dofadr[m.joint("joint1").id]
+  assert m.dof_damping[dof_id] == pytest.approx(0.8)
+
+
+def test_viscous_damping_explicit_overrides_xml(device):
+  """Explicit viscous_damping overrides XML damping."""
+  cfg = BuiltinMotorActuatorCfg(
+    target_names_expr=("joint1",), effort_limit=100.0, viscous_damping=2.5
+  )
+  entity = create_entity_with_actuator(_XML_WITH_JOINT_PROPERTIES, cfg)
+  entity, sim = initialize_entity(entity, device)
+
+  m = sim.mj_model
+  dof_id = m.jnt_dofadr[m.joint("joint1").id]
+  assert m.dof_damping[dof_id] == pytest.approx(2.5)
+
+
+def test_viscous_damping_explicit_zero_overrides_xml(device):
+  """Explicit 0.0 overrides XML damping (distinguishes None from 0.0)."""
+  cfg = BuiltinMotorActuatorCfg(
+    target_names_expr=("joint1",), effort_limit=100.0, viscous_damping=0.0
+  )
+  entity = create_entity_with_actuator(_XML_WITH_JOINT_PROPERTIES, cfg)
+  entity, sim = initialize_entity(entity, device)
+
+  m = sim.mj_model
+  dof_id = m.jnt_dofadr[m.joint("joint1").id]
+  assert m.dof_damping[dof_id] == pytest.approx(0.0)
+
+
+def test_viscous_damping_produces_passive_force(device, robot_xml):
+  """Nonzero viscous_damping produces a dissipative passive force."""
+  damping = 3.0
+  actuator_cfg = BuiltinMotorActuatorCfg(
+    target_names_expr=("joint.*",), effort_limit=100.0, viscous_damping=damping
+  )
+  entity = create_entity_with_actuator(robot_xml, actuator_cfg)
+  entity, sim = initialize_entity(entity, device)
+
+  vel = torch.tensor([[1.0, -2.0]], device=device)
+  entity.write_joint_state_to_sim(
+    position=torch.zeros(1, 2, device=device),
+    velocity=vel,
+  )
+  entity.write_data_to_sim()
+  sim.step()
+
+  qfrc_damper = sim.data.qfrc_damper[0]
+  # f = -b * v
+  assert qfrc_damper[6].item() == pytest.approx(-damping * 1.0, abs=1e-4)
+  assert qfrc_damper[7].item() == pytest.approx(-damping * -2.0, abs=1e-4)
@@ -0,0 +1,215 @@
+"""Tests for BuiltinActuatorGroup."""
+
+import mujoco
+import pytest
+import torch
+from conftest import get_test_device, load_fixture_xml
+
+from mjlab.actuator import (
+  BuiltinMotorActuatorCfg,
+  BuiltinPositionActuatorCfg,
+  IdealPdActuatorCfg,
+)
+from mjlab.entity import Entity, EntityArticulationInfoCfg, EntityCfg
+from mjlab.sim.sim import Simulation, SimulationCfg
+
+ROBOT_XML = load_fixture_xml("floating_base_articulated")
+
+ROBOT_XML_3JOINT = """
+<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="1.0"/>
+      <body name="link1" pos="0 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="0.1"/>
+      </body>
+      <body name="link2" pos="0 0 0">
+        <joint name="joint2" type="hinge" axis="0 0 1" range="-1.57 1.57"/>
+        <geom name="link2_geom" type="box" size="0.1 0.1 0.1" mass="0.1"/>
+      </body>
+      <body name="link3" pos="0 0 0">
+        <joint name="joint3" type="hinge" axis="0 0 1" range="-1.57 1.57"/>
+        <geom name="link3_geom" type="box" size="0.1 0.1 0.1" mass="0.1"/>
+      </body>
+    </body>
+  </worldbody>
+</mujoco>
+"""
+
+
+@pytest.fixture(scope="module")
+def device():
+  return get_test_device()
+
+
+def create_entity(actuator_cfgs, robot_xml=ROBOT_XML, **kwargs):
+  cfg = EntityCfg(
+    spec_fn=lambda: mujoco.MjSpec.from_string(robot_xml),
+    articulation=EntityArticulationInfoCfg(actuators=actuator_cfgs),
+    **kwargs,
+  )
+  return Entity(cfg)
+
+
+def initialize_entity(entity, device, num_envs=1):
+  model = entity.compile()
+  sim_cfg = SimulationCfg()
+  sim = Simulation(num_envs=num_envs, cfg=sim_cfg, model=model, device=device)
+  entity.initialize(model, sim.model, sim.data, device)
+  return entity, sim
+
+
+def test_position_actuator_batched(device):
+  """BuiltinPositionActuator writes position targets via batched path."""
+  actuator_cfg = BuiltinPositionActuatorCfg(
+    target_names_expr=("joint.*",), stiffness=50.0, damping=5.0
+  )
+  entity = create_entity((actuator_cfg,))
+  entity, sim = initialize_entity(entity, device)
+
+  entity.set_joint_position_target(torch.tensor([[0.5, -0.3]], device=device))
+  entity.write_data_to_sim()
+
+  ctrl = sim.data.ctrl[0]
+  assert torch.allclose(ctrl, torch.tensor([0.5, -0.3], device=device))
+
+
+def test_motor_actuator_batched(device):
+  """BuiltinMotorActuator writes effort targets via batched path."""
+  actuator_cfg = BuiltinMotorActuatorCfg(
+    target_names_expr=("joint.*",), effort_limit=100.0
+  )
+  entity = create_entity((actuator_cfg,))
+  entity, sim = initialize_entity(entity, device)
+
+  entity.set_joint_effort_target(torch.tensor([[10.0, -5.0]], device=device))
+  entity.write_data_to_sim()
+
+  ctrl = sim.data.ctrl[0]
+  assert torch.allclose(ctrl, torch.tensor([10.0, -5.0], device=device))
+
+
+def test_mixed_builtin_actuators(device):
+  """Multiple builtin actuator types can coexist and all use batched path."""
+  position_cfg = BuiltinPositionActuatorCfg(
+    target_names_expr=("joint1",), stiffness=50.0, damping=5.0
+  )
+  motor_cfg = BuiltinMotorActuatorCfg(target_names_expr=("joint2",), effort_limit=100.0)
+  entity = create_entity((position_cfg, motor_cfg))
+  entity, sim = initialize_entity(entity, device)
+
+  entity.set_joint_position_target(torch.tensor([[0.5, 0.0]], device=device))
+  entity.set_joint_effort_target(torch.tensor([[0.0, -3.0]], device=device))
+  entity.write_data_to_sim()
+
+  ctrl = sim.data.ctrl[0]
+  assert torch.allclose(ctrl, torch.tensor([0.5, -3.0], device=device))
+
+
+def test_builtin_and_custom_actuators(device):
+  """Builtin actuators use batched path, custom actuators use compute()."""
+  builtin_cfg = BuiltinPositionActuatorCfg(
+    target_names_expr=("joint1",), stiffness=50.0, damping=5.0
+  )
+  custom_cfg = IdealPdActuatorCfg(
+    target_names_expr=("joint2",), effort_limit=100.0, stiffness=50.0, damping=5.0
+  )
+  entity = create_entity((builtin_cfg, custom_cfg))
+  entity, sim = initialize_entity(entity, device)
+
+  entity.write_joint_state_to_sim(
+    position=torch.tensor([[0.0, 0.0]], device=device),
+    velocity=torch.tensor([[0.0, 0.0]], device=device),
+  )
+
+  entity.set_joint_position_target(torch.tensor([[0.5, 0.2]], device=device))
+  entity.set_joint_velocity_target(torch.tensor([[0.0, 0.0]], device=device))
+  entity.set_joint_effort_target(torch.tensor([[0.0, 0.0]], device=device))
+  entity.write_data_to_sim()
+
+  ctrl = sim.data.ctrl[0]
+  # joint1: builtin position -> ctrl = 0.5
+  # joint2: ideal pd -> ctrl = kp * (0.2 - 0.0) = 50.0 * 0.2 = 10.0
+  assert torch.allclose(ctrl, torch.tensor([0.5, 10.0], device=device))
+
+
+def test_builtin_group_mismatched_indices(device):
+  """Controls are written correctly when actuator order differs from joint order."""
+  # Actuators: position on joint2, motor on joint1+joint3.
+  # MuJoCo ctrl order follows declaration: [joint2, joint1, joint3].
+  position_cfg = BuiltinPositionActuatorCfg(
+    target_names_expr=("joint2",), stiffness=50.0, damping=5.0
+  )
+  motor_cfg = BuiltinMotorActuatorCfg(
+    target_names_expr=("joint1", "joint3"), effort_limit=100.0
+  )
+  entity = create_entity((position_cfg, motor_cfg), robot_xml=ROBOT_XML_3JOINT)
+  entity, sim = initialize_entity(entity, device, num_envs=1)
+
+  entity.set_joint_position_target(torch.tensor([[10.0, 20.0, 30.0]], device=device))
+  entity.set_joint_effort_target(torch.tensor([[100.0, 200.0, 300.0]], device=device))
+  entity.write_data_to_sim()
+
+  # ctrl follows declaration order: joint2=20, joint1=100, joint3=300.
+  expected = torch.tensor([20.0, 100.0, 300.0], device=device)
+  assert torch.allclose(sim.data.ctrl[0], expected)
+
+
+def test_sort_actuators(device):
+  """sort_actuators=True reorders ctrl to match joint definition order."""
+  position_cfg = BuiltinPositionActuatorCfg(
+    target_names_expr=("joint2",), stiffness=50.0, damping=5.0
+  )
+  motor_cfg = BuiltinMotorActuatorCfg(
+    target_names_expr=("joint1", "joint3"), effort_limit=100.0
+  )
+  entity = create_entity(
+    (position_cfg, motor_cfg),
+    robot_xml=ROBOT_XML_3JOINT,
+    sort_actuators=True,
+  )
+  entity, sim = initialize_entity(entity, device, num_envs=1)
+
+  entity.set_joint_position_target(torch.tensor([[10.0, 20.0, 30.0]], device=device))
+  entity.set_joint_effort_target(torch.tensor([[100.0, 200.0, 300.0]], device=device))
+  entity.write_data_to_sim()
+
+  # Sorted by joint order: joint1=100, joint2=20, joint3=300.
+  expected = torch.tensor([100.0, 20.0, 300.0], device=device)
+  assert torch.allclose(sim.data.ctrl[0], expected)
+
+
+def test_position_actuator_ctrlrange_matches_forcerange(device):
+  """Commanding to ctrlrange limits produces forces matching forcerange."""
+  actuator_cfg = BuiltinPositionActuatorCfg(
+    target_names_expr=("joint.*",),
+    stiffness=50.0,
+    damping=5.0,
+    effort_limit=100.0,
+  )
+  entity = create_entity((actuator_cfg,), robot_xml=ROBOT_XML_3JOINT)
+  model = entity.compile()
+  sim = Simulation(num_envs=1, cfg=SimulationCfg(), model=model, device=device)
+  entity.initialize(model, sim.model, sim.data, device)
+
+  n = model.nu
+  ctrlrange = torch.tensor(model.actuator_ctrlrange, device=device, dtype=torch.float32)
+  forcerange = torch.tensor(
+    model.actuator_forcerange, device=device, dtype=torch.float32
+  )
+  zeros = torch.zeros((1, n), device=device)
+
+  for col, label in [(1, "max"), (0, "min")]:
+    entity.write_joint_state_to_sim(position=zeros, velocity=zeros)
+    entity.set_joint_position_target(ctrlrange[:, col].unsqueeze(0))
+    entity.set_joint_effort_target(zeros)
+    entity.write_data_to_sim()
+    sim.step()
+
+    forces = sim.data.actuator_force[0]
+    assert torch.allclose(forces, forcerange[:, col], rtol=0.05), (
+      f"Forces at ctrl {label} {forces} not within 5% of "
+      f"forcerange {label} {forcerange[:, col]}"
+    )
@@ -0,0 +1,18 @@
+import mujoco
+import pytest
+
+from mjlab.asset_zoo.robots import get_g1_robot_cfg, get_go1_robot_cfg
+from mjlab.entity import Entity
+
+
+@pytest.mark.parametrize(
+  "robot_name,robot_cfg_fn",
+  [
+    ("G1", get_g1_robot_cfg),
+    ("GO1", get_go1_robot_cfg),
+  ],
+)
+def test_robot_compiles_parametrized(robot_name: str, robot_cfg_fn) -> None:
+  """Tests that all robots in the asset zoo compile without errors."""
+  robot_cfg = robot_cfg_fn()
+  assert isinstance(Entity(robot_cfg).compile(), mujoco.MjModel)
@@ -0,0 +1,328 @@
+"""Tests for the auto_reset config flag."""
+
+import pytest
+import torch
+from conftest import get_test_device
+from test_command_manager import CounterCommand, CounterCommandCfg
+
+from mjlab.envs import ManagerBasedRlEnv
+from mjlab.managers.event_manager import EventTermCfg
+from mjlab.tasks.cartpole.cartpole_env_cfg import cartpole_balance_env_cfg
+
+
+@pytest.fixture(scope="module")
+def device():
+  return get_test_device()
+
+
+def _make_cfg(auto_reset: bool):
+  cfg = cartpole_balance_env_cfg()
+  cfg.episode_length_s = 0.5  # 10 steps at dt=0.05
+  cfg.scene.num_envs = 4
+  cfg.auto_reset = auto_reset
+  return cfg
+
+
+def _step_until_done_env(env):
+  """Step with zero actions until at least one env is done. Return step outputs."""
+  for _ in range(env.max_episode_length + 5):
+    action = torch.zeros((env.num_envs, 1), device=env.device)
+    result = env.step(action)
+    terminated, truncated = result[2], result[3]
+    if (terminated | truncated).any():
+      return result
+  pytest.fail("No env terminated within max_episode_length steps")
+
+
+def test_auto_reset_true_resets_done_envs(device):
+  """With auto_reset=True (default), done envs are reset during step."""
+  env = ManagerBasedRlEnv(cfg=_make_cfg(auto_reset=True), device=device)
+  env.reset()
+  _, _, terminated, truncated, _ = _step_until_done_env(env)
+  done = terminated | truncated
+  done_ids = done.nonzero(as_tuple=False).squeeze(-1)
+
+  # Episode counter was reset to 0 for done envs.
+  assert (env.episode_length_buf[done_ids] == 0).all()
+  env.close()
+
+
+def test_auto_reset_false_preserves_terminal_state(device):
+  """With auto_reset=False, done envs are NOT reset and obs is the terminal state."""
+  env = ManagerBasedRlEnv(cfg=_make_cfg(auto_reset=False), device=device)
+  env.reset()
+  obs, _, terminated, truncated, _ = _step_until_done_env(env)
+  done = terminated | truncated
+  done_ids = done.nonzero(as_tuple=False).squeeze(-1)
+
+  # Episode counter was NOT reset (still at max_episode_length).
+  assert (env.episode_length_buf[done_ids] == env.max_episode_length).all()
+
+  # The returned obs must reflect the current (post-decimation terminal) sim
+  # state. Since no reset ran and the sim wasn't touched after step(), a fresh
+  # observation_manager.compute() on the current sim state must match exactly.
+  # This catches regressions where step() might return stale or post-reset obs.
+  env.observation_manager._obs_buffer = None  # bypass cache
+  fresh_obs = env.observation_manager.compute()
+  for group in obs:
+    returned = obs[group]
+    current = fresh_obs[group]
+    assert isinstance(returned, torch.Tensor) and isinstance(current, torch.Tensor)
+    assert torch.equal(returned, current)
+  env.close()
+
+
+def test_auto_reset_false_explicit_reset_works(device):
+  """After auto_reset=False, calling reset(env_ids=...) resets those envs."""
+  env = ManagerBasedRlEnv(cfg=_make_cfg(auto_reset=False), device=device)
+  env.reset()
+  _, _, terminated, truncated, _ = _step_until_done_env(env)
+  done = terminated | truncated
+  done_ids = done.nonzero(as_tuple=False).squeeze(-1)
+
+  # Manually reset done envs.
+  env.reset(env_ids=done_ids)
+  assert (env.episode_length_buf[done_ids] == 0).all()
+
+  # Can continue stepping after manual reset.
+  action = torch.zeros((env.num_envs, 1), device=env.device)
+  obs, reward, _, _, _ = env.step(action)
+  assert obs is not None
+  assert reward is not None
+  env.close()
+
+
+def test_auto_reset_false_requires_manual_reset_before_next_step(device):
+  """Raw env should reject another step until done envs are explicitly reset."""
+  env = ManagerBasedRlEnv(cfg=_make_cfg(auto_reset=False), device=device)
+  env.reset()
+  _step_until_done_env(env)
+
+  action = torch.zeros((env.num_envs, 1), device=env.device)
+  with pytest.raises(RuntimeError, match="must be reset via reset"):
+    env.step(action)
+
+  env.close()
+
+
+def _slice_obs(obs: dict, ids: torch.Tensor) -> dict[str, torch.Tensor]:
+  """Return a new obs dict containing only the rows at ``ids`` (per group)."""
+  return {k: v[ids] for k, v in obs.items() if isinstance(v, torch.Tensor)}
+
+
+def test_auto_reset_false_user_loop_pattern(device):
+  """Example: run your own training loop against an auto_reset=False env.
+
+  The pattern is:
+    1. After step(), derive done_ids from terminated | truncated.
+    2. Slice obs[done_ids] to get the true terminal observation and use it for
+      bootstrap / target computation.
+    3. Call env.reset(env_ids=done_ids) to reset only the done envs.
+    4. Continue stepping with the full batch.
+  """
+  env = ManagerBasedRlEnv(cfg=_make_cfg(auto_reset=False), device=device)
+  obs, _ = env.reset(seed=0)
+  episode_count = torch.zeros(env.num_envs, dtype=torch.long, device=env.device)
+  last_terminal_obs: dict[str, torch.Tensor] | None = None
+
+  action = torch.zeros((env.num_envs, 1), device=env.device)
+  for _ in range((env.max_episode_length + 2) * 3):
+    obs, _, terminated, truncated, _ = env.step(action)
+    done = terminated | truncated
+    if not done.any():
+      continue
+
+    done_ids = done.nonzero(as_tuple=False).squeeze(-1)
+    last_terminal_obs = _slice_obs(obs, done_ids)  # feed this to your critic/replay
+
+    episode_count[done_ids] += 1
+    obs, _ = env.reset(env_ids=done_ids)
+    if (episode_count >= 2).all():
+      break
+
+  assert (episode_count >= 2).all()
+  assert last_terminal_obs is not None
+  env.close()
+
+
+def test_auto_reset_false_obs_differs_from_auto_reset_true(device):
+  """Terminal obs (auto_reset=False) differs from post-reset obs (auto_reset=True)."""
+  # Run with auto_reset=True, capture post-reset obs for done envs.
+  env_on = ManagerBasedRlEnv(cfg=_make_cfg(auto_reset=True), device=device)
+  env_on.reset(seed=42)
+  obs_on, _, _, _, _ = _step_until_done_env(env_on)
+  env_on.close()
+
+  # Run with auto_reset=False with the same seed, capture terminal obs.
+  env_off = ManagerBasedRlEnv(cfg=_make_cfg(auto_reset=False), device=device)
+  env_off.reset(seed=42)
+  obs_off, _, _, _, _ = _step_until_done_env(env_off)
+  env_off.close()
+
+  # The observations should differ: one is post-reset, the other is terminal.
+  for group in obs_on:
+    on_val = obs_on[group]
+    off_val = obs_off[group]
+    assert isinstance(on_val, torch.Tensor) and isinstance(off_val, torch.Tensor)
+    assert not torch.equal(on_val, off_val)
+
+
+def test_partial_reset_leaves_other_envs_obs_buffers_untouched(device):
+  """reset(env_ids=...) must not advance other envs' history/delay buffers."""
+  cfg = _make_cfg(auto_reset=False)
+  cfg.observations["actor"].terms["cart_pos"].history_length = 4
+  cfg.observations["actor"].terms["cart_vel"].delay_min_lag = 2
+  cfg.observations["actor"].terms["cart_vel"].delay_max_lag = 2
+  env = ManagerBasedRlEnv(cfg=cfg, device=device)
+  env.reset()
+  action = torch.zeros((env.num_envs, 1), device=env.device)
+  for _ in range(3):
+    env.step(action)
+
+  om = env.observation_manager
+  hist = om._group_obs_term_history_buffer["actor"]["cart_pos"]
+  delay = om._group_obs_term_delay_buffer["actor"]["cart_vel"]
+  h_before = hist.buffer[0].clone()
+  d_before = delay.peek()[0].clone()
+
+  env.reset(env_ids=torch.tensor([1], dtype=torch.int64, device=env.device))
+
+  # Env 0 was not reset: history window and delayed obs are untouched.
+  assert torch.equal(hist.buffer[0], h_before)
+  assert torch.equal(delay.peek()[0], d_before)
+
+  # Env 1 was reset: history is backfilled with its single post-reset frame.
+  h1 = hist.buffer[1]
+  assert torch.all(h1 == h1[0])
+  assert hist.current_length[1].item() == 1
+  env.close()
+
+
+# Section: parity between auto-reset and the explicit reset() flow.
+
+_COMMAND_T = 7.0
+_INTERVAL_T = 5.0
+_MARKER_VEL = 37.0
+
+
+def _noop_event(env, env_ids) -> None:
+  del env, env_ids
+
+
+def _write_marker_velocity(env, env_ids) -> None:
+  """Overwrite joint velocities with a recognizable marker value."""
+  asset = env.scene["cartpole"]
+  if env_ids is None:
+    env_ids = torch.arange(env.num_envs, device=env.device)
+  joint_pos = asset.data.joint_pos[env_ids]
+  joint_vel = torch.full_like(asset.data.joint_vel[env_ids], _MARKER_VEL)
+  asset.write_joint_state_to_sim(joint_pos, joint_vel, env_ids=env_ids)
+
+
+def _make_parity_cfg(auto_reset: bool):
+  """Cartpole cfg with fixed-interval command and event timers.
+
+  Fixed ranges make timer values deterministic, so assertions are exact and
+  independent of RNG state and physics nondeterminism.
+  """
+  cfg = _make_cfg(auto_reset)
+  cfg.commands = {
+    "counter": CounterCommandCfg(resampling_time_range=(_COMMAND_T, _COMMAND_T))
+  }
+  cfg.events["probe"] = EventTermCfg(
+    func=_noop_event, mode="interval", interval_range_s=(_INTERVAL_T, _INTERVAL_T)
+  )
+  return cfg
+
+
+def _stagger_env0(env, steps_until_reset: int) -> torch.Tensor:
+  """Advance env 0's episode clock so it times out after the given steps."""
+  env.reset()
+  env.episode_length_buf[0] = env.max_episode_length - steps_until_reset
+  return torch.zeros((env.num_envs, 1), device=env.device)
+
+
+def test_auto_reset_preserves_fresh_command_timer(device):
+  """A command timer resampled by an auto-reset is not decremented that step."""
+  env = ManagerBasedRlEnv(cfg=_make_parity_cfg(auto_reset=True), device=device)
+  action = _stagger_env0(env, steps_until_reset=2)
+  env.step(action)
+  env.step(action)  # Env 0 times out and auto-resets here.
+  assert env.episode_length_buf[0].item() == 0
+
+  term = env.command_manager.get_term("counter")
+  assert isinstance(term, CounterCommand)
+  time_left = term.time_left
+  expected_running = _COMMAND_T - 2 * env.step_dt
+  assert torch.allclose(time_left[0], torch.tensor(_COMMAND_T, device=env.device))
+  assert torch.allclose(
+    time_left[1:], torch.tensor(expected_running, device=env.device)
+  )
+  env.close()
+
+
+def test_auto_reset_preserves_fresh_interval_event_timer(device):
+  """An interval event timer resampled by an auto-reset is not decremented."""
+  env = ManagerBasedRlEnv(cfg=_make_parity_cfg(auto_reset=True), device=device)
+  action = _stagger_env0(env, steps_until_reset=2)
+  env.step(action)
+  env.step(action)  # Env 0 times out and auto-resets here.
+  assert env.episode_length_buf[0].item() == 0
+
+  timer = env.event_manager._interval_term_time_left[0]
+  expected_running = _INTERVAL_T - 2 * env.step_dt
+  assert torch.allclose(timer[0], torch.tensor(_INTERVAL_T, device=env.device))
+  assert torch.allclose(timer[1:], torch.tensor(expected_running, device=env.device))
+  env.close()
+
+
+def test_interval_event_acts_on_pre_reset_state(device):
+  """An interval event firing on a reset step hits the terminal state, so the
+  freshly reset env comes out clean, as in the explicit reset flow."""
+  cfg = _make_cfg(auto_reset=True)
+  cfg.events["kick"] = EventTermCfg(
+    func=_write_marker_velocity, mode="interval", interval_range_s=(0.0, 0.0)
+  )
+  env = ManagerBasedRlEnv(cfg=cfg, device=device)
+  action = _stagger_env0(env, steps_until_reset=1)
+  env.step(action)  # Kick fires for all envs; env 0 then auto-resets.
+  assert env.episode_length_buf[0].item() == 0
+
+  joint_vel = env.scene["cartpole"].data.joint_vel
+  # Env 0 was reset after the kick: its velocity is the reset distribution's,
+  # not the marker. Env 1 was not reset and still carries the marker.
+  assert joint_vel[0].abs().max().item() < 1.0
+  assert torch.allclose(joint_vel[1], torch.full_like(joint_vel[1], _MARKER_VEL))
+  env.close()
+
+
+def test_auto_reset_matches_manual_reset_timers(device):
+  """After a reset, both flows leave identical command/event timer state."""
+  auto_env = ManagerBasedRlEnv(cfg=_make_parity_cfg(auto_reset=True), device=device)
+  manual_env = ManagerBasedRlEnv(cfg=_make_parity_cfg(auto_reset=False), device=device)
+  auto_env.reset(seed=0)
+  manual_env.reset(seed=0)
+
+  action = torch.zeros((auto_env.num_envs, 1), device=auto_env.device)
+  done = torch.zeros(auto_env.num_envs, dtype=torch.bool, device=auto_env.device)
+  for _ in range(auto_env.max_episode_length):
+    auto_env.step(action)
+    _, _, terminated, truncated, _ = manual_env.step(action)
+    done = terminated | truncated
+  done_ids = done.nonzero(as_tuple=False).squeeze(-1)
+  assert len(done_ids) == manual_env.num_envs  # Time-out is synchronized.
+  manual_env.reset(env_ids=done_ids)
+
+  for env in (auto_env, manual_env):
+    term = env.command_manager.get_term("counter")
+    assert isinstance(term, CounterCommand)
+    interval_timer = env.event_manager._interval_term_time_left[0]
+    assert torch.all(env.episode_length_buf == 0)
+    assert torch.allclose(term.time_left, torch.full_like(term.time_left, _COMMAND_T))
+    assert torch.allclose(interval_timer, torch.full_like(interval_timer, _INTERVAL_T))
+    # Stateful command advance: resample zeroed the counter, then exactly one
+    # post-reset update ran in both flows.
+    assert torch.all(term.ticks == 1)
+
+  auto_env.close()
+  manual_env.close()
@@ -0,0 +1,653 @@
+"""Tests for BuiltinDcMotorActuator.
+
+Covers wiring of MuJoCo's native ``<dcmotor>`` element through mjlab: the
+three input modes (voltage / position / velocity), torque saturation,
+config validation, and DR integration.
+"""
+
+import math
+from unittest.mock import Mock
+
+import mujoco
+import pytest
+import torch
+from conftest import (
+  create_entity_with_actuator,
+  get_test_device,
+  initialize_entity,
+  load_fixture_xml,
+)
+
+from mjlab.actuator import (
+  BuiltinDcMotorActuator,
+  BuiltinDcMotorActuatorCfg,
+  DcMotorDatasheetParams,
+  DcMotorInputMode,
+  DcMotorPhysicalParams,
+)
+from mjlab.actuator.actuator import TransmissionType
+from mjlab.entity import Entity, EntityArticulationInfoCfg, EntityCfg
+from mjlab.envs.mdp import dr
+from mjlab.managers.scene_entity_config import SceneEntityCfg
+from mjlab.scene import Scene, SceneCfg
+from mjlab.sim.sim import Simulation, SimulationCfg
+
+ROBOT_XML = load_fixture_xml("floating_base_articulated")
+
+# Motor characterization used throughout (resolves to K=0.24, R=2.88).
+V_NOM, TAU_STALL, OMEGA_NL = 24.0, 2.0, 100.0
+K = V_NOM / OMEGA_NL
+R = K * V_NOM / TAU_STALL
+
+
+@pytest.fixture(scope="module")
+def device():
+  return get_test_device()
+
+
+DATASHEET = DcMotorDatasheetParams(
+  nominal_voltage=V_NOM, stall_torque=TAU_STALL, no_load_speed=OMEGA_NL
+)
+
+
+def _make_cfg(
+  *,
+  mode: DcMotorInputMode = DcMotorInputMode.POSITION,
+  stiffness=5.0,
+  damping=0.5,
+  voltage_limit=24.0,
+  **extra,
+) -> BuiltinDcMotorActuatorCfg:
+  """Build a cfg with sensible PID defaults. ``extra`` forwards any other
+  BuiltinDcMotorActuatorCfg kwarg (effort_limit, integral_gain, thermal,
+  delay_*, etc.)."""
+  return BuiltinDcMotorActuatorCfg(
+    target_names_expr=("joint.*",),
+    mode=mode,
+    motor_params=DATASHEET,
+    stiffness=stiffness,
+    damping=damping,
+    voltage_limit=voltage_limit,
+    **extra,
+  )
+
+
+def _make_entity(**kwargs) -> Entity:
+  return create_entity_with_actuator(ROBOT_XML, _make_cfg(**kwargs))
+
+
+def _make_initialized(device, **kwargs):
+  """Build entity from cfg kwargs and initialize it through the sim."""
+  return initialize_entity(_make_entity(**kwargs), device)
+
+
+def _drive(
+  entity: Entity,
+  sim,
+  device: str,
+  *,
+  pos_target=None,
+  vel_target=None,
+  effort_target=None,
+  q0=None,
+  qd0=None,
+) -> None:
+  zero = torch.zeros(1, 2, device=device)
+  entity.write_joint_state_to_sim(
+    position=q0 if q0 is not None else zero,
+    velocity=qd0 if qd0 is not None else zero,
+  )
+  entity.set_joint_position_target(pos_target if pos_target is not None else zero)
+  entity.set_joint_velocity_target(vel_target if vel_target is not None else zero)
+  entity.set_joint_effort_target(effort_target if effort_target is not None else zero)
+  entity.write_data_to_sim()
+  sim.forward()
+
+
+# Wiring sanity.
+
+
+def test_kr_packed_into_gainprm(device):
+  """The XML compiler derives K and R from the nominal triplet."""
+  _, sim = initialize_entity(_make_entity(effort_limit=1.5), device)
+  m = sim.mj_model
+  for i in range(2):
+    assert m.actuator_gainprm[i, 0] == pytest.approx(R, abs=1e-6)
+    assert m.actuator_gainprm[i, 1] == pytest.approx(K, abs=1e-6)
+    assert m.actuator_gainprm[i, 4] == pytest.approx(5.0)  # kp
+    assert m.actuator_gainprm[i, 6] == pytest.approx(0.5)  # kd
+    assert m.actuator_gainprm[i, 7] == pytest.approx(24.0)  # Vmax
+    assert m.actuator_gainprm[i, 8] == pytest.approx(1.0)  # input_mode=position
+    assert m.actuator_gaintype[i] == mujoco.mjtGain.mjGAIN_DCMOTOR
+    assert m.actuator_biastype[i] == mujoco.mjtBias.mjBIAS_DCMOTOR
+    # No activation state: ki=0, no inductance, no thermal/lugre/slew.
+    assert m.actuator_actnum[i] == 0
+
+
+def test_motor_const_path(device):
+  """Physical params pack K = sqrt(Kt*Ke) and R verbatim."""
+  cfg = BuiltinDcMotorActuatorCfg(
+    target_names_expr=("joint.*",),
+    mode=DcMotorInputMode.VOLTAGE,
+    motor_params=DcMotorPhysicalParams(kt=0.1, ke=0.05, resistance=2.0),
+  )
+  _, sim = initialize_entity(create_entity_with_actuator(ROBOT_XML, cfg), device)
+  m = sim.mj_model
+  for i in range(2):
+    assert m.actuator_gainprm[i, 0] == pytest.approx(2.0, abs=1e-6)
+    assert m.actuator_gainprm[i, 1] == pytest.approx((0.1 * 0.05) ** 0.5, abs=1e-6)
+
+
+# Stateless motor physics.
+
+
+def test_voltage_mode_steady_state(device):
+  """At rest, ctrl = V -> tau = K * V / R."""
+  cfg = BuiltinDcMotorActuatorCfg(
+    target_names_expr=("joint.*",),
+    mode=DcMotorInputMode.VOLTAGE,
+    motor_params=DATASHEET,
+  )
+  entity, sim = initialize_entity(create_entity_with_actuator(ROBOT_XML, cfg), device)
+  V = torch.tensor([[10.0, -5.0]], device=device)
+  _drive(entity, sim, device, effort_target=V)
+  v_adr = entity.indexing.joint_v_adr
+  expected = K * V[0] / R
+  assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], expected, atol=1e-4)
+
+
+def test_voltage_mode_voltage_limit_zero_is_noop(device):
+  """Docstring promises ``voltage_limit=0`` disables clamping. Verify against
+  MuJoCo's ``dcmotor_voltage`` (which only clamps when ``Vmax > 0``)."""
+  cfg = BuiltinDcMotorActuatorCfg(
+    target_names_expr=("joint.*",),
+    mode=DcMotorInputMode.VOLTAGE,
+    motor_params=DATASHEET,
+    voltage_limit=0.0,
+  )
+  entity, sim = initialize_entity(create_entity_with_actuator(ROBOT_XML, cfg), device)
+  V = torch.tensor([[1000.0, 0.0]], device=device)  # absurdly high voltage.
+  _drive(entity, sim, device, effort_target=V)
+  v_adr = entity.indexing.joint_v_adr
+  expected = K * V[0] / R
+  assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], expected, atol=1e-2)
+
+
+def test_back_emf_reduces_torque_at_velocity(device):
+  """Same V, joint moving at omega: tau = K * (V - K * omega) / R."""
+  cfg = BuiltinDcMotorActuatorCfg(
+    target_names_expr=("joint.*",),
+    mode=DcMotorInputMode.VOLTAGE,
+    motor_params=DATASHEET,
+  )
+  entity, sim = initialize_entity(create_entity_with_actuator(ROBOT_XML, cfg), device)
+  V = torch.tensor([[10.0, 0.0]], device=device)
+  omega0 = torch.tensor([[2.0, 0.0]], device=device)
+  _drive(entity, sim, device, effort_target=V, qd0=omega0)
+  v_adr = entity.indexing.joint_v_adr
+  expected = K * (V[0] - K * omega0[0]) / R
+  assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], expected, atol=1e-4)
+
+
+def test_position_mode_pid_at_rest(device):
+  """kd=0, no Vmax clamp: tau = K * kp * (target - q) / R."""
+  # voltage_limit must be >0 (cfg invariant), pick it big enough not to clamp.
+  entity, sim = initialize_entity(
+    _make_entity(damping=0.0, voltage_limit=1000.0), device
+  )
+  pos = torch.tensor([[0.1, -0.05]], device=device)
+  _drive(entity, sim, device, pos_target=pos)
+  v_adr = entity.indexing.joint_v_adr
+  expected = K * 5.0 * pos[0] / R
+  assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], expected, atol=1e-4)
+
+
+def test_position_mode_voltage_clamp(device):
+  """Huge position error -> PID voltage saturates at Vmax."""
+  entity, sim = initialize_entity(
+    _make_entity(stiffness=100.0, damping=0.0, voltage_limit=2.0),
+    device,
+  )
+  # kp * err = 100 * 0.5 = 50 V, well above Vmax=2.
+  pos = torch.tensor([[0.5, 0.0]], device=device)
+  _drive(entity, sim, device, pos_target=pos)
+  v_adr = entity.indexing.joint_v_adr
+  qfrc = sim.data.qfrc_actuator[0, v_adr]
+  expected_first = K * 2.0 / R  # tau at clamped V.
+  assert qfrc[0].item() == pytest.approx(expected_first, abs=1e-4)
+  assert qfrc[1].item() == pytest.approx(0.0, abs=1e-4)
+
+
+def test_velocity_mode_pid(device):
+  """P-only velocity tracking: tau = K * kp * (target - qdot) / R."""
+  entity, sim = initialize_entity(
+    _make_entity(mode=DcMotorInputMode.VELOCITY, damping=0.0, voltage_limit=1000.0),
+    device,
+  )
+  qd0 = torch.tensor([[1.0, 0.0]], device=device)
+  vel_target = torch.tensor([[3.0, 0.0]], device=device)
+  _drive(entity, sim, device, vel_target=vel_target, qd0=qd0)
+  v_adr = entity.indexing.joint_v_adr
+  # back-EMF subtracts K*omega; this is folded into the dcmotor bias.
+  # voltage = kp*(target - qdot); tau = K*(voltage - K*omega)/R.
+  voltage = 5.0 * (vel_target[0] - qd0[0])
+  expected = K * (voltage - K * qd0[0]) / R
+  assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], expected, atol=1e-4)
+
+
+def test_effort_limit_clamps_torque(device):
+  """forcerange clamps the algebraic torque output."""
+  entity, sim = initialize_entity(
+    _make_entity(stiffness=100.0, damping=0.0, voltage_limit=1000.0, effort_limit=0.1),
+    device,
+  )
+  m = sim.mj_model
+  for i in range(2):
+    assert m.actuator_forcelimited[i] == 1
+    assert m.actuator_forcerange[i, 0] == pytest.approx(-0.1)
+    assert m.actuator_forcerange[i, 1] == pytest.approx(0.1)
+
+  # Unclamped tau would be K * 100 * 0.5 / R ~= K*50/R, well above 0.1.
+  pos = torch.tensor([[0.5, 0.0]], device=device)
+  _drive(entity, sim, device, pos_target=pos)
+  v_adr = entity.indexing.joint_v_adr
+  qfrc = sim.data.qfrc_actuator[0, v_adr]
+  assert qfrc[0].item() == pytest.approx(0.1, abs=1e-4)
+  assert qfrc[1].item() == pytest.approx(0.0, abs=1e-4)
+
+
+# Cogging.
+
+
+def test_cogging_packed_into_biasprm(device):
+  """``cogging=(A, Np, phi)`` packs into ``biasprm[0:3]``."""
+  cfg = BuiltinDcMotorActuatorCfg(
+    target_names_expr=("joint.*",),
+    mode=DcMotorInputMode.VOLTAGE,
+    motor_params=DATASHEET,
+    cogging=(0.5, 4.0, 0.1),
+  )
+  _, sim = initialize_entity(create_entity_with_actuator(ROBOT_XML, cfg), device)
+  m = sim.mj_model
+  for i in range(2):
+    assert m.actuator_biasprm[i, 0] == pytest.approx(0.5)
+    assert m.actuator_biasprm[i, 1] == pytest.approx(4.0)
+    assert m.actuator_biasprm[i, 2] == pytest.approx(0.1)
+
+
+def test_cogging_contributes_torque(device):
+  """At ctrl=0 (no electromagnetic torque), qfrc_actuator equals the cogging
+  term ``A * sin(Np * q + phi)`` evaluated at the joint angle."""
+  A, Np, phi = 0.5, 4.0, 0.1
+  cfg = BuiltinDcMotorActuatorCfg(
+    target_names_expr=("joint.*",),
+    mode=DcMotorInputMode.VOLTAGE,
+    motor_params=DATASHEET,
+    cogging=(A, Np, phi),
+  )
+  entity, sim = initialize_entity(create_entity_with_actuator(ROBOT_XML, cfg), device)
+  q0, q1 = 0.3, -0.2
+  _drive(
+    entity,
+    sim,
+    device,
+    q0=torch.tensor([[q0, q1]], device=device),
+    effort_target=torch.zeros(1, 2, device=device),
+  )
+  v_adr = entity.indexing.joint_v_adr
+  qfrc = sim.data.qfrc_actuator[0, v_adr]
+  assert qfrc[0].item() == pytest.approx(A * math.sin(Np * q0 + phi), abs=1e-5)
+  assert qfrc[1].item() == pytest.approx(A * math.sin(Np * q1 + phi), abs=1e-5)
+
+
+def test_cogging_bypasses_effort_limit(device):
+  """Cogging is added *after* the forcerange clamp (MuJoCo's intentional
+  model: ``effort_limit`` bounds electromagnetic torque, cogging is
+  mechanical). Total torque can exceed ``effort_limit`` by up to the
+  cogging amplitude."""
+  A, Np, phi = 0.5, 0.0, math.pi / 2  # sin(pi/2)=1, so cogging = A at any q.
+  cfg = BuiltinDcMotorActuatorCfg(
+    target_names_expr=("joint.*",),
+    mode=DcMotorInputMode.VOLTAGE,
+    motor_params=DATASHEET,
+    cogging=(A, Np, phi),
+    effort_limit=0.05,  # An order of magnitude below A.
+  )
+  entity, sim = initialize_entity(create_entity_with_actuator(ROBOT_XML, cfg), device)
+  # Pick a voltage large enough that the electromagnetic torque alone
+  # would saturate forcerange at +/- 0.05.
+  V = torch.tensor([[100.0, 0.0]], device=device)
+  _drive(entity, sim, device, effort_target=V)
+  v_adr = entity.indexing.joint_v_adr
+  qfrc = sim.data.qfrc_actuator[0, v_adr]
+  # joint1: electromagnetic clamped to +0.05, plus cogging A=0.5.
+  assert qfrc[0].item() == pytest.approx(0.05 + A, abs=1e-5)
+  # joint2: zero voltage, electromagnetic=0, only cogging.
+  assert qfrc[1].item() == pytest.approx(A, abs=1e-5)
+
+
+# Optional stateful extensions (integral, slew, inductance, thermal, LuGre).
+# Each behavior check compares against a baseline with the feature disabled
+# so that removing the wiring in edit_spec causes the comparison to fail.
+
+
+def _step_n(entity, sim, device, n: int, *, pos_target=None, eff_target=None):
+  zero = torch.zeros(1, 2, device=device)
+  entity.write_joint_state_to_sim(position=zero, velocity=zero)
+  for _ in range(n):
+    entity.set_joint_position_target(pos_target if pos_target is not None else zero)
+    entity.set_joint_velocity_target(zero)
+    entity.set_joint_effort_target(eff_target if eff_target is not None else zero)
+    entity.write_data_to_sim()
+    sim.step()
+
+
+def _qfrc(entity, sim) -> torch.Tensor:
+  return sim.data.qfrc_actuator[0, entity.indexing.joint_v_adr].clone()
+
+
+def test_integral_gain_ramps_torque(device):
+  """Integrator in position mode ramps torque over time even with ``kp``
+  and ``kd`` near zero."""
+  # stiffness must be > 0 (validation); choose tiny so ki dominates.
+  base = dict(
+    mode=DcMotorInputMode.POSITION, stiffness=1e-4, damping=0.0, voltage_limit=24.0
+  )
+  ent_off, sim_off = _make_initialized(device, **base, integral_gain=0.0)
+  ent_on, sim_on = _make_initialized(device, **base, integral_gain=10.0)
+
+  target = torch.tensor([[0.5, 0.0]], device=device)
+  for sim, ent in ((sim_off, ent_off), (sim_on, ent_on)):
+    _step_n(ent, sim, device, n=20, pos_target=target)
+  assert _qfrc(ent_on, sim_on)[0].abs() > 100.0 * _qfrc(ent_off, sim_off)[0].abs()
+
+
+def test_slew_rate_limits_voltage(device):
+  """``slew_rate`` rate-limits ``ctrl``: after one step, effective voltage
+  is far below the requested input."""
+  base = dict(
+    mode=DcMotorInputMode.VOLTAGE, stiffness=0.0, damping=0.0, voltage_limit=0.0
+  )
+  ent_off, sim_off = _make_initialized(device, **base, slew_rate=0.0)
+  ent_on, sim_on = _make_initialized(device, **base, slew_rate=10.0)
+
+  V = torch.tensor([[100.0, 0.0]], device=device)
+  for sim, ent in ((sim_off, ent_off), (sim_on, ent_on)):
+    _step_n(ent, sim, device, n=1, eff_target=V)
+  assert _qfrc(ent_off, sim_off)[0] > 100.0 * _qfrc(ent_on, sim_on)[0]
+
+
+def test_inductance_lags_current(device):
+  """Large ``inductance`` (te >> dt) suppresses early-step torque."""
+  base = dict(
+    mode=DcMotorInputMode.VOLTAGE, stiffness=0.0, damping=0.0, voltage_limit=0.0
+  )
+  ent_off, sim_off = _make_initialized(device, **base, inductance=0.0)
+  ent_on, sim_on = _make_initialized(device, **base, inductance=1.0)
+
+  V = torch.tensor([[10.0, 0.0]], device=device)
+  for sim, ent in ((sim_off, ent_off), (sim_on, ent_on)):
+    _step_n(ent, sim, device, n=2, eff_target=V)
+  assert _qfrc(ent_off, sim_off)[0].abs() > 10.0 * _qfrc(ent_on, sim_on)[0].abs()
+
+
+def test_thermal_decays_torque(device):
+  """I^2R heating raises T, which raises effective resistance and decays
+  torque over time."""
+  # Params chosen for visible effect in a handful of steps without going
+  # numerically unstable: small C (fast heating) and modest alpha.
+  base = dict(
+    mode=DcMotorInputMode.VOLTAGE, stiffness=0.0, damping=0.0, voltage_limit=0.0
+  )
+  ent_off, sim_off = _make_initialized(device, **base)
+  ent_on, sim_on = _make_initialized(
+    device, **base, thermal=(1.0, 0.1, 0.0, 0.01, 0.0, 0.0)
+  )
+
+  V = torch.tensor([[100.0, 0.0]], device=device)
+  for sim, ent in ((sim_off, ent_off), (sim_on, ent_on)):
+    _step_n(ent, sim, device, n=5, eff_target=V)
+  assert _qfrc(ent_on, sim_on)[0].abs() < 0.5 * _qfrc(ent_off, sim_off)[0].abs()
+
+
+def test_lugre_subtracts_friction(device):
+  """LuGre friction subtracts a velocity-dependent force after the
+  ``effort_limit`` clamp (mechanical, like cogging)."""
+  # Static comparison at v>0, ctrl=0; avoids feedback between LuGre slowing
+  # the joint and back-EMF easing off under sim.step().
+  #   no LuGre:  qfrc = -K^2 * v / R              (back-EMF only)
+  #   w/ LuGre:  qfrc = -K^2 * v / R - sigma1*v - ...
+  base = dict(
+    mode=DcMotorInputMode.VOLTAGE, stiffness=0.0, damping=0.0, voltage_limit=0.0
+  )
+  ent_off, sim_off = _make_initialized(device, **base)
+  ent_on, sim_on = _make_initialized(
+    device, **base, lugre=(1e4, 100.0, 0.1, 0.15, 0.01)
+  )
+
+  zero = torch.zeros(1, 2, device=device)
+  v0 = torch.tensor([[1.0, 0.0]], device=device)
+  for sim, ent in ((sim_off, ent_off), (sim_on, ent_on)):
+    ent.write_joint_state_to_sim(position=zero, velocity=v0)
+    ent.set_joint_position_target(zero)
+    ent.set_joint_velocity_target(zero)
+    ent.set_joint_effort_target(zero)
+    ent.write_data_to_sim()
+    sim.forward()
+  assert abs(_qfrc(ent_on, sim_on)[0]) > 100.0 * abs(_qfrc(ent_off, sim_off)[0])
+
+
+# Config validation.
+
+
+def test_pid_mode_requires_gains():
+  with pytest.raises(ValueError, match="stiffness"):
+    BuiltinDcMotorActuatorCfg(
+      target_names_expr=("j",),
+      mode=DcMotorInputMode.POSITION,
+      motor_params=DATASHEET,
+      voltage_limit=1.0,
+    )
+  with pytest.raises(ValueError, match="voltage_limit"):
+    BuiltinDcMotorActuatorCfg(
+      target_names_expr=("j",),
+      mode=DcMotorInputMode.POSITION,
+      motor_params=DATASHEET,
+      stiffness=1.0,
+    )
+
+
+def test_voltage_mode_rejects_pid_gains():
+  with pytest.raises(ValueError, match="VOLTAGE"):
+    BuiltinDcMotorActuatorCfg(
+      target_names_expr=("j",),
+      mode=DcMotorInputMode.VOLTAGE,
+      motor_params=DATASHEET,
+      stiffness=1.0,
+    )
+
+
+def test_site_rejected():
+  with pytest.raises(ValueError, match="SITE"):
+    BuiltinDcMotorActuatorCfg(
+      target_names_expr=("j",),
+      motor_params=DATASHEET,
+      stiffness=1.0,
+      voltage_limit=1.0,
+      transmission_type=TransmissionType.SITE,
+    )
+
+
+# Joint-level passthrough.
+
+
+def test_armature_applied(device):
+  _, sim = initialize_entity(_make_entity(armature=0.7), device)
+  m = sim.mj_model
+  for jname in ("joint1", "joint2"):
+    dof_id = m.jnt_dofadr[m.joint(jname).id]
+    assert m.dof_armature[dof_id] == pytest.approx(0.7)
+
+
+# Domain randomization.
+
+
+def _scene_env(
+  device,
+  num_envs=2,
+  mode: DcMotorInputMode = DcMotorInputMode.POSITION,
+):
+  def spec_fn():
+    spec = mujoco.MjSpec.from_string(ROBOT_XML)
+    for a in list(spec.actuators):
+      spec.delete(a)
+    return spec
+
+  entity_cfg = EntityCfg(
+    spec_fn=spec_fn,
+    articulation=EntityArticulationInfoCfg(
+      actuators=(
+        BuiltinDcMotorActuatorCfg(
+          target_names_expr=("joint.*",),
+          mode=mode,
+          motor_params=DATASHEET,
+          stiffness=5.0 if mode != DcMotorInputMode.VOLTAGE else 0.0,
+          damping=0.5 if mode != DcMotorInputMode.VOLTAGE else 0.0,
+          voltage_limit=24.0 if mode != DcMotorInputMode.VOLTAGE else 0.0,
+          effort_limit=50.0,
+        ),
+      )
+    ),
+  )
+  scene_cfg = SceneCfg(num_envs=num_envs, entities={"robot": entity_cfg})
+  scene = Scene(scene_cfg, device)
+  model = scene.compile()
+  sim = Simulation(num_envs=num_envs, cfg=SimulationCfg(), model=model, device=device)
+  scene.initialize(model, sim.model, sim.data)
+
+  env = Mock()
+  env.num_envs = num_envs
+  env.device = device
+  env.scene = {"robot": scene["robot"]}
+  env.sim = sim
+  return env
+
+
+@pytest.mark.parametrize(
+  "operation, kp_in, kd_in, kp_expected, kd_expected",
+  [
+    # scale: multiplies the configured defaults (kp=5.0, kd=0.5).
+    ("scale", 2.0, 3.0, 2.0 * 5.0, 3.0 * 0.5),
+    # abs: writes the value directly.
+    ("abs", 10.0, 2.0, 10.0, 2.0),
+  ],
+)
+def test_dr_pd_gains_position_mode(
+  device, operation, kp_in, kd_in, kp_expected, kd_expected
+):
+  env = _scene_env(device)
+  robot = env.scene["robot"]
+  act = robot.actuators[0]
+  assert isinstance(act, BuiltinDcMotorActuator)
+  ctrl_ids = act.global_ctrl_ids
+  env.sim.expand_model_fields(("actuator_gainprm", "actuator_biasprm"))
+
+  dr.pd_gains(
+    env,
+    env_ids=torch.tensor([0], device=device),
+    kp_range=(kp_in, kp_in),
+    kd_range=(kd_in, kd_in),
+    asset_cfg=SceneEntityCfg("robot"),
+    operation=operation,
+  )
+
+  m = env.sim.model
+  n = len(ctrl_ids)
+  assert torch.allclose(
+    m.actuator_gainprm[0, ctrl_ids, 4], torch.full((n,), kp_expected, device=device)
+  )
+  assert torch.allclose(
+    m.actuator_gainprm[0, ctrl_ids, 6], torch.full((n,), kd_expected, device=device)
+  )
+  # Other env untouched (cfg defaults).
+  assert torch.allclose(m.actuator_gainprm[1, ctrl_ids, 4], torch.tensor(5.0))
+  assert torch.allclose(m.actuator_gainprm[1, ctrl_ids, 6], torch.tensor(0.5))
+
+
+def test_dr_pd_gains_voltage_mode_rejected(device):
+  env = _scene_env(device, mode=DcMotorInputMode.VOLTAGE)
+  env.sim.expand_model_fields(("actuator_gainprm", "actuator_biasprm"))
+  with pytest.raises(ValueError, match="VOLTAGE"):
+    dr.pd_gains(
+      env,
+      env_ids=torch.tensor([0], device=device),
+      kp_range=(1.0, 1.0),
+      kd_range=(1.0, 1.0),
+      asset_cfg=SceneEntityCfg("robot"),
+    )
+
+
+def test_dr_effort_limits_writes_forcerange(device):
+  env = _scene_env(device)
+  robot = env.scene["robot"]
+  act = robot.actuators[0]
+  ctrl_ids = act.global_ctrl_ids
+  env.sim.expand_model_fields(
+    ("actuator_forcerange", "jnt_actfrcrange", "tendon_actfrcrange")
+  )
+
+  dr.effort_limits(
+    env,
+    env_ids=torch.tensor([0], device=device),
+    effort_limit_range=(123.0, 123.0),
+    asset_cfg=SceneEntityCfg("robot"),
+    operation="abs",
+  )
+
+  m = env.sim.model
+  n = len(ctrl_ids)
+  assert torch.allclose(
+    m.actuator_forcerange[0, ctrl_ids, 0],
+    torch.full((n,), -123.0, device=device),
+  )
+  assert torch.allclose(
+    m.actuator_forcerange[0, ctrl_ids, 1],
+    torch.full((n,), 123.0, device=device),
+  )
+  # Env 1 keeps the configured default of 50.
+  assert torch.allclose(m.actuator_forcerange[1, ctrl_ids, 1], torch.tensor(50.0))
+
+
+# Delay.
+
+
+def test_delay_position_mode(device):
+  """A 2-step lag should make position-mode torque reference step-0 target."""
+  entity, sim = initialize_entity(
+    _make_entity(
+      stiffness=10.0,
+      damping=0.0,
+      voltage_limit=1000.0,
+      delay_min_lag=2,
+      delay_max_lag=2,
+    ),
+    device,
+  )
+  zero = torch.zeros(1, 2, device=device)
+  entity.write_joint_state_to_sim(position=zero, velocity=zero)
+  targets = [
+    torch.tensor([[0.1, 0.0]], device=device),
+    torch.tensor([[0.3, 0.0]], device=device),
+    torch.tensor([[0.5, 0.0]], device=device),
+  ]
+  for p in targets:
+    entity.set_joint_position_target(p)
+    entity.set_joint_velocity_target(zero)
+    entity.set_joint_effort_target(zero)
+    entity.write_data_to_sim()
+    sim.forward()
+
+  v_adr = entity.indexing.joint_v_adr
+  # With lag=2 and three writes, the effective target is targets[0].
+  expected = K * 10.0 * targets[0][0] / R
+  assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], expected, atol=1e-4)
@@ -0,0 +1,182 @@
+"""Tests for BuiltinActuatorGroup (unified delayed + non-delayed)."""
+
+import mujoco
+import pytest
+import torch
+from conftest import get_test_device, load_fixture_xml
+
+from mjlab.actuator import (
+  BuiltinPositionActuatorCfg,
+  IdealPdActuatorCfg,
+)
+from mjlab.entity import Entity, EntityArticulationInfoCfg, EntityCfg
+from mjlab.sim.sim import Simulation, SimulationCfg
+
+ROBOT_XML = load_fixture_xml("floating_base_articulated")
+
+
+@pytest.fixture(scope="module")
+def device():
+  return get_test_device()
+
+
+def make_entity(actuator_cfgs, num_envs, device):
+  cfg = EntityCfg(
+    spec_fn=lambda: mujoco.MjSpec.from_string(ROBOT_XML),
+    articulation=EntityArticulationInfoCfg(actuators=actuator_cfgs),
+  )
+  entity = Entity(cfg)
+  model = entity.compile()
+  sim = Simulation(num_envs=num_envs, cfg=SimulationCfg(), model=model, device=device)
+  entity.initialize(model, sim.model, sim.data, device)
+  return entity, sim
+
+
+def test_fused_group_created(device):
+  """Delayed builtins with same config are fused into one group."""
+  cfg1 = BuiltinPositionActuatorCfg(
+    target_names_expr=("joint1",),
+    stiffness=50.0,
+    damping=5.0,
+    delay_min_lag=1,
+    delay_max_lag=3,
+  )
+  cfg2 = BuiltinPositionActuatorCfg(
+    target_names_expr=("joint2",),
+    stiffness=50.0,
+    damping=5.0,
+    delay_min_lag=1,
+    delay_max_lag=3,
+  )
+  entity, _ = make_entity((cfg1, cfg2), num_envs=2, device=device)
+
+  assert len(entity._builtin_group._delayed_groups) == 1
+  assert len(entity._custom_actuators) == 0
+
+
+def test_different_delay_configs_separate_groups(device):
+  """Delayed builtins with different delay configs get separate groups."""
+  cfg1 = BuiltinPositionActuatorCfg(
+    target_names_expr=("joint1",),
+    stiffness=50.0,
+    damping=5.0,
+    delay_min_lag=1,
+    delay_max_lag=3,
+  )
+  cfg2 = BuiltinPositionActuatorCfg(
+    target_names_expr=("joint2",),
+    stiffness=50.0,
+    damping=5.0,
+    delay_min_lag=5,
+    delay_max_lag=10,
+  )
+  entity, _ = make_entity((cfg1, cfg2), num_envs=2, device=device)
+
+  assert len(entity._builtin_group._delayed_groups) == 2
+
+
+def test_non_builtin_delayed_not_fused(device):
+  """Builtin delay fusion does not absorb ideal PD actuators.
+
+  The ideal PD actuator is fused separately into the PD group, not into the
+  builtin group and not left on the per-actuator custom path.
+  """
+  delayed_builtin = BuiltinPositionActuatorCfg(
+    target_names_expr=("joint1",),
+    stiffness=50.0,
+    damping=5.0,
+    delay_min_lag=1,
+    delay_max_lag=3,
+  )
+  custom = IdealPdActuatorCfg(
+    target_names_expr=("joint2",),
+    stiffness=50.0,
+    damping=5.0,
+    effort_limit=100.0,
+  )
+  entity, _ = make_entity((delayed_builtin, custom), num_envs=2, device=device)
+
+  assert len(entity._builtin_group._delayed_groups) == 1
+  assert len(entity._fused_actuator_group._groups) == 1
+  assert len(entity._custom_actuators) == 0
+
+
+def test_delayed_controls_written(device):
+  """Fused delayed builtins write controls to sim correctly."""
+  cfg = BuiltinPositionActuatorCfg(
+    target_names_expr=("joint.*",),
+    stiffness=50.0,
+    damping=5.0,
+    delay_min_lag=0,
+    delay_max_lag=0,
+  )
+  entity, sim = make_entity((cfg,), num_envs=1, device=device)
+
+  target = torch.tensor([[0.5, -0.3]], device=device)
+  entity.set_joint_position_target(target)
+  entity.write_data_to_sim()
+
+  # With lag=0, output equals input.
+  assert torch.allclose(sim.data.ctrl[0], target[0])
+
+
+def test_delay_actually_delays(device):
+  """With nonzero fixed lag, a new target takes lag steps to appear."""
+  cfg = BuiltinPositionActuatorCfg(
+    target_names_expr=("joint.*",),
+    stiffness=50.0,
+    damping=5.0,
+    delay_min_lag=2,
+    delay_max_lag=2,
+  )
+  entity, sim = make_entity((cfg,), num_envs=1, device=device)
+
+  target_a = torch.tensor([[1.0, 2.0]], device=device)
+  target_b = torch.tensor([[5.0, 6.0]], device=device)
+
+  # Fill the buffer with target_a so lag clamp doesn't mask the test.
+  for _ in range(3):
+    entity.set_joint_position_target(target_a)
+    entity.write_data_to_sim()
+  assert torch.allclose(sim.data.ctrl[0], target_a[0])
+
+  # Now send target_b. With lag=2, it takes 2 more steps to arrive.
+  entity.set_joint_position_target(target_b)
+  entity.write_data_to_sim()
+  assert torch.allclose(sim.data.ctrl[0], target_a[0])  # still old
+
+  entity.set_joint_position_target(target_b)
+  entity.write_data_to_sim()
+  assert torch.allclose(sim.data.ctrl[0], target_a[0])  # still old
+
+  entity.set_joint_position_target(target_b)
+  entity.write_data_to_sim()
+  assert torch.allclose(sim.data.ctrl[0], target_b[0])  # now arrives
+
+
+def test_reset_clears_buffer(device):
+  """Reset clears delay buffers so old values don't leak through."""
+  cfg = BuiltinPositionActuatorCfg(
+    target_names_expr=("joint.*",),
+    stiffness=50.0,
+    damping=5.0,
+    delay_min_lag=2,
+    delay_max_lag=2,
+  )
+  entity, sim = make_entity((cfg,), num_envs=1, device=device)
+
+  target_a = torch.tensor([[1.0, 2.0]], device=device)
+  target_b = torch.tensor([[5.0, 6.0]], device=device)
+
+  # Fill buffer with target_a.
+  for _ in range(3):
+    entity.set_joint_position_target(target_a)
+    entity.write_data_to_sim()
+  assert torch.allclose(sim.data.ctrl[0], target_a[0])
+
+  # Reset, then fill with target_b. Old target_a must not appear.
+  entity.reset(torch.tensor([0], device=device))
+  for _ in range(3):
+    entity.set_joint_position_target(target_b)
+    entity.write_data_to_sim()
+  assert torch.allclose(sim.data.ctrl[0], target_b[0])
@@ -0,0 +1,468 @@
+"""Tests for BuiltinPdActuator.
+
+Covers the unique surface of the actuator: paired <position>/<velocity>
+elements per target, joint/tendon-level actfrcrange sum-clamp, DR for both
+gains and effort limits, delay synchronization, and the ordering invariant
+that DR depends on.
+"""
+
+from unittest.mock import Mock
+
+import mujoco
+import pytest
+import torch
+from conftest import (
+  create_entity_with_actuator,
+  get_test_device,
+  initialize_entity,
+  load_fixture_xml,
+)
+
+from mjlab.actuator import BuiltinPdActuator, BuiltinPdActuatorCfg
+from mjlab.actuator.actuator import TransmissionType
+from mjlab.entity import Entity, EntityArticulationInfoCfg, EntityCfg
+from mjlab.envs.mdp import dr
+from mjlab.managers.scene_entity_config import SceneEntityCfg
+from mjlab.scene import Scene, SceneCfg
+from mjlab.sim.sim import Simulation, SimulationCfg
+
+ROBOT_XML = load_fixture_xml("floating_base_articulated")
+KP = 100.0
+KD = 10.0
+
+
+@pytest.fixture(scope="module")
+def device():
+  return get_test_device()
+
+
+def _make_entity(
+  *,
+  effort_limit: float | None = 50.0,
+  armature: float | None = None,
+  delay_max_lag: int = 0,
+  delay_min_lag: int = 0,
+  delay_hold_prob: float = 0.0,
+) -> Entity:
+  cfg = BuiltinPdActuatorCfg(
+    target_names_expr=("joint.*",),
+    stiffness=KP,
+    damping=KD,
+    effort_limit=effort_limit,
+    armature=armature,
+    delay_min_lag=delay_min_lag,
+    delay_max_lag=delay_max_lag,
+    delay_hold_prob=delay_hold_prob,
+  )
+  return create_entity_with_actuator(ROBOT_XML, cfg)
+
+
+def _at_rest_with_targets(
+  entity: Entity,
+  sim,
+  device: str,
+  pos_target: torch.Tensor,
+  vel_target: torch.Tensor,
+) -> None:
+  entity.write_joint_state_to_sim(
+    position=torch.zeros(1, 2, device=device),
+    velocity=torch.zeros(1, 2, device=device),
+  )
+  entity.set_joint_position_target(pos_target)
+  entity.set_joint_velocity_target(vel_target)
+  entity.set_joint_effort_target(torch.zeros(1, 2, device=device))
+  entity.write_data_to_sim()
+  sim.forward()
+
+
+# ---------------------------------------------------------------------------
+# Structural invariants
+# ---------------------------------------------------------------------------
+
+
+def test_two_ctrls_per_target_with_pos_then_vel_layout(device):
+  """Each target gets one <position> + one <velocity>, in halves."""
+  entity, sim = initialize_entity(_make_entity(), device)
+  act = entity.actuators[0]
+  assert isinstance(act, BuiltinPdActuator)
+
+  n = act.num_targets
+  assert n == len(act.target_names) == 2
+  assert len(act.ctrl_ids) == 2 * n
+  assert len(act.global_ctrl_ids) == 2 * n
+
+  names = [sim.mj_model.actuator(i).name for i in act.global_ctrl_ids.tolist()]
+  assert names[:n] == [f"{name}_pd_pos" for name in act.target_names]
+  assert names[n:] == [f"{name}_pd_vel" for name in act.target_names]
+
+
+def test_site_transmission_rejected():
+  with pytest.raises(ValueError, match="SITE"):
+    BuiltinPdActuatorCfg(
+      target_names_expr=("x",),
+      stiffness=1.0,
+      damping=1.0,
+      transmission_type=TransmissionType.SITE,
+    )
+
+
+def test_armature_applied_once(device):
+  """Joint armature must come from the position element only; double-applying
+  would silently double dof_armature."""
+  _, sim = initialize_entity(_make_entity(armature=0.7), device)
+  m = sim.mj_model
+  for jname in ("joint1", "joint2"):
+    dof_id = m.jnt_dofadr[m.joint(jname).id]
+    assert m.dof_armature[dof_id] == pytest.approx(0.7)
+
+
+# ---------------------------------------------------------------------------
+# Force computation
+# ---------------------------------------------------------------------------
+
+
+def test_position_only(device):
+  """Zero vel target: qfrc = kp * pos_target."""
+  entity, sim = initialize_entity(_make_entity(effort_limit=None), device)
+  pos = torch.tensor([[0.1, -0.05]], device=device)
+  _at_rest_with_targets(entity, sim, device, pos, torch.zeros(1, 2, device=device))
+  v_adr = entity.indexing.joint_v_adr
+  assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], KP * pos[0], atol=1e-4)
+
+
+def test_velocity_only(device):
+  """Zero pos target, joint at rest: qfrc = kd * vel_target."""
+  entity, sim = initialize_entity(_make_entity(effort_limit=None), device)
+  vel = torch.tensor([[0.3, -0.2]], device=device)
+  _at_rest_with_targets(entity, sim, device, torch.zeros(1, 2, device=device), vel)
+  v_adr = entity.indexing.joint_v_adr
+  assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], KD * vel[0], atol=1e-4)
+
+
+def test_pd_superposition(device):
+  """Both targets nonzero: qfrc = kp * pos_target + kd * vel_target."""
+  entity, sim = initialize_entity(_make_entity(effort_limit=None), device)
+  pos = torch.tensor([[0.1, -0.05]], device=device)
+  vel = torch.tensor([[0.2, -0.1]], device=device)
+  _at_rest_with_targets(entity, sim, device, pos, vel)
+  v_adr = entity.indexing.joint_v_adr
+  expected = KP * pos[0] + KD * vel[0]
+  assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], expected, atol=1e-4)
+
+
+def test_actfrcrange_sum_clamp(device):
+  """A pos error big enough to make kp*err exceed effort_limit must be
+  clamped at the joint, not allowed to ride through the unbounded element."""
+  entity, sim = initialize_entity(_make_entity(effort_limit=5.0), device)
+  # kp * 10.0 = 1000, well over the 5.0 clamp.
+  pos = torch.tensor([[10.0, 0.0]], device=device)
+  _at_rest_with_targets(entity, sim, device, pos, torch.zeros(1, 2, device=device))
+  v_adr = entity.indexing.joint_v_adr
+  qfrc = sim.data.qfrc_actuator[0, v_adr]
+  assert qfrc[0].item() == pytest.approx(5.0, abs=1e-4)
+  assert qfrc[1].item() == pytest.approx(0.0, abs=1e-4)
+
+
+def test_effort_limit_none_leaves_joint_unlimited(device):
+  """effort_limit=None: jnt_actfrclimited stays 0 on the targeted joints."""
+  _, sim = initialize_entity(_make_entity(effort_limit=None), device)
+  m = sim.mj_model
+  for jname in ("joint1", "joint2"):
+    jid = m.joint(jname).id
+    assert m.jnt_actfrclimited[jid] == 0
+
+
+def test_actuator_forcerange_not_set(device):
+  """We deliberately leave per-element forcerange unset; the limit lives on
+  the joint. Inspection of actuator_force[i] thus shows the unclamped value."""
+  entity, sim = initialize_entity(_make_entity(effort_limit=5.0), device)
+  m = sim.mj_model
+  for ctrl_id in entity.actuators[0].global_ctrl_ids.tolist():
+    assert m.actuator_forcelimited[ctrl_id] == 0
+
+
+# ---------------------------------------------------------------------------
+# Delay synchronization
+# ---------------------------------------------------------------------------
+
+
+def test_delay_syncs_pos_and_vel(device):
+  """The shared delay buffer must lag pos and vel together."""
+  entity, sim = initialize_entity(
+    _make_entity(effort_limit=None, delay_min_lag=2, delay_max_lag=2),
+    device,
+  )
+  pos_targets = [
+    torch.tensor([[0.1, 0.0]], device=device),
+    torch.tensor([[0.3, 0.0]], device=device),
+    torch.tensor([[0.5, 0.0]], device=device),
+  ]
+  vel_targets = [
+    torch.tensor([[1.0, 0.0]], device=device),
+    torch.tensor([[2.0, 0.0]], device=device),
+    torch.tensor([[3.0, 0.0]], device=device),
+  ]
+  entity.write_joint_state_to_sim(
+    position=torch.zeros(1, 2, device=device),
+    velocity=torch.zeros(1, 2, device=device),
+  )
+  for p, v in zip(pos_targets, vel_targets, strict=True):
+    entity.set_joint_position_target(p)
+    entity.set_joint_velocity_target(v)
+    entity.set_joint_effort_target(torch.zeros(1, 2, device=device))
+    entity.write_data_to_sim()
+    sim.forward()
+
+  v_adr = entity.indexing.joint_v_adr
+  # With lag=2, both halves should reference step-0 values.
+  expected = KP * pos_targets[0][0] + KD * vel_targets[0][0]
+  assert torch.allclose(sim.data.qfrc_actuator[0, v_adr], expected, atol=1e-4)
+
+
+def test_reset_clears_delay_buffer(device):
+  entity, _ = initialize_entity(_make_entity(delay_min_lag=1, delay_max_lag=3), device)
+  act = entity.actuators[0]
+  assert act._delay_buffer is not None
+  entity.set_joint_position_target(torch.full((1, 2), 0.5, device=device))
+  entity.set_joint_velocity_target(torch.zeros(1, 2, device=device))
+  entity.write_data_to_sim()
+
+  entity.reset(torch.tensor([0], device=device))
+  assert act._delay_buffer.current_lags[0] == 0
+
+
+# ---------------------------------------------------------------------------
+# Domain randomization
+# ---------------------------------------------------------------------------
+
+
+def _scene_env(device, transmission=TransmissionType.JOINT, num_envs=2):
+  """Build a real scene/sim with one BuiltinPd-driven entity for DR tests."""
+  if transmission == TransmissionType.JOINT:
+    xml = ROBOT_XML
+    targets = ("joint.*",)
+  else:
+    xml = load_fixture_xml("tendon_finger")
+    # tendon_finger ships with motor/position/velocity actuators; we need a
+    # bare spec so BuiltinPd can attach to the tendon without name clashes.
+    targets = ("finger_tendon",)
+
+  def spec_fn():
+    spec = mujoco.MjSpec.from_string(xml)
+    # Strip any pre-existing actuators so BuiltinPd's added elements own ctrl.
+    for a in list(spec.actuators):
+      spec.delete(a)
+    return spec
+
+  entity_cfg = EntityCfg(
+    spec_fn=spec_fn,
+    articulation=EntityArticulationInfoCfg(
+      actuators=(
+        BuiltinPdActuatorCfg(
+          target_names_expr=targets,
+          stiffness=KP,
+          damping=KD,
+          effort_limit=50.0,
+          transmission_type=transmission,
+        ),
+      )
+    ),
+  )
+  scene_cfg = SceneCfg(num_envs=num_envs, entities={"robot": entity_cfg})
+  scene = Scene(scene_cfg, device)
+  model = scene.compile()
+  sim = Simulation(num_envs=num_envs, cfg=SimulationCfg(), model=model, device=device)
+  scene.initialize(model, sim.model, sim.data)
+
+  env = Mock()
+  env.num_envs = num_envs
+  env.device = device
+  env.scene = {"robot": scene["robot"]}
+  env.sim = sim
+  return env
+
+
+def test_dr_pd_gains_scales_halves_independently(device):
+  env = _scene_env(device)
+  robot = env.scene["robot"]
+  act = robot.actuators[0]
+  assert isinstance(act, BuiltinPdActuator)
+  n = act.num_targets
+  pos_ids = act.global_ctrl_ids[:n]
+  vel_ids = act.global_ctrl_ids[n:]
+
+  # Expand fields so DR can write per-env.
+  env.sim.expand_model_fields(("actuator_gainprm", "actuator_biasprm"))
+
+  dr.pd_gains(
+    env,
+    env_ids=torch.tensor([0], device=device),
+    kp_range=(2.0, 2.0),
+    kd_range=(3.0, 3.0),
+    asset_cfg=SceneEntityCfg("robot"),
+    operation="scale",
+  )
+
+  m = env.sim.model
+  # Position half: gainprm[0] and biasprm[1] both scaled by kp=2, biasprm[2]
+  # must stay zero (no kd injection).
+  assert torch.allclose(
+    m.actuator_gainprm[0, pos_ids, 0],
+    torch.full((n,), 2.0 * KP, device=device),
+  )
+  assert torch.allclose(
+    m.actuator_biasprm[0, pos_ids, 1],
+    torch.full((n,), -2.0 * KP, device=device),
+  )
+  assert torch.allclose(
+    m.actuator_biasprm[0, pos_ids, 2], torch.zeros(n, device=device)
+  )
+  # Velocity half: gainprm[0] and biasprm[2] both scaled by kd=3, biasprm[1]
+  # stays zero (no kp injection).
+  assert torch.allclose(
+    m.actuator_gainprm[0, vel_ids, 0],
+    torch.full((n,), 3.0 * KD, device=device),
+  )
+  assert torch.allclose(
+    m.actuator_biasprm[0, vel_ids, 2],
+    torch.full((n,), -3.0 * KD, device=device),
+  )
+  assert torch.allclose(
+    m.actuator_biasprm[0, vel_ids, 1], torch.zeros(n, device=device)
+  )
+  # The other env must be untouched.
+  assert torch.allclose(m.actuator_gainprm[1, pos_ids, 0], torch.tensor(KP))
+  assert torch.allclose(m.actuator_gainprm[1, vel_ids, 0], torch.tensor(KD))
+
+
+def test_dr_pd_gains_abs_writes_correct_columns(device):
+  env = _scene_env(device)
+  robot = env.scene["robot"]
+  act = robot.actuators[0]
+  assert isinstance(act, BuiltinPdActuator)
+  n = act.num_targets
+  pos_ids = act.global_ctrl_ids[:n]
+  vel_ids = act.global_ctrl_ids[n:]
+  env.sim.expand_model_fields(("actuator_gainprm", "actuator_biasprm"))
+
+  dr.pd_gains(
+    env,
+    env_ids=torch.tensor([0, 1], device=device),
+    kp_range=(200.0, 200.0),
+    kd_range=(25.0, 25.0),
+    asset_cfg=SceneEntityCfg("robot"),
+    operation="abs",
+  )
+
+  m = env.sim.model
+  assert torch.allclose(
+    m.actuator_gainprm[:, pos_ids, 0],
+    torch.full((env.num_envs, n), 200.0, device=device),
+  )
+  assert torch.allclose(
+    m.actuator_biasprm[:, pos_ids, 1],
+    torch.full((env.num_envs, n), -200.0, device=device),
+  )
+  assert torch.allclose(
+    m.actuator_biasprm[:, pos_ids, 2],
+    torch.zeros(env.num_envs, n, device=device),
+  )
+  assert torch.allclose(
+    m.actuator_gainprm[:, vel_ids, 0],
+    torch.full((env.num_envs, n), 25.0, device=device),
+  )
+  assert torch.allclose(
+    m.actuator_biasprm[:, vel_ids, 2],
+    torch.full((env.num_envs, n), -25.0, device=device),
+  )
+
+
+def test_dr_effort_limits_writes_jnt_actfrcrange(device):
+  env = _scene_env(device, transmission=TransmissionType.JOINT)
+  robot = env.scene["robot"]
+  act = robot.actuators[0]
+  assert isinstance(act, BuiltinPdActuator)
+  env.sim.expand_model_fields(
+    ("actuator_forcerange", "jnt_actfrcrange", "tendon_actfrcrange")
+  )
+
+  joint_ids = robot.indexing.joint_ids[act.target_ids]
+  pre_forcerange = env.sim.model.actuator_forcerange.clone()
+
+  dr.effort_limits(
+    env,
+    env_ids=torch.tensor([0], device=device),
+    effort_limit_range=(123.0, 123.0),
+    asset_cfg=SceneEntityCfg("robot"),
+    operation="abs",
+  )
+
+  m = env.sim.model
+  # The joint sum-clamp was rewritten on env 0 only.
+  assert torch.allclose(
+    m.jnt_actfrcrange[0, joint_ids],
+    torch.tensor([[-123.0, 123.0]] * len(joint_ids), device=device),
+  )
+  assert torch.allclose(
+    m.jnt_actfrcrange[1, joint_ids],
+    torch.tensor([[-50.0, 50.0]] * len(joint_ids), device=device),
+  )
+  # Per-element actuator_forcerange must be untouched for BuiltinPd: that
+  # field belongs to the existing single-element actuator semantic.
+  assert torch.allclose(m.actuator_forcerange, pre_forcerange)
+
+
+def test_dr_effort_limits_scale_multiplies_default(device):
+  """``scale`` multiplies the configured ``effort_limit`` (50.0) by the sample."""
+  env = _scene_env(device, transmission=TransmissionType.JOINT)
+  robot = env.scene["robot"]
+  act = robot.actuators[0]
+  assert isinstance(act, BuiltinPdActuator)
+  env.sim.expand_model_fields(
+    ("actuator_forcerange", "jnt_actfrcrange", "tendon_actfrcrange")
+  )
+  joint_ids = robot.indexing.joint_ids[act.target_ids]
+
+  dr.effort_limits(
+    env,
+    env_ids=torch.tensor([0], device=device),
+    effort_limit_range=(2.0, 2.0),
+    asset_cfg=SceneEntityCfg("robot"),
+    operation="scale",
+  )
+
+  m = env.sim.model
+  # Default is [-50, 50], scaled by 2 -> [-100, 100].
+  assert torch.allclose(
+    m.jnt_actfrcrange[0, joint_ids],
+    torch.tensor([[-100.0, 100.0]] * len(joint_ids), device=device),
+  )
+
+
+def test_dr_effort_limits_writes_tendon_actfrcrange(device):
+  env = _scene_env(device, transmission=TransmissionType.TENDON)
+  robot = env.scene["robot"]
+  act = robot.actuators[0]
+  assert isinstance(act, BuiltinPdActuator)
+  env.sim.expand_model_fields(
+    ("actuator_forcerange", "jnt_actfrcrange", "tendon_actfrcrange")
+  )
+  tendon_ids = robot.indexing.tendon_ids[act.target_ids]
+
+  dr.effort_limits(
+    env,
+    env_ids=torch.tensor([0], device=device),
+    effort_limit_range=(77.0, 77.0),
+    asset_cfg=SceneEntityCfg("robot"),
+    operation="abs",
+  )
+
+  m = env.sim.model
+  assert torch.allclose(
+    m.tendon_actfrcrange[0, tendon_ids],
+    torch.tensor([[-77.0, 77.0]] * len(tendon_ids), device=device),
+  )
+  assert torch.allclose(
+    m.tendon_actfrcrange[1, tendon_ids],
+    torch.tensor([[-50.0, 50.0]] * len(tendon_ids), device=device),
+  )
@@ -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"
@@ -0,0 +1,439 @@
+"""Tests for camera_sensor.py and the unified sense pipeline."""
+
+from __future__ import annotations
+
+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 import CameraSensorCfg, CameraSensorData
+from mjlab.sim.sim import Simulation, SimulationCfg
+
+
+@pytest.fixture(scope="module")
+def device():
+  return get_test_device()
+
+
+SCENE_WITH_CAMERA_XML = """
+  <mujoco>
+    <worldbody>
+      <light pos="0 0 3" dir="0 0 -1"/>
+      <geom name="floor" type="plane" size="10 10 0.1" pos="0 0 0"
+            rgba="0.5 0.5 0.5 1"/>
+      <geom name="red_box" type="box" size="0.5 0.5 0.5" pos="0 0 0.5"
+            rgba="1 0 0 1"/>
+      <camera name="overhead_cam" pos="0 0 3" quat="1 0 0 0"
+              fovy="45" resolution="32 24"/>
+    </worldbody>
+  </mujoco>
+"""
+
+
+def _make_scene_and_sim(
+  device: str,
+  sensors: tuple,
+  xml: str = SCENE_WITH_CAMERA_XML,
+  num_envs: int = 2,
+) -> tuple[Scene, Simulation]:
+  entity_cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(xml))
+  scene_cfg = SceneCfg(
+    num_envs=num_envs,
+    env_spacing=5.0,
+    entities={"world": entity_cfg},
+    sensors=sensors,
+  )
+  scene = Scene(scene_cfg, device)
+  model = scene.compile()
+  sim = Simulation(num_envs=num_envs, cfg=SimulationCfg(), model=model, device=device)
+  scene.initialize(sim.mj_model, sim.model, sim.data)
+  if scene.sensor_context is not None:
+    sim.set_sensor_context(scene.sensor_context)
+  return scene, sim
+
+
+def test_camera_rgb_shape(device):
+  """RGB data has correct shape [num_envs, height, width, 3]."""
+  cam_cfg = CameraSensorCfg(
+    name="test_cam",
+    camera_name="world/overhead_cam",
+    width=32,
+    height=24,
+    data_types=("rgb",),
+  )
+  scene, sim = _make_scene_and_sim(device, sensors=(cam_cfg,))
+
+  sim.forward()
+  sim.sense()
+  sensor = scene["test_cam"]
+  data = sensor.data
+
+  assert isinstance(data, CameraSensorData)
+  assert data.rgb is not None
+  assert data.rgb.shape == (2, 24, 32, 3)
+  assert data.rgb.dtype == torch.uint8
+  assert data.depth is None
+
+
+def test_camera_depth_shape(device):
+  """Depth data has correct shape [num_envs, height, width]."""
+  cam_cfg = CameraSensorCfg(
+    name="test_cam",
+    camera_name="world/overhead_cam",
+    width=32,
+    height=24,
+    data_types=("depth",),
+  )
+  scene, sim = _make_scene_and_sim(device, sensors=(cam_cfg,))
+
+  sim.forward()
+  sim.sense()
+  data = scene["test_cam"].data
+
+  assert isinstance(data, CameraSensorData)
+  assert data.depth is not None
+  assert data.depth.shape == (2, 24, 32, 1)
+  assert data.depth.dtype == torch.float32
+  assert data.rgb is None
+
+
+def test_camera_rgb_and_depth(device):
+  """Both RGB and depth can be requested simultaneously."""
+  cam_cfg = CameraSensorCfg(
+    name="test_cam",
+    camera_name="world/overhead_cam",
+    width=32,
+    height=24,
+    data_types=("rgb", "depth"),
+  )
+  scene, sim = _make_scene_and_sim(device, sensors=(cam_cfg,))
+
+  sim.forward()
+  sim.sense()
+  data = scene["test_cam"].data
+
+  assert data.rgb is not None
+  assert data.depth is not None
+  assert data.rgb.shape == (2, 24, 32, 3)
+  assert data.depth.shape == (2, 24, 32, 1)
+
+
+def test_camera_create_new(device):
+  """Camera sensor can create a new camera (not wrap existing)."""
+  # XML without any cameras.
+  xml = """
+    <mujoco>
+      <worldbody>
+        <light pos="0 0 3" dir="0 0 -1"/>
+        <geom name="floor" type="plane" size="10 10 0.1"
+              rgba="0.5 0.5 0.5 1"/>
+      </worldbody>
+    </mujoco>
+  """
+  cam_cfg = CameraSensorCfg(
+    name="my_camera",
+    pos=(0.0, 0.0, 3.0),
+    quat=(1.0, 0.0, 0.0, 0.0),
+    fovy=45.0,
+    width=16,
+    height=12,
+    data_types=("rgb",),
+  )
+  scene, sim = _make_scene_and_sim(device, sensors=(cam_cfg,), xml=xml)
+
+  sim.forward()
+  sim.sense()
+  data = scene["my_camera"].data
+
+  assert data.rgb is not None
+  assert data.rgb.shape == (2, 12, 16, 3)
+
+
+def test_camera_rgb_not_all_zeros(device):
+  """RGB data should contain non-zero values when scene has objects."""
+  cam_cfg = CameraSensorCfg(
+    name="test_cam",
+    camera_name="world/overhead_cam",
+    width=32,
+    height=24,
+    data_types=("rgb",),
+  )
+  scene, sim = _make_scene_and_sim(device, sensors=(cam_cfg,))
+
+  sim.forward()
+  sim.sense()
+  data = scene["test_cam"].data
+
+  assert data.rgb is not None
+  # At least some pixels should be non-zero (scene has colored objects).
+  assert data.rgb.any(), "RGB data is all zeros - rendering may have failed"
+
+
+def test_camera_wrap_existing_overrides_fovy(device):
+  """Wrapping an existing camera can override its fovy."""
+  cam_cfg = CameraSensorCfg(
+    name="test_cam",
+    camera_name="world/overhead_cam",
+    width=32,
+    height=24,
+    fovy=90.0,  # Override the camera's fovy
+    data_types=("rgb",),
+  )
+  scene, sim = _make_scene_and_sim(device, sensors=(cam_cfg,))
+
+  # Verify the camera's fovy was updated.
+  cam_id = sim.mj_model.camera("world/overhead_cam").id
+  assert sim.mj_model.cam_fovy[cam_id] == pytest.approx(90.0)
+
+  sim.forward()
+  sim.sense()
+  data = scene["test_cam"].data
+  assert data.rgb is not None
+
+
+@pytest.mark.skipif(
+  not torch.cuda.is_available(), reason="CUDA required for sense graph"
+)
+def test_sense_graph_captured(device):
+  """Verify sense_graph is captured on CUDA."""
+  cam_cfg = CameraSensorCfg(
+    name="test_cam",
+    camera_name="world/overhead_cam",
+    width=16,
+    height=12,
+    data_types=("rgb",),
+  )
+  scene, sim = _make_scene_and_sim(device, sensors=(cam_cfg,))
+
+  # On CUDA with mempool, sense_graph should be captured.
+  if sim.use_cuda_graph:
+    assert sim.sense_graph is not None
+
+
+def test_camera_clone_data(device):
+  """clone_data=True returns independent tensors."""
+  cam_cfg = CameraSensorCfg(
+    name="test_cam",
+    camera_name="world/overhead_cam",
+    width=16,
+    height=12,
+    data_types=("rgb",),
+    clone_data=True,
+  )
+  scene, sim = _make_scene_and_sim(device, sensors=(cam_cfg,))
+
+  sim.forward()
+  sim.sense()
+  data1 = scene["test_cam"].data
+  rgb1 = data1.rgb
+
+  # Invalidate cache and re-read.
+  scene["test_cam"]._invalidate_cache()
+  data2 = scene["test_cam"].data
+  rgb2 = data2.rgb
+
+  # With clone_data=True, the tensors should be different objects.
+  assert rgb1 is not None and rgb2 is not None
+  assert rgb1.data_ptr() != rgb2.data_ptr()
+
+
+def test_camera_sensor_context_created(device):
+  """Scene creates SensorContext when camera sensors are present."""
+  cam_cfg = CameraSensorCfg(
+    name="test_cam",
+    camera_name="world/overhead_cam",
+    width=16,
+    height=12,
+    data_types=("rgb",),
+  )
+  scene, _ = _make_scene_and_sim(device, sensors=(cam_cfg,))
+
+  assert scene.sensor_context is not None
+  assert scene.sensor_context.has_cameras
+
+
+INTRINSIC_CAM_XML = """
+  <mujoco>
+    <worldbody>
+      <light pos="0 0 5" dir="0 0 -1"/>
+      <geom name="floor" type="plane" size="10 10 0.1" pos="0 0 0"
+            rgba="0.5 0.5 0.5 1"/>
+      <geom name="small_box" type="box" size="0.1 0.1 0.1" pos="0 0 0.1"
+            rgba="1 0 0 1"/>
+      <camera name="intrinsic_cam" pos="0 0 5" quat="1 0 0 0"
+              resolution="32 24" sensorsize="0.01 0.01"
+              focal="0.005 0.005"/>
+    </worldbody>
+  </mujoco>
+"""
+
+
+def test_expand_cam_intrinsic_disables_precomputed_rays(device):
+  """Expanding cam_intrinsic disables precomputed rays on the render context."""
+  cam_cfg = CameraSensorCfg(
+    name="test_cam",
+    camera_name="world/intrinsic_cam",
+    width=32,
+    height=24,
+    data_types=("rgb",),
+  )
+  scene, sim = _make_scene_and_sim(device, sensors=(cam_cfg,), xml=INTRINSIC_CAM_XML)
+  assert scene.sensor_context is not None
+
+  # Before expansion: precomputed rays enabled.
+  rc = scene.sensor_context.render_context
+  assert rc.use_precomputed_rays is True
+
+  sim.expand_model_fields(("cam_intrinsic",))
+
+  # After expansion: precomputed rays disabled.
+  rc = scene.sensor_context.render_context
+  assert rc.use_precomputed_rays is False
+
+
+def test_cam_intrinsic_dr_changes_rendered_image(device):
+  """Per-env cam_intrinsic DR produces different images across envs."""
+  cam_cfg = CameraSensorCfg(
+    name="test_cam",
+    camera_name="world/intrinsic_cam",
+    width=32,
+    height=24,
+    data_types=("rgb",),
+  )
+  scene, sim = _make_scene_and_sim(
+    device, sensors=(cam_cfg,), xml=INTRINSIC_CAM_XML, num_envs=2
+  )
+  sim.expand_model_fields(("cam_intrinsic",))
+
+  # Render baseline.
+  sim.forward()
+  sim.sense()
+  baseline = scene["test_cam"].data.rgb.clone()
+
+  # Set env 0 to a very different focal length (5x zoom).
+  cam_id = sim.mj_model.camera("world/intrinsic_cam").id
+  sim.model.cam_intrinsic[0, cam_id, :2] *= 5.0
+
+  sim.forward()
+  sim.sense()
+  after = scene["test_cam"].data.rgb
+
+  # Env 0 should differ (zoomed in), env 1 should be unchanged.
+  assert not torch.equal(after[0], baseline[0]), (
+    "Env 0 image unchanged after cam_intrinsic modification"
+  )
+  assert torch.equal(after[1], baseline[1]), "Env 1 image changed unexpectedly"
+
+
+def test_camera_segmentation_shape(device):
+  """Segmentation data has correct shape [num_envs, height, width, 2]."""
+  cam_cfg = CameraSensorCfg(
+    name="test_cam",
+    camera_name="world/overhead_cam",
+    width=32,
+    height=24,
+    data_types=("segmentation",),
+  )
+  scene, sim = _make_scene_and_sim(device, sensors=(cam_cfg,))
+
+  sim.forward()
+  sim.sense()
+  data = scene["test_cam"].data
+
+  assert isinstance(data, CameraSensorData)
+  assert data.segmentation is not None
+  assert data.segmentation.shape == (2, 24, 32, 2)
+  assert data.segmentation.dtype == torch.int32
+  assert data.rgb is None
+  assert data.depth is None
+
+
+def test_camera_segmentation_content(device):
+  """Segmentation contains typed MuJoCo ids for geom hits and background."""
+  cam_cfg = CameraSensorCfg(
+    name="test_cam",
+    camera_name="world/overhead_cam",
+    width=32,
+    height=24,
+    data_types=("segmentation",),
+  )
+  scene, sim = _make_scene_and_sim(device, sensors=(cam_cfg,))
+
+  sim.forward()
+  sim.sense()
+  seg = scene["test_cam"].data.segmentation
+  assert seg is not None
+
+  seg_np = seg.cpu().numpy()
+  obj_ids = seg_np[..., 0]
+  obj_types = seg_np[..., 1]
+
+  geom_mask = obj_types == int(mujoco.mjtObj.mjOBJ_GEOM)
+  background_mask = obj_types == -1
+
+  assert geom_mask.any(), "Expected at least one geom hit"
+  assert (obj_ids[geom_mask] >= 0).all()
+  assert len(set(obj_ids[geom_mask].flatten().tolist())) > 1
+  if background_mask.any():
+    assert (obj_ids[background_mask] == -1).all()
+
+
+def test_camera_rgb_depth_segmentation(device):
+  """All three data types can be requested simultaneously."""
+  cam_cfg = CameraSensorCfg(
+    name="test_cam",
+    camera_name="world/overhead_cam",
+    width=32,
+    height=24,
+    data_types=("rgb", "depth", "segmentation"),
+  )
+  scene, sim = _make_scene_and_sim(device, sensors=(cam_cfg,))
+
+  sim.forward()
+  sim.sense()
+  data = scene["test_cam"].data
+
+  assert data.rgb is not None
+  assert data.depth is not None
+  assert data.segmentation is not None
+  assert data.rgb.shape == (2, 24, 32, 3)
+  assert data.depth.shape == (2, 24, 32, 1)
+  assert data.segmentation.shape == (2, 24, 32, 2)
+
+
+def test_camera_create_on_parent_body(device):
+  """Camera sensor can create a new camera attached to a body."""
+  xml = """
+    <mujoco>
+      <worldbody>
+        <light pos="0 0 3" dir="0 0 -1"/>
+        <geom name="floor" type="plane" size="10 10 0.1"
+              rgba="0.5 0.5 0.5 1"/>
+        <body name="arm_link" pos="0 0 1">
+          <geom type="box" size="0.1 0.1 0.1" rgba="1 0 0 1"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  """
+  cam_cfg = CameraSensorCfg(
+    name="wrist_cam",
+    parent_body="world/arm_link",
+    pos=(0.0, 0.0, 0.2),
+    quat=(1.0, 0.0, 0.0, 0.0),
+    fovy=45.0,
+    width=16,
+    height=12,
+    data_types=("rgb",),
+  )
+  scene, sim = _make_scene_and_sim(device, sensors=(cam_cfg,), xml=xml)
+
+  sim.forward()
+  sim.sense()
+  data = scene["wrist_cam"].data
+
+  assert isinstance(data, CameraSensorData)
+  assert data.rgb is not None
+  assert data.rgb.shape == (2, 12, 16, 3)
@@ -0,0 +1,268 @@
+"""Tests for CircularBuffer."""
+
+import pytest
+import torch
+from conftest import get_test_device
+
+from mjlab.utils.buffers import CircularBuffer
+
+
+@pytest.fixture
+def device():
+  """Test device fixture."""
+  return get_test_device()
+
+
+def test_circular_buffer_basic_append(device):
+  """Basic append and chronological retrieval (oldest -> newest)."""
+  buffer = CircularBuffer(max_len=3, batch_size=2, device=device)
+
+  buffer.append(torch.tensor([[1.0, 2.0], [3.0, 4.0]], device=device))
+  buffer.append(torch.tensor([[5.0, 6.0], [7.0, 8.0]], device=device))
+  buffer.append(torch.tensor([[9.0, 10.0], [11.0, 12.0]], device=device))
+
+  result = buffer.buffer
+  assert result.shape == (2, 3, 2)
+  # Oldest to newest.
+  assert torch.allclose(
+    result[0], torch.tensor([[1.0, 2.0], [5.0, 6.0], [9.0, 10.0]], device=device)
+  )
+  assert torch.allclose(
+    result[1], torch.tensor([[3.0, 4.0], [7.0, 8.0], [11.0, 12.0]], device=device)
+  )
+
+
+def test_circular_buffer_overwrite(device):
+  """Overwrites oldest once capacity reached."""
+  buffer = CircularBuffer(max_len=2, batch_size=1, device=device)
+
+  buffer.append(torch.tensor([[1.0]], device=device))
+  buffer.append(torch.tensor([[2.0]], device=device))
+  buffer.append(torch.tensor([[3.0]], device=device))  # Overwrites first.
+
+  result = buffer.buffer
+  assert result.shape == (1, 2, 1)
+  assert torch.allclose(result[0], torch.tensor([[2.0], [3.0]], device=device))
+
+
+def test_circular_buffer_reset_single_batch(device):
+  """Reset clears values and counters for specified batch rows."""
+  buffer = CircularBuffer(max_len=2, batch_size=3, device=device)
+
+  buffer.append(torch.tensor([[1.0], [2.0], [3.0]], device=device))
+  buffer.append(torch.tensor([[4.0], [5.0], [6.0]], device=device))
+
+  # Reset only batch index 1.
+  buffer.reset(batch_ids=torch.tensor([1], device=device))
+
+  result = buffer.buffer
+  # Oldest-to-newest for each batch row.
+  assert result[0, 0, 0] == 1.0
+  assert result[1, 0, 0] == 0.0  # Reset backfilled to zeros for that row.
+  assert result[2, 0, 0] == 3.0
+
+  # current_length reflects reset: rows 0 and 2 had 2 pushes, row 1 is 0.
+  cl = buffer.current_length
+  assert torch.equal(cl.cpu(), torch.tensor([2, 0, 2]))
+
+
+def test_circular_buffer_first_append_fills(device):
+  """First append back-fills whole history for each batch row."""
+  buffer = CircularBuffer(max_len=3, batch_size=2, device=device)
+  buffer.append(torch.tensor([[1.0], [2.0]], device=device))
+
+  result = buffer.buffer
+  assert torch.allclose(result[0], torch.tensor([[1.0], [1.0], [1.0]], device=device))
+  assert torch.allclose(result[1], torch.tensor([[2.0], [2.0], [2.0]], device=device))
+
+  # And current_length reflects valid frames so far.
+  cl = buffer.current_length
+  assert torch.equal(cl.cpu(), torch.tensor([1, 1]))
+
+
+def test_current_length_counts_and_clamps(device):
+  """current_length counts per-batch valid frames and clamps to max_len."""
+  buffer = CircularBuffer(max_len=4, batch_size=3, device=device)
+
+  # Two appends -> length 2 everywhere.
+  for _ in range(2):
+    buffer.append(torch.arange(3, dtype=torch.float32, device=device).unsqueeze(-1))
+
+  assert torch.equal(buffer.current_length.cpu(), torch.tensor([2, 2, 2]))
+
+  # Reset middle row -> it becomes 0.
+  buffer.reset(batch_ids=[1])
+  assert torch.equal(buffer.current_length.cpu(), torch.tensor([2, 0, 2]))
+
+  # One more append -> rows [0,2] become 3; row 1 becomes 1.
+  buffer.append(torch.arange(3, dtype=torch.float32, device=device).unsqueeze(-1))
+  assert torch.equal(buffer.current_length.cpu(), torch.tensor([3, 1, 3]))
+
+  # Fill beyond capacity -> clamp to max_len.
+  for _ in range(5):
+    buffer.append(torch.arange(3, dtype=torch.float32, device=device).unsqueeze(-1))
+  assert torch.equal(buffer.current_length.cpu(), torch.tensor([4, 4, 4]))
+
+
+def test_reset_all_none_path(device):
+  """reset(None) zeros the entire buffer and counters without indexing."""
+  buffer = CircularBuffer(max_len=3, batch_size=2, device=device)
+  buffer.append(torch.tensor([[1.0], [2.0]], device=device))
+  buffer.append(torch.tensor([[3.0], [4.0]], device=device))
+
+  buffer.reset()  # None -> reset all.
+
+  # Counters are zero.
+  assert torch.equal(buffer.current_length.cpu(), torch.tensor([0, 0]))
+  # Buffer zeros (safe to read even after reset because storage exists).
+  result = buffer.buffer
+  assert torch.count_nonzero(result) == 0
+
+
+def test_getitem_lifo_and_clamp(device):
+  """__getitem__ returns lagged frames per-batch (LIFO), clamping when needed."""
+  buffer = CircularBuffer(max_len=3, batch_size=2, device=device)
+
+  buffer.append(torch.tensor([[1.0], [10.0]], device=device))  # t0
+  buffer.append(torch.tensor([[2.0], [20.0]], device=device))  # t1
+  buffer.append(torch.tensor([[3.0], [30.0]], device=device))  # t2
+
+  # Lag 0 for batch 0 (-> 3), lag 2 for batch 1 (-> oldest 10).
+  out = buffer[torch.tensor([0, 2], device=device)]
+  assert torch.allclose(out, torch.tensor([[3.0], [10.0]], device=device))
+
+  # Clamp: huge lag for batch 0 -> oldest (1), lag 1 for batch 1 -> 20.
+  out = buffer[torch.tensor([99, 1], device=device)]
+  assert torch.allclose(out, torch.tensor([[1.0], [20.0]], device=device))
+
+
+def test_backfill_after_per_batch_reset(device):
+  """After resetting a row, the next append back-fills its entire history for that row."""
+  buffer = CircularBuffer(max_len=3, batch_size=2, device=device)
+
+  buffer.append(torch.tensor([[1.0], [10.0]], device=device))  # t0
+  buffer.append(torch.tensor([[2.0], [20.0]], device=device))  # t1
+
+  # Reset only batch row 1; row 0 remains with 2 valid frames.
+  buffer.reset(batch_ids=[1])
+  assert torch.equal(buffer.current_length.cpu(), torch.tensor([2, 0]))
+
+  # Next append: row 0 gets new value; row 1 is "first push" -> back-filled.
+  buffer.append(torch.tensor([[3.0], [99.0]], device=device))  # t2
+
+  hist = buffer.buffer  # shape (2, 3, 1)
+  # Row 0 keeps real chronology [1, 2, 3].
+  assert torch.allclose(
+    hist[0].squeeze(-1), torch.tensor([1.0, 2.0, 3.0], device=device)
+  )
+  # Row 1 is back-filled to all 99s.
+  assert torch.allclose(
+    hist[1].squeeze(-1), torch.tensor([99.0, 99.0, 99.0], device=device)
+  )
+
+
+def test_backfill_multidimensional_data(device):
+  """First-push backfill works for >2D data whose trailing dims differ from batch.
+
+  Regression for #1045: the backfill broadcast hard-coded a 3D condition shape,
+  which misaligned the batch axis with a trailing axis when num_targets != batch.
+  Here batch_size=4 but the trailing axis is 2, so the buggy broadcast raised a
+  4-vs-2 mismatch.
+  """
+  batch_size, num_targets, num_channels = 4, 2, 3
+  buffer = CircularBuffer(max_len=3, batch_size=batch_size, device=device)
+
+  data = torch.arange(
+    batch_size * num_targets * num_channels, dtype=torch.float32, device=device
+  ).reshape(batch_size, num_targets, num_channels)
+  buffer.append(data)
+
+  hist = buffer.buffer  # (batch, max_len, num_targets, num_channels)
+  assert hist.shape == (batch_size, 3, num_targets, num_channels)
+  # First push backfills every history slot with the same frame per batch row.
+  for t in range(3):
+    assert torch.allclose(hist[:, t], data)
+
+
+def test_errors_and_types(device):
+  """Error paths: wrong batch, pre-append access, and bad key size."""
+  buffer = CircularBuffer(max_len=2, batch_size=2, device=device)
+
+  # Wrong batch size on append.
+  with pytest.raises(ValueError):
+    buffer.append(torch.tensor([[1.0]], device=device))  # batch_size=1 wrong
+
+  # buffer property before first append.
+  with pytest.raises(RuntimeError):
+    _ = CircularBuffer(max_len=1, batch_size=1, device=device).buffer
+
+  # __getitem__ before any valid pushes.
+  with pytest.raises(RuntimeError):
+    _ = buffer[torch.tensor([0, 0], device=device)]
+
+  # Now append once so storage exists and counters > 0.
+  buffer.append(torch.tensor([[1.0], [2.0]], device=device))
+
+  # __getitem__ with wrong key length.
+  with pytest.raises(ValueError):
+    _ = buffer[torch.tensor([0], device=device)]
+
+
+def test_dtype_and_device_preserved(device):
+  """Buffer preserves dtype and device."""
+  buffer = CircularBuffer(max_len=2, batch_size=2, device=device)
+  x = torch.tensor([[1.0], [2.0]], dtype=torch.float32, device=device)
+  buffer.append(x)
+  assert buffer.buffer.dtype == torch.float32
+  assert buffer.buffer.device.type == torch.device(device).type
+
+
+def test_backfill_only_touches_given_rows(device):
+  """backfill writes one frame into all slots of given rows, no pointer advance."""
+  buffer = CircularBuffer(max_len=3, batch_size=3, device=device)
+  buffer.append(torch.tensor([[1.0], [10.0], [100.0]], device=device))
+  buffer.append(torch.tensor([[2.0], [20.0], [200.0]], device=device))
+  buffer.append(torch.tensor([[3.0], [30.0], [300.0]], device=device))
+
+  buffer.reset(batch_ids=torch.tensor([1], device=device))
+  data = torch.tensor([[-1.0], [99.0], [-1.0]], device=device)
+  buffer.backfill(data, torch.tensor([1], device=device))
+
+  result = buffer.buffer
+  assert torch.allclose(
+    result[0].flatten(), torch.tensor([1.0, 2.0, 3.0], device=device)
+  )
+  assert torch.allclose(
+    result[1].flatten(), torch.tensor([99.0, 99.0, 99.0], device=device)
+  )
+  assert torch.allclose(
+    result[2].flatten(), torch.tensor([100.0, 200.0, 300.0], device=device)
+  )
+  assert buffer.current_length.tolist() == [3, 1, 3]
+
+  # Subsequent appends treat the backfilled row normally.
+  buffer.append(torch.tensor([[4.0], [40.0], [400.0]], device=device))
+  result = buffer.buffer
+  assert torch.allclose(
+    result[1].flatten(), torch.tensor([99.0, 99.0, 40.0], device=device)
+  )
+  assert buffer.current_length.tolist() == [3, 2, 3]
+
+
+def test_backfill_uninitialized_raises(device):
+  buffer = CircularBuffer(max_len=2, batch_size=2, device=device)
+  with pytest.raises(RuntimeError, match="not initialized"):
+    buffer.backfill(torch.zeros(2, 1, device=device), torch.tensor([0], device=device))
+
+
+def test_getitem_lag_clamps_to_oldest_after_wrap(device):
+  buffer = CircularBuffer(max_len=2, batch_size=1, device=device)
+  for v in [1.0, 2.0, 3.0]:
+    buffer.append(torch.tensor([[v]], device=device))
+
+  # Retained frames: [2, 3] (num_pushes exceeds max_len).
+  assert buffer[torch.tensor([0], device=device)].flatten().item() == 3.0
+  assert buffer[torch.tensor([1], device=device)].flatten().item() == 2.0
+  # A lag beyond the retained history clamps to the oldest frame instead of
+  # wrapping around to a newer one.
+  assert buffer[torch.tensor([2], device=device)].flatten().item() == 2.0
@@ -0,0 +1,179 @@
+"""Tests for command manager."""
+
+from dataclasses import dataclass
+from unittest.mock import Mock
+
+import pytest
+import torch
+from conftest import get_test_device
+
+from mjlab.envs import ManagerBasedRlEnv
+from mjlab.managers.command_manager import CommandTerm, CommandTermCfg
+from mjlab.tasks.cartpole.cartpole_env_cfg import cartpole_balance_env_cfg
+from mjlab.tasks.tracking.mdp.commands import MotionCommand
+
+
+@pytest.fixture(scope="module")
+def device():
+  return get_test_device()
+
+
+class CounterCommand(CommandTerm):
+  """Stateful command term: a per-env counter ticked by _update_command."""
+
+  def __init__(self, cfg, env):
+    super().__init__(cfg, env)
+    self.ticks = torch.zeros(self.num_envs, dtype=torch.long, device=self.device)
+
+  @property
+  def command(self) -> torch.Tensor:
+    return self.ticks.unsqueeze(-1).float()
+
+  def _update_metrics(self) -> None:
+    pass
+
+  def _resample_command(self, env_ids: torch.Tensor) -> None:
+    self.ticks[env_ids] = 0
+
+  def _update_command(self, env_ids: torch.Tensor | None = None) -> None:
+    if env_ids is None:
+      self.ticks += 1
+    else:
+      self.ticks[env_ids] += 1
+
+
+@dataclass(kw_only=True)
+class CounterCommandCfg(CommandTermCfg):
+  resampling_time_range: tuple[float, float] = (1e9, 1e9)
+
+  def build(self, env) -> CounterCommand:
+    return CounterCommand(self, env)
+
+
+@pytest.fixture
+def counter_env(device):
+  cfg = cartpole_balance_env_cfg()
+  cfg.scene.num_envs = 4
+  cfg.commands = {"counter": CounterCommandCfg()}
+  env = ManagerBasedRlEnv(cfg=cfg, device=device)
+  yield env
+  env.close()
+
+
+def test_partial_reset_does_not_advance_other_envs(counter_env):
+  env = counter_env
+  term = env.command_manager.get_term("counter")
+  assert isinstance(term, CounterCommand)
+
+  # Full reset: every env is resampled (counter zeroed) then ticked once.
+  env.reset()
+  assert term.ticks.tolist() == [1, 1, 1, 1]
+
+  action = torch.zeros((env.num_envs, 1), device=env.device)
+  env.step(action)
+  assert term.ticks.tolist() == [2, 2, 2, 2]
+
+  # Partial reset: only env 1 is resampled and ticked. Before the fix for
+  # issue #1138 the other envs advanced too (to 3).
+  env.reset(env_ids=torch.tensor([1], dtype=torch.int64, device=env.device))
+  assert term.ticks.tolist() == [2, 1, 2, 2]
+
+  # The next step advances everyone by exactly one.
+  env.step(action)
+  assert term.ticks.tolist() == [3, 2, 3, 3]
+
+
+def test_old_style_update_command_raises(counter_env):
+  """Terms with the old zero-arg _update_command fail fast at construction."""
+
+  class OldStyleCommand(CounterCommand):
+    def _update_command(self) -> None:  # type: ignore[override]
+      self.ticks += 1
+
+  @dataclass(kw_only=True)
+  class OldStyleCommandCfg(CounterCommandCfg):
+    def build(self, env) -> "OldStyleCommand":
+      return OldStyleCommand(self, env)
+
+  with pytest.raises(TypeError, match="env_ids"):
+    OldStyleCommandCfg().build(counter_env)
+
+
+def _make_motion_command_stub(time_steps, total, sampling_mode="uniform"):
+  """A stub with just enough state to drive MotionCommand._update_command."""
+  cmd = Mock()
+  cmd.time_steps = torch.tensor(time_steps, dtype=torch.long)
+  cmd.motion = Mock()
+  cmd.motion.time_step_total = total
+  cmd.cfg = Mock()
+  cmd.cfg.sampling_mode = sampling_mode
+  cmd.cfg.adaptive_alpha = 0.5
+  cmd.bin_failed_count = torch.tensor([1.0, 1.0])
+  cmd._current_bin_failed = torch.tensor([4.0, 4.0])
+  cmd._pending_forward = False
+  cmd._resample_command = Mock(
+    side_effect=lambda ids: setattr(cmd, "_pending_forward", True)
+  )
+  return cmd
+
+
+def test_motion_command_update_scopes_time_advance():
+  cmd = _make_motion_command_stub([2, 5, 7], total=100)
+  MotionCommand._update_command(cmd, env_ids=torch.tensor([1]))
+  assert cmd.time_steps.tolist() == [2, 6, 7]
+  cmd._resample_command.assert_not_called()
+  cmd.update_relative_body_poses.assert_called_once()
+
+  MotionCommand._update_command(cmd, env_ids=None)
+  assert cmd.time_steps.tolist() == [3, 7, 8]
+
+
+def test_motion_command_update_resamples_on_wraparound():
+  cmd = _make_motion_command_stub([2, 9], total=10)
+  MotionCommand._update_command(cmd, env_ids=torch.tensor([1]))
+  # Env 1 wrapped past the end of the motion and must be resampled; env 0
+  # is untouched.
+  (wrap_ids,), _ = cmd._resample_command.call_args
+  assert wrap_ids.tolist() == [1]
+  assert cmd.time_steps[0].item() == 2
+  cmd._env.sim.forward.assert_called_once()
+
+
+def test_motion_command_ema_folds_only_on_step_update():
+  cmd = _make_motion_command_stub([2, 5], total=100, sampling_mode="adaptive")
+  MotionCommand._update_command(cmd, env_ids=torch.tensor([0]))
+  # Reset-scoped update: EMA untouched, pending failure counts preserved.
+  assert cmd.bin_failed_count.tolist() == [1.0, 1.0]
+  assert cmd._current_bin_failed.tolist() == [4.0, 4.0]
+
+  MotionCommand._update_command(cmd, env_ids=None)
+  # Per-step update: EMA folds the counts and clears them.
+  assert cmd.bin_failed_count.tolist() == [2.5, 2.5]
+  assert cmd._current_bin_failed.tolist() == [0.0, 0.0]
+
+
+def test_motion_command_gui_reset_forwards_before_pose_update():
+  """apply_gui_reset must refresh kinematics between the state write and
+  update_relative_body_poses (viewer forwards only after it returns)."""
+  calls = []
+  cmd = Mock()
+  cmd._scrubber_handles = (Mock(value=5),)
+  cmd.reset_to_frame = lambda ids, frame: calls.append("reset_to_frame")
+  cmd._env.sim.forward = lambda: calls.append("forward")
+  cmd.update_relative_body_poses = lambda: calls.append("update_poses")
+
+  assert MotionCommand.apply_gui_reset(cmd, torch.tensor([0])) is True
+  assert calls == ["reset_to_frame", "forward", "update_poses"]
+
+
+def test_motion_command_timer_resample_triggers_forward():
+  """A timer-expiry resample (flag set before _update_command) forwards."""
+  cmd = _make_motion_command_stub([2, 5], total=100)
+  cmd._pending_forward = True  # As set by a compute-path _resample_command.
+  MotionCommand._update_command(cmd, env_ids=None)
+  cmd._env.sim.forward.assert_called_once()
+  assert cmd._pending_forward is False
+
+  cmd._env.sim.forward.reset_mock()
+  MotionCommand._update_command(cmd, env_ids=None)
+  cmd._env.sim.forward.assert_not_called()
@@ -0,0 +1,41 @@
+"""Tests for CurriculumManager."""
+
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from mjlab.managers.curriculum_manager import CurriculumManager, CurriculumTermCfg
+
+
+@pytest.fixture
+def mock_env():
+  env = Mock()
+  env.num_envs = 2
+  return env
+
+
+def test_get_active_iterable_terms_handles_dict_and_scalar_state(mock_env):
+  """Dict- and scalar-shaped curriculum states both yield flat value lists.
+
+  Regression: the dict branch previously indexed `terms` (a list) by term
+  name, raising TypeError. Only observable through callers that invoke
+  get_active_iterable_terms, which no in-tree caller currently does.
+  """
+
+  def dict_state_func(env, env_ids):
+    return {"a": torch.tensor(1.5), "b": 2.0}
+
+  def scalar_state_func(env, env_ids):
+    return torch.tensor(7.0)
+
+  cfg = {
+    "dict_term": CurriculumTermCfg(func=dict_state_func, params={}),
+    "scalar_term": CurriculumTermCfg(func=scalar_state_func, params={}),
+  }
+  manager = CurriculumManager(cfg, mock_env)
+  manager.compute()
+
+  terms = dict(manager.get_active_iterable_terms(0))
+  assert terms["dict_term"] == [1.5, 2.0]
+  assert terms["scalar_term"] == [7.0]
@@ -0,0 +1,284 @@
+"""Tests for DC motor actuator torque-speed curve."""
+
+import pytest
+import torch
+from conftest import (
+  create_entity_with_actuator,
+  get_test_device,
+  initialize_entity,
+  load_fixture_xml,
+)
+
+from mjlab.actuator import DcMotorActuatorCfg
+
+
+@pytest.fixture(scope="module")
+def device():
+  return get_test_device()
+
+
+@pytest.fixture(scope="module")
+def robot_xml():
+  return load_fixture_xml("floating_base_articulated")
+
+
+def test_dc_motor_stall_torque(device, robot_xml):
+  """DC motor produces full saturation_effort at zero velocity."""
+  kp = 100.0
+  kd = 10.0
+  saturation_effort = 20.0
+  velocity_limit = 30.0
+
+  entity = create_entity_with_actuator(
+    robot_xml,
+    DcMotorActuatorCfg(
+      target_names_expr=("joint.*",),
+      effort_limit=saturation_effort,  # Set to saturation to not constrain.
+      stiffness=kp,
+      damping=kd,
+      saturation_effort=saturation_effort,
+      velocity_limit=velocity_limit,
+    ),
+  )
+
+  entity, sim = initialize_entity(entity, device)
+
+  # Zero velocity, large position error to produce high PD torque.
+  joint_pos = torch.tensor([[0.0]], device=device)
+  joint_vel = torch.tensor([[0.0]], device=device)
+  entity.write_joint_state_to_sim(joint_pos, joint_vel)
+
+  # Set target to produce positive torque demand >> saturation_effort.
+  entity.set_joint_position_target(torch.tensor([[2.0]], device=device))
+  entity.set_joint_velocity_target(torch.zeros(1, 1, device=device))
+  entity.set_joint_effort_target(torch.zeros(1, 1, device=device))
+  entity.write_data_to_sim()
+
+  # At zero velocity, should be clipped to saturation_effort.
+  ctrl = sim.data.ctrl[0]
+  assert torch.allclose(ctrl, torch.tensor([saturation_effort], device=device))
+
+
+def test_dc_motor_zero_torque_at_max_velocity(device, robot_xml):
+  """DC motor produces zero torque at maximum velocity."""
+  kp = 100.0
+  kd = 0.0
+  saturation_effort = 20.0
+  velocity_limit = 30.0
+
+  entity = create_entity_with_actuator(
+    robot_xml,
+    DcMotorActuatorCfg(
+      target_names_expr=("joint.*",),
+      effort_limit=saturation_effort,  # Set to saturation to not constrain.
+      stiffness=kp,
+      damping=kd,
+      saturation_effort=saturation_effort,
+      velocity_limit=velocity_limit,
+    ),
+  )
+
+  entity, sim = initialize_entity(entity, device)
+
+  # At max velocity, large position error to produce high PD torque.
+  joint_pos = torch.tensor([[0.0]], device=device)
+  joint_vel = torch.tensor([[velocity_limit]], device=device)
+  entity.write_joint_state_to_sim(joint_pos, joint_vel)
+
+  # Set target to produce positive torque demand >> saturation_effort.
+  entity.set_joint_position_target(torch.tensor([[2.0]], device=device))
+  entity.set_joint_velocity_target(torch.zeros(1, 1, device=device))
+  entity.set_joint_effort_target(torch.zeros(1, 1, device=device))
+  entity.write_data_to_sim()
+
+  # At max velocity, should produce zero torque.
+  ctrl = sim.data.ctrl[0]
+  assert torch.allclose(ctrl, torch.tensor([0.0], device=device), atol=1e-5)
+
+
+def test_dc_motor_linear_torque_speed_curve(device, robot_xml):
+  """DC motor torque varies linearly between zero and max velocity."""
+  kp = 100.0
+  kd = 0.0
+  saturation_effort = 20.0
+  velocity_limit = 30.0
+
+  entity = create_entity_with_actuator(
+    robot_xml,
+    DcMotorActuatorCfg(
+      target_names_expr=("joint.*",),
+      effort_limit=saturation_effort,  # Set to saturation to not constrain.
+      stiffness=kp,
+      damping=kd,
+      saturation_effort=saturation_effort,
+      velocity_limit=velocity_limit,
+    ),
+  )
+
+  entity, sim = initialize_entity(entity, device)
+
+  # Test at half max velocity: should produce half saturation_effort.
+  joint_pos = torch.tensor([[0.0]], device=device)
+  joint_vel = torch.tensor([[velocity_limit / 2.0]], device=device)
+  entity.write_joint_state_to_sim(joint_pos, joint_vel)
+
+  # Large position error to produce high PD torque.
+  entity.set_joint_position_target(torch.tensor([[2.0]], device=device))
+  entity.set_joint_velocity_target(torch.zeros(1, 1, device=device))
+  entity.set_joint_effort_target(torch.zeros(1, 1, device=device))
+  entity.write_data_to_sim()
+
+  # At half velocity, should produce half saturation_effort.
+  ctrl = sim.data.ctrl[0]
+  expected = saturation_effort * 0.5
+  assert torch.allclose(ctrl, torch.tensor([expected], device=device), rtol=1e-4)
+
+
+def test_dc_motor_effort_limit_constrains_output(device, robot_xml):
+  """Continuous effort_limit constrains output below saturation_effort."""
+  kp = 100.0
+  kd = 0.0
+  saturation_effort = 20.0
+  effort_limit = 5.0
+  velocity_limit = 30.0
+
+  entity = create_entity_with_actuator(
+    robot_xml,
+    DcMotorActuatorCfg(
+      target_names_expr=("joint.*",),
+      effort_limit=effort_limit,
+      stiffness=kp,
+      damping=kd,
+      saturation_effort=saturation_effort,
+      velocity_limit=velocity_limit,
+    ),
+  )
+
+  entity, sim = initialize_entity(entity, device)
+
+  # Zero velocity: would produce saturation_effort without effort_limit.
+  joint_pos = torch.tensor([[0.0]], device=device)
+  joint_vel = torch.tensor([[0.0]], device=device)
+  entity.write_joint_state_to_sim(joint_pos, joint_vel)
+
+  # Large position error to produce high PD torque.
+  entity.set_joint_position_target(torch.tensor([[2.0]], device=device))
+  entity.set_joint_velocity_target(torch.zeros(1, 1, device=device))
+  entity.set_joint_effort_target(torch.zeros(1, 1, device=device))
+  entity.write_data_to_sim()
+
+  # Should be clamped to effort_limit, not saturation_effort.
+  ctrl = sim.data.ctrl[0]
+  assert torch.allclose(ctrl, torch.tensor([effort_limit], device=device))
+
+
+def test_dc_motor_negative_velocity_behavior(device, robot_xml):
+  """DC motor handles negative velocities symmetrically."""
+  kp = 100.0
+  kd = 0.0
+  saturation_effort = 20.0
+  velocity_limit = 30.0
+
+  entity = create_entity_with_actuator(
+    robot_xml,
+    DcMotorActuatorCfg(
+      target_names_expr=("joint.*",),
+      effort_limit=saturation_effort,  # Set to saturation to not constrain
+      stiffness=kp,
+      damping=kd,
+      saturation_effort=saturation_effort,
+      velocity_limit=velocity_limit,
+    ),
+  )
+
+  entity, sim = initialize_entity(entity, device)
+
+  # Negative velocity, target produces negative torque demand.
+  joint_pos = torch.tensor([[0.0]], device=device)
+  joint_vel = torch.tensor([[-velocity_limit / 2.0]], device=device)
+  entity.write_joint_state_to_sim(joint_pos, joint_vel)
+
+  # Set target to produce negative torque demand.
+  entity.set_joint_position_target(torch.tensor([[-2.0]], device=device))
+  entity.set_joint_velocity_target(torch.zeros(1, 1, device=device))
+  entity.set_joint_effort_target(torch.zeros(1, 1, device=device))
+  entity.write_data_to_sim()
+
+  # Should produce -half saturation_effort.
+  ctrl = sim.data.ctrl[0]
+  expected = -saturation_effort * 0.5
+  assert torch.allclose(ctrl, torch.tensor([expected], device=device), rtol=1e-4)
+
+
+def test_dc_motor_corner_velocity_transition(device, robot_xml):
+  """DC motor transitions correctly at corner velocity where curves intersect."""
+  kp = 100.0
+  kd = 0.0
+  saturation_effort = 20.0
+  effort_limit = 10.0
+  velocity_limit = 30.0
+
+  entity = create_entity_with_actuator(
+    robot_xml,
+    DcMotorActuatorCfg(
+      target_names_expr=("joint.*",),
+      effort_limit=effort_limit,
+      stiffness=kp,
+      damping=kd,
+      saturation_effort=saturation_effort,
+      velocity_limit=velocity_limit,
+    ),
+  )
+
+  entity, sim = initialize_entity(entity, device)
+
+  # Corner velocity: where torque-speed curve intersects effort_limit.
+  # vel_corner = velocity_limit * (1 - effort_limit / saturation_effort)
+  vel_corner = velocity_limit * (1.0 - effort_limit / saturation_effort)
+
+  joint_pos = torch.tensor([[0.0]], device=device)
+  joint_vel = torch.tensor([[vel_corner]], device=device)
+  entity.write_joint_state_to_sim(joint_pos, joint_vel)
+
+  # Large position error to produce high PD torque.
+  entity.set_joint_position_target(torch.tensor([[2.0]], device=device))
+  entity.set_joint_velocity_target(torch.zeros(1, 1, device=device))
+  entity.set_joint_effort_target(torch.zeros(1, 1, device=device))
+  entity.write_data_to_sim()
+
+  # At corner velocity, should produce exactly effort_limit.
+  ctrl = sim.data.ctrl[0]
+  assert torch.allclose(ctrl, torch.tensor([effort_limit], device=device), rtol=1e-4)
+
+
+def test_dc_motor_warns_when_effort_limit_is_inf():
+  """DcMotorActuatorCfg warns when effort_limit is inf."""
+  import warnings
+
+  # inf triggers both warnings (is inf + exceeds saturation), so catch both.
+  with warnings.catch_warnings(record=True) as w:
+    warnings.simplefilter("always")
+    DcMotorActuatorCfg(
+      target_names_expr=("joint.*",),
+      stiffness=100.0,
+      damping=10.0,
+      saturation_effort=20.0,
+      velocity_limit=30.0,
+      # effort_limit intentionally not set (defaults to inf).
+    )
+    # Should trigger the "is inf" warning.
+    assert len(w) >= 1
+    assert any("effort_limit is set to inf" in str(warning.message) for warning in w)
+
+
+def test_dc_motor_warns_when_effort_limit_exceeds_saturation():
+  """DcMotorActuatorCfg warns when effort_limit > saturation_effort."""
+  with pytest.warns(UserWarning, match="effort_limit.*exceeds saturation_effort"):
+    DcMotorActuatorCfg(
+      target_names_expr=("joint.*",),
+      stiffness=100.0,
+      damping=10.0,
+      saturation_effort=20.0,
+      velocity_limit=30.0,
+      effort_limit=25.0,  # > saturation_effort.
+    )
@@ -0,0 +1,327 @@
+"""Tests for DelayBuffer."""
+
+import pytest
+import torch
+from conftest import get_test_device
+
+from mjlab.utils.buffers import DelayBuffer
+
+
+@pytest.fixture
+def device():
+  """Test device fixture."""
+  return get_test_device()
+
+
+def make_gen(seed: int, device: str) -> torch.Generator:
+  """Create a seeded generator for reproducible tests."""
+  gen = torch.Generator(device=device)
+  gen.manual_seed(seed)
+  return gen
+
+
+##
+# Basic behavior.
+##
+
+
+def test_delay_buffer_zero_lag(device):
+  """Zero lag returns current observation."""
+  buffer = DelayBuffer(min_lag=0, max_lag=0, batch_size=2, device=device)
+
+  buffer.append(torch.tensor([[1.0], [2.0]], device=device))
+  result = buffer.compute()
+  assert torch.allclose(result, torch.tensor([[1.0], [2.0]], device=device))
+
+  buffer.append(torch.tensor([[3.0], [4.0]], device=device))
+  result = buffer.compute()
+  assert torch.allclose(result, torch.tensor([[3.0], [4.0]], device=device))
+
+
+def test_delay_buffer_constant_lag(device):
+  """Constant lag returns observation from N steps ago."""
+  buffer = DelayBuffer(min_lag=2, max_lag=2, batch_size=1, device=device)
+
+  buffer.append(torch.tensor([[1.0]], device=device))
+  result = buffer.compute()
+  assert torch.allclose(result, torch.tensor([[1.0]], device=device))
+
+  buffer.append(torch.tensor([[2.0]], device=device))
+  result = buffer.compute()
+  assert torch.allclose(result, torch.tensor([[1.0]], device=device))
+
+  buffer.append(torch.tensor([[3.0]], device=device))
+  result = buffer.compute()
+  assert torch.allclose(result, torch.tensor([[1.0]], device=device))
+
+  buffer.append(torch.tensor([[4.0]], device=device))
+  result = buffer.compute()
+  assert torch.allclose(result, torch.tensor([[2.0]], device=device))
+
+
+def test_value_matches_lag_per_env(device):
+  """Delayed values correctly match current lags."""
+  B = 3
+  buf = DelayBuffer(
+    0, 2, batch_size=B, per_env=True, device=device, generator=make_gen(1234, device)
+  )
+  for t in range(6):
+    buf.append(torch.full((B, 1), float(t), device=device))
+    y = buf.compute()
+    lags = buf.current_lags
+    for e in range(B):
+      eff = int(min(lags[e].item(), t))
+      assert y[e].item() == float(t - eff)
+
+
+##
+# Lag sampling modes.
+##
+
+
+def test_delay_buffer_per_env_lags(device):
+  """Per-env mode allows different lags per environment."""
+  buffer = DelayBuffer(min_lag=0, max_lag=3, batch_size=4, per_env=True, device=device)
+
+  for i in range(10):
+    buffer.append(torch.full((4, 1), float(i), device=device))
+    buffer.compute()
+
+  lags = buffer.current_lags
+  assert torch.all(lags >= 0)
+  assert torch.all(lags <= 3)
+
+
+def test_delay_buffer_shared_lags(device):
+  """Shared mode uses same lag across all environments."""
+  buffer = DelayBuffer(min_lag=0, max_lag=3, batch_size=4, per_env=False, device=device)
+
+  for i in range(10):
+    buffer.append(torch.full((4, 1), float(i), device=device))
+    buffer.compute()
+
+  lags = buffer.current_lags
+  assert torch.all(lags == lags[0])
+
+
+##
+# Hold probability.
+##
+
+
+def test_hold_prob_always_hold(device):
+  """hold_prob=1.0 keeps lag frozen."""
+  buf = DelayBuffer(
+    0,
+    3,
+    batch_size=1,
+    device=device,
+    update_period=1,
+    hold_prob=1.0,
+    generator=make_gen(7, device),
+  )
+  for t in range(2):
+    buf.append(torch.tensor([[float(t)]], device=device))
+    buf.compute()
+  first = buf.current_lags.item()
+  for t in range(2, 20):
+    buf.append(torch.tensor([[float(t)]], device=device))
+    buf.compute()
+    assert buf.current_lags.item() == first
+
+
+def test_hold_prob_never_hold(device):
+  """hold_prob=0.0 allows lag to change."""
+  buf = DelayBuffer(
+    0,
+    3,
+    batch_size=1,
+    device=device,
+    update_period=1,
+    hold_prob=0.0,
+    generator=make_gen(42, device),
+  )
+  for t in range(2):
+    buf.append(torch.tensor([[float(t)]], device=device))
+    buf.compute()
+  prev, changed = buf.current_lags.item(), False
+  for t in range(2, 20):
+    buf.append(torch.tensor([[float(t)]], device=device))
+    buf.compute()
+    cur = buf.current_lags.item()
+    if cur != prev:
+      changed = True
+      break
+    prev = cur
+  assert changed
+
+
+##
+# Update period.
+##
+
+
+def test_delay_buffer_update_period(device):
+  """Update period controls how often lags are resampled."""
+  buffer = DelayBuffer(
+    min_lag=0,
+    max_lag=3,
+    batch_size=1,
+    update_period=3,
+    per_env_phase=False,
+    device=device,
+  )
+
+  lag_history = []
+  for i in range(12):
+    buffer.append(torch.tensor([[float(i)]], device=device))
+    buffer.compute()
+    lag_history.append(buffer.current_lags[0].item())
+
+  assert len(lag_history) == 12
+
+
+def test_update_period_changes_only_on_schedule(device):
+  """Lags update only at scheduled intervals."""
+  buf = DelayBuffer(
+    0,
+    10,  # Wider range makes value changes more likely
+    batch_size=1,
+    update_period=3,
+    per_env_phase=False,
+    device=device,
+    generator=make_gen(123, device),
+  )
+
+  # Track when lags are updated by checking step_count timing
+  lag_values = []
+  for t in range(12):
+    buf.append(torch.tensor([[float(t)]], device=device))
+    buf.compute()
+    lag_values.append(buf.current_lags.item())
+
+  # Verify lag stays constant between update periods
+  # Update happens at step 0, 3, 6, 9
+  # So: [0-2] same, [3-5] same, [6-8] same, [9-11] same
+  assert lag_values[0] == lag_values[1] == lag_values[2]
+  assert lag_values[3] == lag_values[4] == lag_values[5]
+  assert lag_values[6] == lag_values[7] == lag_values[8]
+  assert lag_values[9] == lag_values[10] == lag_values[11]
+
+
+##
+# Reset behavior.
+##
+
+
+def test_delay_buffer_reset_all(device):
+  """Reset clears all environments."""
+  buffer = DelayBuffer(min_lag=1, max_lag=2, batch_size=2, device=device)
+
+  for i in range(5):
+    buffer.append(torch.full((2, 1), float(i), device=device))
+    buffer.compute()
+
+  buffer.reset()
+
+  assert torch.all(buffer.current_lags == 0)
+  assert torch.all(buffer._step_count == 0)
+
+
+def test_delay_buffer_reset_partial(device):
+  """Reset with batch_ids only resets specified environments."""
+  buffer = DelayBuffer(min_lag=0, max_lag=3, batch_size=4, device=device)
+
+  for _ in range(5):
+    buffer.append(torch.arange(4, device=device).unsqueeze(1).float())
+    buffer.compute()
+
+  lags_before = buffer.current_lags.clone()
+
+  buffer.reset(batch_ids=torch.tensor([0, 2], device=device))
+
+  lags_after = buffer.current_lags
+
+  assert lags_after[0] == 0
+  assert lags_after[2] == 0
+  assert lags_after[1] == lags_before[1]
+  assert lags_after[3] == lags_before[3]
+
+
+def test_partial_reset_then_backfill(device):
+  """Partial reset returns zeros until next append backfills."""
+  B = 3
+  buf = DelayBuffer(1, 2, batch_size=B, device=device, generator=make_gen(9, device))
+  for t in range(3):
+    x = torch.arange(1, B + 1, device=device).float().unsqueeze(1) * 10 + t
+    buf.append(x)
+    buf.compute()
+  buf.reset(batch_ids=torch.tensor([1], device=device))
+  y = buf.compute()
+  assert torch.allclose(y[1], torch.zeros_like(y[1]))
+  x_new = torch.tensor([[111.0], [999.0], [333.0]], device=device)
+  buf.append(x_new)
+  y2 = buf.compute()
+  assert torch.allclose(y2[1], torch.tensor([[999.0]], device=device))
+
+
+##
+# Error handling.
+##
+
+
+def test_delay_buffer_not_initialized(device):
+  """Compute before append raises error."""
+  buffer = DelayBuffer(min_lag=0, max_lag=3, batch_size=1, device=device)
+
+  with pytest.raises(RuntimeError, match="Buffer not initialized"):
+    buffer.compute()
+
+
+def test_append_wrong_batch_raises(device):
+  """Appending wrong batch size raises error."""
+  buf = DelayBuffer(0, 1, batch_size=2, device=device)
+  with pytest.raises(ValueError):
+    buf.append(torch.zeros(3, 1, device=device))
+
+
+def test_delay_buffer_validation(device):
+  """Input validation catches invalid parameters."""
+  with pytest.raises(ValueError, match="min_lag must be >= 0"):
+    DelayBuffer(min_lag=-1, max_lag=3, batch_size=1, device=device)
+
+  with pytest.raises(ValueError, match="max_lag.*must be >= min_lag"):
+    DelayBuffer(min_lag=5, max_lag=3, batch_size=1, device=device)
+
+  with pytest.raises(ValueError, match="hold_prob must be in"):
+    DelayBuffer(min_lag=0, max_lag=3, batch_size=1, hold_prob=1.5, device=device)
+
+  with pytest.raises(ValueError, match="update_period must be >= 0"):
+    DelayBuffer(min_lag=0, max_lag=3, batch_size=1, update_period=-1, device=device)
+
+
+def test_backfill_and_peek_leave_others_untouched(device):
+  """backfill serves the reset row's new frame; peek advances no schedules."""
+  B = 3
+  buf = DelayBuffer(1, 2, batch_size=B, device=device, generator=make_gen(9, device))
+  for t in range(4):
+    x = torch.arange(1, B + 1, device=device).float().unsqueeze(1) * 10 + t
+    buf.append(x)
+    buf.compute()
+
+  out_before = buf.peek()
+  lags_before = buf.current_lags.clone()
+  steps_before = buf._step_count.clone()
+
+  reset_ids = torch.tensor([1], device=device)
+  buf.reset(batch_ids=reset_ids)
+  buf.backfill(torch.tensor([[-1.0], [999.0], [-1.0]], device=device), reset_ids)
+  out = buf.peek()
+
+  # Reset row serves its backfilled frame (lag zeroed, one valid frame).
+  assert torch.allclose(out[1], torch.tensor([999.0], device=device))
+  # Other rows: same delayed frame, lags and schedule untouched.
+  keep = torch.tensor([0, 2], device=device)
+  assert torch.allclose(out[keep], out_before[keep])
+  assert torch.equal(buf.current_lags[keep], lags_before[keep])
+  assert torch.equal(buf._step_count[keep], steps_before[keep])
@@ -0,0 +1,363 @@
+"""Tests for delayed actuators."""
+
+import mujoco
+import pytest
+import torch
+from conftest import get_test_device, load_fixture_xml
+
+from mjlab.actuator import (
+  BuiltinPositionActuatorCfg,
+  IdealPdActuatorCfg,
+)
+from mjlab.entity import Entity, EntityArticulationInfoCfg, EntityCfg
+from mjlab.sim.sim import Simulation, SimulationCfg
+
+ROBOT_XML = load_fixture_xml("floating_base_articulated")
+
+
+@pytest.fixture(scope="module")
+def device():
+  return get_test_device()
+
+
+def create_entity_with_delayed_builtin(delay_min_lag=0, delay_max_lag=3):
+  cfg = EntityCfg(
+    spec_fn=lambda: mujoco.MjSpec.from_string(ROBOT_XML),
+    articulation=EntityArticulationInfoCfg(
+      actuators=(
+        BuiltinPositionActuatorCfg(
+          target_names_expr=("joint.*",),
+          effort_limit=100.0,
+          stiffness=80.0,
+          damping=10.0,
+          delay_min_lag=delay_min_lag,
+          delay_max_lag=delay_max_lag,
+        ),
+      )
+    ),
+  )
+  return Entity(cfg)
+
+
+def create_entity_with_delayed_ideal(delay_min_lag=0, delay_max_lag=3):
+  cfg = EntityCfg(
+    spec_fn=lambda: mujoco.MjSpec.from_string(ROBOT_XML),
+    articulation=EntityArticulationInfoCfg(
+      actuators=(
+        IdealPdActuatorCfg(
+          target_names_expr=("joint.*",),
+          effort_limit=100.0,
+          stiffness=80.0,
+          damping=10.0,
+          delay_min_lag=delay_min_lag,
+          delay_max_lag=delay_max_lag,
+        ),
+      )
+    ),
+  )
+  return Entity(cfg)
+
+
+def initialize_entity(entity, device, num_envs=1):
+  model = entity.compile()
+  sim_cfg = SimulationCfg()
+  sim = Simulation(num_envs=num_envs, cfg=sim_cfg, model=model, device=device)
+  entity.initialize(model, sim.model, sim.data, device)
+  return entity, sim
+
+
+def test_delayed_builtin_applies_constant_delay(device):
+  """Test that delayed builtin actuator delays position targets."""
+  entity = create_entity_with_delayed_builtin(delay_min_lag=2, delay_max_lag=2)
+  entity, sim = initialize_entity(entity, device)
+
+  # Set position targets for 3 steps.
+  targets = [
+    torch.tensor([[0.1, 0.2]], device=device),
+    torch.tensor([[0.3, 0.4]], device=device),
+    torch.tensor([[0.5, 0.6]], device=device),
+  ]
+
+  joint_vel = torch.zeros(1, 2, device=device)
+
+  for target in targets:
+    entity.set_joint_position_target(target)
+    entity.set_joint_velocity_target(joint_vel)
+    entity.set_joint_effort_target(torch.zeros(1, 2, device=device))
+    entity.write_data_to_sim()
+
+  # After 3 steps with lag=2, the output should be the target from step 0.
+  ctrl = sim.data.ctrl[0]
+  # With constant lag=2, after 3 appends, we expect target from step 0.
+  assert torch.allclose(ctrl, targets[0][0], atol=1e-5)
+
+
+def test_delayed_ideal_applies_delay(device):
+  """Test that delayed ideal actuator delays position targets."""
+  entity = create_entity_with_delayed_ideal(delay_min_lag=2, delay_max_lag=2)
+  entity, sim = initialize_entity(entity, device)
+
+  joint_pos = torch.zeros(1, 2, device=device)
+  joint_vel = torch.zeros(1, 2, device=device)
+  entity.write_joint_state_to_sim(joint_pos, joint_vel)
+
+  # Set position targets for 3 steps.
+  targets = [
+    torch.tensor([[0.1, 0.2]], device=device),
+    torch.tensor([[0.3, 0.4]], device=device),
+    torch.tensor([[0.5, 0.6]], device=device),
+  ]
+
+  for target in targets:
+    entity.set_joint_position_target(target)
+    entity.set_joint_velocity_target(joint_vel)
+    entity.set_joint_effort_target(torch.zeros(1, 2, device=device))
+    entity.write_data_to_sim()
+    sim.forward()  # Compute actuator forces
+
+  # The computed torque should use the delayed target from step 0.
+  joint_v_adr = entity.indexing.joint_v_adr
+  qfrc = sim.data.qfrc_actuator[0, joint_v_adr]
+
+  # Expected torque: kp * (delayed_target - joint_pos) + kd * (0 - joint_vel) + 0
+  # = 80.0 * targets[0] + 0 = 80.0 * [0.1, 0.2]
+  expected_torque = 80.0 * targets[0][0]
+  assert torch.allclose(qfrc, expected_torque, atol=1e-4)
+
+
+def test_delayed_ideal_delays_velocity(device):
+  """Velocity targets share the same delay as position targets.
+
+  Regression test: the velocity reference used to bypass the delay buffer, so
+  the damping term consumed the latest target instead of the delayed one.
+  """
+  entity = create_entity_with_delayed_ideal(delay_min_lag=2, delay_max_lag=2)
+  entity, sim = initialize_entity(entity, device)
+
+  joint_pos = torch.zeros(1, 2, device=device)
+  joint_vel = torch.zeros(1, 2, device=device)
+  entity.write_joint_state_to_sim(joint_pos, joint_vel)
+
+  # Only the velocity target varies; position and effort stay zero.
+  vel_targets = [
+    torch.tensor([[0.1, 0.2]], device=device),
+    torch.tensor([[0.3, 0.4]], device=device),
+    torch.tensor([[0.5, 0.6]], device=device),
+  ]
+
+  for vel_target in vel_targets:
+    entity.set_joint_position_target(joint_pos)
+    entity.set_joint_velocity_target(vel_target)
+    entity.set_joint_effort_target(torch.zeros(1, 2, device=device))
+    entity.write_data_to_sim()
+    sim.forward()
+
+  joint_v_adr = entity.indexing.joint_v_adr
+  qfrc = sim.data.qfrc_actuator[0, joint_v_adr]
+
+  # With lag=2, the damping term uses the velocity target from step 0:
+  # kd * (delayed_vel_target - 0) = 10.0 * [0.1, 0.2].
+  expected_torque = 10.0 * vel_targets[0][0]
+  assert torch.allclose(qfrc, expected_torque, atol=1e-4)
+
+
+def test_delayed_ideal_delays_effort(device):
+  """Feedforward effort targets share the same delay as position targets."""
+  entity = create_entity_with_delayed_ideal(delay_min_lag=2, delay_max_lag=2)
+  entity, sim = initialize_entity(entity, device)
+
+  joint_pos = torch.zeros(1, 2, device=device)
+  joint_vel = torch.zeros(1, 2, device=device)
+  entity.write_joint_state_to_sim(joint_pos, joint_vel)
+
+  effort_targets = [
+    torch.tensor([[1.0, 2.0]], device=device),
+    torch.tensor([[3.0, 4.0]], device=device),
+    torch.tensor([[5.0, 6.0]], device=device),
+  ]
+
+  for effort_target in effort_targets:
+    entity.set_joint_position_target(joint_pos)
+    entity.set_joint_velocity_target(joint_vel)
+    entity.set_joint_effort_target(effort_target)
+    entity.write_data_to_sim()
+    sim.forward()
+
+  joint_v_adr = entity.indexing.joint_v_adr
+  qfrc = sim.data.qfrc_actuator[0, joint_v_adr]
+
+  # With lag=2, the feedforward term uses the effort target from step 0.
+  expected_torque = effort_targets[0][0]
+  assert torch.allclose(qfrc, expected_torque, atol=1e-4)
+
+
+def test_delayed_actuator_reset(device):
+  """Test that reset clears the delay buffer."""
+  entity = create_entity_with_delayed_builtin(delay_min_lag=1, delay_max_lag=3)
+  entity, _ = initialize_entity(entity, device, num_envs=2)
+
+  # Set some targets to fill the buffer.
+  entity.set_joint_position_target(torch.ones(2, 2, device=device) * 0.5)
+  entity.set_joint_velocity_target(torch.zeros(2, 2, device=device))
+  entity.set_joint_effort_target(torch.zeros(2, 2, device=device))
+  entity.write_data_to_sim()
+
+  # Reset env 0.
+  entity.reset(torch.tensor([0], device=device))
+
+  # Check that delay buffer was reset for env 0.
+  actuator = entity.actuators[0]
+  assert actuator.has_delay
+  assert actuator._delay_buffer is not None
+  assert actuator._delay_buffer.current_lags[0] == 0
+
+
+def test_delayed_actuator_set_lags(device):
+  """Test that set_lags sets lag values on all delay buffers."""
+  entity = create_entity_with_delayed_builtin(delay_min_lag=0, delay_max_lag=5)
+  entity, _ = initialize_entity(entity, device, num_envs=4)
+
+  actuator = entity.actuators[0]
+  assert actuator.has_delay
+
+  # Set lags for all environments.
+  lags = torch.tensor([1, 2, 3, 4], device=device)
+  actuator.set_lags(lags)
+
+  # Check that lags were set.
+  buffer = actuator._delay_buffer
+  assert buffer is not None
+  assert torch.equal(buffer.current_lags, lags)
+
+
+def test_delayed_actuator_set_lags_subset(device):
+  """Test that set_lags can set lag values for a subset of environments."""
+  entity = create_entity_with_delayed_builtin(delay_min_lag=0, delay_max_lag=5)
+  entity, _ = initialize_entity(entity, device, num_envs=4)
+
+  actuator = entity.actuators[0]
+  assert actuator.has_delay
+
+  # Set lags for envs 1 and 3 only.
+  env_ids = torch.tensor([1, 3], device=device)
+  lags = torch.tensor([4, 5], device=device)
+  actuator.set_lags(lags, env_ids)
+
+  # Check that only specified envs were updated.
+  buffer = actuator._delay_buffer
+  assert buffer is not None
+  assert buffer.current_lags[0] == 0  # Unchanged (initial value)
+  assert buffer.current_lags[1] == 4
+  assert buffer.current_lags[2] == 0  # Unchanged
+  assert buffer.current_lags[3] == 5
+
+
+def test_delayed_actuator_set_lags_clamps_to_range(device):
+  """Test that set_lags clamps values to the configured lag range."""
+  entity = create_entity_with_delayed_builtin(delay_min_lag=1, delay_max_lag=3)
+  entity, _ = initialize_entity(entity, device, num_envs=2)
+
+  actuator = entity.actuators[0]
+  assert actuator.has_delay
+
+  # Try to set lags outside the valid range.
+  lags = torch.tensor([0, 10], device=device)  # 0 < min_lag, 10 > max_lag
+  actuator.set_lags(lags)
+
+  # Lags should be clamped to [1, 3].
+  buffer = actuator._delay_buffer
+  assert buffer is not None
+  assert buffer.current_lags[0] == 1  # Clamped from 0
+  assert buffer.current_lags[1] == 3  # Clamped from 10
+
+
+def test_delayed_actuator_set_lags_affects_delay(device):
+  """Test that setting lags actually changes the delay behavior."""
+  # Use hold_prob=1.0 to prevent automatic lag resampling.
+  cfg = EntityCfg(
+    spec_fn=lambda: mujoco.MjSpec.from_string(ROBOT_XML),
+    articulation=EntityArticulationInfoCfg(
+      actuators=(
+        BuiltinPositionActuatorCfg(
+          target_names_expr=("joint.*",),
+          effort_limit=100.0,
+          stiffness=80.0,
+          damping=10.0,
+          delay_min_lag=0,
+          delay_max_lag=5,
+          delay_hold_prob=1.0,  # Prevent automatic resampling
+        ),
+      )
+    ),
+  )
+  entity = Entity(cfg)
+  entity, sim = initialize_entity(entity, device, num_envs=1)
+
+  actuator = entity.actuators[0]
+  assert actuator.has_delay
+
+  # Set lag to 1.
+  actuator.set_lags(torch.tensor([1], device=device))
+
+  # Fill the buffer with known targets.
+  targets = [
+    torch.tensor([[0.1, 0.2]], device=device),
+    torch.tensor([[0.3, 0.4]], device=device),
+    torch.tensor([[0.5, 0.6]], device=device),
+  ]
+
+  for target in targets:
+    entity.set_joint_position_target(target)
+    entity.set_joint_velocity_target(torch.zeros(1, 2, device=device))
+    entity.set_joint_effort_target(torch.zeros(1, 2, device=device))
+    entity.write_data_to_sim()
+
+  # With lag=1, after 3 steps, ctrl should use target from step 1 (index 1).
+  ctrl = sim.data.ctrl[0]
+  assert torch.allclose(ctrl, targets[1][0], atol=1e-5)
+
+
+def test_delayed_actuator_set_lags_overwritten_without_hold_prob(device):
+  """Test that set_lags gets overwritten when delay_hold_prob < 1.0."""
+  # Use min_lag=max_lag=2 so resampling always produces 2.
+  cfg = EntityCfg(
+    spec_fn=lambda: mujoco.MjSpec.from_string(ROBOT_XML),
+    articulation=EntityArticulationInfoCfg(
+      actuators=(
+        BuiltinPositionActuatorCfg(
+          target_names_expr=("joint.*",),
+          effort_limit=100.0,
+          stiffness=80.0,
+          damping=10.0,
+          delay_min_lag=2,
+          delay_max_lag=2,
+          delay_hold_prob=0.0,  # Always resample
+        ),
+      )
+    ),
+  )
+  entity = Entity(cfg)
+  entity, sim = initialize_entity(entity, device, num_envs=1)
+
+  actuator = entity.actuators[0]
+  assert actuator.has_delay
+  buffer = actuator._delay_buffer
+  assert buffer is not None
+
+  # Set lag to 2 (the only valid value, so set_lags won't clamp it).
+  actuator.set_lags(torch.tensor([2], device=device))
+  assert buffer.current_lags[0] == 2
+
+  # Now manually set _current_lags to 0 to simulate "we want 0".
+  # This bypasses clamping to test the resampling behavior.
+  buffer._current_lags[0] = 0
+  assert buffer.current_lags[0] == 0
+
+  # After compute, with hold_prob=0.0, it resamples to [min_lag, max_lag] = 2.
+  entity.set_joint_position_target(torch.zeros(1, 2, device=device))
+  entity.set_joint_velocity_target(torch.zeros(1, 2, device=device))
+  entity.set_joint_effort_target(torch.zeros(1, 2, device=device))
+  entity.write_data_to_sim()
+
+  # Lag should have been resampled back to 2.
+  assert buffer.current_lags[0] == 2
@@ -0,0 +1,423 @@
+"""Tests for differential IK action space."""
+
+from unittest.mock import Mock
+
+import mujoco
+import pytest
+import torch
+from conftest import get_test_device
+
+from mjlab.actuator.actuator import TransmissionType
+from mjlab.actuator.builtin_actuator import BuiltinPositionActuatorCfg
+from mjlab.entity import Entity, EntityArticulationInfoCfg, EntityCfg
+from mjlab.envs import ManagerBasedRlEnv
+from mjlab.envs.mdp.actions import DifferentialIKAction, DifferentialIKActionCfg
+from mjlab.sim.sim import MujocoCfg, Simulation, SimulationCfg
+from mjlab.utils.lab_api.math import axis_angle_from_quat, quat_mul
+
+ARM_XML = """\
+<mujoco>
+  <worldbody>
+    <body name="base" pos="0 0 0.5">
+      <geom name="base_geom" type="cylinder" size="0.05 0.02" mass="1.0"
+            contype="0" conaffinity="0"/>
+      <body name="link1" pos="0 0 0.02">
+        <joint name="joint1" type="hinge" axis="0 1 0" range="-3.14 3.14"/>
+        <geom name="link1_geom" type="capsule" fromto="0 0 0 0 0 0.3"
+              size="0.02" mass="0.5" contype="0" conaffinity="0"/>
+        <body name="link2" pos="0 0 0.3">
+          <joint name="joint2" type="hinge" axis="0 1 0" range="-3.14 3.14"/>
+          <geom name="link2_geom" type="capsule" fromto="0 0 0 0 0 0.3"
+                size="0.02" mass="0.5" contype="0" conaffinity="0"/>
+          <body name="ee" pos="0 0 0.3">
+            <joint name="joint3" type="hinge" axis="0 1 0"
+                   range="-3.14 3.14"/>
+            <geom name="ee_geom" type="sphere" size="0.03" mass="0.1"
+                  contype="0" conaffinity="0"/>
+            <site name="ee_site" pos="0 0 0.05"/>
+          </body>
+        </body>
+      </body>
+    </body>
+  </worldbody>
+</mujoco>
+"""
+
+TIGHT_ARM_XML = """\
+<mujoco>
+  <worldbody>
+    <body name="base" pos="0 0 0.5">
+      <geom name="base_geom" type="cylinder" size="0.05 0.02" mass="1.0"
+            contype="0" conaffinity="0"/>
+      <body name="link1" pos="0 0 0.02">
+        <joint name="joint1" type="hinge" axis="0 1 0" range="-0.5 0.5"/>
+        <geom name="link1_geom" type="capsule" fromto="0 0 0 0 0 0.3"
+              size="0.02" mass="0.5" contype="0" conaffinity="0"/>
+        <body name="link2" pos="0 0 0.3">
+          <joint name="joint2" type="hinge" axis="0 1 0" range="-0.5 0.5"/>
+          <geom name="link2_geom" type="capsule" fromto="0 0 0 0 0 0.3"
+                size="0.02" mass="0.5" contype="0" conaffinity="0"/>
+          <body name="ee" pos="0 0 0.3">
+            <joint name="joint3" type="hinge" axis="0 1 0"
+                   range="-0.5 0.5"/>
+            <geom name="ee_geom" type="sphere" size="0.03" mass="0.1"
+                  contype="0" conaffinity="0"/>
+          </body>
+        </body>
+      </body>
+    </body>
+  </worldbody>
+</mujoco>
+"""
+
+NUM_ENVS = 4
+
+
+@pytest.fixture
+def device():
+  return get_test_device()
+
+
+def _make_entity(device, xml=ARM_XML):
+  cfg = EntityCfg(
+    spec_fn=lambda: mujoco.MjSpec.from_string(xml),
+    articulation=EntityArticulationInfoCfg(
+      actuators=(
+        BuiltinPositionActuatorCfg(
+          target_names_expr=("joint.*",),
+          transmission_type=TransmissionType.JOINT,
+          stiffness=1.0,
+          damping=1.0,
+          effort_limit=1.0,
+        ),
+      )
+    ),
+  )
+  entity = Entity(cfg)
+  model = entity.compile()
+  sim_cfg = SimulationCfg(mujoco=MujocoCfg(gravity=(0, 0, 0)))
+  sim = Simulation(num_envs=NUM_ENVS, cfg=sim_cfg, model=model, device=device)
+  entity.initialize(model, sim.model, sim.data, device)
+  return entity, sim
+
+
+def _make_env(entity, sim, device):
+  env = Mock(spec=ManagerBasedRlEnv)
+  env.num_envs = NUM_ENVS
+  env.device = device
+  env.scene = {"robot": entity}
+  env.sim = sim
+  return env
+
+
+@pytest.mark.parametrize(
+  "orientation_weight,relative,expected_dim",
+  [
+    (0.0, True, 3),
+    (0.0, False, 3),
+    (1.0, True, 6),
+    (1.0, False, 7),
+  ],
+)
+def test_action_dim(device, orientation_weight, relative, expected_dim):
+  entity, sim = _make_entity(device)
+  env = _make_env(entity, sim, device)
+  cfg = DifferentialIKActionCfg(
+    entity_name="robot",
+    actuator_names=("joint.*",),
+    frame_name="ee",
+    frame_type="body",
+    orientation_weight=orientation_weight,
+    use_relative_mode=relative,
+  )
+  action = cfg.build(env)
+  assert action.action_dim == expected_dim
+
+
+@pytest.mark.parametrize(
+  "frame_type,frame_name",
+  [
+    ("body", "ee"),
+    ("site", "ee_site"),
+    ("geom", "ee_geom"),
+  ],
+)
+def test_frame_types(device, frame_type, frame_name):
+  entity, sim = _make_entity(device)
+  env = _make_env(entity, sim, device)
+  cfg = DifferentialIKActionCfg(
+    entity_name="robot",
+    actuator_names=("joint.*",),
+    frame_name=frame_name,
+    frame_type=frame_type,
+    orientation_weight=0.0,
+  )
+  action: DifferentialIKAction = cfg.build(env)
+  assert action.action_dim == 3
+
+  pos, quat = action._get_frame_pose()
+  assert pos.shape == (NUM_ENVS, 3)
+  assert quat.shape == (NUM_ENVS, 4)
+
+
+def test_ik_convergence_position(device):
+  """Verify IK converges kinematically by writing q+dq directly to qpos."""
+  entity, sim = _make_entity(device)
+  env = _make_env(entity, sim, device)
+  cfg = DifferentialIKActionCfg(
+    entity_name="robot",
+    actuator_names=("joint.*",),
+    frame_name="ee",
+    frame_type="body",
+    orientation_weight=0.0,
+    use_relative_mode=False,
+    damping=0.05,
+  )
+  action: DifferentialIKAction = cfg.build(env)
+
+  initial_pos, _ = action._get_frame_pose()
+  target = initial_pos.clone()
+  target[:, 0] += 0.05
+
+  action.process_actions(target)
+  for _ in range(10):
+    dq = action.compute_dq()
+    q = entity.data.joint_pos[:, action._joint_ids] + dq
+    entity.write_joint_position_to_sim(q, joint_ids=action._joint_ids)
+    sim.forward()
+
+  final_pos, _ = action._get_frame_pose()
+  error = (final_pos - target).norm(dim=-1)
+  assert (error < 0.003).all(), f"Position error too large: {error}"
+
+
+def test_reset_clears_buffers(device):
+  entity, sim = _make_entity(device)
+  env = _make_env(entity, sim, device)
+  cfg = DifferentialIKActionCfg(
+    entity_name="robot",
+    actuator_names=("joint.*",),
+    frame_name="ee",
+    frame_type="body",
+  )
+  action: DifferentialIKAction = cfg.build(env)
+
+  actions = torch.randn(NUM_ENVS, 6, device=device)
+  action.process_actions(actions)
+
+  reset_ids = torch.tensor([0, 2], device=device)
+  action.reset(env_ids=reset_ids)
+
+  assert torch.all(action.raw_action[[0, 2]] == 0.0)
+  assert not torch.all(action.raw_action[[1, 3]] == 0.0)
+
+
+def test_joint_limits_respected(device):
+  """IK with joint_limit_weight keeps joints within bounds."""
+  entity, sim = _make_entity(device, xml=TIGHT_ARM_XML)
+  env = _make_env(entity, sim, device)
+
+  action_no_lim: DifferentialIKAction = DifferentialIKActionCfg(
+    entity_name="robot",
+    actuator_names=("joint.*",),
+    frame_name="ee",
+    frame_type="body",
+    orientation_weight=0.0,
+    use_relative_mode=False,
+    damping=0.05,
+  ).build(env)
+  action_lim: DifferentialIKAction = DifferentialIKActionCfg(
+    entity_name="robot",
+    actuator_names=("joint.*",),
+    frame_name="ee",
+    frame_type="body",
+    orientation_weight=0.0,
+    use_relative_mode=False,
+    damping=0.05,
+    joint_limit_weight=10.0,
+  ).build(env)
+
+  initial_pos, _ = action_no_lim._get_frame_pose()
+  target = initial_pos.clone()
+  target[:, 0] += 0.15
+
+  limits = entity.data.soft_joint_pos_limits[:, action_lim._joint_ids]
+  for action in (action_no_lim, action_lim):
+    entity.write_joint_position_to_sim(
+      torch.zeros(NUM_ENVS, 3, device=device),
+      joint_ids=action._joint_ids,
+    )
+    sim.forward()
+    action.process_actions(target)
+    for _ in range(10):
+      dq = action.compute_dq()
+      q = entity.data.joint_pos[:, action._joint_ids] + dq
+      entity.write_joint_position_to_sim(q, joint_ids=action._joint_ids)
+      sim.forward()
+
+  q_lim = entity.data.joint_pos[:, action_lim._joint_ids]
+  upper_viol = (q_lim - limits[..., 1]).clamp(min=0).sum()
+  lower_viol = (limits[..., 0] - q_lim).clamp(min=0).sum()
+  total_violation = upper_viol + lower_viol
+  assert total_violation < 0.01, f"Joint limit violation too large: {total_violation}"
+
+
+def _quat_conjugate(q: torch.Tensor) -> torch.Tensor:
+  return torch.cat([q[..., :1], -q[..., 1:]], dim=-1)
+
+
+def _central_fd_position(action, entity, sim, q0, eps):
+  """Central finite-difference position Jacobian."""
+  nj = action._num_joints
+  jacp_fd = torch.zeros(q0.shape[0], 3, nj, device=q0.device)
+  for i in range(nj):
+    q_plus = q0.clone()
+    q_plus[:, i] += eps
+    entity.write_joint_position_to_sim(q_plus, joint_ids=action._joint_ids)
+    sim.forward()
+    pos_plus = action._get_frame_pose()[0].clone()
+
+    q_minus = q0.clone()
+    q_minus[:, i] -= eps
+    entity.write_joint_position_to_sim(q_minus, joint_ids=action._joint_ids)
+    sim.forward()
+    pos_minus = action._get_frame_pose()[0].clone()
+
+    jacp_fd[:, :, i] = (pos_plus - pos_minus) / (2 * eps)
+  return jacp_fd
+
+
+def _central_fd_orientation(action, entity, sim, q0, quat0, eps):
+  """Central finite-difference rotation Jacobian via axis-angle."""
+  nj = action._num_joints
+  jacr_fd = torch.zeros(q0.shape[0], 3, nj, device=q0.device)
+  for i in range(nj):
+    q_plus = q0.clone()
+    q_plus[:, i] += eps
+    entity.write_joint_position_to_sim(q_plus, joint_ids=action._joint_ids)
+    sim.forward()
+    _, quat_plus = action._get_frame_pose()
+    q_diff_plus = quat_mul(quat_plus, _quat_conjugate(quat0))
+    aa_plus = axis_angle_from_quat(q_diff_plus)
+
+    q_minus = q0.clone()
+    q_minus[:, i] -= eps
+    entity.write_joint_position_to_sim(q_minus, joint_ids=action._joint_ids)
+    sim.forward()
+    _, quat_minus = action._get_frame_pose()
+    q_diff_minus = quat_mul(quat_minus, _quat_conjugate(quat0))
+    aa_minus = axis_angle_from_quat(q_diff_minus)
+
+    jacr_fd[:, :, i] = (aa_plus - aa_minus) / (2 * eps)
+  return jacr_fd
+
+
+# Float32 central FD: truncation O(eps^2) ~ 1e-6, round-off O(u/eps) ~ 1e-4.
+_FD_EPS = 1e-3
+_FD_ATOL = 1e-3
+
+
+def test_jacobian_finite_difference_position(device):
+  """Verify position Jacobian against central finite differences."""
+  entity, sim = _make_entity(device)
+  env = _make_env(entity, sim, device)
+  cfg = DifferentialIKActionCfg(
+    entity_name="robot",
+    actuator_names=("joint.*",),
+    frame_name="ee",
+    frame_type="body",
+    orientation_weight=0.0,
+    use_relative_mode=False,
+    damping=0.05,
+  )
+  action: DifferentialIKAction = cfg.build(env)
+
+  q0 = 0.3 * torch.tensor([[0.5, -0.3, 0.8]] * NUM_ENVS, device=device)
+  entity.write_joint_position_to_sim(q0, joint_ids=action._joint_ids)
+  sim.forward()
+
+  pos0 = action._get_frame_pose()[0].clone()
+  action._point_torch[:] = pos0
+  action._compute_jacobian()
+  jacp_analytic = action._jacp_torch[:, :, action._joint_dof_ids].clone()
+
+  jacp_fd = _central_fd_position(action, entity, sim, q0, _FD_EPS)
+
+  assert torch.allclose(jacp_analytic, jacp_fd, atol=_FD_ATOL), (
+    f"Position Jacobian mismatch.\n"
+    f"Max error: {(jacp_analytic - jacp_fd).abs().max():.2e}"
+  )
+
+
+def test_jacobian_finite_difference_orientation(device):
+  """Verify rotation Jacobian against central finite differences."""
+  entity, sim = _make_entity(device)
+  env = _make_env(entity, sim, device)
+  cfg = DifferentialIKActionCfg(
+    entity_name="robot",
+    actuator_names=("joint.*",),
+    frame_name="ee",
+    frame_type="body",
+    use_relative_mode=False,
+    damping=0.05,
+  )
+  action: DifferentialIKAction = cfg.build(env)
+
+  q0 = 0.3 * torch.tensor([[0.5, -0.3, 0.8]] * NUM_ENVS, device=device)
+  entity.write_joint_position_to_sim(q0, joint_ids=action._joint_ids)
+  sim.forward()
+
+  _, quat0 = action._get_frame_pose()
+  quat0 = quat0.clone()
+
+  pos0 = action._get_frame_pose()[0].clone()
+  action._point_torch[:] = pos0
+  action._compute_jacobian()
+  jacr_analytic = action._jacr_torch[:, :, action._joint_dof_ids].clone()
+
+  jacr_fd = _central_fd_orientation(action, entity, sim, q0, quat0, _FD_EPS)
+
+  assert torch.allclose(jacr_analytic, jacr_fd, atol=_FD_ATOL), (
+    f"Rotation Jacobian mismatch.\n"
+    f"Max error: {(jacr_analytic - jacr_fd).abs().max():.2e}"
+  )
+
+
+def test_posture_regularization(device):
+  """Posture weight biases joints toward target when task error is zero."""
+  entity, sim = _make_entity(device)
+  env = _make_env(entity, sim, device)
+
+  cfg = DifferentialIKActionCfg(
+    entity_name="robot",
+    actuator_names=("joint.*",),
+    frame_name="ee",
+    frame_type="body",
+    orientation_weight=0.0,
+    use_relative_mode=False,
+    damping=0.05,
+    posture_weight=0.1,
+    posture_target={".*": 0.0},
+  )
+  action: DifferentialIKAction = cfg.build(env)
+
+  q0 = torch.tensor([[0.3, -0.2, 0.5]] * NUM_ENVS, device=device)
+  entity.write_joint_position_to_sim(q0, joint_ids=action._joint_ids)
+  sim.forward()
+
+  # Target = current frame position, so only posture term drives motion.
+  frame_pos, _ = action._get_frame_pose()
+  action.process_actions(frame_pos.clone())
+
+  for _ in range(10):
+    dq = action.compute_dq()
+    q = entity.data.joint_pos[:, action._joint_ids] + dq
+    entity.write_joint_position_to_sim(q, joint_ids=action._joint_ids)
+    sim.forward()
+
+  q_final = entity.data.joint_pos[:, action._joint_ids]
+  assert q_final.abs().mean() < q0.abs().mean(), (
+    f"Posture regularization did not pull toward target.\n"
+    f"Initial mean |q|: {q0.abs().mean():.4f}, "
+    f"Final mean |q|: {q_final.abs().mean():.4f}"
+  )
@@ -0,0 +1,304 @@
+"""Tests for encoder bias simulation.
+
+We simulate this by storing `encoder_bias` per joint in EntityData:
+  - Observations: policy sees (true_position + bias)
+  - Actions: position commands are converted via (command - bias) before simulation
+
+This ensures joint limits apply to the true physical position, not biased readings.
+"""
+
+from collections.abc import Callable
+from functools import partial
+
+import mujoco
+import pytest
+import torch
+from conftest import get_test_device
+
+from mjlab.actuator import BuiltinPositionActuatorCfg
+from mjlab.entity import EntityArticulationInfoCfg, EntityCfg
+from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg, mdp
+from mjlab.managers.event_manager import EventTermCfg
+from mjlab.managers.observation_manager import ObservationGroupCfg, ObservationTermCfg
+from mjlab.managers.scene_entity_config import SceneEntityCfg
+from mjlab.scene import SceneCfg
+from mjlab.sim import MujocoCfg, SimulationCfg
+from mjlab.terrains import TerrainEntityCfg
+
+# =============================================================================
+# Test fixtures and helpers
+# =============================================================================
+
+SLIDING_MASS_XML = """
+<mujoco>
+  <option timestep="0.002"/>
+  <worldbody>
+    <body name="mass" pos="0 0 0">
+      <joint name="slide" type="slide" axis="1 0 0" range="-1 1" limited="true"/>
+      <geom name="mass_geom" type="sphere" size="0.1" mass="1.0"/>
+    </body>
+  </worldbody>
+  <sensor>
+    <jointpos name="slide_pos" joint="slide"/>
+  </sensor>
+</mujoco>
+"""
+
+
+def _make_robot_cfg():
+  return EntityCfg(
+    spec_fn=lambda: mujoco.MjSpec.from_string(SLIDING_MASS_XML),
+    articulation=EntityArticulationInfoCfg(
+      actuators=(
+        BuiltinPositionActuatorCfg(
+          target_names_expr=(".*",), stiffness=1000.0, damping=100.0
+        ),
+      )
+    ),
+  )
+
+
+def _make_env_cfg(
+  obs_func: Callable | None = None,
+  num_envs: int = 2,
+  events: dict | None = None,
+) -> ManagerBasedRlEnvCfg:
+  if obs_func is None:
+    obs_func = partial(mdp.joint_pos_rel, biased=True)
+  return ManagerBasedRlEnvCfg(
+    scene=SceneCfg(
+      terrain=TerrainEntityCfg(terrain_type="plane"),
+      num_envs=num_envs,
+      extent=1.0,
+      entities={"robot": _make_robot_cfg()},
+    ),
+    observations={
+      "actor": ObservationGroupCfg(
+        terms={"obs": ObservationTermCfg(func=obs_func)},
+      ),
+    },
+    actions={
+      "joint_pos": mdp.JointPositionActionCfg(
+        entity_name="robot", actuator_names=(".*",), scale=1.0
+      )
+    },
+    events=events or {},
+    sim=SimulationCfg(mujoco=MujocoCfg(timestep=0.002, iterations=1)),
+    decimation=1,
+    episode_length_s=10.0,
+  )
+
+
+@pytest.fixture(scope="module")
+def device():
+  return get_test_device()
+
+
+# =============================================================================
+# EntityData.encoder_bias field tests
+# =============================================================================
+
+
+def test_encoder_bias_initialized_to_zero(device):
+  """encoder_bias should be zeros with shape (num_envs, num_joints)."""
+  env = ManagerBasedRlEnv(cfg=_make_env_cfg(num_envs=4), device=device)
+  env.reset()
+
+  robot = env.scene["robot"]
+  assert robot.data.encoder_bias.shape == (4, 1)
+  assert (robot.data.encoder_bias == 0).all()
+
+  env.close()
+
+
+def test_encoder_bias_can_be_set_per_env(device):
+  """encoder_bias can be set differently per environment."""
+  env = ManagerBasedRlEnv(cfg=_make_env_cfg(), device=device)
+  env.reset()
+
+  robot = env.scene["robot"]
+  robot.data.encoder_bias[0, 0] = 0.1
+  robot.data.encoder_bias[1, 0] = -0.2
+
+  assert robot.data.encoder_bias[0, 0].item() == pytest.approx(0.1)
+  assert robot.data.encoder_bias[1, 0].item() == pytest.approx(-0.2)
+
+  env.close()
+
+
+def test_joint_pos_biased_equals_joint_pos_plus_bias(device):
+  """joint_pos_biased property should return joint_pos + encoder_bias."""
+  env = ManagerBasedRlEnv(cfg=_make_env_cfg(), device=device)
+  env.reset()
+
+  robot = env.scene["robot"]
+  robot.data.encoder_bias[:, 0] = 0.25
+
+  torch.testing.assert_close(
+    robot.data.joint_pos_biased, robot.data.joint_pos + robot.data.encoder_bias
+  )
+
+  env.close()
+
+
+# =============================================================================
+# Observation tests
+# =============================================================================
+
+
+def test_joint_pos_rel_includes_encoder_bias(device):
+  """joint_pos_rel observation should include encoder_bias."""
+  env = ManagerBasedRlEnv(cfg=_make_env_cfg(), device=device)
+  env.reset()
+
+  robot = env.scene["robot"]
+  bias = 0.5
+
+  obs_before = env.observation_manager.compute()["actor"]
+  assert isinstance(obs_before, torch.Tensor)
+  obs_before = obs_before.clone()
+  robot.data.encoder_bias[:, 0] = bias
+  env.observation_manager._obs_buffer = None  # Invalidate cache.
+  obs_after = env.observation_manager.compute()["actor"]
+
+  # Observation should increase by bias amount.
+  torch.testing.assert_close(obs_after, obs_before + bias, atol=1e-5, rtol=0)
+
+  env.close()
+
+
+def test_joint_vel_rel_ignores_encoder_bias(device):
+  """joint_vel_rel should NOT include encoder_bias (bias is constant, d/dt = 0)."""
+  env = ManagerBasedRlEnv(cfg=_make_env_cfg(obs_func=mdp.joint_vel_rel), device=device)
+  env.reset()
+
+  robot = env.scene["robot"]
+
+  obs_before = env.observation_manager.compute()["actor"]
+  assert isinstance(obs_before, torch.Tensor)
+  obs_before = obs_before.clone()
+  robot.data.encoder_bias[:, 0] = 0.5
+  env.observation_manager._obs_buffer = None
+  obs_after = env.observation_manager.compute()["actor"]
+
+  torch.testing.assert_close(obs_before, obs_after, atol=1e-6, rtol=0)
+
+  env.close()
+
+
+# =============================================================================
+# Action tests
+# =============================================================================
+
+
+def test_position_action_subtracts_encoder_bias(device):
+  """JointPositionAction should subtract encoder_bias from commands."""
+  env = ManagerBasedRlEnv(cfg=_make_env_cfg(), device=device)
+  env.reset()
+
+  robot = env.scene["robot"]
+  robot.data.encoder_bias[0, 0] = 0.0
+  robot.data.encoder_bias[1, 0] = 0.3
+
+  # Same command to both envs.
+  env.step(torch.tensor([[0.5], [0.5]], device=device))
+
+  # Actual targets should differ by bias.
+  assert robot.data.joint_pos_target[0, 0].item() == pytest.approx(0.5, abs=1e-5)
+  assert robot.data.joint_pos_target[1, 0].item() == pytest.approx(0.2, abs=1e-5)
+
+  env.close()
+
+
+# =============================================================================
+# End-to-end consistency test
+# =============================================================================
+
+
+def test_bias_compensation_produces_identical_physical_behavior(device):
+  """Envs with different biases should have identical physics when compensated.
+
+  If env0 has bias=0 and env1 has bias=0.3, and we want both to reach physical
+  position 0.4, we command:
+    - env0: 0.4 (in encoder frame)
+    - env1: 0.7 (in encoder frame, which becomes 0.4 after bias subtraction)
+
+  Both should reach the same physical position. Their observations should differ
+  by the bias amount.
+  """
+  env = ManagerBasedRlEnv(cfg=_make_env_cfg(), device=device)
+  env.reset()
+
+  robot = env.scene["robot"]
+  bias_env0, bias_env1 = 0.0, 0.3
+  robot.data.encoder_bias[0, 0] = bias_env0
+  robot.data.encoder_bias[1, 0] = bias_env1
+
+  # Command same physical target via encoder frame.
+  target_physical = 0.4
+  action = torch.tensor(
+    [[target_physical + bias_env0], [target_physical + bias_env1]], device=device
+  )
+
+  for _ in range(100):
+    env.step(action)
+
+  # Physical positions should match.
+  pos_env0 = robot.data.joint_pos[0, 0].item()
+  pos_env1 = robot.data.joint_pos[1, 0].item()
+  assert pos_env0 == pytest.approx(pos_env1, abs=1e-4)
+
+  # Observations should differ by bias.
+  env.observation_manager._obs_buffer = None
+  obs = env.observation_manager.compute()["actor"]
+  assert isinstance(obs, torch.Tensor)
+  obs_diff = obs[1, 0].item() - obs[0, 0].item()
+  assert obs_diff == pytest.approx(bias_env1 - bias_env0, abs=1e-4)
+
+  env.close()
+
+
+# =============================================================================
+# Event randomization test
+# =============================================================================
+
+
+def test_randomize_encoder_bias_event(device):
+  """randomize_encoder_bias should sample values within specified range."""
+  env_cfg = ManagerBasedRlEnvCfg(
+    scene=SceneCfg(
+      terrain=TerrainEntityCfg(terrain_type="plane"),
+      num_envs=100,
+      extent=10.0,
+      entities={"robot": _make_robot_cfg()},
+    ),
+    observations={
+      "actor": ObservationGroupCfg(
+        terms={"obs": ObservationTermCfg(func=partial(mdp.joint_pos_rel, biased=True))},
+      ),
+    },
+    actions={
+      "joint_pos": mdp.JointPositionActionCfg(
+        entity_name="robot", actuator_names=(".*",), scale=1.0
+      )
+    },
+    events={
+      "randomize_bias": EventTermCfg(
+        func=mdp.dr.encoder_bias,
+        mode="startup",
+        params={"bias_range": (-0.1, 0.1), "asset_cfg": SceneEntityCfg("robot")},
+      )
+    },
+    sim=SimulationCfg(mujoco=MujocoCfg(timestep=0.002, iterations=1)),
+    decimation=1,
+    episode_length_s=10.0,
+  )
+
+  env = ManagerBasedRlEnv(cfg=env_cfg, device=device)
+  env.reset()
+
+  biases = env.scene["robot"].data.encoder_bias[:, 0]
+  assert (biases >= -0.1).all() and (biases <= 0.1).all()
+  assert biases.std() > 0.01  # Not all identical.
+
+  env.close()
@@ -0,0 +1,870 @@
+"""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)
@@ -0,0 +1,581 @@
+"""Tests for EntityData."""
+
+import math
+from pathlib import Path
+
+import mujoco
+import numpy as np
+import pytest
+import torch
+from conftest import (
+  create_entity_with_actuator,
+  get_test_device,
+  initialize_entity,
+  load_fixture_xml,
+)
+
+from mjlab.actuator import BuiltinMotorActuatorCfg
+from mjlab.actuator.actuator import TransmissionType
+from mjlab.entity import Entity, EntityArticulationInfoCfg, EntityCfg
+from mjlab.sim.sim import Simulation, SimulationCfg
+from mjlab.utils.lab_api.math import quat_apply, quat_apply_inverse
+
+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>
+"""
+
+FLOATING_BASE_COM_OFFSET_XML = """
+<mujoco>
+  <worldbody>
+    <body name="object" pos="0 0 1">
+      <freejoint name="free_joint"/>
+      <inertial pos="0.1 0.05 0" mass="1" diaginertia="0.01 0.01 0.01"/>
+      <geom type="box" size="0.1 0.1 0.1"/>
+    </body>
+  </worldbody>
+</mujoco>
+"""
+
+
+@pytest.fixture(scope="module")
+def device():
+  """Test device fixture."""
+  return get_test_device()
+
+
+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 initialize_entity_with_sim(entity, device, num_envs=1):
+  """Initialize an entity with a simulation."""
+  model = entity.compile()
+  sim_cfg = SimulationCfg()
+  sim = Simulation(num_envs=num_envs, cfg=sim_cfg, model=model, device=device)
+  entity.initialize(model, sim.model, sim.data, device)
+  return entity, sim
+
+
+def test_root_velocity_world_frame_roundtrip(device):
+  """Test reading and writing root velocity is a no-op (world frame).
+
+  Verifies that the API is consistent: if you write a velocity, read it
+  back, and write it again, you get the same result. This ensures no
+  unintended transformations happen in the read/write cycle.
+  """
+  entity = create_floating_base_entity()
+  entity, sim = initialize_entity_with_sim(entity, device)
+
+  pose = torch.tensor([0.0, 0.0, 1.0, 0.6, 0.2, 0.3, 0.7141], device=device).unsqueeze(
+    0
+  )
+  entity.write_root_link_pose_to_sim(pose)
+
+  vel_w = torch.tensor([1.0, 0.5, 0.0, 0.0, 0.3, 0.1], device=device).unsqueeze(0)
+  entity.write_root_link_velocity_to_sim(vel_w)
+  sim.forward()
+
+  vel_w_read = entity.data.root_link_vel_w.clone()
+  assert torch.allclose(vel_w_read, vel_w, atol=1e-4)
+
+  entity.write_root_link_velocity_to_sim(vel_w_read)
+  sim.forward()
+  vel_w_after = entity.data.root_link_vel_w
+
+  assert torch.allclose(vel_w_after, vel_w_read, atol=1e-4)
+
+
+def test_root_velocity_frame_conversion(device):
+  """Test angular velocity converts from world to body frame internally.
+
+  The API accepts angular velocity in world frame, but MuJoCo's qvel
+  stores it in body frame. This test verifies the conversion happens
+  correctly by checking qvel directly.
+  """
+  entity = create_floating_base_entity()
+  entity, sim = initialize_entity_with_sim(entity, device)
+
+  quat_w = torch.tensor([0.6, 0.2, 0.3, 0.7141], device=device).unsqueeze(0)
+  pose = torch.cat([torch.zeros(1, 3, device=device), quat_w], dim=-1)
+  entity.write_root_link_pose_to_sim(pose)
+
+  lin_vel_w = torch.tensor([1.0, 0.5, 0.2], device=device).unsqueeze(0)
+  ang_vel_w = torch.tensor([0.1, 0.2, 0.3], device=device).unsqueeze(0)
+  vel_w = torch.cat([lin_vel_w, ang_vel_w], dim=-1)
+  entity.write_root_link_velocity_to_sim(vel_w)
+
+  v_slice = entity.data.indexing.free_joint_v_adr
+  qvel = sim.data.qvel[:, v_slice]
+
+  assert torch.allclose(qvel[:, :3], lin_vel_w, atol=1e-5)
+
+  expected_ang_vel_b = quat_apply_inverse(quat_w, ang_vel_w)
+  assert torch.allclose(qvel[:, 3:], expected_ang_vel_b, atol=1e-5)
+
+
+def test_write_velocity_uses_qpos_not_xquat(device):
+  """Test write_root_velocity uses qpos (not stale xquat).
+
+  Writing pose then velocity without forward() must work. This would fail
+  if write_root_velocity used xquat (stale) instead of qpos (current).
+  """
+  entity = create_floating_base_entity()
+  entity, sim = initialize_entity_with_sim(entity, device)
+
+  initial_pose = torch.tensor(
+    [0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 0.0], device=device
+  ).unsqueeze(0)
+  entity.write_root_link_pose_to_sim(initial_pose)
+  sim.forward()  # xquat now has identity orientation.
+
+  # Write different orientation without forward() - xquat stale, qpos current.
+  new_pose = torch.tensor(
+    [0.0, 0.0, 1.0, 0.707, 0.0, 0.707, 0.0], device=device
+  ).unsqueeze(0)
+  vel_w = torch.tensor([1.0, 0.0, 0.0, 0.0, 1.0, 0.0], device=device).unsqueeze(0)
+
+  entity.write_root_link_pose_to_sim(new_pose)
+  entity.write_root_link_velocity_to_sim(vel_w)
+
+  sim.forward()
+  vel_w_read = entity.data.root_link_vel_w
+
+  assert torch.allclose(vel_w_read, vel_w, atol=1e-4)
+
+
+def test_write_root_com_velocity(device):
+  """COM velocity write must produce the same qvel as manual conversion."""
+  num_envs = 4
+  cfg = EntityCfg(
+    spec_fn=lambda: mujoco.MjSpec.from_string(FLOATING_BASE_COM_OFFSET_XML)
+  )
+  entity = Entity(cfg)
+  entity, sim = initialize_entity_with_sim(entity, device, num_envs=num_envs)
+
+  # Give each env a different non-identity orientation.
+  pose = torch.zeros(num_envs, 7, device=device)
+  pose[:, 2] = 1.0
+  quats = torch.tensor(
+    [
+      [1.0, 0.0, 0.0, 0.0],
+      [0.707, 0.707, 0.0, 0.0],
+      [0.5, 0.5, 0.5, 0.5],
+      [0.0, 0.707, 0.0, 0.707],
+    ],
+    device=device,
+  )
+  pose[:, 3:7] = quats
+  entity.write_root_link_pose_to_sim(pose)
+
+  com_vel = torch.tensor(
+    [[1.0, 0.5, -0.2, 0.1, 0.2, 0.3], [-0.5, 1.0, 0.0, -0.1, 0.0, 0.4]],
+    device=device,
+  )
+  env_ids = torch.tensor([1, 3], device=device)
+
+  # Write COM velocity via the API.
+  entity.write_root_com_velocity_to_sim(com_vel, env_ids=env_ids)
+  qvel_from_api = sim.data.qvel.clone()
+
+  # Manually convert COM velocity to link velocity and write that instead.
+  com_offset_b = sim.model.body_ipos[:, entity.indexing.root_body_id]
+  com_offset_w = quat_apply(quats[env_ids], com_offset_b[env_ids])
+  lin_vel_link = com_vel[:, :3] - torch.cross(com_vel[:, 3:], com_offset_w, dim=-1)
+  link_vel = torch.cat([lin_vel_link, com_vel[:, 3:]], dim=-1)
+  entity.write_root_link_velocity_to_sim(link_vel, env_ids=env_ids)
+  qvel_from_manual = sim.data.qvel.clone()
+
+  assert torch.allclose(qvel_from_api, qvel_from_manual, atol=1e-5)
+
+  # Untouched envs should have zero velocity.
+  v_adr = entity.indexing.free_joint_v_adr
+  assert torch.all(qvel_from_api[0, v_adr] == 0.0)
+  assert torch.all(qvel_from_api[2, v_adr] == 0.0)
+
+
+def test_read_requires_forward_to_be_current(device):
+  """Test read properties are stale until forward() is called.
+
+  Demonstrates why event order matters and why forward() is needed
+  between writes and reads.
+  """
+  entity = create_floating_base_entity()
+  entity, sim = initialize_entity_with_sim(entity, device)
+
+  sim.forward()
+  initial_pose = entity.data.root_link_pose_w.clone()
+
+  new_pose = torch.tensor(
+    [1.0, 2.0, 3.0, 0.707, 0.0, 0.707, 0.0], device=device
+  ).unsqueeze(0)
+  entity.write_root_link_pose_to_sim(new_pose)
+
+  stale_pose = entity.data.root_link_pose_w
+  assert torch.allclose(stale_pose, initial_pose, atol=1e-5)
+
+  sim.forward()
+  current_pose = entity.data.root_link_pose_w
+  assert torch.allclose(current_pose, new_pose, atol=1e-4)
+  assert not torch.allclose(current_pose, initial_pose, atol=1e-4)
+
+
+@pytest.mark.parametrize(
+  "property_name,expected_shape",
+  [
+    # Root properties.
+    ("root_link_pose_w", (1, 7)),
+    ("root_link_pos_w", (1, 3)),
+    ("root_link_quat_w", (1, 4)),
+    ("root_link_vel_w", (1, 6)),
+    ("root_link_lin_vel_w", (1, 3)),
+    ("root_link_ang_vel_w", (1, 3)),
+    ("root_com_pose_w", (1, 7)),
+    ("root_com_pos_w", (1, 3)),
+    ("root_com_quat_w", (1, 4)),
+    ("root_com_vel_w", (1, 6)),
+    ("root_com_lin_vel_w", (1, 3)),
+    ("root_com_ang_vel_w", (1, 3)),
+    # Body properties (we only have 1 body in this test).
+    ("body_link_pose_w", (1, 1, 7)),
+    ("body_link_pos_w", (1, 1, 3)),
+    ("body_link_quat_w", (1, 1, 4)),
+    ("body_link_vel_w", (1, 1, 6)),
+    ("body_com_pose_w", (1, 1, 7)),
+    ("body_com_vel_w", (1, 1, 6)),
+  ],
+)
+def test_entity_data_properties_accessible(device, property_name, expected_shape):
+  """Test that all EntityData properties can be accessed without errors."""
+  entity = create_floating_base_entity()
+  entity, sim = initialize_entity_with_sim(entity, device)
+
+  sim.forward()
+
+  value = getattr(entity.data, property_name)
+  assert value.shape == expected_shape, f"{property_name} has unexpected shape"
+
+
+def test_entity_data_reset_clears_all_targets(device):
+  """Test that EntityData.clear_state() zeros out all target buffers."""
+  xml_path = Path(__file__).parent / "fixtures" / "tendon_finger.xml"
+
+  cfg = EntityCfg(
+    spec_fn=lambda: mujoco.MjSpec.from_file(str(xml_path)),
+    articulation=EntityArticulationInfoCfg(
+      actuators=(
+        BuiltinMotorActuatorCfg(
+          target_names_expr=("finger_tendon",),
+          transmission_type=TransmissionType.TENDON,
+          effort_limit=10.0,
+        ),
+      )
+    ),
+  )
+  entity = Entity(cfg)
+  entity, sim = initialize_entity_with_sim(entity, device, num_envs=4)
+
+  entity.data.joint_pos_target[:] = 1.0
+  entity.data.joint_vel_target[:] = 2.0
+  entity.data.joint_effort_target[:] = 3.0
+  entity.data.tendon_len_target[:] = 4.0
+  entity.data.tendon_vel_target[:] = 5.0
+  entity.data.tendon_effort_target[:] = 6.0
+
+  entity.data.clear_state()
+
+  assert torch.all(entity.data.joint_pos_target == 0.0)
+  assert torch.all(entity.data.joint_vel_target == 0.0)
+  assert torch.all(entity.data.joint_effort_target == 0.0)
+  assert torch.all(entity.data.tendon_len_target == 0.0)
+  assert torch.all(entity.data.tendon_vel_target == 0.0)
+  assert torch.all(entity.data.tendon_effort_target == 0.0)
+
+
+def test_entity_data_reset_partial_envs(device):
+  """Test that EntityData.clear_state() can reset specific environments."""
+  xml_path = Path(__file__).parent / "fixtures" / "tendon_finger.xml"
+
+  cfg = EntityCfg(
+    spec_fn=lambda: mujoco.MjSpec.from_file(str(xml_path)),
+    articulation=EntityArticulationInfoCfg(
+      actuators=(
+        BuiltinMotorActuatorCfg(
+          target_names_expr=("finger_tendon",),
+          transmission_type=TransmissionType.TENDON,
+          effort_limit=10.0,
+        ),
+      )
+    ),
+  )
+  entity = Entity(cfg)
+  entity, sim = initialize_entity_with_sim(entity, device, num_envs=4)
+
+  entity.data.tendon_len_target[:] = 7.0
+  entity.data.tendon_vel_target[:] = 8.0
+  entity.data.tendon_effort_target[:] = 9.0
+
+  env_ids = torch.tensor([1, 3], device=device)
+  entity.data.clear_state(env_ids)
+
+  assert torch.all(entity.data.tendon_len_target[0] == 7.0)
+  assert torch.all(entity.data.tendon_len_target[1] == 0.0)
+  assert torch.all(entity.data.tendon_len_target[2] == 7.0)
+  assert torch.all(entity.data.tendon_len_target[3] == 0.0)
+
+  assert torch.all(entity.data.tendon_vel_target[0] == 8.0)
+  assert torch.all(entity.data.tendon_vel_target[1] == 0.0)
+  assert torch.all(entity.data.tendon_vel_target[2] == 8.0)
+  assert torch.all(entity.data.tendon_vel_target[3] == 0.0)
+
+  assert torch.all(entity.data.tendon_effort_target[0] == 9.0)
+  assert torch.all(entity.data.tendon_effort_target[1] == 0.0)
+  assert torch.all(entity.data.tendon_effort_target[2] == 9.0)
+  assert torch.all(entity.data.tendon_effort_target[3] == 0.0)
+
+
+# Joint limits tests.
+
+MIXED_LIMITS_XML = """
+<mujoco>
+  <worldbody>
+    <body name="base">
+      <body name="cart" pos="0 0 1">
+        <joint name="slider" type="slide" axis="1 0 0" limited="true" range="-2 2"/>
+        <geom type="box" size="0.1 0.1 0.1" mass="1"/>
+        <body name="pole">
+          <joint name="hinge" type="hinge" axis="0 1 0"/>
+          <geom type="capsule" fromto="0 0 0 0 0 0.5" size="0.02" mass="0.1"/>
+        </body>
+      </body>
+    </body>
+  </worldbody>
+</mujoco>
+"""
+
+ALL_LIMITED_XML = """
+<mujoco>
+  <worldbody>
+    <body name="base">
+      <body name="link1" pos="0 0 0.1">
+        <joint name="j1" type="hinge" axis="0 0 1" limited="true" range="-1.5 1.5"/>
+        <geom type="box" size="0.05 0.05 0.1" mass="1"/>
+        <body name="link2" pos="0 0 0.2">
+          <joint name="j2" type="hinge" axis="0 1 0" limited="true" range="-1 1"/>
+          <geom type="box" size="0.05 0.05 0.1" mass="0.5"/>
+        </body>
+      </body>
+    </body>
+  </worldbody>
+</mujoco>
+"""
+
+ALL_UNLIMITED_XML = """
+<mujoco>
+  <worldbody>
+    <body name="base">
+      <body name="link1" pos="0 0 0.1">
+        <joint name="j1" type="hinge" axis="0 0 1"/>
+        <geom type="box" size="0.05 0.05 0.1" mass="1"/>
+        <body name="link2" pos="0 0 0.2">
+          <joint name="j2" type="hinge" axis="0 1 0"/>
+          <geom type="box" size="0.05 0.05 0.1" mass="0.5"/>
+        </body>
+      </body>
+    </body>
+  </worldbody>
+</mujoco>
+"""
+
+
+def test_soft_limits_unlimited_joints_are_infinite(device):
+  """Unlimited joints must have [-inf, inf] soft limits, not [0, 0]."""
+  cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(MIXED_LIMITS_XML))
+  entity = Entity(cfg)
+  entity, _ = initialize_entity_with_sim(entity, device)
+
+  limits = entity.data.soft_joint_pos_limits[0]
+  # slider (index 0) is limited: finite limits.
+  assert torch.isfinite(limits[0, 0]) and torch.isfinite(limits[0, 1])
+  assert limits[0, 0] == -2.0
+  assert limits[0, 1] == 2.0
+  # hinge (index 1) is unlimited: infinite limits.
+  assert limits[1, 0] == float("-inf")
+  assert limits[1, 1] == float("inf")
+
+
+def test_soft_limits_all_limited(device):
+  """All-limited joints should have finite soft limits."""
+  cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(ALL_LIMITED_XML))
+  entity = Entity(cfg)
+  entity, _ = initialize_entity_with_sim(entity, device)
+
+  limits = entity.data.soft_joint_pos_limits[0]
+  for j in range(limits.shape[0]):
+    assert torch.isfinite(limits[j, 0]) and torch.isfinite(limits[j, 1])
+
+
+def test_soft_limits_all_unlimited(device):
+  """All-unlimited joints should have infinite soft limits."""
+  cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(ALL_UNLIMITED_XML))
+  entity = Entity(cfg)
+  entity, _ = initialize_entity_with_sim(entity, device)
+
+  limits = entity.data.soft_joint_pos_limits[0]
+  for j in range(limits.shape[0]):
+    assert limits[j, 0] == float("-inf")
+    assert limits[j, 1] == float("inf")
+
+
+def test_joint_pos_limits_match_soft_limits_for_unlimited(device):
+  """joint_pos_limits and default_joint_pos_limits should also be inf for unlimited."""
+  cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(MIXED_LIMITS_XML))
+  entity = Entity(cfg)
+  entity, _ = initialize_entity_with_sim(entity, device)
+
+  # Hinge is joint index 1, unlimited.
+  for limits_tensor in (
+    entity.data.joint_pos_limits,
+    entity.data.default_joint_pos_limits,
+    entity.data.soft_joint_pos_limits,
+  ):
+    assert limits_tensor[0, 1, 0] == float("-inf")
+    assert limits_tensor[0, 1, 1] == float("inf")
+
+
+def test_reset_joints_by_offset_respects_unlimited(device):
+  """reset_joints_by_offset must not clamp unlimited joints to [0, 0]."""
+  cfg = EntityCfg(
+    spec_fn=lambda: mujoco.MjSpec.from_string(MIXED_LIMITS_XML),
+    init_state=EntityCfg.InitialStateCfg(
+      joint_pos={"slider": 0.0, "hinge": math.pi},
+      joint_vel={".*": 0.0},
+    ),
+  )
+  entity = Entity(cfg)
+  entity, _ = initialize_entity_with_sim(entity, device)
+
+  # Simulate what reset_joints_by_offset does.
+  default_pos = entity.data.default_joint_pos
+  soft_limits = entity.data.soft_joint_pos_limits
+  joint_pos = default_pos.clone()
+  joint_pos = joint_pos.clamp_(soft_limits[..., 0], soft_limits[..., 1])
+
+  # Hinge should stay at pi, not get clamped to 0.
+  assert torch.allclose(
+    joint_pos[0, 1], torch.tensor(math.pi, device=device), atol=1e-5
+  )
+
+
+# Generalized force accessor tests.
+
+TWO_LINK_ARM_XML = """
+<mujoco>
+  <worldbody>
+    <body name="base">
+      <body name="upper" pos="0 0 0.5">
+        <joint name="shoulder" type="hinge" axis="0 1 0"/>
+        <geom type="capsule" fromto="0 0 0 0 0 0.4" size="0.03" mass="1.0"/>
+        <body name="lower" pos="0 0 0.4">
+          <joint name="elbow" type="hinge" axis="0 1 0"/>
+          <geom type="capsule" fromto="0 0 0 0 0 0.3" size="0.02" mass="0.5"/>
+        </body>
+      </body>
+    </body>
+  </worldbody>
+</mujoco>
+"""
+
+
+def test_qfrc_actuator_slices_joint_dofs_only(device):
+  """qfrc_actuator should expose only articulated joint DoFs."""
+  xml = load_fixture_xml("floating_base_articulated")
+  cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(xml))
+  entity = Entity(cfg)
+  entity, sim = initialize_entity_with_sim(entity, device)
+
+  joint_v_adr = entity.indexing.joint_v_adr
+  free_v_adr = entity.indexing.free_joint_v_adr
+  nv = sim.data.qvel.shape[1]
+
+  values = torch.arange(1, nv + 1, device=device, dtype=torch.float32)
+  sim.data.qfrc_actuator[0] = values
+  expected = values[joint_v_adr]
+  actual = entity.data.qfrc_actuator[0]
+  assert torch.equal(actual, expected)
+
+  assert free_v_adr.numel() > 0
+  assert entity.data.qfrc_actuator.shape[-1] == joint_v_adr.numel()
+  assert entity.data.qfrc_actuator.shape[-1] != free_v_adr.numel()
+
+
+def test_qfrc_actuator_matches_motor_command_with_gear(device):
+  """qfrc_actuator should equal actuator_force projected through joint gear."""
+  xml = load_fixture_xml("fixed_base_articulated")
+  entity = create_entity_with_actuator(
+    xml,
+    BuiltinMotorActuatorCfg(
+      target_names_expr=("joint.*",),
+      effort_limit=100.0,
+      gear=3.0,
+    ),
+  )
+  entity, sim = initialize_entity(entity, device)
+
+  commanded = torch.tensor([[2.0, -1.0]], device=device)
+  entity.set_joint_effort_target(commanded)
+  entity.write_data_to_sim()
+  sim.forward()
+
+  assert torch.allclose(entity.data.actuator_force, commanded, atol=1e-6)
+  assert torch.allclose(entity.data.qfrc_actuator, 3.0 * commanded, atol=1e-6)
+
+
+def test_qfrc_external_matches_mj_applyFT(device):
+  """qfrc_external should match J^T * xfrc_applied computed by CPU MuJoCo."""
+  cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(TWO_LINK_ARM_XML))
+  entity = Entity(cfg)
+  entity, sim = initialize_entity_with_sim(entity, device)
+
+  entity.write_joint_state_to_sim(
+    torch.tensor([[0.5, -0.3]], device=device),
+    torch.tensor([[0.0, 0.0]], device=device),
+  )
+  sim.forward()
+
+  lower_body_id = int(entity.indexing.body_ids[1].item())
+  force = torch.tensor([1.0, 0.0, -2.0], device=device)
+  torque = torch.tensor([0.0, 0.5, 0.0], device=device)
+  sim.data.xfrc_applied[0, lower_body_id, 0:3] = force
+  sim.data.xfrc_applied[0, lower_body_id, 3:6] = torque
+  sim.forward()
+
+  mjd = mujoco.MjData(sim.mj_model)
+  mjd.qpos[:] = sim.data.qpos[0].cpu().numpy()
+  mjd.qvel[:] = sim.data.qvel[0].cpu().numpy()
+  mujoco.mj_forward(sim.mj_model, mjd)
+
+  qfrc_reference = np.zeros(sim.mj_model.nv)
+  mujoco.mj_applyFT(
+    sim.mj_model,
+    mjd,
+    force.cpu().numpy().astype(np.float64),
+    torque.cpu().numpy().astype(np.float64),
+    mjd.xipos[lower_body_id],
+    lower_body_id,
+    qfrc_reference,
+  )
+  expected = torch.from_numpy(qfrc_reference).to(device=device, dtype=torch.float32)
+
+  joint_v_adr = entity.indexing.joint_v_adr
+  assert torch.allclose(
+    entity.data.qfrc_external, expected[joint_v_adr].unsqueeze(0), atol=1e-5
+  )
@@ -0,0 +1,275 @@
+"""Tests for reward_curriculum and termination_curriculum."""
+
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from mjlab.envs.mdp.curriculums import reward_curriculum, termination_curriculum
+from mjlab.managers.curriculum_manager import CurriculumTermCfg
+from mjlab.managers.reward_manager import RewardTermCfg
+from mjlab.managers.termination_manager import TerminationTermCfg
+
+
+def _reward_func(env):
+  return torch.ones(env.num_envs)
+
+
+def _termination_func(env):
+  return torch.zeros(env.num_envs, dtype=torch.bool)
+
+
+def _make_reward_cfg(
+  weight: float = 1.0,
+  params: dict | None = None,
+) -> RewardTermCfg:
+  return RewardTermCfg(
+    func=_reward_func,
+    weight=weight,
+    params=params if params is not None else {"std": 0.5, "scale": 1.0},
+  )
+
+
+def _make_termination_cfg(
+  params: dict | None = None,
+) -> TerminationTermCfg:
+  return TerminationTermCfg(
+    func=_termination_func,
+    params=params if params is not None else {"threshold": float("inf")},
+  )
+
+
+def _build_reward(env, reward_name, stages):
+  params = {"reward_name": reward_name, "stages": stages}
+  cfg = CurriculumTermCfg(func=reward_curriculum, params=params)
+  instance = reward_curriculum(cfg, env)
+  return instance(env, env_ids=torch.tensor([0, 1]), **params)
+
+
+def _build_termination(env, termination_name, stages):
+  params = {"termination_name": termination_name, "stages": stages}
+  cfg = CurriculumTermCfg(func=termination_curriculum, params=params)
+  instance = termination_curriculum(cfg, env)
+  return instance(env, env_ids=torch.tensor([0, 1]), **params)
+
+
+def _make_reward_env(step_counter, reward_cfg):
+  env = Mock()
+  env.common_step_counter = step_counter
+  env.reward_manager.get_term_cfg.return_value = reward_cfg
+  return env
+
+
+def _make_termination_env(step_counter, term_cfg):
+  env = Mock()
+  env.common_step_counter = step_counter
+  env.termination_manager.get_term_cfg.return_value = term_cfg
+  return env
+
+
+# Reward: weight
+
+
+def test_reward_weight_unchanged_before_threshold():
+  rc = _make_reward_cfg()
+  env = _make_reward_env(0, rc)
+  _build_reward(env, "r", [{"step": 100, "weight": 2.0}])
+  assert rc.weight == pytest.approx(1.0)
+
+
+def test_reward_weight_applied_at_threshold():
+  rc = _make_reward_cfg()
+  env = _make_reward_env(100, rc)
+  _build_reward(env, "r", [{"step": 100, "weight": 2.0}])
+  assert rc.weight == pytest.approx(2.0)
+
+
+def test_reward_weight_later_stage_wins():
+  rc = _make_reward_cfg()
+  env = _make_reward_env(500, rc)
+  _build_reward(
+    env,
+    "r",
+    [
+      {"step": 0, "weight": 0.5},
+      {"step": 100, "weight": 1.5},
+      {"step": 400, "weight": 3.0},
+    ],
+  )
+  assert rc.weight == pytest.approx(3.0)
+
+
+def test_reward_weight_partial_application():
+  rc = _make_reward_cfg()
+  env = _make_reward_env(150, rc)
+  _build_reward(
+    env,
+    "r",
+    [
+      {"step": 100, "weight": 2.0},
+      {"step": 200, "weight": 4.0},
+    ],
+  )
+  assert rc.weight == pytest.approx(2.0)
+
+
+def test_step_zero_applies_immediately():
+  rc = _make_reward_cfg()
+  env = _make_reward_env(0, rc)
+  _build_reward(env, "r", [{"step": 0, "weight": 9.0}])
+  assert rc.weight == pytest.approx(9.0)
+
+
+# Reward: params
+
+
+def test_reward_params_updated():
+  rc = _make_reward_cfg()
+  env = _make_reward_env(200, rc)
+  _build_reward(env, "r", [{"step": 100, "params": {"std": 0.2}}])
+  assert rc.params["std"] == 0.2
+
+
+def test_reward_params_unchanged_before_threshold():
+  rc = _make_reward_cfg()
+  env = _make_reward_env(0, rc)
+  _build_reward(env, "r", [{"step": 100, "params": {"std": 0.2}}])
+  assert rc.params["std"] == 0.5
+
+
+def test_reward_multiple_params_updated():
+  rc = _make_reward_cfg()
+  env = _make_reward_env(200, rc)
+  _build_reward(env, "r", [{"step": 100, "params": {"std": 0.2, "scale": 2.0}}])
+  assert rc.params["std"] == 0.2
+  assert rc.params["scale"] == 2.0
+
+
+# Reward: combined weight + params
+
+
+def test_reward_weight_and_params_in_same_stage():
+  rc = _make_reward_cfg()
+  env = _make_reward_env(200, rc)
+  _build_reward(env, "r", [{"step": 100, "weight": 5.0, "params": {"std": 0.1}}])
+  assert rc.weight == pytest.approx(5.0)
+  assert rc.params["std"] == 0.1
+
+
+# Termination: params
+
+
+def test_termination_params_updated():
+  tc = _make_termination_cfg()
+  env = _make_termination_env(200, tc)
+  _build_termination(env, "energy", [{"step": 100, "params": {"threshold": 500.0}}])
+  assert tc.params["threshold"] == 500.0
+
+
+def test_termination_params_unchanged_before_threshold():
+  tc = _make_termination_cfg()
+  env = _make_termination_env(0, tc)
+  _build_termination(env, "energy", [{"step": 100, "params": {"threshold": 500.0}}])
+  assert tc.params["threshold"] == float("inf")
+
+
+def test_termination_later_stage_wins():
+  tc = _make_termination_cfg()
+  env = _make_termination_env(500, tc)
+  _build_termination(
+    env,
+    "energy",
+    [
+      {"step": 0, "params": {"threshold": 1000.0}},
+      {"step": 100, "params": {"threshold": 700.0}},
+      {"step": 400, "params": {"threshold": 400.0}},
+    ],
+  )
+  assert tc.params["threshold"] == 400.0
+
+
+# Validation: shared engine
+
+
+def test_unknown_reward_param_raises():
+  rc = _make_reward_cfg()
+  env = _make_reward_env(0, rc)
+  params = {"reward_name": "r", "stages": [{"step": 0, "params": {"stdd": 0.2}}]}
+  cfg = CurriculumTermCfg(func=reward_curriculum, params=params)
+  with pytest.raises(KeyError, match="unknown param"):
+    reward_curriculum(cfg, env)
+
+
+def test_unknown_termination_param_raises():
+  tc = _make_termination_cfg()
+  env = _make_termination_env(0, tc)
+  params = {
+    "termination_name": "energy",
+    "stages": [{"step": 0, "params": {"thresholddd": 1.0}}],
+  }
+  cfg = CurriculumTermCfg(func=termination_curriculum, params=params)
+  with pytest.raises(KeyError, match="unknown param"):
+    termination_curriculum(cfg, env)
+
+
+def test_unsorted_stages_raise():
+  rc = _make_reward_cfg()
+  env = _make_reward_env(0, rc)
+  params = {
+    "reward_name": "r",
+    "stages": [
+      {"step": 200, "weight": 1.0},
+      {"step": 100, "weight": 2.0},
+    ],
+  }
+  cfg = CurriculumTermCfg(func=reward_curriculum, params=params)
+  with pytest.raises(ValueError, match="nondecreasing"):
+    reward_curriculum(cfg, env)
+
+
+def test_duplicate_steps_allowed():
+  rc = _make_reward_cfg()
+  env = _make_reward_env(200, rc)
+  _build_reward(
+    env,
+    "r",
+    [
+      {"step": 100, "weight": 2.0},
+      {"step": 100, "params": {"std": 0.1}},
+    ],
+  )
+  assert rc.weight == pytest.approx(2.0)
+  assert rc.params["std"] == 0.1
+
+
+# Logging keys
+
+
+def test_reward_logs_only_staged_keys():
+  rc = _make_reward_cfg()
+  env = _make_reward_env(200, rc)
+  result = _build_reward(
+    env, "r", [{"step": 100, "weight": 5.0, "params": {"std": 0.2}}]
+  )
+  assert result["weight"].item() == pytest.approx(5.0)
+  assert result["std"].item() == pytest.approx(0.2)
+  assert "scale" not in result  # Not in any stage.
+
+
+def test_reward_omits_weight_when_not_staged():
+  rc = _make_reward_cfg()
+  env = _make_reward_env(200, rc)
+  result = _build_reward(env, "r", [{"step": 100, "params": {"std": 0.2}}])
+  assert "weight" not in result
+  assert "std" in result
+
+
+def test_termination_log_keys():
+  tc = _make_termination_cfg()
+  env = _make_termination_env(200, tc)
+  result = _build_termination(
+    env, "energy", [{"step": 100, "params": {"threshold": 500.0}}]
+  )
+  assert "threshold" in result
+  assert result["threshold"].item() == pytest.approx(500.0)
+  assert "weight" not in result  # No weight for termination.
@@ -0,0 +1,159 @@
+"""Tests for flat patch sampling on heightfield terrains."""
+
+import numpy as np
+import pytest
+import torch
+from conftest import get_test_device
+
+from mjlab.terrains.heightfield_terrains import HfRandomUniformTerrainCfg
+from mjlab.terrains.terrain_entity import TerrainEntity, TerrainEntityCfg
+from mjlab.terrains.terrain_generator import (
+  FlatPatchSamplingCfg,
+  TerrainGeneratorCfg,
+)
+from mjlab.terrains.utils import find_flat_patches_from_heightfield
+
+
+@pytest.fixture
+def rng() -> np.random.Generator:
+  return np.random.default_rng(42)
+
+
+def test_patches_avoid_step_boundary(rng: np.random.Generator):
+  """Patches must not land near a sharp height discontinuity."""
+  heights = np.zeros((80, 80), dtype=np.float64)
+  # Step along column axis → appears on MuJoCo x-axis.
+  heights[:, 40:] = 1.0
+  patch_radius = 0.3
+  cfg = FlatPatchSamplingCfg(
+    num_patches=50, patch_radius=patch_radius, max_height_diff=0.05
+  )
+  patches = find_flat_patches_from_heightfield(
+    heights=heights, horizontal_scale=0.1, z_offset=0.0, cfg=cfg, rng=rng
+  )
+  assert patches.shape == (50, 3)
+  for i in range(len(patches)):
+    x, z = patches[i, 0], patches[i, 2]
+    # Step is at x=4.0; no patch center should be within patch_radius of it.
+    assert x < 3.7 or x >= 4.3, f"Patch {i} at x={x:.2f} too close to step"
+    # z must be consistent with the side the patch landed on.
+    if x < 3.7:
+      assert abs(z) < 0.1, f"Low-x patch should have z~0, got {z}"
+    else:
+      assert abs(z - 1.0) < 0.1, f"High-x patch should have z~1, got {z}"
+
+
+def test_range_constraints(rng: np.random.Generator):
+  """x/y range filters restrict where patches can land."""
+  heights = np.zeros((80, 80), dtype=np.float64)
+  cfg = FlatPatchSamplingCfg(
+    num_patches=10,
+    patch_radius=0.3,
+    max_height_diff=0.1,
+    x_range=(2.0, 6.0),
+    y_range=(2.0, 6.0),
+  )
+  patches = find_flat_patches_from_heightfield(
+    heights=heights, horizontal_scale=0.1, z_offset=0.0, cfg=cfg, rng=rng
+  )
+  assert np.all(patches[:, 0] >= 2.0) and np.all(patches[:, 0] <= 6.0)
+  assert np.all(patches[:, 1] >= 2.0) and np.all(patches[:, 1] <= 6.0)
+
+
+def test_fallback_when_no_valid_patches(rng: np.random.Generator):
+  """When nothing is flat, all patches fall back to the terrain center."""
+  # Linear ramp along rows — every footprint spans a large height range.
+  heights = np.arange(80).reshape(80, 1).repeat(80, axis=1).astype(np.float64)
+  cfg = FlatPatchSamplingCfg(num_patches=5, patch_radius=0.3, max_height_diff=0.001)
+  patches = find_flat_patches_from_heightfield(
+    heights=heights, horizontal_scale=0.1, z_offset=0.0, cfg=cfg, rng=rng
+  )
+  center_x = 80 * 0.1 / 2.0
+  center_y = 80 * 0.1 / 2.0
+  np.testing.assert_allclose(patches[:, 0], center_x)
+  np.testing.assert_allclose(patches[:, 1], center_y)
+
+
+def test_patches_respect_edge_margin(rng: np.random.Generator):
+  """No patch center should be within patch_radius of the heightfield edge."""
+  heights = np.zeros((80, 80), dtype=np.float64)
+  patch_radius = 0.5
+  h_scale = 0.1
+  cfg = FlatPatchSamplingCfg(
+    num_patches=200, patch_radius=patch_radius, max_height_diff=0.1
+  )
+  patches = find_flat_patches_from_heightfield(
+    heights=heights, horizontal_scale=h_scale, z_offset=0.0, cfg=cfg, rng=rng
+  )
+  terrain_x = 80 * h_scale
+  terrain_y = 80 * h_scale
+  margin = patch_radius
+  # Every patch center must be at least patch_radius from each edge.
+  assert np.all(patches[:, 0] >= margin - h_scale)
+  assert np.all(patches[:, 0] <= terrain_x - margin + h_scale)
+  assert np.all(patches[:, 1] >= margin - h_scale)
+  assert np.all(patches[:, 1] <= terrain_y - margin + h_scale)
+
+
+def test_grid_resolution(rng: np.random.Generator):
+  """Finer grid_resolution still produces correct flat patches on a step."""
+  heights = np.zeros((80, 80), dtype=np.float64)
+  heights[:, 40:] = 1.0
+  cfg = FlatPatchSamplingCfg(
+    num_patches=30,
+    patch_radius=0.3,
+    max_height_diff=0.05,
+    grid_resolution=0.025,
+  )
+  patches = find_flat_patches_from_heightfield(
+    heights=heights, horizontal_scale=0.1, z_offset=0.0, cfg=cfg, rng=rng
+  )
+  assert patches.shape == (30, 3)
+  for i in range(len(patches)):
+    x, z = patches[i, 0], patches[i, 2]
+    assert x < 3.7 or x >= 4.3, f"Patch {i} at x={x:.2f} too close to step"
+    if x < 3.7:
+      assert abs(z) < 0.1
+    else:
+      assert abs(z - 1.0) < 0.1
+
+
+def test_terrain_importer_end_to_end():
+  """Flat patches flow through generator → importer with correct shape and bounds."""
+  device = get_test_device()
+  terrain_size = (4.0, 4.0)
+  num_patches = 5
+  patch_cfg = FlatPatchSamplingCfg(
+    num_patches=num_patches, patch_radius=0.3, max_height_diff=0.1
+  )
+  gen_cfg = TerrainGeneratorCfg(
+    seed=123,
+    size=terrain_size,
+    num_rows=2,
+    num_cols=2,
+    sub_terrains={
+      "rough": HfRandomUniformTerrainCfg(
+        proportion=1.0,
+        noise_range=(0.02, 0.08),
+        noise_step=0.01,
+        flat_patch_sampling={"spawn": patch_cfg},
+      ),
+    },
+  )
+  importer_cfg = TerrainEntityCfg(
+    terrain_type="generator", terrain_generator=gen_cfg, num_envs=4
+  )
+  importer = TerrainEntity(importer_cfg, device=device)
+
+  assert "spawn" in importer.flat_patches
+  patches = importer.flat_patches["spawn"]
+  assert patches.shape == (2, 2, num_patches, 3)
+  assert patches.dtype == torch.float
+
+  # Every patch should be within the world-space extent of the full terrain grid.
+  # Grid is centered at origin → spans [-total/2, +total/2].
+  half_x = gen_cfg.num_rows * terrain_size[0] / 2
+  half_y = gen_cfg.num_cols * terrain_size[1] / 2
+  pts = patches.cpu().numpy().reshape(-1, 3)
+  assert np.all(pts[:, 0] >= -half_x) and np.all(pts[:, 0] <= half_x)
+  assert np.all(pts[:, 1] >= -half_y) and np.all(pts[:, 1] <= half_y)
@@ -0,0 +1,230 @@
+"""Tests for per-foot height sensor and terrain-aware rewards."""
+
+from __future__ import annotations
+
+import pytest
+import torch
+from conftest import get_test_device, make_scene_and_sim
+
+from mjlab.sensor import ObjRef, RingPatternCfg, TerrainHeightSensorCfg
+from mjlab.sensor.terrain_height_sensor import TerrainHeightSensor
+
+# Platform top at z=0.5 (box center z=0.25, half-height 0.25).
+# Body at z=1.0, left_foot at z=0.8 (offset -0.2), right_foot at z=0.6 (offset -0.4).
+# Expected: left 0.3m above platform, right 0.1m above platform.
+# Both within max_distance=1.0.
+TWO_FEET_ABOVE_PLATFORM_XML = """
+  <mujoco>
+    <worldbody>
+      <geom name="platform" type="box" size="5 5 0.25" pos="0 0 0.25"/>
+      <body name="base" pos="0 0 1">
+        <freejoint name="free_joint"/>
+        <geom name="base_geom" type="sphere" size="0.05" mass="1.0"/>
+        <site name="left_foot" pos="-0.1 0 -0.2"/>
+        <site name="right_foot" pos="0.1 0 -0.4"/>
+      </body>
+    </worldbody>
+  </mujoco>
+"""
+
+# Stepped terrain: step top at z=0.4 for x<0, ground at z=0 for x>0.
+# Body at x=0, z=0.8. left_foot at x=-0.5 (over step), right_foot at x=0.5 (over ground).
+# Expected: left 0.4m above step, right 0.8m above ground.
+STEPPED_TERRAIN_XML = """
+  <mujoco>
+    <worldbody>
+      <geom name="ground" type="plane" size="5 5 0.1" pos="0 0 0"/>
+      <geom name="step" type="box" size="5 5 0.2" pos="-5 0 0.2"/>
+      <body name="base" pos="0 0 0.8">
+        <freejoint name="free_joint"/>
+        <geom name="base_geom" type="sphere" size="0.05" mass="1.0"/>
+        <site name="left_foot" pos="-0.5 0 0"/>
+        <site name="right_foot" pos="0.5 0 0"/>
+      </body>
+    </worldbody>
+  </mujoco>
+"""
+
+
+def _foot_sensor_cfg() -> TerrainHeightSensorCfg:
+  """Match the shipped config: yaw alignment, max_distance=1.0, group 0."""
+  return TerrainHeightSensorCfg(
+    name="foot_height_scan",
+    frame=(
+      ObjRef(type="site", name="left_foot", entity="robot"),
+      ObjRef(type="site", name="right_foot", entity="robot"),
+    ),
+    ray_alignment="yaw",
+    pattern=RingPatternCfg.single_ring(radius=0.04, num_samples=4),
+    max_distance=1.0,
+    exclude_parent_body=True,
+    include_geom_groups=(0,),
+  )
+
+
+@pytest.fixture(scope="module")
+def device():
+  return get_test_device()
+
+
+class _FakeEnv:
+  def __init__(self, scene):
+    self.scene = scene
+
+
+def test_foot_height_flat_platform(device):
+  """Two feet at different heights above a flat platform."""
+  cfg = _foot_sensor_cfg()
+  scene, sim = make_scene_and_sim(device, TWO_FEET_ABOVE_PLATFORM_XML, (cfg,))
+  sim.step()
+  sim.sense()
+
+  sensor: TerrainHeightSensor = scene["foot_height_scan"]
+  heights = sensor.data.heights
+
+  assert heights.shape == (1, 2)
+  # Left foot: z=0.8, platform top at z=0.5, height = 0.3.
+  assert heights[0, 0].item() == pytest.approx(0.3, abs=0.05)
+  # Right foot: z=0.6, platform top at z=0.5, height = 0.1.
+  assert heights[0, 1].item() == pytest.approx(0.1, abs=0.05)
+
+
+def test_foot_height_stepped_terrain(device):
+  """Feet over different terrain heights give different readings."""
+  cfg = _foot_sensor_cfg()
+  scene, sim = make_scene_and_sim(device, STEPPED_TERRAIN_XML, (cfg,))
+  sim.step()
+  sim.sense()
+
+  sensor: TerrainHeightSensor = scene["foot_height_scan"]
+  heights = sensor.data.heights
+
+  assert heights.shape == (1, 2)
+  # Left foot at x=-0.5, z=0.8, over step top at z=0.4 -> height ~0.4.
+  assert heights[0, 0].item() == pytest.approx(0.4, abs=0.1)
+  # Right foot at x=0.5, z=0.8, over ground at z=0 -> height ~0.8.
+  assert heights[0, 1].item() == pytest.approx(0.8, abs=0.1)
+
+
+def test_foot_height_observation(device):
+  """foot_height observation delegates to sensor.data.heights."""
+  from mjlab.tasks.velocity.mdp.observations import foot_height
+
+  cfg = _foot_sensor_cfg()
+  scene, sim = make_scene_and_sim(device, TWO_FEET_ABOVE_PLATFORM_XML, (cfg,))
+  sim.step()
+  sim.sense()
+
+  env = _FakeEnv(scene)
+  obs = foot_height(env, "foot_height_scan")  # type: ignore[invalid-argument-type]
+  sensor: TerrainHeightSensor = scene["foot_height_scan"]
+  direct = sensor.data.heights
+
+  assert torch.allclose(obs, direct)
+
+
+def test_foot_height_multi_env(device):
+  """Sensor works correctly across multiple environments."""
+  cfg = _foot_sensor_cfg()
+  scene, sim = make_scene_and_sim(
+    device, TWO_FEET_ABOVE_PLATFORM_XML, (cfg,), num_envs=4
+  )
+  sim.step()
+  sim.sense()
+
+  sensor: TerrainHeightSensor = scene["foot_height_scan"]
+  heights = sensor.data.heights
+
+  assert heights.shape == (4, 2)
+  # All envs should report same heights (identical geometry).
+  for i in range(4):
+    assert heights[i, 0].item() == pytest.approx(0.3, abs=0.05)
+    assert heights[i, 1].item() == pytest.approx(0.1, abs=0.05)
+
+
+def test_foot_height_miss_returns_max_distance(device):
+  """Feet beyond max_distance report max_distance, not -1."""
+  # Body at z=3 with feet at z=2.8 and z=2.6.
+  # max_distance=1.0, ground at z=0, so both feet are >1m above ground.
+  miss_xml = """
+    <mujoco>
+      <worldbody>
+        <geom name="floor" type="plane" size="5 5 0.1" pos="0 0 0"/>
+        <body name="base" pos="0 0 3">
+          <freejoint name="free_joint"/>
+          <geom name="base_geom" type="sphere" size="0.05" mass="1.0"/>
+          <site name="left_foot" pos="-0.1 0 -0.2"/>
+          <site name="right_foot" pos="0.1 0 -0.4"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  """
+  cfg = _foot_sensor_cfg()
+  scene, sim = make_scene_and_sim(device, miss_xml, (cfg,))
+  sim.step()
+  sim.sense()
+
+  sensor: TerrainHeightSensor = scene["foot_height_scan"]
+  heights = sensor.data.heights
+
+  # Both feet are >1m above ground, beyond max_distance=1.0.
+  assert heights[0, 0].item() == pytest.approx(1.0, abs=0.01)
+  assert heights[0, 1].item() == pytest.approx(1.0, abs=0.01)
+
+
+def test_foot_penetration_plane(device):
+  """Foot below a ground plane should report near-zero, not max_distance."""
+  xml = """
+    <mujoco>
+      <worldbody>
+        <geom name="ground" type="plane" size="5 5 0.1" pos="0 0 0"/>
+        <body name="base" pos="0 0 0.05">
+          <freejoint name="free_joint"/>
+          <geom name="base_geom" type="sphere" size="0.02" mass="1.0" group="1"/>
+          <site name="left_foot" pos="0 0 -0.04"/>
+          <site name="right_foot" pos="0 0 -0.06"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  """
+  # left_foot at z=0.01 (above), right_foot at z=-0.01 (below).
+  cfg = _foot_sensor_cfg()
+  scene, sim = make_scene_and_sim(device, xml, (cfg,))
+  sim.step()
+  sim.forward()
+  sim.sense()
+
+  heights = scene["foot_height_scan"].data.heights[0]
+  assert heights[0].item() < 0.1
+  assert heights[1].item() < 0.5
+
+
+def test_foot_penetration_box(device):
+  """Foot inside box terrain should report near-zero, not box thickness.
+
+  Regression: rays inside a box hit the bottom face, producing a bogus
+  height equal to the box thickness.
+  """
+  xml = """
+    <mujoco>
+      <worldbody>
+        <geom name="terrain" type="box" size="5 5 0.5" pos="0 0 0.5"/>
+        <body name="base" pos="0 0 1.05">
+          <freejoint name="free_joint"/>
+          <geom name="base_geom" type="sphere" size="0.02" mass="1.0" group="1"/>
+          <site name="left_foot" pos="0 0 -0.04"/>
+          <site name="right_foot" pos="0 0 -0.08"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  """
+  # Terrain top at z=1.0. left_foot at z=1.01, right_foot at z=0.97.
+  cfg = _foot_sensor_cfg()
+  scene, sim = make_scene_and_sim(device, xml, (cfg,))
+  sim.step()
+  sim.forward()
+  sim.sense()
+
+  heights = scene["foot_height_scan"].data.heights[0]
+  assert heights[0].item() < 0.1
+  assert heights[1].item() < 0.5
@@ -0,0 +1,364 @@
+"""Tests for FusedActuatorGroup (fused stateless-control-law actuators)."""
+
+import mujoco
+import pytest
+import torch
+from conftest import get_test_device, load_fixture_xml
+
+from mjlab.actuator import (
+  DcMotorActuator,
+  DcMotorActuatorCfg,
+  IdealPdActuator,
+  IdealPdActuatorCfg,
+  LearnedMlpActuator,
+)
+from mjlab.actuator.actuator import TransmissionType
+from mjlab.actuator.dc_actuator import dc_motor_clip
+from mjlab.entity import Entity, EntityArticulationInfoCfg, EntityCfg
+from mjlab.sim.sim import Simulation, SimulationCfg
+
+ROBOT_XML = load_fixture_xml("floating_base_articulated")
+TENDON_XML = load_fixture_xml("tendon_finger")
+
+
+@pytest.fixture(scope="module")
+def device():
+  return get_test_device()
+
+
+def make_entity(actuator_cfgs, num_envs, device):
+  cfg = EntityCfg(
+    spec_fn=lambda: mujoco.MjSpec.from_string(ROBOT_XML),
+    articulation=EntityArticulationInfoCfg(actuators=actuator_cfgs),
+  )
+  entity = Entity(cfg)
+  model = entity.compile()
+  sim = Simulation(num_envs=num_envs, cfg=SimulationCfg(), model=model, device=device)
+  entity.initialize(model, sim.model, sim.data, device)
+  return entity, sim
+
+
+def test_idealpd_actuators_fused(device):
+  """Ideal PD actuators with matching delay config fuse into one group."""
+  cfg1 = IdealPdActuatorCfg(target_names_expr=("joint1",), stiffness=50.0, damping=5.0)
+  cfg2 = IdealPdActuatorCfg(target_names_expr=("joint2",), stiffness=30.0, damping=3.0)
+  entity, _ = make_entity((cfg1, cfg2), num_envs=2, device=device)
+
+  assert len(entity._fused_actuator_group._groups) == 1
+  assert entity._fused_actuator_group._groups[0].target_ids.numel() == 2
+  assert len(entity._custom_actuators) == 0
+
+
+def test_different_delay_configs_separate_groups(device):
+  """Ideal PD actuators with different delay configs get separate groups."""
+  cfg1 = IdealPdActuatorCfg(
+    target_names_expr=("joint1",), stiffness=50.0, damping=5.0, delay_max_lag=3
+  )
+  cfg2 = IdealPdActuatorCfg(
+    target_names_expr=("joint2",), stiffness=50.0, damping=5.0, delay_max_lag=5
+  )
+  entity, _ = make_entity((cfg1, cfg2), num_envs=2, device=device)
+
+  assert len(entity._fused_actuator_group._groups) == 2
+
+
+def test_fusable_detection_by_compute():
+  """Fusability is keeping the shared stateless-law compute, not a flag."""
+  from mjlab.actuator import BuiltinPositionActuator
+
+  # IdealPd defines the shared law-applying compute; DcMotor inherits it.
+  assert DcMotorActuator.compute is IdealPdActuator.compute
+  # Custom compute (learned network, built-in field passthrough) opts out, with
+  # no flag and immune to the DcMotor -> LearnedMlp inheritance trap.
+  assert LearnedMlpActuator.compute is not IdealPdActuator.compute
+  assert BuiltinPositionActuator.compute is not IdealPdActuator.compute
+
+
+def test_dcmotor_fuses_separately_and_matches(device):
+  """DcMotor fuses into its own group (distinct law) and matches per-actuator."""
+  ideal = IdealPdActuatorCfg(target_names_expr=("joint1",), stiffness=50.0, damping=5.0)
+  dc = DcMotorActuatorCfg(
+    target_names_expr=("joint2",),
+    stiffness=50.0,
+    damping=5.0,
+    effort_limit=20.0,
+    saturation_effort=40.0,
+    velocity_limit=30.0,
+  )
+  entity, _ = make_entity((ideal, dc), num_envs=4, device=device)
+  fused = entity._fused_actuator_group
+
+  # Different control laws (clamp vs torque-speed curve) -> separate groups,
+  # nothing left on the custom path.
+  assert len(fused._groups) == 2
+  assert len(entity._custom_actuators) == 0
+
+  data = entity.data
+  # Large targets so the DC torque-speed clip is actually active.
+  data.joint_pos_target[:] = 10.0 * torch.randn_like(data.joint_pos_target)
+  data.joint_vel_target[:] = torch.randn_like(data.joint_vel_target)
+
+  fused.apply_controls(data)
+  for group in fused._groups:
+    got = data.data.ctrl[:, data.indexing.ctrl_ids[group.ctrl_ids]]
+    ref = torch.cat(
+      [act.compute(act.get_command(data)) for act in group.absorbed_actuators], dim=1
+    )
+    assert torch.equal(got, ref)
+
+
+def test_fused_matches_per_actuator(device):
+  """Fused control output is identical to per-actuator compute (no delay)."""
+  cfg1 = IdealPdActuatorCfg(
+    target_names_expr=("joint1",), stiffness=50.0, damping=5.0, effort_limit=100.0
+  )
+  cfg2 = IdealPdActuatorCfg(
+    target_names_expr=("joint2",), stiffness=30.0, damping=3.0, effort_limit=80.0
+  )
+  entity, _ = make_entity((cfg1, cfg2), num_envs=4, device=device)
+  data = entity.data
+  group = entity._fused_actuator_group._groups[0]
+
+  data.joint_pos_target[:] = torch.randn_like(data.joint_pos_target)
+  data.joint_vel_target[:] = torch.randn_like(data.joint_vel_target)
+
+  entity._fused_actuator_group.apply_controls(data)
+  fused_ctrl = data.data.ctrl[:, data.indexing.ctrl_ids[group.ctrl_ids]].clone()
+
+  reference = torch.cat(
+    [act.compute(act.get_command(data)) for act in group.absorbed_actuators], dim=1
+  )
+  assert torch.equal(fused_ctrl, reference)
+
+
+def test_set_gains_writes_through_view(device):
+  """Per-actuator set_gains mutates the fused gain tensor and its output."""
+  cfg1 = IdealPdActuatorCfg(
+    target_names_expr=("joint1",), stiffness=50.0, damping=5.0, effort_limit=1e6
+  )
+  cfg2 = IdealPdActuatorCfg(
+    target_names_expr=("joint2",), stiffness=30.0, damping=3.0, effort_limit=1e6
+  )
+  entity, _ = make_entity((cfg1, cfg2), num_envs=4, device=device)
+  data = entity.data
+  group = entity._fused_actuator_group._groups[0]
+  act0 = group.absorbed_actuators[0]
+  assert isinstance(act0, IdealPdActuator)
+  n0 = act0.target_ids.numel()
+
+  env_ids = torch.arange(4, device=device)
+  new_kp = torch.full((4, n0), 123.0, device=device)
+  act0.set_gains(env_ids, kp=new_kp)
+
+  # The view aliases the fused tensor in place.
+  assert torch.equal(group.params["stiffness"][:, :n0], new_kp)
+
+  data.joint_pos_target[:] = torch.randn_like(data.joint_pos_target)
+  entity._fused_actuator_group.apply_controls(data)
+  fused_ctrl = data.data.ctrl[:, data.indexing.ctrl_ids[group.ctrl_ids]]
+  expected0 = act0.compute(act0.get_command(data))
+  assert torch.equal(fused_ctrl[:, :n0], expected0)
+
+
+def test_fused_delay_applies_lag(device):
+  """A fused group with constant lag returns the command from `lag` steps ago."""
+  lag = 3
+  cfg = IdealPdActuatorCfg(
+    target_names_expr=("joint.*",),
+    stiffness=0.0,  # isolate the feedforward effort term.
+    damping=0.0,
+    effort_limit=1e9,
+    delay_min_lag=lag,
+    delay_max_lag=lag,
+  )
+  entity, _ = make_entity((cfg,), num_envs=2, device=device)
+  data = entity.data
+  group = entity._fused_actuator_group._groups[0]
+  assert group.delay_buffer is not None
+
+  # First append happened during a prior apply; drive a ramp and read it back.
+  data.joint_effort_target[:] = 0.0
+  entity._fused_actuator_group.apply_controls(data)  # seed history with 0.
+  seen = []
+  for step in range(1, 6):
+    data.joint_effort_target[:] = float(step)
+    entity._fused_actuator_group.apply_controls(data)
+    seen.append(data.data.ctrl[:, data.indexing.ctrl_ids[group.ctrl_ids]][0, 0].item())
+
+  # ctrl reflects the effort target from `lag` steps earlier; history starts at
+  # the seeded 0 and the buffer fills before the ramp shows through.
+  assert seen == [0.0, 0.0, 0.0, 1.0, 2.0]
+
+
+def test_no_idealpd_actuators_empty_group(device):
+  """With no ideal PD actuators the fused group is empty and harmless."""
+  from mjlab.actuator import BuiltinPositionActuatorCfg
+
+  cfg = BuiltinPositionActuatorCfg(
+    target_names_expr=("joint.*",), stiffness=50.0, damping=5.0
+  )
+  entity, _ = make_entity((cfg,), num_envs=2, device=device)
+  assert len(entity._fused_actuator_group._groups) == 0
+  entity._fused_actuator_group.apply_controls(entity.data)  # no-op, must not raise.
+
+
+def make_tendon_entity(actuator_cfgs, num_envs, device):
+  cfg = EntityCfg(
+    spec_fn=lambda: mujoco.MjSpec.from_string(TENDON_XML),
+    articulation=EntityArticulationInfoCfg(actuators=actuator_cfgs),
+  )
+  entity = Entity(cfg)
+  model = entity.compile()
+  sim = Simulation(num_envs=num_envs, cfg=SimulationCfg(), model=model, device=device)
+  entity.initialize(model, sim.model, sim.data, device)
+  return entity, sim
+
+
+def test_tendon_transmission_fused_matches_per_actuator(device):
+  """TENDON-transmission ideal PD actuators fuse and gather the right fields."""
+  cfg = IdealPdActuatorCfg(
+    target_names_expr=("finger_tendon",),
+    transmission_type=TransmissionType.TENDON,
+    stiffness=50.0,
+    damping=5.0,
+    effort_limit=100.0,
+  )
+  entity, _ = make_tendon_entity((cfg,), num_envs=4, device=device)
+  data = entity.data
+  group = entity._fused_actuator_group._groups[0]
+  assert group.transmission_type == TransmissionType.TENDON
+
+  data.tendon_len_target[:] = torch.randn_like(data.tendon_len_target)
+  data.tendon_vel_target[:] = torch.randn_like(data.tendon_vel_target)
+  data.tendon_effort_target[:] = torch.randn_like(data.tendon_effort_target)
+
+  entity._fused_actuator_group.apply_controls(data)
+  fused_ctrl = data.data.ctrl[:, data.indexing.ctrl_ids[group.ctrl_ids]].clone()
+
+  act0 = group.absorbed_actuators[0]
+  reference = act0.compute(act0.get_command(data))
+  assert torch.equal(fused_ctrl, reference)
+
+
+def test_fused_group_shares_one_lag_per_env(device):
+  """Fused actuators sharing a delay config share one lag per env, not one each."""
+  cfg1 = IdealPdActuatorCfg(
+    target_names_expr=("joint1",),
+    stiffness=0.0,
+    damping=0.0,
+    effort_limit=1e9,
+    delay_min_lag=0,
+    delay_max_lag=5,
+  )
+  cfg2 = IdealPdActuatorCfg(
+    target_names_expr=("joint2",),
+    stiffness=0.0,
+    damping=0.0,
+    effort_limit=1e9,
+    delay_min_lag=0,
+    delay_max_lag=5,
+  )
+  entity, _ = make_entity((cfg1, cfg2), num_envs=8, device=device)
+  group = entity._fused_actuator_group._groups[0]
+  assert len(group.absorbed_actuators) == 2
+  act0, act1 = group.absorbed_actuators
+
+  # Both actuators alias the same shared DelayBuffer.
+  assert group.delay_buffer is not None
+  assert act0._delay_buffer is act1._delay_buffer is group.delay_buffer
+
+  # Setting lags through one actuator's handle moves the whole group's lags,
+  # since set_lags reaches into the shared buffer.
+  env_ids = torch.arange(8, device=device)
+  lags = torch.tensor([0, 1, 2, 3, 4, 5, 3, 2], device=device)
+  act0.set_lags(lags, env_ids)
+  assert torch.equal(group.delay_buffer.current_lags, lags)
+  zeros = torch.zeros(8, dtype=torch.long, device=device)
+  act1.set_lags(zeros, env_ids)
+  assert torch.equal(group.delay_buffer.current_lags, zeros)
+
+
+def test_fused_group_reset_clears_shared_buffer(device):
+  """Resetting one absorbed actuator's env_ids clears the shared delay buffer."""
+  cfg1 = IdealPdActuatorCfg(
+    target_names_expr=("joint1",),
+    stiffness=0.0,
+    damping=0.0,
+    effort_limit=1e9,
+    delay_min_lag=2,
+    delay_max_lag=2,
+  )
+  cfg2 = IdealPdActuatorCfg(
+    target_names_expr=("joint2",),
+    stiffness=0.0,
+    damping=0.0,
+    effort_limit=1e9,
+    delay_min_lag=2,
+    delay_max_lag=2,
+  )
+  entity, _ = make_entity((cfg1, cfg2), num_envs=4, device=device)
+  data = entity.data
+  group = entity._fused_actuator_group._groups[0]
+  assert group.delay_buffer is not None
+
+  data.joint_effort_target[:] = 7.0
+  for _ in range(3):
+    entity._fused_actuator_group.apply_controls(data)
+  assert torch.all(group.delay_buffer._buffer.current_length == 3)
+
+  # Resetting env 1 through one absorbed actuator zeros the shared buffer's
+  # history and counter for that env, visible to the other absorbed actuator.
+  reset_ids = torch.tensor([1], device=device)
+  group.absorbed_actuators[0].reset(reset_ids)
+  assert group.delay_buffer._buffer.current_length[1] == 0
+  assert group.delay_buffer._buffer.current_length[0] == 3
+
+
+def test_dcmotor_effort_limit_randomization_updates_torque_speed_curve(device):
+  """DC motor torque-speed clip tracks force_limit after it is randomized.
+
+  Regression: the corner velocity of the torque-speed curve used to be cached
+  once at initialize() time, so calling set_effort_limit (as domain
+  randomization does) silently left the clip using the stale, pre-randomized
+  force_limit.
+  """
+  cfg = DcMotorActuatorCfg(
+    target_names_expr=("joint1",),
+    stiffness=0.0,
+    damping=0.0,
+    effort_limit=20.0,
+    saturation_effort=40.0,
+    velocity_limit=30.0,
+  )
+  entity, _ = make_entity((cfg,), num_envs=2, device=device)
+  data = entity.data
+  group = entity._fused_actuator_group._groups[0]
+  act0 = group.absorbed_actuators[0]
+  assert isinstance(act0, DcMotorActuator)
+
+  env_ids = torch.arange(2, device=device)
+  new_limit = torch.full((2, 1), 5.0, device=device)
+  act0.set_effort_limit(env_ids, effort_limit=new_limit)
+
+  data.joint_effort_target[:] = 100.0
+  data.joint_vel_target[:] = 0.0
+  current_vel = act0.get_command(data).vel.clone()  # actual current joint vel.
+  entity._fused_actuator_group.apply_controls(data)
+  fused_ctrl = data.data.ctrl[:, data.indexing.ctrl_ids[group.ctrl_ids]]
+
+  assert act0.saturation_effort is not None
+  assert act0.velocity_limit_motor is not None
+  assert act0.force_limit is not None
+  expected = dc_motor_clip(
+    torch.full_like(fused_ctrl, 100.0),
+    act0.saturation_effort,
+    act0.velocity_limit_motor,
+    act0.force_limit,
+    current_vel,
+  )
+  assert torch.equal(fused_ctrl, expected)
+  # The randomized (not the original) force_limit bounds the clipped torque:
+  # with vel == 0 the corner velocity is never engaged, so the clamp is exactly
+  # +/- force_limit.
+  assert torch.all(current_vel == 0.0)
+  assert torch.allclose(fused_ctrl, new_limit)
@@ -0,0 +1,141 @@
+"""Tests for g1_constants.py."""
+
+import re
+
+import mujoco
+import numpy as np
+import pytest
+
+from mjlab.asset_zoo.robots.unitree_g1 import g1_constants
+from mjlab.entity import Entity
+from mjlab.utils.string import resolve_expr
+
+
+@pytest.fixture(scope="module")
+def g1_entity() -> Entity:
+  return Entity(g1_constants.get_g1_robot_cfg())
+
+
+@pytest.fixture(scope="module")
+def g1_model(g1_entity: Entity) -> mujoco.MjModel:
+  return g1_entity.spec.compile()
+
+
+# fmt: off
+@pytest.mark.parametrize(
+  "actuator_config,stiffness,damping",
+  [
+    (g1_constants.G1_ACTUATOR_5020, g1_constants.STIFFNESS_5020, g1_constants.DAMPING_5020),
+    (g1_constants.G1_ACTUATOR_7520_14, g1_constants.STIFFNESS_7520_14, g1_constants.DAMPING_7520_14),
+    (g1_constants.G1_ACTUATOR_7520_22, g1_constants.STIFFNESS_7520_22, g1_constants.DAMPING_7520_22),
+    (g1_constants.G1_ACTUATOR_4010, g1_constants.STIFFNESS_4010, g1_constants.DAMPING_4010),
+    (g1_constants.G1_ACTUATOR_WAIST, g1_constants.STIFFNESS_5020 * 2, g1_constants.DAMPING_5020 * 2),
+    (g1_constants.G1_ACTUATOR_ANKLE, g1_constants.STIFFNESS_5020 * 2, g1_constants.DAMPING_5020 * 2),
+  ],
+)
+# fmt: on
+def test_actuator_parameters(g1_model, actuator_config, stiffness, damping):
+  for i in range(g1_model.nu):
+    actuator = g1_model.actuator(i)
+    actuator_name = actuator.name
+    matches = any(
+      re.match(pattern, actuator_name) for pattern in actuator_config.target_names_expr
+    )
+    if matches:
+      assert actuator.gainprm[0] == stiffness
+      assert actuator.biasprm[1] == -stiffness
+      assert actuator.biasprm[2] == -damping
+      assert actuator.forcerange[0] == -actuator_config.effort_limit
+      assert actuator.forcerange[1] == actuator_config.effort_limit
+
+
+def test_keyframe_base_position(g1_model) -> None:
+  data = mujoco.MjData(g1_model)
+  mujoco.mj_resetDataKeyframe(g1_model, data, 0)
+  mujoco.mj_forward(g1_model, data)
+  np.testing.assert_array_equal(data.qpos[:3], g1_constants.KNEES_BENT_KEYFRAME.pos)
+  np.testing.assert_array_equal(data.qpos[3:7], g1_constants.KNEES_BENT_KEYFRAME.rot)
+
+
+def test_keyframe_joint_positions(g1_entity, g1_model) -> None:
+  """Test that keyframe joint positions match the configuration."""
+  key = g1_model.key("init_state")
+  expected_joint_pos = g1_constants.KNEES_BENT_KEYFRAME.joint_pos
+  assert expected_joint_pos is not None
+  expected_values = resolve_expr(expected_joint_pos, g1_entity.joint_names, 0.0)
+  for joint_name, expected_value in zip(
+    g1_entity.joint_names, expected_values, strict=True
+  ):
+    joint = g1_model.joint(joint_name)
+    qpos_idx = joint.qposadr[0]
+    actual_value = key.qpos[qpos_idx]
+    np.testing.assert_allclose(
+      actual_value,
+      expected_value,
+      rtol=1e-5,
+      err_msg=f"Joint {joint_name} position mismatch: "
+      f"expected {expected_value}, got {actual_value}",
+    )
+
+
+def test_foot_collision_geoms(g1_model) -> None:
+  foot_pattern = r"^(left|right)_foot[1-7]_collision$"
+  for i in range(g1_model.ngeom):
+    geom = g1_model.geom(i)
+    geom_name = geom.name
+    # Foot collision geoms should have condim=3, priority=1, and friction=0.6.
+    if re.match(foot_pattern, geom_name):
+      assert geom.condim == 3
+      assert geom.priority == 1
+      assert geom.friction[0] == 0.6
+
+
+def test_non_foot_collision_geoms(g1_model) -> None:
+  foot_pattern = r"^(left|right)_foot[1-7]_collision$"
+  for i in range(g1_model.ngeom):
+    geom = g1_model.geom(i)
+    geom_name = geom.name
+    if "_collision" not in geom_name:
+      continue
+    # Non-foot collision geoms should have condim=1
+    if not re.match(foot_pattern, geom_name):
+      assert geom.condim == 1
+
+
+def test_collision_geom_count(g1_model) -> None:
+  # There should be 7 geoms (capsules) per foot.
+  collision_geoms = [
+    g1_model.geom(i).name
+    for i in range(g1_model.ngeom)
+    if "_collision" in g1_model.geom(i).name
+  ]
+  assert len(collision_geoms) > 0
+  foot_geoms = [
+    name
+    for name in collision_geoms
+    if re.match(r"^(left|right)_foot[1-7]_collision$", name)
+  ]
+  assert len(foot_geoms) == 14
+
+
+def test_g1_entity_creation(g1_entity) -> None:
+  assert g1_entity.num_actuators == 29
+  assert g1_entity.num_joints == 29
+  assert g1_entity.is_actuated
+  assert not g1_entity.is_fixed_base
+
+
+def test_g1_actuators_configured_correctly(g1_model):
+  """Verify that all G1 actuators have correct control and force limiting.
+
+  All 29 G1 actuators should have ctrllimited=False (allowing setpoints beyond
+  joint limits) and forcelimited=True (limiting forces to effort limits).
+  """
+  for i in range(g1_model.nu):
+    actuator_name = g1_model.actuator(i).name
+    assert g1_model.actuator_ctrllimited[i] == 0, (
+      f"Actuator '{actuator_name}' has ctrllimited=True, expected False"
+    )
+    assert g1_model.actuator_forcelimited[i] == 1, (
+      f"Actuator '{actuator_name}' has forcelimited=False, expected True"
+    )
@@ -0,0 +1,114 @@
+"""Tests for go1_constants.py."""
+
+import re
+
+import mujoco
+import numpy as np
+import pytest
+
+from mjlab.asset_zoo.robots.unitree_go1 import go1_constants
+from mjlab.entity import Entity
+from mjlab.utils.string import resolve_expr
+
+
+@pytest.fixture(scope="module")
+def go1_entity() -> Entity:
+  return Entity(go1_constants.get_go1_robot_cfg())
+
+
+@pytest.fixture(scope="module")
+def go1_model(go1_entity: Entity) -> mujoco.MjModel:
+  return go1_entity.spec.compile()
+
+
+# fmt: off
+@pytest.mark.parametrize(
+  "actuator_config,stiffness,damping",
+  [
+    (go1_constants.GO1_HIP_ACTUATOR_CFG, go1_constants.STIFFNESS_HIP, go1_constants.DAMPING_HIP),
+    (go1_constants.GO1_KNEE_ACTUATOR_CFG, go1_constants.STIFFNESS_KNEE, go1_constants.DAMPING_KNEE),
+  ],
+)
+# fmt: on
+def test_actuator_parameters(go1_model, actuator_config, stiffness, damping):
+  for i in range(go1_model.nu):
+    actuator = go1_model.actuator(i)
+    actuator_name = actuator.name
+    matches = any(
+      re.match(pattern, actuator_name) for pattern in actuator_config.target_names_expr
+    )
+    if matches:
+      assert actuator.gainprm[0] == stiffness
+      assert actuator.biasprm[1] == -stiffness
+      assert actuator.biasprm[2] == -damping
+      assert actuator.forcerange[0] == -actuator_config.effort_limit
+      assert actuator.forcerange[1] == actuator_config.effort_limit
+
+
+def test_keyframe_joint_positions(go1_entity, go1_model) -> None:
+  """Test that keyframe joint positions match the configuration."""
+  key = go1_model.key("init_state")
+  expected_joint_pos = go1_constants.INIT_STATE.joint_pos
+  assert expected_joint_pos is not None
+  expected_values = resolve_expr(expected_joint_pos, go1_entity.joint_names, 0.0)
+  for joint_name, expected_value in zip(
+    go1_entity.joint_names, expected_values, strict=True
+  ):
+    joint = go1_model.joint(joint_name)
+    qpos_idx = joint.qposadr[0]
+    actual_value = key.qpos[qpos_idx]
+    np.testing.assert_allclose(
+      actual_value,
+      expected_value,
+      rtol=1e-5,
+      err_msg=f"Joint {joint_name} position mismatch: "
+      f"expected {expected_value}, got {actual_value}",
+    )
+
+
+def test_foot_collision_geoms(go1_model) -> None:
+  """Foot collision geoms should have specific properties."""
+  foot_pattern = r"^[FR][LR]_foot_collision$"
+  for i in range(go1_model.ngeom):
+    geom = go1_model.geom(i)
+    if re.match(foot_pattern, geom.name):
+      assert geom.condim == 6
+      assert geom.priority == 1
+      assert geom.friction[0] == 1.0
+
+
+def test_all_collision_geoms_enabled(go1_model) -> None:
+  """Every *_collision geom, including numbered ones, must stay collidable.
+
+  The go1 XML numbers its secondary collision capsules (FR_thigh_collision1,
+  ...), which a pattern anchored at "_collision" would silently drop, letting
+  disable_other_geoms zero their contype. Guard against that regression with
+  an oracle independent of the config's own regex resolution.
+  """
+  collision_geoms = [
+    go1_model.geom(i)
+    for i in range(go1_model.ngeom)
+    if "_collision" in go1_model.geom(i).name
+  ]
+  assert len(collision_geoms) == 30
+  for geom in collision_geoms:
+    assert geom.contype == 1, f"{geom.name} lost its contype"
+    assert geom.conaffinity == 1, f"{geom.name} lost its conaffinity"
+
+
+def test_collision_geom_count(go1_model) -> None:
+  """Go1 should have 4 foot collision geoms."""
+  foot_pattern = r"^[FR][LR]_foot_collision$"
+  foot_geoms = [
+    g.name for g in [go1_model.geom(i) for i in range(go1_model.ngeom)]
+    if re.match(foot_pattern, g.name)
+  ]
+  assert len(foot_geoms) == 4
+
+
+def test_go1_entity_creation(go1_entity) -> None:
+  """Test basic Go1 entity properties."""
+  assert go1_entity.num_actuators == 12
+  assert go1_entity.num_joints == 12
+  assert go1_entity.is_actuated
+  assert not go1_entity.is_fixed_base
@@ -0,0 +1,128 @@
+"""Tests for GPU selection utilities."""
+
+import os
+
+import pytest
+
+from mjlab.utils.gpu import select_gpus
+
+
+@pytest.fixture(autouse=True)
+def clean_cuda_env():
+  """Clean CUDA_VISIBLE_DEVICES before each test."""
+  original = os.environ.get("CUDA_VISIBLE_DEVICES")
+  yield
+  # Restore original value after test.
+  if original is None:
+    os.environ.pop("CUDA_VISIBLE_DEVICES", None)
+  else:
+    os.environ["CUDA_VISIBLE_DEVICES"] = original
+
+
+def test_select_gpus_with_visible_devices_set():
+  """Respects CUDA_VISIBLE_DEVICES when selecting GPUs."""
+  os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"
+
+  # Select first two GPUs (indices 0, 1 -> physical GPUs 0, 1).
+  selected, num = select_gpus([0, 1])
+  assert selected == [0, 1]
+  assert num == 2
+
+
+def test_select_gpus_with_non_contiguous_devices():
+  """Handles non-contiguous CUDA_VISIBLE_DEVICES correctly."""
+  os.environ["CUDA_VISIBLE_DEVICES"] = "1,3,5"
+
+  # Index 0 -> physical GPU 1.
+  selected, num = select_gpus([0])
+  assert selected == [1]
+  assert num == 1
+
+  # Index 1 -> physical GPU 3.
+  selected, num = select_gpus([1])
+  assert selected == [3]
+  assert num == 1
+
+  # Indices 0, 2 -> physical GPUs 1, 5.
+  selected, num = select_gpus([0, 2])
+  assert selected == [1, 5]
+  assert num == 2
+
+
+def test_select_gpus_all_option():
+  """Selects all GPUs when 'all' is specified."""
+  os.environ["CUDA_VISIBLE_DEVICES"] = "2,4,6"
+
+  selected, num = select_gpus("all")
+  assert selected == [2, 4, 6]
+  assert num == 3
+
+
+def test_select_gpus_single_gpu():
+  """Correctly selects a single GPU."""
+  os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
+
+  selected, num = select_gpus([0])
+  assert selected == [0]
+  assert num == 1
+
+
+def test_select_gpus_with_spaces():
+  """Handles CUDA_VISIBLE_DEVICES with spaces correctly."""
+  os.environ["CUDA_VISIBLE_DEVICES"] = " 0 , 1 , 2 "
+
+  selected, num = select_gpus([1, 2])
+  assert selected == [1, 2]
+  assert num == 2
+
+
+def test_select_gpus_empty_segments():
+  """Handles empty segments in CUDA_VISIBLE_DEVICES."""
+  # This can happen with malformed environment variables.
+  os.environ["CUDA_VISIBLE_DEVICES"] = "0,,2"
+
+  selected, num = select_gpus([0, 1])
+  assert selected == [0, 2]
+  assert num == 2
+
+
+def test_select_gpus_respects_user_order():
+  """Preserves the order specified by the user."""
+  os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"
+
+  # User requests in reverse order.
+  selected, num = select_gpus([3, 2, 1, 0])
+  assert selected == [3, 2, 1, 0]
+  assert num == 4
+
+
+def test_select_gpus_cpu_mode_explicit():
+  """Selects CPU mode when None is specified."""
+  os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"
+
+  selected, num = select_gpus(None)
+  assert selected is None
+  assert num == 0
+
+
+def test_select_gpus_cpu_mode_empty_cuda_visible_devices():
+  """Selects CPU mode when CUDA_VISIBLE_DEVICES is empty."""
+  os.environ["CUDA_VISIBLE_DEVICES"] = ""
+
+  # Should return CPU mode (None, 0) since no GPUs are visible.
+  selected, num = select_gpus([0])
+  assert selected is None
+  assert num == 0
+
+
+def test_select_gpus_mig_uuids():
+  """Handles MIG GPU UUIDs in CUDA_VISIBLE_DEVICES."""
+  os.environ["CUDA_VISIBLE_DEVICES"] = "MIG-GPU-abc-123,MIG-GPU-def-456"
+
+  selected, num = select_gpus("all")
+  assert selected == ["MIG-GPU-abc-123", "MIG-GPU-def-456"]
+  assert num == 2
+
+  selected, num = select_gpus([0])
+  assert selected == ["MIG-GPU-abc-123"]
+  assert num == 1
@@ -0,0 +1,26 @@
+"""Tests for mjlab.utils.lab_api.math module."""
+
+import pytest
+import torch
+from conftest import get_test_device
+
+from mjlab.utils.lab_api.math import apply_delta_pose
+
+
+@pytest.fixture
+def device():
+  return get_test_device()
+
+
+def test_apply_delta_pose_zero_rotation_is_finite_and_identity(device):
+  """Zero rotation delta should return finite values and preserve input pose."""
+  source_pos = torch.zeros(2, 3, device=device)
+  source_rot = torch.tensor([[1.0, 0.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0]], device=device)
+  delta_pose = torch.zeros(2, 6, device=device)
+
+  target_pos, target_rot = apply_delta_pose(source_pos, source_rot, delta_pose)
+
+  assert torch.isfinite(target_pos).all()
+  assert torch.isfinite(target_rot).all()
+  assert torch.allclose(target_pos, source_pos)
+  assert torch.allclose(target_rot, source_rot)
@@ -0,0 +1,466 @@
+"""Tests for learned MLP actuator."""
+
+import mujoco
+import pytest
+import torch
+from conftest import get_test_device
+
+from mjlab.actuator import LearnedMlpActuator, LearnedMlpActuatorCfg
+from mjlab.entity import Entity, EntityArticulationInfoCfg, EntityCfg
+from mjlab.sim.sim import Simulation, SimulationCfg
+
+ROBOT_XML = """
+<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="1.0"/>
+      <body name="link1" pos="0 0 0">
+        <joint name="joint1" type="hinge" axis="0 0 1" range="-3.14 3.14"/>
+        <geom name="link1_geom" type="box" size="0.1 0.1 0.1" mass="0.1"/>
+      </body>
+    </body>
+  </worldbody>
+</mujoco>
+"""
+
+ROBOT_XML_TWO_JOINTS = """
+<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="1.0"/>
+      <body name="link1" pos="0 0 0">
+        <joint name="joint1" type="hinge" axis="0 0 1" range="-3.14 3.14"/>
+        <geom name="link1_geom" type="box" size="0.1 0.1 0.1" mass="0.1"/>
+        <body name="link2" pos="0.2 0 0">
+          <joint name="joint2" type="hinge" axis="0 0 1" range="-3.14 3.14"/>
+          <geom name="link2_geom" type="box" size="0.1 0.1 0.1" mass="0.1"/>
+        </body>
+      </body>
+    </body>
+  </worldbody>
+</mujoco>
+"""
+
+
+@pytest.fixture(scope="module")
+def device():
+  return get_test_device()
+
+
+@pytest.fixture(scope="module")
+def identity_network_file(tmp_path_factory):
+  """Create a simple identity network: output = sum of all inputs."""
+
+  class IdentityNet(torch.nn.Module):
+    def forward(self, x):
+      # Sum all inputs and return as single output per joint.
+      return x.sum(dim=1, keepdim=True)
+
+  net = IdentityNet()
+  tmp_dir = tmp_path_factory.mktemp("networks")
+  network_path = tmp_dir / "identity_net.pt"
+  torch.jit.script(net).save(str(network_path))
+  return str(network_path)
+
+
+@pytest.fixture(scope="module")
+def constant_network_file(tmp_path_factory):
+  """Create a network that returns constant value regardless of input."""
+
+  class ConstantNet(torch.nn.Module):
+    def forward(self, x):
+      # Return constant 10.0 for each joint.
+      return torch.full((x.shape[0], 1), 10.0, device=x.device)
+
+  net = ConstantNet()
+  tmp_dir = tmp_path_factory.mktemp("networks")
+  network_path = tmp_dir / "constant_net.pt"
+  torch.jit.script(net).save(str(network_path))
+  return str(network_path)
+
+
+@pytest.fixture(scope="module")
+def subtract_network_file(tmp_path_factory):
+  """Create a network that subtracts second half from first half of inputs.
+
+  For input [a, b, c, d], returns (a + b) - (c + d).
+  This is sensitive to input order, unlike sum-based networks.
+  """
+
+  class SubtractNet(torch.nn.Module):
+    def forward(self, x):
+      # Split input in half and subtract second half from first half.
+      mid = x.shape[1] // 2
+      first_half = x[:, :mid].sum(dim=1, keepdim=True)
+      second_half = x[:, mid:].sum(dim=1, keepdim=True)
+      return first_half - second_half
+
+  net = SubtractNet()
+  tmp_dir = tmp_path_factory.mktemp("networks")
+  network_path = tmp_dir / "subtract_net.pt"
+  torch.jit.script(net).save(str(network_path))
+  return str(network_path)
+
+
+def create_entity_with_actuator(actuator_cfg):
+  cfg = EntityCfg(
+    spec_fn=lambda: mujoco.MjSpec.from_string(ROBOT_XML),
+    articulation=EntityArticulationInfoCfg(actuators=(actuator_cfg,)),
+  )
+  return Entity(cfg)
+
+
+def initialize_entity(entity, device, num_envs=1):
+  model = entity.compile()
+  sim_cfg = SimulationCfg()
+  sim = Simulation(num_envs=num_envs, cfg=sim_cfg, model=model, device=device)
+  entity.initialize(model, sim.model, sim.data, device)
+  return entity, sim
+
+
+def test_learned_mlp_network_loads(device, identity_network_file):
+  """Verify network loads from TorchScript file and initializes."""
+  entity = create_entity_with_actuator(
+    LearnedMlpActuatorCfg(
+      target_names_expr=("joint.*",),
+      network_file=identity_network_file,
+      pos_scale=1.0,
+      vel_scale=1.0,
+      torque_scale=1.0,
+      input_order="pos_vel",
+      history_length=1,
+      saturation_effort=100.0,
+      velocity_limit=30.0,
+      effort_limit=50.0,
+    )
+  )
+
+  entity, sim = initialize_entity(entity, device)
+
+  # Verify network was loaded.
+  actuator = entity.actuators[0]
+  assert isinstance(actuator, LearnedMlpActuator)
+  assert actuator.network is not None
+  assert actuator._pos_error_history is not None
+  assert actuator._vel_history is not None
+
+
+def test_learned_mlp_history_indexing(device, identity_network_file):
+  """Test that history_length uses consecutive timesteps."""
+  entity = create_entity_with_actuator(
+    LearnedMlpActuatorCfg(
+      target_names_expr=("joint.*",),
+      network_file=identity_network_file,
+      pos_scale=1.0,
+      vel_scale=0.0,  # Zero out velocity contribution.
+      torque_scale=1.0,
+      input_order="pos_vel",
+      history_length=3,  # Use current + 2 previous timesteps.
+      saturation_effort=100.0,
+      velocity_limit=30.0,
+      effort_limit=100.0,
+    )
+  )
+
+  entity, sim = initialize_entity(entity, device)
+
+  # Set up state at rest.
+  joint_pos = torch.tensor([[0.0]], device=device)
+  joint_vel = torch.tensor([[0.0]], device=device)
+  entity.write_joint_state_to_sim(joint_pos, joint_vel)
+
+  # Set position targets to create known position errors.
+  # Step 0: pos_error = 1.0
+  # Step 1: pos_error = 2.0
+  # Step 2: pos_error = 3.0
+  targets = [
+    torch.tensor([[1.0]], device=device),
+    torch.tensor([[2.0]], device=device),
+    torch.tensor([[3.0]], device=device),
+  ]
+
+  for target in targets:
+    entity.set_joint_position_target(target)
+    entity.set_joint_velocity_target(joint_vel)
+    entity.set_joint_effort_target(torch.zeros(1, 1, device=device))
+    entity.write_data_to_sim()
+
+  # After 3 steps with history_length=3:
+  # Network input: [current, -1 step, -2 steps] = [3.0, 2.0, 1.0]
+  # Identity network sums: 3.0 + 2.0 + 1.0 = 6.0
+  ctrl = sim.data.ctrl[0]
+  assert torch.allclose(ctrl, torch.tensor([6.0], device=device), atol=1e-4)
+
+
+def test_learned_mlp_input_order_pos_vel(device, subtract_network_file):
+  """Verify input_order='pos_vel' concatenates position then velocity."""
+  entity = create_entity_with_actuator(
+    LearnedMlpActuatorCfg(
+      target_names_expr=("joint.*",),
+      network_file=subtract_network_file,
+      pos_scale=1.0,
+      vel_scale=1.0,
+      torque_scale=1.0,
+      input_order="pos_vel",
+      history_length=1,
+      saturation_effort=100.0,
+      velocity_limit=30.0,
+      effort_limit=100.0,
+    )
+  )
+
+  entity, sim = initialize_entity(entity, device)
+
+  # Set state: pos_error = 2.0, vel = 3.0.
+  joint_pos = torch.tensor([[0.0]], device=device)
+  joint_vel = torch.tensor([[3.0]], device=device)
+  entity.write_joint_state_to_sim(joint_pos, joint_vel)
+
+  entity.set_joint_position_target(torch.tensor([[2.0]], device=device))
+  entity.set_joint_velocity_target(torch.zeros(1, 1, device=device))
+  entity.set_joint_effort_target(torch.zeros(1, 1, device=device))
+  entity.write_data_to_sim()
+
+  # Network input: [pos_error, vel] = [2.0, 3.0]
+  # Subtract network: 2.0 - 3.0 = -1.0
+  ctrl = sim.data.ctrl[0]
+  assert torch.allclose(ctrl, torch.tensor([-1.0], device=device), atol=1e-4)
+
+
+def test_learned_mlp_input_order_vel_pos(device, subtract_network_file):
+  """Verify input_order='vel_pos' concatenates velocity then position."""
+  entity = create_entity_with_actuator(
+    LearnedMlpActuatorCfg(
+      target_names_expr=("joint.*",),
+      network_file=subtract_network_file,
+      pos_scale=1.0,
+      vel_scale=1.0,
+      torque_scale=1.0,
+      input_order="vel_pos",
+      history_length=1,
+      saturation_effort=100.0,
+      velocity_limit=30.0,
+      effort_limit=100.0,
+    )
+  )
+
+  entity, sim = initialize_entity(entity, device)
+
+  # Set state: pos_error = 2.0, vel = 3.0.
+  joint_pos = torch.tensor([[0.0]], device=device)
+  joint_vel = torch.tensor([[3.0]], device=device)
+  entity.write_joint_state_to_sim(joint_pos, joint_vel)
+
+  entity.set_joint_position_target(torch.tensor([[2.0]], device=device))
+  entity.set_joint_velocity_target(torch.zeros(1, 1, device=device))
+  entity.set_joint_effort_target(torch.zeros(1, 1, device=device))
+  entity.write_data_to_sim()
+
+  # Network input: [vel, pos_error] = [3.0, 2.0]
+  # Subtract network: 3.0 - 2.0 = 1.0 (opposite of pos_vel!)
+  ctrl = sim.data.ctrl[0]
+  assert torch.allclose(ctrl, torch.tensor([1.0], device=device), atol=1e-4)
+
+
+def test_learned_mlp_scaling_applied(device, identity_network_file):
+  """Test pos_scale, vel_scale, torque_scale are applied correctly."""
+  entity = create_entity_with_actuator(
+    LearnedMlpActuatorCfg(
+      target_names_expr=("joint.*",),
+      network_file=identity_network_file,
+      pos_scale=2.0,
+      vel_scale=3.0,
+      torque_scale=4.0,
+      input_order="pos_vel",
+      history_length=1,
+      saturation_effort=100.0,
+      velocity_limit=30.0,
+      effort_limit=100.0,
+    )
+  )
+
+  entity, sim = initialize_entity(entity, device)
+
+  # Set state: pos_error = 1.0, vel = 1.0.
+  joint_pos = torch.tensor([[0.0]], device=device)
+  joint_vel = torch.tensor([[1.0]], device=device)
+  entity.write_joint_state_to_sim(joint_pos, joint_vel)
+
+  entity.set_joint_position_target(torch.tensor([[1.0]], device=device))
+  entity.set_joint_velocity_target(torch.zeros(1, 1, device=device))
+  entity.set_joint_effort_target(torch.zeros(1, 1, device=device))
+  entity.write_data_to_sim()
+
+  # Network input: [1.0 * 2.0, 1.0 * 3.0] = [2.0, 3.0]
+  # Identity network sums: 2.0 + 3.0 = 5.0
+  # Output scaled: 5.0 * 4.0 = 20.0
+  ctrl = sim.data.ctrl[0]
+  assert torch.allclose(ctrl, torch.tensor([20.0], device=device), atol=1e-4)
+
+
+def test_learned_mlp_reset_clears_history(device, identity_network_file):
+  """Test reset zeroes history buffers for specified environments."""
+  entity = create_entity_with_actuator(
+    LearnedMlpActuatorCfg(
+      target_names_expr=("joint.*",),
+      network_file=identity_network_file,
+      pos_scale=1.0,
+      vel_scale=1.0,
+      torque_scale=1.0,
+      input_order="pos_vel",
+      history_length=2,
+      saturation_effort=100.0,
+      velocity_limit=30.0,
+      effort_limit=100.0,
+    )
+  )
+
+  entity, sim = initialize_entity(entity, device, num_envs=2)
+
+  # Set state and targets to fill buffers.
+  joint_pos = torch.tensor([[0.0], [0.0]], device=device)
+  joint_vel = torch.tensor([[1.0], [1.0]], device=device)
+  entity.write_joint_state_to_sim(joint_pos, joint_vel)
+
+  entity.set_joint_position_target(torch.tensor([[5.0], [5.0]], device=device))
+  entity.set_joint_velocity_target(torch.zeros(2, 1, device=device))
+  entity.set_joint_effort_target(torch.zeros(2, 1, device=device))
+  entity.write_data_to_sim()
+
+  # Reset only env 0.
+  entity.reset(torch.tensor([0], device=device))
+
+  # Check history buffers.
+  actuator = entity.actuators[0]
+  assert isinstance(actuator, LearnedMlpActuator)
+  assert actuator._pos_error_history is not None
+  assert actuator._vel_history is not None
+  assert actuator._pos_error_history.current_length[0] == 0
+  assert actuator._pos_error_history.current_length[1] > 0
+  assert actuator._vel_history.current_length[0] == 0
+  assert actuator._vel_history.current_length[1] > 0
+
+
+def test_learned_mlp_inherits_dc_motor_limits(device, constant_network_file):
+  """Test that DC motor saturation_effort and velocity limits apply."""
+  saturation_effort = 20.0
+  velocity_limit = 30.0
+
+  entity = create_entity_with_actuator(
+    LearnedMlpActuatorCfg(
+      target_names_expr=("joint.*",),
+      network_file=constant_network_file,  # Always outputs 10.0.
+      pos_scale=1.0,
+      vel_scale=1.0,
+      torque_scale=100.0,  # Scale to 1000.0 to exceed limits.
+      input_order="pos_vel",
+      history_length=1,
+      saturation_effort=saturation_effort,
+      velocity_limit=velocity_limit,
+      effort_limit=saturation_effort,  # Set to saturation to not constrain.
+    )
+  )
+
+  entity, sim = initialize_entity(entity, device)
+
+  # At zero velocity, should be clipped to saturation_effort.
+  joint_pos = torch.tensor([[0.0]], device=device)
+  joint_vel = torch.tensor([[0.0]], device=device)
+  entity.write_joint_state_to_sim(joint_pos, joint_vel)
+
+  entity.set_joint_position_target(torch.tensor([[1.0]], device=device))
+  entity.set_joint_velocity_target(torch.zeros(1, 1, device=device))
+  entity.set_joint_effort_target(torch.zeros(1, 1, device=device))
+  entity.write_data_to_sim()
+
+  # Network outputs 10.0, scaled by 100.0 = 1000.0.
+  # Should be clipped to saturation_effort = 20.0.
+  ctrl = sim.data.ctrl[0]
+  assert torch.allclose(ctrl, torch.tensor([saturation_effort], device=device))
+
+
+def test_learned_mlp_dc_motor_zero_torque_at_max_velocity(
+  device, constant_network_file
+):
+  """Test that DC motor produces zero torque at max velocity."""
+  saturation_effort = 20.0
+  velocity_limit = 30.0
+
+  entity = create_entity_with_actuator(
+    LearnedMlpActuatorCfg(
+      target_names_expr=("joint.*",),
+      network_file=constant_network_file,  # Always outputs 10.0.
+      pos_scale=1.0,
+      vel_scale=1.0,
+      torque_scale=100.0,  # Scale to 1000.0.
+      input_order="pos_vel",
+      history_length=1,
+      saturation_effort=saturation_effort,
+      velocity_limit=velocity_limit,
+      effort_limit=saturation_effort,  # Set to saturation to not constrain.
+    )
+  )
+
+  entity, sim = initialize_entity(entity, device)
+
+  # At max velocity, should produce zero torque.
+  joint_pos = torch.tensor([[0.0]], device=device)
+  joint_vel = torch.tensor([[velocity_limit]], device=device)
+  entity.write_joint_state_to_sim(joint_pos, joint_vel)
+
+  entity.set_joint_position_target(torch.tensor([[1.0]], device=device))
+  entity.set_joint_velocity_target(torch.zeros(1, 1, device=device))
+  entity.set_joint_effort_target(torch.zeros(1, 1, device=device))
+  entity.write_data_to_sim()
+
+  # At max velocity, DC motor should clip to zero.
+  ctrl = sim.data.ctrl[0]
+  assert torch.allclose(ctrl, torch.tensor([0.0], device=device), atol=1e-5)
+
+
+def test_learned_mlp_multi_joint_reshaping(device, identity_network_file):
+  """Test that multi-joint reshaping works correctly."""
+  cfg = EntityCfg(
+    spec_fn=lambda: mujoco.MjSpec.from_string(ROBOT_XML_TWO_JOINTS),
+    articulation=EntityArticulationInfoCfg(
+      actuators=(
+        LearnedMlpActuatorCfg(
+          target_names_expr=("joint.*",),
+          network_file=identity_network_file,
+          pos_scale=1.0,
+          vel_scale=1.0,
+          torque_scale=1.0,
+          input_order="pos_vel",
+          history_length=1,
+          saturation_effort=100.0,
+          velocity_limit=30.0,
+          effort_limit=100.0,
+        ),
+      )
+    ),
+  )
+
+  entity = Entity(cfg)
+  model = entity.compile()
+  sim_cfg = SimulationCfg()
+  sim = Simulation(num_envs=1, cfg=sim_cfg, model=model, device=device)
+  entity.initialize(model, sim.model, sim.data, device)
+
+  # Set different states for each joint.
+  # Joint 1: pos_error = 1.0, vel = 2.0
+  # Joint 2: pos_error = 3.0, vel = 4.0
+  joint_pos = torch.tensor([[0.0, 0.0]], device=device)
+  joint_vel = torch.tensor([[2.0, 4.0]], device=device)
+  entity.write_joint_state_to_sim(joint_pos, joint_vel)
+
+  entity.set_joint_position_target(torch.tensor([[1.0, 3.0]], device=device))
+  entity.set_joint_velocity_target(torch.zeros(1, 2, device=device))
+  entity.set_joint_effort_target(torch.zeros(1, 2, device=device))
+  entity.write_data_to_sim()
+
+  # Network input for joint 1: [1.0, 2.0] → sum = 3.0
+  # Network input for joint 2: [3.0, 4.0] → sum = 7.0
+  ctrl = sim.data.ctrl[0]
+  assert torch.allclose(ctrl, torch.tensor([3.0, 7.0], device=device), atol=1e-4)
@@ -0,0 +1,172 @@
+"""Tests for manager config immutability.
+
+Managers should not mutate the original config objects passed to them.
+This allows the same config to be reused across multiple manager instantiations.
+"""
+
+import inspect
+from unittest.mock import Mock
+
+import mujoco
+import pytest
+import torch
+from conftest import get_test_device
+
+from mjlab.entity import Entity, EntityCfg
+from mjlab.managers.observation_manager import (
+  ObservationGroupCfg,
+  ObservationManager,
+  ObservationTermCfg,
+)
+from mjlab.managers.reward_manager import RewardManager, RewardTermCfg
+from mjlab.managers.scene_entity_config import SceneEntityCfg
+from mjlab.sim.sim import Simulation, SimulationCfg
+
+
+@pytest.fixture
+def device():
+  return get_test_device()
+
+
+@pytest.fixture
+def mock_env(device):
+  env = Mock()
+  env.num_envs = 2
+  env.device = device
+  env.step_dt = 0.02
+  env.max_episode_length_s = 10.0
+  env.scene = {}
+  return env
+
+
+class ClassObsTerm:
+  """Class-based observation term that triggers func mutation."""
+
+  def __init__(self, cfg, env):
+    self.device = env.device
+
+  def __call__(self, env):
+    return torch.zeros((env.num_envs, 3), device=self.device)
+
+
+class ClassRewardTerm:
+  """Class-based reward term that triggers func mutation."""
+
+  def __init__(self, cfg, env):
+    self.device = env.device
+
+  def __call__(self, env):
+    return torch.ones(env.num_envs, device=self.device)
+
+
+def test_manager_preserves_class_func(mock_env):
+  """Managers should not mutate class-based func to instance.
+
+  When term_cfg.func is a class, _resolve_common_term_cfg instantiates it.
+  Without copying term_cfg first, this mutates the original config.
+  """
+  term_cfg = ObservationTermCfg(func=ClassObsTerm, params={})
+
+  assert inspect.isclass(term_cfg.func), "precondition: func should be a class"
+
+  cfg = {"actor": ObservationGroupCfg(terms={"obs1": term_cfg})}
+  _ = ObservationManager(cfg, mock_env)
+
+  assert inspect.isclass(term_cfg.func), "func was mutated from class to instance"
+
+
+def test_observation_shared_terms_between_groups(mock_env):
+  """Terms shared between policy and critic groups get independent settings."""
+  term_cfg = ObservationTermCfg(func=ClassObsTerm, params={})
+
+  cfg = {
+    "actor": ObservationGroupCfg(terms={"obs": term_cfg}, history_length=5),
+    "critic": ObservationGroupCfg(terms={"obs": term_cfg}),  # no history
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+
+  policy_obs = manager.compute()["actor"]
+  critic_obs = manager.compute()["critic"]
+
+  # Policy has history (5 * 3 = 15), critic doesn't (3).
+  assert isinstance(policy_obs, torch.Tensor)
+  assert isinstance(critic_obs, torch.Tensor)
+  assert policy_obs.shape[1] == 15  # 3 features * 5 history
+  assert critic_obs.shape[1] == 3  # 3 features, no history
+
+
+def test_get_term_cfg_returns_resolved_config(mock_env):
+  """get_term_cfg should return resolved config with func as instance."""
+  term_cfg = ObservationTermCfg(func=ClassObsTerm, params={})
+
+  # Original config has func as a class.
+  assert inspect.isclass(term_cfg.func), "precondition: func should be a class"
+
+  cfg = {"actor": ObservationGroupCfg(terms={"obs1": term_cfg})}
+  manager = ObservationManager(cfg, mock_env)
+
+  # Original config should remain unchanged.
+  assert inspect.isclass(term_cfg.func), "original config should not be mutated"
+  assert inspect.isclass(manager.cfg["actor"].terms["obs1"].func), (
+    "manager.cfg should preserve original class"
+  )
+
+  # Resolved config should have func as an instance.
+  resolved_cfg = manager.get_term_cfg("actor", "obs1")
+  assert not inspect.isclass(resolved_cfg.func), (
+    "get_term_cfg should return resolved config with func as instance"
+  )
+  assert isinstance(resolved_cfg.func, ClassObsTerm), (
+    f"Expected ClassObsTerm instance, got {type(resolved_cfg.func)}"
+  )
+
+
+def test_scene_entity_cfg_in_params_not_mutated(device):
+  """SceneEntityCfg in params should not be mutated when resolved."""
+  # NOTE: 2 joints so selecting one won't optimize to slice(None).
+  xml = """
+  <mujoco>
+    <worldbody>
+      <body name="body1" pos="0 0 1">
+        <joint name="joint1" type="hinge" axis="0 0 1"/>
+        <geom type="sphere" size="0.1" mass="1"/>
+        <body name="body2" pos="0.2 0 0">
+          <joint name="joint2" type="hinge" axis="0 0 1"/>
+          <geom type="sphere" size="0.1" mass="1"/>
+        </body>
+      </body>
+    </worldbody>
+  </mujoco>
+  """
+  entity_cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(xml))
+  entity = Entity(entity_cfg)
+  model = entity.compile()
+
+  sim = Simulation(num_envs=2, cfg=SimulationCfg(), model=model, device=device)
+  entity.initialize(model, sim.model, sim.data, device)
+
+  env = Mock()
+  env.num_envs = 2
+  env.device = device
+  env.max_episode_length_s = 10.0
+  env.scene = {"robot": entity}
+
+  # Select only joint1 --> joint_ids will mutate from slice(None) to [0].
+  asset_cfg = SceneEntityCfg(name="robot", joint_names=("joint1",))
+  original_joint_ids = asset_cfg.joint_ids  # slice(None) before resolve
+
+  def reward_func(env, asset_cfg):
+    del asset_cfg  # Unused.
+    return torch.ones(env.num_envs, device=env.device)
+
+  term_cfg = RewardTermCfg(
+    func=reward_func, params={"asset_cfg": asset_cfg}, weight=1.0
+  )
+
+  _ = RewardManager({"reward1": term_cfg}, env)
+
+  assert asset_cfg.joint_ids == original_joint_ids, (
+    f"SceneEntityCfg.joint_ids was mutated from {original_joint_ids} to "
+    f"{asset_cfg.joint_ids}"
+  )
@@ -0,0 +1,93 @@
+"""Tests for lifting command resample kinematics freshness.
+
+A timer-expiry resample runs inside command compute, after the step's single
+sim.forward(). The commands must refresh kinematics themselves so observations
+and the reward-path cache see post-teleport state.
+"""
+
+from types import SimpleNamespace
+from typing import TYPE_CHECKING, cast
+
+import pytest
+import torch
+from conftest import get_test_device, make_scene_and_sim
+
+from mjlab.tasks.manipulation.mdp.commands import (
+  LiftingCommandCfg,
+  MultiCubeLiftingCommandCfg,
+)
+
+if TYPE_CHECKING:
+  from mjlab.envs import ManagerBasedRlEnv
+
+CUBE_XML = """
+<mujoco>
+  <worldbody>
+    <body name="cube">
+      <freejoint/>
+      <geom name="cube_geom" type="box" size="0.02 0.02 0.02" mass="0.1"/>
+    </body>
+  </worldbody>
+</mujoco>
+"""
+
+
+@pytest.fixture(scope="module")
+def device():
+  return get_test_device()
+
+
+def _make_env(device, entity_names):
+  scene, sim = make_scene_and_sim(
+    device, {name: CUBE_XML for name in entity_names}, sensors=(), num_envs=2
+  )
+  env = SimpleNamespace(scene=scene, sim=sim, num_envs=scene.num_envs, device=device)
+  return cast("ManagerBasedRlEnv", env)
+
+
+def _assert_kinematics_fresh(scene, sim, entity_name):
+  """xpos-derived root pos must match qpos after a compute-path resample."""
+  ent = scene[entity_name]
+  q_adr = ent.indexing.free_joint_q_adr
+  qpos_pos = sim.data.qpos[:, q_adr[:3]]
+  assert torch.allclose(ent.data.root_link_pos_w, qpos_pos, atol=1e-6)
+
+
+def test_lifting_resample_refreshes_kinematics(device):
+  env = _make_env(device, ("cube0",))
+  cfg = LiftingCommandCfg(entity_name="cube0", resampling_time_range=(0.001, 0.001))
+  term = cfg.build(env)
+
+  term.reset(env_ids=torch.arange(env.num_envs, device=device))
+  env.sim.forward()
+
+  counter = term.command_counter.clone()
+  term.compute(dt=1.0)  # Timer (1 ms) expires: teleports the object.
+  assert (term.command_counter > counter).all()
+
+  _assert_kinematics_fresh(env.scene, env.sim, "cube0")
+
+
+def test_multi_cube_resample_refreshes_kinematics_and_cache(device):
+  names = ("cube0", "cube1")
+  env = _make_env(device, names)
+  cfg = MultiCubeLiftingCommandCfg(
+    entity_names=names, resampling_time_range=(0.001, 0.001)
+  )
+  term = cfg.build(env)
+
+  term.reset(env_ids=torch.arange(env.num_envs, device=device))
+  env.sim.forward()
+
+  counter = term.command_counter.clone()
+  term.compute(dt=1.0)  # Timer (1 ms) expires: teleports all cubes.
+  assert (term.command_counter > counter).all()
+
+  for name in names:
+    _assert_kinematics_fresh(env.scene, env.sim, name)
+
+  # Reward-path cache serves the post-teleport position of the target cube.
+  all_pos = torch.stack([env.scene[n].data.root_link_pos_w for n in names])
+  arange = torch.arange(env.num_envs, device=device)
+  expected = all_pos[term.target_selection, arange]
+  assert torch.allclose(term.target_object_pos(), expected, atol=1e-6)
@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+from types import SimpleNamespace
+from typing import TYPE_CHECKING, cast
+
+import mujoco
+import torch
+from conftest import get_test_device
+
+from mjlab.sensor import CameraSensorData
+from mjlab.tasks.manipulation.mdp.commands import MultiCubeLiftingCommand
+from mjlab.tasks.manipulation.mdp.observations import (
+  camera_segmentation,
+  camera_target_cube_mask,
+)
+
+if TYPE_CHECKING:
+  from mjlab.envs import ManagerBasedRlEnv
+
+
+def _make_env(segmentation: torch.Tensor, target_geom_ids: torch.Tensor):
+  sensor = SimpleNamespace(data=CameraSensorData(segmentation=segmentation))
+
+  command = object.__new__(MultiCubeLiftingCommand)
+  command._padded_geom_ids = target_geom_ids
+  command.target_selection = torch.arange(
+    target_geom_ids.shape[0], device=target_geom_ids.device
+  )
+
+  command_manager = SimpleNamespace(get_term=lambda _: command)
+  return SimpleNamespace(
+    scene={"seg_cam": sensor},
+    command_manager=command_manager,
+  )
+
+
+def test_camera_segmentation_returns_bchw():
+  device = get_test_device()
+  geom = int(mujoco.mjtObj.mjOBJ_GEOM)
+  seg = torch.tensor(
+    [
+      [[[1, geom], [2, geom], [-1, -1]], [[3, geom], [4, geom], [-1, -1]]],
+      [[[5, geom], [6, geom], [-1, -1]], [[7, geom], [8, geom], [-1, -1]]],
+    ],
+    dtype=torch.int32,
+    device=device,
+  )
+  env = _make_env(seg, torch.tensor([[1], [7]], dtype=torch.int32, device=device))
+  env = cast("ManagerBasedRlEnv", env)
+
+  obs = camera_segmentation(env, "seg_cam")
+
+  assert obs.shape == (2, 2, 2, 3)
+  assert obs.dtype == torch.int32
+  assert torch.equal(obs[:, 0], seg[..., 0])
+  assert torch.equal(obs[:, 1], seg[..., 1])
+
+
+def test_camera_target_cube_mask_filters_to_geom_hits():
+  device = get_test_device()
+  geom = int(mujoco.mjtObj.mjOBJ_GEOM)
+  flex = int(mujoco.mjtObj.mjOBJ_FLEX)
+  seg = torch.tensor(
+    [
+      [[[3, geom], [3, flex], [0, geom]], [[-1, -1], [4, geom], [3, geom]]],
+      [[[5, geom], [7, flex], [5, geom]], [[7, geom], [-1, -1], [0, geom]]],
+    ],
+    dtype=torch.int32,
+    device=device,
+  )
+  env = _make_env(seg, torch.tensor([[3], [7]], dtype=torch.int32, device=device))
+  env = cast("ManagerBasedRlEnv", env)
+
+  mask = camera_target_cube_mask(env, "seg_cam", "lift_height")
+
+  expected = torch.tensor(
+    [
+      [[[1.0, 0.0, 0.0], [0.0, 0.0, 1.0]]],
+      [[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]],
+    ],
+    dtype=torch.float32,
+    device=device,
+  )
+  assert mask.shape == (2, 1, 2, 3)
+  assert torch.equal(mask, expected)
@@ -0,0 +1,359 @@
+"""Tests for metrics manager functionality."""
+
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from mjlab.managers.metrics_manager import (
+  MetricsManager,
+  MetricsTermCfg,
+  NullMetricsManager,
+)
+
+
+class SimpleTestMetric:
+  """A class-based metric that tracks state."""
+
+  def __init__(self, cfg: MetricsTermCfg, env):
+    self.call_count = torch.zeros(env.num_envs, device=env.device)
+
+  def __call__(self, env, **kwargs):
+    self.call_count += 1
+    return torch.ones(env.num_envs, device=env.device) * 0.5
+
+  def reset(self, env_ids: torch.Tensor | None = None, env=None):
+    if env_ids is not None and len(env_ids) > 0:
+      self.call_count[env_ids] = 0
+
+
+@pytest.fixture
+def mock_env():
+  env = Mock()
+  env.num_envs = 4
+  env.device = "cpu"
+  env.scene = {"robot": Mock()}
+  return env
+
+
+def test_episode_averages_and_reset(mock_env):
+  """Compute for N steps, reset a subset, verify averages and zeroing."""
+  cfg = {
+    "term": MetricsTermCfg(
+      func=lambda env: torch.ones(env.num_envs, device=env.device) * 0.5,
+      params={},
+    )
+  }
+  manager = MetricsManager(cfg, mock_env)
+
+  for _ in range(10):
+    manager.compute()
+
+  info = manager.reset(env_ids=torch.tensor([0, 1]))
+
+  # Each env: sum=5.0, count=10, avg=0.5. Mean across 2 reset envs = 0.5.
+  assert info["Episode_Metrics/term"].item() == pytest.approx(0.5)
+  # Reset envs zeroed; non-reset envs untouched.
+  assert manager._episode_sums["term"][0] == 0.0
+  assert manager._step_count[0] == 0
+  assert manager._episode_sums["term"][2] == pytest.approx(5.0)
+  assert manager._step_count[2] == 10
+
+
+def test_early_termination_uses_per_env_step_count(mock_env):
+  """Envs with different episode lengths get correct per-step averages."""
+  step = [0]
+
+  def step_dependent_metric(env):
+    step[0] += 1
+    return torch.full((env.num_envs,), float(step[0]), device=env.device)
+
+  cfg = {"m": MetricsTermCfg(func=step_dependent_metric, params={})}
+  manager = MetricsManager(cfg, mock_env)
+
+  # 4 steps for all envs: values are 1, 2, 3, 4.
+  for _ in range(4):
+    manager.compute()
+  # Env 0: sum=10, count=4. Reset it (env 1 keeps accumulating).
+  manager.reset(env_ids=torch.tensor([0]))
+
+  # 2 more steps: values are 5, 6.
+  for _ in range(2):
+    manager.compute()
+  # Env 0: sum=11, count=2, avg=5.5.
+  # Env 1: sum=21, count=6, avg=3.5.
+  info = manager.reset(env_ids=torch.tensor([0, 1]))
+  # Mean of [5.5, 3.5] = 4.5.
+  assert info["Episode_Metrics/m"].item() == pytest.approx(4.5)
+
+
+def test_class_based_metric_reset_targets_correct_envs(mock_env):
+  """Class-based term's reset() is called with the correct env_ids."""
+  cfg = {"term": MetricsTermCfg(func=SimpleTestMetric, params={})}
+  manager = MetricsManager(cfg, mock_env)
+  term = manager._class_term_cfgs[0].func
+
+  for _ in range(10):
+    manager.compute()
+
+  manager.reset(env_ids=torch.tensor([0, 2]))
+
+  assert term.call_count[0] == 0
+  assert term.call_count[1] == 10
+  assert term.call_count[2] == 0
+  assert term.call_count[3] == 10
+
+
+def test_null_metrics_manager(mock_env):
+  """NullMetricsManager doesn't crash and returns empty dict on reset."""
+  manager = NullMetricsManager()
+  manager.compute()
+  assert manager.reset(env_ids=torch.tensor([0])) == {}
+
+
+def test_none_terms_are_skipped(mock_env):
+  """None terms in config are skipped without error."""
+  cfg: dict[str, MetricsTermCfg | None] = {
+    "valid": MetricsTermCfg(
+      func=lambda env: torch.ones(env.num_envs, device=env.device),
+      params={},
+    ),
+    "skipped": None,
+  }
+  manager = MetricsManager(cfg, mock_env)  # type: ignore[arg-type]
+  assert manager._term_names == ["valid"]
+
+
+def test_per_substep_averaging(mock_env):
+  """Substep terms are averaged across substeps; step terms evaluated once."""
+  substep_call = [0]
+
+  def rising_metric(env):
+    substep_call[0] += 1
+    return torch.full((env.num_envs,), float(substep_call[0]), device=env.device)
+
+  cfg = {
+    "step_term": MetricsTermCfg(
+      func=lambda env: torch.ones(env.num_envs, device=env.device) * 2.0,
+      params={},
+    ),
+    "substep_term": MetricsTermCfg(func=rising_metric, params={}, per_substep=True),
+  }
+  manager = MetricsManager(cfg, mock_env)
+
+  for _ in range(2):  # 2 env steps
+    for _ in range(4):  # 4 substeps each
+      manager.compute_substep()
+    manager.compute()
+
+  info = manager.reset(env_ids=torch.tensor([0]))
+  assert info["Episode_Metrics/step_term"].item() == pytest.approx(2.0)
+  # Substeps: step 0 -> [1,2,3,4] avg=2.5; step 1 -> [5,6,7,8] avg=6.5.
+  # Episode avg = (2.5 + 6.5) / 2 = 4.5.  A broken single-eval path
+  # would give (1 + 2) / 2 = 1.5.
+  assert info["Episode_Metrics/substep_term"].item() == pytest.approx(4.5)
+
+
+def test_substep_reset_zeroes_accumulators(mock_env):
+  """Reset zeroes substep accumulators for the given env_ids."""
+  cfg = {
+    "sub": MetricsTermCfg(
+      func=lambda env: torch.ones(env.num_envs, device=env.device),
+      params={},
+      per_substep=True,
+    ),
+  }
+  manager = MetricsManager(cfg, mock_env)
+
+  # Accumulate some substep values without calling compute().
+  for _ in range(3):
+    manager.compute_substep()
+
+  # Reset env 0 -- its substep accum should be zeroed.
+  manager.reset(env_ids=torch.tensor([0]))
+  assert manager._substep_accum[0][0].item() == 0.0
+  # Env 1 still has its accumulated value.
+  assert manager._substep_accum[0][1].item() == pytest.approx(3.0)
+
+
+def test_reduce_last_reports_final_step_value(mock_env):
+  """reduce='last' reports the last step's value, not the episode average."""
+  step = [0]
+
+  def rising_metric(env):
+    step[0] += 1
+    return torch.full((env.num_envs,), float(step[0]), device=env.device)
+
+  cfg = {"term": MetricsTermCfg(func=rising_metric, params={}, reduce="last")}
+  manager = MetricsManager(cfg, mock_env)
+
+  for _ in range(4):
+    manager.compute()
+
+  info = manager.reset(env_ids=torch.tensor([0]))
+
+  # Values were 1, 2, 3, 4. Mean would be 2.5; last should be 4.
+  assert info["Episode_Metrics/term"].item() == pytest.approx(4.0)
+
+
+def test_reduce_mean_and_last_coexist(mock_env):
+  """Mixed reduce modes in the same manager report correctly."""
+  step = [0]
+
+  def rising_metric(env):
+    step[0] += 1
+    return torch.full((env.num_envs,), float(step[0]), device=env.device)
+
+  cfg = {
+    "mean_term": MetricsTermCfg(func=rising_metric, params={}, reduce="mean"),
+    "last_term": MetricsTermCfg(func=rising_metric, params={}, reduce="last"),
+  }
+  manager = MetricsManager(cfg, mock_env)
+
+  for _ in range(3):
+    manager.compute()
+
+  info = manager.reset(env_ids=torch.tensor([0]))
+
+  # mean_term sees steps 1, 3, 5 (odd calls); episode avg = 3.0.
+  # last_term sees steps 2, 4, 6 (even calls); last value = 6.0.
+  assert info["Episode_Metrics/mean_term"].item() == pytest.approx(3.0)
+  assert info["Episode_Metrics/last_term"].item() == pytest.approx(6.0)
+
+
+def test_metrics_step_shape_validation_rejects_bad_compute_output(mock_env):
+  """Step metrics must return one scalar per environment."""
+  cfg = {"bad": MetricsTermCfg(func=lambda env: torch.ones(env.num_envs, 1))}
+  manager = MetricsManager(cfg, mock_env)
+  with pytest.raises(ValueError, match="MetricsManager term 'bad'.*expected \\(4,\\)"):
+    manager.compute()
+
+
+def test_metrics_substep_shape_validation_rejects_bad_compute_output(mock_env):
+  """Substep metrics must return one scalar per environment."""
+  cfg = {
+    "bad": MetricsTermCfg(
+      func=lambda env: torch.ones(env.num_envs, 1), per_substep=True
+    )
+  }
+  manager = MetricsManager(cfg, mock_env)
+  with pytest.raises(ValueError, match="MetricsManager term 'bad'.*expected \\(4,\\)"):
+    manager.compute_substep()
+
+
+def test_reduce_max_reports_episode_peak(mock_env):
+  """reduce='max' reports the highest value seen during the episode."""
+  step = [0]
+
+  def rising_then_falling(env):
+    step[0] += 1
+    val = 5.0 - abs(step[0] - 3)  # values: 3, 4, 5, 4, 3
+    return torch.full((env.num_envs,), val, device=env.device)
+
+  cfg = {"term": MetricsTermCfg(func=rising_then_falling, params={}, reduce="max")}
+  manager = MetricsManager(cfg, mock_env)
+
+  for _ in range(5):
+    manager.compute()
+
+  info = manager.reset(env_ids=torch.tensor([0]))
+
+  assert info["Episode_Metrics/term"].item() == pytest.approx(5.0)
+
+
+def test_reduce_max_reset_clears_to_neg_inf(mock_env):
+  """After reset, max tracking restarts from -inf."""
+  cfg = {
+    "term": MetricsTermCfg(
+      func=lambda env: torch.ones(env.num_envs, device=env.device) * 10.0,
+      params={},
+      reduce="max",
+    )
+  }
+  manager = MetricsManager(cfg, mock_env)
+
+  manager.compute()
+  manager.reset(env_ids=torch.tensor([0]))
+
+  # After reset, the max buffer for env 0 should be -inf.
+  assert manager._episode_max["term"][0].item() == float("-inf")
+  # Env 1 was not reset, so it keeps its max.
+  assert manager._episode_max["term"][1].item() == pytest.approx(10.0)
+
+
+def test_reduce_max_coexists_with_mean_and_last(mock_env):
+  """All three reduce modes work correctly in the same manager."""
+  step = [0]
+
+  def rising_metric(env):
+    step[0] += 1
+    return torch.full((env.num_envs,), float(step[0]), device=env.device)
+
+  cfg = {
+    "mean_term": MetricsTermCfg(func=rising_metric, params={}, reduce="mean"),
+    "last_term": MetricsTermCfg(func=rising_metric, params={}, reduce="last"),
+    "max_term": MetricsTermCfg(func=rising_metric, params={}, reduce="max"),
+  }
+  manager = MetricsManager(cfg, mock_env)
+
+  for _ in range(3):
+    manager.compute()
+
+  info = manager.reset(env_ids=torch.tensor([0]))
+
+  # mean_term sees 1, 4, 7; avg = 4.0
+  assert info["Episode_Metrics/mean_term"].item() == pytest.approx(4.0)
+  # last_term sees 2, 5, 8; last = 8.0
+  assert info["Episode_Metrics/last_term"].item() == pytest.approx(8.0)
+  # max_term sees 3, 6, 9; max = 9.0
+  assert info["Episode_Metrics/max_term"].item() == pytest.approx(9.0)
+
+
+def test_reduce_max_with_per_substep(mock_env):
+  """reduce='max' with per_substep tracks the max of step-averaged values."""
+  substep_call = [0]
+
+  def rising_metric(env):
+    substep_call[0] += 1
+    return torch.full((env.num_envs,), float(substep_call[0]), device=env.device)
+
+  cfg = {
+    "sub_max": MetricsTermCfg(
+      func=rising_metric, params={}, per_substep=True, reduce="max"
+    ),
+  }
+  manager = MetricsManager(cfg, mock_env)
+
+  # Step 0: substeps [1, 2] -> avg = 1.5
+  for _ in range(2):
+    manager.compute_substep()
+  manager.compute()
+
+  # Step 1: substeps [3, 4] -> avg = 3.5
+  for _ in range(2):
+    manager.compute_substep()
+  manager.compute()
+
+  info = manager.reset(env_ids=torch.tensor([0]))
+
+  # Max of step averages: max(1.5, 3.5) = 3.5
+  assert info["Episode_Metrics/sub_max"].item() == pytest.approx(3.5)
+
+
+def test_no_substep_terms_no_overhead(mock_env):
+  """When no per_substep terms exist, compute_substep is a no-op."""
+  cfg = {
+    "step_only": MetricsTermCfg(
+      func=lambda env: torch.ones(env.num_envs, device=env.device),
+      params={},
+    ),
+  }
+  manager = MetricsManager(cfg, mock_env)
+
+  # Should not raise or change state.
+  manager.compute_substep()
+  assert manager._substep_count == 0
+
+  manager.compute()
+  info = manager.reset(env_ids=torch.tensor([0]))
+  assert info["Episode_Metrics/step_only"].item() == pytest.approx(1.0)
@@ -0,0 +1,409 @@
+"""Tests for NaN guard functionality."""
+
+import tempfile
+from pathlib import Path
+
+import mujoco
+import numpy as np
+import pytest
+import torch
+from conftest import get_test_device
+
+from mjlab.sim.sim import Simulation, SimulationCfg
+from mjlab.utils.nan_guard import NanGuardCfg
+
+
+@pytest.fixture
+def simple_model():
+  """Create a simple MuJoCo model for testing."""
+  xml = """
+  <mujoco>
+    <worldbody>
+      <body>
+        <freejoint/>
+        <geom type="box" size="0.1 0.1 0.1"/>
+      </body>
+    </worldbody>
+  </mujoco>
+  """
+  spec = mujoco.MjSpec.from_string(xml)
+  return spec.compile()
+
+
+@pytest.fixture
+def mocap_model():
+  """Create a model with a mocap body and a freejoint body."""
+  xml = """
+  <mujoco>
+    <worldbody>
+      <body name="mocap_body" mocap="true">
+        <geom type="sphere" size="0.05"/>
+      </body>
+      <body name="free_body">
+        <freejoint/>
+        <geom type="box" size="0.1 0.1 0.1"/>
+      </body>
+    </worldbody>
+  </mujoco>
+  """
+  spec = mujoco.MjSpec.from_string(xml)
+  return spec.compile()
+
+
+def test_nan_guard_disabled_by_default(simple_model):
+  """NaN guard should be disabled by default with no overhead."""
+  cfg = SimulationCfg()
+  sim = Simulation(num_envs=2, cfg=cfg, model=simple_model, device=get_test_device())
+
+  assert not sim.nan_guard.enabled
+  sim.step()  # Should not trigger any capture.
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Likely bug on CPU MjWarp")
+def test_nan_guard_captures_and_dumps_on_nan(simple_model):
+  """NaN guard should capture states and dump when NaN detected."""
+  with tempfile.TemporaryDirectory() as tmpdir:
+    cfg = SimulationCfg(
+      nan_guard=NanGuardCfg(
+        enabled=True,
+        buffer_size=5,
+        output_dir=tmpdir,
+        max_envs_to_dump=2,
+      )
+    )
+    sim = Simulation(num_envs=4, cfg=cfg, model=simple_model, device=get_test_device())
+
+    # Run a few steps to populate buffer.
+    for _ in range(3):
+      sim.step()
+
+    # Inject NaN into environment 1.
+    sim.data.qpos[1, 0] = float("nan")
+
+    # Next step should trigger dump.
+    sim.step()
+
+    # Check that timestamped dump file was created.
+    dump_files = [
+      f for f in Path(tmpdir).glob("nan_dump_*.npz") if "latest" not in f.name
+    ]
+    assert len(dump_files) == 1
+
+    # Load and inspect the dump.
+    dump = np.load(dump_files[0], allow_pickle=True)
+    metadata = dump["_metadata"].item()
+
+    assert metadata["num_envs_total"] == 4
+    assert metadata["num_envs_dumped"] == 1
+    assert 1 in metadata["nan_env_ids"]
+    assert metadata["buffer_size"] == 4
+
+    # Check that states were captured.
+    assert "states_step_000000" in dump
+    assert "states_step_000001" in dump
+    assert "states_step_000002" in dump
+    assert "states_step_000003" in dump
+
+    # Verify state shape: (num_envs_dumped, state_size).
+    state = dump["states_step_000000"]
+    assert state.shape[0] == 1
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Likely bug on CPU MjWarp")
+def test_nan_guard_detects_correct_env_ids(simple_model):
+  """NaN guard should correctly identify which environments have NaN/Inf."""
+  with tempfile.TemporaryDirectory() as tmpdir:
+    cfg = SimulationCfg(
+      nan_guard=NanGuardCfg(enabled=True, buffer_size=5, output_dir=tmpdir)
+    )
+    sim = Simulation(num_envs=10, cfg=cfg, model=simple_model, device=get_test_device())
+
+    # Run a few steps to populate buffer.
+    for _ in range(3):
+      sim.step()
+
+    # Inject NaN/Inf into environments 2, 5, and 7 (in state variables).
+    sim.data.qpos[2, 0] = float("nan")
+    sim.data.qvel[5, 1] = float("nan")
+    sim.data.qvel[7, 2] = float("inf")
+
+    # Next step should trigger dump.
+    sim.step()
+
+    # Load and inspect the dump.
+    dump_files = [
+      f for f in Path(tmpdir).glob("nan_dump_*.npz") if "latest" not in f.name
+    ]
+    assert len(dump_files) == 1
+
+    dump = np.load(dump_files[0], allow_pickle=True)
+    metadata = dump["_metadata"].item()
+
+    # Should detect exactly the environments with NaN/Inf.
+    nan_env_ids = set(metadata["nan_env_ids"])
+    assert nan_env_ids == {2, 5, 7}, f"Expected {{2, 5, 7}}, got {nan_env_ids}"
+
+
+def test_nan_guard_saves_model(simple_model):
+  """NaN guard should save model file alongside state dump."""
+  with tempfile.TemporaryDirectory() as tmpdir:
+    cfg = SimulationCfg(
+      nan_guard=NanGuardCfg(enabled=True, buffer_size=5, output_dir=tmpdir)
+    )
+    sim = Simulation(num_envs=2, cfg=cfg, model=simple_model, device=get_test_device())
+
+    # Inject NaN and trigger dump.
+    sim.data.qpos[0, 0] = float("nan")
+    sim.step()
+
+    # Check that both dump and model files were created.
+    dump_files = [
+      f for f in Path(tmpdir).glob("nan_dump_*.npz") if "latest" not in f.name
+    ]
+    model_files = [
+      f for f in Path(tmpdir).glob("model_*.mjb") if "latest" not in f.name
+    ]
+    assert len(dump_files) == 1
+    assert len(model_files) == 1
+
+    # Verify model file can be loaded.
+    dump = np.load(dump_files[0], allow_pickle=True)
+    metadata = dump["_metadata"].item()
+    model_path = Path(tmpdir) / metadata["model_file"]
+    assert model_path.exists()
+
+    # Load model and verify it's valid.
+    loaded_model = mujoco.MjModel.from_binary_path(str(model_path))
+    assert loaded_model.nq == simple_model.nq
+    assert loaded_model.nv == simple_model.nv
+
+
+@pytest.mark.slow
+def test_nan_guard_with_complex_model():
+  """NaN guard should work with complex robot model."""
+  from mjlab.scene import Scene
+  from mjlab.tasks.velocity.config.go1.env_cfgs import unitree_go1_rough_env_cfg
+
+  scene = Scene(unitree_go1_rough_env_cfg().scene, device=get_test_device())
+  model = scene.compile()
+
+  with tempfile.TemporaryDirectory() as tmpdir:
+    cfg = SimulationCfg(
+      nan_guard=NanGuardCfg(enabled=True, buffer_size=3, output_dir=tmpdir)
+    )
+    sim = Simulation(num_envs=2, cfg=cfg, model=model, device=get_test_device())
+
+    # Run a few steps.
+    for _ in range(2):
+      sim.step()
+
+    # Inject NaN and trigger dump.
+    sim.data.qpos[0, 0] = float("nan")
+    sim.step()
+
+    # Verify dump and model files were created.
+    dump_files = [
+      f for f in Path(tmpdir).glob("nan_dump_*.npz") if "latest" not in f.name
+    ]
+    model_files = [
+      f for f in Path(tmpdir).glob("model_*.mjb") if "latest" not in f.name
+    ]
+    assert len(dump_files) == 1
+    assert len(model_files) == 1
+
+    # Load the saved model and verify it matches.
+    dump = np.load(dump_files[0], allow_pickle=True)
+    metadata = dump["_metadata"].item()
+    model_path = Path(tmpdir) / metadata["model_file"]
+
+    loaded_model = mujoco.MjModel.from_binary_path(str(model_path))
+    assert loaded_model.nq == model.nq
+    assert loaded_model.nv == model.nv
+    assert loaded_model.nu == model.nu  # Check actuators too
+
+    # Verify we can create MjData and restore a state.
+    loaded_data = mujoco.MjData(loaded_model)
+    state = dump["states_step_000000"][0]
+    state_spec = metadata.get("state_spec", mujoco.mjtState.mjSTATE_PHYSICS.value)
+    mujoco.mj_setState(loaded_model, loaded_data, state, state_spec)
+    mujoco.mj_forward(loaded_model, loaded_data)
+
+    # Data should be valid (no NaN in derived quantities after forward).
+    assert not np.isnan(loaded_data.qpos).any()
+
+
+def test_nan_guard_only_dumps_once(simple_model):
+  """NaN guard should only dump once per training run."""
+  with tempfile.TemporaryDirectory() as tmpdir:
+    cfg = SimulationCfg(
+      nan_guard=NanGuardCfg(enabled=True, buffer_size=5, output_dir=tmpdir)
+    )
+    sim = Simulation(num_envs=2, cfg=cfg, model=simple_model, device=get_test_device())
+
+    # Inject NaN.
+    sim.data.qpos[0, 0] = float("nan")
+    sim.step()
+
+    # Should have exactly one timestamped dump.
+    dump_files = [
+      f for f in Path(tmpdir).glob("nan_dump_*.npz") if "latest" not in f.name
+    ]
+    assert len(dump_files) == 1
+
+    # Inject another NaN.
+    sim.data.qpos[1, 0] = float("nan")
+    sim.step()
+
+    # Should still have only one timestamped dump.
+    dump_files = [
+      f for f in Path(tmpdir).glob("nan_dump_*.npz") if "latest" not in f.name
+    ]
+    assert len(dump_files) == 1
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Likely bug on CPU MjWarp")
+def test_nan_guard_respects_buffer_size(simple_model):
+  """NaN guard should only keep last K states in buffer."""
+  with tempfile.TemporaryDirectory() as tmpdir:
+    cfg = SimulationCfg(
+      nan_guard=NanGuardCfg(enabled=True, buffer_size=3, output_dir=tmpdir)
+    )
+    sim = Simulation(num_envs=2, cfg=cfg, model=simple_model, device=get_test_device())
+
+    # Run 10 steps.
+    for _ in range(10):
+      sim.step()
+
+    # Inject NaN.
+    sim.data.qpos[0, 0] = float("nan")
+    sim.step()
+
+    # Load dump.
+    dump_files = [
+      f for f in Path(tmpdir).glob("nan_dump_*.npz") if "latest" not in f.name
+    ]
+    dump = np.load(dump_files[0], allow_pickle=True)
+
+    # Should only have 3 states (buffer size).
+    state_keys = [k for k in dump.keys() if k.startswith("states_step_")]
+    assert len(state_keys) == 3
+
+    # Should be the last 3 steps (steps 8, 9, 10 where step 10 has NaN).
+    assert "states_step_000008" in dump
+    assert "states_step_000009" in dump
+    assert "states_step_000010" in dump
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Likely bug on CPU MjWarp")
+def test_nan_guard_captures_high_indexed_envs(simple_model):
+  """NaN guard should capture NaN in high-indexed environments beyond max_envs_to_dump."""
+  with tempfile.TemporaryDirectory() as tmpdir:
+    cfg = SimulationCfg(
+      nan_guard=NanGuardCfg(
+        enabled=True, buffer_size=5, output_dir=tmpdir, max_envs_to_dump=3
+      )
+    )
+    sim = Simulation(num_envs=10, cfg=cfg, model=simple_model, device=get_test_device())
+
+    for _ in range(3):
+      sim.step()
+
+    sim.data.qpos[7, 0] = float("nan")
+    sim.step()
+
+    dump_files = [
+      f for f in Path(tmpdir).glob("nan_dump_*.npz") if "latest" not in f.name
+    ]
+    assert len(dump_files) == 1
+
+    dump = np.load(dump_files[0], allow_pickle=True)
+    metadata = dump["_metadata"].item()
+
+    assert 7 in metadata["nan_env_ids"]
+    assert 7 in metadata["dumped_env_ids"]
+    assert metadata["num_envs_total"] == 10
+    assert metadata["num_envs_dumped"] == 1
+
+    state = dump["states_step_000000"]
+    assert state.shape[0] == 1
+
+
+def test_nan_guard_creates_latest_symlinks(simple_model):
+  """NaN guard should create latest symlinks that work correctly."""
+  with tempfile.TemporaryDirectory() as tmpdir:
+    cfg = SimulationCfg(
+      nan_guard=NanGuardCfg(enabled=True, buffer_size=5, output_dir=tmpdir)
+    )
+    sim = Simulation(num_envs=2, cfg=cfg, model=simple_model, device=get_test_device())
+
+    sim.data.qpos[0, 0] = float("nan")
+    sim.step()
+
+    latest_dump = Path(tmpdir) / "nan_dump_latest.npz"
+    latest_model = Path(tmpdir) / "model_latest.mjb"
+
+    # Verify symlinks exist.
+    assert latest_dump.exists()
+    assert latest_model.exists()
+    assert latest_dump.is_symlink()
+    assert latest_model.is_symlink()
+
+    # Load via symlink and verify it works.
+    dump = np.load(latest_dump, allow_pickle=True)
+    metadata = dump["_metadata"].item()
+
+    # Metadata should reference the timestamped model file.
+    assert metadata["model_file"].startswith("model_")
+    assert metadata["model_file"].endswith(".mjb")
+    assert (Path(tmpdir) / metadata["model_file"]).exists()
+
+    # Loading via symlink should work.
+    loaded_model = mujoco.MjModel.from_binary_path(str(latest_model))
+    assert loaded_model.nq == simple_model.nq
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Likely bug on CPU MjWarp")
+def test_nan_guard_captures_mocap_state(mocap_model):
+  """NaN guard should capture and restore mocap body poses."""
+  with tempfile.TemporaryDirectory() as tmpdir:
+    cfg = SimulationCfg(
+      nan_guard=NanGuardCfg(enabled=True, buffer_size=5, output_dir=tmpdir)
+    )
+    sim = Simulation(num_envs=2, cfg=cfg, model=mocap_model, device=get_test_device())
+
+    # Set a known mocap position before stepping.
+    sim.data.mocap_pos[0, 0] = torch.tensor([1.0, 2.0, 3.0], device=get_test_device())
+
+    # Run a few steps.
+    for _ in range(3):
+      sim.step()
+
+    # Inject NaN to trigger dump.
+    sim.data.qpos[0, 0] = float("nan")
+    sim.step()
+
+    # Load dump.
+    dump_files = [
+      f for f in Path(tmpdir).glob("nan_dump_*.npz") if "latest" not in f.name
+    ]
+    assert len(dump_files) == 1
+
+    dump = np.load(dump_files[0], allow_pickle=True)
+    metadata = dump["_metadata"].item()
+
+    # Metadata should include state_spec with mocap flags.
+    assert "state_spec" in metadata
+    state_spec = metadata["state_spec"]
+    assert state_spec & mujoco.mjtState.mjSTATE_MOCAP_POS.value
+    assert state_spec & mujoco.mjtState.mjSTATE_MOCAP_QUAT.value
+
+    # State size should be larger than physics-only.
+    physics_size = mujoco.mj_stateSize(mocap_model, mujoco.mjtState.mjSTATE_PHYSICS)
+    assert metadata["state_size"] > physics_size
+
+    # Restore a state and verify mocap data round-trips.
+    loaded_data = mujoco.MjData(mocap_model)
+    state = dump["states_step_000000"][0]
+    mujoco.mj_setState(mocap_model, loaded_data, state, state_spec)
+    assert loaded_data.mocap_pos.shape == (1, 3)
+    np.testing.assert_allclose(loaded_data.mocap_pos[0], [1.0, 2.0, 3.0], atol=1e-5)
@@ -0,0 +1,40 @@
+"""Tests for native viewer custom action dispatch."""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+from mjlab.viewer.base import ViewerAction
+from mjlab.viewer.native.viewer import NativeMujocoViewer
+
+
+def _make_viewer(num_envs: int = 3) -> NativeMujocoViewer:
+  env = MagicMock()
+  env.cfg.viewer.env_idx = 0
+  env.unwrapped.num_envs = num_envs
+  return NativeMujocoViewer(env, MagicMock())
+
+
+def test_toggle_actions_dispatch_by_action_enum():
+  v = _make_viewer()
+
+  assert not v._show_plots
+  assert v._handle_custom_action(ViewerAction.TOGGLE_PLOTS, None)
+  assert v._show_plots
+
+  assert v._show_debug_vis
+  assert v._handle_custom_action(ViewerAction.TOGGLE_DEBUG_VIS, None)
+  assert not v._show_debug_vis
+
+  assert not v._show_all_envs
+  assert v._handle_custom_action(ViewerAction.TOGGLE_SHOW_ALL_ENVS, None)
+  assert v._show_all_envs
+
+
+def test_prev_next_env_actions_wrap_and_succeed():
+  v = _make_viewer(num_envs=3)
+  v.env_idx = 0
+  assert v._handle_custom_action(ViewerAction.PREV_ENV, None)
+  assert v.env_idx == 2
+  assert v._handle_custom_action(ViewerAction.NEXT_ENV, None)
+  assert v.env_idx == 0
@@ -0,0 +1,127 @@
+"""Validate that notebook cells import only symbols that exist in mjlab.
+
+The cartpole tutorial notebook writes Python files via ``%%writefile`` cells
+and imports directly from ``mjlab`` in others. When a public API is renamed
+or removed, those references silently rot until a user runs the notebook.
+This test parses each code cell, finds ``from mjlab... import X`` statements,
+and verifies that ``X`` is actually importable.
+
+Modules that the notebook itself creates (via ``%%writefile`` into the mjlab
+source tree) are skipped, since they don't exist until the notebook runs.
+"""
+
+from __future__ import annotations
+
+import ast
+import importlib
+import json
+import re
+from pathlib import Path
+
+import pytest
+
+NOTEBOOKS_DIR = Path(__file__).parent.parent / "notebooks"
+
+# Match "%%writefile /content/mjlab/src/mjlab/foo/bar.py" style targets that
+# land inside the mjlab package.
+_WRITEFILE_MJLAB_RE = re.compile(
+  r"^\s*%%writefile\s+\S*?/src/(mjlab(?:/[\w/]+)+)\.py\s*$"
+)
+
+
+def _extract_python_source(cell_source: str) -> str | None:
+  """Return Python source for a cell, or None if it has no Python to parse.
+
+  Handles ``%%writefile path.py`` cells by returning their body (only when the
+  target is a ``.py`` file; XML/YAML/etc. writefiles are skipped). Strips
+  shell escapes (``!...``) and line magics (``%...``).
+  """
+  lines = cell_source.splitlines()
+  if lines and lines[0].lstrip().startswith("%%writefile"):
+    target = lines[0].split(maxsplit=1)[1] if len(lines[0].split()) > 1 else ""
+    if not target.endswith(".py"):
+      return None
+    return "\n".join(lines[1:])
+  # Skip other cell magics (%%bash, %%html, etc.) entirely.
+  if lines and lines[0].lstrip().startswith("%%"):
+    return None
+  kept = [ln for ln in lines if not ln.lstrip().startswith(("!", "%"))]
+  return "\n".join(kept)
+
+
+def _iter_mjlab_imports(source: str):
+  """Yield (module, name) for every ``from mjlab... import name`` in source."""
+  tree = ast.parse(source)
+  for node in ast.walk(tree):
+    if isinstance(node, ast.ImportFrom) and node.module:
+      if node.module == "mjlab" or node.module.startswith("mjlab."):
+        for alias in node.names:
+          yield node.module, alias.name
+
+
+def _collect_written_modules(nb: dict) -> set[str]:
+  """Return the set of mjlab modules the notebook creates via %%writefile."""
+  written: set[str] = set()
+  for cell in nb.get("cells", []):
+    if cell.get("cell_type") != "code":
+      continue
+    source = "".join(cell.get("source", []))
+    first_line = source.splitlines()[0] if source else ""
+    match = _WRITEFILE_MJLAB_RE.match(first_line)
+    if match:
+      written.add(match.group(1).replace("/", "."))
+  return written
+
+
+def _collect_notebook_imports(path: Path):
+  """Return (imports, written_modules) where imports is a list of
+  (cell_index, module, name) tuples for every ``from mjlab... import`` in the
+  notebook, and written_modules is the set of mjlab modules the notebook
+  creates itself.
+  """
+  nb = json.loads(path.read_text())
+  written = _collect_written_modules(nb)
+  imports = []
+  for idx, cell in enumerate(nb.get("cells", [])):
+    if cell.get("cell_type") != "code":
+      continue
+    source = "".join(cell.get("source", []))
+    python_src = _extract_python_source(source)
+    if not python_src or not python_src.strip():
+      continue
+    try:
+      cell_imports = list(_iter_mjlab_imports(python_src))
+    except SyntaxError as e:
+      pytest.fail(f"Cell {idx} in {path.name} has invalid Python syntax: {e}")
+    for module, name in cell_imports:
+      imports.append((idx, module, name))
+  return imports, written
+
+
+@pytest.mark.parametrize("notebook_path", sorted(NOTEBOOKS_DIR.glob("*.ipynb")))
+def test_notebook_mjlab_imports_resolve(notebook_path: Path) -> None:
+  """Every ``from mjlab... import X`` in a notebook must resolve."""
+  imports, written_modules = _collect_notebook_imports(notebook_path)
+  failures = []
+  for cell_idx, module, name in imports:
+    # Skip modules the notebook writes itself; they don't exist until runtime.
+    if module in written_modules or any(
+      module.startswith(w + ".") for w in written_modules
+    ):
+      continue
+    try:
+      mod = importlib.import_module(module)
+    except ImportError as e:
+      failures.append(
+        f"cell {cell_idx}: `from {module} import {name}` failed to import module: {e}"
+      )
+      continue
+    if not hasattr(mod, name):
+      failures.append(
+        f"cell {cell_idx}: `from {module} import {name}` — "
+        f"'{name}' is not exposed by '{module}'"
+      )
+  if failures:
+    pytest.fail(
+      f"Stale mjlab imports in {notebook_path.name}:\n  " + "\n  ".join(failures)
+    )
@@ -0,0 +1,660 @@
+"""Tests for observation delay functionality."""
+
+from unittest.mock import Mock
+
+import pytest
+import torch
+from conftest import get_test_device
+
+from mjlab.managers.observation_manager import (
+  ObservationGroupCfg,
+  ObservationManager,
+  ObservationTermCfg,
+)
+
+
+@pytest.fixture
+def device():
+  """Test device fixture."""
+  return get_test_device()
+
+
+@pytest.fixture
+def mock_env(device):
+  """Create a mock environment."""
+  env = Mock()
+  env.num_envs = 4
+  env.device = device
+  env.step_dt = 0.02
+  return env
+
+
+@pytest.fixture
+def simple_obs_func(device):
+  """Create a simple observation function that returns a counter."""
+  counter = {"value": 0}
+
+  def obs_func(env):
+    counter["value"] += 1
+    return torch.full((env.num_envs, 3), float(counter["value"]), device=device)
+
+  return obs_func
+
+
+##
+# Basic delay tests.
+##
+
+
+def test_no_delay_by_default(mock_env, simple_obs_func):
+  """Test that observations work without delay (default behavior)."""
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs1": ObservationTermCfg(func=simple_obs_func, params={}),
+      }
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+  assert manager.group_obs_dim["actor"] == (3,)
+
+  obs = manager.compute(update_history=True)
+  policy_obs = obs["actor"]
+  assert isinstance(policy_obs, torch.Tensor)
+  assert policy_obs.shape == (4, 3)
+
+
+def test_constant_delay(mock_env, simple_obs_func, device):
+  """Test observation with constant delay (min_lag = max_lag = 2)."""
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs1": ObservationTermCfg(
+          func=simple_obs_func, params={}, delay_min_lag=2, delay_max_lag=2
+        ),
+      }
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+  # Note: counter is incremented during _prepare_terms (value=1).
+
+  # First compute uses value=2.
+  # Delay: lag=2 sampled, buffer only has 1 frame, clamped to lag=0, returns 2.
+  obs1 = manager.compute(update_history=True)
+  policy_obs1 = obs1["actor"]
+  assert isinstance(policy_obs1, torch.Tensor)
+  assert torch.allclose(policy_obs1[0], torch.full((3,), 2.0, device=device))
+
+  # Second compute uses value=3.
+  # Delay: lag=2, buffer has 2 frames, clamped to lag=1, returns 2.
+  obs2 = manager.compute(update_history=True)
+  policy_obs2 = obs2["actor"]
+  assert isinstance(policy_obs2, torch.Tensor)
+  assert torch.allclose(policy_obs2[0], torch.full((3,), 2.0, device=device))
+
+  # Third compute uses value=4.
+  # Delay: lag=2, buffer full (3 frames), returns value from 2 steps ago = 2.
+  obs3 = manager.compute(update_history=True)
+  policy_obs3 = obs3["actor"]
+  assert isinstance(policy_obs3, torch.Tensor)
+  assert torch.allclose(policy_obs3[0], torch.full((3,), 2.0, device=device))
+
+  # Fourth compute uses value=5.
+  # Delay: lag=2, returns value from 2 steps ago = 3.
+  obs4 = manager.compute(update_history=True)
+  policy_obs4 = obs4["actor"]
+  assert isinstance(policy_obs4, torch.Tensor)
+  assert torch.allclose(policy_obs4[0], torch.full((3,), 3.0, device=device))
+
+  # Fifth compute uses value=6.
+  # Delay: lag=2, returns value from 2 steps ago = 4.
+  obs5 = manager.compute(update_history=True)
+  policy_obs5 = obs5["actor"]
+  assert isinstance(policy_obs5, torch.Tensor)
+  assert torch.allclose(policy_obs5[0], torch.full((3,), 4.0, device=device))
+
+
+def test_zero_delay_returns_current(mock_env, simple_obs_func, device):
+  """Test that delay with min_lag=max_lag=0 returns current observation."""
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs1": ObservationTermCfg(
+          func=simple_obs_func, params={}, delay_min_lag=0, delay_max_lag=0
+        ),
+      }
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+
+  # First compute uses value=2 (counter at 1 after _prepare_terms).
+  obs1 = manager.compute(update_history=True)
+  policy_obs1 = obs1["actor"]
+  assert isinstance(policy_obs1, torch.Tensor)
+  assert torch.allclose(policy_obs1[0], torch.full((3,), 2.0, device=device))
+
+  # Second compute uses value=3.
+  obs2 = manager.compute(update_history=True)
+  policy_obs2 = obs2["actor"]
+  assert isinstance(policy_obs2, torch.Tensor)
+  assert torch.allclose(policy_obs2[0], torch.full((3,), 3.0, device=device))
+
+
+def test_lag_one_delay(mock_env, simple_obs_func, device):
+  """Test lag=1 returns previous observation."""
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs1": ObservationTermCfg(
+          func=simple_obs_func, params={}, delay_min_lag=1, delay_max_lag=1
+        ),
+      }
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+
+  # Build sequence: 2, 3, 4, 5, 6
+  # Expected with lag=1: 2, 2, 3, 4, 5
+
+  obs1 = manager.compute(update_history=True)  # value=2, lag constrained, returns 2
+  assert isinstance(obs1["actor"], torch.Tensor)
+  assert torch.allclose(obs1["actor"][0], torch.full((3,), 2.0, device=device))
+
+  obs2 = manager.compute(update_history=True)  # value=3, lag=1, returns 2
+  assert isinstance(obs2["actor"], torch.Tensor)
+  assert torch.allclose(obs2["actor"][0], torch.full((3,), 2.0, device=device))
+
+  obs3 = manager.compute(update_history=True)  # value=4, lag=1, returns 3
+  assert isinstance(obs3["actor"], torch.Tensor)
+  assert torch.allclose(obs3["actor"][0], torch.full((3,), 3.0, device=device))
+
+  obs4 = manager.compute(update_history=True)  # value=5, lag=1, returns 4
+  assert isinstance(obs4["actor"], torch.Tensor)
+  assert torch.allclose(obs4["actor"][0], torch.full((3,), 4.0, device=device))
+
+  obs5 = manager.compute(update_history=True)  # value=6, lag=1, returns 5
+  assert isinstance(obs5["actor"], torch.Tensor)
+  assert torch.allclose(obs5["actor"][0], torch.full((3,), 5.0, device=device))
+
+
+##
+# Delay with other features.
+##
+
+
+def test_delay_with_history(mock_env, simple_obs_func, device):
+  """Test that delay is applied before history."""
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs1": ObservationTermCfg(
+          func=simple_obs_func,
+          params={},
+          delay_min_lag=1,
+          delay_max_lag=1,
+          history_length=2,
+          flatten_history_dim=False,
+        ),
+      }
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+  assert manager.group_obs_dim["actor"] == (2, 3)
+
+  # First compute: value=2, delay gives 2, history [2, 2].
+  obs1 = manager.compute(update_history=False)
+  policy_obs1 = obs1["actor"]
+  assert isinstance(policy_obs1, torch.Tensor)
+  assert torch.allclose(
+    policy_obs1[0],
+    torch.stack(
+      [torch.full((3,), 2.0, device=device), torch.full((3,), 2.0, device=device)]
+    ),
+  )
+
+  # Second compute: value=3, delay gives 2 (lag=1), history updated to [2, 2].
+  obs2 = manager.compute(update_history=True)
+  policy_obs2 = obs2["actor"]
+  assert isinstance(policy_obs2, torch.Tensor)
+  assert torch.allclose(
+    policy_obs2[0],
+    torch.stack(
+      [torch.full((3,), 2.0, device=device), torch.full((3,), 2.0, device=device)]
+    ),
+  )
+
+  # Third compute: value=4, delay gives 3 (lag=1), history updated to [2, 3].
+  obs3 = manager.compute(update_history=True)
+  policy_obs3 = obs3["actor"]
+  assert isinstance(policy_obs3, torch.Tensor)
+  assert torch.allclose(
+    policy_obs3[0],
+    torch.stack(
+      [torch.full((3,), 2.0, device=device), torch.full((3,), 3.0, device=device)]
+    ),
+  )
+
+
+def test_delay_with_scale(mock_env, simple_obs_func, device):
+  """Test that scaling is applied before delay."""
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs1": ObservationTermCfg(
+          func=simple_obs_func,
+          params={},
+          scale=2.0,
+          delay_min_lag=1,
+          delay_max_lag=1,
+        ),
+      }
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+
+  # First compute: value=2, scaled to 4.
+  obs1 = manager.compute(update_history=True)
+  assert isinstance(obs1["actor"], torch.Tensor)
+  assert torch.allclose(obs1["actor"][0], torch.full((3,), 4.0, device=device))
+
+  # Second compute: value=3, scaled to 6, delay returns 4 (lag=1).
+  obs2 = manager.compute(update_history=True)
+  assert isinstance(obs2["actor"], torch.Tensor)
+  assert torch.allclose(obs2["actor"][0], torch.full((3,), 4.0, device=device))
+
+  # Third compute: value=4, scaled to 8, delay returns 6 (lag=1).
+  obs3 = manager.compute(update_history=True)
+  assert isinstance(obs3["actor"], torch.Tensor)
+  assert torch.allclose(obs3["actor"][0], torch.full((3,), 6.0, device=device))
+
+
+##
+# Reset tests.
+##
+
+
+def test_reset_clears_delay_buffer(mock_env, simple_obs_func, device):
+  """Test that reset clears delay buffer and restarts lag constraint."""
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs1": ObservationTermCfg(
+          func=simple_obs_func, params={}, delay_min_lag=2, delay_max_lag=2
+        ),
+      }
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+
+  # Build up delay buffer with several steps.
+  obs = manager.compute(update_history=True)
+  for _ in range(4):
+    obs = manager.compute(update_history=True)
+  # At this point, delay should be working (lag=2).
+  # Value should be from 2 steps ago.
+  assert isinstance(obs["actor"], torch.Tensor)
+  last_val = obs["actor"][0, 0].item()
+
+  # Reset all environments.
+  manager.reset()
+
+  # After reset, buffer and lags are cleared.
+  # Next compute should return current (lag constrained to 0).
+  obs = manager.compute(update_history=True)
+  policy_obs = obs["actor"]
+  assert isinstance(policy_obs, torch.Tensor)
+  # Value should be current, not delayed.
+  current_val = policy_obs[0, 0].item()
+  # Verify it's different from before reset (counter continues).
+  assert current_val != last_val
+
+
+def test_reset_partial_envs_with_verification(mock_env, device):
+  """Test partial reset actually resets only specified environments."""
+  # Use per-env counters to track each env independently.
+  counters = torch.zeros(4, dtype=torch.long, device=device)
+
+  def per_env_obs_func(env):
+    counters[:] += 1
+    # Return different values per env so we can track them.
+    return counters.unsqueeze(1).repeat(1, 3).float()
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs1": ObservationTermCfg(
+          func=per_env_obs_func, params={}, delay_min_lag=1, delay_max_lag=1
+        ),
+      }
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+
+  # Build up some history.
+  for _ in range(3):
+    manager.compute(update_history=True)
+
+  # Get obs for env 1 and 3 before reset.
+  obs_before = manager.compute(update_history=True)
+  assert isinstance(obs_before["actor"], torch.Tensor)
+  env1_before = obs_before["actor"][1, 0].item()
+  env3_before = obs_before["actor"][3, 0].item()
+
+  # Reset only envs 0 and 2.
+  manager.reset(env_ids=torch.tensor([0, 2], device=device))
+
+  # Next compute.
+  obs_after = manager.compute(update_history=True)
+
+  # Envs 1 and 3 should have continuous delayed values.
+  # Env 0 and 2 should have restarted (returning current due to lag constraint).
+  assert isinstance(obs_after["actor"], torch.Tensor)
+  env1_after = obs_after["actor"][1, 0].item()
+  env3_after = obs_after["actor"][3, 0].item()
+
+  # Non-reset envs should have changed (counter incremented).
+  assert env1_after != env1_before
+  assert env3_after != env3_before
+
+
+##
+# Shared vs per-env delay.
+##
+
+
+def test_shared_delay_actual_verification(mock_env, device):
+  """Verify that per_env=False gives same lag across all envs."""
+
+  # Use unique values per env to distinguish them.
+  def unique_obs_func(env):
+    # Return env index as the observation value.
+    return torch.arange(env.num_envs, device=device).unsqueeze(1).repeat(1, 3).float()
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs1": ObservationTermCfg(
+          func=unique_obs_func,
+          params={},
+          delay_min_lag=0,
+          delay_max_lag=2,
+          delay_per_env=False,  # Same lag for all envs.
+        ),
+      }
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+
+  # Run several steps and track the delay buffer.
+  delay_buffer = manager._group_obs_term_delay_buffer["actor"]["obs1"]
+
+  for _ in range(10):
+    manager.compute(update_history=True)
+
+  # Check that all envs have the same lag.
+  lags = delay_buffer.current_lags
+  assert torch.all(lags == lags[0]), f"Expected same lag for all envs, got {lags}"
+
+
+##
+# Update period tests.
+##
+
+
+def test_update_period_actual_verification(mock_env, device):
+  """Verify that update_period controls lag update frequency."""
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs1": ObservationTermCfg(
+          func=lambda env: torch.zeros((env.num_envs, 3), device=device),
+          params={},
+          delay_min_lag=0,
+          delay_max_lag=3,
+          delay_update_period=3,
+          delay_per_env_phase=False,  # All envs update together.
+        ),
+      }
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+  delay_buffer = manager._group_obs_term_delay_buffer["actor"]["obs1"]
+
+  # Track lag changes over time.
+  lag_history = []
+  for _ in range(12):
+    manager.compute(update_history=True)
+    lag_history.append(delay_buffer.current_lags[0].item())
+
+  # Lags should only change every 3 steps (at steps 0, 3, 6, 9).
+  assert len(lag_history) == 12
+
+
+##
+# Dimension verification tests.
+##
+
+
+def test_delay_preserves_dimensions(mock_env, simple_obs_func):
+  """Test that delay preserves observation dimensions."""
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs1": ObservationTermCfg(
+          func=simple_obs_func, params={}, delay_min_lag=1, delay_max_lag=3
+        ),
+      }
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+  assert manager.group_obs_dim["actor"] == (3,)
+
+  obs = manager.compute(update_history=True)
+  policy_obs = obs["actor"]
+  assert isinstance(policy_obs, torch.Tensor)
+  assert policy_obs.shape == (4, 3)
+
+
+def test_mixed_delay_and_no_delay_terms(mock_env, simple_obs_func, device):
+  """Test group with both delayed and non-delayed terms."""
+
+  counter = {"value": 0}
+
+  def obs_func2(env):
+    counter["value"] += 1
+    return torch.full((env.num_envs, 2), float(counter["value"]) * 10, device=device)
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs_with_delay": ObservationTermCfg(
+          func=simple_obs_func, params={}, delay_min_lag=1, delay_max_lag=1
+        ),
+        "obs_no_delay": ObservationTermCfg(func=obs_func2, params={}),
+      }
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+
+  # Should concatenate: 3 + 2 = 5.
+  assert manager.group_obs_dim["actor"] == (5,)
+
+  # First compute.
+  obs1 = manager.compute(update_history=True)
+  assert isinstance(obs1["actor"], torch.Tensor)
+  assert obs1["actor"].shape == (4, 5)
+  # First 3 values should be delayed obs (value 2).
+  # Last 2 values should be non-delayed obs (value 2*10=20).
+  assert torch.allclose(obs1["actor"][0, :3], torch.full((3,), 2.0, device=device))
+  assert torch.allclose(obs1["actor"][0, 3:], torch.full((2,), 20.0, device=device))
+
+  # Second compute.
+  obs2 = manager.compute(update_history=True)
+  assert isinstance(obs2["actor"], torch.Tensor)
+  # Delayed obs should be 2 (lag=1), non-delayed should be 3*10=30.
+  assert torch.allclose(obs2["actor"][0, :3], torch.full((3,), 2.0, device=device))
+  assert torch.allclose(obs2["actor"][0, 3:], torch.full((2,), 30.0, device=device))
+
+  # Third compute.
+  obs3 = manager.compute(update_history=True)
+  assert isinstance(obs3["actor"], torch.Tensor)
+  # Delayed obs should be 3 (lag=1), non-delayed should be 4*10=40.
+  assert torch.allclose(obs3["actor"][0, :3], torch.full((3,), 3.0, device=device))
+  assert torch.allclose(obs3["actor"][0, 3:], torch.full((2,), 40.0, device=device))
+
+
+##
+# Cache behavior tests.
+##
+
+
+def test_compute_without_update_returns_cached(mock_env, simple_obs_func):
+  """Test that compute(update_history=False) returns cached result.
+
+  This test verifies the fix for the double-push bug where calling compute()
+  multiple times per control step would push observations twice to the delay
+  buffer, effectively halving the actual delay.
+  """
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs1": ObservationTermCfg(
+          func=simple_obs_func, params={}, delay_min_lag=1, delay_max_lag=1
+        ),
+      }
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+
+  # First compute with update_history=True should compute fresh and push to buffer.
+  obs1 = manager.compute(update_history=True)
+  assert isinstance(obs1["actor"], torch.Tensor)
+  val1 = obs1["actor"][0, 0].item()
+
+  # Second compute with update_history=False should return cached result.
+  obs2 = manager.compute(update_history=False)
+  assert isinstance(obs2["actor"], torch.Tensor)
+  val2 = obs2["actor"][0, 0].item()
+  assert val1 == val2, "compute(update_history=False) should return cached result"
+
+  # Multiple calls without update should all return the same cached value.
+  for _ in range(5):
+    obs = manager.compute(update_history=False)
+    assert isinstance(obs["actor"], torch.Tensor)
+    assert obs["actor"][0, 0].item() == val1
+
+  # Third compute with update_history=True should compute fresh observations.
+  obs3 = manager.compute(update_history=True)
+  assert isinstance(obs3["actor"], torch.Tensor)
+  val3 = obs3["actor"][0, 0].item()
+  # Value should have advanced (previous value due to lag=1).
+  assert val3 == val1  # lag=1 returns previous observation
+
+
+def test_delay_buffer_not_double_pushed(mock_env, simple_obs_func):
+  """Test that delay buffer is not pushed twice per control step.
+
+  This tests the specific bug where calling compute() twice per control step
+  (e.g., in step() and get_observations()) would push observations twice,
+  effectively halving the actual delay.
+  """
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs1": ObservationTermCfg(
+          func=simple_obs_func, params={}, delay_min_lag=2, delay_max_lag=2
+        ),
+      }
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+  delay_buffer = manager._group_obs_term_delay_buffer["actor"]["obs1"]
+
+  # Simulate the training loop pattern:
+  # 1. compute(update_history=True) in step()
+  # 2. compute(update_history=False) in get_observations()
+
+  # First "step": compute with update, buffer should have 1 entry.
+  manager.compute(update_history=True)
+  assert delay_buffer._buffer.current_length[0].item() == 1
+
+  # Call compute without update (like get_observations does).
+  manager.compute(update_history=False)
+  # Buffer should STILL have 1 entry (not 2).
+  assert delay_buffer._buffer.current_length[0].item() == 1
+
+  # Second "step": compute with update, buffer should have 2 entries.
+  manager.compute(update_history=True)
+  assert delay_buffer._buffer.current_length[0].item() == 2
+
+  # Call compute without update multiple times.
+  for _ in range(3):
+    manager.compute(update_history=False)
+  # Buffer should STILL have 2 entries.
+  assert delay_buffer._buffer.current_length[0].item() == 2
+
+  # Third "step": compute with update, buffer should have 3 entries (max for lag=2).
+  manager.compute(update_history=True)
+  assert delay_buffer._buffer.current_length[0].item() == 3
+
+
+def test_cache_invalidated_on_reset(mock_env, simple_obs_func):
+  """Test that observation cache is invalidated when environments are reset."""
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs1": ObservationTermCfg(
+          func=simple_obs_func, params={}, delay_min_lag=1, delay_max_lag=1
+        ),
+      }
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+
+  # Build up some state.
+  for _ in range(3):
+    manager.compute(update_history=True)
+
+  # Get cached observation.
+  obs_before = manager.compute(update_history=False)
+  assert isinstance(obs_before["actor"], torch.Tensor)
+  val_before = obs_before["actor"][0, 0].item()
+
+  # Reset.
+  manager.reset()
+
+  # Next compute should return fresh observations (cache invalidated).
+  obs_after = manager.compute(update_history=True)
+  assert isinstance(obs_after["actor"], torch.Tensor)
+  val_after = obs_after["actor"][0, 0].item()
+
+  # After reset, delay buffer is cleared, so we get the current (fresh) observation.
+  # The value should have advanced since the obs_func counter continues incrementing.
+  assert val_after != val_before
@@ -0,0 +1,612 @@
+"""Tests for observation history functionality."""
+
+from unittest.mock import Mock
+
+import pytest
+import torch
+from conftest import get_test_device
+
+from mjlab.managers.observation_manager import (
+  ObservationGroupCfg,
+  ObservationManager,
+  ObservationTermCfg,
+)
+
+
+@pytest.fixture
+def device():
+  """Test device fixture."""
+  return get_test_device()
+
+
+@pytest.fixture
+def mock_env(device):
+  """Create a mock environment."""
+  env = Mock()
+  env.num_envs = 4
+  env.device = device
+  env.step_dt = 0.02
+  return env
+
+
+@pytest.fixture
+def simple_obs_func(device):
+  """Create a simple observation function that returns a counter."""
+  counter = {"value": 0}
+
+  def obs_func(env):
+    counter["value"] += 1
+    return torch.full((env.num_envs, 3), float(counter["value"]), device=device)
+
+  return obs_func
+
+
+# Basic observation history tests.
+
+
+def test_no_history_by_default(mock_env, simple_obs_func):
+  """Test that observations work without history (default behavior)."""
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs1": ObservationTermCfg(func=simple_obs_func, params={}),
+      }
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+  assert manager.group_obs_dim["actor"] == (3,)
+
+  obs = manager.compute()
+  policy_obs = obs["actor"]
+  assert isinstance(policy_obs, torch.Tensor)
+  assert policy_obs.shape == (4, 3)
+
+
+def test_single_step_history(mock_env, simple_obs_func):
+  """Test observation with history_length=1."""
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs1": ObservationTermCfg(func=simple_obs_func, params={}, history_length=1),
+      }
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+  assert manager.group_obs_dim["actor"] == (3,)
+
+  obs = manager.compute()
+  policy_obs = obs["actor"]
+  assert isinstance(policy_obs, torch.Tensor)
+  assert policy_obs.shape == (4, 3)
+
+
+def test_multi_step_history_flattened(mock_env, simple_obs_func):
+  """Test observation with history_length=3 and flattened."""
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs1": ObservationTermCfg(
+          func=simple_obs_func, params={}, history_length=3, flatten_history_dim=True
+        ),
+      }
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+  assert manager.group_obs_dim["actor"] == (9,)
+
+  obs = manager.compute()
+  policy_obs = obs["actor"]
+  assert isinstance(policy_obs, torch.Tensor)
+  assert policy_obs.shape == (4, 9)
+
+
+def test_multi_step_history_not_flattened(mock_env, simple_obs_func):
+  """Test observation with history_length=3 and not flattened."""
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs1": ObservationTermCfg(
+          func=simple_obs_func, params={}, history_length=3, flatten_history_dim=False
+        ),
+      }
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+  assert manager.group_obs_dim["actor"] == (3, 3)
+
+  obs = manager.compute()
+  policy_obs = obs["actor"]
+  assert isinstance(policy_obs, torch.Tensor)
+  assert policy_obs.shape == (4, 3, 3)
+
+
+# History accumulation tests.
+
+
+def test_history_accumulates_correctly(mock_env, simple_obs_func):
+  """Test that history buffer accumulates observations in correct order."""
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs1": ObservationTermCfg(
+          func=simple_obs_func, params={}, history_length=3, flatten_history_dim=False
+        ),
+      }
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+  device = mock_env.device
+  # Note: counter is incremented during _prepare_terms (value=1).
+
+  # First compute uses value=2 and initializes buffer.
+  obs = manager.compute(update_history=False)
+  policy_obs = obs["actor"]
+  assert isinstance(policy_obs, torch.Tensor)
+  assert torch.allclose(policy_obs[0], torch.full((3, 3), 2.0, device=device))
+
+  # Update with value=3.
+  obs = manager.compute(update_history=True)
+  policy_obs = obs["actor"]
+  assert isinstance(policy_obs, torch.Tensor)
+  # History: [2, 2, 3] (oldest to newest).
+  expected = torch.stack(
+    [
+      torch.full((3,), 2.0, device=device),
+      torch.full((3,), 2.0, device=device),
+      torch.full((3,), 3.0, device=device),
+    ]
+  )
+  assert torch.allclose(policy_obs[0], expected)
+
+  # Update with value=4.
+  obs = manager.compute(update_history=True)
+  policy_obs = obs["actor"]
+  assert isinstance(policy_obs, torch.Tensor)
+  # History: [2, 3, 4].
+  expected = torch.stack(
+    [
+      torch.full((3,), 2.0, device=device),
+      torch.full((3,), 3.0, device=device),
+      torch.full((3,), 4.0, device=device),
+    ]
+  )
+  assert torch.allclose(policy_obs[0], expected)
+
+  # Update with value=5, circular overwrite.
+  obs = manager.compute(update_history=True)
+  policy_obs = obs["actor"]
+  assert isinstance(policy_obs, torch.Tensor)
+  # History: [3, 4, 5].
+  expected = torch.stack(
+    [
+      torch.full((3,), 3.0, device=device),
+      torch.full((3,), 4.0, device=device),
+      torch.full((3,), 5.0, device=device),
+    ]
+  )
+  assert torch.allclose(policy_obs[0], expected)
+
+
+def test_update_history_false_doesnt_modify_buffer(mock_env, simple_obs_func):
+  """Test that update_history=False doesn't modify the buffer."""
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs1": ObservationTermCfg(
+          func=simple_obs_func, params={}, history_length=2, flatten_history_dim=False
+        ),
+      }
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+
+  # Initialize (value=1).
+  obs1 = manager.compute(update_history=False)
+  policy_obs1 = obs1["actor"]
+  assert isinstance(policy_obs1, torch.Tensor)
+
+  # Call without update (value=2, but buffer unchanged).
+  obs2 = manager.compute(update_history=False)
+  policy_obs2 = obs2["actor"]
+  assert isinstance(policy_obs2, torch.Tensor)
+
+  # History should still be [1, 1].
+  assert torch.allclose(policy_obs1, policy_obs2)
+
+
+# Group-level history override tests.
+
+
+def test_group_history_overrides_term(mock_env, simple_obs_func):
+  """Test group history_length overrides term history_length."""
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      history_length=5,  # Group level.
+      flatten_history_dim=False,
+      terms={
+        "obs1": ObservationTermCfg(
+          func=simple_obs_func,
+          params={},
+          history_length=2,  # Overridden.
+        ),
+      },
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+  assert manager.group_obs_dim["actor"] == (5, 3)
+
+  obs = manager.compute()
+  policy_obs = obs["actor"]
+  assert isinstance(policy_obs, torch.Tensor)
+  assert policy_obs.shape == (4, 5, 3)
+
+
+# History reset tests.
+
+
+def test_reset_clears_all_envs(mock_env, simple_obs_func):
+  """Test that reset without env_ids clears all environments."""
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs1": ObservationTermCfg(
+          func=simple_obs_func, params={}, history_length=2, flatten_history_dim=False
+        ),
+      }
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+
+  # Build up history.
+  manager.compute(update_history=True)
+  manager.compute(update_history=True)
+
+  # Reset all.
+  manager.reset()
+
+  # Buffer should be zeroed.
+  obs = manager.compute(update_history=False)
+  policy_obs = obs["actor"]
+  assert isinstance(policy_obs, torch.Tensor)
+  assert torch.allclose(policy_obs, torch.zeros((4, 2, 3), device=mock_env.device))
+
+
+def test_reset_partial_envs(mock_env, simple_obs_func):
+  """Test that reset with specific env_ids only resets those envs."""
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs1": ObservationTermCfg(
+          func=simple_obs_func, params={}, history_length=3, flatten_history_dim=False
+        ),
+      }
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+
+  # Build up history.
+  manager.compute(update_history=True)
+  manager.compute(update_history=True)
+  obs_before = manager.compute(update_history=True)
+  policy_obs_before = obs_before["actor"]
+  assert isinstance(policy_obs_before, torch.Tensor)
+
+  # Reset only envs 0 and 2.
+  manager.reset(env_ids=torch.tensor([0, 2], device=mock_env.device))
+
+  obs_after = manager.compute(update_history=False)
+  policy_obs_after = obs_after["actor"]
+  assert isinstance(policy_obs_after, torch.Tensor)
+
+  # Envs 0 and 2 should be reset (zeros), 1 and 3 unchanged.
+  assert torch.allclose(
+    policy_obs_after[0], torch.zeros((3, 3), device=mock_env.device)
+  )
+  assert torch.allclose(policy_obs_after[1], policy_obs_before[1])
+  assert torch.allclose(
+    policy_obs_after[2], torch.zeros((3, 3), device=mock_env.device)
+  )
+  assert torch.allclose(policy_obs_after[3], policy_obs_before[3])
+
+
+def test_reset_partial_envs_with_backfill(mock_env, simple_obs_func):
+  """Test that reset envs get backfilled on next update."""
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs1": ObservationTermCfg(
+          func=simple_obs_func, params={}, history_length=3, flatten_history_dim=False
+        ),
+      }
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+  device = mock_env.device
+
+  # Build history: [2, 3, 4] for all envs (counter starts at 1 after _prepare_terms).
+  manager.compute(update_history=True)  # value=2
+  manager.compute(update_history=True)  # value=3
+  manager.compute(update_history=True)  # value=4
+
+  # Reset only envs 0 and 2.
+  manager.reset(env_ids=torch.tensor([0, 2], device=device))
+
+  # Next update with value=5.
+  obs = manager.compute(update_history=True)
+  policy_obs = obs["actor"]
+  assert isinstance(policy_obs, torch.Tensor)
+
+  # Env 0: [5, 5, 5] (backfilled after reset).
+  expected_env0 = torch.stack(
+    [
+      torch.full((3,), 5.0, device=device),
+      torch.full((3,), 5.0, device=device),
+      torch.full((3,), 5.0, device=device),
+    ]
+  )
+  assert torch.allclose(policy_obs[0], expected_env0)
+
+  # Env 1: [3, 4, 5] (continues normally).
+  expected_env1 = torch.stack(
+    [
+      torch.full((3,), 3.0, device=device),
+      torch.full((3,), 4.0, device=device),
+      torch.full((3,), 5.0, device=device),
+    ]
+  )
+  assert torch.allclose(policy_obs[1], expected_env1)
+
+  # Env 2: [5, 5, 5] (backfilled after reset).
+  assert torch.allclose(policy_obs[2], expected_env0)
+
+  # Env 3: [3, 4, 5] (continues normally).
+  assert torch.allclose(policy_obs[3], expected_env1)
+
+
+# History with other features tests.
+
+
+def test_history_with_clip(mock_env, simple_obs_func):
+  """Test that clipping is applied before history."""
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs1": ObservationTermCfg(
+          func=simple_obs_func,
+          params={},
+          history_length=2,
+          flatten_history_dim=False,
+          clip=(-0.5, 0.5),
+        ),
+      }
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+
+  obs = manager.compute(update_history=True)
+  policy_obs = obs["actor"]
+  assert isinstance(policy_obs, torch.Tensor)
+  # Values should be clipped.
+  assert torch.all(policy_obs >= -0.5)
+  assert torch.all(policy_obs <= 0.5)
+
+
+def test_history_with_scale(mock_env, simple_obs_func):
+  """Test that scaling is applied before history."""
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs1": ObservationTermCfg(
+          func=simple_obs_func,
+          params={},
+          history_length=2,
+          flatten_history_dim=False,
+          scale=2.0,
+        ),
+      }
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+  device = mock_env.device
+  # Counter at 1 after _prepare_terms.
+
+  # First call uses value=2, scaled to 4, initializes buffer.
+  manager.compute(update_history=False)
+  # Second call uses value=3, scaled to 6, and updates history.
+  obs = manager.compute(update_history=True)
+  policy_obs = obs["actor"]
+  assert isinstance(policy_obs, torch.Tensor)
+
+  # History: [4, 6].
+  expected = torch.stack(
+    [torch.full((3,), 4.0, device=device), torch.full((3,), 6.0, device=device)]
+  )
+  assert torch.allclose(policy_obs[0], expected)
+
+
+# Mixed history terms tests.
+
+
+def test_mixed_terms_concatenated(mock_env, simple_obs_func, device):
+  """Test group with both history and non-history terms concatenated."""
+
+  counter = {"value": 0}
+
+  def obs_func2(env):
+    counter["value"] += 1
+    return torch.full((env.num_envs, 2), float(counter["value"]) * 10, device=device)
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs_with_history": ObservationTermCfg(
+          func=simple_obs_func, params={}, history_length=2, flatten_history_dim=True
+        ),
+        "obs_no_history": ObservationTermCfg(func=obs_func2, params={}),
+      }
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+
+  # Should concatenate: (3*2) + 2 = 8.
+  assert manager.group_obs_dim["actor"] == (8,)
+
+  obs = manager.compute()
+  policy_obs = obs["actor"]
+  assert isinstance(policy_obs, torch.Tensor)
+  assert policy_obs.shape == (4, 8)
+
+
+def test_no_double_append_on_first_call(mock_env, simple_obs_func):
+  """Test that first call with update_history=True only appends once, not twice."""
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs1": ObservationTermCfg(
+          func=simple_obs_func, params={}, history_length=3, flatten_history_dim=False
+        ),
+      }
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+  device = mock_env.device
+  # Counter is at 1 after _prepare_terms.
+
+  # First call with update_history=True (value=2).
+  # This should initialize the buffer AND append once (not twice).
+  obs = manager.compute(update_history=True)
+  policy_obs = obs["actor"]
+  assert isinstance(policy_obs, torch.Tensor)
+
+  # Verify buffer was initialized and backfilled correctly.
+  # All slots should be filled with value 2.
+  expected_first = torch.stack(
+    [
+      torch.full((3,), 2.0, device=device),
+      torch.full((3,), 2.0, device=device),
+      torch.full((3,), 2.0, device=device),
+    ]
+  )
+  assert torch.allclose(policy_obs[0], expected_first)
+
+  # Get the circular buffer and check pointer position.
+  circular_buffer = manager._group_obs_term_history_buffer["actor"]["obs1"]
+  # After one append, pointer should be at 0 (not 1 which would indicate double-append).
+  assert circular_buffer._pointer == 0
+  # And we should have exactly 1 push recorded.
+  assert torch.all(circular_buffer._num_pushes == 1)
+
+  # Second call with update_history=True (value=3).
+  obs = manager.compute(update_history=True)
+  policy_obs = obs["actor"]
+  assert isinstance(policy_obs, torch.Tensor)
+
+  # History should be [2, 2, 3] (oldest to newest).
+  expected_second = torch.stack(
+    [
+      torch.full((3,), 2.0, device=device),
+      torch.full((3,), 2.0, device=device),
+      torch.full((3,), 3.0, device=device),
+    ]
+  )
+  assert torch.allclose(policy_obs[0], expected_second)
+
+  # Pointer should now be at 1.
+  assert circular_buffer._pointer == 1
+  # And we should have exactly 2 pushes.
+  assert torch.all(circular_buffer._num_pushes == 2)
+
+
+def test_term_major_ordering(mock_env, device):
+  """Test that history flattening uses term-major ordering.
+
+  Term-major: [A_t0, A_t1, ..., A_tH-1, B_t0, B_t1, ..., B_tH-1]
+  Time-major: [A_t0, B_t0, ..., A_t1, B_t1, ..., A_tH-1, B_tH-1]
+
+  This test verifies mjlab uses term-major ordering.
+  """
+
+  # Create observation functions with distinct values per term.
+  def obs_A(env):
+    # Returns [100, 101] for easy identification.
+    return torch.tensor([[100.0, 101.0]] * env.num_envs, device=device)
+
+  def obs_B(env):
+    # Returns [200, 201, 202] for easy identification.
+    return torch.tensor([[200.0, 201.0, 202.0]] * env.num_envs, device=device)
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      concatenate_terms=True,
+      terms={
+        "term_A": ObservationTermCfg(
+          func=obs_A, params={}, history_length=3, flatten_history_dim=True
+        ),
+        "term_B": ObservationTermCfg(
+          func=obs_B, params={}, history_length=3, flatten_history_dim=True
+        ),
+      },
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+
+  # Compute observations (history will be backfilled with same values).
+  obs = manager.compute(update_history=False)
+  policy_obs = obs["actor"]
+  assert isinstance(policy_obs, torch.Tensor)
+
+  # Expected shape: (4 envs, 2*3 + 3*3) = (4, 15).
+  assert policy_obs.shape == (4, 15)
+
+  # Check ordering for first environment.
+  result = policy_obs[0].cpu().tolist()
+
+  # Term-major: all A history, then all B history.
+  # A has 2 dims, B has 3 dims, history_length=3.
+  # fmt: off
+  expected_term_major = [
+    100.0, 101.0, 100.0, 101.0, 100.0, 101.0,  # All A history
+    200.0, 201.0, 202.0, 200.0, 201.0, 202.0, 200.0, 201.0, 202.0,  # All B
+  ]
+  # fmt: on
+
+  # Time-major would interleave: [A_t0, B_t0, A_t1, B_t1, A_t2, B_t2].
+  # fmt: off
+  expected_time_major = [
+    100.0, 101.0, 200.0, 201.0, 202.0,  # Frame t0
+    100.0, 101.0, 200.0, 201.0, 202.0,  # Frame t1
+    100.0, 101.0, 200.0, 201.0, 202.0,  # Frame t2
+  ]
+  # fmt: on
+
+  # Verify term-major ordering.
+  assert result == expected_term_major, f"Expected term-major ordering, got {result}"
+  assert result != expected_time_major, "Should not match time-major"
@@ -0,0 +1,63 @@
+"""Tests for ObservationManager behavior."""
+
+from unittest.mock import Mock
+
+import pytest
+import torch
+from conftest import get_test_device
+
+from mjlab.managers.observation_manager import (
+  ObservationGroupCfg,
+  ObservationManager,
+  ObservationTermCfg,
+)
+
+
+@pytest.fixture
+def mock_env():
+  env = Mock()
+  env.num_envs = 4
+  env.device = get_test_device()
+  env.step_dt = 0.02
+  return env
+
+
+def _dummy_term(mock_env):
+  def func(_env, **_kwargs):
+    return torch.zeros(mock_env.num_envs, 3, device=mock_env.device)
+
+  return ObservationTermCfg(func=func)
+
+
+def test_empty_terms_dict_skipped(mock_env):
+  """A group declared with no terms is skipped rather than raising."""
+  cfg = {"actor": ObservationGroupCfg(terms={})}
+  mgr = ObservationManager(cfg, mock_env)
+
+  assert "actor" not in mgr.active_terms
+  assert "actor" not in mgr.group_obs_dim
+
+
+def test_all_terms_none_skipped(mock_env):
+  """A group whose every term is None is skipped rather than raising."""
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={"a": None, "b": None},  # type: ignore[dict-item]
+    ),
+  }
+  mgr = ObservationManager(cfg, mock_env)
+
+  assert "actor" not in mgr.active_terms
+  assert "actor" not in mgr.group_obs_dim
+
+
+def test_empty_group_skipped_alongside_active_group(mock_env):
+  """Active groups coexist with empty ones; only the empty group is dropped."""
+  cfg = {
+    "actor": ObservationGroupCfg(terms={"a": _dummy_term(mock_env)}),
+    "critic": ObservationGroupCfg(terms={}),
+  }
+  mgr = ObservationManager(cfg, mock_env)
+
+  assert "actor" in mgr.active_terms
+  assert "critic" not in mgr.active_terms
@@ -0,0 +1,428 @@
+"""Tests for observation NaN handling functionality."""
+
+from unittest.mock import Mock
+
+import pytest
+import torch
+from conftest import get_test_device
+
+from mjlab.managers.observation_manager import (
+  ObservationGroupCfg,
+  ObservationManager,
+  ObservationTermCfg,
+)
+
+
+@pytest.fixture
+def device():
+  """Test device fixture."""
+  return get_test_device()
+
+
+@pytest.fixture
+def mock_env(device):
+  """Create a mock environment."""
+  env = Mock()
+  env.num_envs = 4
+  env.device = device
+  env.step_dt = 0.02
+  return env
+
+
+def test_nan_disabled_policy_no_check(mock_env, device):
+  """Policy='disabled' should pass through NaNs unchanged."""
+
+  def obs_with_nan(env):
+    obs = torch.ones((env.num_envs, 3), device=device)
+    obs[1, 1] = float("nan")
+    return obs
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={"obs1": ObservationTermCfg(func=obs_with_nan, params={})},
+      nan_policy="disabled",
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+  obs = manager.compute()
+
+  policy_obs = obs["actor"]
+  assert isinstance(policy_obs, torch.Tensor)
+  assert torch.isnan(policy_obs[1, 1])  # NaN should pass through.
+
+
+def test_nan_sanitize_policy_silently_fixes(mock_env, device):
+  """Policy='sanitize' should zero NaNs without warning."""
+
+  def obs_with_nan(env):
+    obs = torch.ones((env.num_envs, 3), device=device)
+    obs[1, 1] = float("nan")
+    obs[2, 0] = float("inf")
+    return obs
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={"obs1": ObservationTermCfg(func=obs_with_nan, params={})},
+      nan_policy="sanitize",
+      nan_check_per_term=True,
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+  obs = manager.compute()
+
+  policy_obs = obs["actor"]
+  assert isinstance(policy_obs, torch.Tensor)
+  assert not torch.isnan(policy_obs).any()
+  assert not torch.isinf(policy_obs).any()
+  assert policy_obs[1, 1] == 0.0
+  assert policy_obs[2, 0] == 0.0
+
+
+def test_nan_warn_policy_logs_and_sanitizes(mock_env, device, capsys):
+  """Policy='warn' should log warning and sanitize."""
+
+  def obs_with_nan(env):
+    obs = torch.ones((env.num_envs, 3), device=device)
+    obs[1, 1] = float("nan")
+    return obs
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={"obs1": ObservationTermCfg(func=obs_with_nan, params={})},
+      nan_policy="warn",
+      nan_check_per_term=True,
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+  obs = manager.compute()
+
+  captured = capsys.readouterr()
+  assert "ObservationManager" in captured.out
+  assert "NaN/Inf" in captured.out
+  assert "actor/obs1" in captured.out
+  assert "Sanitizing" in captured.out
+
+  policy_obs = obs["actor"]
+  assert isinstance(policy_obs, torch.Tensor)
+  assert not torch.isnan(policy_obs).any()
+  assert policy_obs[1, 1] == 0.0
+
+
+def test_nan_error_policy_raises(mock_env, device):
+  """Policy='error' should raise ValueError with context."""
+
+  def obs_with_nan(env):
+    obs = torch.ones((env.num_envs, 3), device=device)
+    obs[1, 1] = float("nan")
+    return obs
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={"obs1": ObservationTermCfg(func=obs_with_nan, params={})},
+      nan_policy="error",
+      nan_check_per_term=True,
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+
+  with pytest.raises(ValueError) as excinfo:
+    manager.compute()
+
+  assert "NaN/Inf detected" in str(excinfo.value)
+  assert "actor/obs1" in str(excinfo.value)
+  assert "1" in str(excinfo.value)  # Environment ID.
+
+
+def test_nan_check_per_term_identifies_source(mock_env, device, capsys):
+  """Per-term checking should identify which term has NaN."""
+
+  def clean_obs(env):
+    return torch.ones((env.num_envs, 3), device=device)
+
+  def obs_with_nan(env):
+    obs = torch.ones((env.num_envs, 3), device=device)
+    obs[1, 1] = float("nan")
+    return obs
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "clean": ObservationTermCfg(func=clean_obs, params={}),
+        "problematic": ObservationTermCfg(func=obs_with_nan, params={}),
+      },
+      nan_policy="warn",
+      nan_check_per_term=True,
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+  obs = manager.compute()
+
+  # Check that only problematic term is logged.
+  captured = capsys.readouterr()
+  assert "actor/problematic" in captured.out
+  assert "actor/clean" not in captured.out
+
+  # Both terms should be sanitized in output.
+  policy_obs = obs["actor"]
+  assert isinstance(policy_obs, torch.Tensor)
+  assert not torch.isnan(policy_obs).any()
+
+
+def test_nan_check_final_only(mock_env, device, capsys):
+  """Final-only checking should check concatenated result."""
+
+  def obs_with_nan(env):
+    obs = torch.ones((env.num_envs, 3), device=device)
+    obs[1, 1] = float("nan")
+    return obs
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={"obs1": ObservationTermCfg(func=obs_with_nan, params={})},
+      nan_policy="warn",
+      nan_check_per_term=False,  # Check only final result.
+      concatenate_terms=True,
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+  obs = manager.compute()
+
+  # Should log group name, not term name.
+  captured = capsys.readouterr()
+  assert "actor" in captured.out
+  assert "actor/obs1" not in captured.out  # Should not show term.
+
+  policy_obs = obs["actor"]
+  assert isinstance(policy_obs, torch.Tensor)
+  assert not torch.isnan(policy_obs).any()
+
+
+def test_nan_check_dict_output(mock_env, device):
+  """NaN checking should work with non-concatenated dict output."""
+
+  def obs_with_nan(env):
+    obs = torch.ones((env.num_envs, 3), device=device)
+    obs[1, 1] = float("nan")
+    return obs
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={"obs1": ObservationTermCfg(func=obs_with_nan, params={})},
+      nan_policy="sanitize",
+      concatenate_terms=False,
+      nan_check_per_term=False,
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+  obs = manager.compute()
+
+  policy: dict[str, torch.Tensor] = obs["actor"]  # type: ignore[assignment]
+  policy_obs = policy["obs1"]
+  assert not torch.isnan(policy_obs).any()
+  assert policy_obs[1, 1] == 0.0
+
+
+def test_nan_before_delay_buffer(mock_env, device):
+  """NaN should be sanitized before entering delay buffer."""
+
+  call_count = [0]
+
+  def obs_with_nan(env):
+    obs = torch.ones((env.num_envs, 3), device=device)
+    # Only inject NaN on first call.
+    if call_count[0] == 0:
+      obs[1, 1] = float("nan")
+    call_count[0] += 1
+    return obs
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs1": ObservationTermCfg(
+          func=obs_with_nan,
+          params={},
+          delay_min_lag=1,
+          delay_max_lag=1,  # 1-step delay.
+        )
+      },
+      nan_policy="sanitize",
+      nan_check_per_term=True,
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+
+  # First compute: NaN should be sanitized before entering delay buffer.
+  obs1 = manager.compute(update_history=True)
+  # First output comes from buffer initialization (zeros).
+  obs1_policy = obs1["actor"]
+  assert isinstance(obs1_policy, torch.Tensor)
+  assert not torch.isnan(obs1_policy).any()
+
+  # Second compute: should get sanitized value from buffer.
+  obs2 = manager.compute(update_history=True)
+  obs2_policy = obs2["actor"]
+  assert isinstance(obs2_policy, torch.Tensor)
+  assert not torch.isnan(obs2_policy).any()
+
+
+def test_nan_before_history_buffer(mock_env, device):
+  """NaN should be sanitized before entering history buffer."""
+
+  call_count = [0]
+
+  def obs_with_nan(env):
+    obs = torch.ones((env.num_envs, 3), device=device)
+    # Inject NaN on first call.
+    if call_count[0] == 0:
+      obs[1, 1] = float("nan")
+    call_count[0] += 1
+    return obs
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs1": ObservationTermCfg(
+          func=obs_with_nan, params={}, history_length=3, flatten_history_dim=True
+        )
+      },
+      nan_policy="sanitize",
+      nan_check_per_term=True,
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+
+  # First compute: NaN should be sanitized.
+  obs1 = manager.compute(update_history=True)
+  obs1_policy = obs1["actor"]
+  assert isinstance(obs1_policy, torch.Tensor)
+  assert not torch.isnan(obs1_policy).any()
+
+  # Subsequent computes should also have no NaN from history.
+  for _ in range(3):
+    obs = manager.compute(update_history=True)
+    obs_policy = obs["actor"]
+    assert isinstance(obs_policy, torch.Tensor)
+    assert not torch.isnan(obs_policy).any()
+
+
+def test_nan_handling_with_multiple_groups(mock_env, device):
+  """NaN handling should work independently for multiple groups."""
+
+  def obs_with_nan(env):
+    obs = torch.ones((env.num_envs, 3), device=device)
+    obs[1, 1] = float("nan")
+    return obs
+
+  def clean_obs(env):
+    return torch.ones((env.num_envs, 3), device=device)
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={"obs1": ObservationTermCfg(func=obs_with_nan, params={})},
+      nan_policy="sanitize",
+    ),
+    "critic": ObservationGroupCfg(
+      terms={"obs1": ObservationTermCfg(func=clean_obs, params={})},
+      nan_policy="disabled",
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+  obs = manager.compute()
+
+  policy_obs = obs["actor"]
+  assert isinstance(policy_obs, torch.Tensor)
+  assert not torch.isnan(policy_obs).any()
+
+  critic_obs = obs["critic"]
+  assert isinstance(critic_obs, torch.Tensor)
+  assert not torch.isnan(critic_obs).any()
+
+
+def test_nan_handling_with_scaling(mock_env, device):
+  """NaN handling should work correctly with scaling applied."""
+
+  def obs_with_nan(env):
+    obs = torch.full((env.num_envs, 3), 2.0, device=device)
+    obs[1, 1] = float("nan")
+    return obs
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={"obs1": ObservationTermCfg(func=obs_with_nan, params={}, scale=0.5)},
+      nan_policy="sanitize",
+      nan_check_per_term=True,
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+  obs = manager.compute()
+
+  policy_obs = obs["actor"]
+  assert isinstance(policy_obs, torch.Tensor)
+  # NaN should be sanitized to 0, other values scaled: 2.0 * 0.5 = 1.0.
+  assert not torch.isnan(policy_obs).any()
+  assert policy_obs[1, 1] == 0.0
+  assert policy_obs[0, 0] == 1.0  # Scaled value.
+
+
+def test_nan_handling_with_clipping(mock_env, device):
+  """NaN handling should work correctly with clipping applied."""
+
+  def obs_with_nan(env):
+    obs = torch.full((env.num_envs, 3), 10.0, device=device)
+    obs[1, 1] = float("nan")
+    return obs
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs1": ObservationTermCfg(func=obs_with_nan, params={}, clip=(-1.0, 1.0))
+      },
+      nan_policy="sanitize",
+      nan_check_per_term=True,
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+  obs = manager.compute()
+
+  policy_obs = obs["actor"]
+  assert isinstance(policy_obs, torch.Tensor)
+  # NaN should be sanitized to 0, other values clipped to 1.0.
+  assert not torch.isnan(policy_obs).any()
+  assert policy_obs[1, 1] == 0.0
+  assert policy_obs[0, 0] == 1.0  # Clipped value.
+
+
+def test_negative_inf_handling(mock_env, device):
+  """Test that negative infinity is also handled."""
+
+  def obs_with_neginf(env):
+    obs = torch.ones((env.num_envs, 3), device=device)
+    obs[1, 1] = float("-inf")
+    return obs
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={"obs1": ObservationTermCfg(func=obs_with_neginf, params={})},
+      nan_policy="sanitize",
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+  obs = manager.compute()
+
+  policy_obs = obs["actor"]
+  assert isinstance(policy_obs, torch.Tensor)
+  assert not torch.isinf(policy_obs).any()
+  assert policy_obs[1, 1] == 0.0
@@ -0,0 +1,315 @@
+"""Tests for observation noise functionality."""
+
+from unittest.mock import Mock
+
+import pytest
+import torch
+from conftest import get_test_device
+
+from mjlab.managers.observation_manager import (
+  ObservationGroupCfg,
+  ObservationManager,
+  ObservationTermCfg,
+)
+from mjlab.utils.noise.noise_cfg import (
+  ConstantNoiseCfg,
+  NoiseModelWithAdditiveBiasCfg,
+)
+
+
+@pytest.fixture
+def device():
+  """Test device fixture."""
+  return get_test_device()
+
+
+@pytest.fixture
+def mock_env(device):
+  """Create a mock environment."""
+  env = Mock()
+  env.num_envs = 4
+  env.device = device
+  env.step_dt = 0.02
+  return env
+
+
+def test_noise_applied_when_corruption_enabled(mock_env, device):
+  """Test that noise is applied when enable_corruption=True."""
+
+  def obs_func(env):
+    return torch.ones((env.num_envs, 3), device=device)
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs1": ObservationTermCfg(
+          func=obs_func,
+          params={},
+          noise=ConstantNoiseCfg(bias=0.5, operation="add"),
+        ),
+      },
+      enable_corruption=True,
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+  obs = manager.compute()
+
+  policy_obs = obs["actor"]
+  assert isinstance(policy_obs, torch.Tensor)
+  # Expect 1.0 + 0.5 = 1.5
+  expected = torch.full((4, 3), 1.5, device=device)
+  assert torch.allclose(policy_obs, expected)
+
+
+def test_noise_not_applied_when_corruption_disabled(mock_env, device):
+  """Test that noise is NOT applied when enable_corruption=False."""
+
+  def obs_func(env):
+    return torch.ones((env.num_envs, 3), device=device)
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs1": ObservationTermCfg(
+          func=obs_func,
+          params={},
+          noise=ConstantNoiseCfg(bias=0.5, operation="add"),
+        ),
+      },
+      enable_corruption=False,
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+  obs = manager.compute()
+
+  policy_obs = obs["actor"]
+  assert isinstance(policy_obs, torch.Tensor)
+  # Noise should NOT be applied, expect original value of 1.0
+  expected = torch.full((4, 3), 1.0, device=device)
+  assert torch.allclose(policy_obs, expected)
+
+
+def test_noise_add_operation(mock_env, device):
+  """Test noise with 'add' operation."""
+
+  def obs_func(env):
+    return torch.full((env.num_envs, 3), 2.0, device=device)
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs1": ObservationTermCfg(
+          func=obs_func,
+          params={},
+          noise=ConstantNoiseCfg(bias=0.3, operation="add"),
+        ),
+      },
+      enable_corruption=True,
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+  obs = manager.compute()
+
+  policy_obs = obs["actor"]
+  assert isinstance(policy_obs, torch.Tensor)
+  # Expect 2.0 + 0.3 = 2.3
+  expected = torch.full((4, 3), 2.3, device=device)
+  assert torch.allclose(policy_obs, expected)
+
+
+def test_noise_scale_operation(mock_env, device):
+  """Test noise with 'scale' operation."""
+
+  def obs_func(env):
+    return torch.full((env.num_envs, 3), 2.0, device=device)
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs1": ObservationTermCfg(
+          func=obs_func,
+          params={},
+          noise=ConstantNoiseCfg(bias=0.5, operation="scale"),
+        ),
+      },
+      enable_corruption=True,
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+  obs = manager.compute()
+
+  policy_obs = obs["actor"]
+  assert isinstance(policy_obs, torch.Tensor)
+  # Expect 2.0 * 0.5 = 1.0
+  expected = torch.full((4, 3), 1.0, device=device)
+  assert torch.allclose(policy_obs, expected)
+
+
+def test_noise_abs_operation(mock_env, device):
+  """Test noise with 'abs' operation (replaces data with bias)."""
+
+  def obs_func(env):
+    return torch.full((env.num_envs, 3), 2.0, device=device)
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs1": ObservationTermCfg(
+          func=obs_func,
+          params={},
+          noise=ConstantNoiseCfg(bias=0.7, operation="abs"),
+        ),
+      },
+      enable_corruption=True,
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+  obs = manager.compute()
+
+  policy_obs = obs["actor"]
+  assert isinstance(policy_obs, torch.Tensor)
+  # Expect data to be replaced with bias = 0.7
+  expected = torch.full((4, 3), 0.7, device=device)
+  assert torch.allclose(policy_obs, expected)
+
+
+def test_noise_with_per_dimension_bias(mock_env, device):
+  """Test noise with per-dimension bias (tuple)."""
+
+  def obs_func(env):
+    return torch.ones((env.num_envs, 3), device=device)
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs1": ObservationTermCfg(
+          func=obs_func,
+          params={},
+          noise=ConstantNoiseCfg(bias=(0.1, 0.2, 0.3), operation="add"),
+        ),
+      },
+      enable_corruption=True,
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+  obs = manager.compute()
+
+  policy_obs = obs["actor"]
+  assert isinstance(policy_obs, torch.Tensor)
+  # Expect 1.0 + [0.1, 0.2, 0.3] = [1.1, 1.2, 1.3]
+  expected = torch.tensor([[1.1, 1.2, 1.3]] * 4, device=device)
+  assert torch.allclose(policy_obs, expected)
+
+
+def test_noise_tensor_caching(device):
+  """Test that tensor conversion is cached across multiple calls."""
+  noise = ConstantNoiseCfg(bias=0.5)
+  data = torch.ones((4, 3), device=device)
+
+  # First call should create the tensor
+  result1 = noise.apply(data)
+
+  # Verify the cache was populated and get reference to cached tensor
+  device_key = str(result1.device)
+  assert device_key in noise._tensor_cache
+  assert "bias" in noise._tensor_cache[device_key]
+  cached_tensor = noise._tensor_cache[device_key]["bias"]
+  cached_data_ptr = cached_tensor.data_ptr()
+
+  # Second call should use the same cached tensor
+  result2 = noise.apply(data)
+
+  # Verify the cached tensor is the same object (same memory address)
+  assert noise._tensor_cache[device_key]["bias"].data_ptr() == cached_data_ptr
+
+  # Both results should be correct
+  expected = torch.full((4, 3), 1.5, device=device)
+  assert torch.allclose(result1, expected)
+  assert torch.allclose(result2, expected)
+
+
+def test_shared_term_name_noise_models_are_per_group(mock_env, device):
+  """Each group owns its own noise model instance when they share a term name.
+
+  ConstantNoiseCfg(op="add") shifts the additive-bias tensor by the configured
+  amount on every reset, so each group's observation drifts by its own value
+  after each manager.reset().
+  """
+
+  def obs_func(env):
+    return torch.zeros((env.num_envs, 3), device=device)
+
+  def group(bias: float) -> ObservationGroupCfg:
+    return ObservationGroupCfg(
+      terms={
+        "obs": ObservationTermCfg(
+          func=obs_func,
+          params={},
+          noise=NoiseModelWithAdditiveBiasCfg(
+            noise_cfg=ConstantNoiseCfg(bias=0.0, operation="add"),
+            bias_noise_cfg=ConstantNoiseCfg(bias=bias, operation="add"),
+          ),
+        ),
+      },
+      enable_corruption=True,
+    )
+
+  manager = ObservationManager({"actor": group(10.0), "critic": group(-1.0)}, mock_env)
+
+  obs = manager.compute()
+  assert isinstance(obs["actor"], torch.Tensor)
+  assert isinstance(obs["critic"], torch.Tensor)
+  assert torch.allclose(obs["actor"], torch.full((4, 3), 10.0, device=device))
+  assert torch.allclose(obs["critic"], torch.full((4, 3), -1.0, device=device))
+
+  manager.reset()
+  obs = manager.compute()
+  assert isinstance(obs["actor"], torch.Tensor)
+  assert isinstance(obs["critic"], torch.Tensor)
+  assert torch.allclose(obs["actor"], torch.full((4, 3), 20.0, device=device))
+  assert torch.allclose(obs["critic"], torch.full((4, 3), -2.0, device=device))
+
+
+def test_multiple_terms_with_different_noise(mock_env, device):
+  """Test multiple observation terms with different noise configs."""
+
+  def obs_func_a(env):
+    return torch.ones((env.num_envs, 2), device=device)
+
+  def obs_func_b(env):
+    return torch.full((env.num_envs, 2), 3.0, device=device)
+
+  cfg = {
+    "actor": ObservationGroupCfg(
+      terms={
+        "obs_a": ObservationTermCfg(
+          func=obs_func_a,
+          params={},
+          noise=ConstantNoiseCfg(bias=0.1, operation="add"),
+        ),
+        "obs_b": ObservationTermCfg(
+          func=obs_func_b,
+          params={},
+          noise=ConstantNoiseCfg(bias=2.0, operation="scale"),
+        ),
+      },
+      enable_corruption=True,
+    ),
+  }
+
+  manager = ObservationManager(cfg, mock_env)
+  obs = manager.compute()
+
+  policy_obs = obs["actor"]
+  assert isinstance(policy_obs, torch.Tensor)
+  # obs_a: 1.0 + 0.1 = 1.1
+  # obs_b: 3.0 * 2.0 = 6.0
+  # Concatenated: [1.1, 1.1, 6.0, 6.0]
+  expected = torch.tensor([[1.1, 1.1, 6.0, 6.0]] * 4, device=device)
+  assert torch.allclose(policy_obs, expected)
@@ -0,0 +1,45 @@
+"""Tests for offscreen renderer environment selection."""
+
+import numpy as np
+import pytest
+
+from mjlab.viewer.offscreen_renderer import OffscreenRenderer
+from mjlab.viewer.viewer_config import ViewerConfig
+
+
+def test_env_ids_clamps_when_fewer_envs_than_requested():
+  # Default max_extra_envs=2 with a single env must not raise (csv_to_npz and
+  # single-env video recording rely on this).
+  cfg = ViewerConfig(max_extra_envs=2)
+  ids = OffscreenRenderer._get_env_ids(cfg, np.zeros((1, 3)))
+  assert ids == (0,)
+
+
+def test_env_ids_primary_is_env_idx_despite_identical_origins():
+  # With identical origins, distance ties must not evict env_idx from the set.
+  cfg = ViewerConfig(env_idx=3, max_extra_envs=2)
+  ids = OffscreenRenderer._get_env_ids(cfg, np.zeros((8, 3)))
+  assert ids[0] == 3
+  assert len(ids) == 3
+  assert len(set(ids)) == 3
+
+
+def test_env_ids_selects_nearest_neighbors():
+  cfg = ViewerConfig(env_idx=0, max_extra_envs=2)
+  origins = np.array(
+    [[0.0, 0.0, 0.0], [10.0, 0.0, 0.0], [1.0, 0.0, 0.0], [2.0, 0.0, 0.0]]
+  )
+  ids = OffscreenRenderer._get_env_ids(cfg, origins)
+  assert ids == (0, 2, 3)
+
+
+def test_env_ids_out_of_range_env_idx_is_clamped():
+  cfg = ViewerConfig(env_idx=99, max_extra_envs=0)
+  ids = OffscreenRenderer._get_env_ids(cfg, np.zeros((4, 3)))
+  assert ids == (3,)
+
+
+def test_env_ids_negative_max_extra_envs_raises():
+  cfg = ViewerConfig(max_extra_envs=-1)
+  with pytest.raises(ValueError, match="max_extra_envs"):
+    OffscreenRenderer._get_env_ids(cfg, np.zeros((4, 3)))
@@ -0,0 +1,164 @@
+"""Tests for PD actuator equivalence."""
+
+import pytest
+import torch
+from conftest import (
+  create_entity_with_actuator,
+  get_test_device,
+  initialize_entity,
+  load_fixture_xml,
+)
+
+from mjlab.actuator import BuiltinPositionActuatorCfg, IdealPdActuatorCfg
+
+
+@pytest.fixture(scope="module")
+def device():
+  return get_test_device()
+
+
+@pytest.fixture(scope="module")
+def robot_xml():
+  return load_fixture_xml("floating_base_articulated")
+
+
+def test_ideal_pd_matches_builtin_at_rest(device, robot_xml):
+  """IdealPd and BuiltinPd produce equivalent actuator forces (qfrc_actuator)."""
+  kp, kv = 80.0, 10.0
+
+  ideal_entity = create_entity_with_actuator(
+    robot_xml,
+    IdealPdActuatorCfg(
+      target_names_expr=("joint.*",), effort_limit=100.0, stiffness=kp, damping=kv
+    ),
+  )
+  builtin_entity = create_entity_with_actuator(
+    robot_xml,
+    BuiltinPositionActuatorCfg(
+      target_names_expr=("joint.*",), effort_limit=100.0, stiffness=kp, damping=kv
+    ),
+  )
+
+  ideal_entity, ideal_sim = initialize_entity(ideal_entity, device)
+  builtin_entity, builtin_sim = initialize_entity(builtin_entity, device)
+
+  joint_pos = torch.tensor([[0.1, -0.05]], device=device)
+  joint_vel = torch.tensor([[0.0, 0.0]], device=device)
+  ideal_entity.write_joint_state_to_sim(joint_pos, joint_vel)
+  builtin_entity.write_joint_state_to_sim(joint_pos, joint_vel)
+
+  pos_target = torch.tensor([[0.5, -0.3]], device=device)
+  vel_target = torch.tensor([[0.2, -0.1]], device=device)
+  ideal_entity.set_joint_position_target(pos_target)
+  ideal_entity.set_joint_velocity_target(vel_target)
+  ideal_entity.set_joint_effort_target(torch.zeros(1, 2, device=device))
+  builtin_entity.set_joint_position_target(pos_target)
+
+  ideal_entity.write_data_to_sim()
+  builtin_entity.write_data_to_sim()
+
+  joint_v_adr = ideal_entity.indexing.joint_v_adr
+  ideal_qfrc = ideal_sim.data.qfrc_actuator[0, joint_v_adr]
+  builtin_qfrc = builtin_sim.data.qfrc_actuator[0, joint_v_adr]
+
+  assert torch.allclose(ideal_qfrc, builtin_qfrc, rtol=1e-4, atol=1e-5)
+
+
+def test_ideal_pd_matches_builtin_with_velocity(device, robot_xml):
+  """IdealPd and BuiltinPd produce equivalent damping forces with nonzero vel."""
+  kp, kv = 80.0, 10.0
+
+  ideal_entity = create_entity_with_actuator(
+    robot_xml,
+    IdealPdActuatorCfg(
+      target_names_expr=("joint.*",), effort_limit=100.0, stiffness=kp, damping=kv
+    ),
+  )
+  builtin_entity = create_entity_with_actuator(
+    robot_xml,
+    BuiltinPositionActuatorCfg(
+      target_names_expr=("joint.*",), effort_limit=100.0, stiffness=kp, damping=kv
+    ),
+  )
+
+  ideal_entity, ideal_sim = initialize_entity(ideal_entity, device)
+  builtin_entity, builtin_sim = initialize_entity(builtin_entity, device)
+
+  joint_pos = torch.tensor([[0.2, -0.1]], device=device)
+  joint_vel = torch.tensor([[0.8, -0.5]], device=device)
+  ideal_entity.write_joint_state_to_sim(joint_pos, joint_vel)
+  builtin_entity.write_joint_state_to_sim(joint_pos, joint_vel)
+
+  pos_target = torch.tensor([[0.5, 0.1]], device=device)
+  vel_target = torch.tensor([[0.3, -0.2]], device=device)
+  ideal_entity.set_joint_position_target(pos_target)
+  ideal_entity.set_joint_velocity_target(vel_target)
+  ideal_entity.set_joint_effort_target(torch.zeros(1, 2, device=device))
+  builtin_entity.set_joint_position_target(pos_target)
+
+  ideal_entity.write_data_to_sim()
+  builtin_entity.write_data_to_sim()
+
+  joint_v_adr = ideal_entity.indexing.joint_v_adr
+  ideal_qfrc = ideal_sim.data.qfrc_actuator[0, joint_v_adr]
+  builtin_qfrc = builtin_sim.data.qfrc_actuator[0, joint_v_adr]
+
+  assert torch.allclose(ideal_qfrc, builtin_qfrc, rtol=1e-4, atol=1e-5)
+
+
+def test_ideal_pd_with_feedforward_effort(device, robot_xml):
+  """IdealPd adds feedforward effort to PD output."""
+  kp, kv = 50.0, 5.0
+
+  ideal_entity = create_entity_with_actuator(
+    robot_xml,
+    IdealPdActuatorCfg(
+      target_names_expr=("joint.*",), effort_limit=100.0, stiffness=kp, damping=kv
+    ),
+  )
+
+  ideal_entity, ideal_sim = initialize_entity(ideal_entity, device)
+
+  joint_pos = torch.tensor([[0.5, 0.0]], device=device)
+  joint_vel = torch.tensor([[0.0, 0.0]], device=device)
+  ideal_entity.write_joint_state_to_sim(joint_pos, joint_vel)
+
+  ideal_entity.set_joint_position_target(joint_pos)
+  ideal_entity.set_joint_velocity_target(joint_vel)
+  ideal_entity.set_joint_effort_target(torch.tensor([[2.0, -1.0]], device=device))
+  ideal_entity.write_data_to_sim()
+
+  ideal_ctrl = ideal_sim.data.ctrl[0]
+  assert torch.allclose(ideal_ctrl, torch.tensor([2.0, -1.0], device=device))
+
+
+def test_ideal_pd_effort_clamping(device, robot_xml):
+  """IdealPd clamps computed torques to [-effort_limit, effort_limit]."""
+  kp, kv = 100.0, 10.0
+  effort_limit = 5.0
+
+  ideal_entity = create_entity_with_actuator(
+    robot_xml,
+    IdealPdActuatorCfg(
+      target_names_expr=("joint.*",),
+      effort_limit=effort_limit,
+      stiffness=kp,
+      damping=kv,
+    ),
+  )
+
+  ideal_entity, ideal_sim = initialize_entity(ideal_entity, device)
+
+  joint_pos = torch.tensor([[0.0, 0.0]], device=device)
+  joint_vel = torch.tensor([[0.0, 0.0]], device=device)
+  ideal_entity.write_joint_state_to_sim(joint_pos, joint_vel)
+
+  ideal_entity.set_joint_position_target(torch.tensor([[1.0, -1.0]], device=device))
+  ideal_entity.set_joint_velocity_target(torch.zeros(1, 2, device=device))
+  ideal_entity.set_joint_effort_target(torch.zeros(1, 2, device=device))
+  ideal_entity.write_data_to_sim()
+
+  ideal_ctrl = ideal_sim.data.ctrl[0]
+  assert torch.allclose(
+    ideal_ctrl, torch.tensor([effort_limit, -effort_limit], device=device)
+  )
@@ -0,0 +1,109 @@
+"""Tests for Perlin noise terrain generation."""
+
+import mujoco
+import numpy as np
+import pytest
+
+from mjlab.terrains.heightfield_terrains import HfPerlinNoiseTerrainCfg
+
+
+@pytest.fixture
+def rng() -> np.random.Generator:
+  return np.random.default_rng(42)
+
+
+def test_perlin_terrain_generation(rng: np.random.Generator):
+  """Verify that Perlin noise terrain generates a valid TerrainOutput."""
+  spec = mujoco.MjSpec()
+  spec.worldbody.add_body(name="terrain")
+
+  cfg = HfPerlinNoiseTerrainCfg(
+    proportion=1.0,
+    size=(10.0, 10.0),
+    height_range=(0.1, 0.5),
+    octaves=4,
+    persistence=0.5,
+    lacunarity=2.0,
+    scale=10.0,
+    horizontal_scale=0.1,
+    resolution=0.1,
+  )
+
+  output = cfg.function(difficulty=0.5, spec=spec, rng=rng)
+
+  # Basic checks for TerrainOutput.
+  assert output.origin.shape == (3,)
+  assert len(output.geometries) == 1
+  assert output.geometries[0].geom is not None
+  assert output.geometries[0].hfield is not None
+
+  # Check hfield size.
+  # size = [x/2, y/2, max_height, base_thickness]
+  hfield = output.geometries[0].hfield
+  assert hfield.size[0] == 5.0
+  assert hfield.size[1] == 5.0
+  assert hfield.size[2] > 0  # max height should be positive
+
+  # Check hfield data.
+  data = np.asarray(hfield.userdata)
+  assert len(data) == 100 * 100  # (10.0/0.1) * (10.0/0.1)
+  assert np.all(data >= 0.0)
+  assert np.all(data <= 1.0)
+  assert np.max(data) > 0.0  # it should not be flat 0
+
+  # Check if height variation is present.
+  assert np.min(data) < np.max(data)
+
+
+def test_perlin_terrain_resolution(rng: np.random.Generator):
+  """Verify that higher resolution increases pixel count."""
+  spec = mujoco.MjSpec()
+  spec.worldbody.add_body(name="terrain")
+
+  cfg = HfPerlinNoiseTerrainCfg(
+    proportion=1.0,
+    size=(10.0, 10.0),
+    height_range=(0.1, 0.5),
+    horizontal_scale=0.1,
+    resolution=0.05,  # 200 pixels
+  )
+
+  output = cfg.function(difficulty=0.5, spec=spec, rng=rng)
+  hfield = output.geometries[0].hfield
+  assert hfield is not None
+  # (10.0 / 0.05) = 200
+  assert hfield.nrow == 200
+  assert hfield.ncol == 200
+  assert len(hfield.userdata) == 200 * 200
+
+
+def test_perlin_terrain_with_border(rng: np.random.Generator):
+  """Verify Perlin noise terrain with a border."""
+  spec = mujoco.MjSpec()
+  spec.worldbody.add_body(name="terrain")
+
+  cfg = HfPerlinNoiseTerrainCfg(
+    proportion=1.0,
+    size=(10.0, 10.0),
+    height_range=(0.1, 0.5),
+    border_width=2.0,
+    horizontal_scale=0.1,
+    resolution=0.1,
+  )
+
+  output = cfg.function(difficulty=0.5, spec=spec, rng=rng)
+  hfield = output.geometries[0].hfield
+  assert hfield is not None
+  data = np.asarray(hfield.userdata).reshape(100, 100)
+
+  # Border pixels: 2.0 / 0.1 = 20
+  border_pixels = 20
+  top_border = data[:border_pixels, :]
+  bottom_border = data[-border_pixels:, :]
+  left_border = data[:, :border_pixels]
+  right_border = data[:, -border_pixels:]
+
+  assert np.all(top_border == top_border[0, 0])
+  assert np.all(bottom_border == bottom_border[0, 0])
+  assert np.all(left_border == left_border[0, 0])
+  assert np.all(right_border == right_border[0, 0])
@@ -0,0 +1,213 @@
+"""Tests for sensor-based projected gravity (framezaxis up-vector sensor).
+
+The shipped robots expose a ``framezaxis`` sensor that outputs the world Z-axis in the
+IMU site frame; negating it gives projected gravity. These tests check the sensor (and
+the ``projected_gravity_from_sensor`` observation that wraps it) against an independent
+ground-truth computation, and verify that -- unlike the entity-data
+``projected_gravity_b`` -- it tracks the IMU site orientation, which is what makes IMU
+mounting domain randomization observable.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import TYPE_CHECKING, cast
+
+import mujoco
+import pytest
+import torch
+from conftest import get_test_device
+
+from mjlab.entity import EntityCfg
+from mjlab.envs.mdp import dr
+from mjlab.envs.mdp.observations import projected_gravity_from_sensor
+from mjlab.managers.scene_entity_config import SceneEntityCfg
+from mjlab.scene import Scene, SceneCfg
+from mjlab.sim.sim import Simulation, SimulationCfg
+
+if TYPE_CHECKING:
+  from mjlab.envs import ManagerBasedRlEnv
+
+# Gravity points along world -Z; projected gravity is this expressed in a body frame.
+_GRAVITY_DIR_W = (0.0, 0.0, -1.0)
+
+
+def _quat_to_mat(q: tuple[float, float, float, float]) -> torch.Tensor:
+  """Rotation matrix from a (w, x, y, z) quaternion. Independent of MuJoCo/mjlab."""
+  w, x, y, z = q
+  return torch.tensor(
+    [
+      [1 - 2 * (y * y + z * z), 2 * (x * y - w * z), 2 * (x * z + w * y)],
+      [2 * (x * y + w * z), 1 - 2 * (x * x + z * z), 2 * (y * z - w * x)],
+      [2 * (x * z - w * y), 2 * (y * z + w * x), 1 - 2 * (x * x + y * y)],
+    ],
+    dtype=torch.float64,
+  )
+
+
+def _expected_projected_gravity(q: tuple[float, float, float, float]) -> torch.Tensor:
+  """Ground-truth projected gravity for a body with world orientation ``q``.
+
+  proj = R(q)^T @ g_world, computed from an explicit rotation matrix so it does not
+  share a code path with the sensor or with ``projected_gravity_b``.
+  """
+  g_w = torch.tensor(_GRAVITY_DIR_W, dtype=torch.float64)
+  return _quat_to_mat(q).T @ g_w
+
+
+class Env:
+  """Minimal env stub for driving observation and dr functions in tests."""
+
+  def __init__(self, scene, sim, device):
+    self.scene = scene
+    self.sim = sim
+    self.num_envs = scene.num_envs
+    self.device = device
+
+
+def _make_env(scene, sim, device) -> ManagerBasedRlEnv:
+  """Build the env stub, typed as the real env for the functions under test."""
+  return cast("ManagerBasedRlEnv", Env(scene, sim, device))
+
+
+@pytest.fixture(scope="module")
+def device():
+  return get_test_device()
+
+
+def _robot_xml(site_euler: str = "0 0 0") -> str:
+  """Free-floating box with an IMU site and the framezaxis up-vector sensor."""
+  return f"""
+    <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="imu" pos="0.05 0 0" euler="{site_euler}"/>
+        </body>
+      </worldbody>
+      <sensor>
+        <framezaxis name="imu_upvector" objtype="body" objname="world"
+                    reftype="site" refname="imu"/>
+      </sensor>
+    </mujoco>
+  """
+
+
+def _build(xml: str, device: str, num_envs: int = 2):
+  entity_cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(xml))
+  scene = Scene(
+    SceneCfg(num_envs=num_envs, env_spacing=3.0, entities={"robot": entity_cfg}),
+    device,
+  )
+  model = scene.compile()
+  sim = Simulation(
+    num_envs=num_envs, cfg=SimulationCfg(njmax=20), model=model, device=device
+  )
+  scene.initialize(sim.mj_model, sim.model, sim.data)
+  return scene, sim
+
+
+def _set_root_quat(robot, q: tuple[float, float, float, float], device: str) -> None:
+  root_state = robot.data.default_root_state.clone()
+  root_state[:, 3:7] = torch.tensor(q, device=device, dtype=root_state.dtype)
+  robot.write_root_state_to_sim(root_state)
+
+
+def test_sensor_matches_ground_truth_when_site_aligned(device):
+  """Sensor and entity both equal hand-computed projected gravity for a tilted base."""
+  scene, sim = _build(_robot_xml(), device)
+  robot = scene["robot"]
+
+  # Compose a 0.6 rad roll with a 0.3 rad pitch into a single root quaternion.
+  ax = (math.cos(0.3), math.sin(0.3), 0.0, 0.0)
+  ay = (math.cos(0.15), 0.0, math.sin(0.15), 0.0)
+  q = (
+    ax[0] * ay[0] - ax[1] * ay[1] - ax[2] * ay[2] - ax[3] * ay[3],
+    ax[0] * ay[1] + ax[1] * ay[0] + ax[2] * ay[3] - ax[3] * ay[2],
+    ax[0] * ay[2] - ax[1] * ay[3] + ax[2] * ay[0] + ax[3] * ay[1],
+    ax[0] * ay[3] + ax[1] * ay[2] - ax[2] * ay[1] + ax[3] * ay[0],
+  )
+  _set_root_quat(robot, q, device)
+  sim.forward()
+
+  expected = _expected_projected_gravity(q).to(device=device, dtype=torch.float32)
+  # Guard against a vacuous pass: the tilt must actually move gravity off straight-down.
+  straight_down = torch.tensor(_GRAVITY_DIR_W, device=device)
+  assert (expected - straight_down).abs().max() > 0.3
+
+  sensor_grav = -scene["robot/imu_upvector"].data
+  entity_grav = robot.data.projected_gravity_b
+  torch.testing.assert_close(sensor_grav[0], expected, atol=1e-5, rtol=0)
+  torch.testing.assert_close(entity_grav[0], expected, atol=1e-5, rtol=0)
+
+
+def test_observation_fn_tracks_site_orientation(device):
+  """The observation fn reflects IMU site tilt; the entity-data version does not.
+
+  With the base upright but the IMU site rolled 30 deg about x, projected gravity in the
+  site frame is (0, -sin30, -cos30). The entity-data version stays straight-down because
+  it uses the root body orientation and is blind to the site.
+  """
+  scene_rot, sim_rot = _build(_robot_xml(site_euler="30 0 0"), device)
+  scene_flat, sim_flat = _build(_robot_xml(site_euler="0 0 0"), device)
+  sim_rot.forward()
+  sim_flat.forward()
+
+  # Drive through the actual shipped observation function, not the raw sensor.
+  env_rot = _make_env(scene_rot, sim_rot, device)
+  env_flat = _make_env(scene_flat, sim_flat, device)
+  grav_rot = projected_gravity_from_sensor(env_rot, "robot/imu_upvector")
+  grav_flat = projected_gravity_from_sensor(env_flat, "robot/imu_upvector")
+
+  expected_rot = torch.tensor(
+    [0.0, -math.sin(math.radians(30)), -math.cos(math.radians(30))], device=device
+  )
+  straight_down = torch.tensor(_GRAVITY_DIR_W, device=device)
+  torch.testing.assert_close(grav_rot[0], expected_rot, atol=1e-5, rtol=0)
+  torch.testing.assert_close(grav_flat[0], straight_down, atol=1e-5, rtol=0)
+
+  # The entity-data version is unchanged by the site rotation (so it cannot be used to
+  # observe IMU mounting randomization), confirming why the sensor path is needed.
+  entity_rot = scene_rot["robot"].data.projected_gravity_b
+  torch.testing.assert_close(entity_rot[0], straight_down, atol=1e-5, rtol=0)
+
+
+@pytest.mark.filterwarnings(
+  "ignore:Use of index_put_ on expanded tensors is deprecated:UserWarning"
+)
+def test_site_quat_randomization_changes_sensor(device):
+  """The full DR path: running ``dr.site_quat`` perturbs the gravity observation.
+
+  This is what the G1 example configs rely on -- randomizing the IMU site orientation
+  must show up in the sensor-based projected gravity, per-environment.
+  """
+  scene, sim = _build(_robot_xml(), device, num_envs=4)
+  sim.expand_model_fields(("site_quat",))
+  env = _make_env(scene, sim, device)
+
+  sim.forward()
+  straight_down = torch.tensor(_GRAVITY_DIR_W, device=device)
+  before = projected_gravity_from_sensor(env, "robot/imu_upvector").clone()
+  # Upright base + identity site quat => straight-down gravity in every env.
+  torch.testing.assert_close(before, straight_down.expand_as(before), atol=1e-5, rtol=0)
+
+  torch.manual_seed(0)
+  dr.site_quat(
+    env,
+    env_ids=None,
+    roll_range=(-0.3, 0.3),
+    pitch_range=(-0.3, 0.3),
+    yaw_range=(-0.3, 0.3),
+    asset_cfg=SceneEntityCfg("robot", site_names=("imu",)),
+  )
+  sim.forward()
+  after = projected_gravity_from_sensor(env, "robot/imu_upvector")
+
+  # Randomization moved the reading off straight-down and made it env-dependent.
+  assert (after - straight_down).abs().max() > 0.05
+  assert not torch.allclose(after, before, atol=1e-3)
+  assert torch.unique(after, dim=0).shape[0] >= 2
+  # The perturbation is a rotation, so gravity stays a unit vector.
+  norms = torch.linalg.norm(after, dim=-1)
+  torch.testing.assert_close(norms, torch.ones_like(norms), atol=1e-5, rtol=0)
@@ -0,0 +1,502 @@
+"""Tests for the recorder manager."""
+
+from unittest.mock import Mock
+
+import mujoco
+import pytest
+import torch
+from conftest import get_test_device
+
+from mjlab.actuator import BuiltinPositionActuatorCfg
+from mjlab.entity import EntityArticulationInfoCfg, EntityCfg
+from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg, mdp
+from mjlab.managers.observation_manager import ObservationGroupCfg, ObservationTermCfg
+from mjlab.managers.recorder_manager import (
+  NullRecorderManager,
+  RecorderManager,
+  RecorderTerm,
+  RecorderTermCfg,
+)
+from mjlab.managers.termination_manager import TerminationTermCfg
+from mjlab.scene import SceneCfg
+from mjlab.sim import MujocoCfg, SimulationCfg
+from mjlab.terrains import TerrainEntityCfg
+
+
+class _CountingRecorder(RecorderTerm):
+  """Records calls to each lifecycle method for inspection."""
+
+  def __init__(self, cfg, env):
+    super().__init__(cfg, env)
+    self.pre_reset_calls: list[torch.Tensor] = []
+    self.post_reset_calls: list[torch.Tensor] = []
+    self.post_step_count = 0
+    self.closed = False
+
+  def record_pre_reset(self, env_ids: torch.Tensor) -> None:
+    self.pre_reset_calls.append(env_ids.clone())
+
+  def record_post_reset(self, env_ids: torch.Tensor) -> None:
+    self.post_reset_calls.append(env_ids.clone())
+
+  def record_post_step(self) -> None:
+    self.post_step_count += 1
+
+  def close(self) -> None:
+    self.closed = True
+
+
+@pytest.fixture
+def mock_env() -> Mock:
+  env = Mock()
+  env.num_envs = 4
+  env.device = "cpu"
+  return env
+
+
+@pytest.fixture
+def manager_and_term(mock_env: Mock) -> tuple[RecorderManager, _CountingRecorder]:
+  cfg = {"recorder": RecorderTermCfg(func=_CountingRecorder, params={})}
+  manager = RecorderManager(cfg, mock_env)
+  term = manager.get_term("recorder")
+  assert isinstance(term, _CountingRecorder)
+  return manager, term
+
+
+# Lifecycle semantics.
+
+
+def test_terminal_action_available_at_record_pre_reset(mock_env: Mock) -> None:
+  """record_pre_reset fires before _reset_idx, so action_manager.action is intact.
+
+  In step(), the sequence is:
+    record_pre_reset(reset_env_ids)   <- action still set
+    _reset_idx(reset_env_ids)         <- zeroes action_manager.action
+    ...
+    record_post_step()                <- action is now 0 for reset envs
+
+  A recorder that wants the terminal action MUST capture it at record_pre_reset.
+  """
+  seen: dict[str, torch.Tensor] = {}
+
+  class ActionCapture(RecorderTerm):
+    def record_pre_reset(self, env_ids: torch.Tensor) -> None:
+      seen["pre_reset_action"] = self._env.action_manager.action[env_ids].clone()
+
+    def record_post_step(self) -> None:
+      seen["post_step_action"] = self._env.action_manager.action.clone()
+
+  mock_env.action_manager = Mock()
+  mock_env.action_manager.action = torch.tensor(
+    [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]]
+  )
+
+  cfg = {"r": RecorderTermCfg(func=ActionCapture, params={})}
+  manager = RecorderManager(cfg, mock_env)
+
+  reset_env_ids = torch.tensor([0, 2])
+
+  # Simulate step(): record_pre_reset fires before _reset_idx zeros actions.
+  manager.record_pre_reset(reset_env_ids)
+  assert torch.equal(
+    seen["pre_reset_action"],
+    torch.tensor([[1.0, 2.0], [5.0, 6.0]]),
+  ), "Terminal action must be intact at record_pre_reset"
+
+  # Simulate _reset_idx: zeros action for reset envs.
+  mock_env.action_manager.action[reset_env_ids] = 0.0
+
+  manager.record_post_step()
+  assert torch.all(seen["post_step_action"][reset_env_ids] == 0.0), (
+    "Action is zeroed for reset envs at record_post_step; use record_pre_reset"
+    " to capture terminal actions"
+  )
+  # Non-reset envs retain their action.
+  non_reset = torch.tensor([1, 3])
+  assert torch.equal(
+    seen["post_step_action"][non_reset],
+    torch.tensor([[3.0, 4.0], [7.0, 8.0]]),
+  )
+
+
+def test_record_pre_reset_obs_buf_is_previous_step(mock_env: Mock) -> None:
+  """obs_buf at record_pre_reset is the previous step's observation, not the post-action
+  terminal state. It is the input obs the agent used to pick the terminal action.
+  """
+  seen: dict[str, torch.Tensor] = {}
+
+  # Use a plain tensor for obs_buf to avoid dict-subscript complexity in the type
+  # checker. The point of this test is lifecycle timing, not data structure.
+  class ObsCapture(RecorderTerm):
+    def record_pre_reset(self, env_ids: torch.Tensor) -> None:
+      obs: torch.Tensor = self._env.obs_buf  # type: ignore[assignment]
+      seen["pre_reset_obs"] = obs[env_ids].clone()
+
+    def record_post_step(self) -> None:
+      obs: torch.Tensor = self._env.obs_buf  # type: ignore[assignment]
+      seen["post_step_obs"] = obs.clone()
+
+  prev_obs = torch.arange(8, dtype=torch.float).reshape(4, 2)
+  mock_env.obs_buf = prev_obs.clone()
+
+  cfg = {"r": RecorderTermCfg(func=ObsCapture, params={})}
+  manager = RecorderManager(cfg, mock_env)
+
+  reset_env_ids = torch.tensor([1])
+  manager.record_pre_reset(reset_env_ids)
+  # obs_buf at pre_reset is unchanged from the previous step.
+  assert torch.equal(seen["pre_reset_obs"], prev_obs[reset_env_ids])
+
+  # Simulate _reset_idx + obs recompute: reset env gets new initial obs.
+  mock_env.obs_buf[reset_env_ids] = torch.zeros(1, 2)
+  manager.record_post_step()
+  assert torch.all(seen["post_step_obs"][reset_env_ids] == 0.0), (
+    "After reset, obs_buf for the reset env holds the new-episode initial obs"
+  )
+
+
+def test_record_post_reset_fires_after_obs_are_fresh(mock_env: Mock) -> None:
+  """record_post_reset is called after obs are computed, so obs_buf already holds the
+  new-episode initial observation for the reset environments.
+  """
+  seen: dict[str, torch.Tensor] = {}
+
+  class ResetObsCapture(RecorderTerm):
+    def record_post_reset(self, env_ids: torch.Tensor) -> None:
+      obs: torch.Tensor = self._env.obs_buf  # type: ignore[assignment]
+      seen["obs"] = obs[env_ids].clone()
+
+  reset_obs = torch.ones(4, 2) * 99.0
+  mock_env.obs_buf = reset_obs
+
+  cfg = {"r": RecorderTermCfg(func=ResetObsCapture, params={})}
+  manager = RecorderManager(cfg, mock_env)
+
+  env_ids = torch.tensor([0, 3])
+  manager.record_post_reset(env_ids)
+  assert torch.equal(seen["obs"], reset_obs[env_ids])
+
+
+def test_record_pre_reset_before_record_post_step_in_step(mock_env: Mock) -> None:
+  """record_pre_reset fires before record_post_step within a single step."""
+  call_log: list[str] = []
+
+  class OrderRecorder(RecorderTerm):
+    def record_pre_reset(self, env_ids: torch.Tensor) -> None:
+      call_log.append("pre_reset")
+
+    def record_post_step(self) -> None:
+      call_log.append("post_step")
+
+  cfg = {"r": RecorderTermCfg(func=OrderRecorder, params={})}
+  manager = RecorderManager(cfg, mock_env)
+
+  manager.record_pre_reset(torch.tensor([0]))
+  manager.record_post_step()
+
+  assert call_log == ["pre_reset", "post_step"]
+
+
+# Manager dispatch.
+
+
+def test_record_post_step_called_each_step(
+  manager_and_term: tuple[RecorderManager, _CountingRecorder],
+) -> None:
+  manager, term = manager_and_term
+  for _ in range(5):
+    manager.record_post_step()
+  assert term.post_step_count == 5
+
+
+def test_record_pre_reset_passes_env_ids(
+  manager_and_term: tuple[RecorderManager, _CountingRecorder],
+) -> None:
+  manager, term = manager_and_term
+  env_ids = torch.tensor([0, 2])
+  manager.record_pre_reset(env_ids)
+  assert len(term.pre_reset_calls) == 1
+  assert torch.equal(term.pre_reset_calls[0], env_ids)
+
+
+def test_record_post_reset_passes_env_ids(
+  manager_and_term: tuple[RecorderManager, _CountingRecorder],
+) -> None:
+  manager, term = manager_and_term
+  env_ids = torch.tensor([1, 3])
+  manager.record_post_reset(env_ids)
+  assert len(term.post_reset_calls) == 1
+  assert torch.equal(term.post_reset_calls[0], env_ids)
+
+
+def test_close_propagates_to_all_terms(mock_env: Mock) -> None:
+  cfg = {
+    "a": RecorderTermCfg(func=_CountingRecorder, params={}),
+    "b": RecorderTermCfg(func=_CountingRecorder, params={}),
+  }
+  manager = RecorderManager(cfg, mock_env)
+  term_a = manager.get_term("a")
+  term_b = manager.get_term("b")
+  assert isinstance(term_a, _CountingRecorder)
+  assert isinstance(term_b, _CountingRecorder)
+
+  manager.close()
+  assert term_a.closed
+  assert term_b.closed
+
+
+def test_terms_called_in_registration_order(mock_env: Mock) -> None:
+  call_log: list[str] = []
+
+  class NamedRecorder(RecorderTerm):
+    def __init__(self, cfg, env):
+      super().__init__(cfg, env)
+      self._name: str = cfg.params["name"]
+
+    def record_post_step(self) -> None:
+      call_log.append(self._name)
+
+  cfg = {
+    "first": RecorderTermCfg(func=NamedRecorder, params={"name": "first"}),
+    "second": RecorderTermCfg(func=NamedRecorder, params={"name": "second"}),
+    "third": RecorderTermCfg(func=NamedRecorder, params={"name": "third"}),
+  }
+  manager = RecorderManager(cfg, mock_env)
+  manager.record_post_step()
+  assert call_log == ["first", "second", "third"]
+
+
+def test_multiple_terms_have_independent_state(mock_env: Mock) -> None:
+  cfg = {
+    "a": RecorderTermCfg(func=_CountingRecorder, params={}),
+    "b": RecorderTermCfg(func=_CountingRecorder, params={}),
+  }
+  manager = RecorderManager(cfg, mock_env)
+  term_a = manager.get_term("a")
+  term_b = manager.get_term("b")
+  assert isinstance(term_a, _CountingRecorder)
+  assert isinstance(term_b, _CountingRecorder)
+
+  manager.record_post_step()
+  manager.record_post_step()
+  manager.record_pre_reset(torch.tensor([0]))
+
+  assert term_a.post_step_count == 2
+  assert term_b.post_step_count == 2
+  assert len(term_a.pre_reset_calls) == 1
+  assert len(term_b.pre_reset_calls) == 1
+
+
+def test_active_terms_matches_registration_order(mock_env: Mock) -> None:
+  cfg = {
+    "first": RecorderTermCfg(func=_CountingRecorder, params={}),
+    "second": RecorderTermCfg(func=_CountingRecorder, params={}),
+  }
+  manager = RecorderManager(cfg, mock_env)
+  assert manager.active_terms == ["first", "second"]
+
+
+def test_none_terms_skipped(mock_env: Mock) -> None:
+  cfg: dict[str, RecorderTermCfg | None] = {
+    "valid": RecorderTermCfg(func=_CountingRecorder, params={}),
+    "skipped": None,
+  }
+  manager = RecorderManager(cfg, mock_env)  # type: ignore[arg-type]
+  assert manager.active_terms == ["valid"]
+  assert len(manager._terms) == 1
+
+
+def test_function_based_term_raises(mock_env: Mock) -> None:
+  cfg = {"bad": RecorderTermCfg(func=lambda env: None, params={})}
+  with pytest.raises(TypeError, match="RecorderTerm subclass"):
+    RecorderManager(cfg, mock_env)
+
+
+def test_cfg_params_accessible_on_term(mock_env: Mock) -> None:
+  cfg = {"t": RecorderTermCfg(func=_CountingRecorder, params={"path": "/tmp/out.csv"})}
+  manager = RecorderManager(cfg, mock_env)
+  term = manager.get_term("t")
+  assert isinstance(term, _CountingRecorder)
+  assert term.cfg.params["path"] == "/tmp/out.csv"
+
+
+def test_deepcopy_protects_original_cfg(mock_env: Mock) -> None:
+  """Creating a manager must not mutate the caller's cfg dict.
+
+  _resolve_common_term_cfg replaces cfg.func with the instantiated object in-place on
+  the internal copy. Without deepcopy the caller's cfg would be silently corrupted,
+  breaking any second use.
+  """
+  original_cfg = {"t": RecorderTermCfg(func=_CountingRecorder, params={})}
+  RecorderManager(original_cfg, mock_env)
+  assert original_cfg["t"].func is _CountingRecorder
+
+
+def test_null_recorder_manager_no_ops() -> None:
+  manager = NullRecorderManager()
+  env_ids = torch.tensor([0, 1])
+  manager.record_pre_reset(env_ids)
+  manager.record_post_reset(env_ids)
+  manager.record_post_step()
+  manager.close()
+  assert manager.active_terms == []
+
+
+# Public term lookup.
+
+
+def test_get_term_returns_registered_instance(mock_env: Mock) -> None:
+  """get_term returns the live term the manager dispatches to, not a copy."""
+  cfg = {"recorder": RecorderTermCfg(func=_CountingRecorder, params={})}
+  manager = RecorderManager(cfg, mock_env)
+  term = manager.get_term("recorder")
+  assert isinstance(term, _CountingRecorder)
+  manager.record_post_step()
+  # The instance returned by get_term is the same one the manager dispatches to.
+  assert term.post_step_count == 1
+
+
+def test_get_term_missing_raises(mock_env: Mock) -> None:
+  """get_term raises KeyError when the requested term name is not registered."""
+  cfg = {"recorder": RecorderTermCfg(func=_CountingRecorder, params={})}
+  manager = RecorderManager(cfg, mock_env)
+  with pytest.raises(KeyError, match="missing"):
+    manager.get_term("missing")
+
+
+def test_contains_for_recorder_manager(mock_env: Mock) -> None:
+  """`in` reflects whether a term name is registered on the manager."""
+  cfg = {"recorder": RecorderTermCfg(func=_CountingRecorder, params={})}
+  manager = RecorderManager(cfg, mock_env)
+  assert "recorder" in manager
+  assert "other" not in manager
+
+
+def test_contains_and_get_term_for_null_manager() -> None:
+  """NullRecorderManager has no terms: `in` is always False and get_term raises."""
+  manager = NullRecorderManager()
+  assert "anything" not in manager
+  with pytest.raises(KeyError):
+    manager.get_term("anything")
+
+
+def test_base_term_all_hooks_are_no_ops(mock_env: Mock) -> None:
+  """Base RecorderTerm hooks do nothing; __call__ raises NotImplementedError."""
+
+  class MinimalRecorder(RecorderTerm):
+    pass
+
+  term = MinimalRecorder(
+    cfg=RecorderTermCfg(func=MinimalRecorder, params={}), env=mock_env
+  )
+  env_ids = torch.tensor([0])
+  term.record_pre_reset(env_ids)
+  term.record_post_reset(env_ids)
+  term.record_post_step()
+  term.close()
+
+  with pytest.raises(NotImplementedError):
+    term()
+
+
+# Integration: real ManagerBasedRlEnv.
+
+
+_SLIDING_MASS_XML = """
+<mujoco>
+  <option timestep="0.002"/>
+  <worldbody>
+    <body name="mass" pos="0 0 0">
+      <joint name="slide" type="slide" axis="1 0 0" range="-1 1" limited="true"/>
+      <geom name="mass_geom" type="sphere" size="0.1" mass="1.0"/>
+    </body>
+  </worldbody>
+</mujoco>
+"""
+
+
+def _make_env_cfg(recorders: dict | None = None) -> ManagerBasedRlEnvCfg:
+  """One-step episode (time_out fires after a single step)."""
+  timestep = 0.002
+  return ManagerBasedRlEnvCfg(
+    scene=SceneCfg(
+      terrain=TerrainEntityCfg(terrain_type="plane"),
+      num_envs=2,
+      extent=1.0,
+      entities={
+        "robot": EntityCfg(
+          spec_fn=lambda: mujoco.MjSpec.from_string(_SLIDING_MASS_XML),
+          articulation=EntityArticulationInfoCfg(
+            actuators=(
+              BuiltinPositionActuatorCfg(
+                target_names_expr=(".*",), stiffness=100.0, damping=10.0
+              ),
+            )
+          ),
+        )
+      },
+    ),
+    observations={
+      "actor": ObservationGroupCfg(
+        terms={"obs": ObservationTermCfg(func=mdp.joint_pos_rel)}
+      )
+    },
+    actions={
+      "joint_pos": mdp.JointPositionActionCfg(
+        entity_name="robot", actuator_names=(".*",), scale=1.0
+      )
+    },
+    terminations={"time_out": TerminationTermCfg(func=mdp.time_out, time_out=True)},
+    events={},
+    sim=SimulationCfg(mujoco=MujocoCfg(timestep=timestep, iterations=1)),
+    decimation=1,
+    episode_length_s=timestep,  # terminates every env after exactly one step
+    recorders=recorders or {},
+  )
+
+
+@pytest.fixture(scope="module")
+def device() -> str:
+  return get_test_device()
+
+
+def test_integration_terminal_action_available_at_record_pre_reset(device: str) -> None:
+  """In a real env.step(), record_pre_reset fires before _reset_idx zeroes
+  action_manager.action for reset envs. This verifies the hook ordering described in
+  the docstring is enforced by the actual step() implementation.
+  """
+  seen: dict[str, torch.Tensor] = {}
+
+  class HookOrderRecorder(RecorderTerm):
+    def record_pre_reset(self, env_ids: torch.Tensor) -> None:
+      # Action must still be non-zero; _reset_idx has not run yet.
+      seen["pre_reset_action"] = self._env.action_manager.action[env_ids].clone()
+
+    def record_post_step(self) -> None:
+      seen["post_step_action"] = self._env.action_manager.action.clone()
+      seen["reset_buf"] = self._env.reset_buf.clone()
+
+  cfg = _make_env_cfg(
+    recorders={"r": RecorderTermCfg(func=HookOrderRecorder, params={})}
+  )
+  env = ManagerBasedRlEnv(cfg=cfg, device=device)
+  env.reset()
+
+  action = torch.ones(2, 1, device=device) * 0.5
+  env.step(action)
+
+  # With episode_length_s = one timestep, both envs time out on the first step.
+  assert seen["reset_buf"].all(), "Both envs should have reset after one step"
+
+  # Terminal action must have been captured intact at record_pre_reset.
+  assert seen["pre_reset_action"].shape == (2, 1)
+  assert torch.allclose(seen["pre_reset_action"], action), (
+    "record_pre_reset must see the applied action before _reset_idx zeroes it"
+  )
+
+  # By record_post_step the action is zeroed for all reset envs.
+  assert torch.all(seen["post_step_action"] == 0.0), (
+    "action_manager.action is zeroed for reset envs by the time record_post_step runs"
+  )
+
+  env.close()
@@ -0,0 +1,34 @@
+"""Tests for the Viser reward bar panel term cap."""
+
+from unittest.mock import Mock
+
+import pytest
+
+from mjlab.viewer.viser.reward_bar_panel import RewardBarPanel
+
+
+def _term_names(n: int) -> list[str]:
+  return [f"term_{i}" for i in range(n)]
+
+
+def test_warns_and_truncates_over_cap():
+  """More terms than the cap warns and keeps only the first ``max_terms``."""
+  server = Mock()
+  with pytest.warns(UserWarning, match="exceed max_terms"):
+    panel = RewardBarPanel(server, _term_names(21), update_dt=1 / 30, max_terms=20)
+  assert panel._term_names == _term_names(21)[:20]
+
+
+def test_no_warning_at_or_below_cap(recwarn):
+  """At or below the cap, all terms are kept and no truncation warning fires."""
+  server = Mock()
+  panel = RewardBarPanel(server, _term_names(20), update_dt=1 / 30, max_terms=20)
+  assert panel._term_names == _term_names(20)
+  assert not [w for w in recwarn if "exceed max_terms" in str(w.message)]
+
+
+def test_raised_cap_shows_all_terms():
+  """Raising ``max_terms`` (the fix for #1079) surfaces the overflow terms."""
+  server = Mock()
+  panel = RewardBarPanel(server, _term_names(24), update_dt=1 / 30, max_terms=32)
+  assert panel._term_names == _term_names(24)
@@ -0,0 +1,363 @@
+"""Tests for reward manager functionality."""
+
+from unittest.mock import Mock
+
+import mujoco
+import pytest
+import torch
+from conftest import get_test_device
+
+from mjlab.actuator import BuiltinPositionActuatorCfg
+from mjlab.entity import Entity, EntityArticulationInfoCfg, EntityCfg
+from mjlab.envs.mdp.rewards import electrical_power_cost, joint_torques_l2
+from mjlab.managers.reward_manager import RewardManager, RewardTermCfg
+from mjlab.managers.scene_entity_config import SceneEntityCfg
+from mjlab.sim.sim import Simulation, SimulationCfg
+
+PARTIALLY_ACTUATED_ROBOT_XML = """
+<mujoco>
+  <worldbody>
+    <body name="base" pos="0 0 0.5">
+      <geom name="base_geom" type="cylinder" size="0.1 0.05" mass="1.0"/>
+      <body name="link1" pos="0 0 0.1">
+        <joint name="actuated_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="0.5"/>
+        <body name="link2" pos="0 0 0.4">
+          <joint name="passive_joint" 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.3"/>
+          <body name="link3" pos="0 0 0.3">
+            <joint name="actuated_joint2" type="hinge" axis="1 0 0" range="-1.57 1.57"/>
+            <geom name="link3_geom" type="box" size="0.05 0.05 0.1" mass="0.2"/>
+          </body>
+        </body>
+      </body>
+    </body>
+  </worldbody>
+</mujoco>
+"""
+
+
+@pytest.fixture(scope="module")
+def device():
+  """Test device fixture."""
+  return get_test_device()
+
+
+class SimpleTestReward:
+  """A simple class-based reward for testing that tracks state."""
+
+  def __init__(self, cfg: RewardTermCfg, env):
+    self.num_envs = env.num_envs
+    self.device = env.device
+    self.current_air_time = torch.zeros((self.num_envs, 1), device=self.device)
+
+  def __call__(self, env, **kwargs):
+    self.current_air_time += 0.01
+    return torch.ones(env.num_envs, device=env.device)
+
+  def reset(self, env_ids: torch.Tensor | None = None, env=None):
+    if env_ids is not None and len(env_ids) > 0:
+      self.current_air_time[env_ids] = 0
+
+
+class StatelessReward:
+  """A stateless class-based reward without reset method."""
+
+  def __init__(self, cfg: RewardTermCfg, env):
+    pass
+
+  def __call__(self, env, **kwargs):
+    return torch.ones(env.num_envs)
+
+
+@pytest.fixture
+def mock_env():
+  """Create a mock environment for testing."""
+  env = Mock()
+  env.num_envs = 4
+  env.device = "cpu"
+  env.step_dt = 0.01
+  env.max_episode_length_s = 10.0
+  robot = Mock()
+  env.scene = {"robot": robot}
+  env.command_manager.get_command = Mock(
+    return_value=torch.tensor([[1.0, 0.0, 0.0]] * 4)
+  )
+  return env
+
+
+@pytest.fixture
+def class_reward_config():
+  """Config with a class-based reward."""
+  return {
+    "term": RewardTermCfg(
+      func=SimpleTestReward,
+      weight=1.0,
+      params={},
+    )
+  }
+
+
+@pytest.fixture
+def function_reward_config():
+  """Config with a function-based reward."""
+  return {
+    "term": RewardTermCfg(
+      func=lambda env: torch.ones(env.num_envs),
+      weight=1.0,
+      params={},
+    )
+  }
+
+
+@pytest.fixture
+def stateless_reward_config():
+  """Config with a stateless class-based reward."""
+  return {
+    "term": RewardTermCfg(
+      func=StatelessReward,
+      weight=1.0,
+      params={},
+    )
+  }
+
+
+def test_class_based_reward_reset(mock_env, class_reward_config):
+  """Test that class-based reward terms are tracked and have reset called."""
+  manager = RewardManager(class_reward_config, mock_env)
+  term = manager._class_term_cfgs[0].func
+
+  for _ in range(10):
+    manager.compute(dt=0.01)
+  assert (term.current_air_time > 0).all()
+
+  manager.reset(env_ids=torch.tensor([0, 2]))
+
+  # Check that only specified envs were reset.
+  assert term.current_air_time[0, 0] == 0
+  assert term.current_air_time[1, 0] > 0
+  assert term.current_air_time[2, 0] == 0
+  assert term.current_air_time[3, 0] > 0
+
+
+def test_function_based_reward_not_tracked(mock_env, function_reward_config):
+  """Test that function-based reward terms are not tracked as class terms."""
+  manager = RewardManager(function_reward_config, mock_env)
+  assert len(manager._class_term_cfgs) == 0
+
+
+def test_stateless_class_reward_no_reset(mock_env, stateless_reward_config):
+  """Test that stateless class-based rewards without reset don't break reset."""
+  manager = RewardManager(stateless_reward_config, mock_env)
+
+  # Stateless rewards without reset method should not be tracked.
+  assert len(manager._class_term_cfgs) == 0
+
+  # Reset should work without errors.
+  manager.reset(env_ids=torch.tensor([0, 2]))
+
+
+def test_electrical_power_cost_partially_actuated(device):
+  """Test electrical_power_cost on robots with passive joints.
+
+  This test verifies that:
+  1. The reward correctly identifies and uses only actuated joints (ignoring passive joints)
+  2. Power cost is computed as sum of max(0, force*velocity) for each actuator,
+     meaning only positive work is penalized (not regenerative braking)
+  """
+  entity_cfg = EntityCfg(
+    spec_fn=lambda: mujoco.MjSpec.from_string(PARTIALLY_ACTUATED_ROBOT_XML),
+    articulation=EntityArticulationInfoCfg(
+      actuators=(
+        BuiltinPositionActuatorCfg(
+          target_names_expr=("actuated_joint1", "actuated_joint2"),
+          effort_limit=10.0,
+          stiffness=100.0,
+          damping=10.0,
+        ),
+      )
+    ),
+  )
+  entity = Entity(entity_cfg)
+  model = entity.compile()
+
+  num_envs = 2
+  sim_cfg = SimulationCfg()
+  sim = Simulation(num_envs=num_envs, cfg=sim_cfg, model=model, device=device)
+  entity.initialize(model, sim.model, sim.data, device)
+
+  env = Mock()
+  env.num_envs = num_envs
+  env.device = device
+  env.scene = {"robot": entity}
+
+  asset_cfg = SceneEntityCfg(name="robot", joint_names=("actuated.*",))
+
+  reward_cfg = RewardTermCfg(
+    func=electrical_power_cost, weight=1.0, params={"asset_cfg": asset_cfg}
+  )
+
+  reward = electrical_power_cost(reward_cfg, env)
+
+  assert len(reward._joint_ids) == 2
+
+  # Test case 1: All forces and velocities aligned (all positive work).
+  # actuated_joint1 (dof 0), actuated_joint2 (dof 2).
+  # Note: dof 1 is the passive joint, not used in power calculation.
+  sim.data.qfrc_actuator[:] = 0.0
+  sim.data.qfrc_actuator[:, 0] = torch.tensor([2.0, 1.0], device=device)
+  sim.data.qfrc_actuator[:, 2] = torch.tensor([3.0, 4.0], device=device)
+  sim.data.qvel[:] = 0.0
+  sim.data.qvel[:, 0] = torch.tensor([1.0, 2.0], device=device)
+  sim.data.qvel[:, 2] = torch.tensor([2.0, 1.0], device=device)
+
+  power_cost = reward(env, asset_cfg)
+
+  # Expected: env0 = 2.0*1.0 + 3.0*2.0 = 8.0, env1 = 1.0*2.0 + 4.0*1.0 = 6.0.
+  expected = torch.tensor([8.0, 6.0], device=device)
+  assert torch.allclose(power_cost, expected)
+
+  # Test case 2: Some negative forces (regenerative braking, not penalized).
+  sim.data.qfrc_actuator[:] = 0.0
+  sim.data.qfrc_actuator[:, 0] = torch.tensor([-2.0, 1.0], device=device)
+  sim.data.qfrc_actuator[:, 2] = torch.tensor([3.0, -4.0], device=device)
+
+  power_cost = reward(env, asset_cfg)
+
+  # Expected: env0 = max(0,-2.0*1.0) + max(0,3.0*2.0) = 0 + 6.0 = 6.0.
+  #           env1 = max(0,1.0*2.0) + max(0,-4.0*1.0) = 2.0 + 0 = 2.0.
+  expected = torch.tensor([6.0, 2.0], device=device)
+  assert torch.allclose(power_cost, expected)
+
+
+def test_reward_manager_handles_nan_values(mock_env):
+  """Test that RewardManager converts NaN/Inf reward values to zero."""
+
+  def nan_reward(env):
+    """Reward function that returns NaN for some environments."""
+    r = torch.ones(env.num_envs, device=env.device)
+    r[1] = float("nan")
+    r[3] = float("inf")
+    return r
+
+  cfg = {"nan_term": RewardTermCfg(func=nan_reward, weight=1.0, params={})}
+  manager = RewardManager(cfg, mock_env)
+
+  rewards = manager.compute(dt=0.01)
+
+  # NaN and Inf should be converted to 0.
+  assert not torch.isnan(rewards).any(), "Reward buffer contains NaN"
+  assert not torch.isinf(rewards).any(), "Reward buffer contains Inf"
+  assert rewards[0] == pytest.approx(0.01)
+  assert rewards[1] == 0.0
+  assert rewards[2] == pytest.approx(0.01)
+  assert rewards[3] == 0.0
+
+
+def test_reward_manager_handles_neginf_values(mock_env):
+  """Test that RewardManager converts negative infinity to zero."""
+
+  def neginf_reward(env):
+    r = torch.ones(env.num_envs, device=env.device)
+    r[2] = float("-inf")
+    return r
+
+  cfg = {"neginf_term": RewardTermCfg(func=neginf_reward, weight=1.0, params={})}
+  manager = RewardManager(cfg, mock_env)
+
+  rewards = manager.compute(dt=0.01)
+
+  assert not torch.isinf(rewards).any()
+  assert rewards[2] == 0.0
+
+
+def test_reward_scaling_enabled(mock_env):
+  """Test that rewards are scaled by dt when scale_by_dt=True (default)."""
+
+  def constant_reward(env):
+    return torch.ones(env.num_envs, device=env.device)
+
+  cfg = {"term": RewardTermCfg(func=constant_reward, weight=2.0, params={})}
+  manager = RewardManager(cfg, mock_env, scale_by_dt=True)
+
+  dt = 0.02
+  rewards = manager.compute(dt=dt)
+
+  # With scaling: reward = raw_value (1.0) * weight (2.0) * dt (0.02) = 0.04
+  expected = 1.0 * 2.0 * dt
+  assert torch.allclose(rewards, torch.full((4,), expected))
+
+  # _step_reward should be unscaled (raw_value * weight)
+  step_reward = manager._step_reward[:, 0]
+  expected_step = 1.0 * 2.0
+  assert torch.allclose(step_reward, torch.full((4,), expected_step))
+
+
+def test_reward_scaling_disabled(mock_env):
+  """Test that rewards are not scaled by dt when scale_by_dt=False."""
+
+  def constant_reward(env):
+    return torch.ones(env.num_envs, device=env.device)
+
+  cfg = {"term": RewardTermCfg(func=constant_reward, weight=2.0, params={})}
+  manager = RewardManager(cfg, mock_env, scale_by_dt=False)
+
+  dt = 0.02
+  rewards = manager.compute(dt=dt)
+
+  # Without scaling: reward = raw_value (1.0) * weight (2.0) = 2.0
+  expected = 1.0 * 2.0
+  assert torch.allclose(rewards, torch.full((4,), expected))
+
+  # _step_reward should still be unscaled (same as reward when not scaling)
+  step_reward = manager._step_reward[:, 0]
+  assert torch.allclose(step_reward, torch.full((4,), expected))
+
+
+def test_reward_scaling_default_is_enabled(mock_env):
+  """Test that scale_by_dt defaults to True for backward compatibility."""
+
+  def constant_reward(env):
+    return torch.ones(env.num_envs, device=env.device)
+
+  cfg = {"term": RewardTermCfg(func=constant_reward, weight=1.0, params={})}
+  # Don't pass scale_by_dt - should default to True
+  manager = RewardManager(cfg, mock_env)
+
+  dt = 0.01
+  rewards = manager.compute(dt=dt)
+
+  # Default (scaling enabled): reward = 1.0 * 1.0 * 0.01 = 0.01
+  assert torch.allclose(rewards, torch.full((4,), 0.01))
+
+
+def test_reward_shape_validation_rejects_bad_compute_output(mock_env):
+  """Reward terms must return one scalar per environment."""
+  cfg = {"bad": RewardTermCfg(func=lambda env: torch.ones(env.num_envs, 1), weight=1.0)}
+  manager = RewardManager(cfg, mock_env)
+  with pytest.raises(ValueError, match="RewardManager term 'bad'.*expected \\(4,\\)"):
+    manager.compute(dt=0.01)
+
+
+def test_joint_torques_l2_with_actuator_ids(mock_env):
+  """Test that joint_torques_l2 only penalizes specified actuators."""
+  mock_env.scene["robot"].data.actuator_force = torch.tensor([[1.0, 2.0, 3.0, 4.0]] * 4)
+
+  asset_cfg = SceneEntityCfg(name="robot", actuator_ids=[0, 2])
+  result = joint_torques_l2(mock_env, asset_cfg)
+
+  # Only actuators 0 and 2: 1^2 + 3^2 = 10.0
+  expected = torch.full((4,), 10.0)
+  assert torch.allclose(result, expected)
+
+
+def test_joint_torques_l2_all_actuators(mock_env):
+  """Test that joint_torques_l2 uses all actuators by default."""
+  mock_env.scene["robot"].data.actuator_force = torch.tensor([[1.0, 2.0, 3.0]] * 4)
+
+  asset_cfg = SceneEntityCfg(name="robot")
+  result = joint_torques_l2(mock_env, asset_cfg)
+
+  # All actuators: 1^2 + 2^2 + 3^2 = 14.0
+  expected = torch.full((4,), 14.0)
+  assert torch.allclose(result, expected)
@@ -0,0 +1,346 @@
+"""Tests for RL exporter utilities."""
+
+import os
+import tempfile
+
+import mujoco
+import onnx
+import pytest
+import torch
+from conftest import get_test_device
+
+from mjlab.actuator import XmlActuatorCfg
+from mjlab.entity import EntityArticulationInfoCfg, EntityCfg
+from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg, mdp
+from mjlab.managers.observation_manager import ObservationGroupCfg, ObservationTermCfg
+from mjlab.rl.exporter_utils import (
+  attach_metadata_to_onnx,
+  get_base_metadata,
+  list_to_csv_str,
+)
+from mjlab.scene import SceneCfg
+from mjlab.sim import MujocoCfg, SimulationCfg
+from mjlab.terrains import TerrainEntityCfg
+
+
+def test_list_to_csv_str():
+  """Test CSV string conversion utility."""
+  # Test with floats.
+  result = list_to_csv_str([1.23456, 2.34567, 3.45678], decimals=3)
+  assert result == "1.235,2.346,3.457"
+
+  # Test with integers.
+  result = list_to_csv_str([1, 2, 3], decimals=2)
+  assert result == "1.00,2.00,3.00"
+
+  # Test with mixed types.
+  result = list_to_csv_str([1.5, "hello", 2.5], decimals=1)
+  assert result == "1.5,hello,2.5"
+
+  # Test custom delimiter.
+  result = list_to_csv_str([1.0, 2.0, 3.0], decimals=1, delimiter=";")
+  assert result == "1.0;2.0;3.0"
+
+  # Test nested sequence entries (e.g. per-term clip ranges or vector scales):
+  # each entry is joined with the sub-delimiter so the outer delimiter stays
+  # unambiguous and round-trips via a plain split.
+  result = list_to_csv_str([[1.0, 2.0], 3.0, [4.0, 5.0]], decimals=1)
+  assert result == "1.0;2.0,3.0,4.0;5.0"
+  entries = result.split(",")
+  assert entries == ["1.0;2.0", "3.0", "4.0;5.0"]
+  assert entries[0].split(";") == ["1.0", "2.0"]
+
+
+def test_attach_metadata_to_onnx():
+  """Test that metadata can be attached to ONNX models."""
+  # Create a dummy ONNX model.
+  with tempfile.TemporaryDirectory() as tmpdir:
+    onnx_path = os.path.join(tmpdir, "test_policy.onnx")
+
+    # Create minimal ONNX model.
+    input_tensor = onnx.helper.make_tensor_value_info(
+      "input", onnx.TensorProto.FLOAT, [1, 2]
+    )
+    output_tensor = onnx.helper.make_tensor_value_info(
+      "output", onnx.TensorProto.FLOAT, [1, 2]
+    )
+    node = onnx.helper.make_node("Identity", ["input"], ["output"])
+    graph = onnx.helper.make_graph(
+      [node], "test_graph", [input_tensor], [output_tensor]
+    )
+    model = onnx.helper.make_model(graph)
+    onnx.save(model, onnx_path)
+
+    # Attach metadata.
+    metadata = {
+      "run_path": "test/run/path",
+      "joint_names": ["joint_a", "joint_b"],
+      "joint_stiffness": [20.0, 10.0],
+      "joint_damping": [1.0, 1.0],
+      "extra_field": "extra_value",
+    }
+    attach_metadata_to_onnx(onnx_path, metadata)
+
+    # Load and verify metadata was attached.
+    loaded_model = onnx.load(onnx_path)
+    metadata_props = {prop.key: prop.value for prop in loaded_model.metadata_props}
+
+    # Check all metadata fields are present.
+    assert "run_path" in metadata_props
+    assert "joint_names" in metadata_props
+    assert "joint_stiffness" in metadata_props
+    assert "extra_field" in metadata_props
+
+    # Check values are correct.
+    assert metadata_props["run_path"] == "test/run/path"
+    assert metadata_props["extra_field"] == "extra_value"
+
+    # Check list was converted to CSV string.
+    joint_names = metadata_props["joint_names"].split(",")
+    assert len(joint_names) == 2
+    assert "joint_a" in joint_names
+    assert "joint_b" in joint_names
+
+    # Check stiffness values are in natural joint order.
+    stiffness_values = [float(x) for x in metadata_props["joint_stiffness"].split(",")]
+    assert stiffness_values == [20.0, 10.0]  # Natural order: joint_a (20), joint_b (10)
+
+
+def test_attach_metadata_to_onnx_nested_clip_and_scale():
+  """Per-term clip ranges and vector scales must round-trip through a plain
+  comma split without corrupting neighboring entries."""
+  with tempfile.TemporaryDirectory() as tmpdir:
+    onnx_path = os.path.join(tmpdir, "test_policy.onnx")
+
+    input_tensor = onnx.helper.make_tensor_value_info(
+      "input", onnx.TensorProto.FLOAT, [1, 2]
+    )
+    output_tensor = onnx.helper.make_tensor_value_info(
+      "output", onnx.TensorProto.FLOAT, [1, 2]
+    )
+    node = onnx.helper.make_node("Identity", ["input"], ["output"])
+    graph = onnx.helper.make_graph(
+      [node], "test_graph", [input_tensor], [output_tensor]
+    )
+    model = onnx.helper.make_model(graph)
+    onnx.save(model, onnx_path)
+
+    metadata = {
+      "observation_terms_clip": [
+        [float("-inf"), float("inf")],
+        [-1.0, 1.0],
+        [float("-inf"), float("inf")],
+      ],
+      "observation_terms_scale": [1.0, [0.5, 1.2, 0.3], 2.0],
+    }
+    attach_metadata_to_onnx(onnx_path, metadata)
+
+    loaded_model = onnx.load(onnx_path)
+    metadata_props = {prop.key: prop.value for prop in loaded_model.metadata_props}
+
+    clip_entries = metadata_props["observation_terms_clip"].split(",")
+    assert clip_entries == ["-inf;inf", "-1.000;1.000", "-inf;inf"]
+    assert [float(x) for x in clip_entries[1].split(";")] == [-1.0, 1.0]
+
+    scale_entries = metadata_props["observation_terms_scale"].split(",")
+    assert scale_entries == ["1.000", "0.500;1.200;0.300", "2.000"]
+    assert [float(x) for x in scale_entries[1].split(";")] == [0.5, 1.2, 0.3]
+
+
+# Robot with 2 joints but only 1 actuator (underactuated).
+ROBOT_XML_UNDERACTUATED = """
+<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="1.0"/>
+      <body name="link1" pos="0 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="0.1"/>
+      </body>
+      <body name="link2" pos="0 0 0">
+        <joint name="joint2" type="hinge" axis="0 0 1" range="-1.57 1.57"/>
+        <geom name="link2_geom" type="box" size="0.1 0.1 0.1" mass="0.1"/>
+      </body>
+    </body>
+  </worldbody>
+  <actuator>
+    <motor name="actuator1" joint="joint2" gear="1.0"/>
+  </actuator>
+</mujoco>
+"""
+
+
+@pytest.fixture(scope="module")
+def device():
+  return get_test_device()
+
+
+def test_get_base_metadata_skips_non_actuated_joints(device):
+  """get_base_metadata handles non-actuated joints without KeyError."""
+  robot_cfg = EntityCfg(
+    spec_fn=lambda: mujoco.MjSpec.from_string(ROBOT_XML_UNDERACTUATED),
+    articulation=EntityArticulationInfoCfg(
+      actuators=(XmlActuatorCfg(target_names_expr=(".*",)),)
+    ),
+  )
+
+  env_cfg = ManagerBasedRlEnvCfg(
+    scene=SceneCfg(
+      terrain=TerrainEntityCfg(terrain_type="plane"),
+      num_envs=1,
+      extent=1.0,
+      entities={"robot": robot_cfg},
+    ),
+    observations={
+      "actor": ObservationGroupCfg(
+        terms={
+          "joint_pos": ObservationTermCfg(
+            func=lambda env: env.scene["robot"].data.joint_pos,
+            history_length=5,
+            scale=2.0,
+          ),
+        },
+      ),
+    },
+    actions={
+      "joint_pos": mdp.JointPositionActionCfg(
+        entity_name="robot", actuator_names=(".*",), scale=1.0
+      )
+    },
+    sim=SimulationCfg(mujoco=MujocoCfg(timestep=0.01, iterations=1)),
+    decimation=1,
+    episode_length_s=1.0,
+  )
+
+  env = ManagerBasedRlEnv(cfg=env_cfg, device=device)
+  metadata = get_base_metadata(env, run_path="dummy/run")
+
+  robot = env.scene["robot"]
+
+  # All joints (including non-actuated) should be listed in joint_names metadata.
+  joint_names_meta = metadata["joint_names"]
+  assert isinstance(joint_names_meta, list)
+  assert joint_names_meta == list(robot.joint_names)
+  assert "joint1" in joint_names_meta
+  assert "joint2" in joint_names_meta
+
+  # Stiffness/damping are only defined for actuated joints, in natural joint order.
+  stiffness_meta = metadata["joint_stiffness"]
+  damping_meta = metadata["joint_damping"]
+  assert isinstance(stiffness_meta, list)
+  assert isinstance(damping_meta, list)
+  assert len(stiffness_meta) == len(robot.spec.actuators)
+  assert len(damping_meta) == len(robot.spec.actuators)
+
+  observation_names = metadata["observation_names"]
+  assert isinstance(observation_names, list)
+  assert "joint_pos" in observation_names
+
+  observation_terms_scale = metadata["observation_terms_scale"]
+  assert isinstance(observation_terms_scale, list)
+  assert len(observation_terms_scale) == len(observation_names)
+  assert observation_terms_scale[0] == 2.0
+
+  observation_history_length = metadata["observation_terms_history_length"]
+  assert isinstance(observation_history_length, list)
+  assert len(observation_history_length) == len(observation_names)
+  assert observation_history_length[0] == 5
+
+  observation_flatten_history_dim = metadata["observation_terms_flatten_history_dim"]
+  assert isinstance(observation_flatten_history_dim, list)
+  assert len(observation_flatten_history_dim) == len(observation_names)
+  # Default flatten_history_dim is 1.0 (True)
+  assert observation_flatten_history_dim[0] is True
+
+  observation_terms_clip = metadata["observation_terms_clip"]
+  assert isinstance(observation_terms_clip, list)
+  assert len(observation_terms_clip) == len(observation_names)
+  # Default clip is [-inf, inf]
+  assert observation_terms_clip[0] == [float("-inf"), float("inf")]
+
+  env.close()
+
+
+def test_get_base_metadata_multiple_terms_scale_clip(device):
+  """get_base_metadata preserves obs term ordering and handles scalar/vector tensor scale, tuple scale, and non-default clip."""
+  robot_cfg = EntityCfg(
+    spec_fn=lambda: mujoco.MjSpec.from_string(ROBOT_XML_UNDERACTUATED),
+    articulation=EntityArticulationInfoCfg(
+      actuators=(XmlActuatorCfg(target_names_expr=(".*",)),)
+    ),
+  )
+
+  env_cfg = ManagerBasedRlEnvCfg(
+    scene=SceneCfg(
+      terrain=TerrainEntityCfg(terrain_type="plane"),
+      num_envs=1,
+      extent=1.0,
+      entities={"robot": robot_cfg},
+    ),
+    observations={
+      "actor": ObservationGroupCfg(
+        terms={
+          "joint_pos": ObservationTermCfg(
+            func=lambda env: env.scene["robot"].data.joint_pos,
+            scale=torch.tensor(
+              2.0
+            ),  # tensor scale -> should be converted via .tolist()
+            history_length=1,
+          ),
+          "joint_vel": ObservationTermCfg(
+            func=lambda env: env.scene["robot"].data.joint_vel,
+            scale=(0.5, 1.5),  # tuple scale -> stored as-is
+            clip=(-1.0, 1.0),  # non-default clip
+            history_length=1,
+          ),
+          "joint_pos_scaled": ObservationTermCfg(
+            func=lambda env: env.scene["robot"].data.joint_pos,
+            scale=torch.tensor([1.0, 2.0]),  # per-element tensor -> list of floats
+            history_length=1,
+          ),
+        },
+      ),
+    },
+    actions={
+      "joint_pos": mdp.JointPositionActionCfg(
+        entity_name="robot", actuator_names=(".*",), scale=1.0
+      )
+    },
+    sim=SimulationCfg(mujoco=MujocoCfg(timestep=0.01, iterations=1)),
+    decimation=1,
+    episode_length_s=1.0,
+  )
+
+  env = ManagerBasedRlEnv(cfg=env_cfg, device=device)
+  metadata = get_base_metadata(env, run_path="dummy/run")
+
+  observation_names = metadata["observation_names"]
+  scales = metadata["observation_terms_scale"]
+  clips = metadata["observation_terms_clip"]
+  assert isinstance(observation_names, list)
+  assert isinstance(scales, list)
+  assert isinstance(clips, list)
+
+  # Terms appear in definition order.
+  assert observation_names == ["joint_pos", "joint_vel", "joint_pos_scaled"]
+
+  # Scalar tensor scale is unwrapped to a plain float.
+  assert scales[0] == 2.0
+
+  # Tuple scale is stored as list.
+  assert scales[1] == [0.5, 1.5]
+
+  # Per-element tensor scale is converted to a list of floats.
+  assert scales[2] == [1.0, 2.0]
+
+  # joint_pos has default clip.
+  assert clips[0] == [float("-inf"), float("inf")]
+
+  # joint_vel has explicit clip converted to list.
+  assert clips[1] == [-1.0, 1.0]
+
+  # joint_pos_scaled has default clip.
+  assert clips[2] == [float("-inf"), float("inf")]
+
+  env.close()
@@ -0,0 +1,661 @@
+"""Tests for MjlabOnPolicyRunner."""
+
+import ast
+import tempfile
+from dataclasses import asdict
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import mujoco
+import onnx
+import pytest
+import torch
+from conftest import get_test_device
+from rsl_rl.models import MLPModel
+from rsl_rl.utils import WandbLogWriter
+from rsl_rl.utils.logger import Logger
+from tensordict import TensorDict
+
+import mjlab.scripts.train as train_mod
+from mjlab.actuator import XmlActuatorCfg
+from mjlab.entity import EntityArticulationInfoCfg, EntityCfg
+from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg, mdp
+from mjlab.managers.observation_manager import ObservationGroupCfg, ObservationTermCfg
+from mjlab.rl import RslRlOnPolicyRunnerCfg, RslRlVecEnvWrapper
+from mjlab.rl.runner import MjlabOnPolicyRunner
+from mjlab.rl.spatial_softmax import SpatialSoftmaxCNNModel
+from mjlab.scene import SceneCfg
+from mjlab.sim import MujocoCfg, SimulationCfg
+from mjlab.tasks.tracking.rl.runner import _OnnxMotionModel
+from mjlab.terrains import TerrainEntityCfg
+from mjlab.utils.os import dump_yaml
+
+
+@pytest.fixture(scope="module")
+def device():
+  return get_test_device()
+
+
+@pytest.fixture(scope="module")
+def env(device):
+  robot_xml = """
+  <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="1.0"/>
+        <body name="link1" pos="0 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="0.1"/>
+        </body>
+      </body>
+    </worldbody>
+    <actuator>
+      <motor name="actuator1" joint="joint1" gear="1.0"/>
+    </actuator>
+  </mujoco>
+  """
+  robot_cfg = EntityCfg(
+    spec_fn=lambda: mujoco.MjSpec.from_string(robot_xml),
+    articulation=EntityArticulationInfoCfg(
+      actuators=(XmlActuatorCfg(target_names_expr=(".*",)),)
+    ),
+  )
+
+  env_cfg = ManagerBasedRlEnvCfg(
+    scene=SceneCfg(
+      terrain=TerrainEntityCfg(terrain_type="plane"),
+      num_envs=2,
+      extent=1.0,
+      entities={"robot": robot_cfg},
+    ),
+    observations={
+      "actor": ObservationGroupCfg(
+        terms={
+          "joint_pos": ObservationTermCfg(
+            func=lambda env: env.scene["robot"].data.joint_pos
+          ),
+        },
+      ),
+      "critic": ObservationGroupCfg(
+        terms={
+          "joint_pos": ObservationTermCfg(
+            func=lambda env: env.scene["robot"].data.joint_pos
+          ),
+        },
+      ),
+    },
+    actions={
+      "joint_pos": mdp.JointPositionActionCfg(
+        entity_name="robot", actuator_names=(".*",), scale=1.0
+      )
+    },
+    sim=SimulationCfg(mujoco=MujocoCfg(timestep=0.01, iterations=1)),
+    decimation=1,
+    episode_length_s=1.0,
+  )
+
+  env = ManagerBasedRlEnv(cfg=env_cfg, device=device)
+  yield env
+  env.close()
+
+
+def test_runner_persists_common_step_counter(env, device, monkeypatch):
+  """MjlabOnPolicyRunner should save and restore common_step_counter."""
+  wrapped_env = RslRlVecEnvWrapper(env)
+  agent_cfg = RslRlOnPolicyRunnerCfg(
+    num_steps_per_env=4, max_iterations=10, save_interval=5
+  )
+
+  with tempfile.TemporaryDirectory() as tmpdir:
+    runner = MjlabOnPolicyRunner(
+      wrapped_env, asdict(agent_cfg), log_dir=tmpdir, device=device
+    )
+    monkeypatch.setattr(runner.logger, "save_model", lambda *args, **kwargs: None)
+
+    wrapped_env.unwrapped.common_step_counter = 12345
+    checkpoint_path = str(Path(tmpdir) / "test_checkpoint.pt")
+    runner.save(checkpoint_path)
+
+    wrapped_env.unwrapped.common_step_counter = 0
+    runner.load(checkpoint_path)
+
+    assert wrapped_env.unwrapped.common_step_counter == 12345
+
+
+def test_runner_handles_old_checkpoints_without_env_state(env, device):
+  """Old checkpoints without env_state should load without crashing."""
+
+  wrapped_env = RslRlVecEnvWrapper(env)
+  agent_cfg = RslRlOnPolicyRunnerCfg(
+    num_steps_per_env=4, max_iterations=10, save_interval=5
+  )
+
+  with tempfile.TemporaryDirectory() as tmpdir:
+    runner = MjlabOnPolicyRunner(
+      wrapped_env, asdict(agent_cfg), log_dir=tmpdir, device=device
+    )
+
+    checkpoint_path = str(Path(tmpdir) / "old_checkpoint.pt")
+    old_checkpoint = {
+      "actor_state_dict": runner.alg.actor.state_dict(),
+      "critic_state_dict": runner.alg.critic.state_dict(),
+      "optimizer_state_dict": runner.alg.optimizer.state_dict(),
+      "iter": 100,
+      "infos": None,
+    }
+    torch.save(old_checkpoint, checkpoint_path)
+
+    wrapped_env.unwrapped.common_step_counter = 999
+    runner.load(checkpoint_path)
+
+    assert wrapped_env.unwrapped.common_step_counter == 999
+
+
+@pytest.mark.filterwarnings("ignore::DeprecationWarning")
+def test_export_policy_to_onnx(env, device):
+  """runner.export_policy_to_onnx() produces a valid ONNX file."""
+  wrapped_env = RslRlVecEnvWrapper(env)
+  agent_cfg = RslRlOnPolicyRunnerCfg(
+    num_steps_per_env=4, max_iterations=10, save_interval=5
+  )
+
+  with tempfile.TemporaryDirectory() as tmpdir:
+    runner = MjlabOnPolicyRunner(
+      wrapped_env, asdict(agent_cfg), log_dir=tmpdir, device=device
+    )
+    runner.export_policy_to_onnx(tmpdir, "test_policy.onnx")
+    onnx_path = Path(tmpdir) / "test_policy.onnx"
+    assert onnx_path.exists()
+    onnx.checker.check_model(str(onnx_path))
+
+
+def _make_actor(obs_dim=8, output_dim=4, obs_normalization=True):
+  obs = TensorDict({"actor": torch.zeros(1, obs_dim)})
+  obs_groups = {"actor": ["actor"]}
+  return MLPModel(
+    obs=obs,
+    obs_groups=obs_groups,
+    obs_set="actor",
+    output_dim=output_dim,
+    hidden_dims=[32, 32],
+    activation="elu",
+    obs_normalization=obs_normalization,
+  )
+
+
+def _train_normalizer(actor, n_batches=50, batch_size=64):
+  actor.train()
+  for _ in range(n_batches):
+    obs = TensorDict({"actor": torch.randn(batch_size, actor.obs_dim) * 5 + 3})
+    actor.update_normalization(obs)
+  actor.eval()
+
+
+def _model_output(actor, x_flat):
+  obs = TensorDict({"actor": x_flat})
+  with torch.no_grad():
+    return actor(obs)
+
+
+def test_onnx_export_matches_actor():
+  """as_onnx() model produces the same output as the full actor with normalization."""
+  actor = _make_actor(obs_normalization=True)
+  _train_normalizer(actor)
+  onnx_model = actor.as_onnx(verbose=False)
+  onnx_model.eval()
+  x = torch.randn(4, actor.obs_dim)
+  model_out = _model_output(actor, x)
+  with torch.no_grad():
+    onnx_out = onnx_model(x)
+  torch.testing.assert_close(model_out, onnx_out, atol=1e-6, rtol=0)
+
+
+def test_onnx_export_without_normalization():
+  """as_onnx() works when normalization is disabled."""
+  actor = _make_actor(obs_normalization=False)
+  onnx_model = actor.as_onnx(verbose=False)
+  onnx_model.eval()
+  x = torch.randn(4, actor.obs_dim)
+  with torch.no_grad():
+    out = onnx_model(x)
+  assert out.shape == (4, 4)
+
+
+@pytest.mark.filterwarnings("ignore::DeprecationWarning")
+def test_onnx_runtime_roundtrip_matches_pytorch():
+  """Exported .onnx file produces the same outputs as PyTorch via onnxruntime."""
+  ort = pytest.importorskip("onnxruntime")
+  actor = _make_actor(obs_normalization=True)
+  _train_normalizer(actor)
+  onnx_model = actor.as_onnx(verbose=False)
+  onnx_model.eval()
+
+  x = torch.randn(4, actor.obs_dim)
+  expected = _model_output(actor, x)
+
+  with tempfile.TemporaryDirectory() as tmpdir:
+    path = Path(tmpdir) / "policy.onnx"
+    torch.onnx.export(
+      onnx_model,
+      (x,),
+      str(path),
+      input_names=onnx_model.input_names,  # pyright: ignore[reportArgumentType]
+      output_names=onnx_model.output_names,  # pyright: ignore[reportArgumentType]
+      opset_version=18,
+      dynamo=False,
+    )
+    sess = ort.InferenceSession(str(path))
+    [actual] = sess.run(None, {"obs": x.numpy()})
+
+  torch.testing.assert_close(torch.from_numpy(actual), expected, atol=1e-5, rtol=0)
+
+
+# CNN (spatial-softmax) ONNX export tests.
+
+_IMG_H, _IMG_W, _IMG_C = 16, 16, 3
+_OBS_DIM_1D = 8
+_OUTPUT_DIM = 4
+
+
+def _make_cnn_actor(obs_normalization=True):
+  obs = TensorDict(
+    {
+      "actor": torch.zeros(1, _OBS_DIM_1D),
+      "camera": torch.zeros(1, _IMG_C, _IMG_H, _IMG_W),
+    }
+  )
+  obs_groups = {"actor": ["actor", "camera"]}
+  cnn_cfg = {
+    "output_channels": [8],
+    "kernel_size": [3],
+    "stride": [1],
+    "spatial_softmax_temperature": 1.0,
+  }
+  return SpatialSoftmaxCNNModel(
+    obs=obs,
+    obs_groups=obs_groups,
+    obs_set="actor",
+    output_dim=_OUTPUT_DIM,
+    cnn_cfg=cnn_cfg,
+    hidden_dims=[32, 32],
+    activation="elu",
+    obs_normalization=obs_normalization,
+  )
+
+
+def _train_cnn_normalizer(actor, n_batches=50, batch_size=64):
+  actor.train()
+  for _ in range(n_batches):
+    obs = TensorDict(
+      {
+        "actor": torch.randn(batch_size, _OBS_DIM_1D) * 5 + 3,
+        "camera": torch.randn(batch_size, _IMG_C, _IMG_H, _IMG_W),
+      }
+    )
+    actor.update_normalization(obs)
+  actor.eval()
+
+
+def _cnn_model_output(actor, x_1d, x_2d):
+  obs = TensorDict({"actor": x_1d, "camera": x_2d})
+  with torch.no_grad():
+    return actor(obs)
+
+
+def test_cnn_onnx_export_matches_actor():
+  """as_onnx() with SpatialSoftmaxCNNModel matches the original model."""
+  actor = _make_cnn_actor(obs_normalization=True)
+  _train_cnn_normalizer(actor)
+
+  onnx_model = actor.as_onnx(verbose=False)
+  onnx_model.eval()
+
+  x_1d = torch.randn(4, _OBS_DIM_1D)
+  x_2d = torch.randn(4, _IMG_C, _IMG_H, _IMG_W)
+
+  expected = _cnn_model_output(actor, x_1d, x_2d)
+  with torch.no_grad():
+    actual = onnx_model(x_1d, x_2d)
+  torch.testing.assert_close(actual, expected, atol=1e-6, rtol=0)
+
+
+@pytest.mark.filterwarnings("ignore::DeprecationWarning")
+def test_cnn_onnx_export_to_file():
+  """SpatialSoftmaxCNNModel exports to a valid ONNX file."""
+  actor = _make_cnn_actor(obs_normalization=False)
+  onnx_model = actor.as_onnx(verbose=False)
+  onnx_model.to("cpu")
+  onnx_model.eval()
+
+  with tempfile.TemporaryDirectory() as tmpdir:
+    onnx_path = Path(tmpdir) / "cnn_policy.onnx"
+    torch.onnx.export(
+      onnx_model,
+      onnx_model.get_dummy_inputs(),  # pyright: ignore[reportCallIssue]
+      str(onnx_path),
+      export_params=True,
+      opset_version=18,
+      input_names=onnx_model.input_names,  # pyright: ignore[reportArgumentType]
+      output_names=onnx_model.output_names,  # pyright: ignore[reportArgumentType]
+      dynamic_axes={},
+      dynamo=False,
+    )
+    assert onnx_path.exists()
+    onnx.checker.check_model(str(onnx_path))
+
+
+@pytest.mark.filterwarnings("ignore::DeprecationWarning")
+def test_cnn_onnx_runtime_roundtrip_matches_pytorch():
+  """Exported CNN .onnx file produces the same outputs as PyTorch via onnxruntime."""
+  ort = pytest.importorskip("onnxruntime")
+  actor = _make_cnn_actor(obs_normalization=True)
+  _train_cnn_normalizer(actor)
+
+  onnx_model = actor.as_onnx(verbose=False)
+  onnx_model.eval()
+
+  x_1d = torch.randn(4, _OBS_DIM_1D)
+  x_2d = torch.randn(4, _IMG_C, _IMG_H, _IMG_W)
+  expected = _cnn_model_output(actor, x_1d, x_2d)
+
+  with tempfile.TemporaryDirectory() as tmpdir:
+    path = Path(tmpdir) / "cnn_policy.onnx"
+    torch.onnx.export(
+      onnx_model,
+      (x_1d, x_2d),
+      str(path),
+      input_names=onnx_model.input_names,  # pyright: ignore[reportArgumentType]
+      output_names=onnx_model.output_names,  # pyright: ignore[reportArgumentType]
+      opset_version=18,
+      dynamo=False,
+    )
+    sess = ort.InferenceSession(str(path))
+    [actual] = sess.run(None, {"obs": x_1d.numpy(), "camera": x_2d.numpy()})
+
+  torch.testing.assert_close(torch.from_numpy(actual), expected, atol=1e-5, rtol=0)
+
+
+def test_get_export_paths():
+  """_get_export_paths resolves the correct dir, filename, and full path."""
+  # Normal case: "model" only appears in the checkpoint filename.
+  export_dir, filename, onnx_path = MjlabOnPolicyRunner._get_export_paths(
+    "/logs/2026-03-30_12-00-00/model_10.pt"
+  )
+  assert export_dir == Path("/logs/2026-03-30_12-00-00")
+  assert filename == "2026-03-30_12-00-00.onnx"
+  assert onnx_path == Path("/logs/2026-03-30_12-00-00/2026-03-30_12-00-00.onnx")
+
+  # Bug case: "model" also appears in a parent directory name — the old
+  # path.split("model")[0] would have truncated to "/tmp/my_".
+  export_dir, filename, onnx_path = MjlabOnPolicyRunner._get_export_paths(
+    "/tmp/my_model_experiment/2026-03-30/model_10.pt"
+  )
+  assert export_dir == Path("/tmp/my_model_experiment/2026-03-30")
+  assert filename == "2026-03-30.onnx"
+  assert onnx_path == Path("/tmp/my_model_experiment/2026-03-30/2026-03-30.onnx")
+
+
+def test_agent_cfg_serializable_after_runner_creation(env, device):
+  """dump_yaml must be called before runner creation.
+
+  The runner mutates agent_cfg in-place (e.g. resolve_symmetry_config injects
+  non-serializable objects). Verify that the train script writes config files before
+  constructing the runner.
+
+  Regression test for https://github.com/mjlab-org/mjlab/issues/764.
+  """
+  wrapped_env = RslRlVecEnvWrapper(env)
+  agent_cfg = asdict(
+    RslRlOnPolicyRunnerCfg(num_steps_per_env=4, max_iterations=10, save_interval=5)
+  )
+
+  # Dump should succeed before runner creation.
+  with tempfile.TemporaryDirectory() as tmpdir:
+    dump_yaml(Path(tmpdir) / "agent.yaml", agent_cfg)
+
+  # Create runner (mutates agent_cfg via resolve_symmetry_config).
+  with tempfile.TemporaryDirectory() as tmpdir:
+    MjlabOnPolicyRunner(wrapped_env, agent_cfg, log_dir=tmpdir, device=device)
+
+  # Confirm that the runner added non-serializable keys to agent_cfg.
+  sym_cfg = agent_cfg.get("algorithm", {}).get("symmetry_cfg")
+  runner_mutated = sym_cfg is not None or "multi_gpu" in agent_cfg
+  assert runner_mutated, "Expected runner to mutate agent_cfg"
+
+  # Verify the train script calls dump_yaml before runner_cls().
+  source = Path(train_mod.__file__).read_text()
+  tree = ast.parse(source)
+
+  dump_yaml_line = None
+  runner_cls_line = None
+  for node in ast.walk(tree):
+    if isinstance(node, ast.Call):
+      func = node.func
+      # Look for dump_yaml(..., agent_cfg)
+      if isinstance(func, ast.Name) and func.id == "dump_yaml":
+        for arg in node.args:
+          if isinstance(arg, ast.Name) and arg.id == "agent_cfg":
+            dump_yaml_line = node.lineno
+      # Look for runner_cls(...)
+      if isinstance(func, ast.Name) and func.id == "runner_cls":
+        runner_cls_line = node.lineno
+
+  assert dump_yaml_line is not None, "dump_yaml(agent_cfg) not found"
+  assert runner_cls_line is not None, "runner_cls() not found"
+  assert dump_yaml_line < runner_cls_line, (
+    f"dump_yaml (line {dump_yaml_line}) must be called before "
+    f"runner_cls (line {runner_cls_line})"
+  )
+
+
+class _MockMotion:
+  """Minimal mock of a motion object with tensor attributes."""
+
+  def __init__(self, num_steps, num_joints=12, num_bodies=5):
+    self.joint_pos = torch.randn(num_steps, num_joints)
+    self.joint_vel = torch.randn(num_steps, num_joints)
+    self.body_pos_w = torch.randn(num_steps, num_bodies, 3)
+    self.body_quat_w = torch.randn(num_steps, num_bodies, 4)
+    self.body_lin_vel_w = torch.randn(num_steps, num_bodies, 3)
+    self.body_ang_vel_w = torch.randn(num_steps, num_bodies, 3)
+
+
+def test_onnx_motion_model_policy_matches_actor():
+  """_OnnxMotionModel actions output matches calling the actor directly."""
+
+  actor = _make_actor(obs_normalization=True)
+  _train_normalizer(actor)
+  motion = _MockMotion(num_steps=50)
+
+  model = _OnnxMotionModel(actor, motion)
+  model.eval()
+
+  x = torch.randn(4, actor.obs_dim)
+  time_step = torch.tensor([[5]], dtype=torch.float32)
+
+  with torch.no_grad():
+    actions, *_ = model(x, time_step)
+  expected = _model_output(actor, x)
+  torch.testing.assert_close(actions, expected, atol=1e-6, rtol=0)
+
+
+def test_onnx_motion_model_returns_correct_motion_frame():
+  """_OnnxMotionModel returns the motion data at the requested time step."""
+
+  actor = _make_actor(obs_normalization=False)
+  motion = _MockMotion(num_steps=50)
+
+  model = _OnnxMotionModel(actor, motion)
+  model.eval()
+
+  x = torch.randn(1, actor.obs_dim)
+  t = 17
+  time_step = torch.tensor([[t]], dtype=torch.float32)
+
+  with torch.no_grad():
+    _, joint_pos, joint_vel, body_pos, body_quat, body_lin_vel, body_ang_vel = model(
+      x, time_step
+    )
+
+  torch.testing.assert_close(joint_pos, motion.joint_pos[t : t + 1])
+  torch.testing.assert_close(joint_vel, motion.joint_vel[t : t + 1])
+  torch.testing.assert_close(body_pos, motion.body_pos_w[t : t + 1])
+  torch.testing.assert_close(body_quat, motion.body_quat_w[t : t + 1])
+  torch.testing.assert_close(body_lin_vel, motion.body_lin_vel_w[t : t + 1])
+  torch.testing.assert_close(body_ang_vel, motion.body_ang_vel_w[t : t + 1])
+
+
+def test_onnx_motion_model_clamps_out_of_bounds_time_step():
+  """_OnnxMotionModel clamps time_step beyond motion length to last frame."""
+
+  num_steps = 20
+  actor = _make_actor(obs_normalization=False)
+  motion = _MockMotion(num_steps=num_steps)
+
+  model = _OnnxMotionModel(actor, motion)
+  model.eval()
+
+  x = torch.randn(1, actor.obs_dim)
+  time_step = torch.tensor([[999]], dtype=torch.float32)
+
+  with torch.no_grad():
+    _, joint_pos, *_ = model(x, time_step)
+
+  torch.testing.assert_close(joint_pos, motion.joint_pos[num_steps - 1 : num_steps])
+
+
+def _make_logger_mock(is_wandb: bool) -> MagicMock:
+  """Mock an rsl-rl Logger, specced off a real (cheap, inert) instance.
+
+  Speccing off an instance rather than the class keeps the instance attributes,
+  so an attribute rsl-rl renames or drops fails the test instead of silently
+  resolving to a truthy mock.
+  """
+  logger = Logger(
+    log_dir=None,
+    cfg={"algorithm": {}},
+    env_cfg={},
+    num_envs=1,
+    is_distributed=False,
+    gpu_world_size=1,
+    gpu_global_rank=0,
+    device="cpu",
+  )
+  mock = MagicMock(spec=logger)
+  mock.writer = MagicMock(spec=WandbLogWriter) if is_wandb else MagicMock()
+  return mock
+
+
+@pytest.mark.parametrize(
+  ("runner_module", "runner_class"),
+  [
+    ("mjlab.tasks.velocity.rl.runner", "VelocityOnPolicyRunner"),
+    ("mjlab.tasks.manipulation.rl.runner", "ManipulationOnPolicyRunner"),
+  ],
+)
+def test_task_runner_uploads_onnx_for_wandb_logger(
+  runner_module, runner_class, monkeypatch, tmp_path
+):
+  """ONNX is uploaded to W&B when the logger's writer is a WandbLogWriter."""
+  import importlib
+
+  runner_mod = importlib.import_module(runner_module)
+  runner_cls = getattr(runner_mod, runner_class)
+  runner = runner_cls.__new__(runner_cls)
+  runner.cfg = {"upload_model": True}
+  runner.logger = _make_logger_mock(is_wandb=True)
+  runner.env = MagicMock()
+  runner.export_policy_to_onnx = MagicMock()
+
+  monkeypatch.setattr(MjlabOnPolicyRunner, "save", lambda *a, **kw: None)
+  monkeypatch.setattr(runner_mod, "get_base_metadata", lambda *a: {})
+  monkeypatch.setattr(runner_mod, "attach_metadata_to_onnx", lambda *a: None)
+
+  checkpoint = tmp_path / "run-dir" / "model_100.pt"
+  checkpoint.parent.mkdir()
+  checkpoint.touch()
+  onnx_path = checkpoint.parent / f"{checkpoint.parent.name}.onnx"
+
+  mock_run = MagicMock()
+  mock_run.name = "test-run"
+  with patch.object(runner_mod, "wandb") as mock_wandb:
+    mock_wandb.run = mock_run
+    runner.save(str(checkpoint))
+
+  mock_wandb.save.assert_called_once_with(
+    str(onnx_path), base_path=str(checkpoint.parent)
+  )
+
+
+def _make_tracking_runner_shell(registry_name, is_wandb, upload_model=True):
+  """Build a MotionTrackingOnPolicyRunner with all heavy parts mocked out."""
+  from mjlab.tasks.tracking.rl.runner import MotionTrackingOnPolicyRunner
+
+  runner = MotionTrackingOnPolicyRunner.__new__(MotionTrackingOnPolicyRunner)
+  runner.registry_name = registry_name
+  runner.cfg = {"upload_model": upload_model}
+  runner.logger = _make_logger_mock(is_wandb)
+
+  mock_motion_term = MagicMock()
+  mock_motion_term.cfg.anchor_body_name = "pelvis"
+  mock_motion_term.cfg.body_names = ["body1"]
+  runner.env = MagicMock()
+  runner.env.unwrapped.command_manager.get_term.return_value = mock_motion_term
+  return runner
+
+
+def test_tracking_runner_registers_artifact_for_wandb_logger(monkeypatch, tmp_path):
+  """use_artifact is called when the logger's writer is a WandbLogWriter."""
+  from mjlab.rl.runner import MjlabOnPolicyRunner
+  from mjlab.tasks.tracking.rl import runner as runner_mod
+
+  runner = _make_tracking_runner_shell("org/motions/motion:latest", is_wandb=True)
+
+  monkeypatch.setattr(MjlabOnPolicyRunner, "save", lambda *a, **kw: None)
+  monkeypatch.setattr(runner_mod, "get_base_metadata", lambda *a: {})
+  monkeypatch.setattr(runner_mod, "attach_metadata_to_onnx", lambda *a: None)
+  monkeypatch.setattr(
+    runner.env.unwrapped.__class__,
+    "export_policy_to_onnx",
+    lambda *a, **kw: None,
+    raising=False,
+  )
+
+  checkpoint = tmp_path / "run-dir" / "model_100.pt"
+  checkpoint.parent.mkdir()
+  checkpoint.touch()
+
+  mock_run = MagicMock()
+  mock_run.name = "test-run"
+
+  with patch.object(runner_mod, "wandb") as mock_wandb:
+    mock_wandb.run = mock_run
+    runner.export_policy_to_onnx = MagicMock()
+    runner.save(str(checkpoint))
+
+  mock_run.use_artifact.assert_called_once_with("org/motions/motion:latest")
+
+
+def test_tracking_runner_does_not_register_artifact_for_tensorboard(
+  monkeypatch, tmp_path
+):
+  """use_artifact is NOT called when using the tensorboard logger."""
+  from mjlab.rl.runner import MjlabOnPolicyRunner
+  from mjlab.tasks.tracking.rl import runner as runner_mod
+
+  runner = _make_tracking_runner_shell("org/motions/motion:latest", is_wandb=False)
+
+  monkeypatch.setattr(MjlabOnPolicyRunner, "save", lambda *a, **kw: None)
+  monkeypatch.setattr(runner_mod, "get_base_metadata", lambda *a: {})
+  monkeypatch.setattr(runner_mod, "attach_metadata_to_onnx", lambda *a: None)
+
+  checkpoint = tmp_path / "run-dir" / "model_100.pt"
+  checkpoint.parent.mkdir()
+  checkpoint.touch()
+
+  with patch.object(runner_mod, "wandb") as mock_wandb:
+    runner.export_policy_to_onnx = MagicMock()
+    runner.save(str(checkpoint))
+
+  mock_wandb.run.use_artifact.assert_not_called()
@@ -0,0 +1,597 @@
+"""Tests for Scene class."""
+
+from unittest.mock import Mock
+
+import mujoco
+import mujoco_warp as mjwarp
+import pytest
+import torch
+from conftest import get_test_device, load_fixture_xml
+
+from mjlab.actuator import BuiltinPositionActuatorCfg
+from mjlab.entity import Entity, EntityArticulationInfoCfg, EntityCfg
+from mjlab.scene import Scene, SceneCfg
+from mjlab.sim.sim import Simulation, SimulationCfg
+from mjlab.sim.sim_data import WarpBridge
+
+# ============================================================================
+# Fixtures
+# ============================================================================
+
+
+@pytest.fixture
+def device():
+  """Test device fixture."""
+  return get_test_device()
+
+
+@pytest.fixture
+def simple_entity_xml():
+  """Simple entity XML for testing."""
+  return load_fixture_xml("fixed_base_box")
+
+
+@pytest.fixture
+def robot_entity_xml():
+  """Robot entity XML for testing."""
+  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="1.0"/>
+          <body name="link1" pos="0.3 0 0">
+            <joint name="joint1" type="hinge" axis="0 0 1" range="0 1.57"/>
+            <geom name="link1_geom" type="box" size="0.1 0.1 0.1" mass="0.1"/>
+          </body>
+        </body>
+      </worldbody>
+    </mujoco>
+    """
+
+
+@pytest.fixture
+def simple_entity_cfg(simple_entity_xml):
+  """Entity config for a simple box."""
+  return EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(simple_entity_xml))
+
+
+@pytest.fixture
+def robot_entity_cfg(robot_entity_xml):
+  """Entity config for a robot."""
+  return EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(robot_entity_xml))
+
+
+@pytest.fixture
+def minimal_scene_cfg():
+  """Minimal scene configuration."""
+  return SceneCfg(
+    num_envs=1,
+    env_spacing=2.0,
+  )
+
+
+@pytest.fixture
+def scene_with_entities_cfg(simple_entity_cfg, robot_entity_cfg):
+  """Scene configuration with multiple entities."""
+  return SceneCfg(
+    num_envs=4,
+    env_spacing=3.0,
+    entities={
+      "box": simple_entity_cfg,
+      "robot": robot_entity_cfg,
+    },
+  )
+
+
+@pytest.fixture
+def entity_with_site_xml():
+  """Entity XML with a site for tendon attachment."""
+  return """
+    <mujoco>
+      <worldbody>
+        <body name="box" pos="0 0 0.5">
+          <freejoint name="free"/>
+          <geom name="box_geom" type="box" size="0.1 0.1 0.1" mass="1.0"/>
+          <site name="hook" pos="0 0 0.1" size="0.01"/>
+        </body>
+      </worldbody>
+    </mujoco>
+    """
+
+
+@pytest.fixture
+def initialized_scene(scene_with_entities_cfg, device):
+  """Create an initialized scene with simulation."""
+  scene = Scene(scene_with_entities_cfg, device)
+  model = scene.compile()
+  data = mujoco.MjData(model)
+  mujoco.mj_resetData(model, data)
+
+  wp_model = mjwarp.put_model(model)
+  wp_data = mjwarp.put_data(model, data, nworld=scene.num_envs)
+  wp_model = WarpBridge(wp_model, nworld=scene.num_envs)
+  wp_data = WarpBridge(wp_data)
+
+  scene.initialize(model, wp_model, wp_data)  # type: ignore
+  return scene, wp_data
+
+
+@pytest.fixture
+def mock_entities():
+  """Create mock entities for testing."""
+  mock_box = Mock(spec=Entity)
+  mock_robot = Mock(spec=Entity)
+  return {"box": mock_box, "robot": mock_robot}
+
+
+# ============================================================================
+# Basic Scene Tests
+# ============================================================================
+
+
+def test_minimal_scene_creation(minimal_scene_cfg, device):
+  """Test creating a minimal scene with no entities."""
+  scene = Scene(minimal_scene_cfg, device)
+
+  assert scene.num_envs == 1
+  assert scene.env_spacing == 2.0
+  assert len(scene.entities) == 0
+  assert scene.terrain is None
+
+
+def test_scene_with_entities(scene_with_entities_cfg, device):
+  """Test creating a scene with multiple entities."""
+  scene = Scene(scene_with_entities_cfg, device)
+
+  assert scene.num_envs == 4
+  assert scene.env_spacing == 3.0
+  assert len(scene.entities) == 2
+  assert "box" in scene.entities
+  assert "robot" in scene.entities
+  assert isinstance(scene.entities["box"], Entity)
+  assert isinstance(scene.entities["robot"], Entity)
+
+
+# ============================================================================
+# Scene Compilation Tests
+# ============================================================================
+
+
+def test_compile_empty_scene(minimal_scene_cfg, device):
+  """Test compiling an empty scene."""
+  scene = Scene(minimal_scene_cfg, device=device)
+  model = scene.compile()
+
+  assert isinstance(model, mujoco.MjModel)
+  assert model.nbody == 1
+  assert model.nq == model.nv == 0
+
+
+def test_compile_scene_with_entities(scene_with_entities_cfg, device):
+  """Test compiling a scene with entities."""
+  scene = Scene(scene_with_entities_cfg, device)
+  model = scene.compile()
+
+  assert isinstance(model, mujoco.MjModel)
+  # Should have world + entity bodies.
+  assert model.nbody > 1
+  # Check that entity names are prefixed
+  body_names = [model.body(i).name for i in range(model.nbody)]
+  assert any("box/" in name for name in body_names)
+  assert any("robot/" in name for name in body_names)
+
+
+def test_write_zip(minimal_scene_cfg, tmp_path, device):
+  """Test exporting scene to zip file."""
+  scene = Scene(minimal_scene_cfg, device)
+  out = tmp_path / "scene_pkg"
+
+  scene.write(out, zip=True)
+  assert out.with_suffix(".zip").exists()
+
+
+def test_write_skips_unreferenced_assets(minimal_scene_cfg, tmp_path, device):
+  """write() only includes assets referenced in the generated XML."""
+  scene = Scene(minimal_scene_cfg, device)
+  scene._spec.assets["unused_mesh.stl"] = b"fake"
+  out = tmp_path / "out"
+  scene.write(out)
+  assert (out / "scene.xml").exists()
+  assets_dir = out / "assets"
+  asset_files = list(assets_dir.rglob("*")) if assets_dir.exists() else []
+  assert not any(f.is_file() for f in asset_files)
+
+
+def test_write_no_traversal_escape(minimal_scene_cfg, tmp_path, device):
+  """Asset keys with path traversal must not escape the output directory."""
+  scene = Scene(minimal_scene_cfg, device)
+  scene._spec.assets["../../assets/robot/mesh.stl"] = b"fake"
+  out = tmp_path / "subdir" / "out"
+  scene.write(out)
+  for f in tmp_path.rglob("*"):
+    if f.is_file() and f.name != "scene.xml":
+      assert str(f).startswith(str(out)), f"File escaped output dir: {f}"
+
+
+# ============================================================================
+# Entity Access Tests
+# ============================================================================
+
+
+def test_entity_dict_access(scene_with_entities_cfg, device):
+  """Test accessing entities through dictionary."""
+  scene = Scene(scene_with_entities_cfg, device)
+
+  box = scene.entities["box"]
+  robot = scene.entities["robot"]
+
+  assert isinstance(box, Entity)
+  assert isinstance(robot, Entity)
+  assert box.is_fixed_base
+  assert not robot.is_fixed_base
+
+
+def test_entity_getitem_access(scene_with_entities_cfg, device):
+  """Test accessing entities through __getitem__."""
+  scene = Scene(scene_with_entities_cfg, device)
+
+  box = scene["box"]
+  robot = scene["robot"]
+
+  assert isinstance(box, Entity)
+  assert isinstance(robot, Entity)
+
+
+def test_invalid_entity_access(scene_with_entities_cfg, device):
+  """Test accessing non-existent entity raises KeyError."""
+  scene = Scene(scene_with_entities_cfg, device)
+
+  with pytest.raises(KeyError, match="Scene element 'invalid' not found"):
+    _ = scene["invalid"]
+
+
+# ============================================================================
+# Scene Initialization Tests
+# ============================================================================
+
+
+def test_scene_initialize(initialized_scene, device):
+  """Test that scene initialization sets up entities."""
+  scene, _ = initialized_scene
+
+  # Check default env origins are set.
+  assert scene._default_env_origins is not None
+  assert scene._default_env_origins.shape == (4, 3)  # 4 envs, 3D positions.
+  assert scene._default_env_origins.device.type == device.split(":")[0]
+
+  # Check entities are initialized.
+  for entity in scene.entities.values():
+    assert hasattr(entity, "data")
+    assert entity.data is not None
+
+
+def test_env_origins_without_terrain(initialized_scene):
+  """Test env_origins property without terrain."""
+  scene, _ = initialized_scene
+
+  origins = scene.env_origins
+  assert origins.shape == (4, 3)
+  assert torch.all(origins == 0)  # Default origins should be zeros.
+
+
+# ============================================================================
+# Scene Operations Tests
+# ============================================================================
+
+
+def test_scene_reset(minimal_scene_cfg, mock_entities, device):
+  """Test that reset calls reset on all entities."""
+  scene = Scene(minimal_scene_cfg, device)
+  scene._entities = mock_entities
+
+  # Reset all environments
+  scene.reset()
+  for entity in mock_entities.values():
+    entity.reset.assert_called_once_with(None)
+
+  # Reset specific environments
+  for entity in mock_entities.values():
+    entity.reset.reset_mock()
+
+  env_ids = torch.tensor([0, 2])
+  scene.reset(env_ids)
+  for entity in mock_entities.values():
+    entity.reset.assert_called_once_with(env_ids)
+
+
+def test_scene_update(minimal_scene_cfg, mock_entities, device):
+  """Test that update calls update on all entities."""
+  scene = Scene(minimal_scene_cfg, device)
+  scene._entities = mock_entities
+
+  dt = 0.01
+  scene.update(dt)
+
+  for entity in mock_entities.values():
+    entity.update.assert_called_once_with(dt)
+
+
+def test_scene_write_data_to_sim(minimal_scene_cfg, mock_entities, device):
+  """Test that write_data_to_sim calls the method on all entities."""
+  scene = Scene(minimal_scene_cfg, device)
+  scene._entities = mock_entities
+
+  scene.write_data_to_sim()
+
+  for entity in mock_entities.values():
+    entity.write_data_to_sim.assert_called_once()
+
+
+# ============================================================================
+# Integration Tests
+# ============================================================================
+
+
+def test_full_scene_lifecycle(robot_entity_cfg, device, tmp_path):
+  """Test complete scene lifecycle from creation to simulation."""
+  scene_cfg = SceneCfg(
+    num_envs=3,
+    env_spacing=2.5,
+    entities={
+      "robot1": robot_entity_cfg,
+      "robot2": robot_entity_cfg,
+    },
+  )
+
+  scene = Scene(scene_cfg, device)
+
+  assert scene.num_envs == 3
+  assert len(scene.entities) == 2
+
+  model = scene.compile()
+  data = mujoco.MjData(model)
+  mujoco.mj_resetData(model, data)
+
+  wp_model = mjwarp.put_model(model)
+  wp_data = mjwarp.put_data(model, data, nworld=scene.num_envs)
+  wp_model = WarpBridge(wp_model, nworld=scene.num_envs)
+  wp_data = WarpBridge(wp_data)
+
+  scene.initialize(model, wp_model, wp_data)  # type: ignore
+
+  scene.reset()
+  scene.update(0.01)
+  scene.write_data_to_sim()
+
+  scene.reset(env_ids=torch.tensor([0, 2]))
+
+  out = tmp_path / "test_scene_pkg"
+  scene.write(out, zip=True)
+  assert out.with_suffix(".zip").exists()
+
+  for entity in scene.entities.values():
+    assert entity.data is not None
+    if not entity.is_fixed_base:
+      assert entity.data.root_link_pose_w.shape == (3, 7)
+
+
+# ============================================================================
+# Scene spec_fn Tests
+# ============================================================================
+
+
+def test_scene_spec_fn_adds_site(device):
+  """Test that spec_fn can add elements to the scene spec."""
+
+  def add_custom_site(spec: mujoco.MjSpec) -> None:
+    spec.worldbody.add_site(name="custom_site", pos=(1, 2, 3))
+
+  cfg = SceneCfg(spec_fn=add_custom_site)
+  scene = Scene(cfg, device)
+  model = scene.compile()
+
+  site_id = model.site("custom_site").id
+  assert site_id >= 0
+  assert tuple(model.site_pos[site_id]) == (1.0, 2.0, 3.0)
+
+
+def test_scene_spec_fn_cross_entity_tendon(entity_with_site_xml, device):
+  """Test that spec_fn can create tendons between two entities."""
+  entity_cfg = EntityCfg(
+    spec_fn=lambda: mujoco.MjSpec.from_string(entity_with_site_xml)
+  )
+
+  def add_tendon_between_entities(spec: mujoco.MjSpec) -> None:
+    # Add a world anchor site.
+    spec.worldbody.add_site(name="anchor", pos=(0, 0, 2), size=(0.01,) * 3)
+    # Create tendon from anchor to entity A.
+    tendon_a = spec.add_tendon(name="rope_a", width=0.005)
+    tendon_a.wrap_site("anchor")
+    tendon_a.wrap_site("entity_a/hook")
+    # Create tendon from anchor to entity B.
+    tendon_b = spec.add_tendon(name="rope_b", width=0.005)
+    tendon_b.wrap_site("anchor")
+    tendon_b.wrap_site("entity_b/hook")
+
+  cfg = SceneCfg(
+    entities={"entity_a": entity_cfg, "entity_b": entity_cfg},
+    spec_fn=add_tendon_between_entities,
+  )
+  scene = Scene(cfg, device)
+  model = scene.compile()
+
+  # Verify tendons exist.
+  assert model.tendon("rope_a").id >= 0
+  assert model.tendon("rope_b").id >= 0
+  # Verify sites are referenced correctly (2 wraps per tendon).
+  assert model.ntendon == 2
+  assert model.nwrap == 4
+
+
+# ============================================================================
+# Keyframe Merging Tests
+# ============================================================================
+
+
+@pytest.fixture
+def floating_box_cfg():
+  """Entity config for a floating box with initial position."""
+  xml = """
+    <mujoco>
+      <worldbody>
+        <body name="box">
+          <freejoint name="box_joint"/>
+          <geom type="box" size="0.1 0.1 0.1"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  """
+  return EntityCfg(
+    init_state=EntityCfg.InitialStateCfg(pos=(1.0, 2.0, 3.0)),
+    spec_fn=lambda: mujoco.MjSpec.from_string(xml),
+  )
+
+
+@pytest.fixture
+def floating_sphere_cfg():
+  """Entity config for a floating sphere with initial position."""
+  xml = """
+    <mujoco>
+      <worldbody>
+        <body name="sphere">
+          <freejoint name="sphere_joint"/>
+          <geom type="sphere" size="0.1"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  """
+  return EntityCfg(
+    init_state=EntityCfg.InitialStateCfg(pos=(4.0, 5.0, 6.0)),
+    spec_fn=lambda: mujoco.MjSpec.from_string(xml),
+  )
+
+
+def test_single_entity_keyframe(floating_box_cfg, device):
+  """Test that a single entity produces one merged keyframe."""
+  cfg = SceneCfg(entities={"box": floating_box_cfg})
+  scene = Scene(cfg, device)
+  model = scene.compile()
+
+  assert model.nkey == 1
+  assert model.key(0).name == "init_state"
+  # qpos: [x, y, z, qw, qx, qy, qz]
+  assert tuple(model.key(0).qpos[:3]) == (1.0, 2.0, 3.0)
+
+
+def test_multiple_entities_merged_keyframe(
+  floating_box_cfg, floating_sphere_cfg, device
+):
+  """Test that multiple entities produce a single merged keyframe."""
+  cfg = SceneCfg(entities={"box": floating_box_cfg, "sphere": floating_sphere_cfg})
+  scene = Scene(cfg, device)
+  model = scene.compile()
+
+  assert model.nkey == 1
+  assert model.key(0).name == "init_state"
+  # Box qpos (0-6), sphere qpos (7-13).
+  qpos = model.key(0).qpos
+  assert tuple(qpos[:3]) == (1.0, 2.0, 3.0)  # box position
+  assert tuple(qpos[7:10]) == (4.0, 5.0, 6.0)  # sphere position
+
+
+# ============================================================================
+# Multi-Entity Actuator Tests
+# ============================================================================
+
+
+def test_two_actuated_entities_write_ctrl(device):
+  """Test that two identical actuated entities write controls to the correct global positions."""
+
+  robot_xml = load_fixture_xml("floating_base_articulated")
+
+  entity_cfg = EntityCfg(
+    spec_fn=lambda: mujoco.MjSpec.from_string(robot_xml),
+    articulation=EntityArticulationInfoCfg(
+      actuators=(
+        BuiltinPositionActuatorCfg(
+          target_names_expr=("joint.*",),
+          effort_limit=100.0,
+          stiffness=80.0,
+          damping=10.0,
+        ),
+      )
+    ),
+  )
+
+  num_envs = 2
+  scene_cfg = SceneCfg(
+    num_envs=num_envs,
+    env_spacing=3.0,
+    entities={"robot_a": entity_cfg, "robot_b": entity_cfg},
+  )
+
+  scene = Scene(scene_cfg, device)
+  model = scene.compile()
+  sim = Simulation(num_envs=num_envs, cfg=SimulationCfg(), model=model, device=device)
+  scene.initialize(model, sim.model, sim.data)
+
+  robot_a = scene["robot_a"]
+  robot_b = scene["robot_b"]
+  assert isinstance(robot_a, Entity)
+  assert isinstance(robot_b, Entity)
+
+  # Set different position targets for each entity.
+  target_a = torch.tensor([[0.1, 0.2]], device=device).expand(num_envs, -1)
+  target_b = torch.tensor([[0.5, 0.6]], device=device).expand(num_envs, -1)
+
+  robot_a.set_joint_position_target(target_a)
+  robot_a.set_joint_velocity_target(torch.zeros(num_envs, 2, device=device))
+  robot_a.set_joint_effort_target(torch.zeros(num_envs, 2, device=device))
+
+  robot_b.set_joint_position_target(target_b)
+  robot_b.set_joint_velocity_target(torch.zeros(num_envs, 2, device=device))
+  robot_b.set_joint_effort_target(torch.zeros(num_envs, 2, device=device))
+
+  scene.write_data_to_sim()
+
+  # Verify that each entity's controls landed in the correct global ctrl positions.
+  global_ctrl_a = robot_a.indexing.ctrl_ids
+  global_ctrl_b = robot_b.indexing.ctrl_ids
+
+  # The two entities should have different global ctrl ranges.
+  assert not torch.equal(global_ctrl_a, global_ctrl_b)
+
+  # Check that the ctrl values in global positions match each entity's expected output.
+  ctrl_a = sim.data.ctrl[0, global_ctrl_a]
+  ctrl_b = sim.data.ctrl[0, global_ctrl_b]
+
+  # Controls should differ since targets differ.
+  assert not torch.allclose(ctrl_a, ctrl_b)
+
+
+def test_entity_with_option_flags_warns(device):
+  """Entity XML <option> flags are not propagated by MjSpec.attach()."""
+  xml_with_option = """
+    <mujoco>
+      <option>
+        <flag contact="disable" />
+      </option>
+      <worldbody>
+        <body name="box" pos="0 0 0.5">
+          <freejoint name="free"/>
+          <geom name="box_geom" type="box" size="0.1 0.1 0.1" mass="1.0"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  """
+  cfg = SceneCfg(
+    num_envs=1,
+    entities={
+      "box": EntityCfg(
+        spec_fn=lambda: mujoco.MjSpec.from_string(xml_with_option),
+      ),
+    },
+  )
+  with pytest.warns(UserWarning, match="disableflags"):
+    Scene(cfg, device)
@@ -0,0 +1,207 @@
+"""Unit tests for SceneEntityCfg resolution logic."""
+
+from dataclasses import dataclass
+
+import pytest
+
+from mjlab.managers.scene_entity_config import SceneEntityCfg
+
+
+@dataclass
+class _FakeEntity:
+  name: str
+
+  joint_names: tuple[str, ...]
+  body_names: tuple[str, ...]
+  geom_names: tuple[str, ...]
+  site_names: tuple[str, ...]
+
+  @property
+  def num_joints(self) -> int:
+    return len(self.joint_names)
+
+  @property
+  def num_bodies(self) -> int:
+    return len(self.body_names)
+
+  @property
+  def num_geoms(self) -> int:
+    return len(self.geom_names)
+
+  @property
+  def num_sites(self) -> int:
+    return len(self.site_names)
+
+  # find_* helpers return (ids, names) similar to Entity API.
+  def _find(
+    self,
+    query_names: tuple[str, ...],
+    pool: tuple[str, ...],
+    preserve_order: bool,
+  ) -> tuple[list[int], list[str]]:
+    # Treat query as exact names (no regex) to keep tests minimal.
+    indices = [list(pool).index(n) for n in query_names]
+    if not preserve_order:
+      indices.sort()
+    names = [list(pool)[i] for i in indices]
+    return indices, names
+
+  def find_joints(self, query_names: tuple[str, ...], preserve_order: bool = False):
+    return self._find(query_names, self.joint_names, preserve_order)
+
+  def find_bodies(self, query_names: tuple[str, ...], preserve_order: bool = False):
+    return self._find(query_names, self.body_names, preserve_order)
+
+  def find_geoms(self, query_names: tuple[str, ...], preserve_order: bool = False):
+    return self._find(query_names, self.geom_names, preserve_order)
+
+  def find_sites(self, query_names: tuple[str, ...], preserve_order: bool = False):
+    return self._find(query_names, self.site_names, preserve_order)
+
+
+@pytest.fixture
+def fake_entity() -> _FakeEntity:
+  names = ("a", "b", "c")
+  return _FakeEntity(
+    name="robot",
+    joint_names=names,
+    body_names=names,
+    geom_names=names,
+    site_names=names,
+  )
+
+
+@pytest.fixture
+def fake_scene(fake_entity: _FakeEntity):
+  # Minimal scene-like mapping.
+  return {fake_entity.name: fake_entity}
+
+
+@pytest.mark.parametrize(
+  "field_names",
+  [
+    ("joint_names", "joint_ids"),
+    ("body_names", "body_ids"),
+    ("geom_names", "geom_ids"),
+    ("site_names", "site_ids"),
+  ],
+)
+def test_names_to_ids_sets_slice_when_all(fake_scene, fake_entity, field_names):
+  """Providing full set of names resolves ids and collapses to slice(None)."""
+  names_attr, ids_attr = field_names
+
+  cfg = SceneEntityCfg(name=fake_entity.name)
+  setattr(cfg, names_attr, getattr(fake_entity, names_attr))
+
+  cfg.resolve(fake_scene)
+
+  ids_value = getattr(cfg, ids_attr)
+  assert isinstance(ids_value, slice), f"{ids_attr} should collapse to slice(None)"
+  assert ids_value == slice(None)
+
+
+@pytest.mark.parametrize(
+  "field_names,ids",
+  [
+    (("joint_names", "joint_ids"), [0, 2]),
+    (("body_names", "body_ids"), [1]),
+    (("geom_names", "geom_ids"), [2, 0]),
+    (("site_names", "site_ids"), [1, 2]),
+  ],
+)
+def test_ids_to_names_resolves_names(fake_scene, fake_entity, field_names, ids):
+  """Providing explicit ids populates the corresponding names list."""
+  names_attr, ids_attr = field_names
+
+  cfg = SceneEntityCfg(name=fake_entity.name)
+  setattr(cfg, ids_attr, ids.copy())
+
+  cfg.resolve(fake_scene)
+
+  names_value = getattr(cfg, names_attr)
+  expected = [getattr(fake_entity, names_attr)[i] for i in ids]
+  assert names_value == expected
+
+
+@pytest.mark.parametrize(
+  "field_names,provided_names,provided_ids",
+  [
+    (("joint_names", "joint_ids"), ["a"], [1]),
+    (("body_names", "body_ids"), ["b"], [2]),
+    (("geom_names", "geom_ids"), ["c"], [0]),
+    (("site_names", "site_ids"), ["a"], [2]),
+  ],
+)
+def test_inconsistent_names_and_ids_raise(
+  fake_scene, field_names, provided_names, provided_ids
+):
+  """When both names and ids are provided but disagree, a ValueError is raised."""
+  names_attr, ids_attr = field_names
+
+  cfg = SceneEntityCfg(name="robot")
+  setattr(cfg, names_attr, provided_names.copy())
+  setattr(cfg, ids_attr, provided_ids.copy())  # Must be list to trigger check.
+
+  with pytest.raises(ValueError):
+    cfg.resolve(fake_scene)
+
+
+def test_tuple_names_with_consistent_ids(fake_scene):
+  """Tuple names with matching IDs should validate without error."""
+  cfg = SceneEntityCfg(name="robot")
+  cfg.joint_names = ("a", "b")
+  cfg.joint_ids = [0, 1]
+
+  cfg.resolve(fake_scene)
+
+  assert cfg.joint_names == ["a", "b"]
+  assert cfg.joint_ids == [0, 1]
+
+
+@pytest.mark.parametrize(
+  "field_names",
+  [
+    ("joint_names", "joint_ids"),
+    ("body_names", "body_ids"),
+    ("geom_names", "geom_ids"),
+    ("site_names", "site_ids"),
+  ],
+)
+def test_names_reordered_to_match_ids_when_not_preserving_order(
+  fake_scene, field_names
+):
+  """With preserve_order=False, names are reordered to internal order."""
+  names_attr, ids_attr = field_names
+
+  # Give a subset in reverse of internal order.
+  cfg = SceneEntityCfg(name="robot")
+  setattr(cfg, names_attr, ["c", "a"])
+
+  cfg.resolve(fake_scene)
+
+  # IDs and names must both be in internal order: a=0, c=2.
+  assert getattr(cfg, ids_attr) == [0, 2]
+  assert getattr(cfg, names_attr) == ["a", "c"]
+
+
+@pytest.mark.parametrize(
+  "field_names",
+  [
+    ("joint_names", "joint_ids"),
+    ("body_names", "body_ids"),
+    ("geom_names", "geom_ids"),
+    ("site_names", "site_ids"),
+  ],
+)
+def test_names_preserve_user_order_when_preserving_order(fake_scene, field_names):
+  """With preserve_order=True, names stay in the user-specified order."""
+  names_attr, ids_attr = field_names
+
+  cfg = SceneEntityCfg(name="robot", preserve_order=True)
+  setattr(cfg, names_attr, ["c", "a"])
+
+  cfg.resolve(fake_scene)
+
+  # User order preserved: c=2, a=0.
+  assert getattr(cfg, ids_attr) == [2, 0]
+  assert getattr(cfg, names_attr) == ["c", "a"]
@@ -0,0 +1,250 @@
+"""Tests for sensor caching behavior."""
+
+from __future__ import annotations
+
+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 import GridPatternCfg, ObjRef, RayCastSensorCfg
+from mjlab.sensor.contact_sensor import ContactMatch, ContactSensorCfg
+from mjlab.sim.sim import Simulation, SimulationCfg
+
+
+@pytest.fixture(scope="module")
+def device():
+  """Test device fixture."""
+  return get_test_device()
+
+
+FALLING_BOX_XML = """
+<mujoco>
+  <worldbody>
+    <body name="ground" pos="0 0 0">
+      <geom name="ground_geom" type="plane" size="5 5 0.1" rgba="0.5 0.5 0.5 1"/>
+    </body>
+    <body name="box" pos="0 0 0.5">
+      <freejoint name="box_joint"/>
+      <geom name="box_geom" type="box" size="0.1 0.1 0.1" rgba="0.8 0.3 0.3 1"
+        mass="1.0"/>
+    </body>
+  </worldbody>
+</mujoco>
+"""
+
+
+ROBOT_WITH_FLOOR_XML = """
+<mujoco>
+  <worldbody>
+    <geom name="floor" type="plane" size="10 10 0.1" pos="0 0 0"/>
+    <body name="base" pos="0 0 2">
+      <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.1"/>
+    </body>
+  </worldbody>
+</mujoco>
+"""
+
+
+def create_contact_sensor_scene(device: str, num_envs: int = 2):
+  """Create a scene with a contact sensor."""
+  entity_cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(FALLING_BOX_XML))
+
+  contact_sensor_cfg = ContactSensorCfg(
+    name="box_contact",
+    primary=ContactMatch(mode="geom", pattern="box_geom", entity="box"),
+    secondary=None,
+    fields=("found", "force"),
+  )
+
+  scene_cfg = SceneCfg(
+    num_envs=num_envs,
+    env_spacing=3.0,
+    entities={"box": entity_cfg},
+    sensors=(contact_sensor_cfg,),
+  )
+
+  scene = Scene(scene_cfg, device)
+  model = scene.compile()
+  sim_cfg = SimulationCfg(njmax=75)
+  sim = Simulation(num_envs=num_envs, cfg=sim_cfg, model=model, device=device)
+  scene.initialize(sim.mj_model, sim.model, sim.data)
+
+  return scene, sim
+
+
+def create_raycast_sensor_scene(device: str, num_envs: int = 2):
+  """Create a scene with a raycast sensor."""
+  entity_cfg = EntityCfg(
+    spec_fn=lambda: mujoco.MjSpec.from_string(ROBOT_WITH_FLOOR_XML)
+  )
+
+  raycast_cfg = RayCastSensorCfg(
+    name="terrain_scan",
+    frame=ObjRef(type="body", name="base", entity="robot"),
+    pattern=GridPatternCfg(
+      size=(0.5, 0.5), resolution=0.25, direction=(0.0, 0.0, -1.0)
+    ),
+    max_distance=10.0,
+  )
+
+  scene_cfg = SceneCfg(
+    num_envs=num_envs,
+    env_spacing=5.0,
+    entities={"robot": entity_cfg},
+    sensors=(raycast_cfg,),
+  )
+
+  scene = Scene(scene_cfg, device)
+  model = scene.compile()
+  sim_cfg = SimulationCfg(njmax=20)
+  sim = Simulation(num_envs=num_envs, cfg=sim_cfg, model=model, device=device)
+  scene.initialize(sim.mj_model, sim.model, sim.data)
+  if scene.sensor_context is not None:
+    sim.set_sensor_context(scene.sensor_context)
+
+  return scene, sim
+
+
+def test_cache_invalidated_by_update(device):
+  """Verify that update() invalidates the cache."""
+  scene, sim = create_contact_sensor_scene(device)
+  sensor = scene["box_contact"]
+
+  sim.step()
+  scene.update(dt=sim.cfg.mujoco.timestep)
+
+  # Access data to populate cache.
+  data1 = sensor.data
+
+  # Call update to invalidate.
+  sensor.update(dt=sim.cfg.mujoco.timestep)
+
+  # Access data again - should be a new object.
+  data2 = sensor.data
+
+  # Should be different objects (cache was invalidated and recomputed).
+  assert data1 is not data2
+
+
+def test_cache_invalidated_by_reset(device):
+  """Verify that reset() invalidates the cache."""
+  scene, sim = create_contact_sensor_scene(device)
+  sensor = scene["box_contact"]
+
+  sim.step()
+  scene.update(dt=sim.cfg.mujoco.timestep)
+
+  # Access data to populate cache.
+  data1 = sensor.data
+
+  # Call reset to invalidate.
+  sensor.reset(None)
+
+  # Access data again - should be a new object.
+  data2 = sensor.data
+
+  # Should be different objects (cache was invalidated and recomputed).
+  assert data1 is not data2
+
+
+def test_compute_data_called_once_per_cache_period(device):
+  """Verify _compute_data is only called once when cache is valid."""
+  scene, sim = create_contact_sensor_scene(device)
+  sensor = scene["box_contact"]
+
+  sim.step()
+  scene.update(dt=sim.cfg.mujoco.timestep)
+
+  # Patch _compute_data to count calls.
+  original_compute = sensor._compute_data
+  call_count = [0]
+
+  def counting_compute():
+    call_count[0] += 1
+    return original_compute()
+
+  sensor._compute_data = counting_compute
+
+  # Access data 5 times without invalidating cache.
+  for _ in range(5):
+    _ = sensor.data
+
+  # _compute_data should have been called only once.
+  assert call_count[0] == 1
+
+
+def test_data_reflects_physics_after_step_and_update(device):
+  """Verify sensor data reflects physics state after step + update cycle."""
+  scene, sim = create_raycast_sensor_scene(device)
+  sensor = scene["terrain_scan"]
+  robot_entity = scene["robot"]
+
+  # Initial position at z=2.
+  sim.step()
+  scene.update(dt=sim.cfg.mujoco.timestep)
+  sim.sense()
+  data1 = sensor.data
+
+  # All rays should hit floor ~2m away.
+  assert torch.allclose(
+    data1.distances, torch.full_like(data1.distances, 2.0), atol=0.1
+  )
+
+  # Move robot higher.
+  root_state = torch.zeros((2, 13), device=device)
+  root_state[:, 2] = 4.0  # Move to z=4.
+  root_state[:, 3] = 1.0  # Unit quaternion.
+  robot_entity.write_root_state_to_sim(root_state)
+
+  sim.step()
+  scene.update(dt=sim.cfg.mujoco.timestep)
+  sim.sense()
+
+  data2 = sensor.data
+
+  # Now rays should hit floor ~4m away.
+  assert torch.allclose(
+    data2.distances, torch.full_like(data2.distances, 4.0), atol=0.1
+  )
+
+  # Data objects should be different (cache was invalidated).
+  assert data1 is not data2
+
+
+def test_stale_cache_without_update(device):
+  """Verify that without update(), cache returns stale data."""
+  scene, sim = create_raycast_sensor_scene(device)
+  sensor = scene["terrain_scan"]
+  robot_entity = scene["robot"]
+
+  # Initial step + update.
+  sim.step()
+  scene.update(dt=sim.cfg.mujoco.timestep)
+  sim.sense()
+
+  # Access data to populate cache.
+  data1 = sensor.data
+  initial_distances = data1.distances.clone()
+
+  # Move robot higher.
+  root_state = torch.zeros((2, 13), device=device)
+  root_state[:, 2] = 4.0  # Move to z=4.
+  root_state[:, 3] = 1.0  # Unit quaternion.
+  robot_entity.write_root_state_to_sim(root_state)
+
+  # Step but DON'T call update - cache should NOT be invalidated.
+  sim.step()
+
+  # Access data - should still be cached (stale).
+  data2 = sensor.data
+
+  # Data should be the same object (cache was not invalidated).
+  assert data1 is data2
+
+  # Values should be the same (stale).
+  assert torch.allclose(data2.distances, initial_distances)
@@ -0,0 +1,188 @@
+"""Tests for sim.py."""
+
+import mujoco
+import mujoco_warp as mjwarp
+import numpy as np
+import pytest
+import torch
+from conftest import get_test_device
+
+from mjlab.sim import MujocoCfg, Simulation, SimulationCfg
+
+
+@pytest.fixture
+def device():
+  """Test device fixture."""
+  return get_test_device()
+
+
+@pytest.fixture
+def robot_xml():
+  """Simple robot with geoms and joints."""
+  return """
+    <mujoco>
+      <worldbody>
+        <body name="base" pos="0 0 1">
+          <freejoint name="free_joint"/>
+          <geom name="base_geom" type="box" size="0.1 0.1 0.1" mass="1.0"
+            friction="0.5 0.01 0.005"/>
+          <body name="foot1" pos="0.2 0 0">
+            <joint name="joint1" type="hinge" axis="0 0 1" range="0 1.57"/>
+            <geom name="foot1_geom" type="box" size="0.05 0.05 0.05" mass="0.1"
+              friction="0.5 0.01 0.005"/>
+          </body>
+          <body name="foot2" pos="-0.2 0 0">
+            <joint name="joint2" type="hinge" axis="0 0 1" range="0 1.57"/>
+            <geom name="foot2_geom" type="box" size="0.05 0.05 0.05" mass="0.1"
+              friction="0.5 0.01 0.005"/>
+          </body>
+        </body>
+      </worldbody>
+    </mujoco>
+    """
+
+
+def test_simulation_config_is_piped(robot_xml, device):
+  """Test that SimulationCfg values are applied to both mj_model and wp_model."""
+  model = mujoco.MjModel.from_xml_string(robot_xml)
+
+  cfg = SimulationCfg(
+    contact_sensor_maxmatch=128,
+    broadphase="sap_tile",
+    broadphase_filter=("plane", "aabb"),
+    mujoco=MujocoCfg(
+      timestep=0.02,
+      integrator="euler",
+      solver="cg",
+      iterations=7,
+      ls_iterations=14,
+      ccd_iterations=20,
+      gravity=(0, 0, 7.5),
+      enableflags=("energy",),
+    ),
+  )
+
+  sim = Simulation(num_envs=1, cfg=cfg, model=model, device=device)
+
+  # MujocoCfg should be applied to mj_model.
+  assert sim.mj_model.opt.timestep == cfg.mujoco.timestep
+  assert sim.mj_model.opt.integrator == mujoco.mjtIntegrator.mjINT_EULER
+  assert sim.mj_model.opt.solver == mujoco.mjtSolver.mjSOL_CG
+  assert sim.mj_model.opt.iterations == cfg.mujoco.iterations
+  assert sim.mj_model.opt.ls_iterations == cfg.mujoco.ls_iterations
+  assert sim.mj_model.opt.ccd_iterations == cfg.mujoco.ccd_iterations
+  assert tuple(sim.mj_model.opt.gravity) == cfg.mujoco.gravity
+  assert sim.mj_model.opt.enableflags & mujoco.mjtEnableBit.mjENBL_ENERGY
+
+  # MujocoCfg should be inherited by wp_model via put_model.
+  np.testing.assert_almost_equal(
+    sim.model.opt.timestep[0].cpu().numpy(), cfg.mujoco.timestep
+  )
+  np.testing.assert_almost_equal(
+    sim.model.opt.gravity[0].cpu().numpy(), cfg.mujoco.gravity
+  )
+  assert sim.model.opt.integrator == mujoco.mjtIntegrator.mjINT_EULER
+  assert sim.model.opt.solver == mujoco.mjtSolver.mjSOL_CG
+  assert sim.model.opt.iterations == cfg.mujoco.iterations
+  assert sim.model.opt.enableflags & mujoco.mjtEnableBit.mjENBL_ENERGY
+
+  # SimulationCfg's warp-only settings should be applied to wp_model.opt.
+  assert sim.wp_model.opt.contact_sensor_maxmatch == cfg.contact_sensor_maxmatch
+  assert sim.wp_model.opt.broadphase == mjwarp.BroadphaseType.SAP_TILE
+  assert sim.wp_model.opt.broadphase_filter == (
+    mjwarp.BroadphaseFilter.PLANE | mjwarp.BroadphaseFilter.AABB
+  )
+
+
+def test_default_broadphase_keeps_put_model_heuristic(robot_xml, device):
+  """Unset broadphase settings should not override put_model's own heuristic."""
+  model = mujoco.MjModel.from_xml_string(robot_xml)
+  heuristic_opt = mjwarp.put_model(model).opt
+
+  sim = Simulation(num_envs=1, cfg=SimulationCfg(), model=model, device=device)
+
+  assert sim.wp_model.opt.broadphase == heuristic_opt.broadphase
+  assert sim.wp_model.opt.broadphase_filter == heuristic_opt.broadphase_filter
+
+
+def test_ls_parallel_is_deprecated():
+  """Setting the removed ls_parallel option warns instead of erroring."""
+  with pytest.warns(DeprecationWarning, match="ls_parallel"):
+    SimulationCfg(ls_parallel=True)
+
+
+def test_sim_reset_restores_initial_state(robot_xml, device):
+  """Test that sim.reset() restores qpos/qvel to initial values."""
+  model = mujoco.MjModel.from_xml_string(robot_xml)
+  sim = Simulation(num_envs=2, cfg=SimulationCfg(), model=model, device=device)
+
+  qpos0 = sim.data.qpos.clone()
+  qvel0 = sim.data.qvel.clone()
+
+  # Run simulation to modify state.
+  for _ in range(10):
+    sim.step()
+
+  assert not torch.allclose(sim.data.qpos, qpos0)
+  assert not torch.allclose(sim.data.qvel, qvel0)
+
+  # Reset should restore initial state.
+  sim.reset()
+
+  torch.testing.assert_close(sim.data.qpos[:], qpos0)
+  torch.testing.assert_close(sim.data.qvel[:], qvel0)
+  # qacc_warmstart should be zeroed.
+  assert (sim.data.qacc_warmstart == 0).all()
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="Likely bug on CPU MjWarp")
+def test_sim_reset_selective(robot_xml, device):
+  """Test that sim.reset() only affects specified environments."""
+  model = mujoco.MjModel.from_xml_string(robot_xml)
+  sim = Simulation(num_envs=4, cfg=SimulationCfg(), model=model, device=device)
+
+  qpos0 = sim.data.qpos.clone()
+
+  # Run simulation to modify state.
+  for _ in range(10):
+    sim.step()
+
+  qpos_after_sim = sim.data.qpos.clone()
+
+  # Reset only env 1 and 3.
+  sim.reset(torch.tensor([1, 3], device=device))
+
+  # Envs 1 and 3 should be reset.
+  torch.testing.assert_close(sim.data.qpos[1], qpos0[1])
+  torch.testing.assert_close(sim.data.qpos[3], qpos0[3])
+  # Envs 0 and 2 should be unchanged.
+  torch.testing.assert_close(sim.data.qpos[0], qpos_after_sim[0])
+  torch.testing.assert_close(sim.data.qpos[2], qpos_after_sim[2])
+
+
+def test_xpos_matches_qpos_after_forward(robot_xml, device):
+  """sim.step() leaves xpos stale; sim.forward() makes it match qpos.
+
+  In MuJoCo, mj_step = mj_step1 (forward kinematics + forces) + mj_step2
+  (integration). After mj_step, qpos/qvel are post-integration but xpos is
+  from the pre-integration forward pass. sim.forward() recomputes xpos from
+  the current qpos.
+  """
+  model = mujoco.MjModel.from_xml_string(robot_xml)
+  cfg = SimulationCfg(mujoco=MujocoCfg(timestep=0.01))  # Large dt for clear signal
+  sim = Simulation(num_envs=2, cfg=cfg, model=model, device=device)
+
+  # Step enough for significant velocity -> large staleness gap.
+  for _ in range(50):
+    sim.step()
+
+  # xpos is stale: reflects pre-integration state of last step.
+  # For the freejoint body (body 1), qpos[:3] is the true position.
+  xpos_stale = sim.data.xpos[:, 1].clone()
+  qpos_pos = sim.data.qpos[:, :3].clone()
+  assert not torch.allclose(xpos_stale, qpos_pos, atol=1e-4)
+
+  # forward() refreshes derived quantities from current qpos.
+  sim.forward()
+  xpos_fresh = sim.data.xpos[:, 1].clone()
+  torch.testing.assert_close(xpos_fresh, qpos_pos, atol=1e-5, rtol=0)
@@ -0,0 +1,100 @@
+"""Tests for sim data bridge."""
+
+from dataclasses import dataclass
+
+import pytest
+import torch
+import warp as wp
+from conftest import get_test_device
+
+from mjlab.sim.sim_data import TorchArray, WarpBridge
+
+
+@dataclass
+class MockData:
+  arr: wp.array
+  val: float = 1.0
+
+
+@pytest.fixture
+def device():
+  """Test device fixture."""
+  return get_test_device()
+
+
+@pytest.fixture
+def mock_data(device):
+  with wp.ScopedDevice(device):
+    return MockData(arr=wp.array([[1.0, 2.0], [3.0, 4.0]], dtype=wp.float32))
+
+
+def test_torch_array_shares_memory(device):
+  """TorchArray modifications affect underlying warp array."""
+  with wp.ScopedDevice(device):
+    wp_arr = wp.array([1.0, 2.0], dtype=wp.float32)
+  torch_arr = TorchArray(wp_arr)
+  torch_arr[0] = 99.0
+  assert wp_arr.numpy()[0] == 99.0
+
+
+def test_torch_array_ops_work(device):
+  """TorchArray supports PyTorch operations."""
+  with wp.ScopedDevice(device):
+    wp_arr = wp.array([1.0, 2.0], dtype=wp.float32)
+  torch_arr = TorchArray(wp_arr)
+  assert torch.allclose(torch_arr * 2, torch.tensor([2.0, 4.0], device=device))
+  assert torch.sum(torch_arr).item() == 3.0  # type: ignore
+
+
+def test_torch_array_expands_batched_singleton_world_dim(device):
+  """Size-1 world dim expands to nworld when mjwarp marks the array batched."""
+  with wp.ScopedDevice(device):
+    wp_arr = wp.array([[1.0, 2.0]], dtype=wp.float32)
+  setattr(wp_arr, "_is_batched", True)  # noqa: B010
+  torch_arr = TorchArray(wp_arr, nworld=4)
+  assert torch_arr.shape == (4, 2)
+
+
+def test_torch_array_keeps_unbatched_singleton_dim(device):
+  """A real size-1 leading dim without the batched marker is not expanded."""
+  with wp.ScopedDevice(device):
+    wp_arr = wp.array([[1.0, 2.0]], dtype=wp.float32)
+  torch_arr = TorchArray(wp_arr, nworld=4)
+  assert torch_arr.shape == (1, 2)
+
+
+def test_bridge_wraps_arrays(mock_data):
+  """WarpBridge wraps warp arrays as TorchArray."""
+  bridge = WarpBridge(mock_data)
+  assert isinstance(bridge.arr, TorchArray)
+  assert bridge.val == 1.0  # Non-arrays pass through.
+
+
+def test_bridge_preserves_memory_on_slice_assign(mock_data, device):
+  """Slice assignment preserves memory addresses (CUDA graph safe)."""
+  bridge = WarpBridge(mock_data)
+  initial_ptr = bridge.arr._tensor.data_ptr()
+
+  bridge.arr[:] = torch.zeros((2, 2), device=device)
+
+  assert bridge.arr._tensor.data_ptr() == initial_ptr
+  assert torch.all(bridge.arr._tensor == 0)
+
+
+def test_bridge_raises_on_setattr(mock_data, device):
+  """Direct assignment raises AttributeError with helpful message."""
+  bridge = WarpBridge(mock_data)
+
+  with pytest.raises(AttributeError, match="Cannot set attribute 'arr' on WarpBridge"):
+    bridge.arr = torch.zeros((2, 2), device=device)
+
+  with pytest.raises(AttributeError, match="Use in-place operations instead"):
+    bridge.val = 42.0
+
+
+def test_bridge_caches_wrappers(mock_data):
+  """WarpBridge caches wrapper objects."""
+  bridge = WarpBridge(mock_data)
+  arr1 = bridge.arr
+  arr2 = bridge.arr
+  assert arr1 is arr2  # Same cached object.
@@ -0,0 +1,622 @@
+"""Tests for spec_config.py."""
+
+import math
+from typing import cast
+
+import mujoco
+import pytest
+
+from mjlab.utils.spec_config import (
+  CameraCfg,
+  CollisionCfg,
+  GeomCfg,
+  LightCfg,
+  MaterialCfg,
+  MeshCfg,
+  TextureCfg,
+)
+
+
+@pytest.fixture
+def simple_robot_xml():
+  """Minimal robot XML for testing."""
+  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="1.0"/>
+          <body name="link1" pos="0 0 0">
+            <joint name="joint1" type="hinge" axis="0 0 1" range="0 1.57"/>
+            <geom name="link1_geom" type="box" size="0.1 0.1 0.1" mass="0.1"/>
+          </body>
+          <body name="link2" pos="0 0 0">
+            <joint name="joint2" type="hinge" axis="0 0 1" range="0 1.57"/>
+            <geom name="link2_geom" type="box" size="0.1 0.1 0.1" mass="0.1"/>
+          </body>
+        </body>
+      </worldbody>
+    </mujoco>
+    """
+
+
+# Collision Tests
+
+
+@pytest.fixture
+def multi_geom_spec():
+  """Spec with multiple geoms for collision testing."""
+  spec = mujoco.MjSpec()
+  body = spec.worldbody.add_body(name="test_body")
+  body.add_geom(
+    name="left_foot1_collision", type=mujoco.mjtGeom.mjGEOM_BOX, size=[0.1, 0.1, 0.1]
+  )
+  body.add_geom(
+    name="right_foot3_collision", type=mujoco.mjtGeom.mjGEOM_BOX, size=[0.1, 0.1, 0.1]
+  )
+  body.add_geom(
+    name="arm_collision", type=mujoco.mjtGeom.mjGEOM_BOX, size=[0.1, 0.1, 0.1]
+  )
+  return spec
+
+
+def test_collision_basic_properties(multi_geom_spec):
+  """CollisionCfg should set basic collision properties."""
+  collision_cfg = CollisionCfg(
+    geom_names_expr=("arm_collision",), contype=2, conaffinity=3, condim=4, priority=1
+  )
+  collision_cfg.edit_spec(multi_geom_spec)
+
+  geom = multi_geom_spec.geom("arm_collision")
+  assert geom.contype == 2
+  assert geom.conaffinity == 3
+  assert geom.condim == 4
+  assert geom.priority == 1
+
+
+def test_collision_regex_matching(multi_geom_spec):
+  """CollisionCfg should support regex pattern matching."""
+  collision_cfg = CollisionCfg(
+    geom_names_expr=(r"^(left|right)_foot\d_collision$",),
+    contype=1,
+    conaffinity=1,
+    condim=3,
+    priority=1,
+    friction=(0.6,),
+    disable_other_geoms=False,
+  )
+  collision_cfg.edit_spec(multi_geom_spec)
+
+  left_foot = multi_geom_spec.geom("left_foot1_collision")
+  assert left_foot.condim == 3
+  assert left_foot.priority == 1
+  assert left_foot.friction[0] == 0.6
+
+  right_foot = multi_geom_spec.geom("right_foot3_collision")
+  assert right_foot.condim == 3
+
+  arm = multi_geom_spec.geom("arm_collision")
+  assert arm.condim == 3  # Default unchanged.
+
+
+def test_collision_dict_field_resolution(multi_geom_spec):
+  """CollisionCfg should support dict-based field resolution."""
+  collision_cfg = CollisionCfg(
+    geom_names_expr=(r".*_foot\d_collision$", "arm_collision"),
+    contype=1,
+    conaffinity=1,
+    condim={r".*_foot\d_collision$": 3, "arm_collision": 1},
+    priority={r".*_foot\d_collision$": 2, "arm_collision": 0},
+  )
+  collision_cfg.edit_spec(multi_geom_spec)
+
+  left_foot = multi_geom_spec.geom("left_foot1_collision")
+  assert left_foot.condim == 3
+  assert left_foot.priority == 2
+
+  arm = multi_geom_spec.geom("arm_collision")
+  assert arm.condim == 1
+  assert arm.priority == 0
+
+
+def test_collision_margin_gap_solmix(multi_geom_spec):
+  """CollisionCfg should set margin, gap, and solmix on matched geoms."""
+  collision_cfg = CollisionCfg(
+    geom_names_expr=("arm_collision",),
+    contype=1,
+    conaffinity=1,
+    condim=3,
+    priority=0,
+    margin=0.01,
+    gap=0.005,
+    solmix=0.5,
+    disable_other_geoms=False,
+  )
+  collision_cfg.edit_spec(multi_geom_spec)
+
+  geom = multi_geom_spec.geom("arm_collision")
+  assert geom.margin == pytest.approx(0.01)
+  assert geom.gap == pytest.approx(0.005)
+  assert geom.solmix == pytest.approx(0.5)
+
+
+def test_collision_margin_gap_solmix_dict(multi_geom_spec):
+  """CollisionCfg should support dict-based margin, gap, solmix overrides."""
+  collision_cfg = CollisionCfg(
+    geom_names_expr=(r".*_foot\d_collision$", "arm_collision"),
+    contype=1,
+    conaffinity=1,
+    condim=3,
+    priority=0,
+    margin={r".*_foot\d_collision$": 0.02, "arm_collision": 0.01},
+    gap={r".*_foot\d_collision$": 0.01, "arm_collision": 0.0},
+    solmix={r".*_foot\d_collision$": 0.8, "arm_collision": 0.2},
+    disable_other_geoms=False,
+  )
+  collision_cfg.edit_spec(multi_geom_spec)
+
+  left_foot = multi_geom_spec.geom("left_foot1_collision")
+  assert left_foot.margin == pytest.approx(0.02)
+  assert left_foot.gap == pytest.approx(0.01)
+  assert left_foot.solmix == pytest.approx(0.8)
+
+  arm = multi_geom_spec.geom("arm_collision")
+  assert arm.margin == pytest.approx(0.01)
+  assert arm.solmix == pytest.approx(0.2)
+
+
+def test_collision_structural_dict_requires_coverage(multi_geom_spec):
+  """Dict-valued structural fields must cover every matched geom."""
+  collision_cfg = CollisionCfg(
+    geom_names_expr=(".*_collision",),
+    contype=1,
+    conaffinity=1,
+    condim={r".*_foot\d_collision$": 3},  # Does not cover arm_collision.
+    priority=0,
+  )
+  with pytest.raises(ValueError, match="condim dict does not match geom"):
+    collision_cfg.edit_spec(multi_geom_spec)
+
+
+def test_collision_structural_fields_are_required():
+  """Structural fields cannot be None."""
+  cfg = CollisionCfg(
+    geom_names_expr=(".*",),
+    contype=1,
+    conaffinity=1,
+    condim=cast(int, None),
+    priority=0,
+  )
+  with pytest.raises(ValueError, match="condim is required"):
+    cfg.validate()
+
+
+def test_collision_disable_other_geoms(multi_geom_spec):
+  """CollisionCfg should disable non-matching geoms when requested."""
+  collision_cfg = CollisionCfg(
+    geom_names_expr=("left_foot1_collision",),
+    contype=2,
+    conaffinity=1,
+    condim=3,
+    priority=0,
+    disable_other_geoms=True,
+  )
+  collision_cfg.edit_spec(multi_geom_spec)
+
+  left_foot = multi_geom_spec.geom("left_foot1_collision")
+  assert left_foot.contype == 2
+
+  right_foot = multi_geom_spec.geom("right_foot3_collision")
+  assert right_foot.contype == 0
+  assert right_foot.conaffinity == 0
+
+  arm = multi_geom_spec.geom("arm_collision")
+  assert arm.contype == 0
+
+
+# fmt: off
+@pytest.mark.parametrize(
+  "param,value,expected_error",
+  [
+    ("condim", -1, "condim must be one of"),
+    ("condim", 2, "condim must be one of"),
+    ("contype", -1, "contype must be non-negative"),
+    ("conaffinity", -1, "conaffinity must be non-negative"),
+    ("priority", -1, "priority must be non-negative"),
+    ("margin", -0.1, "margin must be non-negative"),
+    ("gap", -0.1, "gap must be non-negative"),
+    ("solmix", 1.5, r"solmix must be in \[0, 1\]"),
+  ],
+)
+# fmt: on
+def test_collision_validation(param, value, expected_error):
+  """CollisionCfg should validate parameters."""
+  with pytest.raises(ValueError, match=expected_error):
+    cfg = CollisionCfg(
+      geom_names_expr=("test",),
+      contype=value if param == "contype" else 1,
+      conaffinity=value if param == "conaffinity" else 1,
+      condim=value if param == "condim" else 3,
+      priority=value if param == "priority" else 0,
+      margin=value if param == "margin" else None,
+      gap=value if param == "gap" else None,
+      solmix=value if param == "solmix" else None,
+    )
+    cfg.validate()
+
+
+# Mesh Tests
+
+
+def _sphere_points(n: int) -> list[float]:
+  """Deterministic point cloud on a unit sphere (Fibonacci lattice).
+
+  Every point is on the hull, so the unconstrained hull has n vertices.
+  """
+  phi = math.pi * (3.0 - math.sqrt(5.0))
+  verts: list[float] = []
+  for i in range(n):
+    y = 1.0 - 2.0 * (i + 0.5) / n
+    r = math.sqrt(max(0.0, 1.0 - y * y))
+    theta = phi * i
+    verts += [math.cos(theta) * r, y, math.sin(theta) * r]
+  return verts
+
+
+@pytest.fixture
+def multi_mesh_spec():
+  """Spec with two mesh assets, each a 100-vertex sphere point cloud."""
+  spec = mujoco.MjSpec()
+  for name in ("blob_a", "blob_b"):
+    mesh = spec.add_mesh()
+    mesh.name = name
+    mesh.uservert = _sphere_points(100)
+  return spec
+
+
+def _hull_numvert(spec: mujoco.MjSpec, mesh_name: str) -> int:
+  """Compile and return the convex-hull vertex count for a mesh asset.
+
+  maxhullvert caps the hull stored in mesh_graph (which the narrowphase support
+  function walks), not the raw mesh_vert array, so we read it back from the
+  compiled graph: the first int at mesh_graphadr is the hull vertex count.
+  """
+  body = spec.worldbody.add_body()
+  geom = body.add_geom()
+  geom.type = mujoco.mjtGeom.mjGEOM_MESH
+  geom.meshname = mesh_name
+  model = spec.compile()
+  mesh_id = model.mesh(mesh_name).id
+  return int(model.mesh_graph[model.mesh_graphadr[mesh_id]])
+
+
+def test_mesh_maxhullvert_scalar(multi_mesh_spec):
+  """MeshCfg should set maxhullvert on every matched mesh asset."""
+  MeshCfg(mesh_names_expr=(".*",), maxhullvert=12).edit_spec(multi_mesh_spec)
+
+  assert multi_mesh_spec.mesh("blob_a").maxhullvert == 12
+  assert multi_mesh_spec.mesh("blob_b").maxhullvert == 12
+
+
+def test_mesh_maxhullvert_regex_matching(multi_mesh_spec):
+  """Only meshes matching the regex should be modified; others stay default."""
+  MeshCfg(mesh_names_expr=("blob_a",), maxhullvert=8).edit_spec(multi_mesh_spec)
+
+  assert multi_mesh_spec.mesh("blob_a").maxhullvert == 8
+  assert multi_mesh_spec.mesh("blob_b").maxhullvert == -1  # Unchanged default.
+
+
+def test_mesh_maxhullvert_dict_resolution(multi_mesh_spec):
+  """A dict should resolve per-mesh values by pattern."""
+  MeshCfg(
+    mesh_names_expr=(".*",),
+    maxhullvert={"blob_a": 6, "blob_b": 20},
+  ).edit_spec(multi_mesh_spec)
+
+  assert multi_mesh_spec.mesh("blob_a").maxhullvert == 6
+  assert multi_mesh_spec.mesh("blob_b").maxhullvert == 20
+
+
+def test_mesh_unconstrained_hull_keeps_all_vertices(multi_mesh_spec):
+  """Baseline: without a cap, every sphere point is a hull vertex."""
+  assert _hull_numvert(multi_mesh_spec, "blob_a") == 100
+
+
+def test_mesh_maxhullvert_caps_compiled_hull(multi_mesh_spec):
+  """The cap should actually shrink the compiled convex hull below the baseline."""
+  MeshCfg(mesh_names_expr=("blob_a",), maxhullvert=10).edit_spec(multi_mesh_spec)
+  assert _hull_numvert(multi_mesh_spec, "blob_a") == 10
+
+
+@pytest.mark.parametrize("value", [0, 2, 3, -2])
+def test_mesh_maxhullvert_validation(value):
+  """maxhullvert must be -1 or greater than 3."""
+  with pytest.raises(ValueError, match="maxhullvert must be"):
+    MeshCfg(mesh_names_expr=(".*",), maxhullvert=value).validate()
+
+
+def test_mesh_maxhullvert_dict_validation():
+  """Validation should also cover dict-valued maxhullvert."""
+  with pytest.raises(ValueError, match="maxhullvert must be.*pattern 'blob_a'"):
+    MeshCfg(mesh_names_expr=(".*",), maxhullvert={"blob_a": 2}).validate()
+
+
+def test_mesh_maxhullvert_allows_unlimited():
+  """-1 (unlimited) is valid and is the MuJoCo default."""
+  MeshCfg(mesh_names_expr=(".*",), maxhullvert=-1).validate()
+
+
+# Visual Element Tests
+
+
+def test_texture_cfg():
+  """TextureCfg should add textures to spec."""
+  spec = mujoco.MjSpec()
+  texture_cfg = TextureCfg(
+    name="test_texture",
+    type="2d",
+    builtin="checker",
+    rgb1=(1.0, 0.0, 0.0),
+    rgb2=(0.0, 1.0, 0.0),
+    width=64,
+    height=64,
+  )
+  texture_cfg.edit_spec(spec)
+
+  texture = spec.texture("test_texture")
+  assert texture.name == "test_texture"
+
+
+def test_texture_cfg_image_fields():
+  """TextureCfg should set file, grid, nchannel, flip, and random fields."""
+  spec = mujoco.MjSpec()
+  texture_cfg = TextureCfg(
+    name="test_texture",
+    type="cube",
+    mark="random",
+    random=0.25,
+    file="sky.png",
+    gridsize=(3, 4),
+    gridlayout=".U..LFRB.D..",
+    nchannel=4,
+    hflip=True,
+    vflip=True,
+  )
+  texture_cfg.edit_spec(spec)
+
+  texture = spec.texture("test_texture")
+  assert texture.file == "sky.png"
+  assert tuple(texture.gridsize) == (3, 4)
+  assert texture.nchannel == 4
+  assert texture.hflip
+  assert texture.vflip
+  assert texture.random == pytest.approx(0.25)
+
+
+@pytest.mark.parametrize(
+  "param,value,expected_error",
+  [
+    ("width", 0, "width must be positive"),
+    ("width", -1, "width must be positive"),
+    ("height", 0, "height must be positive"),
+    ("height", -1, "height must be positive"),
+  ],
+)
+def test_texture_cfg_validation(param, value, expected_error):
+  """TextureCfg should validate explicitly set width and height."""
+  with pytest.raises(ValueError, match=expected_error):
+    TextureCfg(
+      name="test_texture",
+      type="2d",
+      builtin="checker",
+      width=value if param == "width" else 64,
+      height=value if param == "height" else 64,
+    ).validate()
+
+
+def test_texture_requires_source():
+  """A texture with no builtin, file, or cubefiles has nothing to render."""
+  with pytest.raises(ValueError, match="must specify a builtin pattern"):
+    TextureCfg(name="test_texture", type="2d").validate()
+
+
+@pytest.mark.parametrize("gridsize", [None, (1, 1)])
+def test_texture_single_file_skybox_validation(gridsize):
+  """A single-image skybox cannot be sampled as six faces by the Warp renderer."""
+  with pytest.raises(ValueError, match="must declare a gridsize"):
+    TextureCfg(
+      name="sky", type="skybox", file="sky.png", gridsize=gridsize
+    ).validate()
+
+
+def test_texture_skybox_allows_grid_and_cubefiles():
+  """A larger grid or explicit cubefiles cover all six faces."""
+  TextureCfg(name="sky", type="skybox", file="sky.png", gridsize=(6, 1)).validate()
+  TextureCfg(
+    name="sky", type="skybox", cubefiles=("a", "b", "c", "d", "e", "f")
+  ).validate()
+
+
+def test_texture_skybox_builtin_overrides_single_file():
+  """A builtin pattern takes precedence over the file, so the grid is irrelevant."""
+  TextureCfg(
+    name="sky", type="skybox", builtin="gradient", file="sky.png"
+  ).validate()
+
+
+def test_material_cfg():
+  """MaterialCfg should add materials to spec."""
+  spec = mujoco.MjSpec()
+  material_cfg = MaterialCfg(
+    name="test_material",
+    texuniform=True,
+    texrepeat=(2, 2),
+    reflectance=0.5,
+  )
+  material_cfg.edit_spec(spec)
+
+  material = spec.material("test_material")
+  assert material.name == "test_material"
+
+
+def test_light_cfg():
+  """LightCfg should add lights to spec."""
+  spec = mujoco.MjSpec()
+  light_cfg = LightCfg(
+    name="test_light",
+    body="world",
+    type="spot",
+    pos=(1.0, 2.0, 3.0),
+    dir=(0.0, 0.0, -1.0),
+  )
+  light_cfg.edit_spec(spec)
+
+  light = spec.light("test_light")
+  assert light.name == "test_light"
+
+
+def test_light_cfg_color_and_attenuation():
+  """LightCfg should set diffuse, specular, ambient, active, attenuation."""
+  spec = mujoco.MjSpec()
+  light_cfg = LightCfg(
+    name="test_light",
+    diffuse=(1.0, 0.0, 0.0),
+    specular=(0.0, 1.0, 0.0),
+    ambient=(0.0, 0.0, 1.0),
+    active=False,
+    attenuation=(1.0, 0.5, 0.25),
+  )
+  light_cfg.edit_spec(spec)
+
+  light = spec.light("test_light")
+  assert tuple(light.diffuse) == pytest.approx((1.0, 0.0, 0.0))
+  assert tuple(light.specular) == pytest.approx((0.0, 1.0, 0.0))
+  assert tuple(light.ambient) == pytest.approx((0.0, 0.0, 1.0))
+  assert not light.active
+  assert tuple(light.attenuation) == pytest.approx((1.0, 0.5, 0.25))
+
+
+def test_camera_cfg():
+  """CameraCfg should add cameras to spec."""
+  spec = mujoco.MjSpec()
+  camera_cfg = CameraCfg(
+    name="test_camera", body="world", fovy=60.0, pos=(0.0, 0.0, 5.0)
+  )
+  camera_cfg.edit_spec(spec)
+
+  camera = spec.camera("test_camera")
+  assert camera.name == "test_camera"
+
+
+def test_material_cfg_geom_assignment():
+  """MaterialCfg.geom_names_expr assigns material to matching geoms."""
+  spec = mujoco.MjSpec()
+  body = spec.worldbody.add_body(name="test_body")
+  body.add_geom(
+    name="link1_visual",
+    type=mujoco.mjtGeom.mjGEOM_BOX,
+    size=[0.1, 0.1, 0.1],
+  )
+  body.add_geom(
+    name="link2_visual",
+    type=mujoco.mjtGeom.mjGEOM_BOX,
+    size=[0.1, 0.1, 0.1],
+  )
+  body.add_geom(
+    name="arm_collision",
+    type=mujoco.mjtGeom.mjGEOM_BOX,
+    size=[0.1, 0.1, 0.1],
+  )
+
+  mat_cfg = MaterialCfg(
+    name="my_mat",
+    texuniform=True,
+    texrepeat=(1, 1),
+    geom_names_expr=(r".*_visual$",),
+  )
+  mat_cfg.edit_spec(spec)
+
+  assert spec.geom("link1_visual").material == "my_mat"
+  assert spec.geom("link2_visual").material == "my_mat"
+  assert spec.geom("arm_collision").material == ""
+
+
+def test_geom_group_basic(multi_geom_spec):
+  """GeomCfg should set the group of matching geoms."""
+  group_cfg = GeomCfg(
+    geom_names_expr=(r"^(left|right)_foot\d_collision$",), group=3
+  )
+  group_cfg.edit_spec(multi_geom_spec)
+
+  assert multi_geom_spec.geom("left_foot1_collision").group == 3
+  assert multi_geom_spec.geom("right_foot3_collision").group == 3
+
+  arm = multi_geom_spec.geom("arm_collision")
+  assert arm.group == 0  # Default unchanged.
+
+
+def test_geom_group_none_is_noop(multi_geom_spec):
+  """GeomCfg should leave the group alone when it is None."""
+  multi_geom_spec.geom("arm_collision").group = 2
+  GeomCfg(geom_names_expr=(".*",)).edit_spec(multi_geom_spec)
+
+  assert multi_geom_spec.geom("arm_collision").group == 2
+
+
+def test_geom_group_dict_field_resolution(multi_geom_spec):
+  """GeomCfg should support dict-based field resolution."""
+  group_cfg = GeomCfg(
+    geom_names_expr=(r".*_foot\d_collision$", "arm_collision"),
+    group={r".*_foot\d_collision$": 3, "arm_collision": 4},
+  )
+  group_cfg.edit_spec(multi_geom_spec)
+
+  assert multi_geom_spec.geom("left_foot1_collision").group == 3
+  assert multi_geom_spec.geom("arm_collision").group == 4
+
+
+def test_geom_group_unnamed_geoms():
+  """GeomCfg should update every matching geom, including unnamed ones."""
+  spec = mujoco.MjSpec()
+  body = spec.worldbody.add_body(name="test_body")
+  for _ in range(3):
+    body.add_geom(type=mujoco.mjtGeom.mjGEOM_BOX, size=[0.1, 0.1, 0.1])
+
+  group_cfg = GeomCfg(geom_names_expr=(".*",), group=3)
+  group_cfg.edit_spec(spec)
+
+  assert [g.group for g in spec.geoms] == [3, 3, 3]
+
+
+@pytest.mark.parametrize("value", [-1, mujoco.mjNGROUP, {"arm_collision": -1}])
+def test_geom_group_validation(value):
+  """GeomCfg should validate the group."""
+  with pytest.raises(ValueError, match="group must be in"):
+    cfg = GeomCfg(geom_names_expr=("test",), group=value)
+    cfg.validate()
+
+
+def test_geom_patches_collision_attrs(multi_geom_spec):
+  """GeomCfg should patch collision attributes, leaving unset ones alone."""
+  multi_geom_spec.geom("left_foot1_collision").contype = 2
+  GeomCfg(
+    geom_names_expr=(r"^(left|right)_foot\d_collision$",),
+    condim=6,
+    friction=(1.0, 5e-3, 5e-4),
+  ).edit_spec(multi_geom_spec)
+
+  left_foot = multi_geom_spec.geom("left_foot1_collision")
+  assert left_foot.condim == 6
+  assert left_foot.friction[0] == pytest.approx(1.0)
+  assert left_foot.contype == 2  # Unset attribute untouched.
+
+  arm = multi_geom_spec.geom("arm_collision")
+  assert arm.condim == 3  # Unmatched geom untouched.
+
+
+def test_geom_collision_attr_validation():
+  """GeomCfg should validate collision attributes like CollisionCfg does."""
+  with pytest.raises(ValueError, match="condim must be one of"):
+    GeomCfg(geom_names_expr=("test",), condim=2).validate()
+  with pytest.raises(ValueError, match="solmix must be in"):
+    GeomCfg(geom_names_expr=("test",), solmix=1.5).validate()
@@ -0,0 +1,214 @@
+"""Tests for spec.py utilities."""
+
+import mujoco
+import numpy as np
+import pytest
+
+from mjlab.utils.spec import create_position_actuator, create_velocity_actuator
+
+
+@pytest.fixture
+def spec_with_limited_joint():
+  """Create a spec with a limited joint for testing."""
+  spec = mujoco.MjSpec()
+  body = spec.worldbody.add_body(name="test_body")
+  joint = body.add_joint(
+    name="test_joint",
+    type=mujoco.mjtJoint.mjJNT_SLIDE,
+    axis=[0, 0, 1],
+    range=[-1.0, 1.0],
+  )
+  joint.limited = mujoco.mjtLimited.mjLIMITED_TRUE
+  body.add_geom(type=mujoco.mjtGeom.mjGEOM_BOX, size=[0.1, 0.1, 0.1], mass=1.0)
+  return spec
+
+
+def test_position_actuator_allows_setpoints_beyond_joint_limits(
+  spec_with_limited_joint,
+):
+  """Verify that position actuators allow commanding setpoints beyond joint limits.
+
+  This is the core fix: position actuators have ctrllimited=False, allowing
+  policies to command positions outside kinematic limits to maximize torque.
+  We verify this by comparing torque when commanding beyond vs at the limit.
+  """
+  stiffness = 100.0
+  create_position_actuator(
+    spec_with_limited_joint,
+    "test_joint",
+    stiffness=stiffness,
+    damping=10.0,
+    effort_limit=500.0,
+  )
+  model = spec_with_limited_joint.compile()
+  data = mujoco.MjData(model)
+
+  # Verify ctrllimited=False.
+  actuator_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_ACTUATOR, "test_joint")
+  assert model.actuator_ctrllimited[actuator_id] == 0
+
+  # Set joint position.
+  current_pos = 0.5
+  data.qpos[0] = current_pos
+  data.qvel[0] = 0.0
+
+  # Command beyond joint upper limit (1.0).
+  target_beyond = 2.0
+  data.ctrl[0] = target_beyond
+  mujoco.mj_forward(model, data)
+  torque_beyond = data.actuator_force[0]
+
+  # Expected: -kp * (q - target) = -100 * (0.5 - 2.0) = 150.0.
+  expected_torque_beyond = -stiffness * (current_pos - target_beyond)
+  np.testing.assert_allclose(torque_beyond, expected_torque_beyond, rtol=1e-5)
+
+  # If control were clipped to limit (1.0), torque would be 50.0.
+  expected_if_clipped = -stiffness * (current_pos - 1.0)
+  assert abs(torque_beyond) == 3.0 * abs(expected_if_clipped)
+
+
+def test_position_actuator_forces_clipped_to_effort_limit(spec_with_limited_joint):
+  """Verify that actuator forces are clipped to effort_limit, not control limits.
+
+  While control setpoints aren't clipped, forces must still be limited to
+  prevent unrealistic torques.
+  """
+  effort_limit = 10.0
+  stiffness = 1000.0
+  create_position_actuator(
+    spec_with_limited_joint,
+    "test_joint",
+    stiffness=stiffness,
+    damping=1.0,
+    effort_limit=effort_limit,
+  )
+  model = spec_with_limited_joint.compile()
+  data = mujoco.MjData(model)
+
+  # Verify forcelimited=True.
+  actuator_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_ACTUATOR, "test_joint")
+  assert model.actuator_forcelimited[actuator_id] == 1
+  np.testing.assert_array_almost_equal(
+    model.actuator_forcerange[actuator_id], [-effort_limit, effort_limit]
+  )
+
+  # Large position error would generate force > effort_limit without clipping.
+  data.qpos[0] = 0.0
+  data.ctrl[0] = 3.0
+  mujoco.mj_step(model, data)
+
+  # Force should saturate at effort_limit.
+  assert abs(data.actuator_force[0]) <= effort_limit + 1e-6
+  assert abs(data.actuator_force[0]) >= effort_limit - 1e-3
+
+
+def test_ctrllimited_true_would_clip_internally():
+  """Demonstrate that ctrllimited=True clips control internally (the bug).
+
+  This test shows what would happen with the bug: commanding at vs beyond
+  the limit generates identical forces, proving internal clipping.
+  """
+  spec = mujoco.MjSpec()
+  body = spec.worldbody.add_body(name="test_body")
+  joint = body.add_joint(
+    name="test_joint",
+    type=mujoco.mjtJoint.mjJNT_SLIDE,
+    axis=[0, 0, 1],
+    range=[-1.0, 1.0],
+  )
+  joint.limited = mujoco.mjtLimited.mjLIMITED_TRUE
+  body.add_geom(type=mujoco.mjtGeom.mjGEOM_BOX, size=[0.1, 0.1, 0.1], mass=1.0)
+
+  # Manually create actuator WITH THE BUG: ctrllimited=True.
+  stiffness = 100.0
+  actuator = spec.add_actuator(name="test_joint", target="test_joint")
+  actuator.trntype = mujoco.mjtTrn.mjTRN_JOINT
+  actuator.dyntype = mujoco.mjtDyn.mjDYN_NONE
+  actuator.gaintype = mujoco.mjtGain.mjGAIN_FIXED
+  actuator.biastype = mujoco.mjtBias.mjBIAS_AFFINE
+  actuator.ctrllimited = True  # THE BUG.
+  actuator.ctrlrange[:] = np.array([-1.0, 1.0])
+  actuator.gainprm[0] = stiffness
+  actuator.biasprm[1] = -stiffness
+  actuator.biasprm[2] = -10.0
+
+  model = spec.compile()
+  data = mujoco.MjData(model)
+
+  data.qpos[0] = 0.5
+  data.qvel[0] = 0.0
+
+  # Command at limit.
+  data.ctrl[0] = 1.0
+  mujoco.mj_forward(model, data)
+  force_at_limit = data.actuator_force[0]
+
+  # Command beyond limit.
+  data.ctrl[0] = 2.0
+  mujoco.mj_forward(model, data)
+  force_beyond = data.actuator_force[0]
+
+  # With the bug, forces are IDENTICAL (both clipped to 1.0 internally).
+  np.testing.assert_allclose(force_beyond, force_at_limit, rtol=1e-10)
+  expected_force = -stiffness * (0.5 - 1.0)
+  np.testing.assert_allclose(force_at_limit, expected_force, rtol=1e-5)
+
+
+@pytest.fixture
+def spec_with_continuous_joint():
+  """Create a spec with a continuous (unlimited) hinge joint."""
+  spec = mujoco.MjSpec()
+  body = spec.worldbody.add_body(name="wheel")
+  body.add_joint(
+    name="wheel_joint",
+    type=mujoco.mjtJoint.mjJNT_HINGE,
+    axis=[0, 0, 1],
+  )
+  body.add_geom(type=mujoco.mjtGeom.mjGEOM_CYLINDER, size=[0.1, 0.05], mass=1.0)
+  return spec
+
+
+def test_velocity_actuator_continuous_joint(spec_with_continuous_joint):
+  """Velocity actuator compiles for continuous joints with no range."""
+  damping = 10.0
+  create_velocity_actuator(
+    spec_with_continuous_joint,
+    "wheel_joint",
+    damping=damping,
+  )
+  model = spec_with_continuous_joint.compile()
+  data = mujoco.MjData(model)
+
+  actuator_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_ACTUATOR, "wheel_joint")
+  assert model.actuator_ctrllimited[actuator_id] == 0
+
+  # Commanding a velocity target should produce force = damping * (ctrl - qvel).
+  data.qvel[0] = 0.0
+  data.ctrl[0] = 5.0
+  mujoco.mj_forward(model, data)
+  expected_force = damping * 5.0
+  np.testing.assert_allclose(data.actuator_force[0], expected_force, rtol=1e-5)
+
+
+def test_velocity_actuator_forces_clipped_to_effort_limit(
+  spec_with_continuous_joint,
+):
+  """Velocity actuator force is bounded by effort_limit."""
+  effort_limit = 20.0
+  create_velocity_actuator(
+    spec_with_continuous_joint,
+    "wheel_joint",
+    damping=1000.0,
+    effort_limit=effort_limit,
+  )
+  model = spec_with_continuous_joint.compile()
+  data = mujoco.MjData(model)
+
+  actuator_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_ACTUATOR, "wheel_joint")
+  assert model.actuator_forcelimited[actuator_id] == 1
+
+  # Large velocity error would exceed effort_limit without clamping.
+  data.qvel[0] = 0.0
+  data.ctrl[0] = 100.0
+  mujoco.mj_forward(model, data)
+  assert abs(data.actuator_force[0]) <= effort_limit + 1e-6
@@ -0,0 +1,144 @@
+"""Generic tests for task config integrity."""
+
+import pytest
+
+from mjlab.envs import ManagerBasedRlEnvCfg
+from mjlab.managers.observation_manager import ObservationGroupCfg
+from mjlab.tasks.registry import list_tasks, load_env_cfg
+
+
+@pytest.fixture(scope="module")
+def all_task_ids() -> list[str]:
+  """Get all registered task IDs."""
+  return list_tasks()
+
+
+def test_all_tasks_loadable(all_task_ids: list[str]) -> None:
+  """All registered tasks should be loadable without errors."""
+  for task_id in all_task_ids:
+    try:
+      cfg = load_env_cfg(task_id)
+      assert isinstance(cfg, ManagerBasedRlEnvCfg), (
+        f"Task {task_id} did not return ManagerBasedRlEnvCfg"
+      )
+    except Exception as e:
+      pytest.fail(f"Failed to load task '{task_id}': {e}")
+
+
+def test_all_tasks_have_play_config(all_task_ids: list[str]) -> None:
+  """All tasks should be loadable in play mode."""
+  for task_id in all_task_ids:
+    try:
+      cfg = load_env_cfg(task_id, play=True)
+      assert isinstance(cfg, ManagerBasedRlEnvCfg), (
+        f"Task {task_id} play mode did not return ManagerBasedRlEnvCfg"
+      )
+    except Exception as e:
+      pytest.fail(f"Failed to load task '{task_id}' in play mode: {e}")
+
+
+def test_play_mode_episode_length(all_task_ids: list[str]) -> None:
+  """Play mode tasks should have infinite episode length."""
+  for task_id in all_task_ids:
+    cfg = load_env_cfg(task_id, play=True)
+    assert cfg.episode_length_s >= 1e9, (
+      f"{task_id} (play mode) episode_length_s={cfg.episode_length_s}, expected >= 1e9"
+    )
+
+
+def test_play_mode_observation_corruption_disabled(all_task_ids: list[str]) -> None:
+  """Play mode tasks should have observation corruption disabled for policy."""
+  for task_id in all_task_ids:
+    cfg = load_env_cfg(task_id, play=True)
+
+    assert "actor" in cfg.observations, (
+      f"Play mode task {task_id} missing 'policy' observation group"
+    )
+
+    policy_obs = cfg.observations["actor"]
+    assert isinstance(policy_obs, ObservationGroupCfg), (
+      f"Play mode task {task_id} policy observation is not ObservationGroupCfg"
+    )
+
+    assert not policy_obs.enable_corruption, (
+      f"Play mode task {task_id} has enable_corruption=True, expected False"
+    )
+
+
+def test_training_mode_observation_corruption_enabled(all_task_ids: list[str]) -> None:
+  """Training mode tasks should have observation corruption enabled for policy."""
+  for task_id in all_task_ids:
+    cfg = load_env_cfg(task_id)
+
+    assert "actor" in cfg.observations, (
+      f"Training task {task_id} missing 'policy' observation group"
+    )
+
+    policy_obs = cfg.observations["actor"]
+    assert isinstance(policy_obs, ObservationGroupCfg), (
+      f"Training task {task_id} policy observation is not ObservationGroupCfg"
+    )
+
+    assert policy_obs.enable_corruption, (
+      f"Training task {task_id} has enable_corruption=False, expected True"
+    )
+
+
+def test_critic_observation_corruption_always_disabled(all_task_ids: list[str]) -> None:
+  """Critic observations should always have corruption disabled."""
+  for task_id in all_task_ids:
+    cfg = load_env_cfg(task_id)
+
+    if "critic" not in cfg.observations:
+      continue
+
+    critic_obs = cfg.observations["critic"]
+    assert isinstance(critic_obs, ObservationGroupCfg), (
+      f"Task {task_id} critic observation is not ObservationGroupCfg"
+    )
+
+    assert not critic_obs.enable_corruption, (
+      f"Task {task_id} has critic enable_corruption=True, expected False"
+    )
+
+
+def test_play_training_observation_structure_match(all_task_ids: list[str]) -> None:
+  """Play and training configs should have matching observation structure."""
+  for task_id in all_task_ids:
+    training_cfg = load_env_cfg(task_id)
+    play_cfg = load_env_cfg(task_id, play=True)
+
+    # Same observation groups.
+    assert set(training_cfg.observations.keys()) == set(play_cfg.observations.keys()), (
+      f"Observation groups mismatch between {task_id} training and play modes"
+    )
+
+    # Same observation terms within each group.
+    for obs_group_name in training_cfg.observations:
+      training_terms = set(training_cfg.observations[obs_group_name].terms.keys())
+      play_terms = set(play_cfg.observations[obs_group_name].terms.keys())
+
+      assert training_terms == play_terms, (
+        f"Observation terms mismatch in group '{obs_group_name}' "
+        f"between {task_id} training and play modes"
+      )
+
+
+def test_play_training_action_structure_match(all_task_ids: list[str]) -> None:
+  """Play and training configs should have matching action structure."""
+  for task_id in all_task_ids:
+    training_cfg = load_env_cfg(task_id)
+    play_cfg = load_env_cfg(task_id, play=True)
+
+    assert set(training_cfg.actions.keys()) == set(play_cfg.actions.keys()), (
+      f"Action structure mismatch between {task_id} training and play modes"
+    )
+
+
+def test_play_mode_disables_push_robot(all_task_ids: list[str]) -> None:
+  """Play mode tasks should disable push_robot event."""
+  for task_id in all_task_ids:
+    cfg = load_env_cfg(task_id, play=True)
+    assert "push_robot" not in cfg.events, (
+      f"Play mode task {task_id} has push_robot event, expected it to be removed"
+    )
@@ -0,0 +1,240 @@
+"""Tests for MDP termination functions."""
+
+from unittest.mock import Mock
+
+import mujoco
+import pytest
+import torch
+from conftest import get_test_device
+
+from mjlab.envs.mdp.terminations import nan_detection
+from mjlab.managers.termination_manager import TerminationManager, TerminationTermCfg
+from mjlab.sim.sim import Simulation, SimulationCfg
+from mjlab.tasks.velocity.mdp.terminations import (
+  out_of_terrain_bounds,
+  terrain_edge_reached,
+)
+
+
+@pytest.fixture
+def simple_model():
+  xml = """
+  <mujoco>
+    <worldbody>
+      <body>
+        <freejoint/>
+        <geom type="box" size="0.1 0.1 0.1"/>
+      </body>
+    </worldbody>
+  </mujoco>
+  """
+  return mujoco.MjModel.from_xml_string(xml)
+
+
+@pytest.fixture
+def mock_env_with_sim(simple_model):
+  env = Mock()
+  env.num_envs = 4
+  env.device = get_test_device()
+  env.episode_length_buf = torch.zeros(4, dtype=torch.long, device=env.device)
+  env.max_episode_length = 1000
+  env.sim = Simulation(
+    num_envs=env.num_envs, cfg=SimulationCfg(), model=simple_model, device=env.device
+  )
+  return env
+
+
+def test_nan_detection_function(mock_env_with_sim):
+  """Test that nan_detection correctly identifies NaN/Inf environments."""
+  env = mock_env_with_sim
+
+  # No NaNs initially.
+  result = nan_detection(env)
+  assert result.shape == (4,)
+  assert not result.any()
+
+  # Inject NaN in qpos for env 1.
+  env.sim.data.qpos[1, 0] = float("nan")
+  result = nan_detection(env)
+  assert result[1] and not result[0] and not result[2] and not result[3]
+
+  # Inject Inf in qacc_warmstart for env 3.
+  env.sim.data.qacc_warmstart[3, 0] = float("-inf")
+  result = nan_detection(env)
+  assert result[1] and result[3] and not result[0] and not result[2]
+
+
+def test_nan_detection_with_termination_manager(mock_env_with_sim):
+  """Test that nan_detection is properly logged by termination manager."""
+  env = mock_env_with_sim
+
+  cfg = {
+    "nan_term": TerminationTermCfg(func=nan_detection, params={}, time_out=False),
+  }
+
+  manager = TerminationManager(cfg, env)
+
+  # No terminations initially.
+  result = manager.compute()
+  assert not result.any()
+  assert not manager.terminated.any()
+  assert not manager.time_outs.any()
+
+  # Inject NaN in env 1.
+  env.sim.data.qpos[1, 0] = float("nan")
+
+  # Should detect termination in env 1.
+  result = manager.compute()
+  assert result[1] and not result[0] and not result[2] and not result[3]
+  assert manager.terminated[1]
+  assert not manager.time_outs[1]
+
+  # Reset should log the termination.
+  reset_info = manager.reset(torch.tensor([1], device=env.device))
+  assert "Episode_Termination/nan_term" in reset_info
+  assert reset_info["Episode_Termination/nan_term"] == 1
+
+  # Inject Inf in multiple envs.
+  env.sim.data.qvel[0, 0] = float("inf")
+  env.sim.data.qacc[2, 0] = float("-inf")
+
+  result = manager.compute()
+  assert result[0] and result[2]
+
+  # Reset should log multiple terminations.
+  reset_info = manager.reset(torch.tensor([0, 2], device=env.device))
+  assert reset_info["Episode_Termination/nan_term"] == 2
+
+
+@pytest.fixture
+def mock_env():
+  """Minimal mock environment for shape validation tests."""
+  env = Mock()
+  env.num_envs = 4
+  env.device = "cpu"
+  env.scene = {"robot": Mock()}
+  return env
+
+
+def test_termination_shape_validation_rejects_bad_compute_output(mock_env):
+  """Termination terms must return one boolean per environment."""
+  cfg = {
+    "bad": TerminationTermCfg(
+      func=lambda env: torch.zeros(env.num_envs, 1, dtype=torch.bool)
+    )
+  }
+  manager = TerminationManager(cfg, mock_env)
+  with pytest.raises(
+    ValueError, match="TerminationManager term 'bad'.*expected \\(4,\\)"
+  ):
+    manager.compute()
+
+
+@pytest.fixture
+def mock_terrain_env():
+  device = get_test_device()
+  num_envs = 4
+
+  env = Mock()
+  env.num_envs = num_envs
+  env.device = device
+  env.episode_length_buf = torch.full((num_envs,), 10, dtype=torch.long, device=device)
+
+  # Terrain with 8x8m sub-terrains.
+  terrain = Mock()
+  terrain.cfg.terrain_type = "generator"
+  terrain.cfg.terrain_generator.size = (8.0, 8.0)
+
+  env.scene.terrain = terrain
+  env.scene.env_origins = torch.zeros(num_envs, 3, device=device)
+
+  # Robot at spawn origin.
+  asset = Mock()
+  asset.data.root_link_pos_w = torch.zeros(num_envs, 3, device=device)
+  env.scene.__getitem__ = Mock(return_value=asset)
+
+  return env, asset
+
+
+def test_terrain_edge_reached_within_bounds(mock_terrain_env):
+  env, asset = mock_terrain_env
+  result = terrain_edge_reached(env, threshold_fraction=0.95)
+  assert not result.any()
+
+
+def test_terrain_edge_reached_at_edge(mock_terrain_env):
+  env, asset = mock_terrain_env
+  # Move env 1 past 95% of half-size (4.0 * 0.95 = 3.8m).
+  asset.data.root_link_pos_w[1, 0] = 3.9
+  result = terrain_edge_reached(env, threshold_fraction=0.95)
+  assert result[1] and not result[0] and not result[2] and not result[3]
+
+
+def test_terrain_edge_reached_skips_early_steps(mock_terrain_env):
+  env, asset = mock_terrain_env
+  asset.data.root_link_pos_w[0, 0] = 5.0
+  env.episode_length_buf[:] = 1  # Too early.
+  result = terrain_edge_reached(env, threshold_fraction=0.95)
+  assert not result.any()
+
+
+def test_terrain_edge_reached_no_terrain(mock_terrain_env):
+  env, asset = mock_terrain_env
+  env.scene.terrain = None
+  asset.data.root_link_pos_w[0, 0] = 100.0
+  result = terrain_edge_reached(env, threshold_fraction=0.95)
+  assert not result.any()
+
+
+@pytest.fixture
+def mock_grid_terrain_env():
+  device = get_test_device()
+  num_envs = 4
+
+  env = Mock()
+  env.num_envs = num_envs
+  env.device = device
+
+  # 10x10 grid of 8x8m sub-terrains → half_x = 40m, limit = 39.7m.
+  terrain = Mock()
+  terrain.cfg.terrain_type = "generator"
+  terrain.cfg.terrain_generator.size = (8.0, 8.0)
+  terrain.cfg.terrain_generator.border_width = 0.0
+  # Config num_rows/num_cols may differ from the effective grid (e.g. curriculum
+  # mode sets num_cols = len(sub_terrains) regardless of cfg.num_cols). The
+  # termination should read terrain_origins, not the config fields.
+  terrain.cfg.terrain_generator.num_rows = 999
+  terrain.cfg.terrain_generator.num_cols = 999
+  terrain.terrain_origins = torch.zeros(10, 10, 3, device=device)
+
+  env.scene.terrain = terrain
+
+  asset = Mock()
+  asset.data.root_link_pos_w = torch.zeros(num_envs, 3, device=device)
+  env.scene.__getitem__ = Mock(return_value=asset)
+
+  return env, asset
+
+
+def test_out_of_terrain_bounds_within(mock_grid_terrain_env):
+  env, asset = mock_grid_terrain_env
+  result = out_of_terrain_bounds(env)
+  assert not result.any()
+
+
+def test_out_of_terrain_bounds_outside(mock_grid_terrain_env):
+  env, asset = mock_grid_terrain_env
+  # Past the grid edge (limit_x = 40 - 0.3 = 39.7m).
+  asset.data.root_link_pos_w[2, 0] = 40.0
+  result = out_of_terrain_bounds(env)
+  assert result[2] and not result[0] and not result[1] and not result[3]
+
+
+def test_out_of_terrain_bounds_includes_border_width(mock_grid_terrain_env):
+  env, asset = mock_grid_terrain_env
+  # 2m flat border extends the footprint: half_x = 40 + 2 = 42m.
+  env.scene.terrain.cfg.terrain_generator.border_width = 2.0
+  asset.data.root_link_pos_w[0, 0] = 41.0  # Inside border, should not fire.
+  asset.data.root_link_pos_w[1, 0] = 42.0  # Past border edge (limit 41.7m).
+  result = out_of_terrain_bounds(env)
+  assert not result[0] and result[1]
@@ -0,0 +1,80 @@
+"""Tests for terrain configuration presets."""
+
+import mujoco
+import numpy as np
+import pytest
+
+import mjlab.terrains as terrain_gen
+from mjlab.terrains.config import (
+  ALL_TERRAIN_PRESETS,
+  ALL_TERRAINS_CFG,
+  ROUGH_TERRAINS_CFG,
+  STAIRS_TERRAINS_CFG,
+  pyramid_stairs,
+  terrain_preset,
+)
+from mjlab.terrains.terrain_generator import SubTerrainCfg
+
+
+def test_all_presets_return_sub_terrain_cfg():
+  for name, fn in ALL_TERRAIN_PRESETS.items():
+    cfg = fn(proportion=1.0)
+    assert isinstance(cfg, SubTerrainCfg), (
+      f"Preset {name!r} returned {type(cfg)}, expected SubTerrainCfg"
+    )
+
+
+def test_preset_overrides():
+  cfg = pyramid_stairs(proportion=0.5, step_width=0.5)
+  assert cfg.proportion == 0.5
+  assert cfg.step_width == 0.5
+  # Default should still apply for unoverridden fields.
+  assert cfg.platform_width == 3.0
+
+
+def test_rough_terrains_cfg_structure():
+  assert ROUGH_TERRAINS_CFG.size == (8.0, 8.0)
+  assert ROUGH_TERRAINS_CFG.num_rows == 10
+  assert ROUGH_TERRAINS_CFG.num_cols == 20
+  assert len(ROUGH_TERRAINS_CFG.sub_terrains) == 7
+  total = sum(c.proportion for c in ROUGH_TERRAINS_CFG.sub_terrains.values())
+  assert abs(total - 1.0) < 1e-6
+
+
+def test_stairs_terrains_cfg_structure():
+  assert STAIRS_TERRAINS_CFG.curriculum is True
+  assert len(STAIRS_TERRAINS_CFG.sub_terrains) == 4
+  assert "flat" in STAIRS_TERRAINS_CFG.sub_terrains
+  assert "easy_stairs" in STAIRS_TERRAINS_CFG.sub_terrains
+
+
+def test_all_terrains_cfg_matches_presets():
+  assert set(ALL_TERRAINS_CFG.sub_terrains.keys()) == set(ALL_TERRAIN_PRESETS.keys())
+  assert ALL_TERRAINS_CFG.num_cols == len(ALL_TERRAIN_PRESETS)
+
+
+def test_terrain_preset_decorator():
+  """Custom preset is registered in ALL_TERRAIN_PRESETS."""
+
+  @terrain_preset
+  def _test_custom(**overrides):
+    return terrain_gen.BoxFlatTerrainCfg(**overrides)
+
+  assert "_test_custom" in ALL_TERRAIN_PRESETS
+  cfg = ALL_TERRAIN_PRESETS["_test_custom"](proportion=0.5)
+  assert isinstance(cfg, SubTerrainCfg)
+  assert cfg.proportion == 0.5
+
+  # Clean up.
+  del ALL_TERRAIN_PRESETS["_test_custom"]
+
+
+@pytest.mark.slow
+def test_all_presets_generate_terrain():
+  """Each preset can generate terrain without error."""
+  for _name, fn in ALL_TERRAIN_PRESETS.items():
+    cfg = fn(proportion=1.0, size=(4.0, 4.0))
+    spec = mujoco.MjSpec()
+    spec.worldbody.add_body(name="terrain")
+    rng = np.random.default_rng(42)
+    cfg.function(difficulty=0.5, spec=spec, rng=rng)
@@ -0,0 +1,201 @@
+"""Tests for proportion-based robot spawning on terrain columns."""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+import torch
+
+from mjlab.terrains.primitive_terrains import BoxFlatTerrainCfg
+from mjlab.terrains.terrain_entity import TerrainEntity, TerrainEntityCfg
+from mjlab.terrains.terrain_generator import TerrainGenerator, TerrainGeneratorCfg
+
+
+@pytest.fixture(scope="module")
+def terrain_entity() -> TerrainEntity:
+  cfg = TerrainEntityCfg(terrain_type="plane", num_envs=1, env_spacing=2.0)
+  return TerrainEntity(cfg, device="cpu")
+
+
+def _make_origins(num_rows: int, num_cols: int) -> torch.Tensor:
+  origins = torch.zeros(num_rows, num_cols, 3)
+  for r in range(num_rows):
+    for c in range(num_cols):
+      origins[r, c, 0] = float(r)
+      origins[r, c, 1] = float(c)
+  return origins
+
+
+def _normalized(proportions: list[float]) -> np.ndarray:
+  p = np.array(proportions)
+  return p / p.sum()
+
+
+def _get_counts(entity: TerrainEntity, num_cols: int) -> list[int]:
+  return [int((entity.terrain_types == col).sum().item()) for col in range(num_cols)]
+
+
+# Fallback (no proportions).
+
+
+def test_even_distribution_without_proportions(terrain_entity: TerrainEntity) -> None:
+  num_envs = 30
+  num_cols = 3
+  origins = _make_origins(5, num_cols)
+
+  terrain_entity._compute_env_origins_curriculum(num_envs, origins, None)
+
+  assert _get_counts(terrain_entity, num_cols) == [10, 10, 10]
+
+
+# Proportional distribution.
+
+
+def test_basic_proportions(terrain_entity: TerrainEntity) -> None:
+  num_envs = 100
+  proportions = _normalized([0.5, 0.3, 0.2])
+  origins = _make_origins(5, len(proportions))
+
+  terrain_entity._compute_env_origins_curriculum(num_envs, origins, proportions)
+
+  assert _get_counts(terrain_entity, len(proportions)) == [50, 30, 20]
+
+
+@pytest.mark.parametrize("num_envs", [3, 7, 32, 64, 100, 512, 1024, 4096])
+def test_counts_sum_to_num_envs(terrain_entity: TerrainEntity, num_envs: int) -> None:
+  """Total terrain_types length must always equal num_envs."""
+  proportions = _normalized([0.6, 0.25, 0.1, 0.05])
+  origins = _make_origins(5, len(proportions))
+
+  terrain_entity._compute_env_origins_curriculum(num_envs, origins, proportions)
+
+  assert len(terrain_entity.terrain_types) == num_envs
+
+
+@pytest.mark.parametrize(
+  "num_envs, proportions",
+  [
+    (3, [0.98, 0.01, 0.01]),
+    (5, [0.9, 0.05, 0.03, 0.01, 0.01]),
+    (100, [0.9, 0.05, 0.05]),
+    (10, [0.97, 0.01, 0.01, 0.01]),
+  ],
+)
+def test_every_column_gets_at_least_one(
+  terrain_entity: TerrainEntity, num_envs: int, proportions: list[float]
+) -> None:
+  """When num_envs >= num_cols, every column gets at least 1 robot."""
+  p = _normalized(proportions)
+  origins = _make_origins(5, len(p))
+
+  terrain_entity._compute_env_origins_curriculum(num_envs, origins, p)
+
+  counts = _get_counts(terrain_entity, len(p))
+  for col, count in enumerate(counts):
+    assert count >= 1, f"Column {col} got 0 envs (counts={counts})"
+
+
+def test_fewer_envs_than_cols(terrain_entity: TerrainEntity) -> None:
+  """When num_envs < num_cols, should still work without errors."""
+  num_envs = 2
+  proportions = _normalized([0.5, 0.3, 0.1, 0.1])
+  origins = _make_origins(5, len(proportions))
+
+  result = terrain_entity._compute_env_origins_curriculum(
+    num_envs, origins, proportions
+  )
+
+  assert len(terrain_entity.terrain_types) == num_envs
+  assert result.shape == (num_envs, 3)
+
+
+def test_equal_proportions_gives_equal_counts(terrain_entity: TerrainEntity) -> None:
+  num_envs = 30
+  proportions = _normalized([1.0, 1.0, 1.0, 1.0])
+  origins = _make_origins(5, len(proportions))
+
+  terrain_entity._compute_env_origins_curriculum(num_envs, origins, proportions)
+
+  counts = _get_counts(terrain_entity, len(proportions))
+  for c in counts:
+    assert 7 <= c <= 8, f"Expected 7 or 8 envs per column, got {c}"
+
+
+def test_env_origins_shape(terrain_entity: TerrainEntity) -> None:
+  num_envs = 20
+  proportions = _normalized([0.5, 0.3, 0.2])
+  origins = _make_origins(5, len(proportions))
+
+  result = terrain_entity._compute_env_origins_curriculum(
+    num_envs, origins, proportions
+  )
+
+  assert result.shape == (num_envs, 3)
+
+
+def test_env_origins_match_terrain_origins(terrain_entity: TerrainEntity) -> None:
+  """Each env_origin must correspond to a valid (level, type) in origins."""
+  num_envs = 50
+  proportions = _normalized([0.6, 0.3, 0.1])
+  origins = _make_origins(5, len(proportions))
+
+  result = terrain_entity._compute_env_origins_curriculum(
+    num_envs, origins, proportions
+  )
+
+  for i in range(num_envs):
+    level = int(terrain_entity.terrain_levels[i].item())
+    ttype = int(terrain_entity.terrain_types[i].item())
+    expected = origins[level, ttype]
+    assert torch.allclose(result[i], expected), (
+      f"env {i}: origin {result[i].tolist()} != "
+      f"origins[{level},{ttype}]={expected.tolist()}"
+    )
+
+
+def test_mismatched_cols_falls_back_to_even(terrain_entity: TerrainEntity) -> None:
+  """When proportions length != num_cols, should fall back to even."""
+  num_envs = 30
+  num_cols = 3
+  # 4 proportions but only 3 columns in origins.
+  proportions = _normalized([0.4, 0.3, 0.2, 0.1])
+  origins = _make_origins(5, num_cols)
+
+  terrain_entity._compute_env_origins_curriculum(num_envs, origins, proportions)
+
+  assert _get_counts(terrain_entity, num_cols) == [10, 10, 10]
+
+
+def test_proportions_are_respected_approximately(
+  terrain_entity: TerrainEntity,
+) -> None:
+  """With large num_envs, actual distribution should track proportions."""
+  num_envs = 10000
+  raw = [0.5, 0.3, 0.15, 0.05]
+  proportions = _normalized(raw)
+  origins = _make_origins(5, len(proportions))
+
+  terrain_entity._compute_env_origins_curriculum(num_envs, origins, proportions)
+
+  counts = _get_counts(terrain_entity, len(proportions))
+  for col, (count, prop) in enumerate(zip(counts, proportions, strict=True)):
+    expected = prop * num_envs
+    assert abs(count - expected) < 10, (
+      f"Column {col}: got {count}, expected ~{expected:.0f}"
+    )
+
+
+def test_generator_does_not_mutate_cfg_num_cols() -> None:
+  """TerrainGenerator must not mutate the caller's cfg.num_cols."""
+  cfg = TerrainGeneratorCfg(
+    size=(4.0, 4.0),
+    curriculum=True,
+    num_cols=1,
+    sub_terrains={
+      "a": BoxFlatTerrainCfg(proportion=0.5, size=(4.0, 4.0)),
+      "b": BoxFlatTerrainCfg(proportion=0.3, size=(4.0, 4.0)),
+      "c": BoxFlatTerrainCfg(proportion=0.2, size=(4.0, 4.0)),
+    },
+  )
+  TerrainGenerator(cfg)
+  assert cfg.num_cols == 1, f"cfg.num_cols was mutated to {cfg.num_cols}"
@@ -0,0 +1,188 @@
+"""Tests for terrain normal fitting from sensor hit points."""
+
+from unittest.mock import MagicMock, PropertyMock
+
+import torch
+
+from mjlab.sensor import RayCastData, RayCastSensor
+from mjlab.tasks.velocity.mdp.terrain_utils import (
+  fit_terrain_normal,
+  terrain_normal_from_sensors,
+)
+
+
+def test_flat_ground():
+  """Points on z=0 plane → normal = [0, 0, 1]."""
+  B, N = 4, 10
+  points = torch.zeros(B, N, 3)
+  torch.manual_seed(0)
+  points[:, :, 0] = torch.randn(B, N)
+  points[:, :, 1] = torch.randn(B, N)
+  valid_mask = torch.ones(B, N, dtype=torch.bool)
+
+  normal = fit_terrain_normal(points, valid_mask)
+
+  assert normal.shape == (B, 3)
+  expected = torch.tensor([0.0, 0.0, 1.0])
+  for b in range(B):
+    torch.testing.assert_close(normal[b], expected, atol=1e-5, rtol=1e-5)
+
+
+def test_tilted_plane():
+  """Points on z = 0.5*x plane → normal perpendicular to that."""
+  B, N = 2, 20
+  torch.manual_seed(42)
+  points = torch.zeros(B, N, 3)
+  x = torch.randn(B, N)
+  y = torch.randn(B, N)
+  points[:, :, 0] = x
+  points[:, :, 1] = y
+  points[:, :, 2] = 0.5 * x
+  valid_mask = torch.ones(B, N, dtype=torch.bool)
+
+  normal = fit_terrain_normal(points, valid_mask)
+
+  # Normal to plane z = 0.5*x is (-0.5, 0, 1) normalized.
+  expected_raw = torch.tensor([-0.5, 0.0, 1.0])
+  expected = expected_raw / expected_raw.norm()
+  assert normal.shape == (B, 3)
+  for b in range(B):
+    torch.testing.assert_close(normal[b], expected, atol=1e-4, rtol=1e-4)
+
+
+def test_partial_misses():
+  """Half invalid points with junk z values → still correct normal."""
+  B, N = 2, 20
+  torch.manual_seed(42)
+  points = torch.zeros(B, N, 3)
+  points[:, :, 0] = torch.randn(B, N)
+  points[:, :, 1] = torch.randn(B, N)
+  valid_mask = torch.ones(B, N, dtype=torch.bool)
+  # Invalidate even-indexed points and put junk in them.
+  valid_mask[:, ::2] = False
+  points[:, ::2, 2] = 999.0
+
+  normal = fit_terrain_normal(points, valid_mask)
+
+  expected = torch.tensor([0.0, 0.0, 1.0])
+  assert normal.shape == (B, 3)
+  for b in range(B):
+    torch.testing.assert_close(normal[b], expected, atol=1e-4, rtol=1e-4)
+
+
+def test_fewer_than_3_valid_fallback():
+  """Fewer than 3 valid points (including zero) falls back to [0, 0, 1]."""
+  B, N = 3, 10
+  points = torch.randn(B, N, 3)
+  valid_mask = torch.zeros(B, N, dtype=torch.bool)
+  # Batch 0: 0 valid, batch 1: 1 valid, batch 2: 2 valid.
+  valid_mask[1, 0] = True
+  valid_mask[2, :2] = True
+
+  normal = fit_terrain_normal(points, valid_mask)
+
+  expected = torch.tensor([0.0, 0.0, 1.0])
+  for b in range(B):
+    torch.testing.assert_close(normal[b], expected, atol=1e-6, rtol=1e-6)
+
+
+def test_collinear_points_fallback():
+  """Collinear points don't define a plane, should fall back to [0, 0, 1]."""
+  B, N = 2, 10
+  points = torch.zeros(B, N, 3)
+  # All points along the X axis.
+  points[:, :, 0] = torch.linspace(0, 1, N)
+  valid_mask = torch.ones(B, N, dtype=torch.bool)
+
+  normal = fit_terrain_normal(points, valid_mask)
+
+  expected = torch.tensor([0.0, 0.0, 1.0])
+  for b in range(B):
+    torch.testing.assert_close(normal[b], expected, atol=1e-6, rtol=1e-6)
+
+
+def test_non_finite_points_fallback():
+  """Non-finite hits (from a diverged env) must not crash eigh; fall back to up."""
+  B, N = 3, 10
+  points = torch.zeros(B, N, 3)
+  # Batch 0: a valid tilted plane.
+  points[0, :, 0] = torch.linspace(0, 1, N)
+  points[0, :, 1] = torch.linspace(0, 1, N)
+  points[0, :, 2] = 0.2 * points[0, :, 0]
+  # Batch 1: all NaN. Batch 2: all Inf.
+  points[1] = float("nan")
+  points[2] = float("inf")
+  valid_mask = torch.ones(B, N, dtype=torch.bool)
+
+  normal = fit_terrain_normal(points, valid_mask)
+
+  assert torch.isfinite(normal).all()
+  expected_up = torch.tensor([0.0, 0.0, 1.0])
+  torch.testing.assert_close(normal[1], expected_up, atol=1e-6, rtol=1e-6)
+  torch.testing.assert_close(normal[2], expected_up, atol=1e-6, rtol=1e-6)
+  # The valid plane is still fit and points upward.
+  assert normal[0, 2] > 0
+  torch.testing.assert_close(normal[0].norm(), torch.tensor(1.0), atol=1e-4, rtol=0)
+
+
+def test_small_valid_plane_still_fits():
+  """A tiny but planar patch should not be rejected as degenerate."""
+  x_vals = torch.linspace(0.0, 1e-4, 4)
+  y_vals = torch.linspace(0.0, 8e-5, 3)
+  xx, yy = torch.meshgrid(x_vals, y_vals, indexing="ij")
+  patch = torch.stack((xx.reshape(-1), yy.reshape(-1)), dim=-1)
+  N = patch.shape[0]
+  B = 2
+  points = torch.zeros(B, N, 3)
+  points[:, :, 0] = patch[:, 0]
+  points[:, :, 1] = patch[:, 1]
+  points[:, :, 2] = 0.5 * patch[:, 0]
+  valid_mask = torch.ones(B, N, dtype=torch.bool)
+
+  normal = fit_terrain_normal(points, valid_mask)
+
+  expected_raw = torch.tensor([-0.5, 0.0, 1.0])
+  expected = expected_raw / expected_raw.norm()
+  for b in range(B):
+    torch.testing.assert_close(normal[b], expected, atol=1e-4, rtol=1e-4)
+
+
+def test_terrain_normal_from_sensors():
+  """Mock env/sensors, verify subsampling + concatenation."""
+  B = 4
+  torch.manual_seed(7)
+
+  # RayCastSensor mock: 100 rays, all hits on z=0.
+  raycast_hit_pos = torch.zeros(B, 100, 3)
+  raycast_hit_pos[:, :, 0] = torch.randn(B, 100)
+  raycast_hit_pos[:, :, 1] = torch.randn(B, 100)
+  raycast_distances = torch.ones(B, 100)  # all valid
+
+  raycast_sensor = MagicMock(spec=RayCastSensor)
+  raycast_data = RayCastData(
+    distances=raycast_distances,
+    normals_w=torch.zeros(B, 100, 3),
+    hit_pos_w=raycast_hit_pos,
+    pos_w=torch.zeros(B, 3),
+    quat_w=torch.zeros(B, 4),
+    frame_pos_w=torch.zeros(B, 1, 3),
+    frame_quat_w=torch.zeros(B, 1, 4),
+  )
+  type(raycast_sensor).data = PropertyMock(return_value=raycast_data)
+
+  # Mock env.
+  sensors = {"raycast": raycast_sensor}
+  env = MagicMock()
+  env.scene.__getitem__ = MagicMock(side_effect=lambda name: sensors[name])
+
+  normal = terrain_normal_from_sensors(
+    env,
+    sensor_names=("raycast",),
+    max_points=16,
+  )
+
+  # 16 subsampled raycast points on z=0.
+  assert normal.shape == (B, 3)
+  expected = torch.tensor([0.0, 0.0, 1.0])
+  for b in range(B):
+    torch.testing.assert_close(normal[b], expected, atol=1e-4, rtol=1e-4)
@@ -0,0 +1,180 @@
+"""Tests for terrain generation."""
+
+from dataclasses import dataclass
+
+import mujoco
+import numpy as np
+import pytest
+
+from mjlab.terrains.config import ALL_TERRAIN_PRESETS
+from mjlab.terrains.primitive_terrains import (
+  _MIN_BORDER_HEIGHT,
+  BoxInvertedPyramidStairsTerrainCfg,
+  BoxPyramidStairsTerrainCfg,
+  BoxSteppingStonesTerrainCfg,
+)
+from mjlab.terrains.terrain_generator import (
+  SubTerrainCfg,
+  TerrainGenerator,
+  TerrainGeneratorCfg,
+  TerrainGeometry,
+  TerrainOutput,
+)
+
+_CFG = BoxSteppingStonesTerrainCfg(
+  proportion=1.0,
+  size=(8.0, 8.0),
+  stone_size_range=(0.2, 0.6),
+  stone_distance_range=(0.05, 0.25),
+  stone_height=0.2,
+  stone_height_variation=0.05,
+  stone_size_variation=0.05,
+  displacement_range=0.1,
+  floor_depth=2.0,
+  platform_width=1.5,
+  border_width=0.25,
+)
+
+
+def _generate_stones(
+  cfg: BoxSteppingStonesTerrainCfg,
+  difficulty: float,
+  rng: np.random.Generator,
+) -> list[tuple[float, float, float, float]]:
+  """Generate terrain and return stone (cx, cy, half_x, half_y) tuples."""
+  spec = mujoco.MjSpec()
+  spec.worldbody.add_body(name="terrain")
+  output = cfg.function(difficulty=difficulty, spec=spec, rng=rng)
+
+  center = cfg.size[0] / 2
+  stones = []
+  for geom_info in output.geometries:
+    geom = geom_info.geom
+    if geom is None:
+      continue
+    pos, size = geom.pos, geom.size
+    # Skip platform, floor, and border geoms. The platform is the geom centered
+    # exactly at the patch center (its size is grid-snapped, not the configured
+    # width, so it is identified by position alone).
+    is_platform = np.isclose(pos[0], center) and np.isclose(pos[1], center)
+    is_full_span = np.isclose(size[0], cfg.size[0] / 2) or np.isclose(
+      size[1], cfg.size[1] / 2
+    )
+    if is_platform or is_full_span:
+      continue
+    stones.append((pos[0], pos[1], size[0], size[1]))
+  return stones
+
+
+def test_no_stone_centers_inside_platform():
+  """No stone center should fall inside the platform."""
+  center = _CFG.size[0] / 2
+  p_half = _CFG.platform_width / 2
+  p_min, p_max = center - p_half, center + p_half
+
+  for difficulty in [0.0, 0.5, 1.0]:
+    stones = _generate_stones(_CFG, difficulty, np.random.default_rng(42))
+    for cx, cy, _, _ in stones:
+      assert not (p_min <= cx <= p_max and p_min <= cy <= p_max), (
+        f"Stone at ({cx:.3f}, {cy:.3f}) inside platform at difficulty={difficulty}"
+      )
+
+
+def test_stone_size_decreases_with_difficulty():
+  """Average stone size should be smaller at higher difficulty."""
+  sizes = {}
+  for difficulty in [0.0, 1.0]:
+    stones = _generate_stones(_CFG, difficulty, np.random.default_rng(42))
+    sizes[difficulty] = np.mean([hx + hy for _, _, hx, hy in stones])
+
+  assert sizes[0.0] > sizes[1.0]
+
+
+@pytest.mark.parametrize(
+  "cfg_cls", [BoxPyramidStairsTerrainCfg, BoxInvertedPyramidStairsTerrainCfg]
+)
+def test_pyramid_stairs_border_present_at_zero_difficulty(cfg_cls):
+  """At difficulty 0 the step height collapses to 0, but the flat border frame
+  must still be generated as solid, non-degenerate geometry (regression for the
+  empty-boundary bug, issue #1033)."""
+  cfg = cfg_cls(
+    size=(8.0, 8.0),
+    step_height_range=(0.0, 0.2),
+    step_width=0.3,
+    platform_width=3.0,
+    border_width=1.0,
+  )
+  spec = mujoco.MjSpec()
+  spec.worldbody.add_body(name="terrain")
+  output = cfg.function(difficulty=0.0, spec=spec, rng=np.random.default_rng(0))
+
+  # The border frame sits below z=0 (top flush at ground level); inner step
+  # boxes are centered at z=0. Identify the frame by its downward offset.
+  border_geoms = [
+    g.geom for g in output.geometries if g.geom is not None and g.geom.pos[2] < -1e-4
+  ]
+  assert len(border_geoms) == 4, "Expected four border frame boxes."
+  for geom in border_geoms:
+    # Each frame box must be solid, not a degenerate zero-height geom, and its
+    # top must be flush with the ground plane at z=0.
+    assert geom.size[2] >= _MIN_BORDER_HEIGHT / 2 - 1e-9
+    assert np.isclose(geom.pos[2] + geom.size[2], 0.0, atol=1e-6)
+
+
+@pytest.mark.parametrize("preset_name", sorted(ALL_TERRAIN_PRESETS))
+@pytest.mark.parametrize("difficulty", [0.0, 1.0])
+def test_preset_compiles_across_difficulty(preset_name, difficulty):
+  """Every terrain preset must generate compilable MuJoCo geometry across the
+  full difficulty range. Difficulty 0 is exercised explicitly because curriculum
+  row 0 lands there deterministically, which previously produced degenerate
+  geometry (zero-height hfields, NaN colors, missing borders)."""
+  cfg = ALL_TERRAIN_PRESETS[preset_name](size=(8.0, 8.0))
+  spec = mujoco.MjSpec()
+  spec.worldbody.add_body(name="terrain")
+  cfg.function(difficulty=difficulty, spec=spec, rng=np.random.default_rng(0))
+  # Compiling validates geom/hfield sizes and rgba values (catches NaNs and
+  # non-positive sizes that MuJoCo rejects).
+  spec.compile()
+
+
+@dataclass(kw_only=True)
+class _BoxesTerrainCfg(SubTerrainCfg):
+  """Test terrain that emits two boxes, named per `names` (None leaves unnamed)."""
+
+  names: tuple[str | None, str | None] = (None, None)
+
+  def function(
+    self, difficulty: float, spec: mujoco.MjSpec, rng: np.random.Generator
+  ) -> TerrainOutput:
+    del difficulty, rng
+    body = spec.body("terrain")
+    geometries = []
+    for i in range(2):
+      geom = body.add_geom(
+        type=mujoco.mjtGeom.mjGEOM_BOX,
+        size=(0.5, 0.5, 0.5),
+        pos=(i, 0, 0),
+        name=self.names[i],
+      )
+      geometries.append(TerrainGeometry(geom=geom))
+    return TerrainOutput(origin=np.zeros(3), geometries=geometries)
+
+
+def test_compile_preserves_custom_names_and_defaults_unnamed_geoms():
+  """Geoms with a custom name from a sub-terrain must be preserved, while
+  unnamed geoms must fall back to the generator's default "terrain_{i}" name."""
+  cfg = TerrainGeneratorCfg(
+    size=(4.0, 4.0),
+    curriculum=True,
+    sub_terrains={
+      "named": _BoxesTerrainCfg(size=(4.0, 4.0), names=("foo_0", "foo_1")),
+      "unnamed": _BoxesTerrainCfg(size=(4.0, 4.0)),
+    },
+  )
+  generator = TerrainGenerator(cfg)
+  spec = mujoco.MjSpec()
+  generator.compile(spec)
+  model = spec.compile()
+
+  names = [model.geom(i).name for i in range(model.ngeom)]
+  assert names == ["foo_0", "foo_1", "terrain_2", "terrain_3"]
@@ -0,0 +1,179 @@
+"""Tests for motion tracking evaluation metrics."""
+
+import math
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from mjlab.tasks.tracking.mdp.metrics import (
+  compute_ee_orientation_error,
+  compute_ee_position_error,
+  compute_joint_velocity_error,
+  compute_mpkpe,
+  compute_root_relative_mpkpe,
+)
+
+
+@pytest.fixture
+def mock_command():
+  """Create a mock MotionCommand for testing."""
+  command = Mock()
+  command.num_envs = 4
+  command.device = "cpu"
+  command.cfg = Mock()
+  command.cfg.body_names = (
+    "pelvis",
+    "left_knee",
+    "right_knee",
+    "left_ankle",
+    "right_ankle",
+    "left_wrist",
+    "right_wrist",
+  )
+  return command
+
+
+def test_mpkpe_zero_when_positions_match(mock_command):
+  """Test MPKPE is zero when global positions are identical."""
+  num_bodies = len(mock_command.cfg.body_names)
+  positions = torch.rand(mock_command.num_envs, num_bodies, 3)
+
+  mock_command.body_pos_w = positions.clone()
+  mock_command.robot_body_pos_w = positions.clone()
+
+  mpkpe = compute_mpkpe(mock_command)
+
+  assert mpkpe.shape == (mock_command.num_envs,)
+  assert torch.allclose(mpkpe, torch.zeros(mock_command.num_envs), atol=1e-6)
+
+
+def test_mpkpe_correct_error(mock_command):
+  """Test MPKPE computes the correct mean global error."""
+  num_bodies = len(mock_command.cfg.body_names)
+
+  mock_command.body_pos_w = torch.zeros(mock_command.num_envs, num_bodies, 3)
+  mock_command.robot_body_pos_w = torch.zeros(mock_command.num_envs, num_bodies, 3)
+  mock_command.robot_body_pos_w[:, :, 0] = 1.0  # 1 unit offset in x
+
+  mpkpe = compute_mpkpe(mock_command)
+
+  assert torch.allclose(mpkpe, torch.ones(mock_command.num_envs), atol=1e-6)
+
+
+def test_mpkpe_uses_global_reference(mock_command):
+  """MPKPE must read the global reference, not the drift-cancelled one.
+
+  Pins issue #1006: setting body_pos_relative_w to match the robot exactly
+  would yield zero error if it were (incorrectly) used; the metric must
+  instead follow body_pos_w.
+  """
+  num_bodies = len(mock_command.cfg.body_names)
+  robot_pos = torch.rand(mock_command.num_envs, num_bodies, 3)
+  mock_command.robot_body_pos_w = robot_pos.clone()
+  mock_command.body_pos_relative_w = robot_pos.clone()  # zero error if misused
+  mock_command.body_pos_w = robot_pos.clone()
+  mock_command.body_pos_w[:, :, 0] += 1.0  # 1 unit of global drift
+
+  mpkpe = compute_mpkpe(mock_command)
+
+  assert torch.allclose(mpkpe, torch.ones(mock_command.num_envs), atol=1e-6)
+
+
+def test_r_mpkpe_zero_when_relative_positions_match(mock_command):
+  """R-MPKPE is zero when re-anchored positions are identical."""
+  num_bodies = len(mock_command.cfg.body_names)
+  positions = torch.rand(mock_command.num_envs, num_bodies, 3)
+
+  mock_command.body_pos_relative_w = positions.clone()
+  mock_command.robot_body_pos_w = positions.clone()
+
+  r_mpkpe = compute_root_relative_mpkpe(mock_command)
+
+  assert r_mpkpe.shape == (mock_command.num_envs,)
+  assert torch.allclose(r_mpkpe, torch.zeros(mock_command.num_envs), atol=1e-6)
+
+
+def test_r_mpkpe_uses_relative_reference(mock_command):
+  """R-MPKPE reads the re-anchored reference, not the global one.
+
+  Setting body_pos_w to match the robot exactly would yield zero error if
+  it were (incorrectly) used; the metric must instead follow
+  body_pos_relative_w.
+  """
+  num_bodies = len(mock_command.cfg.body_names)
+  robot_pos = torch.rand(mock_command.num_envs, num_bodies, 3)
+  mock_command.robot_body_pos_w = robot_pos.clone()
+  mock_command.body_pos_w = robot_pos.clone()  # zero error if misused
+  mock_command.body_pos_relative_w = robot_pos.clone()
+  mock_command.body_pos_relative_w[:, :, 0] += 1.0  # 1 unit of local pose error
+
+  r_mpkpe = compute_root_relative_mpkpe(mock_command)
+
+  assert torch.allclose(r_mpkpe, torch.ones(mock_command.num_envs), atol=1e-6)
+
+
+def test_joint_velocity_error_rms(mock_command):
+  """Joint velocity error is the per-joint RMS of the velocity error."""
+  num_joints = 3
+
+  mock_command.joint_vel = torch.zeros(mock_command.num_envs, num_joints)
+  mock_command.robot_joint_vel = torch.zeros(mock_command.num_envs, num_joints)
+  mock_command.robot_joint_vel[:, 0] = 3.0
+  mock_command.robot_joint_vel[:, 1] = 4.0  # Error [3, 4, 0]
+
+  error = compute_joint_velocity_error(mock_command)
+
+  expected = math.sqrt((3.0**2 + 4.0**2 + 0.0**2) / num_joints)
+  assert torch.allclose(error, torch.ones(mock_command.num_envs) * expected, atol=1e-6)
+
+
+def test_ee_position_error_only_uses_specified_bodies(mock_command):
+  """Test EE position error only uses specified bodies."""
+  num_bodies = len(mock_command.cfg.body_names)
+
+  mock_command.body_pos_relative_w = torch.zeros(mock_command.num_envs, num_bodies, 3)
+  mock_command.robot_body_pos_w = torch.zeros(mock_command.num_envs, num_bodies, 3)
+
+  # Large error for pelvis (not an EE).
+  mock_command.robot_body_pos_w[:, 0, :] = 100.0
+  # Small error for ankles.
+  mock_command.robot_body_pos_w[:, 3, 0] = 1.0  # left_ankle
+  mock_command.robot_body_pos_w[:, 4, 0] = 1.0  # right_ankle
+
+  error = compute_ee_position_error(mock_command, ("left_ankle", "right_ankle"))
+
+  # Should only reflect ankle error, not pelvis.
+  assert torch.allclose(error, torch.ones(mock_command.num_envs), atol=1e-6)
+
+
+def test_ee_orientation_error_detects_rotation(mock_command):
+  """Test EE orientation error detects rotations."""
+  num_bodies = len(mock_command.cfg.body_names)
+
+  identity_quat = torch.tensor([1.0, 0.0, 0.0, 0.0])
+  mock_command.body_quat_relative_w = (
+    identity_quat.view(1, 1, 4).expand(mock_command.num_envs, num_bodies, 4).clone()
+  )
+
+  # 90 degree rotation around z-axis.
+  rotated_quat = torch.tensor([0.7071, 0.0, 0.0, 0.7071])
+  mock_command.robot_body_quat_w = (
+    rotated_quat.view(1, 1, 4).expand(mock_command.num_envs, num_bodies, 4).clone()
+  )
+
+  error = compute_ee_orientation_error(mock_command, ("left_wrist",))
+
+  # Error should be approximately pi/2 radians.
+  expected = torch.ones(mock_command.num_envs) * (3.14159 / 2)
+  assert torch.allclose(error, expected, atol=0.01)
+
+
+def test_ee_metrics_raise_on_unknown_body(mock_command):
+  """Unknown end-effector names raise instead of silently scoring zero."""
+  num_bodies = len(mock_command.cfg.body_names)
+  mock_command.body_pos_relative_w = torch.zeros(mock_command.num_envs, num_bodies, 3)
+  mock_command.robot_body_pos_w = torch.zeros(mock_command.num_envs, num_bodies, 3)
+
+  with pytest.raises(ValueError, match="not tracked"):
+    compute_ee_position_error(mock_command, ("nonexistent_body",))
@@ -0,0 +1,134 @@
+"""Tests specific to motion tracking tasks."""
+
+import pytest
+
+from mjlab.asset_zoo.robots import G1_ACTION_SCALE
+from mjlab.envs.mdp.actions import JointPositionActionCfg
+from mjlab.tasks.registry import list_tasks, load_env_cfg
+from mjlab.tasks.tracking.mdp import MotionCommandCfg
+
+
+@pytest.fixture(scope="module")
+def tracking_task_ids() -> list[str]:
+  """Get all tracking task IDs."""
+  return [t for t in list_tasks() if "Tracking" in t]
+
+
+@pytest.fixture(scope="module")
+def g1_tracking_task_ids(tracking_task_ids: list[str]) -> list[str]:
+  """Get all G1 tracking task IDs."""
+  return [t for t in tracking_task_ids if "G1" in t]
+
+
+def test_tracking_tasks_have_motion_command(tracking_task_ids: list[str]) -> None:
+  """All tracking tasks should have a 'motion' command of type MotionCommandCfg."""
+  for task_id in tracking_task_ids:
+    cfg = load_env_cfg(task_id)
+
+    assert "motion" in cfg.commands, f"Task {task_id} missing 'motion' command"
+
+    motion_cmd = cfg.commands["motion"]
+    assert isinstance(motion_cmd, MotionCommandCfg), (
+      f"Task {task_id} motion command is not MotionCommandCfg"
+    )
+
+
+def test_tracking_tasks_have_self_collision_sensor(
+  tracking_task_ids: list[str],
+) -> None:
+  """All tracking tasks should have a self_collision sensor."""
+  for task_id in tracking_task_ids:
+    cfg = load_env_cfg(task_id)
+
+    assert cfg.scene.sensors is not None, f"Task {task_id} has no sensors"
+
+    sensor_names = {s.name for s in cfg.scene.sensors}
+    assert "self_collision" in sensor_names, (
+      f"Task {task_id} missing self_collision sensor"
+    )
+
+
+def test_tracking_no_state_estimation_observations() -> None:
+  """No-state-estimation tasks remove observations that depend on state estimation."""
+  task_id = "Mjlab-Tracking-Flat-Unitree-G1-No-State-Estimation"
+
+  # Test both training and play modes
+  for play_mode in [False, True]:
+    cfg = load_env_cfg(task_id, play=play_mode)
+    mode_str = "play mode" if play_mode else "training mode"
+
+    assert "actor" in cfg.observations, (
+      f"Task {task_id} ({mode_str}) missing policy observations"
+    )
+    actor_terms = cfg.observations["actor"].terms
+
+    assert "motion_anchor_pos_b" not in actor_terms, (
+      f"Task {task_id} ({mode_str}) has motion_anchor_pos_b in policy, "
+      "expected it to be removed for no-state-estimation variant"
+    )
+    assert "base_lin_vel" not in actor_terms, (
+      f"Task {task_id} ({mode_str}) has base_lin_vel in policy, "
+      "expected it to be removed for no-state-estimation variant"
+    )
+
+
+def test_tracking_play_disables_rsi_randomization() -> None:
+  """Tracking play tasks should disable RSI randomization."""
+  tracking_tasks = [
+    "Mjlab-Tracking-Flat-Unitree-G1",
+    "Mjlab-Tracking-Flat-Unitree-G1-No-State-Estimation",
+  ]
+
+  for task_id in tracking_tasks:
+    cfg = load_env_cfg(task_id, play=True)
+
+    motion_cmd = cfg.commands["motion"]
+    assert isinstance(motion_cmd, MotionCommandCfg), (
+      f"Task {task_id} (play mode) motion command is not MotionCommandCfg"
+    )
+
+    assert motion_cmd.pose_range == {}, (
+      f"Task {task_id} (play mode) has non-empty pose_range={motion_cmd.pose_range}, "
+      "expected empty dict for disabled RSI"
+    )
+    assert motion_cmd.velocity_range == {}, (
+      f"Task {task_id} (play mode) has non-empty velocity_range={motion_cmd.velocity_range}, "
+      "expected empty dict for disabled RSI"
+    )
+
+
+def test_tracking_play_uses_start_sampling_mode() -> None:
+  """Tracking play tasks should use sampling_mode='start'."""
+  tracking_tasks = [
+    "Mjlab-Tracking-Flat-Unitree-G1",
+    "Mjlab-Tracking-Flat-Unitree-G1-No-State-Estimation",
+  ]
+
+  for task_id in tracking_tasks:
+    cfg = load_env_cfg(task_id, play=True)
+
+    motion_cmd = cfg.commands["motion"]
+    assert isinstance(motion_cmd, MotionCommandCfg), (
+      f"Task {task_id} (play mode) motion command is not MotionCommandCfg"
+    )
+
+    assert motion_cmd.sampling_mode == "start", (
+      f"Task {task_id} (play mode) sampling_mode={motion_cmd.sampling_mode}, expected 'start'"
+    )
+
+
+def test_g1_tracking_has_correct_action_scale(g1_tracking_task_ids: list[str]) -> None:
+  """G1 tracking tasks should use G1_ACTION_SCALE."""
+  for task_id in g1_tracking_task_ids:
+    cfg = load_env_cfg(task_id)
+
+    assert "joint_pos" in cfg.actions, f"Task {task_id} missing 'joint_pos' action"
+
+    joint_pos_action = cfg.actions["joint_pos"]
+    assert isinstance(joint_pos_action, JointPositionActionCfg), (
+      f"Task {task_id} joint_pos action is not JointPositionActionCfg"
+    )
+
+    assert joint_pos_action.scale == G1_ACTION_SCALE, (
+      f"Task {task_id} action scale mismatch, expected G1_ACTION_SCALE"
+    )
@@ -0,0 +1,30 @@
+import tyro
+
+from mjlab.envs import ManagerBasedRlEnvCfg
+from mjlab.rl import RslRlOnPolicyRunnerCfg
+from mjlab.scene import SceneCfg
+from mjlab.scripts.train import TrainConfig
+from mjlab.terrains import TerrainEntityCfg
+
+
+def test_tyro_defaults():
+  """Regression test for #677 due to incompatibilities between declared types and
+  default values of the env configs.
+  """
+  tyro.cli(
+    TrainConfig,
+    default=TrainConfig(
+      env=ManagerBasedRlEnvCfg(
+        decimation=5,
+        scene=SceneCfg(
+          terrain=TerrainEntityCfg(),
+        ),
+      ),
+      agent=RslRlOnPolicyRunnerCfg(),
+    ),
+    config=(
+      tyro.conf.AvoidSubcommands,
+      tyro.conf.FlagConversionOff,
+    ),
+    args=[],
+  )
@@ -0,0 +1,76 @@
+"""Test that all tyro.cli() calls include mjlab.TYRO_FLAGS."""
+
+import ast
+from pathlib import Path
+
+import pytest
+
+ROOT = Path(__file__).parent.parent
+SEARCH_DIRS = [ROOT / "src", ROOT / "scripts"]
+
+
+def find_tyro_cli_calls() -> list[tuple[Path, int, bool]]:
+  """Find all tyro.cli() calls and check if they have TYRO_FLAGS.
+
+  Returns a list of (filepath, line_number, has_tyro_flags) tuples.
+  """
+  results = []
+
+  for search_dir in SEARCH_DIRS:
+    if not search_dir.exists():
+      continue
+    for filepath in search_dir.rglob("*.py"):
+      try:
+        source = filepath.read_text()
+        tree = ast.parse(source, filename=str(filepath))
+      except SyntaxError:
+        continue
+
+      for node in ast.walk(tree):
+        if not isinstance(node, ast.Call):
+          continue
+
+        # Check if this is a tyro.cli() call
+        func = node.func
+        is_tyro_cli = False
+
+        if isinstance(func, ast.Attribute) and func.attr == "cli":
+          # tyro.cli(...)
+          if isinstance(func.value, ast.Name) and func.value.id == "tyro":
+            is_tyro_cli = True
+        elif isinstance(func, ast.Name) and func.id == "cli":
+          # from tyro import cli; cli(...)
+          is_tyro_cli = True
+
+        if not is_tyro_cli:
+          continue
+
+        # Check if config= keyword argument contains TYRO_FLAGS
+        has_tyro_flags = False
+        for keyword in node.keywords:
+          if keyword.arg == "config":
+            config_source = ast.unparse(keyword.value)
+            if "TYRO_FLAGS" in config_source:
+              has_tyro_flags = True
+              break
+
+        results.append((filepath, node.lineno, has_tyro_flags))
+
+  return results
+
+
+def test_tyro_cli_calls_have_tyro_flags():
+  """Ensure all tyro.cli() calls include config=mjlab.TYRO_FLAGS."""
+  all_calls = find_tyro_cli_calls()
+  assert all_calls, "No tyro.cli() calls found"
+
+  violations = [
+    (filepath, line) for filepath, line, has_flags in all_calls if not has_flags
+  ]
+
+  if violations:
+    msg = "tyro.cli() calls missing config=mjlab.TYRO_FLAGS:\n"
+    for filepath, line in violations:
+      rel_path = filepath.relative_to(ROOT)
+      msg += f"  {rel_path}:{line}\n"
+    pytest.fail(msg)
@@ -0,0 +1,120 @@
+"""Tests for the velocity command's initial-velocity injection.
+
+The init_velocity_prob path runs inside the reset pipeline, after reset
+events wrote the new pose to qpos but before sim.forward(), so it must not
+read (or write back) derived kinematics.
+"""
+
+from types import SimpleNamespace
+from typing import TYPE_CHECKING, cast
+
+import pytest
+import torch
+from conftest import get_test_device, load_fixture_xml, make_scene_and_sim
+
+from mjlab.tasks.velocity.mdp.velocity_command import UniformVelocityCommandCfg
+
+if TYPE_CHECKING:
+  from mjlab.envs import ManagerBasedRlEnv
+
+
+@pytest.fixture(scope="module")
+def device():
+  return get_test_device()
+
+
+def test_init_velocity_preserves_fresh_reset_pose(device):
+  scene, sim = make_scene_and_sim(
+    device, load_fixture_xml("floating_base_articulated"), sensors=(), num_envs=2
+  )
+  env = cast(
+    "ManagerBasedRlEnv",
+    SimpleNamespace(scene=scene, sim=sim, num_envs=2, device=device),
+  )
+  cfg = UniformVelocityCommandCfg(
+    entity_name="robot",
+    resampling_time_range=(1e9, 1e9),
+    init_velocity_prob=1.0,
+    rel_heading_envs=0.0,
+    ranges=UniformVelocityCommandCfg.Ranges(
+      lin_vel_x=(0.5, 0.5), lin_vel_y=(0.2, 0.2), ang_vel_z=(0.0, 0.0)
+    ),
+  )
+  term = cfg.build(env)
+  robot = scene["robot"]
+  env_ids = torch.arange(2, device=device)
+
+  # Derived kinematics now hold the spawn pose (the "previous episode" state).
+  sim.forward()
+
+  # Emulate a reset event: write a fresh pose to qpos, no forward yet.
+  pose = torch.tensor(
+    [
+      [1.0, 2.0, 1.5, 1.0, 0.0, 0.0, 0.0],
+      [3.0, -1.0, 1.5, 1.0, 0.0, 0.0, 0.0],
+    ],
+    device=device,
+  )
+  robot.write_root_link_pose_to_sim(pose, env_ids=env_ids)
+
+  term.reset(env_ids=env_ids)
+  sim.forward()
+
+  # The fresh reset pose survives. Before the fix, the init-velocity path
+  # wrote the stale pre-reset pose back into the sim.
+  assert torch.allclose(robot.data.root_link_pos_w, pose[:, :3], atol=1e-6)
+
+  # Planar velocity matches the sampled command (identity orientation, so
+  # body frame equals world frame).
+  assert torch.allclose(
+    robot.data.root_link_lin_vel_b[:, :2],
+    term.vel_command_b[:, :2],
+    atol=1e-5,
+  )
+
+
+def test_mid_episode_resample_does_not_write_velocity(device):
+  """Init velocity applies on reset only; a timer-expiry resample runs after
+  step()'s forward and must not write sim state."""
+  scene, sim = make_scene_and_sim(
+    device, load_fixture_xml("floating_base_articulated"), sensors=(), num_envs=2
+  )
+  env = cast(
+    "ManagerBasedRlEnv",
+    SimpleNamespace(scene=scene, sim=sim, num_envs=2, device=device, step_dt=0.02),
+  )
+  cfg = UniformVelocityCommandCfg(
+    entity_name="robot",
+    resampling_time_range=(0.001, 0.001),  # Expires on the first compute.
+    init_velocity_prob=1.0,
+    rel_heading_envs=0.0,
+    ranges=UniformVelocityCommandCfg.Ranges(
+      lin_vel_x=(0.7, 0.7), lin_vel_y=(0.3, 0.3), ang_vel_z=(0.0, 0.0)
+    ),
+  )
+  term = cfg.build(env)
+  robot = scene["robot"]
+  env_ids = torch.arange(2, device=device)
+
+  term.reset(env_ids=env_ids)
+  sim.forward()
+
+  # Mid-episode the robot has decelerated to rest.
+  robot.write_root_link_velocity_to_sim(
+    torch.zeros(2, 6, device=device), env_ids=env_ids
+  )
+  sim.forward()
+
+  # Step's command compute: the 1 ms timer expires and resamples.
+  counter = term.command_counter.clone()
+  term.compute(dt=1.0)
+  assert (term.command_counter > counter).all()
+
+  # Only the command changed; sim velocity is untouched.
+  qvel_lin = sim.data.qvel[:, robot.indexing.free_joint_v_adr[:3]]
+  assert torch.allclose(qvel_lin, torch.zeros_like(qvel_lin), atol=1e-6)
+  assert torch.allclose(
+    robot.data.root_link_lin_vel_b,
+    torch.zeros_like(robot.data.root_link_lin_vel_b),
+    atol=1e-6,
+  )
@@ -0,0 +1,190 @@
+"""Tests for velocity task reward functions."""
+
+from __future__ import annotations
+
+import math
+from unittest.mock import MagicMock, PropertyMock
+
+import torch
+
+from mjlab.managers.reward_manager import RewardTermCfg
+from mjlab.managers.scene_entity_config import SceneEntityCfg
+from mjlab.sensor import RayCastData, RayCastSensor
+from mjlab.tasks.velocity.mdp.rewards import upright
+from mjlab.utils.lab_api.math import quat_from_euler_xyz
+
+
+def _identity_quat(B: int) -> torch.Tensor:
+  """(w, x, y, z) = (1, 0, 0, 0)."""
+  q = torch.zeros(B, 4)
+  q[:, 0] = 1.0
+  return q
+
+
+def _quat_from_roll(roll_rad: float, B: int = 1) -> torch.Tensor:
+  roll = torch.full((B,), roll_rad)
+  zero = torch.zeros(B)
+  return quat_from_euler_xyz(roll, zero, zero)
+
+
+def _quat_from_pitch(pitch_rad: float, B: int = 1) -> torch.Tensor:
+  pitch = torch.full((B,), pitch_rad)
+  zero = torch.zeros(B)
+  return quat_from_euler_xyz(zero, pitch, zero)
+
+
+def _make_env_and_reward(
+  terrain_sensor_names: tuple[str, ...] | None = None,
+  body_quat_w: torch.Tensor | None = None,
+  terrain_hit_z: float = 0.0,
+  terrain_slope_x: float = 0.0,
+):
+  """Build mocked env + upright reward instance.
+
+  Args:
+    terrain_sensor_names: If set, enables terrain-aware mode.
+    body_quat_w: [B, 4] root orientation. Defaults to identity.
+    terrain_hit_z: Z value for flat terrain hits.
+    terrain_slope_x: Slope in X (z = terrain_slope_x * x).
+  """
+  B = 1 if body_quat_w is None else body_quat_w.shape[0]
+  if body_quat_w is None:
+    body_quat_w = _identity_quat(B)
+
+  # Mock asset data. Use explicit asset_cfg with no body_names so
+  # body_ids stays None and the reward uses root_link_quat_w.
+  asset = MagicMock()
+  asset.data.root_link_quat_w = body_quat_w
+  asset.data.root_link_pos_w = torch.zeros(B, 3)
+  asset.data.gravity_vec_w = torch.tensor([0.0, 0.0, -1.0]).expand(B, 3)
+  asset_cfg = SceneEntityCfg("robot", body_names=None, body_ids=[])
+
+  # Mock terrain sensor if needed.
+  sensors: dict = {"robot": asset}
+  if terrain_sensor_names is not None:
+    N = 100
+    torch.manual_seed(0)
+    hit_pos = torch.zeros(B, N, 3)
+    hit_pos[:, :, 0] = torch.randn(B, N)
+    hit_pos[:, :, 1] = torch.randn(B, N)
+    hit_pos[:, :, 2] = terrain_hit_z + terrain_slope_x * hit_pos[:, :, 0]
+
+    raycast_sensor = MagicMock(spec=RayCastSensor)
+    raycast_data = RayCastData(
+      distances=torch.ones(B, N),
+      normals_w=torch.zeros(B, N, 3),
+      hit_pos_w=hit_pos,
+      pos_w=torch.zeros(B, 3),
+      quat_w=torch.zeros(B, 4),
+      frame_pos_w=torch.zeros(B, 1, 3),
+      frame_quat_w=torch.zeros(B, 1, 4),
+    )
+    type(raycast_sensor).data = PropertyMock(return_value=raycast_data)
+    for name in terrain_sensor_names:
+      sensors[name] = raycast_sensor
+
+  env = MagicMock()
+  env.scene.__getitem__ = MagicMock(side_effect=lambda n: sensors[n])
+
+  params: dict = {"std": 1.0, "asset_cfg": asset_cfg}
+  if terrain_sensor_names is not None:
+    params["terrain_sensor_names"] = terrain_sensor_names
+  cfg = MagicMock(spec=RewardTermCfg)
+  cfg.params = params
+
+  reward_fn = upright(cfg, env)
+  return env, reward_fn, params
+
+
+def test_world_up_identity_gives_max_reward():
+  """Perfectly upright robot on flat ground → reward ≈ 1."""
+  env, reward, params = _make_env_and_reward()
+  r = reward(env, std=params["std"], asset_cfg=params["asset_cfg"])
+  assert r.shape == (1,)
+  assert r.item() > 0.99
+
+
+def test_world_up_tilted_gives_lower_reward():
+  """30° roll → reward significantly below 1."""
+  quat = _quat_from_roll(math.radians(30))
+  env, reward, params = _make_env_and_reward(body_quat_w=quat)
+  r = reward(env, std=params["std"], asset_cfg=params["asset_cfg"])
+  assert r.item() < 0.8
+
+
+def test_terrain_aware_aligned_with_slope():
+  """Robot pitched to match a slope → terrain-aware reward ≈ 1."""
+  slope = 0.5  # z = 0.5 * x
+  tilt = math.atan(slope)  # Pitch to match slope in XZ plane.
+  quat = _quat_from_pitch(-tilt)
+  env, reward, params = _make_env_and_reward(
+    terrain_sensor_names=("terrain_scan",),
+    body_quat_w=quat,
+    terrain_slope_x=slope,
+  )
+  r = reward(
+    env,
+    std=params["std"],
+    asset_cfg=params["asset_cfg"],
+    terrain_sensor_names=params["terrain_sensor_names"],
+  )
+  # Should be close to 1 since robot matches terrain.
+  assert r.item() > 0.9
+
+
+def test_terrain_aware_upright_on_slope_penalized():
+  """Robot staying vertical on a slope → terrain-aware reward < 1."""
+  slope = 0.5
+  quat = _identity_quat(1)  # Robot is world-vertical, not matching slope.
+  env, reward, params = _make_env_and_reward(
+    terrain_sensor_names=("terrain_scan",),
+    body_quat_w=quat,
+    terrain_slope_x=slope,
+  )
+  r = reward(
+    env,
+    std=params["std"],
+    asset_cfg=params["asset_cfg"],
+    terrain_sensor_names=params["terrain_sensor_names"],
+  )
+  # Should be penalized since robot doesn't match terrain.
+  assert r.item() < 0.95
+
+
+def test_terrain_aware_flat_ground_matches_world_up():
+  """On flat terrain, terrain-aware and world-up should give same reward."""
+  quat = _quat_from_roll(math.radians(15))
+  env_t, reward_t, params_t = _make_env_and_reward(
+    terrain_sensor_names=("terrain_scan",),
+    body_quat_w=quat,
+  )
+  env_w, reward_w, params_w = _make_env_and_reward(body_quat_w=quat)
+
+  r_terrain = reward_t(
+    env_t,
+    std=params_t["std"],
+    asset_cfg=params_t["asset_cfg"],
+    terrain_sensor_names=params_t["terrain_sensor_names"],
+  )
+  r_world = reward_w(env_w, std=params_w["std"], asset_cfg=params_w["asset_cfg"])
+
+  torch.testing.assert_close(r_terrain, r_world, atol=0.02, rtol=0.02)
+
+
+def test_batch_consistency():
+  """Multiple envs with different orientations get independent rewards."""
+  B = 4
+  quats = torch.zeros(B, 4)
+  quats[:, 0] = 1.0  # All identity.
+  # Tilt env 2 by 45°.
+  quats[2] = _quat_from_roll(math.radians(45))[0]
+
+  env, reward, params = _make_env_and_reward(body_quat_w=quats)
+  r = reward(env, std=params["std"], asset_cfg=params["asset_cfg"])
+
+  assert r.shape == (B,)
+  # Env 0, 1, 3 should be ~1, env 2 should be lower.
+  assert r[0].item() > 0.99
+  assert r[1].item() > 0.99
+  assert r[2].item() < 0.7
+  assert r[3].item() > 0.99
@@ -0,0 +1,199 @@
+"""Tests specific to velocity tasks."""
+
+import pytest
+
+from mjlab.asset_zoo.robots import G1_ACTION_SCALE, GO1_ACTION_SCALE
+from mjlab.envs.mdp.actions import JointPositionActionCfg
+from mjlab.tasks.registry import list_tasks, load_env_cfg
+from mjlab.tasks.velocity.mdp import UniformVelocityCommandCfg
+
+
+@pytest.fixture(scope="module")
+def velocity_task_ids() -> list[str]:
+  """Get all velocity task IDs."""
+  return [t for t in list_tasks() if "Velocity" in t]
+
+
+@pytest.fixture(scope="module")
+def g1_velocity_task_ids(velocity_task_ids: list[str]) -> list[str]:
+  """Get all G1 velocity task IDs."""
+  return [t for t in velocity_task_ids if "G1" in t]
+
+
+@pytest.fixture(scope="module")
+def go1_velocity_task_ids(velocity_task_ids: list[str]) -> list[str]:
+  """Get all Go1 velocity task IDs."""
+  return [t for t in velocity_task_ids if "Go1" in t]
+
+
+@pytest.fixture(scope="module")
+def rough_velocity_task_ids(velocity_task_ids: list[str]) -> list[str]:
+  """Get all rough terrain velocity task IDs."""
+  return [t for t in velocity_task_ids if "Rough" in t]
+
+
+@pytest.fixture(scope="module")
+def flat_velocity_task_ids(velocity_task_ids: list[str]) -> list[str]:
+  """Get all flat terrain velocity task IDs."""
+  return [t for t in velocity_task_ids if "Flat" in t]
+
+
+def test_velocity_tasks_have_twist_command(velocity_task_ids: list[str]) -> None:
+  """All velocity tasks should have a velocity command."""
+  for task_id in velocity_task_ids:
+    cfg = load_env_cfg(task_id)
+
+    assert "twist" in cfg.commands, f"Task {task_id} missing 'twist' command"
+
+    twist_cmd = cfg.commands["twist"]
+    assert isinstance(twist_cmd, UniformVelocityCommandCfg), (
+      f"Task {task_id} twist command is not UniformVelocityCommandCfg"
+    )
+
+
+def test_g1_velocity_has_required_sensors(g1_velocity_task_ids: list[str]) -> None:
+  """G1 velocity tasks should have feet/ground and self collision sensors."""
+  for task_id in g1_velocity_task_ids:
+    cfg = load_env_cfg(task_id)
+
+    assert cfg.scene.sensors is not None, f"Task {task_id} has no sensors"
+
+    sensor_names = {s.name for s in cfg.scene.sensors}
+    assert "feet_ground_contact" in sensor_names, (
+      f"Task {task_id} missing feet_ground_contact sensor"
+    )
+    assert "self_collision" in sensor_names, (
+      f"Task {task_id} missing self_collision sensor"
+    )
+
+
+def test_go1_velocity_has_required_sensors(go1_velocity_task_ids: list[str]) -> None:
+  """Go1 velocity tasks should have feet/ground and collision sensors."""
+  for task_id in go1_velocity_task_ids:
+    cfg = load_env_cfg(task_id)
+
+    assert cfg.scene.sensors is not None, f"Task {task_id} has no sensors"
+
+    sensor_names = {s.name for s in cfg.scene.sensors}
+    assert "feet_ground_contact" in sensor_names, (
+      f"Task {task_id} missing feet_ground_contact sensor"
+    )
+    if "Rough" in task_id:
+      for name in (
+        "self_collision",
+        "thigh_ground_touch",
+        "shank_ground_touch",
+        "trunk_ground_touch",
+      ):
+        assert name in sensor_names, f"Task {task_id} missing {name} sensor"
+
+
+def test_flat_velocity_tasks_have_plane_terrain(
+  flat_velocity_task_ids: list[str],
+) -> None:
+  """Flat velocity tasks should have terrain_type='plane' and no terrain_generator."""
+  for task_id in flat_velocity_task_ids:
+    cfg = load_env_cfg(task_id)
+
+    assert cfg.scene.terrain is not None, f"Task {task_id} has no terrain config"
+    assert cfg.scene.terrain.terrain_type == "plane", (
+      f"Task {task_id} terrain_type={cfg.scene.terrain.terrain_type}, expected 'plane'"
+    )
+    assert cfg.scene.terrain.terrain_generator is None, (
+      f"Task {task_id} has terrain_generator, expected None for flat terrain"
+    )
+
+
+def test_rough_velocity_tasks_have_generator_terrain(
+  rough_velocity_task_ids: list[str],
+) -> None:
+  """Rough velocity tasks should have generator terrain."""
+  for task_id in rough_velocity_task_ids:
+    cfg = load_env_cfg(task_id)
+
+    assert cfg.scene.terrain is not None, f"Task {task_id} has no terrain config"
+    assert cfg.scene.terrain.terrain_type == "generator", (
+      f"Task {task_id} terrain_type={cfg.scene.terrain.terrain_type}, "
+      "expected 'generator'"
+    )
+    assert cfg.scene.terrain.terrain_generator is not None, (
+      f"Task {task_id} has no terrain_generator, expected one for rough terrain"
+    )
+
+
+def test_rough_velocity_training_has_curriculum_enabled() -> None:
+  """Rough velocity training tasks should have terrain curriculum enabled."""
+  rough_training_tasks = [
+    "Mjlab-Velocity-Rough-Unitree-G1",
+    "Mjlab-Velocity-Rough-Unitree-Go1",
+  ]
+
+  for task_id in rough_training_tasks:
+    cfg = load_env_cfg(task_id)
+
+    assert cfg.scene.terrain is not None, f"Task {task_id} has no terrain config"
+    assert cfg.scene.terrain.terrain_generator is not None, (
+      f"Task {task_id} has no terrain_generator"
+    )
+    assert cfg.scene.terrain.terrain_generator.curriculum is True, (
+      f"Task {task_id} curriculum={cfg.scene.terrain.terrain_generator.curriculum}, "
+      "expected True"
+    )
+
+
+def test_rough_velocity_play_has_curriculum_disabled() -> None:
+  """Rough velocity play tasks should have terrain curriculum disabled."""
+  rough_training_tasks = [
+    "Mjlab-Velocity-Rough-Unitree-G1",
+    "Mjlab-Velocity-Rough-Unitree-Go1",
+  ]
+
+  for task_id in rough_training_tasks:
+    cfg = load_env_cfg(task_id, play=True)
+
+    assert cfg.scene.terrain is not None, (
+      f"Task {task_id} (play mode) has no terrain config"
+    )
+    assert cfg.scene.terrain.terrain_generator is not None, (
+      f"Task {task_id} (play mode) has no terrain_generator"
+    )
+    assert cfg.scene.terrain.terrain_generator.curriculum is False, (
+      f"Task {task_id} (play mode) curriculum={cfg.scene.terrain.terrain_generator.curriculum}, "
+      "expected False"
+    )
+
+
+def test_g1_velocity_has_correct_action_scale(g1_velocity_task_ids: list[str]) -> None:
+  """G1 velocity tasks should use G1_ACTION_SCALE."""
+  for task_id in g1_velocity_task_ids:
+    cfg = load_env_cfg(task_id)
+
+    assert "joint_pos" in cfg.actions, f"Task {task_id} missing 'joint_pos' action"
+
+    joint_pos_action = cfg.actions["joint_pos"]
+    assert isinstance(joint_pos_action, JointPositionActionCfg), (
+      f"Task {task_id} joint_pos action is not JointPositionActionCfg"
+    )
+
+    assert joint_pos_action.scale == G1_ACTION_SCALE, (
+      f"Task {task_id} action scale mismatch, expected G1_ACTION_SCALE"
+    )
+
+
+def test_go1_velocity_has_correct_action_scale(
+  go1_velocity_task_ids: list[str],
+) -> None:
+  """Go1 velocity tasks should use GO1_ACTION_SCALE."""
+  for task_id in go1_velocity_task_ids:
+    cfg = load_env_cfg(task_id)
+
+    assert "joint_pos" in cfg.actions, f"Task {task_id} missing 'joint_pos' action"
+
+    joint_pos_action = cfg.actions["joint_pos"]
+    assert isinstance(joint_pos_action, JointPositionActionCfg), (
+      f"Task {task_id} joint_pos action is not JointPositionActionCfg"
+    )
+
+    assert joint_pos_action.scale == GO1_ACTION_SCALE, (
+      f"Task {task_id} action scale mismatch, expected GO1_ACTION_SCALE"
+    )
@@ -0,0 +1,79 @@
+"""Tests for the velocity terrain-level curriculum."""
+
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import torch
+
+from mjlab.managers.scene_entity_config import SceneEntityCfg
+from mjlab.tasks.velocity.mdp.curriculums import terrain_levels_vel
+
+
+class _FakeTerrain:
+  """Minimal terrain stub implementing the real level-update arithmetic."""
+
+  def __init__(self, num_envs: int, size: float = 8.0):
+    self.terrain_levels = torch.zeros(num_envs, dtype=torch.long)
+    self.terrain_types = torch.zeros(num_envs, dtype=torch.long)
+    self.max_terrain_level = 10
+    # [num_rows, num_cols, 3]. One column, enough rows for promotion.
+    self.terrain_origins = torch.zeros(self.max_terrain_level, 1, 3)
+    self.env_origins = torch.zeros(num_envs, 3)
+    self.cfg = SimpleNamespace(
+      terrain_generator=SimpleNamespace(size=(size, size), sub_terrains={})
+    )
+
+  def update_env_origins(self, env_ids, move_up, move_down):
+    # Mirror TerrainEntity.update_env_origins exactly, including the cap-wrap
+    # branch and the env_origins write, so the fake can't pass vacuously.
+    self.terrain_levels[env_ids] += 1 * move_up - 1 * move_down
+    self.terrain_levels[env_ids] = torch.where(
+      self.terrain_levels[env_ids] >= self.max_terrain_level,
+      torch.randint_like(self.terrain_levels[env_ids], self.max_terrain_level),
+      torch.clip(self.terrain_levels[env_ids], 0),
+    )
+    self.env_origins[env_ids] = self.terrain_origins[
+      self.terrain_levels[env_ids], self.terrain_types[env_ids]
+    ]
+
+
+def _make_env(terrain: _FakeTerrain, common_step_counter: int, walked: float):
+  num_envs = terrain.terrain_levels.shape[0]
+  # Robot has "walked" `walked` meters in x from its origin.
+  asset = Mock()
+  asset.data.root_link_pos_w = torch.tensor([[walked, 0.0, 0.0]] * num_envs)
+
+  env = Mock()
+  env.common_step_counter = common_step_counter
+  env.max_episode_length_s = 20.0
+  env.scene.__getitem__ = Mock(return_value=asset)
+  env.scene.terrain = terrain
+  env.scene.env_origins = torch.zeros(num_envs, 3)
+  # Zero command so move_down stays inactive; isolates move_up behavior.
+  env.command_manager.get_command = Mock(return_value=torch.zeros(num_envs, 2))
+  return env
+
+
+def test_first_reset_does_not_promote_levels():
+  """On the initial reset the far spawn/origin gap must not bump levels."""
+  terrain = _FakeTerrain(num_envs=4, size=8.0)
+  # walked >> size/2 (=4.0), which would normally trigger move_up.
+  env = _make_env(terrain, common_step_counter=0, walked=100.0)
+
+  terrain_levels_vel(
+    env, torch.arange(4), command_name="twist", asset_cfg=SceneEntityCfg("robot")
+  )
+
+  assert torch.all(terrain.terrain_levels == 0)
+
+
+def test_subsequent_reset_promotes_on_long_walk():
+  """After stepping, a long walk still promotes the level as before."""
+  terrain = _FakeTerrain(num_envs=4, size=8.0)
+  env = _make_env(terrain, common_step_counter=1, walked=100.0)
+
+  terrain_levels_vel(
+    env, torch.arange(4), command_name="twist", asset_cfg=SceneEntityCfg("robot")
+  )
+
+  assert torch.all(terrain.terrain_levels == 1)
@@ -0,0 +1,79 @@
+"""Tests for video recording with mediapy."""
+
+from pathlib import Path
+from unittest.mock import Mock
+
+import mediapy as media
+import numpy as np
+import torch
+
+
+def _make_mock_env(num_envs: int = 1):
+  """Create a mock environment that produces random RGB frames."""
+  env = Mock()
+  env.render_mode = "rgb_array"
+  env.metadata = {"render_fps": 30}
+  env.render.return_value = np.random.default_rng().integers(
+    0, 255, (num_envs, 64, 64, 3), dtype=np.uint8
+  )
+  env.step.return_value = (
+    torch.zeros(num_envs),  # obs
+    torch.zeros(num_envs),  # reward
+    torch.zeros(num_envs, dtype=torch.bool),  # terminated
+    torch.zeros(num_envs, dtype=torch.bool),  # truncated
+    {},  # info
+  )
+  env.close.return_value = None
+  env.unwrapped = env
+  return env
+
+
+def test_step_trigger_writes_video(tmp_path: Path):
+  """VideoRecorder writes a readable mp4 when the step trigger fires."""
+  from mjlab.utils.wrappers.video_recorder import VideoRecorder
+
+  env = _make_mock_env()
+  recorder = VideoRecorder(
+    env,
+    video_folder=tmp_path,
+    step_trigger=lambda step: step == 0,
+    video_length=5,
+    disable_logger=True,
+  )
+
+  action = torch.zeros(1)
+  for _ in range(6):
+    recorder.step(action)
+
+  recorder.close()
+
+  videos = list(tmp_path.glob("*.mp4"))
+  assert len(videos) == 1
+
+  # Verify the file is a valid video readable by mediapy.
+  frames = media.read_video(str(videos[0]))
+  assert len(frames) == 5
+  assert frames[0].shape == (64, 64, 3)
+
+
+def test_accepts_string_path(tmp_path: Path):
+  """VideoRecorder accepts a string path for video_folder."""
+  from mjlab.utils.wrappers.video_recorder import VideoRecorder
+
+  env = _make_mock_env()
+  folder = str(tmp_path / "vids")
+  recorder = VideoRecorder(
+    env,
+    video_folder=folder,
+    step_trigger=lambda step: step == 0,
+    video_length=3,
+    disable_logger=True,
+  )
+
+  action = torch.zeros(1)
+  for _ in range(4):
+    recorder.step(action)
+
+  recorder.close()
+
+  assert list(Path(folder).glob("*.mp4"))
@@ -0,0 +1,269 @@
+"""Tests for the single-accumulator BaseViewer in base_alt.py."""
+
+from __future__ import annotations
+
+import time
+from unittest.mock import MagicMock
+
+from mjlab.viewer.base import BaseViewer
+
+
+class FakeViewer(BaseViewer):
+  """Minimal concrete viewer for testing."""
+
+  def __init__(self, step_dt: float = 0.01, frame_rate: float = 60.0):
+    env = MagicMock()
+    env.unwrapped.step_dt = step_dt
+    env.cfg.viewer = MagicMock()
+    super().__init__(env, MagicMock(return_value=MagicMock()), frame_rate=frame_rate)
+    self.sim_step_count = 0
+    self.render_count = 0
+    self._last_tick_time = time.perf_counter()
+
+  def setup(self) -> None: ...
+
+  def sync_env_to_viewer(self) -> None:
+    self.render_count += 1
+
+  def sync_viewer_to_env(self) -> None: ...
+  def close(self) -> None: ...
+  def is_running(self) -> bool:
+    return True
+
+  def _execute_step(self) -> bool:
+    self.sim_step_count += 1
+    self._step_count += 1
+    self._stats_steps += 1
+    return True
+
+  def inject_tick(self, dt: float) -> bool:
+    """Call tick() with a controlled dt."""
+    now = time.perf_counter()
+    self._last_tick_time = now - dt
+    return self.tick()
+
+
+# Physics stepping.
+
+
+def test_stepping():
+  """Physics steps match sim-time budget."""
+  v = FakeViewer(step_dt=0.01)
+  v.inject_tick(dt=0.01)
+  assert v.sim_step_count == 1
+
+  v.inject_tick(dt=0.03)
+  assert v.sim_step_count == 4
+
+  v.inject_tick(dt=0.0)
+  assert v.sim_step_count == 4
+
+
+def test_accumulator_carries():
+  """Fractional budget carries across ticks."""
+  v = FakeViewer(step_dt=0.01)
+  v.inject_tick(dt=0.01)
+  assert v.sim_step_count == 1
+
+  v.inject_tick(dt=0.015)
+  assert v.sim_step_count == 2  # 0.015 budget, 1 step, 0.005 carry
+
+  v.inject_tick(dt=0.015)
+  assert v.sim_step_count == 4  # 0.005+0.015=0.02, 2 steps
+
+
+def test_high_speed():
+  v = FakeViewer(step_dt=0.01, frame_rate=60.0)
+  v._speed_index = v.SPEED_MULTIPLIERS.index(8.0)
+  v._time_multiplier = 8.0
+  v.inject_tick(dt=1.0 / 60.0)
+  assert v.sim_step_count == 13
+
+
+def test_slow_speed():
+  v = FakeViewer(step_dt=0.01, frame_rate=60.0)
+  v._speed_index = 0
+  v._time_multiplier = 1 / 32
+  for _ in range(19):
+    v.inject_tick(dt=1.0 / 60.0)
+  assert v.sim_step_count == 0
+  for _ in range(3):
+    v.inject_tick(dt=1.0 / 60.0)
+  assert v.sim_step_count >= 1
+
+
+# Render timing.
+
+
+def test_render_at_frame_rate():
+  v = FakeViewer(step_dt=0.01, frame_rate=60.0)
+  assert v.inject_tick(dt=0.001) is True  # First tick renders.
+  assert v.inject_tick(dt=0.001) is False  # Too soon.
+  assert v.inject_tick(dt=1.0 / 60.0) is True  # Frame time elapsed.
+
+
+# Pause and resume.
+
+
+def test_pause_stops_physics():
+  v = FakeViewer(step_dt=0.01)
+  v.inject_tick(dt=0.02)
+  before = v.sim_step_count
+
+  v.pause()
+  v.inject_tick(dt=0.5)
+  assert v.sim_step_count == before
+
+
+def test_resume_no_burst():
+  v = FakeViewer(step_dt=0.01)
+  v.inject_tick(dt=0.03)
+  count_before = v.sim_step_count
+
+  v.pause()
+  v._last_error = "some error"
+
+  v.resume()
+  assert v._last_error is None
+  assert v._sim_budget == 0.0
+
+  v.inject_tick(dt=0.01)
+  assert v.sim_step_count - count_before == 1
+
+
+# Single step.
+
+
+def test_single_step_while_paused():
+  v = FakeViewer(step_dt=0.01)
+  v.pause()
+
+  v.request_single_step()
+  v.inject_tick(dt=0.0)
+
+  assert v.sim_step_count == 1
+  assert v._is_paused
+
+
+def test_single_step_ignored_when_running():
+  v = FakeViewer(step_dt=0.01)
+  v.inject_tick(dt=0.02)
+  before = v.sim_step_count
+
+  v.request_single_step()
+  v.inject_tick(dt=0.02)
+  assert v.sim_step_count > before
+
+
+# Error recovery.
+
+
+def test_error_pauses():
+  v = FakeViewer(step_dt=0.01)
+  v._execute_step = BaseViewer._execute_step.__get__(v, FakeViewer)  # type: ignore[attr-defined]
+  v.policy = MagicMock(side_effect=RuntimeError("test error"))
+
+  v.inject_tick(dt=0.02)
+
+  assert v._is_paused
+  assert v._last_error is not None
+  assert "test error" in v._last_error
+
+
+def test_reset_clears():
+  v = FakeViewer(step_dt=0.01)
+  v.inject_tick(dt=0.02)
+  v._last_error = "err"
+
+  v.reset_environment()
+  assert v._step_count == 0
+  assert v._sim_budget == 0.0
+  assert v._last_error is None
+
+
+def test_reset_calls_policy_reset():
+  """reset_environment() calls policy.reset() if available."""
+  v = FakeViewer(step_dt=0.01)
+  v.policy = MagicMock()
+  v.policy.reset = MagicMock()
+
+  v.reset_environment()
+  v.policy.reset.assert_called_once()
+
+
+def test_reset_without_policy_reset():
+  """reset_environment() works fine when policy has no reset method."""
+  v = FakeViewer(step_dt=0.01)
+  v.policy = lambda obs: obs  # plain callable, no reset attribute
+
+  v.reset_environment()  # should not raise
+
+
+# Formatting and status.
+
+
+def test_format_speed():
+  assert BaseViewer._format_speed(1.0) == "1x"
+  assert BaseViewer._format_speed(0.5) == "1/2x"
+  assert BaseViewer._format_speed(0.25) == "1/4x"
+  assert BaseViewer._format_speed(1 / 32) == "1/32x"
+
+
+def test_status_snapshot():
+  v = FakeViewer(step_dt=0.01)
+  v._fps = 60.0
+  v._sps = 50.0
+  v._last_error = "err"
+
+  status = v.get_status()
+  assert abs(status.actual_realtime - 0.5) < 1e-10
+  assert status.speed_label == "1x"
+  assert status.last_error == "err"
+
+
+def test_capped_clears_each_tick():
+  """Capped resets to False each tick; only True when deadline hit."""
+  v = FakeViewer(step_dt=0.01)
+  v.inject_tick(dt=0.02)
+  assert not v._was_capped
+
+  # Force capped state, then verify next tick clears it.
+  v._was_capped = True
+  v.inject_tick(dt=0.02)
+  assert not v._was_capped
+
+
+def test_capped_false_when_no_remaining_budget():
+  """Deadline hit with no remaining budget is NOT capped (transient stall)."""
+  v = FakeViewer(step_dt=0.01, frame_rate=60.0)
+  # dt=0.01 at 1x: budget=0.01, exactly 1 step, budget goes to 0.
+  # Even if deadline fires (via GC), no remaining work to drop.
+  v.inject_tick(dt=0.01)
+  assert not v._was_capped
+
+
+# Spiral protection.
+
+
+def test_no_spiral():
+  v = FakeViewer(step_dt=0.01, frame_rate=60.0)
+  v._speed_index = v.SPEED_MULTIPLIERS.index(8.0)
+  v._time_multiplier = 8.0
+
+  original = v._execute_step
+
+  def slow_step() -> bool:
+    time.sleep(0.005)
+    return original()
+
+  v._execute_step = slow_step  # type: ignore[assignment]
+
+  start = time.perf_counter()
+  v.inject_tick(dt=1.0 / 60.0)
+  elapsed = time.perf_counter() - start
+
+  assert elapsed < 0.05
+  assert v.sim_step_count < 13
+  assert v.sim_step_count >= 1
+  # Capped because slow steps + large budget means remaining work was dropped.
+  assert v._was_capped
@@ -0,0 +1,573 @@
+"""Test viser conversion functions with various robot models."""
+
+from pathlib import Path
+
+import mujoco
+
+from mjlab.viewer.viser import (
+  create_site_mesh,
+  get_geom_texture_id,
+  group_geoms_by_visual_compat,
+  merge_geoms,
+  merge_sites,
+  mujoco_mesh_to_trimesh,
+)
+
+
+def load_robot_model(robot_name: str) -> mujoco.MjModel:
+  """Load a robot model from the asset zoo."""
+  base_path = Path(__file__).parent.parent / "src/mjlab/asset_zoo/robots"
+
+  # Map robot names to their XML files.
+  robot_paths = {
+    "unitree_g1": base_path / "unitree_g1/xmls/g1.xml",
+    "unitree_go1": base_path / "unitree_go1/xmls/go1.xml",
+  }
+
+  if robot_name not in robot_paths:
+    raise ValueError(f"Unknown robot: {robot_name}")
+
+  xml_path = robot_paths[robot_name]
+  if not xml_path.exists():
+    raise FileNotFoundError(f"Robot XML not found: {xml_path}")
+
+  return mujoco.MjModel.from_xml_path(str(xml_path))
+
+
+def test_unitree_g1_conversion():
+  """Test conversion with Unitree G1 robot."""
+  model = load_robot_model("unitree_g1")
+
+  mesh_geom_count = 0
+  has_textures = False
+
+  for geom_idx in range(model.ngeom):
+    if model.geom_type[geom_idx] == mujoco.mjtGeom.mjGEOM_MESH:
+      mesh_geom_count += 1
+
+      # Convert to trimesh.
+      mesh = mujoco_mesh_to_trimesh(model, geom_idx)
+
+      # Basic checks.
+      assert mesh is not None, f"Failed to convert geom {geom_idx}"
+      assert len(mesh.vertices) > 0, f"Mesh {geom_idx} has no vertices"
+      assert len(mesh.faces) > 0, f"Mesh {geom_idx} has no faces"
+
+      # Check for textures.
+      if hasattr(mesh.visual, "uv"):
+        has_textures = True
+
+  assert mesh_geom_count > 0, "No mesh geometries found in Unitree G1"
+  print(f"✓ Unitree G1: Successfully converted {mesh_geom_count} mesh geometries")
+  if has_textures:
+    print("  - Found textured meshes")
+
+
+def test_unitree_go1_conversion():
+  """Test conversion with Unitree Go1 robot."""
+  model = load_robot_model("unitree_go1")
+
+  mesh_geom_count = 0
+  primitive_geom_count = 0
+
+  for geom_idx in range(model.ngeom):
+    geom_type = model.geom_type[geom_idx]
+
+    if geom_type == mujoco.mjtGeom.mjGEOM_MESH:
+      mesh_geom_count += 1
+
+      # Convert to trimesh.
+      mesh = mujoco_mesh_to_trimesh(model, geom_idx)
+
+      # Basic checks.
+      assert mesh is not None, f"Failed to convert geom {geom_idx}"
+      assert len(mesh.vertices) > 0, f"Mesh {geom_idx} has no vertices"
+      assert len(mesh.faces) > 0, f"Mesh {geom_idx} has no faces"
+    else:
+      # Count primitive geometries (box, sphere, capsule, etc.).
+      primitive_geom_count += 1
+
+  print(f"✓ Unitree Go1: Successfully converted {mesh_geom_count} mesh geometries")
+  print(f"  - Also has {primitive_geom_count} primitive geometries")
+
+
+def test_texture_extraction():
+  """Test texture extraction with a simple textured model."""
+  # Create a model with procedural texture.
+  xml_string = """
+    <mujoco>
+        <asset>
+            <!-- Procedural checker texture -->
+            <texture name="checker" type="2d" builtin="checker" width="64" height="64"
+                     rgb1="1 0 0" rgb2="0 0 1"/>
+            <material name="checker_mat" texture="checker" rgba="1 1 1 1"/>
+
+            <!-- Simple box mesh with 8 vertices -->
+            <mesh name="box"
+                  vertex="0 0 0  1 0 0  1 1 0  0 1 0  0 0 1  1 0 1  1 1 1  0 1 1"
+                  face="0 1 2  0 2 3  4 5 6  4 6 7"/>
+        </asset>
+
+        <worldbody>
+            <geom name="textured_mesh" type="mesh" mesh="box" material="checker_mat"/>
+        </worldbody>
+    </mujoco>
+    """
+
+  model = mujoco.MjModel.from_xml_string(xml_string)
+
+  # Find the mesh geom.
+  for geom_idx in range(model.ngeom):
+    if model.geom_type[geom_idx] == mujoco.mjtGeom.mjGEOM_MESH:
+      mesh = mujoco_mesh_to_trimesh(model, geom_idx)
+
+      assert mesh is not None, "Failed to convert textured mesh"
+      assert mesh.visual is not None, "Mesh has no visual"
+
+      # Check if texture was extracted (procedural textures should work).
+      matid = model.geom_matid[geom_idx]
+      if matid >= 0 and matid < model.nmat:
+        texid = model.mat_texid[matid]
+        # texid might be an array, get the first element if so.
+        if hasattr(texid, "__len__"):
+          texid = texid[0] if len(texid) > 0 else -1
+        if texid >= 0:
+          # Should have extracted the checker texture.
+          print("✓ Texture extraction: Successfully extracted procedural texture")
+          return
+
+  print("✓ Texture extraction: Tested (no textured meshes in simple model)")
+
+
+def test_material_colors():
+  """Test that material colors are properly applied."""
+  xml_string = """
+    <mujoco>
+        <asset>
+            <material name="red" rgba="1 0 0 1"/>
+            <material name="green" rgba="0 1 0 0.5"/>
+            <!-- Tetrahedron with 4 vertices (minimum for MuJoCo mesh) -->
+            <mesh name="tetra"
+                  vertex="0 0 0  1 0 0  0.5 0.866 0  0.5 0.289 0.816"
+                  face="0 1 2  0 1 3  0 2 3  1 2 3"/>
+        </asset>
+
+        <worldbody>
+            <geom name="red_mesh" type="mesh" mesh="tetra" material="red"/>
+            <geom name="green_mesh" type="mesh" mesh="tetra" material="green" pos="2 0 0"/>
+            <geom name="default_mesh" type="mesh" mesh="tetra" pos="4 0 0"/>
+        </worldbody>
+    </mujoco>
+    """
+
+  model = mujoco.MjModel.from_xml_string(xml_string)
+
+  colors_found = []
+  for geom_idx in range(model.ngeom):
+    if model.geom_type[geom_idx] == mujoco.mjtGeom.mjGEOM_MESH:
+      mesh = mujoco_mesh_to_trimesh(model, geom_idx)
+
+      # Check visual colors.
+      visual = mesh.visual
+      if visual and hasattr(visual, "vertex_colors"):
+        # Get the first vertex color (they should all be the same).
+        color = visual.vertex_colors[0]
+        colors_found.append(tuple(color))
+
+  assert len(colors_found) == 3, "Should have 3 meshes"
+
+  # Check that we got different colors.
+  assert colors_found[0][:3] == (255, 0, 0), "First mesh should be red"
+  assert colors_found[1][:3] == (0, 255, 0), "Second mesh should be green"
+  # Third mesh should have default color.
+
+  print("✓ Material colors: Correctly applied to meshes")
+
+
+def test_performance():
+  """Test conversion performance with a complex model."""
+  import time
+
+  model = load_robot_model("unitree_g1")
+
+  mesh_geoms = []
+  for geom_idx in range(model.ngeom):
+    if model.geom_type[geom_idx] == mujoco.mjtGeom.mjGEOM_MESH:
+      mesh_geoms.append(geom_idx)
+
+  start_time = time.time()
+  for geom_idx in mesh_geoms:
+    mujoco_mesh_to_trimesh(model, geom_idx)
+  elapsed = time.time() - start_time
+
+  avg_time = elapsed / len(mesh_geoms) * 1000  # Convert to ms.
+  print(f"✓ Performance: Converted {len(mesh_geoms)} meshes in {elapsed:.3f}s")
+  print(f"  - Average: {avg_time:.2f}ms per mesh")
+
+  # Warn if it's too slow.
+  if avg_time > 50:
+    print("  ⚠ Warning: Conversion is slow (>50ms per mesh)")
+
+
+def test_mesh_with_texture_coordinates():
+  """Test meshes with texture coordinates (Issue #225)."""
+  xml_string = """
+    <mujoco>
+        <asset>
+            <!-- Simple tetrahedron WITH texture coordinates -->
+            <mesh name="textured_tetra"
+                  vertex="0 0 0  1 0 0  0.5 0.866 0  0.5 0.289 0.816"
+                  texcoord="0 0  1 0  0.5 1  0.5 0.5"
+                  face="0 1 2  0 1 3  0 2 3  1 2 3"/>
+        </asset>
+        <worldbody>
+            <geom type="mesh" mesh="textured_tetra"/>
+        </worldbody>
+    </mujoco>
+    """
+
+  model = mujoco.MjModel.from_xml_string(xml_string)
+
+  # Verify the model has texture coordinates.
+  assert model.mesh_texcoord.shape == (4, 2), "Model should have 4 texture coordinates"
+  assert model.mesh_texcoordnum[0] == 4, "Mesh should have 4 texture coordinates"
+
+  # Convert the mesh.
+  mesh = mujoco_mesh_to_trimesh(model, 0)
+
+  # Verify the mesh was created successfully.
+  assert mesh is not None, "Failed to convert textured mesh"
+  assert len(mesh.vertices) > 0, "Mesh has no vertices"
+  assert len(mesh.faces) > 0, "Mesh has no faces"
+
+  # With texture coordinates, vertices are duplicated per face.
+  # 4 faces * 3 vertices per face = 12 vertices.
+  assert len(mesh.vertices) == 12, (
+    f"Expected 12 duplicated vertices, got {len(mesh.vertices)}"
+  )
+  assert len(mesh.faces) == 4, f"Expected 4 faces, got {len(mesh.faces)}"
+
+  print(
+    "✓ Mesh with texture coordinates: Successfully converted mesh with texcoords (Issue #225 fixed)"
+  )
+
+
+def _make_mixed_visual_model() -> mujoco.MjModel:
+  """Create a model with both textured and color-only geoms on the same body."""
+  xml = """
+    <mujoco>
+      <asset>
+        <texture name="checker" type="2d" builtin="checker" width="64" height="64"
+                 rgb1="1 0 0" rgb2="0 0 1"/>
+        <material name="tex_mat" texture="checker"/>
+        <material name="color_mat" rgba="0 1 0 1"/>
+        <mesh name="textured_tetra"
+              vertex="0 0 0  1 0 0  0.5 0.866 0  0.5 0.289 0.816"
+              texcoord="0 0  1 0  0.5 1  0.5 0.5"
+              face="0 1 2  0 1 3  0 2 3  1 2 3"/>
+        <mesh name="plain_tetra"
+              vertex="2 0 0  3 0 0  2.5 0.866 0  2.5 0.289 0.816"
+              face="0 1 2  0 1 3  0 2 3  1 2 3"/>
+      </asset>
+      <worldbody>
+        <body name="mixed">
+          <geom name="tex_geom" type="mesh" mesh="textured_tetra" material="tex_mat"/>
+          <geom name="color_geom" type="mesh" mesh="plain_tetra" material="color_mat"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  """
+  return mujoco.MjModel.from_xml_string(xml)
+
+
+def test_get_geom_texture_id():
+  """Verify correct texture ID detection for textured vs untextured geoms."""
+  model = _make_mixed_visual_model()
+
+  # Find geom indices by name.
+  tex_geom_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, "tex_geom")
+  color_geom_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, "color_geom")
+
+  tex_id = get_geom_texture_id(model, tex_geom_id)
+  assert tex_id >= 0, "Textured geom should have a non-negative texture ID"
+
+  color_id = get_geom_texture_id(model, color_geom_id)
+  assert color_id == -1, "Color-only geom should have texture ID -1"
+
+
+def test_get_geom_texture_id_primitive():
+  """Primitive (non-mesh) geoms should always return -1."""
+  xml = """
+    <mujoco>
+      <worldbody>
+        <geom name="sphere" type="sphere" size="0.1"/>
+      </worldbody>
+    </mujoco>
+  """
+  model = mujoco.MjModel.from_xml_string(xml)
+  geom_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, "sphere")
+  assert get_geom_texture_id(model, geom_id) == -1
+
+
+def test_group_geoms_by_visual_compat():
+  """Verify geoms are correctly partitioned by visual compatibility."""
+  model = _make_mixed_visual_model()
+
+  tex_geom_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, "tex_geom")
+  color_geom_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, "color_geom")
+
+  groups = group_geoms_by_visual_compat(model, [tex_geom_id, color_geom_id])
+
+  # Should produce two groups: one textured, one untextured.
+  assert len(groups) == 2, f"Expected 2 groups, got {len(groups)}"
+
+  # Each group should have exactly one geom.
+  assert all(len(g) == 1 for g in groups)
+
+  # Flatten to a set and verify all geoms are present.
+  all_ids = {gid for g in groups for gid in g}
+  assert all_ids == {tex_geom_id, color_geom_id}
+
+
+def test_group_geoms_single_type():
+  """All untextured geoms should land in a single group."""
+  xml = """
+    <mujoco>
+      <asset>
+        <material name="red" rgba="1 0 0 1"/>
+        <material name="blue" rgba="0 0 1 1"/>
+        <mesh name="t1"
+              vertex="0 0 0  1 0 0  0.5 0.866 0  0.5 0.289 0.816"
+              face="0 1 2  0 1 3  0 2 3  1 2 3"/>
+        <mesh name="t2"
+              vertex="2 0 0  3 0 0  2.5 0.866 0  2.5 0.289 0.816"
+              face="0 1 2  0 1 3  0 2 3  1 2 3"/>
+      </asset>
+      <worldbody>
+        <body name="b">
+          <geom name="g1" type="mesh" mesh="t1" material="red"/>
+          <geom name="g2" type="mesh" mesh="t2" material="blue"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  """
+  model = mujoco.MjModel.from_xml_string(xml)
+  g1 = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, "g1")
+  g2 = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, "g2")
+
+  groups = group_geoms_by_visual_compat(model, [g1, g2])
+  assert len(groups) == 1, "Two untextured geoms should form a single group"
+  assert set(groups[0]) == {g1, g2}
+
+
+def test_merge_mixed_visual_geoms():
+  """Regression: merging each visual-compat subgroup must preserve visuals."""
+  import trimesh.visual
+
+  model = _make_mixed_visual_model()
+
+  tex_geom_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, "tex_geom")
+  color_geom_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, "color_geom")
+
+  groups = group_geoms_by_visual_compat(model, [tex_geom_id, color_geom_id])
+
+  for group in groups:
+    mesh = merge_geoms(model, group)
+    assert len(mesh.vertices) > 0
+    assert len(mesh.faces) > 0
+
+    tex_id = get_geom_texture_id(model, group[0])
+    if tex_id >= 0:
+      # Textured group should have TextureVisuals, not default gray.
+      assert isinstance(mesh.visual, trimesh.visual.TextureVisuals), (
+        f"Expected TextureVisuals for textured group, got {type(mesh.visual)}"
+      )
+    else:
+      # Color group should have ColorVisuals.
+      assert isinstance(mesh.visual, trimesh.visual.ColorVisuals), (
+        f"Expected ColorVisuals for color group, got {type(mesh.visual)}"
+      )
+      # Verify the color is green (not default gray).
+      vertex_colors = mesh.visual.vertex_colors
+      assert vertex_colors[0][1] == 255, (
+        f"Expected green channel 255, got {vertex_colors[0][1]}"
+      )
+
+
+def _make_site_model() -> mujoco.MjModel:
+  """Create a model with all 5 site types."""
+  xml = """
+    <mujoco>
+      <worldbody>
+        <body name="test_body" pos="0 0 1">
+          <joint type="free"/>
+          <geom type="sphere" size="0.01"/>
+          <site name="sphere_site" type="sphere" size="0.1" rgba="1 0 0 1"/>
+          <site name="capsule_site" type="capsule" size="0.05 0.1" rgba="0 1 0 1"
+                pos="1 0 0"/>
+          <site name="ellipsoid_site" type="ellipsoid" size="0.1 0.05 0.03"
+                rgba="0 0 1 1" pos="2 0 0"/>
+          <site name="cylinder_site" type="cylinder" size="0.05 0.1" rgba="1 1 0 1"
+                pos="3 0 0"/>
+          <site name="box_site" type="box" size="0.1 0.05 0.03" rgba="0 1 1 1"
+                pos="4 0 0"/>
+        </body>
+      </worldbody>
+    </mujoco>
+  """
+  return mujoco.MjModel.from_xml_string(xml)
+
+
+def test_create_site_mesh_sphere():
+  """Test sphere site mesh creation."""
+  model = _make_site_model()
+  site_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE, "sphere_site")
+  mesh = create_site_mesh(model, site_id)
+  assert len(mesh.vertices) > 0
+  assert len(mesh.faces) > 0
+  # Check color is red.
+  assert mesh.visual is not None
+  assert mesh.visual.vertex_colors[0][0] == 255
+  assert mesh.visual.vertex_colors[0][1] == 0
+
+
+def test_create_site_mesh_capsule():
+  """Test capsule site mesh creation."""
+  model = _make_site_model()
+  site_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE, "capsule_site")
+  mesh = create_site_mesh(model, site_id)
+  assert len(mesh.vertices) > 0
+  assert len(mesh.faces) > 0
+  # Check color is green.
+  assert mesh.visual is not None
+  assert mesh.visual.vertex_colors[0][1] == 255
+
+
+def test_create_site_mesh_ellipsoid():
+  """Test ellipsoid site mesh creation."""
+  model = _make_site_model()
+  site_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE, "ellipsoid_site")
+  mesh = create_site_mesh(model, site_id)
+  assert len(mesh.vertices) > 0
+  assert len(mesh.faces) > 0
+  # Check color is blue.
+  assert mesh.visual is not None
+  assert mesh.visual.vertex_colors[0][2] == 255
+
+
+def test_create_site_mesh_cylinder():
+  """Test cylinder site mesh creation."""
+  model = _make_site_model()
+  site_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE, "cylinder_site")
+  mesh = create_site_mesh(model, site_id)
+  assert len(mesh.vertices) > 0
+  assert len(mesh.faces) > 0
+
+
+def test_create_site_mesh_box():
+  """Test box site mesh creation."""
+  model = _make_site_model()
+  site_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE, "box_site")
+  mesh = create_site_mesh(model, site_id)
+  assert len(mesh.vertices) > 0
+  assert len(mesh.faces) > 0
+
+
+def test_create_site_mesh_rgba_fallback():
+  """Test that all-zero RGBA falls back to gray."""
+  xml = """
+    <mujoco>
+      <worldbody>
+        <site name="zero_rgba" type="sphere" size="0.1" rgba="0 0 0 0"/>
+      </worldbody>
+    </mujoco>
+  """
+  model = mujoco.MjModel.from_xml_string(xml)
+  site_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE, "zero_rgba")
+  mesh = create_site_mesh(model, site_id)
+  # Should be gray (0.5 * 255 = 127), not transparent black.
+  assert mesh.visual is not None
+  color = mesh.visual.vertex_colors[0]
+  assert color[0] == 127, f"Expected R=127, got {color[0]}"
+  assert color[1] == 127, f"Expected G=127, got {color[1]}"
+  assert color[2] == 127, f"Expected B=127, got {color[2]}"
+  assert color[3] == 255, f"Expected A=255, got {color[3]}"
+
+
+def test_merge_sites_single():
+  """Test merging a single site returns the site mesh."""
+  model = _make_site_model()
+  site_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE, "sphere_site")
+  mesh = merge_sites(model, [site_id])
+  assert len(mesh.vertices) > 0
+  assert len(mesh.faces) > 0
+
+
+def test_merge_sites_multiple():
+  """Test merging multiple sites combines them correctly."""
+  model = _make_site_model()
+  sphere_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE, "sphere_site")
+  capsule_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE, "capsule_site")
+  box_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SITE, "box_site")
+
+  single_sphere = create_site_mesh(model, sphere_id)
+  single_capsule = create_site_mesh(model, capsule_id)
+  single_box = create_site_mesh(model, box_id)
+
+  merged = merge_sites(model, [sphere_id, capsule_id, box_id])
+
+  # Merged mesh should have at least the sum of individual vertices.
+  expected_min_verts = (
+    len(single_sphere.vertices)
+    + len(single_capsule.vertices)
+    + len(single_box.vertices)
+  )
+  assert len(merged.vertices) >= expected_min_verts, (
+    f"Expected >= {expected_min_verts} vertices, got {len(merged.vertices)}"
+  )
+  assert len(merged.faces) > 0
+
+
+if __name__ == "__main__":
+  # Run all tests
+  print("=" * 60)
+  print("Testing Viser Conversions")
+  print("=" * 60)
+
+  tests = [
+    test_unitree_g1_conversion,
+    test_unitree_go1_conversion,
+    test_texture_extraction,
+    test_material_colors,
+    test_performance,
+    test_mesh_with_texture_coordinates,
+    test_get_geom_texture_id,
+    test_get_geom_texture_id_primitive,
+    test_group_geoms_by_visual_compat,
+    test_group_geoms_single_type,
+    test_merge_mixed_visual_geoms,
+  ]
+
+  failed = []
+  for test in tests:
+    try:
+      test()
+    except Exception as e:
+      test_name = test.__name__ if hasattr(test, "__name__") else str(test)
+      print(f"✗ {test_name}: {e}")
+      import traceback
+
+      traceback.print_exc()
+      failed.append(test_name)
+
+  print("\n" + "=" * 60)
+  if failed:
+    print(f"Failed tests: {', '.join(failed)}")
+    import sys
+
+    sys.exit(1)
+  else:
+    print("All tests passed!")
+    import sys
+
+    sys.exit(0)
@@ -0,0 +1,116 @@
+"""Tests for Viser viewer update-policy helpers."""
+
+from dataclasses import dataclass
+from typing import Any
+from unittest.mock import MagicMock
+
+import mujoco
+from mjviser import ViserMujocoScene
+
+from mjlab.viewer.viser.overlays import ViserContactOverlays, ViserDebugOverlays
+from mjlab.viewer.viser.viewer import ViserPlayViewer
+
+
+@dataclass
+class _DummyScene:
+  env_idx: int
+  debug_visualization_enabled: bool
+  show_contact_points: bool = False
+  show_contact_forces: bool = False
+  needs_update: bool = False
+  clear_count: int = 0
+  clear_debug_count: int = 0
+
+  def clear(self) -> None:
+    self.clear_count += 1
+
+  def clear_debug_all(self) -> None:
+    self.clear_debug_count += 1
+
+
+class _DummyEnv:
+  def __init__(self, unwrapped: Any):
+    self._unwrapped = unwrapped
+
+  @property
+  def unwrapped(self) -> Any:
+    return self._unwrapped
+
+
+def test_should_update_cameras():
+  assert ViserPlayViewer._should_update_cameras(paused=False, has_pending_updates=False)
+  assert not ViserPlayViewer._should_update_cameras(
+    paused=True, has_pending_updates=False
+  )
+  assert ViserPlayViewer._should_update_cameras(paused=True, has_pending_updates=True)
+
+
+def test_should_submit_scene_update():
+  # Odd ticks are skipped to keep scene submits around 30Hz.
+  assert not ViserPlayViewer._should_submit_scene_update(
+    counter=1, paused=False, has_pending_updates=True
+  )
+  # Running: submit on even ticks regardless of pending flags.
+  assert ViserPlayViewer._should_submit_scene_update(
+    counter=2, paused=False, has_pending_updates=False
+  )
+  # Paused: submit only with pending updates.
+  assert not ViserPlayViewer._should_submit_scene_update(
+    counter=2, paused=True, has_pending_updates=False
+  )
+  assert ViserPlayViewer._should_submit_scene_update(
+    counter=2, paused=True, has_pending_updates=True
+  )
+
+
+def test_scene_requires_live_refresh():
+  """Base scene sets needs_update when contact overlays are visible."""
+  model = mujoco.MjModel.from_xml_string("<mujoco><worldbody/></mujoco>")
+  server = MagicMock()
+  scene = ViserMujocoScene(server, model, num_envs=1)
+
+  scene.show_contact_points = False
+  scene.show_contact_forces = False
+  scene.needs_update = False
+  scene._last_body_xpos = None  # No cached data, refresh is a no-op.
+  scene.refresh_visualization()
+  assert not scene.needs_update
+
+  scene.show_contact_points = True
+  assert scene.show_contact_points or scene.show_contact_forces
+
+
+def test_debug_overlays_env_switch_and_queue():
+  """env_switch respects enabled flag; queue clears + calls update_visualizers."""
+  unwrapped = MagicMock(spec=["update_visualizers"])
+  env = _DummyEnv(unwrapped)
+  scene = _DummyScene(env_idx=0, debug_visualization_enabled=False)
+  overlays = ViserDebugOverlays(env=env, scene=scene)
+
+  # Disabled: env switch does not clear.
+  overlays.on_env_switch()
+  assert scene.clear_debug_count == 0
+
+  # Enabled: env switch clears, queue dispatches.
+  scene.debug_visualization_enabled = True
+  overlays.on_env_switch()
+  assert scene.clear_debug_count == 1
+
+  overlays.queue()
+  assert scene.clear_count == 1
+  unwrapped.update_visualizers.assert_called_once_with(scene)
+
+
+def test_contact_overlays_env_switch():
+  """env_switch requests scene update only when contacts are enabled."""
+  scene = _DummyScene(env_idx=0, debug_visualization_enabled=False)
+  overlays = ViserContactOverlays(scene=scene)
+
+  assert not overlays.is_enabled()
+  overlays.on_env_switch()
+  assert not scene.needs_update
+
+  scene.show_contact_points = True
+  assert overlays.is_enabled()
+  overlays.on_env_switch()
+  assert scene.needs_update
@@ -0,0 +1,178 @@
+"""Tests for XML actuator wrappers."""
+
+import mujoco
+import pytest
+from conftest import get_test_device
+
+from mjlab.actuator import XmlActuator, XmlActuatorCfg
+from mjlab.entity import Entity, EntityArticulationInfoCfg, EntityCfg
+from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg, mdp
+from mjlab.managers.observation_manager import ObservationGroupCfg, ObservationTermCfg
+from mjlab.scene import SceneCfg
+from mjlab.sim import MujocoCfg, SimulationCfg
+from mjlab.terrains import TerrainEntityCfg
+
+# Robot with 2 joints but only 1 actuator defined (underactuated).
+ROBOT_XML_UNDERACTUATED = """
+<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="1.0"/>
+      <body name="link1" pos="0 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="0.1"/>
+      </body>
+      <body name="link2" pos="0 0 0">
+        <joint name="joint2" type="hinge" axis="0 0 1" range="-1.57 1.57"/>
+        <geom name="link2_geom" type="box" size="0.1 0.1 0.1" mass="0.1"/>
+      </body>
+    </body>
+  </worldbody>
+  <actuator>
+    <motor name="actuator1" joint="joint2" gear="1.0"/>
+  </actuator>
+</mujoco>
+"""
+
+
+@pytest.fixture(scope="module")
+def device():
+  return get_test_device()
+
+
+def test_xml_actuator_underactuated_with_wildcard():
+  """XmlActuator filters to joints with XML actuators when using wildcard."""
+  cfg = EntityCfg(
+    spec_fn=lambda: mujoco.MjSpec.from_string(ROBOT_XML_UNDERACTUATED),
+    articulation=EntityArticulationInfoCfg(
+      actuators=(XmlActuatorCfg(target_names_expr=(".*",)),)
+    ),
+  )
+  entity = Entity(cfg)
+  entity.compile()
+
+  # Should only control joint2 (which has an XML actuator), not joint1.
+  assert len(entity._actuators) == 1
+  actuator = entity._actuators[0]
+  assert actuator._target_names == ["joint2"]
+
+
+def test_xml_actuator_no_matching_actuators_raises_error():
+  """XmlActuator raises error when no targets have matching XML actuators."""
+  with pytest.raises(
+    ValueError, match="No XML actuators found for any targets matching the patterns"
+  ):
+    cfg = EntityCfg(
+      spec_fn=lambda: mujoco.MjSpec.from_string(ROBOT_XML_UNDERACTUATED),
+      articulation=EntityArticulationInfoCfg(
+        actuators=(XmlActuatorCfg(target_names_expr=("joint1",)),)
+      ),
+    )
+    entity = Entity(cfg)
+    entity.compile()
+
+
+def test_joint_action_underactuated_with_wildcard(device):
+  """JointAction with wildcard pattern matches only actuated joints."""
+  robot_cfg = EntityCfg(
+    spec_fn=lambda: mujoco.MjSpec.from_string(ROBOT_XML_UNDERACTUATED),
+    articulation=EntityArticulationInfoCfg(
+      actuators=(XmlActuatorCfg(target_names_expr=(".*",)),)
+    ),
+  )
+
+  env_cfg = ManagerBasedRlEnvCfg(
+    scene=SceneCfg(
+      terrain=TerrainEntityCfg(terrain_type="plane"),
+      num_envs=1,
+      extent=1.0,
+      entities={"robot": robot_cfg},
+    ),
+    observations={
+      "actor": ObservationGroupCfg(
+        terms={
+          "joint_pos": ObservationTermCfg(
+            func=lambda env: env.scene["robot"].data.joint_pos
+          ),
+        },
+      ),
+    },
+    actions={
+      "joint_effort": mdp.JointEffortActionCfg(
+        entity_name="robot", actuator_names=(".*",), scale=1.0
+      )
+    },
+    sim=SimulationCfg(mujoco=MujocoCfg(timestep=0.01, iterations=1)),
+    decimation=1,
+    episode_length_s=1.0,
+  )
+
+  env = ManagerBasedRlEnv(cfg=env_cfg, device=device)
+  action_term = env.action_manager._terms["joint_effort"]
+  assert isinstance(action_term, mdp.JointEffortAction)
+
+  # Wildcard should resolve to only actuated joint (joint2), not all joints.
+  assert action_term.action_dim == 1
+  assert action_term.target_names == ["joint2"]
+  assert action_term.target_ids.tolist() == [1]
+
+  env.close()
+
+
+# Robot with a "general" actuator (hand-tuned gainprm/biasprm that don't match
+# any standard pattern). Strict auto-detection would reject this.
+ROBOT_XML_GENERAL_ACTUATOR = """
+<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="1.0"/>
+      <body name="link1" pos="0 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="0.1"/>
+      </body>
+    </body>
+  </worldbody>
+  <actuator>
+    <general name="custom_act" joint="joint1"
+             gaintype="fixed" biastype="affine"
+             gainprm="10" biasprm="5 -3 -1"/>
+  </actuator>
+</mujoco>
+"""
+
+
+def test_xml_actuator_explicit_command_field_bypasses_detection():
+  """Explicit command_field allows non-standard XML actuators."""
+  cfg = EntityCfg(
+    spec_fn=lambda: mujoco.MjSpec.from_string(ROBOT_XML_GENERAL_ACTUATOR),
+    articulation=EntityArticulationInfoCfg(
+      actuators=(
+        XmlActuatorCfg(
+          target_names_expr=("joint1",),
+          command_field="effort",
+        ),
+      )
+    ),
+  )
+  entity = Entity(cfg)
+  entity.compile()
+
+  actuator = entity._actuators[0]
+  assert isinstance(actuator, XmlActuator)
+  assert actuator.command_field == "effort"
+  assert actuator._target_names == ["joint1"]
+
+
+def test_xml_actuator_auto_detection_rejects_general_actuator():
+  """Auto-detection raises for non-standard actuators when command_field is None."""
+  with pytest.raises(ValueError, match="Cannot determine command field"):
+    cfg = EntityCfg(
+      spec_fn=lambda: mujoco.MjSpec.from_string(ROBOT_XML_GENERAL_ACTUATOR),
+      articulation=EntityArticulationInfoCfg(
+        actuators=(XmlActuatorCfg(target_names_expr=("joint1",)),)
+      ),
+    )
+    entity = Entity(cfg)
+    entity.compile()
@@ -0,0 +1,116 @@
+"""Tests for mjlab.utils.xml."""
+
+import zipfile
+
+import mujoco
+import numpy as np
+
+from mjlab.utils.xml import fix_spec_xml, strip_buffer_textures
+
+
+def test_strip_buffer_textures():
+  """Stripping buffer textures removes them and enables XML roundtrip."""
+  spec = mujoco.MjSpec()
+  body = spec.worldbody.add_body(name="terrain")
+
+  tex = spec.add_texture(
+    name="hf_tex",
+    type=mujoco.mjtTexture.mjTEXTURE_2D,
+    width=2,
+    height=2,
+  )
+  tex.data = (np.ones((2, 2, 3), dtype=np.uint8) * 128).tobytes()
+
+  mat = spec.add_material(name="hf_mat")
+  mat.textures[mujoco.mjtTextureRole.mjTEXROLE_RGB] = "hf_tex"
+
+  hf = spec.add_hfield(name="hf", nrow=4, ncol=4, size=[1, 1, 0.1, 0.01])
+  hf.userdata = np.zeros(16).tolist()
+  body.add_geom(
+    type=mujoco.mjtGeom.mjGEOM_HFIELD,
+    hfieldname="hf",
+    material="hf_mat",
+  )
+
+  strip_buffer_textures(spec)
+
+  assert len(spec.textures) == 0
+  assert len(spec.materials) == 0
+  assert spec.geoms[0].material == ""
+
+  # XML roundtrip should succeed now.
+  model = mujoco.MjModel.from_xml_string(spec.to_xml())
+  assert model.nhfield == 1
+
+
+def test_fix_spec_xml():
+  """Removes empty defaults and optionally injects meshdir."""
+  xml = (
+    "<mujoco>\n"
+    "  <compiler/>\n"
+    "  <default/>\n"
+    '  <default class="foo/bar"/>\n'
+    "  <worldbody/>\n"
+    "</mujoco>\n"
+  )
+  result = fix_spec_xml(xml, meshdir="assets")
+
+  assert "<default/>" not in result
+  assert 'class="foo/bar"' not in result
+  assert 'meshdir="assets"' in result
+  assert "meshdir" not in fix_spec_xml(xml)
+
+
+def test_rough_terrain_write_xml_roundtrip(tmp_path):
+  """Export terrain with hfield via Entity.write_xml, reload as MjModel."""
+  from mjlab.terrains.heightfield_terrains import (
+    HfRandomUniformTerrainCfg,
+  )
+  from mjlab.terrains.terrain_entity import TerrainEntity, TerrainEntityCfg
+  from mjlab.terrains.terrain_generator import TerrainGeneratorCfg
+
+  terrain = TerrainEntity(
+    TerrainEntityCfg(
+      terrain_type="generator",
+      terrain_generator=TerrainGeneratorCfg(
+        size=(2.0, 2.0),
+        num_rows=1,
+        num_cols=1,
+        sub_terrains={
+          "rough": HfRandomUniformTerrainCfg(
+            proportion=1.0,
+            noise_range=(-0.05, 0.05),
+            noise_step=0.005,
+          ),
+        },
+      ),
+    ),
+    device="cpu",
+  )
+  xml_path = tmp_path / "terrain.xml"
+  terrain.write_xml(xml_path)
+
+  model = mujoco.MjModel.from_xml_path(str(xml_path))
+  assert model.nhfield == 1
+
+
+def test_scene_write_zip(tmp_path):
+  """Scene.write produces a directory; with zip=True, a zip archive."""
+  from mjlab.scene.scene import Scene, SceneCfg
+
+  scene = Scene(SceneCfg(), device="cpu")
+
+  # Directory export.
+  dir_out = tmp_path / "scene_dir"
+  scene.write(dir_out)
+  assert (dir_out / "scene.xml").exists()
+  mujoco.MjModel.from_xml_path(str(dir_out / "scene.xml"))
+
+  # Zip export.
+  zip_out = tmp_path / "scene_pkg"
+  scene.write(zip_out, zip=True)
+  zip_path = zip_out.with_suffix(".zip")
+  assert zip_path.exists()
+  assert not zip_out.exists()
+  with zipfile.ZipFile(zip_path) as zf:
+    assert "scene.xml" in zf.namelist()