64 lines
1.6 KiB
Python
64 lines
1.6 KiB
Python
|
"""Tests for ObservationManager behavior."""
|
||
|
|
|
||
|
|
from unittest.mock import Mock
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
import torch
|
||
|
|
from conftest import get_test_device
|
||
|
|
|
||
|
|
from mjlab.managers.observation_manager import (
|
||
|
|
ObservationGroupCfg,
|
||
|
|
ObservationManager,
|
||
|
|
ObservationTermCfg,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
@pytest.fixture
|
||
|
|
def mock_env():
|
||
|
|
env = Mock()
|
||
|
|
env.num_envs = 4
|
||
|
|
env.device = get_test_device()
|
||
|
|
env.step_dt = 0.02
|
||
|
|
return env
|
||
|
|
|
||
|
|
|
||
|
|
def _dummy_term(mock_env):
|
||
|
|
def func(_env, **_kwargs):
|
||
|
|
return torch.zeros(mock_env.num_envs, 3, device=mock_env.device)
|
||
|
|
|
||
|
|
return ObservationTermCfg(func=func)
|
||
|
|
|
||
|
|
|
||
|
|
def test_empty_terms_dict_skipped(mock_env):
|
||
|
|
"""A group declared with no terms is skipped rather than raising."""
|
||
|
|
cfg = {"actor": ObservationGroupCfg(terms={})}
|
||
|
|
mgr = ObservationManager(cfg, mock_env)
|
||
|
|
|
||
|
|
assert "actor" not in mgr.active_terms
|
||
|
|
assert "actor" not in mgr.group_obs_dim
|
||
|
|
|
||
|
|
|
||
|
|
def test_all_terms_none_skipped(mock_env):
|
||
|
|
"""A group whose every term is None is skipped rather than raising."""
|
||
|
|
cfg = {
|
||
|
|
"actor": ObservationGroupCfg(
|
||
|
|
terms={"a": None, "b": None}, # type: ignore[dict-item]
|
||
|
|
),
|
||
|
|
}
|
||
|
|
mgr = ObservationManager(cfg, mock_env)
|
||
|
|
|
||
|
|
assert "actor" not in mgr.active_terms
|
||
|
|
assert "actor" not in mgr.group_obs_dim
|
||
|
|
|
||
|
|
|
||
|
|
def test_empty_group_skipped_alongside_active_group(mock_env):
|
||
|
|
"""Active groups coexist with empty ones; only the empty group is dropped."""
|
||
|
|
cfg = {
|
||
|
|
"actor": ObservationGroupCfg(terms={"a": _dummy_term(mock_env)}),
|
||
|
|
"critic": ObservationGroupCfg(terms={}),
|
||
|
|
}
|
||
|
|
mgr = ObservationManager(cfg, mock_env)
|
||
|
|
|
||
|
|
assert "actor" in mgr.active_terms
|
||
|
|
assert "critic" not in mgr.active_terms
|