mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-06 09:37:06 +00:00
feat(pi05): integrate RenderMessagesStep for advantage conditioning
Add RenderedMessagesToTaskStep adapter that bridges recipe-rendered chat messages back into PI05's task-string prompt format. When recipe_path is set on PI05Config, the preprocessor inserts RenderMessagesStep + adapter before prompt construction, enabling RECAP advantage text to flow end-to-end through the recipe YAML system.
This commit is contained in:
Khalil Meftah
@@ -87,6 +87,11 @@ class PI05Config(PreTrainedConfig):
|
||||
freeze_vision_encoder: bool = False # Freeze only the vision encoder
|
||||
train_expert_only: bool = False # Freeze entire VLM, train only action expert and projections
|
||||
|
||||
# Language conditioning (e.g. RECAP advantage). When set, RenderMessagesStep
|
||||
# is inserted into the preprocessor to resolve language_persistent rows via
|
||||
# the recipe YAML before prompt construction.
|
||||
recipe_path: str | None = None
|
||||
|
||||
# Optimizer settings: see openpi `AdamW`
|
||||
optimizer_lr: float = 2.5e-5 # see openpi `CosineDecaySchedule: peak_lr`
|
||||
optimizer_betas: tuple[float, float] = (0.9, 0.95)
|
||||
|
||||
@@ -111,9 +111,10 @@ def make_pi05_pre_post_processors(
|
||||
1. Renaming features to match pretrained configurations.
|
||||
2. Normalizing input and output features based on dataset statistics.
|
||||
3. Adding a batch dimension.
|
||||
4. Appending a newline character to the task description for tokenizer compatibility.
|
||||
5. Tokenizing the text prompt using the PaliGemma tokenizer.
|
||||
6. Moving all data to the specified device.
|
||||
4. (Optional) Rendering language annotations via recipe YAML.
|
||||
5. (Optional) Flattening rendered messages into the task string.
|
||||
6. Tokenizing the text prompt using the PaliGemma tokenizer.
|
||||
7. Moving all data to the specified device.
|
||||
|
||||
The post-processing pipeline handles the model's output by:
|
||||
1. Moving data to the CPU.
|
||||
@@ -122,8 +123,6 @@ def make_pi05_pre_post_processors(
|
||||
Args:
|
||||
config: The configuration object for the PI0 policy.
|
||||
dataset_stats: A dictionary of statistics for normalization.
|
||||
preprocessor_kwargs: Additional arguments for the pre-processor pipeline.
|
||||
postprocessor_kwargs: Additional arguments for the post-processor pipeline.
|
||||
|
||||
Returns:
|
||||
A tuple containing the configured pre-processor and post-processor pipelines.
|
||||
@@ -147,16 +146,31 @@ def make_pi05_pre_post_processors(
|
||||
norm_map=config.normalization_mapping,
|
||||
stats=dataset_stats,
|
||||
),
|
||||
Pi05PrepareStateTokenizerProcessorStep(max_state_dim=config.max_state_dim),
|
||||
TokenizerProcessorStep(
|
||||
tokenizer_name="google/paligemma-3b-pt-224",
|
||||
max_length=config.tokenizer_max_length,
|
||||
padding_side="right",
|
||||
padding="max_length",
|
||||
),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
]
|
||||
|
||||
# Insert language rendering steps when a recipe is configured (e.g. RECAP advantage)
|
||||
if config.recipe_path is not None:
|
||||
from lerobot.configs.recipe import load_recipe
|
||||
from lerobot.processor.render_messages_processor import RenderMessagesStep
|
||||
from lerobot.processor.rendered_messages_to_task import RenderedMessagesToTaskStep
|
||||
|
||||
recipe = load_recipe(config.recipe_path)
|
||||
input_steps.append(RenderMessagesStep(recipe=recipe))
|
||||
input_steps.append(RenderedMessagesToTaskStep())
|
||||
|
||||
input_steps.extend(
|
||||
[
|
||||
Pi05PrepareStateTokenizerProcessorStep(max_state_dim=config.max_state_dim),
|
||||
TokenizerProcessorStep(
|
||||
tokenizer_name="google/paligemma-3b-pt-224",
|
||||
max_length=config.tokenizer_max_length,
|
||||
padding_side="right",
|
||||
padding="max_length",
|
||||
),
|
||||
DeviceProcessorStep(device=config.device),
|
||||
]
|
||||
)
|
||||
|
||||
output_steps: list[ProcessorStep] = [
|
||||
UnnormalizerProcessorStep(
|
||||
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
|
||||
|
||||
@@ -0,0 +1,84 @@
|
||||
#!/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.
|
||||
|
||||
"""Adapter step that flattens rendered chat messages back into a task string.
|
||||
|
||||
Bridges RenderMessagesStep (which outputs structured messages) to policies
|
||||
that expect a plain task string in complementary_data["task"] (e.g. PI05).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from lerobot.configs import PipelineFeatureType, PolicyFeature
|
||||
|
||||
from .pipeline import ComplementaryDataProcessorStep, ProcessorStepRegistry
|
||||
|
||||
|
||||
@ProcessorStepRegistry.register(name="rendered_messages_to_task")
|
||||
class RenderedMessagesToTaskStep(ComplementaryDataProcessorStep):
|
||||
"""Extract user-role message content from rendered messages into the task string.
|
||||
|
||||
After RenderMessagesStep renders a recipe into structured messages, this
|
||||
step extracts content from all user-role messages, joins them, and writes
|
||||
the result to complementary_data["task"]. This allows downstream steps
|
||||
(like Pi05PrepareStateTokenizerProcessorStep) to consume the
|
||||
advantage-conditioned prompt without modification.
|
||||
|
||||
No-ops when the "messages" key is absent (backward compatible with
|
||||
pipelines that don't use language annotations).
|
||||
"""
|
||||
|
||||
def complementary_data(self, complementary_data: dict) -> dict:
|
||||
messages = complementary_data.get("messages")
|
||||
if messages is None:
|
||||
return complementary_data
|
||||
|
||||
user_parts = []
|
||||
for msg in messages:
|
||||
if msg.get("role") == "user":
|
||||
content = msg.get("content", "")
|
||||
if isinstance(content, str) and content:
|
||||
user_parts.append(content)
|
||||
elif isinstance(content, list):
|
||||
# HF multimodal blocks: extract text blocks
|
||||
for block in content:
|
||||
if isinstance(block, dict) and block.get("type") == "text":
|
||||
text = block.get("text", "")
|
||||
if text:
|
||||
user_parts.append(text)
|
||||
|
||||
new_complementary_data = dict(complementary_data)
|
||||
|
||||
if user_parts:
|
||||
task = complementary_data.get("task")
|
||||
# Wrap in list if the original task was a list (batched)
|
||||
joined = "\n".join(user_parts)
|
||||
if isinstance(task, list):
|
||||
new_complementary_data["task"] = [joined] * len(task)
|
||||
else:
|
||||
new_complementary_data["task"] = joined
|
||||
|
||||
# Remove consumed rendering outputs
|
||||
new_complementary_data.pop("messages", None)
|
||||
new_complementary_data.pop("message_streams", None)
|
||||
new_complementary_data.pop("target_message_indices", None)
|
||||
|
||||
return new_complementary_data
|
||||
|
||||
def transform_features(
|
||||
self, features: dict[PipelineFeatureType, dict[str, PolicyFeature]]
|
||||
) -> dict[PipelineFeatureType, dict[str, PolicyFeature]]:
|
||||
return features
|
||||
@@ -0,0 +1,186 @@
|
||||
#!/usr/bin/env python
|
||||
|
||||
"""Tests for RenderedMessagesToTaskStep and PI05 pipeline integration with advantage."""
|
||||
|
||||
import pytest
|
||||
|
||||
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.processor.rendered_messages_to_task import RenderedMessagesToTaskStep # noqa: E402
|
||||
from lerobot.types import TransitionKey # noqa: E402
|
||||
|
||||
|
||||
def test_rendered_messages_to_task_noops_without_messages():
|
||||
"""Without messages key, the step is a no-op."""
|
||||
transition = create_transition(complementary_data={"task": "pick up the cup"})
|
||||
step = RenderedMessagesToTaskStep()
|
||||
out = step(transition)
|
||||
data = out[TransitionKey.COMPLEMENTARY_DATA]
|
||||
assert data["task"] == "pick up the cup"
|
||||
|
||||
|
||||
def test_rendered_messages_to_task_extracts_user_content():
|
||||
"""Extracts user-role message content and joins with newline."""
|
||||
transition = create_transition(
|
||||
complementary_data={
|
||||
"task": "original task",
|
||||
"messages": [
|
||||
{"role": "user", "content": "pick up the cup"},
|
||||
{"role": "user", "content": "Advantage: positive"},
|
||||
{"role": "assistant", "content": "reach for cup"},
|
||||
],
|
||||
"message_streams": ["high_level", "high_level", "low_level"],
|
||||
"target_message_indices": [2],
|
||||
}
|
||||
)
|
||||
step = RenderedMessagesToTaskStep()
|
||||
out = step(transition)
|
||||
data = out[TransitionKey.COMPLEMENTARY_DATA]
|
||||
|
||||
assert data["task"] == "pick up the cup\nAdvantage: positive"
|
||||
assert "messages" not in data
|
||||
assert "message_streams" not in data
|
||||
assert "target_message_indices" not in data
|
||||
|
||||
|
||||
def test_rendered_messages_to_task_handles_multimodal_blocks():
|
||||
"""Extracts text from HF multimodal content blocks."""
|
||||
transition = create_transition(
|
||||
complementary_data={
|
||||
"task": "original",
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "image", "image": "placeholder"},
|
||||
{"type": "text", "text": "describe this"},
|
||||
],
|
||||
},
|
||||
{"role": "assistant", "content": "a cup on a table"},
|
||||
],
|
||||
"message_streams": ["high_level", "low_level"],
|
||||
"target_message_indices": [1],
|
||||
}
|
||||
)
|
||||
step = RenderedMessagesToTaskStep()
|
||||
out = step(transition)
|
||||
data = out[TransitionKey.COMPLEMENTARY_DATA]
|
||||
|
||||
assert data["task"] == "describe this"
|
||||
|
||||
|
||||
def test_rendered_messages_to_task_preserves_list_task_format():
|
||||
"""When original task is a list (batched), output is also a list."""
|
||||
transition = create_transition(
|
||||
complementary_data={
|
||||
"task": ["task1", "task2"],
|
||||
"messages": [
|
||||
{"role": "user", "content": "rendered task"},
|
||||
{"role": "assistant", "content": "do it", "target": True},
|
||||
],
|
||||
"message_streams": ["high_level", "low_level"],
|
||||
"target_message_indices": [1],
|
||||
}
|
||||
)
|
||||
step = RenderedMessagesToTaskStep()
|
||||
out = step(transition)
|
||||
data = out[TransitionKey.COMPLEMENTARY_DATA]
|
||||
|
||||
assert data["task"] == ["rendered task", "rendered task"]
|
||||
|
||||
|
||||
def test_full_render_then_flatten_pipeline():
|
||||
"""RenderMessagesStep + RenderedMessagesToTaskStep produces correct task string."""
|
||||
recipe = TrainingRecipe(
|
||||
messages=[
|
||||
MessageTurn(role="user", content="${task}", stream="high_level"),
|
||||
MessageTurn(
|
||||
role="user",
|
||||
content="Advantage: ${advantage}",
|
||||
stream="high_level",
|
||||
if_present="advantage",
|
||||
),
|
||||
MessageTurn(role="assistant", content="${subtask}", stream="low_level", target=True),
|
||||
]
|
||||
)
|
||||
transition = create_transition(
|
||||
complementary_data={
|
||||
"task": "pick up the cup",
|
||||
"timestamp": torch.tensor(0.5),
|
||||
"index": torch.tensor(0),
|
||||
"language_persistent": [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "reach for the cup",
|
||||
"style": "subtask",
|
||||
"timestamp": 0.0,
|
||||
"camera": None,
|
||||
"tool_calls": None,
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "positive",
|
||||
"style": "advantage",
|
||||
"timestamp": 0.1,
|
||||
"camera": None,
|
||||
"tool_calls": None,
|
||||
},
|
||||
],
|
||||
"language_events": [],
|
||||
}
|
||||
)
|
||||
|
||||
# Step 1: Render recipe
|
||||
rendered = RenderMessagesStep(recipe=recipe)(transition)
|
||||
# Step 2: Flatten to task string
|
||||
out = RenderedMessagesToTaskStep()(rendered)
|
||||
data = out[TransitionKey.COMPLEMENTARY_DATA]
|
||||
|
||||
assert "pick up the cup" in data["task"]
|
||||
assert "Advantage: positive" in data["task"]
|
||||
|
||||
|
||||
def test_full_render_advantage_absent_skips_turn():
|
||||
"""When advantage row is absent, the advantage turn is skipped via if_present."""
|
||||
recipe = TrainingRecipe(
|
||||
messages=[
|
||||
MessageTurn(role="user", content="${task}", stream="high_level"),
|
||||
MessageTurn(
|
||||
role="user",
|
||||
content="Advantage: ${advantage}",
|
||||
stream="high_level",
|
||||
if_present="advantage",
|
||||
),
|
||||
MessageTurn(role="assistant", content="${subtask}", stream="low_level", target=True),
|
||||
]
|
||||
)
|
||||
transition = create_transition(
|
||||
complementary_data={
|
||||
"task": "pick up the cup",
|
||||
"timestamp": torch.tensor(0.5),
|
||||
"index": torch.tensor(0),
|
||||
"language_persistent": [
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "reach for the cup",
|
||||
"style": "subtask",
|
||||
"timestamp": 0.0,
|
||||
"camera": None,
|
||||
"tool_calls": None,
|
||||
},
|
||||
],
|
||||
"language_events": [],
|
||||
}
|
||||
)
|
||||
|
||||
rendered = RenderMessagesStep(recipe=recipe)(transition)
|
||||
out = RenderedMessagesToTaskStep()(rendered)
|
||||
data = out[TransitionKey.COMPLEMENTARY_DATA]
|
||||
|
||||
assert data["task"] == "pick up the cup"
|
||||
assert "Advantage" not in data["task"]
|
||||
Reference in New Issue
Block a user