mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-28 12:15:59 +00:00
feat(data): add recipe-driven language supervision
This commit is contained in:
@@ -12,7 +12,9 @@ from lerobot.processor.render_messages_processor import RenderMessagesStep # no
|
||||
from lerobot.types import TransitionKey # noqa: E402
|
||||
|
||||
|
||||
def test_render_messages_step_noops_without_language_columns():
|
||||
def test_render_messages_step_renders_task_fallback_without_language_columns():
|
||||
"""No language columns + a task string → low-level task fallback render,
|
||||
matching what the policy sees at eval time on unannotated observations."""
|
||||
recipe = TrainingRecipe(
|
||||
messages=[
|
||||
MessageTurn(role="user", content="${task}", stream="high_level"),
|
||||
@@ -21,6 +23,24 @@ def test_render_messages_step_noops_without_language_columns():
|
||||
)
|
||||
transition = create_transition(complementary_data={"task": "do it"})
|
||||
|
||||
out = RenderMessagesStep(recipe)(transition)
|
||||
data = out[TransitionKey.COMPLEMENTARY_DATA]
|
||||
|
||||
assert data["messages"] == [{"role": "user", "content": "do it"}]
|
||||
assert data["message_streams"] == ["low_level"]
|
||||
assert data["target_message_indices"] == []
|
||||
assert data["task"] == "do it"
|
||||
|
||||
|
||||
def test_render_messages_step_noops_without_language_columns_or_task():
|
||||
recipe = TrainingRecipe(
|
||||
messages=[
|
||||
MessageTurn(role="user", content="${task}", stream="high_level"),
|
||||
MessageTurn(role="assistant", content="${subtask}", stream="low_level", target=True),
|
||||
]
|
||||
)
|
||||
transition = create_transition(complementary_data={})
|
||||
|
||||
assert RenderMessagesStep(recipe)(transition) == transition
|
||||
|
||||
|
||||
@@ -58,3 +78,70 @@ def test_render_messages_step_renders_and_drops_raw_language():
|
||||
assert data["messages"][-1]["content"] == "reach carefully"
|
||||
assert data["message_streams"] == ["high_level", "low_level"]
|
||||
assert data["target_message_indices"] == [1]
|
||||
|
||||
|
||||
def test_render_messages_step_falls_back_to_low_level_task_when_recipe_misses():
|
||||
recipe = TrainingRecipe(
|
||||
messages=[
|
||||
MessageTurn(
|
||||
role="assistant",
|
||||
content="${subtask}",
|
||||
stream="high_level",
|
||||
target=True,
|
||||
if_present="subtask",
|
||||
),
|
||||
]
|
||||
)
|
||||
transition = create_transition(
|
||||
complementary_data={
|
||||
"task": "pick the cube",
|
||||
"timestamp": torch.tensor(0.0),
|
||||
"index": torch.tensor(7),
|
||||
"language_persistent": [],
|
||||
"language_events": [{"style": "unmatched", "timestamp": 0.0}],
|
||||
}
|
||||
)
|
||||
|
||||
out = RenderMessagesStep(recipe)(transition)
|
||||
data = out[TransitionKey.COMPLEMENTARY_DATA]
|
||||
|
||||
assert data["messages"] == [{"role": "user", "content": "pick the cube"}]
|
||||
assert data["message_streams"] == ["low_level"]
|
||||
assert data["target_message_indices"] == []
|
||||
|
||||
|
||||
def test_render_messages_step_falls_back_per_sample_in_batched_language():
|
||||
recipe = TrainingRecipe(
|
||||
messages=[
|
||||
MessageTurn(
|
||||
role="assistant",
|
||||
content="${subtask}",
|
||||
stream="high_level",
|
||||
target=True,
|
||||
if_present="subtask",
|
||||
),
|
||||
]
|
||||
)
|
||||
transition = create_transition(
|
||||
action=torch.arange(4).reshape(2, 2),
|
||||
complementary_data={
|
||||
"task": ["pick the cube", "open the drawer"],
|
||||
"timestamp": torch.tensor([0.0, 1.0]),
|
||||
"index": torch.tensor([7, 8]),
|
||||
"language_persistent": [[], []],
|
||||
"language_events": [
|
||||
[{"style": "unmatched", "timestamp": 0.0}],
|
||||
[{"style": "unmatched", "timestamp": 1.0}],
|
||||
],
|
||||
},
|
||||
)
|
||||
|
||||
out = RenderMessagesStep(recipe)(transition)
|
||||
data = out[TransitionKey.COMPLEMENTARY_DATA]
|
||||
|
||||
assert data["messages"] == [
|
||||
[{"role": "user", "content": "pick the cube"}],
|
||||
[{"role": "user", "content": "open the drawer"}],
|
||||
]
|
||||
assert data["message_streams"] == [["low_level"], ["low_level"]]
|
||||
assert data["target_message_indices"] == [[], []]
|
||||
|
||||
@@ -25,7 +25,7 @@ import pytest
|
||||
import torch
|
||||
|
||||
from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature
|
||||
from lerobot.processor import DataProcessorPipeline, TokenizerProcessorStep
|
||||
from lerobot.processor import ActionTokenizerProcessorStep, DataProcessorPipeline, TokenizerProcessorStep
|
||||
from lerobot.processor.converters import create_transition, identity_transition
|
||||
from lerobot.types import TransitionKey
|
||||
from lerobot.utils.constants import (
|
||||
@@ -88,6 +88,46 @@ class MockTokenizer:
|
||||
return result
|
||||
|
||||
|
||||
def test_action_tokenizer_config_preserves_token_mapping():
|
||||
processor = object.__new__(ActionTokenizerProcessorStep)
|
||||
processor.trust_remote_code = True
|
||||
processor.max_action_tokens = 384
|
||||
processor.fast_skip_tokens = 64
|
||||
processor.paligemma_tokenizer_name = "custom/paligemma"
|
||||
processor.allow_truncation = False
|
||||
processor.action_tokenizer_name = "custom/fast"
|
||||
processor.action_tokenizer_input_object = None
|
||||
|
||||
assert processor.get_config() == {
|
||||
"trust_remote_code": True,
|
||||
"max_action_tokens": 384,
|
||||
"fast_skip_tokens": 64,
|
||||
"paligemma_tokenizer_name": "custom/paligemma",
|
||||
"allow_truncation": False,
|
||||
"action_tokenizer_name": "custom/fast",
|
||||
}
|
||||
|
||||
|
||||
def test_action_tokenizer_can_reject_truncated_sequences():
|
||||
processor = object.__new__(ActionTokenizerProcessorStep)
|
||||
processor.max_action_tokens = 4
|
||||
processor.fast_skip_tokens = 128
|
||||
processor.allow_truncation = False
|
||||
processor.action_tokenizer = lambda _actions: [1, 2, 3]
|
||||
processor._paligemma_tokenizer = type(
|
||||
"Tokenizer",
|
||||
(),
|
||||
{
|
||||
"vocab_size": 1000,
|
||||
"bos_token_id": 2,
|
||||
"encode": lambda _self, text, **_kwargs: [10, 11] if text == "Action: " else [12, 1],
|
||||
},
|
||||
)()
|
||||
|
||||
with pytest.raises(ValueError, match="max_action_tokens=4"):
|
||||
processor._tokenize_action(torch.zeros(1, 2, 1))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_tokenizer():
|
||||
"""Provide a mock tokenizer for testing."""
|
||||
|
||||
Reference in New Issue
Block a user