Behavioral models improve task detection for better robot control

Task Inference Beyond Least Squares in Behavioral Foundation Models

Machine Learning

Summary

Robots often need to figure out what task they should do by understanding rewards or goals, but current methods can guess wrong and lead to less effective actions. The authors introduced a new way called BLS that balances understanding rewards while making sure the robot's behavior matches what truly works best. This new method helps robots perform better on tasks like walking, moving objects, or human-like motions, and it does so without needing extra computing power. They tested it on several challenging robot control tasks and saw clear improvements.

What this means in practice

  • For robot control engineers: Improve robot task recognition and control accuracy in locomotion and manipulation without costly retraining.
  • For game ai developers: Create more adaptive game agents that infer player goals better and respond naturally in complex control environments.

Authors

Kuan-Hsun Tu, Chien-Sheng Chiang, Hsin-Wei Chen, Ping-Chun Hsieh, Tsung-Wei Ke

Abstract

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