Koopman-Assisted Trajectory Synthesis: A Data Augmentation Framework for Offline Imitation Learning

ICLR 2026 Conference SubmissionAnonymous Authors
Offline Imitation Learning; Offline Reinforcement Learning; Data Augmentation
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.

Disclaimer
This report is AI-GENERATED using Large Language Models and WisPaper (A scholar search engine). It analyzes academic papers' tasks and contributions against retrieved prior work. While this system identifies POTENTIAL overlaps and novel directions, ITS COVERAGE IS NOT EXHAUSTIVE AND JUDGMENTS ARE APPROXIMATE. These results are intended to assist human reviewers and SHOULD NOT be relied upon as a definitive verdict on novelty.
NOTE that some papers exist in multiple, slightly different versions (e.g., with different titles or URLs). The system may retrieve several versions of the same underlying work. The current automated pipeline does not reliably align or distinguish these cases, so human reviewers will need to disambiguate them manually.
If you have any questions, please contact: mingzhang23@m.fudan.edu.cn

Overview

Overall Novelty Assessment

The paper proposes KATS, a framework for generating complete multi-step trajectories in offline imitation learning using Koopman operator theory with state-equivariant assumptions and refined generator matrices. It resides in the 'Dynamics-Based Trajectory Generation' leaf, which contains five papers total including the original work. This leaf sits within the broader 'Trajectory-Level Synthesis and Adaptation' branch, indicating a moderately populated research direction focused on generating full trajectories rather than single-step augmentations. The taxonomy shows this is an active but not overcrowded area, with sibling leaves exploring diffusion-based synthesis and demonstration adaptation as alternative trajectory-level approaches.

The taxonomy reveals several neighboring research directions that contextualize this work. The sibling 'Diffusion-Based Trajectory Synthesis' leaf contains four papers using generative models for trajectory creation, representing an alternative paradigm to dynamics-based methods. Adjacent branches include 'Corrective and Interventional Augmentation' (addressing distribution shift through corrective labels) and 'Model-Based Data Generation' (using world models for synthesis). The scope note for the original leaf explicitly excludes diffusion and stitching methods, positioning KATS within dynamics-model approaches that preserve system constraints. This placement suggests the work bridges classical control theory (Koopman operators) with modern imitation learning, occupying a distinct methodological niche.

Among 25 candidates examined across three contributions, the trajectory-level synthesis contribution shows one refutable candidate from 10 examined, while the state-equivariant representation (0 from 5) and refined generator matrix (0 from 10) appear more novel within this limited search scope. The single refutable case for trajectory synthesis suggests some prior work addresses multi-step generation, though the specific combination of Koopman theory with state-equivariance and error-correction mechanisms may differentiate KATS. The computational efficiency and robustness contributions show no clear refutations among their examined candidates, indicating these technical innovations may represent more distinctive advances within the constrained literature sample.

Based on this limited analysis of 25 semantically similar papers, KATS appears to occupy a specialized position combining established Koopman theory with novel efficiency and robustness mechanisms for trajectory synthesis. The search scope covers top semantic matches but cannot claim exhaustiveness across the broader offline IL literature. The taxonomy structure suggests moderate competition in dynamics-based trajectory generation, with the work's novelty likely residing in its specific technical approach rather than the high-level goal of multi-step synthesis.

Taxonomy

Core-task Taxonomy Papers
50
3
Claimed Contributions
25
Contribution Candidate Papers Compared
1
Refutable Paper

Research Landscape Overview

Core task: data augmentation for offline imitation learning. The field addresses the challenge of learning policies from fixed demonstration datasets by generating additional training data. The taxonomy organizes approaches into several major branches: Trajectory-Level Synthesis creates new state-action sequences through dynamics models or generative processes; Corrective and Interventional Augmentation incorporates failure recovery and human feedback; Visual and Perceptual Augmentation applies transformations to observation spaces; Invariance and Equivariance-Based methods exploit symmetries; Model-Based Generation uses learned world models; Cross-Domain Transfer leverages data from related tasks; Policy-Guided approaches use learned policies to inform augmentation; Application-Specific methods target domains like robotics or autonomous driving; and Evaluation branches establish theoretical guarantees. Representative works span from trajectory synthesis methods like Mimicgen[1] and DemoGen[5] to corrective approaches such as IntervenGen[7] and visual augmentation techniques like RoCoDA[9]. Within Trajectory-Level Synthesis, a particularly active line explores dynamics-based generation, where learned models produce plausible rollouts to expand limited datasets. Koopman Trajectory Synthesis[0] sits squarely in this branch, using Koopman operator theory to generate trajectories that respect system dynamics. This contrasts with nearby works: DemoGen[5] emphasizes task-conditioned generation for robotic manipulation, while Reverse Augmentation[36] explores backward trajectory construction. The tension between model accuracy and generalization appears across these methods—some prioritize faithful dynamics modeling (Offline Trajectory Optimization[40]), while others focus on diversity and coverage. Koopman Trajectory Synthesis[0] distinguishes itself through its linear operator framework for nonlinear dynamics, offering a middle ground between model fidelity and computational tractability that differs from diffusion-based approaches like those in related generative methods.

Claimed Contributions

Trajectory-level synthesis process avoiding compounding errors

The authors introduce a method that generates entire expert trajectories as the base unit for data augmentation, rather than single-step transitions. This approach mitigates the compounding errors common in state-space rollouts and ensures generated trajectories are dynamically consistent within the linear Koopman space.

10 retrieved papers
Can Refute
State-equivariant Koopman representation for computational efficiency

The framework leverages a state-equivariant assumption instead of action-equivariant modeling, which avoids severe computational and memory costs of prior approaches. This design makes KATS highly efficient and scalable for complex tasks by learning only a single operator rather than per-action operators.

5 retrieved papers
Refined generator matrix to counteract approximation errors

The authors design an adaptive symmetric generator matrix that makes the model more robust to the inherent approximation errors of finite-dimensional Koopman representations. This is achieved through an optimization process weighted by the Koopman model's prediction error, improving the quality of synthesized trajectories.

10 retrieved papers

Core Task Comparisons

Comparisons with papers in the same taxonomy category

Contribution Analysis

Detailed comparisons for each claimed contribution

Contribution

Trajectory-level synthesis process avoiding compounding errors

The authors introduce a method that generates entire expert trajectories as the base unit for data augmentation, rather than single-step transitions. This approach mitigates the compounding errors common in state-space rollouts and ensures generated trajectories are dynamically consistent within the linear Koopman space.

Contribution

State-equivariant Koopman representation for computational efficiency

The framework leverages a state-equivariant assumption instead of action-equivariant modeling, which avoids severe computational and memory costs of prior approaches. This design makes KATS highly efficient and scalable for complex tasks by learning only a single operator rather than per-action operators.

Contribution

Refined generator matrix to counteract approximation errors

The authors design an adaptive symmetric generator matrix that makes the model more robust to the inherent approximation errors of finite-dimensional Koopman representations. This is achieved through an optimization process weighted by the Koopman model's prediction error, improving the quality of synthesized trajectories.