# Learning Causally Invariant Reward Functions from Diverse Demonstrations

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in Inverse reinforcement learning, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 29 of its 52 references that are in the map.
- First posted on arXiv 12 Sep 2024.

## Placed in

- Offline RL and imitation learning from data / [Inverse reinforcement learning](https://mlsanity.com/offlinerl/inverse-reinforcement-learning/)
- Deep reinforcement learning methods / [RL for robot locomotion and manipulation](https://mlsanity.com/deeprl/rl-for-robot-locomotion-and-manipulation/)

## Builds on (most cited first, 29 of 29)

- [Generative Adversarial Networks](https://mlsanity.com/paper/2203.00667/) (2022)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Improved Training of Wasserstein GANs](https://mlsanity.com/paper/1704.00028/) (2017)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [Generative Adversarial Imitation Learning](https://mlsanity.com/paper/1606.03476/) (2016)
- [Causal inference using invariant prediction: identification and confidence intervals](https://mlsanity.com/paper/1501.01332/) (2015)
- [Learning Robust Rewards with Adversarial Inverse Reinforcement Learning](https://mlsanity.com/paper/1710.11248/) (2017)
- [Reinforcement Learning and Control as Probabilistic Inference: Tutorial and Review](https://mlsanity.com/paper/1805.00909/) (2018)
- [Guided Cost Learning: Deep Inverse Optimal Control via Policy Optimization](https://mlsanity.com/paper/1603.00448/) (2016)
- [Causal Confusion in Imitation Learning](https://mlsanity.com/paper/1905.11979/) (2019)
- [Invariance Principle Meets Information Bottleneck for Out-of-Distribution Generalization](https://mlsanity.com/paper/2106.06607/) (2021)
- [Invariant Risk Minimization Games](https://mlsanity.com/paper/2002.04692/) (2020)
- [Spectral Norm Regularization for Improving the Generalizability of Deep Learning](https://mlsanity.com/paper/1705.10941/) (2017)
- [Maximum Entropy Deep Inverse Reinforcement Learning](https://mlsanity.com/paper/1507.04888/) (2015)
- [A Connection between Generative Adversarial Networks, Inverse Reinforcement Learning, and Energy-Based Models](https://mlsanity.com/paper/1611.03852/) (2016)
- [Invariant Causal Prediction for Nonlinear Models](https://mlsanity.com/paper/1706.08576/) (2017)
- [Anchor regression: heterogeneous data meets causality](https://mlsanity.com/paper/1801.06229/) (2018)
- [Invariant Causal Prediction for Block MDPs](https://mlsanity.com/paper/2003.06016/) (2020)
- [Observational Overfitting in Reinforcement Learning](https://mlsanity.com/paper/1912.02975/) (2019)
- [Of Moments and Matching: A Game-Theoretic Framework for Closing the Imitation Gap](https://mlsanity.com/paper/2103.03236/) (2021)
- [Off-Policy Imitation Learning from Observations](https://mlsanity.com/paper/2102.13185/) (2021)
- [f-IRL: Inverse Reinforcement Learning via State Marginal Matching](https://mlsanity.com/paper/2011.04709/) (2020)
- [Task-Relevant Adversarial Imitation Learning](https://mlsanity.com/paper/1910.01077/) (2019)
- [Invariant Policy Optimization: Towards Stronger Generalization in Reinforcement Learning](https://mlsanity.com/paper/2006.01096/) (2020)
- [Robust Inverse Reinforcement Learning under Transition Dynamics Mismatch](https://mlsanity.com/paper/2007.01174/) (2020)
- [Causal Imitation Learning under Temporally Correlated Noise](https://mlsanity.com/paper/2202.01312/) (2022)
- [Kernel Exponential Family Estimation via Doubly Dual Embedding](https://mlsanity.com/paper/1811.02228/) (2018)
- [Invariant Causal Imitation Learning for Generalizable Policies](https://mlsanity.com/paper/2311.01489/) (2023)
- [Inverse Reinforcement Learning via Matching of Optimality Profiles](https://mlsanity.com/paper/2011.09264/) (2020)

## Built on by (most cited first, 2 of 2)

- [Factored Causal Representation Learning for Robust Reward Modeling in RLHF](https://mlsanity.com/paper/2601.21350/) (2026)
- [Hallucination-Resistant Relation Extraction via Dependency-Aware Sentence Simplification and Two-tiered Hierarchical Refinement](https://mlsanity.com/paper/2508.14391/) (2025)

## Abstract (arXiv)

Inverse reinforcement learning methods aim to retrieve the reward function of a Markov decision process based on a dataset of expert demonstrations. The commonplace scarcity and heterogeneous sources of such demonstrations can lead to the absorption of spurious correlations in the data by the learned reward function. Consequently, this adaptation often exhibits behavioural overfitting to the expert data set when a policy is trained on the obtained reward function under distribution shift of the environment dynamics. In this work, we explore a novel regularization approach for inverse reinforcement learning methods based on the causal invariance principle with the goal of improved reward function generalization. By applying this regularization to both exact and approximate formulations of the learning task, we demonstrate superior policy performance when trained using the recovered reward functions in a transfer setting