From 68a29d0ab374f3e7e58b8ec8dcabc2e695318635 Mon Sep 17 00:00:00 2001
From: R1ckF <38724381+R1ckF@users.noreply.github.com>
Date: Mon, 17 Sep 2018 21:33:39 +0000
Subject: [PATCH] --play now works with LSTM (#595)

---
 baselines/run.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/baselines/run.py b/baselines/run.py
index cc1a512..b84bb47 100644
--- a/baselines/run.py
+++ b/baselines/run.py
@@ -221,8 +221,11 @@ def main():
         logger.log("Running trained model")
         env = build_env(args)
         obs = env.reset()
+        def initialize_placeholders(nlstm=128,**kwargs):
+            return np.zeros((args.num_env, 2*nlstm)), np.zeros((1))
+        state, dones = initialize_placeholders(**extra_args)
         while True:
-            actions = model.step(obs)[0]
+            actions, _, state, _ = model.step(obs,S=state, M=dones)
             obs, _, done, _ = env.step(actions)
             env.render()
             done = done.any() if isinstance(done, np.ndarray) else done