Add Real-Time Chunking (RTC) support for flow matching models

Implement Real-Time Chunking (RTC) for action chunking policies using flow
matching denoising. RTC enables smooth action transitions between consecutive
chunks by using prefix guidance during denoising.

Key features:
- RTCProcessor class with denoise_step method for RTC guidance
- Tracker system for debug tracking using time-based dictionary storage
- RTCDebugVisualizer with comprehensive visualization utilities
- Integration with SmolVLA policy for flow matching models
- Support for multiple prefix attention schedules (ZEROS, ONES, LINEAR, EXP)
- Configurable execution horizon and max guidance weight
- Example scripts for dataset evaluation and real-time control

Technical details:
- Uses autograd-based gradient computation for RTC corrections
- Time-based tracking eliminates duplicate step issues
- Proxy methods in RTCProcessor for cleaner API
- Full integration with LeRobot's policy and dataset systems

Files added/modified:
- src/lerobot/configs/types.py: Add RTCAttentionSchedule enum
- src/lerobot/policies/rtc/: Core RTC implementation
  - configuration_rtc.py: RTC configuration
  - modeling_rtc.py: RTCProcessor with denoise_step
  - debug_handler.py: Tracker for debug information
  - debug_visualizer.py: Visualization utilities
- src/lerobot/policies/smolvla/modeling_smolvla.py: RTC integration
- examples/rtc/: Example scripts and evaluation tools

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Alexander Soare <alexander.soare159@gmail.com>
Co-Authored-By: Claude <noreply@anthropic.com>

src/lerobot/configs/types.py

@@ -43,3 +43,10 @@ class NormalizationMode(str, Enum):
 class PolicyFeature:
     type: FeatureType
     shape: tuple[int, ...]
+
+
+class RTCAttentionSchedule(str, Enum):
+    ZEROS = "ZEROS"
+    ONES = "ONES"
+    LINEAR = "LINEAR"
+    EXP = "EXP"

src/lerobot/policies/rtc/README.md

@@ -0,0 +1,28 @@
+# Real-Time Chunking (RTC) Module
+
+This module implements Real-Time Chunking and related adaptive inference techniques for robotics policies in LeRobot.
+
+## Overview
+
+Real-Time Chunking (RTC) addresses the challenge of real-time inference in action chunking policies by treating chunk generation as an inpainting problem. It strategically handles overlapping timesteps between action chunks using prefix attention mechanisms.
+
+It is particularly effective for handling long-horizon inference in robotics policies.
+
+## Integration with Policies
+
+RTC can be integrated with any policy that supports flow mathicng for chunking:
+
+- **SmolVLA**: Vision-language-action model with RTC support
+- **Pi0**: Action prediction model with adaptive chunking
+
+## Original Implementation
+
+This implementation is based on Physical Intelligence's Kinetix RTC:
+
+- [Original RTC implementation](https://github.com/Physical-Intelligence/real-time-chunking-kinetix/blob/main/src/model.py#L214)
+- [Kinetix GitHub Repository](https://github.com/Physical-Intelligence/real-time-chunking-kinetix)
+
+## References
+
+- [Real Time Chunking Paper](https://www.physicalintelligence.company/research/real_time_chunking)
+- [Physical Intelligence Kinetix](https://github.com/Physical-Intelligence/real-time-chunking-kinetix)

src/lerobot/policies/rtc/configuration_rtc.py

@@ -0,0 +1,55 @@
+#!/usr/bin/env python
+
+# Copyright 2025 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Real Time Chunking (RTC) and Bidirectional Decoding (BID) configuration classes.
+
+Based on:
+- Real Time Chunking: https://www.physicalintelligence.company/research/real_time_chunking
+"""
+
+from dataclasses import dataclass
+
+from lerobot.configs.types import RTCAttentionSchedule
+
+
+@dataclass
+class RTCConfig:
+    """Configuration for Real Time Chunking (RTC) inference.
+
+    RTC improves real-time inference by treating chunk generation as an inpainting problem,
+    strategically handling overlapping timesteps between action chunks using prefix attention.
+    """
+
+    # Infrastructure
+    enabled: bool = False
+
+    # Core RTC settings
+    # Todo change to exp
+    prefix_attention_schedule: RTCAttentionSchedule = RTCAttentionSchedule.LINEAR
+    max_guidance_weight: float = 1.0
+    execution_horizon: int = 10
+
+    # Debug settings
+    debug: bool = False
+    debug_maxlen: int = 100
+
+    def __post_init__(self):
+        """Validate RTC configuration parameters."""
+        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:
+            raise ValueError(f"debug_maxlen must be positive, got {self.debug_maxlen}")

src/lerobot/policies/rtc/debug_handler.py

@@ -0,0 +1,339 @@
+#!/usr/bin/env python
+
+# Copyright 2025 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Debug information handler for Real-Time Chunking (RTC)."""
+
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch
+from torch import Tensor
+
+
+@dataclass
+class DebugStep:
+    """Container for debug information from a single denoising step.
+
+    Attributes:
+        step_idx (int): Step index/counter.
+        x_t (Tensor | None): Current latent/state tensor.
+        v_t (Tensor | None): Velocity from denoiser.
+        x1_t (Tensor | None): Denoised prediction (x_t - time * v_t).
+        correction (Tensor | None): Correction gradient tensor.
+        err (Tensor | None): Weighted error term.
+        weights (Tensor | None): Prefix attention weights.
+        guidance_weight (float | Tensor | None): Applied guidance weight.
+        time (float | Tensor | None): Time parameter.
+        inference_delay (int | None): Inference delay parameter.
+        execution_horizon (int | None): Execution horizon parameter.
+        metadata (dict[str, Any]): Additional metadata.
+    """
+
+    step_idx: int = 0
+    x_t: Tensor | None = None
+    v_t: Tensor | None = None
+    x1_t: Tensor | None = None
+    correction: Tensor | None = None
+    err: Tensor | None = None
+    weights: Tensor | None = None
+    guidance_weight: float | Tensor | None = None
+    time: float | Tensor | None = None
+    inference_delay: int | None = None
+    execution_horizon: int | None = None
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+    def to_dict(self, include_tensors: bool = False) -> dict[str, Any]:
+        """Convert debug step to dictionary.
+
+        Args:
+            include_tensors (bool): If True, include tensor values. If False, only include
+                tensor statistics (shape, mean, std, min, max).
+
+        Returns:
+            Dictionary representation of the debug step.
+        """
+        result = {
+            "step_idx": self.step_idx,
+            "guidance_weight": (
+                self.guidance_weight.item()
+                if isinstance(self.guidance_weight, Tensor)
+                else self.guidance_weight
+            ),
+            "time": self.time.item() if isinstance(self.time, Tensor) else self.time,
+            "inference_delay": self.inference_delay,
+            "execution_horizon": self.execution_horizon,
+            "metadata": self.metadata.copy(),
+        }
+
+        # Add tensor information
+        tensor_fields = ["x_t", "v_t", "x1_t", "correction", "err", "weights"]
+        for field_name in tensor_fields:
+            tensor = getattr(self, field_name)
+            if tensor is not None:
+                if include_tensors:
+                    result[field_name] = tensor.detach().cpu()
+                else:
+                    result[f"{field_name}_stats"] = {
+                        "shape": tuple(tensor.shape),
+                        "mean": tensor.mean().item(),
+                        "std": tensor.std().item(),
+                        "min": tensor.min().item(),
+                        "max": tensor.max().item(),
+                    }
+
+        return result
+
+
+class Tracker:
+    """Collects and manages debug information for RTC processing.
+
+    This tracker stores debug information from recent denoising steps in a dictionary,
+    using time as the key for efficient lookups and updates.
+
+    Args:
+        enabled (bool): Whether debug collection is enabled.
+        maxlen (int | None): Optional sliding window size. If provided, only the
+            most recent ``maxlen`` debug steps are kept. If ``None``, keeps all.
+    """
+
+    def __init__(self, enabled: bool = False, maxlen: int = 100):
+        self.enabled = enabled
+        self._steps = {} if enabled else None  # Dictionary with time as key
+        self._maxlen = maxlen
+        self._step_counter = 0
+
+    def reset(self) -> None:
+        """Clear all recorded debug information."""
+        if self.enabled and self._steps is not None:
+            self._steps.clear()
+        self._step_counter = 0
+
+    def track(
+        self,
+        time: float | Tensor,
+        x_t: Tensor | None = None,
+        v_t: Tensor | None = None,
+        x1_t: Tensor | None = None,
+        correction: Tensor | None = None,
+        err: Tensor | None = None,
+        weights: Tensor | None = None,
+        guidance_weight: float | Tensor | None = None,
+        inference_delay: int | None = None,
+        execution_horizon: int | None = None,
+        **metadata,
+    ) -> None:
+        """Track debug information for a denoising step at a given time.
+
+        If a step with the given time already exists, it will be updated with the new data.
+        Otherwise, a new step will be created. Only non-None fields are updated/set.
+
+        Args:
+            time (float | Tensor): Time parameter - used as the key to identify the step.
+            x_t (Tensor | None): Current latent/state tensor.
+            v_t (Tensor | None): Velocity from denoiser.
+            x1_t (Tensor | None): Denoised prediction.
+            correction (Tensor | None): Correction gradient tensor.
+            err (Tensor | None): Weighted error term.
+            weights (Tensor | None): Prefix attention weights.
+            guidance_weight (float | Tensor | None): Applied guidance weight.
+            inference_delay (int | None): Inference delay parameter.
+            execution_horizon (int | None): Execution horizon parameter.
+            **metadata: Additional metadata to store.
+        """
+        if not self.enabled:
+            return
+
+        # Convert time to float and round to avoid float precision issues
+        time_value = time.item() if isinstance(time, Tensor) else time
+        time_key = round(time_value, 6)  # Use rounded time as dictionary key
+
+        # Check if step with this time already exists
+        if time_key in self._steps:
+            # Update existing step with non-None fields
+            existing_step = self._steps[time_key]
+            if x_t is not None:
+                existing_step.x_t = x_t.detach().clone()
+            if v_t is not None:
+                existing_step.v_t = v_t.detach().clone()
+            if x1_t is not None:
+                existing_step.x1_t = x1_t.detach().clone()
+            if correction is not None:
+                existing_step.correction = correction.detach().clone()
+            if err is not None:
+                existing_step.err = err.detach().clone()
+            if weights is not None:
+                existing_step.weights = weights.detach().clone()
+            if guidance_weight is not None:
+                existing_step.guidance_weight = guidance_weight
+            if inference_delay is not None:
+                existing_step.inference_delay = inference_delay
+            if execution_horizon is not None:
+                existing_step.execution_horizon = execution_horizon
+            if metadata:
+                existing_step.metadata.update(metadata)
+        else:
+            # Create new step
+            step = DebugStep(
+                step_idx=self._step_counter,
+                x_t=x_t.detach().clone() if x_t is not None else None,
+                v_t=v_t.detach().clone() if v_t is not None else None,
+                x1_t=x1_t.detach().clone() if x1_t is not None else None,
+                correction=correction.detach().clone() if correction is not None else None,
+                err=err.detach().clone() if err is not None else None,
+                weights=weights.detach().clone() if weights is not None else None,
+                guidance_weight=guidance_weight,
+                time=time_value,
+                inference_delay=inference_delay,
+                execution_horizon=execution_horizon,
+                metadata=metadata,
+            )
+
+            # Add to dictionary
+            self._steps[time_key] = step
+            self._step_counter += 1
+
+            # Enforce maxlen if set
+            if self._maxlen is not None and len(self._steps) > self._maxlen:
+                # Remove oldest entry (first key in dict - Python 3.7+ preserves insertion order)
+                oldest_key = next(iter(self._steps))
+                del self._steps[oldest_key]
+
+    def get_recent_steps(self, n: int = 1) -> list[DebugStep]:
+        """Get the n most recent debug steps.
+
+        Args:
+            n (int): Number of recent steps to retrieve.
+
+        Returns:
+            List of DebugStep objects (may be empty if disabled or no steps recorded).
+        """
+        if not self.enabled or self._steps is None:
+            return []
+
+        # Get all values and return the last n
+        all_steps = list(self._steps.values())
+        return all_steps[-n:]
+
+    def get_all_steps(self) -> list[DebugStep]:
+        """Get all recorded debug steps.
+
+        Returns:
+            List of all DebugStep objects (may be empty if disabled).
+        """
+        if not self.enabled or self._steps is None:
+            return []
+
+        return list(self._steps.values())
+
+    def get_step_stats_summary(self) -> dict[str, Any]:
+        """Get summary statistics across all recorded steps.
+
+        Returns:
+            Dictionary containing aggregate statistics.
+        """
+        if not self.enabled or self._steps is None or len(self._steps) == 0:
+            return {"enabled": self.enabled, "total_steps": 0}
+
+        # Aggregate statistics from dictionary values
+        corrections = [s.correction for s in self._steps.values() if s.correction is not None]
+        errors = [s.err for s in self._steps.values() if s.err is not None]
+        guidance_weights = [s.guidance_weight for s in self._steps.values() if s.guidance_weight is not None]
+
+        summary = {
+            "enabled": self.enabled,
+            "total_steps": len(self._steps),
+            "step_counter": self._step_counter,
+        }
+
+        if corrections:
+            correction_norms = torch.tensor([c.norm().item() for c in corrections])
+            summary["correction_norms"] = {
+                "mean": correction_norms.mean().item(),
+                "std": correction_norms.std().item(),
+                "min": correction_norms.min().item(),
+                "max": correction_norms.max().item(),
+            }
+
+        if errors:
+            error_norms = torch.tensor([e.norm().item() for e in errors])
+            summary["error_norms"] = {
+                "mean": error_norms.mean().item(),
+                "std": error_norms.std().item(),
+                "min": error_norms.min().item(),
+                "max": error_norms.max().item(),
+            }
+
+        if guidance_weights:
+            gw_tensor = torch.tensor([gw.item() if isinstance(gw, Tensor) else gw for gw in guidance_weights])
+            summary["guidance_weights"] = {
+                "mean": gw_tensor.mean().item(),
+                "std": gw_tensor.std().item(),
+                "min": gw_tensor.min().item(),
+                "max": gw_tensor.max().item(),
+            }
+
+        return summary
+
+    def export_to_dict(self, include_tensors: bool = False) -> dict[str, Any]:
+        """Export all debug information to a dictionary.
+
+        Args:
+            include_tensors (bool): If True, include full tensor values. If False,
+                only include tensor statistics.
+
+        Returns:
+            Dictionary containing all debug information.
+        """
+        if not self.enabled or self._steps is None:
+            return {"enabled": False, "steps": []}
+
+        return {
+            "enabled": True,
+            "total_steps": len(self._steps),
+            "step_counter": self._step_counter,
+            "steps": [step.to_dict(include_tensors=include_tensors) for step in self._steps.values()],
+        }
+
+    def __len__(self) -> int:
+        """Return the number of recorded debug steps."""
+        if not self.enabled or self._steps is None:
+            return 0
+        return len(self._steps)
+
+    @staticmethod
+    def tensor_stats(tensor: Tensor, name: str = "tensor") -> str:
+        """Generate readable statistics string for a tensor.
+
+        Args:
+            tensor: Input tensor
+            name: Name to display
+
+        Returns:
+            Formatted string with shape and statistics
+        """
+        if tensor is None:
+            return f"{name}: None"
+
+        stats = (
+            f"{name}: shape={tuple(tensor.shape)}, "
+            f"dtype={tensor.dtype}, "
+            f"device={tensor.device}, "
+            f"min={tensor.min().item():.4f}, "
+            f"max={tensor.max().item():.4f}, "
+            f"mean={tensor.mean().item():.4f}, "
+            f"std={tensor.std().item():.4f}"
+        )
+        return stats

src/lerobot/policies/rtc/debug_visualizer.py

@@ -0,0 +1,460 @@
+#!/usr/bin/env python
+
+# Copyright 2025 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Visualization utilities for RTC debug information."""
+
+import matplotlib.pyplot as plt
+import torch
+from matplotlib.figure import Figure
+from torch import Tensor
+
+from lerobot.policies.rtc.debug_handler import Tracker
+
+
+class RTCDebugVisualizer:
+    """Visualizer for RTC debug information.
+
+    This class provides methods to visualize debug information collected by the Tracker,
+    including corrections, errors, weights, and guidance weights over denoising steps.
+    """
+
+    @staticmethod
+    def plot_waypoints(
+        axes,
+        tensor,
+        start_from: int = 0,
+        color: str = "blue",
+        label: str = "",
+        alpha: float = 0.7,
+        linewidth: float = 2,
+        marker: str | None = None,
+        markersize: int = 4,
+    ):
+        """Plot trajectories across multiple dimensions.
+
+        This function plots a tensor's values across time for multiple dimensions,
+        with each dimension plotted on a separate axis.
+
+        Args:
+            axes: Array of matplotlib axes (one for each dimension).
+            tensor: The tensor to plot (can be torch.Tensor or numpy array).
+                   Shape should be (time_steps, num_dims) or (batch, time_steps, num_dims).
+            start_from: Starting index for the x-axis.
+            color: Color for the plot lines.
+            label: Label for the plot legend.
+            alpha: Transparency level for the plot.
+            linewidth: Width of the plot lines.
+            marker: Marker style for data points (e.g., 'o', 's', '^').
+            markersize: Size of the markers.
+        """
+        import numpy as np
+        import torch
+
+        # Handle None tensor
+        if tensor is None:
+            return
+
+        # Convert tensor to numpy if needed
+        tensor_np = tensor.detach().cpu().numpy() if isinstance(tensor, torch.Tensor) else tensor
+
+        # Handle different tensor shapes
+        if tensor_np.ndim == 3:
+            # If batch dimension present, take first batch
+            tensor_np = tensor_np[0]
+        elif tensor_np.ndim == 1:
+            # If 1D, reshape to (time_steps, 1)
+            tensor_np = tensor_np.reshape(-1, 1)
+
+        # Get dimensions
+        time_steps, num_dims = tensor_np.shape
+
+        # Create x-axis indices
+        x_indices = np.arange(start_from, start_from + time_steps)
+
+        # Plot each dimension on its corresponding axis
+        num_axes = len(axes) if hasattr(axes, "__len__") else 1
+        for dim_idx in range(min(num_dims, num_axes)):
+            ax = axes[dim_idx] if hasattr(axes, "__len__") else axes
+
+            # Plot the trajectory
+            if marker:
+                ax.plot(
+                    x_indices,
+                    tensor_np[:, dim_idx],
+                    color=color,
+                    label=label if dim_idx == 0 else "",  # Only show label once
+                    alpha=alpha,
+                    linewidth=linewidth,
+                    marker=marker,
+                    markersize=markersize,
+                )
+            else:
+                ax.plot(
+                    x_indices,
+                    tensor_np[:, dim_idx],
+                    color=color,
+                    label=label if dim_idx == 0 else "",  # Only show label once
+                    alpha=alpha,
+                    linewidth=linewidth,
+                )
+
+            # Add grid and labels if not already present
+            if not ax.xaxis.get_label().get_text():
+                ax.set_xlabel("Step", fontsize=10)
+            if not ax.yaxis.get_label().get_text():
+                ax.set_ylabel(f"Dim {dim_idx}", fontsize=10)
+            ax.grid(True, alpha=0.3)
+
+            # Add legend if label provided and this is the first dimension
+            if label and dim_idx == 0:
+                ax.legend(loc="best", fontsize=8)
+
+    @staticmethod
+    def plot_debug_summary(
+        tracker: Tracker,
+        save_path: str | None = None,
+        show: bool = False,
+        figsize: tuple[int, int] = (16, 12),
+    ) -> Figure:
+        """Create a comprehensive summary plot of debug information.
+
+        Args:
+            tracker (Tracker): Tracker with recorded steps.
+            save_path (str | None): Path to save the figure. If None, figure is not saved.
+            show (bool): Whether to display the figure.
+            figsize (tuple[int, int]): Figure size in inches (width, height).
+
+        Returns:
+            Figure: The matplotlib figure object.
+        """
+        if not tracker.enabled or len(tracker) == 0:
+            print("Tracker is disabled or has no recorded steps.")
+            return None
+
+        steps = tracker.get_all_steps()
+        num_steps = len(steps)
+
+        # Create figure with subplots
+        fig, axes = plt.subplots(3, 2, figsize=figsize)
+        fig.suptitle(f"RTC Debug Summary ({num_steps} steps)", fontsize=16, fontweight="bold")
+
+        # Plot 1: Correction norms over steps
+        ax = axes[0, 0]
+        correction_norms = [step.correction.norm().item() for step in steps if step.correction is not None]
+        if correction_norms:
+            ax.plot(correction_norms, marker="o", linewidth=2, markersize=4)
+            ax.set_xlabel("Step Index", fontsize=12)
+            ax.set_ylabel("Correction Norm", fontsize=12)
+            ax.set_title("Correction Magnitude Over Steps", fontsize=13, fontweight="bold")
+            ax.grid(True, alpha=0.3)
+
+        # Plot 2: Error norms over steps
+        ax = axes[0, 1]
+        error_norms = [step.err.norm().item() for step in steps if step.err is not None]
+        if error_norms:
+            ax.plot(error_norms, marker="o", linewidth=2, markersize=4, color="orange")
+            ax.set_xlabel("Step Index", fontsize=12)
+            ax.set_ylabel("Error Norm", fontsize=12)
+            ax.set_title("Error Magnitude Over Steps", fontsize=13, fontweight="bold")
+            ax.grid(True, alpha=0.3)
+
+        # Plot 3: Guidance weights over steps
+        ax = axes[1, 0]
+        guidance_weights = [
+            step.guidance_weight.item() if isinstance(step.guidance_weight, Tensor) else step.guidance_weight
+            for step in steps
+            if step.guidance_weight is not None
+        ]
+        if guidance_weights:
+            ax.plot(guidance_weights, marker="o", linewidth=2, markersize=4, color="green")
+            ax.set_xlabel("Step Index", fontsize=12)
+            ax.set_ylabel("Guidance Weight", fontsize=12)
+            ax.set_title("Guidance Weight Over Steps", fontsize=13, fontweight="bold")
+            ax.grid(True, alpha=0.3)
+
+        # Plot 4: Time parameter over steps
+        ax = axes[1, 1]
+        times = [
+            step.time.item() if isinstance(step.time, Tensor) else step.time
+            for step in steps
+            if step.time is not None
+        ]
+        if times:
+            ax.plot(times, marker="o", linewidth=2, markersize=4, color="purple")
+            ax.set_xlabel("Step Index", fontsize=12)
+            ax.set_ylabel("Time Parameter", fontsize=12)
+            ax.set_title("Time Parameter Over Steps", fontsize=13, fontweight="bold")
+            ax.grid(True, alpha=0.3)
+
+        # Plot 5: Correction vs Error relationship
+        ax = axes[2, 0]
+        if correction_norms and error_norms:
+            ax.scatter(error_norms, correction_norms, alpha=0.6, s=50)
+            ax.set_xlabel("Error Norm", fontsize=12)
+            ax.set_ylabel("Correction Norm", fontsize=12)
+            ax.set_title("Correction vs Error", fontsize=13, fontweight="bold")
+            ax.grid(True, alpha=0.3)
+
+        # Plot 6: Prefix attention weights visualization (last step)
+        ax = axes[2, 1]
+        last_step = steps[-1]
+        if last_step.weights is not None:
+            weights = last_step.weights.squeeze().cpu().numpy()
+            ax.plot(weights, linewidth=2, marker="o", markersize=4, color="red")
+            ax.set_xlabel("Time Index", fontsize=12)
+            ax.set_ylabel("Weight Value", fontsize=12)
+            ax.set_title("Prefix Attention Weights (Last Step)", fontsize=13, fontweight="bold")
+            ax.grid(True, alpha=0.3)
+            ax.set_ylim(-0.1, 1.1)
+
+        plt.tight_layout()
+
+        if save_path:
+            fig.savefig(save_path, dpi=150, bbox_inches="tight")
+            print(f"Debug summary saved to {save_path}")
+
+        if show:
+            plt.show()
+        else:
+            plt.close(fig)
+
+        return fig
+
+    @staticmethod
+    def plot_correction_heatmap(
+        tracker: Tracker,
+        save_path: str | None = None,
+        show: bool = False,
+        figsize: tuple[int, int] = (14, 8),
+        max_dims: int = 6,
+    ) -> Figure:
+        """Create a heatmap showing correction values across steps and action dimensions.
+
+        Args:
+            tracker (Tracker): Tracker with recorded steps.
+            save_path (str | None): Path to save the figure.
+            show (bool): Whether to display the figure.
+            figsize (tuple[int, int]): Figure size in inches.
+            max_dims (int): Maximum number of action dimensions to visualize.
+
+        Returns:
+            Figure: The matplotlib figure object.
+        """
+        if not tracker.enabled or len(tracker) == 0:
+            print("Tracker is disabled or has no recorded steps.")
+            return None
+
+        steps = tracker.get_all_steps()
+
+        # Collect corrections across steps (shape: [num_steps, time, action_dim])
+        corrections = [step.correction for step in steps if step.correction is not None]
+        if not corrections:
+            print("No corrections found in debug steps.")
+            return None
+
+        # Stack corrections: [num_steps, time, action_dim]
+        # Take mean over time dimension and limit action dims
+        corrections_stacked = torch.stack(corrections)  # [num_steps, batch, time, action_dim]
+        corrections_mean = corrections_stacked.mean(dim=(1, 2))  # [num_steps, action_dim]
+
+        # Limit to max_dims
+        corrections_mean = corrections_mean[:, :max_dims].cpu().numpy()
+
+        fig, ax = plt.subplots(figsize=figsize)
+        im = ax.imshow(corrections_mean.T, aspect="auto", cmap="RdBu_r", interpolation="nearest")
+
+        ax.set_xlabel("Step Index", fontsize=12)
+        ax.set_ylabel("Action Dimension", fontsize=12)
+        ax.set_title("Correction Values Heatmap (averaged over time)", fontsize=14, fontweight="bold")
+
+        # Colorbar
+        cbar = plt.colorbar(im, ax=ax)
+        cbar.set_label("Correction Value", fontsize=12)
+
+        plt.tight_layout()
+
+        if save_path:
+            fig.savefig(save_path, dpi=150, bbox_inches="tight")
+            print(f"Correction heatmap saved to {save_path}")
+
+        if show:
+            plt.show()
+        else:
+            plt.close(fig)
+
+        return fig
+
+    @staticmethod
+    def plot_step_by_step_comparison(
+        tracker: Tracker,
+        step_idx: int = -1,
+        save_path: str | None = None,
+        show: bool = False,
+        figsize: tuple[int, int] = (18, 10),
+        max_dims: int = 6,
+    ) -> Figure:
+        """Plot detailed comparison for a single denoising step.
+
+        Args:
+            tracker (Tracker): Tracker with recorded steps.
+            step_idx (int): Step index to visualize (-1 for last step).
+            save_path (str | None): Path to save the figure.
+            show (bool): Whether to display the figure.
+            figsize (tuple[int, int]): Figure size in inches.
+            max_dims (int): Maximum number of action dimensions to visualize.
+
+        Returns:
+            Figure: The matplotlib figure object.
+        """
+        if not tracker.enabled or len(tracker) == 0:
+            print("Tracker is disabled or has no recorded steps.")
+            return None
+
+        steps = tracker.get_all_steps()
+        step = steps[step_idx]
+
+        fig, axes = plt.subplots(2, 3, figsize=figsize)
+        fig.suptitle(
+            f"Detailed Step Analysis (Step {step.step_idx})",
+            fontsize=16,
+            fontweight="bold",
+        )
+
+        # Get tensors and squeeze batch dimension
+        x_t = step.x_t.squeeze(0).cpu().numpy() if step.x_t is not None else None
+        v_t = step.v_t.squeeze(0).cpu().numpy() if step.v_t is not None else None
+        x1_t = step.x1_t.squeeze(0).cpu().numpy() if step.x1_t is not None else None
+        correction = step.correction.squeeze(0).cpu().numpy() if step.correction is not None else None
+        err = step.err.squeeze(0).cpu().numpy() if step.err is not None else None
+        weights = step.weights.squeeze().cpu().numpy() if step.weights is not None else None
+
+        # Limit to max_dims
+        num_dims = min(max_dims, x_t.shape[1] if x_t is not None else 0)
+
+        # Plot 1: x_t (current state)
+        ax = axes[0, 0]
+        if x_t is not None:
+            for dim in range(num_dims):
+                ax.plot(x_t[:, dim], label=f"Dim {dim}", alpha=0.7)
+            ax.set_title("x_t (Current State)", fontsize=12, fontweight="bold")
+            ax.set_xlabel("Time Index")
+            ax.set_ylabel("Value")
+            ax.legend(fontsize=8)
+            ax.grid(True, alpha=0.3)
+
+        # Plot 2: v_t (velocity)
+        ax = axes[0, 1]
+        if v_t is not None:
+            for dim in range(num_dims):
+                ax.plot(v_t[:, dim], label=f"Dim {dim}", alpha=0.7)
+            ax.set_title("v_t (Velocity)", fontsize=12, fontweight="bold")
+            ax.set_xlabel("Time Index")
+            ax.set_ylabel("Value")
+            ax.legend(fontsize=8)
+            ax.grid(True, alpha=0.3)
+
+        # Plot 3: x1_t (predicted state)
+        ax = axes[0, 2]
+        if x1_t is not None:
+            for dim in range(num_dims):
+                ax.plot(x1_t[:, dim], label=f"Dim {dim}", alpha=0.7)
+            ax.set_title("x1_t (Predicted State)", fontsize=12, fontweight="bold")
+            ax.set_xlabel("Time Index")
+            ax.set_ylabel("Value")
+            ax.legend(fontsize=8)
+            ax.grid(True, alpha=0.3)
+
+        # Plot 4: correction
+        ax = axes[1, 0]
+        if correction is not None:
+            for dim in range(num_dims):
+                ax.plot(correction[:, dim], label=f"Dim {dim}", alpha=0.7)
+            ax.set_title("Correction", fontsize=12, fontweight="bold")
+            ax.set_xlabel("Time Index")
+            ax.set_ylabel("Value")
+            ax.legend(fontsize=8)
+            ax.grid(True, alpha=0.3)
+
+        # Plot 5: error
+        ax = axes[1, 1]
+        if err is not None:
+            for dim in range(num_dims):
+                ax.plot(err[:, dim], label=f"Dim {dim}", alpha=0.7)
+            ax.set_title("Error (Weighted)", fontsize=12, fontweight="bold")
+            ax.set_xlabel("Time Index")
+            ax.set_ylabel("Value")
+            ax.legend(fontsize=8)
+            ax.grid(True, alpha=0.3)
+
+        # Plot 6: prefix weights
+        ax = axes[1, 2]
+        if weights is not None:
+            ax.plot(weights, linewidth=2, marker="o", markersize=4, color="red")
+            ax.set_title("Prefix Attention Weights", fontsize=12, fontweight="bold")
+            ax.set_xlabel("Time Index")
+            ax.set_ylabel("Weight Value")
+            ax.grid(True, alpha=0.3)
+            ax.set_ylim(-0.1, 1.1)
+
+        plt.tight_layout()
+
+        if save_path:
+            fig.savefig(save_path, dpi=150, bbox_inches="tight")
+            print(f"Step-by-step comparison saved to {save_path}")
+
+        if show:
+            plt.show()
+        else:
+            plt.close(fig)
+
+        return fig
+
+    @staticmethod
+    def print_debug_statistics(tracker: Tracker) -> None:
+        """Print summary statistics from the tracker.
+
+        Args:
+            tracker (Tracker): Tracker with recorded steps.
+        """
+        if not tracker.enabled:
+            print("Tracker is disabled.")
+            return
+
+        stats = tracker.get_step_stats_summary()
+
+        print("\n" + "=" * 60)
+        print("RTC Debug Statistics Summary")
+        print("=" * 60)
+        print(f"Enabled: {stats['enabled']}")
+        print(f"Total steps recorded: {stats['total_steps']}")
+        print(f"Step counter: {stats['step_counter']}")
+
+        if "correction_norms" in stats:
+            print("\nCorrection Norms:")
+            for key, value in stats["correction_norms"].items():
+                print(f"  {key}: {value:.6f}")
+
+        if "error_norms" in stats:
+            print("\nError Norms:")
+            for key, value in stats["error_norms"].items():
+                print(f"  {key}: {value:.6f}")
+
+        if "guidance_weights" in stats:
+            print("\nGuidance Weights:")
+            for key, value in stats["guidance_weights"].items():
+                print(f"  {key}: {value:.6f}")
+
+        print("=" * 60 + "\n")

src/lerobot/policies/rtc/latency_tracker.py

@@ -0,0 +1,72 @@
+#!/usr/bin/env python
+
+# Copyright 2025 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Latency tracking utilities for Real-Time Chunking (RTC)."""
+
+from collections import deque
+
+import numpy as np
+
+
+class LatencyTracker:
+    """Tracks recent latencies and provides max/percentile queries.
+
+    Args:
+        maxlen (int | None): Optional sliding window size. If provided, only the
+            most recent ``maxlen`` latencies are kept. If ``None``, keeps all.
+    """
+
+    def __init__(self, maxlen: int = 100):
+        self._values = deque(maxlen=maxlen)
+        self.reset()
+
+    def reset(self) -> None:
+        """Clear all recorded latencies."""
+        self._values.clear()
+        self.max_latency = 0.0
+
+    def add(self, latency: float) -> None:
+        """Add a latency sample (seconds)."""
+        # Ensure numeric and non-negative
+        val = float(latency)
+
+        if val < 0:
+            return
+        self._values.append(val)
+        self.max_latency = max(self.max_latency, val)
+
+    def __len__(self) -> int:
+        return len(self._values)
+
+    def max(self) -> float | None:
+        """Return the maximum latency or None if empty."""
+        return self.max_latency
+
+    def percentile(self, q: float) -> float | None:
+        """Return the q-quantile (q in [0,1]) of recorded latencies or None if empty."""
+        if not self._values:
+            return 0.0
+        q = float(q)
+        if q <= 0.0:
+            return min(self._values)
+        if q >= 1.0:
+            return self.max_latency
+        vals = np.array(list(self._values), dtype=np.float32)
+        return float(np.quantile(vals, q))
+
+    def p95(self) -> float | None:
+        """Return the 95th percentile latency or None if empty."""
+        return self.percentile(0.95)

src/lerobot/policies/rtc/modeling_rtc.py

@@ -0,0 +1,325 @@
+#!/usr/bin/env python
+
+# Copyright 2025 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.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Real-Time Chunking (RTC) implementation for LeRobot.
+
+Based on Physical Intelligence's Kinetix implementation:
+https://github.com/Physical-Intelligence/real-time-chunking-kinetix/blob/main/src/model.py#L214
+"""
+
+import logging
+import math
+
+import torch
+from torch import Tensor
+
+from lerobot.configs.types import RTCAttentionSchedule
+from lerobot.policies.rtc.configuration_rtc import RTCConfig
+from lerobot.policies.rtc.debug_handler import Tracker
+
+logger = logging.getLogger(__name__)
+
+
+class RTCProcessor:
+    """Real-Time Chunking processor for action chunking policies.
+
+    This class implements RTC techniques including velocity calculation,
+    prefix attention, and adaptive chunk processing.
+    """
+
+    def __init__(self, rtc_config: RTCConfig):
+        self.rtc_config = rtc_config
+
+        self.tracker = None
+
+        if rtc_config.debug:
+            self.tracker = Tracker(
+                enabled=rtc_config.debug,
+                maxlen=rtc_config.debug_maxlen,
+            )
+
+    # ====================== Tracker Proxy Methods ======================
+    def track_debug(
+        self,
+        time: float | Tensor,
+        x_t: Tensor | None = None,
+        v_t: Tensor | None = None,
+        x1_t: Tensor | None = None,
+        correction: Tensor | None = None,
+        err: Tensor | None = None,
+        weights: Tensor | None = None,
+        guidance_weight: float | Tensor | None = None,
+        inference_delay: int | None = None,
+        execution_horizon: int | None = None,
+        **metadata,
+    ) -> None:
+        """Proxy method to track debug information.
+
+        If tracker is None or disabled, this method does nothing.
+        Otherwise, it forwards the call to tracker.track().
+        """
+        if self.tracker is not None:
+            self.tracker.track(
+                time=time,
+                x_t=x_t,
+                v_t=v_t,
+                x1_t=x1_t,
+                correction=correction,
+                err=err,
+                weights=weights,
+                guidance_weight=guidance_weight,
+                inference_delay=inference_delay,
+                execution_horizon=execution_horizon,
+                **metadata,
+            )
+
+    def get_tracker_stats(self) -> dict | None:
+        """Get tracker statistics summary.
+
+        Returns None if tracker is disabled or None.
+        """
+        if self.tracker is not None:
+            return self.tracker.get_step_stats_summary()
+        return None
+
+    def get_all_debug_steps(self) -> list:
+        """Get all debug steps from tracker.
+
+        Returns empty list if tracker is disabled or None.
+        """
+        if self.tracker is not None:
+            return self.tracker.get_all_steps()
+        return []
+
+    def get_recent_debug_steps(self, n: int = 1) -> list:
+        """Get recent debug steps from tracker.
+
+        Returns empty list if tracker is disabled or None.
+        """
+        if self.tracker is not None:
+            return self.tracker.get_recent_steps(n)
+        return []
+
+    def is_debug_enabled(self) -> bool:
+        """Check if debug tracking is enabled.
+
+        Returns True if tracker exists and is enabled.
+        """
+        return self.tracker is not None and self.tracker.enabled
+
+    def reset_tracker(self) -> None:
+        """Reset the tracker, clearing all recorded steps.
+
+        Does nothing if tracker is None.
+        """
+        if self.tracker is not None:
+            self.tracker.reset()
+
+    def get_tracker_length(self) -> int:
+        """Get the number of recorded debug steps.
+
+        Returns 0 if tracker is disabled or None.
+        """
+        if self.tracker is not None:
+            return len(self.tracker)
+        return 0
+
+    # ====================== End Tracker Proxy Methods ======================
+
+    def denoise_step(
+        self,
+        x_t,
+        prev_chunk_left_over,
+        inference_delay,
+        time,
+        original_denoise_step_partial,
+        execution_horizon=None,
+    ) -> Tensor:
+        """RTC guidance wrapper around an existing denoiser.
+
+        This method wraps an original denoising callable that only takes ``x_t`` and
+        returns a base denoised velocity ``v_t``. It then applies Real-Time Chunking
+        (RTC) prefix guidance using the leftover prefix from the previous chunk.
+
+        Args:
+            x_t (Tensor): Current latent/state to denoise. Shape ``(B, T, A)`` or ``(T, A)``.
+            prev_chunk_left_over (Tensor | None): Unexecuted prefix from the previous
+                chunk. Shape ``(B, T_prev, A)`` or ``(T_prev, A)``. If ``None``, no guidance
+                is applied and the method returns ``v_t`` from the original denoiser.
+            inference_delay (int): Number of timesteps from the prefix to use for guidance.
+            time (float | Tensor): Scalar in [0, 1] indicating normalized time. Must be
+                broadcastable with ``x_t``.
+            original_denoise_step_partial (Callable[[Tensor], Tensor]): Callable that
+                computes the base denoised velocity given only ``x_t``.
+            execution_horizon (int | None): Horizon used to build prefix weights. If
+                ``None``, defaults to ``self.rtc_config.execution_horizon``.
+
+        Returns:
+            Tensor: Guided velocity with the same shape as ``v_t``.
+
+        Notes:
+            - If inputs are 2D, a batch dimension is temporarily added and removed at the end.
+            - If ``prev_chunk_left_over`` is shorter than the current chunk length ``T``, it is
+              right-padded with zeros to match ``T``.
+            - Prefix weights are constructed via ``get_prefix_weights(inference_delay, execution_horizon, T)``
+              and broadcast to ``(B, T, A)``.
+            - Guidance correction is computed via autograd using ``x1_t = x_t + time * v_t`` and
+              ``error = (prev_chunk_left_over - x1_t) * weights``.
+            - The final guidance weight is clamped by ``max_guidance_weight`` from the config.
+
+        Reference:
+            https://www.physicalintelligence.company/download/real_time_chunking.pdf
+        """
+
+        # In the original implementation, the time goes from 0 to 1 and
+        # In our implementation, the time goes from 1 to 0
+        # So we need to invert the time
+        tau = 1 - time
+
+        x_t = x_t.clone().detach()
+
+        if prev_chunk_left_over is None:
+            # First step, no guidance - return v_t
+            v_t = original_denoise_step_partial(x_t)
+            return v_t
+
+        squeezed = False
+        if len(x_t.shape) < 3:
+            # Add batch dimension
+            x_t = x_t.unsqueeze(0)
+            squeezed = True
+
+        if len(prev_chunk_left_over.shape) < 3:
+            # Add batch dimension
+            prev_chunk_left_over = prev_chunk_left_over.unsqueeze(0)
+
+        if execution_horizon is None:
+            execution_horizon = self.rtc_config.execution_horizon
+
+        # If the previous action chunk is to short then it doesn't make sense to use long execution horizon
+        # because there is nothing to merge
+        if execution_horizon > prev_chunk_left_over.shape[1]:
+            execution_horizon = prev_chunk_left_over.shape[1]
+
+        batch_size = x_t.shape[0]
+        action_chunk_size = x_t.shape[1]
+        action_dim = x_t.shape[2]
+
+        if prev_chunk_left_over.shape[1] < action_chunk_size or prev_chunk_left_over.shape[2] < action_dim:
+            padded = torch.zeros(batch_size, action_chunk_size, action_dim).to(x_t.device)
+            padded[:, : prev_chunk_left_over.shape[1], : prev_chunk_left_over.shape[2]] = prev_chunk_left_over
+            prev_chunk_left_over = padded
+
+        assert prev_chunk_left_over.shape == x_t.shape, (
+            "The padded previous chunk must be the same size as the input tensor"
+        )
+
+        weights = (
+            self.get_prefix_weights(inference_delay, execution_horizon, action_chunk_size)
+            .to(x_t.device)
+            .unsqueeze(0)
+            .unsqueeze(-1)
+        )
+
+        with torch.enable_grad():
+            v_t = original_denoise_step_partial(x_t)
+            x_t.requires_grad_(True)
+
+            x1_t = x_t - time * v_t  # noqa: N806
+            err = (prev_chunk_left_over - x1_t) * weights
+            grad_outputs = err.clone().detach()
+            correction = torch.autograd.grad(x1_t, x_t, grad_outputs, retain_graph=False)[0]
+
+        max_guidance_weight = torch.as_tensor(self.rtc_config.max_guidance_weight)
+        squared_one_minus_tau = (1 - tau) ** 2
+        inv_r2 = (squared_one_minus_tau + tau**2) / (squared_one_minus_tau)
+        c = torch.nan_to_num((1 - tau) / tau, posinf=max_guidance_weight)
+        guidance_weight = torch.nan_to_num(c * inv_r2, posinf=max_guidance_weight)
+        guidance_weight = torch.minimum(guidance_weight, max_guidance_weight)
+
+        result = v_t - guidance_weight * correction
+
+        # Remove the batch dimension if it was added
+        if squeezed:
+            result = result.squeeze(0)
+            correction = correction.squeeze(0)
+            x1_t = x1_t.squeeze(0)
+            err = err.squeeze(0)
+
+        # Record debug information (all params except x_t which is recorded externally)
+        self.track_debug(
+            time=time,
+            v_t=v_t,
+            x1_t=x1_t,
+            correction=correction,
+            err=err,
+            weights=weights,
+            guidance_weight=guidance_weight,
+            inference_delay=inference_delay,
+            execution_horizon=execution_horizon,
+        )
+
+        return result
+
+    def get_prefix_weights(self, start, end, total):
+        start = min(start, end)
+
+        if self.rtc_config.prefix_attention_schedule == RTCAttentionSchedule.ZEROS:
+            weights = torch.zeros(total)
+            weights[:start] = 1.0
+        elif self.rtc_config.prefix_attention_schedule == RTCAttentionSchedule.ONES:
+            weights = torch.ones(total)
+            weights[end:] = 0.0
+        elif self.rtc_config.prefix_attention_schedule == RTCAttentionSchedule.LINEAR:
+            lin_weights = self._linweights(start, end, total)
+            weights = self._add_trailing_zeros(lin_weights, total, end)
+            weights = self._add_leading_ones(weights, start, total)
+        elif self.rtc_config.prefix_attention_schedule == RTCAttentionSchedule.EXP:
+            lin_weights = self._linweights(start, end, total)
+            lin_weights = lin_weights * torch.expm1(lin_weights).div(math.e - 1)
+            weights = self._add_trailing_zeros(lin_weights, total, end)
+            weights = self._add_leading_ones(weights, start, total)
+
+        return weights
+
+    def _linweights(self, start, end, total):
+        skip_steps_at_end = max(total - end, 0)
+
+        linspace_steps = total - skip_steps_at_end - start
+
+        if end <= start or linspace_steps <= 0:
+            return torch.tensor([])
+
+        return torch.linspace(1, 0, linspace_steps + 2)[1:-1]
+
+    def _add_trailing_zeros(self, weights, total, end):
+        zeros_len = total - end
+
+        if zeros_len <= 0:
+            return weights
+
+        zeros = torch.zeros(zeros_len)
+        return torch.cat([weights, zeros])
+
+    def _add_leading_ones(self, weights, start, total):
+        ones_len = min(start, total)
+
+        if ones_len <= 0:
+            return weights
+
+        ones = torch.ones(ones_len)
+        return torch.cat([ones, weights])

src/lerobot/policies/smolvla/modeling_smolvla.py

@@ -55,11 +55,15 @@ policy = SmolVLAPolicy.from_pretrained("lerobot/smolvla_base")
 import math
 from collections import deque
 
+import matplotlib.pyplot as plt
+import numpy as np
 import torch
 import torch.nn.functional as F  # noqa: N812
 from torch import Tensor, nn
 
 from lerobot.policies.pretrained import PreTrainedPolicy
+from lerobot.policies.rtc.debug_visualizer import RTCDebugVisualizer
+from lerobot.policies.rtc.modeling_rtc import RTCProcessor
 from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
 from lerobot.policies.smolvla.smolvlm_with_expert import SmolVLMWithExpertModel
 from lerobot.policies.utils import (
@@ -68,6 +72,9 @@ from lerobot.policies.utils import (
 from lerobot.utils.constants import ACTION, OBS_LANGUAGE_ATTENTION_MASK, OBS_LANGUAGE_TOKENS, OBS_STATE
 from lerobot.utils.utils import get_safe_dtype
 
+# Make plot_waypoints easily accessible
+plot_waypoints = RTCDebugVisualizer.plot_waypoints
+
 
 def create_sinusoidal_pos_embedding(
     time: torch.tensor, dimension: int, min_period: float, max_period: float, device="cpu"
@@ -232,8 +239,8 @@ class SmolVLAPolicy(PreTrainedPolicy):
         super().__init__(config)
         config.validate_features()
         self.config = config
-
-        self.model = VLAFlowMatching(config)
+        self.init_rtc_processor()
+        self.model = VLAFlowMatching(config, rtc_processor=self.rtc_processor)
         self.reset()
 
     def reset(self):
@@ -242,10 +249,27 @@ class SmolVLAPolicy(PreTrainedPolicy):
             ACTION: deque(maxlen=self.config.n_action_steps),
         }
 
+    def init_rtc_processor(self, verbose: bool = False):
+        """Initialize RTC processor with optional verbose logging.
+
+        Args:
+            verbose: Enable verbose debug logging in RTCProcessor (currently unused)
+        """
+        self.rtc_processor = None
+
+        if self.config.rtc_config is not None and self.config.rtc_config.enabled:
+            self.rtc_processor = RTCProcessor(self.config.rtc_config)
+
+            # In case of calling init_rtc_processor after the model is created
+            # We need to set the rtc_processor to the model
+            # During the normal initialization process the model is not created yet
+            if self.model is not None:
+                self.model.rtc_processor = self.rtc_processor
+
     def get_optim_params(self) -> dict:
         return self.parameters()
 
-    def _get_action_chunk(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor:
+    def _get_action_chunk(self, batch: dict[str, Tensor], noise: Tensor | None = None, **kwargs) -> Tensor:
         # TODO: Check if this for loop is needed.
         # Context: In fact, self.queues contains only ACTION field, and in inference, we don't have action in the batch
         # In the case of offline inference, we have the action in the batch
@@ -260,7 +284,9 @@ class SmolVLAPolicy(PreTrainedPolicy):
         lang_tokens = batch[f"{OBS_LANGUAGE_TOKENS}"]
         lang_masks = batch[f"{OBS_LANGUAGE_ATTENTION_MASK}"]
 
-        actions = self.model.sample_actions(images, img_masks, lang_tokens, lang_masks, state, noise=noise)
+        actions = self.model.sample_actions(
+            images, img_masks, lang_tokens, lang_masks, state, noise=noise, **kwargs
+        )
 
         # Unpad actions
         original_action_dim = self.config.action_feature.shape[0]
@@ -278,30 +304,33 @@ class SmolVLAPolicy(PreTrainedPolicy):
         return batch
 
     @torch.no_grad()
-    def predict_action_chunk(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor:
+    def predict_action_chunk(self, batch: dict[str, Tensor], noise: Tensor | None = None, **kwargs) -> Tensor:
         self.eval()
 
         batch = self._prepare_batch(batch)
         self._queues = populate_queues(self._queues, batch, exclude_keys=[ACTION])
 
-        actions = self._get_action_chunk(batch, noise)
+        actions = self._get_action_chunk(batch, noise, **kwargs)
         return actions
 
     @torch.no_grad()
-    def select_action(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor:
+    def select_action(self, batch: dict[str, Tensor], noise: Tensor | None = None, **kwargs) -> Tensor:
         """Select a single action given environment observations.
 
         This method wraps `select_actions` in order to return one action at a time for execution in the
         environment. It works by managing the actions in a queue and only calling `select_actions` when the
         queue is empty.
         """
+
+        assert not self._rtc_enabled(), (
+            "RTC is not supported for select_action, use it with predict_action_chunk"
+        )
+
         self.eval()
         batch = self._prepare_batch(batch)
         self._queues = populate_queues(self._queues, batch, exclude_keys=[ACTION])
 
-        # Action queue logic for n_action_steps > 1. When the action_queue is depleted, populate it by
-        # querying the policy.
-        if len(self._queues[ACTION]) == 0:
+        if self._check_get_actions_condition():
             actions = self._get_action_chunk(batch, noise)
 
             # `self.predict_action_chunk` returns a (batch_size, n_action_steps, action_dim) tensor, but the queue
@@ -310,6 +339,12 @@ class SmolVLAPolicy(PreTrainedPolicy):
 
         return self._queues[ACTION].popleft()
 
+    def _check_get_actions_condition(self) -> bool:
+        return len(self._queues[ACTION]) == 0
+
+    def _rtc_enabled(self) -> bool:
+        return self.config.rtc_config is not None and self.config.rtc_config.enabled
+
     def forward(self, batch: dict[str, Tensor], noise=None, time=None) -> dict[str, Tensor]:
         """Do a full training forward pass to compute the loss"""
         if self.config.adapt_to_pi_aloha:
@@ -471,7 +506,7 @@ class VLAFlowMatching(nn.Module):
     └──────────────────────────────┘
     """
 
-    def __init__(self, config: SmolVLAConfig):
+    def __init__(self, config: SmolVLAConfig, rtc_processor: RTCProcessor | None = None):
         super().__init__()
         self.config = config
 
@@ -485,7 +520,6 @@ class VLAFlowMatching(nn.Module):
             num_vlm_layers=self.config.num_vlm_layers,
             self_attn_every_n_layers=self.config.self_attn_every_n_layers,
             expert_width_multiplier=self.config.expert_width_multiplier,
-            device=self.config.device,
         )
         self.state_proj = nn.Linear(
             self.config.max_state_dim, self.vlm_with_expert.config.text_config.hidden_size
@@ -510,6 +544,12 @@ class VLAFlowMatching(nn.Module):
         self.add_image_special_tokens = self.config.add_image_special_tokens
         self.image_end_token = torch.tensor([self.fake_image_token], dtype=torch.long)
         self.prefix_length = self.config.prefix_length
+        self.rtc_processor = rtc_processor
+
+        # For visualization of x_t during denoising
+        self.denoise_step_counter = 0
+        self.viz_fig = None
+        self.viz_axs = None
 
     def set_requires_grad(self):
         for params in self.state_proj.parameters():
@@ -706,11 +746,25 @@ class VLAFlowMatching(nn.Module):
         losses = F.mse_loss(u_t, v_t, reduction="none")
         return losses
 
-    def sample_actions(self, images, img_masks, lang_tokens, lang_masks, state, noise=None) -> Tensor:
-        """Do a full inference forward and compute the action (batch_size x num_steps x num_motors)"""
+    def sample_actions(
+        self, images, img_masks, lang_tokens, lang_masks, state, noise=None, **kwargs
+    ) -> Tensor:
+        """Do a full inference forward and compute the action (batch_size x num_steps x num_motors)
+
+        Args:
+            viz_xt_axs: Optional matplotlib axes for plotting x_t trajectories (array of 6 axes)
+            viz_vt_axs: Optional matplotlib axes for plotting v_t trajectories (array of 6 axes)
+            viz_x1t_axs: Optional matplotlib axes for plotting x1_t predicted state and error (array of 6 axes)
+                         When RTC is enabled, plots both x1_t (solid line) and error (orange dashed line)
+        """
         bsize = state.shape[0]
         device = state.device
 
+        # Extract visualization axes from kwargs
+        viz_xt_axs = kwargs.pop("viz_xt_axs", None)
+        viz_vt_axs = kwargs.pop("viz_vt_axs", None)
+        viz_x1t_axs = kwargs.pop("viz_x1t_axs", None)
+
         if noise is None:
             actions_shape = (bsize, self.config.chunk_size, self.config.max_action_dim)
             noise = self.sample_noise(actions_shape, device)
@@ -734,17 +788,167 @@ class VLAFlowMatching(nn.Module):
 
         x_t = noise
         time = torch.tensor(1.0, dtype=torch.float32, device=device)
+        correction = None
+        x1_t = None
+        error = None
+        use_provided_axes = viz_xt_axs is not None and viz_vt_axs is not None
+
         while time >= -dt / 2:
             expanded_time = time.expand(bsize)
-            v_t = self.denoise_step(
-                prefix_pad_masks,
-                past_key_values,
-                x_t,
-                expanded_time,
-            )
+
+            # Define a closure function to properly capture expanded_time
+            # This avoids the lambda expression (E731) and loop variable binding (B023) issues
+            def denoise_step_partial_call(input_x_t, current_timestep=expanded_time):
+                return self.denoise_step(
+                    x_t=input_x_t,
+                    prefix_pad_masks=prefix_pad_masks,
+                    past_key_values=past_key_values,
+                    timestep=current_timestep,
+                )
+
+            if self.config.rtc_config is not None and self.config.rtc_config.enabled:
+                inference_delay = kwargs.get("inference_delay")
+                prev_chunk_left_over = kwargs.get("prev_chunk_left_over")
+                execution_horizon = kwargs.get("execution_horizon", self.config.rtc_config.execution_horizon)
+
+                v_t = self.rtc_processor.denoise_step(
+                    x_t=x_t,
+                    prev_chunk_left_over=prev_chunk_left_over,
+                    inference_delay=inference_delay,
+                    time=time,
+                    original_denoise_step_partial=denoise_step_partial_call,
+                    execution_horizon=execution_horizon,
+                )
+            else:
+                v_t = denoise_step_partial_call(x_t)
+
             # Euler step
             x_t += dt * v_t
             time += dt
+
+            # Record x_t after Euler step (other params are recorded in rtc_processor.denoise_step)
+            if (
+                self.config.rtc_config is not None
+                and self.config.rtc_config.enabled
+                and correction is not None
+            ):
+                self.rtc_processor.track_debug(time=time, x_t=x_t)
+
+            # Visualize x_t using plot_waypoints - accumulate all denoise steps
+            # Use provided axes or create new ones
+            if not use_provided_axes:
+                if self.viz_fig is None:
+                    # Create figure once on first denoise step
+                    self.viz_fig, self.viz_axs = plt.subplots(6, 1, figsize=(12, 12))
+                    self.viz_v_fig, self.viz_v_axs = plt.subplots(6, 1, figsize=(12, 12))
+                xt_axs = self.viz_axs
+                vt_axs = self.viz_v_axs
+            else:
+                xt_axs = viz_xt_axs
+                vt_axs = viz_vt_axs
+
+            # Define colors for different denoise steps (using a colormap)
+            colors = plt.cm.viridis(np.linspace(0, 1, self.config.num_steps))
+            color = colors[self.denoise_step_counter % len(colors)]
+
+            # Plot this denoise step
+            plot_waypoints(xt_axs, x_t, start_from=0, color=color, label=f"Step {self.denoise_step_counter}")
+
+            # Plot this denoise step
+            plot_waypoints(vt_axs, v_t, start_from=0, color=color, label=f"Step {self.denoise_step_counter}")
+
+            if correction is not None:
+                plot_waypoints(
+                    vt_axs,
+                    correction,
+                    start_from=0,
+                    color="red",
+                    label=f"Step corr {self.denoise_step_counter}",
+                )
+
+            # Plot x1_t if axes provided and RTC is enabled
+            if viz_x1t_axs is not None and x1_t is not None:
+                plot_waypoints(
+                    viz_x1t_axs,
+                    x1_t,
+                    start_from=0,
+                    color=color,
+                    label=f"x1_t Step {self.denoise_step_counter}",
+                )
+
+                # Plot error on the same axes with different color
+                if error is not None:
+                    # Use orange color for error
+                    # Handle batch dimension if present
+                    error_chunk = error[0].cpu().numpy() if len(error.shape) == 3 else error.cpu().numpy()
+
+                    num_dims = min(error_chunk.shape[-1], 6)
+                    for j in range(num_dims):
+                        viz_x1t_axs[j].plot(
+                            np.arange(0, error_chunk.shape[0]),
+                            error_chunk[:, j],
+                            color="orange",
+                            linestyle="--",
+                            alpha=0.7,
+                            label=f"error Step {self.denoise_step_counter}",
+                        )
+
+            self.denoise_step_counter += 1
+
+        # Save visualization of x_t denoise steps (only if using internal figures)
+        if not use_provided_axes and self.viz_fig is not None:
+            plt.figure(self.viz_fig.number)
+
+            xt_name = "smolvla_x_t_denoise_steps.png"
+            v_name = "smolvla_v_denoise_steps.png"
+
+            if self.config.rtc_config is not None and self.config.rtc_config.enabled:
+                xt_name = "smolvla_x_t_with_rtc_denoise_steps.png"
+                v_name = "smolvla_v_with_rtc_denoise_steps.png"
+
+                prev_chunk_left_over = kwargs.get("prev_chunk_left_over")
+
+                if prev_chunk_left_over is not None:
+                    plot_waypoints(
+                        self.viz_axs, prev_chunk_left_over, start_from=0, color="red", label="Ground truth"
+                    )
+
+            plt.savefig(xt_name)
+            plt.close(self.viz_fig)
+
+            # Reset for next inference
+            self.viz_fig = None
+            self.viz_axs = None
+            self.denoise_step_counter = 0
+
+            plt.figure(self.viz_v_fig.number)
+            plt.savefig(v_name)
+            plt.close(self.viz_v_fig)
+
+            self.viz_v_fig = None
+            self.viz_v_axs = None
+
+        # Plot ground truth on provided axes if available
+        if use_provided_axes:
+            prev_chunk_left_over = kwargs.get("prev_chunk_left_over")
+            if (
+                prev_chunk_left_over is not None
+                and self.config.rtc_config is not None
+                and self.config.rtc_config.enabled
+            ):
+                plot_waypoints(
+                    viz_xt_axs, prev_chunk_left_over, start_from=0, color="red", label="Ground truth"
+                )
+                # Also plot ground truth on x1_t axes if provided
+                if viz_x1t_axs is not None:
+                    plot_waypoints(
+                        viz_x1t_axs, prev_chunk_left_over, start_from=0, color="red", label="Ground truth"
+                    )
+
+        # Reset counter when using provided axes (for next call)
+        if use_provided_axes:
+            self.denoise_step_counter = 0
+
         return x_t
 
     def denoise_step(