Invariant Causal Imitation Learning for Generalizable Policies
This addresses the challenge of deploying imitation learning policies in new, unseen environments, which is crucial for real-world applications like robotics and healthcare, though it appears incremental as it builds on existing causal and invariant learning ideas.
The paper tackles the problem of learning imitation policies that generalize to unseen environments by leveraging data from multiple settings to avoid spurious correlations. It proposes Invariant Causal Imitation Learning (ICIL), which learns an invariant feature representation and matches expert behavior, showing effectiveness in control and healthcare tasks.
Consider learning an imitation policy on the basis of demonstrated behavior from multiple environments, with an eye towards deployment in an unseen environment. Since the observable features from each setting may be different, directly learning individual policies as mappings from features to actions is prone to spurious correlations -- and may not generalize well. However, the expert's policy is often a function of a shared latent structure underlying those observable features that is invariant across settings. By leveraging data from multiple environments, we propose Invariant Causal Imitation Learning (ICIL), a novel technique in which we learn a feature representation that is invariant across domains, on the basis of which we learn an imitation policy that matches expert behavior. To cope with transition dynamics mismatch, ICIL learns a shared representation of causal features (for all training environments), that is disentangled from the specific representations of noise variables (for each of those environments). Moreover, to ensure that the learned policy matches the observation distribution of the expert's policy, ICIL estimates the energy of the expert's observations and uses a regularization term that minimizes the imitator policy's next state energy. Experimentally, we compare our methods against several benchmarks in control and healthcare tasks and show its effectiveness in learning imitation policies capable of generalizing to unseen environments.