Breaking Habits: On The Role Of The Advantage Function In Learning Causal State Representations

Abstract

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.

Stats

• citations0
• S2 citations—
• github stars0
• HF likes0
• heat score0.00
• arxiv keysuau2025breaking

Related papers

• Bad Habits: Policy Confounding And Out-of-trajectory Generalization In RL (2023)0.00
• Direct Advantage Estimation (2021)0.00
• Causal Confusion And Reward Misidentification In Preference-based Reward Learning (2022)0.00
• Learning Nonlinear Causal Reductions To Explain Reinforcement Learning Policies (2025)0.00
• ACE : Off-policy Actor-critic With Causality-aware Entropy Regularization (2024)0.00
• Tackling Non-stationarity In Reinforcement Learning Via Causal-origin Representation (2023)0.00
• Learning By Doing: An Online Causal Reinforcement Learning Framework With Causal-aware Policy (2024)1.56
• From Kicking To Causality: Simulating Infant Agency Detection With A Robust Intrinsic Reward (2025)0.00