fastwam rtc

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 = ""

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))

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)