# Breaking Habits: On the Role of the Advantage Function in Learning Causal State Representations

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Deep RL algorithms and representations, one of 15 directions in Deep reinforcement learning methods.
- Builds on 19 of its 33 references that are in the map.
- First posted on arXiv 13 Jun 2025.

## Placed in

- Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/)

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

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Causal Confusion in Imitation Learning](https://mlsanity.com/paper/1905.11979/) (2019)
- [Revisiting the Arcade Learning Environment: Evaluation Protocols and Open Problems for General Agents](https://mlsanity.com/paper/1709.06009/) (2017)
- [An Introduction to Deep Reinforcement Learning](https://mlsanity.com/paper/1811.12560/) (2018)
- [A Survey of Zero-shot Generalisation in Deep Reinforcement Learning](https://mlsanity.com/paper/2111.09794/) (2021)
- [DARLA: Improving Zero-Shot Transfer in Reinforcement Learning](https://mlsanity.com/paper/1707.08475/) (2017)
- [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)
- [Decoupling Value and Policy for Generalization in Reinforcement Learning](https://mlsanity.com/paper/2102.10330/) (2021)
- [Causal Imitation Learning with Unobserved Confounders](https://mlsanity.com/paper/2208.06267/) (2022)
- [Beyond variance reduction: Understanding the true impact of baselines on policy optimization](https://mlsanity.com/paper/2008.13773/) (2020)
- [The Role of Baselines in Policy Gradient Optimization](https://mlsanity.com/paper/2301.06276/) (2023)
- [Seeing is not Believing: Robust Reinforcement Learning against Spurious Correlation](https://mlsanity.com/paper/2307.07907/) (2023)
- [On Covariate Shift of Latent Confounders in Imitation and Reinforcement Learning](https://mlsanity.com/paper/2110.06539/) (2021)
- [Direct Advantage Estimation](https://mlsanity.com/paper/2109.06093/) (2021)
- [Bad Habits: Policy Confounding and Out-of-Trajectory Generalization in RL](https://mlsanity.com/paper/2306.02419/) (2023)
- [Can Agents Run Relay Race with Strangers? Generalization of RL to Out-of-Distribution Trajectories](https://mlsanity.com/paper/2304.13424/) (2023)
- [Skill or Luck? Return Decomposition via Advantage Functions](https://mlsanity.com/paper/2402.12874/) (2024)
- [What model does MuZero learn?](https://mlsanity.com/paper/2306.00840/) (2023)

## Abstract (arXiv)

Recent work has shown that reinforcement learning agents can develop policies that exploit spurious correlations between rewards and observations. This phenomenon, known as policy confounding, arises because the agent's policy influences both past and future observation variables, creating a feedback loop that can hinder the agent's ability to generalize beyond its usual trajectories. In this paper, we show that the advantage function, commonly used in policy gradient methods, not only reduces the variance of gradient estimates but also mitigates the effects of policy confounding. By adjusting action values relative to the state representation, the advantage function downweights state-action pairs that are more likely under the current policy, breaking spurious correlations and encouraging the agent to focus on causal factors. We provide both analytical and empirical evidence demonstrating that training with the advantage function leads to improved out-of-trajectory performance.