← All publications

Unbiased Policy Gradient with Random Horizon

  • Deep Learning Theory

We use geometrically sampled trajectory horizons to construct unbiased policy-gradient estimators for infinite-horizon discounted reinforcement learning.

Paper EWRL 2024

Rui Yuan1, Andrii Tretynko2, Simone Rossi1, Thomas Hannagan1

  1. Stellantis, France
  2. Vehicle Software Development, Ukraine

Policy gradients for discounted reinforcement learning target rewards arbitrarily far into the future, while practical rollouts stop after finitely many steps. We address the resulting gradient bias by sampling the rollout horizon from a geometric distribution. The resulting unbiased policy gradient (UPG) estimators include new GPOMDP variants and existing random-horizon estimators. Their sampling complexity removes a logarithmic factor in accuracy from the corresponding vanilla policy-gradient bound. Longer expected rollouts reduce the second-moment bound, making trajectory length and estimator variance joint determinants of sampling cost.

The horizon enters the gradient

We consider a Markov decision process with bounded rewards ∣r(s,a)∣≤rmax⁡|r(s,a)|\leq r_{\max}, initial-state distribution ρ\rho, discount γ∈(0,1)\gamma\in(0,1), and differentiable policy πθ\pi_\theta. For trajectories τ\tau sampled under πθ\pi_\theta, write rt=r(st,at)r_t=r(s_t,a_t). The objective and its GPOMDP gradient representation are

J(θ)=Eτ[∑t=0∞γtrt],vt=∑t′=0t∇log⁡πθ(at′∣st′),G(θ)=∇J(θ)=Eτ[∑t=0∞γtrtvt].\begin{aligned} J(\theta)&=\mathbb E_\tau \left[\sum_{t=0}^{\infty}\gamma^t r_t\right],\\ v_t&=\sum_{t'=0}^{t}\nabla\log\pi_\theta(a_{t'}\mid s_{t'}),\\ G(\theta)&=\nabla J(\theta) =\mathbb E_\tau\left[\sum_{t=0}^{\infty}\gamma^t r_t v_t\right]. \end{aligned}

The vector vtv_t accumulates policy scores through time tt. A fixed truncation horizon drops the remaining reward-score terms and introduces bias. Increasing that horizon reduces the error while increasing sampling cost. We instead encode the discount in an independent stopping rule.

Geometric stopping yields unbiased GPOMDP

Sample H−1∼Geom⁡(1−γ)H-1\sim\operatorname{Geom}(1-\gamma) on the nonnegative integers, then accumulate undiscounted rewards and scores:

Pr⁡(H>t)=γt,E[H]=11−γ,G^UGPOMDP(θ)=∑t=0H−1rtvt,E[G^UGPOMDP(θ)]=G(θ).\begin{aligned} \Pr(H>t)&=\gamma^t,\qquad \mathbb E[H]=\frac{1}{1-\gamma},\\ \widehat G_{\mathrm{UGPOMDP}}(\theta) &=\sum_{t=0}^{H-1}r_t v_t,\\ \mathbb E[\widehat G_{\mathrm{UGPOMDP}}(\theta)]&=G(\theta). \end{aligned}

Each time-tt term survives with exactly its original discount weight. Every rollout is finite almost surely, and arbitrarily late rewards remain represented across samples. At each step, we update the score accumulator and add the current reward times that accumulator to the gradient estimate. Independent trajectory estimates can then be averaged into a batch.

Splitting discounting between rewards and survival

The broader α\alpha-UPG family allocates discounting between the stopping distribution and explicit reward weights. For α∈[0,1)\alpha\in[0,1), its GPOMDP estimator is

H−1∼Geom⁡(1−γ1−α),G^α-UGPOMDP(θ)=∑t=0H−1γαtrtvt,E[H]=11−γ1−α.\begin{aligned} H-1&\sim\operatorname{Geom}(1-\gamma^{1-\alpha}),\\ \widehat G_{\alpha\text{-}\mathrm{UGPOMDP}}(\theta) &=\sum_{t=0}^{H-1}\gamma^{\alpha t}r_t v_t,\\ \mathbb E[H]&=\frac{1}{1-\gamma^{1-\alpha}}. \end{aligned}

Multiplying the survival probability γ(1−α)t\gamma^{(1-\alpha)t} by γαt\gamma^{\alpha t} recovers γt\gamma^t, preserving the same gradient expectation. At α=0\alpha=0 we recover UGPOMDP. Increasing α\alpha lengthens the expected rollout; the theoretical second-moment bound decreases. We also derive an action-value branch, α\alpha-QPGT, containing Q-PGT at α=0\alpha=0 and RPG at α=1/2\alpha=1/2. These branches provide a common framework for comparing existing and new estimators.

Sampling cost depends on variance and horizon

For a smooth objective JJ, unbiasedness permits a stochastic-gradient convergence analysis under an expected-smoothness bound on the estimator's second moment. Uniform bounds on the policy's expected squared score and expected log-policy Hessian norm yield an explicit second-moment bound ν\nu. Choosing the batch size and constant step size according to these bounds gives

E∥∇J(θU)∥2≤ϵ2,N=O ⁣(Lν [J⋆−J(θ(0))] E[H]ϵ4).\begin{aligned} \mathbb E\|\nabla J(\theta_U)\|^2&\leq\epsilon^2,\\ N&=\mathcal O\!\left( \frac{L\nu\,[J^\star-J(\theta^{(0)})]\,\mathbb E[H]}{\epsilon^4} \right). \end{aligned}

Here θU\theta_U is a uniformly sampled iterate, θ(0)\theta^{(0)} is the initial policy parameter, and J⋆=sup⁡θJ(θ)J^\star=\sup_\theta J(\theta) is the optimal objective value. The tolerance ϵ\epsilon measures first-order stationarity, LL is the objective's smoothness constant, and NN counts expected environment samples. For fixed problem constants the accuracy dependence is O(ϵ−4)\mathcal O(\epsilon^{-4}), removing the truncation-related logarithmic factor in the vanilla comparison. Choosing α\alpha therefore requires considering νE[H]\nu\mathbb E[H]: the smallest second-moment bound occurs with longer rollouts. For UGPOMDP, the sampling bound scales as O((1−γ)−6ϵ−4)\mathcal O((1-\gamma)^{-6}\epsilon^{-4}).

Discounted return in control benchmarks

We evaluate two-hidden-layer policies with 64 units and Tanh activations, using softmax outputs for discrete actions and Gaussian outputs for continuous actions. The main comparison averages 10 runs with 95% confidence intervals. Both UGPOMDP and 0.50.5-UGPOMDP outperform GPOMDP on CartPole, LunarLander, Acrobot, and Pendulum. We use γ=0.99\gamma=0.99 for the first three and 0.950.95 for Pendulum, with training batches of 64 trajectories. The preferred α\alpha varies across environments; 0.750.75 performs poorly on Acrobot and best among the tested settings on Pendulum.

We also compare the number of environment interactions and repeat the GPOMDP baseline with twice its truncation horizon to examine how discounted return depends on sampling cost. These experiments use Adam and nonlinear neural policies; the convergence bound applies to stochastic-gradient updates under the assumptions above. In terminating environments, the treatment of terminal states can change gradient estimates, particularly for Q-PGT and RPG, and must preserve the reward and stopping weights used in the unbiasedness calculation.