# Task Inference Beyond Least Squares in Behavioral Foundation Models

arXiv: https://arxiv.org/abs/2610.05350 · HTML: https://mlsanity.com/paper/2610.05350/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 20 of its 23 references that are in the map.
- First posted on arXiv 4 Oct 2026.

## Builds on (most cited first, 20 of 20)

- [DeepMind Control Suite](https://mlsanity.com/paper/1801.00690/) (2018)
- [Successor Features for Transfer in Reinforcement Learning](https://mlsanity.com/paper/1606.05312/) (2016)
- [OGBench: Benchmarking Offline Goal-Conditioned RL](https://mlsanity.com/paper/2410.20092/) (2024)
- [Learning One Representation to Optimize All Rewards](https://mlsanity.com/paper/2103.07945/) (2021)
- [Don't Change the Algorithm, Change the Data: Exploratory Data for Offline Reinforcement Learning](https://mlsanity.com/paper/2201.13425/) (2022)
- [Foundation Policies with Hilbert Representations](https://mlsanity.com/paper/2402.15567/) (2024)
- [The Laplacian in RL: Learning Representations with Efficient Approximations](https://mlsanity.com/paper/1810.04586/) (2018)
- [Does Zero-Shot Reinforcement Learning Exist?](https://mlsanity.com/paper/2209.14935/) (2022)
- [Learning Successor States and Goal-Dependent Values: A Mathematical Viewpoint](https://mlsanity.com/paper/2101.07123/) (2021)
- [Reinforcement Learning from Passive Data via Latent Intentions](https://mlsanity.com/paper/2304.04782/) (2023)
- [Zero-Shot Whole-Body Humanoid Control via Behavioral Foundation Models](https://mlsanity.com/paper/2504.11054/) (2025)
- [BFM-Zero: A Promptable Behavioral Foundation Model for Humanoid Control Using Unsupervised Reinforcement Learning](https://mlsanity.com/paper/2511.04131/) (2025)
- [Temporal Difference Flows](https://mlsanity.com/paper/2503.09817/) (2025)
- [TD-JEPA: Latent-predictive Representations for Zero-Shot Reinforcement Learning](https://mlsanity.com/paper/2510.00739/) (2025)
- [Fast Adaptation with Behavioral Foundation Models](https://mlsanity.com/paper/2504.07896/) (2025)
- [Optimistic Task Inference for Behavior Foundation Models](https://mlsanity.com/paper/2510.20264/) (2025)
- [Which Features are Best for Successor Features?](https://mlsanity.com/paper/2502.10790/) (2025)
- [Soft Forward-Backward Representations for Zero-shot Reinforcement Learning with General Utilities](https://mlsanity.com/paper/2602.06769/) (2026)
- [Zero-Shot Off-Policy Learning](https://mlsanity.com/paper/2602.01962/) (2026)
- [From Reward-Free Representations to Preferences: Rethinking Offline Preference-Based Reinforcement Learning](https://mlsanity.com/paper/2606.01123/) (2026)

## Abstract (arXiv)

Behavioral Foundation Models (BFMs) aim to solve a wide range of downstream tasks without test-time policy learning by inferring a task vector from the reward function. While efficient, the retrieved policies are often suboptimal because of how this task vector is inferred, typically with ordinary least squares (OLS). OLS minimizes reward reconstruction error but leaves the ordering of rewards unconstrained, which can bias the successor measure of the retrieved zero-shot policy away from that of the optimal policy. In this work, we propose BLS, an efficient test-time inference method that balances minimizing reward reconstruction error with reducing successor-measure mismatch. Theoretically, we provide a suboptimality gap upper bound characterized by both successor-measure and reward-function residuals. Empirically, we evaluate BLS on top of state-of-the-art BFMs across benchmarks for locomotion, manipulation, and humanoid control. BLS outperforms existing task inference baselines with negligible computational overhead. Project page: https://embodiedai-ntu.github.io/BLS