Reproduce “Task Inference Beyond Least Squares in Behavioral Foundation Models” on OpenResearch

About this paper

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

BLS refines task-vector inference for behavioral foundation models at test time, adjusting ordinary least squares to rank high-reward states above low-reward states while staying near its solution. This ranking objective is a surrogate for aligning the policy’s future-state distribution, not a direct measure. On OGBench with HILP, success across nine tasks rises from 31.1% to 39.1%, averaged over three seeds.

Key insights

• Low average reward reconstruction error does not necessarily preserve the reward ordering that determines which states a policy should visit.
• BLS uses high-versus-low reward separation as a practical surrogate for behavior similarity, rather than directly optimizing successor-measure mismatch.
• The appropriate strength of the trust-region constraint depends on reward structure: more flexibility tended to help sparse-reward tasks, while stronger anchoring suited dense-reward tasks.

Summary by alphaXiv.