verify.py

#!/usr/bin/env python3
"""Reproduce the reported numbers for TinyLunarPolicy.

  python verify.py                          # structural checks only (torch only)
  python verify.py --episodes 100           # + live Gymnasium evaluation
  python verify.py --seeds 60 --episodes 100  # full 60x100 protocol (mean ~247.2)

The reported reference numbers use the 60 evaluation seeds listed in README
(60 seeds x 100 episodes = 6000 episodes): mean 231.23, 99.67% landings, 0.25% crashes, 60/60 solved.
"""
from __future__ import annotations
import argparse
import os

import numpy as np

from model import TinyLunarPolicy, OBS_DIM, N_ACTIONS

# The 60 evaluation seeds used for the reported reference numbers
# (episode seed = seed * 1000 + episode_index).
EVAL_SEEDS = [
    40122, 42940, 51305, 55316, 60944, 61197, 93978, 96327, 103680, 157789,
    184578, 197504, 208678, 250957, 289579, 296198, 296953, 309957, 318961, 323741,
    349217, 350692, 373376, 382648, 383930, 385293, 393420, 400382, 414089, 450829,
    467030, 482346, 490224, 498162, 501685, 515927, 535096, 543079, 544042, 553850,
    555722, 571111, 615137, 634353, 642133, 660500, 665066, 684081, 732008, 743613,
    800434, 816622, 819515, 839392, 845498, 852069, 852799, 867535, 894908, 898516,
]



def structural_checks(path: str) -> TinyLunarPolicy:
    m = TinyLunarPolicy.from_safetensors(path)
    W = m.weights()
    print(f"params                : {m.n_params}   (expected 32)")
    print(f"weights shape (obs,act): {W.shape}   (expected (8, 4))")
    assert m.n_params == 32, m.n_params

    rng = np.random.default_rng(0)
    O = rng.standard_normal((2000, OBS_DIM)).astype(np.float32)
    ref = np.argmax(O @ W, axis=1)                       # obs @ W  -> argmax
    got = np.asarray([m.act(o) for o in O[:500]])         # module argmax
    assert (ref[:500] == got).all(), "module disagrees with obs @ W argmax"
    print(f"argmax(obs @ W) == module.act : {int((ref[:500] == got).sum())}/500  OK")
    return m


def live_eval(m: TinyLunarPolicy, seeds, n_ep: int, env_id: str = "LunarLander-v3"):
    import gymnasium as gym
    env = gym.make(env_id)
    means, lands, crashes = [], 0, 0
    for s in seeds:
        rs = []
        for ep in range(n_ep):
            o, _ = env.reset(seed=int(s) * 1000 + ep)
            done, tot = False, 0.0
            while not done:
                o, r, term, trunc, info = env.step(m.act(o))
                tot += r
                done = term or trunc
            rs.append(tot)
            lands += int(tot > 100)          # heuristic: landed (LunarLander landing ~ +100..300)
            crashes += int(tot < -100)
        means.append(float(np.mean(rs)))
    means = np.asarray(means)
    print(f"seeds={len(seeds)} x {n_ep} ep = {len(seeds) * n_ep} episodes")
    print(f"mean={means.mean():.3f}  std={means.std():.3f}  "
          f"min={means.min():.2f}  max={means.max():.2f}")
    print(f"seeds with mean>=200: {int((means >= 200).sum())}/{len(means)}")
    print(f"landing rate~{lands / (len(seeds) * n_ep):.4f}  crash rate~{crashes / (len(seeds) * n_ep):.4f}")
    return means


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--weights", default=os.path.join(os.path.dirname(__file__), "model.safetensors"))
    ap.add_argument("--episodes", type=int, default=0, help="episodes per seed (0 = skip live eval)")
    ap.add_argument("--seeds", type=int, default=60, help="number of seeds to evaluate")
    ap.add_argument("--env-id", default="LunarLander-v3")
    a = ap.parse_args()

    m = structural_checks(a.weights)
    if a.episodes > 0:
        live_eval(m, EVAL_SEEDS[: a.seeds], a.episodes, a.env_id)


if __name__ == "__main__":
    main()