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

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.

This paper studies how an AI agent can adapt to new tasks after learning only from previously collected data. This is difficult because the agent learns from past behavior data, then must adapt using its own actions, often with only limited feedback. We introduce MetaSTAR, a method that helps the agent learn more reliable task knowledge from offline data. MetaSTAR builds an internal model of how the environment changes over time, allowing the agent to imagine possible future outcomes before acting. It also learns task descriptions that focus on stable information about the task itself, rather than accidental patterns caused by how the training data was collected. We test MetaSTAR on several control tasks, including cases where feedback is sparse and test tasks differ from training tasks. The results show that MetaSTAR often adapts more reliably than existing methods. This suggests that learning internal environment models can help offline-trained agents transfer more effectively to new situations.