Resolving Action Bottleneck: Agentic Reinforcement Learning Informed by Token-Level Energy

Fuente: arXiv
Enregistré dans:
Détails bibliographiques
Auteurs principaux: He, Langzhou, Zhu, Junyou, Zhou, Yue, Gu, Zhengyao, Liu, Junhua, Huang, Wei-Chieh, Zou, Henry Peng, Wipf, David, Yu, Philip S., Wu, Qitian
Format: Preprint
Publié: 2026
Sujets:
Accès en ligne:
Tags: Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
_version_ 1866911685443321856
author He, Langzhou
Zhu, Junyou
Zhou, Yue
Gu, Zhengyao
Liu, Junhua
Huang, Wei-Chieh
Zou, Henry Peng
Wipf, David
Yu, Philip S.
Wu, Qitian
author_facet He, Langzhou
Zhu, Junyou
Zhou, Yue
Gu, Zhengyao
Liu, Junhua
Huang, Wei-Chieh
Zou, Henry Peng
Wipf, David
Yu, Philip S.
Wu, Qitian
contents Agentic reinforcement learning trains large language models using multi-turn trajectories that interleave long reasoning traces with short environment-facing actions. Common policy-gradient methods, such as PPO and GRPO, treat each token in a trajectory equally, leading to uniform credit assignment. In this paper, we critically demonstrate that such uniform credit assignment largely misallocates token-level training signals. From an energy-based modeling perspective, we show that token-level training signals, quantified by their correlations with reward variance of different rollouts sampled from a given prompt, concentrate sharply on action tokens rather than reasoning tokens, even though action tokens account for only a small fraction of the trajectory. We refer to this phenomenon as the Action Bottleneck. Motivated by this observation, we propose an embarrassingly simple token reweighting approach, ActFocus, that downweights gradients on reasoning tokens, along with an additional energy-based redistribution mechanism that further increases the weights on action tokens with higher uncertainty. Across four environments and different model sizes, ActFocus consistently outperforms PPO and GRPO, yielding final-step gains of up to 65.2 and 63.7 percentage points, respectively, without any additional runtime or memory cost.
format Preprint
id arxiv_https___arxiv_org_abs_2605_14558
institution arXiv
publishDate 2026
record_format arxiv
spellingShingle Resolving Action Bottleneck: Agentic Reinforcement Learning Informed by Token-Level Energy
He, Langzhou
Zhu, Junyou
Zhou, Yue
Gu, Zhengyao
Liu, Junhua
Huang, Wei-Chieh
Zou, Henry Peng
Wipf, David
Yu, Philip S.
Wu, Qitian
Machine Learning
Artificial Intelligence
Computation and Language
I.2.6; I.2.7; I.2.8
Agentic reinforcement learning trains large language models using multi-turn trajectories that interleave long reasoning traces with short environment-facing actions. Common policy-gradient methods, such as PPO and GRPO, treat each token in a trajectory equally, leading to uniform credit assignment. In this paper, we critically demonstrate that such uniform credit assignment largely misallocates token-level training signals. From an energy-based modeling perspective, we show that token-level training signals, quantified by their correlations with reward variance of different rollouts sampled from a given prompt, concentrate sharply on action tokens rather than reasoning tokens, even though action tokens account for only a small fraction of the trajectory. We refer to this phenomenon as the Action Bottleneck. Motivated by this observation, we propose an embarrassingly simple token reweighting approach, ActFocus, that downweights gradients on reasoning tokens, along with an additional energy-based redistribution mechanism that further increases the weights on action tokens with higher uncertainty. Across four environments and different model sizes, ActFocus consistently outperforms PPO and GRPO, yielding final-step gains of up to 65.2 and 63.7 percentage points, respectively, without any additional runtime or memory cost.
title Resolving Action Bottleneck: Agentic Reinforcement Learning Informed by Token-Level Energy
topic Machine Learning
Artificial Intelligence
Computation and Language
I.2.6; I.2.7; I.2.8
url https://arxiv.org/abs/2605.14558