# Factored World Models for Zero-Shot Generalization in Robotic Manipulation

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

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Placed in Embodied and robot world models, one of 11 directions in World models.
- Builds on 27 of its 44 references that are in the map.
- First posted on arXiv 10 Feb 2022.

## Placed in

- World models / [Embodied and robot world models](https://mlsanity.com/wm/embodied-and-robot-world-models/)

## Builds on (most cited first, 27 of 27)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Object-Centric Learning with Slot Attention](https://mlsanity.com/paper/2006.15055/) (2020)
- [Object Detection with Deep Learning: A Review](https://mlsanity.com/paper/1807.05511/) (2018)
- [Interaction Networks for Learning about Objects, Relations and Physics](https://mlsanity.com/paper/1612.00222/) (2016)
- [Neural Relational Inference for Interacting Systems](https://mlsanity.com/paper/1802.04687/) (2018)
- [Embed to Control: A Locally Linear Latent Dynamics Model for Control from Raw Images](https://mlsanity.com/paper/1506.07365/) (2015)
- [Visual Reinforcement Learning with Imagined Goals](https://mlsanity.com/paper/1807.04742/) (2018)
- [Graph networks as learnable physics engines for inference and control](https://mlsanity.com/paper/1806.01242/) (2018)
- [On the Binding Problem in Artificial Neural Networks](https://mlsanity.com/paper/2012.05208/) (2020)
- [Contrastive Learning of Structured World Models](https://mlsanity.com/paper/1911.12247/) (2019)
- [Deep Variational Bayes Filters: Unsupervised Learning of State Space Models from Raw Data](https://mlsanity.com/paper/1605.06432/) (2016)
- [Planning with Goal-Conditioned Policies](https://mlsanity.com/paper/1911.08453/) (2019)
- [PHYRE: A New Benchmark for Physical Reasoning](https://mlsanity.com/paper/1908.05656/) (2019)
- [State Representation Learning for Control: An Overview](https://mlsanity.com/paper/1802.04181/) (2018)
- [Entity Abstraction in Visual Model-Based Reinforcement Learning](https://mlsanity.com/paper/1910.12827/) (2019)
- [Towards Practical Multi-Object Manipulation using Relational Reinforcement Learning](https://mlsanity.com/paper/1912.11032/) (2019)
- [Reasoning About Physical Interactions with Object-Oriented Prediction and Planning](https://mlsanity.com/paper/1812.10972/) (2018)
- [Learning Long-term Visual Dynamics with Region Proposal Interaction Networks](https://mlsanity.com/paper/2008.02265/) (2020)
- [Structured Object-Aware Physics Prediction for Video Modeling and Planning](https://mlsanity.com/paper/1910.02425/) (2019)
- [Robust Locally-Linear Controllable Embedding](https://mlsanity.com/paper/1710.05373/) (2017)
- [Efficient Model-Based Deep Reinforcement Learning with Variational State Tabulation](https://mlsanity.com/paper/1802.04325/) (2018)
- [Object-centric Forward Modeling for Model Predictive Control](https://mlsanity.com/paper/1910.03568/) (2019)
- [Prediction, Consistency, Curvature: Representation Learning for Locally-Linear Control](https://mlsanity.com/paper/1909.01506/) (2019)
- [Policy learning in SE(3) action spaces](https://mlsanity.com/paper/2010.02798/) (2020)
- [Action Priors for Large Action Spaces in Robotics](https://mlsanity.com/paper/2101.04178/) (2021)

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

- [Curious Exploration via Structured World Models Yields Zero-Shot Object Manipulation](https://mlsanity.com/paper/2206.11403/) (2022)
- [Structure in Deep Reinforcement Learning: A Survey and Open Problems](https://mlsanity.com/paper/2306.16021/) (2023)
- [Learning Dynamic Attribute-factored World Models for Efficient Multi-object Reinforcement Learning](https://mlsanity.com/paper/2307.09205/) (2023)
- [Binding Actions to Objects in World Models](https://mlsanity.com/paper/2204.13022/) (2022)
- [Regularity as Intrinsic Reward for Free Play](https://mlsanity.com/paper/2312.01473/) (2023)
- [Generalization properties of contrastive world models](https://mlsanity.com/paper/2401.00057/) (2023)

## Abstract (arXiv)

World models for environments with many objects face a combinatorial explosion of states: as the number of objects increases, the number of possible arrangements grows exponentially. In this paper, we learn to generalize over robotic pick-and-place tasks using object-factored world models, which combat the combinatorial explosion by ensuring that predictions are equivariant to permutations of objects. Previous object-factored models were limited either by their inability to model actions, or by their inability to plan for complex manipulation tasks. We build on recent contrastive methods for training object-factored world models, which we extend to model continuous robot actions and to accurately predict the physics of robotic pick-and-place. To do so, we use a residual stack of graph neural networks that receive action information at multiple levels in both their node and edge neural networks. Crucially, our learned model can make predictions about tasks not represented in the training data. That is, we demonstrate successful zero-shot generalization to novel tasks, with only a minor decrease in model performance. Moreover, we show that an ensemble of our models can be used to plan for tasks involving up to 12 pick and place actions using heuristic search. We also demonstrate transfer to a physical robot.