Model predictive task sampling for efficient and robust adaptation

Abstract Foundation models have revolutionized general-purpose problem-solving, offering rapidtask adaptation through pretraining, meta-training, and finetuning. Recent crucial advances in these paradigms reveal the importance of challenging taskprioritized sampling to enhance adaptation robustness under distribution shifts. However, ranking task difficulties over iteration as a preliminary step typically requiresexhaustive task evaluation, which is practically unaffordable in computation and data-annotation. This study provides a novel perspective to illuminate the possibility of leveraging thedual importance of adaptation robustness and learning efficiency, particularly inscenarios where task evaluation is risky or costly, such as iterative agent-environmentinteractions for robotic policy evaluation or computationally intensive inference steps forfinetuning foundation models. Firstly, we introduce Model Predictive Task Sampling (MPTS), a framework that bridgesthe task space and adaptation risk landscape, providing a theoretical foundation forrobust active task sampling. MPTS employs a generative model to characterize the episodic optimization process andpredicts task-specific adaptation risk via posterior inference. The resulting risk learner amortizes the costly evaluation of task adaptationperformance and provably approximates task difficulty rankings. MPTS seamlesslyintegrates into zero-shot, few-shot, and supervised finetuning settings. Empirically, we conduct extensive experiments in pattern recognition using foundationmodels and sequential decision-making. Our results demonstrate that MPTS significantly enhances adaptation robustness for tailor out-of-distribution (OOD) tasks and improves learning efficiency compared to state-of-the-art (SOTA) methods. The code is available at the project site https://github.com/thu-rllab/MPTS.