Learning Compositional Koopman Operators for Model-Based Control

1 INTRODUCTION . Simulating and controlling complex dynamical systems , such as ropes or soft robots , relies on two key features of the dynamics model : first , it needs to be efficient for system identification and motor control ; second , it needs to be generalizable to a complex , constantly evolving environments . In practice , computational models for complex , nonlinear dynamical systems are often not efficient enough for real-time control ( Mayne , 2000 ) . The Koopman operator theory suggests that identifying nonlinear-to-linear coordinate transformations allows efficient linear approximation of nonlinear systems ( Williams et al. , 2015 ; Mauroy & Goncalves , 2016 ) . Fast as they are , however , existing papers on Koopman operators focus on a single dynamical system , making it hard to generalize to cases where there are a variable number of components . In contrast , recent advances in approximating dynamics models with deep nets have demonstrated its power in characterizing complex , generic environments . In particular , a few recent papers have explored the use of graph nets in dynamics modeling , taking into account the state of each object as well as their interactions . This allows their models to generalize to scenarios with a variable number of objects ( Battaglia et al. , 2016 ; Chang et al. , 2017 ) . Despite their strong generalization power , they are not as efficient in system identification and control , because deep nets are heavily over-parameterized , making optimization time-consuming and sample-inefficient . In this paper , we propose compositional Koopman operators , integrating Koopman operators with graph networks for generalizable and efficient dynamics modeling . We build on the idea of encoding states into object-centric embeddings with graph neural networks , which ensures generalization power . But instead of using over-parameterized neural nets to model state transition , we identify the Koopman matrix and control matrix from data as a linear approximation of the nonlinear dynamical system . The linear approximation allows efficient system identification and control synthesis . The main challenge of extending Koopman theory to multi-object systems is scalability . The number of parameters in the Koopman matrix scales quadratically with the number of objects , which harms the learning efficiency and leads to overfitting . To tackle this issue , we exploit the structure of the ∗indicates equal contributions . Our project page : http : //koopman.csail.mit.edu underlying system and use the same block-wise Koopman sub-matrix for object pairs of the same relation . This significantly reduces the number of parameters that need to be identified by making it independent of the size of the system . Our experiments include simulating and controlling ropes of variable lengths and soft robots of different shapes . The compositional Koopman operators are significantly more accurate than the state-of-the-art learned physics engines ( Battaglia et al. , 2016 ; Li et al. , 2019b ) , and faster when adapting to new environments of unknown physical parameters . Our method also outperforms vanilla deep Koopman methods ( Lusch et al. , 2018 ; Morton et al. , 2018 ) and Koopman models with manually-designed basis functions , which shows the advantages of using a structured Koopman matrix and graph neural networks . Please see our project page for demonstrating videos . 2 RELATED WORK . Koopman operators . The Koopman operator formalism of dynamical systems is rooted in the seminal works of Koopman and Von Neumann in the early 1930s ( Koopman , 1931 ; Koopman & Neumann , 1932 ) . The core idea is to map the state of a nonlinear dynamical system to an embedding space , over which we can linearly propagate into the future . Researchers have proposed various algorithms to explore the Koopman spectral properties from data . A large portion of them are in the class of dynamic mode decomposition ( DMD ) ( Rowley et al. , 2009 ; Schmid , 2010 ; Tu et al. , 2014 ; Williams et al. , 2015 ; Arbabi & Mezic , 2017 ) . The linear representation will enable efficient prediction , estimation , and control using tools from linear dynamical systems ( Williams et al. , 2016 ; Proctor et al. , 2018 ; Mauroy & Goncalves , 2019 ; Korda & Mezić , 2018 ) . People have been using hand-designed Koopman observables for various modeling and control tasks ( Brunton et al. , 2016 ; Kaiser et al. , 2017 ; Abraham et al. , 2017 ; Bruder et al. , 2019b ; Arbabi et al. , 2018 ) . Some recent works have applied the method to the real world and successfully control soft robots with great precision ( Bruder et al. , 2019a ; Mamakoukas et al. , 2019 ) . However , hand-crafted basis functions sometimes fail to generalize to more complex environments . Learning these functions from data using neural nets turns out to generate a more expressive invariant subspace ( Lusch et al. , 2018 ; Takeishi et al. , 2017 ) and has achieved successes in fluid control ( Morton et al. , 2018 ) . Morton et al . ( 2019 ) has also extended the framework to account for uncertainty in the system by inferring a distribution over observations . Our model differs by explicitly modeling the compositionality of the underlying system with graph networks . It generalizes better to environments of a variable number of objects or soft robots of different shapes . Learning-based physical simulators . Battaglia et al . ( 2016 ) and Chang et al . ( 2017 ) first explored learning a simulator from data by approximating object interactions with neural networks . These models are no longer bounded to hard-coded physical rules , and can adapt to scenarios where the underlying physics is unknown . Please refer to Battaglia et al . ( 2018 ) for a full review . Recently , Mrowca et al . ( 2018 ) and Li et al . ( 2019a ) extended these models to approximate particle dynamics of deformable shapes and fluids . Flexible as they are , these models become less efficient during model adaptation in complex scenarios , because the optimization of neural networks usually needs a lot of samples and compute , which limits its use in an online setting . Nagabandi et al . ( 2019a ; b ) proposed to use meta-learning for online adaptation , and have shown to be effective in simulated robots and a real legged millirobot . However , it is not clear whether their methods can generalize to systems with variable numbers of instances . The use of graph nets and Koopman operators in our model allows better generalization ability and enables efficient system identification as we only need to identify the transition matrices , which is essentially a least-square problem and can be solved very efficiently . People have also used the learned physics engines for planning and control . Many previous papers in this direction learn a latent dynamics model together with a policy in a model-based reinforcement learning setup ( Racanière et al. , 2017 ; Hamrick et al. , 2017 ; Pascanu et al. , 2017 ; Hafner et al. , 2019 ) ; a few alternatives use the learned model in model-predictive control ( MPC ) ( Sanchez-Gonzalez et al. , 2018 ; Li et al. , 2019b ; Janner et al. , 2019 ) . In this paper , we leverage the fact that the embeddings in the Koopman space are propagating linearly through time , which allows us to formulate the control problem as quadratic programming and optimize the control signals much more efficiently . 3 APPROACH . We first present the basics of Koopman operators : for a nonlinear dynamical system , the Koopman observation functions can map the state space to an embedding space where the dynamics become linear . We then discuss the compositional nature of physical systems and show how graph networks can be used to capture the compositionality . 3.1 THE KOOPMAN OPERATORS . Let xt ∈ X ⊂ Rn be the state vector for the system at time step t. We consider a non-linear discrete-time dynamical system described by xt+1 = F ( xt ) . The Koopman operator ( Koopman , 1931 ) , denoted as K : F → F , is a linear transformation defined by Kg , g ◦ F , where F is the collection of all functions ( also referred to as observables ) that form an infinite-dimensional Hilbert space . For every function g : X → R belonging to F , we have ( Kg ) ( xt ) = g ( F ( xt ) ) = g ( xt+1 ) , ( 1 ) making the function space F invariant under the action of the Koopman operator . Although the theory guarantees the existence of the Koopman operator , its use in practice is limited by its infinite dimensionality . Most often , we assume there is an invariant subspace G of the Koopman operator . It spans by a set of base observation functions { g1 , · · · , gm } and satisfies that Kg ∈ G for any g ∈ G. With a slightly abuse of the notation , we now use g ( xt ) : Rn → Rm to represent [ g1 ( x t ) , · · · , gm ( xt ) ] T . By constraining the Koopman operator on this invariant subspace , we get a finite-dimensional linear operator K ∈ Rm×m that we refer as the Koopman matrix . Traditionally , people hand-craft base observation functions from the knowledge of underlying physics . The system identification problem is then reduced to finding the Koopman matrix K , which can be solved by linear regression given historical data of the system . Recently , researchers have also explored data-driven methods that automatically find the Koopman invariant subspace via representing the base observation functions g ( x ) via deep neural networks . Although the original Koopman theory does not consider the system with external control inputs , researchers have found that linearly injecting control signals to the Koopman observation space can give us good numerical performance ( Brunton et al. , 2016 ; Bruder et al. , 2019a ) . Mathematically , considering a dynamical system , xt+1 = F ( xt , ut ) , with an external control input ut , we aim to find the Koopman observation functions and the linear dynamics model in the form of g ( xt+1 ) = Kg ( xt ) + Lut , ( 2 ) where the coefficient matrix L is referred to as the control matrix . 3.2 COMPOSITIONAL KOOPMAN OPERATORS . The dynamics of a physical system are governed by physical rules , which are usually shared across different subcomponents in the system . Explicitly modeling such compositionality enables more efficient system identification and control synthesis and provides better generalization ability . Motivating example . Consider a system with N balls moving in a 2D plane , each pair connected by a linear spring . Assume all balls have mass 1 and all springs share the same stiffness coefficient k. We denote the i ’ s ball ’ s position as ( xi , yi ) and its velocity as ( ẋi , ẏi ) . For ball i , equation 3 describes its dynamics , where xi , [ xi , yi , ẋi , ẏi ] T denotes ball i ’ s state : ẋi = ẋiẏiẍi ÿi = ẋi ẏi∑N j=1 k ( xj − xi ) ∑N j=1 k ( yj − yi ) = 0 0 1 00 0 0 1k −Nk 0 0 0 0 k −Nk 0 0 ︸ ︷︷ ︸ , A xiyiẋi ẏi +∑ j 6=i 0 0 0 00 0 0 0k 0 0 0 0 k 0 0 ︸ ︷︷ ︸ , B xjyjẋj ẏj . ( 3 ) We can represent the state of the whole system using the union of every ball ’ s state , where x = [ x1 , · · · , xN ] T . Then the transition matrix is essentially a block matrix , where the matrix parameters are shared among the diagonal or off-diagonal blocks as shown in equation 4 : ẋ = ẋ1 ẋ2 ... ẋN = A B · · · B B A · · · B ... ... . . . ... B B · · · A x1 x2 ... xN . ( 4 ) Based on the linear spring system , we make three observations for multi-object systems . • The system state is composed of the state of each individual object . The dimension of the whole system scales linearly with the number of objects . We formulate the system state by concatenating the state of every object , corresponding to an object-centric state representation . • The transition matrix has a block-wise substructure . After assuming an object-centric state representation , the transition matrix naturally has a block-wise structure as shown in equation 4 . • The same physical interactions share the same transition block . The blocks in the transition matrix encode actual interactions and generalize across systems . A and B govern the dynamics of the linear spring system , and are shared by systems with a different number of objects . These observations inspire us to exploit the structure of multi-object systems , instead of learning separate models for systems that contains different numbers of balls . Compositional Koopman operators . Motivated by the linear spring system , we want to inject a good inductive bias to incorporate compositionality when applying the Koopman theory . This allows better generalization ability and more efficient system identification and better controller design . Figure 1 shows an overview of our model . Considering a system with N objects , we denote xt as the system state at time t and xti is the state of the i ’ th object . We further denote gt , g ( xt ) as the embedding of the state in the Koopman invariant space . In the rest of the paper , we call gt the Koopman embedding . Based on the observation we made in the case of linear spring system , we propose the following assumptions on the compositional structure of the Koopman embedding and the Koopman matrix . • The Koopman embedding of the system is composed of the Koopman embedding of every objects . Similar to the decomposition in the state space , we assume the Koopman embedding can be divided into object-centric sub-embeddings , i.e . gt ∈ RNm denoting the concatenation of gt1 , · · · , gtN , where we use gti = gi ( xt ) ∈ Rm as the Koopman embedding for the i ’ th object . • The Koopman matrix has a block-wise structure . It is natural to think the Koopman matrix is composed of block matrices after assuming an object-centric Koopman embeddings . In equation 5 , Kij ∈ Rm×m and Lij ∈ Rm×l are blocks of the Koopman matrix and the control matrix , where l is the dimension of the action for each object and ut ∈ RNl is the concatenation of ut1 , · · · , utN denoting the total control signal at time t : g t+1 1 ... gt+1N = K11 · · · K1N ... . . . ... KN1 · · · KNN g t 1 ... gtN + L11 · · · L1N ... . . . ... LN1 · · · LNN u t 1 ... utN . ( 5 ) As we have seen in the case of linear spring system , those matrix blocks are not independent , but some of them share the same set of values . • The same physical interactions shall share the same transition block . The equivalence between the blocks should reflect the equivalence of the interactions , where we use the same transition sub-matrix for object pairs of the same relation . For example , if the system is composed of N identical objects interacting with the same relation , then , by symmetry , all the diagonal blocks should be the same , while all the off-diagonal blocks should also be the same . The repetitive structure allows us to efficiently identify the values using least squares regression .

Summary: The paper proposes a novel method for modelling dynamical systems over graphs. The main idea investigated by the authors is to combine Graph Neural Networks together with approximate Koopman embedding. The GNN encodes the input graph to what the authors call "object-centric embedding", whose concatenation over all objects is defacto the approximate Koopman embedding of the system.