mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-23 01:41:54 +00:00
Merge branch 'feat/language-columns' into feat/language-annotation-pipeline
Resolve conflicts and pull in the latest PR 1 fixes. Conflicts: - pyproject.toml: PR 1 added `lerobot-rollout` and PR 2 added `lerobot-annotate` to the same `[project.scripts]` block. Kept both. - uv.lock: dropped both sides and regenerated against the merged `pyproject.toml` (PR 2 dropped the `datatrove` dep when distribution moved to HF Jobs; PR 1's lock didn't have it). Test follow-up: - `tests/annotations/test_pipeline_recipe_render.py` — PR 1 deleted `src/lerobot/configs/recipes/pi05_hirobot.yaml` (review feedback: remove the canonical-recipe file; recipes are user-supplied). The cross-PR contract this test guards is "the recipe DSL renders non-empty messages from pipeline output", which doesn't depend on any specific YAML, so the test now builds an inline blend recipe with the same coverage. Passes. Sweep: 82 passed, 2 failed (pre-existing module-impl bugs: `test_module1_attaches_video_block_to_subtask_prompt`, `test_module2_mid_episode_emits_paired_interjection_and_speech`). The PR 1 carryover (`test_emitted_at_raises_on_ambiguous_per_camera_vqa`) is now passing — the merge brought in PR 1's tightened `_select_one` ambiguity check. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Pepijn
@@ -35,14 +35,60 @@ from lerobot.annotations.steerable_pipeline.modules import (
|
||||
)
|
||||
from lerobot.annotations.steerable_pipeline.validator import StagingValidator
|
||||
from lerobot.annotations.steerable_pipeline.writer import LanguageColumnsWriter
|
||||
from lerobot.configs.recipe import TrainingRecipe
|
||||
from lerobot.configs.recipe import MessageTurn, TrainingRecipe
|
||||
from lerobot.datasets.language_render import render_sample
|
||||
|
||||
from ._helpers import make_canned_responder
|
||||
|
||||
_RECIPE_PATH = (
|
||||
Path(__file__).resolve().parents[2] / "src" / "lerobot" / "configs" / "recipes" / "pi05_hirobot.yaml"
|
||||
)
|
||||
|
||||
def _build_pr1_style_blend_recipe() -> TrainingRecipe:
|
||||
"""Inline blend recipe that consumes every style this pipeline produces.
|
||||
|
||||
PR 1 used to ship ``src/lerobot/configs/recipes/pi05_hirobot.yaml`` as
|
||||
a canonical example, but that file was dropped during PR 1 review. The
|
||||
cross-PR contract this test guards is "the recipe DSL can render
|
||||
non-empty messages from pipeline output", which doesn't require a
|
||||
specific YAML — so we build the equivalent blend in code.
|
||||
"""
|
||||
return TrainingRecipe(
|
||||
blend={
|
||||
"low_level_execution": TrainingRecipe(
|
||||
weight=0.35,
|
||||
messages=[
|
||||
MessageTurn(
|
||||
role="user",
|
||||
content="${task}\nPlan: ${plan}\nMemory: ${memory}",
|
||||
stream="high_level",
|
||||
),
|
||||
MessageTurn(role="assistant", content="${subtask}", stream="low_level", target=True),
|
||||
],
|
||||
),
|
||||
"user_interjection_response": TrainingRecipe(
|
||||
weight=0.16,
|
||||
bindings={
|
||||
"speech": "emitted_at(t, role=assistant, tool_name=say)",
|
||||
"interjection": "emitted_at(t, style=interjection)",
|
||||
},
|
||||
messages=[
|
||||
MessageTurn(role="user", content="${task}", stream="high_level"),
|
||||
MessageTurn(
|
||||
role="user",
|
||||
content="${interjection}",
|
||||
stream="high_level",
|
||||
if_present="interjection",
|
||||
),
|
||||
MessageTurn(
|
||||
role="assistant",
|
||||
content="${plan}",
|
||||
stream="high_level",
|
||||
target=True,
|
||||
if_present="plan",
|
||||
tool_calls_from="speech",
|
||||
),
|
||||
],
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _build_executor() -> Executor:
|
||||
@@ -96,13 +142,7 @@ def test_pr1_canonical_recipe_renders_nonempty_from_pipeline_output(
|
||||
events_lists = table.column("language_events").to_pylist()
|
||||
timestamps = table.column("timestamp").to_pylist()
|
||||
|
||||
recipe = TrainingRecipe.from_yaml(_RECIPE_PATH) if hasattr(TrainingRecipe, "from_yaml") else None
|
||||
if recipe is None:
|
||||
# PR 1 may not expose from_yaml; load via PyYAML and TrainingRecipe(**...)
|
||||
import yaml
|
||||
|
||||
loaded = yaml.safe_load(_RECIPE_PATH.read_text(encoding="utf-8"))
|
||||
recipe = TrainingRecipe(**loaded)
|
||||
recipe = _build_pr1_style_blend_recipe()
|
||||
|
||||
rendered_any = False
|
||||
for ts, persistent, events in zip(timestamps, persistent_lists, events_lists, strict=True):
|
||||
|
||||
+1
-1
@@ -1,3 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:c2b8f8532c7a0b776de5e536b8b54e30b1a0c2e3d5cc25a2d86fe43e40ae5e8c
|
||||
oid sha256:8a31653c11eccdd4d80fd3f6a351cd54c49b8a48db1f7e9faf38fddd7900a09f
|
||||
size 515400
|
||||
|
||||
+1
-1
@@ -1,3 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:224b5fa4828aa88171b68c036e8919c1eae563e2113f03b6461eadf5bf8525a6
|
||||
oid sha256:75bf051698b37dcd7517ec8025a896ab5a0551a6dde5f89d0a3d5d50966e83e6
|
||||
size 31672
|
||||
|
||||
+1
-1
@@ -1,3 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:016d2fa8fe5f58017dfd46f4632fdc19dfd751e32a2c7cde2077c6f95546d6bd
|
||||
oid sha256:88e10930a10041d50f2cf369e6813ac14618d13dad1c21bdde1ac7798611c6ba
|
||||
size 68
|
||||
|
||||
+1
-1
@@ -1,3 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:eca0d87a699620e4fec7e68539b0be91e4cc933f6bf12032da52c182ab6f38cf
|
||||
oid sha256:89833a5ccdb7d85c83f717ff8ec68b8e822005cb8803899acaae88c578e2e3ae
|
||||
size 31672
|
||||
|
||||
@@ -1,3 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:54aecbc1af72a4cd5e9261492f5e7601890517516257aacdf2a0ffb3ce281f1b
|
||||
oid sha256:51effd76b73e972f10d31f5084ab906386134b600c87b2668767d30232a902bd
|
||||
size 992
|
||||
|
||||
@@ -1,3 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:88a9c3775a2aa1e90a08850521970070a4fcf0f6b82aab43cd8ccc5cf77e0013
|
||||
size 47424
|
||||
oid sha256:d4d7a16ca67f9adefac0e0620a7b2e9c822f2db42faaaced7a89fbad60e5ead4
|
||||
size 47680
|
||||
|
||||
@@ -1,3 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:91a2635e05a75fe187a5081504c5f35ce3417378813fa2deaf9ca4e8200e1819
|
||||
oid sha256:796c439ee8a64bf9901ff8325e7419bda8bd316360ee95e6304e8e1ae0f4c36c
|
||||
size 68
|
||||
|
||||
@@ -1,3 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:645bff922ac7bea63ad018ebf77c303c0e4cd2c1c0dc5ef3192865281bef3dc6
|
||||
size 47424
|
||||
oid sha256:ad33a8b47c39c2e1374567ff9da43cdb95e2dbe904c1b02a35051346d3043095
|
||||
size 47680
|
||||
|
||||
+149
-13
@@ -1,10 +1,22 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
from pathlib import Path
|
||||
from textwrap import dedent
|
||||
|
||||
import pytest
|
||||
|
||||
from lerobot.configs.recipe import MessageTurn, TrainingRecipe
|
||||
from lerobot.configs.recipe import MessageTurn, TrainingRecipe, load_recipe
|
||||
|
||||
|
||||
def _minimal_message_turn(content: str = "${task}") -> MessageTurn:
|
||||
return MessageTurn(role="user", content=content, stream="high_level")
|
||||
|
||||
|
||||
def _minimal_target_turn() -> MessageTurn:
|
||||
return MessageTurn(role="assistant", content="ok", stream="high_level", target=True)
|
||||
|
||||
|
||||
# ── Message-recipe validation ────────────────────────────────────────
|
||||
|
||||
|
||||
def test_message_recipe_validates_unknown_binding():
|
||||
@@ -12,21 +24,145 @@ def test_message_recipe_validates_unknown_binding():
|
||||
TrainingRecipe(
|
||||
messages=[
|
||||
MessageTurn(role="user", content="${missing}", stream="high_level"),
|
||||
MessageTurn(role="assistant", content="ok", stream="high_level", target=True),
|
||||
_minimal_target_turn(),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def test_canonical_recipe_loads():
|
||||
recipe = TrainingRecipe.from_yaml(Path("src/lerobot/configs/recipes/pi05_hirobot.yaml"))
|
||||
def test_message_turn_requires_a_stream():
|
||||
"""Every turn must declare a stream — None is rejected at construction.
|
||||
|
||||
Previously this only failed at render time (``_validate_rendered``);
|
||||
catching it here means a malformed recipe YAML errors at load instead
|
||||
of at the first training sample.
|
||||
"""
|
||||
with pytest.raises(ValueError, match="missing a stream"):
|
||||
MessageTurn(role="user", content="${task}")
|
||||
|
||||
|
||||
def test_message_recipe_requires_at_least_one_target():
|
||||
with pytest.raises(ValueError, match="target"):
|
||||
TrainingRecipe(
|
||||
messages=[
|
||||
_minimal_message_turn(),
|
||||
MessageTurn(role="assistant", content="no target", stream="high_level"),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def test_recipe_rejects_both_messages_and_blend():
|
||||
with pytest.raises(ValueError, match="only one"):
|
||||
TrainingRecipe(
|
||||
messages=[_minimal_message_turn(), _minimal_target_turn()],
|
||||
blend={"a": TrainingRecipe(weight=1.0, messages=[_minimal_target_turn()])},
|
||||
)
|
||||
|
||||
|
||||
def test_recipe_rejects_neither_messages_nor_blend():
|
||||
with pytest.raises(ValueError, match="must set one"):
|
||||
TrainingRecipe()
|
||||
|
||||
|
||||
# ── Blend validation ─────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_blend_must_be_non_empty():
|
||||
with pytest.raises(ValueError, match="at least one component"):
|
||||
TrainingRecipe(blend={})
|
||||
|
||||
|
||||
def test_blend_component_must_define_weight():
|
||||
with pytest.raises(ValueError, match="weight"):
|
||||
TrainingRecipe(blend={"a": TrainingRecipe(messages=[_minimal_target_turn()])})
|
||||
|
||||
|
||||
def test_blend_component_weight_must_be_positive():
|
||||
with pytest.raises(ValueError, match="positive weight"):
|
||||
TrainingRecipe(blend={"a": TrainingRecipe(weight=0.0, messages=[_minimal_target_turn()])})
|
||||
|
||||
|
||||
def test_blend_component_must_define_messages():
|
||||
# A bare TrainingRecipe(weight=1.0) would itself raise; build it without
|
||||
# going through __post_init__ to exercise the blend-level validator.
|
||||
bad = TrainingRecipe.__new__(TrainingRecipe)
|
||||
bad.messages = None
|
||||
bad.bindings = None
|
||||
bad.blend = None
|
||||
bad.weight = 1.0
|
||||
with pytest.raises(ValueError, match="must define messages"):
|
||||
TrainingRecipe(blend={"a": bad})
|
||||
|
||||
|
||||
def test_blend_components_cannot_themselves_define_a_blend():
|
||||
inner = TrainingRecipe(blend={"x": TrainingRecipe(weight=1.0, messages=[_minimal_target_turn()])})
|
||||
# Force-bypass the inner component's normal validation so the test
|
||||
# exercises the outer blend's "no nested blends" rule directly.
|
||||
nested = TrainingRecipe.__new__(TrainingRecipe)
|
||||
nested.messages = None
|
||||
nested.bindings = None
|
||||
nested.blend = inner.blend
|
||||
nested.weight = 1.0
|
||||
with pytest.raises(ValueError, match="cannot itself define a blend"):
|
||||
TrainingRecipe(blend={"outer": nested})
|
||||
|
||||
|
||||
# ── from_dict / from_yaml round-trips ────────────────────────────────
|
||||
|
||||
|
||||
def test_from_dict_with_nested_blend():
|
||||
recipe = TrainingRecipe.from_dict(
|
||||
{
|
||||
"blend": {
|
||||
"a": {
|
||||
"weight": 1.0,
|
||||
"messages": [
|
||||
{"role": "user", "content": "${task}", "stream": "high_level"},
|
||||
{"role": "assistant", "content": "a", "stream": "high_level", "target": True},
|
||||
],
|
||||
},
|
||||
"b": {
|
||||
"weight": 2.0,
|
||||
"messages": [
|
||||
{"role": "user", "content": "${task}", "stream": "high_level"},
|
||||
{"role": "assistant", "content": "b", "stream": "high_level", "target": True},
|
||||
],
|
||||
},
|
||||
}
|
||||
}
|
||||
)
|
||||
assert recipe.blend is not None
|
||||
assert set(recipe.blend) == {
|
||||
"memory_update",
|
||||
"user_interjection_response",
|
||||
"high_level_subtask",
|
||||
"low_level_execution",
|
||||
"ask_vqa_top",
|
||||
"ask_vqa_wrist",
|
||||
}
|
||||
assert sum(component.weight for component in recipe.blend.values()) == pytest.approx(0.96)
|
||||
assert set(recipe.blend) == {"a", "b"}
|
||||
assert recipe.blend["b"].weight == 2.0
|
||||
# Inner messages were promoted to MessageTurn instances.
|
||||
assert isinstance(recipe.blend["a"].messages[0], MessageTurn)
|
||||
|
||||
|
||||
def test_from_yaml_round_trips_through_load_recipe(tmp_path: Path):
|
||||
yaml_text = dedent(
|
||||
"""
|
||||
bindings:
|
||||
custom: "active_at(t, style=subtask)"
|
||||
messages:
|
||||
- {role: user, content: "${task}: ${custom}", stream: high_level}
|
||||
- {role: assistant, content: "ok", stream: high_level, target: true}
|
||||
"""
|
||||
).strip()
|
||||
path = tmp_path / "recipe.yaml"
|
||||
path.write_text(yaml_text)
|
||||
|
||||
via_classmethod = TrainingRecipe.from_yaml(path)
|
||||
via_helper = load_recipe(path)
|
||||
|
||||
assert via_classmethod.bindings == {"custom": "active_at(t, style=subtask)"}
|
||||
assert via_classmethod.messages[1].target is True
|
||||
# ``load_recipe`` is just a wrapper, but assert the two paths agree
|
||||
# on the structural result so a future divergence is caught here.
|
||||
assert via_helper.bindings == via_classmethod.bindings
|
||||
assert len(via_helper.messages) == len(via_classmethod.messages)
|
||||
|
||||
|
||||
def test_from_yaml_rejects_non_mapping(tmp_path: Path):
|
||||
path = tmp_path / "bad.yaml"
|
||||
path.write_text("- just\n- a\n- list\n")
|
||||
with pytest.raises(ValueError, match="mapping at the top level"):
|
||||
TrainingRecipe.from_yaml(path)
|
||||
|
||||
@@ -113,7 +113,7 @@ def assert_metadata_consistency(aggr_ds, ds_0, ds_1):
|
||||
"""Test that metadata is correctly aggregated."""
|
||||
# Test basic info
|
||||
assert aggr_ds.fps == ds_0.fps == ds_1.fps, "FPS should be the same across all datasets"
|
||||
assert aggr_ds.meta.info["robot_type"] == ds_0.meta.info["robot_type"] == ds_1.meta.info["robot_type"], (
|
||||
assert aggr_ds.meta.info.robot_type == ds_0.meta.info.robot_type == ds_1.meta.info.robot_type, (
|
||||
"Robot type should be the same"
|
||||
)
|
||||
|
||||
@@ -153,8 +153,8 @@ def assert_video_frames_integrity(aggr_ds, ds_0, ds_1):
|
||||
|
||||
video_keys = list(
|
||||
filter(
|
||||
lambda key: aggr_ds.meta.info["features"][key]["dtype"] == "video",
|
||||
aggr_ds.meta.info["features"].keys(),
|
||||
lambda key: aggr_ds.meta.info.features[key]["dtype"] == "video",
|
||||
aggr_ds.meta.info.features.keys(),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@@ -161,7 +161,7 @@ def test_init_loads_existing_metadata(tmp_path, lerobot_dataset_metadata_factory
|
||||
|
||||
assert meta.total_episodes == 3
|
||||
assert meta.total_frames == 150
|
||||
assert meta.fps == info["fps"]
|
||||
assert meta.fps == info.fps
|
||||
|
||||
|
||||
# ── Property accessors ───────────────────────────────────────────────
|
||||
@@ -385,3 +385,84 @@ def test_finalize_flushes_buffered_metadata(tmp_path):
|
||||
assert episodes_dir.exists()
|
||||
parquet_files = list(episodes_dir.rglob("*.parquet"))
|
||||
assert len(parquet_files) > 0
|
||||
|
||||
|
||||
# ── Tools accessor ───────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_tools_falls_back_to_default_when_info_has_no_tools_field(tmp_path):
|
||||
"""meta.tools returns DEFAULT_TOOLS when info.json doesn't declare any."""
|
||||
from lerobot.datasets.language import DEFAULT_TOOLS
|
||||
|
||||
root = tmp_path / "no_tools"
|
||||
meta = LeRobotDatasetMetadata.create(
|
||||
repo_id="test/no_tools",
|
||||
fps=DEFAULT_FPS,
|
||||
features=SIMPLE_FEATURES,
|
||||
root=root,
|
||||
use_videos=False,
|
||||
)
|
||||
|
||||
assert meta.tools == DEFAULT_TOOLS
|
||||
# info.json on disk should NOT include a `tools` key for clean datasets
|
||||
with open(root / INFO_PATH) as f:
|
||||
info_on_disk = json.load(f)
|
||||
assert "tools" not in info_on_disk
|
||||
|
||||
|
||||
def test_tools_reads_declared_tools_from_info_json(tmp_path):
|
||||
"""A `tools` list written into info.json survives load → meta.tools.
|
||||
|
||||
Regression test for the bug where ``DatasetInfo.from_dict`` silently
|
||||
dropped the ``tools`` key (no matching dataclass field), so
|
||||
``meta.tools`` always returned ``DEFAULT_TOOLS`` regardless of
|
||||
what was on disk.
|
||||
"""
|
||||
from lerobot.datasets.io_utils import load_info
|
||||
|
||||
root = tmp_path / "with_tools"
|
||||
meta = LeRobotDatasetMetadata.create(
|
||||
repo_id="test/with_tools",
|
||||
fps=DEFAULT_FPS,
|
||||
features=SIMPLE_FEATURES,
|
||||
root=root,
|
||||
use_videos=False,
|
||||
)
|
||||
|
||||
custom_tool = {
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "record_observation",
|
||||
"description": "Capture a still image.",
|
||||
"parameters": {
|
||||
"type": "object",
|
||||
"properties": {"label": {"type": "string"}},
|
||||
"required": ["label"],
|
||||
},
|
||||
},
|
||||
}
|
||||
info_path = root / INFO_PATH
|
||||
with open(info_path) as f:
|
||||
raw = json.load(f)
|
||||
raw["tools"] = [custom_tool]
|
||||
with open(info_path, "w") as f:
|
||||
json.dump(raw, f)
|
||||
|
||||
# Reload info from disk and rebind it on the metadata object
|
||||
meta.info = load_info(root)
|
||||
assert meta.tools == [custom_tool]
|
||||
|
||||
|
||||
def test_tools_round_trip_through_dataset_info(tmp_path):
|
||||
"""A `tools` list survives DatasetInfo.from_dict / to_dict."""
|
||||
from lerobot.datasets.utils import DatasetInfo
|
||||
|
||||
raw = {
|
||||
"codebase_version": "v3.1",
|
||||
"fps": 30,
|
||||
"features": SIMPLE_FEATURES,
|
||||
"tools": [{"type": "function", "function": {"name": "say"}}],
|
||||
}
|
||||
info = DatasetInfo.from_dict(raw)
|
||||
assert info.tools == raw["tools"]
|
||||
assert info.to_dict()["tools"] == raw["tools"]
|
||||
|
||||
@@ -1,13 +1,17 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
|
||||
from lerobot.datasets import LeRobotDataset
|
||||
from lerobot.datasets.io_utils import write_info
|
||||
from lerobot.datasets.language import (
|
||||
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
|
||||
pytest.importorskip("pandas", reason="pandas is required (install lerobot[dataset])")
|
||||
|
||||
import numpy as np # noqa: E402
|
||||
import pandas as pd # noqa: E402
|
||||
import pyarrow as pa # noqa: E402
|
||||
|
||||
from lerobot.datasets import LeRobotDataset # noqa: E402
|
||||
from lerobot.datasets.io_utils import write_info # noqa: E402
|
||||
from lerobot.datasets.language import ( # noqa: E402
|
||||
EVENT_ONLY_STYLES,
|
||||
LANGUAGE_EVENTS,
|
||||
LANGUAGE_PERSISTENT,
|
||||
@@ -21,7 +25,7 @@ from lerobot.datasets.language import (
|
||||
language_persistent_arrow_type,
|
||||
validate_camera_field,
|
||||
)
|
||||
from lerobot.datasets.utils import DEFAULT_DATA_PATH
|
||||
from lerobot.datasets.utils import DEFAULT_DATA_PATH # noqa: E402
|
||||
|
||||
|
||||
def test_language_arrow_schema_has_expected_fields():
|
||||
|
||||
@@ -1,11 +1,18 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from lerobot.configs.recipe import MessageTurn, TrainingRecipe
|
||||
from lerobot.datasets.language_render import active_at, emitted_at, nth_next, nth_prev, render_sample
|
||||
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
|
||||
|
||||
from lerobot.configs.recipe import MessageTurn, TrainingRecipe # noqa: E402
|
||||
from lerobot.datasets.language_render import ( # noqa: E402
|
||||
EMITTED_AT_TOLERANCE_S,
|
||||
active_at,
|
||||
emitted_at,
|
||||
nth_next,
|
||||
nth_prev,
|
||||
render_sample,
|
||||
)
|
||||
|
||||
|
||||
def persistent_row(role, content, style, timestamp, tool_calls=None, camera=None):
|
||||
@@ -201,92 +208,50 @@ def test_emitted_at_raises_on_ambiguous_per_camera_vqa():
|
||||
)
|
||||
|
||||
|
||||
def test_per_camera_blend_renders_both_views():
|
||||
recipe = TrainingRecipe(
|
||||
blend={
|
||||
"top": TrainingRecipe(
|
||||
weight=1.0,
|
||||
bindings={
|
||||
"vqa_query": (
|
||||
"emitted_at(t, style=vqa, role=user, camera=observation.images.top)"
|
||||
),
|
||||
"vqa": (
|
||||
"emitted_at(t, style=vqa, role=assistant, camera=observation.images.top)"
|
||||
),
|
||||
},
|
||||
messages=[
|
||||
MessageTurn(
|
||||
role="user",
|
||||
content=[
|
||||
{"type": "image", "feature": "observation.images.top"},
|
||||
{"type": "text", "text": "${vqa_query}"},
|
||||
],
|
||||
stream="high_level",
|
||||
if_present="vqa_query",
|
||||
),
|
||||
MessageTurn(
|
||||
role="assistant",
|
||||
content="${vqa}",
|
||||
stream="high_level",
|
||||
target=True,
|
||||
if_present="vqa",
|
||||
),
|
||||
],
|
||||
def _vqa_subrecipe(camera: str) -> TrainingRecipe:
|
||||
return TrainingRecipe(
|
||||
weight=1.0,
|
||||
bindings={
|
||||
"vqa_query": f"emitted_at(t, style=vqa, role=user, camera={camera})",
|
||||
"vqa": f"emitted_at(t, style=vqa, role=assistant, camera={camera})",
|
||||
},
|
||||
messages=[
|
||||
MessageTurn(
|
||||
role="user",
|
||||
content=[{"type": "image", "feature": camera}, {"type": "text", "text": "${vqa_query}"}],
|
||||
stream="high_level",
|
||||
if_present="vqa_query",
|
||||
),
|
||||
"wrist": TrainingRecipe(
|
||||
weight=1.0,
|
||||
bindings={
|
||||
"vqa_query": (
|
||||
"emitted_at(t, style=vqa, role=user, camera=observation.images.wrist)"
|
||||
),
|
||||
"vqa": (
|
||||
"emitted_at(t, style=vqa, role=assistant, camera=observation.images.wrist)"
|
||||
),
|
||||
},
|
||||
messages=[
|
||||
MessageTurn(
|
||||
role="user",
|
||||
content=[
|
||||
{"type": "image", "feature": "observation.images.wrist"},
|
||||
{"type": "text", "text": "${vqa_query}"},
|
||||
],
|
||||
stream="high_level",
|
||||
if_present="vqa_query",
|
||||
),
|
||||
MessageTurn(
|
||||
role="assistant",
|
||||
content="${vqa}",
|
||||
stream="high_level",
|
||||
target=True,
|
||||
if_present="vqa",
|
||||
),
|
||||
],
|
||||
MessageTurn(
|
||||
role="assistant",
|
||||
content="${vqa}",
|
||||
stream="high_level",
|
||||
target=True,
|
||||
if_present="vqa",
|
||||
),
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
rendered_top = render_sample(
|
||||
recipe=recipe.blend["top"],
|
||||
persistent=PERSISTENT,
|
||||
events=EVENTS_AT_3_TWO_CAMERAS,
|
||||
t=3.0,
|
||||
sample_idx=0,
|
||||
)
|
||||
rendered_wrist = render_sample(
|
||||
recipe=recipe.blend["wrist"],
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("camera", "expected_query", "expected_answer"),
|
||||
[
|
||||
("observation.images.top", "how many cups (top)?", '{"count": 3}'),
|
||||
("observation.images.wrist", "how many cups (wrist)?", '{"count": 1}'),
|
||||
],
|
||||
)
|
||||
def test_per_camera_blend_renders_both_views(camera, expected_query, expected_answer):
|
||||
rendered = render_sample(
|
||||
recipe=_vqa_subrecipe(camera),
|
||||
persistent=PERSISTENT,
|
||||
events=EVENTS_AT_3_TWO_CAMERAS,
|
||||
t=3.0,
|
||||
sample_idx=0,
|
||||
)
|
||||
|
||||
assert rendered_top["messages"][0]["content"][0]["feature"] == "observation.images.top"
|
||||
assert rendered_top["messages"][0]["content"][1]["text"] == "how many cups (top)?"
|
||||
assert rendered_top["messages"][1]["content"] == '{"count": 3}'
|
||||
|
||||
assert rendered_wrist["messages"][0]["content"][0]["feature"] == "observation.images.wrist"
|
||||
assert rendered_wrist["messages"][0]["content"][1]["text"] == "how many cups (wrist)?"
|
||||
assert rendered_wrist["messages"][1]["content"] == '{"count": 1}'
|
||||
assert rendered["messages"][0]["content"][0]["feature"] == camera
|
||||
assert rendered["messages"][0]["content"][1]["text"] == expected_query
|
||||
assert rendered["messages"][1]["content"] == expected_answer
|
||||
|
||||
|
||||
def test_resolve_task_picks_rephrasing_deterministically_per_sample():
|
||||
@@ -319,11 +284,19 @@ def test_resolve_task_picks_rephrasing_deterministically_per_sample():
|
||||
assert seen == {r["content"] for r in rephrasings}
|
||||
# Same sample_idx → same pick (determinism).
|
||||
a = render_sample(
|
||||
recipe=recipe, persistent=rephrasings, events=[], t=0.0, sample_idx=42,
|
||||
recipe=recipe,
|
||||
persistent=rephrasings,
|
||||
events=[],
|
||||
t=0.0,
|
||||
sample_idx=42,
|
||||
dataset_ctx={"task": "canonical"},
|
||||
)
|
||||
b = render_sample(
|
||||
recipe=recipe, persistent=rephrasings, events=[], t=0.0, sample_idx=42,
|
||||
recipe=recipe,
|
||||
persistent=rephrasings,
|
||||
events=[],
|
||||
t=0.0,
|
||||
sample_idx=42,
|
||||
dataset_ctx={"task": "canonical"},
|
||||
)
|
||||
assert a["messages"][0]["content"] == b["messages"][0]["content"]
|
||||
@@ -370,9 +343,65 @@ def test_resolve_task_explicit_override_beats_rephrasings():
|
||||
assert rendered["messages"][0]["content"] == "explicit override wins"
|
||||
|
||||
|
||||
def test_canonical_recipe_can_render_low_level_branch():
|
||||
recipe = TrainingRecipe.from_yaml(Path("src/lerobot/configs/recipes/pi05_hirobot.yaml"))
|
||||
low_level = TrainingRecipe(blend={"low": recipe.blend["low_level_execution"]})
|
||||
def test_emitted_at_persistent_tolerates_small_timestamp_drift():
|
||||
"""Persistent ``emitted_at`` should match within EMITTED_AT_TOLERANCE_S
|
||||
so callers that derive ``t`` arithmetically (``frame_idx / fps``) still
|
||||
line up with the parquet-stored timestamp.
|
||||
"""
|
||||
rows = [persistent_row("assistant", "memo", "memory", 1.0)]
|
||||
# Half a tolerance window — bit-different float, comfortably inside
|
||||
inside = emitted_at(1.0 + EMITTED_AT_TOLERANCE_S / 2, persistent=rows, events=[], style="memory")
|
||||
assert inside is not None and inside["content"] == "memo"
|
||||
|
||||
# Just past the window — no match
|
||||
outside = emitted_at(1.0 + EMITTED_AT_TOLERANCE_S * 2, persistent=rows, events=[], style="memory")
|
||||
assert outside is None
|
||||
|
||||
|
||||
def test_render_sample_rejects_non_dict_language_rows():
|
||||
"""``_normalize_rows`` must surface malformed inputs as TypeError.
|
||||
|
||||
A pipeline that hands the renderer a non-dict (e.g. a stray string)
|
||||
is a real upstream bug — silent skipping would let it propagate.
|
||||
"""
|
||||
recipe = TrainingRecipe(
|
||||
messages=[
|
||||
MessageTurn(role="user", content="${task}", stream="high_level"),
|
||||
MessageTurn(role="assistant", content="ok", stream="high_level", target=True),
|
||||
]
|
||||
)
|
||||
with pytest.raises(TypeError, match="must be dictionaries"):
|
||||
render_sample(
|
||||
recipe=recipe,
|
||||
persistent=["not a dict"],
|
||||
events=[],
|
||||
t=0.0,
|
||||
sample_idx=0,
|
||||
task="x",
|
||||
)
|
||||
|
||||
|
||||
def test_low_level_branch_renders_active_subtask():
|
||||
low_level = TrainingRecipe(
|
||||
blend={
|
||||
"low": TrainingRecipe(
|
||||
weight=1.0,
|
||||
messages=[
|
||||
MessageTurn(
|
||||
role="user",
|
||||
content="${task}\nPlan: ${plan}\nMemory: ${memory}",
|
||||
stream="high_level",
|
||||
),
|
||||
MessageTurn(
|
||||
role="assistant",
|
||||
content="${subtask}",
|
||||
stream="low_level",
|
||||
target=True,
|
||||
),
|
||||
],
|
||||
)
|
||||
}
|
||||
)
|
||||
|
||||
rendered = render_sample(
|
||||
recipe=low_level,
|
||||
|
||||
@@ -80,18 +80,18 @@ def _write_dataset_tree(
|
||||
)
|
||||
tasks = tasks_factory(total_tasks=1)
|
||||
episodes = episodes_factory(
|
||||
features=info["features"],
|
||||
fps=info["fps"],
|
||||
features=info.features,
|
||||
fps=info.fps,
|
||||
total_episodes=1,
|
||||
total_frames=3,
|
||||
tasks=tasks,
|
||||
)
|
||||
stats = stats_factory(features=info["features"])
|
||||
stats = stats_factory(features=info.features)
|
||||
hf_dataset = hf_dataset_factory(
|
||||
features=info["features"],
|
||||
features=info.features,
|
||||
tasks=tasks,
|
||||
episodes=episodes,
|
||||
fps=info["fps"],
|
||||
fps=info.fps,
|
||||
)
|
||||
|
||||
create_info(root, info)
|
||||
@@ -416,6 +416,18 @@ def test_create_initial_counts_zero(tmp_path):
|
||||
assert dataset.num_frames == 0
|
||||
|
||||
|
||||
def test_create_propagates_video_files_size_in_mb(tmp_path):
|
||||
"""video_files_size_in_mb passed to create() is reflected in the dataset metadata."""
|
||||
dataset = LeRobotDataset.create(
|
||||
repo_id=DUMMY_REPO_ID,
|
||||
fps=DEFAULT_FPS,
|
||||
features=SIMPLE_FEATURES,
|
||||
root=tmp_path / "ds",
|
||||
video_files_size_in_mb=42.0,
|
||||
)
|
||||
assert dataset.meta.video_files_size_in_mb == 42.0
|
||||
|
||||
|
||||
def test_add_frame_works_in_write_mode(tmp_path):
|
||||
"""add_frame() succeeds on a dataset created via create()."""
|
||||
dataset = LeRobotDataset.create(
|
||||
|
||||
Vendored
+38
-40
@@ -28,7 +28,7 @@ from datasets import Dataset
|
||||
|
||||
from lerobot.datasets.dataset_metadata import CODEBASE_VERSION, LeRobotDatasetMetadata
|
||||
from lerobot.datasets.feature_utils import get_hf_features_from_features
|
||||
from lerobot.datasets.io_utils import hf_transform_to_torch
|
||||
from lerobot.datasets.io_utils import flatten_dict, hf_transform_to_torch
|
||||
from lerobot.datasets.lerobot_dataset import LeRobotDataset
|
||||
from lerobot.datasets.utils import (
|
||||
DEFAULT_CHUNK_SIZE,
|
||||
@@ -36,10 +36,10 @@ from lerobot.datasets.utils import (
|
||||
DEFAULT_DATA_PATH,
|
||||
DEFAULT_VIDEO_FILE_SIZE_IN_MB,
|
||||
DEFAULT_VIDEO_PATH,
|
||||
DatasetInfo,
|
||||
)
|
||||
from lerobot.datasets.video_utils import encode_video_frames
|
||||
from lerobot.utils.constants import DEFAULT_FEATURES
|
||||
from lerobot.utils.utils import flatten_dict
|
||||
from tests.fixtures.constants import (
|
||||
DEFAULT_FPS,
|
||||
DUMMY_CAMERA_FEATURES,
|
||||
@@ -157,33 +157,31 @@ def info_factory(features_factory):
|
||||
total_episodes: int = 0,
|
||||
total_frames: int = 0,
|
||||
total_tasks: int = 0,
|
||||
total_videos: int = 0,
|
||||
chunks_size: int = DEFAULT_CHUNK_SIZE,
|
||||
data_files_size_in_mb: float = DEFAULT_DATA_FILE_SIZE_IN_MB,
|
||||
video_files_size_in_mb: float = DEFAULT_VIDEO_FILE_SIZE_IN_MB,
|
||||
data_files_size_in_mb: int = DEFAULT_DATA_FILE_SIZE_IN_MB,
|
||||
video_files_size_in_mb: int = DEFAULT_VIDEO_FILE_SIZE_IN_MB,
|
||||
data_path: str = DEFAULT_DATA_PATH,
|
||||
video_path: str = DEFAULT_VIDEO_PATH,
|
||||
motor_features: dict = DUMMY_MOTOR_FEATURES,
|
||||
camera_features: dict = DUMMY_CAMERA_FEATURES,
|
||||
use_videos: bool = True,
|
||||
) -> dict:
|
||||
) -> DatasetInfo:
|
||||
features = features_factory(motor_features, camera_features, use_videos)
|
||||
return {
|
||||
"codebase_version": codebase_version,
|
||||
"robot_type": robot_type,
|
||||
"total_episodes": total_episodes,
|
||||
"total_frames": total_frames,
|
||||
"total_tasks": total_tasks,
|
||||
"total_videos": total_videos,
|
||||
"chunks_size": chunks_size,
|
||||
"data_files_size_in_mb": data_files_size_in_mb,
|
||||
"video_files_size_in_mb": video_files_size_in_mb,
|
||||
"fps": fps,
|
||||
"splits": {},
|
||||
"data_path": data_path,
|
||||
"video_path": video_path if use_videos else None,
|
||||
"features": features,
|
||||
}
|
||||
return DatasetInfo(
|
||||
codebase_version=codebase_version,
|
||||
robot_type=robot_type,
|
||||
total_episodes=total_episodes,
|
||||
total_frames=total_frames,
|
||||
total_tasks=total_tasks,
|
||||
chunks_size=chunks_size,
|
||||
data_files_size_in_mb=data_files_size_in_mb,
|
||||
video_files_size_in_mb=video_files_size_in_mb,
|
||||
fps=fps,
|
||||
splits={},
|
||||
data_path=data_path,
|
||||
video_path=video_path if use_videos else None,
|
||||
features=features,
|
||||
)
|
||||
|
||||
return _create_info
|
||||
|
||||
@@ -333,12 +331,12 @@ def create_videos(info_factory, img_array_factory):
|
||||
total_episodes=total_episodes, total_frames=total_frames, total_tasks=total_tasks
|
||||
)
|
||||
|
||||
video_feats = {key: feats for key, feats in info["features"].items() if feats["dtype"] == "video"}
|
||||
video_feats = {key: feats for key, feats in info.features.items() if feats["dtype"] == "video"}
|
||||
for key, ft in video_feats.items():
|
||||
# create and save images with identifiable content
|
||||
tmp_dir = root / "tmp_images"
|
||||
tmp_dir.mkdir(parents=True, exist_ok=True)
|
||||
for frame_index in range(info["total_frames"]):
|
||||
for frame_index in range(info.total_frames):
|
||||
content = f"{key}-{frame_index}"
|
||||
img = img_array_factory(height=ft["shape"][0], width=ft["shape"][1], content=content)
|
||||
pil_img = PIL.Image.fromarray(img)
|
||||
@@ -348,7 +346,7 @@ def create_videos(info_factory, img_array_factory):
|
||||
video_path = root / DEFAULT_VIDEO_PATH.format(video_key=key, chunk_index=0, file_index=0)
|
||||
video_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
# Use the global fps from info, not video-specific fps which might not exist
|
||||
encode_video_frames(tmp_dir, video_path, fps=info["fps"])
|
||||
encode_video_frames(tmp_dir, video_path, fps=info.fps)
|
||||
shutil.rmtree(tmp_dir)
|
||||
|
||||
return _create_video_directory
|
||||
@@ -433,16 +431,16 @@ def lerobot_dataset_metadata_factory(
|
||||
if info is None:
|
||||
info = info_factory()
|
||||
if stats is None:
|
||||
stats = stats_factory(features=info["features"])
|
||||
stats = stats_factory(features=info.features)
|
||||
if tasks is None:
|
||||
tasks = tasks_factory(total_tasks=info["total_tasks"])
|
||||
tasks = tasks_factory(total_tasks=info.total_tasks)
|
||||
if episodes is None:
|
||||
video_keys = [key for key, ft in info["features"].items() if ft["dtype"] == "video"]
|
||||
video_keys = [key for key, ft in info.features.items() if ft["dtype"] == "video"]
|
||||
episodes = episodes_factory(
|
||||
features=info["features"],
|
||||
fps=info["fps"],
|
||||
total_episodes=info["total_episodes"],
|
||||
total_frames=info["total_frames"],
|
||||
features=info.features,
|
||||
fps=info.fps,
|
||||
total_episodes=info.total_episodes,
|
||||
total_frames=info.total_frames,
|
||||
video_keys=video_keys,
|
||||
tasks=tasks,
|
||||
)
|
||||
@@ -503,23 +501,23 @@ def lerobot_dataset_factory(
|
||||
chunks_size=chunks_size,
|
||||
)
|
||||
if stats is None:
|
||||
stats = stats_factory(features=info["features"])
|
||||
stats = stats_factory(features=info.features)
|
||||
if tasks is None:
|
||||
tasks = tasks_factory(total_tasks=info["total_tasks"])
|
||||
tasks = tasks_factory(total_tasks=info.total_tasks)
|
||||
if episodes_metadata is None:
|
||||
video_keys = [key for key, ft in info["features"].items() if ft["dtype"] == "video"]
|
||||
video_keys = [key for key, ft in info.features.items() if ft["dtype"] == "video"]
|
||||
episodes_metadata = episodes_factory(
|
||||
features=info["features"],
|
||||
fps=info["fps"],
|
||||
total_episodes=info["total_episodes"],
|
||||
total_frames=info["total_frames"],
|
||||
features=info.features,
|
||||
fps=info.fps,
|
||||
total_episodes=info.total_episodes,
|
||||
total_frames=info.total_frames,
|
||||
video_keys=video_keys,
|
||||
tasks=tasks,
|
||||
multi_task=multi_task,
|
||||
)
|
||||
if hf_dataset is None:
|
||||
hf_dataset = hf_dataset_factory(
|
||||
features=info["features"], tasks=tasks, episodes=episodes_metadata, fps=info["fps"]
|
||||
features=info.features, tasks=tasks, episodes=episodes_metadata, fps=info.fps
|
||||
)
|
||||
|
||||
# Write data on disk
|
||||
|
||||
Vendored
+8
-8
@@ -62,19 +62,19 @@ def mock_snapshot_download_factory(
|
||||
if info is None:
|
||||
info = info_factory(data_files_size_in_mb=data_files_size_in_mb, chunks_size=chunks_size)
|
||||
if stats is None:
|
||||
stats = stats_factory(features=info["features"])
|
||||
stats = stats_factory(features=info.features)
|
||||
if tasks is None:
|
||||
tasks = tasks_factory(total_tasks=info["total_tasks"])
|
||||
tasks = tasks_factory(total_tasks=info.total_tasks)
|
||||
if episodes is None:
|
||||
episodes = episodes_factory(
|
||||
features=info["features"],
|
||||
fps=info["fps"],
|
||||
total_episodes=info["total_episodes"],
|
||||
total_frames=info["total_frames"],
|
||||
features=info.features,
|
||||
fps=info.fps,
|
||||
total_episodes=info.total_episodes,
|
||||
total_frames=info.total_frames,
|
||||
tasks=tasks,
|
||||
)
|
||||
if hf_dataset is None:
|
||||
hf_dataset = hf_dataset_factory(tasks=tasks, episodes=episodes, fps=info["fps"])
|
||||
hf_dataset = hf_dataset_factory(tasks=tasks, episodes=episodes, fps=info.fps)
|
||||
|
||||
def _mock_snapshot_download(
|
||||
repo_id: str, # TODO(rcadene): repo_id should be used no?
|
||||
@@ -97,7 +97,7 @@ def mock_snapshot_download_factory(
|
||||
DEFAULT_DATA_PATH.format(chunk_index=0, file_index=0),
|
||||
]
|
||||
|
||||
video_keys = [key for key, feats in info["features"].items() if feats["dtype"] == "video"]
|
||||
video_keys = [key for key, feats in info.features.items() if feats["dtype"] == "video"]
|
||||
for key in video_keys:
|
||||
all_files.append(DEFAULT_VIDEO_PATH.format(video_key=key, chunk_index=0, file_index=0))
|
||||
|
||||
|
||||
@@ -0,0 +1,186 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Smoke tests for EO1's public LeRobot policy interface."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
pytest.importorskip("transformers")
|
||||
|
||||
from lerobot.configs.types import FeatureType, PolicyFeature
|
||||
from lerobot.policies.eo1.modeling_eo1 import EO1Policy
|
||||
from lerobot.utils.constants import ACTION, OBS_STATE
|
||||
|
||||
HIDDEN_SIZE = 8
|
||||
STATE_DIM = 4
|
||||
ACTION_DIM = 3
|
||||
CHUNK_SIZE = 3
|
||||
N_ACTION_STEPS = 2
|
||||
MAX_ACTION_DIM = 6
|
||||
STATE_TOKEN_ID = 5
|
||||
ACTION_TOKEN_ID = 6
|
||||
|
||||
|
||||
class DummyVLMBackbone(nn.Module):
|
||||
def __init__(self, hidden_size: int, vocab_size: int = 64):
|
||||
super().__init__()
|
||||
self.embedding = nn.Embedding(vocab_size, hidden_size)
|
||||
self.config = SimpleNamespace(text_config=SimpleNamespace(hidden_size=hidden_size))
|
||||
|
||||
@property
|
||||
def model(self):
|
||||
return self
|
||||
|
||||
def get_input_embeddings(self):
|
||||
return self.embedding
|
||||
|
||||
def get_rope_index(
|
||||
self,
|
||||
input_ids: torch.Tensor,
|
||||
image_grid_thw: torch.Tensor | None = None,
|
||||
attention_mask: torch.Tensor | None = None,
|
||||
mm_token_type_ids: torch.Tensor | None = None,
|
||||
):
|
||||
batch_size, seq_len = input_ids.shape
|
||||
if attention_mask is None:
|
||||
text_positions = torch.arange(seq_len, device=input_ids.device).expand(batch_size, -1)
|
||||
else:
|
||||
text_positions = attention_mask.long().cumsum(-1) - 1
|
||||
text_positions = text_positions.masked_fill(attention_mask == 0, 0)
|
||||
position_ids = text_positions.view(1, batch_size, seq_len).expand(3, batch_size, seq_len)
|
||||
rope_deltas = torch.zeros(batch_size, 1, dtype=torch.long, device=input_ids.device)
|
||||
return position_ids, rope_deltas
|
||||
|
||||
def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
|
||||
return gradient_checkpointing_kwargs
|
||||
|
||||
def gradient_checkpointing_disable(self):
|
||||
return None
|
||||
|
||||
def forward(
|
||||
self,
|
||||
*,
|
||||
input_ids: torch.Tensor | None = None,
|
||||
inputs_embeds: torch.Tensor | None = None,
|
||||
**kwargs,
|
||||
):
|
||||
if inputs_embeds is None:
|
||||
inputs_embeds = self.embedding(input_ids)
|
||||
return SimpleNamespace(
|
||||
last_hidden_state=inputs_embeds,
|
||||
past_key_values=SimpleNamespace(crop=lambda prefix_len: None),
|
||||
)
|
||||
|
||||
|
||||
def make_eo1_config():
|
||||
from lerobot.policies.eo1.configuration_eo1 import EO1Config
|
||||
|
||||
return EO1Config(
|
||||
device="cpu",
|
||||
dtype="float32",
|
||||
vlm_base="dummy-qwen",
|
||||
vlm_config={},
|
||||
chunk_size=CHUNK_SIZE,
|
||||
n_action_steps=N_ACTION_STEPS,
|
||||
max_state_dim=STATE_DIM,
|
||||
max_action_dim=MAX_ACTION_DIM,
|
||||
num_denoise_steps=2,
|
||||
input_features={
|
||||
OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(STATE_DIM,)),
|
||||
"observation.images.image": PolicyFeature(type=FeatureType.VISUAL, shape=(3, 16, 16)),
|
||||
},
|
||||
output_features={
|
||||
ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(ACTION_DIM,)),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def make_policy_batch(include_action: bool) -> dict[str, torch.Tensor | int]:
|
||||
batch_size = 1
|
||||
seq_len = CHUNK_SIZE + 4
|
||||
input_ids = torch.tensor(
|
||||
[[11, STATE_TOKEN_ID, 12, ACTION_TOKEN_ID, ACTION_TOKEN_ID, ACTION_TOKEN_ID, 13]],
|
||||
dtype=torch.long,
|
||||
)
|
||||
assert input_ids.shape == (batch_size, seq_len)
|
||||
|
||||
batch: dict[str, torch.Tensor | int] = {
|
||||
OBS_STATE: torch.randn(batch_size, STATE_DIM, dtype=torch.float32),
|
||||
"input_ids": input_ids,
|
||||
"attention_mask": torch.ones(batch_size, seq_len, dtype=torch.long),
|
||||
"pixel_values": torch.zeros(batch_size, 3, 4, 4, dtype=torch.float32),
|
||||
"image_grid_thw": torch.tensor([[1, 2, 2]], dtype=torch.long),
|
||||
"mm_token_type_ids": torch.zeros(batch_size, seq_len, dtype=torch.int32),
|
||||
"state_token_id": STATE_TOKEN_ID,
|
||||
"action_token_id": ACTION_TOKEN_ID,
|
||||
}
|
||||
if include_action:
|
||||
batch[ACTION] = torch.randn(batch_size, CHUNK_SIZE, ACTION_DIM, dtype=torch.float32)
|
||||
return batch
|
||||
|
||||
|
||||
def test_lerobot_eo1_forward_pass(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
"lerobot.policies.eo1.modeling_eo1.Qwen2_5_VLForConditionalGeneration.from_pretrained",
|
||||
lambda *args, **kwargs: DummyVLMBackbone(HIDDEN_SIZE),
|
||||
)
|
||||
policy = EO1Policy(make_eo1_config())
|
||||
|
||||
loss, metrics = policy.forward(make_policy_batch(include_action=True))
|
||||
|
||||
assert loss.ndim == 0
|
||||
assert torch.isfinite(loss)
|
||||
assert metrics["loss"] == pytest.approx(loss.item())
|
||||
|
||||
|
||||
def test_lerobot_eo1_inference(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
"lerobot.policies.eo1.modeling_eo1.Qwen2_5_VLForConditionalGeneration.from_pretrained",
|
||||
lambda *args, **kwargs: DummyVLMBackbone(HIDDEN_SIZE),
|
||||
)
|
||||
policy = EO1Policy(make_eo1_config())
|
||||
|
||||
sample_calls = {"count": 0}
|
||||
fixed_chunk = torch.tensor(
|
||||
[
|
||||
[
|
||||
[0.1, 0.2, 0.3, 9.0, 9.0, 9.0],
|
||||
[1.1, 1.2, 1.3, 9.0, 9.0, 9.0],
|
||||
[2.1, 2.2, 2.3, 9.0, 9.0, 9.0],
|
||||
]
|
||||
],
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
def fake_sample_actions(**kwargs):
|
||||
sample_calls["count"] += 1
|
||||
return fixed_chunk
|
||||
|
||||
monkeypatch.setattr(policy.model, "sample_actions", fake_sample_actions)
|
||||
|
||||
batch = make_policy_batch(include_action=False)
|
||||
action_0 = policy.select_action(batch)
|
||||
action_1 = policy.select_action(batch)
|
||||
|
||||
torch.testing.assert_close(action_0, fixed_chunk[:, 0, :ACTION_DIM])
|
||||
torch.testing.assert_close(action_1, fixed_chunk[:, 1, :ACTION_DIM])
|
||||
assert sample_calls["count"] == 1
|
||||
@@ -17,9 +17,9 @@
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from lerobot.policies.rtc.action_interpolator import ActionInterpolator
|
||||
from lerobot.policies.rtc.action_queue import ActionQueue
|
||||
from lerobot.policies.rtc.configuration_rtc import RTCConfig
|
||||
from lerobot.utils.action_interpolator import ActionInterpolator
|
||||
|
||||
# ====================== Fixtures ======================
|
||||
|
||||
|
||||
@@ -26,7 +26,7 @@ def test_rtc_config_default_initialization():
|
||||
"""Test RTCConfig initializes with default values."""
|
||||
config = RTCConfig()
|
||||
|
||||
assert config.enabled is False
|
||||
assert config.enabled is True
|
||||
assert config.prefix_attention_schedule == RTCAttentionSchedule.LINEAR
|
||||
assert config.max_guidance_weight == 10.0
|
||||
assert config.execution_horizon == 10
|
||||
|
||||
@@ -22,7 +22,7 @@ from lerobot.configs.types import (
|
||||
PolicyFeature,
|
||||
RTCAttentionSchedule,
|
||||
)
|
||||
from lerobot.processor import TransitionKey, batch_to_transition
|
||||
from lerobot.processor import TransitionKey, batch_to_transition, create_transition
|
||||
from lerobot.processor.normalize_processor import NormalizerProcessorStep, UnnormalizerProcessorStep
|
||||
from lerobot.processor.relative_action_processor import (
|
||||
AbsoluteActionsProcessorStep,
|
||||
@@ -52,6 +52,9 @@ _rtc_debug_mod = _import_rtc_module("lerobot.policies.rtc.debug_tracker", "debug
|
||||
_rtc_mod = _import_rtc_module("lerobot.policies.rtc.modeling_rtc", "modeling_rtc.py")
|
||||
RTCProcessor = _rtc_mod.RTCProcessor
|
||||
|
||||
_rtc_relative_mod = _import_rtc_module("lerobot.policies.rtc.relative", "relative.py")
|
||||
reanchor_relative_rtc_prefix = _rtc_relative_mod.reanchor_relative_rtc_prefix
|
||||
|
||||
ACTION_DIM = 6
|
||||
CHUNK_SIZE = 50
|
||||
EXECUTION_HORIZON = 10
|
||||
@@ -187,7 +190,7 @@ class TestRTCDenoiseWithRelativeLeftovers:
|
||||
|
||||
|
||||
class TestFullPipelineRelativeRTC:
|
||||
"""End-to-end test of the RTC + relative actions pipeline matching eval_with_real_robot.py flow."""
|
||||
"""End-to-end test of the RTC + relative actions pipeline matching lerobot-rollout flow."""
|
||||
|
||||
def test_preprocessor_caches_state_for_postprocessor(self):
|
||||
"""Preprocessor's relative step should cache state so postprocessor can convert back."""
|
||||
@@ -218,7 +221,9 @@ class TestFullPipelineRelativeRTC:
|
||||
|
||||
def test_roundtrip_with_identity_normalization(self):
|
||||
"""Actions → relative → normalize → [model] → unnormalize → absolute should recover originals.
|
||||
Using mean=0, std=1 normalization (identity)."""
|
||||
|
||||
Using mean=0, std=1 normalization (identity).
|
||||
"""
|
||||
relative_step, normalizer, unnormalizer, absolute_step = _make_relative_pipeline()
|
||||
|
||||
state = torch.randn(1, ACTION_DIM)
|
||||
@@ -240,7 +245,7 @@ class TestFullPipelineRelativeRTC:
|
||||
torch.testing.assert_close(recovered, actions, atol=1e-5, rtol=1e-5)
|
||||
|
||||
def test_eval_loop_simulation(self):
|
||||
"""Simulate the eval_with_real_robot.py loop with relative actions.
|
||||
"""Simulate the lerobot-rollout loop with relative actions.
|
||||
|
||||
Iteration 1: No leftovers → model generates relative actions → store for RTC
|
||||
Iteration 2: Use leftovers as RTC guidance → model generates new relative actions
|
||||
@@ -400,13 +405,113 @@ class TestStateRebasingApproximation:
|
||||
assert error_excluded < 1e-6, f"Excluded joint should have zero error, got {error_excluded}"
|
||||
|
||||
|
||||
class TestRTCReanchoringWithStateNormalizer:
|
||||
"""RTC re-anchoring under non-identity OBS_STATE normalization."""
|
||||
|
||||
@staticmethod
|
||||
def _build_normalizer_with_state_stats():
|
||||
"""Build a relative-action preprocessor with non-trivial OBS_STATE stats."""
|
||||
features = {
|
||||
ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(ACTION_DIM,)),
|
||||
OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(ACTION_DIM,)),
|
||||
}
|
||||
norm_map = {
|
||||
FeatureType.ACTION: NormalizationMode.MEAN_STD,
|
||||
FeatureType.STATE: NormalizationMode.MEAN_STD,
|
||||
}
|
||||
stats = {
|
||||
ACTION: {
|
||||
"mean": torch.zeros(ACTION_DIM).numpy(),
|
||||
"std": (0.5 * torch.ones(ACTION_DIM)).numpy(),
|
||||
},
|
||||
OBS_STATE: {
|
||||
"mean": (5.0 * torch.ones(ACTION_DIM)).numpy(),
|
||||
"std": (2.0 * torch.ones(ACTION_DIM)).numpy(),
|
||||
},
|
||||
}
|
||||
relative_step = RelativeActionsProcessorStep(enabled=True)
|
||||
normalizer = NormalizerProcessorStep(features=features, norm_map=norm_map, stats=stats)
|
||||
return relative_step, normalizer
|
||||
|
||||
def test_reanchor_with_raw_state_matches_normalize_of_absolute_minus_state(self):
|
||||
"""Reanchoring with the raw cached state yields ``normalize(prev_actions_absolute - raw_state)``."""
|
||||
relative_step, normalizer = self._build_normalizer_with_state_stats()
|
||||
|
||||
raw_state = torch.tensor([[7.0, 8.0, 9.0, 10.0, 11.0, 12.0]])
|
||||
relative_step(batch_to_transition({OBS_STATE: raw_state.clone()}))
|
||||
|
||||
prev_actions_absolute = torch.tensor([[2.0, 3.0, 4.0, 5.0, 6.0, 7.0]] * 5)
|
||||
|
||||
result = reanchor_relative_rtc_prefix(
|
||||
prev_actions_absolute=prev_actions_absolute,
|
||||
current_state=relative_step.get_cached_state(),
|
||||
relative_step=relative_step,
|
||||
normalizer_step=normalizer,
|
||||
policy_device="cpu",
|
||||
)
|
||||
|
||||
expected_relative = to_relative_actions(prev_actions_absolute, raw_state, [True] * ACTION_DIM)
|
||||
expected = normalizer(create_transition(action=expected_relative))[TransitionKey.ACTION]
|
||||
torch.testing.assert_close(result, expected, atol=1e-5, rtol=1e-5)
|
||||
|
||||
def test_reanchor_with_normalized_state_produces_wrong_result(self):
|
||||
"""Reanchoring with raw vs. normalized state produces meaningfully different outputs."""
|
||||
relative_step, normalizer = self._build_normalizer_with_state_stats()
|
||||
|
||||
raw_state = torch.tensor([[7.0, 8.0, 9.0, 10.0, 11.0, 12.0]])
|
||||
relative_step(batch_to_transition({OBS_STATE: raw_state.clone()}))
|
||||
|
||||
normalized_obs = normalizer(batch_to_transition({OBS_STATE: raw_state.clone()}))
|
||||
normalized_state = normalized_obs[TransitionKey.OBSERVATION][OBS_STATE]
|
||||
assert not torch.allclose(normalized_state, raw_state)
|
||||
|
||||
prev_actions_absolute = torch.tensor([[2.0, 3.0, 4.0, 5.0, 6.0, 7.0]] * 5)
|
||||
|
||||
result_raw = reanchor_relative_rtc_prefix(
|
||||
prev_actions_absolute=prev_actions_absolute,
|
||||
current_state=raw_state,
|
||||
relative_step=relative_step,
|
||||
normalizer_step=normalizer,
|
||||
policy_device="cpu",
|
||||
)
|
||||
result_normalized = reanchor_relative_rtc_prefix(
|
||||
prev_actions_absolute=prev_actions_absolute,
|
||||
current_state=normalized_state,
|
||||
relative_step=relative_step,
|
||||
normalizer_step=normalizer,
|
||||
policy_device="cpu",
|
||||
)
|
||||
|
||||
max_abs_diff = (result_raw - result_normalized).abs().max()
|
||||
assert max_abs_diff > 0.5, (
|
||||
f"Raw and normalized state produced near-identical outputs (max diff {max_abs_diff:.4f}); "
|
||||
"OBS_STATE stats are too close to identity to be sensitive."
|
||||
)
|
||||
|
||||
def test_engine_pipeline_cached_state_is_raw_after_full_preprocess(self):
|
||||
"""``get_cached_state()`` returns raw OBS_STATE after the full preprocessor pipeline runs."""
|
||||
relative_step, normalizer = self._build_normalizer_with_state_stats()
|
||||
|
||||
raw_state = torch.tensor([[7.0, 8.0, 9.0, 10.0, 11.0, 12.0]])
|
||||
|
||||
transition = batch_to_transition({OBS_STATE: raw_state.clone()})
|
||||
transition = relative_step(transition)
|
||||
preprocessed = normalizer(transition)
|
||||
|
||||
cached = relative_step.get_cached_state()
|
||||
torch.testing.assert_close(cached, raw_state, atol=1e-6, rtol=1e-6)
|
||||
|
||||
post_normalize_state = preprocessed[TransitionKey.OBSERVATION][OBS_STATE]
|
||||
assert not torch.allclose(cached, post_normalize_state, atol=1e-3)
|
||||
|
||||
|
||||
def _detect_relative_actions(preprocessor) -> bool:
|
||||
"""Mirror of the helper in eval_with_real_robot.py for testing without importing it."""
|
||||
"""Mirror of the helper in lerobot-rollout for testing without importing it."""
|
||||
return any(isinstance(step, RelativeActionsProcessorStep) and step.enabled for step in preprocessor.steps)
|
||||
|
||||
|
||||
class TestDetectRelativeActions:
|
||||
"""Test the _detect_relative_actions helper logic used by eval_with_real_robot.py."""
|
||||
"""Test the _detect_relative_actions helper logic used by lerobot-rollout."""
|
||||
|
||||
def test_detects_enabled_relative_step(self):
|
||||
class FakePipeline:
|
||||
|
||||
@@ -1,11 +1,15 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
import torch
|
||||
import pytest
|
||||
|
||||
from lerobot.configs.recipe import MessageTurn, TrainingRecipe
|
||||
from lerobot.processor.converters import create_transition
|
||||
from lerobot.processor.render_messages_processor import RenderMessagesStep
|
||||
from lerobot.types import TransitionKey
|
||||
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
|
||||
|
||||
import torch # noqa: E402
|
||||
|
||||
from lerobot.configs.recipe import MessageTurn, TrainingRecipe # noqa: E402
|
||||
from lerobot.processor.converters import create_transition # noqa: E402
|
||||
from lerobot.processor.render_messages_processor import RenderMessagesStep # noqa: E402
|
||||
from lerobot.types import TransitionKey # noqa: E402
|
||||
|
||||
|
||||
def test_render_messages_step_noops_without_language_columns():
|
||||
|
||||
+10
-34
@@ -1,5 +1,3 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
@@ -21,8 +19,6 @@ import pytest
|
||||
import torch
|
||||
|
||||
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
|
||||
from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
|
||||
from lerobot.policies.sac.reward_model.processor_classifier import make_classifier_processor
|
||||
from lerobot.processor import (
|
||||
DataProcessorPipeline,
|
||||
DeviceProcessorStep,
|
||||
@@ -31,6 +27,8 @@ from lerobot.processor import (
|
||||
TransitionKey,
|
||||
)
|
||||
from lerobot.processor.converters import create_transition, transition_to_batch
|
||||
from lerobot.rewards.classifier.configuration_classifier import RewardClassifierConfig
|
||||
from lerobot.rewards.classifier.processor_classifier import make_classifier_processor
|
||||
from lerobot.utils.constants import OBS_IMAGE, OBS_STATE
|
||||
|
||||
|
||||
@@ -42,7 +40,7 @@ def create_default_config():
|
||||
OBS_IMAGE: PolicyFeature(type=FeatureType.VISUAL, shape=(3, 224, 224)),
|
||||
}
|
||||
config.output_features = {
|
||||
"reward": PolicyFeature(type=FeatureType.ACTION, shape=(1,)), # Classifier output
|
||||
"reward": PolicyFeature(type=FeatureType.ACTION, shape=(1,)),
|
||||
}
|
||||
config.normalization_mapping = {
|
||||
FeatureType.STATE: NormalizationMode.MEAN_STD,
|
||||
@@ -90,17 +88,14 @@ def test_classifier_processor_normalization():
|
||||
config = create_default_config()
|
||||
stats = create_default_stats()
|
||||
|
||||
preprocessor, postprocessor = make_classifier_processor(
|
||||
config,
|
||||
stats,
|
||||
)
|
||||
preprocessor, postprocessor = make_classifier_processor(config, stats)
|
||||
|
||||
# Create test data
|
||||
observation = {
|
||||
OBS_STATE: torch.randn(10),
|
||||
OBS_IMAGE: torch.randn(3, 224, 224),
|
||||
}
|
||||
action = torch.randn(1) # Dummy action/reward
|
||||
action = torch.randn(1)
|
||||
transition = create_transition(observation, action)
|
||||
batch = transition_to_batch(transition)
|
||||
|
||||
@@ -120,10 +115,7 @@ def test_classifier_processor_cuda():
|
||||
config.device = "cuda"
|
||||
stats = create_default_stats()
|
||||
|
||||
preprocessor, postprocessor = make_classifier_processor(
|
||||
config,
|
||||
stats,
|
||||
)
|
||||
preprocessor, postprocessor = make_classifier_processor(config, stats)
|
||||
|
||||
# Create CPU data
|
||||
observation = {
|
||||
@@ -132,7 +124,6 @@ def test_classifier_processor_cuda():
|
||||
}
|
||||
action = torch.randn(1)
|
||||
transition = create_transition(observation, action)
|
||||
|
||||
batch = transition_to_batch(transition)
|
||||
|
||||
# Process through preprocessor
|
||||
@@ -158,10 +149,7 @@ def test_classifier_processor_accelerate_scenario():
|
||||
config.device = "cuda:0"
|
||||
stats = create_default_stats()
|
||||
|
||||
preprocessor, postprocessor = make_classifier_processor(
|
||||
config,
|
||||
stats,
|
||||
)
|
||||
preprocessor, postprocessor = make_classifier_processor(config, stats)
|
||||
|
||||
# Simulate Accelerate: data already on GPU
|
||||
device = torch.device("cuda:0")
|
||||
@@ -171,7 +159,6 @@ def test_classifier_processor_accelerate_scenario():
|
||||
}
|
||||
action = torch.randn(1).to(device)
|
||||
transition = create_transition(observation, action)
|
||||
|
||||
batch = transition_to_batch(transition)
|
||||
|
||||
# Process through preprocessor
|
||||
@@ -201,7 +188,6 @@ def test_classifier_processor_multi_gpu():
|
||||
}
|
||||
action = torch.randn(1).to(device)
|
||||
transition = create_transition(observation, action)
|
||||
|
||||
batch = transition_to_batch(transition)
|
||||
|
||||
# Process through preprocessor
|
||||
@@ -231,7 +217,6 @@ def test_classifier_processor_without_stats():
|
||||
}
|
||||
action = torch.randn(1)
|
||||
transition = create_transition(observation, action)
|
||||
|
||||
batch = transition_to_batch(transition)
|
||||
|
||||
processed = preprocessor(batch)
|
||||
@@ -294,7 +279,6 @@ def test_classifier_processor_mixed_precision():
|
||||
}
|
||||
action = torch.randn(1, dtype=torch.float32)
|
||||
transition = create_transition(observation, action)
|
||||
|
||||
batch = transition_to_batch(transition)
|
||||
|
||||
# Process through preprocessor
|
||||
@@ -312,10 +296,7 @@ def test_classifier_processor_batch_data():
|
||||
config = create_default_config()
|
||||
stats = create_default_stats()
|
||||
|
||||
preprocessor, postprocessor = make_classifier_processor(
|
||||
config,
|
||||
stats,
|
||||
)
|
||||
preprocessor, postprocessor = make_classifier_processor(config, stats)
|
||||
|
||||
# Test with batched data
|
||||
batch_size = 16
|
||||
@@ -325,7 +306,6 @@ def test_classifier_processor_batch_data():
|
||||
}
|
||||
action = torch.randn(batch_size, 1)
|
||||
transition = create_transition(observation, action)
|
||||
|
||||
batch = transition_to_batch(transition)
|
||||
|
||||
# Process through preprocessor
|
||||
@@ -343,15 +323,11 @@ def test_classifier_processor_postprocessor_identity():
|
||||
config = create_default_config()
|
||||
stats = create_default_stats()
|
||||
|
||||
preprocessor, postprocessor = make_classifier_processor(
|
||||
config,
|
||||
stats,
|
||||
)
|
||||
preprocessor, postprocessor = make_classifier_processor(config, stats)
|
||||
|
||||
# Create test data for postprocessor
|
||||
reward = torch.tensor([[0.8], [0.3], [0.9]]) # Batch of rewards/predictions
|
||||
reward = torch.tensor([[0.8], [0.3], [0.9]])
|
||||
transition = create_transition(action=reward)
|
||||
|
||||
_ = transition_to_batch(transition)
|
||||
|
||||
# Process through postprocessor
|
||||
+15
-9
@@ -1,5 +1,3 @@
|
||||
# !/usr/bin/env python
|
||||
|
||||
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
@@ -18,8 +16,8 @@ import pytest
|
||||
import torch
|
||||
|
||||
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
|
||||
from lerobot.policies.sac.reward_model.configuration_classifier import RewardClassifierConfig
|
||||
from lerobot.policies.sac.reward_model.modeling_classifier import ClassifierOutput
|
||||
from lerobot.rewards.classifier.configuration_classifier import RewardClassifierConfig
|
||||
from lerobot.rewards.classifier.modeling_classifier import ClassifierOutput
|
||||
from lerobot.utils.constants import OBS_IMAGE, REWARD
|
||||
from tests.utils import skip_if_package_missing
|
||||
|
||||
@@ -42,7 +40,7 @@ def test_classifier_output():
|
||||
reason="helper2424/resnet10 needs to be updated to work with the latest version of transformers"
|
||||
)
|
||||
def test_binary_classifier_with_default_params():
|
||||
from lerobot.policies.sac.reward_model.modeling_classifier import Classifier
|
||||
from lerobot.rewards.classifier.modeling_classifier import Classifier
|
||||
|
||||
config = RewardClassifierConfig()
|
||||
config.input_features = {
|
||||
@@ -86,7 +84,7 @@ def test_binary_classifier_with_default_params():
|
||||
reason="helper2424/resnet10 needs to be updated to work with the latest version of transformers"
|
||||
)
|
||||
def test_multiclass_classifier():
|
||||
from lerobot.policies.sac.reward_model.modeling_classifier import Classifier
|
||||
from lerobot.rewards.classifier.modeling_classifier import Classifier
|
||||
|
||||
num_classes = 5
|
||||
config = RewardClassifierConfig()
|
||||
@@ -128,11 +126,15 @@ def test_multiclass_classifier():
|
||||
reason="helper2424/resnet10 needs to be updated to work with the latest version of transformers"
|
||||
)
|
||||
def test_default_device():
|
||||
from lerobot.policies.sac.reward_model.modeling_classifier import Classifier
|
||||
from lerobot.rewards.classifier.modeling_classifier import Classifier
|
||||
|
||||
config = RewardClassifierConfig()
|
||||
assert config.device == "cpu"
|
||||
assert config.device is None or config.device == "cpu"
|
||||
|
||||
config.input_features = {
|
||||
OBS_IMAGE: PolicyFeature(type=FeatureType.VISUAL, shape=(3, 224, 224)),
|
||||
}
|
||||
config.num_cameras = 1
|
||||
classifier = Classifier(config)
|
||||
for p in classifier.parameters():
|
||||
assert p.device == torch.device("cpu")
|
||||
@@ -143,11 +145,15 @@ def test_default_device():
|
||||
reason="helper2424/resnet10 needs to be updated to work with the latest version of transformers"
|
||||
)
|
||||
def test_explicit_device_setup():
|
||||
from lerobot.policies.sac.reward_model.modeling_classifier import Classifier
|
||||
from lerobot.rewards.classifier.modeling_classifier import Classifier
|
||||
|
||||
config = RewardClassifierConfig(device="cpu")
|
||||
assert config.device == "cpu"
|
||||
|
||||
config.input_features = {
|
||||
OBS_IMAGE: PolicyFeature(type=FeatureType.VISUAL, shape=(3, 224, 224)),
|
||||
}
|
||||
config.num_cameras = 1
|
||||
classifier = Classifier(config)
|
||||
for p in classifier.parameters():
|
||||
assert p.device == torch.device("cpu")
|
||||
@@ -0,0 +1,447 @@
|
||||
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Tests for the reward model base classes and registry."""
|
||||
|
||||
import json
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from lerobot.configs.rewards import RewardModelConfig
|
||||
from lerobot.optim.optimizers import AdamWConfig
|
||||
from lerobot.rewards.pretrained import PreTrainedRewardModel
|
||||
|
||||
|
||||
@RewardModelConfig.register_subclass(name="_dummy_hub_reward")
|
||||
@dataclass
|
||||
class _DummyHubRewardConfig(RewardModelConfig):
|
||||
def get_optimizer_preset(self):
|
||||
return AdamWConfig(lr=1e-4)
|
||||
|
||||
|
||||
class _DummyHubReward(PreTrainedRewardModel):
|
||||
config_class = _DummyHubRewardConfig
|
||||
name = "_dummy_hub_reward"
|
||||
|
||||
def __init__(self, config):
|
||||
super().__init__(config)
|
||||
self.bias = torch.nn.Parameter(torch.zeros(1))
|
||||
|
||||
def compute_reward(self, batch):
|
||||
return self.bias.expand(1)
|
||||
|
||||
|
||||
def test_reward_model_config_registry():
|
||||
"""Verify that classifier and sarm are registered."""
|
||||
known = RewardModelConfig.get_known_choices()
|
||||
assert "reward_classifier" in known
|
||||
assert "sarm" in known
|
||||
|
||||
|
||||
def test_reward_model_config_lookup():
|
||||
"""Verify that we can look up configs by name."""
|
||||
cls = RewardModelConfig.get_choice_class("reward_classifier")
|
||||
from lerobot.rewards.classifier.configuration_classifier import RewardClassifierConfig
|
||||
|
||||
assert cls is RewardClassifierConfig
|
||||
|
||||
|
||||
def test_factory_get_reward_model_class():
|
||||
"""Test the get_reward_model_class factory."""
|
||||
from lerobot.rewards.factory import get_reward_model_class
|
||||
|
||||
cls = get_reward_model_class("sarm")
|
||||
from lerobot.rewards.sarm.modeling_sarm import SARMRewardModel
|
||||
|
||||
assert cls is SARMRewardModel
|
||||
|
||||
|
||||
def test_factory_unknown_raises():
|
||||
"""Unknown name should raise ValueError."""
|
||||
from lerobot.rewards.factory import get_reward_model_class
|
||||
|
||||
with pytest.raises(ValueError, match="not available"):
|
||||
get_reward_model_class("nonexistent_reward_model")
|
||||
|
||||
|
||||
def test_pretrained_reward_model_requires_config_class():
|
||||
"""Subclass without config_class should fail."""
|
||||
with pytest.raises(TypeError, match="must define 'config_class'"):
|
||||
|
||||
class BadModel(PreTrainedRewardModel):
|
||||
name = "bad"
|
||||
|
||||
def compute_reward(self, batch):
|
||||
pass
|
||||
|
||||
|
||||
def test_pretrained_reward_model_requires_name():
|
||||
"""Subclass without name should fail."""
|
||||
with pytest.raises(TypeError, match="must define 'name'"):
|
||||
|
||||
class BadModel(PreTrainedRewardModel):
|
||||
config_class = RewardModelConfig
|
||||
|
||||
def compute_reward(self, batch):
|
||||
pass
|
||||
|
||||
|
||||
def test_non_trainable_forward_raises():
|
||||
"""Non-trainable model should raise on forward()."""
|
||||
from dataclasses import dataclass
|
||||
|
||||
from lerobot.optim.optimizers import AdamWConfig
|
||||
|
||||
@dataclass
|
||||
class DummyConfig(RewardModelConfig):
|
||||
def get_optimizer_preset(self):
|
||||
return AdamWConfig(lr=1e-4)
|
||||
|
||||
class DummyReward(PreTrainedRewardModel):
|
||||
config_class = DummyConfig
|
||||
name = "dummy_test"
|
||||
|
||||
def compute_reward(self, batch):
|
||||
return torch.zeros(1)
|
||||
|
||||
config = DummyConfig()
|
||||
model = DummyReward(config)
|
||||
|
||||
with pytest.raises(NotImplementedError, match="not trainable"):
|
||||
model.forward({"x": torch.zeros(1)})
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Trainable vs zero-shot (general-purpose) reward models.
|
||||
# The proposal explicitly supports models like TOPReward that wrap a pretrained
|
||||
# VLM and produce a reward signal without any training step. These tests pin
|
||||
# the contract that lets such models coexist with trainable ones.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_is_trainable_false_when_forward_not_overridden():
|
||||
"""A reward model that only implements ``compute_reward`` is zero-shot."""
|
||||
model, _ = _make_dummy_reward_model()
|
||||
assert model.is_trainable is False
|
||||
|
||||
|
||||
def test_is_trainable_true_when_forward_overridden():
|
||||
"""Overriding ``forward`` flips ``is_trainable`` to True."""
|
||||
|
||||
class _TrainableReward(_DummyHubReward):
|
||||
name = "_trainable_dummy_reward"
|
||||
|
||||
def forward(self, batch):
|
||||
loss = (self.bias**2).sum()
|
||||
return loss, {}
|
||||
|
||||
# Register a fresh config subclass so the subclass check passes.
|
||||
@RewardModelConfig.register_subclass(name="_trainable_dummy_reward")
|
||||
@dataclass
|
||||
class _TrainableConfig(_DummyHubRewardConfig):
|
||||
pass
|
||||
|
||||
_TrainableReward.config_class = _TrainableConfig
|
||||
model = _TrainableReward(_TrainableConfig())
|
||||
assert model.is_trainable is True
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# RewardModelConfig.from_pretrained
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_reward_model_config_from_pretrained_raises_when_config_missing(tmp_path):
|
||||
"""``from_pretrained`` must surface a clear ``FileNotFoundError`` when the
|
||||
target directory exists but does not contain ``config.json``, instead of
|
||||
crashing later inside ``draccus.parse``.
|
||||
"""
|
||||
# tmp_path exists but has no config.json
|
||||
with pytest.raises(FileNotFoundError, match="config.json not found"):
|
||||
RewardModelConfig.from_pretrained(tmp_path)
|
||||
|
||||
|
||||
def test_reward_model_config_from_pretrained_roundtrip(tmp_path):
|
||||
"""Round-trip: save a RewardClassifierConfig, reload it, fields must match."""
|
||||
from lerobot.rewards.classifier.configuration_classifier import RewardClassifierConfig
|
||||
|
||||
original = RewardClassifierConfig(
|
||||
num_classes=3,
|
||||
hidden_dim=128,
|
||||
latent_dim=64,
|
||||
num_cameras=1,
|
||||
learning_rate=5e-4,
|
||||
)
|
||||
original._save_pretrained(tmp_path)
|
||||
|
||||
loaded = RewardModelConfig.from_pretrained(tmp_path)
|
||||
|
||||
assert isinstance(loaded, RewardClassifierConfig)
|
||||
assert loaded.num_classes == 3
|
||||
assert loaded.hidden_dim == 128
|
||||
assert loaded.latent_dim == 64
|
||||
assert loaded.num_cameras == 1
|
||||
assert loaded.learning_rate == 5e-4
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# TrainPipelineConfig — reward model training path
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_train_pipeline_config_path_fields_includes_reward_model():
|
||||
"""``--reward_model.path=local/dir`` requires ``reward_model`` to be listed
|
||||
as a draccus path-field on ``TrainPipelineConfig``."""
|
||||
from lerobot.configs.train import TrainPipelineConfig
|
||||
|
||||
fields = TrainPipelineConfig.__get_path_fields__()
|
||||
assert "policy" in fields
|
||||
assert "reward_model" in fields
|
||||
|
||||
|
||||
def test_train_pipeline_config_trainable_config_returns_reward_model_when_set():
|
||||
"""When only ``reward_model`` is set, ``trainable_config`` (used by the
|
||||
trainer for e.g. ``.device``) must return it — not ``None`` from ``policy``."""
|
||||
from lerobot.configs.default import DatasetConfig
|
||||
from lerobot.configs.train import TrainPipelineConfig
|
||||
from lerobot.rewards.classifier.configuration_classifier import RewardClassifierConfig
|
||||
|
||||
reward_cfg = RewardClassifierConfig(device="cpu")
|
||||
cfg = TrainPipelineConfig(
|
||||
dataset=DatasetConfig(repo_id="user/repo"),
|
||||
reward_model=reward_cfg,
|
||||
)
|
||||
|
||||
assert cfg.is_reward_model_training is True
|
||||
assert cfg.trainable_config is reward_cfg
|
||||
# This is what lerobot_train.py uses to decide force_cpu; ``cfg.policy.device``
|
||||
# would AttributeError here because policy is None.
|
||||
assert cfg.trainable_config.device == "cpu"
|
||||
|
||||
|
||||
def test_train_pipeline_config_trainable_config_returns_policy_when_set():
|
||||
"""Mirror of the reward-model case: when only ``policy`` is set,
|
||||
``trainable_config`` must return it."""
|
||||
from lerobot.configs.default import DatasetConfig
|
||||
from lerobot.configs.train import TrainPipelineConfig
|
||||
from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
|
||||
|
||||
policy_cfg = DiffusionConfig(device="cpu")
|
||||
cfg = TrainPipelineConfig(
|
||||
dataset=DatasetConfig(repo_id="user/repo"),
|
||||
policy=policy_cfg,
|
||||
)
|
||||
|
||||
assert cfg.is_reward_model_training is False
|
||||
assert cfg.trainable_config is policy_cfg
|
||||
assert cfg.trainable_config.device == "cpu"
|
||||
|
||||
|
||||
def test_train_pipeline_config_from_pretrained_migrates_legacy_rabc_fields(tmp_path):
|
||||
"""Legacy top-level RA-BC fields should be migrated into ``sample_weighting``."""
|
||||
from lerobot.configs.default import DatasetConfig
|
||||
from lerobot.configs.train import TRAIN_CONFIG_NAME, TrainPipelineConfig
|
||||
from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
|
||||
|
||||
cfg = TrainPipelineConfig(
|
||||
dataset=DatasetConfig(repo_id="user/repo"),
|
||||
policy=DiffusionConfig(device="cpu"),
|
||||
)
|
||||
cfg._save_pretrained(tmp_path)
|
||||
|
||||
config_path = tmp_path / TRAIN_CONFIG_NAME
|
||||
with open(config_path) as f:
|
||||
payload = json.load(f)
|
||||
|
||||
payload.pop("sample_weighting", None)
|
||||
payload.update(
|
||||
{
|
||||
"use_rabc": True,
|
||||
"rabc_progress_path": "hf://datasets/user/repo/sarm_progress.parquet",
|
||||
"rabc_kappa": 0.05,
|
||||
"rabc_epsilon": 1e-5,
|
||||
"rabc_head_mode": "dense",
|
||||
}
|
||||
)
|
||||
with open(config_path, "w") as f:
|
||||
json.dump(payload, f)
|
||||
|
||||
loaded = TrainPipelineConfig.from_pretrained(tmp_path)
|
||||
|
||||
assert loaded.sample_weighting is not None
|
||||
assert loaded.sample_weighting.type == "rabc"
|
||||
assert loaded.sample_weighting.progress_path == "hf://datasets/user/repo/sarm_progress.parquet"
|
||||
assert loaded.sample_weighting.kappa == 0.05
|
||||
assert loaded.sample_weighting.epsilon == 1e-5
|
||||
assert loaded.sample_weighting.head_mode == "dense"
|
||||
|
||||
|
||||
def test_train_pipeline_config_from_pretrained_strips_legacy_rabc_when_disabled(tmp_path):
|
||||
"""Legacy RA-BC fields should be ignored when ``use_rabc`` was false."""
|
||||
from lerobot.configs.default import DatasetConfig
|
||||
from lerobot.configs.train import TRAIN_CONFIG_NAME, TrainPipelineConfig
|
||||
from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
|
||||
|
||||
cfg = TrainPipelineConfig(
|
||||
dataset=DatasetConfig(repo_id="user/repo"),
|
||||
policy=DiffusionConfig(device="cpu"),
|
||||
)
|
||||
cfg._save_pretrained(tmp_path)
|
||||
|
||||
config_path = tmp_path / TRAIN_CONFIG_NAME
|
||||
with open(config_path) as f:
|
||||
payload = json.load(f)
|
||||
|
||||
payload.pop("sample_weighting", None)
|
||||
payload.update(
|
||||
{
|
||||
"use_rabc": False,
|
||||
"rabc_progress_path": "hf://datasets/user/repo/sarm_progress.parquet",
|
||||
"rabc_kappa": 0.05,
|
||||
"rabc_epsilon": 1e-5,
|
||||
"rabc_head_mode": "dense",
|
||||
}
|
||||
)
|
||||
with open(config_path, "w") as f:
|
||||
json.dump(payload, f)
|
||||
|
||||
loaded = TrainPipelineConfig.from_pretrained(tmp_path)
|
||||
|
||||
assert loaded.sample_weighting is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# PreTrainedRewardModel hub upload: push_model_to_hub + generate_model_card.
|
||||
# We test the generation side (offline) fully, and the upload side with HfApi
|
||||
# mocked so nothing actually hits the network.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_dummy_reward_model(**config_kwargs):
|
||||
return _DummyHubReward(_DummyHubRewardConfig(**config_kwargs)), _DummyHubRewardConfig
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def _offline_model_card(monkeypatch):
|
||||
"""``ModelCard.validate`` does a live ``POST`` to huggingface.co — bypass it
|
||||
so tests can run offline."""
|
||||
from huggingface_hub import ModelCard
|
||||
|
||||
monkeypatch.setattr(ModelCard, "validate", lambda self, *a, **kw: None)
|
||||
|
||||
|
||||
def test_reward_model_generate_model_card_renders_expected_fields(_offline_model_card):
|
||||
"""``generate_model_card`` must produce a card with the right metadata and
|
||||
body, using the dedicated reward-model template."""
|
||||
model, _ = _make_dummy_reward_model(
|
||||
license="mit",
|
||||
tags=["robot", "sim"],
|
||||
)
|
||||
|
||||
card = model.generate_model_card(
|
||||
dataset_repo_id="user/my_dataset",
|
||||
model_type=model.config.type,
|
||||
license=model.config.license,
|
||||
tags=model.config.tags,
|
||||
)
|
||||
|
||||
# Metadata (YAML header) — ModelCardData fields.
|
||||
assert card.data.license == "mit"
|
||||
assert card.data.library_name == "lerobot"
|
||||
assert card.data.pipeline_tag == "robotics"
|
||||
assert "reward-model" in card.data.tags
|
||||
assert model.config.type in card.data.tags
|
||||
assert card.data.model_name == model.config.type
|
||||
assert card.data.datasets == "user/my_dataset"
|
||||
|
||||
# Body — specific to the reward-model template, NOT the policy one.
|
||||
body = str(card)
|
||||
assert "Reward Model Card" in body
|
||||
assert "This reward model has been trained" in body
|
||||
assert "--reward_model.type=" in body # reward-model-specific usage block
|
||||
|
||||
|
||||
def test_reward_model_generate_model_card_uses_default_license(_offline_model_card):
|
||||
"""When config.license is None the card falls back to apache-2.0."""
|
||||
model, _ = _make_dummy_reward_model()
|
||||
|
||||
card = model.generate_model_card(
|
||||
dataset_repo_id="user/my_dataset",
|
||||
model_type=model.config.type,
|
||||
license=model.config.license,
|
||||
tags=None,
|
||||
)
|
||||
|
||||
assert card.data.license == "apache-2.0"
|
||||
|
||||
|
||||
def test_reward_model_push_model_to_hub_uploads_expected_files(monkeypatch, _offline_model_card):
|
||||
"""``push_model_to_hub`` must:
|
||||
1. create the repo,
|
||||
2. assemble a temp folder with weights + config.json + train_config.json + README.md,
|
||||
3. call ``api.upload_folder`` on that folder.
|
||||
All network calls are mocked.
|
||||
"""
|
||||
from huggingface_hub.constants import CONFIG_NAME
|
||||
|
||||
from lerobot.configs.default import DatasetConfig
|
||||
from lerobot.configs.train import TRAIN_CONFIG_NAME, TrainPipelineConfig
|
||||
|
||||
model, _ = _make_dummy_reward_model(
|
||||
repo_id="user/my_reward",
|
||||
license="apache-2.0",
|
||||
)
|
||||
# Point the reward model's train config at a dummy dataset repo.
|
||||
train_cfg = TrainPipelineConfig(
|
||||
dataset=DatasetConfig(repo_id="user/my_dataset"),
|
||||
reward_model=model.config,
|
||||
)
|
||||
|
||||
uploaded: dict = {}
|
||||
fake_commit_info = SimpleNamespace(repo_url=SimpleNamespace(url="https://huggingface.co/user/my_reward"))
|
||||
|
||||
class _FakeHfApi:
|
||||
def create_repo(self, repo_id, private=None, exist_ok=False):
|
||||
uploaded["create_repo_id"] = repo_id
|
||||
uploaded["create_private"] = private
|
||||
return SimpleNamespace(repo_id=repo_id)
|
||||
|
||||
def upload_folder(self, *, repo_id, repo_type, folder_path, commit_message, **_kwargs):
|
||||
uploaded["upload_repo_id"] = repo_id
|
||||
uploaded["upload_repo_type"] = repo_type
|
||||
uploaded["commit_message"] = commit_message
|
||||
# Snapshot files assembled in the temp folder — this is the real
|
||||
# contract we care about.
|
||||
uploaded["files"] = sorted(p.name for p in Path(folder_path).iterdir())
|
||||
return fake_commit_info
|
||||
|
||||
from lerobot.rewards import pretrained as reward_pretrained
|
||||
|
||||
monkeypatch.setattr(reward_pretrained, "HfApi", lambda *a, **kw: _FakeHfApi())
|
||||
|
||||
model.push_model_to_hub(train_cfg)
|
||||
|
||||
assert uploaded["create_repo_id"] == "user/my_reward"
|
||||
assert uploaded["upload_repo_id"] == "user/my_reward"
|
||||
assert uploaded["upload_repo_type"] == "model"
|
||||
assert uploaded["commit_message"] == "Upload reward model weights, train config and readme"
|
||||
# Minimum required files that must be uploaded with a reward model.
|
||||
assert CONFIG_NAME in uploaded["files"] # config.json
|
||||
assert TRAIN_CONFIG_NAME in uploaded["files"] # train_config.json
|
||||
assert "README.md" in uploaded["files"]
|
||||
assert any(name.endswith(".safetensors") for name in uploaded["files"])
|
||||
@@ -104,8 +104,8 @@ class TestSARMEncodingProcessorStepEndToEnd:
|
||||
def mock_clip_model(self):
|
||||
"""Mock CLIP model to avoid loading real weights."""
|
||||
with (
|
||||
patch("lerobot.policies.sarm.processor_sarm.CLIPModel") as mock_model_cls,
|
||||
patch("lerobot.policies.sarm.processor_sarm.CLIPProcessor") as mock_processor_cls,
|
||||
patch("lerobot.rewards.sarm.processor_sarm.CLIPModel") as mock_model_cls,
|
||||
patch("lerobot.rewards.sarm.processor_sarm.CLIPProcessor") as mock_processor_cls,
|
||||
):
|
||||
# Mock the CLIP model - return embeddings based on input batch size
|
||||
mock_model = MagicMock()
|
||||
@@ -142,7 +142,7 @@ class TestSARMEncodingProcessorStepEndToEnd:
|
||||
@pytest.fixture
|
||||
def processor_with_mocks(self, mock_clip_model):
|
||||
"""Create a processor with mocked CLIP and dataset metadata for dual mode."""
|
||||
from lerobot.policies.sarm.processor_sarm import SARMEncodingProcessorStep
|
||||
from lerobot.rewards.sarm.processor_sarm import SARMEncodingProcessorStep
|
||||
|
||||
# Dual mode config with both sparse and dense annotations
|
||||
config = MockConfig(
|
||||
@@ -256,7 +256,7 @@ class TestSARMEncodingProcessorStepEndToEnd:
|
||||
|
||||
def test_call_with_batched_input(self, mock_clip_model):
|
||||
"""Test processor __call__ with a batched input (multiple frames) in dual mode."""
|
||||
from lerobot.policies.sarm.processor_sarm import SARMEncodingProcessorStep
|
||||
from lerobot.rewards.sarm.processor_sarm import SARMEncodingProcessorStep
|
||||
|
||||
config = MockConfig(
|
||||
n_obs_steps=8,
|
||||
@@ -332,7 +332,7 @@ class TestSARMEncodingProcessorStepEndToEnd:
|
||||
|
||||
def test_targets_increase_with_progress(self, mock_clip_model):
|
||||
"""Test that both sparse and dense targets increase as frame index progresses."""
|
||||
from lerobot.policies.sarm.processor_sarm import SARMEncodingProcessorStep
|
||||
from lerobot.rewards.sarm.processor_sarm import SARMEncodingProcessorStep
|
||||
|
||||
config = MockConfig(
|
||||
n_obs_steps=8,
|
||||
@@ -404,7 +404,7 @@ class TestSARMEncodingProcessorStepEndToEnd:
|
||||
|
||||
def test_progress_labels_exact_values(self, mock_clip_model):
|
||||
"""Test that progress labels (stage.tau) are computed correctly for known positions."""
|
||||
from lerobot.policies.sarm.processor_sarm import SARMEncodingProcessorStep
|
||||
from lerobot.rewards.sarm.processor_sarm import SARMEncodingProcessorStep
|
||||
|
||||
# Simple setup: 2 sparse stages, 4 dense stages, 100 frame episode
|
||||
config = MockConfig(
|
||||
@@ -495,7 +495,7 @@ class TestSARMEncodingProcessorStepEndToEnd:
|
||||
"""Test that rewind augmentation correctly extends sequence and generates targets."""
|
||||
import random
|
||||
|
||||
from lerobot.policies.sarm.processor_sarm import SARMEncodingProcessorStep
|
||||
from lerobot.rewards.sarm.processor_sarm import SARMEncodingProcessorStep
|
||||
|
||||
config = MockConfig(
|
||||
n_obs_steps=8,
|
||||
@@ -587,8 +587,8 @@ class TestSARMEncodingProcessorStepEndToEnd:
|
||||
|
||||
def test_full_sequence_target_consistency(self, mock_clip_model):
|
||||
"""Test that the full sequence of targets is consistent with frame positions."""
|
||||
from lerobot.policies.sarm.processor_sarm import SARMEncodingProcessorStep
|
||||
from lerobot.policies.sarm.sarm_utils import find_stage_and_tau
|
||||
from lerobot.rewards.sarm.processor_sarm import SARMEncodingProcessorStep
|
||||
from lerobot.rewards.sarm.sarm_utils import find_stage_and_tau
|
||||
|
||||
config = MockConfig(
|
||||
n_obs_steps=8,
|
||||
@@ -18,7 +18,7 @@ import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from lerobot.policies.sarm.sarm_utils import (
|
||||
from lerobot.rewards.sarm.sarm_utils import (
|
||||
apply_rewind_augmentation,
|
||||
compute_absolute_indices,
|
||||
compute_tau,
|
||||
@@ -24,10 +24,6 @@ def lerobot_train(args):
|
||||
return run_command(cmd="lerobot-train", module="lerobot_train", args=args)
|
||||
|
||||
|
||||
def lerobot_record(args):
|
||||
return run_command(cmd="lerobot-record", module="lerobot_record", args=args)
|
||||
|
||||
|
||||
def resolve_model_id_for_peft_training(policy_type):
|
||||
"""PEFT training needs pretrained models, this finds the pretrained model of a policy type for PEFT training."""
|
||||
if policy_type == "smolvla":
|
||||
@@ -155,81 +151,3 @@ def test_peft_training_params_are_fewer(policy_type, tmp_path):
|
||||
f"--output_dir={output_dir}",
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
class DummyRobot:
|
||||
name = "dummy"
|
||||
cameras = []
|
||||
action_features = {"foo": 1.0, "bar": 2.0}
|
||||
observation_features = {"obs1": 1.0, "obs2": 2.0}
|
||||
is_connected = True
|
||||
|
||||
def connect(self, *args):
|
||||
pass
|
||||
|
||||
def disconnect(self):
|
||||
pass
|
||||
|
||||
|
||||
def dummy_make_robot_from_config(*args, **kwargs):
|
||||
return DummyRobot()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("policy_type", ["smolvla"])
|
||||
@skip_if_package_missing("peft")
|
||||
def test_peft_record_loads_policy(policy_type, tmp_path):
|
||||
"""Train a policy with PEFT and attempt to load it with `lerobot-record`."""
|
||||
from peft import PeftModel
|
||||
|
||||
output_dir = tmp_path / f"output_{policy_type}"
|
||||
model_id = resolve_model_id_for_peft_training(policy_type)
|
||||
|
||||
lerobot_train(
|
||||
[
|
||||
f"--policy.path={model_id}",
|
||||
"--policy.push_to_hub=false",
|
||||
"--policy.input_features=null",
|
||||
"--policy.output_features=null",
|
||||
"--peft.method=LORA",
|
||||
"--dataset.repo_id=lerobot/pusht",
|
||||
"--dataset.episodes=[0, 1]",
|
||||
"--steps=1",
|
||||
f"--output_dir={output_dir}",
|
||||
]
|
||||
)
|
||||
|
||||
policy_dir = output_dir / "checkpoints" / "last" / "pretrained_model"
|
||||
dataset_dir = tmp_path / "eval_pusht"
|
||||
single_task = "move the table"
|
||||
loaded_policy = None
|
||||
|
||||
def dummy_record_loop(*args, **kwargs):
|
||||
nonlocal loaded_policy
|
||||
|
||||
if "dataset" not in kwargs:
|
||||
return
|
||||
|
||||
dataset = kwargs["dataset"]
|
||||
dataset.add_frame({"task": single_task})
|
||||
loaded_policy = kwargs["policy"]
|
||||
|
||||
with (
|
||||
patch("lerobot.scripts.lerobot_record.make_robot_from_config", dummy_make_robot_from_config),
|
||||
# disable record loop since we're only interested in successful loading of the policy.
|
||||
patch("lerobot.scripts.lerobot_record.record_loop", dummy_record_loop),
|
||||
# disable speech output
|
||||
patch("lerobot.utils.utils.say"),
|
||||
):
|
||||
lerobot_record(
|
||||
[
|
||||
f"--policy.path={policy_dir}",
|
||||
"--robot.type=so101_follower",
|
||||
"--robot.port=/dev/null",
|
||||
"--dataset.repo_id=lerobot/eval_pusht",
|
||||
f'--dataset.single_task="{single_task}"',
|
||||
f"--dataset.root={dataset_dir}",
|
||||
"--dataset.push_to_hub=false",
|
||||
]
|
||||
)
|
||||
|
||||
assert isinstance(loaded_policy, PeftModel)
|
||||
|
||||
@@ -21,8 +21,9 @@ import pytest
|
||||
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
|
||||
pytest.importorskip("deepdiff", reason="deepdiff is required (install lerobot[hardware])")
|
||||
|
||||
from lerobot.configs.dataset import DatasetRecordConfig
|
||||
from lerobot.scripts.lerobot_calibrate import CalibrateConfig, calibrate
|
||||
from lerobot.scripts.lerobot_record import DatasetRecordConfig, RecordConfig, record
|
||||
from lerobot.scripts.lerobot_record import RecordConfig, record
|
||||
from lerobot.scripts.lerobot_replay import DatasetReplayConfig, ReplayConfig, replay
|
||||
from lerobot.scripts.lerobot_teleoperate import TeleoperateConfig, teleoperate
|
||||
from tests.fixtures.constants import DUMMY_REPO_ID
|
||||
|
||||
@@ -0,0 +1,345 @@
|
||||
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Minimal tests for the rollout module's public API."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import dataclasses
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Import smoke tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_rollout_top_level_imports():
|
||||
import lerobot.rollout
|
||||
|
||||
for name in lerobot.rollout.__all__:
|
||||
assert hasattr(lerobot.rollout, name), f"Missing export: {name}"
|
||||
|
||||
|
||||
def test_inference_submodule_imports():
|
||||
import lerobot.rollout.inference
|
||||
|
||||
for name in lerobot.rollout.inference.__all__:
|
||||
assert hasattr(lerobot.rollout.inference, name), f"Missing export: {name}"
|
||||
|
||||
|
||||
def test_strategies_submodule_imports():
|
||||
import lerobot.rollout.strategies
|
||||
|
||||
for name in lerobot.rollout.strategies.__all__:
|
||||
assert hasattr(lerobot.rollout.strategies, name), f"Missing export: {name}"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Config tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_strategy_config_types():
|
||||
from lerobot.rollout import (
|
||||
BaseStrategyConfig,
|
||||
DAggerStrategyConfig,
|
||||
HighlightStrategyConfig,
|
||||
SentryStrategyConfig,
|
||||
)
|
||||
|
||||
assert BaseStrategyConfig().type == "base"
|
||||
assert SentryStrategyConfig().type == "sentry"
|
||||
assert HighlightStrategyConfig().type == "highlight"
|
||||
assert DAggerStrategyConfig().type == "dagger"
|
||||
|
||||
|
||||
def test_dagger_config_invalid_input_device():
|
||||
from lerobot.rollout import DAggerStrategyConfig
|
||||
|
||||
with pytest.raises(ValueError, match="input_device must be 'keyboard' or 'pedal'"):
|
||||
DAggerStrategyConfig(input_device="joystick")
|
||||
|
||||
|
||||
def test_dagger_config_defaults():
|
||||
from lerobot.rollout import DAggerStrategyConfig
|
||||
|
||||
cfg = DAggerStrategyConfig()
|
||||
assert cfg.num_episodes is None
|
||||
assert cfg.record_autonomous is False
|
||||
assert cfg.input_device == "keyboard"
|
||||
|
||||
|
||||
def test_inference_config_types():
|
||||
from lerobot.rollout import RTCInferenceConfig, SyncInferenceConfig
|
||||
|
||||
assert SyncInferenceConfig().type == "sync"
|
||||
|
||||
rtc = RTCInferenceConfig()
|
||||
assert rtc.type == "rtc"
|
||||
assert rtc.queue_threshold == 30
|
||||
assert rtc.rtc is not None
|
||||
|
||||
|
||||
def test_sentry_config_defaults():
|
||||
from lerobot.rollout import SentryStrategyConfig
|
||||
|
||||
cfg = SentryStrategyConfig()
|
||||
assert cfg.upload_every_n_episodes == 5
|
||||
assert cfg.target_video_file_size_mb is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# RolloutRingBuffer
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_ring_buffer_append_and_eviction():
|
||||
from lerobot.rollout.ring_buffer import RolloutRingBuffer
|
||||
|
||||
buf = RolloutRingBuffer(max_seconds=0.5, max_memory_mb=100.0, fps=10.0)
|
||||
# max_frames = 5
|
||||
for i in range(8):
|
||||
buf.append({"val": i})
|
||||
assert len(buf) == 5
|
||||
|
||||
|
||||
def test_ring_buffer_drain():
|
||||
from lerobot.rollout.ring_buffer import RolloutRingBuffer
|
||||
|
||||
buf = RolloutRingBuffer(max_seconds=1.0, max_memory_mb=100.0, fps=10.0)
|
||||
for i in range(3):
|
||||
buf.append({"val": i})
|
||||
frames = buf.drain()
|
||||
assert len(frames) == 3
|
||||
assert len(buf) == 0
|
||||
assert buf.estimated_bytes == 0
|
||||
|
||||
|
||||
def test_ring_buffer_clear():
|
||||
from lerobot.rollout.ring_buffer import RolloutRingBuffer
|
||||
|
||||
buf = RolloutRingBuffer(max_seconds=1.0, max_memory_mb=100.0, fps=10.0)
|
||||
buf.append({"val": 1})
|
||||
buf.clear()
|
||||
assert len(buf) == 0
|
||||
assert buf.estimated_bytes == 0
|
||||
|
||||
|
||||
def test_ring_buffer_tensor_bytes():
|
||||
from lerobot.rollout.ring_buffer import RolloutRingBuffer
|
||||
|
||||
buf = RolloutRingBuffer(max_seconds=1.0, max_memory_mb=100.0, fps=10.0)
|
||||
t = torch.zeros(100, dtype=torch.float32) # 400 bytes
|
||||
buf.append({"tensor": t})
|
||||
assert buf.estimated_bytes >= 400
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ThreadSafeRobot
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_thread_safe_robot_delegates():
|
||||
from lerobot.rollout.robot_wrapper import ThreadSafeRobot
|
||||
from tests.mocks.mock_robot import MockRobot, MockRobotConfig
|
||||
|
||||
robot = MockRobot(MockRobotConfig(n_motors=3))
|
||||
robot.connect()
|
||||
wrapper = ThreadSafeRobot(robot)
|
||||
|
||||
obs = wrapper.get_observation()
|
||||
assert "motor_1.pos" in obs
|
||||
assert "motor_2.pos" in obs
|
||||
assert "motor_3.pos" in obs
|
||||
|
||||
action = {"motor_1.pos": 0.0, "motor_2.pos": 1.0, "motor_3.pos": 2.0}
|
||||
result = wrapper.send_action(action)
|
||||
assert result == action
|
||||
|
||||
robot.disconnect()
|
||||
|
||||
|
||||
def test_thread_safe_robot_properties():
|
||||
from lerobot.rollout.robot_wrapper import ThreadSafeRobot
|
||||
from tests.mocks.mock_robot import MockRobot, MockRobotConfig
|
||||
|
||||
robot = MockRobot(MockRobotConfig(n_motors=3))
|
||||
robot.connect()
|
||||
wrapper = ThreadSafeRobot(robot)
|
||||
|
||||
assert wrapper.name == "mock_robot"
|
||||
assert "motor_1.pos" in wrapper.observation_features
|
||||
assert "motor_1.pos" in wrapper.action_features
|
||||
assert wrapper.is_connected is True
|
||||
assert wrapper.inner is robot
|
||||
|
||||
robot.disconnect()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Strategy factory
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_create_strategy_dispatches():
|
||||
from lerobot.rollout import (
|
||||
BaseStrategy,
|
||||
BaseStrategyConfig,
|
||||
DAggerStrategy,
|
||||
DAggerStrategyConfig,
|
||||
SentryStrategy,
|
||||
SentryStrategyConfig,
|
||||
create_strategy,
|
||||
)
|
||||
|
||||
assert isinstance(create_strategy(BaseStrategyConfig()), BaseStrategy)
|
||||
assert isinstance(create_strategy(SentryStrategyConfig()), SentryStrategy)
|
||||
assert isinstance(create_strategy(DAggerStrategyConfig()), DAggerStrategy)
|
||||
|
||||
|
||||
def test_create_strategy_unknown_raises():
|
||||
from lerobot.rollout import create_strategy
|
||||
|
||||
cfg = MagicMock()
|
||||
cfg.type = "bogus"
|
||||
with pytest.raises(ValueError, match="Unknown strategy type"):
|
||||
create_strategy(cfg)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Inference factory
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_create_inference_engine_sync():
|
||||
from lerobot.rollout import SyncInferenceConfig, SyncInferenceEngine, create_inference_engine
|
||||
|
||||
engine = create_inference_engine(
|
||||
SyncInferenceConfig(),
|
||||
policy=MagicMock(),
|
||||
preprocessor=MagicMock(),
|
||||
postprocessor=MagicMock(),
|
||||
robot_wrapper=MagicMock(robot_type="mock"),
|
||||
hw_features={},
|
||||
dataset_features={},
|
||||
ordered_action_keys=["k"],
|
||||
task="test",
|
||||
fps=30.0,
|
||||
device="cpu",
|
||||
)
|
||||
assert isinstance(engine, SyncInferenceEngine)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Pure functions
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_estimate_max_episode_seconds_no_video():
|
||||
from lerobot.rollout.strategies import estimate_max_episode_seconds
|
||||
|
||||
assert estimate_max_episode_seconds({}, fps=30.0) == 300.0
|
||||
|
||||
|
||||
def test_estimate_max_episode_seconds_with_video():
|
||||
from lerobot.rollout.strategies import estimate_max_episode_seconds
|
||||
|
||||
features = {"cam": {"dtype": "video", "shape": (480, 640, 3)}}
|
||||
result = estimate_max_episode_seconds(features, fps=30.0)
|
||||
assert result > 0
|
||||
# With a real camera, duration should differ from the fallback
|
||||
assert result != 300.0
|
||||
|
||||
|
||||
def test_safe_push_to_hub():
|
||||
from lerobot.rollout.strategies import safe_push_to_hub
|
||||
|
||||
ds = MagicMock()
|
||||
ds.num_episodes = 0
|
||||
assert safe_push_to_hub(ds) is False
|
||||
ds.push_to_hub.assert_not_called()
|
||||
|
||||
ds.num_episodes = 5
|
||||
assert safe_push_to_hub(ds, tags=["test"]) is True
|
||||
ds.push_to_hub.assert_called_once_with(tags=["test"], private=False)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# DAgger state machine
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_dagger_full_transition_cycle():
|
||||
from lerobot.rollout.strategies import DAggerEvents, DAggerPhase
|
||||
|
||||
events = DAggerEvents()
|
||||
assert events.phase == DAggerPhase.AUTONOMOUS
|
||||
|
||||
# AUTONOMOUS -> PAUSED
|
||||
events.request_transition("pause_resume")
|
||||
old, new = events.consume_transition()
|
||||
assert (old, new) == (DAggerPhase.AUTONOMOUS, DAggerPhase.PAUSED)
|
||||
|
||||
# PAUSED -> CORRECTING
|
||||
events.request_transition("correction")
|
||||
old, new = events.consume_transition()
|
||||
assert (old, new) == (DAggerPhase.PAUSED, DAggerPhase.CORRECTING)
|
||||
|
||||
# CORRECTING -> PAUSED
|
||||
events.request_transition("correction")
|
||||
old, new = events.consume_transition()
|
||||
assert (old, new) == (DAggerPhase.CORRECTING, DAggerPhase.PAUSED)
|
||||
|
||||
# PAUSED -> AUTONOMOUS
|
||||
events.request_transition("pause_resume")
|
||||
old, new = events.consume_transition()
|
||||
assert (old, new) == (DAggerPhase.PAUSED, DAggerPhase.AUTONOMOUS)
|
||||
|
||||
|
||||
def test_dagger_invalid_transition_ignored():
|
||||
from lerobot.rollout.strategies import DAggerEvents, DAggerPhase
|
||||
|
||||
events = DAggerEvents()
|
||||
events.request_transition("correction") # Not valid from AUTONOMOUS
|
||||
assert events.consume_transition() is None
|
||||
assert events.phase == DAggerPhase.AUTONOMOUS
|
||||
|
||||
|
||||
def test_dagger_events_reset():
|
||||
from lerobot.rollout.strategies import DAggerEvents, DAggerPhase
|
||||
|
||||
events = DAggerEvents()
|
||||
events.request_transition("pause_resume")
|
||||
events.consume_transition() # -> PAUSED
|
||||
events.upload_requested.set()
|
||||
events.reset()
|
||||
assert events.phase == DAggerPhase.AUTONOMOUS
|
||||
assert not events.upload_requested.is_set()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Context dataclass
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_rollout_context_fields():
|
||||
from lerobot.rollout import RolloutContext
|
||||
|
||||
field_names = {f.name for f in dataclasses.fields(RolloutContext)}
|
||||
assert field_names == {"runtime", "hardware", "policy", "processors", "data"}
|
||||
@@ -1,8 +1,12 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
import torch
|
||||
import pytest
|
||||
|
||||
from lerobot.utils.collate import lerobot_collate_fn
|
||||
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():
|
||||
@@ -34,3 +38,47 @@ def test_lerobot_collate_preserves_messages_and_drops_raw_language():
|
||||
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
|
||||
|
||||
@@ -24,7 +24,7 @@ import pytest
|
||||
|
||||
pytest.importorskip("grpc")
|
||||
|
||||
from lerobot.rl.process import ProcessSignalHandler # noqa: E402
|
||||
from lerobot.utils.process import ProcessSignalHandler # noqa: E402
|
||||
|
||||
|
||||
# Fixture to reset shutdown_event_counter and original signal handlers before and after each test
|
||||
|
||||
@@ -0,0 +1,401 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Tests for the sample weighting infrastructure."""
|
||||
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("pandas", reason="pandas is required (install lerobot[dataset])")
|
||||
|
||||
import torch
|
||||
|
||||
from lerobot.utils.sample_weighting import (
|
||||
SampleWeighter,
|
||||
SampleWeightingConfig,
|
||||
UniformWeighter,
|
||||
make_sample_weighter,
|
||||
)
|
||||
|
||||
# =============================================================================
|
||||
# Fixtures
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_progress_parquet(tmp_path):
|
||||
"""Create a sample progress parquet file for testing."""
|
||||
import pandas as pd
|
||||
|
||||
# Create sample progress data for 2 episodes with 10 frames each
|
||||
data = {
|
||||
"index": list(range(20)),
|
||||
"episode_index": [0] * 10 + [1] * 10,
|
||||
"frame_index": list(range(10)) * 2,
|
||||
"progress_sparse": [i / 10.0 for i in range(10)] * 2,
|
||||
}
|
||||
df = pd.DataFrame(data)
|
||||
parquet_path = tmp_path / "sarm_progress.parquet"
|
||||
df.to_parquet(parquet_path)
|
||||
return parquet_path
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# SampleWeightingConfig Tests
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def test_config_default_values():
|
||||
"""Test default configuration values."""
|
||||
config = SampleWeightingConfig()
|
||||
assert config.type == "rabc"
|
||||
assert config.progress_path is None
|
||||
assert config.head_mode == "sparse"
|
||||
assert config.kappa == 0.01
|
||||
assert config.epsilon == 1e-6
|
||||
assert config.extra_params == {}
|
||||
|
||||
|
||||
def test_config_custom_values():
|
||||
"""Test configuration with custom values."""
|
||||
config = SampleWeightingConfig(
|
||||
type="rabc",
|
||||
progress_path="/path/to/progress.parquet",
|
||||
head_mode="dense",
|
||||
kappa=0.05,
|
||||
epsilon=1e-8,
|
||||
extra_params={"fallback_weight": 0.5},
|
||||
)
|
||||
assert config.type == "rabc"
|
||||
assert config.progress_path == "/path/to/progress.parquet"
|
||||
assert config.head_mode == "dense"
|
||||
assert config.kappa == 0.05
|
||||
assert config.epsilon == 1e-8
|
||||
assert config.extra_params == {"fallback_weight": 0.5}
|
||||
|
||||
|
||||
def test_config_uniform_type():
|
||||
"""Test configuration for uniform weighting."""
|
||||
config = SampleWeightingConfig(type="uniform")
|
||||
assert config.type == "uniform"
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# UniformWeighter Tests
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def test_uniform_weighter_inherits_from_sample_weighter():
|
||||
"""Test that UniformWeighter is a SampleWeighter."""
|
||||
weighter = UniformWeighter(device=torch.device("cpu"))
|
||||
assert isinstance(weighter, SampleWeighter)
|
||||
|
||||
|
||||
def test_uniform_weighter_compute_batch_weights_with_action_key():
|
||||
"""Test weight computation with 'action' key in batch."""
|
||||
weighter = UniformWeighter(device=torch.device("cpu"))
|
||||
batch = {"action": torch.randn(8, 10)}
|
||||
|
||||
weights, stats = weighter.compute_batch_weights(batch)
|
||||
|
||||
assert weights.shape == (8,)
|
||||
assert torch.allclose(weights, torch.ones(8))
|
||||
assert stats["mean_weight"] == 1.0
|
||||
assert stats["type"] == "uniform"
|
||||
|
||||
|
||||
def test_uniform_weighter_compute_batch_weights_with_index_key():
|
||||
"""Test weight computation with 'index' key in batch."""
|
||||
weighter = UniformWeighter(device=torch.device("cpu"))
|
||||
batch = {"index": torch.arange(16)}
|
||||
|
||||
weights, stats = weighter.compute_batch_weights(batch)
|
||||
|
||||
assert weights.shape == (16,)
|
||||
assert torch.allclose(weights, torch.ones(16))
|
||||
|
||||
|
||||
def test_uniform_weighter_compute_batch_weights_no_tensor_keys():
|
||||
"""Test weight computation with no tensor keys (fallback to size 1)."""
|
||||
weighter = UniformWeighter(device=torch.device("cpu"))
|
||||
batch = {"other_key": "some_value"}
|
||||
|
||||
weights, stats = weighter.compute_batch_weights(batch)
|
||||
|
||||
assert weights.shape == (1,)
|
||||
assert torch.allclose(weights, torch.ones(1))
|
||||
|
||||
|
||||
def test_uniform_weighter_compute_batch_weights_empty_batch_raises():
|
||||
"""Test that empty batch raises ValueError."""
|
||||
weighter = UniformWeighter(device=torch.device("cpu"))
|
||||
batch = {}
|
||||
|
||||
with pytest.raises(ValueError, match="empty batch"):
|
||||
weighter.compute_batch_weights(batch)
|
||||
|
||||
|
||||
def test_uniform_weighter_compute_batch_weights_scans_all_keys():
|
||||
"""Test that batch size is determined by scanning all tensor values."""
|
||||
weighter = UniformWeighter(device=torch.device("cpu"))
|
||||
# Batch with non-standard key containing a tensor
|
||||
batch = {"custom_tensor": torch.randn(7, 3)}
|
||||
|
||||
weights, stats = weighter.compute_batch_weights(batch)
|
||||
|
||||
assert weights.shape == (7,)
|
||||
assert torch.allclose(weights, torch.ones(7))
|
||||
|
||||
|
||||
def test_uniform_weighter_compute_batch_weights_on_cuda():
|
||||
"""Test that weights are placed on the correct device."""
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("CUDA not available")
|
||||
|
||||
weighter = UniformWeighter(device=torch.device("cuda"))
|
||||
batch = {"action": torch.randn(4, 10)}
|
||||
|
||||
weights, _ = weighter.compute_batch_weights(batch)
|
||||
|
||||
assert weights.device.type == "cuda"
|
||||
|
||||
|
||||
def test_uniform_weighter_get_stats():
|
||||
"""Test get_stats returns expected structure."""
|
||||
weighter = UniformWeighter(device=torch.device("cpu"))
|
||||
stats = weighter.get_stats()
|
||||
|
||||
assert stats == {"type": "uniform"}
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# make_sample_weighter Factory Tests
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def test_factory_returns_none_for_none_config():
|
||||
"""Test that None config returns None weighter."""
|
||||
policy = Mock()
|
||||
device = torch.device("cpu")
|
||||
|
||||
result = make_sample_weighter(None, policy, device)
|
||||
|
||||
assert result is None
|
||||
|
||||
|
||||
def test_factory_creates_uniform_weighter():
|
||||
"""Test creation of UniformWeighter."""
|
||||
config = SampleWeightingConfig(type="uniform")
|
||||
policy = Mock()
|
||||
device = torch.device("cpu")
|
||||
|
||||
weighter = make_sample_weighter(config, policy, device)
|
||||
|
||||
assert isinstance(weighter, UniformWeighter)
|
||||
assert isinstance(weighter, SampleWeighter)
|
||||
|
||||
|
||||
def test_factory_raises_for_unknown_type():
|
||||
"""Test that unknown type raises ValueError."""
|
||||
config = SampleWeightingConfig(type="unknown_type")
|
||||
policy = Mock()
|
||||
device = torch.device("cpu")
|
||||
|
||||
with pytest.raises(ValueError, match="Unknown sample weighting type"):
|
||||
make_sample_weighter(config, policy, device)
|
||||
|
||||
|
||||
def test_factory_rabc_requires_chunk_size():
|
||||
"""Test that RABC weighter requires chunk_size in policy config."""
|
||||
config = SampleWeightingConfig(
|
||||
type="rabc",
|
||||
progress_path="/path/to/progress.parquet",
|
||||
)
|
||||
policy = Mock()
|
||||
policy.config = Mock()
|
||||
policy.config.chunk_size = None # No chunk_size
|
||||
device = torch.device("cpu")
|
||||
|
||||
with pytest.raises(ValueError, match="chunk_size"):
|
||||
make_sample_weighter(config, policy, device)
|
||||
|
||||
|
||||
def test_factory_rabc_requires_progress_path_or_dataset_info():
|
||||
"""Test that RABC weighter requires progress_path or dataset info for auto-detection."""
|
||||
config = SampleWeightingConfig(
|
||||
type="rabc",
|
||||
progress_path=None, # No progress path
|
||||
)
|
||||
policy = Mock()
|
||||
policy.config = Mock()
|
||||
policy.config.chunk_size = 50
|
||||
device = torch.device("cpu")
|
||||
|
||||
# Should fail when no progress_path AND no dataset info
|
||||
with pytest.raises(ValueError, match="progress_path"):
|
||||
make_sample_weighter(config, policy, device)
|
||||
|
||||
|
||||
def test_factory_rabc_auto_detects_from_dataset_root(sample_progress_parquet):
|
||||
"""Test that RABC weighter auto-detects progress_path from dataset_root."""
|
||||
config = SampleWeightingConfig(
|
||||
type="rabc",
|
||||
progress_path=None, # Not provided, should auto-detect
|
||||
)
|
||||
policy = Mock()
|
||||
policy.config = Mock()
|
||||
policy.config.chunk_size = 5
|
||||
device = torch.device("cpu")
|
||||
|
||||
# The parquet file is at sample_progress_parquet, get its parent directory
|
||||
dataset_root = sample_progress_parquet.parent
|
||||
weighter = make_sample_weighter(
|
||||
config,
|
||||
policy,
|
||||
device,
|
||||
dataset_root=str(dataset_root),
|
||||
)
|
||||
|
||||
assert weighter is not None
|
||||
from lerobot.rewards.sarm.rabc import RABCWeights
|
||||
|
||||
assert isinstance(weighter, RABCWeights)
|
||||
|
||||
|
||||
def test_factory_rabc_auto_detects_from_repo_id():
|
||||
"""Test that RABC weighter constructs HF path from repo_id."""
|
||||
config = SampleWeightingConfig(
|
||||
type="rabc",
|
||||
progress_path=None, # Not provided, should auto-detect
|
||||
)
|
||||
policy = Mock()
|
||||
policy.config = Mock()
|
||||
policy.config.chunk_size = 50
|
||||
device = torch.device("cpu")
|
||||
|
||||
# This will construct the path but fail when trying to load (file doesn't exist)
|
||||
# We just verify it doesn't raise the "progress_path required" error
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
make_sample_weighter(
|
||||
config,
|
||||
policy,
|
||||
device,
|
||||
dataset_repo_id="test-user/test-dataset",
|
||||
)
|
||||
# Should NOT be the "progress_path required" error - it should try to load the file
|
||||
assert (
|
||||
"progress_path" not in str(exc_info.value).lower() or "auto-detection" in str(exc_info.value).lower()
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Integration Tests with RABCWeights
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def test_rabc_weights_is_sample_weighter(sample_progress_parquet):
|
||||
"""Test that RABCWeights inherits from SampleWeighter."""
|
||||
from lerobot.rewards.sarm.rabc import RABCWeights
|
||||
|
||||
weighter = RABCWeights(
|
||||
progress_path=sample_progress_parquet,
|
||||
chunk_size=5,
|
||||
head_mode="sparse",
|
||||
)
|
||||
assert isinstance(weighter, SampleWeighter)
|
||||
|
||||
|
||||
def test_rabc_compute_batch_weights(sample_progress_parquet):
|
||||
"""Test RABCWeights.compute_batch_weights returns correct structure."""
|
||||
from lerobot.rewards.sarm.rabc import RABCWeights
|
||||
|
||||
weighter = RABCWeights(
|
||||
progress_path=sample_progress_parquet,
|
||||
chunk_size=5,
|
||||
head_mode="sparse",
|
||||
device=torch.device("cpu"),
|
||||
)
|
||||
|
||||
batch = {"index": torch.tensor([0, 1, 2, 3])}
|
||||
weights, stats = weighter.compute_batch_weights(batch)
|
||||
|
||||
assert isinstance(weights, torch.Tensor)
|
||||
assert weights.shape == (4,)
|
||||
assert isinstance(stats, dict)
|
||||
assert "mean_weight" in stats
|
||||
|
||||
|
||||
def test_rabc_get_stats(sample_progress_parquet):
|
||||
"""Test RABCWeights.get_stats returns expected structure."""
|
||||
from lerobot.rewards.sarm.rabc import RABCWeights
|
||||
|
||||
weighter = RABCWeights(
|
||||
progress_path=sample_progress_parquet,
|
||||
chunk_size=5,
|
||||
head_mode="sparse",
|
||||
)
|
||||
|
||||
stats = weighter.get_stats()
|
||||
|
||||
assert stats["type"] == "rabc"
|
||||
assert "num_frames" in stats
|
||||
assert "chunk_size" in stats
|
||||
assert stats["chunk_size"] == 5
|
||||
assert "head_mode" in stats
|
||||
assert stats["head_mode"] == "sparse"
|
||||
assert "delta_mean" in stats
|
||||
assert "delta_std" in stats
|
||||
|
||||
|
||||
def test_factory_creates_rabc_weighter(sample_progress_parquet):
|
||||
"""Test factory creates RABCWeights with valid config."""
|
||||
from lerobot.rewards.sarm.rabc import RABCWeights
|
||||
|
||||
config = SampleWeightingConfig(
|
||||
type="rabc",
|
||||
progress_path=str(sample_progress_parquet),
|
||||
head_mode="sparse",
|
||||
kappa=0.01,
|
||||
)
|
||||
policy = Mock()
|
||||
policy.config = Mock()
|
||||
policy.config.chunk_size = 5
|
||||
device = torch.device("cpu")
|
||||
|
||||
weighter = make_sample_weighter(config, policy, device)
|
||||
|
||||
assert isinstance(weighter, RABCWeights)
|
||||
assert isinstance(weighter, SampleWeighter)
|
||||
|
||||
|
||||
def test_rabc_weights_normalization(sample_progress_parquet):
|
||||
"""Test that RABCWeights normalizes weights to sum to batch_size."""
|
||||
from lerobot.rewards.sarm.rabc import RABCWeights
|
||||
|
||||
weighter = RABCWeights(
|
||||
progress_path=sample_progress_parquet,
|
||||
chunk_size=5,
|
||||
head_mode="sparse",
|
||||
device=torch.device("cpu"),
|
||||
)
|
||||
|
||||
batch = {"index": torch.tensor([0, 1, 2, 3])}
|
||||
weights, _ = weighter.compute_batch_weights(batch)
|
||||
|
||||
# Weights should be normalized to sum approximately to batch_size
|
||||
batch_size = 4
|
||||
assert abs(weights.sum().item() - batch_size) < 0.1
|
||||
Reference in New Issue
Block a user