Dynamic Neural Koopman Distillation for Real-Time Robot Control Using Diffusion Models
For roboticists needing real-time control with multimodal trajectory generation, this method reduces inference latency to enable high-frequency closed-loop execution without sacrificing expressivity.
The paper tackles the latency issue of diffusion models for real-time robot control by distilling multistep inference into a single forward pass, achieving millisecond-level inference while outperforming existing one-step distillation methods on D4RL locomotion benchmarks and demonstrating smooth closed-loop control on a physical robot.
Diffusion models excel at generating diverse and multimodal trajectories for robotic planning, yet their iterative denoising process introduces latency that is incompatible with high-frequency closed-loop control. To address this problem, we propose Dynamic Neural Koopman Distillation, a framework that distills multistep diffusion inference into a single forward pass while retaining the multimodal expressivity of the teacher model. Specifically, we introduce a Factorized Dynamic Koopman layer that models the denoising process through a factorized latent transition with state-dependent modal gains. We evaluate the proposed method on standard D4RL MuJoCo locomotion benchmarks and a physical Kinova manipulator, comparing against one-step baselines. The results show that our method significantly outperforms existing one-step distillation approaches on the reported locomotion tasks, and reduces the inference latency to the millisecond regime compared with the teacher policy. Hardware experiments further demonstrate that our method enables smooth and fast closed-loop execution while maintaining task success and comparable accuracy. A project page is available at https://fdkoopman.github.io/.