simple eval

src/lerobot/policies/rlearn/eval_script.py

@@ -0,0 +1,367 @@
+#!/usr/bin/env python
+
+"""
+Standalone evaluation script for RLearN models.
+
+This script evaluates RLearN reward models on episodes from a dataset,
+generating comparison plots between ground truth rewards and model predictions.
+
+Usage:
+    python src/lerobot/policies/rlearn/eval_script.py --model MODEL_NAME --dataset DATASET_REPO --episodes N
+
+Example:
+    python src/lerobot/policies/rlearn/eval_script.py --model pepijn223/rlearn_mse5 --dataset pepijn223/phone_pipeline_pickup1 --episodes 2
+"""
+
+import argparse
+import os
+import sys
+from pathlib import Path
+
+# Add src to path for imports
+sys.path.append(str(Path(__file__).parent.parent.parent.parent))
+
+import warnings
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from scipy.stats import spearmanr
+from tqdm import tqdm
+
+warnings.filterwarnings("ignore")
+
+# LeRobot imports
+from lerobot.constants import OBS_IMAGE, OBS_IMAGES, OBS_LANGUAGE
+from lerobot.datasets.lerobot_dataset import LeRobotDataset
+from lerobot.policies.rlearn.modeling_rlearn import RLearNPolicy
+
+
+def _to_chw_float01(img):
+    """Ensure CHW float in [0,1]."""
+    if isinstance(img, np.ndarray):
+        img = torch.from_numpy(img)
+    # HWC -> CHW if needed
+    if len(img.shape) == 3 and img.shape[-1] in (1, 3, 4):
+        img = img.permute(2, 0, 1)
+    if img.dtype == torch.uint8:
+        img = img.float() / 255.0
+    else:
+        img = img.float()
+    return torch.clamp(img, 0.0, 1.0)
+
+
+def _get_language(frame_data):
+    lang = None
+    if OBS_LANGUAGE in frame_data:
+        lang = frame_data[OBS_LANGUAGE]
+        if isinstance(lang, list) and len(lang) > 0:
+            lang = lang[0]
+    elif "task" in frame_data:
+        lang = frame_data["task"]
+    return lang if isinstance(lang, str) else "No language provided"
+
+
+def _get_ground_truth_reward(frame_data):
+    """Try common keys for ground-truth reward. Return None if unavailable."""
+    for key in ("reward", "rewards", "gt_reward", "progress"):
+        if key in frame_data:
+            r = frame_data[key]
+            # unwrap single-element lists/arrays
+            if isinstance(r, (list, np.ndarray)) and np.array(r).size == 1:
+                r = float(np.array(r).reshape(-1)[0])
+            try:
+                return float(r)
+            except Exception:
+                pass
+    return None
+
+
+def extract_episode_frames_and_gt(dataset, episode_idx):
+    """Load a full episode: frames (T, C, H, W), language (str), gt_rewards (np.ndarray or None)."""
+    ep_start = dataset.episode_data_index["from"][episode_idx].item()
+    ep_end = dataset.episode_data_index["to"][episode_idx].item()
+    T = ep_end - ep_start
+
+    frames = []
+    gt_rewards = []
+    language = None
+
+    for t in range(T):
+        item = dataset[ep_start + t]
+
+        # image(s)
+        if OBS_IMAGES in item:
+            img = item[OBS_IMAGES]
+        elif OBS_IMAGE in item:
+            img = item[OBS_IMAGE]
+        else:
+            # try to find an image-like key
+            img_keys = [k for k in item.keys() if "image" in k.lower()]
+            if not img_keys:
+                continue
+            img = item[img_keys[0]]
+
+        frames.append(_to_chw_float01(img))
+
+        # language once
+        if language is None:
+            language = _get_language(item)
+
+        # ground-truth reward (optional)
+        r = _get_ground_truth_reward(item)
+        gt_rewards.append(r)
+
+    if not frames:
+        return None, None, None
+
+    frames = torch.stack(frames)  # (T, C, H, W)
+
+    # If all GT entries are None, treat as missing
+    if all(r is None for r in gt_rewards):
+        gt_rewards = None
+    else:
+        # Replace None by forward filling
+        arr = np.array([np.nan if r is None else float(r) for r in gt_rewards], dtype=float)
+        # forward/back fill
+        if np.isnan(arr[0]):
+            first_valid = np.flatnonzero(~np.isnan(arr))
+            if len(first_valid) > 0:
+                arr[0] = arr[first_valid[0]]
+            else:
+                arr[0] = 0.0
+        for i in range(1, len(arr)):
+            if np.isnan(arr[i]):
+                arr[i] = arr[i - 1]
+        gt_rewards = arr
+
+    return frames, language or "No language provided", gt_rewards
+
+
+@torch.no_grad()
+def predict_rewards_sliding(model, frames, language, max_seq_len=16, batch_size=64, device="cuda"):
+    """
+    Sliding-window prediction: for each frame i, create a window [max(0, i-L+1) .. i],
+    left-pad by repeating the first frame to length L (<= 16), and take the last-step prediction.
+    Returns np.ndarray of shape (T,).
+    """
+    T = frames.shape[0]
+    L = int(getattr(getattr(model, "config", object()), "max_seq_len", max_seq_len))
+    L = min(L, max_seq_len)  # hard-cap at 16
+
+    # Preprocessed tensor on device
+    frames = frames.to(device)
+
+    windows = []
+    for i in range(T):
+        start = max(0, i - L + 1)
+        window = frames[start : i + 1]  # (len<=L, C, H, W)
+
+        if window.shape[0] < L:
+            pad_needed = L - window.shape[0]
+            pad = window[:1].expand(pad_needed, -1, -1, -1)  # repeat first frame
+            window = torch.cat([pad, window], dim=0)
+
+        windows.append(window)
+
+    preds = np.zeros(T, dtype=float)
+
+    for s in range(0, T, batch_size):
+        e = min(s + batch_size, T)
+        batch_windows = torch.stack(windows[s:e])  # (B, L, C, H, W)
+
+        batch = {OBS_IMAGES: batch_windows, OBS_LANGUAGE: [language] * (e - s)}  # expects (B, L, C, H, W)
+
+        # Model should return (B, L) or (B,) final-step values. We take the last step.
+        values = model.predict_rewards(batch)  # torch.Tensor
+
+        if values.dim() == 2:
+            last = values[:, -1]
+        else:
+            last = values.squeeze(-1)
+
+        preds[s:e] = last.detach().float().cpu().numpy()
+
+    return preds
+
+
+def plot_episode_eval(episode_idx, gt, pred, language, save_path=None, show=False, title_prefix="RLearN Eval"):
+    """Plot GT vs Predicted over time. Saves PNG if save_path is provided."""
+    T = len(pred)
+    x = np.arange(T)
+
+    plt.figure(figsize=(14, 8))
+    plt.plot(x, pred, linewidth=2.5, marker="o", markersize=3, label="Predicted Reward", color="blue")
+
+    if gt is not None:
+        plt.plot(x, gt, linestyle="--", linewidth=2.5, label="Ground-Truth Reward", color="orange")
+        # Correlation between GT and Pred
+        corr, p = spearmanr(gt, pred)
+        corr_str = f"ρ(GT, Pred) = {0.0 if np.isnan(corr) else corr:.3f} (p={0.0 if np.isnan(p) else p:.3f})"
+    else:
+        expected = np.linspace(0, 1, T)
+        plt.plot(x, expected, linestyle="--", linewidth=2.5, label="Expected Progress (0→1)", color="orange")
+        corr, p = spearmanr(x, pred)
+        corr_str = f"VOC-S ρ(t, Pred) = {0.0 if np.isnan(corr) else corr:.3f} (p={0.0 if np.isnan(p) else p:.3f})"
+
+    plt.title(f"{title_prefix} — Episode {episode_idx}\n{language}\n{corr_str}", fontsize=14)
+    plt.xlabel("Frame Index", fontsize=12)
+    plt.ylabel("Reward / Progress", fontsize=12)
+    plt.legend(fontsize=11)
+    plt.grid(True, alpha=0.3)
+    plt.tight_layout()
+
+    if save_path is not None:
+        plt.savefig(save_path, dpi=200, bbox_inches="tight")
+        print(f"Saved eval image to: {save_path}")
+
+    if show:
+        plt.show()
+    else:
+        plt.close()
+
+
+def eval_episode_sliding(
+    episode_idx, dataset, model, save_dir=".", device="cuda", max_seq_len=16, batch_size=64, title_prefix="RLearN Eval"
+):
+    """End-to-end: load episode, predict with sliding 16-frame windows, and save PNG."""
+    frames, language, gt = extract_episode_frames_and_gt(dataset, episode_idx)
+    if frames is None:
+        print(f"[Episode {episode_idx}] No frames found.")
+        return None
+
+    model.eval()
+
+    pred = predict_rewards_sliding(
+        model=model, frames=frames, language=language, max_seq_len=max_seq_len, batch_size=batch_size, device=device
+    )
+
+    # Basic stats
+    print(f"Episode {episode_idx}: T={len(pred)}, pred∈[{pred.min():.3f},{pred.max():.3f}]")
+    if gt is not None:
+        print(f"GT available: gt∈[{np.nanmin(gt):.3f},{np.nanmax(gt):.3f}]")
+
+    save_path = f"{save_dir}/episode_{episode_idx:04d}_eval.png"
+    plot_episode_eval(
+        episode_idx=episode_idx, gt=gt, pred=pred, language=language, save_path=save_path, show=False, title_prefix=title_prefix
+    )
+    return save_path
+
+
+def main():
+    """Main evaluation script for RLearN models."""
+    # Parse command line arguments
+    parser = argparse.ArgumentParser(description="Evaluate RLearN model on episodes with GT vs Predicted rewards")
+    parser.add_argument("--model", type=str, required=True, help="Model name/path (e.g., pepijn223/rlearn_mse5)")
+    parser.add_argument("--dataset", type=str, required=True, help="Dataset repo (e.g., pepijn223/phone_pipeline_pickup1)")
+    parser.add_argument("--episodes", type=int, default=5, help="Number of episodes to evaluate")
+    parser.add_argument("--output", type=str, default="./eval_results", help="Output directory for images")
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu",
+        help="Device to use",
+    )
+    parser.add_argument("--batch_size", type=int, default=32, help="Batch size for sliding window evaluation")
+
+    args = parser.parse_args()
+
+    # Create output directory
+    output_dir = Path(args.output)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    print("🎯 RLearN Model Evaluation")
+    print("=" * 60)
+    print(f"Model: {args.model}")
+    print(f"Dataset: {args.dataset}")
+    print(f"Episodes: {args.episodes}")
+    print(f"Device: {args.device}")
+    print(f"Output: {output_dir}")
+    print("=" * 60)
+
+    try:
+        # Load dataset
+        print("📁 Loading dataset...")
+
+        dataset = LeRobotDataset(
+            repo_id=args.dataset,
+            episodes=list(range(min(args.episodes, 50))),  # Load enough episodes
+            download_videos=True,
+        )
+
+        print(f"✅ Dataset loaded: {dataset.num_episodes} episodes, {dataset.num_frames} frames")
+        print(f"   Features: {list(dataset.features.keys())}")
+        print(f"   FPS: {dataset.fps}")
+
+        # Load model
+        print("\n🤖 Loading model...")
+
+        model = RLearNPolicy.from_pretrained(args.model)
+        model = model.to(args.device)
+        model.eval()
+
+        print(f"✅ Model loaded on {args.device}")
+        print(f"   Parameters: {sum(p.numel() for p in model.parameters()):,}")
+        print(f"   Trainable: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")
+        print(f"   Max sequence length: {model.config.max_seq_len}")
+
+        # Select episodes to evaluate
+        total_available = min(dataset.num_episodes, args.episodes)
+        episode_indices = list(range(total_available))
+
+        print(f"\n📊 Evaluating {len(episode_indices)} episodes...")
+        print("=" * 60)
+
+        # Run sliding window evaluation on each episode
+        saved_paths = []
+        for i, ep_idx in enumerate(episode_indices):
+            print(f"\n[{i+1}/{len(episode_indices)}] Processing Episode {ep_idx}")
+            print("-" * 40)
+
+            try:
+                save_path = eval_episode_sliding(
+                    episode_idx=ep_idx,
+                    dataset=dataset,
+                    model=model,
+                    save_dir=str(output_dir),
+                    device=args.device,
+                    batch_size=args.batch_size,
+                    title_prefix="RLearN Ground Truth vs Predicted",
+                )
+
+                if save_path:
+                    saved_paths.append(save_path)
+
+            except Exception as e:
+                print(f"❌ Error processing episode {ep_idx}: {e}")
+                import traceback
+
+                traceback.print_exc()
+                continue
+
+        # Summary
+        print("\n" + "=" * 60)
+        print("✅ EVALUATION COMPLETE")
+        print(f"📈 Generated {len(saved_paths)} evaluation plots")
+        print(f"📁 Results saved to: {output_dir}")
+        print("\nGenerated files:")
+        for path in saved_paths:
+            print(f"  • {path}")
+
+        if saved_paths:
+            print(f"\n💡 View the plots to compare ground truth vs predicted rewards!")
+            print(f"   Each plot shows the model's sliding 16-frame window predictions")
+            print(f"   against available ground truth rewards over the episode timeline.")
+
+        return 0
+
+    except Exception as e:
+        print(f"❌ Error during evaluation: {e}")
+        import traceback
+
+        traceback.print_exc()
+        return 1
+
+
+if __name__ == "__main__":
+    exit(main())

src/lerobot/policies/rlearn/eval_visualizer.py

@@ -1,511 +0,0 @@
-#!/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 RLearN evaluation during training.
-
-Creates and saves reward prediction visualizations for held-out episodes.
-"""
-
-from __future__ import annotations
-
-import warnings
-from pathlib import Path
-from typing import Any
-
-import matplotlib.pyplot as plt
-import numpy as np
-import torch
-from matplotlib import rcParams
-from scipy.stats import spearmanr
-from torch import Tensor
-
-from lerobot.constants import OBS_IMAGES, OBS_LANGUAGE
-
-# Set matplotlib backend to avoid GUI issues during training
-rcParams['backend'] = 'Agg'
-
-
-class RLearNEvalVisualizer:
-    """
-    Creates visualization plots for RLearN model evaluation during training.
-    
-    Generates reward prediction plots similar to the evaluation notebook but saves
-    them as images for monitoring training progress.
-    """
-    
-    def __init__(self, model, dataset, device: str = "cuda"):
-        """
-        Args:
-            model: RLearN model instance
-            dataset: LeRobot dataset instance
-            device: Device to run evaluation on
-        """
-        self.model = model
-        self.dataset = dataset
-        self.device = device
-        
-    def get_episode_data(self, episode_idx: int, max_frames: int = 64) -> tuple[Tensor | None, str | None, np.ndarray | None, int | None]:
-        """Extract frames, language, and predict rewards for an episode."""
-        try:
-            # Get episode data
-            ep_start = self.dataset.episode_data_index["from"][episode_idx].item()
-            ep_end = self.dataset.episode_data_index["to"][episode_idx].item()
-            episode_length = min(ep_end - ep_start, max_frames)
-
-            # Collect frames and get language
-            frames = []
-            language = None
-
-            for frame_idx in range(episode_length):
-                global_idx = ep_start + frame_idx
-                frame_data = self.dataset[global_idx]
-
-                # Extract image
-                if OBS_IMAGES in frame_data:
-                    img = frame_data[OBS_IMAGES]
-                else:
-                    img_keys = [k for k in frame_data.keys() if "image" in k.lower()]
-                    if img_keys:
-                        img = frame_data[img_keys[0]]
-                    else:
-                        continue
-
-                if isinstance(img, np.ndarray):
-                    img = torch.from_numpy(img)
-
-                # Ensure CHW format
-                if len(img.shape) == 3 and img.shape[-1] in [1, 3, 4]:
-                    img = img.permute(2, 0, 1)
-
-                # Resize to expected input size (224x224 for SigLIP2)
-                if img.shape[-2:] != (224, 224):
-                    import torch.nn.functional as F
-                    img = F.interpolate(
-                        img.unsqueeze(0), size=(224, 224), mode="bilinear", align_corners=False
-                    ).squeeze(0)
-
-                # Normalize to [0, 1] if needed
-                if img.dtype == torch.uint8:
-                    img = img.float() / 255.0
-
-                frames.append(img)
-
-                # Get language
-                if language is None:
-                    if OBS_LANGUAGE in frame_data:
-                        language = frame_data[OBS_LANGUAGE]
-                        if isinstance(language, list):
-                            language = language[0]
-                    elif "task" in frame_data:
-                        language = frame_data["task"]
-                    else:
-                        language = "No language provided"
-
-            if not frames:
-                return None, None, None, None
-
-            frames_tensor = torch.stack(frames)
-            
-            # Predict rewards using the model's evaluation method
-            with torch.no_grad():
-                self.model.eval()
-                rewards = self._predict_episode_rewards(frames_tensor, language)
-
-            return frames_tensor, language, rewards, episode_length
-
-        except Exception as e:
-            warnings.warn(f"Error processing episode {episode_idx}: {e}")
-            return None, None, None, None
-
-    @torch.no_grad()
-    def _predict_episode_rewards(self, frames: Tensor, language: str, batch_size: int = 16) -> np.ndarray:
-        """
-        Predict rewards for a single episode using proper temporal sequences.
-        
-        Args:
-            frames: Video frames tensor of shape (T, C, H, W)
-            language: Language instruction string
-            batch_size: Maximum number of temporal sequences to process at once
-
-        Returns:
-            Predicted progress/rewards array of shape (T,)
-        """
-        T = frames.shape[0]
-        max_seq_len = self.model.config.max_seq_len
-
-        # Create temporal sequences for each frame
-        temporal_sequences = []
-
-        for i in range(T):
-            # Create sequence ending at frame i
-            seq_frames = []
-            for j in range(max(0, i - max_seq_len + 1), i + 1):
-                # Use frame j if available, otherwise repeat the first available frame
-                frame_idx = max(0, min(j, T - 1))
-                seq_frames.append(frames[frame_idx])
-
-            # Pad sequence to max_seq_len by repeating the first frame if needed
-            while len(seq_frames) < max_seq_len:
-                seq_frames.insert(0, seq_frames[0])  # Prepend first frame
-
-            # Take only the last max_seq_len frames if we have too many
-            seq_frames = seq_frames[-max_seq_len:]
-            temporal_sequences.append(torch.stack(seq_frames))  # (max_seq_len, C, H, W)
-
-        # Stack all temporal sequences: (T, max_seq_len, C, H, W)
-        all_sequences = torch.stack(temporal_sequences)
-
-        # Process in batches
-        rewards = []
-        for i in range(0, T, batch_size):
-            end_idx = min(i + batch_size, T)
-            batch_sequences = all_sequences[i:end_idx].to(self.device)  # (B, max_seq_len, C, H, W)
-
-            # Create batch for model
-            batch = {
-                OBS_IMAGES: batch_sequences,  # (B, T, C, H, W) format expected by model
-                OBS_LANGUAGE: [language] * batch_sequences.shape[0],
-            }
-
-            # Predict rewards - model returns (B, T') but we want the last timestep for each sequence
-            values = self.model.predict_rewards(batch)  # (B, T')
-
-            # Take the last timestep prediction for each sequence (represents current frame reward)
-            if values.dim() == 2:
-                batch_rewards = values[:, -1].cpu().numpy()  # (B,) - last timestep
-            else:
-                batch_rewards = values.cpu().numpy()  # (B,) - already single timestep
-
-            rewards.extend(batch_rewards)
-
-        return np.array(rewards[:T])  # Ensure exact length
-
-    def create_episode_grid_visualization(
-        self, 
-        episode_indices: list[int], 
-        save_path: Path, 
-        step: int | None = None,
-        max_frames: int = 64
-    ) -> dict[str, Any]:
-        """
-        Create a 3x3 grid visualization of episode reward predictions.
-        
-        Args:
-            episode_indices: List of 9 episode indices to visualize
-            save_path: Path to save the visualization image
-            step: Training step (for title)
-            max_frames: Maximum frames per episode to process
-            
-        Returns:
-            Dictionary with evaluation metrics
-        """
-        if len(episode_indices) != 9:
-            raise ValueError("Expected exactly 9 episode indices for 3x3 grid")
-            
-        # Create figure with 3x3 subplots
-        fig, axes = plt.subplots(3, 3, figsize=(20, 16))
-        axes = axes.flatten()
-        
-        eval_metrics = {
-            "voc_s_scores": [],
-            "episode_lengths": [],
-            "reward_ranges": [],
-            "languages": []
-        }
-        
-        for i, episode_idx in enumerate(episode_indices):
-            ax = axes[i]
-            
-            frames, language, rewards, episode_length = self.get_episode_data(episode_idx, max_frames)
-            
-            if rewards is None:
-                ax.text(
-                    0.5, 0.5, f"Episode {episode_idx}\nNo data available",
-                    ha="center", va="center", transform=ax.transAxes,
-                    fontsize=12, bbox=dict(boxstyle="round,pad=0.3", facecolor="lightcoral", alpha=0.7)
-                )
-                ax.set_title(f"Episode {episode_idx} - Error", fontsize=12, pad=10)
-                continue
-                
-            # Plot predicted rewards
-            time_steps = range(len(rewards))
-            ax.plot(
-                time_steps, rewards, "b-", linewidth=2.5, marker="o", markersize=5, 
-                label="Predicted Reward", alpha=0.8
-            )
-            
-            # Add expected progress line (ground truth for ReWiND)
-            expected_progress = np.linspace(0, 1, len(rewards))
-            ax.plot(
-                time_steps, expected_progress, "orange", linestyle="--", linewidth=2.5, 
-                label="Expected Progress (0→1)", alpha=0.8
-            )
-            
-            # Compute VOC-S (Value-Order Correlation for Success)
-            frame_indices = np.arange(1, len(rewards) + 1)
-            correlation, p_value = spearmanr(frame_indices, rewards)
-            if np.isnan(correlation):
-                correlation = 0.0
-                
-            eval_metrics["voc_s_scores"].append(correlation)
-            eval_metrics["episode_lengths"].append(len(rewards))
-            eval_metrics["reward_ranges"].append((rewards.min(), rewards.max()))
-            eval_metrics["languages"].append(language)
-            
-            # Format title with language (truncated) and VOC-S
-            title_lang = language[:35] + "..." if len(language) > 35 else language
-            title = f'Episode {episode_idx}\n"{title_lang}"\nVOC-S: {correlation:.3f}'
-            ax.set_title(title, fontsize=10, pad=15)
-            
-            ax.set_xlabel("Frame Index", fontsize=10)
-            ax.set_ylabel("Reward", fontsize=10)
-            ax.legend(fontsize=8, loc='upper left')
-            ax.grid(True, alpha=0.3)
-            
-            # Color-coded trend indicator
-            if correlation > 0.3:
-                trend_text = "↗ Strong+"
-                trend_color = "darkgreen"
-            elif correlation > 0.1:
-                trend_text = "↗ Weak+"
-                trend_color = "green"
-            elif correlation < -0.3:
-                trend_text = "↘ Strong-"
-                trend_color = "darkred"
-            elif correlation < -0.1:
-                trend_text = "↘ Weak-"
-                trend_color = "red"
-            else:
-                trend_text = "→ Flat"
-                trend_color = "gray"
-                
-            ax.text(
-                0.02, 0.98, trend_text, transform=ax.transAxes,
-                verticalalignment="top", fontsize=9, fontweight="bold",
-                bbox=dict(boxstyle="round,pad=0.3", facecolor=trend_color, alpha=0.2),
-                color=trend_color
-            )
-            
-            # Add reward range info
-            ax.text(
-                0.98, 0.02, f"Range: [{rewards.min():.3f}, {rewards.max():.3f}]",
-                transform=ax.transAxes, ha="right", va="bottom", fontsize=8,
-                bbox=dict(boxstyle="round,pad=0.2", facecolor="lightblue", alpha=0.5)
-            )
-        
-        # Add overall title
-        step_text = f" - Step {step}" if step is not None else ""
-        fig.suptitle(
-            f"RLearN Reward Evaluation{step_text}\n"
-            f"Mean VOC-S: {np.mean(eval_metrics['voc_s_scores']):.3f} | "
-            f"Episodes: {len([s for s in eval_metrics['voc_s_scores'] if s != 0])}/9",
-            fontsize=16, y=0.95
-        )
-        
-        plt.tight_layout()
-        plt.subplots_adjust(top=0.90)  # Make room for suptitle
-        
-        # Save the figure
-        save_path.parent.mkdir(parents=True, exist_ok=True)
-        plt.savefig(save_path, dpi=150, bbox_inches='tight', facecolor='white')
-        plt.close()  # Close to free memory
-        
-        # Calculate summary metrics
-        valid_scores = [s for s in eval_metrics["voc_s_scores"] if s != 0]
-        summary = {
-            "mean_voc_s": np.mean(valid_scores) if valid_scores else 0.0,
-            "std_voc_s": np.std(valid_scores) if valid_scores else 0.0,
-            "num_valid_episodes": len(valid_scores),
-            "total_episodes": len(episode_indices),
-            "mean_episode_length": np.mean(eval_metrics["episode_lengths"]) if eval_metrics["episode_lengths"] else 0,
-            "individual_scores": eval_metrics["voc_s_scores"],
-            "episode_languages": eval_metrics["languages"]
-        }
-        
-        return summary
-        
-    def create_comparison_visualization(
-        self,
-        episode_indices: list[int],
-        save_path: Path,
-        step: int | None = None,
-        max_frames: int = 64,
-        mismatch_templates: list[str] | None = None
-    ) -> dict[str, Any]:
-        """
-        Create correct vs incorrect language comparison visualization.
-        
-        Args:
-            episode_indices: List of episode indices to compare (up to 6)
-            save_path: Path to save the visualization image  
-            step: Training step (for title)
-            max_frames: Maximum frames per episode to process
-            mismatch_templates: Custom mismatch templates
-            
-        Returns:
-            Dictionary with detection metrics
-        """
-        if mismatch_templates is None:
-            mismatch_templates = [
-                "kick the ball", "clean the sink", "dance in place", 
-                "wave your hand", "jump up and down", "do nothing"
-            ]
-            
-        # Limit to 6 episodes for 2x3 grid
-        episode_indices = episode_indices[:6]
-        n_episodes = len(episode_indices)
-        
-        # Create figure with 2x3 subplots
-        fig, axes = plt.subplots(2, 3, figsize=(18, 12))
-        axes = axes.flatten()
-        
-        detection_results = {
-            "correct_finals": [],
-            "incorrect_finals": [],
-            "detection_successes": [],
-            "episode_info": []
-        }
-        
-        for i, episode_idx in enumerate(episode_indices):
-            if i >= 6:  # Limit to 6 subplots
-                break
-                
-            ax = axes[i]
-            
-            # Get episode data with correct language
-            frames, correct_language, correct_rewards, episode_length = self.get_episode_data(episode_idx, max_frames)
-            
-            if correct_rewards is None:
-                ax.text(
-                    0.5, 0.5, f"Episode {episode_idx}\nNo data available",
-                    ha="center", va="center", transform=ax.transAxes
-                )
-                ax.set_title(f"Episode {episode_idx} - Error")
-                continue
-                
-            # Generate incorrect language and predict
-            incorrect_language = mismatch_templates[i % len(mismatch_templates)]
-            incorrect_rewards = self._predict_episode_rewards(frames, incorrect_language)
-            
-            # Plot both reward curves
-            time_steps = range(len(correct_rewards))
-            ax.plot(
-                time_steps, correct_rewards, "g-", linewidth=2.5, marker="o", markersize=4,
-                label=f"Correct: '{correct_language[:25]}...'" if len(correct_language) > 25 else f"Correct: '{correct_language}'"
-            )
-            ax.plot(
-                time_steps, incorrect_rewards, "r-", linewidth=2.5, marker="s", markersize=4,
-                label=f"Incorrect: '{incorrect_language}'"
-            )
-            
-            # Calculate detection success
-            final_correct = correct_rewards[-1]
-            final_incorrect = incorrect_rewards[-1]
-            detection_success = final_correct > final_incorrect
-            
-            detection_results["correct_finals"].append(final_correct)
-            detection_results["incorrect_finals"].append(final_incorrect)
-            detection_results["detection_successes"].append(detection_success)
-            detection_results["episode_info"].append({
-                "episode_idx": episode_idx,
-                "correct_language": correct_language,
-                "incorrect_language": incorrect_language,
-                "final_correct": final_correct,
-                "final_incorrect": final_incorrect
-            })
-            
-            # Color-coded title based on detection success
-            success_indicator = "✓" if detection_success else "✗"
-            title_color = "darkgreen" if detection_success else "darkred"
-            ax.set_title(
-                f"Episode {episode_idx} {success_indicator}\nΔ: {final_correct - final_incorrect:.3f}",
-                color=title_color, fontweight="bold", fontsize=11
-            )
-            
-            ax.set_xlabel("Frame Index")
-            ax.set_ylabel("Reward")
-            ax.legend(fontsize=8, loc='upper left')
-            ax.grid(True, alpha=0.3)
-            
-            # Add final reward values as text
-            ax.text(
-                0.98, 0.02, 
-                f"Final: C={final_correct:.3f}, I={final_incorrect:.3f}",
-                transform=ax.transAxes, ha="right", va="bottom", fontsize=9,
-                bbox=dict(boxstyle="round,pad=0.3", facecolor="lightyellow", alpha=0.7)
-            )
-        
-        # Hide unused subplots
-        for i in range(n_episodes, 6):
-            axes[i].axis('off')
-        
-        # Calculate summary metrics
-        detection_accuracy = np.mean(detection_results["detection_successes"]) if detection_results["detection_successes"] else 0.0
-        mean_correct = np.mean(detection_results["correct_finals"]) if detection_results["correct_finals"] else 0.0
-        mean_incorrect = np.mean(detection_results["incorrect_finals"]) if detection_results["incorrect_finals"] else 0.0
-        
-        # Add overall title
-        step_text = f" - Step {step}" if step is not None else ""
-        fig.suptitle(
-            f"RLearN Language Detection{step_text}\n"
-            f"Accuracy: {detection_accuracy:.1%} | Mean Δ: {mean_correct - mean_incorrect:.3f} | "
-            f"Success: {sum(detection_results['detection_successes'])}/{len(detection_results['detection_successes'])}",
-            fontsize=16, y=0.95
-        )
-        
-        plt.tight_layout()
-        plt.subplots_adjust(top=0.90)
-        
-        # Save the figure
-        save_path.parent.mkdir(parents=True, exist_ok=True)
-        plt.savefig(save_path, dpi=150, bbox_inches='tight', facecolor='white')
-        plt.close()
-        
-        summary = {
-            "detection_accuracy": detection_accuracy,
-            "mean_correct_final": mean_correct,
-            "mean_incorrect_final": mean_incorrect,
-            "separation_score": mean_correct - mean_incorrect,
-            "num_episodes": len(detection_results["detection_successes"]),
-            "individual_results": detection_results["episode_info"]
-        }
-        
-        return summary
-
-
-def select_evaluation_episodes(dataset, num_episodes: int = 9, seed: int = 42) -> list[int]:
-    """
-    Select a diverse set of episodes for evaluation holdout.
-    
-    Args:
-        dataset: LeRobot dataset instance
-        num_episodes: Number of episodes to select
-        seed: Random seed for reproducibility
-        
-    Returns:
-        List of episode indices
-    """
-    np.random.seed(seed)
-    
-    total_episodes = dataset.num_episodes
-    if num_episodes >= total_episodes:
-        return list(range(total_episodes))
-    
-    # Select random episodes
-    episode_indices = np.random.choice(total_episodes, num_episodes, replace=False).tolist()
-    
-    return sorted(episode_indices)

src/lerobot/policies/rlearn/evaluation.py

@@ -1,695 +0,0 @@
-#!/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.
-
-"""
-Evaluation metrics for RLearn (Video-Language Conditioned Reward Model).
-
-Key metrics:
-1. VOC-S (Value-Order Correlation for Success): Spearman correlation between frame indices and predicted rewards
-2. Success vs Failure Detection: Model's ability to distinguish between correct and incorrect language conditions
-"""
-
-from __future__ import annotations
-
-import warnings
-from typing import Any
-
-import numpy as np
-import torch
-from scipy.stats import spearmanr
-from torch import Tensor
-from tqdm import tqdm
-
-from lerobot.constants import OBS_IMAGES, OBS_LANGUAGE
-
-
-def compute_voc_s(
-    predicted_rewards: list[np.ndarray], use_interquartile_mean: bool = True
-) -> dict[str, float]:
-    """
-    Compute Value-Order Correlation for Success (VOC-S).
-
-    Measures whether per-frame rewards increase as successful execution unfolds.
-    For each episode, computes Spearman correlation between frame indices [1..T]
-    and predicted rewards [r1..rT].
-
-    Args:
-        predicted_rewards: List of reward arrays, one per episode. Each array has shape (T,)
-        use_interquartile_mean: If True, use IQM instead of mean for aggregation
-
-    Returns:
-        Dictionary with VOC-S metrics:
-        - voc_s_mean: Mean Spearman correlation across episodes
-        - voc_s_std: Standard deviation of correlations
-        - voc_s_iqm: Interquartile mean (if use_interquartile_mean=True)
-        - num_episodes: Number of episodes evaluated
-        - correlations: Individual correlations per episode
-    """
-    if not predicted_rewards:
-        return {"voc_s_mean": 0.0, "voc_s_std": 0.0, "voc_s_iqm": 0.0, "num_episodes": 0, "correlations": []}
-
-    correlations = []
-
-    for episode_rewards in predicted_rewards:
-        if len(episode_rewards) < 2:
-            # Need at least 2 points for correlation
-            continue
-
-        # Frame indices: [1, 2, ..., T]
-        frame_indices = np.arange(1, len(episode_rewards) + 1)
-
-        # Compute Spearman correlation
-        try:
-            correlation, p_value = spearmanr(frame_indices, episode_rewards)
-
-            # Handle NaN correlations (e.g., all rewards are identical)
-            if np.isnan(correlation):
-                correlation = 0.0
-
-            correlations.append(correlation)
-
-        except Exception as e:
-            warnings.warn(f"Failed to compute correlation for episode: {e}")
-            correlations.append(0.0)
-
-    if not correlations:
-        return {"voc_s_mean": 0.0, "voc_s_std": 0.0, "voc_s_iqm": 0.0, "num_episodes": 0, "correlations": []}
-
-    correlations = np.array(correlations)
-
-    # Compute statistics
-    voc_s_mean = float(np.mean(correlations))
-    voc_s_std = float(np.std(correlations))
-
-    # Interquartile mean: mean of values between 25th and 75th percentiles
-    if use_interquartile_mean and len(correlations) >= 4:
-        q25, q75 = np.percentile(correlations, [25, 75])
-        iqm_mask = (correlations >= q25) & (correlations <= q75)
-        voc_s_iqm = float(np.mean(correlations[iqm_mask]))
-    else:
-        voc_s_iqm = voc_s_mean
-
-    return {
-        "voc_s_mean": voc_s_mean,
-        "voc_s_std": voc_s_std,
-        "voc_s_iqm": voc_s_iqm,
-        "num_episodes": len(correlations),
-        "correlations": correlations.tolist(),
-    }
-
-
-def compute_success_failure_detection(
-    correct_rewards: list[np.ndarray], incorrect_rewards: list[np.ndarray], threshold_percentile: float = 50.0
-) -> dict[str, float]:
-    """
-    Compute success vs failure detection accuracy.
-
-    Tests the model's ability to distinguish between correct and incorrect language conditions.
-    For each episode, compares final reward under correct vs incorrect language instruction.
-
-    Args:
-        correct_rewards: List of reward arrays for episodes with correct language
-        incorrect_rewards: List of reward arrays for episodes with incorrect/mismatched language
-        threshold_percentile: Percentile of correct rewards to use as threshold
-
-    Returns:
-        Dictionary with detection metrics:
-        - detection_accuracy: Fraction of episodes where correct > incorrect
-        - mean_correct_final: Mean final reward for correct language
-        - mean_incorrect_final: Mean final reward for incorrect language
-        - separation_score: (mean_correct - mean_incorrect) / (std_correct + std_incorrect)
-        - num_pairs: Number of episode pairs evaluated
-    """
-    if len(correct_rewards) != len(incorrect_rewards):
-        raise ValueError("Must have same number of correct and incorrect reward sequences")
-
-    if not correct_rewards:
-        return {
-            "detection_accuracy": 0.0,
-            "mean_correct_final": 0.0,
-            "mean_incorrect_final": 0.0,
-            "separation_score": 0.0,
-            "num_pairs": 0,
-        }
-
-    # Extract final rewards (last timestep of each episode)
-    correct_finals = []
-    incorrect_finals = []
-
-    for correct_ep, incorrect_ep in zip(correct_rewards, incorrect_rewards, strict=False):
-        if len(correct_ep) > 0 and len(incorrect_ep) > 0:
-            correct_finals.append(correct_ep[-1])  # Final reward
-            incorrect_finals.append(incorrect_ep[-1])  # Final reward
-
-    if not correct_finals:
-        return {
-            "detection_accuracy": 0.0,
-            "mean_correct_final": 0.0,
-            "mean_incorrect_final": 0.0,
-            "separation_score": 0.0,
-            "num_pairs": 0,
-        }
-
-    correct_finals = np.array(correct_finals)
-    incorrect_finals = np.array(incorrect_finals)
-
-    # Detection accuracy: fraction where correct > incorrect
-    detection_accuracy = float(np.mean(correct_finals > incorrect_finals))
-
-    # Statistics
-    mean_correct = float(np.mean(correct_finals))
-    mean_incorrect = float(np.mean(incorrect_finals))
-    std_correct = float(np.std(correct_finals))
-    std_incorrect = float(np.std(incorrect_finals))
-
-    # Separation score: normalized difference (clamp to prevent extreme values)
-    denominator = std_correct + std_incorrect
-    if denominator > 1e-6:  # Prevent division by very small numbers
-        separation_score = (mean_correct - mean_incorrect) / denominator
-        # Clamp to reasonable range
-        separation_score = np.clip(separation_score, -100.0, 100.0)
-    else:
-        separation_score = 0.0
-
-    return {
-        "detection_accuracy": detection_accuracy,
-        "mean_correct_final": mean_correct,
-        "mean_incorrect_final": mean_incorrect,
-        "separation_score": float(separation_score),
-        "num_pairs": len(correct_finals),
-    }
-
-
-def generate_mismatched_languages(
-    original_languages: list[str], mismatch_templates: list[str] | None = None
-) -> list[str]:
-    """
-    Generate mismatched language instructions for failure detection evaluation.
-
-    Args:
-        original_languages: List of original task descriptions
-        mismatch_templates: Custom mismatch templates. If None, uses defaults.
-
-    Returns:
-        List of mismatched language instructions
-    """
-    if mismatch_templates is None:
-        mismatch_templates = ["kick the ball", "walk to the red shoes", "wave", "do nothing"]
-
-    # For each original language, pick a random mismatch
-    mismatched = []
-    np.random.seed(42)  # For reproducibility
-
-    for i, orig_lang in enumerate(original_languages):
-        # Use modulo to cycle through mismatches if we have more episodes than templates
-        mismatch_idx = i % len(mismatch_templates)
-        mismatched.append(mismatch_templates[mismatch_idx])
-
-    return mismatched
-
-
-class RLearnEvaluator:
-    """
-    Comprehensive evaluator for RLearN reward models.
-
-    Provides methods to evaluate VOC-S and success/failure detection on datasets.
-    """
-
-    def __init__(self, model, device: str = "cuda"):
-        """
-        Args:
-            model: RLearN model instance
-            device: Device to run evaluation on
-        """
-        self.model = model
-        self.device = device
-        self.model.eval()
-
-    @torch.no_grad()
-    def predict_episode_rewards(self, frames: Tensor, language: str, batch_size: int = 16) -> np.ndarray:
-        """
-        Predict rewards for a single episode using proper temporal sequences.
-        
-        Note: With ReWiND loss, the model predicts progress values (0-1) across episodes,
-        which serve as dense reward signals for policy learning.
-
-        Args:
-            frames: Video frames tensor of shape (T, C, H, W)
-            language: Language instruction string
-            batch_size: Maximum number of temporal sequences to process at once
-
-        Returns:
-            Predicted progress/rewards array of shape (T,) with values typically in [0, 1]
-        """
-        T = frames.shape[0]
-        max_seq_len = self.model.config.max_seq_len
-
-        # Preprocess frames to match model expectations
-        processed_frames = self._preprocess_frames(frames)
-
-        # Create temporal sequences for each frame
-        # For frame i, we want frames [i-max_seq_len+1, ..., i-1, i]
-        temporal_sequences = []
-
-        for i in range(T):
-            # Create sequence ending at frame i
-            seq_frames = []
-            for j in range(max(0, i - max_seq_len + 1), i + 1):
-                # Use frame j if available, otherwise repeat the first available frame
-                frame_idx = max(0, min(j, T - 1))
-                seq_frames.append(processed_frames[frame_idx])
-
-            # Pad sequence to max_seq_len by repeating the first frame if needed
-            while len(seq_frames) < max_seq_len:
-                seq_frames.insert(0, seq_frames[0])  # Prepend first frame
-
-            # Take only the last max_seq_len frames if we have too many
-            seq_frames = seq_frames[-max_seq_len:]
-
-            temporal_sequences.append(torch.stack(seq_frames))  # (max_seq_len, C, H, W)
-
-        # Stack all temporal sequences: (T, max_seq_len, C, H, W)
-        all_sequences = torch.stack(temporal_sequences)
-
-        # Process in batches
-        rewards = []
-        for i in range(0, T, batch_size):
-            end_idx = min(i + batch_size, T)
-            batch_sequences = all_sequences[i:end_idx].to(self.device)  # (B, max_seq_len, C, H, W)
-
-            # Create batch for model
-            batch = {
-                OBS_IMAGES: batch_sequences,  # (B, T, C, H, W) format expected by model
-                OBS_LANGUAGE: [language] * batch_sequences.shape[0],
-            }
-
-            # Predict rewards - model returns (B, T') but we want the last timestep for each sequence
-            values = self.model.predict_rewards(batch)  # (B, T')
-
-            # Take the last timestep prediction for each sequence (represents current frame reward)
-            if values.dim() == 2:
-                batch_rewards = values[:, -1].cpu().numpy()  # (B,) - last timestep
-            else:
-                batch_rewards = values.cpu().numpy()  # (B,) - already single timestep
-
-            rewards.extend(batch_rewards)
-
-        return np.array(rewards[:T])  # Ensure exact length
-
-    def _preprocess_frames(self, frames: Tensor) -> Tensor:
-        """
-        Preprocess frames to match model expectations.
-
-        Args:
-            frames: Input frames tensor of shape (T, C, H, W)
-
-        Returns:
-            Preprocessed frames tensor of shape (T, C, H', W')
-        """
-        import torch.nn.functional as F
-
-        T, C, H, W = frames.shape
-
-        # Expected input size for DINO v3 is 224x224
-        target_size = 224
-
-        # Resize frames if needed
-        if H != target_size or W != target_size:
-            # Resize using bilinear interpolation
-            frames = F.interpolate(
-                frames, size=(target_size, target_size), mode="bilinear", align_corners=False
-            )
-
-        # Normalize to [0, 1] if needed
-        if frames.dtype == torch.uint8:
-            frames = frames.float() / 255.0
-
-        # Ensure values are in [0, 1] range
-        frames = torch.clamp(frames, 0.0, 1.0)
-
-        return frames
-
-    def evaluate_voc_s(
-        self, dataset, num_episodes: int = 100, use_interquartile_mean: bool = True
-    ) -> dict[str, Any]:
-        """
-        Evaluate VOC-S on a dataset.
-
-        Args:
-            dataset: LeRobot dataset instance
-            num_episodes: Number of episodes to evaluate (randomly sampled)
-            use_interquartile_mean: Whether to compute IQM
-
-        Returns:
-            VOC-S evaluation results
-        """
-        print(f"Evaluating VOC-S on {num_episodes} episodes...")
-
-        # Sample episodes
-        total_episodes = dataset.num_episodes
-        if num_episodes >= total_episodes:
-            episode_indices = list(range(total_episodes))
-        else:
-            np.random.seed(42)
-            episode_indices = np.random.choice(total_episodes, num_episodes, replace=False)
-
-        predicted_rewards = []
-
-        for ep_idx in tqdm(episode_indices, desc="Computing VOC-S"):
-            try:
-                # Get episode data
-                ep_start = dataset.episode_data_index["from"][ep_idx].item()
-                ep_end = dataset.episode_data_index["to"][ep_idx].item()
-                episode_length = ep_end - ep_start
-
-                # Get frames and language for this episode
-                frames = []
-                language = None
-
-                for frame_idx in range(episode_length):
-                    global_idx = ep_start + frame_idx
-                    frame_data = dataset[global_idx]
-
-                    # Extract image (assuming single camera for now)
-                    if OBS_IMAGES in frame_data:
-                        img = frame_data[OBS_IMAGES]
-                    else:
-                        # Try to find image key
-                        img_keys = [k for k in frame_data.keys() if "image" in k.lower()]
-                        if img_keys:
-                            img = frame_data[img_keys[0]]
-                        else:
-                            continue
-
-                    # Convert to tensor if needed
-                    if isinstance(img, np.ndarray):
-                        img = torch.from_numpy(img)
-
-                    # Ensure CHW format
-                    if len(img.shape) == 3 and img.shape[-1] in [1, 3, 4]:
-                        img = img.permute(2, 0, 1)  # HWC -> CHW
-
-                    # Resize to expected input size (224x224 for DINO v3) BEFORE stacking
-                    if img.shape[-2:] != (224, 224):
-                        import torch.nn.functional as F
-
-                        img = F.interpolate(
-                            img.unsqueeze(0), size=(224, 224), mode="bilinear", align_corners=False
-                        ).squeeze(0)
-
-                    # Normalize to [0, 1] if needed
-                    if img.dtype == torch.uint8:
-                        img = img.float() / 255.0
-
-                    frames.append(img)
-
-                    # Get language instruction
-                    if language is None:
-                        if OBS_LANGUAGE in frame_data:
-                            language = frame_data[OBS_LANGUAGE]
-                            if isinstance(language, list):
-                                language = language[0]
-                        elif "task" in frame_data:
-                            language = frame_data["task"]
-                        else:
-                            language = ""  # Default empty language
-
-                if not frames:
-                    continue
-
-                # Stack frames into video tensor
-                frames_tensor = torch.stack(frames)  # (T, C, H, W)
-
-                # Predict rewards
-                episode_rewards = self.predict_episode_rewards(frames_tensor, language)
-                predicted_rewards.append(episode_rewards)
-
-            except Exception as e:
-                warnings.warn(f"Failed to process episode {ep_idx}: {e}")
-                continue
-
-        # Compute VOC-S
-        voc_results = compute_voc_s(predicted_rewards, use_interquartile_mean)
-
-        print("VOC-S Results:")
-        print(f"  Mean correlation: {voc_results['voc_s_mean']:.4f}")
-        print(f"  Std correlation: {voc_results['voc_s_std']:.4f}")
-        print(f"  IQM correlation: {voc_results['voc_s_iqm']:.4f}")
-        print(f"  Episodes evaluated: {voc_results['num_episodes']}")
-
-        return voc_results
-
-    def evaluate_success_failure_detection(
-        self, dataset, num_episodes: int = 100, mismatch_templates: list[str] | None = None
-    ) -> dict[str, Any]:
-        """
-        Evaluate success vs failure detection.
-
-        Args:
-            dataset: LeRobot dataset instance
-            num_episodes: Number of episodes to evaluate
-            mismatch_templates: Custom mismatch language templates
-
-        Returns:
-            Success/failure detection results
-        """
-        print(f"Evaluating success/failure detection on {num_episodes} episodes...")
-
-        # Sample episodes
-        total_episodes = dataset.num_episodes
-        if num_episodes >= total_episodes:
-            episode_indices = list(range(total_episodes))
-        else:
-            np.random.seed(42)
-            episode_indices = np.random.choice(total_episodes, num_episodes, replace=False)
-
-        correct_rewards = []
-        incorrect_rewards = []
-
-        # Get original languages
-        original_languages = []
-        for ep_idx in episode_indices:
-            ep_start = dataset.episode_data_index["from"][ep_idx].item()
-            frame_data = dataset[ep_start]
-
-            if OBS_LANGUAGE in frame_data:
-                lang = frame_data[OBS_LANGUAGE]
-                if isinstance(lang, list):
-                    lang = lang[0]
-            elif "task" in frame_data:
-                lang = frame_data["task"]
-            else:
-                lang = ""
-
-            original_languages.append(lang)
-
-        # Generate mismatched languages
-        mismatched_languages = generate_mismatched_languages(original_languages, mismatch_templates)
-
-        for i, ep_idx in enumerate(tqdm(episode_indices, desc="Computing detection metrics")):
-            try:
-                # Get episode frames (same as VOC-S evaluation)
-                ep_start = dataset.episode_data_index["from"][ep_idx].item()
-                ep_end = dataset.episode_data_index["to"][ep_idx].item()
-                episode_length = ep_end - ep_start
-
-                frames = []
-                for frame_idx in range(episode_length):
-                    global_idx = ep_start + frame_idx
-                    frame_data = dataset[global_idx]
-
-                    # Extract image
-                    if OBS_IMAGES in frame_data:
-                        img = frame_data[OBS_IMAGES]
-                    else:
-                        img_keys = [k for k in frame_data.keys() if "image" in k.lower()]
-                        if img_keys:
-                            img = frame_data[img_keys[0]]
-                        else:
-                            continue
-
-                    if isinstance(img, np.ndarray):
-                        img = torch.from_numpy(img)
-
-                    if len(img.shape) == 3 and img.shape[-1] in [1, 3, 4]:
-                        img = img.permute(2, 0, 1)
-
-                    # Resize to expected input size (224x224 for DINO v3)
-                    if img.shape[-2:] != (224, 224):
-                        import torch.nn.functional as F
-
-                        img = F.interpolate(
-                            img.unsqueeze(0), size=(224, 224), mode="bilinear", align_corners=False
-                        ).squeeze(0)
-
-                    # Normalize to [0, 1] if needed
-                    if img.dtype == torch.uint8:
-                        img = img.float() / 255.0
-
-                    frames.append(img)
-
-                if not frames:
-                    continue
-
-                frames_tensor = torch.stack(frames)
-
-                # Predict with correct language
-                correct_lang = original_languages[i]
-                correct_ep_rewards = self.predict_episode_rewards(frames_tensor, correct_lang)
-
-                # Predict with incorrect language
-                incorrect_lang = mismatched_languages[i]
-                incorrect_ep_rewards = self.predict_episode_rewards(frames_tensor, incorrect_lang)
-
-                correct_rewards.append(correct_ep_rewards)
-                incorrect_rewards.append(incorrect_ep_rewards)
-
-            except Exception as e:
-                warnings.warn(f"Failed to process episode {ep_idx} for detection: {e}")
-                continue
-
-        # Compute detection metrics
-        detection_results = compute_success_failure_detection(correct_rewards, incorrect_rewards)
-
-        print("Success/Failure Detection Results:")
-        print(f"  Detection accuracy: {detection_results['detection_accuracy']:.4f}")
-        print(f"  Mean correct final reward: {detection_results['mean_correct_final']:.4f}")
-        print(f"  Mean incorrect final reward: {detection_results['mean_incorrect_final']:.4f}")
-        print(f"  Separation score: {detection_results['separation_score']:.4f}")
-        print(f"  Episode pairs evaluated: {detection_results['num_pairs']}")
-
-        return detection_results
-
-    def evaluate_rewind_progress(
-        self, dataset, num_episodes: int = 100
-    ) -> dict[str, Any]:
-        """
-        Evaluate ReWiND-specific progress properties.
-        
-        Checks:
-        1. Progress values are in [0, 1] range
-        2. Progress increases monotonically (or mostly)
-        3. First frames have low progress, last frames have high progress
-        """
-        episodes = np.random.choice(len(dataset.meta.episodes), min(num_episodes, len(dataset.meta.episodes)), replace=False)
-        
-        results = {
-            "progress_range_violations": 0,
-            "monotonicity_scores": [],
-            "start_progress_values": [],
-            "end_progress_values": [],
-            "episodes_evaluated": 0
-        }
-        
-        for ep_idx in episodes:
-            try:
-                # Get episode data
-                ep_start = dataset.episode_data_index["from"][ep_idx].item()
-                ep_end = dataset.episode_data_index["to"][ep_idx].item()
-                
-                # Sample some frames from episode
-                sample_indices = np.linspace(ep_start, ep_end-1, min(20, ep_end-ep_start), dtype=int)
-                
-                frames = []
-                for idx in sample_indices:
-                    item = dataset[idx]
-                    if OBS_IMAGES in item:
-                        frames.append(item[OBS_IMAGES])
-                    elif OBS_IMAGE in item:
-                        frames.append(item[OBS_IMAGE])
-                    else:
-                        continue
-                
-                if len(frames) < 2:
-                    continue
-                    
-                frames = torch.stack(frames)
-                language = dataset[ep_start].get("task", "")
-                
-                # Predict rewards/progress
-                progress = self.predict_episode_rewards(frames, language)
-                
-                # Check range violations
-                range_violations = np.sum((progress < 0) | (progress > 1))
-                results["progress_range_violations"] += range_violations
-                
-                # Check monotonicity (should generally increase)
-                if len(progress) > 1:
-                    diffs = np.diff(progress)
-                    monotonicity = np.mean(diffs >= 0)  # Fraction of non-decreasing steps
-                    results["monotonicity_scores"].append(monotonicity)
-                
-                # Record start/end values
-                results["start_progress_values"].append(progress[0])
-                results["end_progress_values"].append(progress[-1])
-                results["episodes_evaluated"] += 1
-                
-            except Exception as e:
-                print(f"Error evaluating episode {ep_idx}: {e}")
-                continue
-        
-        # Summarize results
-        if results["episodes_evaluated"] > 0:
-            results["mean_monotonicity"] = np.mean(results["monotonicity_scores"])
-            results["mean_start_progress"] = np.mean(results["start_progress_values"])
-            results["mean_end_progress"] = np.mean(results["end_progress_values"])
-            results["progress_increase"] = results["mean_end_progress"] - results["mean_start_progress"]
-        
-        return results
-
-    def comprehensive_evaluation(
-        self,
-        dataset,
-        num_episodes: int = 100,
-        use_interquartile_mean: bool = True,
-        mismatch_templates: list[str] | None = None,
-    ) -> dict[str, Any]:
-        """
-        Run comprehensive evaluation including both VOC-S and detection metrics.
-
-        Returns:
-            Combined evaluation results
-        """
-        print("=" * 60)
-        print("COMPREHENSIVE RLEARN EVALUATION")
-        print("=" * 60)
-
-        # VOC-S evaluation
-        voc_results = self.evaluate_voc_s(
-            dataset, num_episodes=num_episodes, use_interquartile_mean=use_interquartile_mean
-        )
-
-        print("\n" + "=" * 40)
-
-        # Success/failure detection
-        detection_results = self.evaluate_success_failure_detection(
-            dataset, num_episodes=num_episodes, mismatch_templates=mismatch_templates
-        )
-
-        # Combined results
-        results = {
-            "voc_s": voc_results,
-            "detection": detection_results,
-            "overall_score": (
-                voc_results["voc_s_iqm"] * 0.6 + detection_results["detection_accuracy"] * 0.4
-            ),  # Weighted combination
-        }
-
-        print("\n" + "=" * 60)
-        print(f"OVERALL EVALUATION SCORE: {results['overall_score']:.4f}")
-        print("=" * 60)
-
-        return results

src/lerobot/policies/rlearn/modeling_rlearn.py

@@ -508,7 +508,7 @@ class RLearNPolicy(PreTrainedPolicy):
                 all_progress = []
                 
                 # DEBUG: Log indexing details for first sample occasionally
-                debug_indexing = torch.rand(1).item() < 0.05  # 5% chance
+                debug_indexing = torch.rand(1).item() < 0.10  # 10% chance - increased for debugging
                 if debug_indexing:
                     print(f"\n=== INDEXING DEBUG ===")
                     print(f"Delta indices: {delta_indices}")

src/lerobot/policies/rlearn/rlearn_plan.md

@@ -81,6 +81,8 @@ _ Open X-Embodiment (OXE)
   - Test rewind (evaluate) [x]
 - benchmark siglip 2 vs this implementation forward pass, debug speed [x]
 - use siglip 2 [x]
+- Fix evaluation bug !!! []
+- Fix sample episode padding bug !!! []
 - Overfit on one episode []
 - Cleanup code? [] + enable language loss
 - Convert python -m lerobot.datasets.v21.convert_dataset_v20_to_v21 --repo-id=IPEC-COMMUNITY/bc_z_lerobot and train on 1 percent
@@ -92,6 +94,7 @@ _ Open X-Embodiment (OXE)
 - Add other datasets from OXE metioned in rewind []
 - Extend evaluation []
 - Ablation for size vision encoder, language encoder, temporal head []
+- Ablation one mlp head per frame or single mlp head []
 - Add other datasets metnioned here []
 - How can we improve spatial aware learning? solve issue of Contrastive learning and position []
 

src/lerobot/scripts/train.py

@@ -213,26 +213,7 @@ def train(cfg: TrainPipelineConfig):
         episode_data_index=episode_data_index,
     )
     
-    # Setup RLearN evaluation visualizations if enabled
+
-    eval_visualizer = None
-    eval_holdout_episodes = None
-    if (getattr(cfg.policy, "type", None) == "rlearn" and 
-        getattr(cfg.policy, "enable_eval_visualizations", False)):
-        
-        try:
-            from lerobot.policies.rlearn.eval_visualizer import RLearNEvalVisualizer, select_evaluation_episodes
-            
-            logging.info("Setting up RLearN evaluation visualizations")
-            eval_visualizer = RLearNEvalVisualizer(policy, dataset, device=str(device))
-            eval_holdout_episodes = select_evaluation_episodes(
-                dataset, 
-                num_episodes=getattr(cfg.policy, "eval_holdout_episodes", 9),
-                seed=getattr(cfg.policy, "eval_visualization_seed", 42)
-            )
-            logging.info(f"Selected {len(eval_holdout_episodes)} holdout episodes for evaluation: {eval_holdout_episodes}")
-        except ImportError as e:
-            logging.warning(f"Could not setup RLearN evaluation visualizations: {e}")
-            eval_visualizer = None
     
     preprocessor, postprocessor = make_processor(
         policy_cfg=cfg.policy, pretrained_path=cfg.policy.pretrained_path, dataset_stats=dataset.meta.stats
@@ -408,8 +389,7 @@ def train(cfg: TrainPipelineConfig):
         is_log_step = cfg.log_freq > 0 and step % cfg.log_freq == 0
         is_saving_step = step % cfg.save_freq == 0 or step == cfg.steps
         is_eval_step = cfg.eval_freq > 0 and step % cfg.eval_freq == 0
-        is_eval_viz_step = (eval_visualizer is not None and 
+
-                           step % getattr(cfg.policy, "eval_visualization_freq", 1000) == 0)
 
         if is_log_step:
             logging.info(train_tracker)
@@ -461,86 +441,7 @@ def train(cfg: TrainPipelineConfig):
                 wandb_logger.log_dict(wandb_log_dict, step, mode="eval")
                 wandb_logger.log_video(eval_info["video_paths"][0], step, mode="eval")
 
-        # RLearN evaluation visualizations
+
-        if is_eval_viz_step:
-            logging.info(f"Creating RLearN evaluation visualizations at step {step}")
-            try:
-                with torch.no_grad():
-                    policy.eval()
-                    
-                    # Create evaluation visualizations directory
-                    eval_viz_dir = cfg.output_dir / "eval_visualizations"
-                    eval_viz_dir.mkdir(parents=True, exist_ok=True)
-                    
-                    # Create reward prediction visualization (3x3 grid)
-                    reward_viz_path = eval_viz_dir / f"reward_predictions_step_{step:06d}.png"
-                    reward_metrics = eval_visualizer.create_episode_grid_visualization(
-                        episode_indices=eval_holdout_episodes,
-                        save_path=reward_viz_path,
-                        step=step,
-                        max_frames=getattr(cfg.policy, "eval_max_frames", 128)
-                    )
-                    
-                    # Log metrics
-                    eval_viz_metrics = {
-                        "eval_viz/mean_voc_s": reward_metrics["mean_voc_s"],
-                        "eval_viz/std_voc_s": reward_metrics["std_voc_s"],
-                        "eval_viz/valid_episodes": reward_metrics["num_valid_episodes"],
-                        "eval_viz/total_episodes": reward_metrics["total_episodes"],
-                        "eval_viz/mean_episode_length": reward_metrics["mean_episode_length"],
-                    }
-                    
-                    logging.info(f"RLearN Evaluation Results at Step {step}:")
-                    logging.info(f"  Mean VOC-S: {reward_metrics['mean_voc_s']:.4f} (±{reward_metrics['std_voc_s']:.4f})")
-                    logging.info(f"  Valid Episodes: {reward_metrics['num_valid_episodes']}/{reward_metrics['total_episodes']}")
-                    logging.info(f"  Mean Episode Length: {reward_metrics['mean_episode_length']:.1f}")
-                    logging.info(f"  Visualizations saved to: {eval_viz_dir}")
-                    
-                    if wandb_logger:
-                        wandb_logger.log_dict(eval_viz_metrics, step, mode="eval_viz")
-                        
-                        # Log the visualization image both as regular image and as artifact
-                        try:
-                            import wandb
-                            
-                            # Log as regular image for immediate viewing in wandb UI
-                            wandb_logger.wandb_run.log({
-                                f"eval_viz/reward_predictions_step_{step}": wandb.Image(str(reward_viz_path)),
-                            }, step=step)
-                            
-                            # Create and upload artifact with reward prediction visualization
-                            artifact_name = f"rlearn_reward_predictions_step_{step:06d}"
-                            artifact = wandb.Artifact(
-                                name=artifact_name,
-                                type="reward_prediction_visualization",
-                                description=f"RLearN reward prediction visualization at training step {step}",
-                                metadata={
-                                    "step": step,
-                                    "mean_voc_s": reward_metrics["mean_voc_s"],
-                                    "std_voc_s": reward_metrics["std_voc_s"],
-                                    "valid_episodes": reward_metrics["num_valid_episodes"],
-                                    "total_episodes": reward_metrics["total_episodes"],
-                                    "mean_episode_length": reward_metrics["mean_episode_length"],
-                                    "holdout_episodes": eval_holdout_episodes,
-                                }
-                            )
-                            
-                            # Add reward prediction visualization to the artifact
-                            artifact.add_file(str(reward_viz_path), name="reward_predictions.png")
-                            
-                            # Upload the artifact
-                            wandb_logger.wandb_run.log_artifact(artifact)
-                            
-                            logging.info(f"Uploaded wandb artifact: {artifact_name}")
-                            
-                        except Exception as e:
-                            logging.warning(f"Could not log visualization image to wandb: {e}")
-                    
-                    policy.train()  # Return to training mode
-                    
-            except Exception as e:
-                logging.error(f"Error during RLearN evaluation visualization: {e}")
-                # Continue training even if evaluation fails
 
     if eval_env:
         eval_env.close()