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

• URL: http://arxiv.org/abs/2509.02737v1
• Date: Tue, 02 Sep 2025 18:33:11 GMT
• Title: Imitate Optimal Policy: Prevail and Induce Action Collapse in Policy Gradient
• Authors: Zhongzhu Zhou, Yibo Yang, Ziyan Chen, Fengxiang Bie, Haojun Xia, Xiaoxia Wu, Robert Wu, Ben Athiwaratkun, Bernard Ghanem, Shuaiwen Leon Song,
• Abstract summary: Policy 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.<n>We show that under certain constraints, a structure resembling neural collapse, which we refer to as Action Collapse (AC), emerges.<n>We propose the Action Collapse Policy Gradient (ACPG) method, which accordingly affixes a synthetic ETF as our action selection layer.
• Score: 61.440209025381016
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• Latent Spherical Flow Policy for Reinforcement Learning with Combinatorial Actions [31.697208397735395]
  Existing approaches embed task-specific value functions into constrained optimization programs or learn deterministic structured policies, sacrificing generality and policy expressiveness.<n>We propose a solver-induced emphlatent spherical flow policy that brings the expressiveness of modern generative policies to the RL while guaranteeing feasibility by design.<n>Our approach outperforms state-of-the-art baselines by an average of 20.6% across a range of challenging RL tasks.
  arXiv  Detail & Related papers  (2026-01-29T18:49:07Z)
• Polychromic Objectives for Reinforcement Learning [63.37185057794815]
  Reinforcement learning fine-tuning (RLFT) is a dominant paradigm for improving pretrained policies for downstream tasks.<n>We introduce an objective for policy methods that explicitly enforces the exploration and refinement of diverse generations.<n>We show how proximal policy optimization (PPO) can be adapted to optimize this objective.
  arXiv  Detail & Related papers  (2025-09-29T19:32:11Z)
• ACT-JEPA: Novel Joint-Embedding Predictive Architecture for Efficient Policy Representation Learning [90.41852663775086]
  ACT-JEPA is a novel architecture that integrates imitation learning and self-supervised learning.<n>We train a policy to predict action sequences and abstract observation sequences.<n>Our experiments show that ACT-JEPA improves the quality of representations by learning temporal environment dynamics.
  arXiv  Detail & Related papers  (2025-01-24T16:41:41Z)
• Reinforcing Language Agents via Policy Optimization with Action Decomposition [36.984163245259936]
  This paper proposes decomposing language agent optimization from the action level to the token level. We then derive the Bellman backup with Action Decomposition (BAD) to integrate credit assignments for both intra-action and inter-action tokens. Implementing BAD within the PPO algorithm, we introduce Policy Optimization with Action Decomposition (POAD)
  arXiv  Detail & Related papers  (2024-05-23T14:01:44Z)
• Extremum-Seeking Action Selection for Accelerating Policy Optimization [18.162794442835413]
  Reinforcement learning for control over continuous spaces typically uses high-entropy policies, such as Gaussian distributions, for local exploration and estimating policy to optimize performance. We propose to improve action selection in this model-free RL setting by introducing additional adaptive control steps based on Extremum-Seeking Control (ESC) Our methods can be easily added in standard policy optimization to improve learning efficiency, which we demonstrate in various control learning environments.
  arXiv  Detail & Related papers  (2024-04-02T02:39:17Z)
• Probabilistic Reach-Avoid for Bayesian Neural Networks [71.67052234622781]
  We show that an optimal synthesis algorithm can provide more than a four-fold increase in the number of certifiable states. The algorithm is able to provide more than a three-fold increase in the average guaranteed reach-avoid probability.
  arXiv  Detail & Related papers  (2023-10-03T10:52:21Z)
• Diverse Policy Optimization for Structured Action Space [59.361076277997704]
  We propose Diverse Policy Optimization (DPO) to model the policies in structured action space as the energy-based models (EBM) A novel and powerful generative model, GFlowNet, is introduced as the efficient, diverse EBM-based policy sampler. Experiments on ATSC and Battle benchmarks demonstrate that DPO can efficiently discover surprisingly diverse policies.
  arXiv  Detail & Related papers  (2023-02-23T10:48:09Z)
• DEFT: Diverse Ensembles for Fast Transfer in Reinforcement Learning [1.111018778205595]
  We present Diverse Ensembles for Fast Transfer in RL (DEFT), a new ensemble-based method for reinforcement learning in highly multimodal environments. The algorithm is broken down into two main phases: training of ensemble members, and synthesis (or fine-tuning) of the ensemble members into a policy that works in a new environment.
  arXiv  Detail & Related papers  (2022-09-26T04:35:57Z)
• A Regularized Implicit Policy for Offline Reinforcement Learning [54.7427227775581]
  offline reinforcement learning enables learning from a fixed dataset, without further interactions with the environment. We propose a framework that supports learning a flexible yet well-regularized fully-implicit policy. Experiments and ablation study on the D4RL dataset validate our framework and the effectiveness of our algorithmic designs.
  arXiv  Detail & Related papers  (2022-02-19T20:22:04Z)
• Evolutionary Action Selection for Gradient-based Policy Learning [6.282299638495976]
  Evolutionary algorithms (EAs) and Deep Reinforcement Learning (DRL) have recently been combined to integrate the advantages of the two solutions for better policy learning. We propose Evolutionary Action Selection-Twin Delayed Deep Deterministic Policy Gradient (EAS-TD3), a novel combination of EA and DRL.
  arXiv  Detail & Related papers  (2022-01-12T03:31:21Z)
• Discrete Action On-Policy Learning with Action-Value Critic [72.20609919995086]
  Reinforcement learning (RL) in discrete action space is ubiquitous in real-world applications, but its complexity grows exponentially with the action-space dimension. We construct a critic to estimate action-value functions, apply it on correlated actions, and combine these critic estimated action values to control the variance of gradient estimation. These efforts result in a new discrete action on-policy RL algorithm that empirically outperforms related on-policy algorithms relying on variance control techniques.
  arXiv  Detail & Related papers  (2020-02-10T04:23:09Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.