Get dictionary of transition data

def get_transitions(self, indexes):
        """ Get dictionary of transition data """
        assert indexes.shape[1] == self.state_buffer.shape[1], \
            "Must have the same number of indexes as there are environments"

        frame_batch_shape = (
                [indexes.shape[0], indexes.shape[1]]
                + list(self.state_buffer.shape[2:-1])
                + [self.state_buffer.shape[-1] * self.frame_history]
        )

        past_frame_buffer = np.zeros(frame_batch_shape, dtype=self.state_buffer.dtype)
        future_frame_buffer = np.zeros(frame_batch_shape, dtype=self.state_buffer.dtype)

        for buffer_idx, frame_row in enumerate(indexes):
            for env_idx, frame_idx in enumerate(frame_row):
                past_frame_buffer[buffer_idx, env_idx], future_frame_buffer[buffer_idx, env_idx] = (
                    self.get_frame_with_future(frame_idx, env_idx)
                )

        actions = take_along_axis(self.action_buffer, indexes)
        rewards = take_along_axis(self.reward_buffer, indexes)
        dones = take_along_axis(self.dones_buffer, indexes)

        transition_tensors = {
            'observations': past_frame_buffer,
            'actions': actions,
            'rewards': rewards,
            'observations_next': future_frame_buffer,
            'dones': dones.astype(np.float32),
        }

        for name in self.extra_data:
            transition_tensors[name] = take_along_axis(self.extra_data[name], indexes)

        return transition_tensors