[Paper Review] Action-depedent Control Variates for Policy Optimization via Stein's Identity

Hao Liu, Yihao Feng+4 more|arXiv (Cornell University)|Oct 30, 2017

Reinforcement Learning in Robotics18 citations

TL;DR

This paper proposes Stein control variates, a novel variance reduction technique for policy gradient reinforcement learning that leverages Stein's identity to construct action- and state-dependent baseline functions. By enabling more flexible, nonlinear baselines while preserving unbiased gradient estimation, the method significantly improves sample efficiency across multiple continuous control environments, outperforming both value-function baselines and Q-prop in PPO and TRPO settings.

ABSTRACT

Policy gradient methods have achieved remarkable successes in solving challenging reinforcement learning problems. However, it still often suffers from the large variance issue on policy gradient estimation, which leads to poor sample efficiency during training. In this work, we propose a control variate method to effectively reduce variance for policy gradient methods. Motivated by the Stein's identity, our method extends the previous control variate methods used in REINFORCE and advantage actor-critic by introducing more general action-dependent baseline functions. Empirical studies show that our method significantly improves the sample efficiency of the state-of-the-art policy gradient approaches.

Motivation & Objective

• To address the high variance in policy gradient estimation that hinders sample efficiency in deep reinforcement learning.
• To develop a general framework for constructing action- and state-dependent control variates with zero analytic expectation under the policy distribution.
• To extend existing control variate methods like REINFORCE, A2C, and Q-prop by enabling more expressive, nonlinear baseline functions.
• To improve the sample efficiency of state-of-the-art policy optimization algorithms such as PPO and TRPO through variance reduction.
• To provide a theoretically grounded method using Stein's identity that generalizes prior approaches while maintaining unbiased gradient estimation.

Proposed method

• Utilizes Stein's identity to derive a broad class of functions with zero expectation under the policy distribution, enabling valid control variates.
• Constructs action- and state-dependent baseline functions that are not restricted to linear or state-only dependencies, as in prior methods.
• Applies the reparameterization trick in conjunction with Stein's identity to enable differentiable and scalable gradient estimation.
• Derives a general control variate form that allows arbitrary baseline functions depending on both states and actions.
• Employs optimization objectives such as MinVar and FitQ to learn the baseline function parameters, minimizing variance in gradient estimates.
• Integrates the Stein control variate into policy optimization algorithms like PPO and TRPO without introducing bias.

Experimental results

Research questions

• RQ1 Can we construct more flexible, action- and state-dependent control variates that reduce gradient variance beyond existing methods?
• RQ2 Does the use of Stein's identity enable a broader class of valid control variates compared to prior approaches like Q-prop or A2C?
• RQ3 Can Stein control variates improve sample efficiency in policy optimization without relying on off-policy data?
• RQ4 How do different baseline function forms (e.g., MLP, quadratic, linear) compare in performance when used with the Stein control variate?
• RQ5 Can the method consistently outperform standard value-function baselines and Q-prop in diverse continuous control environments?

Key findings

• The Stein control variate significantly reduces gradient variance across multiple environments, improving sample efficiency in both PPO and TRPO.
• On Humanoid-v1 and HumanoidStandup-v1, PPO with MinVar+MLP baseline achieved 3847 ± 249.3 average return, outperforming the value function baseline (128765 ± 13440) in the same setting.
• In Hopper-v1 and Walker2d-v1, all Stein control variate variants, including FitQ+Linear and MinVar+Linear, outperformed Q-prop despite using only on-policy data.
• The Quadratic baseline performed worse than expected in policy training, likely due to slower convergence compared to MLP and linear baselines.
• MinVar+MLP consistently outperformed other configurations, suggesting it is a strong default choice for the method.
• The method generalizes well across environments, showing consistent improvements over vanilla PPO with value function baselines in HalfCheetah-v1, Ant-v1, and others.