From 4c5b2cd9a184d0b071660c9902466aeedd8c9e4f Mon Sep 17 00:00:00 2001 From: wwiiiii Date: Tue, 16 Apr 2019 12:36:10 +0900 Subject: [PATCH] Fix indexing LazyFrames Indexing LazyFrames with index i should return the single channel frame --- baselines/common/atari_wrappers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/baselines/common/atari_wrappers.py b/baselines/common/atari_wrappers.py index 2c9b8c69ce..585081fc7f 100644 --- a/baselines/common/atari_wrappers.py +++ b/baselines/common/atari_wrappers.py @@ -221,7 +221,7 @@ def __len__(self): return len(self._force()) def __getitem__(self, i): - return self._force()[i] + return self._force()[..., i] def make_atari(env_id, max_episode_steps=None): env = gym.make(env_id)