Reward-Free Curricula for Training Robust World Models

TL;DR

This paper introduces WAKER, a method for generating curricula in reward-free exploration to train robust world models that generalize across diverse environments, improving adaptability and performance.

Contribution

The paper proposes a novel approach to create curricula for reward-free training, linking minimax regret to model error, and introduces the WAKER algorithm for robust world model learning.

Findings

WAKER outperforms baselines in robustness and generalization.

WAKER improves training efficiency in diverse environments.

Theoretical connection between minimax regret and model error.

Abstract

There has been a recent surge of interest in developing generally-capable agents that can adapt to new tasks without additional training in the environment. Learning world models from reward-free exploration is a promising approach, and enables policies to be trained using imagined experience for new tasks. However, achieving a general agent requires robustness across different environments. In this work, we address the novel problem of generating curricula in the reward-free setting to train robust world models. We consider robustness in terms of minimax regret over all environment instantiations and show that the minimax regret can be connected to minimising the maximum error in the world model across environment instances. This result informs our algorithm, WAKER: Weighted Acquisition of Knowledge across Environments for Robustness. WAKER selects environments for data collection based on the estimated error of the world model for each environment. Our experiments demonstrate that WAKER outperforms several baselines, resulting in improved robustness, efficiency, and generalisation.