()
| 123 | |
| 124 | |
| 125 | def main(): |
| 126 | p = argparse.ArgumentParser(description=__doc__) |
| 127 | p.add_argument("--demo", required=True) |
| 128 | p.add_argument("--env", default="montezuma_goexplore_robust") |
| 129 | p.add_argument("--seed", type=int, default=0) |
| 130 | p.add_argument("--total-frames", type=int, default=None) |
| 131 | p.add_argument("--n-envs", type=int, default=None) |
| 132 | p.add_argument("--device", default="auto") |
| 133 | p.add_argument("--run-dir", default=None) |
| 134 | p.add_argument("--ckpt-every", type=int, default=None) |
| 135 | p.add_argument("--resume", default=None) |
| 136 | p.add_argument("--eval-episodes", type=int, default=50) |
| 137 | # stretch flags (off by default — see SPEC) |
| 138 | p.add_argument("--sil", action="store_true") |
| 139 | p.add_argument("--autoscale", action="store_true") |
| 140 | p.add_argument("--ent-coef", type=float, default=ENT_COEF, |
| 141 | help="entropy bonus coefficient (default keeps the module constant). " |
| 142 | "Raise to fight a competence plateau where the policy commits before " |
| 143 | "mastering the demo suffix under sticky actions.") |
| 144 | args = p.parse_args() |
| 145 | global TOTAL_FRAMES, N_ENVS |
| 146 | if args.total_frames: |
| 147 | TOTAL_FRAMES = args.total_frames |
| 148 | if args.n_envs: |
| 149 | N_ENVS = args.n_envs |
| 150 | |
| 151 | torch.manual_seed(args.seed) |
| 152 | np.random.seed(args.seed) |
| 153 | device = pick_device(args.device) |
| 154 | demo = load_demo(args.demo) |
| 155 | print(f"device {device} demo {len(demo['actions'])} actions score {demo['score']:.0f} " |
| 156 | f"n_envs {N_ENVS} total_frames {TOTAL_FRAMES:,} ent_coef {args.ent_coef:g}", flush=True) |
| 157 | |
| 158 | envs = [ReplayResetEnv(demo, seed=args.seed * 1000 + i, sticky=STICKY) for i in range(N_ENVS)] |
| 159 | mgr = ResetManager(demo, N_ENVS, move_threshold=MOVE_THRESHOLD) |
| 160 | mgr.assign(envs) |
| 161 | n_actions = envs[0].env.action_space.n |
| 162 | |
| 163 | net = GRUActorCritic(n_actions).to(device) |
| 164 | opt = torch.optim.Adam(net.parameters(), lr=LR, eps=ADAM_EPS) |
| 165 | logger = RunLogger(args.run_dir, args.ckpt_every) |
| 166 | |
| 167 | # reset all envs |
| 168 | stacks = np.stack([_stack_init(e.reset()) for e in envs]) # (N,4,105,80) |
| 169 | hx = torch.zeros(N_ENVS, net.gru_dim, device=device) |
| 170 | ep_start_nr = np.array([e.start_nr for e in envs]) |
| 171 | frames = 0 |
| 172 | update = 0 |
| 173 | |
| 174 | def _state(): |
| 175 | # RNG: global numpy + CPU torch + per-env numpy Generators (sticky / start-point |
| 176 | # sampling streams). Restored on resume so the kill→resume contract is faithful |
| 177 | # for the deterministic streams (preflight S2). MPS policy-sampling has no bit |
| 178 | # determinism (mps-pitfalls) — these cover what is reproducible. |
| 179 | return {"net": net.state_dict(), "opt": opt.state_dict(), "frames": frames, |
| 180 | "update": update, "max_starting_point": mgr.max_starting_point, |
| 181 | "success": mgr.success, "rng": np.random.get_state(), |
| 182 | "torch_rng": torch.get_rng_state(), |
no test coverage detected