"""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 = """
  <mujoco>
    <worldbody>
      <body name="base" pos="0 0 1">
        <freejoint name="free_joint"/>
        <geom name="base_geom" type="box" size="0.2 0.2 0.1" mass="1.0"/>
        <body name="link1" pos="0 0 0">
          <joint name="joint1" type="hinge" axis="0 0 1" range="-1.57 1.57"/>
          <geom name="link1_geom" type="box" size="0.1 0.1 0.1" mass="0.1"/>
        </body>
      </body>
    </worldbody>
    <actuator>
      <motor name="actuator1" joint="joint1" gear="1.0"/>
    </actuator>
  </mujoco>
  """
  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()
