| 148 | return self._frame(), float(r) |
| 149 | |
| 150 | def step(self, action): |
| 151 | frame, raw_r = self._step_raw(action) |
| 152 | info = {"raw_reward": raw_r} |
| 153 | done = False |
| 154 | |
| 155 | # success: as good as the demo from here |
| 156 | if self.extra == 0 and self.score >= self.total_return - self.allowed_score_deficit: |
| 157 | self.extra = int(np.exp(self.rng.random() * 7)) # 1..1096 |
| 158 | if self.extra > 0: |
| 159 | self.extra -= 1 |
| 160 | if self.extra == 0: |
| 161 | done = True |
| 162 | info["random_reset"] = True |
| 163 | info["as_good_as_demo"] = True |
| 164 | |
| 165 | # lag kill: fell behind the demo's pace (windowed-min, deficit-aware) |
| 166 | t = self.action_nr |
| 167 | if not done and t > self.allowed_lag and t < self.n: |
| 168 | lo = max(t - self.allowed_lag, 0) |
| 169 | hi = min(t + self.allowed_lag, self.n) |
| 170 | threshold = float(self.returns[lo:hi].min()) - self.allowed_score_deficit |
| 171 | if self.score < threshold: |
| 172 | done = True |
| 173 | |
| 174 | if self.ale.game_over() or self.action_nr - self.start_nr >= self.max_steps: |
| 175 | done = True |
| 176 | info["increase_entropy"] = (self.action_nr < self.start_nr + self.inc_entropy_threshold) |
| 177 | return frame, np.sign(raw_r), done, info # clipped reward to the agent |
| 178 | |
| 179 | |
| 180 | class ResetManager: |