행동 병목 현상 해결: 토큰 수준 에너지 정보를 활용한 에이전트 기반 강화 학습

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

에이전트 기반 강화 학습은 대규모 언어 모델을 훈련하는 데 사용되며, 이는 긴 추론 과정과 짧은 환경 상호 작용을 번갈아 가며 진행되는 다중 턴 트레이젝토리(trajectory)를 활용합니다. PPO 및 GRPO와 같은 일반적인 정책 경사(policy-gradient) 방법은 트레이젝토리 내의 모든 토큰을 동일하게 취급하여 균일한 신용 할당(credit assignment)을 수행합니다. 본 논문에서는 이러한 균일한 신용 할당이 토큰 수준의 훈련 신호를 크게 잘못 할당한다는 것을 비판적으로 보여줍니다. 에너지 기반 모델링 관점에서, 특정 프롬프트에서 샘플링된 다양한 롤아웃(rollout)의 보상 분산과 상관 관계를 통해 정량화되는 토큰 수준의 훈련 신호가 추론 토큰이 아닌 행동 토큰에 집중되는 경향이 있으며, 이는 트레이젝토리에서 행동 토큰이 차지하는 비율이 매우 작음에도 불구하고 나타나는 현상입니다. 우리는 이 현상을 '행동 병목(Action Bottleneck)' 현상이라고 명명합니다. 이러한 관찰에 기반하여, 우리는 추론 토큰의 그래디언트를 감소시키고, 더 높은 불확실성을 가진 행동 토큰의 가중치를 더욱 증가시키는 추가적인 에너지 기반 재분배 메커니즘과 함께, 매우 간단한 토큰 재가중(reweighting) 접근 방식인 ActFocus를 제안합니다. 4가지 환경 및 다양한 모델 크기를 대상으로 ActFocus는 PPO 및 GRPO보다 일관되게 우수한 성능을 보였으며, 각각 최대 65.2% 및 63.7%의 최종 단계 성능 향상을 달성했습니다. 이러한 성능 향상은 추가적인 런타임 또는 메모리 비용 없이 이루어졌습니다.

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.