Return batch of frames for given indexes

def _get_transitions(self, probs, indexes, tree_idxs, batch_info, forward_steps=1, discount_factor=1.0):
        """ Return batch of frames for given indexes """
        if forward_steps > 1:
            transition_arrays = self.backend.get_transitions_forward_steps(indexes, forward_steps, discount_factor)
        else:
            transition_arrays = self.backend.get_transitions(indexes)

        priority_weight = self.priority_weight.value(batch_info['progress'])

        # Normalize by sum of all probs
        probs = probs / np.array([s.total() for s in self.backend.segment_trees], dtype=float).reshape(1, -1)
        capacity = self.backend.current_size
        weights = (capacity * probs) ** (-priority_weight)
        weights = weights / weights.max(axis=0, keepdims=True)

        transition_arrays['weights'] = weights
        transition_tensors = {k: torch.from_numpy(v) for k, v in transition_arrays.items()}

        transitions = Trajectories(
            num_steps=indexes.shape[0],
            num_envs=indexes.shape[1],
            environment_information=None,
            transition_tensors=transition_tensors,
            rollout_tensors={},
            extra_data={
                'tree_idxs': tree_idxs
            }
        )

        return transitions.to_transitions()