[논문리뷰] Reinforcement World Model Learning for LLM-based Agents

링크: 논문 PDF로 바로 열기

저자: Xiao Yu, Baolin Peng, Ruize Xu, Yelong Shen, Pengcheng He, Suman Nath, Nikhil Singh, Jiangfeng Gao, Zhou Yu

핵심 연구 목표

대규모 언어 모델(LLM) 기반 에이전트가 현실 환경에서 행동 결과(action consequences)를 예측하고 환경 역학에 적응하는 데 겪는 어려움을 해결하는 것을 목표로 합니다. 이는 기존 사전 훈련 목표와 에이전트 활용 사례 간의 불일치에서 비롯되며, 에이전트의 세계 모델링 능력을 향상시켜 효과적인 추론 및 계획을 가능하게 하고자 합니다.

핵심 방법론

본 논문은 강화 세계 모델 학습(Reinforcement World Model Learning, RWML) 이라는 자기 지도 학습(self-supervised) 방식을 제안합니다. 이 방법은 에이전트가 생성한 시뮬레이션된 다음 상태와 환경에서 관찰된 실제 다음 상태 간의 불일치를 측정하는 sim-to-real gap rewards 를 사용하며, 이는 사전 훈련된 임베딩 공간(embedding space) 에서의 유사성을 기반으로 합니다. 이 보상은 표준 GRPO(Shao et al., 2024) 알고리즘을 사용하여 최적화되며, 학습 데이터는 대상 모델이 환경에서 수집한 롤아웃(rollouts)을 통해 생성됩니다.

주요 결과

RWML은 ALFWorld와 τ² Bench 벤치마크에서 뛰어난 성능을 보였습니다. 순수 자기 지도 학습 방식으로도 ALFWorld에서 19.6점 , τ² Bench에서 7.9점 의 성능 향상을 달성했습니다. 특히, 작업 성공 보상(task-success reward)을 사용하는 정책 RL과 결합했을 때, ALFWorld에서 6.9점 , τ² Bench에서 5.7점 더 높은 성능을 기록하며 직접적인 작업 성공 보상 RL을 능가했습니다. 또한, 전문가 데이터 훈련과 유사하거나 그 이상의 성능을 보였고, WM SFT 대비 catastrophic forgetting 이 적음을 입증했습니다.

AI 실무자를 위한 시사점

이 연구는 LLM 기반 에이전트의 세계 모델링 능력 을 향상시키는 확장 가능하고 자기 지도 학습적인 방법론 을 제공합니다. 특히 전문가 데이터나 강력한 LLM 없이 환경과의 상호작용 데이터만으로 세계 모델을 효율적으로 학습할 수 있어, 실제 적용 시 데이터 수집 및 라벨링 비용을 크게 절감 할 수 있습니다. Sim-to-real gap rewards 는 토큰 수준의 정확도보다 의미론적 동등성에 초점을 맞춰 모델 붕괴 및 보상 해킹에 대한 강건성 을 높여, 더욱 안정적인 에이전트 개발에 기여할 수 있습니다.

⚠️ 알림: 이 리뷰는 AI로 작성되었습니다.