code: """Bernoulli MLP Regressor based on MLP with Normalized Inputs."""
  from dowel import tabular
  import numpy as np
  import tensorflow as tf
  
  from garage.tf.distributions import Bernoulli
  from garage.tf.misc import tensor_utils
  from garage.tf.models import NormalizedInputMLPModel
  from garage.tf.optimizers import ConjugateGradientOptimizer, LbfgsOptimizer
  from garage.tf.regressors.regressor import StochasticRegressor
  
  
  class BernoulliMLPRegressor(StochasticRegressor):
      """Fits data to a Bernoulli distribution, parameterized by an MLP.
  
      Args:
          input_shape (tuple[int]): Input shape of the training data. Since an
              MLP model is used, implementation assumes flattened inputs. The
              input shape of each data point should thus be of shape (x, ).
          output_dim (int): Output dimension of the model.
          name (str): Model name, also the variable scope.
          hidden_sizes (list[int]): Output dimension of dense layer(s) for
              the MLP for the network. For example, (32, 32) means the MLP
              consists of two hidden layers, each with 32 hidden units.
          hidden_nonlinearity (Callable): Activation function for intermediate
              dense layer(s). It should return a tf.Tensor. Set it to
              None to maintain a linear activation.
          hidden_w_init (Callable): Initializer function for the weight
              of intermediate dense layer(s). The function should return a
              tf.Tensor. Default is Glorot uniform initializer.
          hidden_b_init (Callable): Initializer function for the bias
              of intermediate dense layer(s). The function should return a
              tf.Tensor. Default is zero initializer.
          output_nonlinearity (Callable): Activation function for output dense
              layer. It should return a tf.Tensor. Set it to None to
              maintain a linear activation.
          output_w_init (Callable): Initializer function for the weight
              of output dense layer(s). The function should return a
              tf.Tensor. Default is Glorot uniform initializer.
          output_b_init (Callable): Initializer function for the bias
              of output dense layer(s). The function should return a
              tf.Tensor. Default is zero initializer.
          optimizer (garage.tf.Optimizer): Optimizer for minimizing the negative
              log-likelihood. Defaults to LbsgsOptimizer
          optimizer_args (dict): Arguments for the optimizer. Default is None,
              which means no arguments.
          tr_optimizer (garage.tf.Optimizer): Optimizer for trust region
              approximation. Defaults to ConjugateGradientOptimizer.
          tr_optimizer_args (dict): Arguments for the trust region optimizer.
              Default is None, which means no arguments.
          use_trust_region (bool): Whether to use trust region constraint.
          max_kl_step (float): KL divergence constraint for each iteration.
          normalize_inputs (bool): Bool for normalizing inputs or not.
          layer_normalization (bool): Bool for using layer normalization or not.
  
      """
  
      def __init__(self,
                   input_shape,
                   output_dim,
                   name='BernoulliMLPRegressor',
                   hidden_sizes=(32, 32),
                   hidden_nonlinearity=tf.nn.relu,
                   hidden_w_init=tf.initializers.glorot_uniform(),
                   hidden_b_init=tf.zeros_initializer(),
                   output_nonlinearity=tf.nn.sigmoid,
                   output_w_init=tf.initializers.glorot_uniform(),
                   output_b_init=tf.zeros_initializer(),
                   optimizer=None,
                   optimizer_args=None,
                   tr_optimizer=None,
                   tr_optimizer_args=None,
                   use_trust_region=True,
                   max_kl_step=0.01,
                   normalize_inputs=True,
                   layer_normalization=False):
  
          super().__init__(input_shape, output_dim, name)
          self._use_trust_region = use_trust_region
          self._max_kl_step = max_kl_step
          self._normalize_inputs = normalize_inputs
  
          with tf.compat.v1.variable_scope(self._name, reuse=False) as vs:
              self._variable_scope = vs
              if optimizer_args is None:
                  optimizer_args = dict()
              if tr_optimizer_args is None:
                  tr_optimizer_args = dict()
  
              if optimizer is None:
                  optimizer = LbfgsOptimizer(**optimizer_args)
              else:
                  optimizer = optimizer(**optimizer_args)
  
              if tr_optimizer is None:
                  tr_optimizer = ConjugateGradientOptimizer(**tr_optimizer_args)
              else:
                  tr_optimizer = tr_optimizer(**tr_optimizer_args)
  
              self._optimizer = optimizer
              self._tr_optimizer = tr_optimizer
              self._first_optimized = False
  
          self.model = NormalizedInputMLPModel(
              input_shape,
              output_dim,
              hidden_sizes=hidden_sizes,
              hidden_nonlinearity=hidden_nonlinearity,
              hidden_w_init=hidden_w_init,
              hidden_b_init=hidden_b_init,
              output_nonlinearity=output_nonlinearity,
              output_w_init=output_w_init,
              output_b_init=output_b_init,
              layer_normalization=layer_normalization)
  
          self._dist = Bernoulli(output_dim)
  
          self._initialize()
  
      def _initialize(self):
          input_var = tf.compat.v1.placeholder(tf.float32,
                                               shape=(None, ) +
                                               self._input_shape)
  
          with tf.compat.v1.variable_scope(self._variable_scope):
              self.model.build(input_var)
  
              ys_var = tf.compat.v1.placeholder(dtype=tf.float32,
                                                name='ys',
                                                shape=(None, self._output_dim))
  
              old_prob_var = tf.compat.v1.placeholder(dtype=tf.float32,
                                                      name='old_prob',
                                                      shape=(None,
                                                             self._output_dim))
  
              y_hat = self.model.networks['default'].y_hat
  
              old_info_vars = dict(p=old_prob_var)
              info_vars = dict(p=y_hat)
  
              mean_kl = tf.reduce_mean(
                  self._dist.kl_sym(old_info_vars, info_vars))
  
              loss = -tf.reduce_mean(
                  self._dist.log_likelihood_sym(ys_var, info_vars))
  
              predicted = y_hat >= 0.5
  
              self._f_predict = tensor_utils.compile_function([input_var],
                                                              predicted)
              self._f_prob = tensor_utils.compile_function([input_var], y_hat)
  
              self._optimizer.update_opt(loss=loss,
                                         target=self,
                                         inputs=[input_var, ys_var])
              self._tr_optimizer.update_opt(
                  loss=loss,
                  target=self,
                  inputs=[input_var, ys_var, old_prob_var],
                  leq_constraint=(mean_kl, self._max_kl_step))
  
      def fit(self, xs, ys):
          """Fit with input data xs and label ys.
  
          Args:
              xs (numpy.ndarray): Input data.
              ys (numpy.ndarray): Label of input data.
  
          """
          if self._normalize_inputs:
              # recompute normalizing constants for inputs
              self.model.networks['default'].x_mean.load(
                  np.mean(xs, axis=0, keepdims=True))
              self.model.networks['default'].x_std.load(
                  np.std(xs, axis=0, keepdims=True) + 1e-8)
  
          if self._use_trust_region and self._first_optimized:
              # To use trust region constraint and optimizer
              old_prob = self._f_prob(xs)
              inputs = [xs, ys, old_prob]
              optimizer = self._tr_optimizer
          else:
              inputs = [xs, ys]
              optimizer = self._optimizer
          loss_before = optimizer.loss(inputs)
          tabular.record('{}/LossBefore'.format(self._name), loss_before)
          optimizer.optimize(inputs)
          loss_after = optimizer.loss(inputs)
          tabular.record('{}/LossAfter'.format(self._name), loss_after)
          tabular.record('{}/dLoss'.format(self._name), loss_before - loss_after)
          self._first_optimized = True
  
      def predict(self, xs):
          """Predict ys based on input xs.
  
          Args:
              xs (numpy.ndarray): Input data of shape (samples, input_dim)
  
          Return:
              numpy.ndarray: The deterministic predicted ys (one hot vectors)
                  of shape (samples, output_dim)
  
          """
          return self._f_predict(xs)
  
      def sample_predict(self, xs):
          """Do a Bernoulli sampling given input xs.
  
          Args:
              xs (numpy.ndarray): Input data of shape (samples, input_dim)
  
          Returns:
              numpy.ndarray: The stochastic sampled ys
                  of shape (samples, output_dim)
  
          """
          p = self._f_prob(xs)
          return self._dist.sample(dict(p=p))
  
      def predict_log_likelihood(self, xs, ys):
          """Log likelihood of ys given input xs.
  
          Args:
              xs (numpy.ndarray): Input data of shape (samples, input_dim)
              ys (numpy.ndarray): Output data of shape (samples, output_dim)
  
          Returns:
              numpy.ndarray: The log likelihood of shape (samples, )
  
          """
          p = self._f_prob(xs)
          return self._dist.log_likelihood(ys, dict(p=p))
  
      def log_likelihood_sym(self, x_var, y_var, name=None):
          """Build a symbolic graph of the log-likelihood.
  
          Args:
              x_var (tf.Tensor): Input tf.Tensor for the input data.
              y_var (tf.Tensor): Input tf.Tensor for the one hot label of data.
              name (str): Name of the new graph.
  
          Return:
              tf.Tensor: Output of the symbolic log-likelihood graph.
  
          """
          with tf.compat.v1.variable_scope(self._variable_scope):
              prob, _, _ = self.model.build(x_var, name=name)
  
          return self._dist.log_likelihood_sym(y_var, dict(p=prob))
  
      def dist_info_sym(self, input_var, state_info_vars=None, name=None):
          """Build a symbolic graph of the distribution parameters.
  
          Args:
              input_var (tf.Tensor): Input tf.Tensor for the input data.
              state_info_vars (dict): a dictionary whose values should contain
                  information about the state of the regressor at the time it
                  received the input.
              name (str): Name of the new graph.
  
          Return:
              dict[tf.Tensor]: Output of the symbolic graph of the distribution
                  parameters.
  
          """
          with tf.compat.v1.variable_scope(self._variable_scope):
              prob, _, _ = self.model.build(input_var, name=name)
  
          return dict(prob=prob)
  
      @property
      def recurrent(self):
          """bool: If this module has a hidden state."""
          return False
  
      @property
      def vectorized(self):
          """bool: If this module supports vectorization input."""
          return True
  
      @property
      def distribution(self):
          """garage.tf.distributions.DiagonalGaussian: Distribution."""
          return self.model.networks['default'].dist
  
      def __getstate__(self):
          """Object.__getstate__.
  
          Returns:
              dict: The state to be pickled for the instance.
  
          """
          new_dict = super().__getstate__()
          del new_dict['_f_predict']
          del new_dict['_f_prob']
          return new_dict
  
      def __setstate__(self, state):
          """Object.__setstate__.
  
          Args:
              state (dict): Unpickled state.
  
          """
          super().__setstate__(state)
          self._initialize()
   """First order optimizer."""
  import time
  
  from dowel import logger
  import pyprind
  import tensorflow as tf
  
  from garage.np.optimizers import BatchDataset
  from garage.tf.misc import tensor_utils
  from garage.tf.optimizers.utils import LazyDict
  
  
  class FirstOrderOptimizer:
      """First order optimier.
  
      Performs (stochastic) gradient descent, possibly using fancier methods like
      ADAM etc.
  
      Args:
          tf_optimizer_cls (tf.Optimizer): Optimizer to be used.
          tf_optimizer_args (dict): Optimizer arguments.
          max_epochs (int): Maximum number of epochs for update.
          tolerance (float): Tolerance for difference in loss during update.
          batch_size (int): Batch size for optimization.
          callback (callable): Function to call during each epoch. Default is
              None.
          verbose (bool): If true, intermediate log message will be printed.
          name (str): Name scope of the optimizer.
  
      """
  
      def __init__(self,
                   tf_optimizer_cls=None,
                   tf_optimizer_args=None,
                   max_epochs=1000,
                   tolerance=1e-6,
                   batch_size=32,
                   callback=None,
                   verbose=False,
                   name='FirstOrderOptimizer'):
          self._opt_fun = None
          self._target = None
          self._callback = callback
          if tf_optimizer_cls is None:
              tf_optimizer_cls = tf.compat.v1.train.AdamOptimizer
          if tf_optimizer_args is None:
              tf_optimizer_args = dict(learning_rate=1e-3)
          self._tf_optimizer = tf_optimizer_cls(**tf_optimizer_args)
          self._max_epochs = max_epochs
          self._tolerance = tolerance
          self._batch_size = batch_size
          self._verbose = verbose
          self._input_vars = None
          self._train_op = None
          self._name = name
  
      def update_opt(self, loss, target, inputs, extra_inputs=None, **kwargs):
          """Construct operation graph for the optimizer.
  
          Args:
              loss (tf.Tensor): Loss objective to minimize.
              target (object): Target object to optimize. The object should
                  implemenet `get_params()` and `get_param_values`.
              inputs (list[tf.Tensor]): List of input placeholders.
              extra_inputs (list[tf.Tensor]): List of extra input placeholders.
              kwargs (dict): Extra unused keyword arguments. Some optimizers
                  have extra input, e.g. KL constraint.
  
          """
          del kwargs
          with tf.name_scope(self._name):
              self._target = target
  
              self._train_op = self._tf_optimizer.minimize(
                  loss, var_list=target.get_params())
  
              if extra_inputs is None:
                  extra_inputs = list()
              self._input_vars = inputs + extra_inputs
              self._opt_fun = LazyDict(
                  f_loss=lambda: tensor_utils.compile_function(
                      inputs + extra_inputs, loss), )
  
      def loss(self, inputs, extra_inputs=None):
          """The loss.
  
          Args:
              inputs (list[numpy.ndarray]): List of input values.
              extra_inputs (list[numpy.ndarray]): List of extra input values.
  
          Returns:
              float: Loss.
  
          Raises:
              Exception: If loss function is None, i.e. not defined.
  
          """
          if self._opt_fun is None:
              raise Exception(
                  'Use update_opt() to setup the loss function first.')
          if extra_inputs is None:
              extra_inputs = tuple()
          return self._opt_fun['f_loss'](*(tuple(inputs) + extra_inputs))
  
          # pylint: disable=too-many-branches
      def optimize(self, inputs, extra_inputs=None, callback=None):
          """Perform optimization.
  
          Args:
              inputs (list[numpy.ndarray]): List of input values.
              extra_inputs (list[numpy.ndarray]): List of extra input values.
              callback (callable): Function to call during each epoch. Default
                  is None.
  
          Raises:
              NotImplementedError: If inputs are invalid.
              Exception: If loss function is None, i.e. not defined.
  
          """
          if not inputs:
              # Assumes that we should always sample mini-batches
              raise NotImplementedError('No inputs are fed to optimizer.')
          if self._opt_fun is None:
              raise Exception(
                  'Use update_opt() to setup the loss function first.')
  
          f_loss = self._opt_fun['f_loss']
  
          if extra_inputs is None:
              extra_inputs = tuple()
  
          last_loss = f_loss(*(tuple(inputs) + extra_inputs))
  
          start_time = time.time()
  
          dataset = BatchDataset(inputs,
                                 self._batch_size,
                                 extra_inputs=extra_inputs)
  
          sess = tf.compat.v1.get_default_session()
  
          for epoch in range(self._max_epochs):
              if self._verbose:
                  logger.log('Epoch {}'.format(epoch))
                  progbar = pyprind.ProgBar(len(inputs[0]))
  
              for batch in dataset.iterate(update=True):
                  sess.run(self._train_op,
                           dict(list(zip(self._input_vars, batch))))
                  if self._verbose:
                      progbar.update(len(batch[0]))
  
              if self._verbose:
                  if progbar.active:
                      progbar.stop()
  
              new_loss = f_loss(*(tuple(inputs) + extra_inputs))
  
              if self._verbose:
                  logger.log('Epoch: {} | Loss: {}'.format(epoch, new_loss))
              if self._callback or callback:
                  elapsed = time.time() - start_time
                  callback_args = dict(
                      loss=new_loss,
                      params=self._target.get_param_values()
                      if self._target else None,
                      itr=epoch,
                      elapsed=elapsed,
                  )
                  if self._callback:
                      self._callback(callback_args)
                  if callback:
                      callback(**callback_args)
  
              if abs(last_loss - new_loss) < self._tolerance:
                  break
              last_loss = new_loss
  
      def __getstate__(self):
          """Object.__getstate__.
  
          Returns:
              dict: The state to be pickled for the instance.
  
          """
          new_dict = self.__dict__.copy()
          del new_dict['_opt_fun']
          del new_dict['_tf_optimizer']
          del new_dict['_train_op']
          del new_dict['_input_vars']
          return new_dict
  
      def __setstate__(self, state):
          """Object.__setstate__.
  
          Args:
              state (dict): Unpickled state.
  
          """
          obj = type(self)()
          self.__dict__.update(obj.__dict__)
          self.__dict__.update(state)
   """This module implements a Hindsight Experience Replay (HER).
  
  See: https://arxiv.org/abs/1707.01495.
  """
  import inspect
  
  import numpy as np
  
  from garage.replay_buffer.replay_buffer import ReplayBuffer
  
  
  def make_her_sample(replay_k, reward_fun):
      """Generate a transition sampler for HER ReplayBuffer.
  
      Args:
          replay_k (float): Ratio between HER replays and regular replays
          reward_fun (callable): Function to re-compute the reward with
              substituted goals
  
      Returns:
          callable: A function that returns sample transitions for HER.
  
      """
      future_p = 1 - (1. / (1 + replay_k))
  
      def _her_sample_transitions(episode_batch, sample_batch_size):
          """Generate a dictionary of transitions.
  
          Args:
              episode_batch (dict): Original transitions which
                  transitions[key] has shape :math:`(N, T, S^*)`.
              sample_batch_size (int): Batch size per sample.
  
          Returns:
              dict[numpy.ndarray]: Transitions.
  
          """
          # Select which episodes to use
          time_horizon = episode_batch['action'].shape[1]
          rollout_batch_size = episode_batch['action'].shape[0]
          episode_idxs = np.random.randint(rollout_batch_size,
                                           size=sample_batch_size)
          # Select time steps to use
          t_samples = np.random.randint(time_horizon, size=sample_batch_size)
          transitions = {
              key: episode_batch[key][episode_idxs, t_samples]
              for key in episode_batch.keys()
          }
  
          her_idxs = np.where(
              np.random.uniform(size=sample_batch_size) < future_p)
          future_offset = np.random.uniform(
              size=sample_batch_size) * (time_horizon - t_samples)
          future_offset = future_offset.astype(int)
          future_t = (t_samples + future_offset)[her_idxs]
  
          future_ag = episode_batch['achieved_goal'][episode_idxs[her_idxs],
                                                     future_t]
          transitions['goal'][her_idxs] = future_ag
  
          achieved_goals = episode_batch['achieved_goal'][episode_idxs[her_idxs],
                                                          t_samples[her_idxs]]
          transitions['achieved_goal'][her_idxs] = achieved_goals
  
          # Re-compute reward since we may have substituted the goal.
          reward_params_keys = inspect.signature(reward_fun).parameters.keys()
          reward_params = {
              rk: transitions[k]
              for k, rk in zip(['next_achieved_goal', 'goal'],
                               list(reward_params_keys)[:-1])
          }
          reward_params['info'] = {}
          transitions['reward'] = reward_fun(**reward_params)
  
          transitions = {
              k: transitions[k].reshape(sample_batch_size,
                                        *transitions[k].shape[1:])
              for k in transitions.keys()
          }
  
          goals = transitions['goal']
          next_inputs = np.concatenate((transitions['next_observation'], goals,
                                        transitions['achieved_goal']),
                                       axis=-1)
          inputs = np.concatenate(
              (transitions['observation'], goals, transitions['achieved_goal']),
              axis=-1)
          transitions['observation'] = inputs
          transitions['next_observation'] = next_inputs
  
          assert transitions['action'].shape[0] == sample_batch_size
          return transitions
  
      return _her_sample_transitions
  
  
  class HerReplayBuffer(ReplayBuffer):
      """Replay buffer for HER (Hindsight Experience Replay).
  
      It constructs hindsight examples using future strategy.
  
      Args:
          replay_k (float): Ratio between HER replays and regular replays
          reward_fun (callable): Function to re-compute the reward with
              substituted goals
          env_spec (garage.envs.EnvSpec): Environment specification.
          size_in_transitions (int): total size of transitions in the buffer
          time_horizon (int): time horizon of rollout.
  
      """
  
      def __init__(self, replay_k, reward_fun, env_spec, size_in_transitions,
                   time_horizon):
          self._env_spec = env_spec
          self._sample_transitions = make_her_sample(replay_k, reward_fun)
          self._replay_k = replay_k
          self._reward_fun = reward_fun
          super().__init__(env_spec, size_in_transitions, time_horizon)
  
      def sample(self, batch_size):
          """Sample a transition of batch_size.
  
          Args:
              batch_size (int): Batch size to sample.
  
          Return:
              dict[numpy.ndarray]: Transitions which transitions[key] has the
                  shape of :math:`(N, S^*)`. Keys include [`observation`,
                  `action`, `goal`, `achieved_goal`, `terminal`,
                  `next_observation`, `next_achieved_goal` and `reward`].
  
          """
          buffer = {}
          for key in self._buffer:
              buffer[key] = self._buffer[key][:self._current_size]
  
          transitions = self._sample_transitions(buffer, batch_size)
  
          for key in (['reward', 'next_observation', 'next_achieved_goal'] +
                      list(self._buffer.keys())):
              assert key in transitions, 'key %s missing from transitions' % key
  
          return transitions
  
      def __getstate__(self):
          """Object.__getstate__.
  
          Returns:
              dict: The state to be pickled for the instance.
  
          """
          new_dict = self.__dict__.copy()
          del new_dict['_sample_transitions']
          return new_dict
  
      def __setstate__(self, state):
          """Object.__setstate__.
  
          Args:
              state (dict): Unpickled state.
  
          """
          self.__dict__ = state
          replay_k = state['_replay_k']
          reward_fun = state['_reward_fun']
          self._sample_transitions = make_her_sample(replay_k, reward_fun)
  
      def add_transitions(self, **kwargs):
          """Add multiple transitions into the replay buffer.
  
          A transition contains one or multiple entries, e.g.
          observation, action, reward, terminal and next_observation.
          The same entry of all the transitions are stacked, e.g.
          {'observation': [obs1, obs2, obs3]} where obs1 is one
          numpy.ndarray observation from the environment.
  
          Args:
              kwargs (dict(str, [numpy.ndarray])): Dictionary that holds
                  the transitions.
  
          """
          obses = kwargs['observation']
          obs = [obs['observation'] for obs in obses]
          d_g = [obs['desired_goal'] for obs in obses]
          a_g = [obs['achieved_goal'] for obs in obses]
          next_obses = kwargs['next_observation']
          super().add_transitions(
              observation=obs,
              action=kwargs['action'],
              goal=d_g,
              achieved_goal=a_g,
              terminal=kwargs['terminal'],
              next_observation=[
                  next_obs['observation'] for next_obs in next_obses
              ],
              next_achieved_goal=[
                  next_obs['achieved_goal'] for next_obs in next_obses
              ],
          )
  
apis: tensorflow.compat.v1.placeholder tensorflow.zeros_initializer tensorflow.compat.v1.variable_scope numpy.std tensorflow.initializers.glorot_uniform numpy.mean
apis: tensorflow.name_scope tensorflow.compat.v1.get_default_session
apis: numpy.random.uniform numpy.concatenate numpy.random.randint