override gym env constraint

rll · Nov 12, 2016 · 4fd25fe · 4fd25fe
1 parent 86fc7bb
commit 4fd25fe
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/rllab/envs/gym_env.py b/rllab/envs/gym_env.py
@@ -30,7 +30,6 @@ def __call__(self, count):
 
 
 class FixedIntervalVideoSchedule(object):
-
     def __init__(self, interval):
         self.interval = interval
 
@@ -89,6 +88,9 @@ def horizon(self):
         return self._horizon
 
     def reset(self):
+        if hasattr(self.env, 'monitor'):
+            if hasattr(self.env.monitor, 'stats_recorder'):
+                self.env.monitor.stats_recorder.done = True
         return self.env.reset()
 
     def step(self, action):