TreeOPSD, 트리 구조와 사후 증류로 장기 에이전트 강화학습 개선

배경

언어 모델 에이전트가 검색, 추론, 도구 호출을 여러 차례 수행하는 장기 작업에서는 최종 답변의 정확성만으로 학습하기 어렵습니다. 검증 가능한 보상이 작업 종료 시점에만 주어지면, 성공을 이끈 중간 행동이 무엇인지 알기 힘들기 때문입니다. 여기에 하나의 rollout 배치를 여러 mini-batch에 걸쳐 재사용하면, 현재 학습 중인 정책이 데이터를 생성한 행동 정책에서 점점 멀어지는 문제도 발생합니다.

TreeOPSD(Tree-Aware Posterior Distillation)는 이 두 문제를 비평가 모델 없이 함께 다루려는 프레임워크입니다. 별도의 가치 네트워크, 턴별 보상 주석, 더 강한 외부 교사를 추가하는 대신, 에이전트 궤적에 나타나는 분기 구조를 학습 신호로 활용합니다.

핵심 아이디어

• 상호작용 턴을 트리 노드로 구성. 같은 접두사를 공유하는 상태에서 여러 후속 행동을 샘플링합니다. 이를 통해 서로 다른 전체 궤적만 비교하는 것이 아니라, 유사한 문맥에서 선택된 행동들을 비교할 수 있습니다.
• 종료 보상을 재귀적으로 역전파. 트리의 말단에서 얻은 결과를 상위 노드로 전달해 턴 단위의 암묵적 기여도를 계산합니다. 종료 보상만 있는 환경에서도 초기 행동에 더 촘촘한 신호를 제공하는 방식입니다.
• 제한된 교사 분포 생성. 각 턴의 기여도와 고정된 행동 정책 대비 우도비를 이용해 폐쇄형 교사 분포를 만듭니다. 높은 보상으로 이어진 행동을 선호하지만, 특정 행동의 확률을 무한히 키우도록 유도하지는 않습니다.
• 정책을 교사에 증류. 학생 정책과 교사 분포 사이의 정방향 KL 발산을 최소화합니다. 행동 정책 주변에서는 이 목적의 접선 방향이 트리 정책 그래디언트와 일치하며, 반복 업데이트에서는 유한한 목표를 향하도록 설계됩니다.

결과와 의미

실험은 다중 홉 질의응답과 심층 연구 벤치마크에서 진행됐으며, 1.7B부터 30B-A3B 규모의 모델을 포함합니다. 동일한 rollout 예산을 맞춘 비교에서 TreeOPSD는 GRPO, GSPO, 고정 기여도를 사용하는 트리 정책 그래디언트 기준선을 일관되게 앞섰다고 보고됩니다. 심층 연구 설정에서 Qwen3-30B-A3B-Thinking-2507은 평균 4.97점 향상됐다는 결과가 제시됐습니다.

TreeOPSD의 핵심적인 의미는 신용 할당과 정책 드리프트를 하나의 목표 설계 안에서 연결했다는 데 있습니다. 중요도 샘플링과 비율 클리핑은 행동 정책과 현재 정책 사이의 차이를 줄이고 업데이트 폭을 제한할 수 있지만, 반복 업데이트를 어디에서 멈출지에 대한 유한한 목표를 직접 제시하지는 않습니다. TreeOPSD는 제한된 교사 분포를 통해 학습이 향해야 할 도착점을 명시하려고 합니다.

다만 여러 분기를 생성하면 추가 계산 비용이 발생하며, 분기 수와 접두사 공유 방식이 효율성에 영향을 줄 수 있습니다. 또한 검증 가능한 종료 보상에 의존한다는 전제도 있습니다. 그럼에도 에이전트 궤적을 단순한 정책 그래디언트 샘플이 아니라 구조화된 자기 증류의 재료로 활용한다는 발상은 장기 에이전트 강화학습의 중요한 연구 방향을 제시합니다.

출처: Hugging Face Daily Papers