# Policy Architectures for Compositional Generalization in Control

arXiv: https://arxiv.org/abs/2203.05960 · HTML: https://mlsanity.com/paper/2203.05960/

## Key facts

- Cited by 23 papers in the map (data through 14 Oct 2026).
- Builds on 37 of its 49 references that are in the map.
- First posted on arXiv 10 Mar 2022.

## Builds on (most cited first, 30 of 37)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Semi-Supervised Classification with Graph Convolutional Networks](https://mlsanity.com/paper/1609.02907/) (2016)
- [Densely Connected Convolutional Networks](https://mlsanity.com/paper/1608.06993/) (2016)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Continuous control with deep reinforcement learning](https://mlsanity.com/paper/1509.02971/) (2015)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [Deep Sets](https://mlsanity.com/paper/1703.06114/) (2017)
- [Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gauges](https://mlsanity.com/paper/2104.13478/) (2021)
- [Hindsight Experience Replay](https://mlsanity.com/paper/1707.01495/) (2017)
- [Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation](https://mlsanity.com/paper/2108.12409/) (2021)
- [Object-Centric Learning with Slot Attention](https://mlsanity.com/paper/2006.15055/) (2020)
- [Offline Reinforcement Learning as One Big Sequence Modeling Problem](https://mlsanity.com/paper/2106.02039/) (2021)
- [Data-Efficient Hierarchical Reinforcement Learning](https://mlsanity.com/paper/1805.08296/) (2018)
- [The Option-Critic Architecture](https://mlsanity.com/paper/1609.05140/) (2016)
- [MONet: Unsupervised Scene Decomposition and Representation](https://mlsanity.com/paper/1901.11390/) (2019)
- [Graph networks as learnable physics engines for inference and control](https://mlsanity.com/paper/1806.01242/) (2018)
- [Recurrent Independent Mechanisms](https://mlsanity.com/paper/1909.10893/) (2019)
- [Contrastive Learning of Structured World Models](https://mlsanity.com/paper/1911.12247/) (2019)
- [Rearrangement: A Challenge for Embodied AI](https://mlsanity.com/paper/2011.01975/) (2020)
- [Language as an Abstraction for Hierarchical Deep Reinforcement Learning](https://mlsanity.com/paper/1906.07343/) (2019)
- [Entity Abstraction in Visual Model-Based Reinforcement Learning](https://mlsanity.com/paper/1910.12827/) (2019)
- [Relational Deep Reinforcement Learning](https://mlsanity.com/paper/1806.01830/) (2018)
- [Learning Multi-Object Dynamics with Compositional Neural Radiance Fields](https://mlsanity.com/paper/2202.11855/) (2022)
- [Neural Task Graphs: Generalizing to Unseen Tasks from a Single Video Demonstration](https://mlsanity.com/paper/1807.03480/) (2018)
- [Neural Dynamic Policies for End-to-End Sensorimotor Learning](https://mlsanity.com/paper/2012.02788/) (2020)
- [Towards Practical Multi-Object Manipulation using Relational Reinforcement Learning](https://mlsanity.com/paper/1912.11032/) (2019)
- [Structured agents for physical construction](https://mlsanity.com/paper/1904.03177/) (2019)
- [The Sensory Neuron as a Transformer: Permutation-Invariant Neural Networks for Reinforcement Learning](https://mlsanity.com/paper/2109.02869/) (2021)

## Built on by (most cited first, 23 of 23)

- [An Investigation into Pre-Training Object-Centric Representations for Reinforcement Learning](https://mlsanity.com/paper/2302.04419/) (2023)
- [MoCoDA: Model-based Counterfactual Data Augmentation](https://mlsanity.com/paper/2210.11287/) (2022)
- [Masked Trajectory Models for Prediction, Representation, and Control](https://mlsanity.com/paper/2305.02968/) (2023)
- [Efficient Data Collection for Robotic Manipulation via Compositional Generalization](https://mlsanity.com/paper/2403.05110/) (2024)
- [Factored Adaptation for Non-Stationary Reinforcement Learning](https://mlsanity.com/paper/2203.16582/) (2022)
- [Entity-Centric Reinforcement Learning for Object Manipulation from Pixels](https://mlsanity.com/paper/2404.01220/) (2024)
- [When does compositional structure yield compositional generalization? A kernel theory](https://mlsanity.com/paper/2405.16391/) (2024)
- [Learning to Extrapolate: A Transductive Approach](https://mlsanity.com/paper/2304.14329/) (2023)
- [Structure in Deep Reinforcement Learning: A Survey and Open Problems](https://mlsanity.com/paper/2306.16021/) (2023)
- [Composing Pre-Trained Object-Centric Representations for Robotics From "What" and "Where" Foundation Models](https://mlsanity.com/paper/2404.13474/) (2024)
- [Learning Dynamic Attribute-factored World Models for Efficient Multi-object Reinforcement Learning](https://mlsanity.com/paper/2307.09205/) (2023)
- [A System for Morphology-Task Generalization via Unified Representation and Behavior Distillation](https://mlsanity.com/paper/2211.14296/) (2022)
- [A Survey on Compositional Generalization in Applications](https://mlsanity.com/paper/2302.01067/) (2023)
- [Entity-based Reinforcement Learning for Autonomous Cyber Defence](https://mlsanity.com/paper/2410.17647/) (2024)
- [Learning Interactive World Model for Object-Centric Reinforcement Learning](https://mlsanity.com/paper/2511.02225/) (2025)
- [Hierarchical Entity-centric Reinforcement Learning with Factored Subgoal Diffusion](https://mlsanity.com/paper/2602.02722/) (2026)
- [Zero-Shot Policy Transfer with Disentangled Task Representation of Meta-Reinforcement Learning](https://mlsanity.com/paper/2210.00350/) (2022)
- [Learning Causal Dynamics Models in Object-Oriented Environments](https://mlsanity.com/paper/2405.12615/) (2024)
- [Language-guided Task Adaptation for Imitation Learning](https://mlsanity.com/paper/2301.09770/) (2023)
- [Efficient Exploration and Discriminative World Model Learning with an Object-Centric Abstraction](https://mlsanity.com/paper/2408.11816/) (2024)
- [UniCO: Towards a Unified Model for Combinatorial Optimization Problems](https://mlsanity.com/paper/2505.06290/) (2025)
- [Neural Operators for Multi-Task Control and Adaptation](https://mlsanity.com/paper/2604.03449/) (2026)
- [Iterative Compositional Data Generation for Robot Control](https://mlsanity.com/paper/2512.10891/) (2025)

## Abstract (arXiv)

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.