text: # coding=utf-8
  # Copyright 2020 The TF-Agents Authors.
  #
  # 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
  #
  #     https://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.
  
  """Implements the EXP3 bandit algorithm.
  
  Implementation based on
  
  "Bandit Algorithms"
    Lattimore and Szepesvari, 2019
    https://tor-lattimore.com/downloads/book/book.pdf
  """
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  from typing import Optional, Text
  
  import gin
  import tensorflow as tf
  
  from tf_agents.agents import data_converter
  from tf_agents.agents import tf_agent
  from tf_agents.bandits.policies import categorical_policy
  from tf_agents.policies import utils as policy_utilities
  from tf_agents.trajectories import policy_step
  from tf_agents.typing import types
  from tf_agents.utils import common
  
  
  def selective_sum(values: types.Tensor, partitions: types.Int,
                    num_partitions: int) -> types.Tensor:
    """Sums entries in `values`, partitioned using `partitions`.
  
    For example,
  
    ```python
       # Returns `[0 + 4 + 5, 2 + 3 + 4]` i.e. `[9, 6]`.
       selective_sum(values=[0, 1, 2, 3, 4, 5],
                     partitions=[0, 1, 1, 1, 0, 0]),
                     num_partitions=2)
    ```
  
    Args:
      values: a `Tensor` with numerical type.
      partitions: an integer `Tensor` with the same shape as `values`. Entry
        `partitions[i]` indicates the partition to which `values[i]` belongs.
      num_partitions: the number of partitions. All values in `partitions` must
        lie in `[0, num_partitions)`.
    Returns:
      A vector of size `num_partitions` with the same dtype as `values`. Entry `i`
      is the sum of all entries in `values` belonging to partition `i`.
    """
    partitioned_values = tf.dynamic_partition(values, partitions, num_partitions)
    return tf.stack([tf.reduce_sum(partition)
                     for partition in partitioned_values])
  
  
  def exp3_update_value(reward: types.Float,
                        log_prob: types.Float) -> types.Float:
    return 1. - (1. - reward) / tf.exp(log_prob)
  
  
  @gin.configurable
  class Exp3Agent(tf_agent.TFAgent):
    """An agent implementing the EXP3 bandit algorithm.
  
    Implementation based on
  
    "Bandit Algorithms"
      Lattimore and Szepesvari, 2019
      https://tor-lattimore.com/downloads/book/book.pdf
    """
  
    def __init__(self,
                 time_step_spec: types.TimeStep,
                 action_spec: types.BoundedTensorSpec,
                 learning_rate: float,
                 name: Optional[Text] = None):
      """Initialize an instance of `Exp3Agent`.
  
      Args:
        time_step_spec: A `TimeStep` spec describing the expected `TimeStep`s.
        action_spec: A scalar `BoundedTensorSpec` with `int32` or `int64` dtype
          describing the number of actions for this agent.
        learning_rate: A float valued scalar. A higher value will force the agent
          to converge on a single action more quickly. A lower value will
          encourage more exploration. This value corresponds to the
          `inverse_temperature` argument passed to `CategoricalPolicy`.
        name: a name for this instance of `Exp3Agent`.
      """
      tf.Module.__init__(self, name=name)
      common.tf_agents_gauge.get_cell('TFABandit').set(True)
      self._num_actions = policy_utilities.get_num_actions_from_tensor_spec(
          action_spec)
      self._weights = tf.compat.v2.Variable(
          tf.zeros(self._num_actions), name='weights')
      self._learning_rate = tf.compat.v2.Variable(
          learning_rate, name='learning_rate')
      policy = categorical_policy.CategoricalPolicy(
          weights=self._weights,
          time_step_spec=time_step_spec,
          action_spec=action_spec,
          inverse_temperature=self._learning_rate)
      # TODO(b/127462472): consider policy=GreedyPolicy(collect_policy).
      super(Exp3Agent, self).__init__(time_step_spec=time_step_spec,
                                      action_spec=policy.action_spec,
                                      policy=policy,
                                      collect_policy=policy,
                                      train_sequence_length=None)
      self._as_trajectory = data_converter.AsTrajectory(
          self.data_context, sequence_length=None)
  
    @property
    def num_actions(self):
      return self._num_actions
  
    @property
    def weights(self):
      return tf.identity(self._weights)
  
    @property
    def learning_rate(self):
      return tf.identity(self._learning_rate)
  
    @learning_rate.setter
    def learning_rate(self, learning_rate):
      return tf.compat.v1.assign(self._learning_rate, learning_rate)
  
    def _initialize(self):
      tf.compat.v1.variables_initializer(self.variables)
  
    def _train(self, experience, weights=None):
      """Updates the policy based on the data in `experience`.
  
      Note that `experience` should only contain data points that this agent has
      not previously seen. If `experience` comes from a replay buffer, this buffer
      should be cleared between each call to `train`.
  
      Args:
        experience: A batch of experience data in the form of a `Trajectory`.
        weights: Unused.
  
      Returns:
        A `LossInfo` containing the loss *before* the training step is taken.
          Note that the loss does not depend on policy state and comes directly
          from the experience (and is therefore not differentiable).
  
          In most cases, if `weights` is provided, the entries of this tuple will
          have been calculated with the weights.  Note that each Agent chooses
          its own method of applying weights.
      """
      del weights  # unused
      experience = self._as_trajectory(experience)
      reward = experience.reward
      log_prob = policy_step.get_log_probability(experience.policy_info)
      action = experience.action
      update_value = exp3_update_value(reward, log_prob)
      weight_update = selective_sum(values=update_value,
                                    partitions=action,
                                    num_partitions=self.num_actions)
      tf.compat.v1.assign_add(self._weights, weight_update)
  
      batch_size = tf.cast(tf.size(reward), dtype=tf.int64)
      self._train_step_counter.assign_add(batch_size)
  
      return tf_agent.LossInfo(loss=-tf.reduce_sum(experience.reward), extra=())
  