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

2026-05-16

AI Summary

Existing behavioral foundation models exhibit insufficient robustness under environmental dynamics shifts, such as changes in friction, actuation, or sensor noise. This work proposes a novel framework that formulates task inference as a robust minimax optimization problem, enabling adaptation to worst-case dynamics perturbations using only offline data from a single nominal environment—without requiring retraining or additional online interaction. By explicitly incorporating dynamics robustness into the task inference stage for the first time, the method significantly outperforms standard behavioral foundation models and current robust offline imitation learning approaches, demonstrating both the effectiveness and superiority of the proposed paradigm.

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.

Problem

• Behavior Foundation Models
• Offline Imitation Learning
• Dynamics Shifts
• Robustness
• Task Inference

Innovation

• Behavior Foundation Models
• robust imitation learning
• dynamics shift
• minimax optimization
• offline reinforcement learning