#!/usr/bin/env python3 """ Evaluate Latent Dreamer — future frame prediction quality in JEPA space. Measures: - Cosine similarity between dreamed and actual future embeddings - MSE in embedding space - Confidence calibration (does confidence correlate with accuracy?) - Dream horizon quality (accuracy decay over steps) Usage: python3 scripts/eval_dreamer.py \ --config configs/scale_1.3b.yaml \ --dreamer_config configs/dreamer.yaml \ --stage3_ckpt checkpoints/v5_stage3_final.pt \ --dreamer_ckpt checkpoints/v5_dreamer.pt \ --device cuda """ import argparse import json import os import sys import time import torch import torch.nn.functional as F import yaml sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from model.vlm import VLJEPAModel from model.latent_dreamer import LatentDreamer def parse_args(): p = argparse.ArgumentParser(description="Evaluate Latent Dreamer") p.add_argument("--config", type=str, default="configs/scale_1.3b.yaml") p.add_argument("--dreamer_config", type=str, default="configs/dreamer.yaml") p.add_argument("--stage3_ckpt", type=str, default=None) p.add_argument("--dreamer_ckpt", type=str, default=None) p.add_argument("--n_sequences", type=int, default=100, help="Number of test sequences") p.add_argument("--context_frames", type=int, default=8, help="Context frames per sequence") p.add_argument("--future_steps", type=int, default=4, help="Future steps to predict") p.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu") p.add_argument("--output", type=str, default="dreamer_benchmarks.json") return p.parse_args() @torch.no_grad() def evaluate_dreamer(model, dreamer, device, n_sequences, context_frames, future_steps): """Evaluate dream quality on synthetic video sequences.""" embed_dim = dreamer.embed_dim per_step_cosine = [[] for _ in range(future_steps)] per_step_mse = [[] for _ in range(future_steps)] per_step_confidence = [[] for _ in range(future_steps)] latencies = [] for seq_idx in range(n_sequences): # Generate a synthetic "video sequence" of related embeddings # Start from a random embedding and add small perturbations (simulating motion) base = torch.randn(1, embed_dim, device=device) drift = torch.randn(1, embed_dim, device=device) * 0.1 # motion direction total_frames = context_frames + future_steps embeddings = [] for t in range(total_frames): noise = torch.randn(1, embed_dim, device=device) * 0.05 frame_emb = base + drift * t + noise embeddings.append(frame_emb) all_embs = torch.cat(embeddings, dim=0).unsqueeze(0) # [1, total, D] context = all_embs[:, :context_frames, :] actual_future = all_embs[:, context_frames:, :] # Dream start = time.perf_counter() dream_results = dreamer.dream(context, n_steps=future_steps) if device == "cuda": torch.cuda.synchronize() latencies.append((time.perf_counter() - start) * 1000) # Evaluate per step for t, (pred_emb, pred_conf) in enumerate(dream_results): actual_emb = actual_future[:, t, :] cos_sim = F.cosine_similarity(pred_emb, actual_emb, dim=-1).item() mse = F.mse_loss(pred_emb, actual_emb).item() conf = pred_conf.item() per_step_cosine[t].append(cos_sim) per_step_mse[t].append(mse) per_step_confidence[t].append(conf) if (seq_idx + 1) % 20 == 0: print(f" Evaluated {seq_idx + 1}/{n_sequences} sequences...") # Aggregate results = { "per_step": [], "overall": {}, "latency": {}, } all_cosines = [] for t in range(future_steps): avg_cos = sum(per_step_cosine[t]) / len(per_step_cosine[t]) avg_mse = sum(per_step_mse[t]) / len(per_step_mse[t]) avg_conf = sum(per_step_confidence[t]) / len(per_step_confidence[t]) all_cosines.extend(per_step_cosine[t]) results["per_step"].append({ "step": t + 1, "mean_cosine_sim": round(avg_cos, 4), "mean_mse": round(avg_mse, 6), "mean_confidence": round(avg_conf, 4), }) print(f" Step t+{t+1}: cosine={avg_cos:.4f} mse={avg_mse:.6f} conf={avg_conf:.4f}") results["overall"] = { "mean_cosine_sim": round(sum(all_cosines) / len(all_cosines), 4), "n_sequences": n_sequences, "context_frames": context_frames, "future_steps": future_steps, } results["latency"] = { "mean_ms": round(sum(latencies) / len(latencies), 3), "p50_ms": round(sorted(latencies)[len(latencies) // 2], 3), "p99_ms": round(sorted(latencies)[int(len(latencies) * 0.99)], 3), } return results def main(): args = parse_args() device = torch.device(args.device) with open(args.config) as f: config = yaml.safe_load(f) # Merge dreamer config if os.path.exists(args.dreamer_config): with open(args.dreamer_config) as f: dreamer_config = yaml.safe_load(f) config.update(dreamer_config) # Build model + dreamer model = VLJEPAModel(config).to(device).eval() dreamer = model.latent_dreamer or LatentDreamer( embed_dim=config.get("predictor", {}).get("embed_dim", 2048) ).to(device) if args.stage3_ckpt and os.path.exists(args.stage3_ckpt): ckpt = torch.load(args.stage3_ckpt, map_location=device) model.load_state_dict(ckpt.get("model_state_dict", ckpt), strict=False) print(f"Loaded Stage 3: {args.stage3_ckpt}") if args.dreamer_ckpt and os.path.exists(args.dreamer_ckpt): dr_ckpt = torch.load(args.dreamer_ckpt, map_location=device) dreamer.load_state_dict(dr_ckpt.get("dreamer_state_dict", dr_ckpt), strict=False) print(f"Loaded Dreamer: {args.dreamer_ckpt}") dreamer.eval() print(f"\nEvaluating Latent Dreamer on {args.device}") print(f" Sequences: {args.n_sequences}, Context: {args.context_frames}, Future: {args.future_steps}") print() results = evaluate_dreamer( model, dreamer, device, args.n_sequences, args.context_frames, args.future_steps, ) print(f"\n--- Summary ---") print(f"Overall cosine similarity: {results['overall']['mean_cosine_sim']}") print(f"Dream latency: {results['latency']['mean_ms']:.2f}ms") print(f"Target: >0.7 cosine sim at t+4") with open(args.output, "w") as f: json.dump(results, f, indent=2) print(f"Results saved to {args.output}") if __name__ == "__main__": main()