Original: L EARNING C OMPOSITIONAL K OOPMAN O PERATORS FOR M ODEL -B ASED C ONTROL
Revised: L EARNING C OMPOSITIONAL K OOPMAN O PERATORS FOR M ODEL -B ASED C ONTROL

Original: Anonymous authors
Revised: Yunzhu Li MIT CSAIL Hao He ∗ MIT CSAIL

Original: Paper under double-blind review
Revised: Jiajun Wu MIT CSAIL Dina Katabi MIT CSAIL

Revised: Antonio Torralba

Revised: MIT CSAIL

Original: A BSTRACT
Revised: A BSTRACT

Original: Finding an embedding space for a linear approximation of a nonlinear dynamical system enables efﬁcient system identiﬁcation and control synthesis.
Revised: Finding an embedding space for a linear approximation of a nonlinear dynamical system enables efﬁcient system identiﬁcation and control synthesis.

Original: The Koopman operator theory lays the foundation for identifying the nonlinear-to-linear coordinate transformations with data-driven methods.
Revised: The Koopman operator theory lays the foundation for identifying the nonlinear-to-linear coordinate transformations with data-driven methods.

Original: Recently, researchers have proposed to use deep neural networks as a more expressive class of basis functions for calculating the Koopman operators.
Revised: Recently, researchers have proposed to use deep neural networks as a more expressive class of basis functions for calculating the Koopman operators.

Original: These approaches, however, assume a ﬁxed dimensional state space; they are therefore not applicable to scenarios with a variable number of objects.
Revised: These approaches, however, assume a ﬁxed dimensional state space; they are therefore not applicable to scenarios with a variable number of objects.

Original: In this paper, we propose to learn compositional
Revised: In this paper, we propose to learn compositional

Original: 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.
Revised: 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.

Original: The learned dynamics can quickly adapt to new environments of unknown physical parameters and produce control signals to achieve a speciﬁed goal.
Revised: The learned dynamics can quickly adapt to new environments of unknown physical parameters and produce control signals to achieve a speciﬁed goal.

Original: Our experiments on manipulating ropes and controlling soft robots show that the proposed method has better efﬁciency and generalization ability than existing baselines.
Revised: Our experiments on manipulating ropes and controlling soft robots show that the proposed method has better efﬁciency and generalization ability than existing baselines.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Simulating and controlling complex dynamical systems such as ropes or soft robots rely on the dynamics model’s two key features: ﬁrst, it needs to be efﬁcient for system identiﬁcation and motor control; second, it needs to be generalizable to a complex, constantly evolving environments.
Revised: Simulating and controlling complex dynamical systems such as ropes or soft robots rely on the dynamics model’s two key features: ﬁrst, it needs to be efﬁcient for system identiﬁcation and motor control; second, it needs to be generalizable to a complex, constantly evolving environments.

Original: In practice, computational models for complex, nonlinear dynamical systems are often not efﬁcient enough for real-time control (Mayne, 2000).
Revised: In practice computational models for complex, nonlinear dynamical systems are often not efﬁcient enough for real-time control (Mayne, 2000).

Original: The Koopman operator theory suggests that identifying nonlinear-to-linear coordinate transformations allows efﬁcient linear approximation of nonlinear systems (Williams et al., 2015; Mauroy & Goncalves, 2016).
Revised: The Koopman operator theory suggests that identifying nonlinear-to-linear coordinate transformations allows efﬁcient linear approximation of nonlinear systems (Williams et al., 2015; Mauroy & Goncalves, 2016).

Original: Fast as they are; however, existing papers on Koopman operators focus on a single dynamical system, making it hard for them to generalize to cases where there are a variable number of components.
Revised: Fast as they are; however, existing papers on Koopman operators focus on a single dynamical system, making it hard for them to generalize to cases where there are a variable number of components.

Original: In contrast, recent advances in approximating dynamics models with deep nets have demonstrated its power in characterizing complex, generic environments.
Revised: In contrast, recent advances in approximating dynamics models with deep nets have demonstrated its power in characterizing complex, generic environments.

Original: 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.
Revised: 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.

Original: This allows their models to generalize to scenarios with a variable number of objects (Battaglia et al., 2016;
Revised: This allows their models to generalize to scenarios with a variable number of objects (Battaglia et al., 2016;

Original: Chang et al., 2017).
Revised: Chang et al., 2017).

Original: Despite their strong generalization power, they are not as efﬁcient in system identiﬁcation and control, because deep nets are heavily over-parameterized, making optimization time-consuming and sample-inefﬁcient.
Revised: Despite their strong generalization power, they are not as efﬁcient in system identiﬁcation and control, because deep nets are heavily over-parameterized, making optimization time-consuming and sample-inefﬁcient.

Original: In this paper, we propose compositional Koopman operators, integrating Koopman operators with graph networks for generalizable and efﬁcient dynamics modeling.
Revised: In this paper, we propose compositional Koopman operators, integrating Koopman operators with graph networks for generalizable and efﬁcient dynamics modeling.

Original: We build on the idea of encoding states into object-centric embeddings with graph neural networks, which ensures generalization power.
Revised: We build on the idea of encoding states into object-centric embeddings with graph neural networks, which ensures generalization power.

Original: But instead of using over-parameterized neural nets to model state transition, we identify the
Revised: But instead of using over-parameterized neural nets to model state transition, we identify the

Original: Koopman matrix and control matrix from data as a linear approximation of the nonlinear dynamical system.
Revised: Koopman matrix and control matrix from data as a linear approximation of the nonlinear dynamical system.

Original: The linear approximation allows efﬁcient system identiﬁcation and control synthesis.
Revised: The linear approximation allows efﬁcient system identiﬁcation and control synthesis.

Original: The main challenge of extending Koopman theory to multi-object systems is scalability.
Revised: The main challenge of extending Koopman theory to multi-object systems is scalability.

Original: The number of parameters in the Koopman matrix scales quadratically with the number of objects, which harms the learning efﬁciency and leads to overﬁtting.
Revised: The number of parameters in the Koopman matrix scales quadratically with the number of objects, which harms the learning efﬁciency and leads to overﬁtting.

Original: To tackle this issue, we exploit the structure of the underlying system and use the same block-wise Koopman sub-matrix for object pairs of the same relation. This signiﬁcantly reduces the number of parameters that need to be identiﬁed by making it independent of the size of the system.
Revised: To tackle this issue, we exploit the structure of the [Table] +b6PDCSymPE004nj4UJgxqAfMqYEAlwZqNDUFYUpMV4iGSCGtTWNmU4M5+eZ60alX3tFq7OavUL4s6SuAAHIJj4IJzUAfXoAGaAINH8AxewZv1ZL1Y79bHdHTBKnb2wB9Ynz8WlpWg</latexit>

Original: Our experiments include simulating and controlling ropes of variable lengths and soft robots of different shapes.
Revised: Our experiments include simulating and controlling ropes of variable lengths and soft robots of different shapes.

Original: The compositional Koopman operators are signiﬁcantly more accurate than the state-of-the-art learned physics engines (Battaglia et al., 2016; Li et al., 2019b), andis 20 times faster when adapting to new environments of unknown physical parameters.
Revised: The compositional Koopman operators are signiﬁcantly 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.

Original: 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.
Revised: 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.

Original: Please see our supplement video for an in-depth demonstration ∗ .
Revised: Please see our project page for demonstrating videos.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Koopman operators.
Revised: Koopman operators.

Original: 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 &
Revised: 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 &

Original: Neumann, 1932).
Revised: Neumann, 1932).

Original: 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.
Revised: 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.

Original: Researchers have proposed various algorithms to explore the Koopman spectral properties from data.
Revised: Researchers have proposed various algorithms to explore the Koopman spectral properties from data.

Original: A large portion of them are in the class of dynamic mode decomposition (DMD) (Rowley et al., 2009; Schmid, 2010;
Revised: A large portion of them are in the class of dynamic mode decomposition (DMD) (Rowley et al., 2009; Schmid, 2010;

Original: Tu et al., 2014; Williams et al., 2015; Arbabi & Mezic, 2017).
Revised: Tu et al., 2014; Williams et al., 2015; Arbabi & Mezic, 2017).

Original: The linear representation will enable efﬁcient prediction, estimation, and control using tools from linear dynamical systems (Williams et al., 2016;
Revised: The linear representation will enable efﬁcient prediction, estimation, and control using tools from linear dynamical systems (Williams et al., 2016;

Original: Proctor et al., 2018; Mauroy & Goncalves, 2017; Korda & Mezi´c, 2018).
Revised: Proctor et al., 2018; Mauroy & Goncalves, 2017; Korda & Mezi´c, 2018).

Original: People have been using hand-designed Koopman observables for various modeling and control tasks (Brunton et al., 2016;
Revised: People have been using hand-designed Koopman observables for various modeling and control tasks (Brunton et al., 2016;

Original: Kaiser et al., 2017; Abraham et al., 2017; Bruder et al., 2018; Arbabi et al., 2018).
Revised: Kaiser et al., 2017; Abraham et al., 2017; Bruder et al., 2018; Arbabi et al., 2018).

Original: Some recent works have applied the method to the real world and successfully control soft robots with great precision (Bruder et al., 2019; Mamakoukas et al., 2019).
Revised: Some recent works have applied the method to the real world and successfully control soft robots with great precision (Bruder et al., 2019; Mamakoukas et al., 2019).

Original: However, hand-crafted basis functions sometimes fail to generalize to more complex environments.
Revised: However, hand-crafted basis functions sometimes fail to generalize to more complex environments.

Original: 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 ﬂuid control (Morton et al., 2018).
Revised: 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 ﬂuid control (Morton et al., 2018).

Original: Morton et al.
Revised: Morton et al.

Original: (2019) has also extended the framework to account for uncertainty in the system by inferring a distribution over observations.
Revised: (2019) has also extended the framework to account for uncertainty in the system by inferring a distribution over observations.

Original: Our model differs by explicitly modeling the compositionality of the underlying system with graph networks.
Revised: Our model differs by explicitly modeling the compositionality of the underlying system with graph networks.

Original: It generalizes better to environments of a variable number of objects or soft robots of different shapes.
Revised: It generalizes better to environments of a variable number of objects or soft robots of different shapes.

Original: Learning-based physical simulators.
Revised: Learning-based physical simulators.

Original: Battaglia et al.
Revised: Battaglia et al.

Original: and Chang et al.
Revised: and Chang et al.

Original: ﬁrst explored learning a simulator from data by approximating object interactions with neural networks.
Revised: ﬁrst explored learning a simulator from data by approximating object interactions with neural networks.

Original: These models are no longer bounded to hard-coded physical rules, and can adapt to scenarios where the underlying physics is unknown.
Revised: These models are no longer bounded to hard-coded physical rules, and can adapt to scenarios where the underlying physics is unknown.

Original: Please refer to Battaglia et al.
Revised: Please refer to Battaglia et al.

Original: (2018) for a full review.
Revised: (2018) for a full review.

Original: Recently,
Revised: Recently,

Original: Mrowca et al.
Revised: Mrowca et al.

Original: (2018) and Li et al.
Revised: (2018) and Li et al.

Original: (2019a) extended these models to approximate particle dynamics of deformable shapes and ﬂuids.
Revised: (2019a) extended these models to approximate particle dynamics of deformable shapes and ﬂuids.

Original: Flexible as they are, these models become less efﬁcient 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.
Revised: Flexible as they are, these models become less efﬁcient 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.

Original: Nagabandi et al.
Revised: Nagabandi et al.

Original: (2018a;b) proposed to use meta-learning for online adaptation, and have shown to be effective in simulated robots and a real legged millirobot.
Revised: (2018a;b) proposed to use meta-learning for online adaptation, and have shown to be effective in simulated robots and a real legged millirobot.

Original: However, it is not clear whether their methods can generalize to systems with variable numbers of instances.
Revised: However, it is not clear whether their methods can generalize to systems with variable numbers of instances.

Original: The use of graph nets and Koopman operators in our model allows better generalization ability and enables efﬁcient system identiﬁcation as we only need to identify the transition matrices, which is essentially a least-square problem and can be solved very efﬁciently.
Revised: The use of graph nets and Koopman operators in our model allows better generalization ability and enables efﬁcient system identiﬁcation as we only need to identify the transition matrices, which is essentially a least-square problem and can be solved very efﬁciently.

Original: People have also used the learned physics engines for planning and control.
Revised: People have also used the learned physics engines for planning and control.

Original: Many previous papers in this direction learn a latent dynamics model together with a policy in a model-based reinforcement learning setup (Racani`ere 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;
Revised: Many previous papers in this direction learn a latent dynamics model together with a policy in a model-based reinforcement learning setup (Racani`ere 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;

Original: Li et al., 2019b; Janner et al., 2019).
Revised: Li et al., 2019b; Janner et al., 2019).

Original: 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 efﬁciently.
Revised: 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 efﬁciently.

Original:  A PPROACH
Revised:  A PPROACH

Original: We ﬁrst 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 [Table]
Revised: We ﬁrst 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.

Original: +b6PDCSymPE004nj4UJgxqAfMqYEAlwZqNDUFYUpMV4iGSCGtTWNmU4M5+eZ60alX3tFq7OavUL4s6SuAAHIJj4IJzUAfXoAGaAINH8AxewZv1ZL1Y79bHdHTBKnb2wB9Ynz8WlpWg</latexit> linear.

Original: We then discuss the compositional nature of  and show  networks can be used to capture the compositionality
Revised: We then discuss the compositional nature of physical systems and show how graph networks can be used to capture the compositionality.

Original: 3. T HE K OOPMAN O PERATORS
Revised: 3. T HE K OOPMAN O PERATORS

Original: Let x t ∈ X ⊂ R n be the state vector for the system at time step t .
Revised: Let x t ∈ X ⊂ R n be the state vector for the system at time step t .

Original: We consider a non-linear discrete-time dynamical system described by x t +1 = F ( x t ) .
Revised: We consider a non-linear discrete-time dynamical system described by x t +1 = F ( x t ) .

Original: We call any function g : X → R an observable of the system, and we note that the set of all observables forms an inﬁnite-dimensional

Original: Hilbert space.

Original: The Koopman operator (Koopman, 1931), denoted by K  , is a linear transformation on this vector space given by K g (cid:44) g ◦ F , so that K g ( x t ) = g (
Revised: The Koopman operator (Koopman, 1931), denoted by K : F → F , is a linear transformation deﬁned by K g (cid:44) g ◦ F , i.e.,

Original: [Equation]
Revised: [Equation]

Revised: for every function g : X → R belonging to F .

Revised: F is the collection of all functions (also referred to as observables ) which forms an inﬁnite-dimensional Hilbert space.

Revised: Note that the function space

Revised: F is invariant under the action of the Koopman operator.

Original: Although the theory guarantees the existence of the Koopman operator, its use in practice is limited by its inﬁnite dimensionality.
Revised: Although the theory guarantees the existence of the Koopman operator, its use in practice is limited by its inﬁnite dimensionality.

Original: Most often, we assume there is an invariant subspace G of the Koopman operator.
Revised: Most often, we assume there is an invariant subspace G of the Koopman operator.

Original: It spanned by a set of base observation functions { g 1 , · · · , g m } and satisﬁes that K g ∈ G for any g ∈ G .
Revised: It spanned by a set of base observation functions { g 1 , · · · , g m } and satisﬁes that K g ∈ G for any g ∈ G .

Original: With a slightly abuse of the notation, we now use g ( x t ) :
Revised: With a slightly abuse of the notation, we now use g ( x t ) :

Original: R n → R m to represent [ g 1 ( x t ) , · · · , g m ( x t )] T .
Revised: R n → R m to represent [ g 1 ( x t ) , · · · , g m ( x t )] T .

Original: By constraining the Koopman operator on this invariant subspace, we get a ﬁnite-dimensional linear operator K ∈ R m × m that we refer as the Koopman matrix .
Revised: By constraining the Koopman operator on this invariant subspace, we get a ﬁnite-dimensional linear operator K ∈ R m × m that we refer as the Koopman matrix .

Original: Traditionally, people hand-craft base observation functions from the knowledge of underlying physics.
Revised: Traditionally, people hand-craft base observation functions from the knowledge of underlying physics.

Original: The system identiﬁcation problem is then reduced to ﬁnding the Koopman matrix K , which can be solved by linear regression given historical data of the system.
Revised: The system identiﬁcation problem is then reduced to ﬁnding the Koopman matrix K , which can be solved by linear regression given historical data of the system.

Original: Recently, researchers have also explored data-driven methods that automatically ﬁnd the Koopman invariant subspace via representing the base observation functions g ( x ) via deep neural networks.
Revised: Recently, researchers have also explored data-driven methods that automatically ﬁnd the Koopman invariant subspace via representing the base observation functions g ( x ) via deep neural networks.

Original: 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., 2019).
Revised: 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., 2019).

Original: Mathematically, considering a dynamical system with an external control input u t , x t +1 = F ( x t , u t ) ,  we aim to ﬁnd the Koopman observation functions and the linear dynamics model in the form ofg ( x t +1 )
Revised: Mathematically, considering a dynamical system, x t +1 = F ( x t , u t ) , with an external control input u t , we aim to ﬁnd the Koopman observation functions and the linear dynamics model in the form of

Original: =
Revised: [Equation]

Original: Kg ( x t ) + L u t , where the coefﬁcient matrix L is referred to as the control matrix .
Revised: where the coefﬁcient matrix L is referred to as the control matrix .

Original: 3. C OMPOSITIONAL K OOPMAN O PERATORS
Revised: 3. C OMPOSITIONAL K OOPMAN O PERATORS

Original: The dynamics of a physical system are governed by physical rules, which are usually shared across different subcomponents in the system.
Revised: The dynamics of a physical system are governed by physical rules, which are usually shared across different subcomponents in the system.

Original: Explicitly modeling such compositionality enables more efﬁcient system identiﬁcation and control synthesis and provides better generalization ability.
Revised: Explicitly modeling such compositionality enables more efﬁcient system identiﬁcation and control synthesis and provides better generalization ability.

Original: Motivating example.
Revised: Motivating example.

Original: Consider a system with N balls moving in a 2D plane, each pair connected by a linear spring.
Revised: Consider a system with N balls moving in a 2D plane, each pair connected by a linear spring.

Original: Assume all balls have mass 1 and all springs share the same stiffness coefﬁcient k .
Revised: Assume all balls have mass 1 and all springs share the same stiffness coefﬁcient k .

Original: We denote the i ’s ball’s position as ( x i , y i ) and its velocity as ( ˙ x i , ˙ y i ) .
Revised: We denote the i ’s ball’s position as ( x i , y i ) and its velocity as ( ˙ x i , ˙ y i ) .

Original: For ball i , equation 1 describes its dynamics, where x i (cid:44) [ x i , y i , ˙ x i , ˙ y i ]
Revised: For ball i , equation 3 describes its dynamics, where x i (cid:44) [ x i , y i , ˙ x i , ˙ y i ]

Original: T denotes ball
Revised: T denotes ball

Original: i ’s state:
Revised: i ’s state:

Original: [Equation]

Original: [Equation]

Original: [-Nk 
-]
Revised: [-Nk 
-]

Original: [Equation]
Revised: [Equation]

Original: 0

Original: 

Original: [Equation]
Revised: [Equation]

Original: (cid:125) (cid:44) A   x i y i ˙

Original: x i ˙ y i   + (cid:88) j (cid:54) = i    0  0k 0 k 0  (cid:124) (cid:123)(cid:122) (cid:125) (cid:44) B   x i y i ˙
Revised: [Equation]

Original: x i ˙ y i 

Original:  .

Original: [Equation]
Revised: [Equation]

Original: We can represent the state of the whole system using the union of every ball’s state, where x = [ x 1 , · · · , x N ] T .
Revised: We can represent the state of the whole system using the union of every ball’s state, where x = [ x 1 , · · · , x N ] T .

Original: 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 2:
Revised: 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:

Original: [Equation]
Revised: [Equation]

Revised: B · · · B B

Revised: A · · · B ..

Revised: B

Revised: B · · ·

Revised: A       x 1 x... x N    .

Original: Based on the linear spring system, we make three observations for multi-object systems.
Revised: Based on the linear spring system, we make three observations for multi-object systems.

Original: • The system state is composed of the state of each individual object.
Revised: • The system state is composed of the state of each individual object.

Original: The dimension of the whole system scales linearly with the number of objects.
Revised: The dimension of the whole system scales linearly with the number of objects.

Original: We formulate the system state by concatenating the state of every object, corresponding to an object-centric state representation.
Revised: We formulate the system state by concatenating the state of every object, corresponding to an object-centric state representation.

Original: • The transition matrix has a block-wise substructure.
Revised: • The transition matrix has a block-wise substructure.

Original: After assuming an object-centric state representation, the transition matrix naturally has a block-wise structure as shown in equation 2.
Revised: After assuming an object-centric state representation, the transition matrix naturally has a block-wise structure as shown in equation 4.

Original: • The same physical interactions share the same transition block.
Revised: • The same physical interactions share the same transition block.

Original: The blocks in the transition matrix encode actual interactions and generalize across systems.
Revised: The blocks in the transition matrix encode actual interactions and generalize across systems.

Original: A and B govern the dynamics of the linear spring system, and are shared by systems with a different number of objects.
Revised: A and B govern the dynamics of the linear spring system, and are shared by systems with a different number of objects.

Original: These observations motivate us to exploit the structure of multi-object systems, instead of learning separate models for systems that contains different numbers of balls.
Revised: 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.

Original: Compositional Koopman operators.
Revised: Compositional Koopman operators.

Original: Motivated by the linear spring system, we want to inject a good inductive bias to incorporate compositionality when applying the Koopman theory.
Revised: Motivated by the linear spring system, we want to inject a good inductive bias to incorporate compositionality when applying the Koopman theory.

Original: This allows better generalization ability and more efﬁcient system identiﬁcation and better controller design.
Revised: This allows better generalization ability and more efﬁcient system identiﬁcation and better controller design.

Original: Figure 1 shows an overview of our model.
Revised: Figure 1 shows an overview of our model.

Original: Considering a system with N objects, we denote x t as the system state at time t and x ti is the state of the i ’th object.
Revised: Considering a system with N objects, we denote x t as the system state at time t and x ti is the state of the i ’th object.

Original: We further denote g t (cid:44) g ( x t ) as the embedding of the state in the Koopman invariant space.
Revised: We further denote g t (cid:44) g ( x t ) as the embedding of the state in the Koopman invariant space.

Original: In the rest of the paper, we call g t the Koopman embedding .
Revised: In the rest of the paper, we call g t the Koopman embedding .

Original: 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.
Revised: 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.

Original: • The Koopman embedding of the system is composed of the Koopman embedding of every objects.
Revised: • The Koopman embedding of the system is composed of the Koopman embedding of every objects.

Original: 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 =
Revised: Similar to the decomposition in the state space, we assume the Koopman embedding can be divided into object-centric sub-embeddings,

Original: [ g t 1(cid:62) , · · · , g tN(cid:62) ] (cid:62) ∈ R Nm , where weoverload the notation and use g ti = g i ( x t ) ∈
Revised: i.e. g t ∈ R Nm denoting the concatenation of g t 1 , · · · , g tN , where we use g ti = g i ( x t ) ∈

Original: R m as the Koopman embedding for the i ’th object.
Revised: R m as the Koopman embedding for the i ’th object.

Original: • The Koopman matrix has a block-wise structure.
Revised: • The Koopman matrix has a block-wise structure.

Original: It is natural to think the Koopman matrix is composed of block matrices after assuming an object-centric Koopman embeddings.
Revised: It is natural to think the Koopman matrix is composed of block matrices after assuming an object-centric Koopman embeddings.

Original: In equation 3, K ij ∈ R m × m and L ij ∈ R m ×| u i | are blocks of the Koopman matrix and the control matrix, while u t =
Revised: 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

Original: [ u t 1(cid:62) , · · · , u tN (cid:62) ] (cid:62) is the  control signal at time t :
Revised: Nl is the concatenation of u t 1 , · · · , u t N denoting the total control signal at time t :

Original: [Equation]
Revised: [Equation]

Original: As we have seen in

Original: the

Original: case

Original:  of linear spring system those matrix blocks are not independent but 
Revised: 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.

Original: some of them share the same set of values. • The same physical interactions shall share the same transition block.
Revised: • The same physical interactions shall share the same transition block.

Original: The equivalence between the blocks should reﬂect the equivalence of the objects, where we use the same transition submatrix for object pairs of the same relation.
Revised: The equivalence between the blocks should reﬂect the equivalence of the objects, where we use the same transition submatrix for object pairs of the same relation.

Original: For example, if the system is composed of N identical objects, then, by symmetry, all the diagonal blocks should be the same while all the off-diagonal blocks should also be the same.
Revised: For example, if the system is composed of N identical objects, then, by symmetry, all the diagonal blocks should be the same while all the off-diagonal blocks should also be the same.

Original: The repetitive structure allows us to efﬁciently identify the values using least square regression.
Revised: The repetitive structure allows us to efﬁciently identify the values using least squares regression.

Original: 3. L EARNING THE K OOPMAN E MBEDDINGS U SING G RAPH N EURAL N ETWORKS
Revised: 3. L EARNING THE K OOPMAN E MBEDDINGS U SING G RAPH N EURAL N ETWORKS

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

Original: =
Revised: =

Original: ( O t , R ) , where vertices O t = { o ti }
Revised: ( O t , R ) , where vertices O t = { o ti }

Original: Ni =1 represent objects and edges R = { r k } N 2 k
Revised: Ni =1 represent objects and edges R = { r k } N 2 k

Original: =1 represent pair-wise relations.
Revised: =1 represent pair-wise relations.

Original: Speciﬁcally, o ti =
Revised: Speciﬁcally, o ti =

Original: ( x ti , a oi ) , where x ti is the state of object i and a oi is a one-hot vector indicating the object type, e.g., ﬁxed or movable.
Revised: ( x ti , a oi ) , where x ti is the state of object i and a oi is a one-hot vector indicating the object type, e.g., ﬁxed or movable.

Original: For relation, we have r k = ( u k , v k , a rk ) , 1 ≤ u k , v k ≤ N , where u k and v k are integers denoting the end points of this directed edge, and a rk is a one-hot vector denoting the type of the relation k .
Revised: For relation, we have r k = ( u k , v k , a rk ) , 1 ≤ u k , v k ≤ N , where u k and v k are integers denoting the end points of this directed edge, and a rk is a one-hot vector denoting the type of the relation k .

Original: We use a graph neural network similar to Interaction Networks (IN) (Battaglia et al., 2016) to generate object-centric Koopman embeddings.
Revised: We use a graph neural network similar to Interaction Networks (IN) (Battaglia et al., 2016) to generate object-centric Koopman embeddings.

Original: IN deﬁnes an object function f O and a relation function f R to model objects and their relations in a compositional way.
Revised: IN deﬁnes an object function f O and a relation function f R to model objects and their relations in a compositional way.

Original: Similar to a message passing procedure, we calculate the edge effect e tk = f R ( o tu k , o tv k , a rk ) k
Revised: Similar to a message passing procedure, we calculate the edge effect e tk = f R ( o tu k , o tv k , a rk ) k

Original: =1 ...N 2 , and node effect g ti =
Revised: =1 ...N 2 , and node effect g ti =

Original: f O ( o ti , (cid:80) k ∈N i e tk ) i =1 ...N , where N i denotes the relations that related to  object i and { g ti } are the derived Koopman embeddings.
Revised: f O ( o ti , (cid:80) k ∈N i e tk ) i =1 ...N , where N i denotes the relations that point to the object i and { g ti } are the derived Koopman embeddings.

Original: We use this graph neural network, denoted as φ , to represent our Koopman observation function.
Revised: We use this graph neural network, denoted as φ , to represent our Koopman observation function.

Original: System identiﬁcation.
Revised: System identiﬁcation.

Original: For a sequence of observations (cid:101) x
Revised: For a sequence of observations (cid:101) x

Original: = [ x 1 , · · · , x T ] from time 1 to time T , we ﬁrst map them to the Koopman space as (cid:101) g
Revised: = [ x 1 , · · · , x T ] from time 1 to time T , we ﬁrst map them to the Koopman space as (cid:101) g

Original: = [ g 1 , · · · , g T ] using the graph encoder φ , where g t = φ ( x t ) .
Revised: = [ g 1 , · · · , g T ] using the graph encoder φ , where g t = φ ( x t ) .

Original: We use g i : j to denote the sub-sequence [ g i , · · · , g j ] .
Revised: We use g i : j to denote the sub-sequence [ g i , · · · , g j ] .

Original: To identify the Koopman matrix, we solve the linear regression min K (cid:107) K g 1: T − 1 − g 2: T (cid:107) 2 .
Revised: To identify the Koopman matrix, we solve the linear regression min K (cid:107) K g 1: T − 1 − g 2: T (cid:107) 2 .

Original: As a result, K = g 2: T ( g 1: T − 1 ) † will asymptotically approach the Koopman operator K with an increasing T .
Revised: As a result, K = g 2: T ( g 1: T − 1 ) † will asymptotically approach the Koopman operator K with an increasing T .

Original: For cases where there are control inputs (cid:101) u
Revised: For cases where there are control inputs (cid:101) u

Original: = [ u 1 , · · · , u T − 1 ] , the calculation of the Koopman matrix and the control matrix is essentially solving a least square problem w.r.t.
Revised: = [ 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.

Original: the objective
Revised: the objective

Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: 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., (cid:101) g ∈ R Nm × T and K ∈
Revised: 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., (cid:101) g ∈ R Nm × T and K ∈

Original: R Nm × Nm .
Revised: R Nm × Nm .

Original: If we do not enforce any structure on the Koopman matrix K , we will have to identify N 2 m 2 parameters.
Revised: If we do not enforce any structure on the Koopman matrix K , we will have to identify N 2 m 2 parameters.

Original: Instead, we can signiﬁcantly reduce the number by leveraging the assumption on the structure of K .
Revised: Instead, we can signiﬁcantly reduce the number by leveraging the assumption on the structure of K .

Original: 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
Revised: 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

Original: h × m × m .
Revised: h × m × m .

Original: Then, the number of parameter to be identiﬁed reduce to hm 2 .
Revised: Then, the number of parameter to be identiﬁed reduce to hm 2 .

Original: Usually, h does not depend on N , and is much smaller than N 2 .
Revised: Usually, h does not depend on N , and is much smaller than N 2 .

Original: 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  .
Revised: 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 .

Original: Finally, as shown in equation 5, we represent the Koopman matrix as the product of the index tensor σ and the parameter tensor ˆ K :[Equation]  , where
Revised: Finally, as shown in equation 7, we represent the Koopman matrix as the product of the index tensor σ and the parameter tensor ˆ K :

Original: [Equation]
Revised: [Equation]

Original: Similar to the Koopman matrix, we assume the same block structure in the control matrix L and denote its parameter as ˆ L ∈ R h [Equation] N ×| u i | .
Revised: 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 .

Original: The least square problem of identifying ˆ K and ˆ L becomes
Revised: The least squares problem of identifying ˆ K and ˆ L becomes

Original: [Equation]
Revised: [Equation]

Revised: where σ ⊗ ˆ K ∈ R Nm × Nm , σ ⊗ ˆ L ∈

Revised: R

Revised: Nm × Nl , g 1: T − 1 ∈ R Nm × ( T − 1) and (cid:101)

Original:  Since the linear least square problems described in equation 4 and equation 6 have analytical solutions, performing system identiﬁcation using our method is very efﬁcient.
Revised: u ∈ R Nl × ( T − 1) . Since the linear least squares problems described in equation 6 and equation 8 have analytical solutions, performing system identiﬁcation using our method is very efﬁcient.

Original: Training GNN models.
Revised: Training GNN models.

Original: To make predictions on the states, we use a graph decoder ψ to map the
Revised: To make predictions on the states, we use a graph decoder ψ to map the

Original: Koopman embeddings back to the original state space.
Revised: Koopman embeddings back to the original state space.

Original: In total, we have three losses to train the graph encoder and decoder.
Revised: In total, we have three losses to train the graph encoder and decoder.

Original: The ﬁrst term is the auto-encoding lossL ae = 1 T (cid:80) i (cid:107) ψ
Revised: The ﬁrst term is the auto-encoding loss

Original: ( φ ( x i ))
Revised: [Equation]

Original: − x i (cid:107) 2 .

Original: The second term is the prediction loss.
Revised: The second term is the prediction loss.

Original: To calculate it, we rollout in the Koopman space and denote the embeddings as ˆ g 1 = g 1 , and ˆ g t +1 =
Revised: To calculate it, we rollout in the Koopman space and denote the embeddings as ˆ g 1 = g 1 , and ˆ g t +1 =

Original: K ˆ g t + L u t , for t = 1 , · · · , T − 1 .
Revised: K ˆ g t + L u t , for t = 1 , · · · , T − 1 .

Original: The prediction loss is deﬁned as the difference between the decoded states and the actual states, i.e.,L pred = 1 T (cid:80)
Revised: The prediction loss is deﬁned as the difference between the decoded states and the actual states, i.e.,

Original: Ti =1 (cid:107) ψ (ˆ g i ) − x i (cid:107) 2 .
Revised: [Equation]

Original: Third, we employ a metric loss to encourage the Koopman embeddings preserving the distance in the original state space.
Revised: Third, we employ a metric loss to encourage the Koopman embeddings preserving the distance in the original state space.

Original: The loss is deﬁned as the absolute error between the distances measured in the Koopman space and that in the original space, i.e.,L metric = (cid:80) ij (cid:12)(cid:12)
Revised: The loss is deﬁned as the absolute error between the distances measured in the Koopman space and that in the original space, i.e.,

Original: [Table]
Revised: [Equation]

Original: Having Koopman embeddings that perserves the distance in the state space is important as we are using the distance in the Koopman space to deﬁne the cost function for downstream control tasks.
Revised: Having Koopman embeddings that perserves the distance in the state space is important as we are using the distance in the Koopman space to deﬁne the cost function for downstream control tasks.

Original: The ultimate training loss is simply the combination of all the terms above:
Revised: The ultimate training loss is simply the combination of all the terms above:

Original: L =
Revised: L =

Original: L ae + L pred + L metric .
Revised: L ae + L pred + L metric .

Original: We then minimize the loss
Revised: We then minimize the loss

Original: L by optimizing the parameters in the graph encoder φ and graph decoder ψ using stochastic gradient descent.
Revised: L by optimizing the parameters in the graph encoder φ and graph decoder ψ using stochastic gradient descent.

Original: Once the model is trained, it can be used for system identiﬁcation, future prediction, and control synthesis.
Revised: Once the model is trained, it can be used for system identiﬁcation, future prediction, and control synthesis.

Original: 3. C ONTROL
Revised: 3. C ONTROL

Original: For a control task, the goal is to synthesize a sequence of control inputs u 1: T that minimize C = (cid:80)
Revised: For a control task, the goal is to synthesize a sequence of control inputs u 1: T that minimize C = (cid:80)

Original: Tt =1 c t ( x t , u t ) , the total incurred cost, where c t ( x t , u t ) is the instantaneous cost.
Revised: Tt =1 c t ( x t , u t ) , the total incurred cost, where c t ( x t , u t ) is the instantaneous cost.

Original: 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 ) =
Revised: 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 ) =

Original: [ t = T ] (cid:107) x
Revised: [ t = T ] (cid:107) x

Original: t − x ∗ (cid:107) 22 + λ (cid:107) u t (cid:107) 22 .
Revised: t − x ∗ (cid:107) 22 + λ (cid:107) u t (cid:107) 22 .

Original: The ﬁrst term promotes the control sequence that matches the state to the goal, while the second term regularizes the control signals.
Revised: The ﬁrst term promotes the control sequence that matches the state to the goal, while the second term regularizes the control signals.

Original: Open-loop control via quadratic programming (QP).
Revised: Open-loop control via quadratic programming (QP).

Original: Our model maps the original nonlinear dynamics to a linear dynamical system.
Revised: Our model maps the original nonlinear dynamics to a linear dynamical system.

Original: We can then solve the control task by solving a linear control problem.
Revised: We can then solve the control task by solving a linear control problem.

Original: With the assumption that the Koopman embeddings preserve the distance measure, we deﬁne the control cost as c t ( g t , u t ) = 

Revised: With the assumption that the Koopman embeddings preserve the distance measure, we deﬁne the control cost as c t ( g t , u t ) =

Original: [ t = T ] (cid:107) g
Revised: [ t = T ] (cid:107) g

Original: t − g ∗ (cid:107) 

Revised: t − g ∗ (cid:107)

Original: + λ (cid:107) u t (cid:107) 22 .
Revised: + λ (cid:107) u t (cid:107) 22 .

Original: As a result, we reduce the problem to minimizing a quadratic cost function C = (cid:80) Tt =1 c t ( g t , u t ) over variables { g t , u t } Tt = under linear constrains g t +

Revised: As a result, we reduce the problem to minimizing a quadratic cost function C = (cid:80) T t =1 c t ( g t , u t ) over variables { g t , u t } T t = under linear constrains g t +

Original: =
Revised: =

Original: K g t + L u t , where g 1 = φ ( x 1 ) and g ∗ =
Revised: K g t + L u t , where g 1 = φ ( x 1 ) and g ∗ =

Original: φ ( x ∗ ) .
Revised: φ ( x ∗ ) .

Original: Model predictive control (MPC).
Revised: Model predictive control (MPC).

Original: Solving the QP gives us control signals, which might not be good enough for long-term control as the prediction error accumulates.
Revised: Solving the QP gives us control signals, which might not be good enough for long-term control as the prediction error accumulates.

Original: We can combine it with
Revised: We can combine it with

Original: Model Predictive Control, assuming feedback from the environment every τ steps.
Revised: Model Predictive Control, assuming feedback from the environment every τ steps.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: Environments.
Revised: Environments.

Original: We evaluate our method by assessing how well it can simulate and control ropes and soft robots.
Revised: We evaluate our method by assessing how well it can simulate and control ropes and soft robots.

Original: Speciﬁcally, we consider three environments.
Revised: Speciﬁcally, we consider three environments.

Original: (1) Rope (Figure 2a): the top mass of a rope is ﬁxed to a speciﬁc height.
Revised: (1) Rope (Figure 2a): the top mass of a rope is ﬁxed to a speciﬁc height.

Original: We apply force to the top mass to move it in a horizontal line.
Revised: We apply force to the top mass to move it in a horizontal line.

Original: The rest of the masses are free to move according to internal force and gravity.
Revised: The rest of the masses are free to move according to internal force and gravity.

Original: (2) Soft (Figure 2b): we aim to control a soft robot that is consist of soft blocks.
Revised: (2) Soft (Figure 2b): we aim to control a soft robot that is consist of soft blocks.

Original: Blocks in dark grey as rigid and those in light blue are soft blocks.
Revised: Blocks in dark grey as rigid and those in light blue are soft blocks.

Original: Each one of the dark blue blocks is soft but have an actuator inside that can contract or expand the block.
Revised: Each one of the dark blue blocks is soft but have an actuator inside that can contract or expand the block.

Original: One of the blocks is pinned to the ground, as shown using the red dots.
Revised: One of the blocks is pinned to the ground, as shown using the red dots.

Original: (3) Swim (Figure 2c): instead of pinning the soft robot to the ground, we let the robot swim in ﬂuids.
Revised: (3) Swim (Figure 2c): instead of pinning the soft robot to the ground, we let the robot swim in ﬂuids.

Original: The colors shown in this environment have the same meaning as in Soft .
Revised: The colors shown in this environment have the same meaning as in Soft .

Revised: [Figure]

Revised: ( a ) R o p e ( b ) S o f t ( c ) S w i m ( a ) R op e

Revised: ( b ) S o f t ( c ) S w i m

Original: Observation space.
Revised: Time Prediction GT Prediction

Revised: GT

Revised: Prediction

Revised: GT

Revised: S i m u l a ti on

Revised: C on t r o l

Revised: Observation space.

Original: In the Rope environment, each mass on the rope is considered as an object.
Revised: In the Rope environment, each mass on the rope is considered as an object.

Original: The observation of each mass is its position and velocity in the 2D plane, which has a dimension of 4 .
Revised: The observation of each mass is its position and velocity in the 2D plane, which has a dimension of 4 .

Original: In total, a rope with N masses has an observation space of dimension 4 N .
Revised: In total, a rope with N masses has an observation space of dimension 4 N .

Original: In both the Soft and the Swim environments, each quadrilateral is considered as an object.
Revised: In both the Soft and the Swim environments, each quadrilateral is considered as an object.

Original: For each quadrilateral, we have access to the positions and velocities of the four corners.
Revised: For each quadrilateral, we have access to the positions and velocities of the four corners.

Original: Thus for a soft robot containing N quadrilaterals, we have a 4 × 4 × N = 

Revised: Thus for a soft robot containing N quadrilaterals, we have a 4 × 4 × N =

Original: N dimensional observation.
Revised: N dimensional observation.

Original: Baselines.
Revised: Baselines.

Original: 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 ).
Revised: 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 ).

Original: IN and PN are the state-of-the-art learning-based physical simulators.
Revised: IN and PN are the state-of-the-art learning-based physical simulators.

Original: For control, we ﬁrst ﬁnetune the parameters in IN and PN to adapt to the test environment of unknown physical parameters.
Revised: For control, we ﬁrst ﬁnetune the parameters in IN and PN to adapt to the test environment of unknown physical parameters.

Original: We then apply gradient descent to the control signals by minimizing the distance between the prediction from IN/PN and the target, following Li et al.
Revised: We then apply gradient descent to the control signals by minimizing the distance between the prediction from IN/PN and the target, following Li et al.

Original: (2019b).
Revised: (2019b).

Original: The generated control sequence is fed to the original simulator to evaluate the performance.
Revised: The generated control sequence is fed to the original simulator to evaluate the performance.

Original: Similar to our method, KPM ﬁts a linear dynamics in the Koopman space.
Revised: Similar to our method, KPM ﬁts a linear dynamics in the Koopman space.

Original: Instead of learning Koopman observations from data, KPM uses polynomials of the original states as the basis functions.
Revised: Instead of learning Koopman observations from data, KPM uses polynomials of the original states as the basis functions.

Original: 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.
Revised: 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.

Original: Data generation.
Revised: Data generation.

Original: We generate 10 , 000 samples for Rope and 50 , 000 samples for Soft and Swim.
Revised: We generate 10 , 000 samples for Rope and 50 , 000 samples for Soft and Swim.

Original: Among them, 90% are used for training and the rest for testing.
Revised: Among them, 90% are used for training, and the rest for testing.

Original: Each data sample has 100 time steps.
Revised: Each data sample has 100 time steps.

Revised: In the dataset, the physical systems have a various number of objects from 5 to 10 , i.e. the ropes have 5 to 10 masses while the soft robots in Soft and Swim environments have 5 to 10 quadrilaterals.

Revised: To

Revised: [Figure]

Revised: [Figure]

Original:  Training and evaluation protocols. 
Revised: [Table] +8SpqVsndRrtQvS9WbLI48nMApnIMHV1CFO6hBAxggPMMrvDmPzovz7nwsWnNONnMMf+B8/gCk14zU</latexit> 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 15 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 .

Revised: For both our model and the baselines, we apply 400 K iterations of gradient steps in the Rope environment and 580 K iterations in the Soft and

Revised: Swim environment.

Original: Our model is trained on the sub-sequence of length 64 from the training set.
Revised: Our model is trained on the sub-sequence of length 64 from the training set, and

Original: For evaluation, we use two metrics: simulation error and control error . For a given data sample, the simulation error at time step t is deﬁned as the mean squared error between the model predictionˆ x t and the ground truth x t .
Revised: 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.

Original: For control, we pick the t 0 ’th frame x t 0 and the t 0 + t ’th frame x t 0 + t from a data sample.
Revised: Our model uses 100 consecutive frames to identify the transition matrix via least-square regression.

Original: Then we ask the model to generate a control sequence of length t to transfer the system from the initial state x t 0 to thetarget state x t 0 + t . The control error is deﬁned as the mean squared distance between the target state and the state of the system at time t .
Revised: 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.

Original: For our experiments we have t = 64 .[Figure] [Figure]
Revised: For evaluation, we use two metrics: simulation error and control error .

Original: R ope ( [Figure] S o ft ( c ) S w i m ( a )
Revised: For a given data sample, the simulation error at time step t is deﬁned as the mean squared error between the model prediction ˆ x t and the ground truth x t .

Original: R ope ( b ) S o ft ( c ) S w i m Time
Revised: For control, we pick the t 0 ’th frame x t 0 and the t 0 + t ’th frame x t 0 + t from a data sample.

Original: Prediction GT Prediction GT
Revised: Then we ask the model to generate a control sequence of length t to transfer the system from the initial state x t 0 to the target state x t 0 + t .

Original: Prediction GT S i m u l a t i on C on t r o l
Revised: The control error is deﬁned as the mean squared distance between the target state and the state of the system at time t .

Original: 4. S IMULATION
Revised: 4. S IMULATION

Original: Figure 2 shows qualitative results on simulation.
Revised: Figure 2 shows qualitative results on simulation.

Original: Our model accurately predicts system dynamics for more than 100 steps.
Revised: Our model accurately predicts system dynamics for more than 100 steps.

Original: 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.
Revised: 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.

Original: For Soft, our model captures the interaction between the body parts and generates accurate prediction over the global movements of the robot.
Revised: For Soft, our model captures the interaction between the body parts and generates accurate prediction over the global movements of the robot.

Original: The error mainly comes from the misalignment of some local components.
Revised: The error mainly comes from the misalignment of some local components.

Original: Figure 3 shows quantitative results.
Revised: Figure 3 shows quantitative results.

Original: IN and PN do not work well  due to poor system identiﬁcation ability.IN and PN rely on gradient descent to adapt to the environments of unknown physical parameters and on average takes over 10 seconds. Our model takes around 0 . 43 seconds per instance in the identiﬁcation stage which is about 20 times faster.
Revised: IN and PN do not work well in the Rope and Swim environments due to insufﬁcient system identiﬁcation ability.

Original: The KPM baseline with simple hand-crafted
Revised: The KPM baseline performs poorly in the Rope and

Original: Koopman basis works reasonably well in the Rope and Swim environment, partly because KPM uses the same system identiﬁcation algorithm as our model, leveraging the prior knowledge on the structure of the system.
Revised: Soft environments indicating the limited power of polynomial Koopman base functions.

Original: Our model signiﬁcantly outperforms all the baselines except in the Swim environment, where we are on par with KPM.
Revised: Our model signiﬁcantly outperforms all the baselines.

Original: 4. C ONTROL
Revised: 4. C ONTROL

Original: As the simulation errors of IN and PN are too large for control, we compare our model with KPM.
Revised: We compare our model with KPM, the Koopman baseline using polynomial basis.

Original: In Rope, we ask the models to perform open-loop control where it only solves the QP once at the beginning.
Revised: In Rope, we ask the models to perform open-loop control where it only solves the QP once at the beginning.

Original: The length of the control sequence is 40 .
Revised: The length of the control sequence is 40 .

Original: 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 aftersteps.
Revised: 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.

Original: Thus every model has a second chance to correct its control sequence by solving the QP again at the time step 32 .[Figure] [Figure] [Table]
Revised: Thus every model has a second chance to correct its control sequence by solving the QP again at the time step 32 .

Original: +8SpqVsndRrtQvS9WbLI48nMApnIMHV1CFO6hBAxggPMMrvDmPzovz7nwsWnNONnMMf+B8/gCk14zU</latexit>

Original: As shown in Figure 2, our model leverages the inertia of the rope and matches the target state accurately.
Revised: As shown in Figure 2, our model leverages the inertia of the rope and matches the target state accurately.

Original: 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.
Revised: 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.

Original: The small control error comes from the slight misalignment of the orientation and the size of the body parts.
Revised: The small control error comes from the slight misalignment of the orientation and the size of the body parts.

Original: Figure 4 shows that quantitatively our model outperforms KPM, too.
Revised: Figure 4 shows that quantitatively our model outperforms KPM, too.

Original: 4. A BLATION S TUDY
Revised: 4. A BLATION S TUDY

Original: Structure of the Koopman matrix.
Revised: Structure of the Koopman matrix.

Original: We explore three different structures of the Koopman matrix,
Revised: We explore three different structures of the Koopman matrix,

Original: Block , Diag and None , to understand its effect on the learned dynamics.
Revised: Block , Diag and None , to understand its effect on the learned dynamics.

Original: None assumes no structure in the Koopman matrix.
Revised: None assumes no structure in the Koopman matrix.

Original: Diag assumes a diagonal block structure of K : all off-diagonal blocks ( K ij where i (cid:54) =
Revised: Diag assumes a diagonal block structure of K : all off-diagonal blocks ( K ij where i (cid:54) =

Original: j ) are zeros and all diagonal blocks share the same values.
Revised: j ) are zeros and all diagonal blocks share the same values.

Original: Block predeﬁnes a block-wise structure, decided by the relation between the objects as introduced in Section 3.3.
Revised: Block predeﬁnes a block-wise structure, decided by the relation between the objects as introduced in Section 3.3.

Revised: simulation,

Original: [Table]
Revised: [Table]

Original: Table 1 includes our model’s simulation error and control error with different Koopman matrix structures in
Revised: Table 1 includes our model’s simulation error and control error with different Koopman matrix structures in

Original: Rope.
Revised: Rope.

Original: All models are trained in the Rope environment with 5 to 9 masses.
Revised: All models are trained in the Rope environment with 5 to 9 masses.

Original: Besides the result on the test set, we also report models’ extrapolation performance in parentheses, where the model is evaluated on system with more masses than training, i.e., 10 to 14 masses.
Revised: 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.

Original: Our model with Block structure consistently achieves a smaller error in all settings.
Revised: Our model with Block structure consistently achieves a smaller error in all settings.

Original: Diag assumes an overly simpliﬁed structure, leading to larger errors and failing to make reasonable controls.
Revised: Diag assumes an overly simpliﬁed structure, leading to larger errors and failing to make reasonable controls.

Original: None has comparable simulation errors but larger control errors.
Revised: None has comparable simulation errors but larger control errors.

Original: Without structure in the Koopman matrix, it overﬁts the data and makes the resulting linear dynamics less amiable to the control.
Revised: Without structure in the Koopman matrix, it overﬁts the data and makes the resulting linear dynamics less amiable to the control.

Original: Hyperparameters.
Revised: Hyperparameters.

Original: In our main experiments, we set the dimension of the Koopman embedding to m = 32 per object.
Revised: In our main experiments, we set the dimension of the Koopman embedding to m = 32 per object.

Original: Online system identiﬁcation requires 800 data samples for each training/test case.
Revised: Online system identiﬁcation requires 800 data samples for each training/test case.

Original: 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 identiﬁcation from 200 to 1 , 600 .
Revised: 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 identiﬁcation from 200 to 1 , 600 .

Original: Figure 4d shows that more data for system identiﬁcation leads to better simulation results.
Revised: Figure 4d shows that more data for system identiﬁcation leads to better simulation results.

Original: Figure 4e shows that dimension 16 gives the best results on simulation.
Revised: Figure 4e shows that dimension 16 gives the best results on simulation.

Original: It may suggest that the intrinsic dimension of the Koopman invariant space of the Rope system is around 16 per object.
Revised: It may suggest that the intrinsic dimension of the Koopman invariant space of the Rope system is around 16 per object.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: Compositionality is common in our daily life.
Revised: Compositionality is common in our daily life.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: By leveraging the compositional structure in the Koopman operator via graph neural nets, our model can efﬁciently manipulate deformable objects such as ropes and soft robots and generalize to systems with variable numbers of instances.
Revised: By leveraging the compositional structure in the Koopman operator via graph neural nets, our model can efﬁciently manipulate deformable objects such as ropes and soft robots, and generalize to systems with variable numbers of instances.

Original: We hope this work could encourage more endeavor in modeling larger and more complex systems by integrating the power of the Koopman theory and the expressiveness of neural networks.
Revised: We hope this work could encourage more endeavor in modeling larger and more complex systems by integrating the power of the Koopman theory and the expressiveness of neural networks.

Original: A E NVIRONMENT AND M ODEL D ETAILS
Revised: A E NVIRONMENT AND M ODEL D ETAILS

Original: Interaction types.
Revised: Interaction types.

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

Original: In the Rope environment, the top mass has a ﬁxed height and is considered differently from the other masses.
Revised: In the Rope environment, the top mass has a ﬁxed height and is considered differently from the other masses.

Original: Thus, we have 2 types of self-interactions for the top mass and the non-top masses.
Revised: Thus, we have 2 types of self-interactions for the top mass and the non-top masses.

Original: In addition, we have 8 types of interactions between different objects.
Revised: In addition, we have 8 types of interactions between different objects.

Original: The objects on a relation could be either top mass or non-top mass.
Revised: The objects on a relation could be either top mass or non-top mass.

Original: It is a combination of 4 .
Revised: It is a combination of 4 .

Original: And the interaction may happen between two nearby masses or masses that are two-hop away.
Revised: And the interaction may happen between two nearby masses or masses that are two-hop away.

Original: In total, the number of interactions between different objects is 4 × 2 = 8 .
Revised: In total, the number of interactions between different objects is 4 × 2 = 8 .

Original: In the Soft and the Swim environments, there are three types of quadrilaterals: rigid, soft and actuated.
Revised: In the Soft and the Swim environments, there are three types of quadrilaterals: rigid, soft and actuated.

Original: We have three types of self-interactions correspondingly.
Revised: We have three types of self-interactions correspondingly.

Original: For the interactions between objects, we add edges between two quadrilaterals only if they are connected by a point or edge.
Revised: For the interactions between objects, we add edges between two quadrilaterals only if they are connected by a point or edge.

Original: We are using different relation types if the types of the connected objects are different.
Revised: Connection from different directions are considered as different relations.

Revised: There are 8 different directions, up , down , left , right , up-left , down-left , up-right , down-right .

Revised: The relation types also encode the type of receiver object.

Original: In total, there are 3 × 3 × 2 =types of relations between different objects.
Revised: Thus, in total, there are (8 + 1) × 3 = 27 types of relations between different objects.

Original: B A DDITIONAL E XPERIMENTS
Revised: B A DDITIONAL E XPERIMENTS

Original: Ablation study on the sample efﬁciency.

Original: To estimate our model’s data-efﬁciency, we train our model with smaller training sets.

Original: In Table 2, we report relative errors, which are our model’s simulation errors divided by PN (trained with standard training set)’s error.

Original: As we can see, even if with only 5% of training data, our model still outperforms PN trained on 100% data, which demonstrates that our model is more sample efﬁcient.

Original: [Table]

Original: Comparison with a classical physical simulator optimized using back-box optimization.
Revised: Comparison with a classical physical simulator optimized using back-box optimization.

Original: 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.
Revised: 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.

Original: If we assume that we know the ground truth model, where we only need to identify relevant physical parameters during the system identiﬁcation stage, Bayesian Optimization (Snoek et al., 2012) (BO) can give us a reasonable estimate of the physical parameters.
Revised: If we assume that we know the ground truth model, where we only need to identify relevant physical parameters during the system identiﬁcation stage, Bayesian Optimization (Snoek et al., 2012) (BO) can give us a reasonable estimate of the physical parameters.

Original: However, BO requires much more time to achieve comparable performance with our method in the Rope environment: 0 .
Revised: However, BO requires much more time to achieve comparable performance with our method in the Rope environment: 0 .

Original: 43 vs. 180 seconds averaged over 100 trails (Ours vs. BO).
Revised: 43 vs. 180 seconds averaged over 100 trails (Ours vs. BO).

Original: 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 might not work as well.
Revised: 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 might not work as well.

Original: In our additional experiments, we approximate the Rope environment using a chained spring-mass system, say n masses and n −springs.
Revised: In our additional experiments, we approximate the Rope environment using a chained spring-mass system, say n masses and n −springs.

Original: While taking much more time, BO still cannot give us a satisfying result: simulation error 0 . 046 vs. 0 .
Revised: While taking much more time, BO still cannot give us a satisfying result: simulation error 0 . 046 vs. 0 .

Original: 084 and control error 0 .
Revised: 084 and control error 0 .

Original: 854 vs. 2 . 

Revised: 854 vs. 2 .

Original: (Ours vs. BO).
Revised: (Ours vs. BO).

Revised: [Figure]

Original: Ablation study on the effectiveness of the metric loss.
Revised: Ablation study on the effectiveness of the metric loss.

Original: The internal linear structure allows us to solve the control problem using quadratic programming, where the objective function for control is deﬁned in the embedding space (Section 3.4); hence, it is desirable to have Koopman embeddings that keep the distance in the original state space.
Revised: The internal linear structure allows us to solve the control problem using quadratic programming, where the objective function for control is deﬁned in the embedding space (Section 3.4); hence, it is desirable to have Koopman embeddings that keep the distance in the original state space.

Original: In Section 3.3, we introduce a metric loss to promote learning a Koopman embedding that preserves the distance measurement.
Revised: In Section 3.3, we introduce a metric loss to promote learning a Koopman embedding that preserves the distance measurement.

Original: [Figure]
Revised: [Figure]

Original: r obab ili t y D en [Figure] ean
Revised: 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.

Original: [Figure]

Original: ed

Original: [Figure] [Figure] [Figure] To demonstrate the effect of the metric loss, we compare the models trained with and without the metric loss.
Revised: To demonstrate the effect of the metric loss, we compare the models trained with and without the metric loss.

Original: We establish the comparison using two measurements, the distance preservation and the prediction accuracy .
Revised: We establish the comparison using two measurements, the distance preservation and the prediction accuracy .

Original: 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 (cid:16) (cid:107)g
Revised: 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 (cid:16) (cid:107)

Original: i − g j (cid:107) 2 (cid:107) x
Revised: [Figure] [Equation] [Figure]

Original: i − x j (cid:107) 2 (cid:17) .
Revised: [Equation] [Figure] (cid:17) .

Original: For the model prediction accuracy, we show the simulation errors.
Revised: For the model prediction accuracy, we show the simulation errors.

Original: -We perform the experiments in the Rope environment and show the result in Figure 5.
Revised: We perform the experiments in the Rope environment and show the result in Figure 5.

Original: On the left, we show the ratio of distance in the learned Koopman space and the distance in the original state space.
Revised: On the left, we show the ratio of distance in the learned Koopman space and the distance in the original state space.

Original: The model trained with metric loss has a log distance ratio that significantly more concentrates on 0 .
Revised: The model trained with metric loss has a log distance ratio that signiﬁcantly more concentrates on 0 .

Original: It means the metric loss effectively regularizes the model to preserve the distance.
Revised: It means the metric loss effectively regularizes the model to preserve the distance.

Original: On the right, we show the simulation errors of the two models, which indicate that two models have comparable prediction performance.
Revised: On the right, we show the simulation errors of the two models, which indicate that two models have comparable prediction performance.

Original: Metric loss effectively enhances the property of distance-preserving while not making a big sacriﬁce on the accuracy of the dynamics modeling.
Revised: Metric loss effectively enhances the property of distance-preserving while not making a big sacriﬁce on the accuracy of the dynamics modeling.

Original: Experiments in a known physical parameter setting.
Revised: Experiments in a known physical parameter setting.

Original: 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.
Revised: 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.

Original: Instead, the parameters are embedded in the transition matrices during the system identiﬁcation stage (Section 3.3).
Revised: Instead, the parameters are embedded in the transition matrices during the system identiﬁcation stage (Section 3.3).

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.