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; Wenhao Ji; Haiyuan Gui; Meng Han; Xiaojian Liao

cs.LGOct 8, 2026

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

Authors: Xinfei Wang, Shanchen Pang, Chenhao Zhang, Shudong Wang, Wenhao Ji, Haiyuan Gui, Meng Han, Xiaojian Liao

Abstract

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.

Figures & tables

Figure 1: (a) Fixed-length prompt sweep on Reach and Coffee-Push (MT50) shows each task’s "optimal" K . (b) On Coffee-Push , dynamic prompting outperforms fixed K =5 and empirically tuned K =30. (c) Comparison of different prompt selection mechanisms in PromptDT on MT30 (curves are smoothed). (d) Cosine similarity between selected prompts and the current inference trajectory under different prompt selection methods.

Panel of Figure 2 — Figure 2: Task-dependent attention under long prompts ( $K{=}30$ , $L{=}20$ ; average over 8 heads). Reach (left) shows attention concentrated on recent tokens; Basketball relies on broader prompt history.

Figure 3: Trajectory stitching and a DT limitation. Left: Complementary suboptimal trajectories: $A$ succeeds early but fails later; $B$ is the opposite. Right: Optimal stitching ( $A_{1}{\to}A_{2}{\to}B_{3}$ ) is blocked by return-conditioned DTs that cannot adjust RTG at the junction.

Figure 4: Overview of DaCe-DT. ( A ) RAPC constructs a per-task offline index from top- $M$ high-return trajectory slices via the shared GPT-2 embedding. ( B ) Stage 1 trains a task-conditioned value function with per-task adaptive expectile regression. ( C ) Stage 2 jointly trains the policy with adaptive prompt length controlled by a dynamic attention mask. ( D ) At inference, the value function corrects the return target, which guides retrieval and prompt assembly before the policy outputs $a_{t}$ .

Method	Meta-World 5 Tasks		Meta-World 30 Tasks		Meta-World 50 Tasks
	Near-optimal	Sub-optimal	Near-optimal	Sub-optimal	Near-optimal	Sub-optimal
MTDIFF	$\textbf{100.0}\pm 0.0$	$66.30\pm 2.31$	$67.52\pm 0.35$	$54.21\pm 1.10$	$61.32\pm 0.89$	$48.94\pm 0.95$
MTDT	$\textbf{100.0}\pm 0.0$	$64.67\pm 5.25$	$71.89\pm 0.95$	$49.33\pm 2.05$	$65.80\pm 1.02$	$42.33\pm 1.89$
Prompt-DT*	$\textbf{100.0}\pm 0.0$	$67.00\pm 2.33$	$71.30\pm 0.67$	$54.33\pm 0.78$	$71.60\pm 1.40$	$51.40\pm 0.35$
HarmoDT-R*	$\textbf{100.0}\pm 0.0$	$68.00\pm 2.33$	$80.10\pm 3.12$	$59.00\pm 3.66$	$74.24\pm 2.66$	$52.45\pm 1.17$
HarmoDT-M*	$\textbf{100.0}\pm 0.0$	$72.12\pm 3.48$	$79.67\pm 1.46$	$\textbf{62.33}\pm 0.66$	$\textbf{78.80}\pm 0.67$	$56.20\pm 0.56$

Table 1: Performance comparison on Meta-World with 5, 30, and 50 randomly sampled tasks under near-optimal and sub-optimal datasets. Each result is averaged over three random seeds. To ensure fair comparison, we adopt the same random seeds as those used in the previous baseline methods.

Method	Near-opt.	Sub-opt.
CARE (online)	$46.12_{\pm 1.30}$	–
PaCo (online)	$54.31_{\pm 1.32}$	–
Soft-M (online)	$53.41_{\pm 0.72}$	–
D2R (online)	$63.53_{\pm 1.22}$	–
MOORE (online)	$72.90_{\pm 3.30}$	–
ARS-LoRA (online)	$\textbf{83.20}_{\pm 1.80}$	–

Table 2: Average success rate across 3 seeds on MT50 with random goals (MT50-rand) under both near-optimal and sub-optimal cases.

Figure 5: Comparison of training reward trajectories between DaCe-DT and HarmoDT (current offline MTRL state-of-the-art) across MT30 and MT50 configurations of Meta-World.

Setting	MTDT	Prompt-DT	HarmoDT-F	Ours
Cheetah-dir	$662.4_{\pm 1.3}$	$935.3_{\pm 2.6}$	$958.5_{\pm 1.5}$	$\mathbf{966.9_{\pm 2.2}}$
Cheetah-vel	$-170.1_{\pm 5.7}$	$-127.7_{\pm 9.9}$	$-66.5_{\pm 1.2}$	$\mathbf{-22.0_{\pm 1.2}}$
Ant-dir	$165.3_{\pm 0.4}$	$278.7_{\pm 38.7}$	$298.3_{\pm 1.0}$	$\mathbf{355.9_{\pm 1.7}}$
Average	$219.2$	$362.1$	$396.8$	$\mathbf{433.6}_{9.27\%\uparrow}$

Table 3: Generalization to unseen tasks. Methods trained on a task subset and evaluated without fine-tuning.

Components			MT30
LGPM	RAPC	VARC	Succ.	Return
✗	✗	✗	71.3	2436.8
✓	✗	✗	83.7	2864.5
✗	✓	✗	85.3	2961.4
✗	✗	✓	82.6	2735.6
✓	✓	✓	93.3	3185.5

Table 4: Ablation on MT30. Impact of LGPM, RAPC and VARC on Success Rate and Return.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Meta-World task visualizations. Representative snapshots of the manipulation tasks used in our experiments (e.g., Assembly, Basketball, Button-Press variants, Coffee variants, Door variants, Drawer variants, Faucet variants, Hammer, etc.), illustrating the diverse objects and interaction dynamics in the Meta-World benchmark.

Parameter	Value
Number of Layers	6
Number of Attention Heads	8
Embedding Dimension	256
Nonlinearity Function	RELU
Context Length K	20
Dropout	0.1

Appendix

Table 5: DaCe-DT training hyperparameters

Parameter	Value
Sparsity	0.2
Batch Size	8
Mask Changing Interval	5000
Minimum of Mask Changing $\eta_{min}$	0
Maximum of Mask Changing $\eta_{max}$	100
Maximum Threshold in Mask	25

Appendix

Table 6: HarmoDT training hyperparameters

Parameter	Value
K (length of context $\tau$ )	20
Batch Size	8
Number of Evaluation Episodes for Each Task	20
Learning Rate	1e-4
Learning Rate Decay Weight	1e-4
Number of Layers	6

Appendix

Table 7: Prompt-DT training hyperparameters

Figure 7: The 8 maze layouts used in our multi-task Maze2D benchmark. Layouts range from $5{\times}5$ to $8{\times}10$ grids and vary in wall topology, yielding 8 navigation tasks with distinct structural complexity.

Task	Ours	Prompt-DT
maze2d-1	106.5	48.0
maze2d-2	95.1	22.0
maze2d-3	0.0	12.0
maze2d-4	32.9	7.5
maze2d-5	4.5	41.1
maze2d-6	22.6	14.9

Appendix

Table 8: Performance on sparse-reward Maze2D with 8 custom maze layouts. VARC is disabled; only RAPC and LGPM are used. Score is cumulative reward per episode.

Task	Return	Succ.	Task	Return	Succ.	Task	Return	Succ.
basketball-v2	2198.46	1.0	door-close-v2	4389.19	1.0	handle-press-v2	4359.66	1.0
bin-picking-v2	4350.81	1.0	door-lock-v2	3152.24	0.9	handle-pull-side-v2	3646.42	1.0
button-press-topdown-v2	3669.39	1.0	door-open-v2	3715.46	1.0	handle-pull-v2	4264.39	1.0
button-press-v2	2608.85	1.0	door-unlock-v2	3840.95	1.0	lever-pull-v2	1880.16	1.0
button-press-wall-v2	2304.72	1.0	hand-insert-v2	3591.48	1.0	peg-insert-side-v2	4506.66	1.0
coffee-button-v2	1658.16	1.0	drawer-close-v2	4831.95	1.0	pick-place-wall-v2	3032.87	0.9

Appendix

Table 9: Per-task episodic return and success rate of DaCe-DT on MetaWorld MT30 (near-optimal, seed 123). Mean return: 3185.45; mean success: 0.933.

Task	Return	Succ.	Task	Return	Succ.	Task	Return	Succ.
basketball-v2	2337.14	1.0	faucet-open-v2	4574.86	1.0	plate-slide-back-side-v2	652.54	1.0
bin-picking-v2	3958.31	1.0	faucet-close-v2	4663.15	1.0	soccer-v2	1043.56	0.3
button-press-topdown-v2	3482.01	0.9	handle-press-side-v2	4731.69	1.0	push-wall-v2	3789.07	0.8
button-press-v2	2396.76	1.0	handle-press-v2	4315.96	0.9	shelf-place-v2	1633.87	0.9
button-press-wall-v2	2155.54	1.0	handle-pull-side-v2	3092.48	1.0	sweep-into-v2	4176.40	0.9
coffee-button-v2	1695.16	0.9	handle-pull-v2	4424.69	1.0	sweep-v2	2449.86	0.7

Appendix

Table 10: Per-task episodic return and success rate of DaCe-DT on MetaWorld MT50 (near-optimal, seed 123). Mean return: 2983.47; mean success: 0.884.

Task	MT5	MT30 (near-opt.)	MT30 (sub-opt.)
basketball-v2	18	24	36
reach-v2	8	6	10
button-press-topdown-v2	10	20	26
button-press-v2	20	10	19
button-press-wall-v2	10	21	27

Appendix

Table 11: Prompt lengths learned by LGPM per task under different configurations. The same task converges to different lengths depending on the task combination, reflecting LGPM’s adaptivity to task heterogeneity.

Figure 8: Overview of prompt-based offline multi-task learning. The prompt disambiguates mixed offline trajectories by encoding task identity, behavioral priors, and goal/return targets, guiding the Transformer to produce task-specific behaviors (e.g., Grasp vs. Pull).

Explore similar work

Distributionally Robust Multi-Task Reinforcement Learning via Adaptive Task Sampling
May 14, 2026Nicholas E. Corrado, Wenyuan Huang, Josiah P. HannaDistributionally Robust RL Multi-Task RL
Computer Sciences Department · University of Wisconsin – Madison
CDCP: Conditional Diffusion Model with Contextual Prompts for Multi-task Offline Safe Reinforcement Learning
Jul 4, 2026Jiayi Guan, Tianle Zhang, Li Shen +8Reinforcement Learning Constrained RL
Tongji University, Shanghai, China · JD Explore Academy, Beijing, China · School of Cyber Science and Technology, Shenzhen Campus of Sun Yat-sen University, Shenzhen, China +2
Learning to Adapt: Representation-Based Reinforcement Learning for Multi-Task Skill Transfer
Jun 11, 2026Aryan Naveen, Haitong Ma, Haldun Balim +1RL for Robotics Reinforcement Learning
Massachusetts Institute of Technology · Harvard School of Engineering and Applied Sciences

May 14, 2026, cs.LGJ/K move · Enter open · S save

Distributionally Robust Multi-Task Reinforcement Learning via Adaptive Task Sampling

Nicholas E. Corrado, Wenyuan Huang, Josiah P. Hanna

Computer Sciences Department · University of Wisconsin – Madison

Distributionally Robust RL Multi-Task RL Robust RL Data-Efficient Learning

Open paper AskPDF arXiv

Multi-task reinforcement learning (MTRL) aims to train a single agent to efficiently optimize performance across multiple tasks simultaneously. However, jointly optimizing all tasks often yields imbalanced learning: agents quickly solve easy tasks but learn slowly on harder ones. While prior work primarily attributes this imbalance to conflicting task gradients and proposes gradient manipulation or specialized architectures to address it, we instead focus on a distinct and under-explored challenge: imbalanced data allocation. Standard MTRL allocates an equal number of environment interactions to each task, which over-allocates data to easy tasks that require relatively few interactions to solve and under-allocates data to hard tasks that require substantially more experience to solve. To address this challenge, we introduce Distributionally Robust Adaptive Task Sampling (DRATS), an algorithm that adaptively prioritizes sampling tasks furthest from being solved. We derive DRATS by formalizing MTRL as a feasibility problem from which we derive a minimax objective for minimizing the worst-case return gap, the difference between a desired target return and the agent's return on a task. In benchmarks like MetaWorld-MT10 and MT50, DRATS improves data efficiency and increases worst-task performance compared to existing task sampling algorithms.

Abstract

Explore similar work

Distributionally Robust Multi-Task Reinforcement Learning via Adaptive Task Sampling

CDCP: Conditional Diffusion Model with Contextual Prompts for Multi-task Offline Safe Reinforcement Learning

Learning to Adapt: Representation-Based Reinforcement Learning for Multi-Task Skill Transfer