TAPO: LLM 에이전트를 위한 전환 인식 정책 최적화

TAPO: Transition-Aware Policy Optimization for LLM Agents

최근, 강화 학습(RL)은 대규모 언어 모델(LLM) 에이전트의 추가 훈련을 위한 중요한 패러다임으로 부상했습니다. 그러나 기존 방법들은 주로 정책 최적화를 위해 희소한 작업 보상에 의존하며, 온라인 상호 작용 중에 자연스럽게 존재하는 또 다른 형태의 밀집형 지도 신호인 행동 실행 후 발생하는 환경 피드백을 충분히 활용하지 못합니다. 최근의 이론 연구에서는 다단계, 목표 지향적인 작업에서 일반화가 환경 결과에 대한 예측적 지식에 달려 있다는 점이 밝혀졌습니다. 이러한 영감을 받아, 우리는 정책 최적화와 전환 감독을 번갈아 수행하는 통합 훈련 프레임워크인 LLM 에이전트를 위한 전환 인식 정책 최적화(TAPO)를 제안합니다. TAPO는 표준적인 강화 학습 업데이트 외에도, 롤아웃 데이터를 재활용하여 공유된 기반 모델에 행동 조건부 다음 관측치 예측 감독을 적용합니다. 이 접근 방식은 모델의 환경 전이 역학 및 행동 결과에 대한 민감도를 향상시키면서 동시에 정책을 최적화합니다. TAPO는 기존 에이전트 강화 학습 알고리즘에 추가적인 전문가 데이터, 샘플링 비용 또는 추론 시간 오버헤드가 필요 없는 경량의 플러그 앤 플레이 확장 모듈입니다. 우리는 다양한 크기의 기초 모델과 다양한 정책 최적화 알고리즘을 사용하여 WebShop 및 ALFWorld에서 체계적인 실험을 수행했습니다. 경험적 결과는 TAPO가 순수한 정책 최적화 기준보다 일관되게 작업 성능을 향상시킨다는 것을 보여줍니다.

Recently, Reinforcement Learning (RL) has emerged as a crucial paradigm for the post-training of Large Language Model (LLM) agents. However, existing methods predominantly rely on sparse task rewards for policy optimization, failing to fully exploit another class of inherently dense supervisory signals naturally present during online interaction: environmental feedback following action execution. Recent theoretical studies suggest that generalization in multi-step, goal-oriented tasks hinges on predictive knowledge of environmental consequences. Inspired by this, we propose TAPO: Transition-Aware Policy Optimization for LLM Agents, a unified training framework that alternates between policy optimization and transition supervision. Beyond standard RL updates, TAPO repurposes rollout data to apply action-conditioned next-observation prediction supervision on a shared backbone model. This approach enhances the model's sensitivity to environmental transition dynamics and action consequences while concurrently optimizing the policy. It serves as a computationally lightweight, plug-and-play enhancement module for existing agent RL algorithms, requiring no additional expert data, extra sampling costs, or inference-time overhead. We conduct systematic experiments on WebShop and ALFWorld, integrating foundation models of various scales with different policy optimization algorithms. Empirical results demonstrate that TAPO consistently improves task performance over pure policy optimization baselines.