From c395286f3c6d4638195311d7b1a4497bd3261245 Mon Sep 17 00:00:00 2001 From: Pepijn Date: Thu, 16 Jul 2026 12:30:40 +0200 Subject: [PATCH] fix(rtc): validate trained prefix capacity --- docs/source/inference.mdx | 2 +- docs/source/pi052.mdx | 5 +-- docs/source/rtc.mdx | 8 +++++ examples/rtc/eval_dataset.py | 1 + src/lerobot/policies/pi05/modeling_pi05.py | 5 ++- src/lerobot/policies/rtc/modeling_rtc.py | 6 +++- src/lerobot/rollout/context.py | 42 ++++++++++++++++------ src/lerobot/rollout/inference/rtc.py | 28 +++++++++++++-- tests/policies/rtc/test_modeling_rtc.py | 13 +++++++ tests/test_rollout.py | 33 +++++++++++++++++ 10 files changed, 125 insertions(+), 18 deletions(-) diff --git a/docs/source/inference.mdx b/docs/source/inference.mdx index 573431a5f..0244f4bb8 100644 --- a/docs/source/inference.mdx +++ b/docs/source/inference.mdx @@ -280,7 +280,7 @@ lerobot-rollout \ | `--inference.rtc.mode` | `guided` (default) or trained-prefix `trained` for compatible Pi052 checkpoints | | `--inference.rtc.max_guidance_weight` | Consistency enforcement strength (default: varies by policy) | | `--inference.rtc.prefix_attention_schedule` | Blend schedule: `LINEAR`, `EXP`, `ONES`, `ZEROS` | -| `--inference.queue_threshold` | Max queue size before backpressure (default: 30) | +| `--inference.queue_threshold` | Backpressure threshold; trained RTC requires at least its maximum delay | See the [Real-Time Chunking](./rtc) guide for details on tuning RTC parameters. diff --git a/docs/source/pi052.mdx b/docs/source/pi052.mdx index 625a92ce8..75bb0a8c8 100644 --- a/docs/source/pi052.mdx +++ b/docs/source/pi052.mdx @@ -152,8 +152,9 @@ than `chunk_size`. Choose it to cover the largest inference latency expected at deployment: at 50 Hz, for example, 10 steps correspond to 200 ms. A delay of zero is included in the uniform sampling distribution, so the checkpoint also continues to receive ordinary flow-matching examples. Set rollout's -`inference.rtc.execution_horizon` to at least this maximum so the previous -chunk cache retains enough actions to construct every supported prefix. +`inference.rtc.execution_horizon` and `inference.queue_threshold` to at least +this maximum so inference starts early enough and the previous chunk retains +every action needed for the committed prefix. Run the resulting checkpoint with the asynchronous `lerobot-rollout` backend and select the trained-prefix path explicitly: diff --git a/docs/source/rtc.mdx b/docs/source/rtc.mdx index d92c9dc76..566e8f6fe 100644 --- a/docs/source/rtc.mdx +++ b/docs/source/rtc.mdx @@ -97,6 +97,10 @@ for step in range(num_steps): - `guided` (default) applies the original Jacobian guidance during denoising and works with ordinary flow-matching checkpoints. - `trained` hard-inpaints the previous chunk's prefix with per-action flow timesteps. It currently requires a Pi052 checkpoint trained with `policy.rtc_training_max_delay > 0` and avoids the guidance backward pass. +For trained mode, both `execution_horizon` and the rollout backend's +`inference.queue_threshold` must be at least the checkpoint's +`rtc_training_max_delay`; rollout validates this before connecting the robot. + **`execution_horizon`**: How many timesteps from the previous chunk to maintain consistency with. Higher values mean smoother transitions but potentially less reactivity. Typical values: 8-12 steps @@ -129,6 +133,10 @@ python examples/rtc/eval_dataset.py \ --device=cuda ``` +Add `--rtc.mode=trained` when evaluating a compatible training-time RTC Pi052 +checkpoint. Unsupported policies reject trained mode instead of falling back to +guided RTC. + The script generates a visualization of the denoising process, comparing standard generation (left) with RTC (right). In the RTC plots, you can see how the first few steps (blue/purple lines) are guided to match the red ground truth trajectory (previous chunk's tail), ensuring a smooth transition between chunks.

diff --git a/examples/rtc/eval_dataset.py b/examples/rtc/eval_dataset.py index bacafead1..acf06323f 100644 --- a/examples/rtc/eval_dataset.py +++ b/examples/rtc/eval_dataset.py @@ -306,6 +306,7 @@ class RTCEvaluator: # Configure RTC rtc_config = RTCConfig( enabled=rtc_enabled, + mode=self.cfg.rtc.mode, execution_horizon=self.cfg.rtc.execution_horizon, max_guidance_weight=self.cfg.rtc.max_guidance_weight, prefix_attention_schedule=self.cfg.rtc.prefix_attention_schedule, diff --git a/src/lerobot/policies/pi05/modeling_pi05.py b/src/lerobot/policies/pi05/modeling_pi05.py index 17d76622f..6e1888a68 100644 --- a/src/lerobot/policies/pi05/modeling_pi05.py +++ b/src/lerobot/policies/pi05/modeling_pi05.py @@ -1284,7 +1284,10 @@ class PI05Policy(PreTrainedPolicy): # Create processor if config provided # If RTC is not enabled - we can still track the denoising data if self.config.rtc_config is not None: - self.rtc_processor = RTCProcessor(self.config.rtc_config) + self.rtc_processor = RTCProcessor( + self.config.rtc_config, + trained_mode_supported=int(getattr(self.config, "rtc_training_max_delay", 0)) > 0, + ) model_value = getattr(self, "model", None) if model_value is not None: diff --git a/src/lerobot/policies/rtc/modeling_rtc.py b/src/lerobot/policies/rtc/modeling_rtc.py index c1aeed328..fe5077c96 100644 --- a/src/lerobot/policies/rtc/modeling_rtc.py +++ b/src/lerobot/policies/rtc/modeling_rtc.py @@ -42,7 +42,11 @@ class RTCProcessor: prefix attention, and adaptive chunk processing. """ - def __init__(self, rtc_config: RTCConfig): + 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." + ) self.rtc_config = rtc_config self.tracker = None diff --git a/src/lerobot/rollout/context.py b/src/lerobot/rollout/context.py index c0dc4393e..b1632ea1f 100644 --- a/src/lerobot/rollout/context.py +++ b/src/lerobot/rollout/context.py @@ -60,6 +60,35 @@ from .robot_wrapper import ThreadSafeRobot logger = logging.getLogger(__name__) +def _validate_trained_rtc_rollout_config(policy_config, inference_config: RTCInferenceConfig) -> None: + """Fail fast when rollout cannot retain every trained RTC prefix.""" + rtc = inference_config.rtc + if not rtc.enabled or rtc.mode != "trained": + return + if policy_config.type != "pi052": + raise ValueError( + "--inference.rtc.mode=trained currently requires a Pi052 checkpoint; " + f"got policy type {policy_config.type!r}." + ) + + training_max_delay = int(getattr(policy_config, "rtc_training_max_delay", 0)) + if training_max_delay <= 0: + raise ValueError( + "--inference.rtc.mode=trained requires a checkpoint trained with " + "--policy.rtc_training_max_delay > 0." + ) + if rtc.execution_horizon < training_max_delay: + raise ValueError( + f"--inference.rtc.execution_horizon ({rtc.execution_horizon}) must be at least the " + f"checkpoint's rtc_training_max_delay ({training_max_delay})." + ) + if inference_config.queue_threshold < training_max_delay: + raise ValueError( + f"--inference.queue_threshold ({inference_config.queue_threshold}) must be at least the " + f"checkpoint's rtc_training_max_delay ({training_max_delay})." + ) + + def _resolve_action_key_order( policy_action_names: list[str] | None, dataset_action_names: list[str] ) -> list[str]: @@ -178,17 +207,8 @@ def build_rollout_context( policy_config = cfg.policy policy_class = get_policy_class(policy_config.type) - if is_rtc and cfg.inference.rtc.enabled and cfg.inference.rtc.mode == "trained": - if policy_config.type != "pi052": - raise ValueError( - "--inference.rtc.mode=trained currently requires a Pi052 checkpoint; " - f"got policy type {policy_config.type!r}." - ) - if int(getattr(policy_config, "rtc_training_max_delay", 0)) <= 0: - raise ValueError( - "--inference.rtc.mode=trained requires a checkpoint trained with " - "--policy.rtc_training_max_delay > 0." - ) + if is_rtc: + _validate_trained_rtc_rollout_config(policy_config, cfg.inference) if hasattr(policy_config, "compile_model"): policy_config.compile_model = cfg.use_torch_compile diff --git a/src/lerobot/rollout/inference/rtc.py b/src/lerobot/rollout/inference/rtc.py index 83758aa1d..8a8ffb47a 100644 --- a/src/lerobot/rollout/inference/rtc.py +++ b/src/lerobot/rollout/inference/rtc.py @@ -57,10 +57,18 @@ _RTC_MAX_CONSECUTIVE_ERRORS: int = 10 _RTC_JOIN_TIMEOUT_S: float = 3.0 -class _TrainedRTCDelayExceededError(RuntimeError): +class _FatalRTCInferenceError(RuntimeError): + """Base class for RTC errors that cannot become valid after a retry.""" + + +class _TrainedRTCDelayExceededError(_FatalRTCInferenceError): """Raised when measured latency exceeds a trained RTC checkpoint's support.""" +class _TrainedRTCPrefixUnavailableError(_FatalRTCInferenceError): + """Raised when the queue cannot provide the prefix used for conditioning.""" + + # --------------------------------------------------------------------------- # RTC helpers # --------------------------------------------------------------------------- @@ -96,6 +104,16 @@ def _trained_rtc_chunk_can_merge( return not has_previous_actions or measured_delay <= conditioned_delay +def _validate_trained_rtc_prefix_available(*, conditioned_delay: int, available_steps: int) -> None: + """Reject hard-prefix inference when the real queue is shorter than its delay.""" + if conditioned_delay > available_steps: + raise _TrainedRTCPrefixUnavailableError( + f"Trained RTC needs {conditioned_delay} committed prefix actions, but the queue has " + f"only {available_steps}. Increase --inference.queue_threshold and " + "--inference.rtc.execution_horizon." + ) + + # --------------------------------------------------------------------------- # RTCInferenceEngine # --------------------------------------------------------------------------- @@ -296,6 +314,12 @@ class RTCInferenceEngine(InferenceEngine): latency = latency_tracker.max() delay = math.ceil(latency / time_per_chunk) if latency else 0 + if self._rtc_config.mode == "trained" and delay > 0: + available_steps = 0 if prev_actions is None else prev_actions.shape[0] + _validate_trained_rtc_prefix_available( + conditioned_delay=delay, + available_steps=available_steps, + ) obs_batch = build_dataset_frame(self._hw_features, obs, prefix="observation") obs_batch = prepare_observation_for_inference( @@ -372,7 +396,7 @@ class RTCInferenceEngine(InferenceEngine): logger.debug("RTC inference latency=%.2fs, queue=%d", new_latency, queue.qsize()) - except _TrainedRTCDelayExceededError: + except _FatalRTCInferenceError: raise except Exception as e: consecutive_errors += 1 diff --git a/tests/policies/rtc/test_modeling_rtc.py b/tests/policies/rtc/test_modeling_rtc.py index e7fdc09c6..6607ea59a 100644 --- a/tests/policies/rtc/test_modeling_rtc.py +++ b/tests/policies/rtc/test_modeling_rtc.py @@ -93,6 +93,19 @@ def test_rtc_processor_initialization_without_debug(rtc_config_debug_disabled): assert processor.tracker is None +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"): + RTCProcessor(config) + + processor = RTCProcessor(config, trained_mode_supported=True) + assert processor.rtc_config.mode == "trained" + + disabled = RTCProcessor(RTCConfig(enabled=False, mode="trained")) + assert disabled.rtc_config.enabled is False + + # ====================== Tracker Proxy Methods Tests ====================== diff --git a/tests/test_rollout.py b/tests/test_rollout.py index 275694848..b6ff54276 100644 --- a/tests/test_rollout.py +++ b/tests/test_rollout.py @@ -17,6 +17,7 @@ from __future__ import annotations import dataclasses +from types import SimpleNamespace from unittest.mock import MagicMock import pytest @@ -130,6 +131,38 @@ def test_trained_rtc_rejects_measured_delay_above_checkpoint_support(): ) +def test_trained_rtc_rejects_prefix_shorter_than_conditioned_delay(): + from lerobot.rollout.inference.rtc import ( + _TrainedRTCPrefixUnavailableError, + _validate_trained_rtc_prefix_available, + ) + + with pytest.raises(_TrainedRTCPrefixUnavailableError, match="only 2"): + _validate_trained_rtc_prefix_available(conditioned_delay=4, available_steps=2) + + +@pytest.mark.parametrize( + ("execution_horizon", "queue_threshold", "match"), + [ + (3, 4, "execution_horizon"), + (4, 3, "queue_threshold"), + ], +) +def test_trained_rtc_rollout_requires_capacity_for_max_delay(execution_horizon, queue_threshold, match): + from lerobot.policies.rtc.configuration_rtc import RTCConfig + from lerobot.rollout.context import _validate_trained_rtc_rollout_config + from lerobot.rollout.inference import RTCInferenceConfig + + policy_config = SimpleNamespace(type="pi052", rtc_training_max_delay=4) + inference_config = RTCInferenceConfig( + rtc=RTCConfig(mode="trained", execution_horizon=execution_horizon), + queue_threshold=queue_threshold, + ) + + with pytest.raises(ValueError, match=match): + _validate_trained_rtc_rollout_config(policy_config, inference_config) + + def test_sentry_config_defaults(): from lerobot.rollout import SentryStrategyConfig