🤖 machine learning

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

This paper proposes a robust minimax optimization framework for Behavior Foundation Models that enables effective offline imitation learning and adaptation to unseen dynamics shifts using only nominal environment data, significantly outperforming existing baselines without requiring pretraining modifications.

Original authors: Rishabh Agrawal, Rahul Jain, Ashutosh Nayyar

Published 2026-05-19

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Rishabh Agrawal, Rahul Jain, Ashutosh Nayyar

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a robot to walk, run, or jump. Traditionally, you'd show the robot a video of an expert doing the task, and the robot tries to copy it. This is called Imitation Learning.

However, there's a catch: robots are often trained in a perfect, frictionless simulation (like a video game), but when they go to the real world, things change. The floor might be slippery, the robot's joints might be rusty, or its motors might be weaker. If the robot was only taught to copy the "perfect" video, it will likely fall over the moment reality gets messy.

This paper introduces a new way to teach robots that makes them much tougher against these real-world changes, without needing to retrain them from scratch.

Here is the breakdown using simple analogies:

1. The Problem: The "Perfect Practice" Trap

The authors use something called Behavior Foundation Models (BFMs). Think of a BFM like a master chef who has tasted thousands of different dishes and learned the fundamental "flavor profiles" of cooking.

How it usually works: The chef learns these flavors in a perfect kitchen with perfect ingredients. When asked to cook a new dish, they can quickly figure out the recipe by looking at just a few ingredients.
The flaw: If the chef tries to cook in a kitchen where the stove is broken, the water is salty, or the ingredients are old, the "perfect kitchen" recipe fails. The chef assumes the kitchen is still perfect, so the dish turns out terrible.

In robotics, this means if the robot's environment changes (like a change in gravity or friction), the standard robot fails because it never learned to handle "broken" kitchens.

2. The Solution: The "Worst-Case" Chef

The authors propose a new method called RBFM (Robust Behavior Foundation Model). Instead of just copying the expert, they teach the robot to prepare for the worst possible version of the environment while it's learning the task.

Imagine the chef is now training with a troublemaker standing next to them.

The Game: The chef tries to figure out the recipe (the task). The troublemaker tries to sabotage the kitchen by making the stove hotter, the floor slippery, or the ingredients stale (simulating "dynamics shifts").
The Goal: The chef must find a recipe that still tastes good even when the troublemaker is doing their worst.
The Result: The chef learns a "robust" recipe. When they go to the real world, even if the stove is slightly broken or the floor is wet, the dish still works because they practiced for those exact scenarios.

3. The Magic Trick: No New Training Needed

Usually, to teach a robot to handle broken stoves, you would need to show it videos of cooking in broken kitchens. This is expensive and hard to do.

The paper's big breakthrough is that you don't need new videos.

They take the robot that was already trained on the "perfect kitchen" data.
They only change the final step (called "Task Inference").
When the robot is asked to do a new job, it runs a quick mental simulation: "If the floor is slippery, how should I move my feet? If the motor is weak, how much force should I use?"
It calculates the best move by assuming the environment might be slightly broken, all while using the same old data it learned from.

4. Two Versions of the Robot

The paper offers two ways to do this "worst-case" thinking, trading off speed for safety:

RBFM-Light (The Quick Thinker): This version is fast. It makes a quick guess about how bad the environment might be and adjusts its plan slightly. It's like a driver who sees a puddle and slows down a bit, just in case. It's very fast to compute.
RBFM-Heavy (The Careful Planner): This version is slower but more thorough. It doesn't just guess; it mathematically proves that its plan will work even if the environment changes in specific, complex ways. It's like a driver who checks the map, the weather, and the road conditions before taking a single step. It takes a bit more time but is safer.

5. The Results: Winning in the Real World

The authors tested this on robots that walk, run, and jump (like a human, a dog, and a cheetah). They messed up the physics in the test:

They made gravity stronger.
They made the ground slippery or soft.
They made the robot's joints stiff or weak.

The Outcome:

Standard Robots (FB-IL): When the environment changed, they fell apart immediately.
Old "Robust" Methods: These worked well but were incredibly slow, taking hours to figure out a single task.
The New RBFM Robots: They handled the changes much better than the standard robots and were thousands of times faster than the old robust methods. They could adapt to a new, broken environment in minutes, not hours.

Summary

The paper says: "Don't just teach the robot to copy the expert; teach it to copy the expert while assuming the world might be broken."

By doing this "worst-case" planning at the very last step of learning, they created a system that is both super fast (like a foundation model) and super tough (like a robot that can handle real-world chaos), all without needing to collect new, messy data.

Technical Summary: When Dynamics Shift, Robust Task Inference Wins

Problem Statement

Behavior Foundation Models (BFMs) have emerged as a scalable solution for imitation learning (IL), enabling agents to rapidly adapt to new tasks by pretraining task-agnostic representations on diverse offline data. However, standard BFMs operate under a critical assumption: that the environment dynamics during pretraining and deployment are identical. In real-world scenarios, this assumption is frequently violated due to sim-to-real discrepancies, mechanical wear, sensor noise, and environmental variations (e.g., changes in friction, actuation limits, or gravity).

Existing remedies for this fragility, such as dynamics-conditioned belief encoders or memory-augmented architectures, typically require pretraining data spanning a wide variety of dynamic regimes. This requirement is often expensive, time-consuming, and infeasible in practical settings where only offline data from a single nominal environment is available. The central question addressed by this work is: Can BFMs be made robust to environment dynamics shifts when offline data from only a single nominal environment is available?

Methodology

The authors propose a framework that shifts the burden of robustness from the pretraining phase to the task-inference phase. Instead of modifying the pretrained BFM or requiring multi-dynamics pretraining data, they formulate the task-inference problem as a robust minimax optimization over an uncertainty set of environment dynamics.

The core approach involves learning a latent task embedding $z$ that minimizes the imitation loss against an adversary. The adversary, in turn, selects a transition model within a specified uncertainty set (defined by a total variation or $f$ -divergence ball around the nominal dynamics) that maximizes this loss. This yields a task embedding that "hedges" against plausible dynamics shifts without altering the pretraining procedure.

Two specific algorithms are introduced to balance computational efficiency and robustness:

RBFM-Light: This variant relaxes the dynamics perturbations into a state-occupancy uncertainty set. By characterizing the sensitivity of the discounted state-occupancy measure to transition kernel perturbations, the authors derive a tractable offline objective. This approach introduces a single robustness parameter optimized jointly with the task embedding $z$ , achieving robustness with minimal computational overhead.
RBFM-Heavy: This variant enforces robustness through explicit Bellman flow constraints. It ensures that the perturbed occupancy measures correspond to valid stationary distributions induced by the expert policy under the underlying MDP structure. By incorporating realizability constraints and using $f$ -divergence (specifically a smooth SoftTV generator), this method avoids overly conservative solutions that might arise from unconstrained occupancy perturbations. It solves a bilevel optimization problem involving the estimation of dual variables (Lagrange multipliers) and the task vector $z$ .

Both methods rely solely on offline data from the nominal environment for pretraining and a small set of expert demonstrations for task inference. No expert data from perturbed environments is required at any stage.

Key Contributions

Robust Formulation: The paper proposes a robust formulation of BFM task inference via minimax optimization over an uncertainty set of environment dynamics. This enables robustness to dynamics shifts without any modification to the BFM pretraining process.
Algorithmic Variants: The authors introduce RBFM-Light and RBFM-Heavy, two algorithms that offer a trade-off between computational efficiency and robustness while preserving the fast task-inference property inherent to BFMs.
Data Efficiency: The approach requires only a single exploratory dataset collected in the nominal setting for pretraining and a small set of expert demonstrations for inference. This contrasts with prior works that rely on data from multiple diverse environments to learn robust representations.
Novelty: To the best of the authors' knowledge, this is the first work to achieve robustness to environment dynamics variation within the BFM framework without requiring pretraining or expert data from multiple dynamics regimes.

Experimental Results

The authors evaluated their methods on the ExORL benchmark across three morphologically diverse environments: Walker, Quadruped, and Cheetah, using datasets generated by unsupervised exploration strategies (RND, APS, PROTO). They introduced various physics perturbations, including gravity changes, body mass scaling, joint friction, tilted gravity, and actuator control range clipping.

Robustness vs. Standard BFM: Across all environments and perturbation types, both RBFM variants consistently outperformed the standard FB-IL baseline. As perturbation severity increased, the performance gap widened, with FB-IL exhibiting the steepest degradation. RBFM-Heavy generally achieved the highest returns, followed by RBFM-Light.
Comparison with Single-Task Robust IL: The proposed methods were compared against single-task distributionally robust offline IL baselines (DRBC and BE-DROIL). While single-task baselines outperformed FB-IL at high perturbation levels, the RBFM framework maintained consistent superiority over both FB-IL and the single-task baselines across the full perturbation sweep. This highlights the advantage of multi-task pretraining combined with robust inference.
Computational Efficiency: A critical finding is the efficiency of the approach. Single-task robust baselines required hours (5–6 hours) per task for optimization. In contrast, RBFM variants reduced task-specific inference time to minutes (5–10 minutes), making them orders of magnitude faster while achieving superior robustness.
Data and Parameter Sensitivity: The study showed that robust task inference critically depends on the quality of the learned FB representations, requiring sufficient pretraining data (e.g., 1M transitions) to become expressive. Additionally, the robustness parameter $\epsilon$ was found to control the trade-off between nominal performance and robustness, with an intermediate value ( $\epsilon=0.8$ ) yielding the best overall trade-off.

Significance and Claims

The paper claims that robust policies can be achieved entirely at task-inference time, significantly improving the practicality of BFMs in dynamic settings. By decoupling robustness from pretraining, the proposed framework allows a single pretrained model to be robustly adapted to new tasks using only offline data from a single nominal environment.

The authors emphasize that this approach addresses a fundamental limitation of current BFMs—their brittleness under dynamics mismatch—without incurring the prohibitive costs of multi-environment pretraining or the computational inefficiency of per-task robust optimization. The results demonstrate that robustness can be "baked in" during the inference stage, offering a scalable path for deploying foundation models in real-world, unstructured environments where dynamics are subject to unpredictable shifts.