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
269 lines
10 KiB
Python
269 lines
10 KiB
Python
"""Tests for CircularBuffer."""
|
|
|
|
import pytest
|
|
import torch
|
|
from conftest import get_test_device
|
|
|
|
from mjlab.utils.buffers import CircularBuffer
|
|
|
|
|
|
@pytest.fixture
|
|
def device():
|
|
"""Test device fixture."""
|
|
return get_test_device()
|
|
|
|
|
|
def test_circular_buffer_basic_append(device):
|
|
"""Basic append and chronological retrieval (oldest -> newest)."""
|
|
buffer = CircularBuffer(max_len=3, batch_size=2, device=device)
|
|
|
|
buffer.append(torch.tensor([[1.0, 2.0], [3.0, 4.0]], device=device))
|
|
buffer.append(torch.tensor([[5.0, 6.0], [7.0, 8.0]], device=device))
|
|
buffer.append(torch.tensor([[9.0, 10.0], [11.0, 12.0]], device=device))
|
|
|
|
result = buffer.buffer
|
|
assert result.shape == (2, 3, 2)
|
|
# Oldest to newest.
|
|
assert torch.allclose(
|
|
result[0], torch.tensor([[1.0, 2.0], [5.0, 6.0], [9.0, 10.0]], device=device)
|
|
)
|
|
assert torch.allclose(
|
|
result[1], torch.tensor([[3.0, 4.0], [7.0, 8.0], [11.0, 12.0]], device=device)
|
|
)
|
|
|
|
|
|
def test_circular_buffer_overwrite(device):
|
|
"""Overwrites oldest once capacity reached."""
|
|
buffer = CircularBuffer(max_len=2, batch_size=1, device=device)
|
|
|
|
buffer.append(torch.tensor([[1.0]], device=device))
|
|
buffer.append(torch.tensor([[2.0]], device=device))
|
|
buffer.append(torch.tensor([[3.0]], device=device)) # Overwrites first.
|
|
|
|
result = buffer.buffer
|
|
assert result.shape == (1, 2, 1)
|
|
assert torch.allclose(result[0], torch.tensor([[2.0], [3.0]], device=device))
|
|
|
|
|
|
def test_circular_buffer_reset_single_batch(device):
|
|
"""Reset clears values and counters for specified batch rows."""
|
|
buffer = CircularBuffer(max_len=2, batch_size=3, device=device)
|
|
|
|
buffer.append(torch.tensor([[1.0], [2.0], [3.0]], device=device))
|
|
buffer.append(torch.tensor([[4.0], [5.0], [6.0]], device=device))
|
|
|
|
# Reset only batch index 1.
|
|
buffer.reset(batch_ids=torch.tensor([1], device=device))
|
|
|
|
result = buffer.buffer
|
|
# Oldest-to-newest for each batch row.
|
|
assert result[0, 0, 0] == 1.0
|
|
assert result[1, 0, 0] == 0.0 # Reset backfilled to zeros for that row.
|
|
assert result[2, 0, 0] == 3.0
|
|
|
|
# current_length reflects reset: rows 0 and 2 had 2 pushes, row 1 is 0.
|
|
cl = buffer.current_length
|
|
assert torch.equal(cl.cpu(), torch.tensor([2, 0, 2]))
|
|
|
|
|
|
def test_circular_buffer_first_append_fills(device):
|
|
"""First append back-fills whole history for each batch row."""
|
|
buffer = CircularBuffer(max_len=3, batch_size=2, device=device)
|
|
buffer.append(torch.tensor([[1.0], [2.0]], device=device))
|
|
|
|
result = buffer.buffer
|
|
assert torch.allclose(result[0], torch.tensor([[1.0], [1.0], [1.0]], device=device))
|
|
assert torch.allclose(result[1], torch.tensor([[2.0], [2.0], [2.0]], device=device))
|
|
|
|
# And current_length reflects valid frames so far.
|
|
cl = buffer.current_length
|
|
assert torch.equal(cl.cpu(), torch.tensor([1, 1]))
|
|
|
|
|
|
def test_current_length_counts_and_clamps(device):
|
|
"""current_length counts per-batch valid frames and clamps to max_len."""
|
|
buffer = CircularBuffer(max_len=4, batch_size=3, device=device)
|
|
|
|
# Two appends -> length 2 everywhere.
|
|
for _ in range(2):
|
|
buffer.append(torch.arange(3, dtype=torch.float32, device=device).unsqueeze(-1))
|
|
|
|
assert torch.equal(buffer.current_length.cpu(), torch.tensor([2, 2, 2]))
|
|
|
|
# Reset middle row -> it becomes 0.
|
|
buffer.reset(batch_ids=[1])
|
|
assert torch.equal(buffer.current_length.cpu(), torch.tensor([2, 0, 2]))
|
|
|
|
# One more append -> rows [0,2] become 3; row 1 becomes 1.
|
|
buffer.append(torch.arange(3, dtype=torch.float32, device=device).unsqueeze(-1))
|
|
assert torch.equal(buffer.current_length.cpu(), torch.tensor([3, 1, 3]))
|
|
|
|
# Fill beyond capacity -> clamp to max_len.
|
|
for _ in range(5):
|
|
buffer.append(torch.arange(3, dtype=torch.float32, device=device).unsqueeze(-1))
|
|
assert torch.equal(buffer.current_length.cpu(), torch.tensor([4, 4, 4]))
|
|
|
|
|
|
def test_reset_all_none_path(device):
|
|
"""reset(None) zeros the entire buffer and counters without indexing."""
|
|
buffer = CircularBuffer(max_len=3, batch_size=2, device=device)
|
|
buffer.append(torch.tensor([[1.0], [2.0]], device=device))
|
|
buffer.append(torch.tensor([[3.0], [4.0]], device=device))
|
|
|
|
buffer.reset() # None -> reset all.
|
|
|
|
# Counters are zero.
|
|
assert torch.equal(buffer.current_length.cpu(), torch.tensor([0, 0]))
|
|
# Buffer zeros (safe to read even after reset because storage exists).
|
|
result = buffer.buffer
|
|
assert torch.count_nonzero(result) == 0
|
|
|
|
|
|
def test_getitem_lifo_and_clamp(device):
|
|
"""__getitem__ returns lagged frames per-batch (LIFO), clamping when needed."""
|
|
buffer = CircularBuffer(max_len=3, batch_size=2, device=device)
|
|
|
|
buffer.append(torch.tensor([[1.0], [10.0]], device=device)) # t0
|
|
buffer.append(torch.tensor([[2.0], [20.0]], device=device)) # t1
|
|
buffer.append(torch.tensor([[3.0], [30.0]], device=device)) # t2
|
|
|
|
# Lag 0 for batch 0 (-> 3), lag 2 for batch 1 (-> oldest 10).
|
|
out = buffer[torch.tensor([0, 2], device=device)]
|
|
assert torch.allclose(out, torch.tensor([[3.0], [10.0]], device=device))
|
|
|
|
# Clamp: huge lag for batch 0 -> oldest (1), lag 1 for batch 1 -> 20.
|
|
out = buffer[torch.tensor([99, 1], device=device)]
|
|
assert torch.allclose(out, torch.tensor([[1.0], [20.0]], device=device))
|
|
|
|
|
|
def test_backfill_after_per_batch_reset(device):
|
|
"""After resetting a row, the next append back-fills its entire history for that row."""
|
|
buffer = CircularBuffer(max_len=3, batch_size=2, device=device)
|
|
|
|
buffer.append(torch.tensor([[1.0], [10.0]], device=device)) # t0
|
|
buffer.append(torch.tensor([[2.0], [20.0]], device=device)) # t1
|
|
|
|
# Reset only batch row 1; row 0 remains with 2 valid frames.
|
|
buffer.reset(batch_ids=[1])
|
|
assert torch.equal(buffer.current_length.cpu(), torch.tensor([2, 0]))
|
|
|
|
# Next append: row 0 gets new value; row 1 is "first push" -> back-filled.
|
|
buffer.append(torch.tensor([[3.0], [99.0]], device=device)) # t2
|
|
|
|
hist = buffer.buffer # shape (2, 3, 1)
|
|
# Row 0 keeps real chronology [1, 2, 3].
|
|
assert torch.allclose(
|
|
hist[0].squeeze(-1), torch.tensor([1.0, 2.0, 3.0], device=device)
|
|
)
|
|
# Row 1 is back-filled to all 99s.
|
|
assert torch.allclose(
|
|
hist[1].squeeze(-1), torch.tensor([99.0, 99.0, 99.0], device=device)
|
|
)
|
|
|
|
|
|
def test_backfill_multidimensional_data(device):
|
|
"""First-push backfill works for >2D data whose trailing dims differ from batch.
|
|
|
|
Regression for #1045: the backfill broadcast hard-coded a 3D condition shape,
|
|
which misaligned the batch axis with a trailing axis when num_targets != batch.
|
|
Here batch_size=4 but the trailing axis is 2, so the buggy broadcast raised a
|
|
4-vs-2 mismatch.
|
|
"""
|
|
batch_size, num_targets, num_channels = 4, 2, 3
|
|
buffer = CircularBuffer(max_len=3, batch_size=batch_size, device=device)
|
|
|
|
data = torch.arange(
|
|
batch_size * num_targets * num_channels, dtype=torch.float32, device=device
|
|
).reshape(batch_size, num_targets, num_channels)
|
|
buffer.append(data)
|
|
|
|
hist = buffer.buffer # (batch, max_len, num_targets, num_channels)
|
|
assert hist.shape == (batch_size, 3, num_targets, num_channels)
|
|
# First push backfills every history slot with the same frame per batch row.
|
|
for t in range(3):
|
|
assert torch.allclose(hist[:, t], data)
|
|
|
|
|
|
def test_errors_and_types(device):
|
|
"""Error paths: wrong batch, pre-append access, and bad key size."""
|
|
buffer = CircularBuffer(max_len=2, batch_size=2, device=device)
|
|
|
|
# Wrong batch size on append.
|
|
with pytest.raises(ValueError):
|
|
buffer.append(torch.tensor([[1.0]], device=device)) # batch_size=1 wrong
|
|
|
|
# buffer property before first append.
|
|
with pytest.raises(RuntimeError):
|
|
_ = CircularBuffer(max_len=1, batch_size=1, device=device).buffer
|
|
|
|
# __getitem__ before any valid pushes.
|
|
with pytest.raises(RuntimeError):
|
|
_ = buffer[torch.tensor([0, 0], device=device)]
|
|
|
|
# Now append once so storage exists and counters > 0.
|
|
buffer.append(torch.tensor([[1.0], [2.0]], device=device))
|
|
|
|
# __getitem__ with wrong key length.
|
|
with pytest.raises(ValueError):
|
|
_ = buffer[torch.tensor([0], device=device)]
|
|
|
|
|
|
def test_dtype_and_device_preserved(device):
|
|
"""Buffer preserves dtype and device."""
|
|
buffer = CircularBuffer(max_len=2, batch_size=2, device=device)
|
|
x = torch.tensor([[1.0], [2.0]], dtype=torch.float32, device=device)
|
|
buffer.append(x)
|
|
assert buffer.buffer.dtype == torch.float32
|
|
assert buffer.buffer.device.type == torch.device(device).type
|
|
|
|
|
|
def test_backfill_only_touches_given_rows(device):
|
|
"""backfill writes one frame into all slots of given rows, no pointer advance."""
|
|
buffer = CircularBuffer(max_len=3, batch_size=3, device=device)
|
|
buffer.append(torch.tensor([[1.0], [10.0], [100.0]], device=device))
|
|
buffer.append(torch.tensor([[2.0], [20.0], [200.0]], device=device))
|
|
buffer.append(torch.tensor([[3.0], [30.0], [300.0]], device=device))
|
|
|
|
buffer.reset(batch_ids=torch.tensor([1], device=device))
|
|
data = torch.tensor([[-1.0], [99.0], [-1.0]], device=device)
|
|
buffer.backfill(data, torch.tensor([1], device=device))
|
|
|
|
result = buffer.buffer
|
|
assert torch.allclose(
|
|
result[0].flatten(), torch.tensor([1.0, 2.0, 3.0], device=device)
|
|
)
|
|
assert torch.allclose(
|
|
result[1].flatten(), torch.tensor([99.0, 99.0, 99.0], device=device)
|
|
)
|
|
assert torch.allclose(
|
|
result[2].flatten(), torch.tensor([100.0, 200.0, 300.0], device=device)
|
|
)
|
|
assert buffer.current_length.tolist() == [3, 1, 3]
|
|
|
|
# Subsequent appends treat the backfilled row normally.
|
|
buffer.append(torch.tensor([[4.0], [40.0], [400.0]], device=device))
|
|
result = buffer.buffer
|
|
assert torch.allclose(
|
|
result[1].flatten(), torch.tensor([99.0, 99.0, 40.0], device=device)
|
|
)
|
|
assert buffer.current_length.tolist() == [3, 2, 3]
|
|
|
|
|
|
def test_backfill_uninitialized_raises(device):
|
|
buffer = CircularBuffer(max_len=2, batch_size=2, device=device)
|
|
with pytest.raises(RuntimeError, match="not initialized"):
|
|
buffer.backfill(torch.zeros(2, 1, device=device), torch.tensor([0], device=device))
|
|
|
|
|
|
def test_getitem_lag_clamps_to_oldest_after_wrap(device):
|
|
buffer = CircularBuffer(max_len=2, batch_size=1, device=device)
|
|
for v in [1.0, 2.0, 3.0]:
|
|
buffer.append(torch.tensor([[v]], device=device))
|
|
|
|
# Retained frames: [2, 3] (num_pushes exceeds max_len).
|
|
assert buffer[torch.tensor([0], device=device)].flatten().item() == 3.0
|
|
assert buffer[torch.tensor([1], device=device)].flatten().item() == 2.0
|
|
# A lag beyond the retained history clamps to the oldest frame instead of
|
|
# wrapping around to a newer one.
|
|
assert buffer[torch.tensor([2], device=device)].flatten().item() == 2.0
|