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

Fuente: arXiv
Enregistré dans:
Détails bibliographiques
Auteurs principaux: Prasanna, Sai, Farid, Karim, Rajan, Raghu, Biedenkapp, André
Format: Preprint
Publié: 2024
Sujets:
Accès en ligne:
Tags: Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
_version_ 1866910553954320384
author Prasanna, Sai
Farid, Karim
Rajan, Raghu
Biedenkapp, André
author_facet Prasanna, Sai
Farid, Karim
Rajan, Raghu
Biedenkapp, André
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2403_10967
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Dreaming of Many Worlds: Learning Contextual World Models Aids Zero-Shot Generalization
Prasanna, Sai
Farid, Karim
Rajan, Raghu
Biedenkapp, André
Machine Learning
Artificial Intelligence
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.
title Dreaming of Many Worlds: Learning Contextual World Models Aids Zero-Shot Generalization
topic Machine Learning
Artificial Intelligence
url https://arxiv.org/abs/2403.10967