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
328 lines
9.2 KiB
Python
328 lines
9.2 KiB
Python
"""Tests for DelayBuffer."""
|
|
|
|
import pytest
|
|
import torch
|
|
from conftest import get_test_device
|
|
|
|
from mjlab.utils.buffers import DelayBuffer
|
|
|
|
|
|
@pytest.fixture
|
|
def device():
|
|
"""Test device fixture."""
|
|
return get_test_device()
|
|
|
|
|
|
def make_gen(seed: int, device: str) -> torch.Generator:
|
|
"""Create a seeded generator for reproducible tests."""
|
|
gen = torch.Generator(device=device)
|
|
gen.manual_seed(seed)
|
|
return gen
|
|
|
|
|
|
##
|
|
# Basic behavior.
|
|
##
|
|
|
|
|
|
def test_delay_buffer_zero_lag(device):
|
|
"""Zero lag returns current observation."""
|
|
buffer = DelayBuffer(min_lag=0, max_lag=0, batch_size=2, device=device)
|
|
|
|
buffer.append(torch.tensor([[1.0], [2.0]], device=device))
|
|
result = buffer.compute()
|
|
assert torch.allclose(result, torch.tensor([[1.0], [2.0]], device=device))
|
|
|
|
buffer.append(torch.tensor([[3.0], [4.0]], device=device))
|
|
result = buffer.compute()
|
|
assert torch.allclose(result, torch.tensor([[3.0], [4.0]], device=device))
|
|
|
|
|
|
def test_delay_buffer_constant_lag(device):
|
|
"""Constant lag returns observation from N steps ago."""
|
|
buffer = DelayBuffer(min_lag=2, max_lag=2, batch_size=1, device=device)
|
|
|
|
buffer.append(torch.tensor([[1.0]], device=device))
|
|
result = buffer.compute()
|
|
assert torch.allclose(result, torch.tensor([[1.0]], device=device))
|
|
|
|
buffer.append(torch.tensor([[2.0]], device=device))
|
|
result = buffer.compute()
|
|
assert torch.allclose(result, torch.tensor([[1.0]], device=device))
|
|
|
|
buffer.append(torch.tensor([[3.0]], device=device))
|
|
result = buffer.compute()
|
|
assert torch.allclose(result, torch.tensor([[1.0]], device=device))
|
|
|
|
buffer.append(torch.tensor([[4.0]], device=device))
|
|
result = buffer.compute()
|
|
assert torch.allclose(result, torch.tensor([[2.0]], device=device))
|
|
|
|
|
|
def test_value_matches_lag_per_env(device):
|
|
"""Delayed values correctly match current lags."""
|
|
B = 3
|
|
buf = DelayBuffer(
|
|
0, 2, batch_size=B, per_env=True, device=device, generator=make_gen(1234, device)
|
|
)
|
|
for t in range(6):
|
|
buf.append(torch.full((B, 1), float(t), device=device))
|
|
y = buf.compute()
|
|
lags = buf.current_lags
|
|
for e in range(B):
|
|
eff = int(min(lags[e].item(), t))
|
|
assert y[e].item() == float(t - eff)
|
|
|
|
|
|
##
|
|
# Lag sampling modes.
|
|
##
|
|
|
|
|
|
def test_delay_buffer_per_env_lags(device):
|
|
"""Per-env mode allows different lags per environment."""
|
|
buffer = DelayBuffer(min_lag=0, max_lag=3, batch_size=4, per_env=True, device=device)
|
|
|
|
for i in range(10):
|
|
buffer.append(torch.full((4, 1), float(i), device=device))
|
|
buffer.compute()
|
|
|
|
lags = buffer.current_lags
|
|
assert torch.all(lags >= 0)
|
|
assert torch.all(lags <= 3)
|
|
|
|
|
|
def test_delay_buffer_shared_lags(device):
|
|
"""Shared mode uses same lag across all environments."""
|
|
buffer = DelayBuffer(min_lag=0, max_lag=3, batch_size=4, per_env=False, device=device)
|
|
|
|
for i in range(10):
|
|
buffer.append(torch.full((4, 1), float(i), device=device))
|
|
buffer.compute()
|
|
|
|
lags = buffer.current_lags
|
|
assert torch.all(lags == lags[0])
|
|
|
|
|
|
##
|
|
# Hold probability.
|
|
##
|
|
|
|
|
|
def test_hold_prob_always_hold(device):
|
|
"""hold_prob=1.0 keeps lag frozen."""
|
|
buf = DelayBuffer(
|
|
0,
|
|
3,
|
|
batch_size=1,
|
|
device=device,
|
|
update_period=1,
|
|
hold_prob=1.0,
|
|
generator=make_gen(7, device),
|
|
)
|
|
for t in range(2):
|
|
buf.append(torch.tensor([[float(t)]], device=device))
|
|
buf.compute()
|
|
first = buf.current_lags.item()
|
|
for t in range(2, 20):
|
|
buf.append(torch.tensor([[float(t)]], device=device))
|
|
buf.compute()
|
|
assert buf.current_lags.item() == first
|
|
|
|
|
|
def test_hold_prob_never_hold(device):
|
|
"""hold_prob=0.0 allows lag to change."""
|
|
buf = DelayBuffer(
|
|
0,
|
|
3,
|
|
batch_size=1,
|
|
device=device,
|
|
update_period=1,
|
|
hold_prob=0.0,
|
|
generator=make_gen(42, device),
|
|
)
|
|
for t in range(2):
|
|
buf.append(torch.tensor([[float(t)]], device=device))
|
|
buf.compute()
|
|
prev, changed = buf.current_lags.item(), False
|
|
for t in range(2, 20):
|
|
buf.append(torch.tensor([[float(t)]], device=device))
|
|
buf.compute()
|
|
cur = buf.current_lags.item()
|
|
if cur != prev:
|
|
changed = True
|
|
break
|
|
prev = cur
|
|
assert changed
|
|
|
|
|
|
##
|
|
# Update period.
|
|
##
|
|
|
|
|
|
def test_delay_buffer_update_period(device):
|
|
"""Update period controls how often lags are resampled."""
|
|
buffer = DelayBuffer(
|
|
min_lag=0,
|
|
max_lag=3,
|
|
batch_size=1,
|
|
update_period=3,
|
|
per_env_phase=False,
|
|
device=device,
|
|
)
|
|
|
|
lag_history = []
|
|
for i in range(12):
|
|
buffer.append(torch.tensor([[float(i)]], device=device))
|
|
buffer.compute()
|
|
lag_history.append(buffer.current_lags[0].item())
|
|
|
|
assert len(lag_history) == 12
|
|
|
|
|
|
def test_update_period_changes_only_on_schedule(device):
|
|
"""Lags update only at scheduled intervals."""
|
|
buf = DelayBuffer(
|
|
0,
|
|
10, # Wider range makes value changes more likely
|
|
batch_size=1,
|
|
update_period=3,
|
|
per_env_phase=False,
|
|
device=device,
|
|
generator=make_gen(123, device),
|
|
)
|
|
|
|
# Track when lags are updated by checking step_count timing
|
|
lag_values = []
|
|
for t in range(12):
|
|
buf.append(torch.tensor([[float(t)]], device=device))
|
|
buf.compute()
|
|
lag_values.append(buf.current_lags.item())
|
|
|
|
# Verify lag stays constant between update periods
|
|
# Update happens at step 0, 3, 6, 9
|
|
# So: [0-2] same, [3-5] same, [6-8] same, [9-11] same
|
|
assert lag_values[0] == lag_values[1] == lag_values[2]
|
|
assert lag_values[3] == lag_values[4] == lag_values[5]
|
|
assert lag_values[6] == lag_values[7] == lag_values[8]
|
|
assert lag_values[9] == lag_values[10] == lag_values[11]
|
|
|
|
|
|
##
|
|
# Reset behavior.
|
|
##
|
|
|
|
|
|
def test_delay_buffer_reset_all(device):
|
|
"""Reset clears all environments."""
|
|
buffer = DelayBuffer(min_lag=1, max_lag=2, batch_size=2, device=device)
|
|
|
|
for i in range(5):
|
|
buffer.append(torch.full((2, 1), float(i), device=device))
|
|
buffer.compute()
|
|
|
|
buffer.reset()
|
|
|
|
assert torch.all(buffer.current_lags == 0)
|
|
assert torch.all(buffer._step_count == 0)
|
|
|
|
|
|
def test_delay_buffer_reset_partial(device):
|
|
"""Reset with batch_ids only resets specified environments."""
|
|
buffer = DelayBuffer(min_lag=0, max_lag=3, batch_size=4, device=device)
|
|
|
|
for _ in range(5):
|
|
buffer.append(torch.arange(4, device=device).unsqueeze(1).float())
|
|
buffer.compute()
|
|
|
|
lags_before = buffer.current_lags.clone()
|
|
|
|
buffer.reset(batch_ids=torch.tensor([0, 2], device=device))
|
|
|
|
lags_after = buffer.current_lags
|
|
|
|
assert lags_after[0] == 0
|
|
assert lags_after[2] == 0
|
|
assert lags_after[1] == lags_before[1]
|
|
assert lags_after[3] == lags_before[3]
|
|
|
|
|
|
def test_partial_reset_then_backfill(device):
|
|
"""Partial reset returns zeros until next append backfills."""
|
|
B = 3
|
|
buf = DelayBuffer(1, 2, batch_size=B, device=device, generator=make_gen(9, device))
|
|
for t in range(3):
|
|
x = torch.arange(1, B + 1, device=device).float().unsqueeze(1) * 10 + t
|
|
buf.append(x)
|
|
buf.compute()
|
|
buf.reset(batch_ids=torch.tensor([1], device=device))
|
|
y = buf.compute()
|
|
assert torch.allclose(y[1], torch.zeros_like(y[1]))
|
|
x_new = torch.tensor([[111.0], [999.0], [333.0]], device=device)
|
|
buf.append(x_new)
|
|
y2 = buf.compute()
|
|
assert torch.allclose(y2[1], torch.tensor([[999.0]], device=device))
|
|
|
|
|
|
##
|
|
# Error handling.
|
|
##
|
|
|
|
|
|
def test_delay_buffer_not_initialized(device):
|
|
"""Compute before append raises error."""
|
|
buffer = DelayBuffer(min_lag=0, max_lag=3, batch_size=1, device=device)
|
|
|
|
with pytest.raises(RuntimeError, match="Buffer not initialized"):
|
|
buffer.compute()
|
|
|
|
|
|
def test_append_wrong_batch_raises(device):
|
|
"""Appending wrong batch size raises error."""
|
|
buf = DelayBuffer(0, 1, batch_size=2, device=device)
|
|
with pytest.raises(ValueError):
|
|
buf.append(torch.zeros(3, 1, device=device))
|
|
|
|
|
|
def test_delay_buffer_validation(device):
|
|
"""Input validation catches invalid parameters."""
|
|
with pytest.raises(ValueError, match="min_lag must be >= 0"):
|
|
DelayBuffer(min_lag=-1, max_lag=3, batch_size=1, device=device)
|
|
|
|
with pytest.raises(ValueError, match="max_lag.*must be >= min_lag"):
|
|
DelayBuffer(min_lag=5, max_lag=3, batch_size=1, device=device)
|
|
|
|
with pytest.raises(ValueError, match="hold_prob must be in"):
|
|
DelayBuffer(min_lag=0, max_lag=3, batch_size=1, hold_prob=1.5, device=device)
|
|
|
|
with pytest.raises(ValueError, match="update_period must be >= 0"):
|
|
DelayBuffer(min_lag=0, max_lag=3, batch_size=1, update_period=-1, device=device)
|
|
|
|
|
|
def test_backfill_and_peek_leave_others_untouched(device):
|
|
"""backfill serves the reset row's new frame; peek advances no schedules."""
|
|
B = 3
|
|
buf = DelayBuffer(1, 2, batch_size=B, device=device, generator=make_gen(9, device))
|
|
for t in range(4):
|
|
x = torch.arange(1, B + 1, device=device).float().unsqueeze(1) * 10 + t
|
|
buf.append(x)
|
|
buf.compute()
|
|
|
|
out_before = buf.peek()
|
|
lags_before = buf.current_lags.clone()
|
|
steps_before = buf._step_count.clone()
|
|
|
|
reset_ids = torch.tensor([1], device=device)
|
|
buf.reset(batch_ids=reset_ids)
|
|
buf.backfill(torch.tensor([[-1.0], [999.0], [-1.0]], device=device), reset_ids)
|
|
out = buf.peek()
|
|
|
|
# Reset row serves its backfilled frame (lag zeroed, one valid frame).
|
|
assert torch.allclose(out[1], torch.tensor([999.0], device=device))
|
|
# Other rows: same delayed frame, lags and schedule untouched.
|
|
keep = torch.tensor([0, 2], device=device)
|
|
assert torch.allclose(out[keep], out_before[keep])
|
|
assert torch.equal(buf.current_lags[keep], lags_before[keep])
|
|
assert torch.equal(buf._step_count[keep], steps_before[keep])
|