Learning Multi-Object Dynamics with Compositional Neural Radiance Fields

Danny Driess, Zhiao Huang, Yunzhu Li, Russ Tedrake, Marc Toussaint

Keywords: Neural Radiance Fields, Dynamics Model Learning, Neural Implicit Representations, Visual Prediction, Graph Neural Networks, Representation Learning, Robotic Manipulation

TL;DR: We present a method to learn compositional multi-object dynamics models from image observations based on implicit object encoders, Neural Radiance Fields (NeRFs), and graph neural networks.

Abstract:

We present a method to learn compositional multi-object dynamics models from image observations based on implicit object encoders, Neural Radiance Fields (NeRFs), and graph neural networks. NeRFs have become a popular choice for representing scenes due to their strong 3D prior. However, most NeRF approaches are trained on a single scene, representing the whole scene with a global model, making generalization to novel scenes, containing different numbers of objects, challenging. Instead, we present a compositional, object-centric auto-encoder framework that maps multiple views of the scene to a set of latent vectors representing each object separately. The latent vectors parameterize individual NeRFs from which the scene can be reconstructed. Based on those latent vectors, we train a graph neural network dynamics model in the latent space to achieve compositionality for dynamics prediction. A key feature of our approach is that the latent vectors are forced to encode 3D information through the NeRF decoder, which enables us to incorporate structural priors in learning the dynamics models, making long-term predictions more stable compared to several baselines. Simulated and real world experiments show that our method can model and learn the dynamics of compositional scenes including rigid and deformable objects.
Video: https://dannydriess.github.io/compnerfdyn/

Student First Author: yes

Website: https://dannydriess.github.io/compnerfdyn/

Community Implementations: [ 2 code implementations](https://www.catalyzex.com/paper/arxiv:2202.11855/code)