Policy Architectures for Compositional Generalization in Control
This addresses the challenge of compositional generalization in control and robotics, enabling policies to extrapolate to different numbers of entities and compose skills in novel ways, which is incremental but impactful for complex task learning.
The paper tackles the problem of learning goal-conditioned policies that struggle with compositional generalization in control tasks, and introduces a framework with policy architectures like Deep Sets and Self Attention that achieve significantly higher success rates with less data on simulated robot manipulation tasks.
Many tasks in control, robotics, and planning can be specified using desired goal configurations for various entities in the environment. Learning goal-conditioned policies is a natural paradigm to solve such tasks. However, current approaches struggle to learn and generalize as task complexity increases, such as variations in number of environment entities or compositions of goals. In this work, we introduce a framework for modeling entity-based compositional structure in tasks, and create suitable policy designs that can leverage this structure. Our policies, which utilize architectures like Deep Sets and Self Attention, are flexible and can be trained end-to-end without requiring any action primitives. When trained using standard reinforcement and imitation learning methods on a suite of simulated robot manipulation tasks, we find that these architectures achieve significantly higher success rates with less data. We also find these architectures enable broader and compositional generalization, producing policies that extrapolate to different numbers of entities than seen in training, and stitch together (i.e. compose) learned skills in novel ways. Videos of the results can be found at https://sites.google.com/view/comp-gen-rl.