wehub-resource-sync
85 lines
2.8 KiB
Python
85 lines
2.8 KiB
Python
#!/usr/bin/env python
|
|
|
|
import pytest
|
|
|
|
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
|
|
|
|
import torch # noqa: E402
|
|
|
|
from lerobot.utils.collate import lerobot_collate_fn # noqa: E402
|
|
|
|
|
|
def test_lerobot_collate_preserves_messages_and_drops_raw_language():
|
|
batch = [
|
|
{
|
|
"index": torch.tensor(0),
|
|
"messages": [{"role": "assistant", "content": "a"}],
|
|
"message_streams": ["low_level"],
|
|
"target_message_indices": [0],
|
|
"language_persistent": [{"content": "raw"}],
|
|
"language_events": [],
|
|
},
|
|
{
|
|
"index": torch.tensor(1),
|
|
"messages": [{"role": "assistant", "content": "b"}],
|
|
"message_streams": ["low_level"],
|
|
"target_message_indices": [0],
|
|
"language_persistent": [{"content": "raw"}],
|
|
"language_events": [],
|
|
},
|
|
]
|
|
|
|
out = lerobot_collate_fn(batch)
|
|
|
|
assert out["index"].tolist() == [0, 1]
|
|
assert out["messages"][0][0]["content"] == "a"
|
|
assert out["messages"][1][0]["content"] == "b"
|
|
assert out["message_streams"] == [["low_level"], ["low_level"]]
|
|
assert out["target_message_indices"] == [[0], [0]]
|
|
assert "language_persistent" not in out
|
|
assert "language_events" not in out
|
|
|
|
|
|
def test_lerobot_collate_passes_through_standard_batch():
|
|
"""On a non-language batch, the collate must match ``default_collate``.
|
|
|
|
Guards against silent regressions: ``lerobot_train.py`` only opts into
|
|
``lerobot_collate_fn`` when the dataset declares language columns, but
|
|
if a future change ever wires it in unconditionally we want the
|
|
behavior to remain a transparent pass-through for ordinary tensor
|
|
batches.
|
|
"""
|
|
from torch.utils.data._utils.collate import default_collate
|
|
|
|
batch = [
|
|
{
|
|
"observation.image": torch.zeros(3, 4, 4),
|
|
"action": torch.tensor([0.0, 1.0]),
|
|
"index": torch.tensor(0),
|
|
},
|
|
{
|
|
"observation.image": torch.ones(3, 4, 4),
|
|
"action": torch.tensor([2.0, 3.0]),
|
|
"index": torch.tensor(1),
|
|
},
|
|
]
|
|
|
|
custom = lerobot_collate_fn(batch)
|
|
expected = default_collate(batch)
|
|
|
|
assert custom.keys() == expected.keys()
|
|
for key in expected:
|
|
assert torch.equal(custom[key], expected[key]), f"key={key} diverged"
|
|
|
|
|
|
def test_lerobot_collate_drops_none_samples():
|
|
"""Recipes that yielded no target message return ``None`` — those samples
|
|
must be filtered out, and an entirely-``None`` batch must collapse to ``None``.
|
|
"""
|
|
batch = [None, {"index": torch.tensor(0)}, None]
|
|
out = lerobot_collate_fn(batch)
|
|
assert out is not None
|
|
assert out["index"].tolist() == [0]
|
|
|
|
assert lerobot_collate_fn([None, None]) is None
|