[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
+2 -4
View File
@@ -99,15 +99,13 @@ class Policy:
t = 0
if save_obs:
obs = []
# TODO(ekl) the squeeze() is needed for Pong-v0, but we should fix
# this in the preprocessor instead
ob = preprocessor.transform(env.reset()).squeeze()
ob = preprocessor.transform(env.reset())
for _ in range(timestep_limit):
ac = self.act(ob[None], random_stream=random_stream)[0]
if save_obs:
obs.append(ob)
ob, rew, done, _ = env.step(ac)
ob = preprocessor.transform(ob).squeeze()
ob = preprocessor.transform(ob)
rews.append(rew)
t += 1
if render: