Table of Contents
Fetching ...

Stabilizing Policy Gradient Methods via Reward Profiling

Shihab Ahmed, El Houcine Bergou, Aritra Dutta, Yue Wang

TL;DR

This work tackles the instability and high variance of policy gradient methods by introducing Reward Profiling, a universal wrapper that selectively accepts updates based on high-confidence performance comparisons. The approach provides Lookback, MixUp, and Three-Points variants, requiring roughly $O( frac{B^2}{2\epsilon^2}\, olinebreak[4] frac{}{} olinebreak[4]\, olinebreak[4] ext{ln}( frac{2T}{\delta})$) extra rollouts and guaranteeing, with high probability, monotonic improvement without slowing convergence. The authors establish a convergence rate of $O(T^{-1/4})$ for the last iterate under standard smoothness assumptions and extend to biased critics. Empirically, Reward Profiling improves convergence speed and reduces return variance across eight continuous-control benchmarks (Box2D, MuJoCo/PyBullet) and in Unity multi-robot tasks, while maintaining broad applicability to TRPO, PPO, and DDPG. The results suggest Reward Profiling as a general, theoretically grounded path to more reliable policy learning in complex environments.

Abstract

Policy gradient methods, which have been extensively studied in the last decade, offer an effective and efficient framework for reinforcement learning problems. However, their performances can often be unsatisfactory, suffering from unreliable reward improvements and slow convergence, due to high variance in gradient estimations. In this paper, we propose a universal reward profiling framework that can be seamlessly integrated with any policy gradient algorithm, where we selectively update the policy based on high-confidence performance estimations. We theoretically justify that our technique will not slow down the convergence of the baseline policy gradient methods, but with high probability, will result in stable and monotonic improvements of their performance. Empirically, on eight continuous-control benchmarks (Box2D and MuJoCo/PyBullet), our profiling yields up to 1.5x faster convergence to near-optimal returns, up to 1.75x reduction in return variance on some setups. Our profiling approach offers a general, theoretically grounded path to more reliable and efficient policy learning in complex environments.

Stabilizing Policy Gradient Methods via Reward Profiling

TL;DR

This work tackles the instability and high variance of policy gradient methods by introducing Reward Profiling, a universal wrapper that selectively accepts updates based on high-confidence performance comparisons. The approach provides Lookback, MixUp, and Three-Points variants, requiring roughly ) extra rollouts and guaranteeing, with high probability, monotonic improvement without slowing convergence. The authors establish a convergence rate of for the last iterate under standard smoothness assumptions and extend to biased critics. Empirically, Reward Profiling improves convergence speed and reduces return variance across eight continuous-control benchmarks (Box2D, MuJoCo/PyBullet) and in Unity multi-robot tasks, while maintaining broad applicability to TRPO, PPO, and DDPG. The results suggest Reward Profiling as a general, theoretically grounded path to more reliable policy learning in complex environments.

Abstract

Policy gradient methods, which have been extensively studied in the last decade, offer an effective and efficient framework for reinforcement learning problems. However, their performances can often be unsatisfactory, suffering from unreliable reward improvements and slow convergence, due to high variance in gradient estimations. In this paper, we propose a universal reward profiling framework that can be seamlessly integrated with any policy gradient algorithm, where we selectively update the policy based on high-confidence performance estimations. We theoretically justify that our technique will not slow down the convergence of the baseline policy gradient methods, but with high probability, will result in stable and monotonic improvements of their performance. Empirically, on eight continuous-control benchmarks (Box2D and MuJoCo/PyBullet), our profiling yields up to 1.5x faster convergence to near-optimal returns, up to 1.75x reduction in return variance on some setups. Our profiling approach offers a general, theoretically grounded path to more reliable and efficient policy learning in complex environments.

Paper Structure

This paper contains 25 sections, 12 theorems, 44 equations, 7 figures, 3 tables, 2 algorithms.

Key Result

Lemma 1

Let the policy $\pi$ have per-step rewards lie in $[0,R_{\max}]$, so that any trajectory return satisfies $0 \;\le\; G(\tau) \;\le\; B \quad\text{with}\quad B \;=\;\frac{R_{\max}}{1-\gamma}.$ Then for any $\epsilon>0$,

Figures (7)

  • Figure 1: Training performance of REINFORCE. The agent converges with stable behavior with our reward profiling. — REINFORCE — REINFORCE+Lookback.
  • Figure 2: Average Return vs. timesteps (up to 100K) for 3 policy‐gradient methods, PPO (top row), TRPO (middle row), and DDPG (bottom row), on 4 continuous‐control benchmarks: LunarLanderContinuous, Ant, Hopper, Walker2D, and under variants: Vanilla, Lookback, Mixup, Three‐Points, with $E=5$ for minimal overhead. Extra rollouts tuning gets convergence gains in complex environments.
  • Figure 3: Multi-agent Reacher task in UnityML: multiple robotic arms must coordinate to reach and manipulate randomly placed targets (spheres) on a raised platform, testing control under interaction.
  • Figure 4: Average return over training in the UnityML Reacher environment, showing profiling-enhanced DDPG not only converges faster but also maintains greater stability compared to its vanilla counterpart.
  • Figure 5: Sensitivity to evaluation rollouts $E$ for DDPG+Three-Points on HopperBulletEnv-v0. Curves show mean selected returns (±1 std). Small $E$(10–20) yields noisy, unstable updates; large $E$(200) improves stability but slows progress; mid-range $E$(50–100) balances both.
  • ...and 2 more figures

Theorems & Definitions (13)

  • Lemma 1: Concentration
  • Lemma 2: Monotonicity
  • Theorem 1: Convergence
  • Remark 1
  • Theorem 2
  • Lemma 3
  • Lemma 4: Smoothness of $J(\theta)$
  • Lemma 5: Descent step
  • Lemma 6
  • Lemma 7
  • ...and 3 more