기술 요약: CAST (Solver Teacher로부터의 신용 할당)

문제 정의

Sokoban, Minesweeper, Rush Hour과 같은 장기적 관점(long-horizon)의 환경에서 범용 에이전트로 동작하도록 대규모 언어 모델(LLM)을 학습시키는 것은 중요한 병목 현상에 직면해 있습니다: 바로 **희소한 보상 신호(sparse reward signals)**입니다. 검증 가능한 보상을 사용하는 강화 학습(RLVR)은 단일 턴 작업에서는 성공적이었으나, 최종 결과 보상(승리/패배)에 의존하기 때문에 어떤 특정 중간 결정이 성공이나 실패를 이끌었는지에 대한 정보를 제공하지 못합니다. 이는 모든 토큰이 동일한 거친 피드로백을 받게 되어, 모델이 장기적인 계획 수립, 목표 지향적 탐색, 그리고 오류 복구 능력을 학습하는 것을 방해하는 근본적인 신용 할당(credit-assignment) 문제를 야기합니다. 이러한 신호를 조밀하게 만들기 위한 기존의 시도들(비싼 비용이 드는 탐색, 학습된 프로세스 보상 모델, 또는 궤적 간 비교 등)은 계산 비용, 감독의 품질, 그리고 신호의 신뢰성 사이에서 트레이드오프를 겪는 경우가 많습니다 도출됩니다.

방법론: CAST

저자들은 도메인 특화 게임 솔버(solver)를 활용하여 LLM 에이전트에 정교한 턴 단위 감독을 제공하는 CAST (Credit Assignment from Solver Teachers) 프레임워크를 제안합니다.

핵심 메커니즘

1. 턴 단위 교사로서의 솔버: 솔버를 단순히 전문가 궤적을 생성하는 용도로만 사용하는 대신(이는 전문가가 방문한 상태로 학습을 제한함), CAST는 솔버를 사용하여 LLM의 온폴리시(on-policy) 행동을 실시간으로 평가합니다.
2. 솔버 어드밴티지 구축:
   • 솔버는 상태 $s$로부터 승리하기 위해 필요한 최소 단계 수를 나타내는 cost-to-go 값 $N(s)$를 계산합니다.
   • LLM이 상태 $s_t$에서 $s_{t+1}$로 전이하기 위해 취한 행동 $a_t$에 대해, 이 방법은 cost-to-go의 변화량을 계산합니다.
   • Shifted solver advantage는 다음과 같이 정의됩니다:
     $$\tilde{A}^{\pi_{Solver}}(s_t, a_t) = N(s_t) - \mathbb{E}[N(s_{t+1})]$$
   • 이 값은 행동이 목표까지의 거리를 줄이면 양수, 중립적이면 0, 거리를 늘리거나 막다른 길(dead end)로 이끌면 음수가 됩니다.
3. 신호 성형(Signal Shaping): 헤비 테일(heavy-tailed) 값(예: 막다른 길에 대한 큰 페널티)으로 인한 학습 불안정성과 게임 간의 스케일 차이를 방지하기 위해, CAST는 두 가지 변환을 적용합니다:
   • Asinh 압축: $g(x) = \text{asinh}(x)$를 사용하여 작은 값은 선형적으로 보존하면서 극단적인 이상치는 로그 단위로 압축합니다.
   • 배치 수준 RMS 정규화: 일관된 스케일을 보장하기 위해 배치 내 루트 평균 제곱(RMS) 크기로 신호를 재조정합니다. 이때 "진전 없음"의 의미인 0을 보존하기 위해 평균을 빼지는 않습니다.
4. RLVR와의 통합: 성형된 솔버 어드밴티지는 표준 궤적 수준 결과 어드밴티지(예: GRPO에서 제공되는)와 결합되어 결합된 어드밴티지 신호를 형성합니다:
   $$\hat{A}_{i,t} = \hat{A}^{outcome}_i + \alpha \cdot h(\tilde{A}^{\pi_{Solver}}_{i,t})$$
   이 결합된 신호는 정책을 업데이트하기 위한 클리핑된 서로게이트 목적 함수(clipped surrogate objective)에 사용됩니다.

이론적 통찰: 로짓 프리(Logit-Free) 온폴리시 증류

본 논문은 솔버 어드밴티지를 최대화하는 것이 **소프트 최적(soft-optimal) 교사로부터의 온폴리시 증류(OPD)**와 수학적으로 동등함을 입증합니다.

• 강력한 솔버가 행동 확률이 $\exp(Q/\tau)$에 비례하는 소프트 최적 정책 역할을 한다는 가정하에, 솔버 어드밴티지는 솔버의 정책에 따른 행동의 로그 확률에 비례합니다.
• 핵심 기여: 이는 솔버의 스칼라 어드밴티지가 본질적으로 교사의 선호도 분포를 인코딩하고 있음을 의미합니다. 결과적으로, CAST는 전체 액션 분포(logits)를 요구하지 않고 오직 스칼라 값만을 사용하여 로짓 프리(logit-free) 증류를 달성합니다. 이를 통해 학생 LLM은 솔버의 가치 함수로부터 효율적으로 학습할 수 있습니다.

주요 기여

1. 솔버 유도 신용(Solver-Derived Credit): 게임 솔버의 상태 가치 변화를 LLM 학습을 위한 미세한 턴 단위 어드밴티지로 변환하는 방법을 도입했습니다.
2. 로짓 프리 증류(Logit-Free Distillation): 솔버 어드밴티지를 최대화하는 것이 전체 교사 분포가 아닌 스칼라 값만을 사용하는 온폴리시 증류와 동등하다는 이론적 증명을 제공하여 효율적인 학습을 가능하게 했습니다.
3. 안정적인 신호 성형: 극단적인 값을 처리하고 도메인 간 스케일 변화를 다루기 위해 asinh 변환과 배치 수준 RMS 정규화를 사용하여 RLVR로의 안정적인 통합을 보장했습니다.
4. 실용성: 정확한 솔버가 추가적인 오버헤드를 거의 발생시키지 않음을 보여주었으며, 정확한 솔버를 사용할 수 없는 경우 학습된 가치 네트워크(예: DQN)가 효과적인 근사치로 기능할 수 있음을 입증했습니다.

실험 결과

이 방법은 Sokoban, Minesweeper, Rush Hour에서 평가되었으며, ALFWorld 및 WebShop에 대한 제로샷 전이 성능도 확인되었습니다.

• 도메인 내 성능: CAST는 세 가지 게임 모두에서 인도메인(in-domain) 및 미보지 난이도(unseen difficulty) 설정 모두에서 모든 훈련된 베이스라인(GRPO, GSPO, DAPO, GiGPO)과 훈련이 필요 없는 프롬프팅 방법(Claude-Opus-4.6과 같은 강력한 폐쇄형 모델 포함)을 능가했습니다.
  • 평균적으로 CAST는 인도메인 작업에서 **62.1%**의 성공률(Avg@4)을 기록했으며, 미보지 난이도에서는 **28.4%**를 기록하여 백본인 DAPO(ID 44.7%, Unseen 18.7%)를 크게 앞질렀습니다.
• 샘플 효율성: CAST는 세 가지 게임 모두에서 DAPO의 피크 검증 성능에 도달하는 데 1.7배에서 2.0배 적은 훈련 단계를 소요했습니다.
• 제로샷 일반화: 추가 파인튜닝 없이, CAST 에이전트는 ALFWorld와 WebShop으로 전이되어 모든 방법 중 가장 높은 평균 제로샷 성능(전체 30.3%)을 달성했으며, 이는 차순위 베이스라인보다 5.6포인트 높았습니다.
• 절제 연구(Ablation Studies):
  • 솔버 어드밴티지 가중치($\alpha$)가 매우 중요했습니다. $\alpha=0.1$이 프로세스 가이드와 결과 목적 사이의 최적의 균형을 제공했습니다.
  • asinh 압축과 RMS 정규화 모두 필수적이었습니다. 둘 중 하나라도 제거하면 성능이 저하되었으며, 특히 asinh를 제거했을 때 극단적인 값으로 인한 그래디언트 불안정성 때문에 가장 큰 성능 하락이 발생했습니다.
• 학습된 솔버: 정확한 솔버를 학습된 DQN 가치 네트워크로 대체해도 대부분의 성능 이점이 유지되었으며, 이는 이 방법이 정확한 솔버가 존재하는 게임을 넘어 적용 가능하다는 것을 보여줍니다.

의의 및 주장

본 논문은 도메인 솔버로부터 유도된 확장 가능하고 미세한 신호를 제공함으로써, 장기적 관점의 RLVR에서 발생하는 근본적인 신용 할당 문제를 해결한다고 주장합니다. 솔버 어드van티지를 최대화하는 것이 로짓 프리 온폴리시 증류와 동등하다는 것을 증명함으로써, CAST는 전체 액션 분포를 생성하는 계산적 부담 없이 전문가 가치 함수와 LLM을 정렬할 수 있는 실용적인 경로를 제시합니다. 결과는 정확한 솔버나 학습된 가치 네트워크를 통한 신뢰할 수 있는 상태 평가가, 희소한 결과로부터 조밀하고 실행 가능한 피드백으로 나아가는 강력한 경로를 제공하며, LLM이 결과 중심의 접근 방식보다 더 효율적으로 복잡한 다단계 의사결정 과제를 마스터할 수 있게 함을 시사합니다.