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

TL;DR

This paper explores diffusion-based sequence models for robot dynamics prediction, demonstrating improved robustness under distribution shifts and real-time operation capabilities compared to deterministic models.

Contribution

It introduces diffusion-based approaches for system identification, showing their advantages over deterministic models in robustness and real-time control scenarios.

Findings

Diffusion models outperform deterministic models under distribution shifts.

Inpainting diffusion achieves the best accuracy in experiments.

Warm-started sampling enables real-time operation of diffusion models.

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.