admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-08-08 17:39:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(sarm): fail fast on unusable episode annotations (#4306)

Browse Source 
* fix(sarm): warn when dense/sparse targets silently collapse to all-zero

In dense_only/dual modes, if meta/episodes/*.parquet has no usable
subtask columns (column absent or NaN), _load_episode_annotations
returns None and find_stage_and_tau yields stage 0 / tau 0 for every
frame. Training "succeeds" but the head silently learns to predict 0
everywhere, with no warning. This complements #2880 (which restored
loading of episodes_df): there the DataFrame is loaded but the
*_subtask_names column is missing/NaN.

Add a one-time validation at processor construction that logs a clear
warning (all episodes missing -> predict-all-zero; some missing ->
partial). Purely additive logging, no change to training math.

Closes #3842

* fix(sarm): fail fast on unusable episode annotations

---------

Co-authored-by: 1thanShih <Smartshithan1620.en12@nycu.edu.tw>
This commit is contained in:
Khalil Meftah
2026-08-03 13:30:00 +02:00
committed by GitHub
parent 3e37269dc6
commit f1efa588b8
2 changed files with 222 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/rewards/test_sarm_processor.py
+137
View File
@@ -692,3 +692,140 @@ class TestSARMEncodingProcessorStepEndToEnd:
assert abs(actual_dense - expected_dense) < 0.01, (
f"Frame {frame}: dense mismatch {actual_dense:.3f} vs expected {expected_dense:.3f}"
)
def test_rejects_missing_dense_annotation_columns(self, mock_clip_model):
"""A multi-stage dense head must reject metadata with no annotation columns."""
from lerobot.rewards.sarm.processor_sarm import SARMEncodingProcessorStep
config = MockConfig(
annotation_mode="dense_only",
dense_subtask_names=["d1", "d2", "d3", "d4"],
dense_temporal_proportions=[0.25, 0.25, 0.25, 0.25],
)
# episodes metadata WITHOUT any dense_subtask_* columns
episodes = [
{"dataset_from_index": 0, "dataset_to_index": 100, "task": "t"},
{"dataset_from_index": 100, "dataset_to_index": 200, "task": "t"},
]
dataset_meta = MockDatasetMeta(episodes)
with pytest.raises(ValueError, match="Training would produce all-zero targets"):
SARMEncodingProcessorStep(config=config, dataset_meta=dataset_meta)
def test_rejects_dense_annotations_when_all_null(self, mock_clip_model):
"""Present-but-null annotation columns must also fail before training."""
from lerobot.rewards.sarm.processor_sarm import SARMEncodingProcessorStep
config = MockConfig(
annotation_mode="dense_only",
dense_subtask_names=["d1", "d2", "d3", "d4"],
dense_temporal_proportions=[0.25, 0.25, 0.25, 0.25],
)
episodes = [
{
"dataset_from_index": 0,
"dataset_to_index": 100,
"task": "t",
"dense_subtask_names": None,
"dense_subtask_start_frames": None,
"dense_subtask_end_frames": None,
},
{
"dataset_from_index": 100,
"dataset_to_index": 200,
"task": "t",
"dense_subtask_names": None,
"dense_subtask_start_frames": None,
"dense_subtask_end_frames": None,
},
]
dataset_meta = MockDatasetMeta(episodes)
with pytest.raises(ValueError, match="none of the 2 episodes have usable annotations"):
SARMEncodingProcessorStep(config=config, dataset_meta=dataset_meta)
def test_rejects_dense_annotations_with_missing_frame_column(self, mock_clip_model):
"""Names alone are not usable when a required frame-boundary column is absent."""
from lerobot.rewards.sarm.processor_sarm import SARMEncodingProcessorStep
config = MockConfig(
annotation_mode="dense_only",
dense_subtask_names=["d1", "d2"],
dense_temporal_proportions=[0.5, 0.5],
)
episodes = [
{
"dataset_from_index": 0,
"dataset_to_index": 100,
"task": "t",
"dense_subtask_names": ["d1", "d2"],
"dense_subtask_start_frames": [0, 50],
}
]
with pytest.raises(ValueError, match="Missing required columns: subtask_end_frames"):
SARMEncodingProcessorStep(config=config, dataset_meta=MockDatasetMeta(episodes))
def test_warns_when_dense_annotations_are_partial(self, mock_clip_model, caplog):
"""Partially annotated datasets remain supported but report exact coverage."""
import logging
from lerobot.rewards.sarm.processor_sarm import SARMEncodingProcessorStep
config = MockConfig(
annotation_mode="dense_only",
dense_subtask_names=["d1", "d2"],
dense_temporal_proportions=[0.5, 0.5],
)
episodes = [
{
"dataset_from_index": 0,
"dataset_to_index": 100,
"task": "t",
"dense_subtask_names": ["d1", "d2"],
"dense_subtask_start_frames": [0, 50],
"dense_subtask_end_frames": [49, 99],
},
{
"dataset_from_index": 100,
"dataset_to_index": 200,
"task": "t",
"dense_subtask_names": None,
"dense_subtask_start_frames": None,
"dense_subtask_end_frames": None,
},
]
with caplog.at_level(logging.WARNING, logger="lerobot.rewards.sarm.processor_sarm"):
SARMEncodingProcessorStep(config=config, dataset_meta=MockDatasetMeta(episodes))
assert "1/2 episodes have unusable annotations" in caplog.text
assert "only the 1 annotated episodes will train the head" in caplog.text
def test_no_warning_when_dense_annotations_present(self, mock_clip_model, caplog):
"""A fully annotated dataset must not emit an annotation-coverage warning."""
import logging
from lerobot.rewards.sarm.processor_sarm import SARMEncodingProcessorStep
config = MockConfig(
annotation_mode="dense_only",
dense_subtask_names=["d1", "d2", "d3", "d4"],
dense_temporal_proportions=[0.25, 0.25, 0.25, 0.25],
)
episodes = [
{
"dataset_from_index": 0,
"dataset_to_index": 100,
"task": "t",
"dense_subtask_names": ["d1", "d2", "d3", "d4"],
"dense_subtask_start_frames": [0, 25, 50, 75],
"dense_subtask_end_frames": [25, 50, 75, 100],
}
]
dataset_meta = MockDatasetMeta(episodes)
with caplog.at_level(logging.WARNING, logger="lerobot.rewards.sarm.processor_sarm"):
SARMEncodingProcessorStep(config=config, dataset_meta=dataset_meta)
assert not any("unusable annotations" in m for m in caplog.messages)