diff --git a/docs/source/inference.mdx b/docs/source/inference.mdx index bdb47e978..573431a5f 100644 --- a/docs/source/inference.mdx +++ b/docs/source/inference.mdx @@ -274,12 +274,13 @@ lerobot-rollout \ --device=cuda ``` -| Flag | Description | -| ------------------------------------------- | -------------------------------------------------------------- | -| `--inference.rtc.execution_horizon` | Steps to blend with previous chunk (default: varies by policy) | -| `--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) | +| Flag | Description | +| ------------------------------------------- | ------------------------------------------------------------------------------- | +| `--inference.rtc.execution_horizon` | Steps to blend with previous chunk (default: varies by policy) | +| `--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) | 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 9308124f8..625a92ce8 100644 --- a/docs/source/pi052.mdx +++ b/docs/source/pi052.mdx @@ -116,12 +116,71 @@ the expected prompt, text target, and action endpoints before scaling up. | `policy.fast_action_loss_weight` | `1.0` | FAST cross-entropy weight | | `policy.knowledge_insulation` | `true` | Blocks action-loss gradients through the VLM K/V path | | `policy.flow_num_repeats` | `5` | Reuses one VLM prefix for independent denoising targets | +| `policy.rtc_training_max_delay` | `0` | Maximum clean-prefix delay; `0` disables training-time RTC | | `policy.lm_head_lr_scale` | `1.0` | Scales language-head learning rate; `1.0` uses the base rate | The loss weights are starting points, not dataset-independent constants. Track flow loss and text/FAST losses separately, and inspect generated subtasks rather than selecting a checkpoint from total loss alone. +### Training-time RTC + +Pi052 optionally supports training-time action conditioning from +[Training-Time Action Conditioning for Efficient Real-Time Chunking](https://arxiv.org/abs/2512.05964). +It simulates inference latency by sampling a clean action prefix for every flow +draw, passing a per-action flow timestep to the action expert, and computing the +flow loss only on the remaining postfix. The default value of `0` leaves the +standard Pi052 objective unchanged. + +```bash +lerobot-train \ + --dataset.repo_id=${HF_USER}/my_language_annotated_dataset \ + --policy.type=pi052 \ + --policy.pretrained_path=lerobot/pi05_base \ + --policy.recipe_path=recipes/subtask_mem.yaml \ + --policy.rtc_training_max_delay=10 \ + --policy.dtype=bfloat16 \ + --policy.device=cuda \ + --batch_size=8 \ + --steps=30000 \ + --output_dir=outputs/pi052_rtc \ + --job_name=pi052_rtc +``` + +`rtc_training_max_delay` is measured in controller steps and must be smaller +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. + +Run the resulting checkpoint with the asynchronous `lerobot-rollout` backend +and select the trained-prefix path explicitly: + +```bash +lerobot-rollout \ + --strategy.type=base \ + --policy.path=outputs/pi052_rtc/checkpoints/last/pretrained_model \ + --inference.type=rtc \ + --inference.rtc.mode=trained \ + --inference.rtc.execution_horizon=10 \ + --robot.type=so100_follower \ + --robot.port=/dev/ttyACM0 \ + --task="pick up the cube" \ + --fps=50 \ + --device=cuda +``` + +The rollout engine measures latency continuously, carries the still-unexecuted +actions from the previous chunk into the next prediction, and discards the +prefix that elapsed during inference. If the measured delay exceeds the +checkpoint's `rtc_training_max_delay`, rollout stops with an explicit error +instead of silently extrapolating beyond the training distribution. Use +`--inference.rtc.mode=guided` for the original Jacobian-guided RTC path; it does +not require a training-time RTC checkpoint but adds backward-pass work during +denoising. + ### Dataset-specific FAST tokenizer The universal FAST tokenizer works out of the box. For a large or diff --git a/docs/source/rtc.mdx b/docs/source/rtc.mdx index eadc34344..d92c9dc76 100644 --- a/docs/source/rtc.mdx +++ b/docs/source/rtc.mdx @@ -1,6 +1,6 @@ # Real-Time Chunking (RTC) -Real-Time Chunking (RTC) is an inference-time method that allows large, flow-matching based robotic policies, such as [Pi0](./pi0), [Pi0.5](./pi05), and [SmolVLA](./smolvla), to produce smooth, continuous, and reactive motion despite having high inference latency. +Real-Time Chunking (RTC) allows large, flow-matching based robotic policies, such as [Pi0](./pi0), [Pi0.5](./pi05), and [SmolVLA](./smolvla), to produce smooth, continuous, and reactive motion despite having high inference latency. LeRobot provides the original inference-time guided mode and, for compatible Pi052 checkpoints, training-time action conditioning with cheap hard-prefix inference. These policies generate chunks of future actions (e.g., 50 steps at a time) instead of single actions. Because the models are large, producing each chunk takes longer than the time it takes the robot to execute it. @@ -92,6 +92,11 @@ for step in range(num_steps): `RTCConfig` has the following parameters to tune: +**`mode`** selects the action-prefix conditioning method: + +- `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. + **`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 @@ -141,6 +146,7 @@ lerobot-rollout \ --strategy.type=base \ --policy.path=${HF_USERNAME}/policy_repo_id \ --inference.type=rtc \ + --inference.rtc.mode=guided \ --inference.rtc.execution_horizon=10 \ --inference.rtc.max_guidance_weight=10.0 \ --robot.type=so100_follower \ @@ -151,6 +157,24 @@ lerobot-rollout \ --device=cuda ``` +For a training-time RTC Pi052 checkpoint, change the mode to `trained`. The +checkpoint records its maximum supported delay, and rollout validates measured +latency against it: + +```bash +lerobot-rollout \ + --strategy.type=base \ + --policy.path=${HF_USERNAME}/pi052_training_rtc \ + --inference.type=rtc \ + --inference.rtc.mode=trained \ + --inference.rtc.execution_horizon=10 \ + --robot.type=so100_follower \ + --robot.port=/dev/tty.usbmodem58FA0834591 \ + --task="Move green small object into the purple platform" \ + --duration=120 \ + --device=cuda +``` + ## How It Differs from the Async Inference in LeRobot Both RTC and [async inference](./async) improve real-time robot control, but they solve different problems. @@ -189,3 +213,5 @@ See `examples/rtc/eval_dataset.py` for a complete example of offline RTC visuali - [Smooth-As-Butter Robot Policies](https://alexander-soare.github.io/robotics/2025/08/05/smooth-as-butter-robot-policies.html) - Excellent technical explanation with real robot results - [Physical Intelligence - Real-Time Chunking](https://www.physicalintelligence.company/research/real_time_chunking) - Original paper and research - [Kinetix RTC Implementation](https://github.com/Physical-Intelligence/real-time-chunking-kinetix) - Reference implementation from Physical Intelligence +- [Training-Time Action Conditioning](https://arxiv.org/abs/2512.05964) - Efficient RTC with clean-prefix conditioning during training +- [RLDX-1](https://github.com/RLWRLD/RLDX-1) - PyTorch reference used for the training-time RTC integration diff --git a/src/lerobot/policies/pi05/modeling_pi05.py b/src/lerobot/policies/pi05/modeling_pi05.py index 863edaca5..76565ff40 100644 --- a/src/lerobot/policies/pi05/modeling_pi05.py +++ b/src/lerobot/policies/pi05/modeling_pi05.py @@ -70,6 +70,54 @@ class ActionSelectKwargs(TypedDict, total=False): execution_horizon: int | None +def _prepare_trained_rtc_prefix( + x_t: Tensor, + prev_chunk_left_over: Tensor | None, + inference_delay: int, + training_max_delay: int, +) -> tuple[Tensor | None, Tensor | None]: + """Pad and validate a hard prefix for training-time RTC inference.""" + if prev_chunk_left_over is None or inference_delay <= 0: + return None, None + if training_max_delay <= 0: + raise ValueError( + "RTC mode='trained' requires a Pi052 checkpoint trained with policy.rtc_training_max_delay > 0." + ) + if inference_delay > training_max_delay: + raise ValueError( + f"Measured RTC inference delay ({inference_delay}) exceeds the checkpoint's " + f"rtc_training_max_delay ({training_max_delay})." + ) + if inference_delay >= x_t.shape[1]: + raise ValueError( + f"RTC inference delay ({inference_delay}) must be smaller than chunk_size ({x_t.shape[1]})." + ) + + previous = prev_chunk_left_over.to(device=x_t.device, dtype=x_t.dtype) + if previous.ndim == 2: + previous = previous.unsqueeze(0) + if previous.ndim != 3: + raise ValueError(f"Expected RTC prefix shape (B, T, A), got {tuple(previous.shape)}") + if previous.shape[0] == 1 and x_t.shape[0] > 1: + previous = previous.expand(x_t.shape[0], -1, -1) + if previous.shape[0] != x_t.shape[0]: + raise ValueError( + f"RTC prefix batch size ({previous.shape[0]}) does not match policy batch ({x_t.shape[0]})." + ) + if previous.shape[1] < inference_delay: + raise ValueError(f"RTC prefix has {previous.shape[1]} steps, but inference_delay={inference_delay}.") + if previous.shape[2] > x_t.shape[2]: + raise ValueError( + f"RTC prefix action dimension ({previous.shape[2]}) exceeds model dimension ({x_t.shape[2]})." + ) + + padded_prefix = torch.zeros_like(x_t) + padded_prefix[:, :inference_delay, : previous.shape[2]] = previous[:, :inference_delay] + prefix_mask = torch.arange(x_t.shape[1], device=x_t.device) < inference_delay + prefix_mask = prefix_mask[None, :, None].expand(x_t.shape[0], -1, x_t.shape[2]) + return padded_prefix, prefix_mask + + _SAFETENSORS_FILE = "model.safetensors" _SAFETENSORS_INDEX = "model.safetensors.index.json" @@ -164,21 +212,21 @@ def get_safe_dtype(target_dtype, device_type): def create_sinusoidal_pos_embedding( # see openpi `create_sinusoidal_pos_embedding` (exact copy) time: torch.Tensor, dimension: int, min_period: float, max_period: float, device="cpu" ) -> Tensor: - """Computes sine-cosine positional embedding vectors for scalar positions.""" + """Computes sine-cosine embeddings for scalar or per-action positions.""" if dimension % 2 != 0: raise ValueError(f"dimension ({dimension}) must be divisible by 2") - if time.ndim != 1: - raise ValueError("The time tensor is expected to be of shape `(batch_size, )`.") + if time.ndim not in (1, 2): + raise ValueError("The time tensor must have shape (batch_size,) or (batch_size, action_horizon).") dtype = get_safe_dtype(torch.float64, device.type) fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=dtype, device=device) period = min_period * (max_period / min_period) ** fraction - # Compute the outer product + # Broadcast the frequency dimension over either (B,) or (B, H). scaling_factor = 1.0 / period * 2 * math.pi - sin_input = scaling_factor[None, :] * time[:, None] - return torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1) + sin_input = time[..., None] * scaling_factor + return torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=-1) def sample_beta(alpha, beta, bsize, device): # see openpi `sample_beta` (exact copy) @@ -948,6 +996,24 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch` ) x_t = noise + rtc_mode = "guided" + trained_prefix = trained_prefix_mask = None + if self._rtc_enabled(): + rtc_mode = self.rtc_processor.rtc_config.mode + if rtc_mode == "trained": + 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 " + "policy.rtc_training_max_delay > 0." + ) + trained_prefix, trained_prefix_mask = _prepare_trained_rtc_prefix( + x_t, + kwargs.get("prev_chunk_left_over"), + int(kwargs.get("inference_delay") or 0), + training_max_delay, + ) + for step in range(num_steps): time = 1.0 + step * dt if times is None: @@ -955,7 +1021,13 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch` else: time_tensor = times[step].expand(bsize) - def denoise_step_partial_call(input_x_t, current_timestep=time_tensor): + denoise_timestep = time_tensor + if trained_prefix is not None: + x_t = torch.where(trained_prefix_mask, trained_prefix, x_t) + denoise_timestep = time_tensor[:, None].expand(bsize, x_t.shape[1]).clone() + denoise_timestep[trained_prefix_mask[..., 0]] = 0.0 + + def denoise_step_partial_call(input_x_t, current_timestep=denoise_timestep): return self.denoise_step( prefix_pad_masks=prefix_pad_masks, past_key_values=past_key_values, @@ -963,7 +1035,7 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch` timestep=current_timestep, ) - if self._rtc_enabled(): + if self._rtc_enabled() and rtc_mode == "guided": inference_delay = kwargs.get("inference_delay") prev_chunk_left_over = kwargs.get("prev_chunk_left_over") execution_horizon = kwargs.get("execution_horizon") @@ -980,6 +1052,8 @@ class PI05Pytorch(nn.Module): # see openpi `PI0Pytorch` v_t = denoise_step_partial_call(x_t) x_t = x_t + dt * v_t + if trained_prefix is not None: + x_t = torch.where(trained_prefix_mask, trained_prefix, x_t) if self.rtc_processor is not None and self.rtc_processor.is_debug_enabled(): self.rtc_processor.track(time=time, x_t=x_t, v_t=v_t) diff --git a/src/lerobot/policies/pi052/configuration_pi052.py b/src/lerobot/policies/pi052/configuration_pi052.py index 63f0f2808..31827de15 100644 --- a/src/lerobot/policies/pi052/configuration_pi052.py +++ b/src/lerobot/policies/pi052/configuration_pi052.py @@ -121,6 +121,15 @@ 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. + 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. + """ + # PaLM-style z-loss stabilizes large-vocabulary CE; 0 disables it. text_ce_z_loss_weight: float = 1e-4 @@ -151,6 +160,11 @@ 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}" diff --git a/src/lerobot/policies/pi052/modeling_pi052.py b/src/lerobot/policies/pi052/modeling_pi052.py index bf344f348..afe50f4d9 100644 --- a/src/lerobot/policies/pi052/modeling_pi052.py +++ b/src/lerobot/policies/pi052/modeling_pi052.py @@ -187,6 +187,54 @@ def _reduce_action_loss(per_sample: Tensor, predict_actions_t: Tensor | None, re return (per_sample * mask).sum() / mask.sum().clamp(min=1.0) +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_per_sample(flow_per_dim: Tensor, prefix_mask: Tensor | None) -> Tensor: + """Average each draw over postfix action positions and dimensions only.""" + 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) + 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 + + # Materialized logits win at VLA token counts; larger dense targets use Liger. _LOGITS_CE_MAX_POSITIONS = 2048 @@ -928,6 +976,7 @@ class PI052Policy(PI05Policy): text_labels is None and predict_actions_t is None and not getattr(self.config, "enable_fast_action_loss", False) + and self.config.rtc_training_max_delay == 0 ): return super().forward(batch, reduction=reduction) @@ -1118,12 +1167,17 @@ class PI052Policy(PI05Policy): noise = self.model.sample_noise(actions.shape, actions.device) time = self.model.sample_time(actions.shape[0], actions.device) - time_expanded = time[:, None, None] - x_t = time_expanded * noise + (1 - time_expanded) * actions + prefix_mask = _sample_training_rtc_prefix_mask( + actions.shape[0], + actions.shape[1], + self.config.rtc_training_max_delay, + actions.device, + ) + x_t, model_time = _build_flow_matching_inputs(actions, noise, time, prefix_mask) u_t = noise - actions # ---- suffix: noisy actions ---------------------------------- - suffix_embs, suffix_pad, suffix_att, adarms_cond = self.model.embed_suffix(x_t, time) + suffix_embs, suffix_pad, suffix_att, adarms_cond = self.model.embed_suffix(x_t, model_time) # ---- bf16 alignment (mirrors PI05Pytorch.forward) ----------- first_layer = self.model.paligemma_with_expert.paligemma.model.language_model.layers[0] @@ -1169,7 +1223,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_per_dim.mean(dim=(1, 2)) + per_sample_flow = _flow_loss_per_sample(flow_per_dim, prefix_mask) flow_loss = _reduce_action_loss(per_sample_flow, predict_actions_t, reduction) return prefix_out, flow_loss @@ -1216,10 +1270,15 @@ class PI052Policy(PI05Policy): noise = model.sample_noise((k * batch_size, *actions.shape[1:]), actions.device) time = model.sample_time(k * batch_size, actions.device) actions_rep = actions.repeat(k, 1, 1) # (k*B, chunk, motor_dim) - time_expanded = time[:, None, None] - x_t = time_expanded * noise + (1 - time_expanded) * actions_rep + prefix_mask_flat = _sample_training_rtc_prefix_mask( + k * batch_size, + chunk, + self.config.rtc_training_max_delay, + actions.device, + ) + x_t, model_time = _build_flow_matching_inputs(actions_rep, noise, time, prefix_mask_flat) u_t = (noise - actions_rep).view(k, batch_size, chunk, -1).transpose(0, 1) # (B, k, chunk, motor) - s_embs, suffix_pad, suffix_att, adarms = model.embed_suffix(x_t, time) + s_embs, suffix_pad, suffix_att, adarms = model.embed_suffix(x_t, model_time) if use_bf16: s_embs = s_embs.to(dtype=torch.bfloat16) suffix_pad = suffix_pad[:batch_size] @@ -1228,9 +1287,10 @@ class PI052Policy(PI05Policy): s_embs.view(k, batch_size, chunk, -1).transpose(0, 1).reshape(batch_size, k * chunk, -1) ) # (B, k*chunk, D) # Broadcast each draw's AdaRMS condition over its action chunk. + if adarms.ndim == 2: + adarms = adarms[:, None, :].expand(-1, chunk, -1) adarms_cond = ( - adarms.view(k, batch_size, 1, adarms.shape[-1]) - .expand(k, batch_size, chunk, adarms.shape[-1]) + adarms.view(k, batch_size, chunk, adarms.shape[-1]) .transpose(0, 1) .reshape(batch_size, k * chunk, adarms.shape[-1]) ) # (B, k*chunk, cond_dim) @@ -1300,7 +1360,10 @@ class PI052Policy(PI05Policy): v_t = model.action_out_proj(suffix_out.to(dtype=torch.float32)) v_t = v_t.view(batch_size, k, chunk, -1) # (B, k, chunk, motor) flow_per_dim = functional.mse_loss(u_t, v_t, reduction="none")[..., :original_action_dim] - per_sample_flow = flow_per_dim.mean(dim=(1, 2, 3)) + 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) return prefix_out, flow_loss diff --git a/src/lerobot/policies/rtc/configuration_rtc.py b/src/lerobot/policies/rtc/configuration_rtc.py index 3d71edf26..3ed9c55a0 100644 --- a/src/lerobot/policies/rtc/configuration_rtc.py +++ b/src/lerobot/policies/rtc/configuration_rtc.py @@ -37,6 +37,10 @@ class RTCConfig: # Infrastructure enabled: bool = True + # ``guided`` is the original inference-time Jacobian guidance. ``trained`` + # hard-inpaints a prefix and requires a compatible training-time RTC checkpoint. + mode: str = "guided" + # Core RTC settings # Todo change to exp prefix_attention_schedule: RTCAttentionSchedule = RTCAttentionSchedule.LINEAR @@ -49,6 +53,8 @@ class RTCConfig: def __post_init__(self): """Validate RTC configuration parameters.""" + if self.mode not in {"guided", "trained"}: + raise ValueError(f"mode must be 'guided' or 'trained', got {self.mode!r}") if self.max_guidance_weight <= 0: raise ValueError(f"max_guidance_weight must be positive, got {self.max_guidance_weight}") if self.debug_maxlen <= 0: diff --git a/src/lerobot/rollout/context.py b/src/lerobot/rollout/context.py index 20a7d715a..c0dc4393e 100644 --- a/src/lerobot/rollout/context.py +++ b/src/lerobot/rollout/context.py @@ -178,6 +178,18 @@ 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 hasattr(policy_config, "compile_model"): policy_config.compile_model = cfg.use_torch_compile diff --git a/tests/policies/pi052/test_pi052_training_time_rtc.py b/tests/policies/pi052/test_pi052_training_time_rtc.py new file mode 100644 index 000000000..347cde9d1 --- /dev/null +++ b/tests/policies/pi052/test_pi052_training_time_rtc.py @@ -0,0 +1,106 @@ +#!/usr/bin/env python + +# Copyright 2026 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. + +"""Unit coverage for Pi052 training-time RTC conditioning.""" + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +pytest.importorskip("transformers") + +from lerobot.policies.pi05.modeling_pi05 import ( # noqa: E402 + _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, +) + + +def test_training_rtc_uses_clean_prefix_and_per_token_time(): + actions = torch.tensor([[[1.0], [2.0], [3.0], [4.0]]]) + noise = torch.tensor([[[10.0], [20.0], [30.0], [40.0]]]) + time = torch.tensor([0.25]) + prefix_mask = torch.tensor([[True, True, False, False]]) + + x_t, model_time = _build_flow_matching_inputs(actions, noise, time, prefix_mask) + + assert model_time.tolist() == [[0.0, 0.0, 0.25, 0.25]] + assert torch.equal(x_t[:, :2], actions[:, :2]) + assert torch.equal(x_t[:, 2:], 0.25 * noise[:, 2:] + 0.75 * actions[:, 2:]) + + +def test_training_rtc_loss_averages_over_postfix_only(): + flow_loss = torch.tensor([[[100.0], [100.0], [2.0], [4.0]]]) + prefix_mask = torch.tensor([[True, True, False, False]]) + + per_sample = _flow_loss_per_sample(flow_loss, prefix_mask) + + assert per_sample.tolist() == [3.0] + + +def test_per_token_time_embedding_preserves_action_axis(): + time = torch.tensor([[0.0, 0.5, 1.0]]) + + embedding = create_sinusoidal_pos_embedding(time, 8, 4e-3, 4.0, time.device) + + assert embedding.shape == (1, 3, 8) + assert not torch.equal(embedding[:, 0], embedding[:, 1]) + + +def test_action_expert_embeds_per_token_flow_times(): + model = PI052Pytorch.__new__(PI052Pytorch) + nn.Module.__init__(model) + model.config = SimpleNamespace(chunk_size=3, min_period=4e-3, max_period=4.0) + model.gradient_checkpointing_enabled = False + model.action_in_proj = nn.Linear(2, 8) + model.time_mlp_in = nn.Linear(8, 8) + model.time_mlp_out = nn.Linear(8, 8) + + suffix, _, _, adarms_cond = model.embed_suffix( + torch.randn(2, 3, 2), + torch.tensor([[0.0, 0.5, 0.5], [0.0, 0.0, 0.5]]), + ) + + assert suffix.shape == (2, 3, 8) + assert adarms_cond.shape == (2, 3, 8) + + +def test_trained_rtc_prefix_is_padded_and_masked(): + x_t = torch.randn(1, 5, 4) + previous = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]) + + prefix, mask = _prepare_trained_rtc_prefix(x_t, previous, inference_delay=2, training_max_delay=3) + + assert prefix.shape == x_t.shape + assert mask.shape == x_t.shape + assert torch.equal(prefix[0, :2, :2], previous[:2]) + assert torch.count_nonzero(prefix[0, :2, 2:]) == 0 + assert mask[0, :2].all() + assert not mask[0, 2:].any() + + +def test_trained_rtc_rejects_delay_outside_training_distribution(): + with pytest.raises(ValueError, match="exceeds the checkpoint"): + _prepare_trained_rtc_prefix( + torch.randn(1, 5, 4), + torch.randn(4, 2), + inference_delay=4, + 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"): + PI052Config(chunk_size=5, n_action_steps=5, rtc_training_max_delay=max_delay) diff --git a/tests/policies/rtc/test_configuration_rtc.py b/tests/policies/rtc/test_configuration_rtc.py index 40d171c0c..88c2f1743 100644 --- a/tests/policies/rtc/test_configuration_rtc.py +++ b/tests/policies/rtc/test_configuration_rtc.py @@ -16,6 +16,8 @@ """Tests for RTC configuration module.""" +import pytest + from lerobot.configs.types import RTCAttentionSchedule from lerobot.policies.rtc.configuration_rtc import RTCConfig @@ -27,6 +29,7 @@ def test_rtc_config_default_initialization(): config = RTCConfig() assert config.enabled is True + assert config.mode == "guided" assert config.prefix_attention_schedule == RTCAttentionSchedule.LINEAR assert config.max_guidance_weight == 10.0 assert config.execution_horizon == 10 @@ -34,10 +37,16 @@ def test_rtc_config_default_initialization(): assert config.debug_maxlen == 100 +def test_rtc_config_rejects_unknown_mode(): + with pytest.raises(ValueError, match="mode must be"): + RTCConfig(mode="unknown") + + def test_rtc_config_custom_initialization(): """Test RTCConfig initializes with custom values.""" config = RTCConfig( enabled=True, + mode="trained", prefix_attention_schedule=RTCAttentionSchedule.EXP, max_guidance_weight=5.0, execution_horizon=20, @@ -46,6 +55,7 @@ def test_rtc_config_custom_initialization(): ) assert config.enabled is True + assert config.mode == "trained" assert config.prefix_attention_schedule == RTCAttentionSchedule.EXP assert config.max_guidance_weight == 5.0 assert config.execution_horizon == 20