charleshsc / harmodt Goto Github PK

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

| 📑 Paper | 🐱 Github Repo |

Shengchao Hu^1,2*, Ziqing Fan^1,2*, Li Shen^3,4, Ya Zhang^1,2, Yanfeng Wang^1,2, Dacheng Tao⁵

¹ Shanghai Jiao Tong University, ² Shanghai AI Laboratory, ³ Sun Yat-sen University, ⁴ JD Explore Academy, ⁵ Nanyang Technological University.

• Overview
• Quick Start
• Citation
• Acknowledgements

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.

Download the dataset MT50 via this Google Drive link.

When your environment is ready, you could run the following script:

python main.py --seed 123 --data_path ./MT50 --prefix_name MT5 # MT30, MT50

If you find this work is relevant with your research or applications, please feel free to cite our work!

@inproceedings{HarmoDT,
    title={HarmoDT: Harmony Multi-Task Decision Transformer for Offline Reinforcement Learning},
    author={Hu, Shengchao and Fan, Ziqing and Shen, Li and Zhang, Ya and Wang, Yanfeng and Tao, Dacheng},
    booktitle={International Conference on Machine Learning},
    year={2024},
}

This repo benefits from DT and MTDiff. Thanks for their wonderful works!