How to Mitigate the Distribution Shift Problem in Robotics Control: A Robust and Adaptive Approach Based on Offline to Online Imitation Learning

Technical Summary: Robust and Adaptive Imitation Learning (RAIL)

Problem Statement
Imitation Learning (IL) in robotics frequently suffers from the distribution shift problem, where an agent fails to execute appropriate actions when encountering states during online inference that were not present in the offline training dataset. This issue arises because expert demonstrations typically cover only a narrow subset of the environment's state-action space. Existing solutions often focus exclusively on either the offline training phase (via dataset augmentation) or the online inference phase (via adaptation), but both approaches exhibit limitations. Offline methods relying on human perturbation are costly, while those using world models are vulnerable to modeling errors. Online methods often struggle to distinguish between novel states caused by distribution shift and those that are simply part of the normal inference process, leading to unstable policy updates or failure to adapt.

Methodology
The authors propose RAIL (Robust Offline to Adaptive Online Imitation Learning), a lifelong, multi-phase framework designed to mitigate distribution shift through two distinct stages:

1. Offline Training Phase (Robustness):

   • Data Utilization: The framework leverages not only expert demonstrations ($D_E$) but also supplementary demonstrations ($D_S$). $D_S$ includes trajectories from novices, suboptimal experts, or automatically generated simulation data. These provide broader state-action coverage.
   • Discriminator with Regularization: To effectively utilize $D_S$, the authors train a discriminator to distinguish between expert and supplementary data. Recognizing that standard discriminators suffer from instability and bias when datasets are imbalanced, they introduce a novel regularization term ($L_{reg}$).
   • Mechanism: This term guides the discriminator to approximate the empirical state-action distributions of both datasets using Gaussian Mixture Models (GMM) and learned policy likelihoods. This stabilizes the discriminator's learning, particularly in early stages, allowing it to more accurately estimate the optimality of training samples.
   • Policy Training: The optimized discriminator output is used to compute weights for Weighted Behavior Cloning (BC), enabling the policy to learn robustly from mixed-quality data.

2. Online Inference Phase (Adaptation):

   • Self-Supervised Adaptation: The framework treats online experiences as supplementary demonstrations. It employs the same discriminator-based learning procedure used in the offline phase to compute self-supervised learning signals.
   • Distribution Shift Detection (UTM): To prevent degrading generalization by learning indiscriminately, the system employs Update Time Management (UTM). It calculates a likelihood metric ($\kappa(s)$) for the current state. Online learning is triggered only if $\kappa(s)$ falls below a threshold for a sustained period ($N_{ds}$), indicating a persistent distribution shift.
   • Adaptive Updates: When a shift is detected, the discriminator and policy are updated using the accumulated online experience, weighted by the degree of distribution shift ($\kappa(s)$). States far from the offline distribution receive more conservative updates.

Key Contributions

1. RAIL Framework: A unified lifelong scheme that addresses distribution shift comprehensively by combining robust offline training with adaptive online inference.
2. Regularized Discriminator: A novel regularization term for the discriminator that stabilizes training and improves the accuracy of optimality estimation for input data samples, particularly when expert and supplementary data are imbalanced.
3. Adaptive Online Learning: A method that computes self-supervised signals using the discriminator and restricts online updates to instances where distribution shift is explicitly detected, ensuring stable adaptation.

Experimental Results
The authors evaluated RAIL in MuJoCo environments (Hopper, Halfcheetah, Walker2d, Ant) under varying levels of Gaussian noise to induce distribution shift.

• Offline Performance: In environments with injected noise, RAIL consistently outperformed baselines (Standard BC, Stable BC, and Model-based BC). The results suggest that expanding state coverage via real supplementary demonstrations is more effective for robustness than policy optimization alone or model-based virtual rollouts.
• Online Performance: During online inference, RAIL demonstrated superior adaptation compared to BC, Beyond Imitation (BI), and GAIL. While BC degraded due to compounding errors and BI struggled with self-supervision signals, RAIL successfully adapted.
• Efficiency: The ablation study showed that removing the UTM technique (RAIL-UTM) resulted in inferior performance and significantly higher training time, highlighting the importance of selective updating. RAIL achieved approximately a 54% reduction in training time compared to RAIL-UTM while maintaining higher final performance.
• Stability: Weighted BC used in RAIL provided more stable policy learning during online adaptation compared to the TRPO-based updates used in GAIL.

Significance and Claims
The paper claims that RAIL offers a practical solution to the distribution shift problem by shifting the focus from solely optimizing policy parameters to expanding state-action coverage through supplementary data. By integrating a stabilized discriminator with a conditional online learning mechanism, the framework enables robots to remain robust against unseen states during training and adapt effectively when such states are encountered in the real world. The authors assert that this approach provides a more stable and efficient path to lifelong imitation learning compared to existing offline-only or online-only strategies.