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

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

Reinforcement Learning Journal, vol. 3, 2024, pp. 1317–1350.

Presented at the Reinforcement Learning Conference (RLC), Amherst Massachusetts, August 9–12, 2024.


Download:

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 Dreamer~(v3) \citep{hafner-arxiv23a}. 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 \url{https://github.com/sai-prasanna/dreaming_of_many_worlds}.


Citation Information:

Sai Prasanna, Karim Farid, Raghu Rajan, and André Biedenkapp. "Dreaming of Many Worlds: Learning Contextual World Models aids Zero-Shot Generalization." Reinforcement Learning Journal, vol. 3, 2024, pp. 1317–1350.

BibTeX:

@article{prasanna2024dreaming,
    title={Dreaming of Many Worlds: {L}earning Contextual World Models aids Zero-Shot Generalization},
    author={Prasanna, Sai and Farid, Karim and Rajan, Raghu and Biedenkapp, Andr{\'{e}}},
    journal={Reinforcement Learning Journal},
    volume={3},
    pages={1317--1350},
    year={2024}
}