Koopman-Assisted Trajectory Synthesis: A Data Augmentation Framework for Offline Imitation Learning
Authors: Jin Wang
Pengcheng He
Ke Jiang
Xiaoyang Tan
Venue: ICLR 2026 Poster (2026)
Keywords: Offline Imitation Learning; Offline Reinforcement Learning; Data Augmentation
Primary area: reinforcement learning

Abstract: Data augmentation plays a pivotal role in offline imitation learning (IL) by alleviating covariate shift, yet existing methods remain constrained. Single-step techniques frequently violate underlying system dynamics, whereas trajectory-level approaches are plagued by compounding errors or scalability limitations. Even recent Koopman-based methods typically function at the single-step level, encountering computational bottlenecks due to action-equivariance requirements and vulnerability to approximation errors. To overcome these challenges, we introduce Koopman-Assisted Trajectory Synthesis (KATS), a novel framework for generating complete, multi-step trajectories. By operating at the trajectory level, KATS effectively mitigates compounding errors. It leverages a state-equivariant assumption to ensure computational efficiency and scalability, while incorporating a refined generator matrix to bolster robustness against Koopman approximation errors. This approach enables a more direct and efficacious mechanism for distribution matching in offline IL. Extensive experiments demonstrate that KATS substantially enhances policy performance and achieves state-of-the-art (SOTA) results, especially in demanding scenarios with narrow expert data distributions.

Published as a conference paper at ICLR 2026
KOOPMAN-ASSISTER TRAJECTORY SYNTHESIS: A DATA AUGMENTATION FRAMEWORK FOR OFFLINE IMITATION LEARNING
Jin Wang, Pengeheng He, Ke Jiang, Xiaoyang Tan" College of Computer Science and Technology Nanjing University of Acronautics and Astronautics
MIIT Key Laboratory of Patten Analysis and Machine Intelligence Nanjing, Ji 210000, China
{jinwl 999, yukina233 ke-jiang, x. (at)aannnnaaae edu, en
ABSTRACT
In offline imitation learning (IL), data augmentation is essential for addressing co- variate shift, yet existing methods face a trade-off: single-step techniques risk dy- namic inconsistency, while trajectory-level approaches suffer from compounding errors or scalability challenges. Recent Koopman-based methods, while promis- ing, have been limited to single-step applications, creating computational bur- dens and vulnerabilities to approximation error. This paper introduces Koopman- Assisted Trajectory Synthesis (KATS), framework designed to resolve these trade-offs. Our primary contribution is a novel trajectory-level synthesis pro- cess that inherently avoids the compounding errors of recursive rollouts. Second, we introduce a state-equivarias equivariint Koopman representation, a key innovation that ensures computational efficiency and scalability, unlike prior action-equivariant models. Third, we bolster robustness by integrating a refined generator matrix to counteract operator approximation errors. Through extensive experimentation, we demonstrate that KATS yields substantial improvements in policy performance and achieves new state-of-the-an results, proving especially effective in challeng- ing,
regimes.
INTRODUCTION
Imitation Learning (IL) enables agents acquire skills from expert demonstrations and has proven highly effective in fields like robotic control and autonomous driving Xie et al.) (2024); Teng et al.) (2022). A primary limitation of classic IL methods, such behavioral cloning (BC), is their suscep- tibility to "covariate "oift" et al. (201T). This issue stems from a mismatch between the training distribution and the states visited during deployment, which triggers a compounding of errors and rapid performance decline Zare et al. (2024). The problem is especially severe in offline imitation fixed dataset and cannot interact with the environment to
learning, as the agent is confined correct its deviations Prudencio et al. (2023).
The challenge of covariate shift is typically addressed through two distinct paradigms: algorithm- centric and data-centric methods. The former centers on enhancing algorithmic robustness, for example, by learning an explicit dynamics model to foresee future states Zhang et al. (2023) or by applying regularization to penalize actions in out-of-distributio (OOD) regions Chang et al. (222)) The latter, in contrast, focuses on improving the training data itself. Data augmentation stands out leading data-centric technique, as it offers cost-effective means enrich the dataset, enhance its diversity, and mitigate covariate shift without requiring further interaction with the environment
Samadi et al. 2024).
The success of data augmentation hinges on generating high-quality, dynamically consistent data Kolev et al. (2024). Common single-step techniques like Mixup ignore temporal context, risking the creation of physically unrealistic states that can bias the model. Synthesizing entire trajectories is
"Corresponding Author

---

Published as a conference paper at ICLR 2026
more robust approach, but it also faces hurdles. Forward-dynamtn models suffer from compounding errors that lead to divergent trajectories, while powerful generative models like GANs Chen et al. (2024) or Diffusion models]anmen
al. (2022) often too data-hun hungry for limited offline datasets and struggle to generate novel yet plausible behaviors without a strong inductive bias.
This reveals the central need in offline IL for framework that can generate novel, consistent tra- jectories without long- nog-horizon error accumulation. The Koopman operator theory is uniquely suited fill this gap. By transforming nonlinear dynamics into a linear evolution in a higher-dimennional space Mezić (2021), it simplifies control and analysis while also allowing for the easy integration of physical priors like symmetries. This feature directly enhances robustness with limited data, addressing a key failure point of other generative models. Consequently, Koopman-basee data aug- mentation has garnered significant interest, spawning methods such as Koopman Q-learning (KFC) Weissenbacher et al. (2022), Koopman Model Predictive Control Wang et al. (2022) and Koopman linear quadratic regulators Shi Meng (2022).
Existing Koopman-based methods, despite their potential, are constrained in their application to complex, high-dimensional control tasks Colbrook et al. (2024). A key issue stems from strategies that operate at the single-step transition level, such as the data augmentation in KFC Wang et al. (2022) Weissenbacher et al. (2022). By focusing on immediate state changes, these methods struggle to model the long- serg-termmmm temporal structures essential for coherent, multi-step task execution Efroni et al. (202T). This limitation highlights clear need to explore the generation of full trajectories within the Koopman framework and to develop the associated theoretical guarantees for such an approach.
A second limitation is the potential for significant computational overhead. In the KFC method, for example, the reliance on an action-equivariance assumption creates linear scaling relationship between the action dimension and both model complexity and inference cost, which constrains its scalability. A final, crucial weakness is the failure of these methods to properly address approxi- mation errors from the Koopman lincarization. Because the reliability of local linear approxima- tions diminishes in areas with minimal or slow dynamics, the integrity of augmented data generated therein is compromised. The introduction of such unreliable data risks biasing the training process, which can guide the learned policy toward non-physical or unstable outcomes.
To overcome these challenges, we introduce Koopman-Assinn novel framework for offline imitation learning. Instead of traditional step-by-step rollouts, KATS generates new data by modeling entire expert trajectories in a learned latent space.
Trajectory Synthesis (KATS), a
Our main contributions are follows.
We introduce trajectory-level augmentation method that uses entire expert demonstra- tions its base unit. This approach mitigates the compounding errors common in state- space rollouts and ensures generated trajectories are dynamically consistent within the lin- ear Koopman space.
Our framework leverages a state- eerivariuuiiiatt assumption, which avoids the severe computa- tional and memory costs of prior equivvaran approaches. This design makes KATS highly efficient and scalable for complex tasks.
We design a refined symmetric generator matrix that makes our model more robust to the inherent approximation errors of dimte-ddieniiooall Koopman representations, improving the quality of synthesized trajectories.
RELATED WORK
Offline Imitation Learning and Distribution Shift. A primary challenge in offline IL. is distribu- tion shift, where policies fail in out-of-distrrbutiio (OOD) states due to compounding errors Xu et al. 2020)) Conservative methods address this by constraining the policy to the expert data distribution Kumar et al. (2020); Chang et al. (2021). Other approaches use learned dynamics models to plan back known regions Shao al. (2024); Janner et al. (2022) explicitly estimate the data support to confine the policy Xu et al. (202bb). Our work complements these by focusing on data-centric mitigation through principled trajectory synthesis.
2

---

Published as a conference paper at ICLR 2026
Data Augmentation in Imitation Learning. Data augmentation aims alleviate distribution shift by expanding the dataset. While simple techniques like noise injection or random visual crops common Laskey et al. (2017): Yang et al. (2022), they often produce dynamically inconsistent samples that can harm learning Gong et al. (2021). Unlike these methods, our work leverages Koopman theory's strong inductive bias for dynamics generate entire trajectories that both physically plausible and behaviorally consistent with the expert.
Koopman Theory in Control and Reinforcement Learning. Koopman operator theory enables the analysis of nonlinear systems via linear representations in a lifted space Mauroy et al. (2020). and has been widely used for learning dynamics Zhao et al. (2023). designing controllers Wang al. (202)); Shi & Meng (2022), and policy learning in RL Weissenbacher et al. (2022). While some works have explored Koopman for data augmentation Jang et al. (2023), our method is distinct. We are the first to systematically develop trajectory-level, state-equivariant Koopman augmentation framework specifically for offline II.. By integrating it with an Inverse Dynamies Model (IDM), generate complete, dynamically consistent, and behaviorally plausible state-action sequences to tackle distribution shift.
PRELIMINARIES
This section provides the necessary background on offline imitation learning, focusing on the chal- lenge of distribution shift and Koopman operator theory, detailing its definition and application in control systems.
3.1 OFFLINE IMITATION LEARNING AND DISTRIBUTION SHIFT
Imitation Learning (IL) aims train policy (a)a)) to mimic expert policy TE using static dataset of expert demonstrations D = {s,, 0;) , A common approach, Behavioral Cloning, frames this supervised learning problem, minimizing loss function such the Mean Squared Error between the policy's actions and the expert's:
LBC = Σ ||: (
The primary challenge in this offline setting is distribution shift. The policy is trained on the expert's state distribution. During deployment, small errors can lead the agent states not seen in D. Since the policy behavior is undefined for these states and no new data can be collected, these errors compound, causing the agent's performance to degrade significantly.
(OOD)
3.2 KOOPMAN-BASHE DATA AUGMENTATION
Koopman operator theory provides powerful framework for linearizing nonlinear dynamics 81+1 F(st, ar) high-dimensional lifted space. Recent methods, such as Koopman Forward Conser- vative Q-learning (KFC) Weissenbacher et al. (2022) Wang et al. (2022). leverage this theory for structured data augmentation, moving beyond simple noise injection.
In particular, the KFC method aims to identify a set of dimensional observable functions, #; g(s)), that transform a system's 's state from its original domain to a higher-dimensional feature space. The objective of this transformation is twofold: to lincarize the system's ' nonlinear dynamics and to ensure the resulting linear model is action rauiii-equiaaiiant In particular, this learned action-equivariant Koopman operator consists of matrices (K,, K2) that approximate the system's 's dynamics in the lifted space as a bilinear model:
21+1
Kg
Σ
Kyat, 2t-
where 2, and 2 are, respectively, two points in the lifted space, corresponding to a pair of states (81,84+1) in the original space.
3

---

Published as a conference paper at ICLR 2026
Then, instead of adding arbitrary noise, KFC finds symmetric generator matrix that represents dynamically consistent enhancements. This is achieved by solving the Sylvester equation, which im- poses that the transformation commutes with the system dynamies. Let K.. Kg + Σ K denote the complete action-dependent Koopman operator for a given action oç.:
(e) Koopman Forward inverse Control (Baseline)
44
Zu
21
72
53
21
Oas
d1
agg
72
(b) Koopman Assisted Trajoctory Synthesis
04
0:
47
Ze
di
02
44
0g
1
Xg
22
62
02
6
0
0g
Figure Overview of KATS framework: An Expert-Guided Koopman Model K learned from expert trajectories. Commuting transformations (a) with * and an inverse dynamics model generate augmented data to train policy my.
Ka Ja. 0., Ka, II
(3)
As illustrated in Figure 1 (a), augment original transition (8 az: 8441, rt), after embedding the states and st+1 into their singh-dimensinall latent representations, 2, and žt+1, the transformation oo, is applied generate augmented counterparts: 2 oa, and 2 aa, 2t+1- One can also use decoder to project the generated new latent vectors back into the original state space, resulting in the augmented states sf and $4+1- Finally, complete, dynamically consistent transition tuple (s. ar. 1. is synthesized by pairing these new states with the original action and reward Weissenbacher et al. (2022).
Despite its promise, the aforementioned action-equivarian/ Koopman-based data augmentation method has been limited to single-step applications. This restriction imposes significant compu- tational burdens and introduces susceptibility to approximation errors. To address these challenges, we proposed policy equivariant Koopman-bald based data augmentation. what follows, we first delve into the theoretical underpinnings of this and subsequently propose practically viable solution.
4 KOOPMAN-ASSISTED TRAJECTORY GENERATION WITH POLICY-EQUIVARIANCE
In this section, we provide the theoretical foundation for our Koopman-Assistee Trajectory Synthe- sis (KATS) framework. Our central argument is that the symmetries of a closed-loop dynamical system, driven by a fixed expert policy, are directly reflected as commutation properties of its as- sociated Koopman operator. We begin by defining this notion of symmetry, then establish formal equivalence to an algebraic property of the operator, and finally show how this equivalence provides rigorous mechanism for generating novel, policy-consiste trajectories.
Definition 1 (Policy-Equue Equivariant Dynamical System). Consider the closed-loop dynamical system described by the discrete-timg time mapping $t+1 = Fz (st). where s, E is the state and π S A fixed policy. Let E, be a group with elements S - S that act the state space. The
4

---

Published as a conference paper at ICLR 2026
system is called E-policy--quuivarine for all E and all E S, the following condition holds: (a = a. (s)). This implies that transforming the state and then applying the policy dynamics is equivalent to applying the dynamics first and then transforming the result.
To analyze these potentially non- non-linear dynamics within a linear framework, we lift the system from the state space to a space of observables using the Koopman operator.
Definition 2 (Koopman Operator for Policy). Let K(()) {g S RR be the space of scalar- valued observable functions on the state space S. The Koopman operator K : KK(S) + K(() associated with the closed-loop system Fs is a linear operator defined by: [K* 2 g(F((s)) In a dimiee-dmmessonnl approximation learned from data, K is represented by a matrix that propagates a latent state vector 2ς E(s() forward in time, i.e., 21+1 Kz
key step connecting the geometric state-space symmetry a to the algebraic operator K₂ is to define how acts on the space of observables. This is achieved through an induced group action.
Lemma 3 (Induced Group Action on Koopman Space). A state space transformation a S + S induces a linear action on the space of observables K(SS), defined by the pullback: [ag](s) 2 g(o (s)).
With these definitions in place, we can now state the central theorem of our analysis. It establishes a formal equivalence between the policy equivariance of the dynamics and a simple commutation relation involving the Koopman operator.
Theorem (Equivalence of Policy Equivariance and Koopman Commutation). Let Fa be closed- loop dynamical system and K be its corresponding Koopman operator. The following two state- ments are equivalent:
1. The system is E-policy-equivvriann::
all o E E
2. The Koopman operator commutes with the induced group action (where a acts on a func- tion g as (sg)(s) g(o=1s)): Kz oK for all E.
Theorem D is not merely a theoretical curiosity; it provides the direct mechanism for our data aug- mentation strategy. The following theorem makes this connection explicit, showing how the com- mutation property can be leveraged to generate new, valid system trajectories.
Theorem 2 (Generating Policy-Consisten Trajectories via Symmetry). equivariant system with Koopman operator K,, and let E E. Let { E(s:)} be the latent representation of trajectory under policy such that 2+41 = Kz Define new trajectory { by ân a 2t. Then this new trajectory evolves according to the same Koopman dynamics: 2141 = K2 Consequently, the decoded trajectory {&g = D(2)) is also a valid trajectory under the closed-loop dynamics of policy N.
Let Fr be a E-policy-
Theorem 3 (Bounded Error for Transformed Trajectories). Let (2 } be latent expert trajectory FF with a sequence of one-step prediction errors { where e = 21+1 K2F Let a be an expert- policy equivariant transformation with commutation error A. Then the one-step prediction error ét for the transformed trajectory (áe = ozxB) } is given by:
é o + AzF.
(4)
Consequently, its norm is bounded as:
|| < |||||| |c + ||||||
(5)
This result formally decomposes the error a synthesized trajectory into two distinet sources:
1. Propagated Model Error (||r)]lee []):
This term reflects the inherent dynamics prediction error 6,, scaled by the magnitude the transformation.
2. Symmetry Violation Error (i) =|=]==/+)== This term stems directly from the imperfect com- mutation between the learned symmetry and the dynamics model.
This decomposition provides a strong theoretical justification for our objective: it is crucial not only to minimize the standard prediction loss (reducing [|<e|] but also explicitly regularize the model to
5

---

Published as a conference paper at ICLR 2026
minimize the commutation residual |||] Please refer to Implication of the Bound in the Appendix for further discussion.
In summary, this analysis formally motivates approach: by learning a Koopman operator K from expert data and then finding linear transformations a that satisfy the commutation relation Ko oK O, are effectively identifying the underlying symmetries of the expert's behavior. Theorem - guarantees that applying these learned symmetries to existing trajectories will produce new data that consistent with the expert's policy dynamics.
5 A PRACTICAL ALGORITHM
Motivated by the above analysis, in this section we detail our practical algorithm Koopman- Assisted Trajectory Synthesis (KATS). designed for augmenting datasets trajectory- offline imitation learning. The primary goal of KATS is to generate novel yet realistic data to improve policy leaming. To this end, the framework synthesizes trajectories with three critical properties: dynamic consistency, behavioral plausibility, and computational efficiency. The KATS pipeline comprises three distinct stages: (1) trajectory synthesis via behavior-driven Koopman operator, (2) robust- ness enhancement using an adaptive symmetry generation scheme, and (3) action inference through a decoupled inverse dynamics model. The entire process is visualized in Figure [
5.1 KOOPMAN-ASSISTA TRAJECTORY SYNTHESIS
The first step of KATS is the construction of generative model for expert trajectories. A key in- novation of our method is the decision to model the closed-loop system dynamics governed by the expert policy, expressed
st+1 s G(s)), as opposed to the conventional open-loop dynamics, S141 = F(t,, a.). Consequently, the expert policy Tg (axl (sx) becomes implicitly embedded within [a
the transition function G. The practical implementation of this principle involves training toencoder (Eg, x to map states into a latent space Z where the dynamics are linear. Within this
linearized space, single, eetion-indeeendent Koopman operator K is then learned to propagate the system state:
(8 ( K
The model is trained by minimizing a combined loss function over the expert dataset Dg, comprising state reconstruction and latent space prediction:
Lrecon Bgw || (Fs(())))
Lkoopman (s., DD (s(++1) (s)
(8)
The proposed action-independent formulation offers two fundamental improvements upon prior methods such as WeCeeessseeaeheer et al. (2022). The first is guarantee of behavioral plausibility: since the operator K is trained directly on expert closed-loop state transitions (st, $t+1), it inherently embeds the expert's 's policy. Since the symmetry transformation commutes with the Koopman oper- ator K, the system dynamics are equívariant under this transformation. Consequently, applying the symmetry generated trajectory yields a new trajectory that remains consistent with the learned policy. The second improvement is dramatic increase in efficiency. Our approach obviates the need for per-action modeling by learning only single operator and its corresponding symmetry basis o, in single offline step. This not only reduces storage complexity from O d²) to a constant but also entirely removes the computational bottleneck associated with per-action processing.
5.2 ADAPTIVE SYMMETRY GENERATION FOR ROBUSTNESS
Any model-based approach is subject to its own approximation emrors. A Koopman model may be inaccurate in regions of the state space where expert data is sparse. Naively generating data in such regions with a flawed model can inject noise and degrade policy performance.
To mitigate this, KATS introduces an adaptive data augmentation scheme. The core idea is to prior- itize data synthesis in regions where the model is most uncertain. Instead of merely finding fixed basis by solving the homogeneous Sylvester equation (Ka ok = 0), we learn the transformation through an optimization process weighted by the Koopman model's prediction error. Specifically, the loss for is defined as:
C. ( Ws, ((((1))) (=)

---

Published as a conference paper at ICLR 2026
where 2; (s(ss and the weight Wa is explicitly tied to the model's local enror:
Ws:
Ka,ll)
(10)
Here, a temperature hyperparameter. This weighting scheme forces the learning of more strongly satisfy the symmetry property in areas where the original operator K struggles, In this way, KATS adaptively generates samples where they are most needed, turning potential weakness into a targeted exploration strategy.
5.3 DECOUPLED ACTION INFERENCE VIA INVERSE DYNAMICS
requirement: the generated symmetric state pairs (s/; 1)41 lack corresponding actions. To address this, we decouple action inference from dynamics modeling by training separate Inverse Dynamics Model (IDM), fIDM S S A. The IDM learns the mapping from state transitions to actions directly from the expert data:
Our action-independent modeling approach creates
LIDM
E 844 DD a fIDM(Sr, st+1)ll
(11)
Once trained, this model infers a plausible action of fIDM (s/1 $441) for each generated state pair, creating complete augmented tuple (s/, a/, 8(41).
5.4 IMPLEMENTATION DETAILS
The complete KATS pipeline is summarized in Algorithm 1 In the final step, we combine the original expert dataset De with our generated data Daug to form a comprehensive training set, Dinal DE U Daug- An imitation learning policy Te(a)s) is then trained this enriched dataset, typically via behavioral cloning, to learn a more robust and generalizable policy:
LBC (0) = -E log To (a]ss)]
(12)
Algorithm Koopman-Assisted Trajectory Synthesis
Require:
(T.). Koopman models (Liş, De, K;) trained via Eq. and Eq. 8 IDM fIDM trained via Eq.
Expert trajectories DE
Ensure: Augmented dataset of trajectories Daug.
1: Initialize Daug + De.
2: Learn symmetry operators {oا)ر/ |= satisfying K, r Ro oK, via the objective in Eq.
3: for each expert trajectory = (so..... ST) in Dg do
4: for each symmetry Jj the basis do
5: Define the end-to-end state transformation Gý(8) = (Ds 002 (x)(s).
6: Generate new state trajectory:
7: (s/,..., s < (Gy(su),... (ST))..
8:
Synthesize the corresponding action trajectory:
9: (a6. a4-1) (JIDM (h) $1),.., fIM (IM 8/)))
10: Assemble and add the new trajectory 7' (5 do ... . to Dang
11: end for
12: end for
13: return Dang.
EXPERIMENTS
In this section, we conduct a comprehensive empirical evaluation our proposed method KATS. Our experiments are designed rigorously assess its effectiveness, versatility, and robustness. We use the challenging D4RL benchmark Fu et al, (2020), which includes diverse locomotion (Mujoco), navigation (Antmaze) and manipulation (Adroit) tasks, serving standard to evaluate offline learning algorithms. All baseline algorithms are configured following the hyperparameters specified in their original publications ensure a fair comparison.

---

Published as a conference paper at ICLR 2026
Table Comparative performance on custom control tasks, Best results for each task in bold.
Domain
Task Name
BC
SRA
MILO
KFC+BC
KATS+BC
antmazze-uaaze
74.0 + 1.2
85.3 1.1
79.1
96.9 0.8
antmaze- caamaaze-umzz--ae diverse
64.0 2.0
81.5
2.1
6.3
90.1
6.1
78.5 2.5
3.3
2.3
82.7
AntMaze
uulmaze-mediunm diveerse
53.7
60.2
2..0
2.2
67.3
untmaae- large-play
2.2
48.3 2.0
1.9
4.5
59.3
antmaaz-latte lare--lrrge--eees
24.9
39.5 1.5
2.0
1.71
44.2 2.7
halffhhctah-mcciiu-eeee
1.82
63.4 - 3.5
1.5
60.9
81.2+
Gym
walker-medium--eyy
107.5+2.00
1041+4.88
9..4+
100.4+1.9
110.4 4.6
meeper-medium-exper
104.5 5.0
5.4
1.3
112.7 3.6
pen- sen-human
3.2
40.2 2.0
1.8
1.5
69.2
3.9
pen-cloned
2..2
4 2.3
2.0
1.7
81.3 1 4.7
hammer-hume
0.8
5.0
5.9
4.6
7.2003
Adroit
hammer-cloned
2.1
2.3+0.4
2.7+0.4
2.8 0.3
4.2 0.5
door-human
9.5
1533=13
2.00
1..1
37.2 2.1
door-cloned
0.2
0.2
2.1 2.1
1.6
4.9
6.1
MAIN RESULTS
We benchmark KATS against suite prominent offline imitation learning algorithms: Behavior
Cloning (BC): The fundamental baseline for imitation learning; SRA [hao et al. (2024): A state-
of-the-art method that employs model-based reverse augmentation; MILO Chang et al. (202T): A
strong model-based algorithm employing policy constraints; KFC+BC Weissenbacher et al. (2022)
baseline that combines Behavioral Cloning (BC) with data augmentation from KFC. Our expen- mental evaluation focuses specifically D4RL. datasets characterized by data sparsity and limited number expert demonstrations, as these scenarios pose the greatest challenge to offline learning algorithms.
Table [ gives the comparative results for the task of offline IL, which unequivocally demonstrate the superiority of our proposed method, KATS, Across all 15 challenging control tasks, KATS+BC significantly outperforming all baselines, including vanilla Behavioral Cloning (BC), prior augmen- tation methods (SRA, MILO), and its direct predecessor, KFC+BC. The most critical insight comes from the direct comparison with KFC+BC. KATS+BC achieves massive performance gains in ev- ery single task, such as improving the score from 68.0 to 96.9 on antmaze-umaze. This significant improvement validates our core contribution: the action-independent formulation. By modeling the system's ' closed-loop dynamics, KATS effectively captures the expert's 's policy. Our empirical results suggest that this approach yields trajectories that more behaviorally consistent than those from cetion-conditioned models like KFC.
Figure 2 gives some illustration of trajectories in the 'maze2d" environment. The visualizations reveal that our method generates paths with greater symmetry at corners and produces a more diverse set of samples in difficult areas, enriching the training data where it is most needed.
O1iinal
skyre
Siyra
Skyre
Styra
Skyra
Styre
Sg
Figure 2: Ablation the Blend Ratio of Symmetric Trajectories in maze2d.
Furthermore, the robust performance of KATS sparse-reward navigation tasks (AntMaze) and noisy, suboptimal human demonstration data (Adroit) highlights its versatility and practical utility. effectively extracts and amplifies the core skills from limited data, critical capability for real-world
8

---

Published as a conference paper at ICLR 2026
Table 2: Performance comparison against offline RL baselines on few-shot D4RL datasets. The best result in each row is highlighted in bold.
Task
Size tratio)
BC
CQL
IQL.
DOGE
POR
TELS
Diffuser
KFC+CQI.
KATS+BC
Hopper
10k (19)
69.1
724
78.529.3
834.1779
87.4.18.
N/A
74,51338
89.127.6
97.1.1004
Hopper-m
10k. (0.5%)
482.16.7
2..9.4.
34.118.7
50.52252
37.916.1
10091668
70.1.6.5
88.2111.3
87.211655
Halfchcetab-e
(1%)
634412211
734-76
57.2417.1
64.7417
86.2152
68.9:41
83,49977
102.7%81
Halfcheetal-me
(0.5%)
4144+9.2
26-5-108
10.5+888
267466
34.7426.
40.7-122
63.288
714+89
75.2+10.4
Walker2dee
10k (35)
76.4117.2
81.211111 73.2--16.1
659188
77.1115.1
N/A
83.316.7
86.111222
109.411.22
2alker2d-me
10K (0.5%)
10.9
19.1=14.4
26.511866
15.3
20.1288.6
87.-13.3
74.2:9.0
81.9669
92.5:18.1
Animaz---d
O.IM (10%)
47.11:19.8
0.5=0.1
34.6118.
41.7::199
42.12:142
60,911699
59.1::1.77
62.1174
10,5
Antmare-u
D.M (10%)
62.3+27.1
0.1-00
65.1+19.4
3-34
61+73
N8.7+7.7
62.1+99
73.1+7.1
X3.1+7.22
Antmaze m-d
0.0M(10%)
36.71220.5
NIA
30.1111.5
NIA
NIA
55.2114.2
51.18..7
56.71100.1
Antmaz-mpp
O (10%)
47.51:16.9
NIA
12515.4
N/A
N/A
62.9117.8
6091122
67.519.3
7131152
Animane--d
DIM(U19)
25.1+12.6
NIA
3.6+4.11
NIA
N/A
39.8+14,1
38.5+109
35.1+104
42.--11.9
Antmaze-4pp
DJM (105)
20.1+15.7
N/A
3.5+411
N/A
N/A
47.3+13.1
47.6+88
49.1+8.1
51.4
applications. In summary, these results confirm that KATS is a highly effective and general-purpose data augmentation technique that significantly pushes the frontier of offline imitation learning.
6.2 RESULTS ON THE OFFLINE RL TASKS
We also validate our method's efficacy in the offline reinforcement learning (RL) setting, which requires assigning rewards to augmented data. Adopting the strategy from KFC, we assign the reward for each synthetic transition by reusing the value from the corresponding step in the source demonstration. Our empirical evaluation includes a direct comparison with competitive offline RL algorithms: TELS Cheng et al. (2025), which employs temporal inverse dynamics regularization; POR Xu et al. (2022a), which leverages state-value guidance and an inverse dynamics model; and DOGE et al. (2022) which focuses on generalizing offline RL by exploiting data geometry.
Table (2 gives the results. One can see from the Table that KATS demonstrates remarkable efficacy in the challenging few-shot offline RL setting, achieving state-of-the-art art results on 10 of the 12 tasks. Our results reveal a crucial insight: augmenting the dataset with high-fidelity, behaviorally consistent trajectories enables simple Behavioral Cloning (BC) policy decisively outperform complex algorithms designed to handle distributional shift.
Notably, these outstanding results were achieved using simple reward assignment heuristic: reusing rewards from the source trajectory. While theoretically suboptimal, the success of this strat- egy strongly implies that the primary driver of performance is the high behavioral fidelity of the generated state-action trajectories. This finding is significant, as it shows that a superior generative model of behavior can allow a simple learning algorithm to surpass complex methods designed for value estimation and policy constraint in offline RL.
In conclusion, the results in rablerJybustty demonstrate that KATS is not merely an imitation learn- ing technique but powerful and versatile data augmentation framework that effectively addresses the challenges of offline RI., especially in the critical and practical context of data scareity.
6.3 COMPATIBILITY AS PLUG-AND-PLAY AUGMENTOR
A key hypothesis is that KATS's ' data augmen- tation is model-agnostic and can benefit any of- fline IL algorithm. To test this, we treat KATS as a pre-processing step and apply it to the data sets used by BC and SRA Shao et al. (2024).
16
For baseline (BC, SRA), we compare its per- formance when trained on the original D4RL dataset versus being trained on the dataset aug- mented by KATS. We denote the augmented versions as 'BC + KATS and 'SRA+ + KATS'.
Figure 3: Performance comparison on offline II. benchmarks with KATS augmentation.
As shown in Figure applying the KATS aug- mentation provides a significant boost in performance for BC and SRA. The improvement is par- ticularly substantial for BC, which suffers from covariate shift, problem that is mitigated by the
9

---

Published as a conference paper at ICLR 2026
richer, symmetrically complete data. The fact that KATS enhances strong baseline like SRA validates its utility as plug-and and-aaay general-purpose purpose module to improve data efficiency in offline imitation learning.
6.4 ABLATION STUDY ON COMPONENT EFFICACY.
To rigorously validate our design choices, we conducted an ablation study on the antmaze and maz mede benchmarks (Table 3) We compared the full KATS model against three variants:
KATS-OAA which replaces the learnable ay-network with fixed analytical matrix;
KATS w/o IDM, which removes the Inverse Dynamics Model the necessity of action correction;
KFC+BC, which employs an action-coupled encoder instead of our decoupled design.
As shown in Table 3 the full KATS model consistently outperforms all variants. The significant performance drop in w/o IDM confirms that the IDM is critical for resolving state-action mismatches in augmented trajectories. Furthermore, KATS surpasses KATS-o, a demonstrating that the adaptive ag getnewwrrk captures more effective symmetry patterns than rigid analytical solutions.
Table 3: Ablation study on antmaze and maze2d.
Domain
Task Name
BC
KFC+BC
KATS w/o IDM
KATS-
KATS (Ours)
antmaze-uman
74.0 1.2
79.1
78.4 4.8
87.3 2.6
96.9 0.8
uutmaz--uaaze-divers.
2.0
1.5
3.6
3.3
9..1
nummaze-medium-play
6.1
2.5
2.8
7..4
82.7
AntMaze
natmaz--mediummdiivrse
4.5
6.8
4.5
67.3 2.5
larmaae-lrggeppaay
35.8
4.5
4.3
5.07+4.0
59.3 1.5
lutmaze-laree-diveree
24.9
2.7
271
2.6
44.2 2.7
maze2d-umaze
89.2 6.0
8..
100.2
113.2 + 5.6
Maze2d
63.7 4.9
66.3 = 12.1
90.8
108.7 1 4.9
maze2d-large
1.3
5..9
41.5 15.8
89.2 12.2
100.1 + 7.2
CONCLUSION AND FUTURE WORK
This paper introduced KATS, a novel data augmentation framework that leverages Koopman the- ory to address the critical distribution shift problem in offline imitation learning. By identifying and applying latent symmetries the system dynamics, KATS generates diverse, yet dynamically consistent data. Our experiments validate that this principled approach effectively enhances the robustness and generalization the learned policy against challenging out-of-diitrrbution states.
Limitations. One limitation of the proposed method is that currently relatively simple Koop- man operator is learnt from data, hence interesting to explore more sophisticated Koopman architectures to capture highly nonlinear dynamics and integrating uncertainty principles.
This work is partially supported by National Natural Science Foundation of China (62476128) and National Key R&D program of China (2021ZD01 13203).
Acknowledgements
Reproducibility. We commit to releasing the supplementary material for review. Our implementation uses PyTorch. All model architectures, hyperparameters, and training procedures are detailed in Appendix D All results averaged over random seeds.
source code upon publication and provide it in
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Jonathan Chang, Masatoshi Uchara, Dhruv Sreenivas, Rahul Kidambi, and Wen Sun. Mitigating covariate shift imitation learning via offline data with partial coverage. Advances in Neural Information Processing Systems, 34:965 979, 2021.
Jiayu Chen, Bhargav Ganguly, Yang Xu, Yongsheng Mei, Tian Lan, and Vaneet Aggarwal. Deep generative models for offline policy learning: Tutorial, survey, and perspectives on future direc- tions. arXiv preprint (xXi::2402.3377 2024.
Peng Cheng, Zhihao Wu, Jianxiong Li, Haoran Xu, Wei Sun, Youfang Lin, Yunxin Liu, and Xi- anyuan Zhan. Pushing the limit of small-efficient offine reinforcement learning, 2025. URL
https::/openrev///
net forum? add-mdddmmmmPP....
Matthew J Colbrook, Igor Mezić, and Alexei Stepanenko. Limits and powers of koopman learning. arXiv preprint arXiv:2407.063: /2, 2024.
Yonathan Efroni, Nadav Merlis, and Shie Mannor. Reinforcement learning with trajectory feedback. In Proceedings of the AAAI conference on artificial intelligence, volume 35, pp. 7288-7295, 2021.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep
data-driven reinforcement learning. arXiv e-prints, pp. arXiv-2004, 2020.
Chengyue Gong, Dilin Wang, Meng Li, Vikas Chandra, and Qiang Liu. Keepaugment: simple information-preservinn data augmentation approach. In 202/ IEEE/CVF Conference Com- puter Vision and Pattern Recognition (CVPR), pp. 1055-1064. IEEE, 2021.
Junwoo Jang, Jungwoo Han, and Jinwhan Kim. K-mixup: Data augmentation for offline rein- forcement learning using mixup in koopman invariant subspace. EXPERT SYSTEMS WITH APPLICATIONS, 225, 2023.
Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine Planning with diffusion for flexible behavior synthesis. In International Conference Machine Learning, pp. 9902-9915. PMLR, 2022.
Victor Kolev, Rafael Rafailov, Kyle Hatch, Jiajun Wu, and Chelsea Finn. Efficient imitation learning with conservative world models. In 6th Annual Learning for Dynamics & Control Conference, pp. 1777-1790. PMLR, 2024.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in neural information processing systems, 33:1 199799991111111 2020.
Michael Laskey, Jonathan Lee, Roy Fox, Anca Dragan, and Ken Goldberg. Dart: Noise injection for robust imitation learning. In Conference on robot learning, pp. 143-156. PMLR, 2017.
Jianxiong Li, Xianyuan Zhan, Haoran Xu, Xiangyu Zhu, Jingjing Liu, and Ya-Qin Zhang. When data geometry meets deep function: Generalizing offline reinforcement learning. arXiv preprint arXiv:2205.11027. 2022.
Alexandre Mauroy, Y Susuki, and Igor Mezic. Koopman operator in systems and control, volume 7. Springer, 2020.
Igor Mezić. Koopman operator, geometry, and learning of dynamical systems. Not. Am. Math. Soc, 68(7):1077-1055, 2021.
Rafael Figueiredo Prudencio, Marcos ROA Maximo, and Esther Luna Colombini. A survey offline reinforcement learning: Taxonomy, review, and open problems. IEEE Transactions on Newral Networks and Learning Systems, 1037 20233
Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. reduction of imitation learning and struc- tured prediction to no-regret online learning. In Proceedings of the fourteenth international ference on artificial intelligence and statistics, pp. 627- 635. JMLR Workshop and Conference Proceedings, 2011.

---

Published as a conference paper at ICLR 2026
Amir Samadi, Konstantinos Koufos, Kurt Debattista, and Mehrdad Dianati. Good data is all imita- tion learning needs. European Conference on Computer Vision, pp. 144-159. Springer, 2024.
Jie-Jing Shao, Hao-Sen Shi, Lan-Zhe Guo, and Yu-Feng Li. Offline imitation learning with model- based reverse augmentation. In Proceedings of the 30th ACM SIGKDD Conference Knowledge Discovery and Data Mining, pp. 2608-2617, 2024.
Haojie Shi and Max Q-H Meng. Deep koopman operator with control for nonlinear systems. IEEE Roboties and Automation Letters, 733):770--7701, 2022.
Siyu Teng, Long Chen, Yunfeng Ai, Yuanye Zhou, Zhe Xuanyuan, and Xuemin Hu. Hicrarchi- cal interpretable imitation learning for end-to-end autonomous driving. IEEE Transactions Intelligent Vehicles, 8 8(1):673-683, 2022.
Meixi Wang. Xuyang Lou, Wei Wu, and Baotong Cui. Koopman-basee mpe with learned dynamics: Hierarchical neural network approach. IEEE Transactions on Neural Networks and Learning Systems, 35(3):3600-3639, 2022.
Matthias Weissenbacher, Samarth Sinha, Animesh Garg, and Kawahara Yoshinobu. Koopman q- learning: Offline reinforcement learning symmetries of dynamics. In International Conference Machine Learming, pp. 23645-23667. PMLR, 2022.
Annie Xie, Lisa Lee, Ted Xiao, and Chelsea Finn. Decomposing the generalization gap in imitation learning for visual robotic manipulation. In 2024 IEEE International Conference on Roboties and Automation (ICRA), pp. 3153-3160. IEFF, 2024.
Haoran Xu, Li Jiang, Li Jianxiong, and Xianyuan Zhan. A policy-guided imitation approach for offline reinforcement learning. Advances in Neural Information Processing Systems, 35:4085- 4098, 2022a.
Haoran Xu, Xianyuan Zhan, Honglei Yin, and Huiling Qin. Discriminator- offline imitation learning from suboptimal demonstrations. In International Conference Machine Learning, pp- 24725 24742. PMLR, 2022b.
Tian Xu, Ziniu Li, and Yang Yu. Error bounds of imitating policies and environments. Advances in Neural Information Processing Systems, 33:15737-15749, 2020.
Suorong Yang, Weikang Xiao, Mengchen Zhang, Suhan Guo, Jian Zhao, and Furao Shen. Image data augmentation for deep learning: A survey. arXiv preprint arXiv:2204.0861 2022.
Maryam Zare, Parham M Kebria, Abbas Khosravi, and Sacid Nahavandi. A survey of imitation learning: Algorithms, recent developments, and challenges. IEEE Transactions on Cybernetics, 2024.
Wenjia Zhang, Haoran Xu, Haoyi Niu, Peng Cheng, Ming Li, Heming Zhang, Guyue Zhou, and Xianyuan Zhan. Discriminator-guide masel-based offline imitation learning. In Conference on robot learning. pp. 1266-1276. PMLR, 2023.
Tianqiao Zhao, Meng Yue, and Jianhui Wang. Deep-learning-boe -aaard koopman modeling for online control synthesis of nonlinear power system transient dynamics. IEEE Transactions on Industrial Informatics, 19110) 2023.
12

---

Published as a conference paper at ICLR 2026
A APPENDIX
This appendix provides the formal proofs and justifications for the definitions, lemmas, and theorems presented in Section
We restate each proposition for clarity before presenting its proof.
A.1 JUSTIFICATION FOR LEMMA 3 (INDUCED GROUP ACTION)
Lemma 4. A state space transformation S + S induces a linear action on the space of observables K(S), defined by the pullback: [s) ((a-1(3))
Justification. This definition, known as the pullback of the function g by the map a, is the standard way transformation a space induces a transformation on functions defined on that space. The intuition is as follows: the new function, ag, when evaluated at a point s, should have the value that the original function had at the point that was mapped to If we denote this original point by s', then (s') = 8, which implies s' = o os)) Therefore, the value of the new function at must be the value of the old function at (s), leading directly to the definition jorg (s) g(o (s)).
The linearity of this induced action can be shown directly. For any scalars a, b E R and observables 91. 92 K(S)):
[ar(01 + bgz)l(s) (agı + bg2)(o
agi ag (s)) + (0 (s))
asogi + (s)
(a(a91) + ((0g)))(s).
Since this holds for all & E S, we have (agi + bg2) a(ag) b(og2).
A.2 PROOF OF THEOREM T (EQUIVALENCE POLICY EQUIVARIANCE AND KOOPMAN COMMUTATION)
Theorem 4. Let Fx be a closed-loop dynamical system and Ky be its corresponding Koopman operator. The following two statements equivalent:
1. The system is --pllicy-euuivvriint:: Fr =)= Fx(8) for all a E E,,8 E.S
2. The Koopman operator commutes with the induced group action: Kzo = oK₃ for all
a E Σ.
Proof. We prove both implications separately.
Part I: Equivariance implies Commutation (1 + 2). Assume the system is E-policy-equivariant, i.e., (z 8) (x). We must show that Kza = a KT. To do this, we apply the operator K to an arbitrary observable g E K(SS) and evaluate at an arbitrary state E
(<)) = (F+)
(by Def. 2 of Kz)
g(0 (FF,
(by Def. 3 of induced action)
o(Fr(o
(by policy equivariance, replacing s with a (s))
[Kng)(o-1(()
(by Def. 2 of Kx)
= (o(K)))](s)
(by Def. 3 of induced action)
Since (sg) (o(K)g)((s) for all g K(S)) and all S, conclude that the operators are equal: K = oK
Part 2: Commutation implies Equivariance (2 1). Assume the Koopman operator commutes with the induced group action, i.e. K2 = a KTz- For any observable y and state s, this means [K= (sg))()) = (s)) We expand both sides using the definitions:
LHS: (og)) (ag) g(o (F(
RHS: (o(K))))(s) (o- -F)
13

---

Published as a conference paper at ICLR 2026
Equating the expanded forms, have:
g ((F
This equality must hold for all observable functions y in K(S). If the space of observables is sufficiently rich (e.g., separating points), this implies that the arguments of g must be equal:
(F+(s)) (s)o
Let s' (s). This implies 8 = a .'s Substituting s' into the equation gives:
F=(5') = (F>
Applying a both sides yields:
a. FF Fg (o 8')
Since this holds for any s' E ES the system is E-policy-cquiivarian.
A.3 PROOF OF THEOREM
2 (GENERATING POLICY-CONSISTER TRAJECTORIES)
Theorem 5. Let Fx. be E-policy -quivv--quivaariant system with Koopman operator Kr and let a E E. Let (z+ = E(s)) be the latent representation of a trajectory under policy n, such that Z441 Kzzt. Define new trujectory { by 2, 2t. Then this new trajectory evolves according to the same Koopman dynamics: 2++1 = K & Consequently, the decoded trajectory {&t = -(22)} also a valid trajectory under the closed-loop dynamics of policy
Proof. We are given expert latent trajectory { that follows the dynamics 21+1 = K Ze. The new trajectory defined { where &
2z. We aim show that 2④41 22
We start with the definition of the next state in the new trajectory, &
2441 a. 7--1
(by definition of the transformed trajectory)
= (Kx21)
(by the dynamics of the original trajectory) (associativity of operator application)
(0KK)22
Since the system is e-pulicy-equivarrint, by Theorem | its Koopman operator commutes with the symmetry action: a K KR We substitute this into our derivation:
2+1 (K
K2 . 21)
(associativity) (by definition &))
K24t
Thus, we have shown that 2 = & This means the transformed latent trajectory { evolves according to the same linear dynamics governed by K2 as the original expert trajectory.
Consequently, the decoder correctly maps latent states back to the state space, the sequence of states (s. D(+)) represents trajectory that is dynamically consistent with the learned model of the expert's closed-loop behavior.
A.4 ROBUSTNESS ANALYSIS: BOUNDED ERROR TRANSFORMED TRAJECTORIES
The theorems in the preceding sections assume an ideal, perfectly commuting Koopman operator and symmetry transformation. In practice, both our learned Koopman operator K and our transformation approximations, meaning their commutation is not perfect. This section provides more general proof that accounts for this imperfection.
We define the commutation error operator, A, as the commutator of and K:
A oK Ko
In an ideal scenario, = In practice, A represents how much our learned symmetry fails to commute with our learned dynamics.
14

---

Published as a conference paper at ICLR 2026
Theorem 6 (Bounded Error for Transformed Trajectories). Let {&{] } be latent expert trajectory
with a sequence of one-step prediction errors {ee). where c = 2 Kef. Let a be an ap- proximately equivarian transformation with commutaaiion error A oK Ko Then the one-step prediction error & for the transformed trajectory (&g = az)) is given by:
& dc + A
(13)
Consequently, its norm is bounded as:
|| S | eell +
|||||=
(14)
Proof. We begin with the definition of the prediction error for the transformed trajectory, &g
- 2411 &
Substitute the definitions of the transformed trajectory, &n ozf and 2441 02
é = - K(oz,)
Now, substitute the expression for the original trajectory's ' dynamics, 2f41 K&f cgi
é = = (K + Ex) Kozg
By the linearity of the operator o, can distribute it over the sum:
& = oı + ae; - Kaz
Rearranging the to group the components related E, and 2F=
é = σee + (aK - Kf
By our definition of the commutation error, A oK Ko, we arrive at the first result:
& 06 + Az
This completes the proof of the first equation.
To derive the bound on the norm, we take a suitable matrix/vector norm (e.g., the L,2 norm) of both sides and apply the triangle inequality (lla bll
S ||a|| + [[6]]):
||ε=]] ||oe, م 1 S ||a + =|A=/11
Finally, applying the property of induced matrix norms (||AT| ||]||| final bound:
each term gives the
| s || + ||||||| |||||
This completes the proof.
IMPLICATIONS OF THE BOUND
This result is highly significant. It shows that the error in a synthesized trajectory has distinct sources:
Propagated Model Error [[ [[): The original model's inability perfectly predict the expert dynamics (c²) is carried over and scaled by the norm of the transformation.
2. Symmetry Violation Error |||+||=++]]]] of the learned symmetry perfectly commute with the learned dynamics.
: A new error term arises directly from the failure
This provides a clear theoretical motivation for not only minimizing the standard reconstruc- tion/prediction loss (which minimizes [|c=]]), but also for explicitly regularizing the model to mini- mize the commutation error ||]
B ROBUSTNESS IN FEW-SHOT SETTINGS
As demonstrated empirically in the Figure our symmetry-based approach exhibits remarkable
robustness to reductions in the size of the expert dataset, regime where standard methods like Behavioral Cloning (BC) often fail. This appendix provides qualitative analysis explaining this superior data efficiency.
15

---

Published as a conference paper at ICLR 2026
Performance vs. Data Availability
1.1
Normalized Score
0.9
0.8
Ours
Behavioral Cloning (BC)
Expert Performance
0.6
10%
25%
50%
100%
Percentage of Expert Data (%)
Figure 4: Few-Shot Performance Comparison on maze2D-umaze.
The Brittleness of Standard Imitation Learning. Standard Behavioral Cloning (BC) directly learns a state-action mapping from the provided expert data a supervised manner. The perfor- mance of BC policy is therefore critically dependent the coverage of the state space by the training data. When the dataset is small, represents a sparse and potentially biased sample of the true expert state distribution. This leads two primary failure modes:
1. Poor Generalization: The learned policy performs well only in the immediate vicinity of the training states. has no knowledge of how to act in the vast "in-between" regions of the state space.
2. Covariate Shift: A small initial error can push the agent into out-of-distributie (OOD) state. Since the policy was not trained on data from this region, its subsequent actions are likely be suboptimal even random, leading to cascade of compounding errors and rapid task failure.
Data Augmentation as the Source of Robustness. In contrast, our method is fundamentally data augmentation technique that addresses the core issue of sparse data coverage. The key advan- tage stems from the learned symmetry transformation, a. Even from a single expert trajectory, our model can generate a continuous family of new, diverse trajectories by applying the transformation pointwise: { = (o (s))))
Crucially, as proven in Appendix A these augmented trajectories guaranteed be dynamically consistent. This process aims to densify the training data distribution, potentially bridging gaps in the state-space coverage that might otherwise hinder learning from sparse datasets. The policy is consequently trained on a much richer and more comprehensive representation of the system's 's behavior. This expanded dataset inoculates the policy against minor deviations, as states that would be OOD for BC are now likely to be in in-distribution for our augmented dataset.
Conclusion. In summary, the performance of BC is directly coupled to the density and quantity of the provided expert samples. Our method decouples this dependency by leveraging the learned ge- ometric and dynamic structure of the system (i.e., its symmetries) to synthetically generate dense, dynamically valid dataset from a sparse source. This intrinsic data amplification is the primary reason for its superior performance and robustness in low-data regimes.
16

---

Published as a conference paper at ICLR 2026
C ANALYSIS OF WEIGHT DISTRIBUTION AND ADAPTIVE SYMMETRY CONTROL
To investigate the interpretability and effectiveness of the learned weights in framework, analyze how these weights adaptively control the symmetry operator. The core objective is to ensure conservative behavior in data-sparse regions (where prediction error is high) while placing trust in augmentation within data-dense regions (where prediction is reliable).
We conducted an empirical validation on the Maze-2D dataset. Specifically, we sampled N 5, 000 state transitions and computed their corresponding weights derived from the Koopman prediction operator. We visualize the relationship between data density, prediction error, and weight magnitude using Principal Component Analysis (PCA).
Observation Manifold (PCAiColored by Latent Diff Probability
Dimendion
Figure 5: PCA visualization of the weight distribution the Maze-2D dataset.
As illustrated in Figure 5 the results corroborate our hypothesis regarding the data-adaptive capa- bility the a training:
Data-Dense Regions: Represented by darker colors in the visualization, these exhibit low prediction errors. Correspondingly, the model assigns small weights, implying that the generated augmentations are reliable and can be trusted.
Sparse Regions: In the central of the maze where data is sparse, observe high prediction errors. The model assigns large weights to these instances, necessitating a con- servative update to prevent distribution shift.
Quantitatively, we measured the deviation of the learned weight distribution from uniform distri- bution. The analysis yields Kullback-Leible (KL) divergence of 0.429 and Total Variation (TV) distance of 0.374. These metrics confirm that the weights non-uniform and effectively capture the underlying geometry and density of the dataset.
D EXPERIMENTAL DETAILS
This section provides a detailed breakdown of the model architectures, hyperparameters, and train- ing procedures used in experiments ensure full reproducibility.
17

---

Published as a conference paper at ICLR 2026
Our method, which refer to Koopman-Assister Trajectory Synthesis for Offline Imitation Learning (KATS), consists of four sequential training stages: Koopman model pre- pre-training, inverse dynamics model training, sigma model training, and finally behavior cloning with data augmenta- tion.
Model Architectures. All neural networks in our framework implemented as Multi-LLayer Perceptrons (MLPs) with Tanh activation functions:
Koopman Model: Consists of an encoder and decoder, both implemented as 3-layer MLPs with hidden dimensions of [512. 512]. The encoder maps from observation space to the Koopman latent space of dimension N₂ = 32, while the decoder reconstructs observations from the latent representation. The linear Koopman operator K E R33232 operates in the latent space.
Inverse Dynamies Model: 3-layer MLP with hidden dimensions of [256, 256], taking concatenated current and next states input and outputting predicted actions.
Sigma Model: 3-layer MLP with hidden dimensions of [256, 256], mapping from latent space to latent space with the same dimensionality N, = 32.
Policy Model: 3-layer MLP with hidden dimensions of [256, 256], mapping from obser- vation space action space. For continuous action spaces, tanh activation applied to the final layer.
Stage 1: Koopman Model Pre-training. The Koopman model, which combines an encoder- decoder architecture with a linear Koopman operator K, is pre-trained to learn the underlying dy- namics in latent space. The model is trained for up 400 epochs with early stopping based on traintto-ess, loss ratio convergence (threshold: 1.0 for both reconstruction and Koopman losses). We use the ADAM optimizer with an initial learning rate of 10-4 44, which halved every 50 epochs after epoch 100. The batch size is set to 256, and 70% of the dataset is used for training with the remaining 30% for validation.
Stage 2: Inverse Dynamics Model Training. The inverse dynamics model is trained to predict actions given state transitions (&t, 8t-1) ) -> az This model essential for generating corresponding actions for the augmented state trajectories. Training is conducted for 50 epochs using the ADAM optimizer with a learning rate of 10-3 and the same train-test split as the forward model.
Stage 3:
Sigma Model Training. The sigma model learns the transformation 2 - Z that enables controlled data augmentation the latent space. The training objective is:
C. B(se [w (((()z (11+1))
where exp(r|l<44 K with 1.5. The model is trained for 30 epochs using the ADAM optimizer with a learning rate of 10-1
Stage 4: Behavior Cloning with Data Augmentation. The final stage trains the policy model using both original and augmented data. The data augmentation process works as follows:
1. Encode original states (st, St+1) into latent space: zt, 7141
2. Apply sigma transformation: 0(21),
3. Decode back to state space: $t, Øt+1
4. Use inverse dynamics model to predict corresponding actions: ā
The policy is trained for 50 epochs using standard mean squared error loss for continuous actions or cross-entropy loss for discrete actions. We use the ADAM optimizer with a learning rate of 5 10 When data augmentation enabled (controlled by use data augment aton flag), the training alternates between original and augmented data with equal weighting.
18

---

Published as a conference paper at ICLR 2026
Hyperparameter Summary.
Latent dimension: N₂ = 32
Batch size: 256 across all training stages
Learning rates: 3 x 10-4 (Koopman model), 10 (invere model), x 10 (policy), 1 10 (sigma model)
Training epochs: Forward model (up to 400 with early stopping), Inverse model (50), Sigma model (30), Policy (50)
Sigma weighting parameter: T 1.5
Data augmentation weight: shift E [0.1, 0.5), controls the weighted fusion be-
tween augmented and original data
Train-test split ratio: 0.7/0.3
Computational Resources. All experiments conducted using PyTorch with CUDA support. The models were trained on NVIDIA GeForce RTX 3090 GPUs with device selection controlled by the cuda device parameter. Each complete training (all four stages) typically requires 2-4 hours depending on the environment complexity and dataset size.
E
USE OF LARGE LANGUAGE MODELS
Large Language Model (LLM) was used to assist in the writing of this manuscript. Its role was exclusively for language polishing, including improving grammar, clarity, and style. The LLM had no role in the research ideation, methodological design, or data interpretation. The human authors solely responsible for the intellectual content of this work.
19


Reviews:
[{"review_id": "xe11B2ewvK", "reviewer": "Reviewer_cuKL", "summary": "This paper proposes a method based on Koopman Theory for generating trajectories from offline data.", "questions": "1. What are the differences and connections between Equations 7–8 and 9–10? In implementation, are they used together or only 9–10?\n\n2. Why does Table 1 compare KATS+BC with KFC+BC, while Table 2 compares KATS+BC with KFC+CQL?\n\n3. Why are the data augmentation methods in Table 1 compared with SRA, MOLI, and KFC+BC, while in Table 2, they are compared with TELS, DOGE, POR, and KFC+CQL?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "The approach appears reasonably sound.", "weaknesses": "1. (I have not read KFC) The authors claim that this work differs from KFC, as KFC only generates single-step data, while KATS generates trajectories. However, judging from Equations 7, 8, 9, and 10, KATS still appears to generate states.\n\n2. Although the experimental results presented by the authors show that KATS performs well, the experiments seem insufficient. For example, there is no ablation study.\n\n3. The baselines compared in Table 1 and Table 2 are inconsistent:\n(1) Table 1 compares KATS+BC and KFC+BC, while Table 2 compares KATS+BC and KFC+CQL. Since the base algorithms of KATS+BC and KFC+CQL are different, the comparison lacks fairness.\n(2) In Table 1, the data augmentation methods compared are SRA, MOLI, and KFC+BC, while in Table 2, the compared methods are TELS, DOGE, POR, and KFC+CQL.", "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": "YbYtKKkWmC", "reviewer": "Reviewer_f7Wk", "summary": "The paper tackles covariate shift in offline imitation learning, where agents are limited to fixed datasets with potentially low diversity. Prior Koopman-based methods operate at the single-step level, causing dynamic inconsistency and high cost. Koopman-Assisted Trajectory Synthesis (KATS) is introduced as a trajectory-level data augmentation method that generates novel yet dynamically consistent expert-like trajectories in a learned state-equivariant linear Koopman latent space. In addition, KATS is adaptive in that it prioritizes data synthesis where the model is uncertain. Theoretical results guarantee that symmetries commuting with the learned Koopman operator yield trajectories consistent with expert policy dynamics. In practice, KATS augments data and then applies simple BC, outperforming more complex offline IL and RL baselines. KATS serves as a plug-and-play augmentation module that enhances existing algorithms through high-fidelity, behaviorally consistent data generation.", "questions": "* Can the authors make it clear what they mean by scalability/scalable? From the phrasing in the paper, it seems that the scalability claimed by the authors is on the complexity of the task (dimension, degrees of freedom). However, since this is an approach for the low-data regime (limited offline dataset), it might be on the how the developed data augmentations impact performance, etc.\n* Koopman theory imposes a strong inductive bias, going against the bitter lesson. Would the authors defend that modeling the temporal relationship between one latent and its successor simply with a linear operator is enough to model system that are more complex than the ones tackled empirically in the paper? In other words, would adding depth to the encoder always be enough to go in a sufficiently “higher-dimensional space” where the dynamics can reasonably be assumed linear?\n* It is unclear in the text why KFC working at the single-step level is costly in compute (L193). Is it clearer costly because the action-equivariant assumption in KFC requires one linear operator per action dimension (Eq. 2), making it costly to scale with action dimensionality?\n* Section 5.1: modeling the closed-loop system dynamics is a design choice rather than an innovation, is it not? Do the other claim it is an innovation because it is unusual?\n* Is there a particular reason why the authors write \"find symmetry basis\" at line 2 of the algorithm, and not \"learn the symmetries\" (\"sigma model training\" only appears in the appendix, L812-813)? By basis or symmetries, do the author mean that the set verifies the properties of basis, then they then use to craft other symmetries as linear combinations of the elementary symmetries of the basis?\n* What would be informative, for the data augmentation methods, is to get the final size of the dataset used to train the policy with BC (compared to the initial size), along with statistics that could give an idea as to how it expands the initial one in terms of diversity.\n* Could the authors include a few words about the baselines, to get the gist of their approach, which would put KATS in perspective; e.g., how does TELS' data augmentation approach differ from KATS?\n* The authors write in the appendix that the \"[policy] training alternates between original and augmented data with equal weighting\". What is the only sampling strategy that the authors tried? Could the authors give the respective dataset sizes for the reader to be able to gauge how likely to overfit the policy is?\n* Have the authors experimented with learning the IDM from the mapping of the expert states into the learned latent space, i.e. from the latents z instead of the states s?\n* What does the weight distribution (in the sigma loss) look like? In other words, how far from uniform is the \"adaptive\" scheme the authors designed?\n* Providing an ablation study comparing KATS with and without any symmetry training and usage would be insightful.\n\nStyle, typos, suggestions:\n* It would be useful to add, in the algorithm, links to the equation according to which the various networks are optimized.\n* [minor] It might be good to mention that that Koopman machinery is used in latent space earlier in the introduction than at the very end. The paper might also benefit from putting KATS in the context of model learning in latent spaces, which the typical RL literature reader might be more familiar with.\n* [minor] L123-124: the comma after \"shift\" should be removed.\n* [minor] L155-156: z_t and z_{t+1} correspond to a transition (s_t, s_{t+1}) or to a pair of states but not a pair of transitions.\n* [minor] L187-188: why use \"aug\" when the figure 1 (a) uses primes to designates the augmentations?\n* [minor] The last sentence of Def. 2 should be emphasized.\n* [minor] L263: properly format the emdash, or use a colon.\n* [minor] End of page 9: \"Limitaitions\" -> Limitations\n* [minor] L810-811: \"Synthesisi\" (extra \"i\"), \"(KATS)\" (missing space prefix).\n* [minor] L289-290: \"Any symmetry transformation that commutes with K is therefore guaranteed to produce trajectories that adhere to this policy.\" I find this to not be clear.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "* Th. 2 is very elegant and clear. Slight suggestion: preamble the section with the gist/a teaser of what the theorem will show.\n* Sec. 5.2: ingenious yet simple.\n* KATS demonstrates that augmenting data and applying behavioral cloning can be a more effective and reliable strategy for imitation, even in data-scarce regimes where traditional apprenticeship learning methods, despite allowing interaction, tend to be fragile and overcomplicated.", "weaknesses": "* Th. 3 should be followed by an \"in-words\" interpretation and description of its consequences, along with a hyperlinked reference to the “Implication of the Bound” section presented in the Appendix.\n* KATS is introduced from the machinery of KFC, but the link is somewhat lost until it comes back in 5.1 line 287, where it is clear.\n* The way the authors present sigma in Fig 2 (“Sigma 0.2, Sigma 0.3” etc.) or as “a symmetry basis” suggests a structured family of symmetry operators or scalars, but in the actual text and appendix, they never describe how that basis or scaling is obtained. The figure’s depiction of multiple sigma’s or scaled sigma’s is not grounded in the described theory or implementation. The authors must explicitly indicate how they operationalize the learned sigma network to obtain their basis.\n* The authors write: \"This dramatic leap provides strong validation for our core contribution: the action-independent formulation.\". It is definitely noteworthy and interesting to observe that the action-independent closed-loop formulation combined with an IDM can yield such results. Stating that the generated trajectories are \"by construction, more behaviorally consistent than those from action-conditioned models like KFC\" might be a bit of a stretch however, but an experimental design could be devised to showcase that further.\n* In the appendix L795-796, the authors write \"This process effectively densifies the training data, filling the gaps in the state-space coverage that a sparse dataset would otherwise leave open.\". That is “potentially” what KATS enables, but the coverage is not proven or showcased as such.\n* The paper could use an additional round of polish to remove the inconsistencies in notations.", "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": "boqIfUbY2q", "reviewer": "Reviewer_AWKq", "summary": "This paper proposes a theoretically principled method to generate augmented, synthetic expert demonstrations for offline imitation learning. With said augmented data, one can perform standard behavioral cloning on the union of the given expert dataset and the augmented dataset. The paper grounds their technique for augmenting trajectories in Koopman theory, done in the latent space of a learned autoencoder over the expert data. In particular, they note that if the learned state latents satisfy certain linear properties (e.g. the transition function under the expert data is linear -- said linear transformation is the Koopman operator), then compounding errors are bounded, leading to useful trajectory generation.\n\nExperiments on MuJoCo tasks, both in IL and RL, validate the hypothesis shown, including strong results compared to prior offline IL baselines that either had to rely on suboptimal offline data (e.g. MILO) and other methods employing data augmentation (e.g. KFC+BC). Furthermore, they also try their method in the offline RL setting, comparing to standard baselines and recent Koopman-focused baselines such as KFC++ and showing strong performance.", "questions": "No fundamental questions from me, but I am curious to know if any large-scale experiments were done for this paper, including with either high-dimensional states or with image observations.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "I really like this area of work, as Koopman theory is well grounded and has been used for quite a while in model-based dynamical system control (although I am not super familiar with the literature). There are many advantages in learning a latent space under which the transition dynamics are linear, which include not even having to do RL directly and employing more stable control theory-focused algorithms.\n\nThe theory seems fine to me, and seems to borrow a lot from the KFC paper (Weissenbacher et al. 2022), leading me to believe that it is sound. The experimental results are also strong, showing strong performance improvement even over other Koopman-based RL and IL algorithms.", "weaknesses": "There are instances where the paper could be written a bit better (e.g. put your citations in parentheses!). I think there are also potential weaknesses to the method empirically, which include (and correct me if I'm wrong) the following:\n\n- In the RL setting, if the reward function is a function of both state and action, then Q learning may be biased. I think that the reward is only the function of the state in MuJoCo and DMC control domains, which means the method is fine there, but in cases where it is not, then I figure due to said bias, learning is difficult even when the latent system is well-learned.\n\n- Generally, these methods seem to work in small-scale tasks such as OpenAI Gym locomotion, while not having been tested on larger-scale domains such as DMC from pixels or larger control tasks. A potentially great use for this method could be to learn such a linear latent system on real robotic datasets, making learning controllers much faster.\n\nThese are not \"make or break\" weaknesses, more so that this seems not to have been tested. In general, for the focus of IL, there are less problems, as for instance removing action conditioning is fine as the Markovian expert policy is embedded into the latent encoding, which is enough for data augmentation.", "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": "1dJHFHBVuO", "reviewer": "Reviewer_qXWS", "summary": "In this paper, a data augmentation framework, namely KATS that leverages Koopman theory to address the critical distribution shift problem in offline imitation learning is introduced. The presented framework can synthesize trajectories for the training data augmentation while avoiding the compounding errors of recursive rollouts and ensuring computational efficiency and scalability. Experimental results show the effectiveness of the proposed approach.", "questions": "Please refer to the weakness points.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well-motivated and aims to address an important issue in the literature, which generates high-quality, dynamically consistent trajectory-level data while avoiding the compounding errors and ensuring computational efficiency. \n2. Experimental results along with theoretical guarantees demonstrate the advantages of KATS, which yields substantial improvements in policy performance on some tasks.", "weaknesses": "1. The presented Koopman-Assisted Trajectory Synthesis (KATS) framework is based on the assumption that the symmetries of a closed-loop dynamical system, driven by a fixed expert policy, are directly reflected as commutation properties of its associated Koopman operator. While the assumption may hold in some cases, it is unclear whether such an argument can be satisfied in a general sense. Can the proposed framework be applied in any type of environment, or is its application limited to some special domains? \n2. In the literature, many works have been proposed for the trajectory-level data augmentation, more recent baselines, especially for diffusion-based approaches, can be added and discussed for the comparison. \n3. In the experiments, the most recent baseline is a rejected paper (TELS) while other baselines were mainly presented two or three years ago. Considering the rapid development of the related research area, it is necessary to adopt more recent baselines to verify the effectiveness of the proposed method.", "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)