[rllib] Fix Preprocessor for ATARI (#1066)

* Removing squeeze, fix atari preprocessing

* nit comment

* comments

* jenkins

* Lint
This commit is contained in:
Richard Liaw
2017-10-03 18:45:02 -07:00
committed by Eric Liang
parent 0dcf36c91e
commit cb6dea94bc
5 changed files with 17 additions and 8 deletions
+3 -2
View File
@@ -22,7 +22,8 @@ class BatchedEnv(object):
def reset(self):
observations = [
self.preprocessor.transform(env.reset()) for env in self.envs]
self.preprocessor.transform(env.reset())[None]
for env in self.envs]
self.shape = observations[0].shape
self.dones = [False for _ in range(self.batchsize)]
return np.vstack(observations)
@@ -43,7 +44,7 @@ class BatchedEnv(object):
break
if render:
self.envs[0].render()
observations.append(self.preprocessor.transform(observation))
observations.append(self.preprocessor.transform(observation)[None])
rewards.append(reward)
self.dones[i] = done
return (np.vstack(observations), np.array(rewards, dtype="float32"),