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

• URL: http://arxiv.org/abs/2606.00780v1
• Date: Sat, 30 May 2026 15:53:03 GMT
• Title: Behavior-Invariant Task Representation Learning with Transformer-based World Models for Offline Meta-Reinforcement Learning
• Abstract summary: We propose a novel framework that integrates information-theoretic task representation learning with a Transformer-based world model.<n>Our method outperforms state-of-the-art approaches, with superior stability and under out-of-distribution and sparse-reward settings.
• Score: 13.491848927208373
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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.

Related papers

• Policy-Based Reinforcement Learning with Action Masking for Dynamic Job Shop Scheduling under Uncertainty: Handling Random Arrivals and Machine Failures [3.2880869992413246]
  We present a novel framework for solving Dynamic Job Shop Scheduling Problems under uncertainty.<n>Our approach follows a model-based paradigm, using Coloured Timed Petri Nets to represent the scheduling environment.<n>We conduct experiments on dynamic JSSP benchmarks, demonstrating that our method consistently outperforms traditional minimization and rule-based approaches in terms of makespan.
  arXiv  Detail & Related papers  (2026-01-14T08:53:46Z)
• MAESTRO: Meta-learning Adaptive Estimation of Scalarization Trade-offs for Reward Optimization [56.074760766965085]
  Group-Relative Policy Optimization has emerged as an efficient paradigm for aligning Large Language Models (LLMs)<n>We propose MAESTRO, which treats reward scalarization as a dynamic latent policy, leveraging the model's terminal hidden states as a semantic bottleneck.<n>We formulate this as a contextual bandit problem within a bi-level optimization framework, where a lightweight Conductor network co-evolves with the policy by utilizing group-relative advantages as a meta-reward signal.
  arXiv  Detail & Related papers  (2026-01-12T05:02:48Z)
• Balance Equation-based Distributionally Robust Offline Imitation Learning [8.607736795429638]
  Imitation Learning (IL) has proven highly effective for robotic and control tasks where manually designing reward functions or explicit controllers is infeasible.<n>Standard IL methods implicitly assume that the environment dynamics remain fixed between training and deployment.<n>We address this challenge through Balance Equation-based Distributionally Robust Offline Learning.<n>We formulate the problem as a distributionally robust optimization over an uncertainty set of transition models, seeking a policy that minimizes the imitation loss under the worst-case transition distribution.
  arXiv  Detail & Related papers  (2025-11-11T07:48:09Z)
• Multi-Modal Manipulation via Multi-Modal Policy Consensus [62.49978559936122]
  We propose a new approach to integrate diverse sensory modalities for robotic manipulation.<n>Our method factorizes the policy into a set of diffusion models, each specialized for a single representation.<n>We evaluate our approach on simulated manipulation tasks in RLBench, as well as real-world tasks such as occluded object picking, in-hand spoon reorientation, and puzzle insertion.
  arXiv  Detail & Related papers  (2025-09-27T19:43:04Z)
• Model-Based Offline Reinforcement Learning with Adversarial Data Augmentation [36.9134885948595]
  We introduce Model-based Offline Reinforcement learning with AdversariaL data augmentation.<n>In MORAL, we replace the fixed horizon rollout by employing adversaria data augmentation to execute alternating sampling with ensemble models.<n>Experiments on D4RL benchmark demonstrate that MORAL outperforms other model-based offline RL methods in terms of policy learning and sample efficiency.
  arXiv  Detail & Related papers  (2025-03-26T07:24:34Z)
• HarmoDT: Harmony Multi-Task Decision Transformer for Offline Reinforcement Learning [72.25707314772254]
  We introduce the Harmony Multi-Task Decision Transformer (HarmoDT), a novel solution designed to identify an optimal harmony subspace of parameters for each task. The upper level of this framework is dedicated to learning a task-specific mask that delineates the harmony subspace, while the inner level focuses on updating parameters to enhance the overall performance of the unified policy.
  arXiv  Detail & Related papers  (2024-05-28T11:41:41Z)
• Policy-Guided Diffusion [30.4597043728046]
  In many real-world settings, agents must learn from an offline dataset gathered by some prior behavior policy. We propose policy-guided diffusion as an alternative to autoregressive offline world models. We show that policy-guided diffusion models a regularized form of the target distribution that balances action likelihood under both the target and behavior policies.
  arXiv  Detail & Related papers  (2024-04-09T14:46:48Z)
• Equivariant Data Augmentation for Generalization in Offline Reinforcement Learning [10.00979536266327]
  We present a novel approach to address the challenge of generalization in offline reinforcement learning (RL) Specifically, we aim to improve the agent's ability to generalize to out-of-distribution goals. We learn a new policy offline based on the augmented dataset, with an off-the-shelf offline RL algorithm.
  arXiv  Detail & Related papers  (2023-09-14T10:22:33Z)
• Offline Reinforcement Learning with Closed-Form Policy Improvement Operators [88.54210578912554]
  Behavior constrained policy optimization has been demonstrated to be a successful paradigm for tackling Offline Reinforcement Learning. In this paper, we propose our closed-form policy improvement operators. We empirically demonstrate their effectiveness over state-of-the-art algorithms on the standard D4RL benchmark.
  arXiv  Detail & Related papers  (2022-11-29T06:29:26Z)
• Distributionally Adaptive Meta Reinforcement Learning [85.17284589483536]
  We develop a framework for meta-RL algorithms that behave appropriately under test-time distribution shifts. Our framework centers on an adaptive approach to distributional robustness that trains a population of meta-policies to be robust to varying levels of distribution shift. We show how our framework allows for improved regret under distribution shift, and empirically show its efficacy on simulated robotics problems.
  arXiv  Detail & Related papers  (2022-10-06T17:55:09Z)
• Latent-Variable Advantage-Weighted Policy Optimization for Offline RL [70.01851346635637]
  offline reinforcement learning methods hold the promise of learning policies from pre-collected datasets without the need to query the environment for new transitions. In practice, offline datasets are often heterogeneous, i.e., collected in a variety of scenarios. We propose to leverage latent-variable policies that can represent a broader class of policy distributions. Our method improves the average performance of the next best-performing offline reinforcement learning methods by 49% on heterogeneous datasets.
  arXiv  Detail & Related papers  (2022-03-16T21:17:03Z)
• A Regularized Implicit Policy for Offline Reinforcement Learning [54.7427227775581]
  offline reinforcement learning enables learning from a fixed dataset, without further interactions with the environment. We propose a framework that supports learning a flexible yet well-regularized fully-implicit policy. Experiments and ablation study on the D4RL dataset validate our framework and the effectiveness of our algorithmic designs.
  arXiv  Detail & Related papers  (2022-02-19T20:22:04Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.