# Imitate Optimal Policy: Prevail and Induce Action Collapse in Policy Gradient

arXiv: https://arxiv.org/abs/2509.02737 · HTML: https://mlsanity.com/paper/2509.02737/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Deep RL algorithms and representations, one of 15 directions in Deep reinforcement learning methods.
- Builds on 33 of its 51 references that are in the map.
- First posted on arXiv 2 Sep 2025.

## Placed in

- Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/)

## Builds on (most cited first, 30 of 33)

- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [Continuous control with deep reinforcement learning](https://mlsanity.com/paper/1509.02971/) (2015)
- [Trust Region Policy Optimization](https://mlsanity.com/paper/1502.05477/) (2015)
- [Asynchronous Methods for Deep Reinforcement Learning](https://mlsanity.com/paper/1602.01783/) (2016)
- [A General Language Assistant as a Laboratory for Alignment](https://mlsanity.com/paper/2112.00861/) (2021)
- [Prevalence of Neural Collapse during the terminal phase of deep learning training](https://mlsanity.com/paper/2008.08186/) (2020)
- [Instruction Tuning with GPT-4](https://mlsanity.com/paper/2304.03277/) (2023)
- [Safe RLHF: Safe Reinforcement Learning from Human Feedback](https://mlsanity.com/paper/2310.12773/) (2023)
- [Is Q-learning Provably Efficient?](https://mlsanity.com/paper/1807.03765/) (2018)
- [A Theoretical Analysis of Deep Q-Learning](https://mlsanity.com/paper/1901.00137/) (2019)
- [Understanding the impact of entropy on policy optimization](https://mlsanity.com/paper/1811.11214/) (2018)
- [Exploring Deep Neural Networks via Layer-Peeled Model: Minority Collapse in Imbalanced Training](https://mlsanity.com/paper/2101.12699/) (2021)
- [Neural Collapse Under MSE Loss: Proximity to and Dynamics on the Central Path](https://mlsanity.com/paper/2106.02073/) (2021)
- [Global Optimality Guarantees For Policy Gradient Methods](https://mlsanity.com/paper/1906.01786/) (2019)
- [Neural collapse with unconstrained features](https://mlsanity.com/paper/2011.11619/) (2020)
- [Global Convergence of Policy Gradient Methods to (Almost) Locally Optimal Policies](https://mlsanity.com/paper/1906.08383/) (2019)
- [Phasic Policy Gradient](https://mlsanity.com/paper/2009.04416/) (2020)
- [Neural Collapse Inspired Feature-Classifier Alignment for Few-Shot Class Incremental Learning](https://mlsanity.com/paper/2302.03004/) (2023)
- [Sample Complexity of Episodic Fixed-Horizon Reinforcement Learning](https://mlsanity.com/paper/1510.08906/) (2015)
- [Non-Asymptotic Gap-Dependent Regret Bounds for Tabular MDPs](https://mlsanity.com/paper/1905.03814/) (2019)
- [Inducing Neural Collapse in Imbalanced Learning: Do We Really Need a Learnable Classifier at the End of Deep Neural Network?](https://mlsanity.com/paper/2203.09081/) (2022)
- [On Linear Convergence of Policy Gradient Methods for Finite MDPs](https://mlsanity.com/paper/2007.11120/) (2020)
- [Neural Collapse with Cross-Entropy Loss](https://mlsanity.com/paper/2012.08465/) (2020)
- [Neural Collapse Inspired Attraction-Repulsion-Balanced Loss for Imbalanced Learning](https://mlsanity.com/paper/2204.08735/) (2022)
- [Linguistic Collapse: Neural Collapse in (Large) Language Models](https://mlsanity.com/paper/2405.17767/) (2024)
- [Regularized Policies are Reward Robust](https://mlsanity.com/paper/2101.07012/) (2021)
- [Inducing Neural Collapse in Deep Long-tailed Learning](https://mlsanity.com/paper/2302.12453/) (2023)
- [TD Convergence: An Optimization Perspective](https://mlsanity.com/paper/2306.17750/) (2023)
- [Inducing Neural Collapse to a Fixed Hierarchy-Aware Frame for Reducing Mistake Severity](https://mlsanity.com/paper/2303.05689/) (2023)

## Abstract (arXiv)

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.