feat(pi05): add training-time RTC

Extend trained RTC conditioning and postfix-only flow loss to plain PI05 so single-task checkpoints can use the same real-time chunking path as PI052.

Co-authored-by: Cursor <cursoragent@cursor.com>

src/lerobot/policies/pi05/configuration_pi05.py

@@ -58,6 +58,8 @@ class PI05Config(PreTrainedConfig):
 
     # Real-Time Chunking (RTC) configuration
     rtc_config: RTCConfig | None = None
+    # Maximum clean action-prefix length sampled during training. Zero disables trained RTC.
+    rtc_training_max_delay: int = 0
 
     image_resolution: tuple[int, int] = (
         DEFAULT_IMAGE_SIZE,
@@ -111,6 +113,11 @@ class PI05Config(PreTrainedConfig):
             raise ValueError(
                 f"n_action_steps ({self.n_action_steps}) cannot be greater than chunk_size ({self.chunk_size})"
             )
+        if not 0 <= self.rtc_training_max_delay < self.chunk_size:
+            raise ValueError(
+                "rtc_training_max_delay must satisfy "
+                f"0 <= delay < chunk_size ({self.chunk_size}), got {self.rtc_training_max_delay}"
+            )
 
         if self.paligemma_variant not in ["gemma_300m", "gemma_2b"]:
             raise ValueError(f"Invalid paligemma_variant: {self.paligemma_variant}")

src/lerobot/policies/pi05/modeling_pi05.py

@@ -81,7 +81,7 @@ def _prepare_trained_rtc_prefix(
         return None, None
     if training_max_delay <= 0:
         raise ValueError(
-            "RTC mode='trained' requires a Pi052 checkpoint trained with policy.rtc_training_max_delay > 0."
+            "RTC mode='trained' requires a checkpoint trained with policy.rtc_training_max_delay > 0."
         )
     if inference_delay > training_max_delay:
         raise ValueError(
@@ -120,6 +120,57 @@ def _prepare_trained_rtc_prefix(
     return padded_prefix, prefix_mask
 
 
+def _sample_training_rtc_prefix_mask(
+    batch_size: int,
+    action_horizon: int,
+    max_delay: int,
+    device: torch.device,
+) -> Tensor | None:
+    """Sample a clean action-prefix length independently for each training example."""
+    if max_delay <= 0:
+        return None
+    delays = torch.randint(0, max_delay + 1, (batch_size,), device=device)
+    positions = torch.arange(action_horizon, device=device)
+    return positions.unsqueeze(0) < delays.unsqueeze(1)
+
+
+def _build_flow_matching_inputs(
+    actions: Tensor,
+    noise: Tensor,
+    time: Tensor,
+    prefix_mask: Tensor | None,
+) -> tuple[Tensor, Tensor]:
+    """Keep the sampled RTC prefix clean while noising the remaining action chunk."""
+    if prefix_mask is None:
+        model_time = time
+        expanded_time = time[:, None, None]
+    else:
+        model_time = time[:, None].expand_as(prefix_mask)
+        model_time = torch.where(prefix_mask, torch.zeros_like(model_time), model_time)
+        expanded_time = model_time.unsqueeze(-1)
+    x_t = expanded_time * noise + (1 - expanded_time) * actions
+    return x_t, model_time
+
+
+def _reduce_training_rtc_loss(
+    losses: Tensor,
+    prefix_mask: Tensor | None,
+    reduction: str,
+) -> Tensor:
+    """Average flow loss over predicted postfix actions, excluding the clean RTC prefix."""
+    if reduction not in {"mean", "none"}:
+        raise ValueError(f"Unsupported loss reduction: {reduction!r}")
+    if prefix_mask is None:
+        return losses.mean() if reduction == "mean" else losses.mean(dim=(1, 2))
+
+    postfix_mask = (~prefix_mask).unsqueeze(-1).expand_as(losses)
+    if reduction == "none":
+        numerator = (losses * postfix_mask).sum(dim=(1, 2))
+        denominator = postfix_mask.sum(dim=(1, 2))
+        return numerator / denominator.clamp(min=1)
+    return (losses * postfix_mask).sum() / postfix_mask.sum().clamp(min=1)
+
+
 _SAFETENSORS_FILE = "model.safetensors"
 _SAFETENSORS_INDEX = "model.safetensors.index.json"
 
@@ -897,14 +948,23 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
 
         return embs, pad_masks, att_masks, adarms_cond
 
-    def forward(self, images, img_masks, tokens, masks, actions, noise, time) -> Tensor:
+    def forward(
+        self,
+        images,
+        img_masks,
+        tokens,
+        masks,
+        actions,
+        noise,
+        time,
+        prefix_mask: Tensor | None = None,
+    ) -> Tensor:
         """Do a full training forward pass and compute the loss."""
-        time_expanded = time[:, None, None]
-        x_t = time_expanded * noise + (1 - time_expanded) * actions
+        x_t, model_time = _build_flow_matching_inputs(actions, noise, time, prefix_mask)
         u_t = noise - actions
 
         prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix(images, img_masks, tokens, masks)
-        suffix_embs, suffix_pad_masks, suffix_att_masks, adarms_cond = self.embed_suffix(x_t, time)
+        suffix_embs, suffix_pad_masks, suffix_att_masks, adarms_cond = self.embed_suffix(x_t, model_time)
 
         if (
             self.paligemma_with_expert.paligemma.model.language_model.layers[0].self_attn.q_proj.weight.dtype
@@ -1423,26 +1483,33 @@ class PI05Policy(PreTrainedPolicy):
 
         noise = self.model.sample_noise(actions.shape, actions.device)
         time = self.model.sample_time(actions.shape[0], actions.device)
+        prefix_mask = _sample_training_rtc_prefix_mask(
+            actions.shape[0],
+            actions.shape[1],
+            self.config.rtc_training_max_delay,
+            actions.device,
+        )
 
         # Compute loss (no separate state needed for PI05)
-        losses = self.model.forward(images, img_masks, tokens, masks, actions, noise, time)
+        losses = self.model.forward(images, img_masks, tokens, masks, actions, noise, time, prefix_mask)
 
         # Truncate losses to actual action dimensions
         original_action_dim = self.config.output_features[ACTION].shape[0]
         losses = losses[:, :, :original_action_dim]
 
-        loss_dict = {
-            "loss_per_dim": losses.mean(dim=[0, 1]).detach().cpu().numpy().tolist(),
-        }
+        if prefix_mask is None:
+            loss_per_dim = losses.mean(dim=(0, 1))
+        else:
+            postfix_mask = (~prefix_mask).unsqueeze(-1).expand_as(losses)
+            loss_per_dim = (losses * postfix_mask).sum(dim=(0, 1)) / postfix_mask.sum(dim=(0, 1)).clamp(min=1)
+        loss_dict = {"loss_per_dim": loss_per_dim.detach().cpu().numpy().tolist()}
 
         if reduction == "none":
-            # Return per-sample losses (B,) by averaging over time and action dims
-            per_sample_loss = losses.mean(dim=(1, 2))
+            per_sample_loss = _reduce_training_rtc_loss(losses, prefix_mask, reduction="none")
             loss_dict["loss"] = per_sample_loss.mean().item()
             return per_sample_loss, loss_dict
-        else:
-            # Default: return scalar mean loss
-            loss = losses.mean()
+
+        loss = _reduce_training_rtc_loss(losses, prefix_mask, reduction="mean")
         loss_dict["loss"] = loss.item()
         return loss, loss_dict
 

src/lerobot/policies/rtc/modeling_rtc.py

@@ -45,7 +45,8 @@ class RTCProcessor:
     def __init__(self, rtc_config: RTCConfig, *, trained_mode_supported: bool = False):
         if rtc_config.enabled and rtc_config.mode == "trained" and not trained_mode_supported:
             raise ValueError(
-                "RTC mode='trained' requires a Pi052 checkpoint trained with rtc_training_max_delay > 0."
+                "RTC mode='trained' requires a PI05-compatible checkpoint trained with "
+                "rtc_training_max_delay > 0."
             )
         self.rtc_config = rtc_config
 

src/lerobot/rollout/context.py

@@ -65,9 +65,9 @@ def _validate_trained_rtc_rollout_config(policy_config, inference_config: RTCInf
     rtc = inference_config.rtc
     if not rtc.enabled or rtc.mode != "trained":
         return
-    if policy_config.type != "pi052":
+    if policy_config.type not in {"pi05", "pi052"}:
         raise ValueError(
-            "--inference.rtc.mode=trained currently requires a Pi052 checkpoint; "
+            "--inference.rtc.mode=trained currently requires a PI05-compatible checkpoint; "
             f"got policy type {policy_config.type!r}."
         )
 

tests/policies/pi0_pi05/test_pi05_training_time_rtc.py

@@ -0,0 +1,45 @@
+#!/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.
+
+import pytest
+import torch
+
+pytest.importorskip("transformers")
+
+from lerobot.policies.pi05.configuration_pi05 import PI05Config  # noqa: E402
+from lerobot.policies.pi05.modeling_pi05 import (  # noqa: E402
+    _build_flow_matching_inputs,
+    _reduce_training_rtc_loss,
+)
+
+
+def test_pi05_training_rtc_uses_clean_prefix_and_per_token_time():
+    actions = torch.tensor([[[1.0], [2.0], [3.0], [4.0]]])
+    noise = torch.tensor([[[10.0], [20.0], [30.0], [40.0]]])
+    time = torch.tensor([0.25])
+    prefix_mask = torch.tensor([[True, True, False, False]])
+
+    x_t, model_time = _build_flow_matching_inputs(actions, noise, time, prefix_mask)
+
+    assert model_time.tolist() == [[0.0, 0.0, 0.25, 0.25]]
+    assert torch.equal(x_t[:, :2], actions[:, :2])
+    assert torch.equal(x_t[:, 2:], 0.25 * noise[:, 2:] + 0.75 * actions[:, 2:])
+
+
+def test_pi05_training_rtc_loss_excludes_clean_prefix():
+    losses = torch.tensor([[[100.0], [100.0], [2.0], [4.0]]])
+    prefix_mask = torch.tensor([[True, True, False, False]])
+
+    loss = _reduce_training_rtc_loss(losses, prefix_mask, reduction="mean")
+
+    assert loss.item() == pytest.approx(3.0)
+
+
+@pytest.mark.parametrize("max_delay", [-1, 5])
+def test_pi05_config_rejects_invalid_training_rtc_delay(max_delay):
+    with pytest.raises(ValueError, match="rtc_training_max_delay"):
+        PI05Config(chunk_size=5, n_action_steps=5, rtc_training_max_delay=max_delay)

tests/policies/rtc/test_modeling_rtc.py

@@ -96,7 +96,7 @@ def test_rtc_processor_initialization_without_debug(rtc_config_debug_disabled):
 def test_rtc_processor_rejects_trained_mode_when_policy_does_not_support_it():
     config = RTCConfig(mode="trained")
 
-    with pytest.raises(ValueError, match="requires a Pi052 checkpoint"):
+    with pytest.raises(ValueError, match="requires a PI05-compatible checkpoint"):
         RTCProcessor(config)
 
     processor = RTCProcessor(config, trained_mode_supported=True)