One raw ALE env that starts episodes from demo states. Not a gym env — the vectorized loop in 3-robustify.py drives it directly.
| 58 | |
| 59 | |
| 60 | class ReplayResetEnv: |
| 61 | """One raw ALE env that starts episodes from demo states. Not a gym env — |
| 62 | the vectorized loop in 3-robustify.py drives it directly.""" |
| 63 | |
| 64 | def __init__(self, demo, seed, *, sticky=0.25, allowed_lag=50, |
| 65 | allowed_score_deficit=0, reset_steps_ignored=0, |
| 66 | inc_entropy_threshold=100, noop_max=30, max_steps=400_000): |
| 67 | import ale_py |
| 68 | import gymnasium as gym |
| 69 | gym.register_envs(ale_py) |
| 70 | self.env = gym.make(demo["env_id"], frameskip=4, |
| 71 | repeat_action_probability=0.0, # we add sticky ourselves |
| 72 | obs_type="grayscale").unwrapped |
| 73 | self.ale = self.env.ale |
| 74 | self.actions = demo["actions"] |
| 75 | self.rewards = demo["rewards"] |
| 76 | self.returns = demo["returns"] # cumulative raw, return-to-here |
| 77 | self.total_return = float(self.returns[-1]) |
| 78 | self.checkpoints = demo["checkpoints"] |
| 79 | self.ckpt_nr = demo["checkpoint_action_nr"] |
| 80 | self.n = len(self.actions) |
| 81 | self.sticky = StickyActionEnv(sticky) if sticky > 0 else None |
| 82 | self.allowed_lag = allowed_lag |
| 83 | self.allowed_score_deficit = allowed_score_deficit |
| 84 | self.reset_steps_ignored = reset_steps_ignored |
| 85 | self.inc_entropy_threshold = inc_entropy_threshold |
| 86 | self.noop_max = noop_max |
| 87 | self.max_steps = max_steps |
| 88 | self.rng = np.random.default_rng(seed) |
| 89 | self.starting_point = self.n - 1 |
| 90 | self.frac_sample = 0.2 |
| 91 | |
| 92 | # --- frame preprocessing: 105x80 grayscale (atari-reset uses RGB; grayscale |
| 93 | # keeps us light and matches the rest of this repo). 4-stack handled in |
| 94 | # the trainer. Returns uint8 (105, 80). --- |
| 95 | def _frame(self): |
| 96 | import cv2 |
| 97 | g = self.ale.getScreenGrayscale() |
| 98 | return cv2.resize(g, (80, 105), interpolation=cv2.INTER_AREA) |
| 99 | |
| 100 | def _restore_to(self, nr): |
| 101 | """Restore the latest checkpoint at or before nr, replay demo actions |
| 102 | up to nr (no sticky — deterministic), return the post-restore frame |
| 103 | from a real act (never a stale post-restore read).""" |
| 104 | ci = int(np.searchsorted(self.ckpt_nr, nr, side="right") - 1) |
| 105 | ci = max(ci, 0) |
| 106 | self.ale.restoreState(pickle.loads(self.checkpoints[ci])) |
| 107 | replay_from = int(self.ckpt_nr[ci]) |
| 108 | last_frame = None |
| 109 | for i in range(replay_from, nr): |
| 110 | self.ale.act(int(self.actions[i])) |
| 111 | last_frame = None # frames during pure replay are not needed |
| 112 | return last_frame |
| 113 | |
| 114 | def reset(self): |
| 115 | # per-episode starting point: 0.8 at the pinned point, 0.2 uniform tail |
| 116 | if self.rng.random() < self.frac_sample: |
| 117 | nr = int(self.rng.integers(self.starting_point, self.n)) |