Pith. sign in

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 →

arxiv 2506.11912 v1 pith:4454ABHN submitted 2025-06-13 cs.LG cs.AI

classification cs.LGcs.AI
keywords policyconfoundingadvantagefunctioncausalstaterepresentationout-of-trajectorygeneralizationgradientsspuriouscorrelationsabstractionreinforcementlearning
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

The paper aims to establish that the advantage function, the usual Q-minus-baseline correction in policy-gradient methods, does more than reduce gradient variance. The author proves that when the action value $Q^\pi(s_t,a_t)$ is evaluated on the full state and the baseline $V^\pi(\Phi(s_t))$ is averaged over the equivalence class induced by a coarse representation $\Phi$, the advantage equals a factor $1-P^\pi(s_t,a_t\mid\Phi(s_t))$ times a contrast term. That factor rescales stochastic policy-gradient updates, shrinking the influence of state-action pairs the policy already visits often and amplifying rare, informative ones, which is exactly what breaks the spurious correlations that the policy itself creates. In three gridworlds, agents trained with advantages generalize to trajectories they never followed, whereas agents trained on raw Q-values fail, and normalizing the advantages removes the benefit. If correct, this turns a standard variance-reduction trick into a mechanism for learning causal rather than habitual state representations.

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.

Watch

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

Editorial extensions of the paper, not claims the author makes directly.

  • 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.
Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

3 major / 3 minor

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)
  1. [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.
  2. [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.
  3. [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)
  1. [Section 7] There is a typo: 'represantion' should be 'representation' in the first paragraph of the Limitations section.
  2. [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.
  3. [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

0 steps flagged · score 0.0 of 10

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 0 free parameters · 3 assumptions · 0 invented entities

No free parameters are fitted in the derivation; the advantage scaling identity is purely algebraic. The experiments use standard hyperparameters but these are not fitted to the OOT claim. The main assumptions are the stationarity of the policy-induced Markov chain, the full-state/coarse-baseline asymmetry, and the identification of Markov representations with causal ones.

assumptions (3)
  • domain assumption The Markov chain induced by policy π on S admits a stationary distribution dπ.
    Definition 6 defines Vπ(Φ(st)) and Qπ(Φ(st), at) as weighted averages over the equivalence class {st}Φ using dπ, so existence is required.
  • ad hoc to paper Qπ(st, at) is estimated with respect to the full state st, while Vπ(Φ(st)) uses only the coarse representation Φ.
    This asymmetry underlies Theorem 1; the author states in Section 7 that Theorem 1 no longer holds if Q is also learned via function approximation subject to Φ. The experiments do not verify this assumption.
  • domain assumption A Markov state representation is equated with a Causal State Representation that includes all true causal factors.
    Used to argue that overcoming policy confounding recovers causal factors; it is a definitional claim, not a proven theorem.

how reviews work

0 comments
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 reproduced from arXiv: 2506.11912 by the authors.

Figure 1
Figure 1. 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 ∈ {0, 1, .., 6} and a binary variable X ∈ {0, 1} indicat￾ing whether or not it has the key. At all states, the agent has two possible actions: mov￾ing left A = 0, or moving right A = 1. The agent receives a rew… view at source ↗
Figure 2
Figure 2. Performance of PPO using Q-values and Advantages in both the training and evaluation [PITH_FULL_IMAGE:figures/full_fig_p008_2.png] view at source ↗
Figure 3
Figure 3. Left: Performance of PPO with and without advantage normalization in the Key2Door evaluation environment. Right: Performance of PPO with different batch sizes in the Key2Door evaluation environment [PITH_FULL_IMAGE:figures/full_fig_p008_3.png] view at source ↗
Figures from the paper (7 more)
Figure 4
Figure 4. Figure 4: KL divergence of action probabilities with and without [PITH_FULL_IMAGE:figures/full_fig_p009_4.png]
Figure 5
Figure 5. Figure 5: Performance of REINFORCE using Q-values and Advantages in both the training and [PITH_FULL_IMAGE:figures/full_fig_p012_5.png]
Figure 6
Figure 6. Figure 6: Performance of PPO with and without advantage normalization in the evaluation variants [PITH_FULL_IMAGE:figures/full_fig_p013_6.png]
Figure 7
Figure 7. Figure 7: Performance of PPO with different batch sizes in the evaluation variants of the Key2Door [PITH_FULL_IMAGE:figures/full_fig_p013_7.png]
Figure 8
Figure 8. Figure 8: KL divergence of action probabilities in the Frozen T-Maze environment at different training [PITH_FULL_IMAGE:figures/full_fig_p014_8.png]
Figure 9
Figure 9. Figure 9: KL divergence of action probabilities in the Diversion environment measured at different [PITH_FULL_IMAGE:figures/full_fig_p014_9.png]
Figure 10
Figure 10. Figure 10: Illustrations of the Frozen T-Maze (left) and Diversion (right) environments. [PITH_FULL_IMAGE:figures/full_fig_p015_10.png]

Discussion (0). Sign in to comment.

Reference graph

Works this paper leans on

34 extracted references · 29 canonical work pages

  1. [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

  2. [2]

    Bakker, B. (2001). Reinforcement learning with long short-term memory. Advances in neural information processing systems , 14

  3. [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

  4. [4]

    C., and Le Roux, N

    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

  5. [5]

    and Vicente, R

    Corcoll, O. and Vicente, R. (2022). Disentangling controlled effects for hierarchical reinforcement learning. In Conference on Causal Learning and Reasoning , pages 178--200. PMLR

  6. [6]

    De Haan, P., Jayaraman, D., and Levine, S. (2019). Causal confusion in imitation learning. Advances in neural information processing systems , 32

  7. [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

  8. [8]

    G., and Pineau, J

    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

Show all 34 references
  1. [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

  2. [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

  3. [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

  4. [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

  5. [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

  6. [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

  7. [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

  8. [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

  9. [17]

    McCallum, A. K. (1995). Reinforcement Learning with Selective Perception and Hidden State . PhD thesis, University of Rochester

  10. [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

  11. [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

  12. [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

  13. [21]

    Pearl, J., Glymour, M., and Jewell, N. P. (2016). Causal inference in statistics: A primer . John Wiley & Sons

  14. [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

  15. [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

  16. [24]

    Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. (2017). Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347

  17. [25]

    Song, X., Jiang, Y., Tu, S., Du, Y., and Neyshabur, B. (2020). Observational overfitting in reinforcement learning. In International Conference on Learning Representations

  18. [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

  19. [27]

    Sutton, R. S. and Barto, A. G. (2018). Reinforcement learning: An introduction . MIT press

  20. [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

  21. [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

  22. [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

  23. [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

  24. [32]

    Williams, R. J. (1992). Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine Learning , 8(3-4):229--256

  25. [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

  26. [34]

    Zhang, J., Kumor, D., and Bareinboim, E. (2020b). Causal imitation learning with unobserved confounders. Advances in neural information processing systems , 33:12263--12274

Pith tools

Reviewed August 7, 2026 · model on record in the stance chip above.