HeteroGrad-DT: Gradient Flow Transformer for heterogeneous multi-task offline reinforcement learning
Year: 2026
Type: article
Abstract: Offline reinforcement learning enables policy learning from static datasets without online exploration, making it particularly suitable for robotic control under safety and data constraints. Decision Transformer recasts offline reinforcement learning as conditional sequence modeling. Extending it to heterogeneous multi-task settings is challenging because task representations and dynamics differ, while joint training can suffer from destructive gradient interference. To address these challenges, we present HeteroGrad-DT, an integrated framework that employs task-specific tokenizers and decoupled prediction heads to retain task-native input and output dimensionality at the model interfaces while learning a shared sequence backbone. It further introduces a Gradient Flow Transformer that periodically coordinates cross-task updates using fixed-dimensional task statistics, compact summaries of shared-backbone gradients accumulated within the current synchronization interval, and an alignment-based coordination objective, rather than relying only on stepwise projection or static regularization. Experiments on the DeepMind Control Suite with up to 30 tasks show that HeteroGrad-DT improves average normalized return over representative multi-task methods. In the 30-task setting, it achieves a 32.3% relative improvement over standard training with a shared backbone and an 8.6% gain over a representative gradient-conflict mitigation baseline using the same task-aware interfaces. Ablation, mechanism, and sensitivity analyses further support complementary benefits from task-aware interfaces and gradient coordination in the evaluated settings.
Source: PLoS ONE
Authors: Ye Wang, Donghui Ma, Jiang Li, Lihong Guo
Language: English
Subfield: Artificial Intelligence
Field: Computer Science
Domain: Physical Sciences
Open Access status: gold
Estimated APC paid: $2,382
OpenAlex ID: W7225783778