RAGEN-2: 에이전트 강화학습에서의 추론 붕괴 현상 분석

멀티턴 AI 에이전트를 강화학습으로 훈련할 때, 모델이 입력에 상관없이 고정된 답변 패턴만 반복하는 '템플릿 붕괴' 현상을 발견하고 이를 해결할 수 있는 새로운 진단 지표와 학습 기법을 제시한다. 보상 분산을 활용해 유의미한 학습 신호가 있는 데이터만 선별함으로써 추론의 질과 작업 성공률을 동시에 높일 수 있다.

강화학습에서 모델은 보상을 최대화하는 방향으로 가중치를 업데이트한다. 이때 '엔트로피'는 모델이 얼마나 다양한 답변을 내놓는지를 측정하지만, 그 답변이 '질문에 맞는 답변'인지는 알려주지 않는다. Transformer 기반 모델이 훈련될 때, 특정 질문에 대한 정답 보상이 명확하지 않으면 모델은 보상을 얻기 위해 노력하기보다 KL Divergence나 엔트로피 보너스 같은 '정규화(Regularization)' 규칙을 지키는 데만 집중하게 된다.

이 과정에서 모델은 '질문이 무엇이든 상관없이' 문법적으로 완벽하고 겉보기에 논리적인 고정된 답변 틀(Template)을 생성하기 시작한다. 이는 마치 시험 문제의 내용을 이해하지 못한 학생이 어떤 질문에도 통용될 법한 모범 답안의 형식만 외워서 써내는 것과 같다. 본 논문은 이러한 현상이 발생하는 근본 원인이 '보상의 차이(분산)'가 너무 작아서 모델이 무엇이 옳은 방향인지 학습할 신호를 찾지 못하기 때문임을 밝혀냈다.

해결책은 간단하면서도 강력하다. 학습 과정에서 여러 번의 시도를 해본 뒤, 시도들 간의 점수 차이가 큰 질문들만 골라서 공부하는 것이다. 점수 차이가 크다는 것은 '어떤 행동이 좋고 나쁜지'가 명확하다는 뜻이며, 이를 통해 모델은 정규화 노이즈에 휘둘리지 않고 실제 문제 해결에 필요한 논리를 학습할 수 있게 된다. 결과적으로 모델은 고정된 틀에서 벗어나 각 입력 상황에 맞는 유연한 추론 능력을 회복한다.

전체적인 접근 방식은 에이전트 강화학습 과정에서 발생하는 그래디언트 불균형을 해결하기 위해 보상 분산 기반의 필터링 루프를 추가하는 것이다. 먼저 상호 정보량(MI)을 추정하기 위해 In-Batch Cross-Scoring 기법을 사용한다. [P개의 프롬프트와 G개의 추론 샘플을 입력으로] → [모든 조합에 대해 모델의 로그 우도(Log-likelihood)를 계산하여] → [자신의 프롬프트에서 나온 확률과 다른 프롬프트에서 나온 확률의 차이를 구하고] → [이 값이 클수록 입력 의존성이 높다고 판단한다].

핵심 메커니즘인 SNR-Aware Filtering은 매 훈련 반복마다 수행된다. [각 프롬프트별로 생성된 G개 궤적의 보상값들을 입력으로] → [표본 분산(Sample Variance)을 계산하여] → [분산이 높은 순서대로 프롬프트를 정렬한 뒤] → [상위 p%의 누적 분산을 가진 프롬프트 그룹만 선택한다]. 이 과정은 수식적으로 |g_task| ≤ sqrt(Var(R|X)) * C 관계에 기반하며, 보상 분산이 0에 가까우면 과업 수행을 위한 그래디언트 신호가 사라지고 정규화 노이즈만 남게 된다는 원리를 이용한다.

학습 전략으로는 Top-p 필터링을 사용하여 분산 분포에 따라 동적으로 학습 데이터 양을 조절한다. 이는 고정된 개수를 뽑는 Top-k 방식보다 유연하며, 전체 배치의 신호가 약할 경우 업데이트를 건너뛰는 안전장치 역할도 수행한다. 구현 시에는 veRL 프레임워크를 기반으로 Qwen2.5-3B 모델을 사용하였으며, PPO, GRPO 등 다양한 RL 알고리즘에 적용 가능함을 입증했다.

메인 벤치마크 결과, Sokoban(계획), FrozenLake(내비게이션), MetaMathQA(수학) 등 7개 환경에서 SNR-Aware Filtering을 적용했을 때 성공률이 일관되게 향상되었다. 특히 Sokoban 환경에서 Qwen2.5-3B 모델 기준, 기본 PPO 대비 성공률이 12.9%에서 28.9%로 약 16%p 상승하는 성과를 거두었다.

상호 정보량(MI) 분석 결과, 기존 엔트로피 지표는 모델의 성능과 상관관계가 거의 없거나 음수(-0.11 ~ -0.14)를 보인 반면, 제안된 MI 프록시는 성능과 높은 양의 상관관계(+0.39)를 보였다. 이는 MI가 추론 품질을 측정하는 훨씬 더 신뢰할 수 있는 지표임을 증명한다.

효율성 분석에서는 그룹 샘플링(G≥2)을 통한 분산 계산 오버헤드가 전체 반복 시간의 0.1% 미만으로 매우 적음을 확인했다. 또한 필터링을 통해 그래디언트 계산 대상이 줄어들면서 단계별 학습 시간(Step Time)이 26%에서 41%까지 단축되는 부수적인 이득도 얻었다.

본 연구는 에이전트 RL의 그래디언트를 과업 관련 성분(g_signal + g_task-noise)과 정규화 성분(g_reg)으로 분해하여 분석한다. 핵심 이론적 기여는 과업 그래디언트의 노름(norm)이 보상 분산의 제곱근에 의해 상한(upper bound)이 결정된다는 것을 증명한 것이다. 보상 분산이 낮아지면 g_signal은 0으로 수렴하지만, KL이나 엔트로피 정규화에 의한 g_reg는 일정하게 유지되어 업데이트 방향이 입력 무관한 방향으로 편향된다.

아키텍처 측면에서는 별도의 보조 모델 없이 훈련 중인 정책 모델 자체를 Scorer로 활용하는 In-Batch Cross-Scoring 방식을 채택했다. 이는 계산 효율성을 극대화하면서도 정보 이론적 지표인 상호 정보량을 효과적으로 근사한다. 특히 MI-ZScore-EMA 프록시는 배치 간 변동성을 안정화하여 훈련 모니터링의 견고함을 높였다.

Prior work인 DAPO나 GRPO가 안정적 훈련을 위해 정규화 항을 조정하는 데 집중했다면, RAGEN-2는 업데이트 신호 자체의 품질을 필터링하는 SNR 제어 메커니즘을 제안했다는 점에서 차별화된다. 이는 정규화 항의 강도와 무관하게 유효 신호가 부족한 업데이트를 원천 차단함으로써 파라미터 드리프트(drift)를 방지한다.

SNR 분해 시 과업 신호와 정규화 노이즈가 완전히 분리된다고 가정했으나 실제로는 그래디언트 누적 과정에서 결합될 가능성이 있다. 또한 모든 실험이 단일 에이전트 환경에서 수행되어 멀티 에이전트 설정에서의 템플릿 붕괴 전파 양상은 아직 미지로 남아있다. 보상 신호가 극도로 희소하거나 노이즈가 심한 환경에서는 보상 분산이 신뢰할 수 있는 프록시 역할을 하기 어려울 수 있다.