Table of Contents
Fetching ...

HarmoDT: Harmony Multi-Task Decision Transformer for Offline Reinforcement Learning

Shengchao Hu, Ziqing Fan, Li Shen, Ya Zhang, Yanfeng Wang, Dacheng Tao

TL;DR

HarmoDT tackles offline multi-task reinforcement learning by learning per-task harmony subspaces through trainable masks within a Transformer-based policy. It casts the problem as a bi-level, gradient-based meta-learning task where the outer level learns task-specific masks and the inner level updates parameters under those masks to maximize collective performance, thereby mitigating gradient conflicts across tasks. Empirically, HarmoDT and its variants surpass strong offline and online baselines on the Meta-World MT50 benchmark, including unseen tasks, and show favorable scalability and robustness across hyper-parameters and model sizes. The work demonstrates that structured parameter sharing via harmony subspaces can unlock efficient, generalizable multi-task policies in offline RL, with practical implications for robotics and other domains requiring offline, multi-task decision making.

Abstract

The purpose of offline multi-task reinforcement learning (MTRL) is to develop a unified policy applicable to diverse tasks without the need for online environmental interaction. Recent advancements approach this through sequence modeling, leveraging the Transformer architecture's scalability and the benefits of parameter sharing to exploit task similarities. However, variations in task content and complexity pose significant challenges in policy formulation, necessitating judicious parameter sharing and management of conflicting gradients for optimal policy performance. In this work, we introduce the Harmony Multi-Task Decision Transformer (HarmoDT), a novel solution designed to identify an optimal harmony subspace of parameters for each task. We approach this as a bi-level optimization problem, employing a meta-learning framework that leverages gradient-based techniques. The upper level of this framework is dedicated to learning a task-specific mask that delineates the harmony subspace, while the inner level focuses on updating parameters to enhance the overall performance of the unified policy. Empirical evaluations on a series of benchmarks demonstrate the superiority of HarmoDT, verifying the effectiveness of our approach.

HarmoDT: Harmony Multi-Task Decision Transformer for Offline Reinforcement Learning

TL;DR

HarmoDT tackles offline multi-task reinforcement learning by learning per-task harmony subspaces through trainable masks within a Transformer-based policy. It casts the problem as a bi-level, gradient-based meta-learning task where the outer level learns task-specific masks and the inner level updates parameters under those masks to maximize collective performance, thereby mitigating gradient conflicts across tasks. Empirically, HarmoDT and its variants surpass strong offline and online baselines on the Meta-World MT50 benchmark, including unseen tasks, and show favorable scalability and robustness across hyper-parameters and model sizes. The work demonstrates that structured parameter sharing via harmony subspaces can unlock efficient, generalizable multi-task policies in offline RL, with practical implications for robotics and other domains requiring offline, multi-task decision making.

Abstract

The purpose of offline multi-task reinforcement learning (MTRL) is to develop a unified policy applicable to diverse tasks without the need for online environmental interaction. Recent advancements approach this through sequence modeling, leveraging the Transformer architecture's scalability and the benefits of parameter sharing to exploit task similarities. However, variations in task content and complexity pose significant challenges in policy formulation, necessitating judicious parameter sharing and management of conflicting gradients for optimal policy performance. In this work, we introduce the Harmony Multi-Task Decision Transformer (HarmoDT), a novel solution designed to identify an optimal harmony subspace of parameters for each task. We approach this as a bi-level optimization problem, employing a meta-learning framework that leverages gradient-based techniques. The upper level of this framework is dedicated to learning a task-specific mask that delineates the harmony subspace, while the inner level focuses on updating parameters to enhance the overall performance of the unified policy. Empirical evaluations on a series of benchmarks demonstrate the superiority of HarmoDT, verifying the effectiveness of our approach.
Paper Structure (31 sections, 11 equations, 6 figures, 6 tables, 2 algorithms)

This paper contains 31 sections, 11 equations, 6 figures, 6 tables, 2 algorithms.

Figures (6)

  • Figure 1: Illustration of a comparative analysis of success rates across various task numbers within the Meta-World benchmark, focusing on prevalent MTRL algorithms. An in-depth exploration of these results refers to Section \ref{['sec:exp']}.
  • Figure 2: Illustration of the harmony degree among trainable weights during training for policies with and without randomly initialized masks (left panel), and the success rates achieved when applying masks with varying sparsity levels (right panel).
  • Figure 3: Illustration of the conflicting problem and the framework of our method HarmoDT to find a harmony subspace for each task. The left panel shows the conflicting phenomenon reflected by divergent task-specific gradients. The middle panel illustrates the procedure to find a harmony subspace for each task via the strategic learning of task masks. The right panel demonstrates the workflow of HarmoDT based on the DT architecture with prompts and learned harmony subspace of weights when handling a task, such as ${\mathcal{T}}_3$.
  • Figure 4: From the left to right, we illustrate the ablation results on the Meta-World benchmark with 50 tasks under the near-optimal case. Default values are listed as $\eta_{min}=0$, $\eta_{max}$ is 100 (about 1e-3% of total weights), $\mathrm{S}=0.2$, $\lambda=10$ and $t_m=5e3$. During each individual ablation, a single parameter is varied, with all other parameters maintained at their default values. Detailed results pertaining to additional settings are comprehensively documented in the Appendix \ref{['sec:moreab']}.
  • Figure 5: The t-SNE visualization of optimal subspace via masks learned by our HarmoDT on the 30 tasks of Meta-World benchmark. The figure illustrates the relational dynamics of task-specific masks, with a focus on 10 representative tasks from the total set.
  • ...and 1 more figures

Theorems & Definitions (4)

  • Definition 3.1: Harmony Score on a Single Weight
  • Definition 3.2: Averaged Harmony Score
  • Definition 4.1: Agreement Score
  • Definition 4.2: Importance Score