Decision Transformer with Task-Aware Representations for Offline Meta-Reinforcement Learning

Xinbo Zhao; Yuhang Liu; Xin Zhang; Lei Le; Zhaohan Xi; Yingxue Zhang

cs.LGNov 2, 2026

Decision Transformer with Task-Aware Representations for Offline Meta-Reinforcement Learning

Authors: Xinbo Zhao, Yuhang Liu, Xin Zhang, Lei Le, Zhaohan Xi, Yingxue Zhang

Organizations: State University of New York at Binghamton · San Diego State University · LinkedIn

Abstract

To enable effective generalization in offline meta-RL, it is essential to accurately infer the underlying task through latent task representations, which provide a compact and transferable summary of task-specific characteristics for guiding policy generation. In this paper, we propose Contrastive Task-Aware Decision Transformer (CoTA-DT), a DT-based offline meta-RL framework that learns task representations by jointly integrating four complementary components: (i) contrastive learning from offline trajectories, (ii) alignment with natural-language task descriptions, (iii) one-step supervision via a task-conditioned world model, and (iv) long-term supervision via a task-conditioned Q function. The learned task representations are used as conditioning signals for a DT-based policy, enabling task-aware meta-policy learning. To further support generalization in both zero-shot and few-shot settings, we design a language-grounded, self-adaptive strategy that enables effective and reliable adaptation to unseen tasks. Notably, in CoTA-DT, the pretrained LLM serves as a shared backbone for language encoding, world modeling, Q-function estimation, and policy learning. This unified design facilitates parameter sharing, improving both efficiency and generalization. Extensive experiments across diverse offline meta-RL benchmarks demonstrate that CoTA-DT consistently outperforms prior methods, particularly in challenging zero-shot generalization scenarios.

Figures & tables

Figure 1: Overview of CoTA-DT. (a) Multi-perspective task representation learning jointly optimizes four signals— (i) contrastive task encoding from offline trajectories, (ii) language alignment between $u_{i}$ and $w_{i}$ , (iii) a task-conditioned world model predicting $\hat{r}_{t},\hat{s}_{t+1}$ , and (iv) a long-horizon Q head predicting $\hat{Q}_{t}$ —all sharing one frozen LLM with a single LoRA adapter. (b) A Decision Transformer policy is trained on the same LLM backbone, conditioned on $z_{t}$ . (c) At zero-shot adaptation, a new task description is encoded into $w_{\star}$ , $z_{0}$ is constructed by aligning with $w_{\star}$ , and the rollout uses self-adaptive $z_{t}$ updates.

Figure 2: Zero-shot test return curves of CoTA-DT against baselines on Mixed datasets.

	BC-Z	BAKU	Prompt-DT	Meta-DT	UNICORN	T2DA-T	CoTA-DT
Pt-Robot	$-13.0$ $\pm 0.7$	$-14.5$ $\pm 0.8$	$-15.2$ $\pm 0.6$	$-12.2$ $\pm 0.5$	$-13.1$ $\pm 0.7$	$-8.7$ $\pm 0.3$	$\mathbf{-7.9}$ $\pm\mathbf{0.4}$
Cheetah-Vel	$-129.2$ $\pm 10.0$	$-112.3$ $\pm 4.2$	$-136.9$ $\pm 17.5$	$-99.7$ $\pm 6.9$	$-101.9$ $\pm 18.0$	$-90.7$ $\pm 5.8$	$\mathbf{-76.5}$ $\pm\mathbf{4.7}$
Ant-Dir	$611.7$ $\pm 30.9$	$695.8$ $\pm 25.5$	$395.8$ $\pm 12.6$	$411.0$ $\pm 13.3$	$426.1$ $\pm 23.8$	$955.0$ $\pm 9.0$	$\mathbf{981.3}$ $\pm\mathbf{16.8}$
Meta-World	$910.3$ $\pm 27.1$	$970.9$ $\pm 36.2$	$867.4$ $\pm 46.5$	$994.2$ $\pm 45.0$	$756.1$ $\pm 60.8$	$1103.5$ $\pm 36.3$	$\mathbf{1292.5}$ $\pm\mathbf{28.2}$

Table 1: Zero-shot test return on the Mixed dataset. Best in bold , second-best underlined .

	BC-Z	BAKU	Prompt-DT	Meta-DT	UNICORN	T2DA-T	CoTA-DT
Pt-Robot	$-12.3$ $\pm 0.8$	$-13.4$ $\pm 0.9$	$-12.0$ $\pm 1.3$	$-10.2$ $\pm 0.6$	$-12.9$ $\pm 1.2$	$-7.6$ $\pm 0.4$	$\mathbf{-6.9}$ $\pm\mathbf{0.2}$
Cheetah-Vel	$-109.2$ $\pm 11.2$	$-97.5$ $\pm 4.3$	$-101.3$ $\pm 10.6$	$-86.5$ $\pm 7.3$	$-89.7$ $\pm 19.8$	$-78.8$ $\pm 4.1$	$\mathbf{-67.3}$ $\pm\mathbf{5.2}$
Ant-Dir	$696.5$ $\pm 39.8$	$817.6$ $\pm 20.1$	$509.0$ $\pm 13.4$	$534.2$ $\pm 14.0$	$459.4$ $\pm 15.5$	$1012.3$ $\pm 9.8$	$\mathbf{1040.2}$ $\pm\mathbf{15.8}$
Meta-World	$1131.7$ $\pm 56.3$	$1129.2$ $\pm 24.8$	$972.8$ $\pm 31.2$	$1083.4$ $\pm 49.5$	$826.8$ $\pm 66.9$	$1189.7$ $\pm 45.3$	$\mathbf{1402.8}$ $\pm\mathbf{30.1}$

Table 2: Few-shot test return on the Mixed dataset. Best in bold , second-best underlined .

Quality		BC-Z	BAKU	Prompt-DT	Meta-DT	T2DA-T	CoTA-DT
Expert	Pt-Robot	$-12.3$ $\pm 0.2$	$-13.7$ $\pm 0.2$	$-15.8$ $\pm 0.3$	$-11.2$ $\pm 0.1$	$-7.3$ $\pm 0.2$	$\mathbf{-6.9}$ $\pm\mathbf{0.4}$
Expert	Meta-World	$875.1$ $\pm 30.3$	$853.3$ $\pm 16.8$	$755.1$ $\pm 38.0$	$830.2$ $\pm 18.9$	$1114.0$ $\pm 59.3$	$\mathbf{1375.7}$ $\pm\mathbf{21.6}$
Medium	Pt-Robot	$-16.1$ $\pm 0.5$	$-14.8$ $\pm 0.6$	$-17.9$ $\pm 0.4$	$-12.6$ $\pm 0.3$	$-12.4$ $\pm 0.5$	$\mathbf{-11.8}$ $\pm\mathbf{0.6}$
Medium	Meta-World	$682.7$ $\pm 54.4$	$679.2$ $\pm 13.1$	$583.4$ $\pm 29.5$	$774.0$ $\pm 35.4$	$1003.5$ $\pm 78.2$	$\mathbf{1055.2}$ $\pm\mathbf{40.3}$

Table 3: Zero-shot return on Expert and Medium datasets. Best in bold , second-best underlined .

Environment	CoTA-DT (full)	w/o $\mathcal{L}_{\mathrm{NCE}}$	w/o $\mathcal{L}_{\mathrm{align}}$	w/o $\mathcal{L}_{\mathrm{rew}}+\mathcal{L}_{\mathrm{dyn}}$	w/o $\mathcal{L}_{Q}$	w/o LLM
Pt-Robot $\uparrow$	$\mathbf{-7.9}\pm\mathbf{0.4}$	$-10.9\pm 0.1$	$-12.3\pm 0.5$	$-10.8\pm 0.3$	$-10.2\pm 0.2$	$-8.4\pm 0.5$
Cheetah-Vel $\uparrow$	$\mathbf{-76.5}\pm\mathbf{4.7}$	$-95.4\pm 6.1$	$-110.8\pm 7.0$	$-99.1\pm 5.8$	$-87.3\pm 6.4$	$-84.6\pm 5.9$
Ant-Dir $\uparrow$	$\mathbf{981.3}\pm\mathbf{16.8}$	$906.7\pm 22.5$	$794.2\pm 26.8$	$874.6\pm 21.0$	$917.1\pm 19.3$	$957.3\pm 14.1$
Meta-World $\uparrow$	$\mathbf{1292.5}\pm\mathbf{28.2}$	$1018.5\pm 25.0$	$868.0\pm 36.7$	$947.3\pm 24.5$	$1106.8\pm 34.9$	$1218.2\pm 37.6$

Table 4: Component and backbone ablations across all four benchmarks (Mixed, zero-shot, mean

\pm

std over 5 seeds).

Figure 3: Task representations on Point-Robot. Each marker is color-coded by goal angle; circles denote task representations, and stars denote text embeddings $w_{i}$ of task descriptions. The top row visualizes aggregated representations $\bar{z}_{i}$ versus $w_{i}$ , while the bottom row shows all $\{z_{i}\}$ learned from sub-trajectories versus $w_{i}$ . Columns compare three types of representations: (a) & (d) show representations learned solely from offline data using the contrastive task encoder; (b) & (e) are language-aligned task representations (task encoding + language alignment); (c) & (f) are the final representations learned by CoTA-DT (task encoding + language alignment + world modeling + Q learning).

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Environment	# Tasks	Train/Test	$d_{s}$	$d_{a}$	$T$	Task variation
Point-Robot	50	45/5	2	2	20	Target $(x,y)$ position
Cheetah-Vel	50	45/5	20	6	200	Target velocity $\in[0.075,3]$
Ant-Dir	50	45/5	27	8	200	Target heading $\in[0,2\pi]$
Meta-World	22	18/4	39	4	200	Diverse manipulation

Appendix

Table 5: Environment statistics.

d_{s}

d_{a}

and

T

are the state, action and episode-length dimensions.

Group	Hyperparameter	Value
Backbone	LLM backbone	LLaMA-3.2 1B (frozen)
	Hidden size	2048
	LoRA rank $r$ / $\alpha$	8 / 16
	LoRA target modules	q_proj , v_proj (all layers)
Context encoder $E_{\phi}$	Architecture	1-layer GRU
	Hidden / output dim	256 / $d_{z}=64$

Appendix

Table 6: CoTA-DT hyperparameters. Values are shared across all four benchmarks unless noted.

Hyperparameter	Value	Test return $\uparrow$
Default (CoTA-DT)	—	$\mathbf{-7.9}$ $\pm\mathbf{0.4}$
$\lambda_{2}$ ( $\mathcal{L}_{\mathrm{align}}$ weight, default $0.5$ )	$0.0$	$-12.3$ $\pm 0.5$
	$0.1$	$-10.1$ $\pm 0.6$
	$1.0$	$-8.4$ $\pm 0.5$
	$2.0$	$-9.7$ $\pm 0.7$
$\lambda_{4}$ ( $\mathcal{L}_{Q}$ weight, default $0.1$ )	$0.0$	$-10.2$ $\pm 0.2$

Appendix

Table 7: Hyperparameter sensitivity of CoTA-DT on Point-Robot (Mixed dataset, zero-shot test return, mean

\pm

std over 5 seeds). Default config:

\lambda_{1}{=}0.1

\lambda_{2}{=}0.5

\lambda_{3}{=}1.0

\lambda_{4}{=}0.1

\alpha{=}1.0

H{=}20

, LoRA rank

r{=}8

\kappa_{\beta}

init

=0.07

Backbone	# params	Hidden dim	Trainable LoRA	Pt-Robot $\uparrow$	Meta-World $\uparrow$
GPT-2 small ( Radford et al., 2019 )	117M	768	0.39M	$-8.5\pm 0.6$	$1186.7\pm 29.6$
Qwen-2.5 0.5B ( Yang and others, 2024 )	0.5B	1024	0.55M	$-8.2\pm 0.5$	$1228.4\pm 35.2$
Qwen-2.5 1.5B ( Yang and others, 2024 )	1.5B	1536	0.95M	$-8.0\pm 0.4$	$1271.1\pm 31.8$
LLaMA-3.2 1B ( Touvron et al., 2023 )	1.2B	2048	0.85M	$\mathbf{-7.9\pm 0.4}$	$\mathbf{1292.5\pm 28.2}$

Appendix

Table 8: Effect of LLM backbone on Point-Robot and Meta-World Mixed (zero-shot test return). Point-Robot results are mean

\pm

std over 5 seeds.

Description style	Test return $\uparrow$	$\Delta$ vs. Standard
Minimal	$-8.8\pm 0.5$	$-0.9$
Standard (default)	$\mathbf{-7.9\pm 0.4}$	—
Noisy	$-11.5\pm 0.5$	$-3.6$

Appendix

Table 9: Robustness of CoTA-DT to description style on Point-Robot Mixed (zero-shot, 5 seeds).

Configuration	Text encoder	Seeds	Pt-Robot $\uparrow$	Meta-World $\uparrow$
T2DA-T (reproduced)	Pretrained CLIP	5	$-8.7\pm 0.3$	$1103.5\pm 36.3$
w/o LLM (small DT)	Separate pretrained CLIP	5	$-8.4\pm 0.5$	$1218.2\pm 37.6$
Random-init 1B + LoRA	Random-init shared backbone	5	$-11.2\pm 0.6$	$916.8\pm 36.0$
CoTA-DT (full)	Pretrained shared backbone	5	$-7.9\pm 0.4$	$1292.5\pm 28.2$

Appendix

Table 10: Decision-backbone controls on Mixed datasets in the zero-shot setting. Results are mean

\pm

std.

	Point-Robot $\uparrow$			Meta-World $\uparrow$
Mode	$k{=}3$	$k{=}5$	$k{=}10$	$k{=}3$	$k{=}5$	$k{=}10$
static	$-9.2$	$-8.6$	$-8.9$	$1281.3$	$1295.1$	$1290.5$
online	$-9.2$	$-8.6$	$-8.9$	$1295.8$	$1318.0$	$1325.2$
inversion only	$-8.4$	$\mathbf{-6.9}$	$\mathbf{-7.2}$	$1349.1$	$1384.7$	$\mathbf{1396.9}$
finetune	$\mathbf{-8.1}$	$\mathbf{-6.9}$	$-7.3$	$\mathbf{1352.6}$	$\mathbf{1402.8}$	$1395.3$

Appendix

Table 11: Few-shot return with different adaptation modes on Point-Robot and Meta-World ( Mixed , mean over 5 seeds). Best per column in bold .

Meta-Reinforcement Learning

Meta-reinforcement learning trains agents across tasks so they can adapt quickly to new environments or objectives.

See all 43 papers

May 30, 2026

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.

Fuyuan Qian, Menglong Zhang, Song Wang +1

Department of Biomedical Engineering, Southern University of Science and Technology, Shenzhen, Guangdong, China · AI Thrust, The Hong Kong University of Science and Technology (Guangzhou), Guangzhou, Guangdong, China · Omni-Intelligence, Shenzhen, Guangdong, China +1

Representation Learning Offline RL Meta-Reinforcement Learning Model-Based RL

Ask PDF

Mar 3, 2026

Temporal Consistency Improves Generalization in Contextual Offline Meta Reinforcement Learning

Offline meta-reinforcement learning seeks to learn a policy that generalizes to new related tasks online. Context-based methods infer a task representation from transition histories, yet learning an effective task representation without supervision remains challenging. Existing methods relying on contrastive learning learn discriminative task representations, but fail to identify task-specific dynamics, while relying on reconstruction can be insufficient to model long-horizon dependencies, limiting generalization to new tasks. We investigate the impact of temporal consistency in latent space on task representation learning, showing that enforcing multi-step predictions in latent space encourages task representations that are able to capture task-dependent dynamics while preventing representation collapse. We provide theoretical analysis characterizing sources of error in value estimation and show through extensive experiments on MuJoCo, Contextual DeepMind Control, and MetaWorld benchmarks that temporal consistency significantly improves both zero-shot and few-shot generalization.

Mohammadreza Nakheai, Aidan Scannell, Kevin Luck +1

Aalto University · Toyota Research Institute · Vrije Universiteit Amsterdam

Representation Learning Offline RL Representation Learning for RL Meta-Reinforcement Learning Generalization in RL Temporal Consistency+2

Ask PDF

Nov 2, 2026

Language Model Augmented Decision Transformer for Multi-Task Offline Reinforcement Learning

Offline multi-task reinforcement learning (MTRL) aims to learn a unified policy from heterogeneous offline datasets that generalizes across tasks. However, existing offline MTRL methods often struggle with two fundamental issues: severe per-task data scarcity and limited cross-task generalization. In this paper, we propose Language-Model Augmented Decision Transformer (LaMA-DT), a unified sequence-modeling framework that leverages a single large language model (LLM) for both trajectory augmentation and multi-task policy learning. LaMA-DT uses natural-language task descriptions to guide task-aligned trajectory augmentation through a set of supervised generation and alignment objectives, and incorporates a scoring-and-filtering mechanism to retain high-quality augmented trajectories. The same LLM is then repurposed as a language-conditioned Decision Transformer and fine-tuned on a mixture of real and augmented data. Experiments on Meta-World and MuJoCo multi-task benchmarks demonstrate that LaMA-DT consistently outperforms competitive baselines, particularly in low-data regimes, showing that language-guided data augmentation substantially improves cross-task generalization.

Xinbo Zhao, Yuhang Liu, Xin Zhang +4

State University of New York at Binghamton · San Diego State University · LinkedIn +1

Multi-Task RL Decision Transformer Large Language Model-Based Data Augmentation

Ask PDF

Abstract

Meta-Reinforcement Learning

Explore similar work

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

Temporal Consistency Improves Generalization in Contextual Offline Meta Reinforcement Learning

Language Model Augmented Decision Transformer for Multi-Task Offline Reinforcement Learning