fix a case when states is None.

2019-03-22 05:28:44 +09:00
parent c320622fe4
commit 23aaf83d79
1 changed files with 2 additions and 1 deletions
--- a/baselines/ppo2/model.py
+++ b/baselines/ppo2/model.py
@@ -136,7 +136,8 @@ class Model(object):
    def step(self, observations, **kwargs):
        kwargs.update({'observations': observations})
        transition = self.act_model.step(**kwargs)
-        return transition['actions'], transition['values'], transition['states'], transition['neglogpacs']
+        states = transition['states'] if 'states' in transition else None
+        return transition['actions'], transition['values'], states, transition['neglogpacs']

    def train(self,
              lr,