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

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

cs.CLNov 2, 2026

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

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

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

Abstract

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.

Figures & tables

Figure 1: LaMA-DT consists of three components: (i) Multi-Objective Fine-Tuning. We adapt the LLM to generate meaningful, task-specific trajectories from limited offline data and natural-language task descriptions by training objective-specific LoRA adapters together with a shared LoRA adapter, enabling both trajectory generation and task–trajectory alignment. (ii) Data Preparation. Candidate trajectories are generated per task using $\text{O}_{3}$ , scored using aggregated multi-objective likelihoods, and the top- $K$ trajectories are retained and merged with the original offline data. (iii) Multi-Task Policy Learning. We repurpose the LLM and shared LoRA adapter as a prompt-conditioned DT, trained on the combined real and augmented dataset, with task descriptions $\mathbf{p}_{i}$ and a provenance token $z_{\text{prov}}$ .

Figure 2: Performance of LaMA-DT with different augmentation ratios $r=|\mathcal{D}^{\text{aug}}_{i}|/|\mathcal{D}_{i}|$ on Meta-World ML30 and MuJoCo Cheetah-dir, across low and standard data regimes, and under sub-optimal and near-optimal data qualities.

Task		CDS	MTBC	MTIQL	MTDIFF	MTDT	Prompt-DT	HarmoDT	LaMA-DT
MW	ML5	50.7 $\pm$ 3.3	46.2 $\pm$ 2.8	61.1 $\pm$ 4.5	76.0 $\pm$ 2.8	74.5 $\pm$ 4.4	76.8 $\pm$ 3.0	78.9 $\pm$ 3.8	83.4 $\pm$ 2.9
MW	ML30	32.9 $\pm$ 2.5	34.5 $\pm$ 1.7	43.3 $\pm$ 2.2	57.5 $\pm$ 1.5	55.2 $\pm$ 2.1	58.7 $\pm$ 1.3	63.4 $\pm$ 1.6	69.8 $\pm$ 1.7
MuJoCo	Cheetah-vel	-134.5 $\pm$ 15.8	-148.4 $\pm$ 11.2	-142.0 $\pm$ 14.4	-186.2 $\pm$ 8.2	-170.1 $\pm$ 17.9	-127.6 $\pm$ 10.9	-86.6 $\pm$ 7.7	-73.5 $\pm$ 12.5
	Cheetah-dir	428.9 $\pm$ 17.5	429.7 $\pm$ 6.6	454.1 $\pm$ 21.3	536.8 $\pm$ 10.7	678.5 $\pm$ 24.3	705.6 $\pm$ 8.3	926.3 $\pm$ 10.4	937.8 $\pm$ 9.6
	Ant-dir	131.0 $\pm$ 10.6	153.8 $\pm$ 15.9	120.6 $\pm$ 9.9	248.0 $\pm$ 16.2	235.7 $\pm$ 15.4	302.5 $\pm$ 17.1	315.6 $\pm$ 5.5	328.5 $\pm$ 8.4

Table 1: The average returns of LaMA-DT vs. baselines on Meta-World (Sub-Optimal) and MuJoCo (Near-Optimal) datasets under standard data regime.

Method	ML30						Cheetah-dir
	Sub-Optimal			Near-Optimal			Sub-Optimal			Near-Optimal
	1%	5%	10%	1%	5%	10%	1%	5%	10%	1%	5%	10%
MTDT	10.9 $\pm$ 3.0	19.2 $\pm$ 2.4	25.6 $\pm$ 2.3	12.5 $\pm$ 2.0	24.1 $\pm$ 2.3	35.6 $\pm$ 1.9	113.4 $\pm$ 20.2	274.5 $\pm$ 18.0	386.3 $\pm$ 21.9	145.2 $\pm$ 23.1	320.5 $\pm$ 19.6	435.6 $\pm$ 20.4
Prompt-DT	12.8 $\pm$ 1.4	24.3 $\pm$ 0.9	36.8 $\pm$ 1.1	15.3 $\pm$ 1.5	28.5 $\pm$ 0.7	41.2 $\pm$ 0.8	191.3 $\pm$ 17.1	378.9 $\pm$ 18.3	507.4 $\pm$ 15.3	256.3 $\pm$ 16.2	474.0 $\pm$ 12.8	544.7 $\pm$ 15.7
HarmoDT	12.1 $\pm$ 1.6	27.0 $\pm$ 1.4	42.6 $\pm$ 1.5	22.4 $\pm$ 1.1	36.1 $\pm$ 0.9	46.8 $\pm$ 0.9	185.2 $\pm$ 12.3	452.2 $\pm$ 17.5	526.8 $\pm$ 13.9	260.7 $\pm$ 11.2	540.3 $\pm$ 9.8	592.0 $\pm$ 10.3
LaMA-DT	15.7 $\pm$ 1.7	39.1 $\pm$ 1.5	51.8 $\pm$ 1.7	38.4 $\pm$ 1.2	52.6 $\pm$ 0.9	61.5 $\pm$ 1.0	225.9 $\pm$ 12.7	521.5 $\pm$ 14.0	654.4 $\pm$ 9.2	450.6 $\pm$ 15.1	720.8 $\pm$ 9.4	803.9 $\pm$ 8.9

Table 2: Performance under low-data regimes (1%, 5%, and 10%) across different data qualities (sub-optimal vs. near-optimal).

Task	w/o $\text{O}_{1}$	w/o $\text{O}_{2}$	w/o $\text{O}_{3}$	w/o $\text{O}_{4}$	w/o $\text{O}_{5}$	w/o $\text{O}_{6}$
ML30	61.7 $\pm$ 1.3	60.5 $\pm$ 2.1	56.4 $\pm$ 2.8	66.3 $\pm$ 1.6	63.8 $\pm$ 2.0	65.0 $\pm$ 2.3
Cheetah-dir	842.3 $\pm$ 14.3	862.4 $\pm$ 10.2	765.9 $\pm$ 11.7	910.2 $\pm$ 16.8	902.6 $\pm$ 15.4	922.5 $\pm$ 26.3
Ant-dir	259.0 $\pm$ 7.0	281.3 $\pm$ 8.5	247.3 $\pm$ 9.6	310.8 $\pm$ 12.3	313.9 $\pm$ 10.4	322.4 $\pm$ 16.2

Table 3: Ablation studies on LaMA-DT.

Source	Dyn. MSE $\downarrow$	State KL vs. real $\downarrow$	L2 Dist. from Dataset	Valid Coverage $\uparrow$
Real expert (reference)	0.0000	0.00	—	100%
MTDIFF ( He et al., 2023 )	0.0107	0.51	0.4552	102%
LaMA-DT (unfiltered)	0.0089	0.42	0.4831	99%
LaMA-DT (filtered)	0.0031	0.18	0.4376	109%

Table 4: Quality of augmented data on Cheetah-dir. We compare against real expert data (reference) and MTDIFF ( He et al., 2023 ) as a generator baseline. Lower values are better for Dyn. MSE, State KL, and L2 Dist., while higher values are better for Valid Coverage. Criteria details are in Appendix E.6 .

Figure 3: t-SNE projection of Cheetah-dir states: unfiltered (left) vs. filtered (right). Gray dots are real states from offline data, red dots are augmented states overlapping the real manifold (in-distribution), red crosses are augmented states lying outside the dataset for exploration, and the dashed contour is the $90\%$ real-data envelope.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

	Cheetah-dir	Cheetah-vel	Ant-dir	Meta-World ML30
Episodes / task	50	50	50	50
RTG target ( forward task)	760	per-task ref. velocity	per-task ref. dir.	4,500
RTG normalization	yes	yes	yes	yes
Provenance at test time	real	real	real	real

Appendix

Table 5: Per-environment evaluation settings used to obtain Tab. 1 and Tab. 2 .

Backbone & tokenization
Parameter	Value
Base LLM backbone	Llama 3.2 1B (frozen)
Number of layers / attention heads	16 / 32
Hidden size	2048
Activation (FFN)	SiLU
Numeric MLPs (state / action / return)	1 layer, Tanh, $\to 2048$

Appendix

Table 6: Hyper-parameters of LaMA-DT in our experiments. All values reflect the configuration that produced Tab. 1 .

Figure 4: Training curves of LaMA-DT and the baselines in Tab. 1 .

Temperature $\mathcal{T}$	ML30	Cheetah-dir
0.0	63.0 $\pm$ 1.1	896.8 $\pm$ 10.8
0.2	67.3 $\pm$ 1.3	917.3 $\pm$ 16.7
0.4 (default)	69.8 $\pm$ 1.7	937.8 $\pm$ 9.6
0.8	59.6 $\pm$ 2.2	874.9 $\pm$ 15.0

Appendix

Table 7: Sensitivity to generation temperature

\mathcal{T}

in trajectory augmentation. We report Meta-World ML30 success rate (%) and MuJoCo Cheetah-dir normalized return. Results are averaged over 3 seeds.

Method	ML30		Cheetah-dir
Method	NO	SO	NO	SO
MTDT	70.9 $\pm$ 1.6	55.2 $\pm$ 2.1	678.5 $\pm$ 24.3	598.1 $\pm$ 22.9
PDT	72.1 $\pm$ 0.9	58.7 $\pm$ 1.3	705.6 $\pm$ 8.3	636.7 $\pm$ 11.2
HarmoDT	82.7 $\pm$ 1.2	63.4 $\pm$ 1.6	926.3 $\pm$ 10.4	872.4 $\pm$ 13.8
LaMA-DT	88.6 $\pm$ 1.5	69.8 $\pm$ 1.7	937.8 $\pm$ 9.6	912.3 $\pm$ 14.9

Appendix

Table 8: Performance across different dataset expertise under standard data regime : Near-Optimal (NO) vs. Sub-Optimal (SO). We report success rate for ML30 and normalized return for Cheetah-dir.

Backbone	Params	Average Return
GPT-2	124M	386.2
Qwen2.5-0.5B	494M	888.0
Qwen2.5-1.5B	1.5B	917.0
Llama 3.2 1B (default)	1.24B	937.8

Appendix

Table 9: LaMA-DT performance with different LLM backbones on Cheetah-dir (50 episodes/task), averaged over the two task directions.

Description variant	ML30 (succ. %)	Cheetah-dir (return)
Full template (default)	69.8	937.8
Minimal (task name only)	55.5	892.2
Drop-expertise field	63.0	920.3
Noisy (typos + grammar errors)	49.6	801.7
Shortened paraphrase	67.9	912.0
Learnable task-ID embedding	61.0	913.4

Appendix

Table 10: Robustness to task-description variations. ML30: average success rate (%); Cheetah-dir: average normalized return.

	Meta-World	MuJoCo
Method	ML30 (un)	C-vel	C-dir	Ant-dir	Average
MTDT	35.9 $\pm$ 4.1	-184.2 $\pm$ 8.1	429.4 $\pm$ 19.5	155.2 $\pm$ 9.5	109.1
PDT	43.5 $\pm$ 2.7	-140.2 $\pm$ 10.9	577.2 $\pm$ 8.5	213.6 $\pm$ 20.4	173.5
LaMA-DT	52.3 $\pm$ 3.2	-94.1 $\pm$ 7.6	668.0 $\pm$ 16.1	237.5 $\pm$ 7.6	216.0

Appendix

Table 11: Unseen-task generalization. Meta-World success rate (%); MuJoCo normalized return (%). Comparison against baselines capable of prompt-conditioning.

Method	Cheetah-dir	ML30
Full LaMA-DT	521.5	39.1
w/o augmentation	359.6	29.8
w/o shared-LoRA initialization	514.9	37.0

Appendix

Table 12: Component analysis with 5% sub-optimal training data. Cheetah-dir reports normalized return; ML30 reports success rate (%). In the first ablation, shared-LoRA initialization is retained. In the second, augmented data are retained.

Multi-Task RL

Multi-task reinforcement learning trains agents to solve related tasks by sharing policies, representations, or experience.

See all 90 papers

DaCe-DT: Data-Centric Offline Multi-Task Reinforcement Learning via Adaptive Prompts and Trajectory Correction for Heterogeneous Tasks
Oct 8, 2026Xinfei Wang, Shanchen Pang, Chenhao Zhang +5Reinforcement Learning Offline RL
Decision Transformer with Task-Aware Representations for Offline Meta-Reinforcement Learning
Nov 2, 2026Xinbo Zhao, Yuhang Liu, Xin Zhang +3Meta-Reinforcement Learning Offline RL
State University of New York at Binghamton · San Diego State University · LinkedIn
DLM: Unified Decision Language Models for Offline Multi-Agent Sequential Decision Making
Apr 26, 2026Zhuohui Zhang, Bin Cheng, Bin HeSequential MARL Offline RL
Department of Control Science and Engineering, Tongji University, Shanghai 201804, China · Shanghai Institute of Intelligent Science and Technology; Shanghai Research Institute for Intelligent Autonomous Systems; State Key Laboratory of Autonomous Intelligent Unmanned Systems, Shanghai 200120, China

Oct 8, 2026J/K move · Enter open · S save

DaCe-DT: Data-Centric Offline Multi-Task Reinforcement Learning via Adaptive Prompts and Trajectory Correction for Heterogeneous Tasks

Xinfei Wang, Shanchen Pang, Chenhao Zhang, Shudong Wang+4

Reinforcement Learning Offline RL

Open paper AskPDF arXiv

Offline multi-task reinforcement learning (Offline MTRL) heavily depends on the quality and distribution of pre-collected data. However, existing methods mainly focus on algorithmic optimization, with less emphasis on data-level improvements to enhance learning ability and generalization performance. This paper, from a data perspective, reveals three key bottlenecks that limit Offline MTRL performance:(i) ineffective utilization of prompts length under diverse task complexities, and (ii) semantic irrelevance of randomly sampled prompt segments, (iii) misleading supervision induced by fragmented and discontinuous trajectories. To address these challenges, we propose DaCe-DT, a robust offline MTRL framework designed to be insensitive to heterogeneous task complexities and data quality, featuring length-gated prompt masking (LGPM), retrieval-augmented prompt construction (RAPC), and value-adaptive return calibration (VARC). Together, these mechanisms enable DaCe-DT to deliver data-centric prompt adaptation and trajectory refinement, resulting in robust multi-task generalization and stable policy learning amid heterogeneous offline data and tasks. Experimental results on Meta-World show that DaCe-DT consistently outperforms state-of-the-art methods, achieving an average improvement of 11.73% on optimal datasets and an improvement of 13.34% on suboptimal datasets, demonstrating its effectiveness in learning stably from imperfect data and improving overall multi-task performance.

Abstract

Multi-Task RL

Explore similar work

DaCe-DT: Data-Centric Offline Multi-Task Reinforcement Learning via Adaptive Prompts and Trajectory Correction for Heterogeneous Tasks

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

DLM: Unified Decision Language Models for Offline Multi-Agent Sequential Decision Making