Import upstream snapshot c19f713c415a699a79d71cd96aa13c3104a05047

Upstream: https://github.com/michaelgillett/mjlab
Upstream-Commit: c19f713c415a699a79d71cd96aa13c3104a05047
Upstream-Branch: main

tests/test_manager_config_immutability.py

@@ -0,0 +1,172 @@
+"""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}"
+  )