()
| 46 | |
| 47 | |
| 48 | def main(): |
| 49 | p = argparse.ArgumentParser(description=__doc__) |
| 50 | p.add_argument("--run-dir", required=True, help="Go-Explore Phase 1 run dir") |
| 51 | p.add_argument("--out", required=True, help="output demo.pkl path") |
| 52 | p.add_argument("--ckpt-every", type=int, default=512, |
| 53 | help="snapshot an ALE restore point every N actions") |
| 54 | p.add_argument("--max-rewards", type=int, default=0, |
| 55 | help="truncate just after the Kth nonzero reward instead of the last " |
| 56 | "(0 = last, default). Use 1 for a first-key-only easy demo — a much " |
| 57 | "shorter horizon for robustification to bootstrap on.") |
| 58 | args = p.parse_args() |
| 59 | |
| 60 | ge = _load_ge() |
| 61 | run_dir = Path(args.run_dir) |
| 62 | |
| 63 | import torch |
| 64 | ckpt_path = run_dir / "ckpt" / "best.pt" |
| 65 | if not ckpt_path.exists(): |
| 66 | ckpt_path = run_dir / "ckpt" / "latest.pt" |
| 67 | ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False) |
| 68 | done = ckpt["archive"]["cells"].get(ge.DONE_KEY) |
| 69 | if not done: |
| 70 | sys.exit("[extract] no DONE cell — run has no end-of-episode trajectory") |
| 71 | archived_score = done["score"] |
| 72 | print(f"[extract] DONE score {archived_score:.0f}, traj_len {done['traj_len']:,}", flush=True) |
| 73 | |
| 74 | explog = ge.ExperienceLog(str(run_dir / "explog")) |
| 75 | explog.load_state(ckpt["explog"]) |
| 76 | actions = explog.reconstruct_actions(done["traj_last"]) |
| 77 | assert len(actions) == done["traj_len"], (len(actions), done["traj_len"]) |
| 78 | |
| 79 | # replay on the exact Phase-1 protocol, capturing rewards + periodic states |
| 80 | import ale_py |
| 81 | import gymnasium as gym |
| 82 | gym.register_envs(ale_py) |
| 83 | env = gym.make("ALE/MontezumaRevenge-v5", frameskip=4, |
| 84 | repeat_action_probability=0.0).unwrapped |
| 85 | env.reset(seed=0) |
| 86 | rewards, checkpoints, ckpt_nr = [], [], [] |
| 87 | score = 0.0 |
| 88 | last_reward_idx = -1 |
| 89 | for i, a in enumerate(actions): |
| 90 | if i % args.ckpt_every == 0: |
| 91 | checkpoints.append(pickle.dumps(env.ale.cloneState())) |
| 92 | ckpt_nr.append(i) |
| 93 | _, r, term, trunc, _ = env.step(int(a)) |
| 94 | rewards.append(float(r)) |
| 95 | score += float(r) |
| 96 | if r != 0: |
| 97 | last_reward_idx = i |
| 98 | if term or trunc: |
| 99 | break |
| 100 | |
| 101 | if score != archived_score: |
| 102 | sys.exit(f"[extract] REPLAY MISMATCH {score} != {archived_score} — " |
| 103 | f"demo is not replayable, refusing to write") |
| 104 | |
| 105 | # truncate just after a reward (papers: start right before a reward; nothing |
no test coverage detected