code: def _sample_batch_prioritized(self, segment_tree, batch_size, history, forward_steps=1):
          
          p_total = segment_tree.total()
          segment = p_total / batch_size
  
          # Get batch of valid samples
          batch = [
              self._get_sample_from_segment(segment_tree, segment, i, history, forward_steps)
              for i in range(batch_size)
          ]
  
          probs, idxs, tree_idxs = zip(*batch)
          return np.array(probs), np.array(idxs), np.array(tree_idxs)
docs: Return indexes of the next sample in from prioritized distribution