Search papers, labs, and topics across Lattice.
This paper introduces Dynamic Neural Koopman Distillation, a method to distill diffusion model policies for real-time robot control by learning a single-step transition model that approximates the iterative denoising process. The core of the approach is a Factorized Dynamic Koopman layer, which models latent transitions with state-dependent modal gains, capturing the multimodal nature of the diffusion policy. Experiments on D4RL benchmarks and a physical Kinova manipulator show that the method outperforms existing one-step distillation techniques and achieves millisecond-level inference latency, enabling smooth closed-loop control.
Unlock real-time robot control with diffusion models: distill complex, multi-step policies into a single, fast forward pass using a novel Koopman operator-inspired architecture.
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/.
CMU ML