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
410 lines
13 KiB
Python
410 lines
13 KiB
Python
"""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)
|