REVIEW 4 major objections 6 minor 22 references
Imitate Optimal Policy: Prevail and Induce Action Collapse in Policy Gradient
T0 review · 4 major / 6 minor · reviewed 2026-08-15 · deepseek-v4-flash
Pith's one-line read The paper claims that optimal policy networks collapse last-layer features into a simplex ETF, and fixing the action layer to that ETF makes the collapse a guaranteed global optimum while speeding up training.
desk verdict The empirical trick of freezing the policy head to an ETF is worth knowing, but the theory proves something about an imitation loss, not the policy gradient objective actually trained, so the induced-collapse claim is unsupported. 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 load-bearing object is the simplex equiangular tight frame (ETF): a set of $K$ vectors with equal norm and pairwise angle $\cos\theta = -1/(K-1)$, the maximal equiangular separation. The paper fixes the action-selection layer $W^*$ to a randomly oriented ETF scaled to norm $\sqrt{E_W}$ and, following the layer-peeled model, treats the last-layer activations as the only variables under the constraint $\|h_s\|_2 \le \sqrt{E_H}$. The proof uses Karush-Kuhn-Tucker conditions; the ETF's defining identities $\sum_k w_k = 0$ and $w_i^T w_j = -E_W/(K-1)$ make the alignment $h = \sqrt{E_H/E_W}\,w_k$ a stationary point, and convexity of the objective makes it the global optimum for every class simultaneously, independent of class sizes.
What would settle it
Take a small finite MDP with $K=2$, unequal class sizes $|S_1| \ne |S_2|$, and fixed ETF weights; solve Eq. (8) exactly and confirm the global optimum is $h_s \propto w_k$. Then compute the actual REINFORCE or PPO objective for those same activations with Monte-Carlo returns and the implied state distribution; if any non-collapsed activation set gives a strictly higher value, the equivalence the proof relies on fails.
Extended reading notes
Core claim
The central claim is Theorem 1: for a layer-peeled policy network with a fixed simplex ETF action-selection layer, every global optimizer $H^*$ of Eq. (8) satisfies $h^{*T}_{s_k} w^*_{k'} = \sqrt{E_H E_W}\left(\frac{K}{K-1}\delta_{k,k'}-\frac{1}{K-1}\right)$ for all $k,k'$ and all states $s\in S_k$. That is, activations collapse to the direction of their optimal action's weight, with length $\sqrt{E_H}$, so Action Collapse is guaranteed by the objective rather than by balanced data or full exploration. The paper also reports the empirical counterpart: in ideal discrete environments, the geometry appears naturally, while in harder environments it does not, and imposing the ETF layer both induces the geometry and improves rewards.
Load-bearing premise
The proof applies to a simplified weighted-log-likelihood objective (Eq. 8), and the paper's own appendix notes that this differs from the general policy-gradient objective without proving equivalence; if the two are not equivalent, Theorem 1 does not cover REINFORCE, PPO, TRPO, or A3C.
Editorial extensions
If this is right
- A discrete policy-gradient agent with a fixed ETF action-selection layer will drive its last-layer activations to the ETF configuration even under partial exploration and imbalanced state sampling, because that geometry is the unique global optimum of the simplified objective.
- ACPG can be layered onto existing discrete policy-gradient algorithms without changing their update rules; in the tested environments it improves best and final rewards and reduces run-to-run variance.
- The fixed ETF gives maximal pairwise angular separation between actions, so the policy's decision boundary is as spread out as possible, which the paper connects to faster convergence and robustness under imbalanced states.
- Because the action-selection layer is frozen, ACPG removes a set of learned parameters and reduces the training objective to aligning activations, which is why the paper reports earlier stopping and lower standard deviation.
Reading between the lines
- A testable extension: if the theorem's mechanism is the binding one, then deliberately limiting backbone capacity should erase ACPG's gains, since the proof assumes activations can be optimized independently.
- The epsilon-greedy results hint that ACPG converts exploration into geometric alignment; an untested corollary is that pairing ACPG with an exploration schedule tuned per environment would widen the reported improvements.
- The same fixed-frame idea could transfer to other finite large action spaces, such as token-selection heads in sequence models, where a frozen ETF might induce a similar collapse; the paper does not test this.
Signed reviews
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper studies the geometry of last-layer features and action-selection weights in discrete policy-gradient (PG) networks. It reports that in idealized environments (fully explored, balanced state subsets, uniform stationary distribution) an 'Action Collapse' (AC) phenomenon occurs: state-action activations and action-selection weights converge to a simplex equiangular tight frame (ETF). To induce this structure in realistic environments, the authors propose Action Collapse Policy Gradient (ACPG), which fixes the action-selection layer to a randomly oriented simplex ETF and trains only the backbone. Theorem 1 claims that any global minimizer of Eq.(8) is a simplex ETF aligned with the fixed weights, irrespective of exploration, balance, or state-subset size. Experiments with REINFORCE, PPO, TRPO, and A3C on several Gym environments report faster and more robust convergence when ACPG is applied.
Significance. The empirical study is broad: four PG algorithms, multiple environments, and 20 seeds, and the proposed intervention is simple and easy to integrate. The paper also identifies a genuine gap, since neural-collapse results from supervised classification do not automatically transfer to reinforcement learning. If Theorem 1 were about the actual objective minimized by PG algorithms, the result would be valuable, giving a provable geometric target for policy networks. However, the theorem optimizes a per-state weighted log-likelihood of the optimal action, which is an imitation/classification loss rather than the PG objective used in the experiments. Appendix B explicitly acknowledges that Eq.(1) and Eq.(2) differ but only asserts that both are valid losses, without proving equivalence of their optimizers. The central conceptual claim, that ACPG 'induces' Action Collapse under PG, is therefore not supported by the presented theory. The empirical improvements may still hold, but the theoretical bridge is missing.
major comments (4)
- [§4.2 (Theorem 1) and Appendix B] Theorem 1 optimizes Eq.(8), which is a weighted log-likelihood of the optimal action a* for each state, i.e., a supervised imitation/classification objective, not the policy-gradient objective in Eq.(1) that REINFORCE, PPO, TRPO, and A3C actually minimize in Section 5. Eq.(1) contains a sum over all actions weighted by πθ(a|s), whereas Eq.(8) contains only the log-probability of a*. Appendix B's final note concedes that Eq.(1) and Eq.(2) differ but merely asserts that both are valid loss functions; it does not show that they share global minimizers or that the algorithms in Section 5 minimize Eq.(8). Moreover, dπ(s) and Ψπ in Eq.(8) are treated as fixed even though in PG they depend on the current policy and hence on H. Consequently, Theorem 1 does not establish that ACPG induces Action Collapse under the PG algorithms evaluated.
- [§2.1 and §4.2] The definition of a* is not operational for PG. In Eq.(2), a* is 'the optimal action predicted by the model', while Theorem 1 treats a* as the true optimal action of state s. In the experiments, the agent does not know optimal actions and samples actions from πθ, so the loss in Eq.(8) is not the loss being minimized. If a* were the model's own argmax, the theorem would describe a self-imitation objective and would still not apply to the sampled-action policy-gradient loss used by the compared algorithms.
- [Appendix B, Eqs. (15)-(16)] The reduction from Eq.(15) to Eq.(16) silently drops the factors dπ(s) and Ψπ and replaces the problem by minimizing Σ_{j≠k} exp(h^T(w_j-w_k)). This equivalence holds only if Ψπ > 0 for all states (and dπ(s) > 0), which is not stated in Theorem 1 and is not guaranteed for returns or advantages in PG. If Ψπ < 0 for some state, the sign of the term flips and the global minimizer of Eq.(8) need not align h with w_k. The theorem therefore requires an explicit positivity assumption or a different argument.
- [§4.2] The theorem's conclusion is largely built into the setup: once the action-selection layer is fixed to a simplex ETF, the unconstrained-features optimum of the weighted softmax loss aligns each activation with the corresponding ETF vector. This is a known property of fixed-ETF classifiers in the neural-collapse literature (e.g., Fang et al. 2021; Yang et al. 2022), not a statement about PG training dynamics. The paper should present Theorem 1 as a property of the chosen target geometry and should not use it to claim that Action Collapse 'emerges' under PG updates.
minor comments (6)
- [§2.1, Eq. (1)] In Eq.(1), Ψπ should be written Ψπ(s,a); as printed, the quantity inside the sum over a appears not to depend on a, which is confusing given the summation index.
- [§2.1 and §4.2] The notation a* is used inconsistently: 'optimal action predicted by the model' in Eq.(2) versus 'true optimal action of state s' in Theorem 1. Please define one convention and use it throughout.
- [§3] It is unclear how Car-Racing, which has a continuous state space, satisfies Condition 3 (dπ(s) = dπ(s′) for all states). Please specify how states are discretized and how the stationary distribution is estimated; without this, the claimed evidence of Action Collapse in Car-Racing is hard to verify.
- [Table 1] Table 1 is difficult to parse: the 'Best Final' columns mix best and final rewards, and some cells contain '--' without explanation. Please restructure the table or add a clear legend explaining each block of columns.
- [Appendix B, Eqs. (13)-(16)] The notation S_k and n_k is introduced in the proof but is not connected to the state subsets S_k defined in Definition 1; also, the summation index in Eq.(14) is inconsistent ('nk' versus 'n_k').
- [Appendix E] There is a typo: 'Theorem 1 1 states...' should be 'Theorem 1 states...'.
Assumptions & free parameters
free parameters (3)
- E_W
- E_H
- learning_rate =
1e-3, 1e-4, 1e-5 or 8e-5, 2.5e-4, 8e-4
assumptions (5)
- domain assumption The last-layer activations are the only optimized variables; the backbone weights are fixed (layer-peeled model).
- ad hoc to paper The objective in Eq.(2) equals the policy-gradient objective or is an equally valid loss for policy training.
- domain assumption The stationary distribution d_pi(s) and return weights Psi are fixed positive constants independent of the activations H.
- ad hoc to paper For each state, minimizing the sum of exponentials is equivalent to minimizing the log-sum-exp objective.
- standard math The simplex ETF weight vectors sum to zero (sum_k w_k = 0).
invented entities (1)
-
Action Collapse (AC)
independent evidence
Cite this review
Pith. "Pith review of Imitate Optimal Policy: Prevail and Induce Action Collapse in Policy Gradient." pith.science (2026). https://pith.science/paper/OVD5QXTF
@misc{pith2026250902737,
author = {Pith},
title = {Pith review of: Imitate Optimal Policy: Prevail and Induce Action Collapse in Policy Gradient},
year = {2026},
howpublished = {\url{https://pith.science/paper/OVD5QXTF}},
note = {Machine review of arXiv:2509.02737}
}
read the original abstract
Policy gradient (PG) methods in reinforcement learning frequently utilize deep neural networks (DNNs) to learn a shared backbone of feature representations used to compute likelihoods in an action selection layer. Numerous studies have been conducted on the convergence and global optima of policy networks, but few have analyzed representational structures of those underlying networks. While training an optimal policy DNN, we observed that under certain constraints, a gentle structure resembling neural collapse, which we refer to as Action Collapse (AC), emerges. This suggests that 1) the state-action activations (i.e. last-layer features) sharing the same optimal actions collapse towards those optimal actions respective mean activations; 2) the variability of activations sharing the same optimal actions converges to zero; 3) the weights of action selection layer and the mean activations collapse to a simplex equiangular tight frame (ETF). Our early work showed those aforementioned constraints to be necessary for these observations. Since the collapsed ETF of optimal policy DNNs maximally separates the pair-wise angles of all actions in the state-action space, we naturally raise a question: can we learn an optimal policy using an ETF structure as a (fixed) target configuration in the action selection layer? Our analytical proof shows that learning activations with a fixed ETF as action selection layer naturally leads to the AC. We thus propose the Action Collapse Policy Gradient (ACPG) method, which accordingly affixes a synthetic ETF as our action selection layer. ACPG induces the policy DNN to produce such an ideal configuration in the action selection layer while remaining optimal. Our experiments across various OpenAI Gym environments demonstrate that our technique can be integrated into any discrete PG methods and lead to favorable reward improvements more quickly and robustly.
Figures
Figures from the paper (1 more)
Reference graph
Works this paper leans on
-
[6]
Contrastive policy gradient: Aligning LLMs on sequence-level scores in a supervised-friendly fashion
Yannis Flet-Berliac, Nathan Grinsztajn, Florian Strub, Eugene Choi, Bill Wu, Chris Cremer, Arash Ahmadian, Yash Chandak, Mohammad Gheshlaghi Azar, Olivier Pietquin, and Matthieu Geist. Contrastive policy gradient: Aligning LLMs on sequence-level scores in a supervised-friendly fashion. In Proceedings of the 2024 Conference on Empirical Methods in Natural ...
work page 2024
-
[9]
Neural collapse under mse loss: Proximity to and dynamics on the central path
XY Han, Vardan Papyan, and David L Donoho. Neural collapse under mse loss: Proximity to and dynamics on the central path. arXiv preprint arXiv:2106.02073,
-
[10]
• End.: Epochs: 100/500, Step per epoch: 100000, Step per collect: 1000, Repeat per collect: 4, Training num: 10, Test num: 10 • Qbe.: Epochs: 100/500, Step per epoch: 100000, Step per collect: 1000, Repeat per collect: 4, Training num: 10, Test num: 10 • Pon.: Epochs: 100/500, Step per epoch: 100000, Step per collect: 1000, Repeat per collect: 4, Trainin...
work page 2024
-
[11]
Inducing Neural Collapse to a Fixed Hierarchy-Aware Frame for Reducing Mistake Severity
Tong Liang and Jim Davis. Inducing neural collapse to a fixed hierarchy-aware frame for reducing mistake severity. arXiv preprint arXiv:2303.05689,
-
[12]
Continuous control with deep reinforcement learning
Timothy P Lillicrap, Jonathan J Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint arXiv:1509.02971,
-
[13]
Neural collapse with cross-entropy loss
Jianfeng Lu and Stefan Steinerberger. Neural collapse with cross-entropy loss. arXiv preprint arXiv:2012.08465,
arXiv 2012
-
[15]
Asynchronous methods for deep reinforcement learning
V olodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In International conference on machine learning, pp. 1928–1937. PMLR,
1928
-
[16]
Baolin Peng, Chunyuan Li, Pengcheng He, Michel Galley, and Jianfeng Gao. Instruction tuning with gpt-4. arXiv preprint arXiv:2304.03277,
Show all 22 references
-
[18]
Liang Xie, Yibo Yang, Deng Cai, and Xiaofei He
URL https://arxiv.org/abs/2405.17767. Liang Xie, Yibo Yang, Deng Cai, and Xiaofei He. Neural collapse inspired attraction–repulsion-balanced loss for imbalanced learning. Neurocomputing, 527:1–14,
-
[19]
Huaqing Xiong, Tengyu Xu, Lin Zhao, Yingbin Liang, and Wei Zhang
doi: 10.1016/j.neucom.2023.01.023. Huaqing Xiong, Tengyu Xu, Lin Zhao, Yingbin Liang, and Wei Zhang. Deterministic policy gradient: Conver- gence analysis. In Uncertainty in Artificial Intelligence, pp. 2159–2169. PMLR, 2022a. Huaqing Xiong, Tengyu Xu, Lin Zhao, Yingbin Liang,...
2023 doi
-
[20]
Convergence and iteration complexity of policy gradient method for infinite-horizon reinforcement learning
Kaiqing Zhang, Alec Koppel, Hao Zhu, and Tamer Ba¸ sar. Convergence and iteration complexity of policy gradient method for infinite-horizon reinforcement learning. In 2019 IEEE 58th Conference on Decision and Control (CDC), pp. 7415–7422. IEEE,
2019
-
[21]
13 Preprint. Imitate Optimal Policy: Prevail and Induce Action Collapse in Policy Gradient Supplementary Material A N EURAL COLLAPSE PHENOMENON Then the neural collapse (NC) phenomenon can be formally described as Papyan et al. (2020): (NC1) Collapse of within-class variabilit...
2020
-
[1998]
On the global convergence of natural actor-critic with two-layer neural network parametrization
Mudit Gaur, Amrit Singh Bedi, Di Wang, and Vaneet Aggarwal. On the global convergence of natural actor-critic with two-layer neural network parametrization. arXiv preprint arXiv:2306.10486,
-
[1999]
Learning classifiers on positive and unlabeled data with policy gradient
Tianyu Li, Chien-Chih Wang, Yukun Ma, Patricia Ortal, Qifang Zhao, Bjorn Stenger, and Yu Hirate. Learning classifiers on positive and unlabeled data with policy gradient. In 2019 IEEE International Conference on Data Mining (ICDM), pp. 399–408. IEEE,
2019
-
[2013]
doi: 10.1613/jair.3912
ISSN 1076-9757. doi: 10.1613/jair.3912. URL http://dx.doi.org/10.1613/jair.3912. Jalaj Bhandari and Daniel Russo. Global optimality guarantees for policy gradient methods. arXiv preprint arXiv:1906.01786,
1906 arXiv
-
[2015]
Proximal policy optimization algorithms
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347,
-
[2017]
Td convergence: An optimization perspective
Kavosh Asadi, Shoham Sabach, Yao Liu, Omer Gottesman, and Rasool Fakoor. Td convergence: An optimization perspective. arXiv preprint arXiv:2306.17750,
-
[2019]
Understanding of a convolutional neural network
Saad Albawi, Tareq Abed Mohammed, and Saad Al-Zawi. Understanding of a convolutional neural network. In 2017 international conference on engineering and technology (ICET), pp. 1–6. Ieee,
2017
-
[2020]
Neural collapse with unconstrained features
Dustin G Mixon, Hans Parshall, and Jianzong Pi. Neural collapse with unconstrained features. arXiv preprint arXiv:2011.11619,
2011 arXiv
-
[2021]
Openai gym
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. arXiv preprint arXiv:1606.01540,
-
[2023]
A general language assistant as a laboratory for alignment
Amanda Askell, Yuntao Bai, Anna Chen, Dawn Drain, Deep Ganguli, Tom Henighan, Andy Jones, Nicholas Joseph, Ben Mann, Nova DasSarma, et al. A general language assistant as a laboratory for alignment. arXiv preprint arXiv:2112.00861,
-
[2024]
doi: 10.18653/v1/2024.emnlp-main.1190
Association for Computational Linguistics. doi: 10.18653/v1/2024.emnlp-main.1190. Matt W Gardner and SR Dorling. Artificial neural networks (the multilayer perceptron)—a review of applications in the atmospheric sciences. Atmospheric environment, 32(14-15):2627–2636,
2024 doi
Reviewed August 15, 2026 · model on record in the stance chip above.
Discussion (0). Continue with ORCID to comment.