Return frame from the buffer together with the next frame

def get_frame_with_future_forward_steps(self, frame_idx, env_idx, forward_steps, discount_factor):
        """ Return frame from the buffer together with the next frame """
        index_array = np.arange(frame_idx, frame_idx+forward_steps) % self.current_size

        if self.current_idx in index_array:
            raise VelException("Cannot provide enough future for the frame")

        past_frame = self.get_frame(frame_idx, env_idx)
        dones_array = self.dones_buffer[index_array, env_idx]

        rewards_array = self.reward_buffer[index_array, env_idx]
        discounted_rewards_array = rewards_array * (discount_factor ** np.arange(forward_steps))

        if dones_array.any():
            # Are we done between current frame and frame + n
            done = True

            dones_shifted = np.zeros_like(dones_array)
            dones_shifted[1:] = dones_array[:-1]
            reward = discounted_rewards_array[~np.logical_or.accumulate(dones_shifted)].sum()
            future_frame = np.zeros_like(past_frame)
        else:
            done = False
            reward = discounted_rewards_array.sum()

            if forward_steps >= self.frame_history:
                frame_indices = (index_array[:self.frame_history] + 1) % self.buffer_capacity
                future_frame = np.moveaxis(self.state_buffer[frame_indices, env_idx], 0, -2).reshape(past_frame.shape)
            else:
                frame_candidate = np.moveaxis(
                    self.state_buffer[(index_array + 1) % self.buffer_capacity, env_idx], 0, -2
                )
                frame_candidate_target_shape = (
                    list(frame_candidate.shape[:-2]) + [frame_candidate.shape[-2] * frame_candidate.shape[-1]]
                )

                future_frame = np.concatenate([
                    past_frame[..., (frame_candidate_target_shape[-1] - past_frame.shape[-1]):],
                    frame_candidate.reshape(frame_candidate_target_shape)
                ], -1)

        return past_frame, future_frame, reward, done