scripts/eval_dreamer.py

#!/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()