Inference-Time Enhancement of Generative Robot Policies via Predictive World Modeling

Han Qi, Haocheng Yin, Aris Zhu, Yilun Du, Heng Yang

Abstract: We present generative predictive control (GPC), a framework for inference-time enhancement of pretrained behavior-cloning policies. Rather than retraining or fine-tuning, GPC augments a frozen diffusion policy at deployment by coupling it with a predictive world model. Concretely, we train an action-conditioned world model on expert demonstrations and random exploration rollouts to forecast the consequences of action proposals produced by the diffusion policy, then perform lightweight online planning that ranks and refines these proposals via model-based look-ahead. This combination of a generative prior with predictive foresight enables test-time adaptation. Across diverse robotic manipulation tasks—state- and vision-based, in simulation and on real hardware—GPC consistently outperforms standard behavior cloning and compares favorably to other inference-time adaptation baselines.