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

arXiv: https://arxiv.org/abs/2605.17017 · HTML: https://mlsanity.com/paper/2605.17017/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Imitation learning methods, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 46 of its 63 references that are in the map.
- First posted on arXiv 16 May 2026.

## Placed in

- Offline RL and imitation learning from data / [Imitation learning methods](https://mlsanity.com/offlinerl/imitation-learning-methods/)

## Builds on (most cited first, 30 of 46)

- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Convex optimization](https://mlsanity.com/paper/2106.01946/) (2021)
- [Generative Adversarial Imitation Learning](https://mlsanity.com/paper/1606.03476/) (2016)
- [Off-Policy Deep Reinforcement Learning without Exploration](https://mlsanity.com/paper/1812.02900/) (2018)
- [Exploration by Random Network Distillation](https://mlsanity.com/paper/1810.12894/) (2018)
- [DeepMind Control Suite](https://mlsanity.com/paper/1801.00690/) (2018)
- [Sim-to-Real Transfer of Robotic Control with Dynamics Randomization](https://mlsanity.com/paper/1710.06537/) (2017)
- [Exploring the Limitations of Behavior Cloning for Autonomous Driving](https://mlsanity.com/paper/1904.08980/) (2019)
- [Successor Features for Transfer in Reinforcement Learning](https://mlsanity.com/paper/1606.05312/) (2016)
- [One-Shot Imitation Learning](https://mlsanity.com/paper/1703.07326/) (2017)
- [One-Shot Visual Imitation Learning via Meta-Learning](https://mlsanity.com/paper/1709.04905/) (2017)
- [DART: Noise Injection for Robust Imitation Learning](https://mlsanity.com/paper/1703.09327/) (2017)
- [Imitation Learning via Off-Policy Distribution Matching](https://mlsanity.com/paper/1912.05032/) (2019)
- [Reinforcement Learning with Prototypical Representations](https://mlsanity.com/paper/2102.11271/) (2021)
- [Learning One Representation to Optimize All Rewards](https://mlsanity.com/paper/2103.07945/) (2021)
- [Imitation Learning from Imperfect Demonstration](https://mlsanity.com/paper/1901.09387/) (2019)
- [APS: Active Pretraining with Successor Features](https://mlsanity.com/paper/2108.13956/) (2021)
- [Universal Successor Features Approximators](https://mlsanity.com/paper/1812.07626/) (2018)
- [Don't Change the Algorithm, Change the Data: Exploratory Data for Offline Reinforcement Learning](https://mlsanity.com/paper/2201.13425/) (2022)
- [Robust Reinforcement Learning using Offline Data](https://mlsanity.com/paper/2208.05129/) (2022)
- [Foundation Policies with Hilbert Representations](https://mlsanity.com/paper/2402.15567/) (2024)
- [Does Zero-Shot Reinforcement Learning Exist?](https://mlsanity.com/paper/2209.14935/) (2022)
- [Learning Successor States and Goal-Dependent Values: A Mathematical Viewpoint](https://mlsanity.com/paper/2101.07123/) (2021)
- [Task-Embedded Control Networks for Few-Shot Imitation Learning](https://mlsanity.com/paper/1810.03237/) (2018)
- [Generalization in Dexterous Manipulation via Geometry-Aware Multi-Task Learning](https://mlsanity.com/paper/2111.03062/) (2021)
- [Distributional Robustness and Regularization in Reinforcement Learning](https://mlsanity.com/paper/2003.02894/) (2020)
- [Watch, Try, Learn: Meta-Learning from Demonstrations and Reward](https://mlsanity.com/paper/1906.03352/) (2019)
- [Robust Imitation Learning from Noisy Demonstrations](https://mlsanity.com/paper/2010.10181/) (2020)

## Abstract (arXiv)

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.