Organizations: University of Applied Science and Arts of Southern Switzerland, Department of Innovative Technologies, IDSIA-SUPSI, Lugano, Switzerland. · Università della Svizzera Italiana, Faculty of Informatics, Lugano, Switzerland. · Politecnico di Milano, Mechanical Department, Milano, Italy.
Accurate modeling of robot dynamics is essential for model-based control, yet remains challenging under distributional shifts and real-time constraints. In this work, we formulate system identification as an in-context meta-learning problem and compare deterministic and generative sequence models for forward dynamics prediction. We take a Transformer-based meta-model, as a strong deterministic baseline, and introduce to this setting two complementary diffusion-based approaches: (i) inpainting diffusion (Diffuser), which learns the joint input-observation distribution, and (ii) conditioned diffusion models (CNN and Transformer), which generate future observations conditioned on control inputs. Through large-scale randomized simulations, we analyze performance across in-distribution and out-of-distribution regimes, as well as computational trade-offs relevant for control. We show that diffusion models significantly improve robustness under distribution shift, with inpainting diffusion achieving the best performance in our experiments. Finally, we demonstrate that warm-started sampling enables diffusion models to operate within real-time constraints, making them viable for control applications. These results highlight generative meta-models as a promising direction for robust system identification in robotics.
Figures & tables
Fig. 1: Comparison of sequence models. Diffusion models ( Inpainting and Conditioned ) are contrasted with a deterministic Transformer ( RoboMorph ). Inpainting Diffusion learns the full trajectory distribution over y1:N from context (the green area). Conditioned Diffusion models system dynamics by generating ym:N conditioned on u1:N , producing future trajectories under input conditioning with smoother denoising updates. Diffusion models require multiple iterative steps (e.g., 20 in this example) to progressively denoise an initial random distribution, refining predictions over time. In contrast, the deterministic Transformer predicts ym:N directly from um:N and context in a single forward pass.
Dataset and Signal
A[Nm]
f[Hz]
D1 – CH
[−4,4]
0.3
D2 – CH
[−4,4]
[0.2,0.4]
D3 – CH
[−4,4]
[0.2,0.6]
D4 – CH
[−4,4]
[0.1,0.7]
D1 – MS
[−30f,30f]
0.15
D2 – MS
[−30f,30f]
[0.05,0.15]
TABLE I: Simulation Parameters Randomization
Fig. 2: Meta-modeling of joint dynamics for chirp and multi-sinusoidal joint torques with fCH=1.0Hz , fMS=0.45Hz master frequencies (OOD), trained on D1 . The results indicate one prediction per signal among 100 randomized scenarios. Diffusion-based architectures seem to better capture global complexity, particularly in OOD signals, such as in this example.
Fig. 3: By selectively covering parts of the domain, it is possible to meta-learn the class of frequency response. For chirp signals of different randomization: fD1=0.30 , fD2=[0.02,0.4] , fD3=[0.2,0.6] , fD4=[0.1,0.7] and sinusoidal signals of different randomization bounds fD1=0.15 , fD2=[0.05,0.15] , fD3=[0.05,0.25] , fD4=[0.01,0.30] the ID and OOD regions shift which effectively enlarges the accuracy of the predictive domain. This improvement on the edges of the domain minimally affects the accuracy observed for the centering points.
Fig. 4: Diffusion-based inference is about 2 order of magnitude larger than non-diffusion-based inference and is a bottleneck of predictive latency. This can be alleviated by warm-starting diffusion models from past trajectories, bringing the inference time to about 40ms with 5 denoising steps.
Fig. 5: As an example, we focus on architectures trained on D2 for chirp and sinusoidal signals with 5 diffusion steps corresponding to about 40ms inference time: Transformer-based models sustain warm-started trajectories relatively well whereas convolution-based models suffer drastically. Among Transformer-based models, diffusion processes provide more versatile representations in modeling high frequency responses whereas for low frequency responses this versatility is practically insignificant.
Diffusion models excel at generating diverse, multimodal trajectories for robotic control, yet their iterative denoising process introduces latency that can constrain real-time closed-loop execution. To address this problem, we propose a Dynamic Neural Koopman (DNK) distillation framework, which distills multistep diffusion inference into a single forward pass conditioned on observations and sampled noise. Specifically, we introduce a Factorized Dynamic Koopman (FDK) layer that models the denoising process through a latent linear transition. The noisy trajectory input is lifted into a latent space, where the FDK layer parameterizes a factorized linear operator with condition-dependent eigenvalues that adapt the noisy-to-denoised transition to the current context. We evaluate DNK across robot-control benchmarks that span locomotion, state-based, and image-based manipulation, comparing against accelerated generative-policy baselines. The results demonstrate that DNK achieves competitive or improved closed-loop performance relative to accelerated generative-policy baselines, while substantially reducing inference latency relative to the strongest high-performing one-step baselines. On long-horizon manipulation tasks, DNK achieves competitive performance with substantially fewer inference parameters than the one-step policy baselines. Hardware experiments on a physical Kinova manipulator further demonstrate millisecond-scale inference and reduced task completion time. A project page is available at https://fdkoopman.github.io/.
Lei Zheng, Peiqi Yu, Zengqi Peng +2
Department of Electrical and Computer Engineering, National University of Singapore, Singapore · Department of Electrical and Computer Engineering, Carnegie Mellon University, USA · Robotics and Autonomous Systems Thrust, The Hong Kong University of Science and Technology, Hong Kong SAR, China
Generative policies such as diffusion models provide expressive data-driven priors for robot behavior, yet the situations a robot faces at execution time rarely match its training data and impose constraints and objectives the policy never learned. Existing inference-time guidance methods either rely on gradients that can push samples away from the learned prior, or on sampling procedures that can be inefficient, collapse modes, or average them into out-of-distribution predictions. We introduce MOSAIC, a mode- and support-aware framework for gradient-free inference-time control of robotic diffusion models. MOSAIC treats the initial noise and all stochastic reverse-step perturbations as a joint input to the pretrained denoising process and adapts a distribution over this input using importance weights computed from fully denoised outputs. Physical feasibility and consistency with the learned support are balanced using adaptive penalties, while mode-wise importance-weighted updates preserve modes and avoid destructive averaging. The same formulation can jointly steer multiple independently trained diffusion models through shared coupling costs, enabling coordinated behavior without requiring a joint model. We evaluate MOSAIC on navigation, dexterous grasping, motion planning with unseen geometry, and collaborative box manipulation with independently trained robot priors. Across these settings, MOSAIC adapts and composes frozen diffusion priors under new deployment constraints without retraining.
Sabrina Bodmer, Rahel Rickenbach, René Zurbrügg +2
Department of Mechanical and Process Engineering, ETH Zürich, Switzerland. · Automatic Control Laboratory, EPFL, Switzerland.
Pretrained generative Diffusion Transformers (DiTs) capture rich pixel-level visual and language-conditioned structure through large-scale image and video generation training. A growing line of robot policies builds on this generative prior, but how it should be transferred to control remains unclear, and existing approaches commonly instantiate this transfer through future visual prediction. We ask a more basic question: what a pretrained generative DiT actually contributes to action learning, and how this prior should be adapted for control. We introduce NowWAM, a future-target-free co-training formulation that denoises the current observation and predicts robot actions from the same visual stream, directly coupling the native generative objective to the action-facing representation across the denoising trajectory. Under matched controlled settings, past and future visual targets perform comparably, while restricting training to the clean endpoint substantially reduces robustness, suggesting that a separate future target is not essential for generative adaptation, while the denoising trajectory remains an effective interface for control. On LIBERO-Plus, NowWAM reaches 87.7% with FLUX2-Klein, improving over the future-target co-training baseline by 6.1 points while halving training visual tokens (784 to 392) and reducing step time from 2.85 s to 1.63 s, a 1.8x speedup. With the pure text-to-image Z-Image backbone, NowWAM further reaches 87.8%, showing that strong control adaptation is not tied to video generation or image-editing backbones.