# Behavior-Invariant Task Representation Learning with Transformer-based World Models for Offline Meta-Reinforcement Learning

arXiv: https://arxiv.org/abs/2606.00780 · HTML: https://mlsanity.com/paper/2606.00780/

## Key facts

- No citing papers in the map yet (data through 14 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.

## Placed in

- World models / [Latent dynamics models and model-based RL](https://mlsanity.com/wm/latent-dynamics-models-and-model-based-rl/)
- Offline RL and imitation learning from data / [Decision transformers and sequence modeling](https://mlsanity.com/offlinerl/decision-transformers-and-sequence-modeling/)

## Builds on (most cited first, 13 of 13)

- [Auto-Encoding Variational Bayes](https://mlsanity.com/paper/1312.6114/) (2013)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [RL2: Fast Reinforcement Learning via Slow Reinforcement Learning](https://mlsanity.com/paper/1611.02779/) (2016)
- [Transformers are Sample-Efficient World Models](https://mlsanity.com/paper/2209.00588/) (2022)
- [VariBAD: A Very Good Method for Bayes-Adaptive Deep RL via Meta-Learning](https://mlsanity.com/paper/1910.08348/) (2019)
- [TransDreamer: Reinforcement Learning with Transformer World Models](https://mlsanity.com/paper/2202.09481/) (2022)
- [Transformer-based World Models Are Happy With 100k Interactions](https://mlsanity.com/paper/2303.07109/) (2023)
- [Training Agents Inside of Scalable World Models](https://mlsanity.com/paper/2509.24527/) (2025)
- [Causal Dynamics Learning for Task-Independent State Abstraction](https://mlsanity.com/paper/2206.13452/) (2022)
- [Learning Transformer-based World Models with Contrastive Predictive Coding](https://mlsanity.com/paper/2503.04416/) (2025)
- [MAMBA: an Effective World Model Approach for Meta-Reinforcement Learning](https://mlsanity.com/paper/2403.09859/) (2024)
- [MuDreamer: Learning Predictive World Models without Reconstruction](https://mlsanity.com/paper/2405.15083/) (2024)
- [FOCAL: Efficient Fully-Offline Meta-Reinforcement Learning via Distance Metric Learning and Behavior Regularization](https://mlsanity.com/paper/2010.01112/) (2020)

## 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.