Proximity-Based Multi-Turn Optimization: Practical Credit Assignment for LLM Agent Training

Yangyi Fang; Jiaye Lin; Xiaoliang Fu; Cong Qin; Haolin Shi; Chang Liu; Peilin Zhao

Source-linked AI summary

Proximity-Based Multi-Turn Optimization: Practical Credit Assignment for LLM Agent Training

Yangyi Fang, Jiaye Lin, Xiaoliang Fu, Cong Qin, Haolin Shi, Chang Liu, Peilin Zhao

arXiv:2602.19225v1Feb 22, 2026cs.AI

TL;DR

Multi-turn agent training must distinguish informative outcomes from stochastic noise, but group-based methods can misallocate credit when task difficulty varies. ProxMO adds success-rate-aware episode modulation and proximity-based step aggregation, and experiments on ALFWorld and WebShop report consistent improvements with independent and synergistic contributions from both mechanisms.

Problem
Existing group-based optimization relies on within-group statistical deviation, overlooking context-dependent informational value when task difficulty varies.
Method
ProxMO combines success-rate-aware modulation for episode-level credit with proximity-based soft aggregation using continuous semantic weighting at the step level.
Results
Experiments on ALFWorld and WebShop demonstrate consistent improvements, while ablations confirm independent and synergistic contributions from both mechanisms.
Takeaways & Limitations
ProxMO provides a practical framework for context-dependent credit assignment in multi-turn reinforcement learning with LLM agents.
Takeaways & Limitations
Evaluation is concentrated on resource-efficient 1.5B and 7B models, leaving generalizability to significantly larger foundation models for future validation.

Abstract

from arXiv · show

Multi-turn LLM agents are becoming pivotal to production systems, spanning customer service automation, e-commerce assistance, and interactive task management, where accurately distinguishing high-value informative signals from stochastic noise is critical for sample-efficient training. In real-world scenarios, a failure in a trivial task may reflect random instability, whereas success in a high-difficulty task signifies a genuine capability breakthrough. Yet, existing group-based policy optimization methods rigidly rely on statistical deviation within discrete batches, frequently misallocating credit when task difficulty fluctuates. To address this issue, we propose Proximity-based Multi-turn Optimization (ProxMO), a practical and robust framework engineered specifically for the constraints of real-world deployment. ProxMO integrates global context via two lightweight mechanisms: success-rate-aware modulation dynamically adapts gradient intensity based on episode-level difficulty, while proximity-based soft aggregation derives baselines through continuous semantic weighting at the step level. Extensive evaluations on ALFWorld and WebShop benchmarks demonstrate that ProxMO yields substantial performance gains over existing baselines with negligible computational cost. Ablation studies further validate the independent and synergistic efficacy of both mechanisms. Crucially, ProxMO offers plug-and-play compatibility with standard GRPO frameworks, facilitating immediate, low-friction adoption in existing industrial training pipelines. Our implementation is available at: \href{https://anonymous.4open.science/r/proxmo-B7E7/README.md}{https://anonymous.4open.science/r/proxmo}.

1 Introduction

Multi-turn agents face context-dependent credit assignment because task difficulty varies, while existing group-based methods rely on within-group statistics. ProxMO addresses this with episode-level success-rate modulation and step-level proximity-based soft aggregation.

Task difficulty makes identical statistical deviations carry different informational value across multi-turn episodes.A failure in a high-success task may be noise, whereas a success in a low-success task may signal a breakthrough.
Existing group-based optimization methods can misallocate learning signals by ignoring context-dependent outcome value.
ProxMO incorporates global context into credit assignment at both episode and step levels.
Success-rate-aware modulation adapts episode-level gradient intensity to task difficulty, while proximity-based aggregation continuously weights states by semantic proximity.
Experiments on ALFWorld and WebShop show effectiveness, with ablations indicating independent and synergistic contributions from both mechanisms.

2 Preliminaries

Multi-turn LLM-agent tasks are modeled as sequential interactions that generate trajectories and returns. GRPO samples trajectories under the same task and estimates advantages from group statistics without explicit value networks.

An LLM policy interacts with a dynamic environment over multiple steps, observing states, generating actions, receiving rewards, and transitioning between states.
A complete episode is represented as a trajectory containing state, action, and reward tuples across its horizon.
The training objective is to find a policy that maximizes expected trajectory return.
GRPO samples N trajectories under the same task instruction and initial state, then derives advantages from group statistics.
GRPO avoids explicit value-function estimation, making the group-based approach memory-efficient and scalable for practical LLM training.

3 Methodology

ProxMO addresses credit-assignment weaknesses at both episode and step levels by adapting advantages to task difficulty and aggregating semantically similar states continuously. It combines these signals into a clipped PPO-style objective for practical training.

Episode-Level: Success-Rate-Aware Advantage Modulation: Standard GRPO z-score normalization assigns identical advantage magnitudes to statistically identical outcomes despite differences in task difficulty.This motivates replacing purely statistical credit assignment with context-aware modulation.
Episode-Level: Success-Rate-Aware Advantage Modulation: ProxMO’s Polarized Signal Controller scales episode-level advantages using the empirical success rate of each trajectory group.The mechanism is designed to adapt credit allocation to episode difficulty.
Episode-Level: Success-Rate-Aware Advantage Modulation: Successes are amplified in low-success groups, while failures are attenuated in high-success groups to strengthen rare breakthroughs and reduce noise penalties.The modulation uses the failure rate, sigmoid activation, and steepness and strength hyperparameters.
Step-Level: Proximity-Based Soft Aggregation: Trajectory-wide advantages cannot distinguish action quality within a trajectory, while hard step-level boundaries can create singleton groups or assign equal weights to dissimilar states.These limitations motivate continuous semantic aggregation rather than discrete partitioning.
Step-Level: Proximity-Based Soft Aggregation: Proximity-Based Soft Aggregation computes step-level baselines from returns weighted by semantic proximity, using temperature-scaled comparisons within task-matched trajectory groups.The method represents states with TF-IDF vectors and uses cosine similarity; temperature controls concentration around similar states.
Unified Training Objective: The framework combines episode-level and step-level advantages through a weighted summation before applying a clipped PPO policy objective.The balance coefficient defaults to ω = 1, while importance sampling ratios and clipping regulate the optimization objective.

4 Experiments

Experiments evaluate ProxMO on ALFWorld and WebShop against closed-source, prompting, and reinforcement-learning baselines, finding consistent gains, robust hyperparameters, synergistic mechanisms, and minimal training overhead.

Experimental Setup: ProxMO is evaluated on ALFWorld and WebShop against closed-source LLMs, prompting agents, GRPO, and GiGPO using Qwen2.5-1.5B/7B-Instruct backbones.ALFWorld contains 3,827 household-task instances, while WebShop contains 1.1M products and 12K user instructions.
Main Results: ProxMO consistently outperforms baselines across model scales and task types, with especially pronounced gains on long-horizon tasks requiring precise credit assignment.The reported long-horizon examples include Look, Cool, and Pick2.
Main Results: Its gains address GRPO’s context-agnostic normalization and GiGPO’s sparse discrete grouping through hierarchical success-rate modulation and proximity-based soft aggregation.The design is reported to provide more informative advantage estimates while maintaining stability in complex environments.
Hyperparameter Sensitivity: Optimal α = 4.0, β = 0.1, and τ = 0.1 remain stable across broad intervals and transfer across model scales and task types.The same configuration is described as near-optimal across ALFWorld and WebShop, reducing the need for extensive tuning.
Ablation Study: Removing either episode-level modulation or step-level aggregation degrades performance, with PSA especially important for long-horizon tasks and PSC for high success-rate variance.The full ProxMO also outperforms the ablated variants and GiGPO, indicating synergistic effects.
Efficiency and Case Study: In a complex ALFWorld task, the ProxMO-trained agent succeeds in 11 steps while GPT-4o fails after 14 steps following goal drift toward an incorrect cabinet action.The example attributes this divergence to proximity-based aggregation enforcing consistency between intermediate decisions and the global target.
Efficiency and Case Study: ProxMO adds only +1.09% training-time overhead versus GRPO across ALFWorld training iterations.The mechanisms use lightweight arithmetic operations without additional neural networks or model forward passes.

5 Conclusion

ProxMO addresses multi-turn credit assignment by incorporating global context at episode and step levels, and experiments report consistent improvements with independent and synergistic contributions from both mechanisms.

Conclusion: ProxMO adapts credit to task difficulty at the episode level and weights states continuously by semantic proximity at the step level.The framework is presented as a context-dependent approach for stable multi-turn reinforcement learning with LLM agents.

Limitations

The paper’s theoretical analysis identifies an informational mismatch in GRPO’s z-score credit allocation, while its empirical scope is concentrated on resource-efficient model scales.

Empirical Scope: The empirical evaluation is primarily limited to resource-efficient 1.5B and 7B models, leaving generalization to significantly larger foundation models for future validation.The stated scope boundary concerns the full spectrum of model capacities.
Theoretical Limitation: GRPO’s z-score normalization allocates advantage magnitudes according to statistical rarity, which can mismatch information value across task success rates.Rare failures in low-success environments may receive excessive penalties, while rare successes may not receive appropriately strong credit.
Theoretical Limitation: ProxMO corrects this mismatch by amplifying signals in low-success regimes and attenuating them in high-success regimes.This success-rate-aware modulation is described as weighting advantages directly as a function of the group success rate p.

A.2 Step-Level Grouping and Singleton Degeneracy

Hard step-level grouping struggles with semantically similar but lexically different states: strict criteria create singleton groups, while loose criteria erase proximity distinctions. On ALFWorld, singleton groups comprise 30–36% of trajectory steps, yielding zero advantage signals for those steps.

Multi-turn states can be semantically similar yet lexically distinct because trajectories differ in location, history, step, or observations.
Strict matching produces rare exact matches and singleton groups, whereas similarity thresholds assign equal weight to states with different semantic proximity.
Singleton groups reduce baseline computation to the state itself, yielding zero advantage and no learning signal.
30–36% of trajectory steps persistently belong to singleton groups across ALFWorld training iterations.
ProxMO replaces discrete grouping with continuous proximity weighting, allowing all states to contribute according to semantic similarity.

B.2 Benchmark Descriptions

ALFWorld evaluates multi-step household-task decision making across six task categories, while WebShop evaluates interactive online shopping through a simulated HTML website. The supplied benchmark description identifies ALFWorld’s task diversity and difficulty variation explicitly.

ALFWorld is an embodied multi-step environment with 3,827 household-task instances across six categories.
ALFWorld difficulty varies substantially across categories, with baseline success rates ranging from 20% for Pick2 to 95% for Pick.
WebShop tests interactive shopping in a simulated HTML website, requiring agents to search, navigate, and purchase items matching user requirements.

B.3 Hyperparameter Analysis

ProxMO’s hyperparameters show broad operational stability: moderate episode-level settings balance sensitivity and variance, while step-level temperature balances precision and discrimination. Optimal configurations remain consistent across ALFWorld and WebShop, reducing tuning demands.

Episode-Level Modulation: Moderate episode steepness balances task-difficulty sensitivity and training stability, while moderate strength avoids ineffective adjustments or excessive gradient variance.
Step-Level Aggregation: Step temperature balances concentrated similarity weights against near-uniform weighting, with gradual degradation at extreme values.
Cross-Task Consistency: Optimal hyperparameter configurations remain consistent across ALFWorld and WebShop despite differences in task structure, episode length, and observations.
Deployment Implications: Wide stability ranges, cross-task consistency, and gradual degradation reduce tuning and adoption barriers for production deployment.

E.1 ALFWorld Example: Complex Multi-Object Task (Pick Two Objects)

The ALFWorld example requires locating two remote controls and placing them in an armchair, demanding exploration and sequential object management. The trajectory illustrates structured planning, object-state management, and adaptive search across the episode.

The task requires locating two remote controls and placing them in an armchair through strategic exploration and sequential object management.
The trajectory demonstrates structured reasoning in which the agent plans globally, explores locations, manages object state, and adapts search based on observations.

E.2 Baseline Comparison: Failed Episode with Execution Error

The failed GPT-4o episode shows that reasonable search can still end in failure when an intermediate placement decision diverges from the task objective. ProxMO addresses this gap with episode-level and step-level credit mechanisms that support consistent multi-step execution.

Baseline failure: GPT-4o pursued a reasonable search strategy but failed after placing the first remote control in a storage area instead of the target armchair.The placement error compromised task completion despite the agent’s initial progress.
Baseline failure: The agent later discovered both remote controls but could not recover after temporarily storing the first remote in a cabinet.The mistake was recognized only after the second remote had been acquired and the episode could not be salvaged.
Credit-assignment failure: Without fine-grained step-level credit assignment, local action rationality can diverge from the global instruction to place both remotes in the armchair.The paper characterizes cabinet storage as locally plausible but inconsistent with the task’s destination requirement.
ProxMO response: ProxMO combines success-rate-aware episode modulation with proximity-based step aggregation to incorporate task difficulty and semantic state similarity into credit assignment.The two mechanisms respectively adjust trajectory-level credit and replace hard state boundaries with continuous semantic weighting.
ProxMO response: The ProxMO-trained agent maintained execution consistency and completed the task successfully within 11 steps.The result is presented as evidence that fine-grained credit assignment supports both efficiency and correctness in multi-step execution.