code: import time
  from contextlib import contextmanager
  from collections import deque
  
  import gym
  from mpi4py import MPI
  import tensorflow as tf
  import numpy as np
  
  import stable_baselines.common.tf_util as tf_util
  from stable_baselines.common.tf_util import total_episode_reward_logger
  from stable_baselines.common import explained_variance, zipsame, dataset, fmt_row, colorize, ActorCriticRLModel, \
      SetVerbosity, TensorboardWriter
  from stable_baselines import logger
  from stable_baselines.common.mpi_adam import MpiAdam
  from stable_baselines.common.cg import conjugate_gradient
  from stable_baselines.common.policies import ActorCriticPolicy
  from stable_baselines.common.misc_util import flatten_lists
  from stable_baselines.common.runners import traj_segment_generator
  from stable_baselines.trpo_mpi.utils import add_vtarg_and_adv
  
  
  class TRPO(ActorCriticRLModel):
      """
      Trust Region Policy Optimization (https://arxiv.org/abs/1502.05477)
  
      :param policy: (ActorCriticPolicy or str) The policy model to use (MlpPolicy, CnnPolicy, CnnLstmPolicy, ...)
      :param env: (Gym environment or str) The environment to learn from (if registered in Gym, can be str)
      :param gamma: (float) the discount value
      :param timesteps_per_batch: (int) the number of timesteps to run per batch (horizon)
      :param max_kl: (float) the Kullback-Leibler loss threshold
      :param cg_iters: (int) the number of iterations for the conjugate gradient calculation
      :param lam: (float) GAE factor
      :param entcoeff: (float) the weight for the entropy loss
      :param cg_damping: (float) the compute gradient dampening factor
      :param vf_stepsize: (float) the value function stepsize
      :param vf_iters: (int) the value function's number iterations for learning
      :param verbose: (int) the verbosity level: 0 none, 1 training information, 2 tensorflow debug
      :param tensorboard_log: (str) the log location for tensorboard (if None, no logging)
      :param _init_setup_model: (bool) Whether or not to build the network at the creation of the instance
      :param policy_kwargs: (dict) additional arguments to be passed to the policy on creation
      :param full_tensorboard_log: (bool) enable additional logging when using tensorboard
          WARNING: this logging can take a lot of space quickly
      :param seed: (int) Seed for the pseudo-random generators (python, numpy, tensorflow).
          If None (default), use random seed. Note that if you want completely deterministic
          results, you must set `n_cpu_tf_sess` to 1.
      :param n_cpu_tf_sess: (int) The number of threads for TensorFlow operations
          If None, the number of cpu of the current machine will be used.
      """
      def __init__(self, policy, env, gamma=0.99, timesteps_per_batch=1024, max_kl=0.01, cg_iters=10, lam=0.98,
                   entcoeff=0.0, cg_damping=1e-2, vf_stepsize=3e-4, vf_iters=3, verbose=0, tensorboard_log=None,
                   _init_setup_model=True, policy_kwargs=None, full_tensorboard_log=False,
                   seed=None, n_cpu_tf_sess=1):
          super(TRPO, self).__init__(policy=policy, env=env, verbose=verbose, requires_vec_env=False,
                                     _init_setup_model=_init_setup_model, policy_kwargs=policy_kwargs,
                                     seed=seed, n_cpu_tf_sess=n_cpu_tf_sess)
  
          self.using_gail = False
          self.timesteps_per_batch = timesteps_per_batch
          self.cg_iters = cg_iters
          self.cg_damping = cg_damping
          self.gamma = gamma
          self.lam = lam
          self.max_kl = max_kl
          self.vf_iters = vf_iters
          self.vf_stepsize = vf_stepsize
          self.entcoeff = entcoeff
          self.tensorboard_log = tensorboard_log
          self.full_tensorboard_log = full_tensorboard_log
  
          # GAIL Params
          self.hidden_size_adversary = 100
          self.adversary_entcoeff = 1e-3
          self.expert_dataset = None
          self.g_step = 1
          self.d_step = 1
          self.d_stepsize = 3e-4
  
          self.graph = None
          self.sess = None
          self.policy_pi = None
          self.loss_names = None
          self.assign_old_eq_new = None
          self.compute_losses = None
          self.compute_lossandgrad = None
          self.compute_fvp = None
          self.compute_vflossandgrad = None
          self.d_adam = None
          self.vfadam = None
          self.get_flat = None
          self.set_from_flat = None
          self.timed = None
          self.allmean = None
          self.nworkers = None
          self.rank = None
          self.reward_giver = None
          self.step = None
          self.proba_step = None
          self.initial_state = None
          self.params = None
          self.summary = None
  
          if _init_setup_model:
              self.setup_model()
  
      def _get_pretrain_placeholders(self):
          policy = self.policy_pi
          action_ph = policy.pdtype.sample_placeholder([None])
          if isinstance(self.action_space, gym.spaces.Discrete):
              return policy.obs_ph, action_ph, policy.policy
          return policy.obs_ph, action_ph, policy.deterministic_action
  
      def setup_model(self):
          # prevent import loops
          from stable_baselines.gail.adversary import TransitionClassifier
  
          with SetVerbosity(self.verbose):
  
              assert issubclass(self.policy, ActorCriticPolicy), "Error: the input policy for the TRPO model must be " \
                                                                 "an instance of common.policies.ActorCriticPolicy."
  
              self.nworkers = MPI.COMM_WORLD.Get_size()
              self.rank = MPI.COMM_WORLD.Get_rank()
              np.set_printoptions(precision=3)
  
              self.graph = tf.Graph()
              with self.graph.as_default():
                  self.set_random_seed(self.seed)
                  self.sess = tf_util.make_session(num_cpu=self.n_cpu_tf_sess, graph=self.graph)
  
                  if self.using_gail:
                      self.reward_giver = TransitionClassifier(self.observation_space, self.action_space,
                                                               self.hidden_size_adversary,
                                                               entcoeff=self.adversary_entcoeff)
  
                  # Construct network for new policy
                  self.policy_pi = self.policy(self.sess, self.observation_space, self.action_space, self.n_envs, 1,
                                               None, reuse=False, **self.policy_kwargs)
  
                  # Network for old policy
                  with tf.variable_scope("oldpi", reuse=False):
                      old_policy = self.policy(self.sess, self.observation_space, self.action_space, self.n_envs, 1,
                                               None, reuse=False, **self.policy_kwargs)
  
                  with tf.variable_scope("loss", reuse=False):
                      atarg = tf.placeholder(dtype=tf.float32, shape=[None])  # Target advantage function (if applicable)
                      ret = tf.placeholder(dtype=tf.float32, shape=[None])  # Empirical return
  
                      observation = self.policy_pi.obs_ph
                      action = self.policy_pi.pdtype.sample_placeholder([None])
  
                      kloldnew = old_policy.proba_distribution.kl(self.policy_pi.proba_distribution)
                      ent = self.policy_pi.proba_distribution.entropy()
                      meankl = tf.reduce_mean(kloldnew)
                      meanent = tf.reduce_mean(ent)
                      entbonus = self.entcoeff * meanent
  
                      vferr = tf.reduce_mean(tf.square(self.policy_pi.value_flat - ret))
  
                      # advantage * pnew / pold
                      ratio = tf.exp(self.policy_pi.proba_distribution.logp(action) -
                                     old_policy.proba_distribution.logp(action))
                      surrgain = tf.reduce_mean(ratio * atarg)
  
                      optimgain = surrgain + entbonus
                      losses = [optimgain, meankl, entbonus, surrgain, meanent]
                      self.loss_names = ["optimgain", "meankl", "entloss", "surrgain", "entropy"]
  
                      dist = meankl
  
                      all_var_list = tf_util.get_trainable_vars("model")
                      var_list = [v for v in all_var_list if "/vf" not in v.name and "/q/" not in v.name]
                      vf_var_list = [v for v in all_var_list if "/pi" not in v.name and "/logstd" not in v.name]
  
                      self.get_flat = tf_util.GetFlat(var_list, sess=self.sess)
                      self.set_from_flat = tf_util.SetFromFlat(var_list, sess=self.sess)
  
                      klgrads = tf.gradients(dist, var_list)
                      flat_tangent = tf.placeholder(dtype=tf.float32, shape=[None], name="flat_tan")
                      shapes = [var.get_shape().as_list() for var in var_list]
                      start = 0
                      tangents = []
                      for shape in shapes:
                          var_size = tf_util.intprod(shape)
                          tangents.append(tf.reshape(flat_tangent[start: start + var_size], shape))
                          start += var_size
                      gvp = tf.add_n([tf.reduce_sum(grad * tangent)
                                      for (grad, tangent) in zipsame(klgrads, tangents)])  # pylint: disable=E1111
                      # Fisher vector products
                      fvp = tf_util.flatgrad(gvp, var_list)
  
                      tf.summary.scalar('entropy_loss', meanent)
                      tf.summary.scalar('policy_gradient_loss', optimgain)
                      tf.summary.scalar('value_function_loss', surrgain)
                      tf.summary.scalar('approximate_kullback-leibler', meankl)
                      tf.summary.scalar('loss', optimgain + meankl + entbonus + surrgain + meanent)
  
                      self.assign_old_eq_new = \
                          tf_util.function([], [], updates=[tf.assign(oldv, newv) for (oldv, newv) in
                                                            zipsame(tf_util.get_globals_vars("oldpi"),
                                                                    tf_util.get_globals_vars("model"))])
                      self.compute_losses = tf_util.function([observation, old_policy.obs_ph, action, atarg], losses)
                      self.compute_fvp = tf_util.function([flat_tangent, observation, old_policy.obs_ph, action, atarg],
                                                          fvp)
                      self.compute_vflossandgrad = tf_util.function([observation, old_policy.obs_ph, ret],
                                                                    tf_util.flatgrad(vferr, vf_var_list))
  
                      @contextmanager
                      def timed(msg):
                          if self.rank == 0 and self.verbose >= 1:
                              print(colorize(msg, color='magenta'))
                              start_time = time.time()
                              yield
                              print(colorize("done in {:.3f} seconds".format((time.time() - start_time)),
                                             color='magenta'))
                          else:
                              yield
  
                      def allmean(arr):
                          assert isinstance(arr, np.ndarray)
                          out = np.empty_like(arr)
                          MPI.COMM_WORLD.Allreduce(arr, out, op=MPI.SUM)
                          out /= self.nworkers
                          return out
  
                      tf_util.initialize(sess=self.sess)
  
                      th_init = self.get_flat()
                      MPI.COMM_WORLD.Bcast(th_init, root=0)
                      self.set_from_flat(th_init)
  
                  with tf.variable_scope("Adam_mpi", reuse=False):
                      self.vfadam = MpiAdam(vf_var_list, sess=self.sess)
                      if self.using_gail:
                          self.d_adam = MpiAdam(self.reward_giver.get_trainable_variables(), sess=self.sess)
                          self.d_adam.sync()
                      self.vfadam.sync()
  
                  with tf.variable_scope("input_info", reuse=False):
                      tf.summary.scalar('discounted_rewards', tf.reduce_mean(ret))
                      tf.summary.scalar('learning_rate', tf.reduce_mean(self.vf_stepsize))
                      tf.summary.scalar('advantage', tf.reduce_mean(atarg))
                      tf.summary.scalar('kl_clip_range', tf.reduce_mean(self.max_kl))
  
                      if self.full_tensorboard_log:
                          tf.summary.histogram('discounted_rewards', ret)
                          tf.summary.histogram('learning_rate', self.vf_stepsize)
                          tf.summary.histogram('advantage', atarg)
                          tf.summary.histogram('kl_clip_range', self.max_kl)
                          if tf_util.is_image(self.observation_space):
                              tf.summary.image('observation', observation)
                          else:
                              tf.summary.histogram('observation', observation)
  
                  self.timed = timed
                  self.allmean = allmean
  
                  self.step = self.policy_pi.step
                  self.proba_step = self.policy_pi.proba_step
                  self.initial_state = self.policy_pi.initial_state
  
                  self.params = tf_util.get_trainable_vars("model") + tf_util.get_trainable_vars("oldpi")
                  if self.using_gail:
                      self.params.extend(self.reward_giver.get_trainable_variables())
  
                  self.summary = tf.summary.merge_all()
  
                  self.compute_lossandgrad = \
                      tf_util.function([observation, old_policy.obs_ph, action, atarg, ret],
                                       [self.summary, tf_util.flatgrad(optimgain, var_list)] + losses)
  
      def learn(self, total_timesteps, callback=None, log_interval=100, tb_log_name="TRPO",
                reset_num_timesteps=True):
  
          new_tb_log = self._init_num_timesteps(reset_num_timesteps)
          callback = self._init_callback(callback)
  
          with SetVerbosity(self.verbose), TensorboardWriter(self.graph, self.tensorboard_log, tb_log_name, new_tb_log) \
                  as writer:
              self._setup_learn()
  
              with self.sess.as_default():
                  callback.on_training_start(locals(), globals())
  
                  seg_gen = traj_segment_generator(self.policy_pi, self.env, self.timesteps_per_batch,
                                                   reward_giver=self.reward_giver,
                                                   gail=self.using_gail, callback=callback)
  
                  episodes_so_far = 0
                  timesteps_so_far = 0
                  iters_so_far = 0
                  t_start = time.time()
                  len_buffer = deque(maxlen=40)  # rolling buffer for episode lengths
                  reward_buffer = deque(maxlen=40)  # rolling buffer for episode rewards
  
                  true_reward_buffer = None
                  if self.using_gail:
                      true_reward_buffer = deque(maxlen=40)
  
                      # Initialize dataloader
                      batchsize = self.timesteps_per_batch // self.d_step
                      self.expert_dataset.init_dataloader(batchsize)
  
                      #  Stats not used for now
                      # TODO: replace with normal tb logging
                      #  g_loss_stats = Stats(loss_names)
                      #  d_loss_stats = Stats(reward_giver.loss_name)
                      #  ep_stats = Stats(["True_rewards", "Rewards", "Episode_length"])
  
                  while True:
                      if timesteps_so_far >= total_timesteps:
                          break
  
                      logger.log("********** Iteration %i ************" % iters_so_far)
  
                      def fisher_vector_product(vec):
                          return self.allmean(self.compute_fvp(vec, *fvpargs, sess=self.sess)) + self.cg_damping * vec
  
                      # ------------------ Update G ------------------
                      logger.log("Optimizing Policy...")
                      # g_step = 1 when not using GAIL
                      mean_losses = None
                      vpredbefore = None
                      tdlamret = None
                      observation = None
                      action = None
                      seg = None
                      for k in range(self.g_step):
                          with self.timed("sampling"):
                              seg = seg_gen.__next__()
  
                          # Stop training early (triggered by the callback)
                          if not seg.get('continue_training', True):  # pytype: disable=attribute-error
                              break
  
                          add_vtarg_and_adv(seg, self.gamma, self.lam)
                          # ob, ac, atarg, ret, td1ret = map(np.concatenate, (obs, acs, atargs, rets, td1rets))
                          observation, action = seg["observations"], seg["actions"]
                          atarg, tdlamret = seg["adv"], seg["tdlamret"]
  
  
                          vpredbefore = seg["vpred"]  # predicted value function before update
                          atarg = (atarg - atarg.mean()) / (atarg.std() + 1e-8)  # standardized advantage function estimate
  
                          # true_rew is the reward without discount
                          if writer is not None:
                              total_episode_reward_logger(self.episode_reward,
                                                          seg["true_rewards"].reshape(
                                                              (self.n_envs, -1)),
                                                          seg["dones"].reshape((self.n_envs, -1)),
                                                          writer, self.num_timesteps)
  
                          args = seg["observations"], seg["observations"], seg["actions"], atarg
                          # Subsampling: see p40-42 of John Schulman thesis
                          # http://joschu.net/docs/thesis.pdf
                          fvpargs = [arr[::5] for arr in args]
  
                          self.assign_old_eq_new(sess=self.sess)
  
                          with self.timed("computegrad"):
                              steps = self.num_timesteps + (k + 1) * (seg["total_timestep"] / self.g_step)
                              run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
                              run_metadata = tf.RunMetadata() if self.full_tensorboard_log else None
                              # run loss backprop with summary, and save the metadata (memory, compute time, ...)
                              if writer is not None:
                                  summary, grad, *lossbefore = self.compute_lossandgrad(*args, tdlamret, sess=self.sess,
                                                                                        options=run_options,
                                                                                        run_metadata=run_metadata)
                                  if self.full_tensorboard_log:
                                      writer.add_run_metadata(run_metadata, 'step%d' % steps)
                                  writer.add_summary(summary, steps)
                              else:
                                  _, grad, *lossbefore = self.compute_lossandgrad(*args, tdlamret, sess=self.sess,
                                                                                  options=run_options,
                                                                                  run_metadata=run_metadata)
  
                          lossbefore = self.allmean(np.array(lossbefore))
                          grad = self.allmean(grad)
                          if np.allclose(grad, 0):
                              logger.log("Got zero gradient. not updating")
                          else:
                              with self.timed("conjugate_gradient"):
                                  stepdir = conjugate_gradient(fisher_vector_product, grad, cg_iters=self.cg_iters,
                                                               verbose=self.rank == 0 and self.verbose >= 1)
                              assert np.isfinite(stepdir).all()
                              shs = .5 * stepdir.dot(fisher_vector_product(stepdir))
                              # abs(shs) to avoid taking square root of negative values
                              lagrange_multiplier = np.sqrt(abs(shs) / self.max_kl)
                              # logger.log("lagrange multiplier:", lm, "gnorm:", np.linalg.norm(g))
                              fullstep = stepdir / lagrange_multiplier
                              expectedimprove = grad.dot(fullstep)
                              surrbefore = lossbefore[0]
                              stepsize = 1.0
                              thbefore = self.get_flat()
                              for _ in range(10):
                                  thnew = thbefore + fullstep * stepsize
                                  self.set_from_flat(thnew)
                                  mean_losses = surr, kl_loss, *_ = self.allmean(
                                      np.array(self.compute_losses(*args, sess=self.sess)))
                                  improve = surr - surrbefore
                                  logger.log("Expected: %.3f Actual: %.3f" % (expectedimprove, improve))
                                  if not np.isfinite(mean_losses).all():
                                      logger.log("Got non-finite value of losses -- bad!")
                                  elif kl_loss > self.max_kl * 1.5:
                                      logger.log("violated KL constraint. shrinking step.")
                                  elif improve < 0:
                                      logger.log("surrogate didn't improve. shrinking step.")
                                  else:
                                      logger.log("Stepsize OK!")
                                      break
                                  stepsize *= .5
                              else:
                                  logger.log("couldn't compute a good step")
                                  self.set_from_flat(thbefore)
                              if self.nworkers > 1 and iters_so_far % 20 == 0:
                                  # list of tuples
                                  paramsums = MPI.COMM_WORLD.allgather((thnew.sum(), self.vfadam.getflat().sum()))
                                  assert all(np.allclose(ps, paramsums[0]) for ps in paramsums[1:])
  
                              for (loss_name, loss_val) in zip(self.loss_names, mean_losses):
                                  logger.record_tabular(loss_name, loss_val)
  
                          with self.timed("vf"):
                              for _ in range(self.vf_iters):
                                  # NOTE: for recurrent policies, use shuffle=False?
                                  for (mbob, mbret) in dataset.iterbatches((seg["observations"], seg["tdlamret"]),
                                                                           include_final_partial_batch=False,
                                                                           batch_size=128,
                                                                           shuffle=True):
                                      grad = self.allmean(self.compute_vflossandgrad(mbob, mbob, mbret, sess=self.sess))
                                      self.vfadam.update(grad, self.vf_stepsize)
  
  
                      # Stop training early (triggered by the callback)
                      if not seg.get('continue_training', True):  # pytype: disable=attribute-error
                          break
  
                      logger.record_tabular("explained_variance_tdlam_before",
                                            explained_variance(vpredbefore, tdlamret))
  
                      if self.using_gail:
                          # ------------------ Update D ------------------
                          logger.log("Optimizing Discriminator...")
                          logger.log(fmt_row(13, self.reward_giver.loss_name))
                          assert len(observation) == self.timesteps_per_batch
                          batch_size = self.timesteps_per_batch // self.d_step
  
                          # NOTE: uses only the last g step for observation
                          d_losses = []  # list of tuples, each of which gives the loss for a minibatch
                          # NOTE: for recurrent policies, use shuffle=False?
                          for ob_batch, ac_batch in dataset.iterbatches((observation, action),
                                                                        include_final_partial_batch=False,
                                                                        batch_size=batch_size,
                                                                        shuffle=True):
                              ob_expert, ac_expert = self.expert_dataset.get_next_batch()
                              # update running mean/std for reward_giver
                              if self.reward_giver.normalize:
                                  self.reward_giver.obs_rms.update(np.concatenate((ob_batch, ob_expert), 0))
  
                              # Reshape actions if needed when using discrete actions
                              if isinstance(self.action_space, gym.spaces.Discrete):
                                  if len(ac_batch.shape) == 2:
                                      ac_batch = ac_batch[:, 0]
                                  if len(ac_expert.shape) == 2:
                                      ac_expert = ac_expert[:, 0]
                              *newlosses, grad = self.reward_giver.lossandgrad(ob_batch, ac_batch, ob_expert, ac_expert)
                              self.d_adam.update(self.allmean(grad), self.d_stepsize)
                              d_losses.append(newlosses)
                          logger.log(fmt_row(13, np.mean(d_losses, axis=0)))
  
                          # lr: lengths and rewards
                          lr_local = (seg["ep_lens"], seg["ep_rets"], seg["ep_true_rets"])  # local values
                          list_lr_pairs = MPI.COMM_WORLD.allgather(lr_local)  # list of tuples
                          lens, rews, true_rets = map(flatten_lists, zip(*list_lr_pairs))
                          true_reward_buffer.extend(true_rets)
                      else:
                          # lr: lengths and rewards
                          lr_local = (seg["ep_lens"], seg["ep_rets"])  # local values
                          list_lr_pairs = MPI.COMM_WORLD.allgather(lr_local)  # list of tuples
                          lens, rews = map(flatten_lists, zip(*list_lr_pairs))
                      len_buffer.extend(lens)
                      reward_buffer.extend(rews)
  
                      if len(len_buffer) > 0:
                          logger.record_tabular("EpLenMean", np.mean(len_buffer))
                          logger.record_tabular("EpRewMean", np.mean(reward_buffer))
                      if self.using_gail:
                          logger.record_tabular("EpTrueRewMean", np.mean(true_reward_buffer))
                      logger.record_tabular("EpThisIter", len(lens))
                      episodes_so_far += len(lens)
                      current_it_timesteps = MPI.COMM_WORLD.allreduce(seg["total_timestep"])
                      timesteps_so_far += current_it_timesteps
                      self.num_timesteps += current_it_timesteps
                      iters_so_far += 1
  
                      logger.record_tabular("EpisodesSoFar", episodes_so_far)
                      logger.record_tabular("TimestepsSoFar", self.num_timesteps)
                      logger.record_tabular("TimeElapsed", time.time() - t_start)
  
                      if self.verbose >= 1 and self.rank == 0:
                          logger.dump_tabular()
  
          callback.on_training_end()
          return self
  
      def save(self, save_path, cloudpickle=False):
          if self.using_gail and self.expert_dataset is not None:
              # Exit processes to pickle the dataset
              self.expert_dataset.prepare_pickling()
          data = {
              "gamma": self.gamma,
              "timesteps_per_batch": self.timesteps_per_batch,
              "max_kl": self.max_kl,
              "cg_iters": self.cg_iters,
              "lam": self.lam,
              "entcoeff": self.entcoeff,
              "cg_damping": self.cg_damping,
              "vf_stepsize": self.vf_stepsize,
              "vf_iters": self.vf_iters,
              "hidden_size_adversary": self.hidden_size_adversary,
              "adversary_entcoeff": self.adversary_entcoeff,
              "expert_dataset": self.expert_dataset,
              "g_step": self.g_step,
              "d_step": self.d_step,
              "d_stepsize": self.d_stepsize,
              "using_gail": self.using_gail,
              "verbose": self.verbose,
              "policy": self.policy,
              "observation_space": self.observation_space,
              "action_space": self.action_space,
              "n_envs": self.n_envs,
              "n_cpu_tf_sess": self.n_cpu_tf_sess,
              "seed": self.seed,
              "_vectorize_action": self._vectorize_action,
              "policy_kwargs": self.policy_kwargs
          }
  
          params_to_save = self.get_parameters()
  
          self._save_to_file(save_path, data=data, params=params_to_save, cloudpickle=cloudpickle)
   import tensorflow as tf
  import numpy as np
  from gym.spaces import Box
  import copy
  from stable_baselines.common.policies import BasePolicy, nature_cnn, register_policy, cnn_1d_extractor
  from stable_baselines.sac.policies import mlp
  from stable_baselines.a2c.utils import lstm, batch_to_seq, seq_to_batch
  
  
  class TD3Policy(BasePolicy):
      """
      Policy object that implements a TD3-like actor critic
  
      :param sess: (TensorFlow session) The current TensorFlow session
      :param ob_space: (Gym Space) The observation space of the environment
      :param ac_space: (Gym Space) The action space of the environment
      :param n_env: (int) The number of environments to run
      :param n_steps: (int) The number of steps to run for each environment
      :param n_batch: (int) The number of batch to run (n_envs * n_steps)
      :param reuse: (bool) If the policy is reusable or not
      :param scale: (bool) whether or not to scale the input
      """
  
      def __init__(self, sess, ob_space, ac_space, n_env=1, n_steps=1, n_batch=None, reuse=False, scale=False,
                   add_action_ph=False):
          super(TD3Policy, self).__init__(sess, ob_space, ac_space, n_env, n_steps, n_batch, reuse=reuse, scale=scale,
                                          add_action_ph=add_action_ph)
          assert isinstance(ac_space, Box), "Error: the action space must be of type gym.spaces.Box"
  
          self.qf1 = None
          self.qf2 = None
          self.q_discrepancy = None
          self.policy = None
  
      def make_actor(self, obs=None, reuse=False, scope="pi"):
          """
          Creates an actor object
  
          :param obs: (TensorFlow Tensor) The observation placeholder (can be None for default placeholder)
          :param reuse: (bool) whether or not to reuse parameters
          :param scope: (str) the scope name of the actor
          :return: (TensorFlow Tensor) the output tensor
          """
          raise NotImplementedError
  
      def make_critics(self, obs=None, action=None, reuse=False,
                       scope="qvalues_fn"):
          """
          Creates the two Q-Values approximator
  
          :param obs: (TensorFlow Tensor) The observation placeholder (can be None for default placeholder)
          :param action: (TensorFlow Tensor) The action placeholder
          :param reuse: (bool) whether or not to reuse parameters
          :param scope: (str) the scope name
          :return: ([tf.Tensor]) Mean, action and log probability
          """
          raise NotImplementedError
  
      def step(self, obs, state=None, mask=None):
          """
          Returns the policy for a single step
  
          :param obs: ([float] or [int]) The current observation of the environment
          :param state: ([float]) The last states (used in recurrent policies)
          :param mask: ([float]) The last masks (used in recurrent policies)
          :return: ([float]) actions
          """
          raise NotImplementedError
  
      def proba_step(self, obs, state=None, mask=None):
          """
          Returns the policy for a single step
  
          :param obs: ([float] or [int]) The current observation of the environment
          :param state: ([float]) The last states (used in recurrent policies)
          :param mask: ([float]) The last masks (used in recurrent policies)
          :return: ([float]) actions
          """
          return self.step(obs, state, mask)
  
  
  class FeedForwardPolicy(TD3Policy):
      """
      Policy object that implements a DDPG-like actor critic, using a feed forward neural network.
  
      :param sess: (TensorFlow session) The current TensorFlow session
      :param ob_space: (Gym Space) The observation space of the environment
      :param ac_space: (Gym Space) The action space of the environment
      :param n_env: (int) The number of environments to run
      :param n_steps: (int) The number of steps to run for each environment
      :param n_batch: (int) The number of batch to run (n_envs * n_steps)
      :param reuse: (bool) If the policy is reusable or not
      :param layers: ([int]) The size of the Neural network for the policy (if None, default to [64, 64])
      :param cnn_extractor: (function (TensorFlow Tensor, ``**kwargs``): (TensorFlow Tensor)) the CNN feature extraction
      :param feature_extraction: (str) The feature extraction type ("cnn" or "mlp")
      :param layer_norm: (bool) enable layer normalisation
      :param act_fun: (tf.func) the activation function to use in the neural network.
      :param kwargs: (dict) Extra keyword arguments for the nature CNN feature extraction
      """
  
      def __init__(self, sess, ob_space, ac_space, n_env=1, n_steps=1, n_batch=None, reuse=False, layers=None,
                   cnn_extractor=nature_cnn, feature_extraction="cnn",
                   layer_norm=False, act_fun=tf.nn.relu, obs_module_indices=None, **kwargs):
          super(FeedForwardPolicy, self).__init__(sess, ob_space, ac_space, n_env, n_steps, n_batch,
                                                  reuse=reuse,
                                                  scale=(feature_extraction == "cnn" and cnn_extractor == nature_cnn))
  
          self._kwargs_check(feature_extraction, kwargs)
          self.layer_norm = layer_norm
          self.feature_extraction = feature_extraction
          self.cnn_kwargs = kwargs
          self.cnn_extractor = cnn_extractor
          self.cnn_vf = self.cnn_kwargs.pop("cnn_vf", True)
          self.reuse = reuse
          if layers is None:
              layers = [64, 64]
          self.layers = layers
          self.obs_module_indices = obs_module_indices
          self.policy_pre_activation = None
  
          assert len(layers) >= 1, "Error: must have at least one hidden layer for the policy."
  
          self.activ_fn = act_fun
  
      def make_actor(self, obs=None, reuse=False, scope="pi"):
          if obs is None:
              obs = self.processed_obs
  
          if self.obs_module_indices is not None:
              obs = tf.gather(obs, self.obs_module_indices["pi"], axis=-1)
          with tf.variable_scope(scope, reuse=reuse):
              if self.feature_extraction == "cnn":
                  pi_h = self.cnn_extractor(obs, name="pi_c1", act_fun=self.activ_fn, **self.cnn_kwargs)
              else:
                  pi_h = tf.layers.flatten(obs)
  
              pi_h = mlp(pi_h, self.layers, self.activ_fn, layer_norm=self.layer_norm)
  
              self.policy_pre_activation = tf.layers.dense(pi_h, self.ac_space.shape[0])
              self.policy = policy = tf.tanh(self.policy_pre_activation)
  
          return policy
  
      def make_critics(self, obs=None, action=None, reuse=False, scope="values_fn", extracted_callback=None):
          if obs is None:
              obs = self.processed_obs
  
          if self.obs_module_indices is not None:
              obs = tf.gather(obs, self.obs_module_indices["vf"], axis=-1)
  
          with tf.variable_scope(scope, reuse=reuse):
              if self.feature_extraction == "cnn" and self.cnn_vf:
                  critics_h = self.cnn_extractor(obs, name="vf_c1", act_fun=self.activ_fn, **self.cnn_kwargs)
              else:
                  critics_h = tf.layers.flatten(obs)
  
              if extracted_callback is not None:
                  critics_h = extracted_callback(critics_h)
  
              # Concatenate preprocessed state and action
              qf_h = tf.concat([critics_h, action], axis=-1)
  
              # Double Q values to reduce overestimation
              with tf.variable_scope('qf1', reuse=reuse):
                  qf1_h = mlp(qf_h, self.layers, self.activ_fn, layer_norm=self.layer_norm)
                  qf1 = tf.layers.dense(qf1_h, 1, name="qf1")
  
              with tf.variable_scope('qf2', reuse=reuse):
                  qf2_h = mlp(qf_h, self.layers, self.activ_fn, layer_norm=self.layer_norm)
                  qf2 = tf.layers.dense(qf2_h, 1, name="qf2")
  
              self.qf1 = qf1
              self.qf2 = qf2
              # TODO: assumes that all qf1 and qf2 can never have opposite signs
              #self.q_discrepancy = tf.square(self.qf1 - self.qf2) / tf.square(tf.maximum(self.qf1, self.qf2))
              #self.q_discrepancy = tf.abs(self.qf1 - self.qf2)
  
          return self.qf1, self.qf2
  
      def step(self, obs, state=None, mask=None):
          return self.sess.run(self.policy, {self.obs_ph: obs})
  
      def get_q_discrepancy(self, obs):
          if isinstance(obs, np.ndarray) and len(obs.shape) == 1: # TODO: check for MLP or CNN policy here
              obs = np.expand_dims(obs, axis=0)
          return self.sess.run(self.q_discrepancy, {self.obs_ph: obs})
  
  
  class RecurrentPolicy(TD3Policy):
      """
          Policy object that implements a DDPG-like actor critic, using a feed forward neural network.
  
          :param sess: (TensorFlow session) The current TensorFlow session
          :param ob_space: (Gym Space) The observation space of the environment
          :param ac_space: (Gym Space) The action space of the environment
          :param n_env: (int) The number of environments to run
          :param n_steps: (int) The number of steps to run for each environment
          :param n_batch: (int) The number of batch to run (n_envs * n_steps)
          :param reuse: (bool) If the policy is reusable or not
          :param layers: ([int]) The size of the Neural network for the policy (if None, default to [64, 64])
          :param cnn_extractor: (function (TensorFlow Tensor, ``**kwargs``): (TensorFlow Tensor)) the CNN feature extraction
          :param feature_extraction: (str) The feature extraction type ("cnn" or "mlp")
          :param layer_norm: (bool) enable layer normalisation
          :param act_fun: (tf.func) the activation function to use in the neural network.
          :param kwargs: (dict) Extra keyword arguments for the nature CNN feature extraction
          """
      recurrent = True
  
      def __init__(self, sess, ob_space, ac_space, layers, n_env=1, n_steps=1, n_batch=None, reuse=False,
                   cnn_extractor=nature_cnn, feature_extraction="mlp", n_lstm=128, share_lstm=False, save_state=False,
                   save_target_state=False, layer_norm=False, act_fun=tf.nn.relu, obs_module_indices=None, **kwargs):
          super(RecurrentPolicy, self).__init__(sess, ob_space, ac_space, n_env, n_steps, n_batch,
                                                  reuse=reuse, add_action_ph=True,
                                                  scale=(feature_extraction == "cnn" and cnn_extractor == nature_cnn))
  
          self._kwargs_check(feature_extraction, kwargs)
          self.layer_norm = layer_norm
          self.feature_extraction = feature_extraction
          self.cnn_kwargs = kwargs
          self.cnn_extractor = cnn_extractor
          self.cnn_vf = self.cnn_kwargs.pop("cnn_vf", True)
          self.reuse = reuse
          self.layers = layers
          self.obs_module_indices = obs_module_indices
  
          self.activ_fn = act_fun
          self.n_lstm = n_lstm
          self.share_lstm = share_lstm
          self._obs_ph = self.processed_obs  # Base class has self.obs_ph as property getting self._obs_ph
          self.obs_tp1_ph = self.processed_obs
  
          assert self.n_batch % self.n_steps == 0, "The batch size must be a multiple of sequence length (n_steps)"
          self._lstm_n_batch = self.n_batch // self.n_steps
  
          self.action_prev = np.zeros((1, *self.ac_space.shape))
  
          self._initial_state = np.zeros((self._lstm_n_batch, self.n_lstm * 2), dtype=np.float32)
          if self.share_lstm:
              self.state = None
          else:
              self.pi_state = None
              self.qf1_state = None
              self.qf2_state = None
  
          with tf.variable_scope("input", reuse=False):
              self.dones_ph = tf.placeholder_with_default(np.zeros((self.n_batch,), dtype=np.float32), (self.n_batch,), name="dones_ph")  # (done t-1)
              if self.share_lstm:
                  self.state_ph = tf.placeholder_with_default(self.initial_state, (self._lstm_n_batch, self.n_lstm * 2), name="state_ph")
              else:
                  self.pi_state_ph = tf.placeholder_with_default(self.initial_state, (self._lstm_n_batch, self.n_lstm * 2), name="pi_state_ph")
                  self.qf1_state_ph = tf.placeholder_with_default(self.initial_state, (self._lstm_n_batch, self.n_lstm * 2), name="qf1_state_ph")
                  self.qf2_state_ph = tf.placeholder_with_default(self.initial_state, (self._lstm_n_batch, self.n_lstm * 2), name="qf2_state_ph")
  
              self.action_prev_ph = tf.placeholder(np.float32, (self.n_batch, *self.ac_space.shape), name="action_prev_ph")
  
          self.save_state = save_state
          self.save_target_state = save_target_state
  
          self.extra_phs = ["action_prev"]
          self.rnn_inputs = ["obs", "action_prev"]
          self.extra_data_names = ["action_prev"]
  
          if self.save_target_state:
              self.extra_data_names = sorted(self.extra_data_names + ["target_action_prev"])
              self.rnn_inputs = sorted(self.rnn_inputs + ["obs_tp1"])
              self.extra_phs = sorted(self.extra_phs + ["target_action_prev"])
  
          if self.save_state:
              state_names = ["state"] if self.share_lstm else ["pi_state", "qf1_state", "qf2_state"]
              if self.save_target_state:
                  state_names.extend(["target_" + state_name for state_name in state_names])
              if self.share_lstm:
                  self.extra_data_names = sorted(self.extra_data_names + state_names)
                  self.extra_phs = sorted(self.extra_phs + state_names)
              else:
                  self.extra_data_names = sorted(self.extra_data_names + state_names)
                  self.extra_phs = sorted(self.extra_phs + state_names)
  
      def _process_phs(self, **phs):
          for ph_name, ph_val in phs.items():
              if ph_val is None:
                  phs[ph_name] = getattr(self, ph_name + "_ph")
              else:
                  try:
                      setattr(self, ph_name + "_ph", ph_val)
                  except AttributeError:
                      setattr(self, "_" + ph_name + "_ph", ph_val)
  
          return phs.values()
  
      def _make_branch(self, branch_name, input_tensor, dones=None, state_ph=None):
          if branch_name == "lstm":
              for i, fc_layer_units in enumerate(self.layers["lstm"]):
                  input_tensor = self.activ_fn(tf.layers.dense(input_tensor, fc_layer_units, name="lstm_fc{}".format(i)))
  
              input_tensor = batch_to_seq(input_tensor, self._lstm_n_batch, self.n_steps)
              masks = batch_to_seq(dones, self._lstm_n_batch, self.n_steps)
              input_tensor, state = lstm(input_tensor, masks, state_ph, "lstm", n_hidden=self.n_lstm,
                                         layer_norm=self.layer_norm)
              input_tensor = seq_to_batch(input_tensor)
  
              return input_tensor, state
          else:
              for i, fc_layer_units in enumerate(self.layers[branch_name]):
                  input_tensor = self.activ_fn(tf.layers.dense(input_tensor, fc_layer_units, name="{}_fc{}".format(branch_name, i)))
  
              return input_tensor
  
      def make_actor(self, ff_phs=None, rnn_phs=None, dones=None, reuse=False, scope="pi"):
          lstm_branch = tf.concat([tf.layers.flatten(ph) for ph in rnn_phs], axis=-1)
          if ff_phs is not None:
              ff_branch = tf.concat([tf.layers.flatten(ph) for ph in ff_phs], axis=-1)
  
          if dones is None:
              dones = self.dones_ph
  
          if self.share_lstm:
              with tf.variable_scope("shared", reuse=tf.AUTO_REUSE):
                  lstm_branch, self.state = self._make_branch("lstm", lstm_branch, dones, self.state_ph)
  
          with tf.variable_scope(scope, reuse=reuse):
              if self.layers["ff"] is not None:
                 ff_branch = self._make_branch("ff", ff_branch)
  
              if not self.share_lstm:
                  lstm_branch, self.pi_state = self._make_branch("lstm", lstm_branch, dones, self.pi_state_ph)
  
              if ff_phs is not None:
                  head = tf.concat([ff_branch, lstm_branch], axis=-1)
              else:
                  head = lstm_branch
  
              head = self._make_branch("head", head)
  
              self.policy_pre_activation = tf.layers.dense(head, self.ac_space.shape[0])
              self.policy = policy = tf.tanh(self.policy_pre_activation)
  
          return policy
  
      def make_critics(self, ff_phs=None, rnn_phs=None, dones=None, reuse=False, scope="values_fn"):
          lstm_branch_in = tf.concat([tf.layers.flatten(ph) for ph in rnn_phs], axis=-1)
          if ff_phs is not None:
              ff_branch_in = tf.concat([tf.layers.flatten(ph) for ph in ff_phs], axis=-1)
  
          if dones is None:
              dones = self.dones_ph
  
          self.qf1, self.qf2 = None, None
          self.qf1_state, self.qf2_state = None, None
  
          if self.share_lstm:
              with tf.variable_scope("shared", reuse=tf.AUTO_REUSE):
                  lstm_branch_s, self.state = self._make_branch("lstm", lstm_branch_in, dones, self.state_ph)
  
          with tf.variable_scope(scope, reuse=reuse):
              # Double Q values to reduce overestimation
              for qf_i in range(1, 3):
                  with tf.variable_scope('qf{}'.format(qf_i), reuse=reuse):
                      lstm_branch = lstm_branch_in
                      if self.layers["ff"] is not None:
                          ff_branch = self._make_branch("ff", ff_branch_in)
                      elif ff_phs is not None:
                          ff_branch = ff_branch_in
  
                      if not self.share_lstm:
                          lstm_branch, state = self._make_branch("lstm", lstm_branch, dones,
                                                                 getattr(self, "qf{}_state_ph".format(qf_i)))
                          setattr(self, "qf{}_state".format(qf_i), state)
                      else:
                          lstm_branch = lstm_branch_s
  
                      if ff_phs is not None:
                          head = tf.concat([ff_branch, lstm_branch], axis=-1)
                      else:
                          head = lstm_branch
  
                      head = self._make_branch("head", head)
  
                      setattr(self, "qf{}".format(qf_i), tf.layers.dense(head, 1, name="qf{}".format(qf_i)))
  
          return self.qf1, self.qf2
  
      def step(self, obs, action_prev=None, state=None, mask=None, feed_dict=None, **kwargs):
          if feed_dict is None:
              feed_dict = {}
          if state is None:
              state = self.initial_state
          if mask is None:
              mask = np.array([False])
          if action_prev is None:
              assert obs.shape[0] == 1
              if mask[0]:
                  self.action_prev = np.zeros((1, *self.ac_space.shape))
              action_prev = self.action_prev
  
          rnn_node = self.state if self.share_lstm else self.pi_state
          state_ph = self.state_ph if self.share_lstm else self.pi_state_ph
  
          feed_dict.update({self.obs_ph: obs, state_ph: state, self.dones_ph: mask,
                                        self.action_prev_ph: action_prev})
  
          action, out_state = self.sess.run([self.policy, rnn_node], feed_dict)
          self.action_prev = action
  
          return action, out_state
  
      @property
      def initial_state(self):
          return self._initial_state
  
      def collect_data(self, _locals, _globals):
          data = {}
          if self.save_state:
              if self.share_lstm:
                  data["state"] = _locals["prev_policy_state"][0, :]
              else:
                  data["pi_state"] = _locals["prev_policy_state"][0, :]
                  if len(_locals["episode_data"]) == 0:
                      qf1_state, qf2_state = self.initial_state, self.initial_state
                  else:
                      qf_feed_dict = {
                          self.qf1_state_ph: _locals["episode_data"][-1]["qf1_state"][None],
                          self.qf2_state_ph: _locals["episode_data"][-1]["qf2_state"][None],
                      }
                      qf_feed_dict.update({getattr(self, data_name + "_ph"): _locals["episode_data"][-1][data_name][None]
                                           for data_name in self.rnn_inputs})
                      qf1_state, qf2_state = self.sess.run([self.qf1_state, self.qf2_state], feed_dict=qf_feed_dict)
                  data["qf1_state"] = qf1_state[0, :]
                  data["qf2_state"] = qf2_state[0, :]
  
          if len(_locals["episode_data"]) == 0:
              data["action_prev"] = np.zeros(*self.ac_space.shape, dtype=np.float32)
          else:
              data["action_prev"] = _locals["episode_data"][-1]["action"]
  
          if self.save_target_state:
              data["target_action_prev_rnn"] = _locals["action"]
  
          return data
  
  
  class DRPolicy(RecurrentPolicy):
      """
          Policy object that implements a DDPG-like actor critic, using a feed forward neural network.
  
          :param sess: (TensorFlow session) The current TensorFlow session
          :param ob_space: (Gym Space) The observation space of the environment
          :param ac_space: (Gym Space) The action space of the environment
          :param n_env: (int) The number of environments to run
          :param n_steps: (int) The number of steps to run for each environment
          :param n_batch: (int) The number of batch to run (n_envs * n_steps)
          :param reuse: (bool) If the policy is reusable or not
          :param layers: ([int]) The size of the Neural network for the policy (if None, default to [64, 64])
          :param cnn_extractor: (function (TensorFlow Tensor, ``**kwargs``): (TensorFlow Tensor)) the CNN feature extraction
          :param feature_extraction: (str) The feature extraction type ("cnn" or "mlp")
          :param layer_norm: (bool) enable layer normalisation
          :param act_fun: (tf.func) the activation function to use in the neural network.
          :param kwargs: (dict) Extra keyword arguments for the nature CNN feature extraction
          """
      recurrent = True
  
      def __init__(self, sess, ob_space, ac_space, goal_size, my_size, n_env=1, n_steps=1, n_batch=None, reuse=False, layers=None,
                   cnn_extractor=nature_cnn, feature_extraction="mlp", n_lstm=128, share_lstm=False,
                   layer_norm=False, act_fun=tf.nn.relu, obs_module_indices=None, **kwargs):
          if layers is None:
              layers = {"ff": [128], "lstm": [128], "head": [128, 128]}
          super().__init__(sess, ob_space, ac_space, layers, n_env, n_steps, n_batch,
                                                  reuse=reuse, cnn_extractor=cnn_extractor,
                                                  feature_extraction=feature_extraction, n_lstm=n_lstm,
                                                  share_lstm=share_lstm, layer_norm=layer_norm, act_fun=act_fun,
                                                  obs_module_indices=obs_module_indices, **kwargs)
  
          with tf.variable_scope("input", reuse=False):
              self.my_ph = tf.placeholder(tf.float32, (None, my_size), name="my_ph")  # the dynamics of the environment
  
          self.goal_size = goal_size
          self.extra_phs = sorted(self.extra_phs + ["my"])
          self.extra_data_names = sorted(self.extra_data_names + ["my"])
  
      def make_actor(self, obs_ff=None, obs_rnn=None, action_prev=None, dones=None, reuse=False, scope="pi"):
          if obs_ff is None:
              obs_ff = self.processed_obs
          if obs_rnn is None:
              obs_rnn = self.processed_obs
          if action_prev is None:
              action_prev = self.action_prev_ph
  
          obs_ff, goal = obs_ff[:, :-self.goal_size], obs_ff[:, -self.goal_size:]
          goal = tf.subtract(goal, obs_ff[:, -self.goal_size:], name="goal_relative")
          obs_rnn = obs_rnn[:, :-self.goal_size]
  
          ff_phs = [obs_ff, goal]
          rnn_phs = [obs_rnn, action_prev]
          return super().make_actor(ff_phs=ff_phs, rnn_phs=rnn_phs, dones=dones, reuse=reuse, scope=scope)
  
      def make_critics(self, obs_ff=None, action_ff=None, my=None, obs_rnn=None, action_prev=None, dones=None, reuse=False, scope="values_fn"):
          if obs_ff is None:
              obs_ff = self.processed_obs
          if action_ff is None:
              action_ff = self.action_ph
          if my is None:
              my = self.my_ph
          if obs_rnn is None:
              obs_rnn = self.processed_obs
          if action_prev is None:
              action_prev = self.action_prev_ph
  
          obs_ff, goal = obs_ff[:, :-self.goal_size], obs_ff[:, -self.goal_size:]
          goal = tf.subtract(goal, obs_ff[:, -self.goal_size:], name="goal_relative")
          obs_rnn = obs_rnn[:, :-self.goal_size]
  
          ff_phs = [obs_ff, goal, my, action_ff]
          rnn_phs = [obs_rnn, action_prev]
          return super().make_critics(ff_phs=ff_phs, rnn_phs=rnn_phs, dones=dones, reuse=reuse, scope=scope)
  
      def collect_data(self, _locals, _globals, **kwargs):
          data = super().collect_data(_locals, _globals)
          if "my" not in _locals or _locals["episode_data"]:
              data["my"] = _locals["self"].env.get_env_parameters()
  
          return data
  
  
  class LstmMlpPolicy(RecurrentPolicy):
      recurrent = True
  
      def __init__(self, sess, ob_space, ac_space, n_env=1, n_steps=1, n_batch=None, reuse=False,
                   layers=None,
                   cnn_extractor=nature_cnn, feature_extraction="mlp", n_lstm=128, share_lstm=False,
                   layer_norm=False, act_fun=tf.nn.relu, obs_module_indices=None, **kwargs):
          if layers is None:
              layers = {"ff": None, "lstm": [64, 64], "head": []}
          else:
              assert layers["ff"] is None
          super().__init__(sess, ob_space, ac_space, layers, n_env, n_steps, n_batch,
                           reuse=reuse, cnn_extractor=cnn_extractor,
                           feature_extraction=feature_extraction, n_lstm=n_lstm,
                           share_lstm=share_lstm, layer_norm=layer_norm, act_fun=act_fun,
                           obs_module_indices=obs_module_indices, **kwargs)
  
      def make_actor(self, obs=None, action_prev=None, dones=None, reuse=False, scope="pi"):
          obs, action_prev, dones = self._process_phs(obs=obs, action_prev=action_prev, dones=dones)
  
          ff_phs = None
          rnn_phs = [obs, action_prev]
          return super().make_actor(ff_phs=ff_phs, rnn_phs=rnn_phs, dones=dones, reuse=reuse, scope=scope)
  
      def make_critics(self, obs=None, action=None, action_prev=None, dones=None, reuse=False, scope="values_fn"):
          obs, action, action_prev, dones = self._process_phs(obs=obs, action=action, action_prev=action_prev, dones=dones)
  
          ff_phs = [action]
          rnn_phs = [obs, action_prev]
          return super().make_critics(ff_phs=ff_phs, rnn_phs=rnn_phs, dones=dones, reuse=reuse, scope=scope)
  
  
  class LstmFFMlpPolicy(RecurrentPolicy):
      recurrent = True
  
      def __init__(self, sess, ob_space, ac_space, n_env=1, n_steps=1, n_batch=None, reuse=False,
                   layers=None,
                   cnn_extractor=nature_cnn, feature_extraction="mlp", n_lstm=128, share_lstm=False,
                   layer_norm=False, act_fun=tf.nn.relu, obs_module_indices=None, **kwargs):
          if layers is None:
              layers = {"ff": [64], "lstm": [64, 64], "head": []}
  
          super().__init__(sess, ob_space, ac_space, layers, n_env, n_steps, n_batch,
                           reuse=reuse, cnn_extractor=cnn_extractor,
                           feature_extraction=feature_extraction, n_lstm=n_lstm,
                           share_lstm=share_lstm, layer_norm=layer_norm, act_fun=act_fun,
                           obs_module_indices=obs_module_indices, **kwargs)
  
      def make_actor(self, obs=None, action_prev=None, dones=None, reuse=False, scope="pi"):
          obs, action_prev, dones = self._process_phs(obs=obs, action_prev=action_prev, dones=dones)
  
          ff_phs = [obs]
          rnn_phs = [obs, action_prev]
          return super().make_actor(ff_phs=ff_phs, rnn_phs=rnn_phs, dones=dones, reuse=reuse, scope=scope)
  
      def make_critics(self, obs=None, action=None, action_prev=None, dones=None, reuse=False, scope="values_fn"):
          obs, action, action_prev, dones = self._process_phs(obs=obs, action=action, action_prev=action_prev, dones=dones)
  
          ff_phs = [obs, action]
          rnn_phs = [obs, action_prev]
          return super().make_critics(ff_phs=ff_phs, rnn_phs=rnn_phs, dones=dones, reuse=reuse, scope=scope)
  
  
  class CnnPolicy(FeedForwardPolicy):
      """
      Policy object that implements actor critic, using a CNN (the nature CNN)
  
      :param sess: (TensorFlow session) The current TensorFlow session
      :param ob_space: (Gym Space) The observation space of the environment
      :param ac_space: (Gym Space) The action space of the environment
      :param n_env: (int) The number of environments to run
      :param n_steps: (int) The number of steps to run for each environment
      :param n_batch: (int) The number of batch to run (n_envs * n_steps)
      :param reuse: (bool) If the policy is reusable or not
      :param _kwargs: (dict) Extra keyword arguments for the nature CNN feature extraction
      """
  
      def __init__(self, sess, ob_space, ac_space, n_env=1, n_steps=1, n_batch=None, reuse=False, **_kwargs):
          super(CnnPolicy, self).__init__(sess, ob_space, ac_space, n_env, n_steps, n_batch, reuse,
                                          feature_extraction="cnn", **_kwargs)
  
  
  class CnnMlpPolicy(FeedForwardPolicy):
      """
      Policy object that implements actor critic, using a CNN (the nature CNN)
  
      :param sess: (TensorFlow session) The current TensorFlow session
      :param ob_space: (Gym Space) The observation space of the environment
      :param ac_space: (Gym Space) The action space of the environment
      :param n_env: (int) The number of environments to run
      :param n_steps: (int) The number of steps to run for each environment
      :param n_batch: (int) The number of batch to run (n_envs * n_steps)
      :param reuse: (bool) If the policy is reusable or not
      :param _kwargs: (dict) Extra keyword arguments for the nature CNN feature extraction
      """
  
      def __init__(self, sess, ob_space, ac_space, n_env=1, n_steps=1, n_batch=None, reuse=False, **_kwargs):
          super(CnnMlpPolicy, self).__init__(sess, ob_space, ac_space, n_env, n_steps, n_batch, reuse,
                                             cnn_extractor=cnn_1d_extractor, feature_extraction="cnn", **_kwargs)
  
  
  class DRCnnMlpPolicy(FeedForwardPolicy):
      """
      Policy object that implements actor critic, using a CNN (the nature CNN)
  
      :param sess: (TensorFlow session) The current TensorFlow session
      :param ob_space: (Gym Space) The observation space of the environment
      :param ac_space: (Gym Space) The action space of the environment
      :param n_env: (int) The number of environments to run
      :param n_steps: (int) The number of steps to run for each environment
      :param n_batch: (int) The number of batch to run (n_envs * n_steps)
      :param reuse: (bool) If the policy is reusable or not
      :param _kwargs: (dict) Extra keyword arguments for the nature CNN feature extraction
      """
  
      def __init__(self, sess, ob_space, ac_space, my_size, n_env=1, n_steps=1, n_batch=None, reuse=False, **_kwargs):
          super(DRCnnMlpPolicy, self).__init__(sess, ob_space, ac_space, n_env, n_steps, n_batch, reuse,
                                             cnn_extractor=cnn_1d_extractor, feature_extraction="cnn", **_kwargs)
  
          with tf.variable_scope("input", reuse=False):
              self.my_ph = tf.placeholder(tf.float32, (self.n_batch, *my_size), name="my_ph")  # (done t-1)
          self.extra_phs = ["my", "target_my"]
          self.extra_data_names = ["my", "target_my"]
  
      def make_critics(self, obs=None, action=None, my=None, reuse=False, scope="values_fn"):
          if my is None:
              my = self.my_ph
  
          return super().make_critics(obs, action, reuse, scope, extracted_callback=lambda x: tf.concat([x, my], axis=-1))
  
      def collect_data(self, _locals, _globals):
          data = []
          for env_i in range(_locals["self"].n_envs):
              d = {}
              if len(_locals["episode_data"][env_i]) == 0 or "my" not in _locals["episode_data"][env_i]:
                  if _locals["self"].n_envs == 1:
                      d["my"] = _locals["self"].env.get_env_parameters()
                  else:
                      d["my"] = _locals["self"].env.env_method("get_env_parameters", indices=env_i)[0]
              else:
                  d["my"] = _locals["episode_data"][env_i][-1]["my"]
  
              d["target_my"] = d["my"]
              data.append(d)
  
          return data
  
  
  class DRMyEstPolicy(FeedForwardPolicy):
      """
      Policy object that implements actor critic, using a CNN (the nature CNN)
  
      :param sess: (TensorFlow session) The current TensorFlow session
      :param ob_space: (Gym Space) The observation space of the environment
      :param ac_space: (Gym Space) The action space of the environment
      :param n_env: (int) The number of environments to run
      :param n_steps: (int) The number of steps to run for each environment
      :param n_batch: (int) The number of batch to run (n_envs * n_steps)
      :param reuse: (bool) If the policy is reusable or not
      :param _kwargs: (dict) Extra keyword arguments for the nature CNN feature extraction
      """
  
      def __init__(self, sess, ob_space, ac_space, my_size, n_env=1, n_steps=1, n_batch=None, reuse=False, loss_weight=1e-3, **_kwargs):
          super().__init__(sess, ob_space, ac_space, n_env, n_steps, n_batch, reuse,
                                             cnn_extractor=cnn_1d_extractor, feature_extraction="mlp", **_kwargs)
  
          self._obs_ph = self.processed_obs  # Base class has self.obs_ph as property getting self._obs_ph
          with tf.variable_scope("input", reuse=False):
              self.my_ph = tf.placeholder(tf.float32, (self.n_batch, *my_size), name="my_ph")  # (done t-1)
              self.action_prev_ph = tf.placeholder(tf.float32, (self.n_batch, *self.ac_space.shape), name="action_prev_ph")
              self.obs_prev_ph = tf.placeholder(tf.float32, (self.n_batch, *self.ob_space.shape), name="obs_prev_ph")
  
          self.loss_weight = loss_weight
          self.obs_prev = np.zeros((1, *self.ob_space.shape))
          self.action_prev = np.zeros((1, *self.ac_space.shape))
          self.my_est_loss_op = None
          self.my_est_op = None
          self.policy_loss = None
          self.my_est = None
          self.extra_phs = ["my", "action_prev", "obs_prev", "target_my", "target_action_prev", "target_obs_prev"]
          self.extra_data_names = ["my", "action_prev", "obs_prev", "target_my", "target_action_prev", "target_obs_prev"]
  
      def _process_phs(self, **phs):
          for ph_name, ph_val in phs.items():
              if ph_val is None:
                  phs[ph_name] = getattr(self, ph_name + "_ph")
              else:
                  try:
                      setattr(self, ph_name + "_ph", ph_val)
                  except AttributeError:
                      setattr(self, "_" + ph_name + "_ph", ph_val)
  
          return phs.values()
  
      def make_actor(self, obs=None, obs_prev=None, action_prev=None, my_gt=None, reuse=False, scope="pi"):
          obs, obs_prev, action_prev, my_gt = self._process_phs(obs=obs, obs_prev=obs_prev, action_prev=action_prev, my=my_gt)
  
          if self.obs_module_indices is not None:
              obs = tf.gather(obs, self.obs_module_indices["pi"], axis=-1)
              obs_prev = tf.gather(obs_prev, self.obs_module_indices["pi"], axis=-1)
  
          with tf.variable_scope(scope + "/my", reuse=reuse):
              my_h = tf.concat([obs, obs_prev, action_prev], axis=-1)
              my_h = mlp(my_h, [64, 64], self.activ_fn, layer_norm=self.layer_norm)
              self.my_est_op = tf.layers.dense(my_h, self.my_ph.shape[-1])
              self.my_est_loss_op = tf.reduce_mean((self.my_est_op - my_gt) ** 2)
              self.policy_loss = self.loss_weight * self.my_est_loss_op
  
          obs = tf.concat([obs, self.my_est_op], axis=-1)
  
          with tf.variable_scope(scope, reuse=reuse):
              if self.feature_extraction == "cnn":
                  pi_h = self.cnn_extractor(obs, name="pi_c1", act_fun=self.activ_fn, **self.cnn_kwargs)
              else:
                  pi_h = tf.layers.flatten(obs)
  
              pi_h = mlp(pi_h, self.layers, self.activ_fn, layer_norm=self.layer_norm)
  
              self.policy_pre_activation = tf.layers.dense(pi_h, self.ac_space.shape[0])
              self.policy = policy = tf.tanh(self.policy_pre_activation)
  
          return policy
  
      def make_critics(self, obs=None, action=None, my=None, reuse=False, scope="values_fn"):
          obs, action, my = self._process_phs(obs=obs, action=action, my=my)
  
          return super().make_critics(obs, action, reuse, scope, extracted_callback=lambda x: tf.concat([x, my], axis=-1))
  
      def collect_data(self, _locals, _globals):
          data = {}
          if "my" not in _locals or _locals["episode_data"]:
              data["my"] = _locals["self"].env.get_env_parameters()
              data["target_my"] = data["my"]
          if len(_locals["episode_data"]) == 0:
              data["obs_prev"] = _locals["obs"]
              data["action_prev"] = _locals["action"]
          else:
              data["obs_prev"] = _locals["episode_data"][-1]["obs"]
              data["action_prev"] = _locals["episode_data"][-1]["action"]
          data["target_obs_prev"] = data["obs_prev"]
          data["target_action_prev"] = data["action_prev"]
  
          return data
      
      def step(self, obs, obs_prev=None, action_prev=None, mask=None):
          if action_prev is None:
              assert obs.shape[0] == 1
              if mask is not None and mask[0]:
                  self.action_prev = np.zeros((1, *self.ac_space.shape))
              action_prev = self.action_prev
          if obs_prev is None:
              if mask is not None and mask[0]:
                  self.obs_prev = np.zeros((1, *self.ob_space.shape))
              obs_prev = self.obs_prev
  
          action, my_est = self.sess.run([self.policy, self.my_est_op], {self.obs_ph: obs,
                                               self.action_prev_ph: action_prev,
                                               self.obs_prev_ph: obs_prev})
          self.action_prev = action
          self.obs_prev = obs
          self.my_est = my_est
  
          #return action, my_est
          return action
  
  
  class LnCnnPolicy(FeedForwardPolicy):
      """
      Policy object that implements actor critic, using a CNN (the nature CNN), with layer normalisation
  
      :param sess: (TensorFlow session) The current TensorFlow session
      :param ob_space: (Gym Space) The observation space of the environment
      :param ac_space: (Gym Space) The action space of the environment
      :param n_env: (int) The number of environments to run
      :param n_steps: (int) The number of steps to run for each environment
      :param n_batch: (int) The number of batch to run (n_envs * n_steps)
      :param reuse: (bool) If the policy is reusable or not
      :param _kwargs: (dict) Extra keyword arguments for the nature CNN feature extraction
      """
  
      def __init__(self, sess, ob_space, ac_space, n_env=1, n_steps=1, n_batch=None, reuse=False, **_kwargs):
          super(LnCnnPolicy, self).__init__(sess, ob_space, ac_space, n_env, n_steps, n_batch, reuse,
                                            feature_extraction="cnn", layer_norm=True, **_kwargs)
  
  
  class MlpPolicy(FeedForwardPolicy):
      """
      Policy object that implements actor critic, using a MLP (2 layers of 64)
  
      :param sess: (TensorFlow session) The current TensorFlow session
      :param ob_space: (Gym Space) The observation space of the environment
      :param ac_space: (Gym Space) The action space of the environment
      :param n_env: (int) The number of environments to run
      :param n_steps: (int) The number of steps to run for each environment
      :param n_batch: (int) The number of batch to run (n_envs * n_steps)
      :param reuse: (bool) If the policy is reusable or not
      :param _kwargs: (dict) Extra keyword arguments for the nature CNN feature extraction
      """
  
      def __init__(self, sess, ob_space, ac_space, n_env=1, n_steps=1, n_batch=None, reuse=False, **_kwargs):
          super(MlpPolicy, self).__init__(sess, ob_space, ac_space, n_env, n_steps, n_batch, reuse,
                                          feature_extraction="mlp", **_kwargs)
  
  
  class LnMlpPolicy(FeedForwardPolicy):
      """
      Policy object that implements actor critic, using a MLP (2 layers of 64), with layer normalisation
  
      :param sess: (TensorFlow session) The current TensorFlow session
      :param ob_space: (Gym Space) The observation space of the environment
      :param ac_space: (Gym Space) The action space of the environment
      :param n_env: (int) The number of environments to run
      :param n_steps: (int) The number of steps to run for each environment
      :param n_batch: (int) The number of batch to run (n_envs * n_steps)
      :param reuse: (bool) If the policy is reusable or not
      :param _kwargs: (dict) Extra keyword arguments for the nature CNN feature extraction
      """
  
      def __init__(self, sess, ob_space, ac_space, n_env=1, n_steps=1, n_batch=None, reuse=False, **_kwargs):
          super(LnMlpPolicy, self).__init__(sess, ob_space, ac_space, n_env, n_steps, n_batch, reuse,
                                            feature_extraction="mlp", layer_norm=True, **_kwargs)
  
  
  register_policy("LstmFFMlpPolicy", LstmFFMlpPolicy)
  register_policy("LstmMlpPolicy", LstmMlpPolicy)
  register_policy("DRPolicy", DRPolicy)
  register_policy("CnnPolicy", CnnPolicy)
  register_policy("LnCnnPolicy", LnCnnPolicy)
  register_policy("MlpPolicy", MlpPolicy)
  register_policy("LnMlpPolicy", LnMlpPolicy)
  register_policy("CnnMlpPolicy", CnnMlpPolicy)
  register_policy("DRCnnMlpPolicy", DRCnnMlpPolicy)
  register_policy("DRMyEstPolicy", DRMyEstPolicy)
  
apis: tensorflow.summary.scalar tensorflow.reshape tensorflow.summary.image tensorflow.variable_scope tensorflow.reduce_sum numpy.isfinite numpy.allclose tensorflow.summary.histogram numpy.set_printoptions tensorflow.Graph numpy.empty_like numpy.mean tensorflow.gradients tensorflow.RunOptions tensorflow.assign tensorflow.placeholder tensorflow.summary.merge_all tensorflow.reduce_mean tensorflow.RunMetadata tensorflow.square numpy.array numpy.concatenate
apis: tensorflow.placeholder_with_default tensorflow.placeholder numpy.zeros tensorflow.subtract tensorflow.layers.flatten tensorflow.tanh tensorflow.reduce_mean tensorflow.variable_scope numpy.expand_dims tensorflow.concat numpy.array tensorflow.gather tensorflow.layers.dense