mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-26 03:06:01 +00:00
fix(rtc): validate trained prefix capacity
This commit is contained in:
@@ -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.
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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.
|
||||
|
||||
<p align="center">
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 ======================
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user