[논문리뷰] Policy and World Modeling Co-Training for Language Agents

링크: 논문 PDF로 바로 열기

메타데이터

저자: Ning Lu, Baijiong Lin, Shengcai Liu, et al.

---

1. Key Terms & Definitions (핵심 용어 및 정의)

• PaW (Policy and World modeling Co-training): On-policy RL 과정 중에 auxiliary world modeling supervision을 결합하여, 추가적인 모델이나 학습 단계 없이 정책 모델의 성능을 향상시키는 프레임워크입니다.
• WM (World Modeling): 에이전트가 수행한 Action이 환경에 미치는 결과(next observation)를 예측함으로써, 단순 보상 학습을 넘어 환경의 동역학(dynamics)을 내재화하는 기술입니다.
• CMAE (Clipped Mean Absolute Error): 고확률 토큰을 무시하고 불확실한 observation 예측에 집중하며, noisy한 데이터로부터 발생하는 gradient 폭주를 방지하기 위해 설계된 WM loss 함수입니다.
• Action-Entropy-based Selection: 에이전트가 결정 지점에서 내린 선택의 엔트로피를 기준으로, 모델 학습에 가장 유익한 Transition 데이터만을 필터링하여 WM supervision의 효율을 높이는 기법입니다.

2. Motivation & Problem Statement (연구 배경 및 문제 정의)

본 논문은 LLM Agent가 표준 RL 학습 과정에서 보상 최적화에만 치중하여 환경의 결과 예측 능력을 결여하는 문제를 해결합니다. 기존 연구들은 별도의 시뮬레이터나 복잡한 다단계 학습, 혹은 추론 시 추가 연산을 요구하여 시스템 복잡도를 높이는 한계가 있었습니다. 저자들은 On-policy RL의 Rollout 데이터 자체가 Action과 그에 따른 next observation이라는 WM supervision을 이미 내포하고 있다는 점에 주목합니다. 이를 활용하여 별도의 추가 과정 없이 정책 학습과 World modeling을 동시에 수행하는 통합적인 프레임워크를 제안합니다 [Figure 1].

3. Method & Key Results (제안 방법론 및 핵심 결과)

PaW는 RL 학습의 base objective에 Reward-adaptive한 WM loss를 추가한 통합 손실 함수를 사용합니다 [Figure 3]. 이 방법론은 Action-entropy를 통해 유익한 데이터를 선별하고, noisy한 observation에 강건한 CMAE loss를 적용하며, 보상 상황에 맞춰 WM loss의 비중(λWM)을 조절하여 안정적인 학습을 도모합니다 [Table 4]. 실험 결과, ALFWorld 및 WebShop 환경에서 GRPO 및 GIGPO 기반 성능을 일관되게 상회하였습니다 [Table 1]. 특히, 보상이 희소한 환경에서 PaW는 Vanilla RL이 실패하는 상황에서도 안정적인 학습 신호를 제공하여 성공률을 비약적으로 높이는 결과를 보여주었습니다 [Figure 4]. 또한, 추가적인 연산 오버헤드는 전체 학습 시간 및 GPU 메모리의 약 2% 수준으로 매우 경제적입니다 [Figure 5].

4. Conclusion & Impact (결론 및 시사점)

본 연구는 On-policy RL rollout을 활용한 PaW 프레임워크를 통해 별도의 인프라 없이도 LLM Agent의 환경 예측 능력을 효과적으로 강화할 수 있음을 입증하였습니다. 이 접근법은 보상 기반 학습이 실패하기 쉬운 복잡한 환경에서 강력한 보조 신호를 제공하여 알고리즘의 견고성을 확보합니다. 향후 다양한 Agentic 작업에서 모델의 효율성과 성능을 동시에 개선하는 표준적인 학습 전략으로 활용될 것으로 기대됩니다.

---

Part 2: 중요 Figure 정보

Figure 1 — 기존 WM 방식 대비 PaW의 통합 구조

Figure 3 — PaW 전체 워크플로우

Figure 4 — 희소 보상 환경에서의 학습 효과

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