Behavior-Invariant Task Representation Learning with Transformer-based World Models for Offline Meta-Reinforcement Learning
On the map: World models → Latent dynamics models and model-based RL
also in 1 more
Abstract (arXiv)
Offline meta-reinforcement learning leverages static datasets to enable agents to generalize to unseen environments by combining offline efficiency with meta-learning adaptability, yet it faces key challenges from context and policy distribution shifts. These issues hinder agents from adapting to online environments, and are further exacerbated under sparse-reward settings. As a result, agents often become trapped in an inherent pattern dilemma, failing to achieve robust generalization. In this work, we propose a novel framework that integrates information-theoretic task representation learning with a Transformer-based stochastic world model. Our approach extracts task-defining latent variables that are invariant to behavior policy, thereby effectively mitigating the context distribution shift. To further handle policy shift and model exploitation, we apply a conservative value penalty to imagination-based rollouts, preventing the policy from exploiting model inaccuracies while maintaining robust adaptation. Extensive evaluations demonstrate that our method outperforms state-of-the-art approaches, with superior stability and generalization under out-of-distribution and sparse-reward settings.
Key facts. No citing papers in the map yet (data through 7 Oct 2026). Placed in Latent dynamics models and model-based RL, one of 11 directions in World models. Builds on 13 of its 48 references that are in the map. First posted on arXiv 30 May 2026.
Builds on
13 of 48 references.
Closest ancestors · 8
Same direction, newest first
- 2025Training Agents Inside of Scalable World Models29 Sep
- 2025Learning Transformer-based World Models with Contrastive Predictive Coding6 Mar
- 2024MuDreamer: Learning Predictive World Models without Reconstruction23 May
- 2024MAMBA: an Effective World Model Approach for Meta-Reinforcement Learning14 Mar
- 2023Transformer-based World Models Are Happy With 100k Interactions13 Mar
- 2022Transformers are Sample-Efficient World Models1 Sep
- 2022TransDreamer: Reinforcement Learning with Transformer World Models19 Feb
- 2018World Models27 Mar
Other references · 5
- 2013Auto-Encoding Variational Bayescited by 3,980
- 2016RL2: Fast Reinforcement Learning via Slow Reinforcement Learningcited by 490
- 2019VariBAD: A Very Good Method for Bayes-Adaptive Deep RL via Meta-Learningcited by 210
- 2022Causal Dynamics Learning for Task-Independent State Abstractioncited by 64
- 2020FOCAL: Efficient Fully-Offline Meta-Reinforcement Learning via Distance Metric Learning and Behavior Regularizationcited by 6
Built on by
Not cited yet.
Similar, not cited
Close in wording, no citation either way
- 2024Meta-DT: Offline Meta-RL as Conditional Sequence Modeling with World Model Disentanglement
- 2024Entropy Regularized Task Representation Learning for Offline Meta-Reinforcement Learning
- 2026Temporal Consistency Improves Generalization in Contextual Offline Meta Reinforcement Learning
- 2022Robust Task Representations for Offline Meta-Reinforcement Learning via Contrastive Learning
- 2024Scrutinize What We Ignore: Reining In Task Representation Shift Of Context-Based Offline Meta Reinforcement Learning
- 2023Generalizable Task Representation Learning for Offline Meta-Reinforcement Learning with Data Limitations
- 2024Disentangling Policy from Offline Task Representation Learning via Adversarial Data Augmentation
- 2024Towards an Information Theoretic Framework of Context-Based Offline Meta-Reinforcement Learning