PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2610.11085v1cs.LG

TL;DR

Offline MTRL seeks to learn across tasks from static datasets, but data heterogeneity, irrelevant prompts, and fragmented trajectories pose data-level challenges. DaCe-DT combines adaptive prompt length, relevant segment retrieval, and return calibration; on Meta-World, it improves over HarmoDT on both optimal and suboptimal datasets.

  • Problem

    Offline MTRL methods largely emphasize algorithmic or architectural changes, while data-level challenges include heterogeneous prompt-length needs, irrelevant prompts, and fragmented trajectories.

    Source: §1 · p. 2

  • Method

    DaCe-DT combines LGPM for task-specific prompt length, RAPC for semantically relevant high-return segments, and VARC for achievable-return calibration and trajectory stitching.

    Source: §1 · p. 2

  • Results

    DaCe-DT improves over HarmoDT by 11.73% on the optimal dataset and 13.34% on the suboptimal dataset in Meta-World.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Experiments on Meta-World show consistent improvements over state-of-the-art methods across datasets of varying quality.

    Source: §6 · p. 9

    DaCe-DT remains bounded by pre-collected data and has been evaluated only on simulation benchmarks with defined task boundaries and reward signals.

    Source: E Limitations · p. 18

Abstract

from arXiv · show

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.

1 Introduction

Offline MTRL seeks cross-task knowledge transfer from static datasets, but existing work largely emphasizes algorithmic or architectural changes over data heterogeneity. DaCe-DT addresses three data-level bottlenecks with adaptive prompts and trajectory refinement, and reports improvements over HarmoDT on Meta-World datasets of differing quality.

1 Introduction

  • Existing methods mainly emphasize architectural or algorithmic changes, while data-centric exploitation of heterogeneous static offline data remains under-addressed.Offline MTRL aims to transfer knowledge across tasks without additional environment interaction.Source: §1 · pp. 1–2
  • On Coffee-Push, an empirically tuned prompt length of 30 improves success rate by 60% over the state-of-the-art baseline, and adaptive prompting raises it to 100%.Across all 50 tasks, the approach improves average success rate by 9.60% over the state of the art.Source: §1 · p. 2
  • LGPM adapts effective prompt length, RAPC retrieves semantically relevant high-return segments, and VARC calibrates achievable returns for reliable stitching.VARC also guides prompt retrieval; RAPC retrieval, VARC filtering, and LGPM prompt trimming form a coordinated design.Source: §1 · p. 2
  • DaCe-DT improves over HarmoDT by 11.73% on the optimal dataset and 13.34% on the suboptimal dataset.The comparison is reported on the Meta-World benchmark.Source: §1 · p. 2
  • The paper identifies prompt-length mismatch, semantically irrelevant sampled prompts, and misleading supervision from fragmented trajectories as three Offline MTRL bottlenecks.The bottlenecks concern heterogeneous task complexity, prompt relevance to the current state, and trajectory continuity.Source: §1 · p. 2

2 Preliminary

Decision Transformer predicts actions from offline trajectories through conditional sequence modeling. Prompt-based DT extends this setup to multiple tasks by adding task-specific context that supports task inference and task-conditioned behavior.

2.1 Decision Transformer (DT)

  • DT trains a causally masked Transformer to predict actions from interleaved return-to-go, state, and action tokens in offline trajectories.Its policy predicts an action autoregressively from the trajectory prefix.Source: §2.1 · p. 3

2.2 Prompt as Task Evidence in Offline MTRL

  • Prompt-based DT prepends task-specific context to the input sequence, enabling task inference and task-conditioned behavior.Prompts serve as task evidence in trajectory space, providing task identification, prior knowledge, and return conditioning.Source: §2.2 · p. 3

3 Rethinking Offline MTRL: A Data-Centric Perspective

The section identifies three data-related constraints on Offline MTRL: task-dependent prompt-length needs, prompt segments that may not match the current state, and fragmented trajectories that impede stitching. It describes how these mismatches can weaken context use or make return-conditioned recombination unreliable.

3.1 Prompt Length Mismatch

  • Simple Reach needs only K=5 for 100% success, whereas complex Coffee-Push needs K=30 to reach 80%; adaptive length adjustment adds 20% over the best fixed baseline.Fixed lengths can under- or over-provision context across tasks and trajectory segments; the same task’s optimal length also varies with task combination and dataset quality.Source: §3.1 · pp. 3–4

3.2 Semantic Irrelevance of Prompts

  • Last-K and random prompt selection rely on recency or sampling rather than semantic relevance to the current decision state.The section describes these approaches as overlooking semantic alignment between prompt segments and current decision context.Source: §3.2 · p. 4
  • RAPC raises mean prompt–trajectory cosine similarity to 0.741, from 0.361 for fixed prompts and 0.462 for random prompts.Fixed and random prompt selection suppress performance during training, while RAPC achieves faster convergence and consistently higher success rates.Source: §3.2 · p. 4

3.3 Sub-Trajectory Fragmentation

  • Complementary locally successful trajectory segments could form a better trajectory, but return-conditioned DTs may lack dataset support for the higher RTG needed at the stitching state.Unsupported return conditioning can induce out-of-distribution generalization, leading DT to follow a lower-return path and limiting recombination.Source: §3.3 · p. 5, Figure 3 · p. 5
  • Direct value-based correction is difficult to apply to prompt-based Offline MTRL because heterogeneous data and return scales can bias reachable-return estimates.Task-specific value functions and stitching modules also add overhead and weaken unified multi-task modeling.Source: §3.3 · p. 5

4 Methodology

DaCe-DT combines task-adaptive prompt length, semantically relevant trajectory retrieval, and value-based return calibration. Its inference pipeline applies return correction, retrieval, and prompt assembly before policy prediction.

4.1 Length-Gated Prompt Masking (LGPM)

  • LGPM learns a task-specific effective prompt length and masks older prompt tokens while preserving causal order.A sparsity regularizer discourages maximum-length prompts, steering each task toward a minimal sufficient context.Source: §4.1 · p. 5

4.2 Retrieval-Augmented Prompt Construction (RAPC)

  • RAPC builds per-task indexes from high-return trajectory slices and retrieves slices matching the current history context.The index uses shared GPT-2 embeddings; retrieval begins after a warm-up period and the index is refreshed periodically.Source: §4.2 · p. 6

4.3 Value-Adaptive Return Calibration (VARC)

  • VARC trains a shared value function with task-specific expectiles, then freezes it for inference-time return calibration.The method adapts expectiles to data noise: noisier tasks use more conservative quantiles, while lower-noise tasks use more aggressive ones.Source: §4.3 · pp. 6–7

4.4 Inference: Return Correction, Retrieval, and Length-Adaptive Prompt Assembly

  • At inference, the value function raises the return target when it detects a higher achievable return, with a tolerance to prevent over-correction.The corrected return and trajectory history guide retrieval; the retrieved slice is cropped or extended to the task’s learned prompt length before policy prediction.Source: §4.4 · p. 7

5 Experiments

Experiments compare DaCe-DT with multi-task offline RL baselines on Meta-World task sets and datasets of varying quality, and examine generalization, sparse-reward robustness, and component ablations.

5.1 Environment and Baselines

  • The evaluation covers MT5, MT30, and MT50, using near-optimal and sub-optimal offline datasets and comparing against representative multi-task offline RL methods.The experiments report average task success rates and include training reward trajectories for DaCe-DT and HarmoDT across MT30 and MT50.Source: §5 · p. 7, §5.1 · p. 8, Figure 5 · p. 8

5.2 Performance Comparison

  • DaCe-DT improves on MT30 by 11.73% on near-optimal and 13.34% on sub-optimal data, and on MT50 by 9.60% and 6.50%, respectively.On MT5, it achieves a 100% success rate on near-optimal data and a 12.96% performance gain on sub-optimal data.Source: Table 1 · p. 8, §5.2 · p. 9

5.3 More Analysis

  • On unseen continuous-control tasks, DaCe-DT outperforms all baselines by 9.27% without fine-tuning.The evaluation follows HarmoDT’s protocol and builds the retrieval index from expert segments using frozen embeddings.Source: §5.3 · p. 9, Table 3 · p. 9
  • With the value function disabled, DaCe-DT scores 36.9 on sparse-reward MAZE2D layouts, compared with 18.2 for Prompt-DT.This evaluation uses eight custom layouts and tests RAPC and LGPM without return calibration.Source: §5.3 · p. 9

5.4 Ablation Study

  • The MT30 ablation compares LGPM-only, RAPC-only, and VARC-only variants with components disabled as specified for each variant.Results are presented for success rate and return; the supplied passage does not give the table values.Source: §5.4 · p. 9

6 Conclusion

The conclusion frames DaCe-DT as a response to three data-centric bottlenecks in offline multi-task reinforcement learning and reports consistent improvements on Meta-World.

6 Conclusion

  • The paper identifies prompt length mismatch, semantically irrelevant prompts, and misleading supervision from fragmented trajectories as barriers to learnability and generalization.Source: §6 · p. 9
  • LGPM adapts context length, RAPC selects semantically aligned behavioral evidence, and VARC calibrates returns for trajectory stitching at inference.Source: §6 · p. 9
  • Meta-World experiments show consistent and significant improvements over state-of-the-art methods across datasets of varying quality.Source: §6 · p. 9

A Detailed Environment and Baselines

The appendix details the Meta-World evaluation setting, offline dataset construction, baseline suite, experimental configuration, and computational resources.

A Detailed Environment and Baselines

  • Meta-World contains 50 manipulation tasks sharing robot dynamics; the study uses random-goal evaluation and averages success across tasks.Offline data come from per-task SAC replay buffers, with near-optimal data spanning training and sub-optimal data limited to the first half of each buffer.Source: A Detailed Environment and Baselines · pp. 11–12

A.1 Baselines

  • The baseline suite includes offline methods such as task-ID-conditioned behavior cloning, IQL, diffusion, and transformer approaches, alongside online multi-task methods.The appendix also describes CARE, PaCO, Soft-M, D2R, MOORE, and ARS-LoRA.Source: §A.1 · p. 12

B.1 Hyper-parameters

  • The policy uses a Transformer implemented from minGPT, with model and baseline hyperparameters reported in tables and experiments run using fixed seeds.Three seeds correspond to distinct randomly sampled task combinations.Source: §B.1 · p. 13, Table 5 · p. 13, Table 6 · p. 14

B.2 Computational Resources

  • Each MT5, MT30, and MT50 configuration is trained for 1 million steps on NVIDIA 4090 GPUs, with three runs per environment.Training time increases with task count and is shorter for the sub-optimal datasets than for the near-optimal datasets.Source: §B.2 · p. 14

C.1 Maze2D Multi-Task Experiments

The Maze2D experiments evaluate multi-task prompt methods on custom sparse-reward layouts, where DaCe-DT outperforms Prompt-DT.

C.1 Maze2D Multi-Task Experiments

  • The multi-task benchmark comprises eight maze layouts, each treated as a task with its own pre-collected offline dataset.Source: §C.1 · p. 14
  • VARC is disabled because it relies on dense reward signals, so the experiment evaluates RAPC and LGPM.Source: §C.1 · p. 14
  • DaCe-DT achieves an average score of 36.9 on sparse-reward Maze2D, compared with 18.2 for Prompt-DT.Source: §C.1 · p. 14

C.2 Per-Task Results on MetaWorld MT30

On near-optimal MetaWorld MT30, DaCe-DT reports an overall mean return of 3185.45 and mean success rate of 0.933.

C.2 Per-Task Results on MetaWorld MT30

  • 3185.45 mean return and 0.933 mean success rate are reported for DaCe-DT on near-optimal MetaWorld MT30.The results use seed 123.Source: §C.2 · p. 14
  • Table 9 reports per-task episodic return and success rate for DaCe-DT on MT30.The dataset is near-optimal, and the reported run uses seed 123.Source: §C.2 · p. 14

C.3 Per-Task Results on MetaWorld MT50

On near-optimal MetaWorld MT50, DaCe-DT reports a mean return of 2983.47 and mean success rate of 0.884; learned prompt lengths vary with task combinations.

C.4 Learned Prompt Lengths (LGPM)

  • The same task can converge to different LGPM prompt lengths under different task combinations.This reflects how co-trained task complexity influences the context window for an individual task.Source: §C.4 · p. 15

D.1 Offline Reinforcement Learning

  • Offline RL learns policies from static datasets without environment interaction, while distributional shift between training data and policy can degrade performance.The cited discussion describes conservative updates and regularization as ways to keep policies close to behavior policies.Source: §D.1 · pp. 15–16

C.3 Per-Task Results on MetaWorld MT50

  • 2983.47 mean return and 0.884 mean success rate are reported for DaCe-DT on near-optimal MetaWorld MT50.The reported run uses seed 123.Source: §C.3 · p. 15, Table 10 · p. 16

D.1 Offline Reinforcement Learning

  • Decision Transformers use conditional sequence modeling to predict actions from historical trajectories, while diffusion-based offline RL models policies or dynamics through iterative denoising.The passage names Diffuser and its variants as examples of diffusion approaches.Source: §D.1 · p. 16

D.2 Multi-task Reinforcement Learning

Multi-task reinforcement learning seeks a unified policy across distinct tasks, but heterogeneous tasks and data can produce gradient conflicts that hinder parameter sharing.

D.2 Multi-task Reinforcement Learning

  • MTRL trains a unified policy across distinct tasks, commonly conditioning on task identifiers or goals or using routing and weight-sharing strategies.These are described as alternative approaches to task-specific behavior and parameter allocation.Source: §D.2 · p. 16
  • Offline MTRL uses task-mixed datasets to transfer knowledge without environment interaction, which is useful when online data collection is expensive or risky.The datasets span multiple tasks.Source: §D.2 · pp. 16–17
  • LGPM’s learned prompt length for a task can differ across configurations, reflecting the influence of the task combination.The table reports learned lengths per task under different configurations.Source: Table 11 · p. 17
  • Task and data heterogeneity can induce gradient conflicts that hinder parameter sharing, and existing offline MTRL approaches use routing or masking to mitigate interference.Parameter-level methods use task-specific binary masks, while routing approaches dynamically allocate network capacity.Source: §D.2 · p. 17

D.3 Prompt-DT

  • Prompt-DT prepends a short demonstration subtrajectory sampled from the target task, and prompt context can encode task identity, behavioral priors, and goal or return targets.The prompt helps disambiguate mixed offline trajectories and guide task-specific behavior.Source: §D.3 · p. 17, Figure 8 · p. 17

E Limitations

This section outlines DaCe-DT’s limitations, including its dependence on offline data quality and diversity, lack of validation beyond simulation benchmarks, and open challenges around distributional gaps and real-world deployment. It also clarifies that LLMs were not part of the core method and were used only for writing assistance.

E Limitations

  • Without active exploration, DaCe-DT cannot recover from gaps in its pre-collected data, so scarce or low-quality datasets may hinder performance.The paper identifies dependence on offline data quality and diversity as a limitation.Source: E Limitations · p. 18, §2 · p. 20
  • Evaluation is limited to simulation benchmarks with defined tasks and rewards; deployment with noisy observations, ambiguous tasks, or sparse or absent rewards remains open.The paper describes this as an unresolved challenge for open-world and real-robot settings.Source: E Limitations · p. 18, §2 · p. 20

2. Limitations

  • LLMs are not a core component of the method; any use was limited to writing assistance and did not affect the scientific methodology.Source: §2 · p. 25
Loading 2610.11085v1…