Generalization of Compositional Tasks with Logical Specification via Implicit Planning

Abstract

In this study, we address the challenge of learning generalizable policies for compositional tasks defined by logical specifications. These tasks consist of multiple temporally extended sub-tasks. Due to the sub-task inter-dependencies and sparse reward issue in long-horizon tasks, existing reinforcement learning (RL) approaches, such as task-conditioned and goal-conditioned policies, continue to struggle with slow convergence and sub-optimal performance in the generalization of compositional tasks. To overcome these limitations, by decomposing the given task into reach-avoid sub-tasks, we introduce a new hierarchical RL framework that trains a high-level planner to select optimal sub-tasks and zero-shot generalizes to other tasks in the sub-task level, which enhances the efficiency and optimality of task generalization. At the high level, we present an implicit planner specifically designed for generalizing compositional tasks. This planner selects the next sub-task and estimates the multi-step return for completing the remaining task from the current state. It learns a latent transition model and performs planning in the latent space to select sub-tasks based on a graph neural network (GNN). Subsequently, the sub-task assigned by the high level guides the low-level module to effectively handle long-horizon tasks, while the estimated return encourages the low-level policy to account for future sub-task dependencies, enhancing its optimality and densifying the sparse rewards. We conduct comprehensive experiments to demonstrate the framework’s advantages over previous methods in terms of both efficiency and optimality.