Dreaming of Many Worlds: Learning Contextual World Models aids Zero-Shot Generalization

Sai Prasanna, Karim Farid, Raghu Rajan, André Biedenkapp

Keywords: Zero-Shot Generalization, Model-based, Dreamer, Contextual Reinforcement Learning

TL;DR: We study zero-shot generalization in MBRL, in particular for Dreamer, by means of using contextual information and devise a novel way of encoding context into the latent space.

Abstract: 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 the 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 experiments show that such systematic incorporation of the context improves the ZSG of the policies trained on the ``dreams'' of the world model. The code for all our experiments is available at https://anonymous.4open.science/r/dreaming_many_worlds.

Already Accepted Paper At Another Venue: already accepted somewhere else

Submission Number: 81