diff --git a/src/lerobot/policies/pi05/configuration_pi05.py b/src/lerobot/policies/pi05/configuration_pi05.py index 124e85cc9..81690ea54 100644 --- a/src/lerobot/policies/pi05/configuration_pi05.py +++ b/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}") diff --git a/src/lerobot/policies/pi05/modeling_pi05.py b/src/lerobot/policies/pi05/modeling_pi05.py index 6e1888a68..8017de28f 100644 --- a/src/lerobot/policies/pi05/modeling_pi05.py +++ b/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,28 +1483,35 @@ 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_dict["loss"] = loss.item() - return loss, loss_dict + + loss = _reduce_training_rtc_loss(losses, prefix_mask, reduction="mean") + loss_dict["loss"] = loss.item() + return loss, loss_dict def _get_default_peft_targets(self) -> dict[str, any]: """Return default PEFT target modules for PI0.5 fine-tuning.""" diff --git a/src/lerobot/policies/rtc/modeling_rtc.py b/src/lerobot/policies/rtc/modeling_rtc.py index fe5077c96..24f0c6a5b 100644 --- a/src/lerobot/policies/rtc/modeling_rtc.py +++ b/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 diff --git a/src/lerobot/rollout/context.py b/src/lerobot/rollout/context.py index b1632ea1f..aa69c8db5 100644 --- a/src/lerobot/rollout/context.py +++ b/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}." ) diff --git a/tests/policies/pi0_pi05/test_pi05_training_time_rtc.py b/tests/policies/pi0_pi05/test_pi05_training_time_rtc.py new file mode 100644 index 000000000..394cf9ea0 --- /dev/null +++ b/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) diff --git a/tests/policies/rtc/test_modeling_rtc.py b/tests/policies/rtc/test_modeling_rtc.py index 6607ea59a..e6a2ea584 100644 --- a/tests/policies/rtc/test_modeling_rtc.py +++ b/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)