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
613 lines
17 KiB
Python
613 lines
17 KiB
Python
"""Tests for observation history functionality."""
|
|
|
|
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 device():
|
|
"""Test device fixture."""
|
|
return get_test_device()
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_env(device):
|
|
"""Create a mock environment."""
|
|
env = Mock()
|
|
env.num_envs = 4
|
|
env.device = device
|
|
env.step_dt = 0.02
|
|
return env
|
|
|
|
|
|
@pytest.fixture
|
|
def simple_obs_func(device):
|
|
"""Create a simple observation function that returns a counter."""
|
|
counter = {"value": 0}
|
|
|
|
def obs_func(env):
|
|
counter["value"] += 1
|
|
return torch.full((env.num_envs, 3), float(counter["value"]), device=device)
|
|
|
|
return obs_func
|
|
|
|
|
|
# Basic observation history tests.
|
|
|
|
|
|
def test_no_history_by_default(mock_env, simple_obs_func):
|
|
"""Test that observations work without history (default behavior)."""
|
|
|
|
cfg = {
|
|
"actor": ObservationGroupCfg(
|
|
terms={
|
|
"obs1": ObservationTermCfg(func=simple_obs_func, params={}),
|
|
}
|
|
),
|
|
}
|
|
|
|
manager = ObservationManager(cfg, mock_env)
|
|
assert manager.group_obs_dim["actor"] == (3,)
|
|
|
|
obs = manager.compute()
|
|
policy_obs = obs["actor"]
|
|
assert isinstance(policy_obs, torch.Tensor)
|
|
assert policy_obs.shape == (4, 3)
|
|
|
|
|
|
def test_single_step_history(mock_env, simple_obs_func):
|
|
"""Test observation with history_length=1."""
|
|
|
|
cfg = {
|
|
"actor": ObservationGroupCfg(
|
|
terms={
|
|
"obs1": ObservationTermCfg(func=simple_obs_func, params={}, history_length=1),
|
|
}
|
|
),
|
|
}
|
|
|
|
manager = ObservationManager(cfg, mock_env)
|
|
assert manager.group_obs_dim["actor"] == (3,)
|
|
|
|
obs = manager.compute()
|
|
policy_obs = obs["actor"]
|
|
assert isinstance(policy_obs, torch.Tensor)
|
|
assert policy_obs.shape == (4, 3)
|
|
|
|
|
|
def test_multi_step_history_flattened(mock_env, simple_obs_func):
|
|
"""Test observation with history_length=3 and flattened."""
|
|
|
|
cfg = {
|
|
"actor": ObservationGroupCfg(
|
|
terms={
|
|
"obs1": ObservationTermCfg(
|
|
func=simple_obs_func, params={}, history_length=3, flatten_history_dim=True
|
|
),
|
|
}
|
|
),
|
|
}
|
|
|
|
manager = ObservationManager(cfg, mock_env)
|
|
assert manager.group_obs_dim["actor"] == (9,)
|
|
|
|
obs = manager.compute()
|
|
policy_obs = obs["actor"]
|
|
assert isinstance(policy_obs, torch.Tensor)
|
|
assert policy_obs.shape == (4, 9)
|
|
|
|
|
|
def test_multi_step_history_not_flattened(mock_env, simple_obs_func):
|
|
"""Test observation with history_length=3 and not flattened."""
|
|
|
|
cfg = {
|
|
"actor": ObservationGroupCfg(
|
|
terms={
|
|
"obs1": ObservationTermCfg(
|
|
func=simple_obs_func, params={}, history_length=3, flatten_history_dim=False
|
|
),
|
|
}
|
|
),
|
|
}
|
|
|
|
manager = ObservationManager(cfg, mock_env)
|
|
assert manager.group_obs_dim["actor"] == (3, 3)
|
|
|
|
obs = manager.compute()
|
|
policy_obs = obs["actor"]
|
|
assert isinstance(policy_obs, torch.Tensor)
|
|
assert policy_obs.shape == (4, 3, 3)
|
|
|
|
|
|
# History accumulation tests.
|
|
|
|
|
|
def test_history_accumulates_correctly(mock_env, simple_obs_func):
|
|
"""Test that history buffer accumulates observations in correct order."""
|
|
|
|
cfg = {
|
|
"actor": ObservationGroupCfg(
|
|
terms={
|
|
"obs1": ObservationTermCfg(
|
|
func=simple_obs_func, params={}, history_length=3, flatten_history_dim=False
|
|
),
|
|
}
|
|
),
|
|
}
|
|
|
|
manager = ObservationManager(cfg, mock_env)
|
|
device = mock_env.device
|
|
# Note: counter is incremented during _prepare_terms (value=1).
|
|
|
|
# First compute uses value=2 and initializes buffer.
|
|
obs = manager.compute(update_history=False)
|
|
policy_obs = obs["actor"]
|
|
assert isinstance(policy_obs, torch.Tensor)
|
|
assert torch.allclose(policy_obs[0], torch.full((3, 3), 2.0, device=device))
|
|
|
|
# Update with value=3.
|
|
obs = manager.compute(update_history=True)
|
|
policy_obs = obs["actor"]
|
|
assert isinstance(policy_obs, torch.Tensor)
|
|
# History: [2, 2, 3] (oldest to newest).
|
|
expected = torch.stack(
|
|
[
|
|
torch.full((3,), 2.0, device=device),
|
|
torch.full((3,), 2.0, device=device),
|
|
torch.full((3,), 3.0, device=device),
|
|
]
|
|
)
|
|
assert torch.allclose(policy_obs[0], expected)
|
|
|
|
# Update with value=4.
|
|
obs = manager.compute(update_history=True)
|
|
policy_obs = obs["actor"]
|
|
assert isinstance(policy_obs, torch.Tensor)
|
|
# History: [2, 3, 4].
|
|
expected = torch.stack(
|
|
[
|
|
torch.full((3,), 2.0, device=device),
|
|
torch.full((3,), 3.0, device=device),
|
|
torch.full((3,), 4.0, device=device),
|
|
]
|
|
)
|
|
assert torch.allclose(policy_obs[0], expected)
|
|
|
|
# Update with value=5, circular overwrite.
|
|
obs = manager.compute(update_history=True)
|
|
policy_obs = obs["actor"]
|
|
assert isinstance(policy_obs, torch.Tensor)
|
|
# History: [3, 4, 5].
|
|
expected = torch.stack(
|
|
[
|
|
torch.full((3,), 3.0, device=device),
|
|
torch.full((3,), 4.0, device=device),
|
|
torch.full((3,), 5.0, device=device),
|
|
]
|
|
)
|
|
assert torch.allclose(policy_obs[0], expected)
|
|
|
|
|
|
def test_update_history_false_doesnt_modify_buffer(mock_env, simple_obs_func):
|
|
"""Test that update_history=False doesn't modify the buffer."""
|
|
|
|
cfg = {
|
|
"actor": ObservationGroupCfg(
|
|
terms={
|
|
"obs1": ObservationTermCfg(
|
|
func=simple_obs_func, params={}, history_length=2, flatten_history_dim=False
|
|
),
|
|
}
|
|
),
|
|
}
|
|
|
|
manager = ObservationManager(cfg, mock_env)
|
|
|
|
# Initialize (value=1).
|
|
obs1 = manager.compute(update_history=False)
|
|
policy_obs1 = obs1["actor"]
|
|
assert isinstance(policy_obs1, torch.Tensor)
|
|
|
|
# Call without update (value=2, but buffer unchanged).
|
|
obs2 = manager.compute(update_history=False)
|
|
policy_obs2 = obs2["actor"]
|
|
assert isinstance(policy_obs2, torch.Tensor)
|
|
|
|
# History should still be [1, 1].
|
|
assert torch.allclose(policy_obs1, policy_obs2)
|
|
|
|
|
|
# Group-level history override tests.
|
|
|
|
|
|
def test_group_history_overrides_term(mock_env, simple_obs_func):
|
|
"""Test group history_length overrides term history_length."""
|
|
|
|
cfg = {
|
|
"actor": ObservationGroupCfg(
|
|
history_length=5, # Group level.
|
|
flatten_history_dim=False,
|
|
terms={
|
|
"obs1": ObservationTermCfg(
|
|
func=simple_obs_func,
|
|
params={},
|
|
history_length=2, # Overridden.
|
|
),
|
|
},
|
|
),
|
|
}
|
|
|
|
manager = ObservationManager(cfg, mock_env)
|
|
assert manager.group_obs_dim["actor"] == (5, 3)
|
|
|
|
obs = manager.compute()
|
|
policy_obs = obs["actor"]
|
|
assert isinstance(policy_obs, torch.Tensor)
|
|
assert policy_obs.shape == (4, 5, 3)
|
|
|
|
|
|
# History reset tests.
|
|
|
|
|
|
def test_reset_clears_all_envs(mock_env, simple_obs_func):
|
|
"""Test that reset without env_ids clears all environments."""
|
|
|
|
cfg = {
|
|
"actor": ObservationGroupCfg(
|
|
terms={
|
|
"obs1": ObservationTermCfg(
|
|
func=simple_obs_func, params={}, history_length=2, flatten_history_dim=False
|
|
),
|
|
}
|
|
),
|
|
}
|
|
|
|
manager = ObservationManager(cfg, mock_env)
|
|
|
|
# Build up history.
|
|
manager.compute(update_history=True)
|
|
manager.compute(update_history=True)
|
|
|
|
# Reset all.
|
|
manager.reset()
|
|
|
|
# Buffer should be zeroed.
|
|
obs = manager.compute(update_history=False)
|
|
policy_obs = obs["actor"]
|
|
assert isinstance(policy_obs, torch.Tensor)
|
|
assert torch.allclose(policy_obs, torch.zeros((4, 2, 3), device=mock_env.device))
|
|
|
|
|
|
def test_reset_partial_envs(mock_env, simple_obs_func):
|
|
"""Test that reset with specific env_ids only resets those envs."""
|
|
|
|
cfg = {
|
|
"actor": ObservationGroupCfg(
|
|
terms={
|
|
"obs1": ObservationTermCfg(
|
|
func=simple_obs_func, params={}, history_length=3, flatten_history_dim=False
|
|
),
|
|
}
|
|
),
|
|
}
|
|
|
|
manager = ObservationManager(cfg, mock_env)
|
|
|
|
# Build up history.
|
|
manager.compute(update_history=True)
|
|
manager.compute(update_history=True)
|
|
obs_before = manager.compute(update_history=True)
|
|
policy_obs_before = obs_before["actor"]
|
|
assert isinstance(policy_obs_before, torch.Tensor)
|
|
|
|
# Reset only envs 0 and 2.
|
|
manager.reset(env_ids=torch.tensor([0, 2], device=mock_env.device))
|
|
|
|
obs_after = manager.compute(update_history=False)
|
|
policy_obs_after = obs_after["actor"]
|
|
assert isinstance(policy_obs_after, torch.Tensor)
|
|
|
|
# Envs 0 and 2 should be reset (zeros), 1 and 3 unchanged.
|
|
assert torch.allclose(
|
|
policy_obs_after[0], torch.zeros((3, 3), device=mock_env.device)
|
|
)
|
|
assert torch.allclose(policy_obs_after[1], policy_obs_before[1])
|
|
assert torch.allclose(
|
|
policy_obs_after[2], torch.zeros((3, 3), device=mock_env.device)
|
|
)
|
|
assert torch.allclose(policy_obs_after[3], policy_obs_before[3])
|
|
|
|
|
|
def test_reset_partial_envs_with_backfill(mock_env, simple_obs_func):
|
|
"""Test that reset envs get backfilled on next update."""
|
|
|
|
cfg = {
|
|
"actor": ObservationGroupCfg(
|
|
terms={
|
|
"obs1": ObservationTermCfg(
|
|
func=simple_obs_func, params={}, history_length=3, flatten_history_dim=False
|
|
),
|
|
}
|
|
),
|
|
}
|
|
|
|
manager = ObservationManager(cfg, mock_env)
|
|
device = mock_env.device
|
|
|
|
# Build history: [2, 3, 4] for all envs (counter starts at 1 after _prepare_terms).
|
|
manager.compute(update_history=True) # value=2
|
|
manager.compute(update_history=True) # value=3
|
|
manager.compute(update_history=True) # value=4
|
|
|
|
# Reset only envs 0 and 2.
|
|
manager.reset(env_ids=torch.tensor([0, 2], device=device))
|
|
|
|
# Next update with value=5.
|
|
obs = manager.compute(update_history=True)
|
|
policy_obs = obs["actor"]
|
|
assert isinstance(policy_obs, torch.Tensor)
|
|
|
|
# Env 0: [5, 5, 5] (backfilled after reset).
|
|
expected_env0 = torch.stack(
|
|
[
|
|
torch.full((3,), 5.0, device=device),
|
|
torch.full((3,), 5.0, device=device),
|
|
torch.full((3,), 5.0, device=device),
|
|
]
|
|
)
|
|
assert torch.allclose(policy_obs[0], expected_env0)
|
|
|
|
# Env 1: [3, 4, 5] (continues normally).
|
|
expected_env1 = torch.stack(
|
|
[
|
|
torch.full((3,), 3.0, device=device),
|
|
torch.full((3,), 4.0, device=device),
|
|
torch.full((3,), 5.0, device=device),
|
|
]
|
|
)
|
|
assert torch.allclose(policy_obs[1], expected_env1)
|
|
|
|
# Env 2: [5, 5, 5] (backfilled after reset).
|
|
assert torch.allclose(policy_obs[2], expected_env0)
|
|
|
|
# Env 3: [3, 4, 5] (continues normally).
|
|
assert torch.allclose(policy_obs[3], expected_env1)
|
|
|
|
|
|
# History with other features tests.
|
|
|
|
|
|
def test_history_with_clip(mock_env, simple_obs_func):
|
|
"""Test that clipping is applied before history."""
|
|
|
|
cfg = {
|
|
"actor": ObservationGroupCfg(
|
|
terms={
|
|
"obs1": ObservationTermCfg(
|
|
func=simple_obs_func,
|
|
params={},
|
|
history_length=2,
|
|
flatten_history_dim=False,
|
|
clip=(-0.5, 0.5),
|
|
),
|
|
}
|
|
),
|
|
}
|
|
|
|
manager = ObservationManager(cfg, mock_env)
|
|
|
|
obs = manager.compute(update_history=True)
|
|
policy_obs = obs["actor"]
|
|
assert isinstance(policy_obs, torch.Tensor)
|
|
# Values should be clipped.
|
|
assert torch.all(policy_obs >= -0.5)
|
|
assert torch.all(policy_obs <= 0.5)
|
|
|
|
|
|
def test_history_with_scale(mock_env, simple_obs_func):
|
|
"""Test that scaling is applied before history."""
|
|
|
|
cfg = {
|
|
"actor": ObservationGroupCfg(
|
|
terms={
|
|
"obs1": ObservationTermCfg(
|
|
func=simple_obs_func,
|
|
params={},
|
|
history_length=2,
|
|
flatten_history_dim=False,
|
|
scale=2.0,
|
|
),
|
|
}
|
|
),
|
|
}
|
|
|
|
manager = ObservationManager(cfg, mock_env)
|
|
device = mock_env.device
|
|
# Counter at 1 after _prepare_terms.
|
|
|
|
# First call uses value=2, scaled to 4, initializes buffer.
|
|
manager.compute(update_history=False)
|
|
# Second call uses value=3, scaled to 6, and updates history.
|
|
obs = manager.compute(update_history=True)
|
|
policy_obs = obs["actor"]
|
|
assert isinstance(policy_obs, torch.Tensor)
|
|
|
|
# History: [4, 6].
|
|
expected = torch.stack(
|
|
[torch.full((3,), 4.0, device=device), torch.full((3,), 6.0, device=device)]
|
|
)
|
|
assert torch.allclose(policy_obs[0], expected)
|
|
|
|
|
|
# Mixed history terms tests.
|
|
|
|
|
|
def test_mixed_terms_concatenated(mock_env, simple_obs_func, device):
|
|
"""Test group with both history and non-history terms concatenated."""
|
|
|
|
counter = {"value": 0}
|
|
|
|
def obs_func2(env):
|
|
counter["value"] += 1
|
|
return torch.full((env.num_envs, 2), float(counter["value"]) * 10, device=device)
|
|
|
|
cfg = {
|
|
"actor": ObservationGroupCfg(
|
|
terms={
|
|
"obs_with_history": ObservationTermCfg(
|
|
func=simple_obs_func, params={}, history_length=2, flatten_history_dim=True
|
|
),
|
|
"obs_no_history": ObservationTermCfg(func=obs_func2, params={}),
|
|
}
|
|
),
|
|
}
|
|
|
|
manager = ObservationManager(cfg, mock_env)
|
|
|
|
# Should concatenate: (3*2) + 2 = 8.
|
|
assert manager.group_obs_dim["actor"] == (8,)
|
|
|
|
obs = manager.compute()
|
|
policy_obs = obs["actor"]
|
|
assert isinstance(policy_obs, torch.Tensor)
|
|
assert policy_obs.shape == (4, 8)
|
|
|
|
|
|
def test_no_double_append_on_first_call(mock_env, simple_obs_func):
|
|
"""Test that first call with update_history=True only appends once, not twice."""
|
|
|
|
cfg = {
|
|
"actor": ObservationGroupCfg(
|
|
terms={
|
|
"obs1": ObservationTermCfg(
|
|
func=simple_obs_func, params={}, history_length=3, flatten_history_dim=False
|
|
),
|
|
}
|
|
),
|
|
}
|
|
|
|
manager = ObservationManager(cfg, mock_env)
|
|
device = mock_env.device
|
|
# Counter is at 1 after _prepare_terms.
|
|
|
|
# First call with update_history=True (value=2).
|
|
# This should initialize the buffer AND append once (not twice).
|
|
obs = manager.compute(update_history=True)
|
|
policy_obs = obs["actor"]
|
|
assert isinstance(policy_obs, torch.Tensor)
|
|
|
|
# Verify buffer was initialized and backfilled correctly.
|
|
# All slots should be filled with value 2.
|
|
expected_first = torch.stack(
|
|
[
|
|
torch.full((3,), 2.0, device=device),
|
|
torch.full((3,), 2.0, device=device),
|
|
torch.full((3,), 2.0, device=device),
|
|
]
|
|
)
|
|
assert torch.allclose(policy_obs[0], expected_first)
|
|
|
|
# Get the circular buffer and check pointer position.
|
|
circular_buffer = manager._group_obs_term_history_buffer["actor"]["obs1"]
|
|
# After one append, pointer should be at 0 (not 1 which would indicate double-append).
|
|
assert circular_buffer._pointer == 0
|
|
# And we should have exactly 1 push recorded.
|
|
assert torch.all(circular_buffer._num_pushes == 1)
|
|
|
|
# Second call with update_history=True (value=3).
|
|
obs = manager.compute(update_history=True)
|
|
policy_obs = obs["actor"]
|
|
assert isinstance(policy_obs, torch.Tensor)
|
|
|
|
# History should be [2, 2, 3] (oldest to newest).
|
|
expected_second = torch.stack(
|
|
[
|
|
torch.full((3,), 2.0, device=device),
|
|
torch.full((3,), 2.0, device=device),
|
|
torch.full((3,), 3.0, device=device),
|
|
]
|
|
)
|
|
assert torch.allclose(policy_obs[0], expected_second)
|
|
|
|
# Pointer should now be at 1.
|
|
assert circular_buffer._pointer == 1
|
|
# And we should have exactly 2 pushes.
|
|
assert torch.all(circular_buffer._num_pushes == 2)
|
|
|
|
|
|
def test_term_major_ordering(mock_env, device):
|
|
"""Test that history flattening uses term-major ordering.
|
|
|
|
Term-major: [A_t0, A_t1, ..., A_tH-1, B_t0, B_t1, ..., B_tH-1]
|
|
Time-major: [A_t0, B_t0, ..., A_t1, B_t1, ..., A_tH-1, B_tH-1]
|
|
|
|
This test verifies mjlab uses term-major ordering.
|
|
"""
|
|
|
|
# Create observation functions with distinct values per term.
|
|
def obs_A(env):
|
|
# Returns [100, 101] for easy identification.
|
|
return torch.tensor([[100.0, 101.0]] * env.num_envs, device=device)
|
|
|
|
def obs_B(env):
|
|
# Returns [200, 201, 202] for easy identification.
|
|
return torch.tensor([[200.0, 201.0, 202.0]] * env.num_envs, device=device)
|
|
|
|
cfg = {
|
|
"actor": ObservationGroupCfg(
|
|
concatenate_terms=True,
|
|
terms={
|
|
"term_A": ObservationTermCfg(
|
|
func=obs_A, params={}, history_length=3, flatten_history_dim=True
|
|
),
|
|
"term_B": ObservationTermCfg(
|
|
func=obs_B, params={}, history_length=3, flatten_history_dim=True
|
|
),
|
|
},
|
|
),
|
|
}
|
|
|
|
manager = ObservationManager(cfg, mock_env)
|
|
|
|
# Compute observations (history will be backfilled with same values).
|
|
obs = manager.compute(update_history=False)
|
|
policy_obs = obs["actor"]
|
|
assert isinstance(policy_obs, torch.Tensor)
|
|
|
|
# Expected shape: (4 envs, 2*3 + 3*3) = (4, 15).
|
|
assert policy_obs.shape == (4, 15)
|
|
|
|
# Check ordering for first environment.
|
|
result = policy_obs[0].cpu().tolist()
|
|
|
|
# Term-major: all A history, then all B history.
|
|
# A has 2 dims, B has 3 dims, history_length=3.
|
|
# fmt: off
|
|
expected_term_major = [
|
|
100.0, 101.0, 100.0, 101.0, 100.0, 101.0, # All A history
|
|
200.0, 201.0, 202.0, 200.0, 201.0, 202.0, 200.0, 201.0, 202.0, # All B
|
|
]
|
|
# fmt: on
|
|
|
|
# Time-major would interleave: [A_t0, B_t0, A_t1, B_t1, A_t2, B_t2].
|
|
# fmt: off
|
|
expected_time_major = [
|
|
100.0, 101.0, 200.0, 201.0, 202.0, # Frame t0
|
|
100.0, 101.0, 200.0, 201.0, 202.0, # Frame t1
|
|
100.0, 101.0, 200.0, 201.0, 202.0, # Frame t2
|
|
]
|
|
# fmt: on
|
|
|
|
# Verify term-major ordering.
|
|
assert result == expected_term_major, f"Expected term-major ordering, got {result}"
|
|
assert result != expected_time_major, "Should not match time-major"
|