support the play.

baselines/run.py

@@ -5,6 +5,7 @@ import gym
 from collections import defaultdict
 import tensorflow as tf
 import numpy as np
+import inspect
 
 from baselines.common.vec_env import VecFrameStack, VecNormalize, VecEnv
 from baselines.common.vec_env.vec_video_recorder import VecVideoRecorder
@@ -218,15 +219,20 @@ def main(args):
         obs = env.reset()
 
         state = model.initial_state if hasattr(model, 'initial_state') else None
-        dones = np.zeros((1,))
+        use_external_memory_management = True if hasattr(model, 'initial_state') else None
+        step_func_args = inspect.getfullargspec(model.step).args
+        done = np.zeros((1,))
 
         episode_rew = 0
         while True:
-            if state is not None:
+            kwargs = {}
-                actions, _, state, _ = model.step(obs,S=state, M=dones)
+            if use_external_memory_management:
-            else:
+                kwargs['S'] = state
-                actions, _, _, _ = model.step(obs)
+                kwargs['M'] = done
+            elif 'done' in step_func_args:
+                kwargs['done'] = done
 
+            actions, _, state, _ = model.step(obs, **kwargs)
             obs, rew, done, _ = env.step(actions)
             episode_rew += rew[0] if isinstance(env, VecEnv) else rew
             env.render()