"""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 = """ """ ROBOT_WITH_FLOOR_XML = """ """ 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)