From 255a01234ed9183e10bb772c3b6f42f41ccd114f Mon Sep 17 00:00:00 2001 From: Maxime Ellerbach Date: Mon, 20 Jul 2026 16:19:08 +0000 Subject: [PATCH] fastwam rtc --- .../policies/fastwam/configuration_fastwam.py | 7 +++ .../policies/fastwam/modeling_fastwam.py | 37 +++++++++++++- src/lerobot/policies/fastwam/wan/modular.py | 51 +++++++++++++++---- 3 files changed, 85 insertions(+), 10 deletions(-) diff --git a/src/lerobot/policies/fastwam/configuration_fastwam.py b/src/lerobot/policies/fastwam/configuration_fastwam.py index 9e76a2eb1..61b52366c 100644 --- a/src/lerobot/policies/fastwam/configuration_fastwam.py +++ b/src/lerobot/policies/fastwam/configuration_fastwam.py @@ -27,6 +27,8 @@ from lerobot.configs import ( from lerobot.optim import AdamWConfig from lerobot.utils.constants import ACTION, OBS_STATE +from ..rtc.configuration_rtc import RTCConfig + WAN22_MODEL_ID = "Wan-AI/Wan2.2-TI2V-5B" WAN22_DIFFUSERS_MODEL_ID = "Wan-AI/Wan2.2-TI2V-5B-Diffusers" FASTWAM_BASE_MODEL_ID = "lerobot/fastwam_base" @@ -214,6 +216,11 @@ class FastWAMConfig(PreTrainedConfig): ) num_inference_steps: int = 10 inference_seed: int | None = 42 + # Real-Time Chunking (RTC): async chunk generation with prefix guidance so a new chunk + # inpaints smoothly onto the still-executing tail of the previous one. `None` disables it + # (default synchronous inference). Consumed by `RTCInferenceEngine`, which calls + # `predict_action_chunk(..., inference_delay=, prev_chunk_left_over=)`. + rtc_config: RTCConfig | None = None rand_device: str = "cpu" text_cfg_scale: float = 1.0 negative_prompt: str = "" diff --git a/src/lerobot/policies/fastwam/modeling_fastwam.py b/src/lerobot/policies/fastwam/modeling_fastwam.py index dbe39e243..25326ecce 100644 --- a/src/lerobot/policies/fastwam/modeling_fastwam.py +++ b/src/lerobot/policies/fastwam/modeling_fastwam.py @@ -22,6 +22,7 @@ import torch from torch import Tensor from lerobot.policies.pretrained import PreTrainedPolicy +from lerobot.policies.rtc.modeling_rtc import RTCProcessor from lerobot.utils.constants import OBS_STATE from lerobot.utils.import_utils import require_package @@ -85,8 +86,23 @@ class FastWAMPolicy(PreTrainedPolicy): for layer in mot.layers: if "video" in layer.blocks: layer.blocks["video"].requires_grad_(False) + self._init_rtc_processor() self.reset() + def _init_rtc_processor(self) -> None: + """Attach a Real-Time Chunking processor to the core model when configured. + + Mirrors the PI0/PI05 pattern: the policy owns the `RTCProcessor` and hands it to + the core `FastWAM` model, which consults it inside `infer_action`'s denoising loop. + """ + self.rtc_processor = None + if self.config.rtc_config is not None: + self.rtc_processor = RTCProcessor(self.config.rtc_config) + self.model.rtc_processor = self.rtc_processor + + def _rtc_enabled(self) -> bool: + return self.config.rtc_config is not None and self.config.rtc_config.enabled + @classmethod def _load_as_safetensor(cls, model, model_file: str, map_location: str, strict: bool): """Shape-aware load that supports cross-embodiment fine-tuning. @@ -227,12 +243,26 @@ class FastWAMPolicy(PreTrainedPolicy): return loss, dict(metrics or {}) @torch.no_grad() - def predict_action_chunk(self, batch: dict[str, Tensor], **_: Any) -> Tensor: + def predict_action_chunk( + self, + batch: dict[str, Tensor], + inference_delay: int | None = None, + prev_chunk_left_over: Tensor | None = None, + execution_horizon: int | None = None, + **_: Any, + ) -> Tensor: """Predict a chunk of actions from the current FastWAM observation. Args: batch (dict[str, Tensor]): Inference batch with `input_image` or image observation keys, plus `context/context_mask` or `prompt`. + inference_delay (int | None): RTC — number of prefix steps assumed already + executed by the time this chunk lands (from measured inference latency). + prev_chunk_left_over (Tensor | None): RTC — the previous chunk's unexecuted + action tail `[T_prev, action_dim]` in model space; guides denoising so the + new chunk inpaints onto it. `None` (default) = plain synchronous inference. + execution_horizon (int | None): RTC — override for the prefix-weight horizon; + `None` falls back to `rtc_config.execution_horizon`. Returns: Tensor: Action chunk with shape `[B, action_horizon, action_dim]`. @@ -249,6 +279,11 @@ class FastWAMPolicy(PreTrainedPolicy): infer_kwargs["prompt"] = None infer_kwargs["context"] = context infer_kwargs["context_mask"] = context_mask + # RTC guidance args flow straight to `infer_action`; they are inert unless an + # RTCProcessor is attached, enabled, and `prev_chunk_left_over` is provided. + infer_kwargs["inference_delay"] = inference_delay + infer_kwargs["prev_chunk_left_over"] = prev_chunk_left_over + infer_kwargs["execution_horizon"] = execution_horizon batch_size = _infer_kwargs_batch_size(infer_kwargs) if batch_size == 1: action = _action_from_model_output(self.model.infer_action(**infer_kwargs)) diff --git a/src/lerobot/policies/fastwam/wan/modular.py b/src/lerobot/policies/fastwam/wan/modular.py index 80cac4be8..25e7760fe 100644 --- a/src/lerobot/policies/fastwam/wan/modular.py +++ b/src/lerobot/policies/fastwam/wan/modular.py @@ -908,6 +908,11 @@ class FastWAM(torch.nn.Module): self.train_scheduler = self.train_video_scheduler self.infer_scheduler = self.infer_video_scheduler + # Optional Real-Time Chunking processor (set by the policy wrapper). When present and + # enabled it guides the action denoising loop in `infer_action` so a freshly generated + # chunk inpaints onto the previous chunk's unexecuted tail. Plain attribute (not an + # nn.Module) — carries no parameters and stays out of state_dict / device moves. + self.rtc_processor = None self.device = torch.device(device) # When pinned (e.g. "cpu"), the frozen text encoder stays on this device instead # of following the model onto the GPU — `_apply` skips it and `encode_prompt` runs @@ -1825,6 +1830,9 @@ class FastWAM(torch.nn.Module): seed: int | None = None, rand_device: str = "cpu", tiled: bool = False, + inference_delay: int | None = None, + prev_chunk_left_over: torch.Tensor | None = None, + execution_horizon: int | None = None, ) -> dict[str, Any]: self.eval() if str(getattr(self.video_expert, "video_attention_mask_mode", "")) != "first_frame_causal": @@ -1877,18 +1885,43 @@ class FastWAM(torch.nn.Module): dtype=latents_action.dtype, shift_override=sigma_shift, ) + rtc_active = ( + self.rtc_processor is not None + and getattr(self.rtc_processor.rtc_config, "enabled", False) + and prev_chunk_left_over is not None + ) + num_train_timesteps = float(self.infer_action_scheduler.num_train_timesteps) for step_t_action, step_delta_action in zip(infer_timesteps_action, infer_deltas_action, strict=True): timestep_action = step_t_action.unsqueeze(0).to(dtype=latents_action.dtype, device=self.device) - pred_action = self._predict_action_noise_with_cache( - latents_action=latents_action, - timestep_action=timestep_action, - context=context, - context_mask=context_mask, - video_kv_cache=video_kv_cache, - attention_mask=attention_mask, - video_seq_len=video_seq_len, - ) + def denoise(x_t, ts=timestep_action): + return self._predict_action_noise_with_cache( + latents_action=x_t, + timestep_action=ts, + context=context, + context_mask=context_mask, + video_kv_cache=video_kv_cache, + attention_mask=attention_mask, + video_seq_len=video_seq_len, + ) + + if rtc_active: + # `time` is the flow-matching noise level sigma in [0, 1]: FastWAM's model + # predicts velocity v = noise - clean, so the clean-action estimate is + # x1 = x_t - sigma * v — exactly RTC's `x1_t = x_t - time * v_t`. + sigma = float(step_t_action.item()) / num_train_timesteps + pred_action = self.rtc_processor.denoise_step( + x_t=latents_action, + prev_chunk_left_over=prev_chunk_left_over.to( + device=latents_action.device, dtype=latents_action.dtype + ), + inference_delay=inference_delay or 0, + time=sigma, + original_denoise_step_partial=denoise, + execution_horizon=execution_horizon, + ) + else: + pred_action = denoise(latents_action) latents_action = self.infer_action_scheduler.step(pred_action, step_delta_action, latents_action)