Learning Compositional Koopman Operators for Model-Based Control
Authors: Yunzhu Li
Hao He
Jiajun Wu
Dina Katabi
Antonio Torralba
Venue:  (2020)
Keywords: Koopman operators
graph neural networks
compositionality

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 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.

Published as a conference paper at ICLR 2020
LEARNING COMPOSITIONAL KOOPMAN OPERATORS FOR MODEL-BASED CONTROL
Yunzhu Li MIT CSAIL
Hao He MIT CSAIL
Jiajun Wu MIT CSAIL
Dina Katabi MIT CSAIL
Antonio Torralba MIT CSAIL
ABSTRACT
Finding an embedding space for a linear approximation of a nonlinear dynamical system enables efficient system identification and control synthesis. The Koop- man 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 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 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 new environments of unknown physical parameters and produce control signals to achieve 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.
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 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 often not efficient enough for real-time control (Mayne, 2000). The Koopman operator theory suggests that identifying nonlinear-to-linee coordinate transformations allows efficient linear approximation of nonlinear systems (Williams et al., 2015; Mauroy & Goncalves, 2016). Fast as they are, however, existing papers Koopman operators focus single dynamical system, making it hard generalize to cases where there are 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 cime-consuming and sample-inefficiae
In this paper, 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 over-parane wer--arameterizeed neural nets to model state transition, we identify the Koopman matrix and control matrix from data as linear approximation of the nonlinear dynamical system. The linear approximation allows efficient system identification and control synthesis.
The main challenge of extending Koopman theory 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:/ http:///ooooma..// il eduu

---

Published as a conference paper at ICLR 2020
Physical System x²
Graph-NN
Graph-NN d
Physical System an"
Object-Centric Embeddings g²
Koopmun Matrix K Control Matrix L
Object-Centric Embeddings g gx11
Figure I: Overview of our model. A graph neural network & takes in the current state of the physical system æ', and generates object-centric representations in the Koopman space g'. 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 the next time step g g Note that K and L, object pairs of the same relation share the same sub-matrix. Another graph neural network vi maps g back to the original state space, i.., x++1, The mapping between g and g 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 be identified by making 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; et al., 2019b), and faster when adapting to new environments of unknown physical parameters. Our method also outperforms vanilla deep Koopman methods (Lusch al., 2018; Morton et al., 2018) and Koopman models with camually-designen basis functions, which shows the advantages of using structured Koopman matrix and graph neural networks. Please see our project page for demonstrating videos.
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 map the state of nonlinear dynamical system 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 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 8; 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), 8). 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 distribution over observations. Our model differs by explicitly modeling the compositionality of the underlying system with graph networks. It generalizes better to environments of 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
2

---

Published as a conference paper at ICLR 2020
samples and compute, which limits its in an online setting. Nagabandi al. (2019a;b) proposed to use meta-learning for online adaptation, and have shown to be effective in simulated robots and a real legged milliroboot. However, it is not clear whether their methods 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 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 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., 201 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 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 E X C R" be the state vector for the system at time step t. We consider a non-linear discrete-time dynamical system described by x²+1 F(w'). The Koopman operator (Koopman, 1931), denoted as K - F, is a linear transformation defined by Kgg - g o 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 RR belonging to F, we have
(K")) g F ((F)(()))))))))))))(()))))))(F g(x²+1),
(1)
making the function space 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 set of base observation functions {gi. I9m } and satisfies that Kg E g for any g E G. With a slightly abuse of the notation, we now use g(x²). IR - Rmm to represent
[g) (x=',, )..,.. 9m ( By constraining the Koopman operator on this invariant subspace, we get dinte-dimeniional linear operator K E mm 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 dynamical system, 2+11 F(w*, u²), with external control input u', we aim to find the Koopman observation functions and the linear dynamics model in the form of
g(x*+1)= Kg(w') + Lui,
where the coefficient matrix L is referred to as the control matrix.
3.2
COMPOSITIONAL KOOPMAN OPERATORS
The dynamics of 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.
3

---

Published as a conference paper at ICLR 2020
Motivating example. Consider system with N balls moving in 2D plane, each pair connected by a linear spring. Assume all balls have mass and all springs share the same stiffiness coefficient We denote the i's ball'ss position as (2), 36) and its velocity as (i,, Mr). For ball i, equation 3 describes its dynamics, where æ;
& [I yis T. út]" denotes ball i's state:
2
i
0
0
1
07
0
0
0
0
ab;
24
y/
Σ
k(ar
mz)
k
Nk:
0
0
0
0
0
1
yi
K
0
0
0
0
I
yj
k(ys v.)
0
Nk
0
0
L
O
0
A
AB
(3)
We can represent the state of the whole system using the union of every ball's state, where [21. _2N'". Then the transition matrix is essentially block matrix, where the matrix parameters are shared among the diagonal or off-diagonal blocks as shown equation 4:
A
B
B
œ1
22
B
22
=
(4)
N
B B
A
[EN
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 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 exploit the structure multi-objes 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 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 shows overview of our model.
Considering a system with N objects, we denote x' as the system state at time f and x is the state of the i'th object. We further denote g" & g(x") as the embedding of the state in the Koopman invariant space. In the rest of the paper, call g the Koopman embedding. Based on the observation 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.c. g" Nm denoting the concatenation of gf. . where we use g = 9: (x=) E I. as the Koopman embedding for the i'th object.
The Koopman matrix has block-wise structure. It natural to think the Koopman matrix is composed of block matrices after assuming an object-centric Koopman embeddings. In equation 5, Ki E p and Lij E [gmxl are blocks of the Koopman matrix and the control matrix, where the dimension of the action for each object and ut E N the concatenation of uf, u denoting the total control signal at time L:
H+1
Ki
KiN
gi
Lu
LIN
+
(5)
.
KNi
KNN
g'x
L
LNN
u'y
4

---

Published as a conference paper at ICLR 2020
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.
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 physical system that contains N objects, we represent the system at time / using directed graph G (O', R), where vertices O = Ti represent objects and edges R {rx) k=1 represent pair-wise relations. Specifically, of = (æ) a.) where of is the state of object i and aj is a one-hot vector indicating the object type, e.g., fixed or movable. For relation, have Th = (ug, Uk. a%) S Uk, Uk I N, where ug and ty2 are integers denoting the end points of this directed edge, and a 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) to generate object-centric Koopman embeddings. IN defines an object funetion fo and relation function fa to model objects and their relations in compositional way. Similar to message passing procedure, we calculate the edge effect e SR(O' d'yy (a)a:a)))= af)R=I... and node effect g = fo EREN; ef.)==1..N. where Mú denotes the relations that point to the object and (gi) are the derived Koopman embeddings. We use this graph neural network, denoted as Φ) to represent our Koopman observation function.
System identification. For sequence of observations ã [x " from time 1 to time T, we first map them to the Koopman space as ğ = [g 9 using the graph encoder ø, where g¹ = o(x'). We use ge to denote the sub-sequence g ... gg). To identify the Koopman matrix, we solve the linear regression ming ||Kg 1:7" |
9 As result, K g (g (g !!)) will asymptotically approach the Koopman operator with an increasing T. For cases where there are control inputs ū = [u², ... u² -," the calculation of the Koopman matrix and the control matrix is essentially solving a least squares problem w.r.t the objective
min [gg I-T + Lu g" 2: 12. K.T.
(6)
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., 9 R NMxT and K E R Nm. If we do not enforce any structure the Koopman matrix K, we will have to identify N2m² parameters. Instead, we can significantly reduce the number by leveraging the assumption on the structure of K. Assume blocks ({ (Ki)) }) of the matrix K are shared and in total there are h different kinds of blocks, which we denote as K E hx Then, the number of parameter to be identified reduce to fum". Usually, does not depend on N, and is much smaller than N², Now, for each block Kij. we have a one-hot vector oi E {0, 11th indicating its type, i.c., K,j oi K E Rm Finally, as shown in equation 7, represent the Koopman matrix as the product of the index tensor and the parameter tensor K:
know some
onk
KINK
Tn
JIN
K K=ook=
where σ
E NxNxh
KNI
NNN
ON1
ONN
Similar to the Koopman matrix, we assume the same block structure in the control matrix and denote its parameter as L. E hxxmi The least squares problem of identifying f and i becomes
min (1 @ igg (o @ i)u g 2:T 122
RLL
where a @ K E Nm o 2 L E NmXNN g 1:T E E NT and ã IT- 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.
5

---

Published as a conference paper at ICLR 2020
Training GNN models. To make predictions the states, we 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 auto-encoding loss
T Σ (v(((())) - æ'1|-
1 T
Le
The second term is the prediction loss. To calculate it, rollout in the Koopman space and denote the embeddings as ğ' g and g"+1 = Kğ¹ + Lu', for t = 1, , -T- The prediction loss is defined as the difference between the decoded states and the actual states, i.e.,
T
pred
llv('') - x
(10)
Third, we employ 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.c..
metric Σ Ellls' - -g'll ||o x
(11)
Having Koopman embeddings that perserves the distance in the state space is important using the distance 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 Lue + t1 fpred X2 meetric We then minimize the loss C 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 I:T that minimize Et= 3 (x', u²), the total incurred cost, where cz u²)
is the instantaneous cost. For example, considering the control task of reaching a desired state x time T, we can design the following instantaneous cost, (x,", u²) = [t= I| _ x 112 + All Al|22111 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 linear dynamical system. We 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 (g u²) | llg g011
112 As result, reduce the problem to minimizing quadratic cost function C = Σt= (g u²) over variables {g', u' IT under linear constrains g² Kg + Luu, where g¹ = e (xx) 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 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 rope is fixed à specific height. We apply force to the top mass 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 control 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 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.

---

Published as a conference paper at ICLR 2020
Prediction
GT
Prediction
GT
Prediction
GT
(a) Rope
Simulation
(c) Swim (b) Soft
Time
Time
Time
(a) Rope
Control
(b) Soft
(c) Swim
Figure Qualitative results. Top: our model prediction matches the ground truth over 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.
Observation space. In the Rope environment, each mass the rope is considered as object. The observation of each mass its position and velocity in the 2D plane, which has a dimension of In total, rope with N masses has 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 soft robot containing N quadrilaterals, we have a 4 x x 16N dimensional observation.
Baselines. We compare model to the following baselines: Interaction Networks (Battaglia et al., 2016) (IN), Propagation Networks (Li et al., 2019b) (PN) and Koopman method with hand- crafted Koopman base functions (KPM). IN and PN the state-of-the-ard 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 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. 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 to quadrilaterals. To evaluate the model's '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 14 objects.
7

---

Published as a conference paper at ICLR 2020
Mean Squared Error
0.3
0.15
IN
0.3
0.2
KPM
0.10
KPM
0.2
KFM
Oum
Ours
Dum
0.1
0.0G
0.1
0.0
0..00
0.0
(a) Rope
(b) Soft
(c) Swim
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.
Control
Ablation Studies
Mcan Squared Error
16
0.15
0.15
-3
15-
0.10
0.10-
10
0.05
0.05
6-
KPM
Ours
KPM
Ours
KPM
Ours
0.00
20
40
80
0.00
20
40
60
80
(d) Number of samples identifying and
(c) The size of object- centric emboddings
(a) Rope
(b) Soft
(c) Swim
Figure 4: Quantitative results control and ablation studies model hyperparameters. Left: box- pox-plots show 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).
Training and evaluation protocols. All models are trained using Adam optimizer (Kingma & Ba, 2015) with learning rate of 10-4 and batch size of 8. Å and X2 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 the Soft and Swim environment. Our model is trained on the sub-seque sub-sequence of length 64 from the training set, and IN/PN aims at minimizing the the ground truth. During test time, the models have to
L1 distance between their prediction adapt to 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 gradient step of length 10-4 for iterations. For evaluation, error and control error. For given episode, the simulation error at time step is defined as the mean squared error between the model prediction á' and the ground truth æ!. For control, pick the initial frame x and the l'th frame x' from a episode. Then we ask the model to generate a control sequence of length to transfer the system from the initial state x to the target state x'. The control error is defined as the mean squared distance between the target state and the state of the system at time
metrics: simulation
4.1 SIMULATION
Figure shows qualitative results on simulation. Our model accurately predicts system dynamies 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 movements of the robot. The error mainly comes from the misalignment of some local components.
the global
We evaluate the models by predicting 100 steps into the future on 500 trajectories and Figure shows quantitative results. IN and PN do not work well in the Rope and Swim environments due insufficient system identification ability. The KPM baseline performs poorly in the Rope and
8

---

Published as a conference paper at ICLR 2020
Soft environments indicating the limited power of polynomial Koopman base functions. Our model significantly outperforms all the baselines.
4.2
CONTROL
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 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 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, 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 ( (Kij where ش 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 includes our model's simulation error and con- trol 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.
Table 1: Ablation study results on the Koop- man 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, show the performance with a horizon of length 40. The numbers in parentheses show the performance on extrapolation.
Our model with Block structure consistently achieves a smaller error 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.
Simulation
Control
Diag None Block
0.133 1744
2.337 (2.809)
0.117 (0.083)
1.522 (1.288)
0.105 (0.075)
0.854 (1.101)
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 '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.
CONCLUSION
Compositionality is common in our daily life. Many ordinary objects contain repetitive subcompo- nents: 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 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.

---

Published as a conference paper at ICLR 2020
REFERENCES
lan Abraham, Gerardo De La Torre, and Todd D Murphey. Model-based control using koopman operators. In RSS, 2017.
Hassan Arbabi and Igor Mezic. Ergodic theory, dynamic mode decomposition, and computation of spectral properties of the koopman operator. SIAM Journal Applied Dynamical Systems, 16(4): 2096-2126, 2017.
Hassan Arbabi, Milan Korda, and Igor Mezic. A data-driven koopman model predictive control framework for nonlinear flows. In CDC, 2018.
Peter W. Battaglia, Razvan Pascanu, Matthew Lai, Danilo Rezende, and Koray Kavukcuoglu. Interac- tion networks for learning about objects, relations and physics. In NeurIPS, 2016.
Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanchez-Gonzalez, Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, et al. Relational inductive biases, deep learning, and graph networks. arXh::1806.0// 00260660016,,...0 2018.
Daniel Bruder, Brent Gillespie, C David Remy, and Ram Vasudevan. Modeling and control soft robots using the koopman operator and model predictive control. In RSS, 2019a.
Daniel Bruder, C David Remy, and Ram Vasudevan. Nonlinear system identification of soft robot dynamics using koopman operator theory. In ICRA, 2019b.
Steven L Brunton, Bingni W Brunton, Joshua L Proctor, and Nathan Kutz. Koopman invariant subspaces and finite linear representations of nonlinear dynamical systems for control. PloS one, 11(2)x0150171, 2016.
Michael B Chang, Tomer Ullman, Antonio Torralba, and Joshua B Tenenbaum. A compositional object-based approach learning physical dynamics. In ICLR, 2017.
Herve Delingette. Toward realistic soft-tissue modeling in medical simulation. Proceedings of the IEEE, 86(3):112-523 1998.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In ICML, 2019.
Jessica Hamrick, Andrew J Ballard, Razvan Pascanu, Oriol Vinyals, Nicolas Heess, and Peter W Battaglia. Metacontrol for adaptive maginationn-aseed optimization. In ICLR, 2017.
Michael Janner, Sergey Levine, William Freeman, Joshua Tenenbaum, Chelsea Finn, and Jiajun Wu. Reasoning about physical interactions with object-oriented prediction and planning. In ICLR, 2019.
Eurika Kaiser, J Nathan Kutz, and Steven L Brunton. Data-driven discovery of koopman eigenfunc- tions for control. arXiv: 146, 2017.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR, 2015.
Bernard O Koopman. Hamiltonian systems and transformation in hilbert space. PNAS, 17(5):315, 1931.
BO Koopman and J Neumann. Dynamical systems of continuous spectra. PNAS, 18(3):255, 1932.
Milan Korda and Igor Mezić. Linear predictors for nonlinear dynamical systems: Koopman operator meets model predictive control. Automatica, 93:149-160. 2018.
Yunzhu Li, Jiajun Wu, Russ Tedrake, Joshua B Tenenbaum, and Antonio Torralba. Learning particle dynamics for manipulating rigid bodies, deformable objects, and fluids. In ICLR, 2019a.
Yunzhu Li, Jiajun Wu, Jun Yan Zhu, Joshua B Tenenbaum, Antonio Torralba, and Russ Tedrake. Propagation networks for model-basee control under partial observation. In ICRA, 2019b.
10

---

Published as a conference paper at ICLR 2020
Bethany Lusch, Nathan Kutz, and Steven L Brunton. Deep learning for universal linear embeddings of nonlinear dynamics. Nature Communications, 91):4955, 2018.
Giorgos Mamakoukas, Maria Castano, Xiaobo Tan, and Todd Murphey. Local koopman operators for data-driven control of robotic systems. RSS, 2019.
Alexandre Mauroy and Jorge Goncalves. Linear identification of nonlinear systems: A lifting technique based on the koopman operator. CDC, 2016.
Alexandre Mauroy and Jorge Goncalves. Koopman-based lifting techniques for nonlinear systems identification. IEEE Transactions on Automatic Control, 2019.
David Mayne, Nonlinear model predictive control: Challenges and opportunities. In Nonlinear Model Predictive Control, pp. 23 44. Springer, 2000.
Jeremy Morton, Freddie D Witherden, Antony Jameson, and Mykel Kochenderfer. Deep dynamical modeling and control of unsteady fluid flows. In NeurIPS, 2018.
Jeremy Morton, Freddie D Witherden, and Mykel Kochenderfer. Deep variational koopman models: Inferring koopman observations for uncertainty-awee dynamics modeling and control. In IJCA,, 2019.
Damian Mrowca, Chengxu Zhuang, Elias Wang, Nick Haber, Li Fei-Fei, Joshua B Tenenbaum, and Daniel LK Yamins. Flexible neural representation for physics prediction. In NeurIPS, 2018.
Anusha Nagabandi, Ignasi Clavera, Simin Liu, Ronald S Fearing, Picter Abbeel, Sergey Levine, and Chelsea Finn. Learning adapt dynamic, real-world environments through meta-reinforcement learning. ICLR, 2019a.
Anusha Nagabandi, Chelsea Finn, and Sergey Levine. Deep online learning via meta-learning: Continual adaptation for model-based
ICLR, 2019b.
Razvan Pascanu, Yujia Li, Oriol Vinyals, Nicolas Heess, Lars Buesing, Sebastien Racanière, David Reichert, Théophane Weber, Daan Wierstra, and Peter Battaglia. Learning model-based planning from scratch. 1rXivv1707.06111 2017.
Joshua L Proctor. Steven L Brunton, and Nathan Kutz. Generalizing koopman theory allow for inputs and control. SIAM Journal on Applied Dynamical Systems, 7(11)9909999) 2018.
Sébastien Racanière, Théophane Weber, David Reichert, Lars Buesing, Arthur Guez, Danilo Jimenez Rezende, Adrià Puigdomènech Badia, Oriol Vinyals, Nicolas Heess, Yujia Li, Razvan Pascanu, Peter Battaglia, David Silver, and Daan Wierstra. forcement learning. In NeurlPS, 2017.
agents for deep rein-
Clarence W Rowley, Igor Mezić, Shervin Bagheri, Philipp Schlatter, and Dan S Henningson. Spectral analysis of nonlinear flows. Journal of Fluid Mechanics, 641:115-127, 2009.
Alvaro Sanchez-Gonzaler Nicolas Heess, Jost Tobias Springenberg, Josh Merel, Martin Riedmiller, Raia Hadsell, and Peter Battaglia. Graph networks as learnable physics engines for inference and control. In ICML, 2018.
Peter Schmid. Dynamic mode decomposition of numerical and experimental data. Journal of Fluid Mechanics, 656:5-28, 2010.
Jasper Snoek, Hugo Larochelle, and Ryan P Adams. Practical bayesian optimization of machine learning algorithms. In NeurIPS, 2012.
Naoya Takeishi, Yoshinobu Kawahara, and Takehisa Yairi. Learning koopman invariant subspaces for dynamic mode decomposition. In NeurlPS, 2017.
Jonathan H Tu, Clarence W Rowley, Dirk M Luchtenburg, Steven Brunton, and J Nathan Kutz. On dynamic mode decomposition: Theory and applications. Journal of Computational Dynamics, (2):391-421, 2014.
11

---

Published as a conference paper at ICLR 2020
Matthew O Williams, loannis G Kevrekidis, and Clarence W Rowley. A data driven approximation of the koopman operator: Extending dynamic mode decomposition. Journal of Nonlinear Science, 25(6):1307-1346, 2015.
Matthew O Williams, Maziar S Hemati, Scott TM Dawson, Ioannis G Kevrekidis, and Clarence W Rowley, Extending data-driven koopman analysis to actuated systems. OFAA-eepersOnLLiee, 49 (18):704-709, 2016.
12

---

Published as a conference paper at ICLR 2020
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 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 either top mass or non-top mass. is a combination of 4. And the interaction may happen between two nearby masses masses that are two-hop away. In total, the number of interactions between different objects is 4 x 2 = 8.
relation could he
In the Soft environments, there are four types of quadrilaterals: rigid, soft, actuated, and fixed. We have four types of self inferretttttsns correspondingly. For the interactions between objects, we add edges between two quadrilaterals only if they are connected by a point edge. Connection from different directions are considered as different relations. There are 8 different directions, up, down, left, right, up-left, down-left, up 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) x 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 optimiza- tion (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 achieve 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 set of points linked by springs and dampers, BO does not work as well. In our additional experiments, we approximate the Rope environment using chained spring-mat mass system, say masses and n 1 springs. While taking much more time, BO still cannot give us a satisfying result: simulation error 0.046 0.084 and control error 0.854 2.547 (Ours 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, is desirable to have Koopman embeddings that preserve the distance in the original state space. In Section 3.3, we introduce 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, compute the distribution of the log-ratio the distance in the Koopman space and in the original state space, i.e., log
"g'-g" For the model prediction accuracy, we show the simulation errors.
|| ||='-=0'11 l
We perform the experiments in the Rope environment and show the result in Figure 5. On the left, we show the ratio of distance the learned Koopman space and the distance the original state space. The model trained with metric loss has 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- preservvn while not making big sacrifice on the accuracy of the dynamics modeling.
13

---

Published as a conference paper at ICLR 2020
w/ metric loss w/o metric loss
Mcan Squared Erron
w/ metric
Mean Squared Error
Probability Density
8-
w/o metric
6-
0.2
0.1
2
0.0
-1.5
-10 -0.5
0.0
0.5
1.0
1.5
20
40
60
80
100
(a) Distance Preservation
(b) Simulation
(c) Control
Figure 5: Ablation study on the metric loss in the Rope environment. (a) shows the distributions
of the logarithm distance ratio, i.e, log Lg'-gg <<='~>/1>
ll2 The model trained with metric loss has 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 on par. (c)
shows that the model trained using the metric loss performs better control.
0.3
Mean Squared Error
IN
0.2
0.0
20
40
60 80
100
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, and PN slightly outperform our method due to the internal linear structure 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 stiffiness, 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 model's 's expressiveness, 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 unseen environments.
14


Reviews:
[{"review_id": "r1eClp0acr", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "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.\nOne 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. \n\n\nI'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.\n\n\nDetailed comments on paper: \n\n1. 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. \n\n2. 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.\n\n3. 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.\n\n\n\nComments on the experiments:\n\n1. In the paper there is no discussion about what are the actual observation space of the environments, could these be clarified better.\n\n2. 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?\n\n3. 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. \n\n4. 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. \n\n\nPS: Please refer to the discussion below with the authors as I have increased my score from 3 to 6.", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "SklAzVndYH", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper introduces an approach to learning compositional koopman operators to efficiently model the dynamics of non-linear systems, consisting of an unspecified number of objects with repetitive dynamics.\n\nThe key contribution of this work is the use of a graph neural network that allows the koopman operator to be learned for systems of  multiple objects, and the incorporation of blockwise structure in the koopman gain and control matrices that improves parameter estimation process and is shown to reduce over-fitting.\n\nResults show the proposed approach is effective, although only simple toy problems are examined and controlled. Nevertheless, this is a useful demonstration of learning for soft robot systems, and the idea of using koopman embeddings is likely to be of value to the ICLR community. \n\nThe paper is well written, and I like the idea of incorporating additional structure into the learning process through the expected blockwise structure.\n\nAs I understand it, the proposed approach is able to exploit this natural blockwise structure due to the assumption that the same physical dynamics are followed by each block (although objects can also be labelled as rigid/moving). How would this approach extend to objects with vary different properties (eg. 5 objects with connected with springs of different stiffness)?\n\nCould you provide more details on the training process, is the model trained in an end-to-end fashion, or in parts?\n\nFor the control experiment, the choice was made to re-evaluate the control after 32 steps. Why was this the case? Is this due to the time taken to simulate/ find controls? MPC would typically re-plan faster than this.\n\nOn a related point, the introduction states that the proposed approach is 20 times faster than baselines, but no evidence of this is provided. How was this assessed, and is this at prediction or training time? How long does it take to replan using the proposed model and SQP?", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper."}, "extra_text": {}}, {"review_id": "H1lOXuJacS", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper is well written and the proposed idea is novel and builds on a sound theoretical framework of Koopman operator theory. A physical system is a represented as a graph; graph neural network is used to encode the current state to an object-centric embedding where the dynamics are assumed to be linear (Koopman operator theory) and modeled as a transition matrix. The key contribution is to recognize that similar physical interactions can be modeled using same parameters which constraints the transition matrix to be block-wise with shared parameters. Furthermore, the model is extended to add a control matrix to model external control. Experiments are conducted on simulations as well as control for 3 different settings - a a hanging rope/string anchored at the top, soft robot with an anchor, and soft robot in fluids.\n\nStrengths\n* Proposed method builds on a sound theoretical basis; although the linear dynamics model appear somewhat limited compared for the complex dynamics, thorough experiments are conducted to demonstrate the effectiveness of the method. The efficiency of the proposed algorithm compared to prior work makes is much practically useful.\n* Well written with examples and illustrative figures.\n* Quantitative analysis together with ablation studies on the structure are insightful.\n\nWeaknesses\n* My primary concern are with the evaluation.\n- Experiments are only conducted on synthetic datasets. While experiments on real datasets are understandably difficult especially for quantitative validation, it would help to map the experiments to real problems to gain a more intuitive understanding and thus cater to a broader community.\n- This paper and several related works are all evaluated on different problems; it would be useful to evaluate on similar tasks; for instance, strings [Battaglia 2016]. It would make it easier to draw comparisons.\n- It's clear that the community would benefit significantly by having a benchmark or a web-based evaluation methodology (similar to OpenAI gym used actively reinforcement learning community). Unfortunately, this paper does not seem to offer a solution to this issue but continue to evaluate in ways similar to previous papers.\n", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper."}, "extra_text": {}}, {"review_id": "H1gWd_J4cS", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper proposes 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 combination of deep Koopman operator with graph neural nets is very novel and interesting. The experiments are also well designed and the results shown in the paper also indicate the effectiveness and efficiency of the algorithm in both simulation and control. In conclusion, I think this work can inspire more wok into modeling larger and more complex systems by integrating the power of the Koopman theory and the expressiveness of neural networks.", "extra_scores": {"experience_assessment": "I do not know much about this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes using object-centered graph neural network embeddings of a dynamical system as approximate Koopman embeddings, and then learning the linear transition matrix to model the dynamics of the system according to the Koopman operator theory. The authors propose adding an inductive bias (a block diagonal structure of the transition matrix with shared components) to limit the number of parameters necessary to learn, which improves the computational efficiency and generalisation of the proposed approach. The authors also propose adding an additional input component that allows for external control of the dynamics of the system. The reviewers initially had concerns about the experimental section, since the approach was only tested on toy domains. The reviewers also asked for more baselines. The authors were able to answer some of the questions raised during the discussion period, and by the end of it all reviewers agreed that this is a solid and novel piece of work that deserves to be accepted. For this reason I recommend acceptance.
Decision: Accept (Spotlight)