#!/usr/bin/env python """ Evaluate Real-Time Chunking (RTC) performance on dataset samples. This script takes two random samples from a dataset: - Uses actions from the first sample as previous chunk - Generates new actions for the second sample with and without RTC It compares action predictions with and without RTC on dataset samples, measuring consistency and ground truth alignment. Usage: python eval_dataset.py \ --policy.path=helper2424/smolvla_check_rtc_last3 \ --dataset.repo_id=helper2424/check_rtc \ --rtc.execution_horizon=8 \ --device=mps """ import logging import os import random import sys from dataclasses import dataclass, field import matplotlib.pyplot as plt import numpy as np import torch from lerobot.configs import parser from lerobot.configs.default import DatasetConfig from lerobot.configs.policies import PreTrainedConfig from lerobot.configs.types import RTCAttentionSchedule from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot.policies.factory import get_policy_class, make_pre_post_processors from lerobot.policies.rtc.configuration_rtc import RTCConfig from lerobot.policies.rtc.debug_visualizer import RTCDebugVisualizer from lerobot.utils.hub import HubMixin logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", force=True, ) logger = logging.getLogger(__name__) # Ensure logs are flushed immediately sys.stdout.reconfigure(line_buffering=True) sys.stderr.reconfigure(line_buffering=True) def set_seed(seed: int): """Set random seed for reproducibility.""" random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed(seed) torch.cuda.manual_seed_all(seed) if torch.backends.mps.is_available(): torch.mps.manual_seed(seed) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False logger.info(f"Random seed set to: {seed}") @dataclass class RTCEvalConfig(HubMixin): """Configuration for RTC evaluation.""" # Policy configuration policy: PreTrainedConfig | None = None # Dataset configuration dataset: DatasetConfig = field(default_factory=DatasetConfig) # RTC configuration rtc: RTCConfig = field( default_factory=lambda: RTCConfig( enabled=True, execution_horizon=20, max_guidance_weight=5.0, prefix_attention_schedule=RTCAttentionSchedule.EXP, debug=True, debug_maxlen=1000, ) ) # Device configuration device: str | None = field( default=None, metadata={"help": "Device to run on (cuda, cpu, mps, auto)"}, ) # Output configuration output_dir: str = field( default="rtc_debug_output", metadata={"help": "Directory to save debug visualizations"}, ) # Seed configuration seed: int = field( default=42, metadata={"help": "Random seed for reproducibility"}, ) inference_delay: int = field( default=4, metadata={"help": "Inference delay for RTC"}, ) def __post_init__(self): # Parse policy path policy_path = parser.get_path_arg("policy") if policy_path: cli_overrides = parser.get_cli_overrides("policy") self.policy = PreTrainedConfig.from_pretrained(policy_path, cli_overrides=cli_overrides) self.policy.pretrained_path = policy_path else: raise ValueError("Policy path is required (--policy.path)") # Auto-detect device if not specified if self.device is None or self.device == "auto": if torch.cuda.is_available(): self.device = "cuda" elif torch.backends.mps.is_available(): self.device = "mps" else: self.device = "cpu" logger.info(f"Auto-detected device: {self.device}") @classmethod def __get_path_fields__(cls) -> list[str]: """This enables the parser to load config from the policy using `--policy.path=local/dir`""" return ["policy"] class RTCEvaluator: """Evaluator for RTC on dataset samples.""" def __init__(self, cfg: RTCEvalConfig): self.cfg = cfg self.device = cfg.device # Load policy logger.info(f"Loading policy from {cfg.policy.pretrained_path}") policy_class = get_policy_class(cfg.policy.type) self.policy = policy_class.from_pretrained(cfg.policy.pretrained_path) self.policy = self.policy.to(self.device) self.policy.eval() # Configure RTC cfg.rtc.enabled = True cfg.rtc.debug = True # Enable debug tracking for visualization self.policy.config.rtc_config = cfg.rtc self.policy.init_rtc_processor() logger.info(f"Policy loaded: {self.policy.name}") logger.info(f"RTC enabled: {cfg.rtc.enabled}") logger.info(f"Execution horizon: {cfg.rtc.execution_horizon}") # Load dataset logger.info(f"Loading dataset: {cfg.dataset.repo_id}") self.dataset = LeRobotDataset(cfg.dataset.repo_id, delta_timestamps={"action": np.arange(50) / 30}) logger.info(f"Dataset loaded: {len(self.dataset)} samples, {self.dataset.num_episodes} episodes") # Create preprocessor/postprocessor self.preprocessor, self.postprocessor = make_pre_post_processors( policy_cfg=cfg.policy, pretrained_path=cfg.policy.pretrained_path, preprocessor_overrides={ "device_processor": {"device": self.device}, }, ) def run_evaluation(self): """Run evaluation on two random dataset samples.""" # Create output directory os.makedirs(self.cfg.output_dir, exist_ok=True) logger.info(f"Output directory: {self.cfg.output_dir}") logger.info("Starting RTC evaluation") logger.info(f"Inference delay: {self.cfg.inference_delay}") data_loader = torch.utils.data.DataLoader(self.dataset, batch_size=1, shuffle=True) loader_iter = iter(data_loader) first_sample = next(loader_iter) second_sample = next(loader_iter) preprocessed_first_sample = self.preprocessor(first_sample) preprocessed_second_sample = self.preprocessor(second_sample) # Don't postprocess the previous chunk prev_chunk_left_over = self.policy.predict_action_chunk( preprocessed_first_sample, )[:, :25, :].squeeze(0) # Sample noise (use same noise for both RTC and non-RTC for fair comparison) noise_size = (1, self.policy.config.chunk_size, self.policy.config.max_action_dim) noise = self.policy.model.sample_noise(noise_size, self.device) noise_clone = noise.clone() # Create side-by-side figures for denoising visualization fig_xt, axs_xt = plt.subplots(6, 2, figsize=(24, 12)) fig_xt.suptitle("x_t Denoising: No RTC (left) vs RTC (right)", fontsize=16) fig_vt, axs_vt = plt.subplots(6, 2, figsize=(24, 12)) fig_vt.suptitle("v_t Denoising: No RTC (left) vs RTC (right)", fontsize=16) fig_x1t, axs_x1t = plt.subplots(6, 2, figsize=(24, 12)) fig_x1t.suptitle("x1_t Predicted State & Error: No RTC (left - empty) vs RTC (right)", fontsize=16) # Generate actions WITHOUT RTC logger.info("Generating actions WITHOUT RTC") self.policy.config.rtc_config.enabled = False with torch.no_grad(): no_rtc_actions = self.policy.predict_action_chunk( preprocessed_second_sample, noise=noise, ) # Plot denoising steps from tracker (no RTC - left column) # Note: No tracker data for non-RTC case since tracking is only done when RTC processor exists # Generate actions WITH RTC logger.info("Generating actions WITH RTC") self.policy.config.rtc_config.enabled = True with torch.no_grad(): rtc_actions = self.policy.predict_action_chunk( preprocessed_second_sample, noise=noise_clone, inference_delay=self.cfg.inference_delay, prev_chunk_left_over=prev_chunk_left_over, execution_horizon=self.cfg.rtc.execution_horizon, ) # Plot denoising steps from tracker (RTC - right column) if self.policy.rtc_processor is not None: num_steps = self.policy.config.num_steps self._plot_denoising_steps_from_tracker( self.policy.rtc_processor.tracker, axs_xt[:, 1], # Right column for x_t axs_vt[:, 1], # Right column for v_t axs_x1t[:, 1], # Right column for x1_t num_steps, ) # Plot ground truth on x_t axes RTCDebugVisualizer.plot_waypoints( axs_xt[:, 1], prev_chunk_left_over, start_from=0, color="red", label="Ground truth" ) # Plot ground truth on x1_t axes RTCDebugVisualizer.plot_waypoints( axs_x1t[:, 1], prev_chunk_left_over, start_from=0, color="red", label="Ground truth" ) # Set titles for denoising plots for ax in axs_xt[:, 0]: ax.set_title("No RTC" if ax == axs_xt[0, 0] else "", fontsize=12) for ax in axs_xt[:, 1]: ax.set_title("RTC" if ax == axs_xt[0, 1] else "", fontsize=12) for ax in axs_vt[:, 0]: ax.set_title("No RTC" if ax == axs_vt[0, 0] else "", fontsize=12) for ax in axs_vt[:, 1]: ax.set_title("RTC" if ax == axs_vt[0, 1] else "", fontsize=12) for ax in axs_x1t[:, 0]: ax.set_title("No RTC (N/A)" if ax == axs_x1t[0, 0] else "", fontsize=12) for ax in axs_x1t[:, 1]: ax.set_title("RTC" if ax == axs_x1t[0, 1] else "", fontsize=12) # Save denoising plots fig_xt.tight_layout() xt_path = os.path.join(self.cfg.output_dir, "denoising_xt_comparison.png") fig_xt.savefig(xt_path, dpi=150) logger.info(f"Saved x_t denoising comparison to {xt_path}") plt.close(fig_xt) fig_vt.tight_layout() vt_path = os.path.join(self.cfg.output_dir, "denoising_vt_comparison.png") fig_vt.savefig(vt_path, dpi=150) logger.info(f"Saved v_t denoising comparison to {vt_path}") plt.close(fig_vt) fig_x1t.tight_layout() x1t_path = os.path.join(self.cfg.output_dir, "denoising_x1t_comparison.png") fig_x1t.savefig(x1t_path, dpi=150) logger.info(f"Saved x1_t predicted state & error comparison to {x1t_path}") plt.close(fig_x1t) # Create side-by-side comparison: No RTC (left) vs RTC (right) fig, axs = plt.subplots(6, 2, figsize=(24, 12)) fig.suptitle("Final Action Comparison: No RTC (left) vs RTC (right)", fontsize=16) # Plot on left column (No RTC) self._plot_actions( axs[:, 0], prev_chunk_left_over[0].cpu().numpy(), no_rtc_actions[0].cpu().numpy(), "No RTC", ) # Plot on right column (RTC) self._plot_actions( axs[:, 1], prev_chunk_left_over[0].cpu().numpy(), rtc_actions[0].detach().cpu().numpy(), "RTC", ) plt.tight_layout() final_path = os.path.join(self.cfg.output_dir, "final_actions_comparison.png") plt.savefig(final_path, dpi=150) logger.info(f"Saved final actions comparison to {final_path}") plt.close(fig) # Visualize debug information if enabled self._visualize_debug_info() logger.info("Evaluation completed successfully") def _plot_actions(self, axs, prev_chunk, predicted_actions, title): """Plot actions comparison on given axes.""" # Ensure arrays are 2D if prev_chunk.ndim == 1: prev_chunk = prev_chunk.reshape(1, -1) if predicted_actions.ndim == 1: predicted_actions = predicted_actions.reshape(1, -1) for j in range(min(prev_chunk.shape[-1], 6)): # Limit to 6 dimensions axs[j].plot( np.arange(prev_chunk.shape[0]), prev_chunk[:, j], color="green", label="Previous Chunk", ) axs[j].plot( np.arange(predicted_actions.shape[0]), predicted_actions[:, j], color="red" if "RTC" in title else "blue", label=title, ) axs[j].set_ylabel("Joint angle", fontsize=14) axs[j].grid() axs[j].legend(loc="upper right", fontsize=14) axs[j].set_title(title if j == 0 else "", fontsize=12) if j == 2: axs[j].set_xlabel("Step #", fontsize=16) def _visualize_debug_info(self): """Visualize debug information from the RTC processor.""" # Use proxy method to check if debug is enabled if not self.policy.rtc_processor.is_debug_enabled(): logger.warning("Debug tracking is disabled. Skipping debug visualization.") return # Get tracker length using proxy method if self.policy.rtc_processor.get_tracker_length() == 0: logger.warning("No debug steps recorded. Skipping debug visualization.") return # Create output directory os.makedirs(self.cfg.output_dir, exist_ok=True) logger.info(f"Saving debug visualizations to {self.cfg.output_dir}") # Still need direct access to tracker for visualization functions # This is acceptable since RTCDebugVisualizer is part of the RTC package tracker = self.policy.rtc_processor.tracker # Print statistics RTCDebugVisualizer.print_debug_statistics(tracker) # Plot debug summary summary_path = os.path.join(self.cfg.output_dir, "debug_summary.png") RTCDebugVisualizer.plot_debug_summary( tracker, save_path=summary_path, show=False, ) # Plot correction heatmap heatmap_path = os.path.join(self.cfg.output_dir, "correction_heatmap.png") RTCDebugVisualizer.plot_correction_heatmap( tracker, save_path=heatmap_path, show=False, ) # Plot step-by-step comparison (last step) step_path = os.path.join(self.cfg.output_dir, "step_comparison_last.png") RTCDebugVisualizer.plot_step_by_step_comparison( tracker, step_idx=-1, save_path=step_path, show=False, ) # Plot step-by-step comparison (first step) step_path_first = os.path.join(self.cfg.output_dir, "step_comparison_first.png") if self.policy.rtc_processor.get_tracker_length() > 0: RTCDebugVisualizer.plot_step_by_step_comparison( tracker, step_idx=0, save_path=step_path_first, show=False, ) logger.info(f"Debug visualizations saved to {self.cfg.output_dir}") def _plot_denoising_steps_from_tracker(self, tracker, xt_axs, vt_axs, x1t_axs, num_steps): """Plot denoising steps from tracker data. Args: tracker: Tracker object containing debug steps xt_axs: Matplotlib axes for x_t plots (array of 6 axes) vt_axs: Matplotlib axes for v_t plots (array of 6 axes) x1t_axs: Matplotlib axes for x1_t plots (array of 6 axes) num_steps: Total number of denoising steps for colormap """ if tracker is None: return debug_steps = tracker.get_all_steps() if not debug_steps: return # Define colors for different denoise steps (using a colormap) colors = plt.cm.viridis(np.linspace(0, 1, num_steps)) for step_idx, debug_step in enumerate(debug_steps): color = colors[step_idx % len(colors)] # Plot x_t if debug_step.x_t is not None: RTCDebugVisualizer.plot_waypoints( xt_axs, debug_step.x_t, start_from=0, color=color, label=f"Step {step_idx}" ) # Plot v_t if debug_step.v_t is not None: RTCDebugVisualizer.plot_waypoints( vt_axs, debug_step.v_t, start_from=0, color=color, label=f"Step {step_idx}" ) # Plot correction in red if debug_step.correction is not None: RTCDebugVisualizer.plot_waypoints( vt_axs, debug_step.correction, start_from=0, color="red", label=f"Step corr {step_idx}", ) # Plot x1_t (predicted state) if x1t_axs is not None and debug_step.x1_t is not None: RTCDebugVisualizer.plot_waypoints( x1t_axs, debug_step.x1_t, start_from=0, color=color, label=f"x1_t Step {step_idx}", ) # Plot error in orange dashed if x1t_axs is not None and debug_step.err is not None: error_chunk = ( debug_step.err[0].cpu().numpy() if len(debug_step.err.shape) == 3 else debug_step.err.cpu().numpy() ) num_dims = min(error_chunk.shape[-1], 6) for j in range(num_dims): x1t_axs[j].plot( np.arange(0, error_chunk.shape[0]), error_chunk[:, j], color="orange", linestyle="--", alpha=0.7, label=f"error Step {step_idx}", ) @parser.wrap() def main(cfg: RTCEvalConfig): """Main entry point for RTC evaluation.""" # Set random seed for reproducibility set_seed(cfg.seed) logger.info("=" * 80) logger.info("RTC Dataset Evaluation") logger.info(f"Config: {cfg}") logger.info("=" * 80) evaluator = RTCEvaluator(cfg) evaluator.run_evaluation() if __name__ == "__main__": main()