From b63a714ae9b04047e6a34e9e83ebc4bd6c9da6fa Mon Sep 17 00:00:00 2001 From: Khalil Meftah Date: Mon, 22 Jun 2026 15:55:39 +0200 Subject: [PATCH] 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. --- .../policies/pi05/configuration_pi05.py | 5 + src/lerobot/policies/pi05/processor_pi05.py | 40 ++-- .../processor/rendered_messages_to_task.py | 84 ++++++++ .../test_rendered_messages_to_task.py | 186 ++++++++++++++++++ 4 files changed, 302 insertions(+), 13 deletions(-) create mode 100644 src/lerobot/processor/rendered_messages_to_task.py create mode 100644 tests/processor/test_rendered_messages_to_task.py diff --git a/src/lerobot/policies/pi05/configuration_pi05.py b/src/lerobot/policies/pi05/configuration_pi05.py index 124e85cc9..06df47b87 100644 --- a/src/lerobot/policies/pi05/configuration_pi05.py +++ b/src/lerobot/policies/pi05/configuration_pi05.py @@ -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) diff --git a/src/lerobot/policies/pi05/processor_pi05.py b/src/lerobot/policies/pi05/processor_pi05.py index 2d015b24f..df5b932e0 100644 --- a/src/lerobot/policies/pi05/processor_pi05.py +++ b/src/lerobot/policies/pi05/processor_pi05.py @@ -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 diff --git a/src/lerobot/processor/rendered_messages_to_task.py b/src/lerobot/processor/rendered_messages_to_task.py new file mode 100644 index 000000000..c4cce25bb --- /dev/null +++ b/src/lerobot/processor/rendered_messages_to_task.py @@ -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 diff --git a/tests/processor/test_rendered_messages_to_task.py b/tests/processor/test_rendered_messages_to_task.py new file mode 100644 index 000000000..552720820 --- /dev/null +++ b/tests/processor/test_rendered_messages_to_task.py @@ -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"]