Model predictive task sampling for efficient and robust adaptation.

Abstract

Foundation models and generalist policies rely on adaptation learning, including meta-training, domain randomization, and supervised finetuning. Challenging task prioritized sampling is essential for adaptation robustness under distribution shifts. However, ranking task difficulty requires exhaustive evaluation, which is prohibitively expensive whether through agent-environment interactions in meta reinforcement learning and domain randomization, or repeated forward passes in large model finetuning. We present Model Predictive Task Sampling (MPTS), a framework that bridges task space and adaptation risk distributions for active task selection. MPTS employs a lightweight generative model to predict task-specific adaptation risk, amortizing costly evaluation while provably ranking task difficulties. It can seamlessly integrate into zero-shot, few-shot, and supervised finetuning settings. We conduct extensive experiments in pattern recognition using foundation models and sequential decision-making. Empirically, MPTS significantly enhances adaptation robustness for tail risk or out-of-distribution tasks compared to superior methods, and improves learning efficiency over subset selection methods such as CVaR_α.