Identifying Selections for Unsupervised Subtask Discovery

TL;DR

This paper introduces a theory and method for identifying subtask selection variables in reinforcement learning data, enabling better subtask discovery and improved generalization in multi-task scenarios.

Contribution

It presents a novel theory to identify selection-based subtasks and develops a sequential NMF method to extract meaningful subgoals from data.

Findings

Learned subtasks improve task generalization in multi-task imitation learning.

The method effectively identifies subgoals as selection variables in complex environments.

Empirical validation on a challenging Kitchen environment demonstrates success.

Abstract

When solving long-horizon tasks, it is intriguing to decompose the high-level task into subtasks. Decomposing experiences into reusable subtasks can improve data efficiency, accelerate policy generalization, and in general provide promising solutions to multi-task reinforcement learning and imitation learning problems. However, the concept of subtasks is not sufficiently understood and modeled yet, and existing works often overlook the true structure of the data generation process: subtasks are the results of a $\textit{selection}$ mechanism on actions, rather than possible underlying confounders or intermediates. Specifically, we provide a theory to identify, and experiments to verify the existence of selection variables in such data. These selections serve as subgoals that indicate subtasks and guide policy. In light of this idea, we develop a sequential non-negative matrix factorization (seq- NMF) method to learn these subgoals and extract meaningful behavior patterns as subtasks. Our empirical results on a challenging Kitchen environment demonstrate that the learned subtasks effectively enhance the generalization to new tasks in multi-task imitation learning scenarios. The codes are provided at https://anonymous.4open.science/r/Identifying\_Selections\_for\_Unsupervised\_Subtask\_Discovery/README.md.