Optimistic Task Inference for Behavior Foundation Models
Thomas Rupf, Marco Bagatella, Marin Vlastelica, Andreas Krause
TL;DR
This paper addresses the data-efficiency bottleneck of Behavior Foundation Models (BFMs) by introducing OpTI-BFM, an online, optimistic method to infer task rewards through environment interaction rather than offline labeling. By leveraging a near-linear relation between rewards, features, and successor features, OpTI-BFM maintains a least-squares estimate of the task embedding and uses confidence sets to select informative task embeddings, achieving sublinear regret in an episodic linear-bandit-style setting. The authors provide a formal regret bound, present practical variants including Thompson Sampling and non-stationary reward adaptation, and demonstrate that OpTI-BFM identifies unseen reward functions within a handful of episodes on standard benchmarks, often matching or surpassing offline reward-inference baselines. The work bridges BFMs and linear-bandit theory, enabling efficient, data-light task inference with potential applicability to pixel-based reward learning and broader online RL scenarios.
Abstract
Behavior Foundation Models (BFMs) are capable of retrieving high-performing policy for any reward function specified directly at test-time, commonly referred to as zero-shot reinforcement learning (RL). While this is a very efficient process in terms of compute, it can be less so in terms of data: as a standard assumption, BFMs require computing rewards over a non-negligible inference dataset, assuming either access to a functional form of rewards, or significant labeling efforts. To alleviate these limitations, we tackle the problem of task inference purely through interaction with the environment at test-time. We propose OpTI-BFM, an optimistic decision criterion that directly models uncertainty over reward functions and guides BFMs in data collection for task inference. Formally, we provide a regret bound for well-trained BFMs through a direct connection to upper-confidence algorithms for linear bandits. Empirically, we evaluate OpTI-BFM on established zero-shot benchmarks, and observe that it enables successor-features-based BFMs to identify and optimize an unseen reward function in a handful of episodes with minimal compute overhead. Code is available at https://github.com/ThomasRupf/opti-bfm.