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

AI summary

Overview

• Research area: Offline multi-task reinforcement learning (Offline MTRL), sequence modeling with Decision Transformers, prompt-based task conditioning, and value-based trajectory stitching.
• Technical level: Advanced. The paper assumes familiarity with return-conditioned sequence models, implicit/expectile value learning, attention masking, retrieval via embedding similarity, and multi-task robotic benchmarks.
• Scope (one sentence): The paper diagnoses three data-level bottlenecks in offline multi-task RL and proposes a single framework, DaCe-DT, whose three mechanisms (LGPM, RAPC, VARC) adapt prompt length, retrieve semantically relevant prompts, and correct return targets, evaluated on Meta-World MT5/MT30/MT50 plus generalization and sparse-reward settings.

What This Paper Is About

Offline multi-task reinforcement learning lets a single policy learn many tasks from pre-collected data with no environment interaction, but most existing work tries to improve it through architecture or algorithm changes. This paper argues the real limiting factor is the data itself: fixed prompt lengths do not fit tasks of different complexity, prompt segments are sampled without regard to their semantic relevance to the current decision state, and offline datasets contain only fragmented, non-stitched high-quality sub-trajectories. The goal is a framework that adapts how context is built and how trajectories are combined, so performance stays stable even when tasks and data quality are highly heterogeneous.

Key Contributions

1. A data-centric diagnosis of Offline MTRL. The authors identify three bottlenecks that limit learnability and generalization: prompt length mismatch under heterogeneous task complexities, semantic irrelevance of randomly or positionally sampled prompt segments, and misleading supervision caused by fragmented, discontinuous trajectories.
2. Three mechanisms mapped one-to-one onto those bottlenecks. Length-Gated Prompt Masking (LGPM) learns a task-specific effective prompt length; Retrieval-Augmented Prompt Construction (RAPC) replaces static prompts with semantically aligned high-return segments retrieved using the model's own GPT-2 embedding layer; Value-Adaptive Return Calibration (VARC) trains a task-conditioned value function with per-task adaptive expectile parameters to calibrate achievable returns at inference.
3. A coordinated pipeline rather than three isolated fixes. RAPC retrieves relevant high-return segments, VARC filters them by achievable-return consistency, and LGPM trims the retrieved prompt to the learned task-specific length, so adaptive context construction and trajectory stitching reinforce each other.
4. Extensive evaluation across data qualities and task counts. Experiments on Meta-World MT5, MT30 and MT50 under near-optimal and sub-optimal datasets, plus zero-shot generalization to unseen continuous-control tasks and a sparse-reward Maze2D test, with ablations isolating each component.

Main Findings

• Consistent gains over state of the art on Meta-World: DaCe-DT reaches average success rates of 100.0 on MT5, 93.33 on MT30 and 88.40 on MT50 under near-optimal data, and 88.00, 75.67 and 64.00 under sub-optimal data. The abstract reports an average improvement of 11.73% on optimal datasets and 13.34% on suboptimal datasets, measured against the state-of-the-art baseline HarmoDT.
• Per-configuration improvements vary: the paper reports gains of 12.96% on sub-optimal MT5, 11.73% and 13.34% on near-optimal and sub-optimal MT30 respectively, and 9.60% and 6.50% on MT50 respectively. Across all 50 tasks the introduction reports a 9.60% overall average success-rate improvement over the state of the art.
• Fixed prompt lengths are genuinely mismatched: on the Reach task, a short prompt with K = 5 already reaches a 100% success rate, while the multi-stage Coffee-Push task needs K = 30 to reach 80%. On Coffee-Push, an empirically tuned fixed length of 30 improves the single-task success rate by 60% over the state-of-the-art baseline; with adaptive prompting the success rate on Coffee-Push rises further to 100%. Adaptive length adjustment yields an additional 20% improvement over the best fixed baseline.
• Prompt selection is the second lever: the mean cosine similarity between prompt and current trajectory is 0.361 for fixed prompts and 0.462 for random prompts, while RAPC raises it to 0.741, relative improvements of 105.6% and 60.3%. RAPC improves the final success rate by 8.4% over random selection and 9.6% over fixed selection, and introducing RAPC alone yields a 9.6% improvement over Prompt-DT.
• Trajectory stitching is the third lever: return-conditioned Decision Transformers cannot adjust return-to-go at a junction, so complementary sub-trajectories (one succeeds early, the other late) cannot be recombined. VARC addresses this at inference by correcting the return target via a frozen value function; the paper notes that direct value-based correction under heterogeneous task data can induce biased reachable-return estimates, OOD conditioning and false stitching points, and that MTIQL's task-specific value functions and stitching modules add considerable overhead and do not fit prompt-based models without explicit task IDs.
• Ablation on MT30 confirms all three components matter: no components gives 71.3 success / 2436.8 return; LGPM only 83.7 / 2864.5; RAPC only 85.3 / 2961.4; VARC only 82.6 / 2735.6; all three together 93.3 / 3185.5.
• Zero-shot generalization to unseen tasks: on Cheetah-dir, Cheetah-vel and Ant-dir, DaCe-DT scores 966.9 ± 2.2, -22.0 ± 1.2 and 355.9 ± 1.7 versus HarmoDT-F's 958.5 ± 1.5, -66.5 ± 1.2 and 298.3 ± 1.0, for an average of 433.6 against 396.8, a 9.27% improvement.
• Robustness under sparse reward: with the value function disabled, on eight custom Maze2D layouts with entirely sparse rewards, the method achieves an average score of 36.9 versus 18.2 for Prompt-DT, indicating that embedding-based retrieval alone improves prompt quality.
• Online baselines on MT50-rand (near-optimal only): CARE 46.12 ± 1.30, PaCo 54.31 ± 1.32, Soft-M 53.41 ± 0.72, D2R 63.53 ± 1.22, MOORE 72.90 ± 3.30, ARS-LoRA 83.20 ± 1.80, compared with DaCe-DT's 88.40.

Methodology in Plain English

The framework builds on the Decision Transformer, which treats offline RL as return-conditioned sequence prediction, and on prompt-based variants that prepend a task-specific context prompt so a shared Transformer can infer the task. DaCe-DT adds three pieces on top.

LGPM (adaptive prompt length). Each task gets a learnable scalar that is squashed through a sigmoid and rounded to give a discrete effective prompt length. A task-specific attention mask is injected into every self-attention layer, blocking prompt tokens older than that length while preserving causal order. Because rounding is not differentiable, gradients flow through it with a straight-through estimator, and a sparsity penalty on the sigmoid values stops tasks from drifting to the maximum length, pushing each task toward a minimal sufficient context.

RAPC (semantic retrieval). For each task, the top-20 trajectories by episodic return are selected, sliced with a sliding window, and encoded using only the model's own GPT-2 embedding projection — return-to-go, state, action and timestep embeddings summed and averaged, then L2-normalized — to build a per-task index. At each decision step, the current history is encoded the same way and the most similar slice is retrieved by maximum inner product search, subject to a filter that the slice's source trajectory return meets or exceeds the current return target. Because embeddings are meaningless early in training, a progressive strategy uses static prompts for a warm-up period and refreshes the index periodically afterwards.

VARC (value-based return calibration). A value function sharing the same GPT-2 architecture is trained with per-task adaptive expectile parameters, so noisy tasks converge toward conservative quantiles and cleaner tasks toward aggressive ones. After convergence the value function and expectile parameters are frozen; they are not used during policy training. At inference, the raw return target is replaced by the value estimate whenever the two differ by more than a tolerance threshold, and the corrected return then guides retrieval and prompt assembly. The retrieved slice is cropped or extended to the task-optimal length learned by LGPM before being fed to the policy.

Training data comes from SAC agents trained per task; the near-optimal dataset aggregates 100M transitions and the sub-optimal dataset uses the first 50% of each task's replay buffer, 50M transitions total.

Why This Matters

The work reframes Offline MTRL as a data problem rather than purely an architecture problem: the same Transformer, fed better-managed context and better-calibrated return targets, outperforms heavier algorithmic alternatives. It also shows that prompt length is a dynamic, task-and-dataset-dependent quantity rather than a hyperparameter to be tuned once, which has direct implications for how prompt-based sequence models are configured.

Real-world applications:

• Robotic manipulation and industrial robotics. The evaluation itself uses a Sawyer robotic arm across 50 manipulation tasks with a 4-dimensional continuous action vector; the same setting applies to warehouse picking, assembly and tool use where a single policy must cover many object interactions.
• Autonomous driving. The paper names autonomous driving among domains that benefit from learning without risky active exploration, where logged driving data is plentiful but heterogeneous and imperfect.
• Healthcare. Also named by the authors as a safety-critical domain where collecting exploratory data is unacceptable, making offline learning from historical records the practical route.
• Any log-based control setting with many related tasks. Whenever a fixed archive of interactions must serve a large family of related goals and only partial high-quality behavior is present in the logs, the retrieval-plus-stitching combination is directly relevant.

Industry relevance: the reported gains are largest under the sub-optimal regime, which is the realistic case for production logs. The ability to generalize zero-shot to unseen tasks at zero training cost, using an index built from a small set of expert segments with frozen embeddings, matters for deployments where retraining per new task is expensive. Comparing favorably with online methods that require environment interaction (ARS-LoRA 83.20 versus 88.40 on MT50-rand near-optimal) is relevant where interaction is costly or unsafe.

Future Directions

The paper does not report an explicit future-work section in the available content; the following are open questions the work raises.

• How far does the retrieval index scale beyond the tested task counts? The index is built from the top-20 trajectories per task; the scaling behavior for substantially larger task families or much longer horizons is not reported.
• How do the adaptive expectile parameters and learned prompt lengths transfer to unseen tasks? For zero-shot evaluation the prompt length is set to the average learned length from training tasks, so task-specific adaptation for genuinely new tasks remains open.
• How sensitive is VARC to the tolerance threshold and to the warm-up and index-refresh schedule? These are the key inference-time knobs introduced by the design, and the available content does not report a sensitivity analysis.
• Can the three mechanisms be simplified or partially merged? The ablation shows each component contributes individually, but whether the coordinated three-stage inference pipeline can be made cheaper for latency-constrained deployment is not addressed.

Target Audience

Researchers and engineers working on offline reinforcement learning, multi-task learning, and sequence-model-based decision making, particularly those using Decision Transformers and prompt-based task conditioning. It is also relevant to practitioners deploying robots or log-driven control systems where a single policy must handle heterogeneous tasks from imperfect, fixed datasets, and to readers interested in data-centric methods as an alternative to architectural or algorithmic complexity.

Authors’ 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.