Introducing video_during_eval hparam.

Author: Błażej O

tensor2tensor/models/research/rl.py

@@ -48,6 +48,7 @@ def ppo_base_v1():
   hparams.add_hparam("epochs_num", 2000)
   hparams.add_hparam("eval_every_epochs", 10)
   hparams.add_hparam("num_eval_agents", 3)
+  hparams.add_hparam("video_during_eval", True)
   return hparams
 
 

tensor2tensor/rl/envs/utils.py

@@ -52,9 +52,8 @@ class EvalVideoWrapper(gym.Wrapper):
     returns last seen observation.
   Videos are only generated during the active runs.
   """
-  def __init__(self, env, directory):
-    super(EvalVideoWrapper, self).__init__(
-      gym.wrappers.Monitor(env, directory, video_callable=lambda i: i % 2 == 0))
+  def __init__(self, env):
+    super(EvalVideoWrapper, self).__init__(env)
     self._reset_counter = 0
     self._active = False
     self._last_returned = None

tensor2tensor/rl/rl_trainer_lib.py

@@ -57,17 +57,21 @@ def define_train(hparams, environment_spec, event_dir):
 
   with tf.variable_scope("eval"):
     eval_env_lambda = env_lambda
-    if event_dir:
-      eval_env_lambda = lambda: utils.EvalVideoWrapper(env_lambda(), event_dir)
+    if event_dir and hparams.video_during_eval:
+      eval_env_lambda = lambda: gym.wrappers.Monitor(
+        env_lambda(), event_dir, video_callable=lambda i: i % 2 == 0)
+    wrapped_eval_env_lambda = lambda: utils.EvalVideoWrapper(eval_env_lambda())
     _, eval_summary = collect.define_collect(
       policy_factory,
-      utils.define_batch_env(eval_env_lambda, hparams.num_eval_agents, xvfb=True),
+      utils.define_batch_env(wrapped_eval_env_lambda, hparams.num_eval_agents,
+                             xvfb=hparams.video_during_eval),
       hparams, eval_phase=True)
   return summary, eval_summary
 
 
 def train(hparams, environment_spec, event_dir=None):
-  train_summary_op, eval_summary_op = define_train(hparams, environment_spec, event_dir)
+  train_summary_op, eval_summary_op = define_train(hparams, environment_spec,
+                                                   event_dir)
 
   if event_dir:
     summary_writer = tf.summary.FileWriter(

tensor2tensor/rl/rl_trainer_lib_test.py

@@ -26,11 +26,13 @@
 class TrainTest(tf.test.TestCase):
 
   def test_no_crash_pendulum(self):
-    hparams = trainer_lib.create_hparams("continuous_action_base", "epochs_num=10")
+    hparams = trainer_lib.create_hparams(
+      "continuous_action_base", "epochs_num=11,video_during_eval=False")
     rl_trainer_lib.train(hparams, "Pendulum-v0")
 
   def test_no_crash_cartpole(self):
-    hparams = trainer_lib.create_hparams("discrete_action_base", "epochs_num=10")
+    hparams = trainer_lib.create_hparams(
+      "discrete_action_base", "epochs_num=11,video_during_eval=False")
     rl_trainer_lib.train(hparams, "CartPole-v0")