| """Online evaluation of state GCBC on cube-single (5 tasks) with Wilson CIs.""" |
| import argparse, os, json, time, statistics |
| import numpy as np, torch, gymnasium as gym, ogbench |
| from common import wilson_ci |
| from train_gcbc import MLPPolicy |
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--ckpt", required=True) |
| ap.add_argument("--env", default="cube-single-v0") |
| ap.add_argument("--rollouts", type=int, default=50) |
| ap.add_argument("--tasks", type=int, default=5) |
| ap.add_argument("--base_seed", type=int, default=20260825) |
| ap.add_argument("--out", default=None) |
| a = ap.parse_args() |
| dev = "cuda" if torch.cuda.is_available() else "cpu" |
| ck = torch.load(a.ckpt, map_location=dev, weights_only=False) |
| mean = torch.tensor(ck["mean"], device=dev); std = torch.tensor(ck["std"], device=dev) |
| model = MLPPolicy(ck["obs_dim"] * 2, ck["act_dim"], ck["hidden"], ck["layers"]).to(dev) |
| model.load_state_dict(ck["model"]); model.eval() |
| env = gym.make(a.env) |
| def act(o, g): |
| o = (torch.tensor(np.asarray(o), dtype=torch.float32, device=dev) - mean) / std |
| g = (torch.tensor(np.asarray(g), dtype=torch.float32, device=dev) - mean) / std |
| with torch.no_grad(): |
| return model(o[None], g[None])[0].cpu().numpy() |
| per_task = {}; all_k = 0; all_n = 0; steps_success = [] |
| t0 = time.time() |
| for task_id in range(1, a.tasks + 1): |
| k = 0 |
| for r in range(a.rollouts): |
| ob, info = env.reset(seed=a.base_seed + 1000 * task_id + r, |
| options=dict(task_id=task_id, render_goal=False)) |
| goal = info["goal"]; succ = 0; sstep = None |
| for t in range(1000): |
| ob, rew, term, trunc, info = env.step(np.clip(act(ob, goal), -1, 1)) |
| if info.get("success"): |
| succ = 1 |
| if sstep is None: sstep = t + 1 |
| if term or trunc: break |
| k += succ; all_n += 1 |
| if succ: steps_success.append(sstep) |
| p, lo, hi = wilson_ci(k, a.rollouts) |
| per_task[task_id] = dict(k=k, n=a.rollouts, rate=p, ci=[lo, hi]) |
| all_k += k |
| print(f" task{task_id}: {k}/{a.rollouts} = {p:.3f} 95% CI [{lo:.3f},{hi:.3f}]", flush=True) |
| p, lo, hi = wilson_ci(all_k, all_n) |
| med = statistics.median(steps_success) if steps_success else None |
| res = dict(env=a.env, rollouts_per_task=a.rollouts, tasks=a.tasks, |
| overall=dict(k=all_k, n=all_n, rate=p, ci=[lo, hi]), |
| per_task=per_task, median_steps_to_success=med, |
| timeout_rate=1 - all_k / all_n, eval_seconds=time.time() - t0, |
| ckpt=a.ckpt) |
| print("OVERALL: %d/%d = %.3f 95%% Wilson CI [%.3f, %.3f] | median steps-to-success=%s | %.0fs" |
| % (all_k, all_n, p, lo, hi, med, time.time() - t0)) |
| out = a.out or os.path.join(os.path.dirname(a.ckpt), "eval.json") |
| json.dump(res, open(out, "w"), indent=2) |
| print("saved ->", out) |
|
|
| if __name__ == "__main__": |
| main() |
|
|
