Optimistic Task Inference for Behavior Foundation Models

AI summary

Overview

Research area: Reinforcement learning, specifically zero-shot RL and Behavior Foundation Models (BFMs) built on Universal Successor Features (USFs), with connections to linear bandits and reward/task inference.

Technical level: Advanced. The paper combines practical deep RL benchmarks with a formal regret analysis grounded in linear contextual bandit theory (LinUCB/OFUL-style proofs), confidence ellipsoids, and successor feature algebra.

Scope in one sentence: The paper proposes OpTI-BFM, an optimistic decision criterion that performs task inference for successor-feature-based BFMs purely through test-time environment interaction, and shows it identifies an unseen reward function within a handful of episodes on DeepMind Control Suite tasks.

What This Paper Is About

BFMs can produce a high-performing policy for any reward function specified at test time, but standard task inference requires a labeled dataset — either a functional form of the reward or a costly set of (state, reward) pairs — and often access to the pre-training data. This paper replaces that offline regression step with an online, interactive alternative: the agent starts from an uninformed prior over task embeddings and actively chooses which behaviors to execute so that the rewards it observes are maximally informative. The goal is to recover near-optimal behavior with few reward labels and no reliance on the original pre-training dataset.

Key Contributions

1. An online task inference framework for BFMs. The paper reformulates task inference as an interactive problem where the agent selects task embeddings (not actions directly), executes the corresponding BFM policy, and observes reward labels from the resulting states, removing the need for the pre-training dataset.

2. OpTI-BFM, an optimistic decision criterion. The method maintains a regularized least-squares estimate of the true task embedding plus a confidence ellipsoid, and selects the task embedding that maximizes an upper-confidence acquisition function combining successor features and the belief over rewards.

3. A regret guarantee. Under four stated assumptions (perfect USF, linear reward with sub-Gaussian noise, exact optimization oracle, bounded norms), an episode-level-update variant of OpTI-BFM incurs expected regret R_n ≤ Õ(d√n) — a sublinear rate derived by direct connection to upper-confidence algorithms for linear bandits.

4. Practical variants and empirical validation. The paper provides a practical implementation using random shooting with 128 candidates, a Thompson Sampling variant (OpTI-BFM-TS), and a non-stationary variant using exponential discounting of old data, evaluated on the ExORL benchmark.

Main Findings

• Oracle-level performance within 5 episodes. On the ExORL environments Walker, Cheetah, and Quadruped (four tasks each), OpTI-BFM recovers Oracle performance within 5 episodes, i.e. 5k environment steps of interaction. The Oracle is obtained by solving the closed-form regression in Eq. 4 with a budget of 50k labeled samples from the pre-training dataset.

• Optimism beats Thompson Sampling on some tasks. A significant gap between OpTI-BFM and its TS variant is observed in Cheetah; the paper notes TS remains promising because it avoids any optimization problem.

• LoLA improves more slowly. LoLA, a blackbox policy search method originally introduced for fast adaptation, makes slower progress, which is more visible over 50 episodes (Fig. 16, Appendix C). This is consistent with existing ablations of LoLA initialized to a random task embedding.

• Active data collection is more informative than passive data. When the task embedding is re-fit offline from the first n transitions of each method, data from actively exploring sources (OpTI-BFM, OpTI-BFM-TS, and RND) outperforms data from the passive Random baseline, and OpTI-BFM and its TS variant tend to be more data-efficient than RND, traced to task-awareness.

• Per-step updates beat episode-level updates. Fixing the task embedding for an entire episode (the version covered by the theory) leads to slower improvement than updating every step, although the episodic variants eventually reach equal performance (Fig. 16).

• Adaptation to non-stationary rewards requires forgetting. On two new Walker tasks, speedup and slowdown, which change the velocity target after a burn-in phase within a single 30k-step episode, OpTI-BFM with ρ = 1.0 struggles to adapt after converging in the first 10k steps. Reducing ρ allows it to track the velocity target more accurately, while a ρ that is too small reduces uncertainty too slowly and can lead to suboptimal behavior.

• Warm-starting from labels works. OpTI-BFM can be warm-started by updating its least-squares estimator n times with provided labeled data; in DMC Cheetah, initial performance increases quickly as n grows, making OpTI-BFM an extension of the traditional offline task inference pipeline.

Methodology in Plain English

The starting point is a pre-trained BFM that supplies a family of policies π_z and successor features ψ, where z is a low-dimensional task embedding. Because rewards are assumed to be (approximately) a linear function of features, i.e. r(s) ≈ φ(s)ᵀz_r, the problem of finding the right policy reduces to estimating the unknown vector z_r from observed rewards.

OpTI-BFM keeps a ridge-regularized least-squares estimate of z_r, updated with each (feature, reward) pair observed online. The same quantities yield a confidence ellipsoid around that estimate. At each step, the algorithm picks the task embedding that would give the highest plausible return, where the "plausibility" is determined by the confidence ellipsoid — the standard optimism-in-the-face-of-uncertainty principle. It then runs the BFM policy conditioned on that embedding, observes the reward, and updates.

Because the criterion involves maximizing over two continuous spaces and a nonlinear map from z to successor features, the authors avoid relying on the BFM being perfectly trained and instead search both the task embedding and the reward weights inside the confidence set. They reformulate it into a single-variable objective (a linear term plus a weighted-norm bonus) and approximate the solution by random shooting with 128 candidate samples, drawn from the ellipsoid via a push-forward of uniform samples from the unit ball. Bookkeeping uses the information vector and the Cholesky factor of the Gram matrix, giving O(d²) updates. A Bayesian reading of the same quantities gives the Thompson Sampling variant, which samples the task embedding from the Gaussian posterior instead of solving the optimization.

Experiments use the Forward-Backward (FB) framework as the BFM backbone, trained and evaluated under the standard FB protocol, on Walker, Cheetah, and Quadruped with four tasks each, reporting min-max intervals over 3 training seeds. Comparisons include OpTI-BFM, OpTI-BFM-TS, LoLA, a Random baseline that samples embeddings uniformly from Z, and an Oracle that runs the optimal policy.

Why This Matters

Impact on research. The paper tightens the link between modern scalable BFMs and classical bandit theory, showing that when a BFM is well-trained, control reduces to a bandit problem over tasks. It also shifts attention from compute efficiency (already a strength of zero-shot RL) to label efficiency, an aspect that determines whether BFMs are deployable outside simulation.

Real-world applications (as motivated or implied by the paper):

• Learning from pixels, where evaluating success from a single image is described as neither straightforward nor cheap, making reward labeling the bottleneck.
• Domains where the pre-training dataset is unavailable or proprietary, since OpTI-BFM never needs to revisit it.
• Settings with changing objectives, such as velocity-target tracking tasks, where the non-stationary variant can follow a reward function that shifts over time.
• Settings with a small amount of existing labeled data, where OpTI-BFM can be warm-started and then refined by interaction.

Industry relevance. Any pipeline that ships a pre-trained generalist policy and adapts it to a customer- or deployment-specific objective faces the same labeling cost; reducing adaptation to a handful of episodes with only in-deployment reward feedback is directly commercially relevant. The O(d²) update cost and the 128-candidate search budget mean the overhead on top of an existing BFM is small.

Future Directions

• Extending the regret bound to per-step updates. The formal guarantee covers only the slower episode-level variant, while the per-step version is empirically stronger; closing this gap is flagged as an important direction for future formal work.

• Studying the structure of the embedding space Z. The paper notes that it makes no significant assumption about the structure of Z and that understanding its elements, formally and practically, remains open.

• Combining with fine-tuning. The authors note that while updating only the task embedding is sample-efficient, fine-tuning additional BFM components may yield better long-run performance, and that OpTI-BFM could be combined with approaches like LoLA.

• Beyond current domains. The authors argue that BFMs' scalability up to complex domains, while retaining linearity in an embedding space, is a ripe field for more theoretically principled approaches that may be impactful beyond regular domains.

Target Audience

Researchers and practitioners working on zero-shot reinforcement learning, behavior foundation models, successor features, and unsupervised RL; bandit theorists interested in connections between linear bandit guarantees and policy learning; and engineers deploying pre-trained generalist policies who need to adapt them to new objectives with minimal reward labeling. Readers should be comfortable with MDP formalism, linear regression, confidence sets, and regret analysis, though the experimental results are readable without the appendix proofs.

Authors’ 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.