mjlab/tests/test_nan_guard.py

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