| """Online evaluation of the pixel GCBC policy in the visual OGBench env.""" |
| import argparse, os, json, time, statistics |
| import numpy as np, torch, gymnasium as gym, ogbench |
| from common import wilson_ci |
| from train_pix import PixelPolicy |
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--ckpt", required=True) |
| ap.add_argument("--env", default="visual-cube-single-v0") |
| ap.add_argument("--rollouts", type=int, default=100) |
| 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" |
| ck = torch.load(a.ckpt, map_location=dev, weights_only=False) |
| model = PixelPolicy(6, ck["act_dim"], ck["width"]).to(dev) |
| model.load_state_dict(ck["model"]); model.eval() |
| env = gym.make(a.env) |
| def act(ob, goal): |
| o = torch.from_numpy(np.asarray(ob)).to(dev).permute(2, 0, 1).float().div_(255.).sub_(0.5) |
| g = torch.from_numpy(np.asarray(goal)).to(dev).permute(2, 0, 1).float().div_(255.).sub_(0.5) |
| with torch.no_grad(): |
| return model(torch.cat([o, g], 0)[None])[0].cpu().numpy() |
| per_task, all_k, all_n, steps_s = {}, 0, 0, [] |
| 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=True)) |
| 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_s.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_s) if steps_s else None |
| res = dict(env=a.env, rollouts_per_task=a.rollouts, 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=%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() |
|
|
