internal merge of PR #1343

PiperOrigin-RevId: 228365154

Author: cclauss

tensor2tensor/data_generators/gym_env.py

@@ -23,6 +23,7 @@
 import itertools
 import os
 import random
+import re
 
 from gym.spaces import Box
 import numpy as np
@@ -641,6 +642,58 @@ def base_env_name(self):
   def num_channels(self):
     return self.observation_space.shape[2]
 
+  @staticmethod
+  def infer_last_epoch_num(data_dir):
+    """Infer highest epoch number from file names in data_dir."""
+    names = os.listdir(data_dir)
+    epochs_str = [re.findall(pattern=r".*\.(-?\d+)$", string=name)
+                  for name in names]
+    epochs_str = sum(epochs_str, [])
+    return max([int(epoch_str) for epoch_str in epochs_str])
+
+  @staticmethod
+  def setup_env_from_hparams(hparams, batch_size, max_num_noops):
+    game_mode = "NoFrameskip-v4"
+    camel_game_name = misc_utils.snakecase_to_camelcase(hparams.game)
+    camel_game_name += game_mode
+    env_name = camel_game_name
+
+    env = T2TGymEnv(base_env_name=env_name,
+                    batch_size=batch_size,
+                    grayscale=hparams.grayscale,
+                    resize_width_factor=hparams.resize_width_factor,
+                    resize_height_factor=hparams.resize_height_factor,
+                    rl_env_max_episode_steps=hparams.rl_env_max_episode_steps,
+                    max_num_noops=max_num_noops, maxskip_envs=True)
+    return env
+
+  @staticmethod
+  def setup_and_load_epoch(hparams, data_dir, which_epoch_data=None):
+    """Load T2TBatchGymEnv with data from one epoch.
+
+    Args:
+      hparams: hparams.
+      data_dir: data directory.
+      which_epoch_data: data from which epoch to load.
+
+    Returns:
+      env.
+    """
+    t2t_env = T2TGymEnv.setup_env_from_hparams(
+        hparams, batch_size=hparams.real_batch_size,
+        max_num_noops=hparams.max_num_noops
+    )
+    # Load data.
+    if which_epoch_data is not None:
+      if which_epoch_data == "last":
+        which_epoch_data = T2TGymEnv.infer_last_epoch_num(data_dir)
+      assert isinstance(which_epoch_data, int), \
+        "{}".format(type(which_epoch_data))
+      t2t_env.start_new_epoch(which_epoch_data, data_dir)
+    else:
+      t2t_env.start_new_epoch(-999)
+    return t2t_env
+
   def _derive_observation_space(self, orig_observ_space):
     height, width, channels = orig_observ_space.shape
     if self.grayscale:

tensor2tensor/models/research/rl.py

@@ -335,7 +335,6 @@ def rlmf_original():
       frame_stack_size=4,
       eval_sampling_temps=[0.0, 0.2, 0.5, 0.8, 1.0, 2.0],
       eval_max_num_noops=8,
-      eval_rl_env_max_episode_steps=1000,
       resize_height_factor=2,
       resize_width_factor=2,
       grayscale=0,

tensor2tensor/rl/evaluator.py

@@ -56,6 +56,7 @@
 import six
 
 from tensor2tensor.bin import t2t_trainer  # pylint: disable=unused-import
+from tensor2tensor.data_generators.gym_env import T2TGymEnv
 from tensor2tensor.rl import player_utils
 from tensor2tensor.rl.envs.simulated_batch_env import PIL_Image
 from tensor2tensor.rl.envs.simulated_batch_env import PIL_ImageDraw
@@ -228,7 +229,7 @@ def main(_):
         directories["data"], directories["world_model"],
         hparams, which_epoch_data=epoch)
   else:
-    env = player_utils.setup_and_load_epoch(
+    env = T2TGymEnv.setup_and_load_epoch(
         hparams, data_dir=directories["data"],
         which_epoch_data=epoch)
     env = FlatBatchEnv(env)

tensor2tensor/rl/player.py

@@ -21,7 +21,6 @@
 
 import copy
 import os
-import re
 
 import gym
 import numpy as np
@@ -40,42 +39,6 @@
 FLAGS = flags.FLAGS
 
 
-def infer_last_epoch_num(data_dir):
-  """Infer highest epoch number from file names in data_dir."""
-  names = os.listdir(data_dir)
-  epochs_str = [re.findall(pattern=r".*\.(-?\d+)$", string=name)
-                for name in names]
-  epochs_str = sum(epochs_str, [])
-  return max([int(epoch_str) for epoch_str in epochs_str])
-
-
-def setup_and_load_epoch(hparams, data_dir, which_epoch_data=None):
-  """Load T2TGymEnv with data from one epoch.
-
-  Args:
-    hparams: hparams.
-    data_dir: data directory.
-    which_epoch_data: data from which epoch to load.
-
-  Returns:
-    env.
-  """
-  t2t_env = rl_utils.setup_env(
-      hparams, batch_size=hparams.real_batch_size,
-      max_num_noops=hparams.max_num_noops
-  )
-  # Load data.
-  if which_epoch_data is not None:
-    if which_epoch_data == "last":
-      which_epoch_data = infer_last_epoch_num(data_dir)
-    assert isinstance(which_epoch_data, int), \
-      "{}".format(type(which_epoch_data))
-    t2t_env.start_new_epoch(which_epoch_data, data_dir)
-  else:
-    t2t_env.start_new_epoch(-999)
-  return t2t_env
-
-
 def make_simulated_gym_env(real_env, world_model_dir, hparams, random_starts):
   """Gym environment with world model."""
   initial_frame_chooser = rl_utils.make_initial_frame_chooser(
@@ -98,7 +61,7 @@ def load_data_and_make_simulated_env(
     data_dir, wm_dir, hparams, which_epoch_data="last", random_starts=True
 ):
   hparams = copy.deepcopy(hparams)
-  t2t_env = setup_and_load_epoch(
+  t2t_env = T2TGymEnv.setup_and_load_epoch(
       hparams, data_dir=data_dir,
       which_epoch_data=which_epoch_data)
   return make_simulated_gym_env(

tensor2tensor/rl/player_utils.py

@@ -28,7 +28,6 @@
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl.dopamine_connector import DQNLearner
 from tensor2tensor.rl.ppo_learner import PPOLearner
-from tensor2tensor.utils import misc_utils
 from tensor2tensor.utils import trainer_lib
 
 import tensorflow as tf
@@ -64,9 +63,8 @@ def evaluate_single_config(
 ):
   """Evaluate the PPO agent in the real environment."""
   eval_hparams = trainer_lib.create_hparams(hparams.base_algo_params)
-  env = setup_env(
-      hparams, batch_size=hparams.eval_batch_size, max_num_noops=max_num_noops,
-      rl_env_max_episode_steps=hparams.eval_rl_env_max_episode_steps
+  env = T2TGymEnv.setup_env_from_hparams(
+      hparams, batch_size=hparams.eval_batch_size, max_num_noops=max_num_noops
   )
   env.start_new_epoch(0)
   env_fn = rl.make_real_env_fn(env)
@@ -100,38 +98,12 @@ def evaluate_all_configs(hparams, agent_model_dir):
   return metrics
 
 
-def summarize_metrics(eval_metrics_writer, metrics, epoch):
-  """Write metrics to summary."""
-  for (name, value) in six.iteritems(metrics):
-    summary = tf.Summary()
-    summary.value.add(tag=name, simple_value=value)
-    eval_metrics_writer.add_summary(summary, epoch)
-  eval_metrics_writer.flush()
-
-
 LEARNERS = {
     "ppo": PPOLearner,
     "dqn": DQNLearner,
 }
 
 
-def setup_env(hparams, batch_size, max_num_noops, rl_env_max_episode_steps=-1):
-  """Setup."""
-  game_mode = "NoFrameskip-v4"
-  camel_game_name = misc_utils.snakecase_to_camelcase(hparams.game)
-  camel_game_name += game_mode
-  env_name = camel_game_name
-
-  env = T2TGymEnv(base_env_name=env_name,
-                  batch_size=batch_size,
-                  grayscale=hparams.grayscale,
-                  resize_width_factor=hparams.resize_width_factor,
-                  resize_height_factor=hparams.resize_height_factor,
-                  rl_env_max_episode_steps=rl_env_max_episode_steps,
-                  max_num_noops=max_num_noops, maxskip_envs=True)
-  return env
-
-
 def update_hparams_from_hparams(target_hparams, source_hparams, prefix):
   """Copy a subset of hparams to target_hparams."""
   for (param_name, param_value) in six.iteritems(source_hparams.values()):

tensor2tensor/rl/rl_utils.py

@@ -37,6 +37,7 @@
 import six
 
 from tensor2tensor.bin import t2t_trainer  # pylint: disable=unused-import
+from tensor2tensor.data_generators.gym_env import T2TGymEnv
 from tensor2tensor.layers import common_video
 from tensor2tensor.models.research import rl
 from tensor2tensor.models.research.rl import make_simulated_env_fn_from_hparams
@@ -377,6 +378,15 @@ def load_metrics(event_dir, epoch):
   return metrics
 
 
+def summarize_metrics(eval_metrics_writer, metrics, epoch):
+  """Write metrics to summary."""
+  for (name, value) in six.iteritems(metrics):
+    summary = tf.Summary()
+    summary.value.add(tag=name, simple_value=value)
+    eval_metrics_writer.add_summary(summary, epoch)
+  eval_metrics_writer.flush()
+
+
 def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
   """Run the main training loop."""
   if report_fn:
@@ -391,10 +401,9 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
 
   epoch = -1
   data_dir = directories["data"]
-  env = rl_utils.setup_env(
+  env = T2TGymEnv.setup_env_from_hparams(
       hparams, batch_size=hparams.real_batch_size,
-      max_num_noops=hparams.max_num_noops,
-      rl_env_max_episode_steps=hparams.rl_env_max_episode_steps
+      max_num_noops=hparams.max_num_noops
   )
   env.start_new_epoch(epoch, data_dir)
 
@@ -484,7 +493,7 @@ def training_loop(hparams, output_dir, report_fn=None, report_metric=None):
         log("World model eval metrics:\n{}".format(pprint.pformat(wm_metrics)))
         metrics.update(wm_metrics)
 
-      rl_utils.summarize_metrics(eval_metrics_writer, metrics, epoch)
+      summarize_metrics(eval_metrics_writer, metrics, epoch)
 
       # Report metrics
       if report_fn:

tensor2tensor/rl/trainer_model_based.py

@@ -84,8 +84,6 @@ def _rlmb_base():
       # Sampling temperatures to try during eval.
       eval_sampling_temps=[0.5, 0.0, 1.0],
       eval_max_num_noops=8,
-      # To speed up the pipeline. Some games want to run forever.
-      eval_rl_env_max_episode_steps=1000,
 
       game="pong",
       # Whether to evaluate the world model in each iteration of the loop to get
@@ -508,7 +506,6 @@ def _rlmb_tiny_overrides():
       resize_width_factor=2,
       wm_eval_rollout_ratios=[1],
       rl_env_max_episode_steps=7,
-      eval_rl_env_max_episode_steps=7,
       simulated_rollout_length=2,
       eval_sampling_temps=[0.0, 1.0],
   )

tensor2tensor/rl/trainer_model_based_params.py

@@ -20,7 +20,7 @@
 python -m tensor2tensor.rl.trainer_model_free \
     --output_dir=$HOME/t2t/rl_v1 \
     --hparams_set=pong_model_free \
-    --hparams='batch_size=15'
+    --loop_hparams='batch_size=15'
 """
 
 from __future__ import absolute_import
@@ -29,6 +29,7 @@
 
 import pprint
 
+from tensor2tensor.data_generators.gym_env import T2TGymEnv
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl import rl_utils
 from tensor2tensor.utils import flags as t2t_flags  # pylint: disable=unused-import
@@ -52,9 +53,9 @@
 
 def initialize_env_specs(hparams):
   """Initializes env_specs using T2TGymEnvs."""
-  env = rl_utils.setup_env(hparams, hparams.batch_size,
-                           hparams.eval_max_num_noops,
-                           hparams.rl_env_max_episode_steps)
+  env = T2TGymEnv.setup_env_from_hparams(
+      hparams, hparams.batch_size, hparams.eval_max_num_noops
+  )
   env.start_new_epoch(0)
 
   # TODO(afrozm): Decouple env_fn from hparams and return both, is there