code: def sample_batch_transitions(self, batch_size, forward_steps=1):
          
          batches = [
              self._sample_batch_prioritized(
                  segment_tree, batch_size, self.deque.frame_history, forward_steps=forward_steps
              )
              for segment_tree in self.segment_trees
          ]
  
          probs, idxs, tree_idxs = zip(*batches)
  
          return np.stack(probs, axis=1).astype(float), np.stack(idxs, axis=1), np.stack(tree_idxs, axis=1)
docs: Return indexes of next sample