Return indexes of next sample

def sample_batch_transitions(self, batch_size, forward_steps=1):
        """ Return indexes of next sample"""
        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)