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

TL;DR

This paper demonstrates that using the advantage function in reinforcement learning reduces policy confounding by breaking spurious correlations, thereby improving the agent's ability to generalize beyond its training trajectories.

Contribution

It reveals that the advantage function not only reduces gradient variance but also mitigates policy confounding, enhancing causal state representation learning.

Findings

Advantage function reduces spurious correlations in observations.

Training with the advantage function improves out-of-trajectory performance.

Analytical and empirical evidence supports the benefits of using the advantage function.

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.