Breaking Habits: On the Role of the Advantage Function in Learning Causal State Representations
Extracted source text · arXiv 2506.11912
\documentclass{article}
% if you need to pass options to natbib, use, e.g.:
% \PassOptionsToPackage{numbers, compress}{natbib}
% before loading neurips_2025
% ready for submission
% \usepackage{neurips_2025}
% to compile a preprint version, e.g., for submission to arXiv, add add the
% [preprint] option:
\usepackage[preprint]{neurips_2025}
% to compile a camera-ready version, add the [final] option, e.g.:
% \usepackage[final]{neurips_2025}
% to avoid loading the natbib package, add option nonatbib:
% \usepackage[nonatbib]{neurips_2025}
\usepackage[utf8]{inputenc} % allow utf-8 input
\usepackage[T1]{fontenc} % use 8-bit T1 fonts
\usepackage{hyperref} % hyperlinks
\usepackage{url} % simple URL typesetting
\usepackage{booktabs} % professional-quality tables
\usepackage{amsfonts} % blackboard math symbols
\usepackage{nicefrac} % compact symbols for 1/2, etc.
\usepackage{microtype} % microtypography
\usepackage{xcolor} % colors
\usepackage{amssymb} % Defines common symbols like \mathbb R
\usepackage{mathtools} % Extends amsmath, providing common math tools
\usepackage{mathrsfs} % Enables \mathscr, which can work in cases that \mathcal does not
%\mathtoolsset{showonlyrefs} % Only number equations that are referenced (optional)
\usepackage{graphicx} % For including images
\usepackage{subcaption} % Allows for the use of subfigures and subcaptions
\usepackage[space]{grffile} % For spaces in image names
\usepackage{url} % For displaying URLs
\usepackage{lipsum} % For placeholder text
\usepackage{xcolor}
\usepackage{colortbl}
\usepackage{thm-restate}
% Define color gradients for positive and negative values
\definecolor{darkblue}{RGB}{0,0,139}
\definecolor{lightblue}{RGB}{173,216,230}
\definecolor{darkred}{RGB}{139,0,0}
\definecolor{lightred}{RGB}{255,204,204}
% Command for positive value coloring (blue scale)
\newcommand{\poscolor}[1]{%
\cellcolor{blue!\the\numexpr100*#1\relax!white} #1%
}
% Command for negative value coloring (red scale)
\newcommand{\negcolor}[1]{%
\cellcolor{red!\the\numexpr-100*#1\relax!white} #1%
}
% Function to map values in a nonlinear way to improve granularity
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% AUTHOR: Fill in the following meta-data
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Enter the title of your paper:
\usepackage{graphicx}
\usepackage{natbib}
\usepackage{mathtools}
\usepackage{amssymb} % Defines common symbols like \mathbb R
\usepackage{mathtools} % Extends amsmath, providing common math tools
\usepackage{mathrsfs} % Enables \mathscr, which can work in cases that \mathcal does not
\mathtoolsset{showonlyrefs} % Only number equations that are referenced (optional)
\usepackage{graphicx} % For including images
\usepackage{subcaption} % Allows for the use of subfigures and subcaptions
\usepackage[space]{grffile} % For spaces in image names
\usepackage{url} % For displaying urls
\usepackage{hyperref}
\usepackage{url}
\usepackage[utf8]{inputenc} % allow utf-8 input
\usepackage[T1]{fontenc} % use 8-bit T1 fonts
\usepackage{url} % simple URL typesetting
\usepackage{booktabs} % professional-quality tables
\usepackage{amsfonts} % blackboard math symbols
\usepackage{nicefrac} % compact symbols for 1/2, etc.
\usepackage{microtype} % microtypography
\usepackage{xcolor} % colors
\usepackage{graphicx}
\usepackage{float}
\usepackage{amssymb}
\usepackage{mathtools}
\usepackage{amsthm}
\usepackage{multirow}
\usepackage{subcaption}
\usepackage{wrapfig}
\usepackage{microtype}
\usepackage{graphicx}
% \usepackage{subfigure}
\usepackage{booktabs} % for professional tables
\usepackage{amsmath}
\newcommand{\ind}{\perp\!\!\!\!\perp}
\DeclareMathOperator*{\supp}{supp}
\DeclareMathOperator*{\dooperator}{do}
\DeclareMathOperator*{\domain}{dom}
% \DeclareMathOperator*{\argmin}{arg\,min}
% \DeclareMathOperator*{\argmax}{arg\,max}
\DeclareMathOperator*{\EX}{\mathbb{E}}
\def\Pr{\mathop{\rm Pr}\nolimits}
\theoremstyle{plain}
\newtheorem{theorem}{Theorem}
\newtheorem{proposition}{Proposition}
\newtheorem{lemma}{Lemma}
\newtheorem{corollary}{Corollary}
\theoremstyle{definition}
\newtheorem{definition}{Definition}
\newtheorem{assumption}{Assumption}
\newtheorem{example}{Example}
\theoremstyle{remark}
\newtheorem{remark}[theorem]{Remark}
\title{Breaking Habits: On the Role of the Advantage \\ Function in Learning Causal State Representations}
% The \author macro works with any number of authors. There are two commands
% used to separate the names and addresses of multiple authors: \And and \AND.
%
% Using \And between authors leaves it to LaTeX to determine where to break the
% lines. Using \AND forces a line break at that point. So, if LaTeX puts 3 of 4
% authors names on the first line, and the last on the second line, try using
% \AND instead of \And before the third author name.
\author{%
Miguel Suau \\
Phaidra \\
\texttt{miguel.suau@phaidra.ai} \\
% examples of more authors
% \And
% Coauthor \\
% Affiliation \\
% Address \\
% \texttt{email} \\
% \AND
% Coauthor \\
% Affiliation \\
% Address \\
% \texttt{email} \\
% \And
% Coauthor \\
% Affiliation \\
% Address \\
% \texttt{email} \\
% \And
% Coauthor \\
% Affiliation \\
% Address \\
% \texttt{email} \\
}
\begin{document}
\maketitle
\begin{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.
\end{abstract}
\section{Introduction}
Imagine a robot trained to perform two tasks: first, navigate from your office to the coffee machine, retrieve a cup of coffee, and return; second, go to the printer room, make copies, and return. There are two possible routes to the coffee machine: one through the printer room and another through a direct corridor. Since the corridor route is shorter, the robot typically avoids the printer room when fetching coffee. However, one day, the corridor is blocked and the robot takes the path through the printer room, unexpectedly returning with a copy of a paper titled \emph{Breaking Habits}.
Why did this happen? After repeatedly performing these tasks, the robot incorrectly associated the printer room with the need to make copies, turning this misassociation into a habit. While this habit works fine under normal conditions, it fails when the robot is forced to deviate from its usual path. This failure mode, referred to as \emph{out-of-trajectory generalization}, was explored by \citet{suau2024bad} in the context of reinforcement learning (RL). The authors showed that such issues arise because the agent’s policy introduces spurious correlations \citep{pearl2016causal} between rewards and observations, a phenomenon they termed \emph{policy confounding}.
\paragraph{Contributions}
In this paper, we observe that the advantage function, commonly used in many policy gradient methods, not only reduces the variance of gradient updates but also plays a crucial role in mitigating this issue. We demonstrate, both analytically and experimentally, that using the advantage function encourages the agent to learn representations that rely more on the true causal factors. It achieves this by scaling state-action pairs according to their probability under a given policy and state representation, effectively breaking the spurious correlations introduced by the agent’s policy. This encourages the agents to focus on causal factors, enabling better out-of-trajectory generalization.
To support the theoretical findings, experiments are conducted in three simple environments. The results indicate that agents trained with $Q$-values fail to generalize beyond their usual trajectories, whereas agents trained with the advantage function exhibit robustness to trajectory deviations. Furthermore, an analysis of the learned state representations reveals that the latter focus on causal factors rather than exploiting spurious correlations.
Finally, we examine the impact of implementation choices, such as batch size and advantage normalization. Through experiments, we show that these factors can significantly affect the learned state representations.
% At this point, the agent begins to exploit the spurious correlations introduced by its policy. This suggests that stopping training early, before the learning content diminishes entirely, can prevent the agent from developing these maladaptive habits. This approach is conceptually similar to early stopping in supervised learning CITATION.
% A straightforward but impractical way to implement early stopping is to periodically evaluate the agent in a test environment and halt training when performance degrades. However, this method poses two significant challenges. First, without a simulator, designing a reliable test environment is difficult and potentially unsafe. Second, even with a simulator, it is impossible to ensure that the test environment is exhaustive enough to detect all possible spurious correlations the policy might introduce.
% Fortunately, there is a more natural and effective alternative: using a value-based advantage function. As we show in this paper, beyond its well-established role in reducing the variance of gradient updates, value-based advantages also serve as a means of quantifying the learning content of a sample relative to the agent's current policy. By leveraging this property, we can inhibit learning when the learning content decreases, effectively addressing the policy confounding issue without relying on external test environments.
% In order to ensure that value-based advantage functions are effective for this second objective we need to force our value estimates to accurately estimate the true state-value under the current policy rather than the state value under the mixture of the most recent policies as it is the case in most policy gradient implementations. This reduces to reinitializing and retraining the value network (or value head) every time we perform a policy update.
% \section{Stuff}
% \color{red}
% SUTTON AND BARTO:
% The former makes
% sense because it causes the parameter to move most in the directions that favor actions
% that yield the highest return. The latter makes sense because otherwise actions that are
% selected frequently are at an advantage (the updates will be more often in their direction)
% and might win out even if they do not yield the highest return.
% \color{black}
% \color{blue}
% MARLOS: BEYOND VARIANCE REDUCTION ICML 2021
% This research paper investigates the impact of baselines on policy optimization in reinforcement learning, challenging the conventional understanding that their primary role is variance reduction. The authors demonstrate that baselines influence learning dynamics beyond variance, leading to "committal" or "non-committal" behavior depending on their value, even with identical variance.
% \color{black}
% \color{blue}
% Csaba: The Role of Baselines in Policy Gradient Optimization
% The key role of the value baseline is to reduce update aggressiveness to achieve
% a particular effect on long-term sampling, rather than simply reduce variance
% In stochastic policy optimization, the two most common techniques for improving the basic algorithm
% are to include on-policy importance sampling (IS) and subtract a baseline. Including on-policy
% IS provides unbiased gradient estimates, but introduces high variance when an action’s sampling
% probability is close to 0.
% These findings show that a key role of the value baseline is to automatically
% ensure “sufficient exploration” during on-policy stochastic optimization. Next, we provide a detailed
% understanding of how baselines modulate the circular interaction between stochastic action sampling
% and updating. Although a baseline has no effect on exact gradients, it can play a major role in
% stochastic gradients.
% Too aggressive update vs. good enough updates.
% "Thus, the
% effect of baseline is to modify the sampling to lie exactly on the boundary of being good enough.
% From this argument the key role of the value baseline is to reduce update aggressiveness to achieve
% a particular effect on long-term sampling, rather than simply reduce variance."
% Can we show that the estimate of the policy gradient that uses the advantage is always lower than the true gradient.
% \color{black}
% ADVANTAGE FUNCTION HIGHLIGHTS UNSUAL EVENTS WHILE DOWNWEIGHTS EVENTS THAT ARE MORE COMMON
% NO CORRELATION BETWEEN POLICY AND ADVANTAGE
% ADVANTAGE SHOWS THE DIFFERENCE IN VALUE BETWEEN THE CURRENT POLICY AND A POLICY THAT IS EXACTLY THE SAME BUT WHICH SELECTS ACTION A AT STATE S WITH PROBABILITY 1.
% AS THE POLICY STARTS EXPLOITING CERTAIN CORRELATIONS THE DECREASES. L=6 IS CORRELATED WITH K=1 EVEN WITH RANDOM POLICY, AS THE POLICY STARTS EXPLOITING THIS CORRELATION THE EXPECTED ADVANTAGE DECREASES WHILE THE EXPECTED VALUE CONTINUES TO INCREASE
% EXPLORE
% \color{blue}
% Mathematical Formalization
% To formalize this, one could define the degree of policy confounding as a measure of how much the expected reward or state transitions deviate when intervening on the state representation (do-operator) versus when not intervening.◦
% Then, one would need to show that the updates based on the advantage function lead to a decrease in this measure of policy confounding, as compared to updates based on the Q-value alone. This would be done by explicitly analyzing the update rule of policy gradient methods in conjunction with the advantage function.
% % The update rule for the policy gradient using the advantage is given by ∇θJ(πθ) = Est,at∼πθ [∇θ log πθ(at | st)Aπ(st, at)]. This update encourages the agent to take actions with high advantage more often, and, according to the previous discussion, the advantage function assigns higher magnitudes to less probable actions under the current policy. This would need to be formally expressed in the proof by deriving an expression for the gradient of the policy when the advantage function is used in the policy gradient method.
% A key step in the proof would be to show how the advantage function alters the update direction, so that the updated policy mitigates the spurious correlations that cause policy confounding. Specifically, the proof would need to show that the update direction using the advantage function leads to a policy that causes the do-operator in the policy confounding definition to approach the normal (non-interventional) case, thereby decreasing policy confounding.
% 7.
% Limitations of the Proof:
% It is important to recognize that this proof would likely be under certain assumptions about the environment, the policy class, and the learning algorithm. This means that the proof would be for a limited set of conditions.
% The effectiveness of the advantage function would have to be demonstrated in scenarios involving specific types of MDPs (Markov Decision Processes), and might not hold in non-Markovian situations.
% It should also be recognized that while the advantage function can break spurious correlations, it does not guarantee the discovery of the absolute best causal representation.
% Conclusion:
% While the sources do not contain a formal proof, the information provided strongly suggests that the advantage function can mitigate policy confounding. The core idea is that by downweighting actions that are frequently taken under the current policy, the advantage function reduces the impact of spurious correlations and encourages the agent to focus on causal features. A formal proof would need to mathematically demonstrate that the advantage function reduces the measure of policy confounding and would also need to provide conditions that constrain the situations under which the proof would hold.
% \color{black}
% \color{red}
% This might involve showing that the advantage function leads to smaller updates for commonly followed trajectories, whereas Q-values might overemphasize those same trajectories, thus being more susceptible to picking up spurious correlations
% Q-value Updates: When using the Q-value directly in policy gradient updates, the magnitude of the update is determined by Q(s, a). There is no adjustment for how likely the action a is under the current policy. If a policy often takes an action a in state s, the update would be based on Q(s, a), and the update will overemphasize this action because it is not balanced with respect to the action probability. This can cause the agent to become overly reliant on the specific trajectories it has seen and can make the agent susceptible to picking up spurious correlations between the observations and rewards in the specific trajectories it commonly follows.
% Focus on Causal Features: By reducing the impact of common trajectories, the agent is incentivized to explore less common trajectories and, potentially, discover the true causal relationships between observations and rewards. The advantage function encourages the agent to focus on causal features by giving more weight to less probable actions.
% In summary, the advantage function scales the update by a term $1 - \pi(a_t | s_t)$ that is smaller for commonly followed actions, whereas Q-values do not incorporate this scaling. This means that the advantage function leads to smaller updates for common trajectories compared to Q-values. The advantage function thereby mitigates policy confounding by reducing the influence of frequently visited trajectories and their spurious correlations while focusing more on less common trajectories that can help the agent find the true causal relations.
% \color{black}
% Note that using the advantage function rather than the $Q$-value does not change the expected gradient. This implies that both should have similar values if the batch size is large enough. However, as the variance of the advantage version is much smaller, the number of samples needed is much smaller. Importantly, when using the advantage, the advantage value of every sample is uncorrelated with the state value, whereas when using the $Q$-value, this is only true in expectation. Hence, unless the batch size is large enough, the policy would be susceptible of picking up on spurious correlations. Note that the number of samples needed increases as the policy entropy decreases.
% In the extreme, when the policy is deterministic the advantage is 0 but the Q value is not. To compensate for a deterministic policy you would need infinitely many samples.
\section{Example: Key2Door}\label{sec:example}
We will use the following example throughout the paper to illustrate the ideas presented.
Figure \ref{fig:key2door} depicts a gridworld environment. The agent's objective is to collect a key placed at the beginning of the corridor and then open the door at the end. The agent's observation consists of its current location $L \in \{0, 1, .., 6\}$ and a binary variable $X \in \{0, 1\}$ indicating whether or not it has the key. At all states, the agent has two possible actions: moving left $A=0$, or moving right $A=1$. The agent receives a reward of $+1$ if it reaches the door after collecting the key, and $0$ if it reaches the door without the key. Additionally, there is a $-0.01$ penalty per timestep to incentivize the agent to take the most direct path.
\begin{wrapfigure}{r}{0.6\textwidth}
\begin{center}
\vspace{-15pt}
\includegraphics[width=0.6\textwidth]{figures/key2door.pdf}
\end{center}
\vspace{-10pt}
\caption{Key2Door environment.}
\vspace{-5pt}
\label{fig:key2door}
\end{wrapfigure}
During training, the agent always starts the episode in location 2. This implies that once the agent has collected the key and is moving toward the door, it can disregard the key once it reaches location 3. The agent can learn to use its location to infer whether it has the key. However, this only works well if the agent consistently follows the optimal policy, as depicted by the purple arrow. If the agent moves randomly through the corridor, it cannot reliably determine whether the key has been collected. Moreover, an agent that relies solely on its location will fail when evaluated in the same task but starting from location 6. The agent will attempt to move right to open the door, even though it has not collected the key.
\section{Related work}
The term policy confounding was introduced by \citet{suau2024bad} to describe a phenomenon in which policies, by influencing both past and future observation variables, induce spurious correlations. These correlations can lead the agent to develop shortcuts, referred to as habits, that are effective only within the trajectories the agent typically follows but fail to generalize out-of-trajectory. Out-of-trajectory (OOT) generalization is a specific instance of the broader problem of out-of-distribution (OOD) generalization \citep{kirk2023survey}. Unlike OOD generalization, which focuses on adapting to environments with different rewards \citep{Taylor09ICML}, observations \citep{mandlekar2017adversarially,zhang2020invariant}, or transition dynamics \citep{higgins2017darla}, OOT generalization seeks to enable agents to generalize across alternative trajectories within the same environment. Empirical evidence of policy confounding has been reported in several studies \citep{machado2018revisiting, Song2020Observational, lan2023can, he2024model, weltevrede2024training}. Similarly, prior work has examined confounding in the context of imitation learning \citep{de2019causal, zhang2020causal, tennenholtz2021covariate, ding2023seeing}. However, unlike policy confounding, these works focus on cases where the agent does not actively contribute to the formation of spurious correlations.
The connection between the advantage function and causality has been explored in prior work. \citet{corcoll2022disentangling} use the advantage function to differentiate between agent-driven effects and external environmental factors, constructing a hierarchy of transformations that the agent can perform on its environment. \citet{pan2022direct} argue that the advantage function can be interpreted as a measure of an action’s causal effect on return and introduce a method for learning advantages directly from on-policy data without relying on a value function. This approach was later extended to support off-policy data \citep{pan2024skill}. While not explicitly framing the advantage function in causal terms, \citet{raileanu2021decoupling} describe it as a measure of the expected additional return gained by selecting a particular action over following the current policy. Their experiments suggest that the advantage function is less prone to overfitting to certain environment-specific idiosyncrasies.
Our work is strongly influenced by \citet{chung2021beyond}, who showed that the choice of baseline in policy gradient updates can sometimes lead to overly aggressive updates, causing what the authors term committal behavior, which may result in suboptimal policies. \citet{mei2022role} further analyzed the theoretical properties of state-value baselines (i.e., advantages), demonstrating that they moderate update aggressiveness and ensure convergence to the optimal policy. As we discuss in the following sections, these insights are closely connected to the findings presented in this work.
% Seeing is not Believing:
% LEARNING CAUSAL STATE REPRESENTATIONS: \citep{zhang2020invariant}
% Passive: Passive methods try to uncover the causal structure of a system directly from observational data (Mooij et al., 2016). In this setting, there is normally data from a few different distributions (resulting from different interventions on the system), and the task is to learn a model that generalizes to a set of test distributions. Perhaps the most popular methods in this category are those that exploit prediction invariance (Peters et al., 2016; Arjovsky et al., 2019; Mooij et al., 2020; Krueger et al., 2021). These are based on the observation that causal models should be invariant to interventions on non-causal variables. Hence, they force the model to be invariant across the training distributions with the hope that this would be sufficient to find the causal predictors. It turns out that this is also the case in RL; causal state representations should be invariant across policies. Translating this observation into a practical algorithm implies that from the agents’ viewpoint, everything would stay the same. That is, agents would explore the environment with the sole intention of maximizing reward. However, the learning objectives would be different. Constraints or penalties would need to be added to the objective to make sure that the learned representations are not biased toward a particular policy. Invariance-based methods have already been applied in the RL context (Zhang et al., 2020). Nonetheless, the objective for Zhang et al. is again generalization to alternative environments. Hence, the set of training distributions is effectively a set of environments that have been intervened differently. In contrast, in our case, to mitigate the effects of policy confounding, the different training distributions are induced by the different policies the agent follows throughout the course of training.
% Active: In this setting, we are allowed to intervene in the system to try to figure out which variables are causal predictors and which others are just statistically associated. There are different types of interventions (Peters et al., 2017): we can, for instance, fix the value of a particular variable (hard intervention) or modify how a variable affects another by adding some noise (soft intervention). The RL setting offers a natural way to intervene in the environment through the agent’s actions. As the host argues in a recent paper (Schölkopf et al., 2021), “(in RL), agents can reflect on their decisions and formulate hypotheses that can be empirically verified”. Hence, CRL can be added as one of the agent’s exploration objectives. This is in contrast to the standard RL formulation, where agents simply explore to maximize the return. Exploration bonuses directed at CRL have already been explored in the literature for sample efficiency in robot manipulation tasks (Seitzer et al., 2021) and zero-shot generalization to other environments (Sontakke et al., 2021). Our goal, again, would be to guarantee that the representations learned by the agent are useful when following alternative trajectories. Hence, during training, the agent will be encouraged to move away from its preferred trajectories such that the spurious correlations introduced by policy confounding can be erased.
\section{Preliminaries}
% \subsection{Problem Formulation}
We focus on problems where states are represented by a set of observation variables, or factors \citep{boutilier1999decision}. This representation is common in tasks that involve modeling policies and value functions using function approximators \citep{franccois2018introduction}. These observation variables typically describe features of the agent's state in the environment.
\begin{definition}[MDP] A Markov decision process (MDP) is a tuple $\langle \mathcal{S}, \mathcal{A}, T, R \rangle$, where $\mathcal{S}$ represents the set of states, $\mathcal{A}$ denotes the set of actions available to the agent, $T : \mathcal{S} \times \mathcal{A} \to \Delta (\mathcal{S})$ is the transition function, and $R : \mathcal{S} \times \mathcal{A} \to \mathbb{R}$ is the reward function. \end{definition}
\begin{definition}[FMDP] A Factored Markov decision process (FMDP) is an MDP where the set of states is defined by a set of observation variables, or factors, $F = {F^1, \dots, F^{|F|}}$. Each variable $F^i$ can take any value from its domain $\mathcal{F}^i$. Consequently, each state $s$ corresponds to a unique combination of values for the variables in $F$, i.e., $s = \langle f^1, \dots, f^{|F|} \rangle \in \times \mathcal{F} = \mathcal{S}$. \end{definition}
\subsection{State representations}
The agent's objective is to learn a policy $\pi: \mathcal{S} \to \Delta (\mathcal{A})$ that maximizes the expected discounted sum of rewards \citep{sutton2018reinforcement}.
However, learning a policy that conditions on every observation variable might be impractical, especially in scenarios with a large number of variables. Fortunately, in many problems, not all variables are essential, and compact state representations can be found that are sufficient for solving the task at hand \citep{McCallum95PhD}. This is where function approximators, such as neural networks, come into play. When used to model policies and value functions, they learn to ignore certain observation variables in $F$ if those variables are deemed unnecessary for estimating values and optimal actions.
\begin{definition}[State representation] A state representation is a function $\Phi: \mathcal{S} \to \bar{\mathcal{S}}$, where $\mathcal{S} = \times \mathcal{F}$, $\bar{\mathcal{S}} = \times \bar{\mathcal{F}}$, and $\bar{F} \subseteq F$. \label{def:state_representation} \end{definition}
Intuitively, a state representation $\Phi(s_t)$ is a projection of a state $s \in \mathcal{S} = \times \mathcal{F}$ onto a lower-dimensional space $\bar{S} = \times \bar{\mathcal{F}}$, defined by a subset of its variables, $\bar{F} \subseteq F$. We use $\{s\}^\Phi = \{s' \in \mathcal{S}: \Phi(s') = \Phi(s)\}$ to denote the equivalence class of $s$ under $\Phi$.
In the Key2Door example a potential state representation could be $\Phi(s_t) = \langle l_t \rangle$ for all $s_t \in \mathcal{S}$. This representation retains only the agent's current location, ignoring all other variables in $F$. Therefore, all states sharing the same value for $L_t$ belong to the same equivalence class.
Not all state representations are sufficient to learn the optimal policy. Some representations may exclude variables that carry valuable information for solving the task. In the Key2Door example, knowing whether the key has been collected is crucial to selecting the next action.
\begin{definition}[Markov state representation] A state representation $\Phi(s_t)$ is said to be Markov if, for all $s_t, s_{t+1} \in \mathcal{S}$ and $a_t \in \mathcal{A}$,
\begin{equation*} R(s_t , a_t) = R(\Phi(s_t), a_t) \quad \text{and} \quad \sum_{s'_{t+1} \in \{s_{t+1}\}^\Phi} T(s'_{t+1} \mid s_t, a_t) = P(\Phi(s_{t+1}) \mid \Phi(s_t), a_t), \end{equation*}
\normalsize
where $R(\Phi(s_t), a_t)$ denotes the reward $R(s'_t, a_t)$ for any $s'_t \in \{s_t\}^\Phi$. \label{def:markov_state} \end{definition}
% \begin{definition}[$\pi$-Markov state representation] A state representation $\Phi^\pi$ is said to be $\pi$-Markov if, for all $s_t, s_{t+1} \in \mathcal{S}^\pi$, $a_t \in \supp(\pi( \cdot \mid s_t))$,
% \begin{equation*}
% R(s_t , a_t) = R^\pi(\Phi^\pi(s_t), a_t)
% \quad \text{and} \quad \sum_{s'_{t+1} \in \{s_{t+1}\}^\Phi_\pi} T(s'_{t+1} \mid s_t, a_t) = P^\pi(\Phi^\pi(s_{t+1}) \mid \Phi^\pi(s_t), a_t),
% % R_h(h_t, a_t) = \EX_\pi[r_t \mid \Phi^\pi(h_t), a_t]
% \end{equation*}
% where $S^\pi \subseteq S$ denotes the set of states visited under $\pi$, $R^\pi(\Phi^\pi(s_t), a_t)$ is the reward $R(s'_t, a_t)$ at any $s'_t \in \{s_t\}^\Phi_\pi$, with $\{s\}^\Phi_\pi =\{s' \in S^\pi: \Phi^\pi(s') = \Phi^\pi(s)\} $, and $P^\pi$ is probability under $\pi$.
% \end{definition}
\subsection{Policy confounding}
The phenomenon of policy confounding plays a critical role in the search for simpler state representations. The policy can mislead the function approximator into forming simpler state representations that rely on spurious correlations, such as having the key when being at location 6, rather than true causal factors.
\begin{definition}[Policy Confounding] A state representation $\Phi : \mathcal{S} \to \bar{\mathcal{S}}$ is said to be confounded by a policy $\pi$ if, for some $s_t, s_{t+1} \in \mathcal{S}$ and $a_t \in \mathcal{A}$,
\begin{equation*}
R^\pi(\Phi(s_t), a_t) \neq R^\pi(\dooperator(\Phi(s_t)), a_t)
\end{equation*}
or
\begin{equation*}
P^\pi(\Phi(s_{t+1}) \mid \Phi(s_t), a_t) \neq P^\pi(\Phi(s_{t+1}) \mid \dooperator(\Phi(s_t)), a_t).
\end{equation*}
\normalsize
where $R^\pi(\Phi^\pi(s_t), a_t)$ is the reward $R(s'_t, a_t)$ at any $s'_t \in \{s_t\}^\Phi_\pi$, with $\{s\}^\Phi_\pi =\{s' \in S^\pi: \Phi^\pi(s') = \Phi^\pi(s)\} $, and $P^\pi$ is probability under $\pi$.
\end{definition}
The operator $\dooperator(\cdot)$ is known as the do-operator, and it is used to represent physical interventions in a system \citep{pearl2016causal}. These interventions are meant to distinguish cause-effect relations from mere statistical associations. In our case, $\dooperator(\Phi(s_t))$ means setting the variables forming the state representation $\Phi(s_t)$ to a particular value and considering all possible states in the equivalence class, $s'_t \in \{s_t\}^\Phi$, (i.e., all states that share the same value for the observation variables that are selected by $\Phi$; independently of whether these are visited by the policy being followed). For instance, in the Key2Door example, $R^{\pi^*}(L = 6) = +1$ when following the optimal policy $\pi^*$ (purple path in Figure \ref{fig:key2door}) since we know that the agent has the key, while $R^{\pi^*}(\dooperator(L = 6)) = \pm 1$ since the agent may or may not have the key.
Strictly speaking, the only representation that satisfies the above conditions is the Markov state representation. This representation is invariant to the agent's policy, making it robust to any deviations in the trajectory. In the following, we refer to Markov state representations as \textit{Causal State Representations}, as they necessarily include all the causal factors governing rewards and transitions.
% \section{Preliminaries}
% \begin{definition}[MDP]
% A Markov decision process (MDP) is a tuple $\langle \mathcal{S}, \mathcal{A}, T, R \rangle$, where $\mathcal{S}$ represents the set of states, $\mathcal{A}$ denotes the set of actions available to the agent, $T : \mathcal{S} \times \mathcal{A} \to \Delta (\mathcal{S})$ is the transitions function, and $R : \mathcal{S} \times \mathcal{A} \to \mathbb{R}$ is the reward function.
% \end{definition}
% In particular, we focus on problems where states are represented by a set of observation variables, or factors \citep{boutilier1999decision}. This representation is common when modeling policies and value functions using function approximators \citep{sutton2018reinforcement}. These observation variables typically describe features of the agent's state in the environment.
% \begin{definition}[FMDP]
% A Factored Markov decision process (FMDP) is an MDP where
% the set of states is defined by a set of observation variables, or factors, $F = \{F^1, \dots, F^{|F|}\}$. Each variable $F^i$ can take any of the values in its domain $\mathcal{F}^i$. Consequently, each state $s$ corresponds to a unique combination of values for the variables in $F$, $s = \langle f^1, \dots, f^{|F|} \rangle \in \times \mathcal{F} = \mathcal{S}$.
% \end{definition}
% The agent's objective is to learn a policy $\pi: \mathcal{S} \to \Delta (\mathcal{A})$ that maximizes the expected discounted sum of rewards \citep{sutton2018reinforcement}. However, learning a policy that conditions on every observation variable might be impractical, particularly in scenarios with a large number of variables. Fortunately, in many problems, not all variables are strictly essential, and compact state representations can be found that are sufficient for solving the task at hand \citep{McCallum95PhD}. This is where function approximators, such as neural networks, come into the picture \citep{franccois2018introduction}. If we use them to model policies and value functions, they will learn to ignore certain observation variables in $F$ if deemed unnecessary for estimating values and optimal actions.
% \begin{wrapfigure}{r}{0.5\textwidth}
% \begin{center}
% \vspace{-20pt}
% \includegraphics[width=0.48\textwidth]{figures/key2door.pdf}
% \end{center}
% \vspace{-5pt}
% \caption{An illustration of the Key2Door environment.}
% \vspace{-5pt}
% \label{fig:key2door}
% \end{wrapfigure}
% For ease of exposition we will make use of the below example to illustrate the phenomenon of policy confounding.
% \begin{example}{\textbf{Key2Door.}}
% \label{ex:key2door}
% Figure \ref{fig:key2door} shows a gridworld environment. The agent's goal is to collect a key placed at the beginning of the corridor in (left) and then open the door at the end. The agent's observation consists of its current location and a binary variable indicating whether or not it has got the key. During training, the agent always starts the episode at the second cell from the left. This means that when moving toward the door, the agent can forget about the key (i.e., ignore the binary variable) once it has reached the third cell. The agent can build the habit of using its own location to tell whether it has or has not got the key yet. This, can only occur when the agent consistently follows the optimal policy, depicted by the purple arrow. Otherwise, if the agent moves randomly through the corridor, it is impossible to tell whether the key has or has not been collected. In contrast, in the evaluation environment, the agent always starts at the second to last cell, this confuses the agent, which is used to already having the key by the time it reaches said cell.
% \end{example}
% \begin{definition}[State representation]
% A state representation is a function $\Phi: \mathcal{S} \to \bar{\mathcal{S}}$, where $\mathcal{S} = \times \mathcal{F}$, $\bar{\mathcal{S}} = \times \bar{\mathcal{F}}$, and $\bar{F} \subseteq F$.
% \label{def:state_representation}
% \end{definition}
% Intuitively, a state representation $\Phi(s_t)$ is a state-specific projection of a state $s \in \mathcal{S} = \times \mathcal{F}$ onto a lower-dimensional space $\bar{S} = \times \bar{\mathcal{F}}$ defined by a subset of its variables, $\bar{F} \subseteq F$. We use $\{s\}^\Phi = \{s' \in \mathcal{S}: \Phi(s') = \Phi(s)\}$ to denote the equivalence class of $s$ under $\Phi$. In Example \ref{ex:frozen}, a potential state representation could be $\Phi(s_t) = \langle l_0, x_0 \rangle$ for all $s_t \in \mathcal{S}$. This representation retains only $L_0$ and $X_0$, ignoring all other variables in $F$. Hence, all states that share the same values for $L_0$ and $X_0$ belong to the same equivalence class.
% Not all state representations are sufficient to learn the optimal policy; some, like the one discussed in the above paragraph, may exclude variables that carry valuable information for the task at hand.
% \begin{definition}[Markov state representation] A state representation $\Phi(s_t)$ is said to be Markov if, for all $s_t, s_{t+1} \in \mathcal{S}$, $a_t \in \mathcal{A}$,
% \begin{equation*}
% R(s_t , a_t) = R(\Phi(s_t), a_t) \quad
% \text{and} \quad \sum_{s'_{t+1} \in \{s_{t+1}\}^\Phi} T(s'_{t+1} \mid s_t, a_t) = \Pr(\Phi(s_{t+1}) \mid \Phi(s_t), a_t),
% \end{equation*}
% where $R(\Phi(s_t), a_t)$ denotes the reward $R(s'_t, a_t)$ at any $s'_t \in \{s_t\}^\Phi$.
% % where $H'_{t+1} = \{h'_{t+1} \in H_{t+1} : \Phi(h'_{t+1}) = \Phi(h_{t+1})\}$.
% \label{def:markov_state}
% \end{definition}
% The phenomenon of policy confounding plays a fundamental role in this quest for simpler state representations, tricking the function approximator into forming state representations that are based on mere spurious correlations.
% \begin{definition}[Policy Confounding] A state representation $\Phi : \mathcal{S} \to \bar{\mathcal{S}}$ is said to be confounded by a policy $\pi$ if, for some $s_t, s_{t+1} \in \mathcal{S}$, $a_t \in \mathcal{A}$,
% \begin{equation*}
% R^\pi(\Phi(s_t), a_t) \neq R^\pi(\dooperator(\Phi(s_t)), a_t)
% \quad \text{or} \quad
% P^\pi(\Phi(s_{t+1}) \mid \Phi(s_t), a_t) \neq P^\pi(\Phi(s_{t+1}) \mid \dooperator(\Phi(s_t)), a_t).
% % $}
% \end{equation*}
% \normalsize
% \end{definition}
% The operator $\dooperator(\cdot)$ is known as the do-operator, and it is used to represent physical interventions in a system \citep{pearl2016causal}. These interventions are meant to distinguish cause-effect relations from mere statistical associations. In our case, $\dooperator(\Phi(s_t))$ means setting the variables forming the state representation $\Phi(s_t)$ to a particular value and considering all possible states in the equivalence class, $s'_t \in \{s_t\}^\Phi$, (i.e., all states that share the same value for the observation variables that are `selected' by $\Phi$; independently of whether these are visited by the policy being followed). For instance, in the above example, $R^{\pi^*}(L = 6) = +1$ when following $\pi^*$ since we know that $X = 1$, while $R^{\pi^*}(\dooperator(L = 6)) = \pm 1$ since $X$ can be either $0$ or $1$.
% \citep{suau2024bad} demonstrated that the phenomenon of policy confounding is particularly problematic when training agents with on-policy methods. This arises because, during policy updates, on-policy methods rely solely on trajectories collected using the current policy. This contrasts with off-policy methods, where the agent is trained on trajectories generated by multiple policies, broadening the trajectory distribution and reducing the risk of the agent learning spurious correlations present in specific trajectories.
% A popular family of on-policy methods includes those that directly optimize the policy by using the gradient of the expected reward with respect to the policy parameters, \(\theta\). The policy gradient theorem formalizes this as:
% \begin{equation}
% \nabla_\theta J(\pi_\theta) = \mathbb{E}_{s_t, a_t \sim \pi_\theta} \left[ \nabla_\theta \log \pi_\theta(a_t \mid s_t) Q^\pi(s_t, a_t) \right],
% \end{equation}
% where \(Q^\pi(s_t, a_t)\) is the action-value function under policy \(\pi\) \citep{sutton2018reinforcement}.
% Following this update rule increases the likelihood of sampling actions that lead to high returns while reducing the probability of actions leading to lower returns. However, this approach is associated with high variance. The \(Q\)-value reflects not only the immediate effectiveness of action \(a_t\) but also its impact on subsequent actions, complicating credit assignment and slowing learning. A common strategy to mitigate variance is to subtract the state value function \(V^\pi(s_t)\) from the \(Q\)-value, leading to the definition of the advantage function:
% \begin{equation}
% A^\pi(s_t, a_t) = Q^\pi(s_t, a_t) - V^\pi(s_t).
% \end{equation}
% This adjustment does not introduce bias, as the value function is independent of the action, but it significantly reduces the variance of the gradient estimation.
% \section{The (New) Role of the Advantage Function}
% \citep{suau2024bad} showed that the phenomenon of policy confounding is particularly troublesome when training agents with on-policy methods. This is because, when updating the policy, on-policy methods are restricted to using trajectories collected with the current policy. This in contrast with off-policy methods where the agent is trained on trajectories generated by multiple policies, hence broading the trajectory distribution and reducing the risk that the agent will pick up on the spurious correlations present in some of the trajectories.
% A popular family of on-policy methods, are those which directly optimize the policy using the gradient of the expected reward with respect to the policy parameters, $\theta$,
% FORMULA
% Following this update rule results in increasing the likelihood of sampling actions leading to high returns while decreasing the probability for those leading to lower returns. The main problem with this approach is its high variance. The $Q$ value contains information not only about how effective action $a$ is in getting high returns but also about subsequent actions. This can severely complicate credit assignment and hence slow down learning. A common solution is to subtract from the use the state value function to reduce the variance.
% EQUATION
% Since the value function does not depend on the action the expected gradient for that action remains the same (no bias) but its variance may be greatly reduced.
% However, as discussed in the introduction, the key message we want to convey in this paper is that the role of the advantage function is more profound than simply reducing the variance of the updates. The first thing to notice is that by subtracting the state value from the Q value we are effectively comparing the value of taking a particular action against the average value when we sample from the current policy.
% Secondly, $A^\pi(s_t, a_t)$ is proportional to the complement of the probability of $a_t$ under the current policy, $1 - \pi(a_t \mid s_t)$:
% \begin{equation}
% \begin{split}
% A^\pi(s_t, a_t) &= Q^\pi(s_t, a_t) - V^\pi(s_t) \\
% &= Q^\pi(s_t, a_t) - \sum_{a' \in \mathcal{A}} \pi(a'_t \mid s_t) Q^\pi(s_t, a'_t) \\
% &= (1 - \pi(a_t \mid s_t))Q^\pi(s_t, a_t) - \sum_{\substack{a' \in \mathcal{A}, a' \neq a}} \pi(a'_t \mid s_t) Q^\pi(s_t, a'_t) \\
% &= (1 - \pi(a_t \mid s_t)) Q^\pi(s_t, a_t) - (1 - \pi(a_t \mid s_t)) \tilde{Q}^\pi(s_t, a_t) \\
% &= (1 - \pi(a_t \mid s_t))\Delta Q^\pi(s_t, a_t)
% †\end{split}
% \end{equation}
% where
% \begin{equation}
% \tilde{Q}^\pi(s_t, a_t) = \frac{\sum_{\substack{a' \in \mathcal{A}, a' \neq a}} \pi(a'_t \mid s_t) Q^\pi(s_t, a'_t)}{\sum_{\substack{a' \in \mathcal{A}, a' \neq a}} \pi(a'_t \mid s_t)} = \frac{\sum_{\substack{a' \in \mathcal{A}, a' \neq a}} \pi(a'_t \mid s_t) Q^\pi(s_t, a'_t)}{1 - \pi(a_t \mid s_t)}
% \end{equation}
% is the average $Q$ value of all actions except $a_t$ weighted by their probabilities and normalized, and $\Delta Q^\pi(s_t, a_t) = Q^\pi(s_t, a_t) - \tilde{Q}^\pi(s_t, a_t)$.
% The above is important because it shows that as the probability of an action under $\pi$ decreases its advantage increases and vice versa.\footnote{Note that $\Delta Q^\pi$ stays constant even when we change the probability of $a_t$, assuming the probabilities for the other actions change proportionally}. This makes intuitive sense since, as mentioned before, the advantage function compares the current action value with the state value (i.e., the average action value under the current policy). Hence, actions that deviate a lot from said average (low probability under $\pi$) will have high advantages. Further, the importance of each sample in the gradient computation is balanced according to how likely it is under the current policy. An action that is very likely does not influence the gradient as much as a very unlikely action since the information content of the former is much lower. In the extreme, when the policy is deterministic, the samples have effectively zero information content, and as such, the advantage also becomes zero, which completely inhibits learning.
% Despite all the above, the question of how this property of the advantage function helps the agent learn causal state representations still remains. Let us illustrate this effect using the Key2Door example (Figure \ref{fig:key2door}).
% But how does all this help the agent learn more robust (causal) representations?
% \input{table}
% Table \ref{table:tab1} shows the $Q$-values and advantages when the agent is at location 6 and has or has not got the key under four different policies, which differ in the probabilities of choosing the optimal action $0.5$, $0.6$, $0.7$, and $0.9$. The last two columns show the state visitation probabilities $\mu^\pi$. The probabilities are normalized to show the relative frequency of being at location 6 and not having the key, versus being at location 6 and having the key.\footnote{Both $Q$ and $A$ are estimated doing value iteration via dynamic programming, whilst the state visitation are estimated by running the policy for multiple episodes and computing the frequency of each of the two states.}
% In each of the four cases we compare the average $Q$-value and advantage $A$ the agent would get if it chose the optimal action (optimal strategy; $Q^\pi_\text{opt}, A^\pi_\text{opt}$), and if it ignored the key and simply went right (naive strategy; $Q^\pi_\text{naive}, A^\pi_\text{naive}$) when being at location 6. To do so we multiply the normalized visitation probabilities, denoted by $\mu^\pi$ (last two columns in Table \ref{table:tab1}), when having and not having the key by the corresponding $Q$ values and advantages.
% For instance, under $\pi(a^* \mid s) = 0.6$ the average $Q$-value for the naive strategy $$Q^\pi_\text{naive} =\mu^\pi(\text{no key})Q^\pi(\text{no key}, a=\text{right}) + \mu^\pi(\text{key})Q^\pi(\text{key}, a=\text{right}) = 0.069 \times 0 + 0.931 \times 1 = \mathbf{0.931},$$ whereas the average $Q$-value for the optimal strategy is $$Q^\pi_\text{opt} = \mu^\pi(\text{no key})Q^\pi(\text{no key}, a=\text{left}) + \mu^\pi(\text{key})Q^\pi(\text{key}, a=\text{right}) = 0.069\times 0.274 + 0.931 \times 1 = \mathbf{0.949}.$$ If we do the same with the advantages we get $A^\pi_\text{naive} = \mathbf{0.048}$ and $A^\pi_\text{opt} =\mathbf{0.067}$. The optimal strategy yields a $\mathbf{39.583\%}$ performance increase with respect to the naive strategy under the advantage function, which contrasts with the $\mathbf{1.933 \%}$ performance increase under the $Q$-value. This implies that the agent will experience a much higher pressure for including the signal in its state representation (so as to know when to go right and left) when optimizing advantages instead of $Q$-values. The same quantities for the other policies are shown in Table $\ref{table:performance}$. We see how the $Q$-value percentage increase stays the same or even decreases as the policy improves, while the advantage percentage increase grows. Interestingly, when the probability of choosing the correct action reaches $0.9$, the percentage increase becomes $0$ for both. This is because the estimated probability of being at location 6 and not having the key is approximately zero (last row, second to last column in Table \ref{table:performance}). Hence, in both scenarios the agent may opt for ignoring the key as there is effectively no difference between the naive and optimal strategies. However, notice the magnitudes of the $Q$-value ($1.000$) and the advatange $0.005$. This implies that the agent will keep updating its state representation at the same pace when optimizing the $Q$-value, which entails a high risk of ignoring the key, while the magnitude of the updates will be significantly smaller with the advantage function, which will eventually become 0 as the agent converges to the optimal policy.
% \input{table2}
\subsection{Policy gradient and advantage function}\label{sec:policy_gradient}
\citet{suau2024bad} demonstrated that the phenomenon of policy confounding is particularly problematic when training agents with on-policy methods. This is because, when updating the policy, on-policy methods rely solely on trajectories collected using the current policy. This contrasts with off-policy methods, where the agent is trained on trajectories generated by multiple policies, thus broadening the trajectory distribution and reducing the risk of the agent picking up on spurious correlations present in specific trajectories.
A popular family of on-policy methods includes those that directly optimize the policy by following the gradient of the expected return with respect to the policy parameters, \(\theta\). The policy gradient theorem \citep{marbach1999simulation, sutton1999policy} formalizes this as:
\begin{equation}
\nabla_\theta J(\pi_\theta) = \mathbb{E}_{s_t, a_t \sim \pi_\theta} \left[ \frac{\nabla_\theta \pi_\theta(a_t \mid s_t)}{\pi_\theta(a_t \mid s_t)} Q^\pi(s_t, a_t) \right],
\end{equation}
\normalsize
where \(Q^\pi(s_t, a_t)\) is the action-value function under policy \(\pi\).
Following this gradient increases the likelihood of sampling actions that lead to high returns while reducing the probability of actions leading to lower returns. However,
in practice, computing the exact gradient is infeasible. Instead, we approximate the gradient using sample estimates from trajectories collected by the policy, which introduces high variance.
A common strategy to reduce variance is to subtract the state value function \(V^\pi(s_t)\) from the \(Q\)-value \citep{Baird94ICoNN, greensmith2001variance}, leading to the definition of the advantage function:
\begin{equation}
A^\pi(s_t, a_t) := Q^\pi(s_t, a_t) - V^\pi(s_t).
\end{equation}
\normalsize
This adjustment does not introduce bias, as the value function is independent of the action, but it significantly reduces the variance of the gradient estimation.
\section{The role of the advantage in learning causal state representations}\label{sec:advantage_role}
In this section, we demonstrate that, beyond its well-known role in variance reduction, the advantage function plays a crucial role in mitigating policy confounding. Specifically, it helps break any spurious correlations that the policy may introduce, thus preventing the agent from forming habits that fail to generalize beyond its typical trajectories. Proofs for the theoretical
results in this section can be found in Appendix \ref{ap:proofs}.
\subsection{Value and advantage function under a state representation}\label{sec:advantage_phi}
Computing the advantage usually involves fitting a model to approximate the value function. Conceptually, the value function can be decomposed into a state representation function $\Phi(s)$, which projects the states into a lower-dimensional space defined by the subset of variables relevant to the task, and a function $V^\pi(\Phi(s))$ that maps the state representation to a scalar value.
% \begin{definition}[$V^\pi$ and $Q^\pi$ under $\Phi$]
% Given a state representation $\Phi : \mathcal{S} \to \mathcal{\bar{S}}$ (Definition~\ref{def:state_representation}), the state value $V^\pi: \bar{\mathcal{S}} \to \mathbb{R}$ and state-action value $Q^\pi: \bar{\mathcal{S}} \times \mathcal{A} \to \mathbb{R}$ under $\Phi$ are defined as:
% \begin{equation}
%
% V^\pi(\Phi(s_t)) = \sum_{s_t' \in \{s_t\}^\Phi} P^\pi(s_t' \mid \Phi(s_t))V^\pi(s_t') \enspace \text{and} \enspace Q^\pi(\Phi(s_t), a_t) = \sum_{s_t' \in \{s_t\}^\Phi} P^\pi(s_t' \mid \Phi(s_t))Q^\pi(s_t', a_t).
% \end{equation}
% with
% \[
% P(s'_t \mid \Phi(s_t)) := \frac{d^\pi(s_t)}{\sum_{s_t' \in \{s_t\}^\Phi} d^\pi(s_t')}.
% \]
% \end{definition}
\begin{definition}[$V^\pi$ and $Q^\pi$ under $\Phi$]
Let $\Phi : \mathcal{S} \to \bar{\mathcal{S}}$ be a state representation (Definition~\ref{def:state_representation}) that induces an equivalence class $\{s\}^\Phi := \{s' \in \mathcal{S} : \Phi(s') = \Phi(s)\}$. Assume the Markov chain under policy $\pi$ admits a stationary distribution $d^\pi$ over $\mathcal{S}$. Then the state value $V^\pi: \bar{\mathcal{S}} \to \mathbb{R}$ and state-action value $Q^\pi: \bar{\mathcal{S}} \times \mathcal{A} \to \mathbb{R}$ under $\Phi$ are defined as:
\[
V^\pi(\Phi(s_t)) := \sum_{s_t' \in \{s_t\}^\Phi} P^\pi(s_t' \mid \Phi(s_t)) \, V^\pi(s_t')
\]
and
\[
Q^\pi(\Phi(s_t), a_t) := \sum_{s_t' \in \{s_t\}^\Phi} P^\pi(s_t' \mid \Phi(s_t)) \, Q^\pi(s_t', a_t),
\]
\normalsize
where the distribution $P^\pi(s_t' \mid \Phi(s_t))$ is defined as:
\[
P^\pi(s_t' \mid \Phi(s_t)) := \frac{d^\pi(s_t')}{\sum_{s \in \{s_t\}^\Phi} d^\pi(s)} \quad \text{for all } s_t' \in \{s_t\}^\Phi.
\]
\end{definition}
\begin{lemma}
Let $\Phi$ be a Causal (Markov) State Representation (Definition \ref{def:markov_state}), then
\begin{equation}
Q^\pi(s_t, a_t) = Q^\pi(s'_t, a_t) = Q^\pi(\Phi(s_t), a_t) \qquad \text{and} \qquad V^\pi(s_t) = V^\pi(s'_t) = V^\pi(\Phi(s_t))
\end{equation}
\normalsize
for all $a \in \mathcal{A}$, $s \in \mathcal{S}$, and, $s' \in \{s\}^\Phi$.
\label{lm:v_q_phi}
\end{lemma}
\begin{definition}[$A_\Phi^\pi$ under $\Phi$]\label{def:advantage_phi}
Given a state representation $\Phi : \mathcal{S} \to \mathcal{\bar{S}}$, the advantage function $A_\Phi^\pi: \mathcal{S} \to \mathbb{R}$ under $\Phi$ is defined as:
\begin{equation}
A^\pi_\Phi(s_t, a_t) := Q^\pi(s_t, a_t) - V^\pi(\Phi(s_t)).
\end{equation}
\normalsize
\end{definition}
Despite the agent's state representation being $\Phi(s_t)$, Definition \ref{def:advantage_phi} shows that the advantage function $A_\Phi^\pi$ still depends on the full state $s_t$. This is because $Q^\pi(s_t, a_t)$ is typically estimated directly for each state $s_t$ and action $a_t$ from Monte Carlo rollouts, making it independent of $\Phi(s_t)$. In contrast, $V^\pi(\Phi(s_t))$ averages over the value of all states equivalent under $\Phi$, effectively marginalizing out the variables not preserved by $\Phi$. As a result, $A_\Phi^\pi$ measures the deviation of a particular state-action pair from the expected value across the equivalence class $\{s_t\}^\Phi$.
\subsection{The scaling effect of the advantage function}
\begin{restatable}{theorem}{theoremone}
Let $\Phi$ be an arbitrary state representation. The advantage function can be expressed as:
\begin{equation}
A^\pi_\Phi(s_t, a_t) = (1 - P^\pi(s_t, a_t \mid \Phi(s_t)))( Q^\pi(s_t, a_t) - \tilde{Q}^\pi(\neg \langle s_t, a_t\rangle))
\end{equation}
\normalsize
where
\begin{equation}
\tilde{Q}_\Phi^\pi(\neg \langle s_t, a_t\rangle) = \frac{\sum_{ s'_t, a'_t \neq s_t, a_t } P^\pi(s'_t, a'_t \mid \Phi(s_t)) Q^\pi(s'_t, a'_t)}{\sum_{ s'_t, a'_t \neq s_t, a_t } P^\pi(s'_t, a'_t \mid \Phi(s_t))}
\end{equation}
\label{thm:advantage_phi}
\end{restatable}
This result highlights that the advantage function is proportional to the complement of the joint probability of a specific state-action pair conditioned on the state representation, i.e., $1 - P^\pi(s_t, a_t \mid \Phi(s_t))$. This ensures that the advantage of less probable state-action pairs is larger in magnitude, and helps offset the overrepresentation of more frequent pairs under $\pi$.
This scaling effect is crucial for mitigating policy confounding as it breaks any spurious correlations induced by the policy. By amplifying rare but informative samples, the advantage function encourages the agent to attend to underlying causal factors rather than simply reinforcing habits. Intuitively, it reweights each sample according to its causal content, preventing high-probability pairs from dominating the gradient updates during training.
Moreover, the difference \( Q^\pi(s_t, a_t) - \tilde{Q}^\pi_\Phi(\neg \langle s_t, a_t \rangle) \) reflects how much better (or worse) a given pair is relative to the alternative pairs. Notably, \( \tilde{Q}^\pi_\Phi(\neg \langle s_t, a_t \rangle) \) is a weighted average over state-action pairs in the same equivalence class, excluding $\langle s_t, a_t \rangle$, and since it does not depend on $P^\pi(s_t, a_t \mid \Phi(s_t))$, it provides a stable, representation-conditional baseline for comparison.
\begin{corollary}
Let $\Phi$ be a Causal (Markov) State Representation. The advantage function can be expressed as:
\begin{equation}
A^\pi_\Phi(s_t, a_t) = (1 - \pi(a_t \mid \Phi(s_t)))( Q^\pi(s_t, a_t) - \tilde{Q}^\pi(s_t, \neg a_t))
\end{equation}
\normalsize
where
\begin{equation}
\tilde{Q}^\pi(s_t, \neg a_t) = \frac{\sum_{a'_t \neq a_t} \pi(a'_t \mid \Phi(s_t)) Q^\pi(s_t, a'_t)}{\sum_{a'_t \neq a_t} \pi(a'_t \mid \Phi(s_t))}
\end{equation}
\end{corollary}
The above result follows directly from Theorem \ref{thm:advantage_phi} and Lemma \ref{lm:v_q_phi}. Note that when the state representation $\Phi$ is causal, there is no need to scale the sample gradients according to the likelihood of the states given $\Phi$, as $\Phi$ already accounts for the true factors governing rewards and transitions.
% However we still need to adjust for the likelihood of the actions... WHY?
% In the extreme case of a deterministic policy, the advantage becomes zero, which halts learning entirely.
\subsection{Impact on policy gradients}
None of the above would matter if the updates strictly followed the exact policy gradient, as the advantage function does not alter the gradient \citep{sutton2018reinforcement}. However, as discussed in Section \ref{sec:policy_gradient}, policy gradients are typically based on sample estimates (i.e., stochastic gradients).
As a result, depending on the agent's policy and the batch size, it is often the case that for a given $\Phi(s_t)$, all samples in a training batch correspond to a single state $s_t$, rather than including all states $s'_t \in \{s\}^\Phi$. Consequently, when using $Q$-values instead of advantages $A^\pi_\Phi$, this can lead to overly aggressive gradient updates \citep{chung2021beyond}, causing the agent to develop habits based on spurious correlations. This issue is further exacerbated as the policy reinforces the likelihood of frequently observed state-action pairs, creating a vicious circle \citep{mei2022role}. In the extreme, when the policy becomes deterministic, certain state-action pairs may have their probabilities driven to zero, requiring an infeasibly large number of samples to correct for the overestimation of the state-action pairs that are actually visited.
In contrast, the advantage function $A^\pi_\Phi$ scales the gradients by adjusting their magnitude based on the probability of each state-action pair (Theorem \ref{thm:advantage_phi}). This moderates the aggressiveness of gradients for frequently sampled state-action pairs while boosting gradients for less common ones, thereby helping the agent break bad habits.
\subsection{Practical considerations}\label{sec:practical_considerations}
Given the points above, an alternative to using advantages could be to employ large batch sizes. However, the batch size may need to be significantly increased depending on the environment for this approach to be effective. Moreover, we find no compelling reason to avoid using advantages.
It is also important to note that the results reported by \citet{suau2024bad} indicate that PPO struggles with out-of-trajectory generalization, despite the default setting utilizing advantages. We hypothesize that this issue arises because, in most PPO implementations, the advantages in each training batch are normalized before the network update. This normalization removes the scaling effect of the advantage function on the gradients. Both of these aspects are further analyzed in the experiments section.
\subsection{Example: Key2Door}\label{sec:advatange_example}
To illustrate the insights discussed in the previous sections, we revisit the Key2Door example (Section \ref{sec:example}). Suppose the agent's state representation consists solely of the location variable, $\Phi(s) = l$. Table \ref{table:tab1} shows the $Q$-values and corresponding advantages $A^\pi_\Phi$ when the agent is at location 6, both with and without the key, under five different policies. These policies differ in the probability of selecting the optimal action: $0.5$, $0.6$, $0.7$, $0.8$, and $0.9$. The last two columns indicate the probability that the agent has, $P^\pi(X=1 \mid L=6)$, or has not, $P^\pi(X=0\mid L=6)$, the key at location 6.\footnote{$Q$-values and advantages are computed using value iteration. $P^\pi(X \mid L=6)$ is estimated by running the policy over multiple episodes and calculating the frequency of each state.
}
\input{table_phi}
As revealed by Theorem \ref{thm:advantage_phi}, the magnitudes of the advantages depend on the joint probability of visiting a specific state-action pair. In particular, the advantage decreases (increases) in magnitude as the probability $P^\pi(X, A \mid L =6)$ increases (decreases). For example, the advantage of moving right \emph{with} the key is much larger when the probability of having the key is $0.832$ compared to when it is $0.999$ ($0.305$ vs. $0.006$). This difference is amplified by the increasing probability of taking the right action ($0.5$ vs. $0.9$). In contrast, the corresponding $Q$-values remain constant at $1$.
Similarly, the advantage of moving right \emph{without} the key is smaller (in magnitude) when the probability of not having the key is $0.168$ than when the probability is $0.000$ ($-0.695$ vs. $-0.994$), even though the $Q$-value remains $0$.\footnote{Note that the magnitude of the advantage of moving left \emph{with} the key first increases and then decreases across rows. This occurs because, although the probability of having the key increases, the probability of moving left decreases.}
The above is key because, as the policy improves, the probability of visiting state-action pairs such $\langle L=6, X=1, A=1 \rangle$ increases, and hence the training batches start to fill with many such pairs. Meanwhile, less frequent but informative pairs, such as $\langle L = 6, X = 0, A = 1 \rangle$, become underrepresented. As a result, training on raw $Q$-values can lead the agent to disregard the key variable $X$, effectively exploiting a spurious correlation between location 6 and having the key. This occurs because updates to the policy treat all state-action pairs equally, regardless of how often they appear in the training data. Notably, even under a fully random policy (top row of Table~\ref{table:tab1}), the probability of having the key at location 6 is already $0.832$, meaning the agent often receives a reward of $1$ simply by moving right at location 6, without explicitly reasoning about whether it has the key.
In contrast, agents trained with the advantage function are less likely to ignore the key variable. When the policy is random, the advantage of moving right \emph{without} the key has a large negative value ($-0.695$), whereas the advantage of moving right \emph{with} the key is just $0.305$, despite the corresponding Q-values being $0$ and $1$ respectively. These differences in magnitude help counteract the overrepresentation of states such as $\langle L=6, X=1, A=1 \rangle$, and ensure that the agent does not develop the habit of moving solely based on location.
% Furthermore, there is little incentive for the agent to actually listen to the key variable as moving left when not having the key only gives the agent an extra value of $0.038$ over the alternative of simply moving right. The value of going left when not having the key does increase as the agent gets better but at the same time, the likelihood of not having the key decreases.
% However, the magnitudes of the $Q$-value ($1.000$) and the advantage ($0.005$) remain distinct. This suggests that when optimizing the $Q$-value, the agent will continue updating its state representation at the same rate, posing a significant risk of learning to disregard the key. In contrast, the advantage function leads to much smaller update magnitudes, which will progressively approach zero as the agent converges to the optimal policy.
% Table \ref{table:tab2} shows a comparison between the advantages the agent perceives when the state representation consists solely of the agent's location $\Phi(s) = l$ and when the state representation is the causal state representation $\Phi(s) = \langle l, x\rangle$. Perhaps unsurprisingly, the magnitudes of the advantages are lower for the causal state representation. This is because, as shown by equation \ref{eq:causal_advantage}, when the representation is causal, the value difference $\Delta Q^\pi(s_t,a_t)$ is only scaled by the complement of the action probability under the current policy, $1 - \pi(a_t \mid s_t)$. This makes sense since the representation is already causal there is no need for rescaling the effect of the sampled states based on their probability given the agent's state representation.
% \input{table}
% In each of the four cases, we compare the average $Q$-value and advantage \(A\) the agent would obtain by choosing the optimal action (optimal strategy; \(Q^\pi_\text{opt}, A^\pi_\text{opt}\)) versus simply ignoring the key and going right (naive strategy; \(Q^\pi_\text{naive}, A^\pi_\text{naive}\)) while at location 6. To do this, we multiply the normalized visitation probabilities \(\mu^\pi\) (from the last two columns in Table \ref{table:tab1}) for both "no key" and "key" by their corresponding $Q$-values and advantages.
% For example, under the policy \(\pi(a^* \mid s) = 0.6\), the average $Q$-value for the naive strategy is calculated as:
% \begin{equation}
% Q^\pi_\text{naive} = \mu^\pi(\text{no key})Q^\pi(\text{no key}, a=\text{right}) + \mu^\pi(\text{key})Q^\pi(\text{key}, a=\text{right}) = 0.069 \times 0 + 0.931 \times 1 = \mathbf{0.931}.
% \end{equation}
% Meanwhile, the average $Q$-value for the optimal strategy is:
% \begin{equation}
% Q^\pi_\text{opt} = \mu^\pi(\text{no key})Q^\pi(\text{no key}, a=\text{left}) + \mu^\pi(\text{key})Q^\pi(\text{key}, a=\text{right}) = 0.069 \times 0.274 + 0.931 \times 1 = \mathbf{0.949}.
% \end{equation}
% For the advantages, we obtain:
% \begin{equation}
% A^\pi_\text{naive} = \mu^\pi(\text{no key})A^\pi(\text{no key}, a=\text{right}) + \mu^\pi(\text{key})A^\pi(\text{key}, a=\text{right}) = 0.069 \times -0.882 + 0.931 \times 0.118 = \mathbf{0.049}.
% \end{equation}
% \begin{equation}
% \begin{aligned}
% A^\pi_\text{opt} &= \mu^\pi(\text{no key})A^\pi(\text{no key}, a=\text{left}) + \mu^\pi(\text{key})A^\pi(\text{key}, a=\text{right}) \\
% &= 0.069 \times 0.109 + 0.931 \times 0.064 = 0.069 \times -0.608 + 0.931 \times 0.118 =\mathbf{0.067}.
% \end{aligned}
% \end{equation}
% $A^\pi_\text{naive} = \mathbf{0.048} \quad \text{and} \quad A^\pi_\text{opt} = \mathbf{0.067}$.
% The optimal strategy results in a \(\frac{0.067-0.048}{0.048} = \mathbf{39.583\%}\) performance increase compared to the naive strategy when optimizing the advantage function, which sharply contrasts with the \(\mathbf{1.933 \%}\) performance increase under the $Q$-value. This indicates that the agent is significantly more incentivized to incorporate the key into its state representation (i.e., to distinguish when to go right or left) when optimizing the advantage function, as opposed to optimizing the $Q$-value. The same metrics for the other policies are shown in Table \ref{table:performance}. Notably, we observe that while the $Q$-value percentage increase remains stable or even decreases as the policy improves, the advantage percentage increase continues to grow.
% Interestingly, when the probability of selecting the correct action reaches $0.9$, the percentage increase for both the $Q$-value and the advantage becomes zero. This is because the estimated probability of being at location 6 without the key is almost 0. In this case, the agent may effectively ignore the key, as the difference between the naive and optimal strategies becomes minimal. However, the magnitudes of the $Q$-value ($1.000$) and the advantage ($0.005$) remain distinct. This suggests that when optimizing the $Q$-value, the agent will continue updating its state representation at the same rate, posing a significant risk of learning to disregard the key. In contrast, the advantage function leads to much smaller update magnitudes, which will progressively approach zero as the agent converges to the optimal policy.
% \input{table2_phi}
% EXPLORATION AND HOW THE POLICY CONVERGES TO 50/50 FOR ACTIONS WITH EQUALLY OPTIMAL VALUES.
% Finally, since the expected value of the advantage function is zero by construction, we can write:
% \begin{equation}
% \begin{split}
% 0 &= \mathbb{E}{a \sim \pi}[A^\pi(s_t, a)] \\
% &= \sum_{a \in \mathcal{A}} \pi(a|s_t)A^\pi(s_t, a) \\
% &= \sum_{a \in \mathcal{A}} \pi(a|s_t)(1 - \pi(a|s_t))\Delta Q^\pi(s_t, a)
% \end{split}
% \end{equation}
% This relationship becomes particularly interesting when we consider how the advantage function is used in policy gradient methods. When we use the advantage function as a critic, we are effectively performing a form of inverse probability weighting. To see this, let's examine the policy gradient update:
% The key insight is that the advantage function naturally corrects for the sampling distribution through the $(1 - \pi(a|s_t))$ term, achieving the same effect as if we had sampled uniformly and used importance weights. This is particularly useful because:
% It allows us to sample from the current policy (which is often more efficient)
% It automatically adjusts the learning signal based on action probabilities
% It doesn't require explicit importance sampling weights, which can have high variance
% This relationship helps explain why advantage-based methods are so effective in practice - they implicitly implement a form of variance reduction that would otherwise require uniform sampling.
% The problem with the above update rule is that it reinforces already learned behaviors.
% Notice how for a fixed value of Q the advantage of actions that are less likely under a given policy is higher than those which are less likely.
% As discussed in the previous section, policies may inadvertently create feedback loops between past and future observation variables as a result of influencing both. These feedback loops can cause spurious correlations that the agent may pick up on.
% \subsection{Reward Re-weighting}
% Reward should be 0 if the policy is completely deterministic. Rewards should be higher the lower the probability of a state/trajectory is.
% An action is rewarded relative to how much that action is taken by the current policy.
% Initially, the probability of opening the door when at location 0 -- and thus obtaining a $+1$ reward, is $\approx 50 \%$. However, as the agent improves, that probability also increases. We need a way of balancing out that effect, and preventing that location, instead of having the key, is credited for the high reward. such that the agent does not confuse being in that location with having the key. The way to do so is through the advantage function. We kind of . We can down-weigh the samples where the agent is at location 6 and has the key while up-weighing the samples where the agent is at location 6 and does not have the key. This way the agent can more easily figure out that it is indeed the key, which is responsible for the high reward.
% EQUIVALENT TO WEIGHTING THE SAMPLES ACCORDING TO THEIR PROBABILITY?
% EXPECTED VALUE OF BEING IN LOCATION 6 IS DIFFERENT THAN EXPECTED VALUE DO(LOCATION=6) DEPENDS ON THE ENVIRONMENT TRANSTIONS THE STARTING STATE AND THE AGENT'S POLICY.
% THE EXTEND TO WHICH A PARTICULAR TRAJECTORY IS INFORMATIVE DEPENDS ON THE EXTENT TO WHICH THE TRAJECTORY DEVIATES FROM THE USUAL TRAJECTORIES DICTATED BY THE AGENT'S POLICY AND THE ENVIRONMENT.
% STATE REPRESENTATION THAT ARE INVARIANT TO THE POLICY AND THE EXOGENOUS VARIABLES.
% PAN AND SCHOLKOPF: "by taking the difference of the two, we essentially
% remove the exogenous component, as in the case of the advantage function"
% SUAU: Not only the exogenous component but also the component that is due to the policy being followed.
% No normalization
% \section{Proof}
% \input{proof}
\section{Experiments}
The experiments aim to verify whether the insights discussed in the previous section hold in practice. Specifically, we seek to demonstrate that training on the advantage function, rather than raw $Q$-values, helps agents develop state representations that better capture causal factors and thus generalize out-of-trajectory. To test this, we conduct experiments on three gridworld environments: the Key2Door environment described in Section \ref{sec:example}, as well as the Frozen T-Maze and Diversion environments introduced by \citet{suau2024bad}.
We evaluate the agent's performance in both the training environments and in modified versions, referred to as the evaluation environments, where, like in the Key2Door environment, the agent is forced to deviate from its usual trajectory. Furthermore, we analyze the effects of advantage normalization and batch size, which, as discussed in section \ref{sec:practical_considerations}, can influence out-of-trajectory generalization. Finally, we inspect the state representations learned by the agents by measuring the KL divergence of the policies between various state observations. Details about the T-Maze and Diversion environments are provided in Appendix \ref{ap:environments}.
\subsection{Experimental setup}\label{sec:experimental_setup}
Agents are trained using two different on-policy policy-gradient methods, REINFORCE \citep{Williams92ML} and PPO \citep{schulman2017proximal} to maximize either the advantage or the $Q$-value. We implement policies and value functions as feedforward neural networks and use a stack of past observations as input in environments that require memory. The results are averaged over 10 random seeds. We report the average return as a function of the number of training steps. The shaded areas show the standard error of the mean. Training is interleaved with periodic evaluations in the original environments and their variants. Further details about the experimental setup are provided in Appendix \ref{ap:experimental_setup}.
\subsection{Results}
\begin{figure}[t]
% \vspace{-15pt}
\centering
% First plot
\begin{subfigure}[t]{0.32\textwidth}
\centering
\includegraphics[width=\textwidth]{figures/keydoor_plot.pdf}
\end{subfigure}
\hfill
% Second plot
\begin{subfigure}[t]{0.32\textwidth}
\centering
\includegraphics[width=\textwidth]{figures/tmaze_plot.pdf}
\end{subfigure}
\hfill
% Third plot
\begin{subfigure}[t]{0.32\textwidth}
\centering
\includegraphics[width=\textwidth]{figures/diversion_plot.pdf}
\end{subfigure}
\hfill
% \vspace{-10pt}
\caption{Performance of PPO using Q-values and Advantages in both the training and evaluation variants of the Key2Door (first plot), Frozen T-Maze (second plot), and Diversion (third plot).
}
\vspace{-10pt}
\label{fig:qvsa}
\end{figure}
\begin{wrapfigure}{r}{0.65\textwidth}
\vspace{-25pt}
\begin{center}
\begin{subfigure}[t]{0.325\textwidth}
\centering
\includegraphics[width=\textwidth]{figures/keydoor_normalization.pdf}
% \vspace{-10pt}
\end{subfigure}
% Second plot
\hspace{-5pt}
\begin{subfigure}[t]{0.325\textwidth}
\centering
\includegraphics[width=\textwidth]{figures/keydoor_batchsize.pdf}
\end{subfigure}
\end{center}
\caption{\textbf{Left:} Performance of PPO with and without advantage normalization in the Key2Door evaluation environment. \textbf{Right:} Performance of PPO with different batch sizes in the Key2Door evaluation environment.}
\vspace{-5pt}
\label{fig:normalize_batchsize}
\end{wrapfigure}
Figure \ref{fig:qvsa} shows the performance curves in all three environments for PPO. Agents trained using the advantage function (green) perform well in both the training and evaluation environments. In contrast, agents trained on the $Q$-value (magenta) perform poorly in the evaluation environments. Similar results using REINFORCE are reported in Appendix \ref{ap:reinforce}.
The plot on the left of Figure \ref{fig:normalize_batchsize} reveals how, as discussed in Section \ref{sec:practical_considerations}, normalizing the advantages removes their scaling effect and results in agents being unable to perform well on the evaluation environment. The plot on the right, on the other hand, shows how the performance of policies trained on the $Q$ value improves as we increase the batch size, suggesting that the problem of state-action pair imbalance can sometimes be partly mitigated by using larger batch sizes. Results for the other two environments are provided in Appendices \ref{ap:advantage_normalization} and \ref{ap:batch_size}.
The heatmaps in Figure \ref{fig:heatmaps} show the KL divergence of action probabilities between the agent having the key and not having the key at each of the six locations, measured at different training steps (top: 10k steps, middle: 50k steps, bottom: 100k steps).
\begin{wrapfigure}{r}{0.65\textwidth}
\vspace{-10pt}
\begin{center}
\begin{subfigure}[t]{0.325\textwidth}
\centering
\includegraphics[width=\textwidth]{figures/kl_heatmaps_q.pdf}
% \vspace{-10pt}
\end{subfigure}
% Second plot
\hspace{-5pt}
\begin{subfigure}[t]{0.325\textwidth}
\centering
\includegraphics[width=\textwidth]{figures/kl_heatmaps_a.pdf}
\end{subfigure}
\end{center}
\vspace{-5pt}
\caption{KL divergence of action probabilities with and without the key, measured at different training steps (top 10k steps, middle 50k steps, and bottom 100k steps) for agents trained on the $Q$-value (left) and the Advantage function (right).}
% \vspace{-15pt}
\label{fig:heatmaps}
\end{wrapfigure}The heatmap on the left corresponds to agents trained on the $Q$-value function, while the one on the right corresponds to agents trained on the advantage. Higher KL divergence indicates that the policy places more importance on the key feature. The results reveal that agents trained to optimize $Q$-value may learn policies that do not necessarily rely on the causal factors due to policy confounding, while policies trained on advantage function may be less affected by this phenomenon. Similar results for the Frozen T-Maze and Diversion environments are provided in Appendix \ref{ap:heatmaps}.
\section{Limitations}
Our analysis assumes discrete state and action spaces, although the results could be extended to continuous domains. In Section \ref{sec:advantage_phi}, when defining the advantage function under a given state representation $\Phi$ (Definition \ref{def:advantage_phi}), we assume that the value function $Q(s_t, a_t)$ is estimated with respect to the full state, while $V^\pi(\Phi(s))$ is defined for the specific state represantion. This assumption is justified in settings where $Q^\pi$ is at least partly estimated from Monte Carlo rollouts, as is common in many actor-critic methods. However, if $Q$ were instead learned entirely via function approximation, then it too would be subject to $\Phi$, and Theorem \ref{thm:advantage_phi} would no longer hold in its current form.
The experiments were conducted in the same three environments introduced by \citet{suau2024bad}, which were specifically designed to expose the phenomenon of policy confounding. While these environments are deliberately simple to facilitate analysis, this simplicity limits the generalizability of our findings. As such, we draw no conclusions about the effectiveness of the advantage function for learning causal representations in more complex or high-dimensional domains. Addressing this question would require further empirical investigation and is left for future work.
Finally, throughout the paper, we have taken care not to make strong claims about the effectiveness of the advantage function in learning causal state representations. Rather, we argue that the advantage function can mitigate policy confounding and support the formation of more causally grounded representations. However, its use provides no guarantees that the resulting representations will be truly causal.
\section{Conclusion}
In this paper, we analyzed the role of the advantage function in helping agents learn causal state representations. We showed that the advantage function scales the gradients by the complement of the probability of the corresponding state-action pair. This increases the magnitude of the gradients for state-action pairs that are less likely under the current policy while decreasing it for those that are more likely. As a result, it downweights the impact of state-action pairs that are overrepresented in training batches while amplifying the impact of those that are underrepresented. This helps break spurious correlations introduced by the policy, allowing agents to focus on the true causal factors. Section \ref{sec:advatange_example} provides a detailed numerical example illustrating this effect.
Our experiments on the Key2Door, Frozen T-Maze, and Diversion environments confirmed that training on advantages leads to more robust agents that generalize better out-of-trajectory. Furthermore, as explained in Section \ref{sec:practical_considerations}, our empirical results reveal how implementation choices, such as batch size and advantage normalization, affect the learned representations. Finally, the KL-divergence analysis of the action probabilities further demonstrates that using the advantage function makes agents more reliant on causal factors.
% \section{Discussion}
% Hence, our contribution is to emphasize that the advantage function plays a more profound role than just variance reduction. It provides a mechanism that influences the balance of gradient contributions, which is crucial for efficient learning dynamics, especially when dealing with sparse or highly probable actions.
% As discussed in the previous section, policies may inadvertently create feedback loops between past and future observation variables as a result of influencing both. These feedback loops can cause spurious correlations that the agent may pick up on.
% The key insight is that the advantage function naturally corrects for the sampling distribution through the $1 - \pi(a|s_t)$ term, automatically adjusting the learning signal based on action probabilities. The problem with the Q value is that it reinforces already learned behaviors. Notice how for a fixed value of Q the advantage of actions that are less likely under a given policy is higher than those which are less likely.
\bibliography{main}
\bibliographystyle{apalike}
\newpage
% \section*{NeurIPS Paper Checklist}
% \begin{enumerate}
% \item {\bf Claims}
% \item[] Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
% \item[] Answer: \answerYes{}%, \answerNo{}, or \answerNA{}.
% \item[] Justification:
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the abstract and introduction do not include the claims made in the paper.
% \item The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
% \item The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
% \item It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
% \end{itemize}
% \item {\bf Limitations}
% \item[] Question: Does the paper discuss the limitations of the work performed by the authors?
% \item[] Answer: \answerYes{}
% \item[] Justification:
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
% \item The authors are encouraged to create a separate "Limitations" section in their paper.
% \item The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
% \item The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
% \item The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
% \item The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
% \item If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
% \item While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
% \end{itemize}
% \item {\bf Theory assumptions and proofs}
% \item[] Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
% \item[] Answer: \answerYes{}
% \item[] Justification:
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the paper does not include theoretical results.
% \item All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.
% \item All assumptions should be clearly stated or referenced in the statement of any theorems.
% \item The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
% \item Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
% \item Theorems and Lemmas that the proof relies upon should be properly referenced.
% \end{itemize}
% \item {\bf Experimental result reproducibility}
% \item[] Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
% \item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification:
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the paper does not include experiments.
% \item If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
% \item If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
% \item Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
% \item While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
% \begin{enumerate}
% \item If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
% \item If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
% \item If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
% \item We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
% \end{enumerate}
% \end{itemize}
% \item {\bf Open access to data and code}
% \item[] Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?
% \item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification: The code is included in the supplemental materials folder.
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that paper does not include experiments requiring code.
% \item Please see the NeurIPS code and data submission guidelines (\url{https://nips.cc/public/guides/CodeSubmissionPolicy}) for more details.
% \item While we encourage the release of code and data, we understand that this might not be possible, so “No” is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
% \item The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (\url{https://nips.cc/public/guides/CodeSubmissionPolicy}) for more details.
% \item The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
% \item The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.
% \item At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
% \item Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
% \end{itemize}
% \item {\bf Experimental setting/details}
% \item[] Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
% \item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification: Details about the experimental setup are provided in Section \ref{sec:experimental_setup} and Appendix \ref{ap:experimental_setup}.
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the paper does not include experiments.
% \item The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
% \item The full details can be provided either with the code, in appendix, or as supplemental material.
% \end{itemize}
% \item {\bf Experiment statistical significance}
% \item[] Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
% \item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification:
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the paper does not include experiments.
% \item The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
% \item The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
% \item The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
% \item The assumptions made should be given (e.g., Normally distributed errors).
% \item It should be clear whether the error bar is the standard deviation or the standard error of the mean.
% \item It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96\% CI, if the hypothesis of Normality of errors is not verified.
% \item For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
% \item If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
% \end{itemize}
% \item {\bf Experiments compute resources}
% \item[] Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
% \item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification: Information on the compute resources is provided in Appendix \ref{ap:experimental_setup}.
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the paper does not include experiments.
% \item The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
% \item The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
% \item The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
% \end{itemize}
% \item {\bf Code of ethics}
% \item[] Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics \url{https://neurips.cc/public/EthicsGuidelines}?
% \item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification:
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
% \item If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
% \item The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).
% \end{itemize}
% \item {\bf Broader impacts}
% \item[] Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
% \item[] Answer: \answerNA{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification: Foundational research
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that there is no societal impact of the work performed.
% \item If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
% \item Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
% \item The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
% \item The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
% \item If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
% \end{itemize}
% \item {\bf Safeguards}
% \item[] Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
% \item[] Answer: \answerNA{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification:
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the paper poses no such risks.
% \item Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
% \item Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
% \item We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.
% \end{itemize}
% \item {\bf Licenses for existing assets}
% \item[] Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
% \item[] Answer: \answerNA{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification:
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the paper does not use existing assets.
% \item The authors should cite the original paper that produced the code package or dataset.
% \item The authors should state which version of the asset is used and, if possible, include a URL.
% \item The name of the license (e.g., CC-BY 4.0) should be included for each asset.
% \item For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
% \item If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, \url{paperswithcode.com/datasets} has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.
% \item For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
% \item If this information is not available online, the authors are encouraged to reach out to the asset's creators.
% \end{itemize}
% \item {\bf New assets}
% \item[] Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
% \item[] Answer: \answerNA{}% Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification:
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the paper does not release new assets.
% \item Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
% \item The paper should discuss whether and how consent was obtained from people whose asset is used.
% \item At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
% \end{itemize}
% \item {\bf Crowdsourcing and research with human subjects}
% \item[] Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
% \item[] Answer: \answerNA{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification:
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
% \item Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
% \item According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
% \end{itemize}
% \item {\bf Institutional review board (IRB) approvals or equivalent for research with human subjects}
% \item[] Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
% \item[] Answer: \answerNA{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification:
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
% \item Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
% \item We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
% \item For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
% \end{itemize}
% \item {\bf Declaration of LLM usage}
% \item[] Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
% %this research?
% \item[] Answer: \answerNA{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification:
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
% \item Please refer to our LLM policy (\url{https://neurips.cc/Conferences/2025/LLM}) for what should or should not be described.
% \end{itemize}
% \end{enumerate}
\newpage
\appendix
\section{Proofs}\label{ap:proofs}
\theoremone*
\begin{proof}
\begin{equation}
\begin{aligned}
A^\pi_\Phi(s_t, a_t) &= Q^\pi(s_t, a_t) - V^\pi(\Phi(s_t)) \\
&= Q^\pi(s_t, a_t) - \sum_{s'_t \in \{s_t\}^\Phi} P^\pi(s'_t \mid \Phi(s_t)) V^\pi(s'_t)\\
&= Q^\pi(s_t, a_t) - \sum_{s'_t \in \{s_t\}^\Phi} P^\pi(s'_t \mid \Phi(s_t))\sum_{a' \in \mathcal{A}} \pi(a' \mid s'_t) Q^\pi(s'_t, a') \\
& = Q^\pi(s_t, a_t) - \sum_{\substack{s'_t \in \{s_t\}^\Phi, a'_t \in \mathcal{A}}} P^\pi(s'_t, a'_t \mid \Phi(s_t)) Q^\pi(s'_t, a'_t) \\
& = (1 - P^\pi(s_t, a_t \mid \Phi(s_t))) Q^\pi(s_t, a_t) - \sum_{ s'_t, a'_t \neq s_t, a_t } P^\pi(s'_t, a'_t \mid \Phi(s_t)) Q^\pi(s'_t, a'_t) \\
&= (1 - P^\pi(s_t, a_t \mid \Phi(s_t))) Q^\pi(s_t, a_t) - (1 - P^\pi(s_t, a_t \mid \Phi(s_t))) \tilde{Q}^\pi(\neg \langle s_t, a_t \rangle)) \\
&= (1 - P^\pi(s_t, a_t \mid \Phi(s_t)))( Q^\pi(s_t, a_t) - \tilde{Q}^\pi(\neg \langle s_t, a_t \rangle)).
\end{aligned}
\end{equation}
\normalsize
since
\begin{equation}
\begin{aligned}
\tilde{Q}^\pi(\neg \langle s_t, a_t \rangle)) &= \frac{\sum_{ s'_t, a'_t \neq s_t, a_t } P^\pi(s'_t, a'_t \mid \Phi(s_t)) Q^\pi(s'_t, a'_t)}{\sum_{ s'_t, a'_t \neq s_t, a_t } P^\pi(s'_t, a'_t \mid \Phi(s_t))} \\
&= \frac{\sum_{ s'_t, a'_t \neq s_t, a_t } P^\pi(s'_t, a'_t \mid \Phi(s_t))Q^\pi(s'_t, a'_t)}{1 - P^\pi(s_t, a_t \mid \Phi(s_t))}.
\end{aligned}
\end{equation}
\end{proof}
\section{Experimental results}\label{ap:experimental_results}
\subsection{Results using REINFORCE}\label{ap:reinforce}
Figure \ref{fig:qvsa_reinforce} shows the performance using REINFORCE in the training and evaluation variants of the three environments.
\begin{figure}[H]
% \vspace{-10pt}
\centering
% First plot
\begin{subfigure}[t]{0.32\textwidth}
\centering
\includegraphics[width=\textwidth]{figures/key2door_reinforce_plot.pdf}
\end{subfigure}
\hfill
% Second plot
\begin{subfigure}[t]{0.32\textwidth}
\centering
\includegraphics[width=\textwidth]{figures/tmaze_plot_reinforce.pdf}
\end{subfigure}
\hfill
% Third plot
\begin{subfigure}[t]{0.32\textwidth}
\centering
\includegraphics[width=\textwidth]{figures/diversion_plot_reinforce.pdf}
\end{subfigure}
\hfill
% \vspace{-10pt}
\caption{Performance of REINFORCE using Q-values and Advantages in both the training and evaluation variants of the Key2Door (first plot), Frozen T-Maze (second plot), and Diversion (third plot) environments.
}
\vspace{-10pt}
\label{fig:qvsa_reinforce}
\end{figure}
\subsection{Advantage normalization}\label{ap:advantage_normalization}
Figure \ref{fig:normalize_all} compares the performance of PPO with and without advantage normalization in the evaluation variants of the three environments.
\begin{figure}[H]
% \vspace{-10pt}
\centering
% First plot
\begin{subfigure}[t]{0.32\textwidth}
\centering
\includegraphics[width=\textwidth]{figures/keydoor_normalization.pdf}
\end{subfigure}
% \hfill
% Second plot
\begin{subfigure}[t]{0.32\textwidth}
\centering
\includegraphics[width=\textwidth]{figures/tmaze_normalization.pdf}
\end{subfigure}
% Third plot
\begin{subfigure}[t]{0.32\textwidth}
\centering
\includegraphics[width=\textwidth]{figures/diversion_normalization.pdf}
\end{subfigure}
\hfill
% \vspace{-10pt}
\caption{Performance of PPO with and without advantage normalization in the evaluation variants of the Key2Door (first plot), Frozen T-Maze (second plot), and Diversion (third plot) environments.}
\vspace{-10pt}
\label{fig:normalize_all}
\end{figure}
\subsection{Batch size}\label{ap:batch_size}
Figure \ref{fig:batch_size_all} compares the performance of PPO with different batch sizes in the evaluation variants of the three environments. While increasing the batch size seems to help in the Key2Door and Diversion environments, it has little effect in the Frozen T-Maze environment.
\begin{figure}[ht]
% \vspace{-10pt}
\centering
% First plot
\begin{subfigure}[t]{0.32\textwidth}
\centering
\includegraphics[width=\textwidth]{figures/keydoor_batchsize.pdf}
\end{subfigure}
\hfill
% Second plot
\begin{subfigure}[t]{0.32\textwidth}
\centering
\includegraphics[width=\textwidth]{figures/tmaze_batchsize.pdf}
\end{subfigure}
\hfill
\begin{subfigure}[t]{0.32\textwidth}
\centering
\includegraphics[width=\textwidth]{figures/diversion_batchsize.pdf}
\label{fig:plot2}
\end{subfigure}
% \vspace{-10pt}
\caption{Performance of PPO with different batch sizes in the evaluation variants of the Key2Door (first plot), Frozen T-Maze (second plot), and Diversion (third plot) environments.}
\vspace{-10pt}
\label{fig:batch_size_all}
\end{figure}
\subsection{Policy KL divergence}\label{ap:heatmaps}
\paragraph{Frozen T-Maze}
Figure \ref{fig:heatmaps_tmaze} shows the KL divergence of action probabilities in the Frozen T-Maze environment, produced by the policy when the signal is either purple or green, measured at different training steps (top: 10k steps, middle: 50k steps, bottom: 100k steps). To compute these divergences, we take the observation stack received by the agent and query the policy network twice: once with the original stack, and once with the signal bit flipped to the opposite value. The KL divergence at each cell is the average over 100 evaluation episodes.
The left heatmaps reveal that after 100k training steps, the agent trained on the $Q$-value largely ignores the signal, except at the starting location. In contrast, the agent trained on the advantage function conditions its action choices on the signal value throughout the maze.
Interestingly, the left heatmaps also reveal that as training progresses, the trajectories followed by the $Q$-value-trained agent become increasingly deterministic. By the end of training, the agent consistently chooses the top path when the signal is green and the bottom path when the signal is purple. In contrast, the agent trained on the advantage function continues to follow a diverse set of trajectories and does not exhibit a strong preference for any particular path. Note that multiple optimal paths exist for each signal value.
\begin{figure}[H]
% \vspace{-10pt}
\centering
% First plot
\begin{subfigure}[t]{0.45\textwidth}
\centering
\includegraphics[width=\textwidth]{figures/kl_heatmaps_tmaze_q.pdf}
\end{subfigure}
% \hfill
% Second plot
\begin{subfigure}[t]{0.45\textwidth}
\centering
\includegraphics[width=\textwidth]{figures/kl_heatmaps_tmaze_a.pdf}
\end{subfigure}
\vspace{-10pt}
\caption{KL divergence of action probabilities in the Frozen T-Maze environment at different training steps (top 10k steps, middle 50k steps, and bottom 100k steps) for agents trained on the $Q$-value (left) and the Advantage function (right).}
% \vspace{-10pt}
\label{fig:heatmaps_tmaze}
\end{figure}
\paragraph{Diversion}
Figure \ref{fig:KL_diversion} shows the policy KL divergence across each column in the Diversion environment, measured at different training steps (top: 3k steps, middle: 10k steps, bottom: 20k steps). For each column, the divergence is computed by comparing the action probabilities output by the agent’s policy when positioned in the top versus bottom row.
The left heatmaps indicate that the agent trained on the $Q$-value learns to ignore the bit indicating the row after 10k training steps. In contrast, the agent trained on the advantage function continues to use the row information when deciding which action to take.
\begin{figure}[H]
% \vspace{-10pt}
\centering
% First plot
\begin{subfigure}[t]{0.45\textwidth}
\centering
\includegraphics[width=\textwidth]{figures/kl_heatmaps_diversion_q.pdf}
\end{subfigure}
% \hfill
% Second plot
\begin{subfigure}[t]{0.45\textwidth}
\centering
\includegraphics[width=\textwidth]{figures/kl_heatmaps_diversion_a.pdf}
\end{subfigure}
% \vspace{-10pt}
\caption{KL divergence of action probabilities in the Diversion environment measured at different training steps (top 3k steps, middle 10k steps, and bottom 20k steps) for agents trained on the $Q$-value (left) and the Advantage function (right).}
% \vspace{-10pt}
\label{fig:KL_diversion}
\end{figure}
\section{Environments}\label{ap:environments}
\begin{figure}[H]
% \vspace{pt}
\centering
\begin{subfigure}[b]{0.33\textwidth}
\centering \includegraphics[width=\textwidth]{figures/frozen_tmaze.pdf}
% \vspace{0.8pt}
\end{subfigure}
\hspace{40pt}
\begin{subfigure}[b]{0.45\textwidth}
\centering
\includegraphics[width=0.95\textwidth]{figures/diversion.pdf}
\vspace{22pt}
\end{subfigure}
\caption{Illustrations of the Frozen T-Maze (left) and Diversion (right) environments.}
\label{fig:frozen_offtrack}
\end{figure}
\paragraph{Frozen T-Maze}
This environment is a variant of the popular T-Maze setting \citep{bakker2001reinforcement}. At the starting location, the agent receives a binary signal: either green or purple. Its task is to navigate to the right and reach the correct goal at the end of the maze. The agent receives a reward of $+1$ for reaching the green (purple) goal when the green (purple) signal was observed, and a penalty of $-1$ otherwise. Additionally, a time penalty of $-0.01$ is applied at each timestep to encourage the agent to take the most direct path to the goal.
At every timestep, the agent observes its current location within the maze, represented by a one-hot-encoded vector. However, the initial signal, represented by a binary variable, is only provided to the agent at the starting location. Crucially, the agent is capable of remembering past observations. When moving randomly, it must retain the initial signal throughout its trajectory. However, once it learns the shortest path to each goal (illustrated by the green and purple arrows), the agent can safely disregard the initial signal. This is because the agent can infer the signal based on its location: if it is on the green (purple) path, it must have received the green (purple) signal. Note that the two paths highlighted in Figure~\ref{fig:frozen_offtrack} are not the only optimal ones. However, for the agent to disregard the initial signal, the paths must not overlap.
We train agents in the original environment and evaluate them in a modified version where an icy surface (shown in blue) is introduced in the middle of the maze. This ice causes the agent to slip between the upper and lower cells.
\paragraph{Diversion}
In this environment, the agent must move from the start state to the goal state shown in Figure~\ref{fig:frozen_offtrack} (right). A reward of $+1$ is given for reaching the goal, and a penalty of $-1$ is incurred if the agent reaches the red cell instead. As in the other environments, there is a per-timestep penalty of $-0.01$. Observations are 8-dimensional binary vectors: the first 7 elements indicate the column where the agent is located, and the last element indicates the row.
After the agent learns the optimal policy (shown by the green arrow), it can ignore the last element of the observation vector. This is because the optimal policy never visits the bottom row. We train the agent in the original environment and evaluate it in a modified version containing a yellow diversion sign in the middle of the maze, which forces the agent to move to the bottom row.
\section{Experimental setup}\label{ap:experimental_setup}
The experiments were run on a laptop equipped with an Apple M2 Pro processor (12 cores) and 16 GB of RAM. Each run took less than 5 minutes and used at most 2\% of the total RAM.
Agents were trained using Stable Baselines3 \citep{stable-baselines3}. The hyperparameters are listed in Tables~\ref{tab:ppo} (PPO) and~\ref{tab:reinforce} (REINFORCE). For PPO, we adopted the hyperparameters used by \citet{schulman2017proximal} in their Atari experiments, except for the learning rate, which we increased to 1.0e-3 to accelerate convergence. We implemented a minimal version of the REINFORCE algorithm with only three hyperparameters (learning rate, discount factor, and entropy coefficient), for which we used the same values as in PPO. For the Frozen T-Maze, we used a stack of the past 30 observations as input, since solving the task requires memory.
\input{hyperparameters}
\end{document}
\begin{table}[!ht]
\centering
\caption{$Q$-values, advantages, and probabilities of key and no key when the agent is at location 6 under five different policies. The agent's state representation consists solely of the location variable.}
\resizebox{\textwidth}{!}{%
\begin{tabular}{lcccccccccc}
\toprule
& \multicolumn{4}{c}{\textbf{$Q$-value}}
& \multicolumn{4}{c}{\textbf{Advantage $\Phi(s) = l$}}
& \multicolumn{2}{c}{$P^\pi(X \mid L=6)$} \\
\cmidrule(r){2-5} \cmidrule(r){6-9} \cmidrule(r){10-11}
& \multicolumn{2}{c}{\textbf{No Key}} & \multicolumn{2}{c}{\textbf{Key}}
& \multicolumn{2}{c}{\textbf{No Key}} & \multicolumn{2}{c}{\textbf{Key}}
& \textbf{No Key} & \textbf{Key} \\
\cmidrule(r){2-3} \cmidrule(r){4-5} \cmidrule(r){6-7} \cmidrule(r){8-9}
& \textbf{Left} & \textbf{Right} & \textbf{Left} & \textbf{Right}
& \textbf{Left} & \textbf{Right} & \textbf{Left} & \textbf{Right}
& & \\
\midrule
$\pi(a^* \mid s) = 0.5$ & 0.038 & 0 & 0.662 & 1 & -0.657 & -0.695 & -0.033 & 0.305 & 0.168 & 0.832 \\
$\pi(a^* \mid s) = 0.6$ & 0.274 & 0 & 0.839 & 1 & -0.608 & -0.882 & -0.043 & 0.118 & 0.069 & 0.931 \\
$\pi(a^* \mid s) = 0.7$ & 0.504 & 0 & 0.905 & 1 & -0.456 & -0.960 & -0.055 & 0.040 & 0.018 & 0.982 \\
$\pi(a^* \mid s) = 0.8$ & 0.664 & 0 & 0.934 & 1 & -0.321 & -0.985 & -0.051 & 0.015 & 0.003 & 0.997 \\
$\pi(a^* \mid s) = 0.9$ & 0.759 & 0 & 0.950 & 1 & -0.235 & -0.994 & -0.044 & 0.006 & 0.000 & 0.999 \\
\bottomrule
\end{tabular}
}
\vspace{-5pt}
\label{table:tab1}
\end{table}
\begin{center}
%
\vspace{-10pt}
\begin{table}[h]
\centering
\caption{PPO hyperparameters.}
\vspace{5pt}
\resizebox{0.7\textwidth}{!}{
\begin{tabular}{ p{5cm}|p{4cm}}
% \hline
% \multicolumn{2}{|c|}{PPO} \\
% \hline
% \hline
Rollout steps & 128 \\
Batch size & 32 \\
Learning rate & 1.0e-3 \\
Number epoch & 3 \\
Discount $\gamma$ & 0.99 \\
GAE $\lambda$ & 0.95 \\
Entropy coefficient & 1.0e-2 \\
Clip range & 0.1 \\
Value coefficient & 1 \\
Number Neurons 1st layer & 128 \\
Number Neurons 2nd layer & 128 \\
% \hline
\end{tabular}
}
\vspace{-20pt}
\label{tab:ppo}
\end{table}
\end{center}
\begin{center}
%
\vspace{-10pt}
\begin{table}[ht]
\centering
\caption{REINFORCE hyperparameters.}
\vspace{5pt}
\resizebox{0.7\textwidth}{!}{
\begin{tabular}{ p{5cm}|p{4cm}}
% \hline
% \multicolumn{2}{|c|}{PPO} \\
% \hline
% \hline
Learning rate & 1.0e-3 \\
Discount $\gamma$ & 0.99 \\
Entropy coefficient & 1.0e-2 \\
Number Neurons 1st layer & 128 \\
Number Neurons 2nd layer & 128 \\
% \hline
\end{tabular}
}
\vspace{-20pt}
\label{tab:reinforce}
\end{table}
\end{center}
\providecommand{\noopsort}[1]{}\providecommand{\noopsort}[1]{}\providecommand{\noopsort}[1]{}\providecommand{\noopsort}[1]{}
\begin{thebibliography}{}
\bibitem[Baird, 1994]{Baird94ICoNN}
Baird, L.~C. (1994).
\newblock Reinforcement learning in continuous time: advantage updating.
\newblock In {\em Neural Networks, 1994. IEEE World Congress on Computational Intelligence., 1994 IEEE International Conference on}, volume~4, pages 2448--2453 vol.4.
\bibitem[Bakker, 2001]{bakker2001reinforcement}
Bakker, B. (2001).
\newblock Reinforcement learning with long short-term memory.
\newblock {\em Advances in neural information processing systems}, 14.
\bibitem[Boutilier et~al., 1999]{boutilier1999decision}
Boutilier, C., Dean, T., and Hanks, S. (1999).
\newblock {Decision-theoretic planning: Structural assumptions and computational leverage}.
\newblock {\em Journal of Artificial Intelligence Research}, 11:1--94.
\bibitem[Chung et~al., 2021]{chung2021beyond}
Chung, W., Thomas, V., Machado, M.~C., and Le~Roux, N. (2021).
\newblock Beyond variance reduction: Understanding the true impact of baselines on policy optimization.
\newblock In {\em International Conference on Machine Learning}, pages 1999--2009. PMLR.
\bibitem[Corcoll and Vicente, 2022]{corcoll2022disentangling}
Corcoll, O. and Vicente, R. (2022).
\newblock Disentangling controlled effects for hierarchical reinforcement learning.
\newblock In {\em Conference on Causal Learning and Reasoning}, pages 178--200. PMLR.
\bibitem[De~Haan et~al., 2019]{de2019causal}
De~Haan, P., Jayaraman, D., and Levine, S. (2019).
\newblock Causal confusion in imitation learning.
\newblock {\em Advances in neural information processing systems}, 32.
\bibitem[Ding et~al., 2023]{ding2023seeing}
Ding, W., Shi, L., Chi, Y., and Zhao, D. (2023).
\newblock Seeing is not believing: Robust reinforcement learning against spurious correlation.
\newblock {\em Advances in Neural Information Processing Systems}, 36:66328--66363.
\bibitem[Fran{\c{c}}ois-Lavet et~al., 2018]{franccois2018introduction}
Fran{\c{c}}ois-Lavet, V., Henderson, P., Islam, R., Bellemare, M.~G., and Pineau, J. (2018).
\newblock An introduction to deep reinforcement learning.
\newblock {\em Foundations and Trends in Machine Learning}.
\bibitem[Greensmith et~al., 2001]{greensmith2001variance}
Greensmith, E., Bartlett, P., and Baxter, J. (2001).
\newblock Variance reduction techniques for gradient estimates in reinforcement learning.
\newblock {\em Advances in Neural Information Processing Systems}, 14.
\bibitem[He et~al., 2024]{he2024model}
He, J., Moerland, T.~M., de~Vries, J.~A., and Oliehoek, F.~A. (2024).
\newblock What model does muzero learn?
\newblock {\em 27th European Conference on Artificial Intelligence}.
\bibitem[Higgins et~al., 2017]{higgins2017darla}
Higgins, I., Pal, A., Rusu, A., Matthey, L., Burgess, C., Pritzel, A., Botvinick, M., Blundell, C., and Lerchner, A. (2017).
\newblock Darla: Improving zero-shot transfer in reinforcement learning.
\newblock In {\em International Conference on Machine Learning}, pages 1480--1490. PMLR.
\bibitem[Kirk et~al., 2023]{kirk2023survey}
Kirk, R., Zhang, A., Grefenstette, E., and Rockt{\"a}schel, T. (2023).
\newblock A survey of zero-shot generalisation in deep reinforcement learning.
\newblock {\em Journal of Artificial Intelligence Research}, 76:201--264.
\bibitem[Lan et~al., 2023]{lan2023can}
Lan, L.-C., Zhang, H., and Hsieh, C.-J. (2023).
\newblock Can agents run relay race with strangers? generalization of {RL} to out-of-distribution trajectories.
\newblock In {\em The Eleventh International Conference on Learning Representations}.
\bibitem[Machado et~al., 2018]{machado2018revisiting}
Machado, M.~C., Bellemare, M.~G., Talvitie, E., Veness, J., Hausknecht, M., and Bowling, M. (2018).
\newblock Revisiting the arcade learning environment: Evaluation protocols and open problems for general agents.
\newblock {\em Journal of Artificial Intelligence Research}, 61:523--562.
\bibitem[Mandlekar et~al., 2017]{mandlekar2017adversarially}
Mandlekar, A., Zhu, Y., Garg, A., Fei-Fei, L., and Savarese, S. (2017).
\newblock Adversarially robust policy learning: Active construction of physically-plausible perturbations.
\newblock In {\em 2017 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)}, pages 3932--3939. IEEE.
\bibitem[Marbach and Tsitsiklis, 1999]{marbach1999simulation}
Marbach, P. and Tsitsiklis, J.~N. (1999).
\newblock Simulation-based optimization of markov reward processes: Implementation issues.
\newblock In {\em Proceedings of the 38th IEEE Conference on Decision and Control (Cat. No. 99CH36304)}, volume~2, pages 1769--1774. IEEE.
\bibitem[McCallum, 1995]{McCallum95PhD}
McCallum, A.~K. (1995).
\newblock {\em Reinforcement Learning with Selective Perception and Hidden State}.
\newblock PhD thesis, University of Rochester.
\bibitem[Mei et~al., 2022]{mei2022role}
Mei, J., Chung, W., Thomas, V., Dai, B., Szepesvari, C., and Schuurmans, D. (2022).
\newblock The role of baselines in policy gradient optimization.
\newblock {\em Advances in Neural Information Processing Systems}, 35:17818--17830.
\bibitem[Pan et~al., 2022]{pan2022direct}
Pan, H.-R., G{\"u}rtler, N., Neitz, A., and Sch{\"o}lkopf, B. (2022).
\newblock Direct advantage estimation.
\newblock {\em Advances in Neural Information Processing Systems}, 35:11869--11880.
\bibitem[Pan and Sch{\"o}lkopf, 2024]{pan2024skill}
Pan, H.-R. and Sch{\"o}lkopf, B. (2024).
\newblock Skill or luck? return decomposition via advantage functions.
\newblock In {\em The Twelfth International Conference on Learning Representations}.
\bibitem[Pearl et~al., 2016]{pearl2016causal}
Pearl, J., Glymour, M., and Jewell, N.~P. (2016).
\newblock {\em Causal inference in statistics: A primer}.
\newblock John Wiley \& Sons.
\bibitem[Raffin et~al., 2021]{stable-baselines3}
Raffin, A., Hill, A., Gleave, A., Kanervisto, A., Ernestus, M., and Dormann, N. (2021).
\newblock Stable-baselines3: Reliable reinforcement learning implementations.
\newblock {\em Journal of Machine Learning Research}, 22(268):1--8.
\bibitem[Raileanu and Fergus, 2021]{raileanu2021decoupling}
Raileanu, R. and Fergus, R. (2021).
\newblock Decoupling value and policy for generalization in reinforcement learning.
\newblock In {\em International Conference on Machine Learning}, pages 8787--8798. PMLR.
\bibitem[Schulman et~al., 2017]{schulman2017proximal}
Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. (2017).
\newblock Proximal policy optimization algorithms.
\newblock {\em arXiv preprint arXiv:1707.06347}.
\bibitem[Song et~al., 2020]{Song2020Observational}
Song, X., Jiang, Y., Tu, S., Du, Y., and Neyshabur, B. (2020).
\newblock Observational overfitting in reinforcement learning.
\newblock In {\em International Conference on Learning Representations}.
\bibitem[Suau et~al., 2024]{suau2024bad}
Suau, M., Spaan, M. T.~J., and Oliehoek, F.~A. (2024).
\newblock Bad habits: Policy confounding and out-of-trajectory generalization in {RL}.
\newblock In {\em Reinforcement Learning Conference}.
\bibitem[Sutton and Barto, 2018]{sutton2018reinforcement}
Sutton, R.~S. and Barto, A.~G. (2018).
\newblock {\em Reinforcement learning: An introduction}.
\newblock MIT press.
\bibitem[Sutton et~al., 1999]{sutton1999policy}
Sutton, R.~S., McAllester, D., Singh, S., and Mansour, Y. (1999).
\newblock Policy gradient methods for reinforcement learning with function approximation.
\newblock {\em Advances in neural information processing systems}, 12.
\bibitem[Taylor and Parr, 2009]{Taylor09ICML}
Taylor, G. and Parr, R. (2009).
\newblock Kernelized value function approximation for reinforcement learning.
\newblock In {\em Proc.\ of the Twenty-Sixth International Conference on Machine learning}, page 128.
\bibitem[Tennenholtz et~al., 2021]{tennenholtz2021covariate}
Tennenholtz, G., Hallak, A., Dalal, G., Mannor, S., Chechik, G., and Shalit, U. (2021).
\newblock On covariate shift of latent confounders in imitation and reinforcement learning.
\newblock {\em arXiv preprint arXiv:2110.06539}.
\bibitem[Weltevrede et~al., 2024]{weltevrede2024training}
Weltevrede, M., Horsch, C., Spaan, M.~T., and B{\"o}hmer, W. (2024).
\newblock Training on more reachable tasks for generalisation in reinforcement learning.
\newblock {\em arXiv preprint arXiv:2410.03565}.
\bibitem[Williams, 1992]{Williams92ML}
Williams, R.~J. (1992).
\newblock Simple statistical gradient-following algorithms for connectionist reinforcement learning.
\newblock {\em Machine Learning}, 8(3-4):229--256.
\bibitem[Zhang et~al., 2020a]{zhang2020invariant}
Zhang, A., Lyle, C., Sodhani, S., Filos, A., Kwiatkowska, M., Pineau, J., Gal, Y., and Precup, D. (2020a).
\newblock Invariant causal prediction for block mdps.
\newblock In {\em International Conference on Machine Learning}, pages 11214--11224. PMLR.
\bibitem[Zhang et~al., 2020b]{zhang2020causal}
Zhang, J., Kumor, D., and Bareinboim, E. (2020b).
\newblock Causal imitation learning with unobserved confounders.
\newblock {\em Advances in neural information processing systems}, 33:12263--12274.
\end{thebibliography}