Skills Regularized Task Decomposition for Multi-task Offline Reinforcement Learning

TL;DR

This paper introduces a skill-based multi-task offline reinforcement learning method that decomposes tasks into shared subtasks using a Wasserstein auto-encoder, effectively leveraging heterogeneous datasets of varying quality.

Contribution

It proposes a novel task decomposition technique with quality-weighted regularization and dataset augmentation to improve multi-task offline RL performance on complex robotic and navigation tasks.

Findings

Outperforms state-of-the-art algorithms on robotic manipulation tasks

Robust to datasets with mixed quality levels

Enhances policy learning with skill-based dataset augmentation

Abstract

Reinforcement learning (RL) with diverse offline datasets can have the advantage of leveraging the relation of multiple tasks and the common skills learned across those tasks, hence allowing us to deal with real-world complex problems efficiently in a data-driven way. In offline RL where only offline data is used and online interaction with the environment is restricted, it is yet difficult to achieve the optimal policy for multiple tasks, especially when the data quality varies for the tasks. In this paper, we present a skill-based multi-task RL technique on heterogeneous datasets that are generated by behavior policies of different quality. To learn the shareable knowledge across those datasets effectively, we employ a task decomposition method for which common skills are jointly learned and used as guidance to reformulate a task in shared and achievable subtasks. In this joint learning, we use Wasserstein auto-encoder (WAE) to represent both skills and tasks on the same latent space and use the quality-weighted loss as a regularization term to induce tasks to be decomposed into subtasks that are more consistent with high-quality skills than others. To improve the performance of offline RL agents learned on the latent space, we also augment datasets with imaginary trajectories relevant to high-quality skills for each task. Through experiments, we show that our multi-task offline RL approach is robust to the mixed configurations of different-quality datasets and it outperforms other state-of-the-art algorithms for several robotic manipulation tasks and drone navigation tasks.