Return frame from the buffer together with the next frame

def get_frame_with_future(self, frame_idx, history_length=1):
        """ Return frame from the buffer together with the next frame """
        if frame_idx == self.current_idx:
            raise VelException("Cannot provide enough future for the frame")

        past_frame = self.get_frame(frame_idx, history_length)

        if history_length > 1:
            assert self.state_buffer.shape[-1] == 1, \
                "State buffer must have last dimension of 1 if we want frame history"

        if not self.dones_buffer[frame_idx]:
            next_idx = (frame_idx + 1) % self.buffer_capacity
            next_frame = self.state_buffer[next_idx]
        else:
            next_idx = (frame_idx + 1) % self.buffer_capacity
            next_frame = np.zeros_like(self.state_buffer[next_idx])

        if history_length > 1:
            future_frame = np.concatenate([
                past_frame.take(indices=np.arange(1, past_frame.shape[-1]), axis=-1), next_frame
            ], axis=-1)
        else:
            future_frame = next_frame

        return past_frame, future_frame