VideoProblem API with dynamic frame composition, video modality and making gym problem and one model work.

PiperOrigin-RevId: 192798025

Author: Lukasz Kaiser

tensor2tensor/data_generators/all_problems.py

@@ -20,45 +20,46 @@
 
 import importlib
 
+
 modules = [
-    'tensor2tensor.data_generators.algorithmic',
-    'tensor2tensor.data_generators.algorithmic_math',
-    'tensor2tensor.data_generators.audio',
-    'tensor2tensor.data_generators.celeba',
-    'tensor2tensor.data_generators.cifar',
-    'tensor2tensor.data_generators.cipher',
-    'tensor2tensor.data_generators.cnn_dailymail',
-    'tensor2tensor.data_generators.desc2code',
-    'tensor2tensor.data_generators.fsns',
-    'tensor2tensor.data_generators.gene_expression',
-    'tensor2tensor.data_generators.gym',
-    'tensor2tensor.data_generators.ice_parsing',
-    'tensor2tensor.data_generators.imagenet',
-    'tensor2tensor.data_generators.imdb',
-    'tensor2tensor.data_generators.librispeech',
-    'tensor2tensor.data_generators.lm1b',
-    'tensor2tensor.data_generators.mnist',
-    'tensor2tensor.data_generators.mscoco',
-    'tensor2tensor.data_generators.multinli',
-    'tensor2tensor.data_generators.ocr',
-    'tensor2tensor.data_generators.problem_hparams',
-    'tensor2tensor.data_generators.ptb',
-    'tensor2tensor.data_generators.snli',
-    'tensor2tensor.data_generators.squad',
-    'tensor2tensor.data_generators.translate_encs',
-    'tensor2tensor.data_generators.translate_ende',
-    'tensor2tensor.data_generators.translate_enfr',
-    'tensor2tensor.data_generators.translate_enmk',
-    'tensor2tensor.data_generators.translate_envi',
-    'tensor2tensor.data_generators.translate_enzh',
-    'tensor2tensor.data_generators.twentybn',
-    'tensor2tensor.data_generators.wiki',
-    'tensor2tensor.data_generators.wsj_parsing',
+    "tensor2tensor.data_generators.algorithmic",
+    "tensor2tensor.data_generators.algorithmic_math",
+    "tensor2tensor.data_generators.audio",
+    "tensor2tensor.data_generators.celeba",
+    "tensor2tensor.data_generators.cifar",
+    "tensor2tensor.data_generators.cipher",
+    "tensor2tensor.data_generators.cnn_dailymail",
+    "tensor2tensor.data_generators.desc2code",
+    "tensor2tensor.data_generators.fsns",
+    "tensor2tensor.data_generators.gene_expression",
+    "tensor2tensor.data_generators.gym",
+    "tensor2tensor.data_generators.ice_parsing",
+    "tensor2tensor.data_generators.imagenet",
+    "tensor2tensor.data_generators.imdb",
+    "tensor2tensor.data_generators.librispeech",
+    "tensor2tensor.data_generators.lm1b",
+    "tensor2tensor.data_generators.mnist",
+    "tensor2tensor.data_generators.mscoco",
+    "tensor2tensor.data_generators.multinli",
+    "tensor2tensor.data_generators.ocr",
+    "tensor2tensor.data_generators.problem_hparams",
+    "tensor2tensor.data_generators.ptb",
+    "tensor2tensor.data_generators.snli",
+    "tensor2tensor.data_generators.squad",
+    "tensor2tensor.data_generators.translate_encs",
+    "tensor2tensor.data_generators.translate_ende",
+    "tensor2tensor.data_generators.translate_enfr",
+    "tensor2tensor.data_generators.translate_enmk",
+    "tensor2tensor.data_generators.translate_envi",
+    "tensor2tensor.data_generators.translate_enzh",
+    "tensor2tensor.data_generators.twentybn",
+    "tensor2tensor.data_generators.wiki",
+    "tensor2tensor.data_generators.wsj_parsing",
 ]
 
 
 for module in modules:
   try:
     importlib.import_module(module)
-  except ImportError:
-    pass
+  except ImportError as error:
+    print("Did not import module: %s; Cause: %s" % (module, str(error)))

tensor2tensor/data_generators/gym.py

@@ -28,23 +28,152 @@
 
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
+from tensor2tensor.data_generators import video_utils
+
 from tensor2tensor.models.research import rl
 from tensor2tensor.rl import collect
 from tensor2tensor.rl.envs import tf_atari_wrappers as atari
 from tensor2tensor.rl.envs.utils import batch_env_factory
+
 from tensor2tensor.utils import registry
 
 import tensorflow as tf
 
-from tensorflow.contrib.training import HParams
-
 
 flags = tf.flags
 FLAGS = flags.FLAGS
 
 flags.DEFINE_string("agent_policy_path", "", "File with model for pong")
 
 
+class GymDiscreteProblem(video_utils.VideoProblem):
+  """Gym environment with discrete actions and rewards."""
+
+  def __init__(self, *args, **kwargs):
+    super(GymDiscreteProblem, self).__init__(*args, **kwargs)
+    self._env = None
+
+  @property
+  def num_input_frames(self):
+    """Number of frames to batch on one input."""
+    return 2
+
+  @property
+  def num_target_frames(self):
+    """Number of frames to batch on one target."""
+    return 1
+
+  @property
+  def extra_reading_spec(self):
+    """Additional data fields to store on disk and their decoders."""
+    data_fields = {
+        "action": tf.FixedLenFeature([1], tf.int64),
+        "reward": tf.FixedLenFeature([1], tf.int64)
+    }
+    decoders = {
+        "action": tf.contrib.slim.tfexample_decoder.Tensor(tensor_key="action"),
+        "reward": tf.contrib.slim.tfexample_decoder.Tensor(tensor_key="reward"),
+    }
+    return data_fields, decoders
+
+  @property
+  def is_generate_per_split(self):
+    """Whether we have a train/test split or just hold out data."""
+    return False  # Just hold out some generated data for evals.
+
+  @property
+  def env_name(self):
+    """This is the name of the Gym environment for this problem."""
+    raise NotImplementedError()
+
+  @property
+  def env(self):
+    if self._env is None:
+      self._env = gym.make(self.env_name)
+    return self._env
+
+  @property
+  def num_actions(self):
+    raise NotImplementedError()
+
+  @property
+  def num_rewards(self):
+    raise NotImplementedError()
+
+  @property
+  def num_steps(self):
+    raise NotImplementedError()
+
+  @property
+  def min_reward(self):
+    raise NotImplementedError()
+
+  def get_action(self, observation=None):
+    return self.env.action_space.sample()
+
+  def hparams(self, defaults, unused_model_hparams):
+    p = defaults
+    p.input_modality = {"inputs": ("video", 256),
+                        "input_reward": ("symbol", self.num_rewards),
+                        "input_action": ("symbol", self.num_actions)}
+    p.target_modality = ("video", 256)
+    p.input_space_id = problem.SpaceID.IMAGE
+    p.target_space_id = problem.SpaceID.IMAGE
+
+  def generate_samples(self, data_dir, tmp_dir, unused_dataset_split):
+    self.env.reset()
+    action = self.get_action()
+    for _ in range(self.num_steps):
+      observation, reward, done, _ = self.env.step(action)
+      action = self.get_action(observation)
+      yield {"frame": observation,
+             "action": [action],
+             "done": [done],
+             "reward": [int(reward - self.min_reward)]}
+
+
+@registry.register_problem
+class GymPongRandom5k(GymDiscreteProblem):
+  """Pong game, random actions."""
+
+  @property
+  def env_name(self):
+    return "PongDeterministic-v4"
+
+  @property
+  def frame_height(self):
+    return 210
+
+  @property
+  def frame_width(self):
+    return 160
+
+  @property
+  def num_actions(self):
+    return 4
+
+  @property
+  def min_reward(self):
+    return -1
+
+  @property
+  def num_rewards(self):
+    return 3
+
+  @property
+  def num_steps(self):
+    return 5000
+
+
+@registry.register_problem
+class GymPongRandom50k(GymPongRandom5k):
+  """Pong game, random actions."""
+
+  @property
+  def num_steps(self):
+    return 50000
+
+
 def moviepy_editor():
   """Access to moviepy that fails gracefully without a moviepy install."""
   try:
@@ -55,11 +184,11 @@ def moviepy_editor():
 
 
 @registry.register_problem
-class GymDiscreteProblem(problem.Problem):
+class GymDiscreteProblemWithAgent(problem.Problem):
   """Gym environment with discrete actions and rewards."""
 
   def __init__(self, *args, **kwargs):
-    super(GymDiscreteProblem, self).__init__(*args, **kwargs)
+    super(GymDiscreteProblemWithAgent, self).__init__(*args, **kwargs)
     self.num_channels = 3
     self.history_size = 2
 
@@ -68,16 +197,17 @@ def __init__(self, *args, **kwargs):
     self.in_graph_wrappers = [(atari.MaxAndSkipWrapper, {"skip": 4})]
     self.collect_hparams = rl.atari_base()
     self.num_steps = 1000
-    self.movies = True
+    self.movies = False
     self.movies_fps = 24
     self.simulated_environment = None
     self.warm_up = 70
 
   def _setup(self):
     in_graph_wrappers = [(atari.ShiftRewardWrapper, {"add_value": 2}),
                          (atari.MemoryWrapper, {})] + self.in_graph_wrappers
-    env_hparams = HParams(in_graph_wrappers=in_graph_wrappers,
-                          simulated_environment=self.simulated_environment)
+    env_hparams = tf.contrib.training.HParams(
+        in_graph_wrappers=in_graph_wrappers,
+        simulated_environment=self.simulated_environment)
 
     generator_batch_env = batch_env_factory(
         self.environment_spec, env_hparams, num_agents=1, xvfb=False)
@@ -234,19 +364,19 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
 
 
 @registry.register_problem
-class GymSimulatedDiscreteProblem(GymDiscreteProblem):
+class GymSimulatedDiscreteProblemWithAgent(GymDiscreteProblemWithAgent):
   """Simulated gym environment with discrete actions and rewards."""
 
   def __init__(self, *args, **kwargs):
-    super(GymSimulatedDiscreteProblem, self).__init__(*args, **kwargs)
+    super(GymSimulatedDiscreteProblemWithAgent, self).__init__(*args, **kwargs)
     # TODO(lukaszkaiser): pull it outside
     self.in_graph_wrappers = [(atari.TimeLimitWrapper, {"timelimit": 150}),
                               (atari.MaxAndSkipWrapper, {"skip": 4})]
     self.simulated_environment = True
     self.movies_fps = 2
 
   def restore_networks(self, sess):
-    super(GymSimulatedDiscreteProblem, self).restore_networks(sess)
+    super(GymSimulatedDiscreteProblemWithAgent, self).restore_networks(sess)
 
     # TODO(lukaszkaiser): adjust regexp for different models
     env_model_loader = tf.train.Saver(tf.global_variables(".*basic_conv_gen.*"))

tensor2tensor/data_generators/image_utils.py

@@ -158,7 +158,7 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
         self.dev_filepaths(data_dir, self.dev_shards, shuffled=False))
 
 
-def _encoded_images(images):
+def encode_images_as_png(images):
   if context.in_eager_mode():
     for image in images:
       yield tf.image.encode_png(image).numpy()
@@ -195,7 +195,7 @@ def image_generator(images, labels):
   if not images:
     raise ValueError("Must provide some images for the generator.")
   width, height, _ = images[0].shape
-  for (enc_image, label) in zip(_encoded_images(images), labels):
+  for (enc_image, label) in zip(encode_images_as_png(images), labels):
     yield {
         "image/encoded": [enc_image],
         "image/format": ["png"],

tensor2tensor/data_generators/problem.py

@@ -781,9 +781,7 @@ def define_shapes(example):
         batch_size_means_tokens = False
       else:
         tf.logging.warning(
-            "Shapes are not fully defined. Assuming batch_size means tokens. "
-            "Override batch_size_means_tokens() "
-            "in your problem subclass if this is undesired behavior.")
+            "Shapes are not fully defined. Assuming batch_size means tokens.")
         batch_size_means_tokens = True
 
     # Batching