fix(rtc): align trained mode with RLDX

src/lerobot/policies/pi05/modeling_pi05.py

@@ -94,6 +94,8 @@ def _prepare_trained_rtc_prefix(
         )
 
     previous = prev_chunk_left_over.to(device=x_t.device, dtype=x_t.dtype)
+    if not torch.isfinite(previous).all():
+        raise ValueError("RTC prefix contains NaN or Inf values.")
     if previous.ndim == 2:
         previous = previous.unsqueeze(0)
     if previous.ndim != 3:

src/lerobot/policies/pi052/configuration_pi052.py

@@ -123,12 +123,7 @@ class PI052Config(PI05Config):
 
     # Training-time RTC (arXiv:2512.05964). Zero preserves standard flow matching.
     rtc_training_max_delay: int = 0
-    """Largest clean action-prefix length sampled during training.
+    """Maximum clean-prefix delay sampled for training-time RTC; zero disables it."""
-
-    A value greater than zero enables training-time action conditioning. Each
-    flow draw samples a delay uniformly from ``[0, rtc_training_max_delay]``;
-    the corresponding action prefix stays clean and is excluded from the loss.
-    """
 
     # PaLM-style z-loss stabilizes large-vocabulary CE; 0 disables it.
     text_ce_z_loss_weight: float = 1e-4

src/lerobot/policies/pi052/modeling_pi052.py

@@ -177,16 +177,6 @@ def _enable_hf_kernels() -> None:
     logger.info("PI052: HF kernels (Liger) enabled — rope, geglu fused.")
 
 
-def _reduce_action_loss(per_sample: Tensor, predict_actions_t: Tensor | None, reduction: str) -> Tensor:
-    """Mask non-action samples and apply the requested batch reduction."""
-    if predict_actions_t is None:
-        return per_sample if reduction == "none" else per_sample.mean()
-    mask = predict_actions_t.to(per_sample.dtype)
-    if reduction == "none":
-        return per_sample * mask
-    return (per_sample * mask).sum() / mask.sum().clamp(min=1.0)
-
-
 def _sample_training_rtc_prefix_mask(
     batch_size: int,
     action_horizon: int,
@@ -224,15 +214,42 @@ def _build_flow_matching_inputs(
     return x_t, model_time
 
 
-def _flow_loss_per_sample(flow_per_dim: Tensor, prefix_mask: Tensor | None) -> Tensor:
+def _flow_loss_components(flow_per_dim: Tensor, prefix_mask: Tensor | None) -> tuple[Tensor, Tensor]:
-    """Average each draw over postfix action positions and dimensions only."""
+    """Return each sample's postfix loss sum and valid-element count."""
     if prefix_mask is None:
-        return flow_per_dim.flatten(start_dim=1).mean(dim=1)
+        flattened = flow_per_dim.flatten(start_dim=1)
-    postfix = (~prefix_mask).unsqueeze(-1).expand_as(flow_per_dim)
+        numerator = flattened.sum(dim=1)
+        denominator = torch.full_like(numerator, flattened.shape[1])
+        return numerator, denominator
+
+    postfix = (~prefix_mask).unsqueeze(-1).expand_as(flow_per_dim).to(flow_per_dim.dtype)
     reduce_dims = tuple(range(1, flow_per_dim.ndim))
     numerator = (flow_per_dim * postfix).sum(dim=reduce_dims)
-    denominator = postfix.sum(dim=reduce_dims).clamp(min=1)
+    denominator = postfix.sum(dim=reduce_dims)
-    return numerator / denominator
+    return numerator, denominator
+
+
+def _flow_loss_per_sample(flow_per_dim: Tensor, prefix_mask: Tensor | None) -> Tensor:
+    """Average each sample over postfix action positions and dimensions only."""
+    numerator, denominator = _flow_loss_components(flow_per_dim, prefix_mask)
+    return numerator / denominator.clamp(min=1)
+
+
+def _reduce_flow_loss(
+    flow_per_dim: Tensor,
+    prefix_mask: Tensor | None,
+    predict_actions_t: Tensor | None,
+    reduction: str,
+) -> Tensor:
+    """Apply action routing and RLDX-compatible postfix normalization."""
+    numerator, denominator = _flow_loss_components(flow_per_dim, prefix_mask)
+    if predict_actions_t is not None:
+        action_mask = predict_actions_t.to(numerator.dtype)
+        numerator = numerator * action_mask
+        denominator = denominator * action_mask
+    if reduction == "none":
+        return numerator / denominator.clamp(min=1)
+    return numerator.sum() / denominator.sum().clamp(min=1)
 
 
 # Materialized logits win at VLA token counts; larger dense targets use Liger.
@@ -1223,8 +1240,7 @@ class PI052Policy(PI05Policy):
         # internally to max_action_dim).
         original_action_dim = self.config.output_features[ACTION].shape[0]
         flow_per_dim = flow_per_dim[:, :, :original_action_dim]
-        per_sample_flow = _flow_loss_per_sample(flow_per_dim, prefix_mask)
+        flow_loss = _reduce_flow_loss(flow_per_dim, prefix_mask, predict_actions_t, reduction)
-        flow_loss = _reduce_action_loss(per_sample_flow, predict_actions_t, reduction)
         return prefix_out, flow_loss
 
     def _ki_forward_kwargs(self, suppress_prefix_grads: bool = False, flex_masks=None) -> dict[str, Any]:
@@ -1363,8 +1379,7 @@ class PI052Policy(PI05Policy):
         prefix_mask = None
         if prefix_mask_flat is not None:
             prefix_mask = prefix_mask_flat.view(k, batch_size, chunk).transpose(0, 1)
-        per_sample_flow = _flow_loss_per_sample(flow_per_dim, prefix_mask)
+        flow_loss = _reduce_flow_loss(flow_per_dim, prefix_mask, predict_actions_t, reduction)
-        flow_loss = _reduce_action_loss(per_sample_flow, predict_actions_t, reduction)
         return prefix_out, flow_loss
 
     def _prefix_ce_losses(

src/lerobot/rollout/inference/rtc.py

@@ -57,6 +57,10 @@ _RTC_MAX_CONSECUTIVE_ERRORS: int = 10
 _RTC_JOIN_TIMEOUT_S: float = 3.0
 
 
+class _TrainedRTCDelayExceededError(RuntimeError):
+    """Raised when measured latency exceeds a trained RTC checkpoint's support."""
+
+
 # ---------------------------------------------------------------------------
 # RTC helpers
 # ---------------------------------------------------------------------------
@@ -76,6 +80,22 @@ def _normalize_prev_actions_length(prev_actions: torch.Tensor, target_steps: int
     return padded
 
 
+def _trained_rtc_chunk_can_merge(
+    *,
+    conditioned_delay: int,
+    measured_delay: int,
+    training_max_delay: int,
+    has_previous_actions: bool,
+) -> bool:
+    """Check that a trained RTC chunk covers the overlap observed during inference."""
+    if measured_delay > training_max_delay:
+        raise _TrainedRTCDelayExceededError(
+            f"Measured RTC inference delay ({measured_delay}) exceeds the checkpoint's "
+            f"rtc_training_max_delay ({training_max_delay})."
+        )
+    return not has_previous_actions or measured_delay <= conditioned_delay
+
+
 # ---------------------------------------------------------------------------
 # RTCInferenceEngine
 # ---------------------------------------------------------------------------
@@ -272,6 +292,7 @@ class RTCInferenceEngine(InferenceEngine):
                         current_time = time.perf_counter()
                         idx_before = queue.get_action_index()
                         prev_actions = queue.get_left_over()
+                        has_previous_actions = prev_actions is not None and prev_actions.numel() > 0
 
                         latency = latency_tracker.max()
                         delay = math.ceil(latency / time_per_chunk) if latency else 0
@@ -321,6 +342,24 @@ class RTCInferenceEngine(InferenceEngine):
                         else:
                             latency_tracker.add(new_latency)
 
+                        if not is_warmup and self._rtc_config.mode == "trained":
+                            training_max_delay = int(
+                                getattr(self._policy.config, "rtc_training_max_delay", 0)
+                            )
+                            if not _trained_rtc_chunk_can_merge(
+                                conditioned_delay=delay,
+                                measured_delay=new_delay,
+                                training_max_delay=training_max_delay,
+                                has_previous_actions=has_previous_actions,
+                            ):
+                                logger.warning(
+                                    "Discarding trained RTC chunk: measured delay %d exceeded "
+                                    "conditioned delay %d; retrying with updated latency",
+                                    new_delay,
+                                    delay,
+                                )
+                                continue
+
                         queue.merge(original, processed, new_delay, idx_before)
 
                         if (
@@ -333,6 +372,8 @@ class RTCInferenceEngine(InferenceEngine):
 
                         logger.debug("RTC inference latency=%.2fs, queue=%d", new_latency, queue.qsize())
 
+                    except _TrainedRTCDelayExceededError:
+                        raise
                     except Exception as e:
                         consecutive_errors += 1
                         logger.error(

tests/policies/pi052/test_pi052_training_time_rtc.py

@@ -24,6 +24,7 @@ from lerobot.policies.pi052.modeling_pi052 import (  # noqa: E402
     PI05Pytorch as PI052Pytorch,
     _build_flow_matching_inputs,
     _flow_loss_per_sample,
+    _reduce_flow_loss,
 )
 
 
@@ -49,6 +50,25 @@ def test_training_rtc_loss_averages_over_postfix_only():
     assert per_sample.tolist() == [3.0]
 
 
+def test_training_rtc_mean_loss_uses_global_postfix_normalization():
+    flow_loss = torch.tensor(
+        [
+            [[1.0], [1.0], [1.0], [1.0]],
+            [[9.0], [9.0], [9.0], [9.0]],
+        ]
+    )
+    prefix_mask = torch.tensor(
+        [
+            [False, False, False, False],
+            [True, True, True, False],
+        ]
+    )
+
+    loss = _reduce_flow_loss(flow_loss, prefix_mask, predict_actions_t=None, reduction="mean")
+
+    assert loss.item() == pytest.approx(13 / 5)
+
+
 def test_per_token_time_embedding_preserves_action_axis():
     time = torch.tensor([[0.0, 0.5, 1.0]])
 
@@ -100,6 +120,20 @@ def test_trained_rtc_rejects_delay_outside_training_distribution():
         )
 
 
+@pytest.mark.parametrize("non_finite", [float("nan"), float("inf")])
+def test_trained_rtc_rejects_non_finite_prefix(non_finite):
+    previous = torch.randn(4, 2)
+    previous[0, 0] = non_finite
+
+    with pytest.raises(ValueError, match="NaN or Inf"):
+        _prepare_trained_rtc_prefix(
+            torch.randn(1, 5, 4),
+            previous,
+            inference_delay=2,
+            training_max_delay=3,
+        )
+
+
 @pytest.mark.parametrize("max_delay", [-1, 5])
 def test_pi052_config_rejects_invalid_training_rtc_delay(max_delay):
     with pytest.raises(ValueError, match="rtc_training_max_delay"):

tests/test_rollout.py

@@ -98,6 +98,38 @@ def test_inference_config_types():
     assert rtc.rtc is not None
 
 
+def test_trained_rtc_retries_chunk_when_measured_delay_exceeds_conditioning():
+    from lerobot.rollout.inference.rtc import _trained_rtc_chunk_can_merge
+
+    assert not _trained_rtc_chunk_can_merge(
+        conditioned_delay=2,
+        measured_delay=3,
+        training_max_delay=4,
+        has_previous_actions=True,
+    )
+    assert _trained_rtc_chunk_can_merge(
+        conditioned_delay=2,
+        measured_delay=3,
+        training_max_delay=4,
+        has_previous_actions=False,
+    )
+
+
+def test_trained_rtc_rejects_measured_delay_above_checkpoint_support():
+    from lerobot.rollout.inference.rtc import (
+        _trained_rtc_chunk_can_merge,
+        _TrainedRTCDelayExceededError,
+    )
+
+    with pytest.raises(_TrainedRTCDelayExceededError, match="rtc_training_max_delay"):
+        _trained_rtc_chunk_can_merge(
+            conditioned_delay=3,
+            measured_delay=5,
+            training_max_delay=4,
+            has_previous_actions=True,
+        )
+
+
 def test_sentry_config_defaults():
     from lerobot.rollout import SentryStrategyConfig