Learning Compositional Koopman Operators for Model-Based Control

Yunzhu Li*      Hao He*      Jiajun Wu      Dina Katabi      Antonio Torralba

(* indicate equal contribution)


Finding an embedding space for a linear approximation of a nonlinear dynamical system enables efficient system identification and control synthesis. The Koopman operator theory lays the foundation for identifying the nonlinear-to-linear coordinate transformations with data-driven methods. Recently, researchers have proposed to use deep neural networks as a more expressive class of basis functions for calculating the Koopman operators. These approaches, however, assume a fixed dimensional state space; they are therefore not applicable to scenarios with a variable number of objects. In this paper, we propose to learn compositional Koopman operators, using graph neural networks to encode the state into object-centric embeddings and using a block-wise linear transition matrix to regularize the shared structure across objects. The learned dynamics can quickly adapt to new environments of unknown physical parameters and produce control signals to achieve a specified goal. Our experiments on manipulating ropes and controlling soft robots show that the proposed method has better efficiency and generalization ability than existing baselines.


Yunzhu Li*, Hao He*, Jiajun Wu, Dina Katabi, and Antonio Torralba
Learning Compositional Koopman Operators for Model-Based Control
ICLR 2020, [Project] [Paper] [BibTex]
Spotlight Presentation
Abridged in NeurIPS 2019 workshop on Graph Representation Learning [Link]

Control Results
(Goal configuration is shown in red.)

Rope Manipulation

Soft Robot Swing
(Red means contract, green indicates expand)

Soft Robot Swim
(Red means contract, green indicates expand)


Related Work

Bethany Lusch, J. Nathan Kutz, Steven L. Brunton
Deep learning for universal linear embeddings of nonlinear dynamics
Nature communications, 4950 (2018)

Yunzhu Li, Jiajun Wu, Jun-Yan Zhu, Joshua B. Tenenbaum, Antonio Torralba, and Russ Tedrake
Propagation Networks for Model-Based Control Under Partial Observation
ICRA 2019, [Project]

Peter W. Battaglia, Razvan Pascanu, Matthew Lai, Danilo Rezende, Koray Kavukcuoglu
Interaction Networks for Learning about Objects, Relations and Physics
NeurIPS 2016