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.

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)

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.
+    4. (Optional) Rendering language annotations via recipe YAML.
-    5. Tokenizing the text prompt using the PaliGemma tokenizer.
+    5. (Optional) Flattening rendered messages into the task string.
-    6. Moving all data to the specified device.
+    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

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

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"]