"""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 = """ """ 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]