강화 기반 세계 모델 학습으로 LLM 에이전트 능력 향상

초록

본 논문은 LLM 기반 에이전트가 환경 변화를 예측하고 행동 결과를 사전 추론하도록 돕는 ‘Reinforcement World Model Learning(RWML)’을 제안한다. RWML은 시뮬레이션된 다음 상태와 실제 관찰된 다음 상태를 사전 학습된 임베딩 공간에서 코사인 유사도로 비교해 보상으로 활용한다. 토큰 수준의 정확성을 강요하는 기존 SFT 방식과 달리 의미적 일관성을 목표로 하며, 전문가 데이터나 별도의 성공 보상 없이도 자체 수집한 롤아웃으로 학습한다. ALFWorld와 τ² Bench에서 베이스 모델 대비 각각 19.6점·7.9점 향상했으며, 성공 보상과 결합했을 때 기존 정책 RL보다 5~7점 앞섰다. 또한, 모델 망각이 적고 LLM‑as‑a‑judge 방식보다 해킹에 강인함을 보였다.

상세 분석

RWML은 LLM을 “행동‑조건부 세계 모델”로 전환시키기 위해 강화학습 프레임워크를 차용한다. 핵심 아이디어는 에이전트가 현재 상태와 행동 히스토리를 입력받아 ‘reason’ 토큰을 생성한 뒤, 다음 상태(ˆsₜ₊₁)를 예측하도록 학습시키는 것이다. 예측된 텍스트와 실제 환경에서 관측된 텍스트 사이의 차이는 사전 학습된 임베딩 모델(E)로 인코딩한 뒤 코사인 거리로 측정한다. 거리 d가 사전 정의된 임계값 τ_d 이하이면 보상 r_WM=1, 그렇지 않으면 0을 부여한다. 이렇게 이진 보상을 사용하면 보상 해킹 위험을 크게 낮출 수 있다.

보상 최적화는 최신 정책 최적화 기법인 GRPO(Generalized PPO)를 적용한다. 중요도 비율 ρ와 KL 정규화 항을 포함해 안정적인 업데이트를 보장한다. 또한, “너무 쉬운” 샘플을 제외하는 서브샘플링 전략을 도입한다. 구체적으로, 별도 SFT 모델(π′)을 이용해 전체 데이터의 10%만으로 초기 학습을 수행하고, 나머지 90%에 대해 π′가 높은 보상을 지속적으로 얻는 경우(≥K번 시도) 이를 ‘easy’로 판단해 확률 p=0.1만큼만 학습에 포함한다. 이 과정은 데이터 다양성을 유지하면서도 모델이 의미 있는 전이 학습을 하도록 만든다.

실험에서는 두 가지 벤치마크를 사용했다. ALFWorld는 텍스트 기반 가상 가정 환경으로, 에이전트는 “sidetable 1에 가라”와 같은 명령을 수행하고 결과 텍스트를 관찰한다. τ² Bench은 도구 호출과 사용자와의 대화를 포함한 복합 환경으로, 상태는 JSON 형식 도구 응답이나 자연어 응답이 될 수 있다. 두 환경 모두 장기 계획과 정확한 상태 추론이 필수적이다.

RWML은 전문가 롤아웃이나 강력한 LLM을 전혀 사용하지 않고, 오직 자체 수집한 롤아웃만으로도 베이스 모델 대비 19.6점(ALFWorld)·7.9점(τ² Bench) 상승을 기록했다. 성공 보상(Policy RL)과 결합했을 때는 기존 정책 기반 강화학습보다 각각 6.9점·5.7점 더 높은 점수를 얻었으며, 전문가 데이터 기반 SFT와 거의 동등한 성능을 보였다.

추가 분석에서는 모델 망각(forgetting) 현상을 평가했다. MMLU‑Redux, IFEval, GSM8k 등 다양한 일반·수학·코딩 벤치마크에 대해 RWML이 WM‑SFT보다 거의 모든 경우에서 더 적은 성능 저하를 보였다. 이는 온라인 RL이 지속적인 정책 업데이트를 통해 기존 지식을 보존하는 특성 때문으로 해석된다.

마지막으로, LLM‑as‑a‑judge 방식과 비교했을 때 RWML은 보상 설계가 단순히 임베딩 유사도에 기반하므로 “보상 해킹”에 대한 취약성이 현저히 낮다. 이는 에이전트가 텍스트 표면 형태가 아닌 의미적 일관성을 학습하도록 유도한다는 점에서 중요한 장점이다.

요약하면, RWML은 (1) 완전 자기 지도(self‑supervised) 방식, (2) 의미 기반 보상 설계, (3) 쉬운 샘플 필터링을 통한 효율적 데이터 활용, (4) 기존 정책 RL과의 시너지 효과라는 네 가지 핵심 요소를 결합해 LLM 기반 에이전트의 세계 모델링 능력을 크게 향상시킨다.