def step(self, action):
obs, reward, done, info = self.env.step(action)
self.obss.append(obs)
return LazyFrames(list(self.obss)), reward, done, info
def reset(self, **kwargs):
obs = self.env.reset(**kwargs)
for _ in range(self._stack):
self.obss.append(obs)
return LazyFrames(list(self.obss))