Dreaming of Many Worlds: Learning Contextual World Models Aids Zero-Shot Generalization
This work addresses the challenge of creating generally capable embodied agents by improving zero-shot generalization in contextual RL, though it is incremental as it builds on existing Dreamer methods.
The paper tackles zero-shot generalization to unseen dynamics in contextual reinforcement learning by proposing a contextual recurrent state-space model (cRSSM) that incorporates context into the world model of Dreamer (v3). The approach improves zero-shot generalization on tasks from the CARL benchmark suite, with qualitative results showing disentanglement of latent state from context for extrapolation to unseen contexts.
Zero-shot generalization (ZSG) to unseen dynamics is a major challenge for creating generally capable embodied agents. To address the broader challenge, we start with the simpler setting of contextual reinforcement learning (cRL), assuming observability of the context values that parameterize the variation in the system's dynamics, such as the mass or dimensions of a robot, without making further simplifying assumptions about the observability of the Markovian state. Toward the goal of ZSG to unseen variation in context, we propose the contextual recurrent state-space model (cRSSM), which introduces changes to the world model of Dreamer (v3) (Hafner et al., 2023). This allows the world model to incorporate context for inferring latent Markovian states from the observations and modeling the latent dynamics. Our approach is evaluated on two tasks from the CARL benchmark suite, which is tailored to study contextual RL. Our experiments show that such systematic incorporation of the context improves the ZSG of the policies trained on the "dreams" of the world model. We further find qualitatively that our approach allows Dreamer to disentangle the latent state from context, allowing it to extrapolate its dreams to the many worlds of unseen contexts. The code for all our experiments is available at https://github.com/sai-prasanna/dreaming_of_many_worlds.