탐색 인지형 정책 최적화를 통한 에이전트 주도 추론 확장 학습

에이전트가 환경 정보를 효율적으로 수집하려면 탐색과 실행을 구분하고 탐색의 필요 시점을 학습으로 판단해야 한다. 본 연구는 exploration-aware 보상과 memory를 통한 외부화된 기억으로 test-time에서 정보 획득과 의사결정을 분리하고, 불확실성이 큰 상태에서만 탐색하는 전략을 학습한다. 4개 벤치마크에서 기존 방법 대비 일관된 성능 향상을 보여주고, 2B 규모 모델로도 상당한 성능을 달성한다.

출발점: 에이전트-강화학습에서 목적은 보상을 최대화하는 정책 학습이다. 기존의 test-time 탐색 접근은 탐색과 action selection을 하나의 정책으로 묶는 경우가 많아 환경의 불확실성에 따른 탐색 필요성을 구분하기 어렵다. 본 논문은 탐색과 기억의 분리된 모듈을 도입하고, Bayesian 기반의 탐색 유용성 보상(R_explore)을 학습 가능한 q_phi로 추정한다. 이를 통해 정보 획득의 가치를 평가하고, visitation-depth를 활용한 그룹화로 학습 신호를 명확히 구분한다. 실험에서는 ALFworld, WebShop, AndroidWorld, OSWorld에서 20-60%의 성능 향상과 2B-크기의 모델이 더 큰 모델 대비 우수한 성능을 달성함을 확인한다. 또한 SFT rollback과 두 단계 학습, 탐색-정보-실행 간 credit assignment의 구분이 안정적 학습과 탐색 효율성을 높임을 보인다.

전체 접근은 두 축으로 구성된다. 1) 탐색-메모리 모드: st = [g; s_t; e_{t-1}; m_{t-1}] 형태의 입력에서 탐색(et)과 기억(mt)을 함께 생성하고, 최종 실행은 at를 포함한다. 수식적으로 at = π_theta(·|st)로 정의한다. 2) 보상 모델과 학습 프로토콜: 총 보상은 R = R_task + α1 R_format + α2 R_explore로 구성되며, R_explore는 베이지안 후방의 정보가치에 의존한다. q_phi(e_{t-1}, m_{t-1}|s)를 추정하여 KL 발산을 최소화하는 방식으로 학습한다. R_explore은 Eq. (9) 및 Eq. (10)에서 정의된다. 학습은 두 단계로 수행되며, 5.1절의 rollback 학습(SFT)과 5.2절의 탐색-인지 정책 최적화로 구성된다. 롤백 데이터셋 D를 구성하고 LSFT를 최소화하여 rollback 능력을 학습한 뒤, τ1...τG를 생성하여 현 상태-방문 깊이 κ에 따라 G(s,ν)로 그룹화한다. 각 그룹에서 Ã_{i,j}를 계산하여 정책 업데이트를 진행한다. 알고리즘은 Algorithm 1에 명시된 흐름을 따른다.

주요 벤치마크에서 EAPO가 baselines보다 일관되게 우수한 성능을 보인다. 텍스트 기반 ALFworld, WebShop, AndroidWorld, OSWorld에서 EAPO가 20%~60%의 절대 성능 향상을 달성하며, 2B 규모 모델이 더 큰 일반/에이전트 모델을 능가하는 경우가 다수이다. 그룹 크기 G를 증가시킬수록 안정성이 증가하지만 수렴 이점은 포화되며, KL 계수 λ은 중간 값에서 가장 잘 작동하고 지나치게 큰 경우 제약이 커져 성능이 떨어진다. γ=0.9로 설정하면 탐색과 실행 간의 균형이 적절해져 탐색 효율성과 실행 성능의 트레이드오프를 잘 맞춘다. SFT 롤백 도입으로 학습 안정성이 개선되며, exploration-aware grouping이 도입되면 탐색Degree와 성능의 상승이 유지된다. OSWorld에 대해 AndroidWorld에서 학습한 모델의 일반화가 확인되었고, 추론 시간 증가도 탐색의 유의미한 증가에 국한된다. 전체 실험은 8-bit GPU 기반 인프라에서 수행되었으며 훈련 오버헤드는 기존 방법 대비 대체로 제한적이다.

아키텍처는 탐색-메모리 모드와 보상 학습 모듈로 구성된다. s~t 입력은 [g; s_t; e_{t-1}; m_{t-1}]이며, 출력은 [e_t; m_t; a_t] 형태의 결합으로 생성된다. 보상은 R_task, R_format, R_explore의 가중합으로 구성되며, R_explore는 Eq. (9) 및 (10)에서 정의된다. R_explore의 추정은 q_phi(e, m|s)를 사용한 변분 추정으로, KL(q_phi || p) 최소화를 목표로 한다. SFT 롤백은 LSFT 손실로 학습되며, 4단계 그룹화(G(s,ν))를 통해 같은 초기 상태와 방문 깊이에서의 트랜지션을 묶어 상대 가치를 추정한다. κ(s_i) = ∑{k<t} I[s_i^k = s_i^t]로 정의되며, 그룹 내 이점은 Ã{i,j} = (R(...)-mean(R_G))/std(R_G)로 계산한다. 실험 설정은 4개 도메인(ALFWorld, WebShop, AndroidWorld, OSWorld)에서 Qwen3 또는 Qwen3-VL 기반 모델 크기로 수행되며 γ=0.9, 그룹 크기 G=16, α1=0.5, α2=1.0, KL λ=0.01, 학습률 1e-4를 기본값으로 사용한다.