Rethinking Policy Diversity in Ensemble Policy Gradient in Large-Scale Reinforcement Learning

AI summary

Overview

Research area: Reinforcement learning (RL), specifically ensemble policy gradient methods trained in massively parallel GPU physics simulators for robotic control.

Technical level: Intermediate. The core idea is accessible, but the paper leans on importance sampling (IS), KL divergence, PPO clipping bias, and effective sample size, with formal proofs deferred to appendices.

Scope: The paper studies how the degree of diversity between policies in a leader-follower ensemble affects learning stability and sample efficiency, and proposes a method, Coupled Policy Optimization (CPO), that regulates that diversity with KL constraints plus an adversarial reward.

What This Paper Is About

Scaling RL to tens of thousands of parallel environments does not automatically improve learning, because a single policy explores in a narrow way. Ensemble methods such as SAPG address this by having one leader and several followers collect data in separate blocks of environments, but it was an open question whether more diversity between those policies is always better. This paper argues that excessive divergence between followers and the leader degrades both training stability and sample efficiency, and proposes to constrain diversity rather than simply maximize it.

Key Contributions

1. A theoretical analysis showing that excessive inter-policy diversity in ensemble policy gradient methods degrades training stability and sample efficiency. Three propositions link IS-ratio deviation from 1 to reduced effective sample size (ESS), link it to larger bias in the PPO-clipped gradient estimate, and show the deviation is upper bounded by the KL divergence between follower and leader policies.
2. CPO, a leader-follower method that adds a KL divergence constraint during follower updates plus an adversarial reward (a discriminator that predicts agent identity from state-action pairs) to prevent policies from overconcentrating.
3. Empirical evidence across six dexterous manipulation, two gripper-based manipulation, and two locomotion tasks that CPO outperforms PPO, DexPBT, and SAPG in sample efficiency and final performance, with KL constraints keeping IS ratios close to one and increasing ESS.
4. An inter-policy KL divergence analysis showing that CPO induces a structured formation in which followers are distributed around the leader, avoiding the policy misalignment observed in the prior method.

Main Findings

• Excessive diversity harms learning stability: Propositions 1-3 formalize the intuition that when followers drift from the leader, the IS ratio deviates from 1, the effective sample size falls, and the bias induced by the PPO clipping operator grows. Proposition 3 shows the expected absolute deviation of the IS ratio from 1 is upper bounded by the square root of twice the KL divergence between the follower and leader policies (via Pinsker's inequality).
• Final performance on dexterous manipulation (after 2 × 10^10 environment steps): On ShadowHand, CPO reached 13762 ± 414 versus SAPG 12882 ± 343, PBT 10294 ± 1728, and PPO 10661 ± 1050. On AllegroHand, CPO reached 14421 ± 885 versus PBT 13239 ± 239, SAPG 11989 ± 817, and PPO 10439 ± 1282. On Reorientation (AllegroKuka Reorientation), CPO reached 43.75 ± 0.65 versus SAPG 38.79 ± 1.66 and PBT 2.92 ± 4.27. On Two-Arms Reorientation, CPO reached 35.30 ± 2.77 versus PBT 26.43 ± 11.12, SAPG 5.11 ± 3.41, and PPO 1.41 ± 0.80.
• Tasks where the gain was not significant: On Regrasping (AllegroKuka Regrasping), CPO scored 37.44 ± 1.21 versus SAPG 37.20 ± 0.65 and PBT 35.26 ± 2.82. On Throw, CPO scored 21.69 ± 2.44 versus SAPG 22.51 ± 1.15. The paper reports no significant improvement over SAPG in these two tasks.
• Sample efficiency: In many tasks CPO reached SAPG's final performance with approximately half the number of environment steps.
• Locomotion: Performance differences across algorithms were smaller. PBT converged fastest, which the authors attribute to broad parallel exploration being advantageous in simpler tasks, while CPO converged slightly faster than SAPG.
• IS ratio and ESS (at 5 × 10^9 environment steps): On ShadowHand, mean IS ratio deviation was 0.889 with ESS rate 0.0223 for SAPG, versus 0.403 and 0.763 for CPO with λ_f = 0.5, 0.297 and 0.871 with λ_f = 0.2, 0.222 and 0.923 with λ_f = 0.1, and 0.187 and 0.941 with λ_f = 0.05. On AllegroKuka Reorientation, SAPG showed 0.608 and 0.110, versus 0.420 and 0.721, 0.276 and 0.888, 0.214 and 0.929, and 0.199 and 0.938 for the same λ_f values. Stronger KL constraints produced smaller deviation and higher ESS, consistent with Proposition 1.
• KL divergence structure: SAPG showed severely misaligned followers in ShadowHand and AllegroKuka Reorientation, whereas CPO maintained stable inter-agent distances. In CPO the leader remained the closest agent to every follower, indicating followers naturally distribute around the leader without overconcentration. In AllegroKuka Regrasping, where performance was similar, SAPG followers did not diverge noticeably.
• Ablation robustness: The KL coefficient ablation in ShadowHand and AllegroKuka Reorientation showed CPO was robust across a wide range of λ_f values and consistently outperformed SAPG, with β in the combined objective fixed at 0.001. The paper suggests starting with a weak constraint (λ_f = 0.5) and gradually strengthening it.
• SAPG entropy ablation: The authors' preliminary ablation on SAPG's entropy regularization found that it promotes exploration and can improve sample efficiency, but increases follower-leader KL divergence and often leads to severe misalignment (details in Appendix A.7).

Methodology in Plain English

The authors start from SAPG, which splits a large number of parallel environments into blocks, assigns one leader policy and several follower policies to them, and lets the leader learn from both its own on-policy data and the followers' off-policy data through importance sampling. All agents share policy and value networks conditioned on identification vectors.

They first analyze what happens when followers wander too far from the leader. Using the effective sample size formula from importance sampling, they show that samples from badly misaligned followers carry little useful weight, and using the structure of PPO's clipped objective, they show that the gradient estimate acquires more bias. They then prove that the expected absolute deviation of the IS ratio from 1 is bounded by the KL divergence between follower and leader, which motivates directly constraining that KL divergence.

CPO implements this as a constrained follower update: maximize the follower's advantage subject to the KL divergence to the leader being at most ε_KL. Following AWAC, this constraint has a closed-form non-parametric solution, which they approximate with a neural network and add to the standard SAPG follower loss, with a temperature λ_f and a scaling coefficient β.

Because a strict KL constraint also pulls followers toward each other and reduces coverage, they add an intrinsic adversarial reward. A discriminator predicts which policy index produced a given state-action pair, and the log-probability of the correct identity becomes a bonus reward for followers. This bonus is deliberately not given to the leader: when the leader is updated from follower samples, only true environment rewards are used, preserving the leader-follower asymmetry.

Experiments use Isaac Gym with N = 24,576 parallel environments, M = 6 blocks, five random seeds, and compare against PPO, DexPBT (PBT), and SAPG on ten tasks: six dexterous manipulation, two gripper-based manipulation (FrankaCubePush and Stack), and two locomotion tasks, all with dense rewards.

Why This Matters

Impact on research: The paper reframes diversity in ensemble RL as something to be regulated rather than maximized, and supplies a concrete, closed-form-friendly mechanism (KL constraint plus identity-discriminator reward) grounded in IS theory. It also documents the failure mode of an existing state-of-the-art method, SAPG, and shows that a smaller IS-ratio deviation and higher ESS are measurable proxies for the improvement. The limitation acknowledged by the authors is that the method still relies on a fixed number of policies and environments per policy.

Real-world applications:

• Dexterous robot hand control, including ShadowHand and AllegroHand style tasks where high-dimensional action spaces and difficult exploration dominate.
• Gripper-based manipulation such as cube pushing and stacking in tabletop settings.
• Locomotion control for legged robots trained in simulation before deployment.
• Large-scale simulated training pipelines where data collection is cheap but useful exploration is the bottleneck.

Industry relevance: Companies using GPU-accelerated simulators for robot learning can apply CPO as a drop-in modification of an existing leader-follower ensemble, gaining sample efficiency without discarding data. The reported ability to reach a baseline's final performance with approximately half the environment steps translates directly into compute savings, and the ESS diagnostics offer a practical monitoring signal for whether an ensemble's policies are collaborating or sabotaging the leader's update.

Future Directions

• Automatically adjusting the number of policies and the number of environments per policy, which the authors identify as a limitation and a promising future direction, since the effective exploration range varies with task and training stage.
• Applying the approach to tasks with high-dimensional action spaces and demanding exploration requirements, which the authors highlight as where massively parallel environments hold the most untapped potential.
• Investigating whether the KL-constraint mechanism extends beyond the leader-follower framework to other agent ensemble schemes such as population-based training.
• Testing the tuning heuristic (start weak at λ_f = 0.5, then strengthen) systematically across a broader range of tasks and constraint schedules.

Target Audience

RL researchers working on policy gradient methods, ensemble or multi-agent training, and importance sampling; robotics engineers training manipulation and locomotion policies in mass-parallel simulators such as Isaac Gym; and practitioners or graduate students already comfortable with PPO, KL regularization, and off-policy correction who want to understand the diversity-stability trade-off in ensemble RL.

Authors’ abstract

Scaling reinforcement learning to tens of thousands of parallel environments requires overcoming the limited exploration capacity of a single policy. Ensemble-based policy gradient methods, which employ multiple policies to collect diverse samples, have recently been proposed to promote exploration. However, merely broadening the exploration space does not always enhance learning capability, since excessive exploration can reduce exploration quality or compromise training stability. In this work, we theoretically analyze the impact of inter-policy diversity on learning efficiency in policy ensembles, and propose Coupled Policy Optimization which regulates diversity through KL constraints between policies. The proposed method enables effective exploration and outperforms strong baselines such as SAPG, PBT, and PPO across multiple tasks, including challenging dexterous manipulation, in terms of both sample efficiency and final performance. Furthermore, analysis of policy diversity and effective sample size during training reveals that follower policies naturally distribute around the leader, demonstrating the emergence of structured and efficient exploratory behavior. Our results indicate that diverse exploration under appropriate regulation is key to achieving stable and sample-efficient learning in ensemble policy gradient methods. Project page at https://naoki04.github.io/paper-cpo/ .