From Embedding to Control: Representations for Stochastic Multi-Object Systems
Authors: Xiaoyuan Cheng
Yiming Yang
Wei Jiang
Chenyang Yuan
Zhuo Sun
Yukun Hu
Venue: ICLR 2026 Poster (2026)
Keywords: Embedding to Control
Controllable Embedding
Graph Representation for Linear Control
Primary area: applications to robotics, autonomy, planning
TL;DR: The paper explores to construct efficient representations for stochastic multi-object systems.

Abstract: This paper studies how to achieve accurate modeling and effective control in stochastic nonlinear dynamics with multiple interacting objects. However, non-uniform interactions and random topologies make this task challenging. We address these challenges by proposing Graph Controllable Embeddings (GCE), a general framework to learn stochastic multi-object dynamics for linear control. Specifically, GCE is built on Hilbert space embeddings, allowing direct embedding of probability distributions of controlled stochastic dynamics into a reproducing kernel Hilbert space (RKHS), which enables linear operations in its RKHS while retaining nonlinear expressiveness. We provide theoretical guarantees on the existence, convergence, and applicability of GCE. Notably, a mean field approximation technique is adopted to efficiently capture inter-object dependencies and achieves provably low sample complexity. By integrating graph neural networks, we construct data-dependent kernel features which are capable of adapting to dynamic interaction patterns and generalizing to even unseen topologies with only limited training instances. GCE scales seamlessly to multi-object systems of varying sizes and topologies. Leveraging the linearity of Hilbert spaces, GCE also supports simple yet effective control algorithms for synthesizing optimal sequences. Experiments on physical systems, robotics, and power grids validate GCE and demonstrate consistent performance improvement over various competitive embedding methods in both in-distribution and few-shot tests.

Published as a conference paper at ICLR 2026
FROM EMBEDDING TO CONTROL: REPRESENTATIONS FOR STOCHASTIC MULTI-OBJECT SYSTEMS
Xiaoyuan Cheng Yiming Yang"
Wei Jiang
Chenyang Yuan
Zhuo Sun Yukun Hust
" DDynamic Systems Lab, University College London, UK
"Independent Researcher, Hong Kong, China "Electrical and Electronic Engineering, University of Sheffield, UK "Statisties and Data Science, Shanghai University of Finance and Economies, China
ABSTRACT
This paper studies how to achieve accurate modeling and effective control stochastic nonlinear dynamics with multiple interacting objects. However, non- uniform interactions and random topologies make this task challenging. We ad- dress these challenges by proposing Graph Controllable Embeddings (GCE), a general framework to learn stochastic multi-object dynamics for linear control. Specificallyy, GCE built on Hilbert space embeddings, allowing direct embed- ding of probability distributions of controlled stochastic dynamics into a repro- ducing kernel Hilbert space (RKHS), which enables linear operations its RKHS while retaining nonlinear expressiveness. We provide theoretical guarantees on the existence, convergence, and applicability of GCE. Notably, a mean field ap- proximation technique is adopted to efficiently capture inter-object dependencies and achieve provably low sample complexity. By integrating graph neural net- works, construct data-dependent kernel features that are capable of adapting to dynamic interaction patterns and generalizing to even unseen topologies with only limited training instances. GCE scales seamlessly of varying sizes and topologies. Leveraging the linearity of Hilbert spaces, GCE also supports simple yet effective control algorithms for synthesizing optimal se- physical systems, robotics, and power grids validate GCE and demonstrate consistent performance improvement over various compet- itive embedding methods in both in-distribution and few-shot tests.
multi-object systems
quences. Experiments
INTRODUCTION
Controlling nonlinear dynamics with continuous state and action spaces across multiple interacting
objects is challenging in domains such network systems (Qin et al. 2022). robotics (Yoneda
let al. 2021), and autonomous agent systems (Gelada et al. 2019). Model based control algorithms
offer a promising solution this problem (Jacobson & Mayne 1970 Todorov & Li 2005). typically
approximating the nonlinear control problem through global or local linearization techniques. While
these methods have demonstrated strong performance (Tassa et al. 2007 Levine Koltun 2013), they generally assume access to known system model and a carefully designed, low-dimensional state representation. In practice, however, system dynamics are often unknown and governed by complex interactions among multiple objects, which significantly complicates control design (Bullo al., 2018). To address these challenges, we propose a framework, termed graph controllable embedding (GCE), which learns controllable embeddings of stochastic multi-object dynamics for efficient control. GCE consists of: (1) Modeling, where the embeddings capture system dynamics and interactions directly from different stochastic graph representations; and (2) Control, where simple linear control algorithms are deployed within the learned embedding space.
Email to first author: 4essxxc4@uc..cc
I Email to corresponding author: yukun. hu@ucl auuucclac..u ac..kk
A controllable embedding is representation space where system dynamics modeled to allow linear locally linear control synthesis, following (Banijamali al. 2018)

---

Published as a conference paper at ICLR 2026
Controllable Embedding. A common approach to handling complex, unknown dynamics is to learn a latent space (c.g., smooth manifold or function space) in which the system evolution becomes easier to model and control (Ha & Schmidhuber 2018 Watter et al., 2015). The goal of constructing such latent spaces is to simplity the dynamics into an approximately linear or globally/locally lin- one, enabling the use of linear control methods (Mauroy & Goncalves] 2016 [Banjjamali al., The mainstream for learning globally linear embeddings is Koopman theory (Brunton et al.,
2018).
2021 Koopman 1931 Mezic| 2020). lifts system states into an infinite-dimensional function space where the dynamics evolve linearly (Korda & Mezić, 2018 Mauroy et al.) 2020] Bevanda al., 2021] Cheng al.] 2023). However, applying Koopman-basse based methods to stochastic multi- object systems is challenging due to two key limitations. First, conventional Koopman operators are originally formulated for deterministic dynamics, making their extension to stochastic settings non-trivial (Brunton] (009). Second, most Koopman-based models treat the system as a single en- tity, thereby neglecting the relational topology among interacting objects (Brunton et al. (02I). As result, they often exhibit poor generalization for multi-object environments and require param- eterizations that scale quadratically with the number of objects, increasing the risk of overfitting. Another stream focuses learning locally linear dynamics (Levine et al.) 2020 MMhmmeedi et al.) 2020 Klushyn et al., 2021), 1), often leveraging variational autoencoders (VAE) for constructing low- dimensional manifolds (Farenga al.) 2024 Mudrik al., 2024). While effective for reconstructing high-dimenn high-dimensional observations such as images, they also neglect object-level interactions, making them less suitable for multi-object systems and ultimately limiting their control performance.
Graph Representation. Graph neural networks (GNNs) offer natural framework for modeling interactions among multiple objects. Early workstreams (Battaglia et al., 2016 Chang et 2016
Poli et al. 2021) laid the foundation for data-driven physics simulators by modeling multi- object
dynamics using GNN-based architectures. Building on this, subsequent studies (Li et al. 2018
2019 Sanehez-Gonzalez et al., 2020, Han et al., 2022, Luo et al.) [2023] Poli et 2021) adopted message-passing networks to learn object-centric dynamics over graphs. However, these methods primarily target prediction tasks rather than control. Consequently, the learned embeddings do not have a linear or locally linear dynamics, which can be easy for downstream control. Then, either additional local linearization techniques difficult nonlinear control methods complicates the problem studied in this paper. Another workstream focuses on model- based multi- agent reinforcement learning for multi-object systems (Jiang et al.] 2018 Liu al.) 2020 Zhang ctt al. 2022] Haramati et al., 2024). While these approaches leverage learned dynamics to improve sample efficiency and value function approximation, the resulting dynamics are neither linear nor locally linear in the embedding space, making them unsuitable for synthesizing optimal sequences.
needed, which
A natural idea for the studied problem in this paper to integrate the controllable embedding and graph representation together. For instance, Li et al. (2020) proposed a compositional Koopman method for deterministic multi-object systems. Specifically, GNN is first used to encode each ob- ject's state into latent representation. Then, shared Koopman operator is constructed to linearly evolve all latent states in a common embedding space. While these kinds of methods are empirically effective, several challenges remain. (1) Theoretical gap in stochastic setting: existing methods for stochastic multi-objeet dynamies lack rigorous theoretical foundations with no formal theories to guide the design of controllable embeddings. (2) Unrealistic relational assumptions: they often assume uniform interaction among all neighbors, a simplification that is typically misspecifieed in probabilistic sense; (3) Limited scalability and generalization: they lack thorough evaluation of scalability and generalization to large-scale or random graphs in stochastic multi-object systems.
Motivated by these findings, GCE in Figure 2 is proposed here to address the above challenges. In detail, we study how to embed stochastic nonlinear dynamics into a general reproducing kernel Hilbert space (RKHS) where the multi-object dynamics become linear and linear control methods can be easily implemented. The choice of Hilbert space embeddings is inspired by their compati- bility with two properties underlying controllable embeddings (Song et al., 2009 2010, Fukumizu et al., 2013); Uehara et al. (202):: (1) Modeling: they provide a flexible and expressive representa- tion to capture complex dynamics with multi-object interactions: (2) Control: their inherent linearity supports simple and analytically tractable control policies. Our key contributions are as follows:
(1) Generalized framework: A theoretical framework for learning controllable embeddings in multi-object systems with formal consistency and convergence guarantees. Unlike prior determinis- tic methods, GCE generalizes representation, prediction, and control synthesis within a probabilis- tically interpretable setting; (2) Adaptive interaction modeling: One specific embedding based on
2

---

Published as a conference paper at ICLR 2026
mean field approximation, supported by theoretical guarantees. This breaks the uniform neighbor influence assumption by adaptively capturing non non-uniform interactions and achieves provably low sample complexity. Moreover, further test multiple forms of potential energy functions, demon- strating the flexibility and effectiveness of our framework. (3) Scalability and generalization: GCE naturally scales large systems with random graphs and remains robust under uncertainty and noise, enabling effective control in challenging environments.
PRELIMINARIES
Notation and Setup. Blackboard bold letters (e.g., O, A. HD) denote compact spaces, capital letters (e.g., OL, Ac. H₂) denote random variables, and lowercase letters (c.g., Ot, ae, h₂) denote their realizations. At time t, O2 and A, are the observation and action random variables with realizations Or E O and a₁ A, respectively. The history random variable H, summarizes past information (here H₂ = 1.-1) with realization h4 E H.. The future observation O follows the conditional distribution P(O, At- ae, He, = h₂) abbrevvaate as P(O, ar he)) hereafter-given the executed action or and the history h2. We use Greek letters & and &ú to denote feature maps into RKHSs, with do. ΦA. and ou mapping into Ho. Hy. and Hy. respectively; v/P == Φn (O.) and ot == do represent the mapped random variable and its realization, respectively. The symbols &b, OD, and represent the tensor product, the direct sum, and the element-wise product, respectively. To clarify indexing, subscripts indicate time steps, while superscripts indicate object indices. For example, of denotes the observation of the i-th object at time step t. Finally, we define the shorthand [M] ;= (1....,N N} for the consecutive integer set. please see more notations in Tables 6 and 7 in Appendix A
Definition (Hilbert Space Embedding of Conditional Distributions (Sriperumbudur et al. 2010 For the conditional distribution P(O, a. he) of the future observation random variable O2 given the action at and history hts its Hilbert space embedding is defined as the conditional expectation of the feature map in a RKHS Ho:
o4, ht]] = { vi P(o. ae, ha)dor = Co AH [1t @ vg], 00
E E(p0
where wp ΦE (he) and us = valae) are the characteristic kernel features in their respective RKHSs Hz and HAi COJAH Hy @ HA + Ho is the conditional embedding operator.
This embedding yields a nonparametric repre- sentation of the conditional distribution in the RKHS (see more details about RHKS in Ap- pendix C. thereby avoiding density estimation (see formal analysis in Appendix and D.2). With characteristic feature maps, the embed- ding uniquely specifies the conditional distribu- tion (Sriperunbudur et al. 2010). Intuitively, the tensor product J.. & th. represents the joint features of the conditioning variables (as, ht), allowing the operator CoJAh to act linearly the joint input space and represent the con- ditional expectation in closed form. In prac- tice, CojAn enables efficient one-step predic- tion from history and action features, and can be applied recursively to perform multi-step rollouts and facilitate control planning (shown in Figure
Pot Aval
Hiour Fauce
1+2
Figure I: Hilbert space embedding of conditional distributions. Left: Stochastic nonlinear dynamics as evolving conditional distributions P(O,la,, hy) (red curves: probability distributions, black dots: realizations). Right: After embedding into RKHS, dynamics become linear under Co (red dots: expectations).
Problem Formulation. We consider a stochastic dynamics composed of N interacting objects, represented as graph G, = (Vr, at time l. Fach one of E Vo == [ corresponds to the observation of the i-th object, and at denotes the control action applicd to it. The full set of actions at time is denoted by V = { The relational structure is encoded by fixed binary adjacency matrix& & = [ Fig=i. where ei = 1 indicates an edge from object i to j. and ei 0 otherwise.
For each object i, we denote its (inclusive) neighborhood by E(i = Eilehs 1) U {i). Our goal
3

---

Published as a conference paper at ICLR 2026
t-slep observation
embedding r-step features
model and control
(t++ 1)-slep observation
Hibbrt space embedding via GNN
approximation
pulbbck
distributions
RXCHS features
disentiongeed hiistory and action featuras
unigue distribtioon ambardiings characteridie fooduana
Figure 2: An example illustration of the GCE framework. The robot with multiple interconnected objects is initialized on top of a thin pole. At time step t. the observation embedded into an RKHS via characteristic feature maps. A mean field approximation is then applied: the history features are computed via an element-wise product to reduce computation, while the action features preserve their structure for optimization (see math in Equation 9 The predicted features mapped back to predict the (t + 1)-step observation, enabling the robot to stabilize on the pole.
to construct controllable embedding to model the stochastic multi-object dynamics governed by the unknown transition function F:
+++1 F(V?, V7,
where FF captures the time evolution of object observations based on actions and interactions. Af- ter system identification, the second goal is enable effective control in the learned embedding space. Specifically, we seek representation in which the dynamics are amenable to linear control that minimizes cumulative cost over M steps: algorithms. The objective is to find actions (Ve )
M J(V Vi), MM
min
s.t. V4+1 = F(VE, V2,
(3)
where J 1 Va R cost function defined over graph representations. As shown in Figure Hilbert space embeddings offer general mechanism to represent stochastic dynamics in form amenable to linear control. This motivates the extension of Hilbert space embeddings to GCE, which captures stochastic multi-object dynamics while enabling efficient linear control. In the next section, we detail how to construct such embeddings to jointly support the modeling objective in Equation] 2 and the control objective in Equation 3
FRAMEWORK OF GRAPH CONTROLLABLE EMBEDDINGS
In this section, we first establish the theoretical foundation of GCE. Then, we propose mean field- based instantiation for capturing non-uniform interaction weights among multiple objects. Finally, compare different embeddings within the GCE framework, analyzing their properties across various settings and deriving sample complexity guarantees.
3.1 FOUNDATION OF GCE: EMBEDDING AND CONSISTENCY
We formalize the Hilbert space embedding of stochastic multi-object system as an instantiation of GCE. Using the feature maps in Definition let s o and u5ft denote the realization features of the i-th object's history, future observation and action, respectively. For each pair of connected objects i and j, we define conditional embedding operator Co (oil HJ that models the contribution of object j's history and action to the conditional expectation of object i's future observation:
Elv O...
af, hR] = ConAAh hi @ 1
Vi,j E [N],
(4)
where wi w.i. indicates the component of the future observation feature of object attributable to object j's influence. The overall conditional expectation of object i's observation feature is:
Ely 0.e (a). M
Σ JEE(()
Co @ v..
(5)
where Cо.AА'И: is the zero operator whenever j E(t). Each Co A implicitly encodes the relative importance of the j-th object's influence on the i-th object, as these effects are generally not equally distributed. More probabilistic analysis and formal derivation referred Appendix D.3
4

---

Published as a conference paper at ICLR 2026
Theorem 1 (Consistency and Convergence of GCE). Let COMAAS be the empirical estima- tor of the embedding operator Co constructed from T i.i.d. samples using character- istic kernel features. For any input set {at, bz define the estimated embedding of ob- ject i's future observation as: Do. EjeE()) CoyAsns ( @ v. Then, as ' + sc. + O. Since the kernel feature is characteristic, this convergence implies that the conditional distribution P(O (at, (()) induced by CoqAAH also converges: P(O) {aL, 1 I-0,, P(O) {at, (o].
10 - E [ (( Ne hii }3=1 -
Theorem shows that, given sufficient samples, the conditional embedding operators in GCE exist and converge consistently, ensuring probabilistic consistency of the induced conditional distributions (see proof in Appendix[F]] [..I). Although gqquation[ppoviiieee of stochastic multi-object dynamics, its direct application for control poses two difficulties. First, estimating Co is computationally intensive due to the high dimensionality introduced by the
a principled formulation for embeddings
tensor product 2/. @ d which in turn increases the required sample size. Second, the entan- gled representation of actions and histories in the tensor space makes sequential action optimization intractable. Inspired by the decomposition of joint distributions in exponential families (Zhang et al. 2008 Altun et al. 2012), we address these difficulties by disentangling the tensorized feature v/
& Vi
into simplified form:
Coq [Vt @ of
Coopar,Coaaallll [[ Coqnuet" +Cousus Vi.j E [N],
(6)
where [- || represents the concatenation operation, COrJHS and Coijas denote neighbor-specife
linear operator that captures the distinct influence of object history and action on object i, respec- tively. this reformulation, the tensor map in Equation i replaced by concatenation ( || v// This concatenation approximates linear relationships in the probability distributions while neglect- ing higher-order interactions the original tensor product. Consequently, the computational com- plexity is significantly reduced, and the action and history representations are disentangled, making sequential action optimization tractable. See more formal analysis in Appendix D.4
3.2 EFFICIENT EMBEDDING: ADAPTIVE MEAN FIELD APPROXIMATION
Although valid controllable embeddings have been formally defined in Equation 6 estimating all linear conditional operators CoIJA>H; remains computationally prohibitive. The number of such operators grows with the number of edges and can scale quadratically with the number of objects in dense graphs. i.e., O(N2) in the worst case. However, modeling all pair-wise interactions is often unnecessary. To mitigate this issue, we leverage mean field approximation, which significantly reduces the sample complexity by summarizing the collective influence of neighbors.
Probabilistically, the mean field approximation avoids dealing with the full joint conditional distri- bution over all pair-wise vistery-observation interactions. It replaces this with aggregated form that captures the combined influence of neighbors, while keeping their contributions unequal. A probabilistic derivation of this factorization, including its separation into history and action terms, is provided in Appendix D.5 In such case, the interaction weight of node interacting with neighbor j at time / approximated
exp(f(W)) , ( Σ EkEero exp(f(()) (il wh
1.j
04
where o,3 [0, 1] denotes the Boltzmann-Gibee weight, f is pair-wise negative potential energy function, and the denominator Σtee(e) exp(/ +/(+))) serves as the partition function, normal- izing the probability one. ff can be either predefined or learned via neural networks. Based on the mean field approximation, the expected feature of object can be computed as the weighted aggrega- tion from its neighbors: Ejerto 0, v. Accordingly, Ejee(e) Co w
h.j can be approximated
5

---

Published as a conference paper at ICLR 2026
by applying shared conditional operator to this aggregated feature:
Σ JEE()
Commi bd Coijar
Σ JEE(1)
where CO-HH is shared approximation operator derived under the homogeneity setting to fur- thermore reduce computational complexity. This formulation captures the essence of the mean field approximation by replacing heterogeneous, neighbor-speciffc interactions with unified, averaged effect, to reduce per-object computation to constant time, independent of the neighbor number.
Unlike H we keep Co. unchanged facilitate the optimization of sequential actions. Finally, the overall conditional expectation of object i's observation feature in Equation 5 changes to:
E h2 ) Co.1 (
Σ JEE()
at
+
Σ Co JEE())
Notably, the cotisertooooss block is often sparse in real-world scenarios. Thus, the com- plexity is dominated by the history-to-observation block whose computational complexity can be reduced by making all conditional linear operators Co-H share the same parameters, leveraging the homogeneous nature of the relational structure (since the transition probability of each object is the same). Under this representation, the total computation time in the history part can be reduced to O(N).
3.3 COMPARISON DIFFERENT EMBEDDINGS: PROPERTIES AND SAMPLE COMPLEXITY
To generalize our framework, we compare four embedding formulations existing in GCE, each for- mally defined in the main text and derived in detail from the probabilistic view in the appendices: (1) Tensor form Tensor, Equation 5 Appendix D.3): full tensor- tensor-product embedding without control- lable structure constraints. This form is probabilistically consistent (see Theorem TT. as it retains all high-order interaction terms; (2) Dense form (Dense, Equation Appendix (D.]): controllable em- bedding without homogeneity or mean field approximation. It remains consistent for all pair-wise terms but ignores higher-order interactions in Hilbert space embedding; (3) Homogeneous form (Hom; cf. Li al. (2020): assumes graph homogeneity and enforces uniform neighbor weights, leading to misspecified product form from probabilistic viewpoint; (4) Homogemeous form with mean field approximation (Hom + Mean, Equation 9 Appendix D.3): leverages homogeneity and adaptive neighbor weighting via Boltzmann- imation to the true joint distribution over the entire graph, retaining unequal influence from each neighbor while reducing complexity. Table sumaarizes their properties. including satisfaction of controllable embedding, applicable scenarios, sample complexity, computation time, ability to in- corporate adaptive weights, and generalizability to random graphs. Overall, Hom + Mean strikes a balance between sample efficiency and expressiveness by approximating the probabilistic structure with adaptive weights, while avoiding the misspecification in Hom and the combinatorial complex- ity in Tensor and Dense. The complete quantitative analysis is given in Appendix E with proved sample complexity and generalization on graphs in Appendix
weights. The product form is an approx-
Table 1: Comparison of four different embeddings in GCE.
Property
Tensor
Dense
Hom
Hom + Mean
Controllable embedding
X
Scenario
heterngeneity
heterogeneity high high
homogeneity medium high
homogeneity low
Sample Complexity
very high very high / (implicit) hard
Computation Time
low
Adaptive Weight
(implicit) hard
X(misspacifiec) medium
((explicit)
Generalize to Random Graphs
casy
From a sample complexity perspective, Hom + Mcan enjoys two key advantages: the mean field approximation replaces O(N2) neighbor-specific operators with single shared operator, greatly
II our setting, "homogeneity" means all nodes share the same operator mapping from history features to observation features, while neighbor influencing weights vary. The shared operator represents one universal law of interaction. The local interaction rule between neighboring objects is the same, e.g., in Rope environment in Figure ??. Since the embedding feature for a "top mass" and "regular mass" are different, applying the same operator to them will naturally produce different dynamic evolutions.

---

Published as a conference paper at ICLR 2026
reducing the number of parameters estimate, while adaptive weights preserve variability across neighbors without explicit pair-wise modeling. These properties enable efficient learning and few- shot adaptation to new random graphs, unlike Tensor and Dense (poor transfer) or Hom (misspecified weights). Using random matrix theory, we show that the required sample size of Hom + Mean scales with an effective operator dimension rather than the total number of graph edges, yielding much smaller requirements for training data than pair-wise conditional embedding operators on graphs (see Theorems a b in Appendix F2
4 IMPLEMENTATION
As shown in Figure the implimentation of GCE is as follows: (1) encoder that maps observa- tions into RKHSs, (2) various ways for estimating adaptive Boltzmann- Gibbss weights, (3) tailored loss functions for controllable embedding, and (4) linear quadratic regulator (LQR) controller in
the RKHS. Pseudo code is given in Appendix H
Mapping to RKHSs. A message passing GNN encoder projects the i-th object's history and future observations into RKHSS whr E Rdd w/t E Rd. The observation feature is wh fv (D fe (o) of ed)), where fv and fe are neural networks extracting deep kernel fea- tures; and action of uses a fixed linear projection to its feature v/2 IRd for easier optimization (see Appendix H).
Boltzman--Gibbs Weights. To com- pute adaptive interaction weights, we em- ploy Nadaraya-Watson kernel estima- tor, which ensures a stable approxima- tion of Boltzmann-Gibbs weights. For each object i, the weight given by a.j G(v,, //)/ Exeet)) G(vi, v where G(+,-) is an exponential kernelized energy function of the form G(v vz))
Table 2: Parameterizations of the potential function ((,)) in Equation 7
(a, X, k) are hyperparameters.
Type
Name
Potential
1 of
Geussian
Kernel bused
Laplace
w/ )T( 10y 11 ne
vMF
We evaluate both kernel-based and neural parameterizations of as shown in Table 2 please refer more de-
exp (( ....)))
Neural-bart based
MLP
tails to Appendix D.5).
Loss Functions.
The GNN block is then utilized to encode this feature in RKHS, the forward loss Liwd in the embedding feature space can be represented as
1
Coyl
Σ
wh at
COMAAN COMAAN
E
C
Lib
CON
NJ at
CONIA'
Liba
mS
(10)
where μ is sampled data distribution of realizations [c)0 N and ||- denotes the Hilbert- Schmidt norm. All conditional operators Čо* share the same learnable parameters. To improve latent control quality, we can achieve M-step open-loop control by optimizing M-step sequential prediction loss. This is implemented by iteratively generating -tsppep sequences in the feature space through autoregressive prediction. After learning the observation features, we utilize the same message passing block as the decoder. The loss function of reconstruction becomes
M
Lroc
E,
...., (10O1....,00NN))))
(o].
NN I
(11)
where t indicates the decoder function, which pulls estimated features 3.00 back to the observation space. Our total loss function has two parts, the forward and reconstruction loss, and we train the encoder-decoder GNN blocks and linear conditional operators based on Equations 10 and 11

---

Published as a conference paper at ICLR 2026
LQR Control. As presented in Equation 3 we directly solve quadratic cost function in feature space over M steps:
M
min
E
10.1
,0.., - -V ".. "2 11, + 118 ..... 1t 1T1..
(12)
(V2) LA
1=1
where Q1 and Q2 are two pre-defined positive definite matrices to measure quadratic costs, and is pre-defined target observation feature which should be achieved in M steps. We minimize the control cost function Equation 12 by using convex optimization.
(1s :.... 2he "
NUMERICAL EXPERIMENTS
We evaluate method four control environments: (1) Rope- chain of masses with the top mass controlled horizontally; (2) Soft a soft robot made interconnected objects; (3) Swim - the soft robot swims freely in fluid (Li et al. 2020); and (4) Power-Grid
high-dimensional system with random topologies aiming to stabilize node voltages. Robustness is tested with additive white noise 2%, 5%, 10%, and 20% of observation variance (see more details in Appendix
Metrics and Baselines. We use open-loop control with 40-step LQR in Rope, while Soft and Swim generate 64-step control with feedback from step 32. Power-Grid uses a 100-step horizon with feedback at step 50. Per- formance is evaluated via control cost (Equation 12 and control error [ / averaged over 200 runs. Few-shot validation on unseen graphs tests the general- ization of the algorithm. The baselines are controllable embedding methods without relational structures, includ- ing VAE (Banjjamali et al. 2018) and Prediction, Consis- tency, Curvature (PCC) (Levine al. 22200,, and graph representation methods, including Koopman Polynomial Model (KPM) (Li et al. 2020). Compositional Koopman Operator (CKO) (Li et al. (020). and GraphODE (Luo al. 2023), where CKO represents the current state-of- the-art in graph embedding control tasks (see control an- imations Appendix [.7).
Rope
36f
Figure 3: The 3D demonstration of con- trol tasks.
Table 3: Performance comparison for in-distribution and few-shot validation. Control cost and error
are shown as mean + standard deviation, with the best and second best results highlighted in green and blue, respectively. Due page limit, for additional Rope and Soft comparison, see Table
In-Distribution Validation
Few-Shor Validation
Methods VAE PCC GrahhODE KPM
Environments
Control cost
Control error
Control cost
Control error
$73.1 z 108.7
0.73 a 0.19
835.4 113.2
0.92 2 0.15
513.3 = 92.5
0.68 2 0.15
732.8 94.5
0.102
87.9
0.52
693.5
0.58
385.5 + 75.2
0.44 = 0.06
523.4 22.8
0.61
0.42
421.0
Ours (vMF)
392.7 = 73.1
0.45 = 0.09
452.3
0.13
Ours (Laplace)
403. 68.3
0.46 0.13
435.7
0.10
Ours (Gusssian)
383.7 + 77.8
0.41 = 0.08
404.3 74,2
0.41 = 0.09
Table 4: Evaluation under varying noise levels in Power-Grid on random graphs with 100-150 ob-
jects. "NaN" means the unstable control. For additional control cost comparison, see Table 10
Method
Noiseless
2%
5%
10%
20%
GraphODE
0.58 0.044
0.073
NaN NaN
N&N NaN
NaN NaN
Control Ervor
KPM
0.42 + 0.008
0.50 + 0.022
CKO
0.47 + 0.031
0.48 0.034
0.51 0.027
0.65 + 0.051 0.63 a 0.037
0.85 0.055
Ours (Gaussian)
0.21 a 0.005
0.27 2 0.018
0.024
0.83 0.M1
5.1 RESULT ANALYSIS
(1) Is it necessary to design specific controllable embeddings for multi-object dynamics? Yes. In Ta- ble general-purposs, controllable embeddings such as VAE and PCC, while capable of capturing
8

---

Published as a conference paper at ICLR 2026
individual trajectories, fail to provide structured features for control. Although GraphODE incorpo- rates relational structures, it lacks an explicit controllable design and depends on local linearization via auttodifferentiation,, resulting in suboptimal performance. This highlights that merely encoding relational dependencies is not enough; the embedding must also be tailored to control objectives.
(2) How well do the empirical findings validate the theoretical predictions of our framework? We divide the answer into four points. (a) Analysis of baseline methods. CKO is an instance of the Hom subclass in our framework, with misspecfied uniform weighting. As shown in Figure 9 (Ap- pendix I.5), the empirical prediction error accumulates faster than our methods in both deterministic and stochastic environments. As result, CKO performance is degraded across all tasks as shown in Table 3 (b) Failure cases and theoretical justification. KPM fails under high-noise conditions (see Table This is primarily because Hilbert space embeddings require characteristic features to uniquely represent probability distributions. The polynomial features in KPM are not characteristic, and therefore cannot faithfully embed the underlying distributions (see Theorem I (c) Analysis sample complexity of different embeddings. To validate the properties summarized in Table [ we compared: Dense, Hom, Hom + Mean and excluded Tensor due to its lack of both local and global linearity. In few-shot settings, we varied the number of training trajectories from 1 to 32 to assess sample efficiency and expressiveness. As shown in Table [S] Hom + Mean consis- tently outperforms Dense and Hom with limited samples, owing to the introduction of mean field approximation. As sample size increases, all methods improve due to reduced operator error, align- ing with Theorems a While Dense eventually has similar control error with Hom Mean, Hom + Mean always maintains lower control costs. (d) Gaussian Kernel leads to stable mean field approximation. Among the three variants for estimating the energy function f. Gaussian consistently outperforms vMF and Laplace as shown in Table 3 Its advantage lies in providing a smoother and more stable mean field approximation, which better captures both central tenden- cies and uncertainty in RKHSs. In contrast, vMF emphasizes directional alignment, and Laplace encourages more uniform energy with slow decay rates, making them less robust across diverse en- vironments. We also experimented with neural network, but it was unstable for the energy function f in RKHSs, further underscoring the benefit of analytically tractable kernels such as Gaussian.
(3) To what extent does the non-uniform weighting from the mean field approximation enhance control performance? As shown in Table our method consistently outperforms baselines across all three tasks in terms of both control cost error. While CKO achieves comparable performance to ours on in-distribution validation for Soft and Swim tasks, it shows a larger performance gap during few- shot generalization. This discrepancy stems from the limitations of misspecified representation (see Figure 12]. which reduces model expressiveness. To further validate this observation, we compare the 100-step prediction errors and sample complexity of CKO and our method in Figure and I0 (Appendix 1.5p. The results validate that our method yields lower error and sample complexity.
(4) Is the proposed framework capable of generalizing to large-scale multi-object systems with ran- dom topologies in noisy environments? Yes. As shown in Table E Power-Grid presents a more challenging setting due its scale, graph complexity, and random topology. GraphODE and KPM fail consistently above 2% noise, producing unstable (NaN) results. At low noise levels, our method significantly outperforms CKO in both control error and cost. While the performance gap narrows with increasing noise, our approach remains more robust and generalizes better across conditions.
Table 5: Control error under different numbers of training trajectories used for few-shot adaptation in the Rope environment. The table reports performance using varying numbers of demonstration trajectories, referred to as "fitting number". For additional control cost comparison, see Table
Method
8
16
32
Dense
0.79 + 0.25
0.41 = 0.11
0.36 + 0.12
0.28 + 0.08
0.26 + 0.10
Contoo Emor Hom
0.74 + 0.25
0.32 0.08
0.30 + 0.06
0.30 :: 0.06
0.30 0.06
Hom+ + Mean
0.51 = 0.12
0.29 = 0.09
0.26 = 0.08
0.25 + 0.08
0.23 + 0.09
Ablation study on Gaussian Kernel. The Gaussian kernel leads to the best performance across all tasks. The Gaussian kernel bandwidth a in Table can be interpreted as a temperature factor that determines how local features are mixed with those of their neighbors. To investigate its impact, we evaluated {0.1, 1, 2, 4, 5, 10) during shot validation in Rope and Soft environments. Due to the quadratic effect of the bandwidth, when a 2 10, the attention weights become nearly uniform across neighboring nodes, which is the reason we do not evaluate with value larger than

---

Published as a conference paper at ICLR 2026
10. We found in Figure that achieves the best performance in both control error and control cost. This can be interpreted as physical temperatures: when o is small (Temp + 0), node dynamics nearly self-determined due to strong penalties for deviation. When is large (Temp oo), all neighbors exert nearly equal influence, reflecting maximal thermal mixing. An ablation on feature dimension is provided in Appendix L.6 with dimension 32 achieving the best performance.
Control
Rope
Control Rope
SSoft
0.4-
100
1500
100
1250
Eror
1000
0.2
160
0.1
0.2
25
7.5
10.0
2.5
7.5
2.5
75
16.0
2.5
5.0
7.5
10.0
Bandwidth
Bandwidth
Bandwidth
Bandwidth
Figure 4: Control error and control cost for varying bandwidth a values in Rope and Soft environ- ments. Results are reported mean and standard deviation over multiple runs.
6 CONCLUSION
We proposed Graph Controllable Embeddings (GCE), a general framework for modeling stochas- tic multi-object dynamics and synthesizing effective control sequences. By introducing the mean field approximation, GCE efficiently captures inter-object dependencies and achieves provably low sample complexity. Leveraging graph neural networks, our method adapts to dynamic interaction patterns and generalizes unseen topologies with limited training data. A limitation of our frame- work is that it currently focuses on pair-wise relations; extending it to richer relational structures, such hypergraphs and incorporating attention- techniques for modeling interaction weights remains unexplored.
ACKNOWLEDGEMENT
The corresponding author would like to thank the financial support provided by the Royal Academy of Engineering through the Industrial Fellowship scheme (IF-2425-19-A1155)
ETHICS STATEMENT AND REPRODUCIBILITY STATEMENT
This work raises specific ethical concerns beyond standard practices in machine learning research. All methods, datasets, and hyperparameters are described in detail, and the core code released in the supplementary materials.
REFERENCES
Yasemin Altun, Alex Smola, and Thomas Hofmann. Exponential families for conditional random fields. arXiv preprint arXiv: 1XX.::200.443.. 2012.
Nachman Aronszajn. Theory of reproducing kernels. Transactions of the American mathematical society, 68(3):337-404 1950.
Ershad Banijamali, Rui Shu, Hung Bui, Ali Ghodsi, et al. Robust locally-linear controllable embed- ding. In International Conference on Artificial Intelligence and Statistics, pp. 1751-1759, PMLR, 2018.
Peter Battaglia, Razvan Pascanu, Matthew Lai, Danilo Jimenez Rezende, et al. Interaction networks for learning about objects, relations and physics. Advances in neural information processing systems, 29, 2016.
Petar Bevanda, Stefan Sosnowski, and Sandra Hirche. Koopman operator dynamical models: Learn- ing, analysis and control. Annual Reviews in Control, 52:197-212, 2021.
10

---

Published as a conference paper at ICLR 2026
Steven L Brunton. Notes koopman operator theory. Universität von Washington, Department of Mechanical Engineering. Zugriff. 30, 2019.
Steven L Brunton, Marko Budišić, Eurika Kaiser, and Nathan Kutz. Modern koopman theory for dynamical systems. arXiv preprint arXiv:2102.12088 2021.
Francesco Bullo, Jorge Cortés, Florian Dörfter, and Sonia Martínez. Lectures on network systems, volume 1. CreateSpace PlatformScotts Valley, 2018.
Michael B Chang, Tomer Ullman, Antonio Torralba, and Joshua B Tenenbaum. A compositional object-based approach to learning physical dynamics. arXiv preprint arXiv: 2016.
Xiaoyuan Cheng, Yiming Yang, Wei Jiang, and Yukun Hu. KEEC: Embed to control equiv- ariant geometry. arXiv preprint arXiv:2312.01542 2023.
Chun Yuan Deng. A generalization of the sherman--morrisnn-wwoadyy formula. Applied Mathe- matics Letters, 24(9):161-1156 2011.
Nicola Farenga, Stefania Fresca, Simone Brivio, and Andrea Manzoni. On latent dynamics learning in nonlinear reduced order modeling. arXiv preprini arXiv:2408. 2024.
Kenji Fukumizu, Le Song, and Arthur Gretton. Kernel bayes' rule. Advances in neural information processing systems, 24, 2011.
Kenji Fukumizu, Song, and Arthur Gretton. Kernel bayes' rule: Bayesian inference with positive definite kernels. The Journal of Machine Learning Research, 14(1):3753-3783, 2013.
Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Mare G Bellemare. Deepmdp: Learning continuous latent space models for representation learning. International conference machine learning, pp. 2170-2179 PMLR, 2019.
David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv: 100. 2018.
Jiaqi Han, Wenbing Huang. Hengbo Ma, Jiachen Li, Josh Tenenbaum, and Chuang Gan. Learning physical dynamies with subequivariant graph neural networks. Advances in Neural Information Processing Systems, 35:26256-26266 2022.
Dan Haramati, Tal Daniel, and Aviv Tamar. Entity-centric reinforcement learning for object manip- ulation from pixels. arXiv preprint arXiv:2404.01220 2024.
Thomas Hofmann, Bernhard Schölkopf, and Alexander Smola. A tutorial review of rkhs methods in machine learning. Technical Report, 2005.
David H Jacobson and David Q Mayne. Differential dynamic programming. Elsevier, 1970.
Jiechuan Jiang, Chen Dun, Tiejun Huang, and Zongqing Lu. Graph convolutional reinforcement learning, arXiv preprint arXiv:1810.09200. 2018
Tosio Kato. Perturbation theory for linear operators, volume 132. Springer Science Business Media, 2013.
Alexej Klushyn, Richard Kurle, Maximilian Soelch, Botond Cseke, and Patrick der Smagt. Latent matters: Learning deep state-space models. Advances in Neural Information Processing Systems, 34:10234-10255, 2021.
Bernard O Koopman. Hamiltonian systems and transformation in hilbert space. Proceedings of the National Academy of Sciences, 1755)331--118, 1931.
Milan Korda and Igor Mezić. Linear predictors for nonlinear dynamical systems: Koopman operator meets model predictive control. Automatica, 93:149-160 2018.
Nir Levine, Yinlam Chow, Rui Shu, Ang Li. Mohammad Ghavamzadeh, and Hung Bui. Predic- tion, consistency, curvature: Representation learning for locally-liness control. arXiv preprint arXiv: 109.001900 2020.

---

Published as a conference paper at ICLR 2026
Sergey Levine and Vladlen Koltun. Variational policy search via trajectory optimization. Advances in neural information processing systems, 26, 2013.
Yunzhu Li, Jiajun Wu, Russ Tedrake, Joshua B Tenenbaum, and Antonio Torralba. Learning par- ticle dynamics for manipulating rigid bodies, deformable objects, and fluids. arXiv preprint arXiv: 2018.
Yunzhu Li, Jiajun Wu, Jun-Yan Zhu, Joshua B Tenenbaum, Antonio Torralba, and Russ Tedrake. Propagation networks for model-based control under partial observation. 2019 International Conference on Robotics and Automation (ICRA), pp. 1205-1211 IEEE, 2019.
Yunzhu Li, Hao He, Jiajun Wu, Dina Katabi, and Antonio Torralba. Learning compositional koop- man operators for model-based control. arXiv preprint arXiv: 19119998002.4.. 2020.
Yong Liu, Weixun Wang, Yujing Hu, Jianye Hao, Xingguo Chen, and Yang Gao. Multi-agent game abstraction via graph attention neural network. In Proceedings of the AAAI conference on artificial intelligence, volume 34, pp. 721 7211-7218, 2020.
Xiao Luo, Jingyang Yuan, Zijie Huang, Huiyu Jiang, Yifang Qin, Wei Ju, Ming Zhang, and Yizhou Sun. Hope: High-order graph ode for modeling interacting dynamics. International conference on machine learning, pp. 23124 23139... PMLR, 2023.
Alexandre Mauroy and Jorge Goncalves. Linear identification of nonlinear systems: A lifting tech- nique based on the koopman operator. In 2016 IEEE 55th Conference on Decision and Control (CDC), pp. 6500-6505. IEEE, 2016.
Alexandre Mauroy, Y Susuki, and Igor Mezic. Koopman operator in systems and control. Springer, 2020.
Igor Mezic. Koopman operator, geometry, and learning, arXiv preprint arXiv:2010.05377, 2020.
Zakaria Mhammedi, Dylan Foster, Max Simchowitz, Dipendra Misra, Wen Sun, Akshay Krish- namurthy, Alexander Rakhlin, and John Langford. Learning the linear quadratic regulator from nonlinear observations. Advances in Neural Information Processing Systems, 33:14532-14543, 2020.
Krikamol Muandet, Kenji Fukumizu, Bharath Sriperumbudur, Bernhard Schölkopf. et al. Kernel mean embedding of distributions: A review and beyond. Foundations and Trends@ in Machine Learning, 10(1-2):1-141, 2017.
Noga Mudrik, Yenho Chen, Eva Yezerets, Christopher Rozell, and Adam S Charles. Decomposed linear dynamical systems (dlds) for learning the latent components of neural dynamics. Journal of Machine Learning Research, 25(59):1 2024.
Michael Poli, Stefano Massaroli, Junyoung Park, Atsushi Yamashita, Hajime Asama, and Jinkyoo Park. Graph neural ordinary differential equations. arXiv preprint urXiv:/911.0755: 2021.
Zengyi Qin, Dawei Sun, and Chuchu Fan. Sablas: Learning safe control for black-box dynamical systems. IEEE Robotics and Automation Letters, 00011928-1933 2022.
Alvaro Sanchez-Gonzalez, Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovee, and Peter
Battaglia. Learning to simulate complex physics with graph networks. In International conference on machine learning, pp. 8459-8468. PMLR, 2020.
Le Song, Jonathan Huang, Alex Smola, and Kenji Fukumizu. Hilbert space embeddings of condi- tional distributions with applications to dynamical systems. In Proceedings of the 26th Annual International Conference on Machine Learning, pp. 9611968. 2009.
Le Song, Byron Boots, Sajid Siddiqi, Geoffrey Gordon, and Alex Smola. Hilbert space embed- dings of hidden markov models. 2010.
Bharath K Sriperumbudur, Arthur Gretton, Kenji Fukumizu, Bernhard Schölkopf, and Gert RG Lanckriet. Hilbert space embeddings and metrics on probability measures. The Journal of Ma- chine Learning Research, 11:1517-1561, 2010.
12

---

Published as a conference paper at ICLR 2026
Bharath K Sriperumbudur, Kenji Fukumizu, and Gert RG Lanckriet. Universality, characteristic kemnlss and rkhs embedding of measures. Journal of Machine Learning Research, 12(7), 2011.
Terence Tao. Topics in random matrix theory, volume 132. American Mathematical Soc., 2012.
Yuval Tassa, Tom Erez, and William Smart. Receding horizon differential dynamic programming. Advances in neural information processing systems, 20, 2007.
Emanuel Todorov and Weiwei A generalized iterative lag method locally-optimal feedback control of constrained nonlinear stochastic systems. In Proceedings of the 2005, American Control Conference, 2005., pp. 300-306. IEEE, 2005.
Joel A Tropp et al. An introduction to matrix concentration inequalities. Foundations and Trends@ in Machine Learning, 8(1-2):1-230, 2015.
Masatoshi Uehara, Ayush Sekhari, Jason Lee, Nathan Kallus, and Wen Sun. Provably efficient inforcement learning in partially observable dynamical systems. Advances in Neural Information Processing Systems, 35:578-592, 2022.
Ramon Van Handel. Probability high dimension. Lecture Notes (Princeton University), 2(3):2-3, 2014.
Roman Vershynin. High-dimensional probability: An introduction with applications in data science, volume 47. Cambridge university press, 2018.
Manuel Watter, Jost Springenberg, Joschka Boedecker, and Martin Riedmiller. Embed to control: A locally linear latent dynamics model for control from raw images. Advances neural information processing systems, 28, 2015.
Takuma Yoneda, Ge Yang. Matthew R Walter, and Bradly Stadie. Invariance through latent align- ment. arXiv preprint arXiv:2112.08526 2021.
Xinhua Zhang, Le Song, Arthur Gretton, and Alex Smola. Kernel measures of independence for non-iid data. Advances neural information processing systems, 21, 2008
Zhi Zhang, Zhuoran Yang, Han Liu, Pratap Tokekar, and Furong Huang. Reinforcement learning under a multi-agent predictive state representation model: Method and theory. In The Tenth International Conference on Learning Representations (ICLR 2022), 2022.
13

---

Published as a conference paper at ICLR 2026
THE USE OF LARGE LANGUAGE MODELS (LLMS)
During the preparation of this manuscript, the authors used ChatGPT to polish the writing (e.g.. improving grammar, readability, and clarity), The content, technical contributions, and conclusions of the paper were developed entirely by the authors, who take full responsibility for all ideas and results presented.
A NOTATIONS
Table 6: Basic Notations of Hilbert Space Embedding
observation
action
history
compact spaces
O O₂ Ot (.,00))) Coo
A At as ka(,,a))
H He
random variables
realizations
he k (t,he) Cни
kernels
covariance operators
CAA
feature maps of realizations
1. (ot)
1 = (at)
w/t = su (hr)
feature maps of random variables RKHS
4 = (O,)
4/ = On (or)
v/ = du (le)
Ho
HA
Hy
Table 7: Notations in the Main Text
Notations
Mcaning
Cojah ComAs' Cooa Comn &
general embedding operator
embedding operator that models object j's effect to i's observation
the disentangled embedding operators
shared embedding operator binary adjacency matrix
E() F
(inclusive) neighborhood
unknown transition function for multi-object dynamics
ff
pair-wise negative potential energy
fe
neural networks in message passing GNN kernelized energy function
G
J
cost function defined over graph representations
A, = O2; H, = h4)
conditional distribution given action a and history h
Q1. Q2
-ositive-definite matrices action of graph at time step t observation of graph at time step Gibbs- Boobs-Boltzzann weight hyperparameters kernel potentials
V
a
(a,
1 1-norm
2-norm
H
norm over Hilbert space norm
14

---

Published as a conference paper at ICLR 2026
B NAVIGATION FOR READING
To help the reader navigate the Appendix, we briefly summarize its structure and content as follows:
Preliminaries: Basics RKHS in C We begin with short introduction of kernel meth- ods and reproducing kernel Hilbert spaces (RKHS), which provides the necessary back- ground for our embedding framework.
Theoretical Foundations of GCE Dn We then construct the Graph Controllable Embed- ding (GCE) under the Hilbert space embedding framework, presenting the Tensor, Dense, and Hom+Mean forms. Each embedding is also interpreted from a probabilistic perspec- tive.
Discussion of Various Embeddings in E Next, we compare the four forms: Tensor, Dense, Hom, and Hom+Mean in terms their properties, sample complexity, and com- putational trade-offs. This section also provides explicit derivations and calculations.
Theoretical Analysis in After that, we present formal proofs, including probabilistic consistency, convergence results, and sample complexity analysis for different embedding forms.
Experimental Details in L Finally, we provide additional experimental details, implemen- tation setups, and extra results that complement the main text.
This organization mirrors the logical flow from foundations, to construction, to comparison, to theo- retical analysis, and finally to empirical support. Readers may choose follow the entire sequence or jump to the section most relevant to their interests.
15

---

Published as a conference paper at ICLR 2026
C PRELIMINARIES: BASICS OF RKHS
In this section, by largely following the research Fukuuizu et al. 2013), the Hilbert spaces are assumed to be separable and dense in L space.
C.1 POSITIVE DEFINITE KERNELS
Let X be a set, and consider symmetric function : kernel, satisfying the condition:
1XX + R, known as a positive definite
n
T,==1
for any finite points TT, ,1,...,37 X and real coefficients CI.-.., Cr. The matrix (M(1,,;)) 1==1 is called the covariance operator.
By the Mloore-Aronszal theorem (Aronszajn 1950), every positive definite kernel on X defines a unique reproducing kernel Hilbert space (RKHS) H with an inner product structure (), comprising functions on such that:
1. For any X, 2() H.
2. The span Span(k(;; x) I X} is dense in H.
3. For all E X and E H. (f. k(,,))) f(x).
The RKHS H. thus characterized by the kernel k with k(.,z) serving as the reproducing kernel. kernel is termed bounded if there exists M 0 such that k(x,,) x)) S M for all X.
C.2 KERNEL MEANS AND COVARIANCE OPERATORS
Let Bx) ) be a measurable space and a random variable taking values in X with distribution Px. Suppose & is a measurable positive definite kernel on X such that 3[k(X,X)) co. The associated RKHS is denoted by H.
The kernel mean embedding mx X in H is defined by:
mx Elk(;
2(x2))xx((1
Using the reproducing property, the kernel mean satisfies:
(f,mx) [x](()),, for all H.
Now consider two measurable spaces (X, Bx) and (Y, By). and let (X, Y) be pair of random variables taking values in X and Y, respectively, with joint distribution Pxy over X x Y. Let kx and ky be measurable positive definite kernels with associated RKHSs Hx and Hy, satisfying Elkx (X, < co and E(ky (Y,Y)] xc.
The (uncentered) covariance operator Cyx Hx Hy is linear operator defined as: (g. (g.Cyxf)ms = g[)()))(Y)), for all f E Hx and E Hy. The integral representations are given by: (Cyxf)(v)= = f ky (y. (v. d)x) M) (Cxxf)(=)= f (x())
Under the bounded assumption, the kernel mean and covariance operators are well-defined for given probability distribution. The (cross-) jovvaiico operators can be written as
Cyx Ejky (: Ax (x Cxx = [E]xxx,, @ xx(,
The conditional mean embedding of (Y = x:) in Hy is then
mylz Elky (-, y) z| = Cy xxxxxxxxx =xxxxxxxxxxxx =x,,,( Cyxxkx (;,2),
(13)
where Cyix: Hx + Hy is the conditional embedding operator and C x is the inverse (or regular- ized inverse) of the covariance operator of Cxx-
16

---

Published as a conference paper at ICLR 2026
D THEORETICAL FOUNDATION OF GCE
we first introduce the essential theoretical foundation of Hilbert space embedding, followed by the construction of GCE.
Following the definitions from Appendix C
D.1 INTRODUCTION TO KERNEL BAYES' RULE AND HILBERT SPACE EMBEDDING
Before presenting the main theoretical analysis, we recall the kernel mean map and its application Bayes' rule, known as the Kernel Bayes' Rule (KBR) (Fukumizu et al.) 2011).
Let YK be random variables taking values in measurable spaces (x, Bx), (Y, By), and (Z, B2), respectively, with joint distribution P. We seek kernel-based analogue of the conditional distribution
= P
P(Z Y = yy=yxx=---
p(y=y|X==>
where x E EX y E Y, and = E are realizations of random variables X, Y. and Z, respectively. To obtain this kernel representation, we require the kernel mean embedding of the conditional joint distribution P(Z,Y X 2).
Let li,, ky: kz be measurable positive definite kernels on X, V. and Z, with associated RKHSs Hx. Hy, and Hz. The canonical feature maps for realizations and random variables as
(x(x)) 2x(,,)),
(5(4) = (.,)), x2(2) = k2),,2),
and
0% (X) = xx(,X), (Y (., 2z = kz (-
We assume each kernel is characteristic, ensuring that the kernel mean embedding defines an injec- tive mapping from probability measures over the domain to elements in the corresponding RKHS Sriperumbudur et al. T)
Step 1. The conditional joint distribution P(Z,Y Y X = 2)) admits the kernel mean embedding
mzy == (dz(Z) (Y) X -| E Hg Hy.
Equivalently, it can be expressed via the conditional embedding operator C2Y1x 7xx + Hz @ Hy:
mzv = C2y (x) = Czy Cxx (x)),,
(14)
where
Czyx := E (Z) @ Φy (Y) 8 ox (X)]
the (uncentered) cross-covariance operator between (Z, Y) and X, and Cxx is the covariance operator of X.
Step 2. Once we have the conditional joint embedding mzylas KBR gives the conditional mean embedding of Z give Y = y and X
mZz.a E(oz(Z) y, y. x] Czr Cr (r(1)
(15)
where
Czri2 = (oz (
E ( x].
Expressing these conditional operators in terms of unconditional covariance operators yields
CzYl2 = =C2xxxxxxx(1) Cyrl Cryx(xx
such that
(crXX(xtxxx)))) Φ2(v).
E(oz(Z) Y y, = x]
(16)
Two Steps in One Formula. Observing the properties of Equations 14 and 15 the conditional embedding of Z given joint distribution X and Y can be written single formula in the tensor space:
E[oz(Z) Y y. z] =ZIxxlox(x) 2 xr(v) = Czy x Cob @ dv(v)] (17)
17

---

Published as a conference paper at ICLR 2026
where tensor covariance operator x(x))( and conditional embedding operator Czjx x
C(xx)xx) = El((v(Y) @ Φx (x())) & (dv((Y) @ xx
and
Czyx = C2xx (XXXXX)
Remark 1. The tensor formulation in Equation 17 is equivalent apply KBR in two steps (XX - (Z,Y), then - Z), but it avoids intermediate conditioning and directly represents P(Z X)
in one operator by Cz
D.2 HILBERT SPACE EMBEDDING OF STOCHASTIC DYNAMICS
In this paper, our goal is to embed the conditional probability
P(O:AA = at, H. h4)
(18)
of the future observation Og given history h2 E HH and E A. Following Definition 1 we denote
the kernel feature mappings of their realizations and random variables in their respective RKHSs Ha. Ha and Hp
o8 = =a(a)), 1ºa oh))),
ve co(o)),
(19)
and
v/" = HH(H),
2 = (a(,),
vf = (c,)
(20)
According the properties of KBR (see Equation 17 in Appendix D.1 the conditional mean embed- ding of O2 given (he, az) is expressed directly in tensor form as
ElU A, a,, H, = h4] = Co (AH)A) @ oil = Co Am ( @ ef
(21)
where COAH and C()H)(H)) are denoted
COAh = Elw @ vo 8 (b))) C(AH)(H)) = Bl(Ve so wit) 6 (v s w
The operator COJAH Hn 8 HA - Ho the conditional embedding operator, mapping joint history-action features to observation features in the RKHS. Figure ① illustrates the core idea of Hilbert space embedding of stochastic dynamics. Left: The nonlinear stochastic dynamics are rep- resented by the evolving conditional distributions P(O,AA, a,, H, h₂), shown as red curves. Red dots indicate possible samples from these distributions, and black dots denote realized obser- vations. Right: After embedding into RKHS, the dynamics become linear under the conditional embedding operator CoJAH, with red dots indicating the corresponding conditional mean estimates in the RKHS.
D.3 DERIVATION OF GCE BASED ON HILBERT SPACE EMBEDDING
The Hilbert space embedding of stochastic dynamics has been established in Appendix D.2 now show how the GCE can be constructed via Hilbert space embedding.
We
Consider a stochastic system of N interacting objects, where each object i is associated with an observation 0', action space A and history space HHI' At time I, object has history hi HP' action a; A and future observation of O'. Also. denote the H; Ai and Oj as the corresponding random variables. Applying Definition [ the corresponding RKHS feature maps are
te = OE (h(),
wia = des (a:);
v : (o.),
in the corresponding RKHSS Hgr H4 and Ho, respectively.
18

---

Published as a conference paper at ICLR 2026
Recalling the principles of KBR in Appendix D.1 can be written as
the joint conditional distribution of all N objects
P( 12
-d,
(Bayes' rule)
P((A) H/
(22)
P(O), - - N II =,, -
(ohject wise Factorization)
In the RKHS setting, this corresponds a global conditional embedding operator Cojarr acting on the pair-wise joint history-acti action feature [u] @ &b ..... A.,NN & v. While this global form is valid, does not exploit the relational structure of the interaction graph, where each object interacts only with subset of other objects (e.g., neighbors).
Decomposition into pair-wise contributions. From the tensorized formulation in Equation 21 the global conditional embedding of all objects can be written
VE
E
0.1 We
O.N the
(A} af. HÇ n2) COJAH
where [22 ................................ v ONN stacks the observation features of all N objects and COJAB is a global oper- ator from the joint history-action RKHS to the joint observation RKHS.
By the linearity of CoJAh, this global operator can be represented in block form:
Co'AA'
COIJANHN
COIAH
CON
CON ANHN
where each block CosJARи н: maps the history-action feature of object j to its contribution in the conditional expectation of object i's observation.
For the i-th object, the i-th row of this block matrix yields:
EM Oi (aj M1
CoyJNNHN ComAEnilo @ J. JEE())
NNN va..N
(23)
Interactions occur only with neighbors j E(); non-neighbors (j & E()),, we set Co. A = 0. This induces a localized form of the global conditional embedding, fully aligned with the entire topology.
factorization of the global conditional embedding operator induced by KBR, where only the operators Co corresponding to edges in the inter- action graph active. This factorization preserves the theoretical grounding of the Hilbert space embedding while enabling efficient modeling and control through localized graph structure.
Interpretation. In this view, Equation 23
Limitations. While the embedding offers powerful representation of distributions, the tensorized formulation in Equation 23 still suffers from two limitations:
Sample complexity. Based on the random matrix theory (Tao 2012), the convergence of the covariance operator grows significantly with the modes of the tensor.
19

---

Published as a conference paper at ICLR 2026
Optimization. The intertwining of history and action features within the tensor prod- uct complicates the optimization of sequential actions. Since control algorithms like LQR/iLQR assume additive disentanglementt between history and action features, the ten- sorized embedding in Equation 23 violates this assumption.
To concretize the controllable embedding within the Hilbert space emhedding framework, the ten- sorized features in Equation 23 must be disentangled to meet the requirements of controllable em- bedding.
D.4 DISENTANGLEMENT OF HISTORY AND ACTION FEATURES FOR OPTIMIZATION
Building on previous studies (Altun et al. 2012 Zhang al. 2008), efficient approach to address the aforementioned problems involves simplifying the tensorized formulation into a concatenated feature representation, [e
hJ || 1/ This decomposition simplifies the problem by factorizing the action and history random variables into two product conditional probability distributions, as exem- plified in (Song et al. (2009] Altun et al.) 2012). By adopting the concatenated feature representation, Equation 23 can be simplified as:
EE 1
0,
"" COMANNN]] (I12 ..... A.N 11
Σ Cos [v] 1113 jEE())
(24)
Σ )EE(0)
Ho
hhd Cor
where the linear conditional operator is simplified to the form [Co no. Co Here, this formula- tion in Equation c2n be regarded an approximation of Equation 23 neglecting the high-order interaction of history and action features. The formulation directly brings two advantages:
Reduced Sample Complexity. Based on the concentration inequalities in random matrices (e.g., matrix Hoeffding and Bernstein inequalities (Van Handel (2014)) the sample com- plexity can be significantly reduced with lower dimension. Then 3-mode tensor operator Cojan in Equation 21 can lead to a higher sample complexity.
Simple Optimization. When optimizing the quadratic cost function constrained by Equa- 3i the only undetermined variable is the fully decomposed action feature 1o instead of the entangled feature db @ v/e.
D.5 EFFICIENT EMBEDDING BY MEAN FIELD APPROXIMATION
In systems with multiple interacting objects, each object directly influences its neighboring objects through their relational structure. For multi-object systems with many objects and dense edges, mod- eling all these pair-wise interactions explicitly becomes computationally expensive, scaling with the number of edges in the graph. The core idea of using mean field approximation is simplify the modeling of complex interactions in systems with many interacting components by approximating average or "mean" effeett, rather than modeling every pair-wise interaction explicitly. Thanks to the disentangled structure Equa- tion 24 we can separately model how history influences the future observation probability and how actions affect the observation, combining them multiplicatively in the joint conditional distribution.
the influence of all other components
given component
To disentangle history and action effects, we assume conditional independence between actions and histories given the observation (A II H O). This allows us to factorize the conditional distribution into pair-wise history potentials and an action-dependent likelihood term as
P( ( 11) {A} a(bees)))
2 (o% = =b:) -P( afbeecnl::). (25)
(pair-wise factorization on history elfect)
II (EE())
X
where
20

---

Published as a conference paper at ICLR 2026
h is pair-wise history potential encoding the interaction between object i's obser- vation and object j's history.
P ( at bes l0:) a disentangled action-dependens likelihood term, reflecting the statistical dependency between the neighbors' actions and object i's observation.
Based on the disentangled probability in Equation 25 the naive assumption of equally weighted potentials Φid for all j is generally invalid, rendering the representation (Li et al. 2020) misspec-
ified.
Mean field normalization for history. Based on Equation 25 only the history part requires mean field normalization. We convert 6/," into normalized influence coefficients by taking its expectation under the current mean field approximation (O)):
E (Oo log ) (0 nz))
exp
i EKEe(H) exp
with JEE(1)
1.
(26)
E (ooog)oogg / (0) ht]]
Here the expectation Boon (under variational distribution ((O))) ensures that the influence weights do not directly depend on the unknown variable O but rather its current mean field approximation.
GCE with mean field approximation. Substituting these weights into the GCE formulation in Equation] 24 and applying the principles of KBR yields:
Σ SEE(1)
CoYHs hos + Coya 15
(27)
=Coq
Σ JEE(()
2
+
Σ JEE(()
Corlss vf.
Here, Coa shared conditional embedding operator used in the homogeneous case, while the non-uniformity of ad directly reflects the unequal pair-wise potentials Φ
Practical Computation. Directly estimating the distribution (O;) is infeasible, as the multi-object dynamics evolves in continuous observation space. Instead, we approximate the influence weight at using the history Af of_1. which temporally close to the future observation of Since the forward dynamics modeled in RKHS. we compute the weights from the history features via a Gibbs measure:
exp(s( (aatt) exp(s(vb bit v )
af
(28)
kEE())
where f f is a negative pair-wise potential function between history features. The normalization in Equation ensures ΣjE(()) a. 1.
We consider the following four choices for to instantiate the pair-wise potential:
1. Gaussian potential. Let shi wi
bid Rdd,, Define
||
SRBP hi whi)
2
with bandwidth a > 0. This corresponds to a positive definite Gaussian kernel and is robust small perturbations.
'ere, exp ().((())) can be regarded practical proxy measure exp [Enoo (ollooo///// (.))
this holds
since there exist some functions that pull features the observation space.
21

---

Published as a conference paper at ICLR 2026
2. Laplace potential.
11
h.3
m- 111
fLap (v) .
with scale > 0. Compared the quadratic potential energy in Gaussian kernel, the Laplace form induces heavier tails due to the nature of I,a norm.
3. von Mises-Fisher (vMF) potential.
Vt
1 with concentration 2 O. This favors alignment on the unit sphere and is natural for directional embeddings.
fNMF( W ll2
1
2 1
4. Neural potential. Let fo Rd Rd R be neural scorer:
fo( ( MLP (
where [iI-]] denotes concatenation. This yields flexible, learnable potential while preserv- ing nonnegativity via the exponential link.
In all cases, the Gibbs normalization in Equation 7 ensures EjE(() a. 1. Hyperparameters (a, X, x)) or the neural parameters 0 can be learned by the joint training.
22

---

Published as a conference paper at ICLR 2026
E DISCUSSION OF VARIOUS EMBEDDINGS IN GCE
This section introduces four embedding formulations used in our framework and evaluates their properties for controllable embedding design. We denote the dimensions of the history, action, and observation features by da, da, and d,, respectively.
E.1 FORMULATION OF VARIOUS EMBEDDINGS IN GCE
Tensor Representation. Originating from the inherent structure of Kernel Bayes' Rule (KBR), the tensor representation directly follows Equation 21 and can be expressed as:
E 1/ Haf, [
Co
COMANHH
w & v
(29)
E
[ Ha
CONIALHI
CONIANHO
J.N & v..
Each sub-operator Cо.AАAи: Rd.x(dxxd.)) acts on the tensorized feature WhS X Vit (dda xd,)x1 entangling history and action information. This entanglement makes it not suitable for controllable embedding, as the representation is neither locally nor globally linear. The compu- tational complexity is ONNd per forward pass, making it computationally expensive for N with large value. For simplicity, we refer to this form as Tensor.
Dense Representation. The dense representation removes the tensor product entanglement by decomposing the history and action features, as Equation 24
WE Hag
0.1
E[
ON
Ha
(30)
Coym
Coynw
CoyAr
Coyaa
COMIHI
CON
CONJA
CONJAN
Here, CoqHs Rdoxdd and Cor Rdo xd making the representation locally linear in each feature type. The complexity is (dddd + doda)) per forward pass, lower than the tensor form but is still quadratic in N if with dense connections. For simplicity, name this form as Dense.
Homogeneous Form with Uniform Weight. This form generalizes the compositional Koopman operator (Li et al. 2020), assuming equal-weight influence from all neighbors:
B Two 0.1 fag, 1/ [
E
(31)
CoYH=
CoynN
CoyA
CoIAAN
+
COMIAT
COMHNN
COMIAI
CONIAN
The sub-operators have the same dimensions as in the dense form, and the complexity remains (d,',,, d,da)). In the setting of (Li et al. 20200), Co иs shared for the entire graph. However, the uniform-weine assumption is generally inaccurate in probabilistic settings (see Equa- tion 25
We call this form Hom.
23

---

Published as a conference paper at ICLR 2026
Homogeneous Form with Mean Field Approximation. To address the limitation of uniform weighting, we propose to use a mean field approximation with adaptive Boltzmann- Gibbs weights:
E
{a.. 1
(32)
Co
Σ ΣjeE(1) oz 1
Coolal
CoulAN
COMHH
EjEE(N) aj o
CONIAL
CONIAN
where denotes the element product with shared operator Co.HH- We refer to this form as Hom + Mean. This form offers:
1. Lower complexity: The history term costs O(Nd,da) and the action term costs O dNNadd))) reducing cost compared to dense and tensor forms.
2. Fewer parameters: Requires fewer learnable parameters than dense tensor forms.
3. Adaptíve weighting: Captures non-uniform neighbor influence without quadratic parame- ter growth.
Comparison of Embeddings.
Table 8
summarizes the properties of the four embeddings.
Table 8: Comparison of embedding formulations in GCE.
Embedding
Complexity
Parameters
Adaptive weights
Tensor
d(Nddddd (N"(d + doda)) O (N + doda))
Quadratic
Yes (implicit)
Dense
Quadratic
Yes (implicit)
Hom
Constant
No (uniform)
Hom + Mean
O O(Noda + N2dodd))
Constant
Yes (explicit)
E.2 EMPIRICAL ESTIMATION OF CONDITIONAL EMBEDDING OPERATORS
The empirical computation of the four forms in GCE is listed follows. Given i.i.d. sample
{o,,,,).... a!....,,,, a .......,,y! AN)) Jt= The empirical covariance and cross-covariance operators are estimated as:
T
CO.D
Σ
E oi
&
hJj
T
t=1
1 T t=1
CO.A)
Σ
jE [N], We & was,
T
CHIHI T
Σ
T
CA = T
Σ
j [N
1. Tensor Form (Tensor). From Equation
29 the empirical block-operator is
COIAAAA T
JEIN,
and the conditional embedding operator computed
COYAAH COYAA (CHHAA))HAA)) (C
+X)
(33)
24

---

Published as a conference paper at ICLR 2026
where AI is the Tikhonov regularization and
T ( s vhdd @ (ve2 @ w.)))
Σ
2. Dense Form (Dense). From Equation 30 the history and action parts are separated:
ČоTи = Co.Ho (CHi + +X)= j E [N].
(34)
BojAs=Cov (C)++M))+ je [N].
(35)
3. Homogeneous Form (Hom). In Equation
31
operators Coijus for all i shared:
N EECCn co (Crm + AT) j je [M].
Coym=
(36)
N
8=1
while Corjas follows the empirical estimation of the Dense form.
Homogeneous with Mean Field Approximation (Hom Mcan). From Equation 32 the
history part uses adaptive weights a!:::
COTNH =Co H + +
1
@ Σ
(
i3 at vi
Σ de ( +
NT
JEE(1)
jEE(()
jEE(1)
(37)
with
N
CorH
1
Σ
at af (h))) )
(38)
NT
JEE(1)
1 NT
CHH
Σ
(39)
(ES(()
JEE(S)
Here, CovT H is shared and computed as in Equation tion of the Dense form.
37 while Cor follows the empirical estima-
25

---

Published as a conference paper at ICLR 2026
F THEORETICAL ANALYSIS
F.1 PROOF OF THEOREM
Proof 1. Step I: Convergence of empirical operators. Based the practical computation in Equation 33 in Appendix E.2 for all [N]
Coywms COYAAS ((NAA))HAOO+++
where
CO.AHH
1 7
Σ
Vi @ (v) @ w.)).
CINHHMAAH -
(of 2 v @ (v. @ v)
Since the feature maps are bounded, continuous, and the samples are ii.dd, the Hilbert space strong law of large numbers (Muandet et al.) 2017) implies
AOAA
a.S.
CO C(ASH)((AHH))
C(AHH)Y(A5H))
HS
HS operator perturbation results yield
in Hilberl-Schmidi norm. With À Ar + 0 and TAT
Kato 2013)
Cooaam - ConAAsH
0.8, 0. loop
Step 2: Pointwise convergence on fixed inputs. For fixed {az
1 define
20.3 Σ EOyAA
1/2
JEE(()
0,i 11
COiJASHE
To [
@
=E 0, {ag. no 11
JEE(1)
Then
20.6 1
<
Σ JEE(1)
CosjARN - Com 11. -1o0t @ 2
S C|E(1)) max JEE(S)
[COYASHS - ConlAsH
where C <, 0o is constant due the boundedness features in the compact space 0' The right-hand side converges to in probability, hence
30.
O.is
T T-000 0.
1/1
H
Step 3: From embedding convergence to weak convergence. If the kernel feature is character- istic (Sriperumbudur et al. 2011), then
MMD (ptot (af, k/)), P(O) (af (()) | 30. 0,
where MMD is the maximum mean discrepancy and thus the above convergence in H implies
P( Aa (2) T+x P(O) (at M
"The maximum mean discrepancy (MMD) between two distributions P and defined as MMD(P.Q) = [mp- mallH where mp and mQ are their kernel mean embeddings in the RKHS H(Hommann et al. 2005).
26

---

Published as a conference paper at ICLR 2026
F.2
PROOF OF ERROR BOUNDS OF VARIOUS EMBEDDINGS IN GCE
The main steps to prove the error bounds for Equation T0 is outlined as follows:
Procedure I: Combine the well-established matrix concentration inequalities in Appendix
to determine the sample complexity the sub-operators Co.u and Co.jai
Procedure 2: Establish the union error bound for the Bulbert-Schhidt norm in the forward
loss function shown in Equation 10
Procedure Prove the sample complexity for other embeddings: Hom and Dense.
F.3 PROCEDURE 1: SAMPLE COMPLEXITY OF OUB-OPPRATORS
The regression in Equations 35 and 37 arises from two primary sources: (1) the estimation of
cross-variance operators Co' A CO+H- H as well covariance operators CAI A CHT hs based finite i..d. 7 samples; (2) the influence of the perturbed term AI. We proceed to derive the non-asymptotic convergence rate step-by-step using concentration results for random matrices.
and Here, we first give an instance proof for the C. and this proof can be plugged into the other cases.
AI)-1, where Co A Rdoodd and ČASAr E Rdaad are two estimated adjoint and self-adjoint operators. respectively. Assuming both two operators are from a finite i..dd T samples in Equation 33 we have probability at least 35 with 5 E (0, 1) and T > clog(2d./8
Proposition For the i-th object in graph, let Čo Co
Amin (A3 A2)
||Co|a - Coy S E + E.
where
Edg < V A/max(Coro)) (Coros)
T.Emin (CArAi) + x))
clog(2da/0)
Amiin(CAiA)
Amiin Manin(CAAAA +
and
2 llog((d, dd))/8j
2 log
+
F <
e clog(2d/5)
Amin (CAiAi)
T'mmn (C
Proof. According Lemmas and operators and errors as
decompose empirical operators two terms: target
= fo-a=Corat + AO'A
(40)
= CALAi AAiAS
where Aot A and A As be regarded the residual enror terms. Our target is derive how the error bounds of A and AAi change with the sample number.
By differencing two operators, obtain
Co.N4 - Cor
-Com AiAi (C +XY-
-Com C (CoLs Aorai)(Caa AA -
(41)
=COiAl CANN (CANA ANIA +AT)-1) AoiAs (CAIAs + NASAI + XI part
part
We divide the derivation of the error bounds into two steps, corresponding to parts I and 2, denoted by E and E respectively.
27

---

Published as a conference paper at ICLR 2026
Step 1. Deriving the Error bound E
Based on the Woodbury identity becomes
(a.k.a. matrix inversion lemma) (Deng 2011 the first error term
CoiAi (Cabas -CCAAA +AN
-Co A (CCres -CAsAt+ +CNA + Anis + +1 (CAsAD AAEn + +A1=
(42)
=CoNCC ALAI (ANIN + A + ANN + A1)
Plugging the Lemmas
and 4 into Equation 42
we get
CO + A +AVA A +1
Coo C-1/22 A AL A MA
(A + AT)(CAA + A + +
(43)
(2log(2d//0)
Sv Symoor(Cooss)
T'mmin (CA)A)) +
Vemin (C
Amin (CA + X)
where the first and second terms are derived using singular value decomposition techniques, the third term is based on the error bound for self-adjoint matrices with probability at least 1 - 5 when T > clog(2d,,/8 (8) (a direct result from matrix Hoeffding inequality in Lemma 4 and the last term Amin(AAAA") A2)
leverages fundamental properties of matrix theory. The primary contributor to the error bound the statistical deviation term A A AE, while the regularization term A7 is typically small in practice.
Step 2. Deriving the Error bound E
To derive the second error term E we consider the bound
|
AC + AAAA + 1|
Amiu (CAi + AAiAS AI)
| ||
log (2dc/8) T TAmin' A5))
Amin (C (CA
(44)
2ogg(2 (2, t..) 1/8)0
2log(2(d, + +d.)
T
(log(2d./6)) Thmin (5)) Amin (CAi +
which holds with probability least 28. The first inequality follows from the standard sub- multiplicative property of the operator norm. The second inequality is a consequence of 4 which applies matrix Hoeffding inequality to the self-adjoint operator CALAS + AAAAI The final inequality uses Lemma 2 which applies the matrix Bernstein inequality to bound the spectral norm of the empirical deviation Aot As with high probability.
Applying union bound over the two error terms E and E the total error is bounded with probability least 1 35.
Proposition 2. For the i-th node in homogeneous graph with N nodes, let Co Co ir ( + Ar)-1, where Čни E BRd xd, and ČO= E xds are two estimated adjoint and self-adjoint op- erators, respectively. Assuming both two operators from a finite i.i.d. T samples in Equation 37 we have probability at least 35 with (0, and T > clog(26,,18 18)
NAmsuHHH
|Coo - Co < El F
where
clog clog(2dn/0) + +2)
Ei.1 S V Amaz (Co
Dinin (CHH)
NT>min (CHH))
Amin AminCCHu + X)
'It is shorten A - (A+ B)
A- B(A A
BY
28

---

Published as a conference paper at ICLR 2026
and
2 log((d, + do d/1//8pe
+ 2 2log((d, (dd 1/8]L
3NT
2,22 S
clog(2ds,] (5)
Amin (CHH) +
NTmmmn (CHH)
Due to the homogeneity of the graph, each node contributes T i.i.d. samples, resulting in total of NT samples. This significantly improves the convergence rate of the estimated operator Cj compared single-node setting.
Remark 2.
F.4 PROCEDURE 2: SAMPLE COMPLEXITY OF Hom + Mean IN LOsS FUNCTION T0
Theorem a (Sample Complexity of Hom + Mean). Considering a homogeneous graph with N nodes, the graph dynamics satisfy the formulation Equation 2 For Hilbert space embedding, history, action, and observation feature dimensions are dn, da and do respectively. For finite i.i.d. sample
number T > max;
clog(2d,, (8)
alog clog(2d,, /6)
we have a probability least 1- 36 with 8
(A
N NAmin (HH)
(0, satisfying
Coyn
15 10d Σ,
Coupal
COyAN
CON
Ejat Vt
N.J. Rid ai
CONPA
CONIAN
HS
<O(N) max(E), Ei,, + 2), O(N2) max(E
where Fod Foi F and F are defined in Propositions Mand
2
Proof. By applying the properties of Cilbert-Schmidh norm and triangle inequality, have when M =1
N Σ
-Co
Σ
L
Vi
jEE(()
fml
Co
Σ
ae (. bij
Coy1
Σ
Oe
SEE(K)
JEE(()
(45)
(onANnN CoN
i==1
Corl CoyN
Σ
Σ W==
Comws-
Coyias
|||||sso|
at
JEE())
All features are well-defined in the Hilbert space, and thus both
Ejeeto of Sh
and a.
are uniformly bounded by a constant (see the properties of kernel functions in Appendix C). By combining with Proposition [ the error bound in the left-hand side scales with the number of nodes as (E),, + E,,2)
Combining Proposition 2 the error bound of right-hand-ae scales O(N2) maxi,, E
Here maxi, (F + F controls the upper error bound of sub-operators actuation components.
Since the two parts hold under T > max, Alog(2d.// /8) Amin (F)
and T >
Allog(2d./8) NAmi (HHY- we rewrite the con-
dition as T max,
Combining the two parts together, we obtain the union error bound for the forward loss function as (NN max, (F%, + max, (F4) +
clog((dd, (8) V. alog(2d, /8)
Xmmnn AT) NAmsa (HH)
29

---

Published as a conference paper at ICLR 2026
F.5 PROCEDURE 3: SAMPLE COMPLEXITY OF OTHER EMBEDDINGS
Theorem b (Sample Complexity of Hom), Considering homogeneous graph with N nodes, the graph dynamics satisfy the formulation Equation 2 For Hilbert space embedding, history, action, and observation feature dimensions are di da and do respectively. For finite iid sample number
> max, Am
clog(2d./15 18)
V
Nomas(HH)) we have a probability at least 36 with d (0.1)
clog (2dn (8)
satisfying
Whil
CooAAl
1i ComlaN
: Lige
CONIH
CONINN
WhiN LiVe
CONIAL
CON
((22) Maax(E, + 2),2) + ON22) (ax
E.2)
where Eaii E.. Ej,1 and Ei are defined in Propositions
and
Proof. The proof of this theorem follows the same idea with Theorem a the main difference is that here need estimate each Cosjra since the implicit weight influenced by the j-th node is unknown. Thus, the upper error bound is scaling quadratically with the number nodes, since there 22 operators in the first block.
Theorem e (Sample Complexity of Dense). Considering graph with N nodes, the graph dy- namics satisfy the formulation Equation 2 For Hilbert space embedding, history, action, and observation feature dimensions da, des and do. respectively. For finite i.i.d. sample number clog(2d..8)) V clog(2d,,/6)
T > maxi. (A-A
we have a probability at least 1 - 35 with E (0,1)
satisfying
Comm
Coynx
Coya
ComJaN
188
CONIHI
CONIAN
Li
CoNIAi
CONIAN
$ ( max((F; + Fiid) O(N2) 1.
max 1.
Eid
where Fig and F are defined Proposition
I and Fi and Figy are defined as
ptd S \mmoo(Coio))
cloo((2d// Tdmin (C
+ +x)
V Domin (CHSH))
Xmin (CHEns + A)
and
2 logg(dd. +d.))
2 +do
F 22 S
37
clog(2da/8
Aman (CHT H3) + A
T Thmin (Cuins)
Proof. The error bound difference between the homogeneous graph and inhomogeneous graph is that each operator CoijHs should be estimated independently (see the explanation in Equations 30 and 31). Thus, each emror bound of Fi. and F can be proved by following Proposition
Following the same idea shown in Theorems a and 2 the union upper bound can be derived as
(N22 max,, (FA 1 F ((N2) max,, E.. E
30

---

Published as a conference paper at ICLR 2026
G AUXILIARY LEMMAS
The auxiliary lemmas in this section are primarily based on the monographs (Vershynin] 2018; Tropp al. (2015). These four lemmas are focused on handling error bounds for the approximated adjoint and self-adjoint operators that arise in regression. The derivation of these error bounds relies on three fundamental matrix inequalities: Matrix Bernstein, Chernoff, and Hoeffding, which serve as key tools for controlling the embedding errors.
Lemma 1 (Matrix Bernstein Inequality (Vershynin) 2018 Tropp et al. 2015)). Consider a finite i..d sequence {Xx j==1 of independent, random matrices with finile dimension di X d2. Assume that X3 is centered at 0 and bounded by a constant for any index k as
[[Xx]= =0 and || S L.
Let = Ek= Xg. Define the matrix variance statistic:
1
v (Z)) max (1)
E x
X I1I:
Σ B[X, Xx]l].
(46)
Then have
-12
P |
| 21) S (d1 d2) exp(
u(Z) + Lt/3'
The introduction of the term v arises from the symmetrization technique used the proof, where the operator norm of the sum is related to the maximum of the operator norms of the variance terms both the original and transposed spaces.
Lemma (Error Bound of Pross-Covarianee Operator). When Lemma 7 holds, we have a proba- bility at least 1 -
||exx- - -yxll 2log d2)/8)u
2log((d: L2)
N
3N
where
Crx Rdixds
(Orx Cxx) - Cxx)
4
02 |(yx - -Cxx) (oi)((Cxxx (Cyx- - Cxx)
v ==max(er, 02).
Remark 3. Lemma 2 is a direct result based on Lemma Following from finite i.i.d. samples, the convergence rate with the number of samples is scaling as O( log d22)) ). The ervor bound
N
established in Lemma is particularly valuable for conducting non-asymptotic analyses of cross- variance operators.
Lemma 3 (Matrix Chernoff Bound Vershynin 2018)). Consider a finite i..dd. sequence (Xa) of random and symmetric matrices with a finite dimension do- Assuming that the eigenvalue of Xg is bounded by a constant for all index k as
0 S Amin (xx)
and Amaz (Xx) S T..
Introduce the random matrix
Z Σ X
Define
Mmnnn Amin (E())).
Then, for any e. E [0, 1), we have
P Amin (7) S )))mmin
S 2d, exp(- (C) L)
(47)
31

---

Published as a conference paper at ICLR 2026
Lemma 4 (Error Bound of Covariance Operator (Tropp et al. (015)). Consider the random variable
Čxx Rddd we have probability least ô with (0, when N (log(2//8)
Amiin
Amin (Cxx) 2 (1- clog(2d/6)
Amin (Cxx);
Nmmin (Cxx)
where
c N with c> O. Amar (xxx)
Remark 4. Lemma can be derived from the matrix Chernoff bound through the Laplace transform technique, which transforms the multiplicative bound into an additive form of matrix concentration inequality. The resulting evoor bound exhibits a convergence rate of O (los with respect to the sample size N. This concentration inequality is particularly effective for analyzing covariance operators, especially the context of empirical covariance operators in Hilbert space embeddingg The primary source of estimation error arises from the finite sampling, which is quantified by this probabilistic bound.
32

---

Published as a conference paper at ICLR 2026
H
PSEUDO CODE
Algorithm Graph Controllable Embedding (Hom+Mean)
Modeling
Require: Data D (T,)) M Naaple
s-step transition tuple 77 (ot:N aat hith) with N
11=0
objects; learning rate o: number of training epoches K: four different pair-wise potential energy energy functions f(.,)) -) in Table 2
1: Initialize GNN encoder and decoder vb and vet
2: for Training epoch & O,... K do
3:
(Encode inputs: observations and history via neural embeddings; actions via a fixed linear projection (ensuring direct recovery back to the original action space))
4: Embedding graph history h. and future observations of v. w...:
@a(o, !:N) and
vh v. (h, action features are obtained via a fixed linear projection: w/ @A (a )
5: Estimate Boltzmann-Giibe weights a exp(sCc
Σ exp (S
6EE())
6: Compute the ensemble features Ejee(h)
7: Estimate (Cot i=1:N and (Coqw) JiJ=I:N by solving Equation by minimizing Hilbert- Schimit norm
8: Compute the loss Liee defined by Equation IT
9: Update encoder and decoder vb and et with stochastic gradient descent
10: end for
Control
Require: target VP; initial observation as history V8: control horizon M; GNN encoder and decoder
d and zit
11: Construct the control objective Equation
12: Obtain the optimal control sequences (V V&r) from minimizing Equation 12
13: Apply (Vr,.., Vi,
".,....,7
33

---

Published as a conference paper at ICLR 2026
1 ADDITIONAL EXPERIMENT DETAILS
1.1 CONTROL ENVIRONMENT DETAILS
In the numerical experiments, we adopt the control environments from (Li et al. 2020), where inter- actions between different objects are considered differently depending on the types of connections and the physical properties the objects involved. The environments are designed to capture diverse interaction dynamics, detailed below:
Rope. In the Rope environment, the top mass has a fixed height and is considered different from the other masses. Thus, there are kinds of self-interactions for the top mass and the non-top masses, respectively. Additionally, we have 8 kinds of interactions between different objects. The objects a relation could be either top mass or non-top mass. It is combination of 4. The interaction may happen between two adjacent masses or masses that are two-hop away. In total, the number of interactions between different objects is 4 2 8. The in-distribution test includes 5 9 objects, and few-shot test is from 10 14 objects. Examples of rope systems are demonstrated in Figure
Soft. In the Soft environments, there are four types of quadrilaterals: rigid, soft, actuated, and fixed. There are four types of self-interactions, respectively. For the interactions be- tween objects, there is an edge between two quadrilaterals only if they are connected by point or edge. Connections from different directions are considered as different relations. There are 8 different directions: up, down, left, right, up-left, down-left, up-right, down- right. The relation types also encode the type of receiver object. Thus, in total, there are (8 + 1) 4 = 36 types of relations between different objects. The in-distribution test includes 9 quadrilaterals, and few-shot test from 10 14 quadrilaterals. Examples of soft robotics are demonstrated in Figure
Swim. In the Swim environment, there are three types of quadrilaterals: rigid, soft, and actuated. Similarly to the Soft environment, different edge types are specified for different connection directions. The number of edge types is (8+1) x = 27. The in in-distribution test includes - 9 quadrilaterals, and few-shot test is from 10 14 quadrilaterals. Examples of swim roboties demonstrated in Figure
Power grid. The primary control objective in this environment is to stabilize the voltage magnitudes at all nodes around a predefined reference voltage value, Vrat = 1 p.u.. Each node has a state vector [V7 V. where V is the voltage magnitude at node i at time step t, and V is its rate of change. This stabilization is achieved by adjusting the reactive power generation Qgen generator nodes in response to both internal system dynamics and external disturbances (i.e., noise) generated by the loads.
The undertying grid topology is modeled ring or grid structures) or randomly generated using an Erdös-Rényí model (Examples are illustrated in Figure B). The graph is assumed be connected, and nodes are randomly designated as either generators or loads. maintaining a generator ratio between 20% and 50%. Generator nodes are controllable via reactive power inputs, while load nodes are subject to time-varying noise disturbances. To learn the system dynamics, we generate 1000 randomized instances with topologies ranging from 50 to 100 nodes and evaluate the control performance on larger graphs with over 100 nodes. The control problem is formulated as the following cost minimization over a time horizon M:
undirected graph, which may be fixed (e.g.,
M
Σl - vi +
+ 1 Q2
where Qi and Q2 are positive definite weighting matrices, and u' is the control input time step L.
I.2 DATA GENERATION
We generate 10,000 episodes for the Rope environment and 50,000 episodes for the Soft and Swim environments, with 90% of the episodes used for training and the remaining 10% for testing. Each
34

---

Published as a conference paper at ICLR 2026
episode consists of 100 time steps. The dataset includes physical systems with varying numbers of objects ranging from to 9. In the Rope environment, ropes consist of 5 to 9 masses, where each mass is treated as object. The observation for each mass includes its position and velocity in 2DD space, resulting in a 4-dimensional observation per mass. In general, a rope with N masses has a AN-diimesionnl observation space. In the Soft and Swim environments, soft robots consist of 5 to quadrilaterals, each treated as an object. For each quadrilateral, we observe the positions and veloc- ities of its four corners, leading to 16-dimensional observation (4 positions 4 velocities). Thus, for a soft robot with N quadrilaterals, the observation has dimensionality of 16N. Additionally, we generate few-shot set with 200 episodes for each environment evaluate the generalization of our graph controllable embedding. The Rope few-shot set includes ropes with 10-14 masses, while the Soft and Swim sets include robots with 10-14 quadrilaterals, more than the number of objects in the training sets. The graph topology is generated with random structure and random number of nodes, ranging from 50 to 150. To ensure the controllability of the random graph, the graph must be connected. This is guaranteed if the edge connection probability satisfies
(1 (1+6)lnn
p
where n is the number of nodes and 5 > O is small constant. In our setup, we set the connection probability to = 0.15.
1.3 TRAINING DETAILS
Our methods and baseline models trained using the Adam optimizer with learning rate of 1e-1 and batch size of The learning rate scheduler follows linear schedule with decay rate of 50% per 100K gradient-descent steps until reaching minimum learning rate of le 66 For all models, we apply 400K iterations of gradient-descent steps in the Rope environment and 580K iterations in the Soft and Swim environments, respectively. The number of future steps in the loss function Equation 10 is set to M = 16. All experiments were run on single RTX 4090 GPU.
1.4 BASELINE ALGORITIIMS
The baseline algorithms are trained on the same dataset and device. We compare the control perfor- mance of our model with two categories of baselines:
* Controllable embedding methods without relational structures:
VAE-based controllable embeddings:
Prediction, Consistency, Curvature
(PCC) https: com/
VinA rchheeeeec////CC-tttrrrr
Graph representation methods:
Koopman Polynomial Model (KPM): https://githbb.com cow.uu/huuco//
Composition innnal http:///ww....aoooooaoo.aeeaaa/a
Koopmanoperrtors..py Conpppeesssooaaloooooooaaeeerrrr
Composition
Compositional Koopman Operator (CKO):
https://github.co com/
35

---

Published as a conference paper at ICLR 2026
Figure 5: Examples of the random parametrized shape of the rope system. The red dots in each sub-figure represent the control targets.
Figure 6: Examples of the random parametrized shape of soft robotics. The red boxes are repre- sented as the control target of each soft robot.
Figure 7: Examples of the random parametrized shape of swim robotics. The red boxes are repre- sented as the control target of each swim robot.
36

---

Published as a conference paper at ICLR 2026
Srdis-Reryiiooo Graph (n=1
Erdiis Erdis-Rérryi Graph 2 (n=94)
Erdis--iére Craph (1=
EriMs--héryj Graph 4 (1=
Figure 8: Examples of the random graph of power grid. The red node represents the generator, and the blue node represents the load.
37

---

Published as a conference paper at ICLR 2026
1.5 ADDITIONAL EXPERIMENTAL RESULTS
Noisy vs. Noiseless Conditions. We evaluate the predictive performance of our models by fore- casting 100 steps into the future in the Rope and Soft environments, with and without observation noise, and compare with the state-of-the-art CKO method. Figure 9 presents quantitative results averaged over 20 trajectories from different initial conditions: shaded regions indicate one standard deviation. Across both environments and noise settings, our method consistently achieves lower NRMSE and standard deviations than CKO, demonstrating improved robustness and accuracy in long-horizon dynamics prediction.
Preciction ecror Rope (noiseless)
Prediction enror Rope (noisel
Prediction ertor Soft (noiseless)
Prediction error Soft (noise]
Gul
Curt
27.5
KRASE (%)
12
25.0
22.5
15:0
12.5
40 Time step
%
80
100
20
40
80
80
102
20
40
55
130
Time step
Time step
Time
stop
Figure 9: Comparison of normalized root-mean-square error (NRMSE) over 100 rollout steps. The left two figures show results for the Rope environment with and without noise; the right figures show results for the Soft environment, Additive noise is aero-mean Gaussian with standard deviation equal to 10% of the standard deviation of the observation data.
Sample Efficiency. To understand the sample efficiency of our model, we vary the number of data samples used for system identification from to 32, and compare with CKO shown in Figure TO This empirical trend aligns with our theoretical analysis: according to Theorems a bound scales proportionally to VITT, where T is the number of samples. Given that constants such do da, and N remain fixed for chosen feature dimension, the sample size T becomes the primary factor influencing the bound. As the fitting number increases (i.e., T), the prediction eπтоr consistently decreases, confirming the theoretical rate. Moreover, as shown in aaddr our method (Hom + Mean) achieves error bound that scales O(N) for history-to-obserratio part, in contrast CKO (the Hom framework) with a bound scaling as (N Consequently, when both methods are evaluated with the same fitting number, our method consistently yields smaller error bound. These results demonstrate that our theoretical predictions are well supported by empirical observations, validating the improved sample efficiency of our approach.
and [b] the error
16
Fitting number (Rope) Ours -uurs
Ritting number (Ropc) OXD
Attting number (Soft) Ours
Fitting number (Soft) CKD
100-
(50 (16g)
NRMSE 198
15
NAMSE 19G
28 15
20
WRASE
70
65 Time step
60
50
130
20
43
60
80
100
40
60
80
100
20
10
130
Time step
Time step
Time step
Figure 10: Prediction enror (NRMSE) over 100 rollout steps under varying fitting number used for system identification, ranging from 1 to 32, indicated the legend. Results shown for both the Rope and Soft environments. The comparison includes our method and the CKO baseline. As number of samples increases, our method achieves lower prediction error with faster convergence, demonstrating improved sample efficiency consistent with the theoretical rate predicted in Theorems and and
Evaluation Time of Different Embeddings. We measure the computation cost of different em- beddings by recording the time required to predict 100 steps. As shown in Figure [I] Hom+Mean achieves the fastest runtime, outperforming both Dense and Hom. This observation is consistent with the theoretical computation complexity analysis summarized in Table
38

---

Published as a conference paper at ICLR 2026
Rape
1.2
Swim
Evaluation Time (second)
0.8
0.6
0.2
Moa+
Figure 11:
Evaluation time (seconds) for predicting 100 steps with different embeddings.
Mom+Mean achieves the lowest runtime, consistent with the theoretical complexity in Table
9
8
7
T=30
Nodes
6
002
Node:
Weights
5
0,4
4
0.8
4
0.8
1.0
0.50
0.56
0.60
0.86
0.70
0.76
2
8
9
Nodes
Nodes
9
9
8
8
7
7
0.0
T=60
0.2
6
Weights
5
24
2
2
8
0.6
0.0
Nodes
Nodes
Figure 12: Visualization of a 9-object Rope. The left column shows observations at t = 30 and 60, with shading color representing pair-wist wair-wise 2 distances. The middle column depicts non-normalized adaptive Boltzmann-Gibbe weights, highlighting non-uniform neighbor influence (ours). The right column illustrates the misspecified uniform weights, where all neighbors exert equal, normalized influence.
39

---

Published as a conference paper at ICLR 2026
Table 9: Full performance comparison. Results are reported in-distribution and few-shot vali- dation. Control cost and control error are shown as mean + standard deviation, with the best and second- best results highlighted in green and blue, respectively.
In. In-Disrihation Validation
Sew-Shot Validation
Methods
Environments
Control
Control
Control
Control
259.7 =
0.65
282.4
0.69
PCC GraphODE KPM CKO
212.6 64.2
0.58 + 0.12
246.7 71.5
187.2
148.5 39.3
0.36 0.09
0.09
206.1 43.5
0.21
Rope
135.9 28.0
0.00
168.5 45.4
0.04
Ours (vMF)
137.5 + 30.2
0.27 0.09
146.3 49.7
0.13
Ours (Laplace)
130.6 32.7
0.32 0.06
162.8 46.2
0.17 0.05
Ours (Gaussian)
122.6 19.2
0.26 0.08
VAE.
37..7
0.54 0.25
768.1 137.0
0.49
2 49.7
0.51 a 0.22
698.3 102.6
0.45 0.21
GraphODE CKO
194.9
0.33 0.15
472.9=84.1
0.29 0.13
164.6
0.09
388.5
0.20
Soft
160.8 + 26.4
0,17 0.06
357.6
0.20 0.07
(yMF)
Ours (Laplace)
169.5 + 19.1
0.18
382.5 70.9
0.19 0.08
Ours (Gaussion)
158.6 21.7
0.13
344.2 56.3
0.12
VAE
573.1 108.7
0.19
835 113.2
0.92 0.15
PCC
513.3 92.5
0.68 0.15
7328
0.80 0.12
GraphODE KPM CKO Ours (vMF)
417.8 87.9
0.117
693.5 58.2
0.58 0.09
385.5 + 75.2
0.44
5
0.61 0.11
Swim
389.1 = 76.9
0,42 0.13
421.0=
0.44 0.08
392.7 = 73.1
0.45 0.09
152.36629
0.13
(Laplace)
403.
435.7
Ours (Gaussian)
383,7 + 77.8
0.41
404.3 74.2
0,41 0.09
Table 10: Full evaluation under varying noise levels in Power-Grid. With the on random graphs with 100-150 objects. Additive noise is introduced with standard deviations equal to 2%, 5%, 10%, 20% of the standard deviation of observations. "NaN" means the unstable control optimization.
Method
Noiseless
2%
5%
10%
20%
GraphODE
0,58 0.043
0.62 a 0.073
NaN
NaN
NaN NaN 0.85 0.055
Control Error
KPM
0.42 + 0.028
0.50
CKO
0.47 + 0.031
0.48 0.034
0.51 0.027
0.65 + 0.051
(Gaussian)
GraphODE
13:29
1.06
KPM
9.19 1.66
11,70 = 1.68
NaN
NaN 2.27 7,81 1.13
NaN 11.43 4 1.14
Control Cost
11.22
Ours (Gaussian)
4.77 0.61
4,99 0.69
5.72
Table 11: Control error and control cost under different numbers of training trajectories used for few-shot adaptation in the Rope environment. The table reports performance using varying numbers of demonstration trajectories, referred to as "fitting number".
Method
4
16
32
Dense
0.79 + 0.25 0.25
0.41 +: 0.11
0.36 0.12
0.28 + 0.08
0.26 + 0.10
Control Hom
0.74
0.30 0.06
0.30 0.06
0.06
Hom Mean
0.51 + 0.12
0.29 1 0.09
0.26 0.08
0.25 - 0.08
0.23 :: 0.09
Dense
183.8 + 38.6
163.9 42,9
161.6 42.9
160.9
159.6 39.6
Control
147.3
135,9 +:
131.6
35.6
Hom Mean
135.3 31.2
119.8 17.5
122.7 = 19.2
118.7 20.8
111.7 = 24,9
40

---

Published as a conference paper at ICLR 2026
1.6 ABLATION STUDY ABOUT FEATURE DIMENSION
We conduct an ablation study to investigate the impact of feature dimension on control performance in Power-Grid. As shown in Table 12 increasing the feature dimension from 32 consistently improves both control accuracy and control cost in both noiseless and noisy settings. The best per- formance is achieved Dim 32, which yields the lowest control eπrоr and cost, indicating an optimal trade-off between representation capacity and learning complexity. Interestingly, while in- creasing the dimension to 64 slightly improves performance in the noiseless case, leads notable degradation under noise, particularly in terms of control cost. This suggests diminishing returns from higher dimensions and potential overfitting. Moreover, the curse of dimensionality in opti- mization may play role, making convergence more difficult and requiring more iterations for the optimizer to find an effective solution.
Table 12:
Comparison of control error and cost across different feature dimensions without and with 2% standard deviation noise in Power- Grid environment.
Noiseless
2%
Control Error
Ours (Dim 8)
0.24 + 0.008
0.29 + 0.011
Ours (Dim = 16)
0.23 0.007
0.29
0.009
Ours (Dim = 32)
0.21 0.005
0.27 + 0.018
Ours (Dim = 64)
0.20 0.008
0.28 0.013
Control Cost
Ours (Dim 8)
8.74 + 2.78
9.80 + 2.84
Ours (Dim = 16)
6.74 + 1.64
7.44 + 0.44
Ours (Dim 32)
4.77 0.61
4.99 + 0.69
Ours (Dim = 64)
10.45 + 2.44
13.56 + 1.56
41

---

Published as a conference paper at ICLR 2026
1.7 ILLUSTRATION OF CONTROL RESULTS
Time direction
Figure 13: Examples of control performance visualization for Rope, Soft, and Swim. From left to right: the first snapshot shows the initial state, followed by the progression towards the target within a fixed time horizon (indicated by the orange circle in Rope, and the orange box in Soft and Swim).
42

---

Published as a conference paper at ICLR 2026
Timestep
Mode Types
Timestep
Generator
Load
Timestop 120
Timestop 180
Figure 14: Control visualization of 102-node power grid over 4 temporal slices (0.05s per step). Blue and red nodes represent loads and generators, respectively, with values indicating voltage mag- nitude (p.u... Under the control strategy, voltage magnitudes gradually stabilize towards the refer- ence value of p.u.
3.0
Reference Voltage (1.0 p.u.)
Node o Voltage
Node 1 Voltage
2.5
Node 2 Voltage
Node 3 Voltage
Node 4 Voltage
2.0
Node Voltage (p.u.)
1.5
1.0
0.5
0.0
-0.5
-1.0
25
100
125
150
175
200
Time Step (0.05s per step)
Figure 15: Node-wise power grid voltage trajectories over time for the visualized graph topology, showing the stabilization of voltage magnitudes towards the reference value of 1.0 p.u.
43


Reviews:
[{"review_id": "6bExHRlJiN", "reviewer": "Reviewer_8izv", "summary": "This paper proposes a framework named Graph Controllable Embeddings (GCE) to tackle the modeling and controlling problems of stochastic nonlinear dynamics in multi-object systems. It utilizes the linear properties of RKHS to transform complex nonlinear stochastic dynamics into linear relationships within RKHS, allowing the application of well-established linear control methods. GCE also integrates graph neural networks (GNNs) to adapt to dynamic interaction patterns and mean field approximation to reduce sample complexity. The evaluation results validate that GCE outperforms competitive methods.", "questions": "No.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The idea is novel. It leverages the linear properties of RKHS to map originally complex nonlinear stochastic dynamics into RKHS, converting them into linear relationships. This further enables the adaptation of mature linear control methods, effectively overcoming the traditional challenges of nonlinear stochastic control for multi-body systems.\n2. The theoretical analysis is detailed and comprehensive. Rigorous derivations are provided for aspects including the existence and convergence of the GCE framework, the sample complexity of mean field approximation, and the property comparison of different embedding forms (such as Tensor, Dense, and Hom).\n3. The experimental results are comprehensive. It covers multiple scenarios including physical systems (Rope), robotics (Soft, Swim), and power grids (Power-Grid). Additionally, tests are conducted on noise robustness, few-shot generalization, and the effects of different components (e.g., kernel function types, feature dimensions).\n4. The presentation is good. The paper is easy to understand.", "weaknesses": "1. Regarding the theoretical analysis in this paper, although convergence and consistency are mentioned as a key contribution in the Introduction section, multi-step approximate derivations (such as Equation (6) only using first-order approximation and adaptive mean field approximation) may undermine the aforementioned theoretical characteristics, leading to a significant deviation between practical implementation and theoretical guarantees.\n2. Regarding scalability and generalization, although the authors claim that the research can scale to larger graphs and more random topologies, this conclusion relies on the strong assumption of the \"shared approximation operator\" mentioned in Line 264 of the paper. The rationality of this inductive bias remains to be verified, which results in a lack of guarantees for the effectiveness of the method’s scalability.", "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": "", "extra_scores": {}}, {"review_id": "CKMPygCVQs", "reviewer": "Reviewer_UzG7", "summary": "This paper introduces Graph Controllable Embeddings (GCE), a general framework for modeling and controlling stochastic multi-object systems with unknown nonlinear dynamics. The core idea is to leverage Hilbert space embeddings to represent the system's stochastic dynamics in a Reproducing Kernel Hilbert Space (RKHS), where the evolution becomes linear and amenable to control. The framework innovatively combines Graph Neural Networks to encode relational structure with a mean-field approximation. This allows for adaptive, non-uniform weighting of inter-object interactions, overcoming a key limitation of prior work. The method is supported by rigorous theoretical guarantees and demonstrates superior performance, especially in few-shot and noisy settings.", "questions": "- The history is defined as the observation at the previous moment, which is a standard first-order Markov assumption. Is this simple historical representation sufficient for systems that require longer historical dependencies to make accurate predictions (for example, systems with momentum or hysteresis effects)?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The paper's contribution is well-supported by its technical components. The proposed framework is built upon the established theory of Kernel Bayes' Rule and provides a principled way to unify several graph-based embedding approaches under a single lens. Its core technical contribution is the introduction of a mean-field approximation to model adaptive, non-uniform interactions, which addresses a specific, documented limitation of prior methods like CKO. The paper includes theoretical proofs for the convergence and sample complexity of the proposed estimators. The claims are further substantiated through a series of experiments on diverse control tasks, including comparisons against relevant baselines.", "weaknesses": "- The most efficient and recommended model variant, `Hom+Mean`, relies on a strong homogeneity assumption where all objects share the same underlying dynamics operator $C_{O|H}$. This may limit its applicability to highly heterogeneous multi-agent systems, where different types of agents (e.g., a mix of ground robots and aerial drones) possess fundamentally different transition models. The experiments, while strong, do not feature such a deeply heterogeneous environment to test this boundary.\n\n- The framework is fundamentally limited to modeling interactions as a composition of pair-wise relationships. While using a multi-layer GNN allows information to propagate across multiple hops, it cannot capture true, non-decomposable higher-order interactions (e.g., physical effects governed by the angle between three bodies, which is not reducible to a sum of pairs). This restricts the method's utility in domains where many-body physics or complex group dynamics are dominant. The authors rightly acknowledge this as a limitation for future work.\n\n- The current experiment seems to lack a direct contrast with a powerful nonlinear Model-Based RL baseline, and such a contrast might more powerfully demonstrate the necessity of the design choice of forced linearization.", "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": "", "extra_scores": {}}, {"review_id": "vmuf98BVz4", "reviewer": "Reviewer_imNB", "summary": "This paper introduces Graph Controllable Embeddings (GCE) to model and control stochastic, nonlinear, multi-object systems. The core idea is to use Hilbert space embeddings to represent the system's stochastic dynamics (probability distributions) within a Reproducing Kernel Hilbert Space (RKHS). The central insight is that in this space, the complex dynamics become linear, permitting simple LQR control. To implement this efficiently, GCE integrates Graph Neural Networks (GNNs) with a mean field approximation. This novel approach uses adaptive, non-uniform weights to capture complex inter-object dependencies, breaking the \"uniform neighbor\" assumption of prior work and achieving provably low sample complexity. Experiments on robotics and large-scale, random-topology power grids demonstrate GCE is scalable, robust to noise, generalizes to unseen topologies, and outperforms baselines in both in-distribution and few-shot tests.", "questions": "1. How is the \"homogeneity\" assumption (a single shared history operator) reconciled with the clear object-type heterogeneity present in the 'Soft' and 'Rope' environments?\n\n2. The action term's complexity remains quadratic, which appears to be the true scalability bottleneck. Could you elaborate on this design choice and whether a mean-field approximation for actions was also considered?\n\n3. Could you provide more insight into the \"unstable\" behavior of the neural potential function? Does this suggest a fundamental difficulty in learning such energy functions in RKHS, or was it a more straightforward optimization challenge?\n\n4. The framework's theoretical guarantees hinge on the use of characteristic kernels. How is this property ensured, or at least encouraged, when the kernel features are learned by the GNN encoder rather than being predefined?\n5. The linearization is strongly motivated by its use with LQR, which assumes a quadratic cost. How much of the framework's value is retained for tasks with non-quadratic costs, where LQR is inapplicable and other methods (like MPC) must be used?\n6. Eq. (6) replaces ψ_h ⊗ ψ_a with concatenation and neglects higher-order interactions. Can you bound or empirically assess the accuracy–complexity trade-off?”", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The paper is exceptionally clear and well-written. It provides a strong motivation by precisely identifying the limitations of existing deterministic and non-controllable methods. The proposed framework is derived logically, making the theoretical foundations accessible.\n\n* The work is built on a strong theoretical foundation. It provides a principled framework for stochastic multi-object control by extending Hilbert space embeddings (RKHS) to handle probability distributions of system dynamics, a non-trivial extension of prior deterministic approaches.\n\n* The proposed \"Hom + Mean\" model is a significant innovation. By introducing adaptive Boltzmann-Gibbs weights to a mean-field approximation, the method effectively captures non-uniform neighbor interactions. This is a well-motivated improvement over prior methods that rely on misspecified uniform weighting.\n\n* The experimental validation is thorough and compelling. The method is tested across diverse environments, including a challenging large-scale, random-topology, and noisy power grid simulation. The results convincingly demonstrate the model's scalability and superior robustness compared to baselines that fail under noise. Ablation studies on sample efficiency further validate the theoretical benefits of the model.", "weaknesses": "* The \"homogeneity\" assumption, which uses a single shared operator for history dynamics, appears to contradict its application to environments with explicitly heterogeneous object types (like 'Soft' and 'Rope'). It is unclear how one operator can model the distinct dynamics of different object types.\n\n* The scalability claims are weakened by the model's design. While the mean-field approximation reduces the history term's complexity, the action term remains dense and scales quadratically with the number of objects, creating a potential bottleneck for large-scale systems.\n\n* The adaptive weighting mechanism seems sensitive to implementation choices. The paper reports that a flexible, neural-network-based potential function was \"unstable,\" and the best performance relied on a specific kernel choice (Gaussian). This suggests potential fragility and limits the generality of the adaptive component.", "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": "", "extra_scores": {}}, {"review_id": "75LkFGJrUi", "reviewer": "Reviewer_5oTz", "summary": "The paper considers the problem of controlling stochastic nonlinear dynamics with interacting objects, and proposes a graph controllable embedding framework to map the dynamics distribution to RKHS. The authors prove asymptotic convergence given sufficient samples, and show improved control cost/error over existing methods on a set of numerical examples.", "questions": "1. Why do you choose a simple LQR for control in the embedding space? How would MPC perform in this space?\n\n2. How does the proposed method compare with [1-4] in weakness 2?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The idea of extending RKHS for nonlinear control to stochastic and multi-object settings is novel. However, such contribution is not very clear given the presentation of the paper and lack of discussion and comparison with closely related work (see weakness 2)\n\n2. The paper is in general comprehensive, with theoretical analysis of the convergence of the proposed method, and empirical studies compared to several baselines.", "weaknesses": "1. The claim on provably reduced sample complexity is not rigorous. It seems that the reduced complexity is through mean field approximation in the algorithm design. There is no formal proof that such an approximation will not degrade performance.\n\n2. RKHS for nonlinear control has been studied in the community, and the paper seems to miss discussion and comparison with those closely related work [1-4]. \n\n3. The explanation of the benchmark problems is missing key details. The paper considers stochastic control, but only one example has noise (power grid), and how noise comes into the dynamics is not given. For benchmark problems in [5], it seems no noises are in the dynamics, unless the authors adapted the setting, though not mentioned in the paper.\n\n[1] Thorpe, Adam J., and Meeko MK Oishi. \"Stochastic optimal control via Hilbert space embeddings of distributions.\" 2021 60th IEEE Conference on Decision and Control (CDC). IEEE, 2021.\n\n[2] Romao, Licio, Ashish R. Hota, and Alessandro Abate. \"Distributionally robust optimal and safe control of stochastic systems via kernel conditional mean embedding.\" 2023 62nd IEEE Conference on Decision and Control (CDC). IEEE, 2023.\n\n[3] Rawlik, Konrad, Marc Toussaint, and Sethu Vijayakumar. \"Path integral control by reproducing kernel Hilbert space embedding.\" arXiv preprint arXiv:1208.2523 (2012).\n\n[4] Bevanda, Petar, et al. \"Nonparametric control Koopman operators.\" arXiv preprint arXiv:2405.07312 (2024).\n\n[5] Li, Yunzhu, et al. \"Learning compositional koopman operators for model-based control.\" arXiv preprint arXiv:1910.08264 (2019).", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)