Title: LEARNING COMPOSITIONAL KOOPMAN OPERATORS FOR MODEL-BASED CONTROL

Abstract: 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 objectcentric 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.

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 Object-Centric Embeddings Physical System x t+1 < l a t e x i t s h a 1 _ b a s e 6 4 = " W y z T h W I y L 1 C 3 Z 5 9 V V 7 E D P Z / E 6 G g = " > A A A B / X i c b V B L S w M x G M z 6 r P W 1 P m 5 e g k U Q h L J b B T 0 W v X i s Y B / Q r i W b z b a h 2 W R J s m J d F v + K F w + K e P V / e P P f m G 3 3 o K 0 D I c P M 9 5 H J + D G j S j v O t 7 W w u L S 8 s l p a K 6 9 v b G 5 t 2 z u 7 L S U S i U k T C y Z k x 0 e K M M p J U 1 P N S C e W B E U + I 2 1 / d J X 7 7 X s i F R X 8 V o 9 j 4 k V o w G l I M d J G 6 t v 7 P V + w Q I 0 j c 6 U P 2 V 2 q T 9 y s b 1 e c q j M B n C d u Q S q g Q K N v f / U C g Z O I c I 0 Z U q r r O r H 2 U i Q 1 x Y x k 5 V 6 i S I z w C A 1 I 1 1 C O I q K 8 d J I + g 0 d G C W A o p D l c w 4 n 6 e y N F k c o D m s k I 6 a G a 9 X L x P 6 + b 6 P D C S y m P E 0 0 4 n j 4 U J g x q A f M q Y E A l w Z q N D U F Y U p M V 4 i G S C G t T W N m U 4 M 5 + e Z 6 0 a l X 3 t F q 7 O a v U L 4 s 6 S u A A H I J j 4 I J z U A f X o A G a A I N H 8 A x e w Z v 1 Z L 1 Y 7 9 b H d H T B K n b 2 w B 9 Y n z 8 W l p W g < / l a t e x i t > Figure 1 : Overview of our model. A graph neural network φ takes in the current state of the physical system x t , and generates object-centric representations in the Koopman space g t . We then use the block-wise Koopman matrix K and control matrix L identified from equation 6 or equation 8 to predict the Koopman embeddings in the next time step g t+1 . Note that in K and L, object pairs of the same relation share the same sub-matrix. Another graph neural network ψ maps g t+1 back to the original state space, i.e., x t+1 . The mapping between g t and g t+1 is linear and is shared across all time steps, where we can iteratively apply K and L to the Koopman embeddings and roll multiple steps into the future. The formulation enables efficient system identification and control synthesis.

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 & Mezić, 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 (Racaniè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 x t ∈ X ⊂ R n be the state vector for the system at time step t. We consider a non-linear discrete-time dynamical system described by x t+1 = F (x t ). 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

EQUATION

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 {g 1 , • • • , g m } and satisfies that Kg ∈ G for any g ∈ G. With a slightly abuse of the notation, we now use g(x t ) :

R n → R m to represent [g 1 (x t ), • • • , g m (x t )] T .

By constraining the Koopman operator on this invariant subspace, we get a finite-dimensional linear operator K ∈ R m×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, x t+1 = F (x t , u t ), with an external control input u t , we aim to find the Koopman observation functions and the linear dynamics model in the form of

EQUATION

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 (x i , y i ) and its velocity as ( ẋi , ẏi ). For ball i, equation 3 describes its dynamics, where x i [x i , y i , ẋi , ẏi ] T denotes ball i's state:

ẋi =    ẋi ẏi ẍi ÿi    =     ẋi ẏi N j=1 k(x j -x i ) N j=1 k(y j -y i )     =    0 0 1 0 0 0 0 1 k -N k 0 0 0 0 k -N k 0 0    A    x i y i ẋi ẏi   + j =i    0 0 0 0 0 0 0 0 k 0 0 0 0 k 0 0    B    x j y j ẋj ẏj    .

(3) We can represent the state of the whole system using the union of every ball's state, where

x = [x 1 , • • • , x N ] 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:

EQUATION

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 x t as the system state at time t and x t i is the state of the i'th object. We further denote g t g(x t ) as the embedding of the state in the Koopman invariant space. In the rest of the paper, we call g t 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. g t ∈ R N m denoting the concatenation of

g t 1 , • • • , g t N ,

where we use g t i = g i (x t ) ∈ R m 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, K ij ∈ R m×m and L ij ∈ R m×l are blocks of the Koopman matrix and the control matrix, where l is the dimension of the action for each object and

u t ∈ R N l is the concatenation of u t 1 , • • • , u t N

denoting the total control signal at time t:

EQUATION

)

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.

3.3. LEARNING THE KOOPMAN EMBEDDINGS USING GRAPH NEURAL NETWORKS:

For a physical system that contains N objects, we represent the system at time t using a directed graph 

G t = (O t ,

r k = (u k , v k , a r k ), 1 ≤ u k , v k ≤ N

, where u k and v k are integers denoting the end points of this directed edge, and a r k is a one-hot vector denoting the type of the relation k. We use a graph neural network similar to Interaction Networks (IN) (Battaglia et al., 2016) System identification. For a sequence of observations x = [x 1 , • • • , x T ] from time 1 to time T , we first map them to the Koopman space as g = [g 1 , • • • , g T ] using the graph encoder φ, where g t = φ(x t ). We use g i:j to denote the sub-sequence

[g i , • • • , g j ].

To identify the Koopman matrix, we solve the linear regression min K Kg 1:T -1g 2:T 2 . As a result, K = g 2:T (g 1:T -1 ) † will asymptotically approach the Koopman operator K with an increasing T . For cases where there are control inputs u = [u 1 , • • • , u T -1 ], the calculation of the Koopman matrix and the control matrix is essentially solving a least squares problem w.r.t. the objective

EQUATION

)

As we mentioned in the Section 3.2, the dimension of the Koopman space is linear to the number of objects in the system, i.e., g ∈ R N m×T and K ∈ R N m×N m . If we do not enforce any structure on the Koopman matrix K, we will have to identify N 2 m 2 parameters. Instead, we can significantly reduce the number by leveraging the assumption on the structure of K. Assume we know some blocks ({K ij }) of the matrix K are shared and in total there are h different kinds of blocks, which we denote as K ∈ R h×m×m . Then, the number of parameter to be identified reduce to hm 2 . Usually, h does not depend on N , and is much smaller than N 2 . Now, for each block K ij , we have a one-hot vector σ ij ∈ {0, 1} h indicating its type, i.e., K ij = σ ij K ∈ R m×m . Finally, as shown in equation 7, we represent the Koopman matrix as the product of the index tensor σ and the parameter tensor K:

K = σ ⊗ K =    σ 11 K • • • σ 1N K . . . . . . . . . σ N 1 K • • • σ N N K   , where σ =    σ 11 • • • σ 1N . . . . . . . . . σ N 1 • • • σ N N    ∈ R N ×N ×h . (7)

Similar to the Koopman matrix, we assume the same block structure in the control matrix L and denote its parameter as L ∈ R h×m×l . The least squares problem of identifying K and L becomes min

EQUATION

where -1) . Since the linear least squares problems described in equation 6 and equation 8 have analytical solutions, performing system identification using our method is very efficient.

σ ⊗ K ∈ R N m×N m , σ ⊗ L ∈ R N m×N l , g 1:T -1 ∈ R N m×(T -1) and u ∈ R N l×(T

Training GNN models. To make predictions on the states, we use a graph decoder ψ to map the Koopman embeddings back to the original state space. In total, we have three losses to train the graph encoder and decoder. The first term is the auto-encoding loss

EQUATION

)

The second term is the prediction loss. To calculate it, we rollout in the Koopman space and denote the embeddings as ĝ1 = g 1 , and

ĝt+1 = K ĝt + Lu t , for t = 1, • • • , T -1.

The prediction loss is defined as the difference between the decoded states and the actual states, i.e.,

EQUATION

)

Third, we employ a metric loss to encourage the Koopman embeddings preserving the distance in the original state space. The loss is defined as the absolute error between the distances measured in the Koopman space and that in the original space, i.e.,

EQUATION

)

Having Koopman embeddings that perserves the distance in the state space is important as we are using the distance in the Koopman space to define the cost function for downstream control tasks.

The final training loss is simply the combination of all the terms above: L = L ae +λ 1 L pred +λ 2 L metric .

We then minimize the loss L by optimizing the parameters in the graph encoder φ and graph decoder ψ using stochastic gradient descent. Once the model is trained, it can be used for system identification, future prediction, and control synthesis.

3.4. CONTROL:

For a control task, the goal is to synthesize a sequence of control inputs u 1:T that minimize C = T t=1 c t (x t , u t ), the total incurred cost, where c t (x t , u t ) is the instantaneous cost. For example, considering the control task of reaching a desired state x * at time T , we can design the following instantaneous cost, c t (x t , u t ) = 1 [t=T ] x tx * 2 2 + λ u t 2 2 . The first term promotes the control sequence that matches the state to the goal, while the second term regularizes the control signals.

Open-loop control via quadratic programming (QP). Our model maps the original nonlinear dynamics to a linear dynamical system. We can then solve the control task by solving a linear control problem. With the assumption that the Koopman embeddings preserve the distance measure, we define the control cost as c t (g t , u t ) = 1 [t=T ] g tg * 2 2 + λ u t 2 2 . As a result, we reduce the problem to minimizing a quadratic cost function C = T t=1 c t (g t , u t ) over variables {g t , u t } T t=1 under linear constrains g t+1 = Kg t + Lu t , where g 1 = φ(x 1 ) and g * = φ(x * ).

Model predictive control (MPC).:

Solving the QP gives us control signals, which might not be good enough for long-term control as the prediction error accumulates. We can combine it with Model Predictive Control, assuming feedback from the environment every τ steps.

4. EXPERIMENTS:

Environments. We evaluate our method by assessing how well it can simulate and control ropes and soft robots. Specifically, we consider three environments. (1) Rope (Figure 2a ): the top mass of a rope is fixed to a specific height. We apply force to the top mass to move it in a horizontal line. The rest of the masses are free to move according to internal force and gravity. (2) Soft (Figure 2b ): we aim to control a soft robot that is consist of soft blocks. Blocks in dark grey are rigid and those in light blue are soft blocks. Each one of the dark blue blocks is soft but have an actuator inside that can contract or expand the block. One of the blocks is pinned to the ground, as shown using the red dots.

(3) Swim (Figure 2c ): instead of pinning the soft robot to the ground, we let the robot swim in fluids. The colors shown in this environment have the same meaning as in Soft. Observation space. In the Rope environment, each mass on the rope is considered as an object. The observation of each mass is its position and velocity in the 2D plane, which has a dimension of 4. In total, a rope with N masses has an observation space of dimension 4N . In both the Soft and the Swim environments, each quadrilateral is considered as an object. For each quadrilateral, we have access to the positions and velocities of the four corners. Thus for a soft robot containing N quadrilaterals, we have a 4 × 4 × N = 16N dimensional observation.

Baselines. We compare our model to the following baselines: Interaction Networks (Battaglia et al., 2016) (IN), Propagation Networks (Li et al., 2019b) (PN) and Koopman method with handcrafted Koopman base functions (KPM). IN and PN are the state-of-the-art learning-based physical simulators, and we evaluate their adaptation ability by finetuning their parameters on a small sequence of observations from the testing environment. Similar to our method, KPM fits a linear dynamics in the Koopman space. Instead of learning Koopman observations from data, KPM uses polynomials of the original states as the basis functions. In our setting, we set the maximum order of the polynomials to be three to make the dimension of the hand-crafted Koopman embeddings match our model's.

Data generation. We generate 10,000 episodes for Rope and 50,000 episodes for Soft and Swim. Among them, 90% are used for training, and the rest for testing. Each episode has 100 time steps. In the dataset, the physical systems have a various number of objects from 5 to 9, i.e. the ropes have 5 to 9 masses while the soft robots in Soft and Swim environments have 5 to 9 quadrilaterals. To evaluate the model's extrapolating generalization ability, for each environment, we generate an extra dataset with the same size as the test set while containing systems consist of 10 to 14 objects. Training and evaluation protocols. All models are trained using Adam optimizer (Kingma & Ba, 2015) with a learning rate of 10 -4 and a batch size of 8. λ 1 and λ 2 are 1.0 and 0.3, respectively, for our model. For both our model and the baselines, we apply 400K iterations of gradient steps in the Rope environment and 580K iterations in the Soft and Swim environment. Our model is trained on the sub-sequence of length 64 from the training set, and IN/PN aims at minimizing the L1 distance between their prediction and the ground truth. During test time, the models have to adapt to a new environment of unknown physical parameters, where they have access to a short sequence of observations and the opportunity to adjust their models' parameters. Our model uses 8 episodes to identify the transition matrix via least-square regression. IN/PN update the model's parameters by minimizing the distance between the model's prediction and the actual observation using a gradient step of length 10 -4 for 5 iterations. For evaluation, we use two metrics: simulation error and control error. For a given episode, the simulation error at time step t is defined as the mean squared error between the model prediction xt and the ground truth x t . For control, we pick the initial frame x 0 and the t'th frame x t from a episode. Then we ask the model to generate a control sequence of length t to transfer the system from the initial state x 0 to the target state x t . The control error is defined as the mean squared distance between the target state and the state of the system at time t.

K < l a t e x i t s h a 1 _ b a s e 6 4 = " H Q f a d D Q r o z r 2 P p v G 9 2 8 H Q m j c O j I = " > A A A B 6 H i c b V D L S g N B E O y N r x h f U Y 9 e B o P g K e x G Q Y 9 B L 4 K X B M w D k i X M T n q T M b O z y 8 y s E E K + w I s H R b z 6 S d 7 8 G y f J H j S x o K G o 6 q a 7 K 0 g E 1 8 Z 1 v 5 3 c 2 v r G 5 l Z + u 7 C z u 7 d / U D w 8 a u o 4 V Q w b L B a x a g d U o + A S G 4 Y b g e 1 E I Y 0 C g a 1 g d D v z W 0 + o N I / l g x k n 6 E d 0 I H n I G T V W q t / 3 i i W 3 7 M 5 B V o m X k R J k q P W K X 9 1 + z N I I p W G C a t 3 x 3 M T 4 E 6 o M Z w K n h W 6 q M a F s R A f Y s V T S C L U / m R 8 6 J W d W 6 Z M w V r a k I X P 1 9 8 S E R l q P o 8 B 2 R t Q M 9 b I 3 E / / z O q k J r / 0 J l 0 l q U L L F o j A V x M R k 9 j X p c 4 X M i L E l l C l u b y V s S B V l x m Z T s C F 4 y y + v k m a l 7 F 2 U K / X L U v U m i y M P J 3 A K 5 + D B F V T h D m r Q A A Y I z / A K b 8 6 j 8 + K 8 O x + L 1 p y T z R z D H z i f P 6 N T j N M = < / l a t e x i t > L < l a t e x i t s h a 1 _ b a s e 6 4 = " Q E 9 o D X O S e R c a v W m Z Z H N Q H B 2 S 4 v k = " > A A A B 6 H i c b V A 9 S w N B E J 2 L X z F + R S 1 t F o N g F e 6 i o G X Q x s I i A f M B y R H 2 N n P J m r 2 9 Y 3 d P C C G / w M Z C E V t / k p 3 / x k 1 y h S Y + G H i 8 N 8 P M v C A R X B v X / X Z y a + s b m 1 v 5 7 c L O 7 t 7 + Q f H w q K n j V D F s s F j E q h 1 Q j Y J L b B h u B L Y T h T Q K B L a C 0 e 3 M b z 2 h 0 j y W D 2 a c o B / R g e Q h Z 9 R Y q X 7 f K 5 b c s j s H W S V e R k q Q o d Y r f n X 7 M U s j l I Y J q n X H c x P j T 6 g y n A m c F r q p x o S y E R 1 g x 1 J J I 9 T + Z H 7 o l J x Z p U / C W N m S h s z V 3 x M T G m k 9 j g L b G V E z 1 M v e T P z P 6 6 Q m v P Y n X C a p Q c k W i 8 J U E B O T 2 d e k z x U y I 8 a W U K a 4 v Z W w I V W U G Z t N w Y b g L b + 8 S p q V s n d R r t Q v S 9 W b L I 4 8 n M A p n I M H V 1 C F O 6 h B A x g g P M M r v D m P z o v z 7 n w s W n N O N n M M f + B 8 / g C k 1 4 z U < / l a t e x i t >

4.1. SIMULATION:

Figure 2 shows qualitative results on simulation. Our model accurately predicts system dynamics for more than 100 steps. For Rope, the small prediction error comes from the slight delay of the force propagation inside the rope; hence, the tail of the rope usually has a larger error. For Soft, our model captures the interaction between the body parts and generates accurate prediction over the global movements of the robot. The error mainly comes from the misalignment of some local components.

We evaluate the models by predicting 100 steps into the future on 500 trajectories and Figure 3 shows quantitative results. IN and PN do not work well in the Rope and Swim environments due to insufficient system identification ability. The KPM baseline performs poorly in the Rope and Soft environments indicating the limited power of polynomial Koopman base functions. Our model significantly outperforms all the baselines.

We compare our model with KPM, the Koopman baseline using polynomial basis. In Rope, we ask the models to perform open-loop control where it only solves the QP once at the beginning. The length of the control sequence is 40. When it comes to Soft/Swim, each model is asked to generate control signals of 64 steps, and we allow the model to receive feedback after 32 steps. Thus every model has a second chance to correct its control sequence by solving the QP again at the time step 32.

As shown in Figure 2 , our model leverages the inertia of the rope and matches the target state accurately. As for controlling a soft body swinging on the ground or swimming in the water, our model can move each part (the boxes) of the body to the exact target position. The small control error comes from the slight misalignment of the orientation and the size of the body parts. Figure 4 shows that quantitatively our model outperforms KPM, too.

4.3. ABLATION STUDY:

Structure of the Koopman matrix. We explore three different structures of the Koopman matrix, Block, Diag and None, to understand its effect on the learned dynamics. None assumes no structure in the Koopman matrix. Diag assumes a diagonal block structure of K: all off-diagonal blocks (K ij where i = j) are zeros and all diagonal blocks share the same values. Block predefines a block-wise structure, decided by the relation between the objects as introduced in Section 3.3.

Table 1 : Ablation study results on the Koopman matrix structure (Rope environment).

For simulation, we show the Mean Squared Error between the prediction and the ground truth at T = 100, whereas for control, we show the performance with a horizon of length 40. The numbers in parentheses show the performance on extrapolation. Table 1 includes our model's simulation error and control error with different Koopman matrix structures in Rope. All models are trained in the Rope environment with 5 to 9 masses. Besides the result on the test set, we also report models' extrapolation performance in parentheses, where the model is evaluated on systems with more masses than training, i.e., 10 to 14 masses.

Our model with Block structure consistently achieves a smaller error in all settings. Diag assumes an overly simplified structure, leading to larger errors and failing to make reasonable controls. None has comparable simulation errors but larger control errors. Without the structure in the Koopman matrix, it overfits the data and makes the resulting linear dynamics less amiable to the control.

Hyperparameters. In our main experiments, we set the dimension of the Koopman embedding to m = 32 per object. Online system identification requires 800 data samples for each training/test case. To understand our model's performance under different hyperparameters, we vary the dimension of the Koopman embedding from 8 to 64 and the number of data samples used for system identification from 200 to 1,600. Figure 4d shows that more data for system identification leads to better simulation results. Figure 4e shows that dimension 16 gives the best results on simulation. It may suggest that the intrinsic dimension of the Koopman invariant space of the Rope system is around 16 per object.

5. CONCLUSION:

Compositionality is common in our daily life. Many ordinary objects contain repetitive subcomponents: ropes and soft robots, as shown in this paper, granular materials such as coffee beans and lego blocks, and deformable objects such as cloth and modeling clay. These objects are known to be very challenging for manipulation using traditional methods, while our formulation opens up a new direction by combining deep Koopman operators with graph neural networks. By leveraging the compositional structure in the Koopman operator via graph neural nets, our model can efficiently manipulate deformable objects such as ropes and soft robots, and generalize to systems with variable numbers of components. We hope this work could encourage more endeavors in modeling larger and more complex systems by integrating the power of the Koopman theory and the expressiveness of neural networks.

A ENVIRONMENT AND MODEL DETAILS:

Interaction types. In our experiments, interactions are considered different if the types are different or the objects involved have different physical properties.

In the Rope environment, the top mass has a fixed height and is considered differently from the other masses. Thus, we have 2 types of self-interactions for the top mass and the non-top masses. In addition, we have 8 types of interactions between different objects. The objects on a relation could be either top mass or non-top mass. It is a combination of 4. And the interaction may happen between two nearby masses or masses that are two-hop away. In total, the number of interactions between different objects is 4 × 2 = 8.

In the Soft environments, there are four types of quadrilaterals: rigid, soft, actuated, and fixed. We have four types of self-interactions correspondingly. For the interactions between objects, we add edges between two quadrilaterals only if they are connected by a point or edge. Connection from different directions are considered as different relations. There are 8 different directions, up, down, left, right, up-left, down-left, up-right, down-right. The relation types also encode the type of receiver object. Thus, in total, there are (8 + 1) × 4 = 36 types of relations between different objects.

In the Swim environment, there are three types of quadrilaterals: rigid, soft, and actuated. Similar to the Soft environment, we use different edge types for different connecting directions; hence, the number of edge type is (8 + 1) × 3 = 27.

B ADDITIONAL EXPERIMENTS:

Comparison with a classical physical simulator optimized using back-box optimization. We have performed comparisons with a classical physical simulator optimized using black-box optimization (Delingette, 1998) by assuming different levels of knowledge over the ground truth model.

If we assume that we know the ground truth model, where we only need to identify relevant physical parameters during the system identification stage, Bayesian Optimization (Snoek et al., 2012) (BO) can give us a reasonable estimate of the physical parameters. However, BO requires much more time to achieve a comparable performance with our method in the Rope environment: 0.43 vs. 180 seconds averaged over 100 trails (Ours vs. BO).

If we are unsure about the ground truth model and we approximate the system using a set of points linked by springs and dampers, BO does not work as well. In our additional experiments, we approximate the Rope environment using a chained spring-mass system, say n masses and n -1 springs. While taking much more time, BO still cannot give us a satisfying result: simulation error 0.046 vs. 0.084 and control error 0.854 vs. 2.547 (Ours vs. BO).

Ablation study on the effectiveness of the metric loss. The internal linear structure allows us to solve the control problem using quadratic programming, where the objective function for control is defined in the embedding space (Section 3.4); hence, it is desirable to have Koopman embeddings that preserve the distance in the original state space. In Section 3.3, we introduce a metric loss to promote learning a Koopman embedding that keeps the distance measurement.

To demonstrate the effect of the metric loss, we compare the models trained with and without the metric loss. We establish the comparison using two measurements, the distance preservation and the prediction accuracy. To evaluate how well the Koopman embeddings preserve the distance, we compute the distribution of the log-ratio of the distance in the Koopman space and in the original state space, i.e., log g i -g j 2 x i -x j 2 . For the model prediction accuracy, we show the simulation errors.

We perform the experiments in the Rope environment and show the result in Figure 5 . On the left, we show the ratio of distance in the learned Koopman space and the distance in the original state space. The model trained with metric loss has a log distance ratio that significantly more concentrates on 0. It means the metric loss effectively regularizes the model to preserve the distance. On the right, we show the simulation errors of the two models, which indicate that two models have comparable prediction performance. Metric loss effectively enhances the property of distance-preserving while not making a big sacrifice on the accuracy of the dynamics modeling. x i -x j 2 . The model trained with metric loss has a distance ratio much more concentrated to 1, which indicates it preserves the distance much better than the counterpart. (b) illustrates the simulation error of two models, where their performance is on par. (c) shows that the model trained using the metric loss performs better control. 

Mean Squared Error:

Figure 6 : Modeling rope with known physical parameters. We show the comparison between our model and IN/PN in scenarios where we have access to the ground truth physical parameters. In this case, IN and PN slightly outperform our method due to the internal linear structure in our model. However, in the real world, we do not always know the physical parameters and their values, which makes our method preferable when adapting to new environments.

Experiments in a known physical parameter setting. Our setting is different from the settings in the original IN and PN papers that we do not assume we know the physical parameters and their values, such as stiffness, mass, and gravity. Instead, the parameters are embedded in the transition matrices during the system identification stage (Section 3.3). If the model has access to the underlying physical parameters, as expected, IN and PN slightly outperform our method as the internal linear structure limits our model's expressiveness, as shown in Figure 6 . In the real world, however, the underlying physical parameters are not always known, which makes our model a better choice when adapting to unseen environments.

:



Figures and Tables:

Figure 2: Qualitative results. Top: our model prediction matches the ground truth over a long period. Bottom: for control, we use red dots or frames to indicate the goal. We apply the control signals generated from our identified model to the original simulator, which allows the agent to achieve the goal accurately. Please refer to our supplementary video for more results.

Figure 3: Quantitative results on simulation. The x axis shows time steps. The solid lines indicate medians and the transparent regions are the interquartile ranges of simulation errors. Our method significantly outperforms the baselines in all testing environments.

Figure 4: Quantitative results on control and ablation studies on model hyperparameters. Left: box-plots the distributions of control errors. The yellow line in the box indicates the median. Our model consistently achieves smaller errors in all environments against KPM. Right: our model's simulation errors with different amount of data for system identification (d) and different dimensions of the Koopman space (e).

Figure 5: Ablation study on the loss in the Rope environment. (a) shows the distributions of the logarithm distance ratio, i.e, log g i -g j 2

R), where vertices O t = {o t i } N i=1 represent objects and edges R = {r k } N 2

Table content: <table><tr><td>k=1 i is the state of object i and i ) , where x t i , a o i = (x t represent pair-wise relations. Specifically, o t a o i is a one-hot vector indicating the object type, e.g., fixed or movable. For relation, we have</td></tr></table>

to generate object-centric Koopman embeddings. IN defines an object function f O and a relation function f R to model objects and their relations in a compositional way. Similar to a message passing procedure, we calculate the edge effect e t k = f R (o t u k , o t v k , a r k ) k=1...N 2 , and node effect g t i = f O (o t i , k∈Ni e t k ) i=1...N , where N i denotes the relations that point to the object i and {g t i } are the derived Koopman embeddings. We use this graph neural network, denoted as φ, to represent our Koopman observation function.

Table content: <table/>

Peer review (iclr2020):
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. One of the key contributions is the reduction in parameters, by assuming that the interactions between different objects in the Koopman space are limited to some fixed number of types, or in other words given the object-centric embedding the Koopman matrix is a block matrix, where each block can only be one of K matrices. In this way the number of parameters is fixed and does not scale with the number of objects, compared to the naive way where it will scale as N^2. In addition to the dynamical modelling the paper adds an extra linear-*control* input in the Koopman embedding space which to affect the dynamics of the system and allow for modelling systems where there is external control being applied. The models are than compared on three small scale tasks, showing better results in mean squared error prediction compared to the three baseline approaches. Additionally, when used for controls on the environments the methods outperforms the one baseline method it is compared to. I*m quite borderline on whether the paper should be accepted or rejected, but currently I*m leaning towards a rejection. The main reason for this decision is that in my opinion the experiments presented are somewhat limited with respect to the baselines used and I have some reservations regarding the results presented for IN and PN discussed below. Detailed comments on paper: 1. I personally like the main idea of the paper, which is to use previous results from approximating the Koopman operator and combining it with GNNs for more accurate physical modelling of object-object interactions. Additionally, the idea of reducing the parameters is quite important. 2. Linear control theory - although it is quite natural to add the control as a linear affect in the latent space, and this has been done numerous times before in the literature, I don*t recall there to be any theory on Koopman embedding when there is a control signal. Additionally, if the policy that has been used in practice is stochastic, the resulting *induced* dynamical system then also becomes stochastic, and to my knowledge, at least in theory, learning Koopman embedding for such systems has more challenges are requires certain assumptions about the true system, such as co-diagonalization and few others. I think this should be discussed in more detailed by the authors (and please do correct me if I*m wrong on any of these statements) as currently for the readers who are not too familiar I think the text might come across as though that the Koopman theory extends naively to these scenarios as well, which I do not think is the case. 3. It is not very clear how does the *metric* loss affects the solution. I would encourage the authors to provide comparison (only in terms of dynamical modelling, without active control) of whether this metric helps, or have some negative effects on the prediction. I think that for instance if the GNN has some form of weight regularization than this indeed would have some non-trivial effect on the resulting representation. Also, it would useful to have plots of how accurately does the embedding preserve the distance to the true states in order to understand this better. Comments on the experiments: 1. In the paper there is no discussion about what are the actual observation space of the environments, could these be clarified better. 2. The block diagonal structure approach in general has been presented as working with multiple types of interactions. However, in practice it seems that the authors have only used two types of interaction -> object-same-object and object-other-object interaction. This however, has never been discussed and is maybe false. Could you clarify these details? 3. The results showed in Figure 3 are somewhat in contrast than the results in the original PN paper, specifically the PN paper states that it can achieve MSE of 7.85 for 1000 time steps, and from figure 6 of that paper it shows about 0.05 MSE over 100 steps on a similar rope environment. These results compared to the one presented here in Figure 3 makes me wonder how well did the authors actually managed to reimperilment the IN and PN paper? Could there be any comments on this as this makes many claims of the proposed method being better questionable and hard to understand its significance in relation to previous work. 4. For the control tasks, it would have been useful to have more than just the single baseline used. There are plenty of algorithms for Reinforcement Learning that could have been used in order to put the method in perspective. E.g. one can apply MPC with ground-truth model (e.g. the simulator) to show the discrepancy with an ideal case. In the RL literature there are plenty of methods for solving smaller problems, parametric and non-parametric: Q-learning, PPO etc... I think this is very important from the reader perspective. PS: Please refer to the discussion below with the authors as I have increased my score from 3 to 6.