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
662 lines
21 KiB
Python
662 lines
21 KiB
Python
"""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()
|