Balance Equation-based Distributionally Robust Offline Imitation Learning

TL;DR

This paper introduces a novel distributionally robust offline imitation learning framework that learns policies resilient to environment dynamics shifts using only expert demonstrations, without additional environment interaction.

Contribution

It formulates a robust optimization approach based on balance equations that is fully offline and data-driven, improving robustness against transition model uncertainties.

Findings

Outperforms existing offline IL methods under environment shifts

Achieves higher robustness and generalization in continuous control tasks

Enables offline learning without environment interaction

Abstract

Imitation Learning (IL) has proven highly effective for robotic and control tasks where manually designing reward functions or explicit controllers is infeasible. However, standard IL methods implicitly assume that the environment dynamics remain fixed between training and deployment. In practice, this assumption rarely holds where modeling inaccuracies, real-world parameter variations, and adversarial perturbations can all induce shifts in transition dynamics, leading to severe performance degradation. We address this challenge through Balance Equation-based Distributionally Robust Offline Imitation Learning, a framework that learns robust policies solely from expert demonstrations collected under nominal dynamics, without requiring further environment interaction. We formulate the problem as a distributionally robust optimization over an uncertainty set of transition models, seeking a policy that minimizes the imitation loss under the worst-case transition distribution. Importantly, we show that this robust objective can be reformulated entirely in terms of the nominal data distribution, enabling tractable offline learning. Empirical evaluations on continuous-control benchmarks demonstrate that our approach achieves superior robustness and generalization compared to state-of-the-art offline IL baselines, particularly under perturbed or shifted environments.