Thirty heterogeneous control tasks share one sequence backbone: HeteroGrad-DT lifts average normalized return by 32.3%
Synopsis
In heterogeneous multi-task offline reinforcement learning, tasks differ in state dimension, action space and dynamics, so a directly shared backbone hits both interface mismatch and gradient conflict. HeteroGrad-DT gives each task its own input tokenizer and prediction head around a shared sequence backbone, and a Gradient Flow Transformer pools per-task gradients over a synchronization interval and learns fusion weights to coordinate the shared update. On the DeepMind Control Suite with up to 30 tasks, its average normalized return exceeds representative multi-task methods, with a 32.3% relative gain over standard shared-backbone training and 8.6% over a PCGrad baseline using the same task-aware interfaces.
Interpretation
Multi-task offline reinforcement learning can now let control tasks with different state and action dimensions share one sequence backbone without first padding them to a common input and output size. Earlier multi-task Decision Transformer methods relied mainly on a shared backbone, task prompts or parameter masks, and still forced heterogeneous state-action spaces into one interface, where zero-padding introduced dimensions with no semantic content. On the 30-task DeepMind Control Suite setting, HeteroGrad-DT reaches an average normalized return of 42.40, 10.35 points or about 32.3% above the directly shared-backbone MTDT, over five random seeds.
The shared backbone is no longer updated only by per-step gradient projection or static regularization; every training interval, gradients are pooled per task and then coordinated together. PCGrad-style methods correct instantaneous conflicts at each step, whereas HeteroGrad-DT averages gradients accumulated within a synchronization interval and feeds them, with task statistics, into a Gradient Flow Transformer that learns fusion weights. With the same task-aware interfaces, it scores 3.35 points, about 8.6%, above HAP plus PCGrad on 30 tasks; removing the Gradient Flow Transformer lowers the 5-, 10- and 30-task scores from 83.46, 77.99 and 42.40 to 80.35, 75.24 and 40.69.
Each task uses its own linear projection to turn states and actions into tokens of a common dimension, and its own prediction head to map the shared representation back to that task's action space. A zero-padded unified interface lets the shared model treat structurally incomparable dimensions as aligned, and entangles task identity with padding patterns. On the fixed MT10 validation split, the task-aware interface generally lowers action-prediction mean squared error and absolute mean error, but the saturation-gap results do not improve uniformly across tasks.
Perspective
Evaluation is limited to state-based simulated control on the DeepMind Control Suite, with at most 30 tasks, medium-replay datasets of 2,000 trajectories per task, five random seeds and 1M training iterations. The training-cost comparison runs on a single NVIDIA RTX 4090, where HeteroGrad-DT has 7.36M parameters, 11928 MiB peak memory, 43.54 iterations per second throughput and 6.38 hours of training.
Performance under physical robot deployment remains to be validated, since real systems bring sensing noise, actuation delay, hardware constraints, dynamics mismatch and safety requirements absent from the current evaluation. How the shared backbone and the gradient descriptor behave as task collections grow well beyond 30 also remains to be examined, and the fixed-dimensional task statistics and interval-mean gradient digest deliberately compress information. Future work will evaluate physical robots or higher-fidelity sim-to-real benchmarks, extend to visual observations, and study adaptive descriptor construction and synchronization policies.