text: import numpy as np
  import random
  
  from baselines.common.segment_tree import SumSegmentTree, MinSegmentTree
  
  
  class ReplayBuffer(object):
      def __init__(self, size):
          """Create Replay buffer.
  
          Parameters
          ----------
          size: int
              Max number of transitions to store in the buffer. When the buffer
              overflows the old memories are dropped.
          """
          self._storage = []
          self._maxsize = size
          self._next_idx = 0
  
      def __len__(self):
          return len(self._storage)
  
      def add(self, obs_t, action, reward, obs_tp1, done):
          data = (obs_t, action, reward, obs_tp1, done)
  
          if self._next_idx >= len(self._storage):
              self._storage.append(data)
          else:
              self._storage[self._next_idx] = data
          self._next_idx = (self._next_idx + 1) % self._maxsize
  
      def _encode_sample(self, idxes):
          obses_t, actions, rewards, obses_tp1, dones = [], [], [], [], []
          for i in idxes:
              data = self._storage[i]
              obs_t, action, reward, obs_tp1, done = data
              obses_t.append(np.array(obs_t, copy=False))
              actions.append(np.array(action, copy=False))
              rewards.append(reward)
              obses_tp1.append(np.array(obs_tp1, copy=False))
              dones.append(done)
          return np.array(obses_t), np.array(actions), np.array(rewards), np.array(obses_tp1), np.array(dones)
  
      def sample(self, batch_size):
          """Sample a batch of experiences.
  
          Parameters
          ----------
          batch_size: int
              How many transitions to sample.
  
          Returns
          -------
          obs_batch: np.array
              batch of observations
          act_batch: np.array
              batch of actions executed given obs_batch
          rew_batch: np.array
              rewards received as results of executing act_batch
          next_obs_batch: np.array
              next set of observations seen after executing act_batch
          done_mask: np.array
              done_mask[i] = 1 if executing act_batch[i] resulted in
              the end of an episode and 0 otherwise.
          """
          idxes = [random.randint(0, len(self._storage) - 1) for _ in range(batch_size)]
          return self._encode_sample(idxes)
  
  
  class PrioritizedReplayBuffer(ReplayBuffer):
      def __init__(self, size, alpha=0.8):
          """Create Prioritized Replay buffer.
  
          Parameters
          ----------
          size: int
              Max number of transitions to store in the buffer. When the buffer
              overflows the old memories are dropped.
          alpha: float
              how much prioritization is used
              (0 - no prioritization, 1 - full prioritization)
  
          See Also
          --------
          ReplayBuffer.__init__
          """
          super(PrioritizedReplayBuffer, self).__init__(size)
          assert alpha >= 0
          self._alpha = alpha
  
          it_capacity = 1
          while it_capacity < size:
              it_capacity *= 2
  
          self._it_sum = SumSegmentTree(it_capacity)
          self._it_min = MinSegmentTree(it_capacity)
          self._max_priority = 1.0
  
      def add(self, *args, **kwargs):
          """See ReplayBuffer.store_effect"""
          idx = self._next_idx
          super().add(*args, **kwargs)
          self._it_sum[idx] = self._max_priority ** self._alpha
          self._it_min[idx] = self._max_priority ** self._alpha
  
      def _sample_proportional(self, batch_size):
          res = []
          p_total = self._it_sum.sum(0, len(self._storage) - 1)
          every_range_len = p_total / batch_size
          for i in range(batch_size):
              mass = random.random() * every_range_len + i * every_range_len
              idx = self._it_sum.find_prefixsum_idx(mass)
              res.append(idx)
          return res
  
      def sample(self, batch_size, beta=None):
          """Sample a batch of experiences.
  
          compared to ReplayBuffer.sample
          it also returns importance weights and idxes
          of sampled experiences.
  
  
          Parameters
          ----------
          batch_size: int
              How many transitions to sample.
          beta: float
              To what degree to use importance weights
              (0 - no corrections, 1 - full correction)
  
          Returns
          -------
          obs_batch: np.array
              batch of observations
          act_batch: np.array
              batch of actions executed given obs_batch
          rew_batch: np.array
              rewards received as results of executing act_batch
          next_obs_batch: np.array
              next set of observations seen after executing act_batch
          done_mask: np.array
              done_mask[i] = 1 if executing act_batch[i] resulted in
              the end of an episode and 0 otherwise.
          weights: np.array
              Array of shape (batch_size,) and dtype np.float32
              denoting importance weight of each sampled transition
          idxes: np.array
              Array of shape (batch_size,) and dtype np.int32
              idexes in buffer of sampled experiences
          """
          if beta is None:
              beta = self.alpha
  
          assert beta > 0
  
          idxes = self._sample_proportional(batch_size)
  
          weights = []
          p_min = self._it_min.min() / self._it_sum.sum()
          max_weight = (p_min * len(self._storage)) ** (-beta)
  
          for idx in idxes:
              p_sample = self._it_sum[idx] / self._it_sum.sum()
              weight = (p_sample * len(self._storage)) ** (-beta)
              weights.append(weight / max_weight)
          weights = np.array(weights)
          encoded_sample = self._encode_sample(idxes)
          # return tuple(list(encoded_sample) + [weights, idxes])
          return tuple(list(encoded_sample))
  
      def update_priorities(self, idxes, priorities):
          """Update priorities of sampled transitions.
  
          sets priority of transition at index idxes[i] in buffer
          to priorities[i].
  
          Parameters
          ----------
          idxes: [int]
              List of idxes of sampled transitions
          priorities: [float]
              List of updated priorities corresponding to
              transitions at the sampled idxes denoted by
              variable `idxes`.
          """
          assert len(idxes) == len(priorities)
          for idx, priority in zip(idxes, priorities):
              assert priority > 0
              assert 0 <= idx < len(self._storage)
              self._it_sum[idx] = priority ** self._alpha
              self._it_min[idx] = priority ** self._alpha
  
              self._max_priority = max(self._max_priority, priority)
  