Owns the shared curriculum across N envs. The trainer calls assign() once to stagger starting points, and update() each time a batch of episodes finishes to march max_starting_point backward.
| 178 | |
| 179 | |
| 180 | class ResetManager: |
| 181 | """Owns the shared curriculum across N envs. The trainer calls assign() once |
| 182 | to stagger starting points, and update() each time a batch of episodes |
| 183 | finishes to march max_starting_point backward.""" |
| 184 | |
| 185 | def __init__(self, demo, n_envs, *, move_threshold=0.1, nudge=100, window=None): |
| 186 | self.n = len(demo["actions"]) |
| 187 | self.n_envs = n_envs |
| 188 | self.move_threshold = move_threshold |
| 189 | self.nudge = nudge |
| 190 | # window = the span of staggered starting points (atari-reset nrstartsteps). |
| 191 | # The move target is move_threshold*window of cumulative success mass. |
| 192 | self.window = window or max(n_envs, 32) |
| 193 | self.max_starting_point = self.n - 1 |
| 194 | self.max_max = self.n - 1 |
| 195 | # latest success-rate per starting-point index |
| 196 | self.success = np.zeros(self.n + 1, dtype=np.float64) |
| 197 | |
| 198 | def assign(self, envs): |
| 199 | """Stagger envs across a window below max_starting_point.""" |
| 200 | per = max(self.window // max(self.n_envs, 1), 1) |
| 201 | for i, e in enumerate(envs): |
| 202 | e.starting_point = max(self.max_starting_point - i * per, 0) |
| 203 | |
| 204 | def record(self, starting_point, success): |
| 205 | # exponential-ish freshening: latest wins (atari-reset keeps last rate) |
| 206 | self.success[min(starting_point, self.n)] = float(success) |
| 207 | |
| 208 | def update(self, envs): |
| 209 | """Move rule (atari-reset ResetManager.proc_infos): forward-cumsum the |
| 210 | per-index success rates from index 0; the new max starting point is the |
| 211 | FIRST index where the cumulative mass reaches move_threshold*window — |
| 212 | i.e. march back as far as the practiced success band supports, no |
| 213 | further. If the mass is never reached (success collapsed), nudge the |
| 214 | curriculum forward (easier) by `nudge`.""" |
| 215 | tail = self.success[: self.max_starting_point + 1] |
| 216 | csum = np.cumsum(tail) # forward: mass accumulated up to each index |
| 217 | hits = np.argwhere(csum >= self.move_threshold * self.window) |
| 218 | if len(hits): |
| 219 | new_max = int(hits[0][0]) # earliest index reaching the mass |
| 220 | self.max_starting_point = max(min(new_max, self.max_starting_point), 0) |
| 221 | else: |
| 222 | self.max_starting_point = min(self.max_starting_point + self.nudge, self.max_max) |
| 223 | self.assign(envs) |
| 224 | return self.max_starting_point |
| 225 | |
| 226 | |
| 227 | def load_demo(path): |