TAPO: Transition-Aware Policy Optimization for LLM Agents

2026-07-30

AI Summary

This work addresses the limitation of existing large language model agents in reinforcement learning, which often rely on sparse task rewards and fail to effectively exploit dense supervisory signals from environmental feedback. To overcome this, the authors propose a unified training framework that alternates between policy optimization and action-conditional next-state prediction, applying self-supervision from environment transition dynamics to a shared backbone. This enhances the agent’s awareness of action consequences and environmental dynamics. The approach is implemented as a lightweight, plug-and-play module compatible with existing reinforcement learning algorithms, requiring no additional expert data, sampling overhead, or inference cost. Experiments on WebShop and ALFWorld benchmarks demonstrate consistent and significant improvements over pure policy optimization baselines across diverse base model scales and policy learning algorithms.

Abstract

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.

Problem

• Reinforcement Learning
• Large Language Model Agents
• Sparse Rewards
• Environmental Feedback
• Transition Dynamics

Innovation

• Transition-aware learning
• Policy optimization
• Reinforcement learning for LLMs
• Action-conditioned prediction
• Dense supervisory signals