text: # Copyright 2017 The TensorFlow Authors All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  
  """Policy neural network.
  
  Implements network which takes in input and produces actions
  and log probabilities given a sampling distribution parameterization.
  """
  
  import tensorflow as tf
  import numpy as np
  
  
  class Policy(object):
    def __init__(self, env_spec, internal_dim,
                 fixed_std=True, recurrent=True,
                 input_prev_actions=True):
      self.env_spec = env_spec
      self.internal_dim = internal_dim
      self.rnn_state_dim = self.internal_dim
      self.fixed_std = fixed_std
      self.recurrent = recurrent
      self.input_prev_actions = input_prev_actions
  
      self.matrix_init = tf.truncated_normal_initializer(stddev=0.01)
      self.vector_init = tf.constant_initializer(0.0)
  
    @property
    def input_dim(self):
      return (self.env_spec.total_obs_dim +
              self.env_spec.total_sampled_act_dim * self.input_prev_actions)
  
    @property
    def output_dim(self):
      return self.env_spec.total_sampling_act_dim
  
    def get_cell(self):
      """Get RNN cell."""
      self.cell_input_dim = self.internal_dim // 2
      cell = tf.contrib.rnn.LSTMCell(self.cell_input_dim,
                                     state_is_tuple=False,
                                     reuse=tf.get_variable_scope().reuse)
  
      cell = tf.contrib.rnn.OutputProjectionWrapper(
          cell, self.output_dim,
          reuse=tf.get_variable_scope().reuse)
  
      return cell
  
    def core(self, obs, prev_internal_state, prev_actions):
      """Core neural network taking in inputs and outputting sampling
      distribution parameters."""
      batch_size = tf.shape(obs[0])[0]
      if not self.recurrent:
        prev_internal_state = tf.zeros([batch_size, self.rnn_state_dim])
  
      cell = self.get_cell()
  
      b = tf.get_variable('input_bias', [self.cell_input_dim],
                          initializer=self.vector_init)
      cell_input = tf.nn.bias_add(tf.zeros([batch_size, self.cell_input_dim]), b)
  
      for i, (obs_dim, obs_type) in enumerate(self.env_spec.obs_dims_and_types):
        w = tf.get_variable('w_state%d' % i, [obs_dim, self.cell_input_dim],
                            initializer=self.matrix_init)
        if self.env_spec.is_discrete(obs_type):
          cell_input += tf.matmul(tf.one_hot(obs[i], obs_dim), w)
        elif self.env_spec.is_box(obs_type):
          cell_input += tf.matmul(obs[i], w)
        else:
          assert False
  
      if self.input_prev_actions:
        if self.env_spec.combine_actions:  # TODO(ofir): clean this up
          prev_action = prev_actions[0]
          for i, action_dim in enumerate(self.env_spec.orig_act_dims):
            act = tf.mod(prev_action, action_dim)
            w = tf.get_variable('w_prev_action%d' % i, [action_dim, self.cell_input_dim],
                                initializer=self.matrix_init)
            cell_input += tf.matmul(tf.one_hot(act, action_dim), w)
            prev_action = tf.to_int32(prev_action / action_dim)
        else:
          for i, (act_dim, act_type) in enumerate(self.env_spec.act_dims_and_types):
            w = tf.get_variable('w_prev_action%d' % i, [act_dim, self.cell_input_dim],
                                initializer=self.matrix_init)
            if self.env_spec.is_discrete(act_type):
              cell_input += tf.matmul(tf.one_hot(prev_actions[i], act_dim), w)
            elif self.env_spec.is_box(act_type):
              cell_input += tf.matmul(prev_actions[i], w)
            else:
              assert False
  
      output, next_state = cell(cell_input, prev_internal_state)
  
      return output, next_state
  
    def sample_action(self, logits, sampling_dim,
                      act_dim, act_type, greedy=False):
      """Sample an action from a distribution."""
      if self.env_spec.is_discrete(act_type):
        if greedy:
          act = tf.argmax(logits, 1)
        else:
          act = tf.reshape(tf.multinomial(logits, 1), [-1])
      elif self.env_spec.is_box(act_type):
        means = logits[:, :sampling_dim / 2]
        std = logits[:, sampling_dim / 2:]
        if greedy:
          act = means
        else:
          batch_size = tf.shape(logits)[0]
          act = means + std * tf.random_normal([batch_size, act_dim])
      else:
        assert False
  
      return act
  
    def entropy(self, logits,
                sampling_dim, act_dim, act_type):
      """Calculate entropy of distribution."""
      if self.env_spec.is_discrete(act_type):
        entropy = tf.reduce_sum(
            -tf.nn.softmax(logits) * tf.nn.log_softmax(logits), -1)
      elif self.env_spec.is_box(act_type):
        means = logits[:, :sampling_dim / 2]
        std = logits[:, sampling_dim / 2:]
        entropy = tf.reduce_sum(
            0.5 * (1 + tf.log(2 * np.pi * tf.square(std))), -1)
      else:
        assert False
  
      return entropy
  
    def self_kl(self, logits,
                sampling_dim, act_dim, act_type):
      """Calculate KL of distribution with itself.
  
      Used layer only for the gradients.
      """
  
      if self.env_spec.is_discrete(act_type):
        probs = tf.nn.softmax(logits)
        log_probs = tf.nn.log_softmax(logits)
        self_kl = tf.reduce_sum(
            tf.stop_gradient(probs) *
            (tf.stop_gradient(log_probs) - log_probs), -1)
      elif self.env_spec.is_box(act_type):
        means = logits[:, :sampling_dim / 2]
        std = logits[:, sampling_dim / 2:]
        my_means = tf.stop_gradient(means)
        my_std = tf.stop_gradient(std)
        self_kl = tf.reduce_sum(
            tf.log(std / my_std) +
            (tf.square(my_std) + tf.square(my_means - means)) /
            (2.0 * tf.square(std)) - 0.5,
            -1)
      else:
        assert False
  
      return self_kl
  
    def log_prob_action(self, action, logits,
                        sampling_dim, act_dim, act_type):
      """Calculate log-prob of action sampled from distribution."""
      if self.env_spec.is_discrete(act_type):
        act_log_prob = tf.reduce_sum(
            tf.one_hot(action, act_dim) * tf.nn.log_softmax(logits), -1)
      elif self.env_spec.is_box(act_type):
        means = logits[:, :sampling_dim / 2]
        std = logits[:, sampling_dim / 2:]
        act_log_prob = (- 0.5 * tf.log(2 * np.pi * tf.square(std))
                        - 0.5 * tf.square(action - means) / tf.square(std))
        act_log_prob = tf.reduce_sum(act_log_prob, -1)
      else:
        assert False
  
      return act_log_prob
  
    def sample_actions(self, output, actions=None, greedy=False):
      """Sample all actions given output of core network."""
      sampled_actions = []
      logits = []
      log_probs = []
      entropy = []
      self_kl = []
  
      start_idx = 0
      for i, (act_dim, act_type) in enumerate(self.env_spec.act_dims_and_types):
        sampling_dim = self.env_spec.sampling_dim(act_dim, act_type)
        if self.fixed_std and self.env_spec.is_box(act_type):
          act_logits = output[:, start_idx:start_idx + act_dim]
  
          log_std = tf.get_variable('std%d' % i, [1, sampling_dim // 2])
          # fix standard deviations to variable
          act_logits = tf.concat(
              [act_logits,
               1e-6 + tf.exp(log_std) + 0 * act_logits], 1)
        else:
          act_logits = output[:, start_idx:start_idx + sampling_dim]
  
        if actions is None:
          act = self.sample_action(act_logits, sampling_dim,
                                   act_dim, act_type,
                                   greedy=greedy)
        else:
          act = actions[i]
  
        ent = self.entropy(act_logits, sampling_dim, act_dim, act_type)
        kl = self.self_kl(act_logits, sampling_dim, act_dim, act_type)
  
        act_log_prob = self.log_prob_action(
            act, act_logits,
            sampling_dim, act_dim, act_type)
  
        sampled_actions.append(act)
        logits.append(act_logits)
        log_probs.append(act_log_prob)
        entropy.append(ent)
        self_kl.append(kl)
  
        start_idx += sampling_dim
  
      assert start_idx == self.env_spec.total_sampling_act_dim
  
      return sampled_actions, logits, log_probs, entropy, self_kl
  
    def get_kl(self, my_logits, other_logits):
      """Calculate KL between one policy output and another."""
      kl = []
      for i, (act_dim, act_type) in enumerate(self.env_spec.act_dims_and_types):
        sampling_dim = self.env_spec.sampling_dim(act_dim, act_type)
        single_my_logits = my_logits[i]
        single_other_logits = other_logits[i]
        if self.env_spec.is_discrete(act_type):
          my_probs = tf.nn.softmax(single_my_logits)
          my_log_probs = tf.nn.log_softmax(single_my_logits)
          other_log_probs = tf.nn.log_softmax(single_other_logits)
          my_kl = tf.reduce_sum(my_probs * (my_log_probs - other_log_probs), -1)
        elif self.env_spec.is_box(act_type):
          my_means = single_my_logits[:, :sampling_dim / 2]
          my_std = single_my_logits[:, sampling_dim / 2:]
          other_means = single_other_logits[:, :sampling_dim / 2]
          other_std = single_other_logits[:, sampling_dim / 2:]
          my_kl = tf.reduce_sum(
              tf.log(other_std / my_std) +
              (tf.square(my_std) + tf.square(my_means - other_means)) /
              (2.0 * tf.square(other_std)) - 0.5,
              -1)
        else:
          assert False
  
        kl.append(my_kl)
  
      return kl
  
    def single_step(self, prev, cur, greedy=False):
      """Single RNN step.  Equivalently, single-time-step sampled actions."""
      prev_internal_state, prev_actions, _, _, _, _ = prev
      obs, actions = cur  # state observed and action taken at this time step
  
      # feed into RNN cell
      output, next_state = self.core(
          obs, prev_internal_state, prev_actions)
  
      # sample actions with values and log-probs
      (actions, logits, log_probs,
       entropy, self_kl) = self.sample_actions(
          output, actions=actions, greedy=greedy)
  
      return (next_state, tuple(actions), tuple(logits), tuple(log_probs),
              tuple(entropy), tuple(self_kl))
  
    def sample_step(self, obs, prev_internal_state, prev_actions, greedy=False):
      """Sample single step from policy."""
      (next_state, sampled_actions, logits, log_probs,
       entropies, self_kls) = self.single_step(
          (prev_internal_state, prev_actions, None, None, None, None),
          (obs, None), greedy=greedy)
      return next_state, sampled_actions
  
    def multi_step(self, all_obs, initial_state, all_actions):
      """Calculate log-probs and other calculations on batch of episodes."""
      batch_size = tf.shape(initial_state)[0]
      time_length = tf.shape(all_obs[0])[0]
      initial_actions = [act[0] for act in all_actions]
      all_actions = [tf.concat([act[1:], act[0:1]], 0)
                     for act in all_actions]  # "final" action is dummy
  
      (internal_states, _, logits, log_probs,
       entropies, self_kls) = tf.scan(
          self.single_step,
          (all_obs, all_actions),
          initializer=self.get_initializer(
              batch_size, initial_state, initial_actions))
  
      # remove "final" computations
      log_probs = [log_prob[:-1] for log_prob in log_probs]
      entropies = [entropy[:-1] for entropy in entropies]
      self_kls = [self_kl[:-1] for self_kl in self_kls]
  
      return internal_states, logits, log_probs, entropies, self_kls
  
    def get_initializer(self, batch_size, initial_state, initial_actions):
      """Get initializer for RNN."""
      logits_init = []
      log_probs_init = []
      for act_dim, act_type in self.env_spec.act_dims_and_types:
        sampling_dim = self.env_spec.sampling_dim(act_dim, act_type)
        logits_init.append(tf.zeros([batch_size, sampling_dim]))
        log_probs_init.append(tf.zeros([batch_size]))
      entropy_init = [tf.zeros([batch_size]) for _ in self.env_spec.act_dims]
      self_kl_init = [tf.zeros([batch_size]) for _ in self.env_spec.act_dims]
  
      return (initial_state,
              tuple(initial_actions),
              tuple(logits_init), tuple(log_probs_init),
              tuple(entropy_init),
              tuple(self_kl_init))
  
    def calculate_kl(self, my_logits, other_logits):
      """Calculate KL between one policy and another on batch of episodes."""
      batch_size = tf.shape(my_logits[0])[1]
      time_length = tf.shape(my_logits[0])[0]
  
      reshaped_my_logits = [
          tf.reshape(my_logit, [batch_size * time_length, -1])
          for my_logit in my_logits]
      reshaped_other_logits = [
          tf.reshape(other_logit, [batch_size * time_length, -1])
          for other_logit in other_logits]
  
      kl = self.get_kl(reshaped_my_logits, reshaped_other_logits)
      kl = [tf.reshape(kkl, [time_length, batch_size])
            for kkl in kl]
      return kl
  
  
  class MLPPolicy(Policy):
    """Non-recurrent policy."""
  
    def get_cell(self):
      self.cell_input_dim = self.internal_dim
  
      def mlp(cell_input, prev_internal_state):
        w1 = tf.get_variable('w1', [self.cell_input_dim, self.internal_dim])
        b1 = tf.get_variable('b1', [self.internal_dim])
  
        w2 = tf.get_variable('w2', [self.internal_dim, self.internal_dim])
        b2 = tf.get_variable('b2', [self.internal_dim])
  
        w3 = tf.get_variable('w3', [self.internal_dim, self.internal_dim])
        b3 = tf.get_variable('b3', [self.internal_dim])
  
        proj = tf.get_variable(
            'proj', [self.internal_dim, self.output_dim])
  
        hidden = cell_input
        hidden = tf.tanh(tf.nn.bias_add(tf.matmul(hidden, w1), b1))
        hidden = tf.tanh(tf.nn.bias_add(tf.matmul(hidden, w2), b2))
  
        output = tf.matmul(hidden, proj)
  
        return output, hidden
  
      return mlp
  
    def single_step(self, obs, actions, prev_actions, greedy=False):
      """Single step."""
      batch_size = tf.shape(obs[0])[0]
      prev_internal_state = tf.zeros([batch_size, self.internal_dim])
  
      output, next_state = self.core(
          obs, prev_internal_state, prev_actions)
  
      # sample actions with values and log-probs
      (actions, logits, log_probs,
       entropy, self_kl) = self.sample_actions(
          output, actions=actions, greedy=greedy)
  
      return (next_state, tuple(actions), tuple(logits), tuple(log_probs),
              tuple(entropy), tuple(self_kl))
  
    def sample_step(self, obs, prev_internal_state, prev_actions, greedy=False):
      """Sample single step from policy."""
      (next_state, sampled_actions, logits, log_probs,
       entropies, self_kls) = self.single_step(obs, None, prev_actions,
                                               greedy=greedy)
      return next_state, sampled_actions
  
    def multi_step(self, all_obs, initial_state, all_actions):
      """Calculate log-probs and other calculations on batch of episodes."""
      batch_size = tf.shape(initial_state)[0]
      time_length = tf.shape(all_obs[0])[0]
  
      # first reshape inputs as a single batch
      reshaped_obs = []
      for obs, (obs_dim, obs_type) in zip(all_obs, self.env_spec.obs_dims_and_types):
        if self.env_spec.is_discrete(obs_type):
          reshaped_obs.append(tf.reshape(obs, [time_length * batch_size]))
        elif self.env_spec.is_box(obs_type):
          reshaped_obs.append(tf.reshape(obs, [time_length * batch_size, obs_dim]))
  
      reshaped_act = []
      reshaped_prev_act = []
      for i, (act_dim, act_type) in enumerate(self.env_spec.act_dims_and_types):
        act = tf.concat([all_actions[i][1:], all_actions[i][0:1]], 0)
        prev_act = all_actions[i]
        if self.env_spec.is_discrete(act_type):
          reshaped_act.append(tf.reshape(act, [time_length * batch_size]))
          reshaped_prev_act.append(
              tf.reshape(prev_act, [time_length * batch_size]))
        elif self.env_spec.is_box(act_type):
          reshaped_act.append(
              tf.reshape(act, [time_length * batch_size, act_dim]))
          reshaped_prev_act.append(
              tf.reshape(prev_act, [time_length * batch_size, act_dim]))
  
      # now inputs go into single step as one large batch
      (internal_states, _, logits, log_probs,
       entropies, self_kls) = self.single_step(
           reshaped_obs, reshaped_act, reshaped_prev_act)
  
      # reshape the outputs back to original time-major format
      internal_states = tf.reshape(internal_states, [time_length, batch_size, -1])
      logits = [tf.reshape(logit, [time_length, batch_size, -1])
                for logit in logits]
      log_probs = [tf.reshape(log_prob, [time_length, batch_size])[:-1]
                   for log_prob in log_probs]
      entropies = [tf.reshape(ent, [time_length, batch_size])[:-1]
                   for ent in entropies]
      self_kls = [tf.reshape(self_kl, [time_length, batch_size])[:-1]
                  for self_kl in self_kls]
  
      return internal_states, logits, log_probs, entropies, self_kls
  