scripts/plot_hamiltonian_slice_flow.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import copy
import csv
import json
import os
import sys
from dataclasses import asdict, dataclass
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE")
MPLCONFIGDIR = REPO_ROOT / ".cache" / "matplotlib"
MPLCONFIGDIR.mkdir(parents=True, exist_ok=True)
os.environ.setdefault("MPLCONFIGDIR", str(MPLCONFIGDIR))
os.environ.setdefault("MUJOCO_GL", "egl")

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import torch

from common import ensure_repo_on_path

ensure_repo_on_path()

from hamworld.runtime import make_env
from hamworld.world_model import CanonicalDynamicsWorldModel, infer_checkpoint_step


@dataclass(frozen=True)
class PairScore:
    q_index: int
    p_index: int
    kind: str
    score: float
    coverage: float
    swirl: float
    speed: float
    occupancy: float
    q_span: float
    p_span: float


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Render 2D Hamiltonian slice flow maps from a HaM-World checkpoint.")
    parser.add_argument("--checkpoint", required=True, help="Path to checkpoint_*.pt.")
    parser.add_argument("--trace-path", required=True, help="Path to dynamics trace .npz file.")
    parser.add_argument("--output-dir", required=True, help="Directory for figures and pair summaries.")
    parser.add_argument("--device", default="cpu", help="Torch device, e.g. cpu or cuda:0.")
    parser.add_argument("--grid-size", type=int, default=81, help="Grid size for slice evaluation.")
    parser.add_argument("--same-k", type=int, default=4, help="How many same-index (q_i, p_i) pairs to render.")
    parser.add_argument("--cross-k", type=int, default=2, help="How many cross-index (q_i, p_j) pairs to render.")
    parser.add_argument("--trace-episodes", type=int, default=3, help="How many trace episodes to overlay.")
    parser.add_argument("--trace-quantile-lo", type=float, default=2.0, help="Low percentile for slice bounds.")
    parser.add_argument("--trace-quantile-hi", type=float, default=98.0, help="High percentile for slice bounds.")
    return parser.parse_args()


def _load_checkpoint_world_model(checkpoint_path: Path, device: torch.device) -> tuple[CanonicalDynamicsWorldModel, dict]:
    payload = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
    config = copy.deepcopy(payload["config"])

    env, spec = make_env(config["task"], int(config["experiment"]["seed"]) + 1000)
    close_fn = getattr(env, "close", None)
    if callable(close_fn):
        close_fn()

    world_model = CanonicalDynamicsWorldModel(
        config=config,
        obs_dim=int(spec.observation_shape[0]),
        action_dim=int(spec.action_shape[0]),
    ).to(device)

    agent_state = payload.get("agent_state")
    if isinstance(agent_state, dict) and "model" in agent_state:
        state_dict = agent_state["model"]
    else:
        state_dict = payload.get("model")
    if state_dict is None:
        raise ValueError(f"Checkpoint does not contain a world model state: {checkpoint_path}")

    world_model.load_state_dict(state_dict)
    world_model.set_schedule_step(infer_checkpoint_step(payload, checkpoint_path))
    world_model.eval()
    return world_model, config


def _load_trace(trace_path: Path) -> dict[str, np.ndarray]:
    trace = np.load(trace_path, allow_pickle=True)
    return {key: trace[key] for key in trace.files}


def _flatten_valid_vectors(values: np.ndarray, valid_mask: np.ndarray) -> np.ndarray:
    mask = np.asarray(valid_mask, dtype=bool)
    return np.asarray(values[mask], dtype=np.float32)


def _robust_span(values: np.ndarray, lo: float = 5.0, hi: float = 95.0) -> float:
    low, high = np.percentile(values, [lo, hi])
    return float(high - low)


def _histogram_entropy(x: np.ndarray, y: np.ndarray, bins: int = 20) -> float:
    hist, _, _ = np.histogram2d(x, y, bins=bins)
    probs = hist.ravel()
    total = probs.sum()
    if total <= 0.0:
        return 0.0
    probs = probs[probs > 0.0] / total
    entropy = -(probs * np.log(probs)).sum()
    max_entropy = np.log(float(bins * bins))
    return float(entropy / max(1e-6, max_entropy))


def _normalize(values: list[float]) -> list[float]:
    maximum = max(values) if values else 1.0
    if maximum <= 1e-6:
        return [0.0 for _ in values]
    return [float(value / maximum) for value in values]


def score_pairs(trace: dict[str, np.ndarray]) -> list[PairScore]:
    valid_mask = np.asarray(trace["valid_mask"], dtype=bool)
    q = _flatten_valid_vectors(trace["q"], valid_mask)
    p = _flatten_valid_vectors(trace["p"], valid_mask)
    dH_dq = _flatten_valid_vectors(trace["dH_dq"], valid_mask)
    dH_dp = _flatten_valid_vectors(trace["dH_dp"], valid_mask)

    q_dim = int(np.asarray(trace["q_dim"]).item())
    p_dim = int(np.asarray(trace["p_dim"]).item())

    raw_rows: list[dict[str, float | int | str]] = []
    for q_index in range(q_dim):
        for p_index in range(p_dim):
            qv = q[:, q_index]
            pv = p[:, p_index]
            dq = dH_dp[:, q_index]
            dp = -dH_dq[:, p_index]
            q_center = float(np.median(qv))
            p_center = float(np.median(pv))

            radial_norm = np.hypot(qv - q_center, pv - p_center)
            flow_norm = np.hypot(dq, dp)
            tangential = np.abs((qv - q_center) * dp - (pv - p_center) * dq)
            swirl = float(np.mean(tangential / np.maximum(1e-6, radial_norm * flow_norm)))

            q_span = _robust_span(qv)
            p_span = _robust_span(pv)
            coverage = float(q_span * p_span)
            speed = float(np.mean(flow_norm))
            occupancy = _histogram_entropy(qv, pv)
            raw_rows.append(
                {
                    "q_index": q_index,
                    "p_index": p_index,
                    "kind": "same" if q_index == p_index else "cross",
                    "coverage": coverage,
                    "swirl": swirl,
                    "speed": speed,
                    "occupancy": occupancy,
                    "q_span": q_span,
                    "p_span": p_span,
                }
            )

    coverage_norm = _normalize([float(row["coverage"]) for row in raw_rows])
    swirl_norm = _normalize([float(row["swirl"]) for row in raw_rows])
    speed_norm = _normalize([float(row["speed"]) for row in raw_rows])
    occupancy_norm = _normalize([float(row["occupancy"]) for row in raw_rows])

    scored: list[PairScore] = []
    for index, row in enumerate(raw_rows):
        score = (
            0.35 * coverage_norm[index]
            + 0.30 * swirl_norm[index]
            + 0.20 * speed_norm[index]
            + 0.15 * occupancy_norm[index]
        )
        scored.append(
            PairScore(
                q_index=int(row["q_index"]),
                p_index=int(row["p_index"]),
                kind=str(row["kind"]),
                score=float(score),
                coverage=float(row["coverage"]),
                swirl=float(row["swirl"]),
                speed=float(row["speed"]),
                occupancy=float(row["occupancy"]),
                q_span=float(row["q_span"]),
                p_span=float(row["p_span"]),
            )
        )

    return sorted(scored, key=lambda item: item.score, reverse=True)


def select_pairs(scores: list[PairScore], same_k: int, cross_k: int) -> list[PairScore]:
    same_pairs = [row for row in scores if row.kind == "same"][: max(0, int(same_k))]
    cross_pairs = [row for row in scores if row.kind == "cross"][: max(0, int(cross_k))]
    return same_pairs + cross_pairs


def _slice_bounds(values: np.ndarray, quantile_lo: float, quantile_hi: float) -> tuple[float, float]:
    low, high = np.percentile(values, [quantile_lo, quantile_hi])
    if not np.isfinite(low) or not np.isfinite(high) or abs(high - low) < 1e-5:
        center = float(np.median(values))
        return center - 1.0, center + 1.0
    margin = 0.18 * float(high - low)
    return float(low - margin), float(high + margin)


def evaluate_slice(

    world_model: CanonicalDynamicsWorldModel,

    q_reference: np.ndarray,

    p_reference: np.ndarray,

    pair: PairScore,

    q_values: np.ndarray,

    p_values: np.ndarray,

    grid_size: int,

    quantile_lo: float,

    quantile_hi: float,

    device: torch.device,

) -> dict[str, np.ndarray]:
    x_min, x_max = _slice_bounds(q_values, quantile_lo, quantile_hi)
    y_min, y_max = _slice_bounds(p_values, quantile_lo, quantile_hi)
    q_axis = np.linspace(x_min, x_max, grid_size, dtype=np.float64)
    p_axis = np.linspace(y_min, y_max, grid_size, dtype=np.float64)
    q_grid, p_grid = np.meshgrid(q_axis, p_axis, indexing="xy")

    q_eval = np.repeat(q_reference[None, :], grid_size * grid_size, axis=0)
    p_eval = np.repeat(p_reference[None, :], grid_size * grid_size, axis=0)
    q_eval[:, pair.q_index] = q_grid.reshape(-1)
    p_eval[:, pair.p_index] = p_grid.reshape(-1)

    q_tensor = torch.as_tensor(q_eval, dtype=torch.float32, device=device).requires_grad_(True)
    p_tensor = torch.as_tensor(p_eval, dtype=torch.float32, device=device).requires_grad_(True)
    energy = world_model.energy_head(q_tensor, p_tensor)
    dH_dq, dH_dp = torch.autograd.grad(energy.sum(), (q_tensor, p_tensor), create_graph=False, retain_graph=False)

    H = energy.detach().cpu().numpy().reshape(grid_size, grid_size)
    U = dH_dp.detach().cpu().numpy()[:, pair.q_index].reshape(grid_size, grid_size)
    V = (-dH_dq.detach().cpu().numpy()[:, pair.p_index]).reshape(grid_size, grid_size)
    speed = np.hypot(U, V)
    return {
        "q_axis": q_axis,
        "p_axis": p_axis,
        "H": H,
        "U": U,
        "V": V,
        "speed": speed,
        "q_grid": q_grid,
        "p_grid": p_grid,
    }


def _overlay_trace(axis: plt.Axes, trace: dict[str, np.ndarray], pair: PairScore, max_episodes: int) -> None:
    valid_mask = np.asarray(trace["valid_mask"], dtype=bool)
    q = np.asarray(trace["q"], dtype=np.float32)
    p = np.asarray(trace["p"], dtype=np.float32)
    episode_count = min(int(max_episodes), int(q.shape[0]))
    colors = ["#7dd3fc", "#38bdf8", "#0ea5e9"]
    for episode_index in range(episode_count):
        mask = valid_mask[episode_index]
        if not np.any(mask):
            continue
        x = q[episode_index, mask, pair.q_index]
        y = p[episode_index, mask, pair.p_index]
        axis.plot(x, y, color=colors[episode_index % len(colors)], linewidth=1.1, alpha=0.9)
        axis.scatter(x[0], y[0], color=colors[episode_index % len(colors)], s=14, marker="o", alpha=0.95)


def draw_pair_figure(

    output_path: Path,

    slice_eval: dict[str, np.ndarray],

    trace: dict[str, np.ndarray],

    pair: PairScore,

    q_reference: np.ndarray,

    p_reference: np.ndarray,

    checkpoint_name: str,

) -> None:
    q_axis = slice_eval["q_axis"]
    p_axis = slice_eval["p_axis"]
    H = slice_eval["H"]
    U = slice_eval["U"]
    V = slice_eval["V"]
    speed = slice_eval["speed"]

    fig, axis = plt.subplots(figsize=(6.4, 5.2), dpi=220)
    heat = axis.contourf(q_axis, p_axis, H, levels=24, cmap="magma")
    axis.contour(q_axis, p_axis, H, levels=10, colors="white", linewidths=0.45, alpha=0.55)
    axis.streamplot(
        q_axis,
        p_axis,
        U,
        V,
        density=1.1,
        color=speed,
        cmap="viridis",
        linewidth=0.9,
        arrowsize=0.9,
    )
    _overlay_trace(axis, trace, pair, max_episodes=3)
    axis.scatter(
        [q_reference[pair.q_index]],
        [p_reference[pair.p_index]],
        s=52,
        marker="x",
        linewidths=1.5,
        color="white",
        zorder=5,
    )
    axis.set_xlabel(f"q[{pair.q_index}]")
    axis.set_ylabel(f"p[{pair.p_index}]")
    axis.set_title(
        f"{pair.kind} pair q[{pair.q_index}] / p[{pair.p_index}]  score={pair.score:.3f}\n{checkpoint_name}",
        loc="left",
    )
    cbar = fig.colorbar(heat, ax=axis, fraction=0.046, pad=0.04)
    cbar.set_label("H(q, p)")
    axis.spines["top"].set_visible(False)
    axis.spines["right"].set_visible(False)
    fig.tight_layout()
    fig.savefig(output_path, bbox_inches="tight")
    plt.close(fig)


def draw_overview(

    output_path: Path,

    slices: list[tuple[PairScore, dict[str, np.ndarray]]],

    trace: dict[str, np.ndarray],

    q_reference: np.ndarray,

    p_reference: np.ndarray,

    checkpoint_name: str,

) -> None:
    columns = 3
    rows = max(1, int(np.ceil(len(slices) / columns)))
    fig, axes = plt.subplots(rows, columns, figsize=(5.3 * columns, 4.5 * rows), dpi=220)
    axes_array = np.atleast_1d(axes).reshape(rows, columns)

    for axis, (pair, slice_eval) in zip(axes_array.flatten(), slices):
        heat = axis.contourf(slice_eval["q_axis"], slice_eval["p_axis"], slice_eval["H"], levels=20, cmap="magma")
        axis.contour(slice_eval["q_axis"], slice_eval["p_axis"], slice_eval["H"], levels=8, colors="white", linewidths=0.35, alpha=0.5)
        axis.streamplot(
            slice_eval["q_axis"],
            slice_eval["p_axis"],
            slice_eval["U"],
            slice_eval["V"],
            density=1.0,
            color=slice_eval["speed"],
            cmap="viridis",
            linewidth=0.8,
            arrowsize=0.8,
        )
        _overlay_trace(axis, trace, pair, max_episodes=2)
        axis.scatter(
            [q_reference[pair.q_index]],
            [p_reference[pair.p_index]],
            s=32,
            marker="x",
            linewidths=1.2,
            color="white",
            zorder=5,
        )
        axis.set_xlabel(f"q[{pair.q_index}]")
        axis.set_ylabel(f"p[{pair.p_index}]")
        axis.set_title(f"{pair.kind}  q[{pair.q_index}] / p[{pair.p_index}]  s={pair.score:.3f}", loc="left", fontsize=10)
        axis.spines["top"].set_visible(False)
        axis.spines["right"].set_visible(False)

    for axis in axes_array.flatten()[len(slices) :]:
        axis.axis("off")

    fig.suptitle(f"Hamiltonian slice flow candidates\n{checkpoint_name}", x=0.055, y=0.995, ha="left", fontsize=14)
    fig.tight_layout(rect=[0, 0, 1, 0.97])
    fig.savefig(output_path, bbox_inches="tight")
    plt.close(fig)


def write_scores_csv(output_path: Path, scores: list[PairScore]) -> None:
    fieldnames = list(asdict(scores[0]).keys()) if scores else []
    with output_path.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=fieldnames)
        writer.writeheader()
        for row in scores:
            writer.writerow(asdict(row))


def main() -> int:
    args = parse_args()
    checkpoint_path = Path(args.checkpoint).expanduser().resolve()
    trace_path = Path(args.trace_path).expanduser().resolve()
    output_dir = Path(args.output_dir).expanduser().resolve()
    output_dir.mkdir(parents=True, exist_ok=True)

    device = torch.device(args.device)
    world_model, config = _load_checkpoint_world_model(checkpoint_path, device)
    trace = _load_trace(trace_path)

    valid_mask = np.asarray(trace["valid_mask"], dtype=bool)
    q_flat = _flatten_valid_vectors(trace["q"], valid_mask)
    p_flat = _flatten_valid_vectors(trace["p"], valid_mask)
    q_reference = np.median(q_flat, axis=0).astype(np.float32)
    p_reference = np.median(p_flat, axis=0).astype(np.float32)

    all_scores = score_pairs(trace)
    selected_pairs = select_pairs(all_scores, same_k=args.same_k, cross_k=args.cross_k)
    if not selected_pairs:
        raise RuntimeError("No q/p pairs were selected for rendering.")

    render_dir = output_dir / "pairs"
    render_dir.mkdir(parents=True, exist_ok=True)

    slice_bundle: list[tuple[PairScore, dict[str, np.ndarray]]] = []
    for pair in selected_pairs:
        q_values = q_flat[:, pair.q_index]
        p_values = p_flat[:, pair.p_index]
        slice_eval = evaluate_slice(
            world_model=world_model,
            q_reference=q_reference,
            p_reference=p_reference,
            pair=pair,
            q_values=q_values,
            p_values=p_values,
            grid_size=int(args.grid_size),
            quantile_lo=float(args.trace_quantile_lo),
            quantile_hi=float(args.trace_quantile_hi),
            device=device,
        )
        slice_bundle.append((pair, slice_eval))
        filename = f"hamiltonian_slice_{pair.kind}_q{pair.q_index}_p{pair.p_index}.png"
        draw_pair_figure(
            output_path=render_dir / filename,
            slice_eval=slice_eval,
            trace=trace,
            pair=pair,
            q_reference=q_reference,
            p_reference=p_reference,
            checkpoint_name=checkpoint_path.parent.parent.parent.name,
        )

    write_scores_csv(output_dir / "pair_scores.csv", all_scores)
    with (output_dir / "selected_pairs.json").open("w", encoding="utf-8") as handle:
        json.dump([asdict(pair) for pair in selected_pairs], handle, indent=2)

    draw_overview(
        output_path=output_dir / "hamiltonian_slice_overview.png",
        slices=slice_bundle,
        trace=trace,
        q_reference=q_reference,
        p_reference=p_reference,
        checkpoint_name=checkpoint_path.parent.parent.parent.name,
    )

    summary = {
        "checkpoint": str(checkpoint_path),
        "trace_path": str(trace_path),
        "task": str(np.asarray(trace["task"]).item()),
        "seed": int(np.asarray(trace["seed"]).item()),
        "mode": str(np.asarray(trace["mode"]).item()),
        "q_reference": q_reference.tolist(),
        "p_reference": p_reference.tolist(),
        "selected_pairs": [asdict(pair) for pair in selected_pairs],
    }
    with (output_dir / "summary.json").open("w", encoding="utf-8") as handle:
        json.dump(summary, handle, indent=2)

    print(output_dir / "hamiltonian_slice_overview.png")
    for pair in selected_pairs:
        print(render_dir / f"hamiltonian_slice_{pair.kind}_q{pair.q_index}_p{pair.p_index}.png")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())