mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-07 01:51:47 +00:00
simple eval
This commit is contained in:
Pepijn
File diff suppressed because one or more lines are too long
@@ -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())
|
||||
@@ -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)
|
||||
@@ -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
|
||||
@@ -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}")
|
||||
|
||||
@@ -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 []
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user