AIF-Guided Diffusion Planning for Nonstationary Control

Yulin Li, Mohsen a jafari, Andrea Matta

Keywords: Active Inference, Diffusion planning, Online Adaptation, Nonstationary Control

Abstract: We present AIF-guided diffusion planning for rapid, reward-free adaptation to abrupt, within-episode dynamics shifts in continuous control. Our controller couples Active Inference (AIF)—which maintains a compact belief over latent dynamics and pursues actions that minimize Expected Free Energy (EFE), balancing goal pursuit and information gain—with a conditional trajectory diffusion policy conditioned on state, goal, and a compact embedding of the regime belief. At deployment, prediction surprise under a hazard prior triggers Bayesian changepoint resets and fast belief correction; the updated belief both guides denoising and supplies an EFE-based scorer over candidate plans, executing the first action of the minimum‑EFE proposal in closed loop. Trained offline on regime‑stratified datasets (no cross‑regime transitions), the method adapts without test‑time rewards or ensembles. On nonstationary MuJoCo HalfCheetah with mid‑episode gravity switches, it improves mid‑window reward (steps 100–200) and preserves final return relative to strong offline baselines. An unsupervised variant that discovers regimes from trajectories remains robust, though slower to recover; ablations indicate the belief pathway and EFE guidance drive the gains.

Primary Area: generative models

Submission Number: 23773