Return frame from the buffer together with the next frame

def get_frame_with_future(self, frame_idx, env_idx):
        """ 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, env_idx)

        if not self.dones_buffer[frame_idx, env_idx]:
            # We're not done
            next_idx = (frame_idx + 1) % self.buffer_capacity
            next_frame = self.state_buffer[next_idx, env_idx]

            if self.frame_history > 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
        else:
            # We are done
            future_frame = np.zeros_like(past_frame)

        return past_frame, future_frame