When Dynamics Shift, Robust Task Inference Wins: Offline Imitation Learning with Behavior Foundation Models Revisited
For practitioners deploying imitation learning in real-world environments with changing dynamics, this work enables robust task inference without modifying pretraining or requiring additional data.
Behavior Foundation Models (BFMs) for imitation learning fail under dynamics shifts (e.g., friction, actuation changes). The authors propose a robust minimax optimization for task inference that adapts to worst-case dynamics perturbations using only offline data from a single nominal environment, significantly outperforming baselines.
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.