File size: 3,565 Bytes
adaf847 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | #!/usr/bin/env python3
"""Structural checks and optional live evaluation for the 14-parameter LunarLander policy.
python verify.py # structural checks only (torch)
python verify.py --episodes 100 # + live Gymnasium evaluation (60 seeds x N episodes)
Reference numbers (60 seeds x 100 episodes): mean 264.06, 60/60 seeds >= 200, 0.35% crashes.
"""
from __future__ import annotations
import argparse
import os
import numpy as np
from model import TinyLunarPolicy, OBS_DIM, N_ACTIONS, N_PARAMS
# The 60 evaluation seeds (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.numpy_weights()
print(f"parameters : {m.n_params} (expected {N_PARAMS})")
print(f"materialised W shape : {W.shape} (expected ({OBS_DIM}, {N_ACTIONS}))")
assert m.n_params == 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)
got = np.asarray([m.act(o) for o in O[:500]])
assert (ref[:500] == got).all(), "module disagrees with argmax(obs @ W)"
print(f"argmax(obs @ W) == act : {(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, last, term, trunc = False, 0.0, 0.0, False, False
while not done:
o, r, term, trunc, _ = env.step(m.act(o))
tot += r
last = r
done = term or trunc
rs.append(tot)
lands += int(term and not trunc and last == 100.0)
crashes += int(term and not trunc and last == -100.0)
means.append(float(np.mean(rs)))
means = np.asarray(means)
n = len(seeds) * n_ep
print(f"{len(seeds)} seeds x {n_ep} ep = {n} episodes")
print(f"mean={means.mean():.3f} std(seed means)={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 / n:.4f} crash rate={crashes / n:.4f}")
return means
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--weights", default=os.path.join(os.path.dirname(os.path.abspath(__file__)),
"model.safetensors"))
ap.add_argument("--episodes", type=int, default=0, help="episodes per seed (0 = skip)")
ap.add_argument("--seeds", type=int, default=60)
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()
|