| | from .vec_env import VecEnvWrapper |
| | import numpy as np |
| | from gym import spaces |
| |
|
| |
|
| | class VecFrameStack(VecEnvWrapper): |
| | def __init__(self, venv, nstack): |
| | self.venv = venv |
| | self.nstack = nstack |
| | wos = venv.observation_space |
| | low = np.repeat(wos.low, self.nstack, axis=-1) |
| | high = np.repeat(wos.high, self.nstack, axis=-1) |
| | self.stackedobs = np.zeros((venv.num_envs,) + low.shape, low.dtype) |
| | observation_space = spaces.Box(low=low, high=high, dtype=venv.observation_space.dtype) |
| | VecEnvWrapper.__init__(self, venv, observation_space=observation_space) |
| |
|
| | def step_wait(self): |
| | obs, rews, news, infos = self.venv.step_wait() |
| | self.stackedobs = np.roll(self.stackedobs, shift=-1, axis=-1) |
| | for (i, new) in enumerate(news): |
| | if new: |
| | self.stackedobs[i] = 0 |
| | self.stackedobs[..., -obs.shape[-1]:] = obs |
| | return self.stackedobs, rews, news, infos |
| |
|
| | def reset(self): |
| | obs = self.venv.reset() |
| | self.stackedobs[...] = 0 |
| | self.stackedobs[..., -obs.shape[-1]:] = obs |
| | return self.stackedobs |
| |
|
