From fa477fca5acffdf051cf79a2180bc8a16fb4e81f Mon Sep 17 00:00:00 2001 From: Pepijn Date: Thu, 16 Jul 2026 11:18:38 +0200 Subject: [PATCH] fix(rtc): align trained mode with RLDX --- src/lerobot/policies/pi05/modeling_pi05.py | 2 + .../policies/pi052/configuration_pi052.py | 7 +-- src/lerobot/policies/pi052/modeling_pi052.py | 55 ++++++++++++------- src/lerobot/rollout/inference/rtc.py | 41 ++++++++++++++ .../pi052/test_pi052_training_time_rtc.py | 34 ++++++++++++ tests/test_rollout.py | 32 +++++++++++ 6 files changed, 145 insertions(+), 26 deletions(-) diff --git a/src/lerobot/policies/pi05/modeling_pi05.py b/src/lerobot/policies/pi05/modeling_pi05.py index 76565ff40..17d76622f 100644 --- a/src/lerobot/policies/pi05/modeling_pi05.py +++ b/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: diff --git a/src/lerobot/policies/pi052/configuration_pi052.py b/src/lerobot/policies/pi052/configuration_pi052.py index 31827de15..ad6f22c84 100644 --- a/src/lerobot/policies/pi052/configuration_pi052.py +++ b/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. - - 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. - """ + """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 diff --git a/src/lerobot/policies/pi052/modeling_pi052.py b/src/lerobot/policies/pi052/modeling_pi052.py index afe50f4d9..cd26998fe 100644 --- a/src/lerobot/policies/pi052/modeling_pi052.py +++ b/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: - """Average each draw over postfix action positions and dimensions only.""" +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: - return flow_per_dim.flatten(start_dim=1).mean(dim=1) - postfix = (~prefix_mask).unsqueeze(-1).expand_as(flow_per_dim) + flattened = flow_per_dim.flatten(start_dim=1) + 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) - return numerator / denominator + denominator = postfix.sum(dim=reduce_dims) + 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_action_loss(per_sample_flow, predict_actions_t, reduction) + flow_loss = _reduce_flow_loss(flow_per_dim, prefix_mask, 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_action_loss(per_sample_flow, predict_actions_t, reduction) + flow_loss = _reduce_flow_loss(flow_per_dim, prefix_mask, predict_actions_t, reduction) return prefix_out, flow_loss def _prefix_ce_losses( diff --git a/src/lerobot/rollout/inference/rtc.py b/src/lerobot/rollout/inference/rtc.py index 0eef62cef..83758aa1d 100644 --- a/src/lerobot/rollout/inference/rtc.py +++ b/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( diff --git a/tests/policies/pi052/test_pi052_training_time_rtc.py b/tests/policies/pi052/test_pi052_training_time_rtc.py index 347cde9d1..398203d64 100644 --- a/tests/policies/pi052/test_pi052_training_time_rtc.py +++ b/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"): diff --git a/tests/test_rollout.py b/tests/test_rollout.py index 85a29ff4c..275694848 100644 --- a/tests/test_rollout.py +++ b/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