Diffusion Sequence Models for Generative In-Context Meta-Learning of Robot Dynamics

Angelo Moroncelli; Matteo Rufolo; Gunes Cagin Aydin; Asad Ali Shahid; Loris Roveda

cs.LGApr 15, 2026

Diffusion Sequence Models for Generative In-Context Meta-Learning of Robot Dynamics

Authors: Angelo Moroncelli, Matteo Rufolo, Gunes Cagin Aydin, Asad Ali Shahid, Loris Roveda

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.

Abstract

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 $y_{1:N}$ from context (the green area). Conditioned Diffusion models system dynamics by generating $y_{m:N}$ conditioned on $u_{1: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 $y_{m:N}$ directly from $u_{m:N}$ and context in a single forward pass.

Dataset and Signal	$A[Nm]$	$f[Hz]$
$D_{1}$ – $CH$	$[-4,4]$	$0.3$
$D_{2}$ – $CH$	$[-4,4]$	$[0.2,0.4]$
$D_{3}$ – $CH$	$[-4,4]$	$[0.2,0.6]$
$D_{4}$ – $CH$	$[-4,4]$	$[0.1,0.7]$
$D_{1}$ – $MS$	$[-30f,30f]$	$0.15$
$D_{2}$ – $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

f_{CH}=1.0Hz

f_{MS}=0.45Hz

master frequencies (OOD), trained on

D_{1}

. 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:

f_{D_{1}}=0.30

f_{D_{2}}=[0.02,0.4]

f_{D_{3}}=[0.2,0.6]

f_{D_{4}}=[0.1,0.7]

and sinusoidal signals of different randomization bounds

f_{D_{1}}=0.15

f_{D_{2}}=[0.05,0.15]

f_{D_{3}}=[0.05,0.25]

f_{D_{4}}=[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

D_{2}

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.

Robot Dynamics Identification

Robot dynamics identification estimates the masses, inertias, friction, actuator behavior, and other forces that govern robot motion.

See all 57 papers

Dynamic Neural Koopman Distillation for Fast Robot Control Using Diffusion Models
May 24, 2026Lei Zheng, Peiqi Yu, Zengqi Peng +2Robotic Control Koopman Operator Learning
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
MOSAIC: Mode- and Support-Aware Inference-Time Control of Robot Diffusion Models
Oct 12, 2026Sabrina Bodmer, Rahel Rickenbach, René Zurbrügg +2Inference-Time Guidance Diffusion Guidance
Department of Mechanical and Process Engineering, ETH Zürich, Switzerland. · Automatic Control Laboratory, EPFL, Switzerland.
Beyond Future Prediction: Denoising as Generative Adaptation for Robot Control
Sep 23, 2026Zanyi Wang, Yuheng Lei, Dengyang Jiang +4Robotic Control Robot Learning
UC San Diego · HKU · HKUST +1

May 24, 2026J/K move · Enter open · S save

Dynamic Neural Koopman Distillation for Fast Robot Control Using Diffusion Models

Lei Zheng, Peiqi Yu, Zengqi Peng, Changliu Liu+1

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

Robotic Control Koopman Operator Learning Diffusion Model Distillation Diffusion Policy Diffusion Model Inference Acceleration

Open paper AskPDF arXiv

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/.

Abstract

Robot Dynamics Identification

Explore similar work

Dynamic Neural Koopman Distillation for Fast Robot Control Using Diffusion Models

MOSAIC: Mode- and Support-Aware Inference-Time Control of Robot Diffusion Models

Beyond Future Prediction: Denoising as Generative Adaptation for Robot Control