refactor(rtc): share Pi05 training helpers

src/lerobot/policies/pi05/modeling_pi05.py

@@ -1066,7 +1066,7 @@ class PI05Pytorch(nn.Module):  # see openpi `PI0Pytorch`
                 training_max_delay = int(getattr(self.config, "rtc_training_max_delay", 0))
                 if training_max_delay <= 0:
                     raise ValueError(
-                        "RTC mode='trained' requires a Pi052 checkpoint trained with "
+                        "RTC mode='trained' requires a checkpoint trained with "
                         "policy.rtc_training_max_delay > 0."
                     )
                 trained_prefix, trained_prefix_mask = _prepare_trained_rtc_prefix(

src/lerobot/policies/pi052/configuration_pi052.py

@@ -121,9 +121,7 @@ class PI052Config(PI05Config):
     # Reuse each VLM prefix across independent denoising draws; 1 restores single-draw flow.
     flow_num_repeats: int = 5
 
-    # Training-time RTC (arXiv:2512.05964). Zero preserves standard flow matching.
+    # Training-time RTC configuration is inherited from PI05Config.
-    rtc_training_max_delay: int = 0
-    """Maximum clean-prefix delay sampled for training-time RTC; zero disables it."""
 
     # PaLM-style z-loss stabilizes large-vocabulary CE; 0 disables it.
     text_ce_z_loss_weight: float = 1e-4
@@ -155,11 +153,6 @@ class PI052Config(PI05Config):
             self.train_expert_only = False
         if self.flow_num_repeats < 1:
             raise ValueError(f"flow_num_repeats must be >= 1, got {self.flow_num_repeats}")
-        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.manual_attention_scope not in {"all", "action"}:
             raise ValueError(
                 f"manual_attention_scope must be 'all' or 'action', got {self.manual_attention_scope!r}"

src/lerobot/policies/pi052/modeling_pi052.py

@@ -36,6 +36,8 @@ from ..pi05.modeling_pi05 import (
     ActionSelectKwargs,
     PI05Policy,
     PI05Pytorch as PI05PytorchBase,
+    _build_flow_matching_inputs,
+    _sample_training_rtc_prefix_mask,
     make_att_2d_masks,
 )
 from .configuration_pi052 import PI052Config
@@ -177,43 +179,6 @@ def _enable_hf_kernels() -> None:
     logger.info("PI052: HF kernels (Liger) enabled — rope, geglu fused.")
 
 
-def _sample_training_rtc_prefix_mask(
-    batch_size: int,
-    action_horizon: int,
-    max_delay: int,
-    device: torch.device,
-) -> Tensor | None:
-    """Sample per-draw clean prefixes for training-time RTC."""
-    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]:
-    """Build noisy actions and scalar/per-token flow times.
-
-    LeRobot's PI0.5 flow uses ``t=0`` for clean data and ``t=1`` for
-    noise, the reverse of the notation in arXiv:2512.05964. Consequently,
-    clean RTC prefix tokens receive ``t=0`` here.
-    """
-    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 _flow_loss_components(flow_per_dim: Tensor, prefix_mask: Tensor | None) -> tuple[Tensor, Tensor]:
     """Return each sample's postfix loss sum and valid-element count."""
     if prefix_mask is None:

tests/policies/pi052/test_pi052_training_time_rtc.py

@@ -16,13 +16,13 @@ from torch import nn
 pytest.importorskip("transformers")
 
 from lerobot.policies.pi05.modeling_pi05 import (  # noqa: E402
+    _build_flow_matching_inputs,
     _prepare_trained_rtc_prefix,
     create_sinusoidal_pos_embedding,
 )
 from lerobot.policies.pi052.configuration_pi052 import PI052Config  # noqa: E402
 from lerobot.policies.pi052.modeling_pi052 import (  # noqa: E402
     PI05Pytorch as PI052Pytorch,
-    _build_flow_matching_inputs,
     _flow_loss_per_sample,
     _reduce_flow_loss,
 )