BiPACE: Bisimulation-Guided Policy Optimization with Action Counterfactual Estimation for LLM Agents

2026-06-24

AI Summary

This work addresses a critical flaw in existing group-based reinforcement learning methods, which erroneously assume equivalence among trajectory steps, leading to state-action credit misassignment: overly fine-grained hash partitioning yields numerous singleton groups, while intra-group averaging conflates state values with action-specific credits. To resolve this, the authors propose BiPACE—a novel advantage estimator that requires no additional critics, auxiliary losses, or sampling overhead. BiPACE integrates behavioral equivalence clustering (BiGPO) and action counterfactual advantage estimation (PACE) to refine credit assignment. Its key innovations include an approximate bisimulation clustering based on cosine distance of policy hidden states to reduce singleton rates, and an action-conditional peer baseline enabling non-parametric estimation of Q(s,a)−V(s). Experiments demonstrate significant improvements, boosting success rates from 90.8% to 97.1±0.9% (all seeds >95%) on ALFWorld with Qwen2.5-7B, and consistently outperforming GRPO and GiGPO across Qwen2.5-1.5B, WebShop, and TextCraft, with only an 11.3% increase in training overhead.

Abstract

Stepwise group-based RL is an attractive way to train long-horizon LLM agents without a learned critic: it reuses multiple sampled rollouts to estimate local advantages. Its weakness is less visible but more fundamental: every group-relative estimator assumes that the steps it compares are equivalent for credit assignment. We show that current agentic variants violate this assumption through a state-action credit mismatch. The observation-hash partition is overly fine on the state side, creating singleton groups with zero step-level signal, while a single within-group mean is too coarse on the action side, mixing state-value estimation with action-specific credit. We introduce BiPACE (Bisimulation-Guided Policy Optimization with Action Counterfactual Estimation), a drop-in advantage estimator that fixes both sides without adding a critic, auxiliary loss, or extra rollouts. BiGPO clusters steps by cosine distance in the actor's own hidden-state geometry, an empirical policy-induced proxy for bisimulation that substantially lowers the singleton rate left by observation hashing. PACE then recenters returns within each behavioral cluster using action-conditioned peer baselines; its Q-style instance estimates a local Q(s,a)-V(s) nonparametrically. On ALFWorld/Qwen2.5-7B, BiPACE_Q raises overall validation success from GiGPO's 90.8 to $97.1\pm0.9$ over three seeds, and crosses the 95% threshold on every seed, which GiGPO never does within the same budget. On Qwen2.5-1.5B it reaches $93.5\pm1.2$ versus GiGPO's 86.7, and on WebShop and TextCraft it improves over GRPO and GiGPO at both model scales. The measured BiPACE-specific overhead is 11.3% of a single training-step wall time. Yet it changes the estimator's comparison unit from surface identity to approximate behavioral equivalence plus action-side counterfactuals. The code is available at https://github.com/TianxiangZhao/BiPACE.

Problem

• credit assignment
• bisimulation
• advantage estimation
• LLM agents
• state-action mismatch

Innovation

• Bisimulation
• Action Counterfactual Estimation
• Advantage Estimation
• LLM Agents
• Policy Optimization