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

📄 arXiv: 2605.14558v1 📥 PDF

作者: Langzhou He, Junyou Zhu, Yue Zhou, Zhengyao Gu, Junhua Liu, Wei-Chieh Huang, Henry Peng Zou, David Wipf, Philip S. Yu, Qitian Wu

分类: cs.LG, cs.AI, cs.CL

发布日期: 2026-05-14

备注: Preprint


💡 一句话要点

ActFocus:通过Token级能量分析解决Agentic强化学习中的动作瓶颈问题

🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture) 支柱九:具身大模型 (Embodied Foundation Models)

关键词: Agentic强化学习 动作瓶颈 Token级能量 梯度重加权 策略梯度 大型语言模型 信用分配

📋 核心要点

  1. 现有Agentic强化学习方法对所有token进行统一信用分配,忽略了动作token的重要性,导致训练效率低下。
  2. ActFocus通过token级能量分析,识别并提升关键动作token的权重,降低推理token的权重,从而优化训练信号分配。
  3. 实验表明,ActFocus在多个环境和模型规模下显著优于PPO和GRPO,且无需额外计算或内存开销。

📝 摘要(中文)

Agentic强化学习使用多轮轨迹训练大型语言模型,这些轨迹交织着长推理过程和短环境交互动作。常见的策略梯度方法(如PPO和GRPO)平等地对待轨迹中的每个token,导致统一的信用分配。本文指出,这种统一的信用分配在很大程度上错误地分配了token级别的训练信号。从基于能量建模的角度来看,我们表明,token级别的训练信号(通过它们与给定prompt的不同rollout的奖励方差的相关性来量化)集中在动作token上,而不是推理token上,即使动作token仅占轨迹的一小部分。我们将这种现象称为“动作瓶颈”。受此观察的启发,我们提出了一种简单的token重加权方法ActFocus,该方法降低了推理token上的梯度,以及一种额外的基于能量的重新分配机制,该机制进一步增加了具有较高不确定性的动作token的权重。在四个环境和不同的模型大小上,ActFocus始终优于PPO和GRPO,最终步的收益分别高达65.2和63.7个百分点,而无需任何额外的运行时或内存成本。

🔬 方法详解

问题定义:Agentic强化学习旨在训练能够进行复杂推理并与环境交互的大型语言模型。现有的策略梯度方法,如PPO和GRPO,在训练过程中对轨迹中的每个token赋予相同的权重,这导致了次优的信用分配。推理token通常占据轨迹的大部分,但它们对最终奖励的贡献远小于直接与环境交互的动作token。这种统一的信用分配方式浪费了大量的计算资源,并阻碍了模型的学习效率。

核心思路:论文的核心思路是识别并解决Agentic强化学习中的“动作瓶颈”问题。作者通过能量建模分析发现,token级别的训练信号与奖励方差的相关性高度集中在动作token上。因此,应该更加关注动作token的训练,并减少推理token的干扰。基于此,论文提出了一种token重加权方法,即ActFocus,通过调整不同token的梯度权重来优化训练过程。

技术框架:ActFocus方法主要包含两个阶段:1) Token重加权:降低推理token的梯度权重,提升动作token的梯度权重。2) 基于能量的重新分配:进一步增加具有较高不确定性的动作token的权重。整体流程如下:首先,使用Agentic强化学习框架生成多轮轨迹;然后,计算每个token的能量值,并根据能量值调整梯度权重;最后,使用调整后的梯度进行模型更新。

关键创新:ActFocus的关键创新在于其token级别的梯度重加权机制,该机制能够更有效地分配训练信号,从而提高Agentic强化学习的效率。与现有方法相比,ActFocus能够更加精准地识别对奖励贡献最大的token,并集中资源进行训练。此外,基于能量的重新分配机制进一步提升了对不确定性较高的动作token的关注度,有助于模型更好地探索环境。

关键设计:ActFocus的关键设计包括:1) Token能量的计算方式:通过计算每个token与奖励方差的相关性来量化其能量值。2) 梯度重加权策略:使用一个简单的权重函数,根据token的能量值调整其梯度权重。3) 基于能量的重新分配机制:使用softmax函数对动作token的能量值进行归一化,并将其作为权重进一步调整动作token的梯度。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

ActFocus在四个不同的环境和多种模型规模下进行了实验,结果表明,ActFocus始终优于PPO和GRPO。在最终步的性能上,ActFocus相比PPO和GRPO分别取得了高达65.2和63.7个百分点的提升,且无需额外的运行时或内存开销。这些结果充分证明了ActFocus方法的有效性和实用性。

🎯 应用场景

ActFocus方法可应用于各种需要复杂推理和环境交互的Agentic强化学习任务,例如机器人控制、游戏AI、对话系统等。通过提高训练效率,ActFocus能够帮助开发更智能、更高效的智能体,从而在实际应用中发挥更大的作用。该研究对于推动Agentic强化学习的发展具有重要意义。

📄 摘要(原文)

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.