Sample-efficient diffusion-based control of complex nonlinear systems

We sincerely appreciate the reviewer's feedback. Our responses are as follows.

Tip: Please visit the link(https://drive.google.com/file/d/1VWaCyEv0NPMPPqCdVfgJDPXoTiuN76MV/view?usp=sharing) for new Tables and Figures.

1. New optimal control baseline

We additionally compare SEDC with learning-based Model Predictive Control (MPC), the only data-driven method among suggested optimal control approaches. Results (link:Table I) show MPC achieves higher target losses across all tasks, likely due to error accumulation. MPC also has significantly longer inference times (e.g. >1000 vs. 0.5) that increase with control horizon. SEDC directly maps initial/target states to complete control trajectories, avoiding compounding errors and reducing computation time.

2. The rationale for isolating linear and nonlinear components in DMD

The design decomposes the clean trajectory prediction into linear and nonlinear components to overcome the limitations of single-network approaches that struggle to model both simultaneously with limited data. It is theoretically grounded in the Taylor expansion of vector-valued functions. Please refer to part 1&2 of our response to reviewer YCVy for the explanation of DMD. We compare this to alternative strategies in the ablation studies. Table 1 shows that compared to single-UNet, the dual-mode architecture reduces error by 54-94% when using only 10% of data, and by 47-57% with full data, confirming higher effectiveness under data scarcity. In Table 3, we also verified that as the nonlinearity of the system increases, the performance benefit gained by applying DMD becomes more pronounced.

3. The use of gradient guidance

Our gradient guidance method incorporates control cost optimization directly into the denoising process, which steers each denoising step toward trajectories that minimize an cost function $J$(e.g. control energy), using equation (3) in the paper to guide the sampling process $p_\theta(\mathbf{x}^{k-1} | \mathbf{x}^k, \mathbf{y}_0^*, \mathbf{y}_f)=\mathcal{N}(\mathbf{x}^{k-1}; \mathbf{\mu} _\theta(\mathbf{x}^k, k, \mathbf{y}_0^*, \mathbf{y}_f), \mathbf{\Sigma}^k).$

The gradient term $\nabla_{\mathbf{x}^k}J(\hat{\mathbf{x}}^0(\mathbf{x}^k))$ in equation (3) computes how changes in the current noisy state affect the cost and updates the sampled mean in a gradient-descent like way. The stability is guaranteed by setting appropriate guidance strength $\lambda$ and is proven by numerous previous works like classifier-guided diffusion (Dhariwal & Nichol, 2021) and diffusion-based planning (Janner et al., 2022). We will make the description of gradient guidance clearer in the revised manuscript.

4. Compounding errors and overfitting of GSF

GSF ensures physical consistency by re-simulating the system with generated controls (Section 4.3), ensuring all finetuning pairs $[\mathbf{u}^0_{\text{update}}, \mathbf{y}^0_{\text{update}}]$ follow system dynamics. Complementarily, our diffusion-based method considers entire trajectories holistically, allowing these two mechanisms to work together to avoid compounding errors. Moreover, generated trajectories with updated states differ from the original training data, preventing overfitting, as confirmed by the stable validation loss (link:Figures I,II) which demonstrates no upward trend.

5. Response to the datasets used

Our benchmark selection (Inverted Pendulum, Kuramoto, and Burgers) follows established control systems research practice, chosen for real-world relevance and diverse nonlinearity and complexity:

• Inverted Pendulum: state 2, control 1, timestep 128
• Kuramoto: state 8, control 8, timestep 15
• Burgers: state 128, control 128, timestep 10

We add experiments on the power grid system (please see part 5 of our response to reviewer YCVy) to demonstrate our generalizability on real-world scenarios.

6. The potential trade offs

Computation costs in the form of training/inference times are reported in appendix C.2, following AdaptDiffuser; other baselines do not report computational costs. The loss curve (link:Figure I) demonstrates our method's training stability. Sensitivity tests (link:Figure III) show performance is highly affected by low diffusion steps but yields only 5-10% improvements at higher steps, following typical diminishing returns in diffusion models. These demonstrations will be included in the final version.

7. The definition of the symbol $\mathbf{y}$

In our paper, we assume full state observability throughout, with y representing the complete observable system state vector. This is reasonable for our evaluated systems. We will revise for consistent terminology and explicitly state this assumption.

We sincerely thank Reviewer vifZ for thorough review and strong recommendation. We greatly appreciate your positive assessment of our work, particularly your recognition of our experimental design, ablation studies, and the potential impact of our work.

Response to extending SEDC to stochastic control settings:

Thank you for your question on extending our work to stochastic control settings. This represents an interesting direction for future research. Our present work focuses on deterministic non-linear systems where SEDC demonstrates significant advantages in sample efficiency and control accuracy. Although we believe the diffusion-based nature of our approach provides a conceptual foundation that could potentially be adapted to stochastic settings, this would require substantial theoretical modifications to our framework components (DSD, DMD, and GSF). Extending to stochastic control would involve addressing additional complexities in modeling state transition probabilities and optimizing over distributions rather than deterministic trajectories. This remains an open research question we are interested in exploring. We will add a brief discussion of these potential extensions in our limitations section to acknowledge the current deterministic focus of our work.

We thank the reviewer again for their valuable feedback and encouraging assessment of our contribution.

We sincerely appreciate the reviewer's feedback. Our responses are as follows.

Tip: Please visit https://drive.google.com/file/d/1JmK5ZuMIg0CJCf1L2fQqobK6gtueOgts/view?usp=sharing for new tables and figures.

1.The core innovations

The gradient guidance and in-painting for goal conditioning aren't our core innovations. The innovation is a new data-driven complex system control framework that significantly improves diffusion model sample-efficiency in high-dimensional nonlinear system control. To tackle the challenge of nonlinearity of complex systems, we innovatively propose DMD that decomposes trajectory prediction into linear and nonlinear components for better capturing nonlinearity under data scarcity. The design is novel and theoretically based on the Taylor expansion. To handle high-dimensional complex physical systems, we are the first to incorporate inverse dynamics into diffusion-based control, maintaining physical consistency in trajectory generation with limited training data. Moreover, to address the unique non-optimal data problem in complex physical systems, we novelly propose GSF that enables exploration beyond the initial training data distribution.

2.Explanation of decomposition into modes

We are not decomposing a noise-corrupted state trajectory into modes. Instead, the DMD actually decomposes the prediction of the clean sampled trajectory into linear and nonlinear modes, overcoming the limitations of single-network approaches that struggle to model both simultaneously. The theoretical foundation is as follows: $\mathbf{y}_c$ is the conditional input, a learnable combination of initial and target states as in the paper. Our denoiser is designed to output the clean state trajectory $\hat{\mathbf{x}}^0$, expressed as a vector function $\mathbf{f}(\mathbf{y}_c)$. It admits a vector Taylor expansion at $\mathbf{y}_c=\mathbf{0}$:

$$\hat{\mathbf{x}}^0=\mathbf{f}(\mathbf{y}_c) = \underbrace{\mathbf{C}_1 \mathbf{y}_c} _{\mathbf{O}_1:\text{1st-order}} + \underbrace{\mathbf{y}_c^T \mathbf{C}_2 \mathbf{y}_c} _{\mathbf{O}_2:\text{2nd-order}} + \mathcal{O}(||\mathbf{y}_c||^3)$$

For linear systems, only the first-order term remains. For nonlinear systems, by neglecting higher-order terms for simplicity, we can decompose the prediction into linear and nonlinear quadratic modes. In our dual-Unet architecture, the first UNet learns to produce linear coefficient ($\mathbf{C}_1$) from noisy trajectories $\hat{\mathbf{x}}^k$, while the second extracts nonlinear modes ($\mathbf{C}_2$) that capture higher-order interactions. It is validated that as the nonlinearity of the system increases, the dual-Unet achieves more benefits compared to the single-Unet(Table 3). We will refine the explanation of DMD in the final manuscript.

3.sample efficiency experiment

Our experimental design is fair-our baselines include AdaptDiffuser which uses fine-tuning yet performs worse(e.g. Figure 3 shows our method achieves lower target loss with just 10% of training data). Unlike AdaptDiffuser, we collect generated trajectories for fine-tuning without reward/discriminator filtering, exposing the model to more diverse samples and better balancing exploration/exploitation. Following your advice, we test applying GSF on the most competitive baseline DiffPhyCon. Result(link:Table III) confirms GSF improves baseline performance, validating the generalizability of our GSF framework.

4.Experiment demonstration

Similar to DecisionDiffuser and DiffPhyCon, we report mean/standard deviation metrics(link:Table I) and pareto frontiers(link:Figure I). Results confirm our method maintains competitive performance across all datasets. We will refine this demonstration in the final manuscript.

5.Limitation of the problems chosen

Our benchmark selection follows established research practice(DiffPhyCon,[1][2]), chosen for diverse nonlinear characteristics and real-world relevance. Although the Inverted Pendulum is low-dimensional, our evaluations in Burgers system(128 states/controls) have sufficiently represented our method's performance in high dimensional dynamics.

Since Kuramoto may oversimplify real-world complexity, we conduct additional experiments on swing dynamics [1,3], which models real-world power grid behavior with higher fidelity and complexity(see link:Figure II). Results (link:Table II) show our method achieves lowest target loss, outperforming DecisionDiffuser by 60%, confirming that our approach's benefits extend to practical complex scenarios.

[1] Data-driven control of complex networks. Nature Communications

[2] Closed-loop Diffusion Control of Complex Physical Systems. ICLR25

[3] How dead ends undermine power grid stability. Nature communications

6.Problems in the supplementary materials

We fix the problem of version compatibility error of the codes in the code link. We also report the frequency at which the benchmark systems can be operated using the methods in link:Table IV.

We sincerely thank Reviewer qYwS for positive assessment of our work and thorough review.

Response to extending SEDC to stochastic control problems:

We appreciate the reviewer's insightful suggestion regarding potential extensions to Schrödinger bridge problems and stochastic optimal control. This represents an interesting direction we had not fully explored yet.

Our framework is primarily designed for deterministic systems, and extending it to stochastic settings would require significant theoretical and architectural modifications. We believe our diffusion-based approach might provide a foundation for addressing stochastic control problems. And for further extension to stochastic scenarios, the inverse dynamics and non-linear decomposition components have to be adapted to accommodate stochastic processes, which would require careful investigation. We will add a limitation section in the final paper that acknowledges the current deterministic focus of our work and discusses potential future extensions to stochastic settings as an open research question.

We thank the reviewer for this valuable suggestion that opens up interesting avenues for future work.