Behaviour Policy Optimization: Provably Lower Variance Return Estimates for Off-Policy Reinforcement Learning

AI summary

Overview

Research area: Reinforcement learning, specifically off-policy policy gradient methods, variance reduction in return estimation, and behaviour policy design.

Technical level: Advanced. The paper relies on importance sampling theory, TD(λ) return estimation, variance-optimal sampling distributions, and fitted Q-evaluation.

One-sentence scope: The paper extends a variance-optimal behaviour-policy design from off-policy evaluation into online off-policy reinforcement learning, derives a new unbiased truncated importance-sampled TD(λ) return estimator with a provable variance-reduction guarantee, and instantiates the idea on top of REINFORCE and PPO.

What This Paper Is About

Most reinforcement learning algorithms that improve a policy using return estimates suffer from high variance, which wastes samples and destabilises training. Recent off-policy evaluation work (the ODI algorithm of Liu and Zhang, 2024) showed the surprising result that deliberately designing the policy used to collect data — rather than simply collecting data with the current policy — can give provably lower-variance return estimates. This paper asks whether that insight survives the transition to online RL, where the target policy keeps changing and where full Monte Carlo returns are a poor choice for policy improvement.

Key Contributions

1. A new return estimator. The authors introduce truncated importance-sampled TD(λ) returns, G^TIS,λ, which combine the variance-optimal behaviour policy idea with TD(λ) bootstrapping, giving control over the bias–variance trade-off and fitting naturally into existing policy-gradient code. They prove unbiasedness (Theorem 1) and a variance-reduction guarantee (Theorem 2) for this estimator.

2. Extension of variance-optimal behaviour policy design to online RL. The paper adapts Liu and Zhang's one-step optimal behaviour policy μ̂(a|s) ∝ π(a|s)√q̂_π(s,a) to the discounted, online setting, and addresses the additional technical difficulties this creates, notably the non-stationarity of the target policy.

3. A practical algorithm (BPO) and a method for learning q̂_π. Theorem 3 shows that q̂_π corresponds to the state-action value function of π under a modified reward r̂(s,a) = 2r(s,a)q_π(s,a) − r²(s,a), which makes it learnable with an additional Q-network trained by fitted Q-evaluation.

4. Continuous-action extension. Because the target distribution in the continuous case has no closed form, the authors derive an alternative loss L_μ^cont (Theorem 4) that is minimised exactly when the behaviour policy matches μ̂(a|s) ∝ π_θ(a|s)√Q̂_ζ̂(s,a).

Main Findings

• On-policy data collection is not variance optimal. Building on Liu and Zhang (2024), the paper states that well-designed behaviour policies give provably lower-variance return estimates than collecting data on-policy — a counter-intuitive result that motivates the whole approach.

• Provable variance reduction. Theorem 2 states that for λ = 1 and c̄, ρ̄ = ∞, the truncated IS TD(λ) estimator under behaviour policy μ̂ has variance at most that under the target policy π, minus a non-negative term ε(s). The term c(s) is non-negative by Jensen's inequality, guaranteeing ε(s) ≥ 0, and c(s) = 0 only in the degenerate case where all actions in a state share the same q̂_π.

• Unbiasedness. Theorem 1 shows that with c̄, ρ̄ = ∞, the estimator is unbiased for the target policy value for all μ in the search space Λ, all states, and all λ in [0,1].

• No-worse and often better performance. For both REINFORCE and PPO, the authors report no-worse and often better sample efficiency and final episode returns compared with the base algorithms.

• ShortCorridor results are modest. On the ShortCorridor grid world (a toy example from Sutton and Barto, 2018) with REINFORCE, the authors explicitly state the gains are marginal; the experiment is included mainly to show the method applies across policy-gradient algorithms. Results are averaged over 100 independent runs with standard error bars, on 10 evaluation episodes.

• MuJoCo results. On continuous-control benchmarks the authors report mean returns averaged over 10 independent runs with SE bars (10 evaluation episodes) on Ant-v5 (gSDE), HalfCheetah-v5 (gSDE), Hopper-v5 (default) and Walker2d-v5 (default). The two clearly aligned comparisons are Ant-v5 (gSDE), where PPO reports 1106 ± 111 and BPO (c̄, ρ̄ = 1.5, 1.5) reports 1287 ± 148, and HalfCheetah-v5 (gSDE), where PPO reports 3425 ± 468 and BPO reports 3742 ± 408. The remaining table entries (including 3527 ± 670, 2126 ± 492 and 2091 ± 408) appear in the truncated table under the PPO row, but the provided content cuts off before the full BPO row, so the complete set of results is not reported here.

• Symlog targets matter for stability. The authors report that symmetric-log targets (Hafner et al., 2024) are crucial because the relative magnitudes of the Q_ζ and Q̂_ζ̂ predictions can be large; symlog squashing lets the same hyperparameters (for example, the same learning rate) be used for both networks.

Methodology in Plain English

The usual setup in off-policy RL is: collect data with one policy (the behaviour policy), then correct for the mismatch with importance sampling weights — the ratio of the target policy's action probability to the behaviour policy's action probability. Naive importance sampling explodes in variance when the two policies differ.

The paper's twist is to stop treating the behaviour policy as a given and instead optimise it. Following Liu and Zhang, the variance-optimal behaviour policy is proportional to the target policy multiplied by the square root of a special quantity q̂_π (a one-step lookahead measure of how much variance each action contributes, quantified in equation 4). So the algorithm runs two loops side by side:

1. The target policy is trained as usual (REINFORCE or PPO), but on rollouts collected by the behaviour policy rather than the target policy, using the new truncated IS TD(λ) returns for the policy-gradient and value targets.
2. The behaviour policy is trained to match the variance-optimal distribution, using a cross-entropy loss in discrete action spaces or a log-probability-based loss in continuous ones.

Because the optimal behaviour policy depends on q̂_π, which itself depends on q_π, the authors train two extra Q-networks with fitted Q-evaluation (FQE): one for q_π using the standard one-step Bellman target, and one for q̂_π using the modified reward r̂ and a γ² discount factor. Importance weights in the return estimator are truncated (with parameters c̄ and ρ̄) to guard against variance explosions, in the spirit of V-trace (IMPALA) and Retrace(λ), and λ itself acts as a bias–variance dial. The authors note λ should not stray far from 1, and that a good default is ρ̄ ≥ c̄ ≥ 1.

Why This Matters

Impact on research. The paper challenges a deeply embedded assumption in reinforcement learning — that data should be collected with the policy you are trying to improve. If a hand-designed behaviour policy can provably reduce return variance while remaining unbiased, it opens a new axis for algorithmic improvements that is orthogonal to the usual levers of baselines, advantage estimation, and gradient clipping. It also bridges off-policy evaluation theory and online policy optimisation, which have largely developed as separate literatures.

Real-world applications:

• Robotics and continuous control, where each environment interaction is expensive and sample efficiency directly determines cost.
• Large reasoning model fine-tuning, which the introduction cites as an active RL application with very expensive rollouts.
• Advanced control systems and high-dimensional games, cited as domains where RL requires millions to billions of interactions.
• Offline-pretraining-plus-online-finetuning pipelines, since the paper explicitly notes that practitioners often deploy online to evaluate or tune policies even after offline training.

Industry relevance. Because BPO is implemented as a wrapper around PPO and REINFORCE — the underlying algorithms stay unchanged — it is relatively easy to adopt for teams already running PPO at scale. The authors provide code at https://github.com/sacktock/BPO. The paper acknowledges the trade-off: the method increases algorithmic complexity and introduces additional hyperparameter tuning burden, which is the main practical obstacle to adoption.

Future Directions

• Reducing the added complexity and hyperparameter load. The authors flag the extra Q-networks, the behaviour policy, and the tuning of c̄, ρ̄ and λ as a real cost, and report that stability of the q_π and q̂_π estimates is critical to BPO. Simplifying this stack is an obvious next step.
• Improving results in simple environments. The marginal ShortCorridor gains suggest the benefit depends on task complexity or on the accuracy of q̂_π; characterising when BPO helps most is unresolved.
• Scaling to more and harder benchmarks. The paper reports REINFORCE on one toy grid world and PPO on four MuJoCo tasks; testing on higher-dimensional and long-horizon domains is a natural extension.
• Alternative estimators and truncation schedules. The paper deliberately uses less aggressive truncation than Retrace(λ) or V-trace, and states that λ should stay near 1; how to best set these trade-offs, and whether other estimators fit the framework, remains open.

Target Audience

This paper is aimed at reinforcement learning researchers and graduate students working on policy optimisation, off-policy learning, or variance reduction. It will be most valuable to readers who are already comfortable with importance sampling, TD(λ)/eligibility traces, and actor-critic methods such as PPO — the theory (Theorems 1–4, with proofs in appendices) is not written for beginners. Practitioners who train PPO-based agents and care about sample efficiency may also benefit, particularly those willing to add the auxiliary networks and tuning the method requires.

Authors’ abstract

Many reinforcement learning algorithms, particularly those that rely on return estimates for policy improvement, can suffer from poor sample efficiency and training instability due to high-variance return estimates. In this paper we leverage new results from off-policy evaluation; it has recently been shown that well-designed behaviour policies can be used to collect off-policy data for provably lower variance return estimates. This result is surprising as it means collecting data on-policy is not variance optimal. We extend this key insight to the online reinforcement learning setting, where both policy evaluation and improvement are interleaved to learn optimal policies. Off-policy RL has been well studied (e.g., IMPALA), with correct and truncated importance weighted samples for de-biasing and managing variance appropriately. Generally these approaches are concerned with reconciling data collected from multiple workers in parallel, while the policy is updated asynchronously, mismatch between the workers and policy is corrected in a mathematically sound way. Here we consider only one worker - the behaviour policy, which is used to collect data for policy improvement, with provably lower variance return estimates. In our experiments we extend two policy-gradient methods with this regime, demonstrating better sample efficiency and performance over a diverse set of environments.