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

2026-05-14

AI Summary

This work addresses a critical limitation in conventional policy gradient methods for reinforcement learning, which assign credit uniformly across all tokens in a trajectory and thereby overlook the pivotal role of action tokens, leading to mismatched training signals. The study identifies and formally names this issue the “action bottleneck,” demonstrating that training signals are in fact highly concentrated on action tokens. To mitigate this, the authors propose ActFocus, a novel approach that leverages energy-based modeling to analyze reward variance and introduces a token-level reweighting mechanism—requiring no additional computational overhead—combined with an uncertainty-aware energy redistribution strategy to refine gradient allocation. ActFocus integrates seamlessly into existing frameworks such as PPO and GRPO, achieving substantial performance gains across four environments and multiple model scales, with final-step success rates improving by up to 65.2 and 63.7 percentage points.

Abstract

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.

Problem

• Action Bottleneck
• Agentic Reinforcement Learning
• Token-level Credit Assignment
• Energy-based Modeling
• Large Language Models

Innovation

• Action Bottleneck
• Token-level Energy
• ActFocus
• Agentic Reinforcement Learning
• Credit Assignment