fix video disabling

berkeleydeeprlcourse · Oct 16, 2020 · b5d3498 · b5d3498
1 parent d28416b
commit b5d3498
Showing 1 changed file with 12 additions and 4 deletions.
diff --git a/hw3/cs285/infrastructure/rl_trainer.py b/hw3/cs285/infrastructure/rl_trainer.py
@@ -54,14 +54,22 @@ def __init__(self, params):
         self.env = gym.make(self.params['env_name'])
         if 'env_wrappers' in self.params:
             # These operations are currently only for Atari envs
-            self.env = wrappers.Monitor(self.env, os.path.join(self.params['logdir'], "gym"), force=True)
-            self.env.enabled = (self.params['video_log_freq'] > 0)
+            self.env = wrappers.Monitor(
+                self.env,
+                os.path.join(self.params['logdir'], "gym"),
+                force=True,
+                video_callable=(None if self.params['video_log_freq'] > 0 else False),
+            )
             self.env = params['env_wrappers'](self.env)
             self.mean_episode_reward = -float('nan')
             self.best_mean_episode_reward = -float('inf')
         if 'non_atari_colab_env' in self.params and self.params['video_log_freq'] > 0:
-            self.env = wrappers.Monitor(self.env, os.path.join(self.params['logdir'], "gym"), force=True)
-            self.env.enabled = (self.params['video_log_freq'] > 0)
+            self.env = wrappers.Monitor(
+                self.env,
+                os.path.join(self.params['logdir'], "gym"),
+                force=True,
+                video_callable=(None if self.params['video_log_freq'] > 0 else False),
+            )
             self.mean_episode_reward = -float('nan')
             self.best_mean_episode_reward = -float('inf')