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
173 lines
5.3 KiB
Python
173 lines
5.3 KiB
Python
"""Tests for manager config immutability.
|
|
|
|
Managers should not mutate the original config objects passed to them.
|
|
This allows the same config to be reused across multiple manager instantiations.
|
|
"""
|
|
|
|
import inspect
|
|
from unittest.mock import Mock
|
|
|
|
import mujoco
|
|
import pytest
|
|
import torch
|
|
from conftest import get_test_device
|
|
|
|
from mjlab.entity import Entity, EntityCfg
|
|
from mjlab.managers.observation_manager import (
|
|
ObservationGroupCfg,
|
|
ObservationManager,
|
|
ObservationTermCfg,
|
|
)
|
|
from mjlab.managers.reward_manager import RewardManager, RewardTermCfg
|
|
from mjlab.managers.scene_entity_config import SceneEntityCfg
|
|
from mjlab.sim.sim import Simulation, SimulationCfg
|
|
|
|
|
|
@pytest.fixture
|
|
def device():
|
|
return get_test_device()
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_env(device):
|
|
env = Mock()
|
|
env.num_envs = 2
|
|
env.device = device
|
|
env.step_dt = 0.02
|
|
env.max_episode_length_s = 10.0
|
|
env.scene = {}
|
|
return env
|
|
|
|
|
|
class ClassObsTerm:
|
|
"""Class-based observation term that triggers func mutation."""
|
|
|
|
def __init__(self, cfg, env):
|
|
self.device = env.device
|
|
|
|
def __call__(self, env):
|
|
return torch.zeros((env.num_envs, 3), device=self.device)
|
|
|
|
|
|
class ClassRewardTerm:
|
|
"""Class-based reward term that triggers func mutation."""
|
|
|
|
def __init__(self, cfg, env):
|
|
self.device = env.device
|
|
|
|
def __call__(self, env):
|
|
return torch.ones(env.num_envs, device=self.device)
|
|
|
|
|
|
def test_manager_preserves_class_func(mock_env):
|
|
"""Managers should not mutate class-based func to instance.
|
|
|
|
When term_cfg.func is a class, _resolve_common_term_cfg instantiates it.
|
|
Without copying term_cfg first, this mutates the original config.
|
|
"""
|
|
term_cfg = ObservationTermCfg(func=ClassObsTerm, params={})
|
|
|
|
assert inspect.isclass(term_cfg.func), "precondition: func should be a class"
|
|
|
|
cfg = {"actor": ObservationGroupCfg(terms={"obs1": term_cfg})}
|
|
_ = ObservationManager(cfg, mock_env)
|
|
|
|
assert inspect.isclass(term_cfg.func), "func was mutated from class to instance"
|
|
|
|
|
|
def test_observation_shared_terms_between_groups(mock_env):
|
|
"""Terms shared between policy and critic groups get independent settings."""
|
|
term_cfg = ObservationTermCfg(func=ClassObsTerm, params={})
|
|
|
|
cfg = {
|
|
"actor": ObservationGroupCfg(terms={"obs": term_cfg}, history_length=5),
|
|
"critic": ObservationGroupCfg(terms={"obs": term_cfg}), # no history
|
|
}
|
|
|
|
manager = ObservationManager(cfg, mock_env)
|
|
|
|
policy_obs = manager.compute()["actor"]
|
|
critic_obs = manager.compute()["critic"]
|
|
|
|
# Policy has history (5 * 3 = 15), critic doesn't (3).
|
|
assert isinstance(policy_obs, torch.Tensor)
|
|
assert isinstance(critic_obs, torch.Tensor)
|
|
assert policy_obs.shape[1] == 15 # 3 features * 5 history
|
|
assert critic_obs.shape[1] == 3 # 3 features, no history
|
|
|
|
|
|
def test_get_term_cfg_returns_resolved_config(mock_env):
|
|
"""get_term_cfg should return resolved config with func as instance."""
|
|
term_cfg = ObservationTermCfg(func=ClassObsTerm, params={})
|
|
|
|
# Original config has func as a class.
|
|
assert inspect.isclass(term_cfg.func), "precondition: func should be a class"
|
|
|
|
cfg = {"actor": ObservationGroupCfg(terms={"obs1": term_cfg})}
|
|
manager = ObservationManager(cfg, mock_env)
|
|
|
|
# Original config should remain unchanged.
|
|
assert inspect.isclass(term_cfg.func), "original config should not be mutated"
|
|
assert inspect.isclass(manager.cfg["actor"].terms["obs1"].func), (
|
|
"manager.cfg should preserve original class"
|
|
)
|
|
|
|
# Resolved config should have func as an instance.
|
|
resolved_cfg = manager.get_term_cfg("actor", "obs1")
|
|
assert not inspect.isclass(resolved_cfg.func), (
|
|
"get_term_cfg should return resolved config with func as instance"
|
|
)
|
|
assert isinstance(resolved_cfg.func, ClassObsTerm), (
|
|
f"Expected ClassObsTerm instance, got {type(resolved_cfg.func)}"
|
|
)
|
|
|
|
|
|
def test_scene_entity_cfg_in_params_not_mutated(device):
|
|
"""SceneEntityCfg in params should not be mutated when resolved."""
|
|
# NOTE: 2 joints so selecting one won't optimize to slice(None).
|
|
xml = """
|
|
<mujoco>
|
|
<worldbody>
|
|
<body name="body1" pos="0 0 1">
|
|
<joint name="joint1" type="hinge" axis="0 0 1"/>
|
|
<geom type="sphere" size="0.1" mass="1"/>
|
|
<body name="body2" pos="0.2 0 0">
|
|
<joint name="joint2" type="hinge" axis="0 0 1"/>
|
|
<geom type="sphere" size="0.1" mass="1"/>
|
|
</body>
|
|
</body>
|
|
</worldbody>
|
|
</mujoco>
|
|
"""
|
|
entity_cfg = EntityCfg(spec_fn=lambda: mujoco.MjSpec.from_string(xml))
|
|
entity = Entity(entity_cfg)
|
|
model = entity.compile()
|
|
|
|
sim = Simulation(num_envs=2, cfg=SimulationCfg(), model=model, device=device)
|
|
entity.initialize(model, sim.model, sim.data, device)
|
|
|
|
env = Mock()
|
|
env.num_envs = 2
|
|
env.device = device
|
|
env.max_episode_length_s = 10.0
|
|
env.scene = {"robot": entity}
|
|
|
|
# Select only joint1 --> joint_ids will mutate from slice(None) to [0].
|
|
asset_cfg = SceneEntityCfg(name="robot", joint_names=("joint1",))
|
|
original_joint_ids = asset_cfg.joint_ids # slice(None) before resolve
|
|
|
|
def reward_func(env, asset_cfg):
|
|
del asset_cfg # Unused.
|
|
return torch.ones(env.num_envs, device=env.device)
|
|
|
|
term_cfg = RewardTermCfg(
|
|
func=reward_func, params={"asset_cfg": asset_cfg}, weight=1.0
|
|
)
|
|
|
|
_ = RewardManager({"reward1": term_cfg}, env)
|
|
|
|
assert asset_cfg.joint_ids == original_joint_ids, (
|
|
f"SceneEntityCfg.joint_ids was mutated from {original_joint_ids} to "
|
|
f"{asset_cfg.joint_ids}"
|
|
)
|