REVIEW 3 major objections 3 minor 34 references
Breaking Habits: On the Role of the Advantage Function in Learning Causal State Representations
T0 review · 3 major / 3 minor · reviewed 2026-08-07 · deepseek-v4-flash
Pith's one-line read This paper claims that the advantage function rescales policy-gradient updates by the complement of a state-action pair's probability under the policy, and that this rescaling mitigates policy confounding and improves out-of-trajectory…
desk verdict A clean algebraic identity about advantages under coarse baselines, with an honest empirical story whose causal interpretation needs stronger tests. read the letter →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
What carries the argument
The central object is the advantage function under a coarse state representation, $A^\pi_\Phi(s_t,a_t)=Q^\pi(s_t,a_t)-V^\pi(\Phi(s_t))$, in which $Q^\pi$ sees the full state while $V^\pi$ is aggregated over the equivalence class $\{s_t\}_\Phi$. This asymmetry is deliberate: the baseline marginalizes away exactly the variables that $\Phi$ discards, so the subtraction measures how much a particular full state-action pair deviates from its represented class. Theorem 1 converts that deviation into a product of a probability-complement weight and a contrast term, and that product is what reweights stochastic gradient samples during training. Corollary 1 then shows that for causal representations the same mechanism reduces to reweighting by the complement of the action probability.
What would settle it
Run the Key2Door evaluation variant with PPO using advantages and no normalization; the paper predicts strong out-of-trajectory performance. Then run the identical training with advantages normalized within each batch; the paper predicts a performance collapse. If normalization produces little or no drop, or if unnormalized advantages offer no benefit over Q-values, the scaling effect is not the operative mechanism.
Extended reading notes
Core claim
On the paper's own terms, the discovery is a decomposition identity for the advantage under a state representation. For any $\Phi$, defining $A^\pi_\Phi(s_t,a_t)=Q^\pi(s_t,a_t)-V^\pi(\Phi(s_t))$ yields $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_\Phi(\neg\langle s_t,a_t\rangle))$, where $\tilde Q^\pi_\Phi(\neg\langle s_t,a_t\rangle)$ is the probability-weighted average Q-value over all other state-action pairs in the same equivalence class. Since stochastic gradient estimates sample pairs with frequency $P^\pi(s_t,a_t\mid\Phi(s_t))$, the complement factor systematically downweights over-sampled pairs and boosts under-sampled ones. The paper argues that this breaks the feedback loop in which a policy makes a pair frequent, the frequent pair dominates gradient updates, and the agent accordingly discards causal variables that merely correlate with the frequent route. When $\Phi$ is a causal (Markov) representation, Corollary 1 simplifies the factor to $1-\pi(a_t\mid\Phi(s_t))$, so only action probabilities need reweighting once the representation already captures the true causal factors.
Load-bearing premise
The argument stands or falls on the premise that $Q^\pi(s_t,a_t)$ is estimated from the full state while the baseline $V^\pi(\Phi(s_t))$ is computed from the coarse representation; the paper's limitations section concedes that if $Q$ is learned entirely through function approximation it too is subject to $\Phi$, and the stated theorem no longer holds.
Editorial extensions
If this is right
- Agents trained on raw Q-values can lock onto spurious correlations and fail when forced off their usual trajectories; advantage-trained agents generalize well in all three environments.
- Standard advantage normalization, as used in many PPO implementations, removes the per-sample scaling and reproduces the Q-value failure mode, making implementation details central to whether causal representations are learned.
- Increasing batch size partly mitigates Q-based training failures, but the paper finds no reason to prefer large batches over simply training on advantages.
- If the state representation is already causal, the advantage reweighting reduces to an action-probability factor, so the method does not need to correct for state frequencies.
Reading between the lines
- Beyond the paper: the complement-probability factor from Theorem 1 could be applied explicitly as a per-sample weight in Q-based or off-policy training, which would test whether the scaling effect transfers to settings without a separate advantage estimator.
- Beyond the paper: the same identity suggests a diagnostic for representation quality, comparing $V^\pi(\Phi(s))$ with the class-aggregated Q-values; a large gap would flag variables that $\Phi$ has collapsed but that still carry causal information.
- Beyond the paper: if the mechanism scales to high-dimensional domains, then RL libraries that normalize advantages by default may be trading causal reliability for numerical stability, and a simple reparameterization could preserve both.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper argues that the advantage function, defined as A^π_Φ(s,a) := Q^π(s,a) − V^π(Φ(s)) where V^π(Φ(s)) averages the value over the Φ-equivalence class of s, downweights state-action pairs that are likely under the current policy by a factor 1 − P^π(s,a|Φ(s)), and that this scaling mitigates policy confounding and improves out-of-trajectory generalization. The main theoretical result (Theorem 1) gives an exact algebraic identity for A^π_Φ. The paper then presents experiments in three gridworlds comparing training with advantages versus Q-values, ablations of advantage normalization and batch size, and an analysis of learned policies via KL-divergence heatmaps.
Significance. The algebraic identity in Theorem 1 is correct and is proved cleanly in Appendix A, and it is a straightforward consequence of Definition 6 (the definition of V(Φ) as a conditional expectation). The paper is honest about its scope and explicitly states in Section 7 that the result does not generally extend when Q is a learned function approximator. The experiments are described with enough detail to be reproducible, and the observation that advantage normalization removes the out-of-trajectory benefit is an interesting empirical finding. However, the theoretical mechanism relies on an asymmetry between Q (full-state) and V (coarse representation) that is not instantiated in the reported experiments, and the causal claim is only supported heuristically. If the authors can resolve this gap—either by instantiating the coarse Φ in experiments or by extending the theory—the paper could make a useful contribution to understanding representation learning in policy gradient methods.
major comments (3)
- [Section 6.1 and Section 7] The experiments do not implement the actor-critic asymmetry required by Theorem 1. The theorem defines A_Φ(s,a)=Q(s,a)−V(Φ(s)) with Q evaluated on the full state and V averaging over the Φ-equivalence class; the paper's own Limitations section (Section 7) concedes that if Q is learned entirely via function approximation then 'Theorem 1 would no longer hold in its current form.' In the reported PPO and REINFORCE experiments, both the policy and the value network take the same observation stack as input, no explicit coarse representation Φ is constructed, and the Q-values used for the Q-value condition are Monte Carlo returns, not a separate full-state Q-function. Consequently, the empirical difference between the advantage and Q-value conditions cannot be attributed to the 1 − P scaling of Theorem 1 rather than to generic variance reduction or clipping. The paper should either modify the experiments to explicitly use a coarse Φ (e.g., define Φ by a hand-selected subset of observation variables and train V(Φ(s)) as the average over the class), or extend the theoretical analysis to function approximation.
- [Section 5.2 (Theorem 1)] The 'scaling effect' is a direct algebraic consequence of Definition 6: V^π(Φ(s)) is defined as the conditional expectation of Q over the equivalence class, so A = Q − E[Q] = (1−p)(Q − E[Q | excluded]) is tautological given the definition. The paper should acknowledge more explicitly that this is not an intrinsic property of the standard advantage function A(s,a)=Q(s,a)−V(s), but of the specific coarse-grained baseline V(Φ(s)). Furthermore, the causal claim in Section 5.3—that this scaling breaks spurious correlations and encourages focus on causal factors—is heuristic; no formal result links the algebraic scaling to the behavior of stochastic gradient ascent or to representation learning. Please either add a formal statement (under simplifying assumptions) or clearly label this connection as a conjecture supported by the experiments.
- [Section 6 (Figures 2–3)] The comparison does not isolate the proposed mechanism. Standard advantage baselines are known to reduce variance and often improve performance relative to pure Q-value (return) training; the ablation with advantage normalization (Figure 3, left) is suggestive but only shows that normalization hurts out-of-trajectory performance, not that the 1−P scaling, rather than any value baseline, is the cause. To support the paper's central claim, the experiments should include a condition with the standard full-state baseline A(s,a)=Q(s,a)−V(s) in addition to A_Φ(s,a)=Q(s,a)−V(Φ(s)) for an explicit Φ, and show that the additional coarse-representation baseline changes the learned representation and out-of-trajectory generalization in the predicted direction.
minor comments (3)
- [Section 7] There is a typo: 'represantion' should be 'representation' in the first paragraph of the Limitations section.
- [Section 5.2] The notation for the alternative-state-action baseline is inconsistent: Theorem 1 and the surrounding text write \tilde{Q}^\pi(\neg\langle s_t,a_t\rangle) without a subscript \Phi, while the proof and Definition 6 use \tilde{Q}^\pi_\Phi. Please use one consistent notation.
- [Definition 5] The notation R^\pi(\Phi(s_t), a_t) and P^\pi(\Phi(s_{t+1}) | \Phi(s_t), a_t) is not formally defined; it would help to state explicitly that these denote expected reward and transition probabilities under π given the coarse state, to avoid confusion with the reward function R defined on full states.
Circularity Check
No significant circularity: Theorem 1 is an algebraic consequence of the paper's own definitions, and no fitted input, self-citation chain, or imported uniqueness theorem is load-bearing.
full rationale
The derivation chain is self-contained. Definition 6 defines Vπ(Φ(st)) as a policy-weighted average over the Φ-equivalence class, and Definition 7 defines AπΦ = Qπ(st, at) − Vπ(Φ(st)). The proof of Theorem 1 (Appendix A) expands Vπ(Φ(st)) as Σ Pπ(s', a' | Φ) Qπ(s', a'), isolates the ⟨st, at⟩ term, and obtains (1 − Pπ(st, at | Φ(st)))(Qπ(st, at) − Q̃π). This is a direct algebraic consequence of the authors' definitions, not a fitted parameter, a renamed empirical pattern, or a result imported from prior work. The interpretive claim that this scaling 'downweights state-action pairs that are more likely under the current policy' is a restatement of the same identity and is not circular in the sense of defining the conclusion into the premises. The paper cites Suau et al. (2024) for the prior notion of policy confounding and for the test environments, but the central theorem and its proof do not depend on that citation, so this is not a load-bearing self-citation. The Limitations section (Sec. 7) concedes that Theorem 1 no longer holds if Q is learned entirely via function approximation; that is a domain-mismatch or validity concern about whether the experiments instantiate the theoretical assumptions, not a circularity in the derivation. Overall, no step reduces a prediction to an input by construction.
Assumptions & free parameters
assumptions (3)
- domain assumption The Markov chain induced by policy π on S admits a stationary distribution dπ.
- ad hoc to paper Qπ(st, at) is estimated with respect to the full state st, while Vπ(Φ(st)) uses only the coarse representation Φ.
- domain assumption A Markov state representation is equated with a Causal State Representation that includes all true causal factors.
Cite this review
Pith. "Pith review of Breaking Habits: On the Role of the Advantage Function in Learning Causal State Representations." pith.science (2026). https://pith.science/paper/4454ABHN
@misc{pith2026250611912,
author = {Pith},
title = {Pith review of: Breaking Habits: On the Role of the Advantage Function in Learning Causal State Representations},
year = {2026},
howpublished = {\url{https://pith.science/paper/4454ABHN}},
note = {Machine review of arXiv:2506.11912}
}
read the original 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.
Figures
Figures from the paper (7 more)
Reference graph
Works this paper leans on
-
[1]
Baird, L. C. (1994). Reinforcement learning in continuous time: advantage updating. In Neural Networks, 1994. IEEE World Congress on Computational Intelligence., 1994 IEEE International Conference on , volume 4, pages 2448--2453 vol.4
work page 1994
-
[2]
Bakker, B. (2001). Reinforcement learning with long short-term memory. Advances in neural information processing systems , 14
work page 2001
-
[3]
Boutilier, C., Dean, T., and Hanks, S. (1999). Decision-theoretic planning: Structural assumptions and computational leverage . Journal of Artificial Intelligence Research , 11:1--94
work page 1999
-
[4]
Chung, W., Thomas, V., Machado, M. C., and Le Roux, N. (2021). Beyond variance reduction: Understanding the true impact of baselines on policy optimization. In International Conference on Machine Learning , pages 1999--2009. PMLR
work page 2021
-
[5]
Corcoll, O. and Vicente, R. (2022). Disentangling controlled effects for hierarchical reinforcement learning. In Conference on Causal Learning and Reasoning , pages 178--200. PMLR
work page 2022
-
[6]
De Haan, P., Jayaraman, D., and Levine, S. (2019). Causal confusion in imitation learning. Advances in neural information processing systems , 32
work page 2019
-
[7]
Ding, W., Shi, L., Chi, Y., and Zhao, D. (2023). Seeing is not believing: Robust reinforcement learning against spurious correlation. Advances in Neural Information Processing Systems , 36:66328--66363
work page 2023
-
[8]
Fran c ois-Lavet, V., Henderson, P., Islam, R., Bellemare, M. G., and Pineau, J. (2018). An introduction to deep reinforcement learning. Foundations and Trends in Machine Learning
work page 2018
Show all 34 references
-
[9]
Greensmith, E., Bartlett, P., and Baxter, J. (2001). Variance reduction techniques for gradient estimates in reinforcement learning. Advances in Neural Information Processing Systems , 14
2001
-
[10]
M., de Vries, J
He, J., Moerland, T. M., de Vries, J. A., and Oliehoek, F. A. (2024). What model does muzero learn? 27th European Conference on Artificial Intelligence
2024
-
[11]
Higgins, I., Pal, A., Rusu, A., Matthey, L., Burgess, C., Pritzel, A., Botvinick, M., Blundell, C., and Lerchner, A. (2017). Darla: Improving zero-shot transfer in reinforcement learning. In International Conference on Machine Learning , pages 1480--1490. PMLR
2017
-
[12]
Kirk, R., Zhang, A., Grefenstette, E., and Rockt \"a schel, T. (2023). A survey of zero-shot generalisation in deep reinforcement learning. Journal of Artificial Intelligence Research , 76:201--264
2023
-
[13]
Lan, L.-C., Zhang, H., and Hsieh, C.-J. (2023). Can agents run relay race with strangers? generalization of RL to out-of-distribution trajectories. In The Eleventh International Conference on Learning Representations
2023
-
[14]
C., Bellemare, M
Machado, M. C., Bellemare, M. G., Talvitie, E., Veness, J., Hausknecht, M., and Bowling, M. (2018). Revisiting the arcade learning environment: Evaluation protocols and open problems for general agents. Journal of Artificial Intelligence Research , 61:523--562
2018
-
[15]
Mandlekar, A., Zhu, Y., Garg, A., Fei-Fei, L., and Savarese, S. (2017). Adversarially robust policy learning: Active construction of physically-plausible perturbations. In 2017 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS) , pages 3932--3939. IEEE
2017
-
[16]
and Tsitsiklis, J
Marbach, P. and Tsitsiklis, J. N. (1999). Simulation-based optimization of markov reward processes: Implementation issues. In Proceedings of the 38th IEEE Conference on Decision and Control (Cat. No. 99CH36304) , volume 2, pages 1769--1774. IEEE
1999
-
[17]
McCallum, A. K. (1995). Reinforcement Learning with Selective Perception and Hidden State . PhD thesis, University of Rochester
1995
-
[18]
Mei, J., Chung, W., Thomas, V., Dai, B., Szepesvari, C., and Schuurmans, D. (2022). The role of baselines in policy gradient optimization. Advances in Neural Information Processing Systems , 35:17818--17830
2022
-
[19]
u rtler, N., Neitz, A., and Sch \
Pan, H.-R., G \"u rtler, N., Neitz, A., and Sch \"o lkopf, B. (2022). Direct advantage estimation. Advances in Neural Information Processing Systems , 35:11869--11880
2022
-
[20]
and Sch \"o lkopf, B
Pan, H.-R. and Sch \"o lkopf, B. (2024). Skill or luck? return decomposition via advantage functions. In The Twelfth International Conference on Learning Representations
2024
-
[21]
Pearl, J., Glymour, M., and Jewell, N. P. (2016). Causal inference in statistics: A primer . John Wiley & Sons
2016
-
[22]
Raffin, A., Hill, A., Gleave, A., Kanervisto, A., Ernestus, M., and Dormann, N. (2021). Stable-baselines3: Reliable reinforcement learning implementations. Journal of Machine Learning Research , 22(268):1--8
2021
-
[23]
and Fergus, R
Raileanu, R. and Fergus, R. (2021). Decoupling value and policy for generalization in reinforcement learning. In International Conference on Machine Learning , pages 8787--8798. PMLR
2021
-
[24]
Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. (2017). Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347
2017 arXiv
-
[25]
Song, X., Jiang, Y., Tu, S., Du, Y., and Neyshabur, B. (2020). Observational overfitting in reinforcement learning. In International Conference on Learning Representations
2020
-
[26]
Suau, M., Spaan, M. T. J., and Oliehoek, F. A. (2024). Bad habits: Policy confounding and out-of-trajectory generalization in RL . In Reinforcement Learning Conference
2024
-
[27]
Sutton, R. S. and Barto, A. G. (2018). Reinforcement learning: An introduction . MIT press
2018
-
[28]
S., McAllester, D., Singh, S., and Mansour, Y
Sutton, R. S., McAllester, D., Singh, S., and Mansour, Y. (1999). Policy gradient methods for reinforcement learning with function approximation. Advances in neural information processing systems , 12
1999
-
[29]
and Parr, R
Taylor, G. and Parr, R. (2009). Kernelized value function approximation for reinforcement learning. In Proc.\ of the Twenty-Sixth International Conference on Machine learning , page 128
2009
-
[30]
Tennenholtz, G., Hallak, A., Dalal, G., Mannor, S., Chechik, G., and Shalit, U. (2021). On covariate shift of latent confounders in imitation and reinforcement learning. arXiv preprint arXiv:2110.06539
2021 arXiv
-
[31]
T., and B \"o hmer, W
Weltevrede, M., Horsch, C., Spaan, M. T., and B \"o hmer, W. (2024). Training on more reachable tasks for generalisation in reinforcement learning. arXiv preprint arXiv:2410.03565
2024 arXiv
-
[32]
Williams, R. J. (1992). Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine Learning , 8(3-4):229--256
1992
-
[33]
Zhang, A., Lyle, C., Sodhani, S., Filos, A., Kwiatkowska, M., Pineau, J., Gal, Y., and Precup, D. (2020a). Invariant causal prediction for block mdps. In International Conference on Machine Learning , pages 11214--11224. PMLR
2020
-
[34]
Zhang, J., Kumor, D., and Bareinboim, E. (2020b). Causal imitation learning with unobserved confounders. Advances in neural information processing systems , 33:12263--12274
2020
Reviewed August 7, 2026 · model on record in the stance chip above.
Discussion (0). Continue with ORCID to comment.