text: # coding=utf-8
  # Copyright 2018 The Dopamine 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
  #
  #      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.
  """A sum tree data structure.
  
  Used for prioritized experience replay. See prioritized_replay_buffer.py
  and Schaul et al. (2015).
  """
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import math
  import random
  
  import numpy as np
  
  
  class SumTree(object):
    """A sum tree data structure for storing replay priorities.
  
    A sum tree is a complete binary tree whose leaves contain values called
    priorities. Internal nodes maintain the sum of the priorities of all leaf
    nodes in their subtree.
  
    For capacity = 4, the tree may look like this:
  
                 +---+
                 |2.5|
                 +-+-+
                   |
           +-------+--------+
           |                |
         +-+-+            +-+-+
         |1.5|            |1.0|
         +-+-+            +-+-+
           |                |
      +----+----+      +----+----+
      |         |      |         |
    +-+-+     +-+-+  +-+-+     +-+-+
    |0.5|     |1.0|  |0.5|     |0.5|
    +---+     +---+  +---+     +---+
  
    This is stored in a list of numpy arrays:
    self.nodes = [ [2.5], [1.5, 1], [0.5, 1, 0.5, 0.5] ]
  
    For conciseness, we allocate arrays as powers of two, and pad the excess
    elements with zero values.
  
    This is similar to the usual array-based representation of a complete binary
    tree, but is a little more user-friendly.
    """
  
    def __init__(self, capacity):
      """Creates the sum tree data structure for the given replay capacity.
  
      Args:
        capacity: int, the maximum number of elements that can be stored in this
          data structure.
  
      Raises:
        ValueError: If requested capacity is not positive.
      """
      assert isinstance(capacity, int)
      if capacity <= 0:
        raise ValueError('Sum tree capacity should be positive. Got: {}'.
                         format(capacity))
  
      self.nodes = []
      tree_depth = int(math.ceil(np.log2(capacity)))
      level_size = 1
      for _ in range(tree_depth + 1):
        nodes_at_this_depth = np.zeros(level_size)
        self.nodes.append(nodes_at_this_depth)
  
        level_size *= 2
  
      self.max_recorded_priority = 1.0
  
    def _total_priority(self):
      """Returns the sum of all priorities stored in this sum tree.
  
      Returns:
        float, sum of priorities stored in this sum tree.
      """
      return self.nodes[0][0]
  
    def sample(self, query_value=None):
      """Samples an element from the sum tree.
  
      Each element has probability p_i / sum_j p_j of being picked, where p_i is
      the (positive) value associated with node i (possibly unnormalized).
  
      Args:
        query_value: float in [0, 1], used as the random value to select a
        sample. If None, will select one randomly in [0, 1).
  
      Returns:
        int, a random element from the sum tree.
  
      Raises:
        Exception: If the sum tree is empty (i.e. its node values sum to 0), or if
          the supplied query_value is larger than the total sum.
      """
      if self._total_priority() == 0.0:
        raise Exception('Cannot sample from an empty sum tree.')
  
      if query_value and (query_value < 0. or query_value > 1.):
        raise ValueError('query_value must be in [0, 1].')
  
      # Sample a value in range [0, R), where R is the value stored at the root.
      query_value = random.random() if query_value is None else query_value
      query_value *= self._total_priority()
  
      # Now traverse the sum tree.
      node_index = 0
      for nodes_at_this_depth in self.nodes[1:]:
        # Compute children of previous depth's node.
        left_child = node_index * 2
  
        left_sum = nodes_at_this_depth[left_child]
        # Each subtree describes a range [0, a), where a is its value.
        if query_value < left_sum:  # Recurse into left subtree.
          node_index = left_child
        else:  # Recurse into right subtree.
          node_index = left_child + 1
          # Adjust query to be relative to right subtree.
          query_value -= left_sum
  
      return node_index
  
    def stratified_sample(self, batch_size):
      """Performs stratified sampling using the sum tree.
  
      Let R be the value at the root (total value of sum tree). This method will
      divide [0, R) into batch_size segments, pick a random number from each of
      those segments, and use that random number to sample from the sum_tree. This
      is as specified in Schaul et al. (2015).
  
      Args:
        batch_size: int, the number of strata to use.
      Returns:
        list of batch_size elements sampled from the sum tree.
  
      Raises:
        Exception: If the sum tree is empty (i.e. its node values sum to 0).
      """
      if self._total_priority() == 0.0:
        raise Exception('Cannot sample from an empty sum tree.')
  
      bounds = np.linspace(0., 1., batch_size + 1)
      assert len(bounds) == batch_size + 1
      segments = [(bounds[i], bounds[i+1]) for i in range(batch_size)]
      query_values = [random.uniform(x[0], x[1]) for x in segments]
      return [self.sample(query_value=x) for x in query_values]
  
    def get(self, node_index):
      """Returns the value of the leaf node corresponding to the index.
  
      Args:
        node_index: The index of the leaf node.
      Returns:
        The value of the leaf node.
      """
      return self.nodes[-1][node_index]
  
    def set(self, node_index, value):
      """Sets the value of a leaf node and updates internal nodes accordingly.
  
      This operation takes O(log(capacity)).
      Args:
        node_index: int, the index of the leaf node to be updated.
        value: float, the value which we assign to the node. This value must be
          nonnegative. Setting value = 0 will cause the element to never be
          sampled.
  
      Raises:
        ValueError: If the given value is negative.
      """
      if value < 0.0:
        raise ValueError('Sum tree values should be nonnegative. Got {}'.
                         format(value))
      self.max_recorded_priority = max(value, self.max_recorded_priority)
  
      delta_value = value - self.nodes[-1][node_index]
  
      # Now traverse back the tree, adjusting all sums along the way.
      for nodes_at_this_depth in reversed(self.nodes):
        # Note: Adding a delta leads to some tolerable numerical inaccuracies.
        nodes_at_this_depth[node_index] += delta_value
        node_index //= 2
  
      assert node_index == 0, ('Sum tree traversal failed, final node index '
                               'is not 0.')
  