Merge pull request #701 from deepsense-ai/rl_model_experiment

RL model-based experiment

Author: lukaszkaiser

tensor2tensor/data_generators/gym.py

@@ -24,76 +24,136 @@
 # Dependency imports
 
 import gym
-import numpy as np
+import os
+from tensorflow.contrib.training import HParams
+from collections import deque
 
 from tensor2tensor.data_generators import generator_utils
 from tensor2tensor.data_generators import problem
-
 from tensor2tensor.models.research import rl
-from tensor2tensor.rl import rl_trainer_lib  # pylint: disable=unused-import
-from tensor2tensor.rl.envs import atari_wrappers
-
-from tensor2tensor.utils import metrics
 from tensor2tensor.utils import registry
+from tensor2tensor.rl.envs.utils import batch_env_factory
+from tensor2tensor.rl.envs.tf_atari_wrappers import MemoryWrapper, TimeLimitWrapper
+from tensor2tensor.rl.envs.tf_atari_wrappers import MaxAndSkipWrapper
+from tensor2tensor.rl.envs.tf_atari_wrappers import PongT2TGeneratorHackWrapper
+from tensor2tensor.rl import collect
 
 import tensorflow as tf
 
 
+def moviepy_editor():
+  """Access to moviepy to allow for import of this file without a moviepy install."""
+  try:
+    from moviepy import editor # pylint: disable=g-import-not-at-top
+  except ImportError:
+    raise ImportError("pip install moviepy to record videos")
+  return editor
+
 flags = tf.flags
 FLAGS = flags.FLAGS
 
-flags.DEFINE_string("model_path", "", "File with model for pong")
-
+flags.DEFINE_string("agent_policy_path", "", "File with model for pong")
 
+@registry.register_problem
 class GymDiscreteProblem(problem.Problem):
   """Gym environment with discrete actions and rewards."""
 
   def __init__(self, *args, **kwargs):
     super(GymDiscreteProblem, self).__init__(*args, **kwargs)
-    self._env = None
+    self.num_channels = 3
+    self.history_size = 2
+
+    # defaults
+    self.environment_spec = lambda: gym.make("PongNoFrameskip-v4")
+    self.in_graph_wrappers = [(MaxAndSkipWrapper, {"skip": 4})]
+    self.collect_hparams = rl.atari_base()
+    self.num_steps = 1000
+    self.movies = True
+    self.movies_fps = 24
+    self.simulated_environment = None
+    self.warm_up = 70
+
+  def _setup(self):
+    # TODO: remove PongT2TGeneratorHackWrapper by writing a modality
+
+    in_graph_wrappers = [(PongT2TGeneratorHackWrapper, {"add_value": 2}),
+                         (MemoryWrapper, {})] + self.in_graph_wrappers
+    env_hparams = 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)
+
+    with tf.variable_scope("", reuse=tf.AUTO_REUSE):
+      policy_lambda = self.collect_hparams.network
+      policy_factory = tf.make_template(
+        "network",
+        functools.partial(policy_lambda, self.environment_spec().action_space, self.collect_hparams),
+        create_scope_now_=True,
+        unique_name_="network")
 
-  def example_reading_spec(self, label_repr=None):
+    with tf.variable_scope("", reuse=tf.AUTO_REUSE):
+      sample_policy = lambda policy: 0*policy.sample()
+      # sample_policy = lambda policy: 0
 
+      self.collect_hparams.epoch_length = 10
+      _, self.collect_trigger_op = collect.define_collect(
+        policy_factory, generator_batch_env, self.collect_hparams,
+        eval_phase=False, policy_to_actions_lambda=sample_policy, scope="define_collect")
+
+    self.avilable_data_size_op = MemoryWrapper.singleton._speculum.size()
+    self.data_get_op = MemoryWrapper.singleton._speculum.dequeue()
+    self.history_buffer = deque(maxlen=self.history_size+1)
+
+  def example_reading_spec(self, label_repr=None):
     data_fields = {
-        "inputs": tf.FixedLenFeature([210, 160, 3], tf.int64),
-        "inputs_prev": tf.FixedLenFeature([210, 160, 3], tf.int64),
-        "targets": tf.FixedLenFeature([210, 160, 3], tf.int64),
-        "action": tf.FixedLenFeature([1], tf.int64),
-        "reward": tf.FixedLenFeature([1], tf.int64)
+      "targets_encoded": tf.FixedLenFeature((), tf.string),
+      "image/format": tf.FixedLenFeature((), tf.string),
+      "action": tf.FixedLenFeature([1], tf.int64),
+      "reward": tf.FixedLenFeature([1], tf.int64),
+      # "done": tf.FixedLenFeature([1], tf.int64)
     }
 
-    return data_fields, None
+    for x in range(self.history_size):
+      data_fields["inputs_encoded_{}".format(x)] =  tf.FixedLenFeature((), tf.string)
 
-  def eval_metrics(self):
-    return [metrics.Metrics.ACC, metrics.Metrics.ACC_PER_SEQ,
-            metrics.Metrics.NEG_LOG_PERPLEXITY, metrics.Metrics.IMAGE_SUMMARY]
 
-  @property
-  def env_name(self):
-    # This is the name of the Gym environment for this problem.
-    raise NotImplementedError()
+    data_items_to_decoders = {
+      "targets":
+        tf.contrib.slim.tfexample_decoder.Image(
+          image_key="targets_encoded",
+          format_key="image/format",
+          shape=[210, 160, 3],
+          channels=3),
 
-  @property
-  def env(self):
-    if self._env is None:
-      self._env = gym.make(self.env_name)
-    return self._env
+      #Just do a pass through
+      "action":tf.contrib.slim.tfexample_decoder.Tensor(tensor_key="action"),
+      "reward":tf.contrib.slim.tfexample_decoder.Tensor(tensor_key="reward"),
+    }
 
-  @property
-  def num_channels(self):
-    return 3
+    for x in range(self.history_size):
+      data_items_to_decoders["inputs_{}".format(x)] =  tf.contrib.slim.tfexample_decoder.Image(
+                image_key="inputs_encoded_{}".format(x),
+                format_key="image/format",
+                shape=[210, 160, 3],
+                channels=3)
+
+    return data_fields, data_items_to_decoders
+
+  # def preprocess_example(self, example, mode, hparams):
+  #   if not self._was_reversed:
+  #     for x in range(self.history_size):
+  #       input_name = "inputs_{}".format(x)
+  #       example[input_name] = tf.image.per_image_standardization(example[input_name])
+  #   return example
 
   @property
   def num_actions(self):
-    raise NotImplementedError()
+    return 4
 
   @property
   def num_rewards(self):
-    raise NotImplementedError()
-
-  @property
-  def num_steps(self):
-    raise NotImplementedError()
+    return 2
 
   @property
   def num_shards(self):
@@ -108,35 +168,70 @@ def get_action(self, observation=None):
 
   def hparams(self, defaults, unused_model_hparams):
     p = defaults
-    p.input_modality = {"inputs": ("image", 256),
-                        "inputs_prev": ("image", 256),
-                        "reward": ("symbol", self.num_rewards),
-                        "action": ("symbol", self.num_actions)}
-    p.target_modality = ("image", 256)
+    # hard coded +1 after "symbol" refers to the fact
+    # that 0 is a special symbol meaning padding
+    # when symbols are e.g. 0, 1, 2, 3 we
+    # shift them to 0, 1, 2, 3, 4
+    p.input_modality = {"action": ("symbol:identity", self.num_actions)}
+
+    for x in range(self.history_size):
+      p.input_modality["inputs_{}".format(x)] = ("image", 256)
+
+    p.target_modality = {"targets": ("image", 256),
+                         "reward":  ("symbol", self.num_rewards+1),
+                         # "done": ("symbol", 2+1)
+                         }
+
     p.input_space_id = problem.SpaceID.IMAGE
     p.target_space_id = problem.SpaceID.IMAGE
 
+  def restore_networks(self, sess):
+    model_saver = tf.train.Saver(
+      tf.global_variables(".*network_parameters.*"))
+    if FLAGS.agent_policy_path:
+      model_saver.restore(sess, FLAGS.agent_policy_path)
+
   def generator(self, data_dir, tmp_dir):
-    self.env.reset()
-    action = self.get_action()
-    prev_observation, observation = None, None
-    for _ in range(self.num_steps):
-      prev_prev_observation = prev_observation
-      prev_observation = observation
-      observation, reward, done, _ = self.env.step(action)
-      action = self.get_action(observation)
-      if done:
-        self.env.reset()
-      def flatten(nparray):
-        flat1 = [x for sublist in nparray.tolist() for x in sublist]
-        return [x for sublist in flat1 for x in sublist]
-      if prev_prev_observation is not None:
-        yield {"inputs_prev": flatten(prev_prev_observation),
-               "inputs": flatten(prev_observation),
-               "action": [action],
-               "done": [done],
-               "reward": [int(reward)],
-               "targets": flatten(observation)}
+    self._setup()
+    clip_files = []
+    with tf.Session() as sess:
+      sess.run(tf.global_variables_initializer())
+      self.restore_networks(sess)
+
+      pieces_generated = 0
+      while pieces_generated<self.num_steps + self.warm_up:
+        avilable_data_size = sess.run(self.avilable_data_size_op)
+        if avilable_data_size>0:
+          observ, reward, action, done = sess.run(self.data_get_op)
+          self.history_buffer.append(observ)
+
+          if self.movies==True and pieces_generated>self.warm_up:
+            file_name = os.path.join(tmp_dir,'output_{}.png'.format(pieces_generated))
+            clip_files.append(file_name)
+            with open(file_name, 'wb') as f:
+              f.write(observ)
+
+          if len(self.history_buffer)==self.history_size+1:
+            pieces_generated += 1
+            ret_dict = {
+              "targets_encoded": [observ],
+              "image/format": ["png"],
+              "action": [int(action)],
+              # "done": [bool(done)],
+              "reward": [int(reward)],
+                }
+            for i, v in enumerate(list(self.history_buffer)[:-1]):
+              ret_dict["inputs_encoded_{}".format(i)] = [v]
+            if pieces_generated>self.warm_up:
+              yield ret_dict
+        else:
+          sess.run(self.collect_trigger_op)
+    if self.movies:
+      # print(clip_files)
+      clip = moviepy_editor().ImageSequenceClip(clip_files, fps=self.movies_fps)
+      clip.write_videofile(os.path.join(data_dir, 'output_{}.mp4'.format(self.name)),
+                           fps=self.movies_fps, codec='mpeg4')
+
 
   def generate_data(self, data_dir, tmp_dir, task_id=-1):
     train_paths = self.training_filepaths(
@@ -150,93 +245,23 @@ def generate_data(self, data_dir, tmp_dir, task_id=-1):
 
 
 @registry.register_problem
-class GymPongRandom5k(GymDiscreteProblem):
-  """Pong game, random actions."""
-
-  @property
-  def env_name(self):
-    return "PongDeterministic-v4"
-
-  @property
-  def num_actions(self):
-    return 4
-
-  @property
-  def num_rewards(self):
-    return 2
-
-  @property
-  def num_steps(self):
-    return 5000
-
-
-@registry.register_problem
-class GymPongTrajectoriesFromPolicy(GymDiscreteProblem):
-  """Pong game, loaded actions."""
+class GymSimulatedDiscreteProblem(GymDiscreteProblem):
+  """Simulated gym environment with discrete actions and rewards."""
 
   def __init__(self, *args, **kwargs):
-    super(GymPongTrajectoriesFromPolicy, self).__init__(*args, **kwargs)
-    self._env = None
-    self._last_policy_op = None
-    self._max_frame_pl = None
-    self._last_action = self.env.action_space.sample()
-    self._skip = 4
-    self._skip_step = 0
-    self._obs_buffer = np.zeros((2,) + self.env.observation_space.shape,
-                                dtype=np.uint8)
-
-  def generator(self, data_dir, tmp_dir):
-    env_spec = lambda: atari_wrappers.wrap_atari(  # pylint: disable=g-long-lambda
-        gym.make(self.env_name),
-        warp=False,
-        frame_skip=4,
-        frame_stack=False)
-    hparams = rl.atari_base()
-    with tf.variable_scope("train", reuse=tf.AUTO_REUSE):
-      policy_lambda = hparams.network
-      policy_factory = tf.make_template(
-          "network",
-          functools.partial(policy_lambda, env_spec().action_space, hparams))
-      self._max_frame_pl = tf.placeholder(
-          tf.float32, self.env.observation_space.shape)
-      actor_critic = policy_factory(tf.expand_dims(tf.expand_dims(
-          self._max_frame_pl, 0), 0))
-      policy = actor_critic.policy
-      self._last_policy_op = policy.mode()
-      with tf.Session() as sess:
-        model_saver = tf.train.Saver(
-            tf.global_variables(".*network_parameters.*"))
-        model_saver.restore(sess, FLAGS.model_path)
-        for item in super(GymPongTrajectoriesFromPolicy,
-                          self).generator(data_dir, tmp_dir):
-          yield item
-
-  # TODO(blazej0): For training of atari agents wrappers are usually used.
-  # Below we have a hacky solution which is a workaround to be used together
-  # with atari_wrappers.MaxAndSkipEnv.
-  def get_action(self, observation=None):
-    if self._skip_step == self._skip - 2: self._obs_buffer[0] = observation
-    if self._skip_step == self._skip - 1: self._obs_buffer[1] = observation
-    self._skip_step = (self._skip_step + 1) % self._skip
-    if self._skip_step == 0:
-      max_frame = self._obs_buffer.max(axis=0)
-      self._last_action = int(tf.get_default_session().run(
-          self._last_policy_op,
-          feed_dict={self._max_frame_pl: max_frame})[0, 0])
-    return self._last_action
-
-  @property
-  def env_name(self):
-    return "PongDeterministic-v4"
-
-  @property
-  def num_actions(self):
-    return 4
-
-  @property
-  def num_rewards(self):
-    return 2
-
-  @property
-  def num_steps(self):
-    return 5000
+    super(GymSimulatedDiscreteProblem, self).__init__(*args, **kwargs)
+    #TODO: pull it outside
+    self.in_graph_wrappers = [(TimeLimitWrapper, {"timelimit": 150}), (MaxAndSkipWrapper, {"skip": 4})]
+    self.simulated_environment = True
+    self.movies_fps = 2
+
+  def restore_networks(self, sess):
+    super(GymSimulatedDiscreteProblem, self).restore_networks(sess)
+
+    #TODO: adjust regexp for different models
+    env_model_loader = tf.train.Saver(tf.global_variables(".*basic_conv_gen.*"))
+    sess = tf.get_default_session()
+
+    ckpts = tf.train.get_checkpoint_state(FLAGS.output_dir)
+    ckpt = ckpts.model_checkpoint_path
+    env_model_loader.restore(sess, ckpt)

tensor2tensor/layers/modalities.py

@@ -100,7 +100,6 @@ def bottom_simple(self, x, name, reuse):
         x = tf.squeeze(x, axis=3)
       while len(x.get_shape()) < 3:
         x = tf.expand_dims(x, axis=-1)
-
       var = self._get_weights()
       x = common_layers.dropout_no_scaling(
           x, 1.0 - self._model_hparams.symbol_dropout)