When Dynamics Shift, Robust Task Inference Wins: Offline Imitation Learning with Behavior Foundation Models Revisited

TL;DR

This paper introduces a robust task inference framework for Behavior Foundation Models that adapts to environment dynamics shifts using a minimax optimization approach, enhancing offline imitation learning robustness.

Contribution

It formulates BFM task inference as a minimax problem to handle worst-case dynamics shifts without retraining, a novel approach in the field.

Findings

Outperforms standard BFM under dynamics shifts

Achieves robustness solely with offline data from a single environment

Demonstrates improved practical applicability in dynamic real-world settings

Abstract

Behavior Foundation Models (BFMs) enable scalable imitation learning (IL) by pretraining task-agnostic representations that can be rapidly adapted to new tasks. However, existing BFMs assume fixed environment dynamics, limiting their robustness under real-world shifts such as changes in friction, actuation, or sensor noise. We address this by formulating BFM task-inference as a robust minimax optimization problem, enabling adaptation to worst-case dynamics perturbations without modifying pretraining. To the best of our knowledge, this is the first BFM-based framework that achieves robustness to dynamics shifts while relying solely on offline data from a single nominal environment. Our approach significantly outperforms standard BFM and robust offline IL baselines under dynamics shifts. These results demonstrate that robust policy can be achieved entirely at task-inference time, improving the practicality of BFMs in dynamic settings.