"""Tests for MjlabOnPolicyRunner.""" import ast import tempfile from dataclasses import asdict from pathlib import Path from unittest.mock import MagicMock, patch import mujoco import onnx import pytest import torch from conftest import get_test_device from rsl_rl.models import MLPModel from rsl_rl.utils import WandbLogWriter from rsl_rl.utils.logger import Logger from tensordict import TensorDict import mjlab.scripts.train as train_mod from mjlab.actuator import XmlActuatorCfg from mjlab.entity import EntityArticulationInfoCfg, EntityCfg from mjlab.envs import ManagerBasedRlEnv, ManagerBasedRlEnvCfg, mdp from mjlab.managers.observation_manager import ObservationGroupCfg, ObservationTermCfg from mjlab.rl import RslRlOnPolicyRunnerCfg, RslRlVecEnvWrapper from mjlab.rl.runner import MjlabOnPolicyRunner from mjlab.rl.spatial_softmax import SpatialSoftmaxCNNModel from mjlab.scene import SceneCfg from mjlab.sim import MujocoCfg, SimulationCfg from mjlab.tasks.tracking.rl.runner import _OnnxMotionModel from mjlab.terrains import TerrainEntityCfg from mjlab.utils.os import dump_yaml @pytest.fixture(scope="module") def device(): return get_test_device() @pytest.fixture(scope="module") def env(device): robot_xml = """ """ robot_cfg = EntityCfg( spec_fn=lambda: mujoco.MjSpec.from_string(robot_xml), articulation=EntityArticulationInfoCfg( actuators=(XmlActuatorCfg(target_names_expr=(".*",)),) ), ) env_cfg = ManagerBasedRlEnvCfg( scene=SceneCfg( terrain=TerrainEntityCfg(terrain_type="plane"), num_envs=2, extent=1.0, entities={"robot": robot_cfg}, ), observations={ "actor": ObservationGroupCfg( terms={ "joint_pos": ObservationTermCfg( func=lambda env: env.scene["robot"].data.joint_pos ), }, ), "critic": ObservationGroupCfg( terms={ "joint_pos": ObservationTermCfg( func=lambda env: env.scene["robot"].data.joint_pos ), }, ), }, actions={ "joint_pos": mdp.JointPositionActionCfg( entity_name="robot", actuator_names=(".*",), scale=1.0 ) }, sim=SimulationCfg(mujoco=MujocoCfg(timestep=0.01, iterations=1)), decimation=1, episode_length_s=1.0, ) env = ManagerBasedRlEnv(cfg=env_cfg, device=device) yield env env.close() def test_runner_persists_common_step_counter(env, device, monkeypatch): """MjlabOnPolicyRunner should save and restore common_step_counter.""" wrapped_env = RslRlVecEnvWrapper(env) agent_cfg = RslRlOnPolicyRunnerCfg( num_steps_per_env=4, max_iterations=10, save_interval=5 ) with tempfile.TemporaryDirectory() as tmpdir: runner = MjlabOnPolicyRunner( wrapped_env, asdict(agent_cfg), log_dir=tmpdir, device=device ) monkeypatch.setattr(runner.logger, "save_model", lambda *args, **kwargs: None) wrapped_env.unwrapped.common_step_counter = 12345 checkpoint_path = str(Path(tmpdir) / "test_checkpoint.pt") runner.save(checkpoint_path) wrapped_env.unwrapped.common_step_counter = 0 runner.load(checkpoint_path) assert wrapped_env.unwrapped.common_step_counter == 12345 def test_runner_handles_old_checkpoints_without_env_state(env, device): """Old checkpoints without env_state should load without crashing.""" wrapped_env = RslRlVecEnvWrapper(env) agent_cfg = RslRlOnPolicyRunnerCfg( num_steps_per_env=4, max_iterations=10, save_interval=5 ) with tempfile.TemporaryDirectory() as tmpdir: runner = MjlabOnPolicyRunner( wrapped_env, asdict(agent_cfg), log_dir=tmpdir, device=device ) checkpoint_path = str(Path(tmpdir) / "old_checkpoint.pt") old_checkpoint = { "actor_state_dict": runner.alg.actor.state_dict(), "critic_state_dict": runner.alg.critic.state_dict(), "optimizer_state_dict": runner.alg.optimizer.state_dict(), "iter": 100, "infos": None, } torch.save(old_checkpoint, checkpoint_path) wrapped_env.unwrapped.common_step_counter = 999 runner.load(checkpoint_path) assert wrapped_env.unwrapped.common_step_counter == 999 @pytest.mark.filterwarnings("ignore::DeprecationWarning") def test_export_policy_to_onnx(env, device): """runner.export_policy_to_onnx() produces a valid ONNX file.""" wrapped_env = RslRlVecEnvWrapper(env) agent_cfg = RslRlOnPolicyRunnerCfg( num_steps_per_env=4, max_iterations=10, save_interval=5 ) with tempfile.TemporaryDirectory() as tmpdir: runner = MjlabOnPolicyRunner( wrapped_env, asdict(agent_cfg), log_dir=tmpdir, device=device ) runner.export_policy_to_onnx(tmpdir, "test_policy.onnx") onnx_path = Path(tmpdir) / "test_policy.onnx" assert onnx_path.exists() onnx.checker.check_model(str(onnx_path)) def _make_actor(obs_dim=8, output_dim=4, obs_normalization=True): obs = TensorDict({"actor": torch.zeros(1, obs_dim)}) obs_groups = {"actor": ["actor"]} return MLPModel( obs=obs, obs_groups=obs_groups, obs_set="actor", output_dim=output_dim, hidden_dims=[32, 32], activation="elu", obs_normalization=obs_normalization, ) def _train_normalizer(actor, n_batches=50, batch_size=64): actor.train() for _ in range(n_batches): obs = TensorDict({"actor": torch.randn(batch_size, actor.obs_dim) * 5 + 3}) actor.update_normalization(obs) actor.eval() def _model_output(actor, x_flat): obs = TensorDict({"actor": x_flat}) with torch.no_grad(): return actor(obs) def test_onnx_export_matches_actor(): """as_onnx() model produces the same output as the full actor with normalization.""" actor = _make_actor(obs_normalization=True) _train_normalizer(actor) onnx_model = actor.as_onnx(verbose=False) onnx_model.eval() x = torch.randn(4, actor.obs_dim) model_out = _model_output(actor, x) with torch.no_grad(): onnx_out = onnx_model(x) torch.testing.assert_close(model_out, onnx_out, atol=1e-6, rtol=0) def test_onnx_export_without_normalization(): """as_onnx() works when normalization is disabled.""" actor = _make_actor(obs_normalization=False) onnx_model = actor.as_onnx(verbose=False) onnx_model.eval() x = torch.randn(4, actor.obs_dim) with torch.no_grad(): out = onnx_model(x) assert out.shape == (4, 4) @pytest.mark.filterwarnings("ignore::DeprecationWarning") def test_onnx_runtime_roundtrip_matches_pytorch(): """Exported .onnx file produces the same outputs as PyTorch via onnxruntime.""" ort = pytest.importorskip("onnxruntime") actor = _make_actor(obs_normalization=True) _train_normalizer(actor) onnx_model = actor.as_onnx(verbose=False) onnx_model.eval() x = torch.randn(4, actor.obs_dim) expected = _model_output(actor, x) with tempfile.TemporaryDirectory() as tmpdir: path = Path(tmpdir) / "policy.onnx" torch.onnx.export( onnx_model, (x,), str(path), input_names=onnx_model.input_names, # pyright: ignore[reportArgumentType] output_names=onnx_model.output_names, # pyright: ignore[reportArgumentType] opset_version=18, dynamo=False, ) sess = ort.InferenceSession(str(path)) [actual] = sess.run(None, {"obs": x.numpy()}) torch.testing.assert_close(torch.from_numpy(actual), expected, atol=1e-5, rtol=0) # CNN (spatial-softmax) ONNX export tests. _IMG_H, _IMG_W, _IMG_C = 16, 16, 3 _OBS_DIM_1D = 8 _OUTPUT_DIM = 4 def _make_cnn_actor(obs_normalization=True): obs = TensorDict( { "actor": torch.zeros(1, _OBS_DIM_1D), "camera": torch.zeros(1, _IMG_C, _IMG_H, _IMG_W), } ) obs_groups = {"actor": ["actor", "camera"]} cnn_cfg = { "output_channels": [8], "kernel_size": [3], "stride": [1], "spatial_softmax_temperature": 1.0, } return SpatialSoftmaxCNNModel( obs=obs, obs_groups=obs_groups, obs_set="actor", output_dim=_OUTPUT_DIM, cnn_cfg=cnn_cfg, hidden_dims=[32, 32], activation="elu", obs_normalization=obs_normalization, ) def _train_cnn_normalizer(actor, n_batches=50, batch_size=64): actor.train() for _ in range(n_batches): obs = TensorDict( { "actor": torch.randn(batch_size, _OBS_DIM_1D) * 5 + 3, "camera": torch.randn(batch_size, _IMG_C, _IMG_H, _IMG_W), } ) actor.update_normalization(obs) actor.eval() def _cnn_model_output(actor, x_1d, x_2d): obs = TensorDict({"actor": x_1d, "camera": x_2d}) with torch.no_grad(): return actor(obs) def test_cnn_onnx_export_matches_actor(): """as_onnx() with SpatialSoftmaxCNNModel matches the original model.""" actor = _make_cnn_actor(obs_normalization=True) _train_cnn_normalizer(actor) onnx_model = actor.as_onnx(verbose=False) onnx_model.eval() x_1d = torch.randn(4, _OBS_DIM_1D) x_2d = torch.randn(4, _IMG_C, _IMG_H, _IMG_W) expected = _cnn_model_output(actor, x_1d, x_2d) with torch.no_grad(): actual = onnx_model(x_1d, x_2d) torch.testing.assert_close(actual, expected, atol=1e-6, rtol=0) @pytest.mark.filterwarnings("ignore::DeprecationWarning") def test_cnn_onnx_export_to_file(): """SpatialSoftmaxCNNModel exports to a valid ONNX file.""" actor = _make_cnn_actor(obs_normalization=False) onnx_model = actor.as_onnx(verbose=False) onnx_model.to("cpu") onnx_model.eval() with tempfile.TemporaryDirectory() as tmpdir: onnx_path = Path(tmpdir) / "cnn_policy.onnx" torch.onnx.export( onnx_model, onnx_model.get_dummy_inputs(), # pyright: ignore[reportCallIssue] str(onnx_path), export_params=True, opset_version=18, input_names=onnx_model.input_names, # pyright: ignore[reportArgumentType] output_names=onnx_model.output_names, # pyright: ignore[reportArgumentType] dynamic_axes={}, dynamo=False, ) assert onnx_path.exists() onnx.checker.check_model(str(onnx_path)) @pytest.mark.filterwarnings("ignore::DeprecationWarning") def test_cnn_onnx_runtime_roundtrip_matches_pytorch(): """Exported CNN .onnx file produces the same outputs as PyTorch via onnxruntime.""" ort = pytest.importorskip("onnxruntime") actor = _make_cnn_actor(obs_normalization=True) _train_cnn_normalizer(actor) onnx_model = actor.as_onnx(verbose=False) onnx_model.eval() x_1d = torch.randn(4, _OBS_DIM_1D) x_2d = torch.randn(4, _IMG_C, _IMG_H, _IMG_W) expected = _cnn_model_output(actor, x_1d, x_2d) with tempfile.TemporaryDirectory() as tmpdir: path = Path(tmpdir) / "cnn_policy.onnx" torch.onnx.export( onnx_model, (x_1d, x_2d), str(path), input_names=onnx_model.input_names, # pyright: ignore[reportArgumentType] output_names=onnx_model.output_names, # pyright: ignore[reportArgumentType] opset_version=18, dynamo=False, ) sess = ort.InferenceSession(str(path)) [actual] = sess.run(None, {"obs": x_1d.numpy(), "camera": x_2d.numpy()}) torch.testing.assert_close(torch.from_numpy(actual), expected, atol=1e-5, rtol=0) def test_get_export_paths(): """_get_export_paths resolves the correct dir, filename, and full path.""" # Normal case: "model" only appears in the checkpoint filename. export_dir, filename, onnx_path = MjlabOnPolicyRunner._get_export_paths( "/logs/2026-03-30_12-00-00/model_10.pt" ) assert export_dir == Path("/logs/2026-03-30_12-00-00") assert filename == "2026-03-30_12-00-00.onnx" assert onnx_path == Path("/logs/2026-03-30_12-00-00/2026-03-30_12-00-00.onnx") # Bug case: "model" also appears in a parent directory name — the old # path.split("model")[0] would have truncated to "/tmp/my_". export_dir, filename, onnx_path = MjlabOnPolicyRunner._get_export_paths( "/tmp/my_model_experiment/2026-03-30/model_10.pt" ) assert export_dir == Path("/tmp/my_model_experiment/2026-03-30") assert filename == "2026-03-30.onnx" assert onnx_path == Path("/tmp/my_model_experiment/2026-03-30/2026-03-30.onnx") def test_agent_cfg_serializable_after_runner_creation(env, device): """dump_yaml must be called before runner creation. The runner mutates agent_cfg in-place (e.g. resolve_symmetry_config injects non-serializable objects). Verify that the train script writes config files before constructing the runner. Regression test for https://github.com/mjlab-org/mjlab/issues/764. """ wrapped_env = RslRlVecEnvWrapper(env) agent_cfg = asdict( RslRlOnPolicyRunnerCfg(num_steps_per_env=4, max_iterations=10, save_interval=5) ) # Dump should succeed before runner creation. with tempfile.TemporaryDirectory() as tmpdir: dump_yaml(Path(tmpdir) / "agent.yaml", agent_cfg) # Create runner (mutates agent_cfg via resolve_symmetry_config). with tempfile.TemporaryDirectory() as tmpdir: MjlabOnPolicyRunner(wrapped_env, agent_cfg, log_dir=tmpdir, device=device) # Confirm that the runner added non-serializable keys to agent_cfg. sym_cfg = agent_cfg.get("algorithm", {}).get("symmetry_cfg") runner_mutated = sym_cfg is not None or "multi_gpu" in agent_cfg assert runner_mutated, "Expected runner to mutate agent_cfg" # Verify the train script calls dump_yaml before runner_cls(). source = Path(train_mod.__file__).read_text() tree = ast.parse(source) dump_yaml_line = None runner_cls_line = None for node in ast.walk(tree): if isinstance(node, ast.Call): func = node.func # Look for dump_yaml(..., agent_cfg) if isinstance(func, ast.Name) and func.id == "dump_yaml": for arg in node.args: if isinstance(arg, ast.Name) and arg.id == "agent_cfg": dump_yaml_line = node.lineno # Look for runner_cls(...) if isinstance(func, ast.Name) and func.id == "runner_cls": runner_cls_line = node.lineno assert dump_yaml_line is not None, "dump_yaml(agent_cfg) not found" assert runner_cls_line is not None, "runner_cls() not found" assert dump_yaml_line < runner_cls_line, ( f"dump_yaml (line {dump_yaml_line}) must be called before " f"runner_cls (line {runner_cls_line})" ) class _MockMotion: """Minimal mock of a motion object with tensor attributes.""" def __init__(self, num_steps, num_joints=12, num_bodies=5): self.joint_pos = torch.randn(num_steps, num_joints) self.joint_vel = torch.randn(num_steps, num_joints) self.body_pos_w = torch.randn(num_steps, num_bodies, 3) self.body_quat_w = torch.randn(num_steps, num_bodies, 4) self.body_lin_vel_w = torch.randn(num_steps, num_bodies, 3) self.body_ang_vel_w = torch.randn(num_steps, num_bodies, 3) def test_onnx_motion_model_policy_matches_actor(): """_OnnxMotionModel actions output matches calling the actor directly.""" actor = _make_actor(obs_normalization=True) _train_normalizer(actor) motion = _MockMotion(num_steps=50) model = _OnnxMotionModel(actor, motion) model.eval() x = torch.randn(4, actor.obs_dim) time_step = torch.tensor([[5]], dtype=torch.float32) with torch.no_grad(): actions, *_ = model(x, time_step) expected = _model_output(actor, x) torch.testing.assert_close(actions, expected, atol=1e-6, rtol=0) def test_onnx_motion_model_returns_correct_motion_frame(): """_OnnxMotionModel returns the motion data at the requested time step.""" actor = _make_actor(obs_normalization=False) motion = _MockMotion(num_steps=50) model = _OnnxMotionModel(actor, motion) model.eval() x = torch.randn(1, actor.obs_dim) t = 17 time_step = torch.tensor([[t]], dtype=torch.float32) with torch.no_grad(): _, joint_pos, joint_vel, body_pos, body_quat, body_lin_vel, body_ang_vel = model( x, time_step ) torch.testing.assert_close(joint_pos, motion.joint_pos[t : t + 1]) torch.testing.assert_close(joint_vel, motion.joint_vel[t : t + 1]) torch.testing.assert_close(body_pos, motion.body_pos_w[t : t + 1]) torch.testing.assert_close(body_quat, motion.body_quat_w[t : t + 1]) torch.testing.assert_close(body_lin_vel, motion.body_lin_vel_w[t : t + 1]) torch.testing.assert_close(body_ang_vel, motion.body_ang_vel_w[t : t + 1]) def test_onnx_motion_model_clamps_out_of_bounds_time_step(): """_OnnxMotionModel clamps time_step beyond motion length to last frame.""" num_steps = 20 actor = _make_actor(obs_normalization=False) motion = _MockMotion(num_steps=num_steps) model = _OnnxMotionModel(actor, motion) model.eval() x = torch.randn(1, actor.obs_dim) time_step = torch.tensor([[999]], dtype=torch.float32) with torch.no_grad(): _, joint_pos, *_ = model(x, time_step) torch.testing.assert_close(joint_pos, motion.joint_pos[num_steps - 1 : num_steps]) def _make_logger_mock(is_wandb: bool) -> MagicMock: """Mock an rsl-rl Logger, specced off a real (cheap, inert) instance. Speccing off an instance rather than the class keeps the instance attributes, so an attribute rsl-rl renames or drops fails the test instead of silently resolving to a truthy mock. """ logger = Logger( log_dir=None, cfg={"algorithm": {}}, env_cfg={}, num_envs=1, is_distributed=False, gpu_world_size=1, gpu_global_rank=0, device="cpu", ) mock = MagicMock(spec=logger) mock.writer = MagicMock(spec=WandbLogWriter) if is_wandb else MagicMock() return mock @pytest.mark.parametrize( ("runner_module", "runner_class"), [ ("mjlab.tasks.velocity.rl.runner", "VelocityOnPolicyRunner"), ("mjlab.tasks.manipulation.rl.runner", "ManipulationOnPolicyRunner"), ], ) def test_task_runner_uploads_onnx_for_wandb_logger( runner_module, runner_class, monkeypatch, tmp_path ): """ONNX is uploaded to W&B when the logger's writer is a WandbLogWriter.""" import importlib runner_mod = importlib.import_module(runner_module) runner_cls = getattr(runner_mod, runner_class) runner = runner_cls.__new__(runner_cls) runner.cfg = {"upload_model": True} runner.logger = _make_logger_mock(is_wandb=True) runner.env = MagicMock() runner.export_policy_to_onnx = MagicMock() monkeypatch.setattr(MjlabOnPolicyRunner, "save", lambda *a, **kw: None) monkeypatch.setattr(runner_mod, "get_base_metadata", lambda *a: {}) monkeypatch.setattr(runner_mod, "attach_metadata_to_onnx", lambda *a: None) checkpoint = tmp_path / "run-dir" / "model_100.pt" checkpoint.parent.mkdir() checkpoint.touch() onnx_path = checkpoint.parent / f"{checkpoint.parent.name}.onnx" mock_run = MagicMock() mock_run.name = "test-run" with patch.object(runner_mod, "wandb") as mock_wandb: mock_wandb.run = mock_run runner.save(str(checkpoint)) mock_wandb.save.assert_called_once_with( str(onnx_path), base_path=str(checkpoint.parent) ) def _make_tracking_runner_shell(registry_name, is_wandb, upload_model=True): """Build a MotionTrackingOnPolicyRunner with all heavy parts mocked out.""" from mjlab.tasks.tracking.rl.runner import MotionTrackingOnPolicyRunner runner = MotionTrackingOnPolicyRunner.__new__(MotionTrackingOnPolicyRunner) runner.registry_name = registry_name runner.cfg = {"upload_model": upload_model} runner.logger = _make_logger_mock(is_wandb) mock_motion_term = MagicMock() mock_motion_term.cfg.anchor_body_name = "pelvis" mock_motion_term.cfg.body_names = ["body1"] runner.env = MagicMock() runner.env.unwrapped.command_manager.get_term.return_value = mock_motion_term return runner def test_tracking_runner_registers_artifact_for_wandb_logger(monkeypatch, tmp_path): """use_artifact is called when the logger's writer is a WandbLogWriter.""" from mjlab.rl.runner import MjlabOnPolicyRunner from mjlab.tasks.tracking.rl import runner as runner_mod runner = _make_tracking_runner_shell("org/motions/motion:latest", is_wandb=True) monkeypatch.setattr(MjlabOnPolicyRunner, "save", lambda *a, **kw: None) monkeypatch.setattr(runner_mod, "get_base_metadata", lambda *a: {}) monkeypatch.setattr(runner_mod, "attach_metadata_to_onnx", lambda *a: None) monkeypatch.setattr( runner.env.unwrapped.__class__, "export_policy_to_onnx", lambda *a, **kw: None, raising=False, ) checkpoint = tmp_path / "run-dir" / "model_100.pt" checkpoint.parent.mkdir() checkpoint.touch() mock_run = MagicMock() mock_run.name = "test-run" with patch.object(runner_mod, "wandb") as mock_wandb: mock_wandb.run = mock_run runner.export_policy_to_onnx = MagicMock() runner.save(str(checkpoint)) mock_run.use_artifact.assert_called_once_with("org/motions/motion:latest") def test_tracking_runner_does_not_register_artifact_for_tensorboard( monkeypatch, tmp_path ): """use_artifact is NOT called when using the tensorboard logger.""" from mjlab.rl.runner import MjlabOnPolicyRunner from mjlab.tasks.tracking.rl import runner as runner_mod runner = _make_tracking_runner_shell("org/motions/motion:latest", is_wandb=False) monkeypatch.setattr(MjlabOnPolicyRunner, "save", lambda *a, **kw: None) monkeypatch.setattr(runner_mod, "get_base_metadata", lambda *a: {}) monkeypatch.setattr(runner_mod, "attach_metadata_to_onnx", lambda *a: None) checkpoint = tmp_path / "run-dir" / "model_100.pt" checkpoint.parent.mkdir() checkpoint.touch() with patch.object(runner_mod, "wandb") as mock_wandb: runner.export_policy_to_onnx = MagicMock() runner.save(str(checkpoint)) mock_wandb.run.use_artifact.assert_not_called()