32a241c28f
Some checks failed
nightly / Test against latest dependencies (py3.10) (push) Has been cancelled
nightly / Test against latest dependencies (py3.13) (push) Has been cancelled
tests / tests (3.13, locked) (push) Has been cancelled
tests / tests (3.13, unlocked) (push) Has been cancelled
tests / pyright (3.10) (push) Has been cancelled
tests / lint-format (push) Has been cancelled
tests / tests (3.10, locked) (push) Has been cancelled
tests / tests (3.11, locked) (push) Has been cancelled
tests / tests (3.12, locked) (push) Has been cancelled
tests / pyright (3.11) (push) Has been cancelled
tests / pyright (3.12) (push) Has been cancelled
tests / pyright (3.13) (push) Has been cancelled
tests / ty-check (3.10) (push) Has been cancelled
tests / ty-check (3.11) (push) Has been cancelled
tests / ty-check (3.12) (push) Has been cancelled
tests / ty-check (3.13) (push) Has been cancelled
tests / stubs (push) Has been cancelled
tests / smoke-test (push) Has been cancelled
Docker / check_paths (push) Has been cancelled
docs / build (push) Has been cancelled
Docker / build (push) Has been cancelled
Upstream: https://github.com/michaelgillett/mjlab Upstream-Commit: c19f713c415a699a79d71cd96aa13c3104a05047 Upstream-Branch: main
251 lines
6.6 KiB
Python
251 lines
6.6 KiB
Python
"""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)
|