Buffer consecutive k observations and stack them on a new last axis. The output observation has shape `original_shape + (k, )`.
| 25 | |
| 26 | |
| 27 | class FrameStack(gym.Wrapper): |
| 28 | """ |
| 29 | Buffer consecutive k observations and stack them on a new last axis. |
| 30 | The output observation has shape `original_shape + (k, )`. |
| 31 | """ |
| 32 | def __init__(self, env, k): |
| 33 | gym.Wrapper.__init__(self, env) |
| 34 | self.k = k |
| 35 | self.frames = deque([], maxlen=k) |
| 36 | |
| 37 | def reset(self): |
| 38 | """Clear buffer and re-fill by duplicating the first observation.""" |
| 39 | ob = self.env.reset() |
| 40 | for _ in range(self.k - 1): |
| 41 | self.frames.append(np.zeros_like(ob)) |
| 42 | self.frames.append(ob) |
| 43 | return self.observation() |
| 44 | |
| 45 | def step(self, action): |
| 46 | ob, reward, done, info = self.env.step(action) |
| 47 | self.frames.append(ob) |
| 48 | return self.observation(), reward, done, info |
| 49 | |
| 50 | def observation(self): |
| 51 | assert len(self.frames) == self.k |
| 52 | return np.stack(self.frames, axis=-1) |
| 53 | |
| 54 | |
| 55 | class _FireResetEnv(gym.Wrapper): |