Source-linked AI summary

PC-PG: Policy Cover Directed Exploration for Provable Policy Gradient Learning

Alekh Agarwal, Mikael Henaff, Sham Kakade, Wen Sun

arXiv:2007.08459v2cs.LGcs.AIstat.ML

TL;DR

Policy gradients can lack the state coverage needed for effective exploration, while existing guarantees for model misspecification and function approximation are limited. PC-PG uses a learned policy cover to direct exploration while optimizing policies on-policy. The paper gives sample guarantees for tabular and linear MDPs and state aggregation, and reports success in a chain-lock experiment where PPO fails at longer horizons and PPO+RND succeeds only about half the time.

  • Problem

    Policy-gradient guarantees rely on state coverage, yet starting from a single state can leave coverage inadequate, and robustness to function approximation and model misspecification remains a concern.

    Source: §1 · p. 2
  • Method

    PC-PG alternates policy improvement with a learned policy cover and uses a coverage-based reward bonus to encourage exploration.

    Source: §3 · pp. 8, 11
  • Results

    PC-PG has polynomial sample guarantees for linear MDPs with RKHS features and state aggregation, and in the chain-lock experiment it succeeds across seeds and horizon lengths while PPO fails at longer horizons and PPO+RND succeeds roughly 50% of the time.

    Source: §1.1 · pp. 4–5, §5.1 · p. 22
  • Takeaways & Limitations

    The learned cover provides near-uniform coverage over both chains in the chain-lock experiment, where PC-PG succeeds across tested seeds and horizons.

    Source: §5.1 · p. 22

    The state-aggregation guarantee depends on bounded approximation error under a comparator policy's state distribution, and feature orthogonality is important for a provable guarantee in the partially well-specified setting.

    Source: §4.3 · p. 15, §4.4 · p. 19

Abstract

from arXiv · show

Direct policy gradient methods for reinforcement learning are a successful approach for a variety of reasons: they are model free, they directly optimize the performance metric of interest, and they allow for richly parameterized policies. Their primary drawback is that, by being local in nature, they fail to adequately explore the environment. In contrast, while model-based approaches and Q-learning directly handle exploration through the use of optimism, their ability to handle model misspecification and function approximation is far less evident. This work introduces the the Policy Cover-Policy Gradient (PC-PG) algorithm, which provably balances the exploration vs. exploitation tradeoff using an ensemble of learned policies (the policy cover). PC-PG enjoys polynomial sample complexity and run time for both tabular MDPs and, more generally, linear MDPs in an infinite dimensional RKHS. Furthermore, PC-PG also has strong guarantees under model misspecification that go beyond the standard worst case $\ell_{\infty}$ assumptions; this includes approximation guarantees for state aggregation under an average case error assumption, along with guarantees under a more general assumption where the approximation error under distribution shift is controlled. We complement the theory with empirical evaluation across a variety of domains in both reward-free and reward-driven settings.

1 Introduction

Policy gradients directly optimize rewards but can fail to explore when their initial state coverage is poor. PC-PG uses a learned policy cover to add strategic exploration while retaining a model-free policy-optimization approach, with guarantees for tabular and linear MDPs and under certain forms of model misspecification.

1 Introduction

  • Poor state coverage can leave policy gradients at low-quality stationary points, while sparse rewards can produce vanishing gradients.The paper motivates exploration methods that preserve policy gradients' direct optimization and ability to handle function approximation and modeling violations.Source: §1 · p. 2

1.1 Our Contributions

  • PC-PG explores with a learned ensemble of policies that forms a state-space cover, and its on-policy approach addresses catastrophic forgetting and Bellman-backup bias.The paper presents it as a direct, model-free policy-optimization alternative to model-based and value-based exploration methods.Source: §1.1 · pp. 2, 4

1.2 Related Work

  • Prior provable exploration methods largely use model-based or value-based approaches; PC-PG adds policy optimization as a model-free alternative.The learned cover lets the algorithm plan through real-world rollouts without storing prior data as an empirical model.Source: §1.1 · p. 4

1.1 Our Contributions

  • PC-PG is sample- and computationally efficient in tabular and linear MDPs, with guarantees extending to infinite-dimensional RKHS features.For linear MDPs, the stated sample bound depends polynomially on the effective dimension measured by maximum information gain and on the maximum RKHS norm of policy Q-functions.Source: §1.1 · p. 4
  • PC-PG also has favorable sample complexity when critic-fitting transfer error is small, a condition described as an average-case alternative to ℓ∞ misspecification assumptions.Transfer error measures prediction error under a comparator policy's state distribution rather than the training distribution.Source: §1.1 · pp. 4–5
  • Under state aggregation, PC-PG's approximation error is weighted by visitation under an optimal policy rather than requiring uniformly accurate aggregation.The informal guarantee gives an ϵ + ϵ_misspec near-optimal policy with sample complexity eO(poly(|Z|, 1/(1 − γ), 1/ϵ)).Source: §1.1 · pp. 4–5
  • Experiments show PC-PG can complement deep-learning approaches and work in settings where prior bonus-based methods such as RND do not recover optimal policies with high probability.The authors characterize this as evidence of viability beyond the regime covered by their sample-complexity bounds.Source: §1.1 · p. 5

1.2 Related Work

  • For tabular and linear MDPs, PC-PG is PAC without a reset-distribution assumption, unlike CPI and NPG under the cited comparison.The paper contrasts this with prior policy-optimization results that require further assumptions on the reset distribution.Source: §1.2 · p. 6

2 Setting

The setting is a discounted MDP in which the agent seeks a policy maximizing expected value from a starting state. The agent can reset to that state and estimate values through rollouts.

2 Setting

  • An MDP is specified by state and action spaces, transition and reward functions, a discount factor, and a starting state; actions are discrete.The reward lies in [0, 1] and the discount factor lies in [0, 1).Source: §2 · p. 7
  • A policy maps each state to a distribution over actions, and the goal is to maximize expected value from the starting state.The paper defines value and state-action value as expected discounted future rewards under a policy.Source: §2 · p. 7
  • The agent is assumed able to reset to the starting state at any point in a trajectory.The text notes this can be replaced by termination at each step with probability 1 − γ.Source: §2 · p. 7
  • The ν-sampler draws an initial state-action pair from ν and follows a policy with geometric episode termination; the Qπ sampler gives an unbiased value estimate from one rollout.These sampling procedures support visitation-measure sampling and Q-value estimation.Source: §2 · p. 8

3 The Policy Cover-Policy Gradient (PC-PG) Algorithm

PC-PG alternates policy improvement with expansion of a policy-cover distribution. It uses a coverage-dependent reward bonus and constrained natural policy-gradient updates to seek underexplored state-action pairs.

3 The Policy Cover-Policy Gradient (PC-PG) Algorithm

  • Ordinary policy-gradient ascent can stall at a poor stationary point when the initial policy has inadequate state coverage.PC-PG therefore improves both the current policy and its coverage distribution across episodes.Source: §3 · pp. 8–9
  • At episode n, PC-PG forms its cover distribution by averaging the state-action visitation measures of the policies learned so far.It adds a bonus to the reward and constrains bonus-positive actions to encourage exploration beyond already covered regions.Source: §3 · p. 9
  • The bonus uses feature covariance under the cover to reward state-action pairs whose features are less explored.In tabular MDPs with one-hot features, the bonus is proportional to the inverse of the number of observations of each pair.Source: §3 · pp. 9–11
  • PC-PG updates policies with natural policy gradient, fitting a value-function critic under the cover and restricting updates at states with positive bonus.The critic fits Qπ(s, a; r + b_n) − b_n before adding the bonus back; positive-bonus actions are chosen uniformly at those states.Source: §3 · pp. 10–11

4 Theory and Examples

Section 4 establishes PC-PG guarantees for RKHS linear MDPs and model-misspecified settings, including state aggregation, bounded transfer error, and partial well-specification. Under partially well-specified models, it also gives guarantees for competing with comparator policies whose values PC-PG realizes exactly, based on feature orthogonality.

4.1 Well specified case: Linear MDPs

  • PC-PG has polynomial sample complexity in the maximum information gain for linear MDPs with possibly infinite-dimensional RKHS features.The paper also states that the result extends to tabular MDPs, replacing information gain with a term depending on the number of states and actions.Source: §4 · p. 11, §4.1 · p. 13

4.2 State-Aggregation under Model Misspecification

  • For state aggregation, PC-PG’s approximation error depends on average misspecification under the comparator’s abstracted state distribution, rather than worst-case global error.The bound may be substantially smaller when this average error is much less than the maximum aggregation error; the stated approximation error also incurs an additional 1/(1 − γ) factor.Source: §4.2 · p. 14

4.3 Agnostic Guarantees with Bounded Transfer Error

  • With bounded transfer error, PC-PG achieves near-optimal performance in polynomial sample complexity when the error is small, without further assumptions on the MDP.Transfer error measures how well an on-policy critic fit under the policy-cover training distribution predicts under a fixed comparator distribution; it is zero for tabular and linear MDPs and bounded by comparator-averaged aggregation error in the state-aggregation case. The assumption was introduced by Agarwal et al. [2019].Source: §4.3 · pp. 15–16
  • Unlike the comparison with Agarwal et al. [2019], PC-PG needs no well-conditioned exploratory initial distribution: its policy cover supplies exploration from the fixed start state.The algorithm’s reward bonus encourages visits where the current cover distribution has a large relative condition number.Source: §4.3 · pp. 16–17

4.4 Robustness to “Delusional Bias” with Partially Well-specified Models

  • For the binary-tree example, PC-PG is guaranteed to find a policy with value greater than 1/2 − ϵ with probability greater than 1 − δ using O(poly(H, d, 1/ϵ, log(1/δ))) samples, because transfer error is zero.Source: §4.4 · p. 18
  • PC-PG may fail to find the globally best policy, but can compete with the best policy in the favorable sub-part satisfying the modeling assumptions.Feature orthogonality is important for a provable guarantee; otherwise, value-fitting errors on the binary subtree can damage estimates on favorable parts.Source: §4.4 · pp. 18–19
  • Bellman-backup approaches face challenges because exact backups can be possible on misspecified regions, while current algorithms lack a guarantee of one-sided value-estimation errors.With arbitrary featurization, the authors say it is unclear why algorithms would avoid estimating a subtree’s value above 1/2 when its true value is below 1/2 − ∆.Source: §4.4 · p. 19
  • State-aggregation features permit any tabular RL algorithm, and PC-PG’s average-case guarantee is distinctive because comparable guarantees are not evident for current tabular methods.The authors attribute the difficulty to those methods’ strong reliance on how they use the Markov property.Source: §4.4 · p. 19
  • PC-PG still competes with a comparator following the leftmost path because it realizes that path’s value functions exactly, without interference from the rest of the MDP.The authors note that making the left path a balanced binary tree with O(H) additional features makes it unclear how to design an initial distribution with a good concentrability coefficient.Source: §4.4 · pp. 19–20

5 Experiments

The experiments evaluate PC-PG’s policy cover in exploration tasks, including reward-free mazes and reward-driven problems. Across the reported tasks, PC-PG combined with PPO outperforms vanilla PPO and PPO+RND on challenging exploration tasks.

5 Experiments

  • PC-PG combined with PPO significantly outperforms vanilla PPO and PPO+RND on challenging exploration tasks.The experiments also test compatibility with neural-network policy optimization.Source: §5 · p. 20
  • The experiments test whether PC-PG’s policy cover widens exploration, finds near-optimal policies in sparse-reward tasks, and avoids catastrophic forgetting; neural policies use a kernel bonus and a rebalanced cover.The rebalancing scheme selects a cover intended to maximize state-space coverage.Source: §5 · p. 20

5.1 Bidirectional Diabolical Combination Lock

  • PC-PG succeeds across seeds and horizons, while PPO+RND succeeds roughly 50% of the time and PPO only at H = 2; PC-PG’s cover nearly uniformly spans both chains.PPO+RND can reach one chain but may become deterministic and forget to explore the other.Source: Figure 3 · p. 21, §5.1 · p. 22

5.2 Reward-free Exploration in Mazes

  • In reward-free mazes, PC-PG quickly visits a significantly higher proportion of locations than PPO and PPO+RND over 10 million steps; cover policies show diverse traces.PPO’s visited proportion stays relatively constant, while PPO+RND explores to some degree.Source: §5.2 · p. 22, Figure 4 · p. 22

5.3 Continuous Control

6 Discussion

The authors report evidence that PC-PG can be combined with neural policy optimization and be effective in practice. They identify learning good features for rich-observation settings as future work.

6 Discussion

  • The experiments provide evidence that PC-PG can be combined with neural policy optimization methods and be effective in practice.Source: §6 · p. 23
  • Future work could combine PC-PG with unsupervised feature learning, including autoencoders or noise-contrastive estimation.Source: §6 · p. 23
  • For rich-observation settings, the authors identify learning a good feature representation as a future direction.Source: §6 · p. 24

A NPG Analysis (Algorithm 3)

The analysis uses an augmented MDP to study NPG’s progress on a bonus-reward objective, then relates that analysis to performance on the original MDP.

A.1 Set up of Augmented MDPs

  • The analysis distinguishes the original MDP, the bonus-reward MDP optimized by NPG, and an augmented MDP used only for analysis.The augmented MDP adds a self-looping action outside the known set, with maximum reward.Source: §A.1 · p. 32

B Relationship between Mn and M

  • The added action absorbs the comparator outside the known set, so its state distribution there does not continue into known states.The comparator matches the original policy on known states; the analysis uses this structure to relate its distributions across MDPs.Source: §A.1 · p. 32, B Relationship between Mn and M · pp. 35–36

A.2 Performance of NPG (Algorithm 3) on the Augmented MDP Mn

  • Mirror-descent analysis bounds NPG’s regret against a comparator policy on the augmented MDP for a fixed episode.The comparator follows its original policy on known states and selects the added action elsewhere.Source: §A.1 · p. 32, §A.2 · pp. 33, 35

B Relationship between Mn and M

  • The analysis uses the relationship between the augmented and bonus-reward MDPs to transfer NPG’s performance bound to the original MDP.A simulation-lemma-style argument links policy performances, while a potential-function argument relates escaping probability to maximum information gain.Source: §A.1 · p. 32, B Relationship between Mn and M · pp. 36–37

C Analysis of PC-PG for the Agnostic Setting (Theorem 4.3)

This section presents PC-PG’s general sample-complexity result under a transfer-bias condition. Its proof bounds exploration episodes and accounts for critic prediction and covariance-estimation errors.

C Analysis of PC-PG for the Agnostic Setting (Theorem 4.3)

  • The section presents the general sample-complexity result for PC-PG under a transfer-bias assumption.The theorem states a high-probability guarantee for any comparator policy in the specified class.Source: C Analysis of PC-PG for the Agnostic… · pp. 37–38
  • The outer-episode bound uses the policy cover’s discovery of unknown states and the limited number of significantly new feature directions.When unknown states are reachable with good probability, NPG finds them because they carry high reward; each discovery adds directions not already well covered.Source: C Analysis of PC-PG for the Agnostic… · p. 38
  • The proof also accounts for critic prediction error through regression analysis and covariance-estimation error through matrix concentration.The section identifies stochastic-gradient regression and matrix concentration inequalities as tools for these bounds.Source: C Analysis of PC-PG for the Agnostic… · p. 39

C.1 Proof of Theorem C.1

The proof assembles the sample-complexity result by controlling critic bias and statistical error, then accounting for regression and covariance-estimation samples.

C.1 Proof of Theorem C.1

  • The critic regresses the bonus-adjusted Q-function onto features, with the bonus subtracted so the fitted target is linear in the linear-MDP setting.The proof describes regression from features to Qπ(s, a; r + bn) − bn(s, a).Source: §C.1 · p. 39
  • Small critic statistical error and feature bias yield low NPG regret against the episode comparator.The proof combines the NPG convergence result with a variance-and-bias tradeoff lemma under the stated comparator-distribution assumption.Source: §C.1 · pp. 39, 41
  • The theorem’s near-optimality conclusion follows by choosing hyperparameters to satisfy the conditions of the preceding regret bound.The proof states that the theorem is verified by substituting its hyperparameter values into the earlier lemma.Source: §C.1 · pp. 41–42
  • Regression samples are required across each NPG iteration and episode, while covariance estimation contributes an additional sample source.The text states that on-policy fitting is repeated inside episodes and that covariance matrices also need to be estimated.Source: §C.1 · pp. 42–43

D Analysis of PC-PG for Linear MDPs (Theorem 4.1)

For linear MDPs, the proof establishes that policy Q-functions are linear in the features and that subtracting the known reward bonus preserves this linear structure for critic fitting.

D Analysis of PC-PG for Linear MDPs (Theorem 4.1)

  • For any policy in a linear MDP, its Q-function is linear in the features, with a bounded coefficient norm under the stated parameter bounds.The bound depends on the model parameter norms and maximum value.Source: D Analysis of PC-PG for Linear MDPs… · p. 44
  • Although adding the reward bonus can make the bonus-reward MDP non-linear under the original features, subtracting that known bonus from Q restores linearity.The resulting function Qπ(s, a; r + bn) − bn(s, a) equals a feature-linear function for every state-action pair.Source: D Analysis of PC-PG for Linear MDPs… · p. 44
  • Under Assumption 4.1, the proof establishes zero critic feature bias for each episode and iteration.This supports the critic choice that regresses the bonus-adjusted Q-function onto the features.Source: D Analysis of PC-PG for Linear MDPs… · p. 44

E Analysis of PC-PG for State-Aggregation (Theorem 4.2)

State aggregation introduces nonzero transfer error, but the analysis relates it to model-misspecification averaged over a fixed comparator’s state distribution. The section defines aggregation error by within-group transition and reward differences and uses a variance-bias tradeoff lemma.

E Analysis of PC-PG for State-Aggregation (Theorem 4.2)

  • The transfer error is related to expected model-misspecification averaged over a fixed comparator’s state distribution.Unlike the linear-MDP analysis, state-aggregation misspecification means the transfer error is not zero.Source: E Analysis of PC-PG for… · pp. 45, 47
  • State aggregation maps state-action pairs to feature groups, with aggregation error defined by within-group transition and reward differences.The state-action feature indicates which group z a pair maps to.Source: E Analysis of PC-PG for… · p. 45
  • The state-aggregation bias-variance lemma replaces the linear-MDP term √Aϵbias with an average model-misspecification term.The replacement is expressed using expected error under the comparator’s state distribution, scaled by (1 − γ)^2.Source: E Analysis of PC-PG for… · pp. 46–47

F Analysis of PC-PG for the Partially Well-specified Models (Corollary 4.1)

For partially well-specified models, the proof constructs an on-policy best fit with zero transfer error. Auxiliary lemmas provide dimension-free least-squares and covariance concentration guarantees.

F Analysis of PC-PG for the Partially Well-specified Models (Corollary 4.1)

  • For any state-action distribution and policy, the proof constructs a best on-policy fit whose error under the comparator distribution is zero.The construction matches Qπ − b on the states visited by the comparator policy.Source: F Analysis of PC-PG for the Partially… · p. 48

G Auxiliary Lemmas

  • The dimension-free least-squares guarantee depends on feature and parameter norms and the bound on y, not the feature dimension.The stated bound uses R = 3(W^2 + WH).Source: G Auxiliary Lemmas · p. 48
  • The auxiliary lemmas establish covariance-matrix concentration and a concentration guarantee for estimates using the inverse covariance matrix.The inverse-covariance result assumes K samples from each of N distributions and specifies a choice of K.Source: G Auxiliary Lemmas · pp. 50–51

H.1 Algorithm Implementation

The experiments implement reward-bonus and reward-free PC-PG variants using NPG with policy-cover restart distributions. The appendix describes their environment setups and implementation choices.

H.1 Algorithm Implementation

  • Both algorithm versions use NPG with restart states sampled by rolling in with a policy from the cover, and experiments use PPO for policy updates.One version adds a reward bonus; the other performs reward-free exploration and can then exploit using the cover as a start distribution.Source: §H.1 · p. 52

H.2.1 Bidirectional Diabolical Combination Lock

  • The bidirectional combination lock requires exploring both locks to find the optimal reward, since pursuing the reward-2 lock yields return 1 while the reward-5 lock yields optimal return 4.The environment includes locally optimal policies that enter dead states or settle for the lower-reward lock.Source: §H.2.1 · p. 54

H.2.2 Mountain Car

  • Mountain Car uses a continuous two-dimensional state space and one-dimensional continuous action space with the reward-bonus algorithm.Source: §H.2.2 · pp. 54–55

H.2.3 Mazes

  • The maze experiments use a 20 × 20 environment with five actions, constant zero reward, and reward-free exploration without exploitation.Source: §H.2.3 · p. 55

H.3 Hyperparameters

  • The appendix reports PPO and PC-PG hyperparameters in separate tables and gives the PC-PG feature choices for each environment.PC-PG uses state features for the combination lock, Random Kitchen Sinks features for Mountain Car, and a randomly initialized convolutional network for mazes.Source: Table 2 · p. 55, §H.3 · p. 56, Table 3 · p. 56
Loading 2007.08459v2…