Occupancy-based Policy Gradient: Estimation, Convergence, and Optimality

Thank you very much for your feedback.

---

re: "Introduce new assumptions in the appendix which appear to influence their main results"

The omission of Asm. 7 and 8 from the preconditions of Theorem 16 was an oversight, rather than a deliberate choice to obscure dependencies. We will include Asm. 7 and 8 in the preconditions of Theorem 16 for the next revision, and space allowing, will move the definitions of Asm. 7 and 8 (and surrounding discussions) into the main text.

We made the difficult decision of not including Asm. 7 and 8 (or Algs. 4 and 5) in the main text due to space constraints, given the breadth of our results in both offline and online policy gradient. Our reasoning was that

• Alg. 4 (requiring Asm. 7) and Alg. 5 (requiring Asm. 8) are algorithms from established papers and have been analyzed therein.
• They are orthogonal to the novelty of our paper, which uses algorithms from the aforementioned works as subroutines.
• Most importantly, we wanted to highlight the distinguishing parts of our analysis given limited space.

In addition, they are weaker in some sense than Asm. 5 (expanded on below), which we mentioned briefly in L310, “We focus on discussing Asm. 5 for the offline gradient function class, which requires a stronger level of expressiveness.” Given extra space, we will expand this discussion based on the following points sourced from the Appendix:

Assumption 8

From L1002-1005: “The weight function class completeness assumption is shown Asm. 8, and is satisfied in low-rank MDPs using linear-over-linear function classes that have pseudo-dimension bounded by MDP rank. It can be seen as a 1-dimensional version of Asm. 5 where $\rho = 1$, and in that sense strictly weaker.”

Assumption 7

This is the standard realizability guarantee for maximum likelihood estimation or supervised learning (it is analogous to [Theorem 21, AKKS20] and also required by [HCJ23]), and only requires that the function class includes the ground-truth function. In comparison, Asm. 5 and 8 involve multiple functions.

---

re: "Paper is not self-contained and the results are rather difficult to verify"

Aside from the above issue (for which our previous response proposes a self-contained fix), we believe that the paper is self-contained and verifiable. For all algorithms and results in the appendix we have included rigorous proofs, and assumptions with justification. However, we empathize with the sentiment of your comment in the sense that our results rely on algorithms from established papers as subroutines, which results in layers of analysis. With more space, we hope to provide more details and results on this in the main body, per the existing descriptions, guarantees, and proofs in Appendices D and E.

---

Questions

Yes, it should be $d^\pi = \mu(s)^\top \psi^\pi$. Thank you for catching these typos. We will correct them in our paper.

---

References

[AKKS20] Alekh Agarwal, Sham Kakade, Akshay Krishnamurthy, and Wen Sun. “Flambe: Structural complexity and representation learning of low rank mdps”. In: Advances in Neural Information Processing Systems (2020).

[HCJ23] Audrey Huang, Jinglin Chen, and Nan Jiang. "Reinforcement learning in low-rank mdps with density features." International Conference on Machine Learning. PMLR, 2023.

Thank you for your detailed and insightful feedback, and for your appreciation of our work. We will refine our presentation per your comments, and include additional discussions according to the points below.

---

Motivation

Yes, these are the main contributions of the online and offline algorithms. We will make them more prominent in our revision.

---

Comparison to previous work

For both the online and offline settings, our results can be divided into two parts, estimation and optimization. Once the estimation and gradient domination conditions (e.g., Eq.(4) for online) are established, the optimization analyses largely follow from existing techniques in the literature.

Overall, the dependencies on sample size $n$, Lipschitz constant $\beta$, and iterations $T$ match what we expect to see based on previous work and the optimization literature [Bec17, KU20, AKLM21, NZJZW22, BR24]. For coverage coefficients, the online version (Definition 3) is analogous to those in existing work (ours is finite-horizon while [AKLM21] is infinite-horizon). For offline gradient estimation, the results in [NZJZW22] pay for coverage of all policies (Assumption 6.1) while ours is single-policy (Theorem 16).

Beyond that, our setting of occupancy-based gradient estimation is novel, so a direct comparison for the gradient estimation bounds is not readily available. In addition, few previous works provide end-to-end analysis incorporating both estimation/statistical error and optimization error (e.g., [AKLM21] establish global convergence with true gradients, while [NZJZW22] estimate off-policy gradients but do not use them for policy gradient).

---

Questions

In Corollary 6, what does it mean to run OccuPG with initial distribution $\nu$?

This means that the initial state is drawn from $\nu$, i.e., $s_1 \sim \nu$. The rest of the trajectory is generated by rolling out the policy in the true MDP. We will revise our paper to make this more explicit.

Definition of $\bar\pi_h$ in L238

We believe our use of the $h$ subscript might have caused some confusion (discussed below), but overall yes. By $\bar\pi_h = \left(\pi \wedge C_h^{\mathsf{a}}\pi^D_h \right)$ we do mean $\bar\pi_h(a_h|s_h) = \min \left\lbrace \pi(a_h|s_h), C_h^{\mathsf{a}} \pi^D_h(a_h|s_h) \right\rbrace$, and we will make this more explicit in L238.

With regards to $\pi(a_h|s_h)$ vs. $\pi_h(a_h|s_h)$, for notational compactness we assumed that the same state cannot appear in multiple timesteps (L60-62). So $\pi(a_h|s_h)$ only refers to the policy’s action at timestep $h$, and is de facto equivalent to $\pi_h(a_h|s_h)$.

In the general offline dataset from Def. 7, where each $h$ has a different dataset, $\pi^D_h$ was intended to indicate the data collection policy for timestep $h$. This might have caused some confusion, and we will clarify our notation accordingly.

Benefit to running OccuPG for expected return (esp. online)?

In the online setting with expected return, we view OccuPG as complementary to value-based PG, in the sense that the former can be beneficial when occupancy modeling is more feasible or aligned with existing inductive biases (e.g., on occupancy representations). However, it is possible that calculating $\nabla\log d^\pi$ is more challenging than auto-differentiating through the value-based policy gradient. This is something that we are interested in exploring experimentally.

More importantly, OccuPG serves as the conceptual framework for our offline PG algorithm Off-OccuPG, and here we do see improvements in sample complexity because, unlike previous works [LSAB19, NZJZW22], we get away with single-policy coverage in gradient estimation.

Tightness of convergence bounds, particularly in the case of optimizing general functionals

This is an interesting question, and currently we are not sure. For occupancy gradient estimation, we do not expect that the rate of $n^{-½}$ can be improved, and our dependence in $\mathsf{p}$ matches similar works on (off-policy) gradient estimation [NZJZW22]. It’s possible that $H$ factors can be reduced by more sophisticated handling of error compounding.

For general functionals, the generality of our bound (obtained through the Lipschitz factor, which doesn’t take into account properties of the objective such as curvature) means that it is likely not tight.

In Lemma 5, when should one expect to have $\mathcal{C}^{\pi^*} < \infty$?

$\mathcal{C}^{\pi^*} < \infty$ if we have access to some distribution over states $\nu \in \Delta(\mathcal{S})$, such that $\nu(s) \gg \sum_h d^{\pi^*}(s)$ for all $s$. In other words, it places nonzero mass on all states visited by the optimal policy.

At the worst, one can always guarantee $\mathcal{C}^{\pi^*}$ is finite (though potentially very large) by setting $\nu = \mathrm{uniform}(\mathcal{S})$ to be uniform over all states. In practice, one may be able to craft a more refined $\nu$ that results in smaller $\mathcal{C}^{\pi^*}$ given some expert or domain knowledge.

---

References

[AKLM21] Alekh Agarwal, Sham M Kakade, Jason D Lee, and Gaurav Mahajan. “On the theory of policy gradient methods: Optimality, approximation, and distribution shift”. In: The Journal of Machine Learning Research 22.1 (2021)

[Bec17] Amir Beck. First-order methods in optimization. SIAM, 2017.

[BR24] Jalaj Bhandari and Daniel Russo. “Global optimality guarantees for policy gradient methods”. In: Operations Research (2024)

[LSAB19] Yao Liu, Adith Swaminathan, Alekh Agarwal, and Emma Brunskill. “Off-policy policy gradient with state distribution correction”

[KU20] Nathan Kallus and Masatoshi Uehara. “Statistically efficient off-policy policy gradients”. In: International Conference on Machine Learning. PMLR. 2020

[NZJZW22] Chengzhuo Ni, Ruiqi Zhang, Xiang Ji, Xuezhou Zhang, and Mengdi Wang. “Optimal Estimation of Policy Gradient via Double Fitted Iteration”. In: International Conference on Machine Learning. PMLR. 2022

Thank you very much for your comments. However, we respectfully disagree that our paper over-claims its results. The reviewer’s comment that “the coverage coefficient… removes the need of exploration” likely results from confusion over the term “exploration”, whose meaning in the PG literature is different from that in PAC-RL. Our use of the term is aligned with the literature, and our results’ dependence on the coverage coefficient is also standard. Results of a similar nature are commonly found in the literature, and we elaborate on the details below.

---

re: "The convergence depends on a coverage coefficient, which removes the need for exploration"

In (online) policy gradient analysis, the coverage coefficient $\mathcal{C}^{\pi^*}$ (Lemma 5) expresses the difficulty of the exploration problem in policy optimization. L184-188 discusses this in-depth.

There may be some confusion over terminology because our algorithm, along with routine PG methods such as REINFORCE and PPO, do not actively explore in the sense of PAC exploration (e.g., UCB and other optimism-in-face-of-uncertainty algorithms). Rather, exploration here refers to the difficulty of finding a good policy within the policy class from online interactions with the environment.

This usage of the phrase “exploration” matches foundational works in the PG literature, for which $\mathcal{C}^{\pi^*}$ is also the standard notion of online PG complexity. They are identical to Definition 3.1 from the seminal work of [AKLM21], and the paragraph above it (modulo differences in infinite and finite horizon MDPs). Definition 3 of [BR24] is another point of comparison. All of these coefficients depend on how well an initial state distribution covers the optimal policy’s occupancy.

That said, we should have defined “exploration” more explicitly in the abstract/introduction to reduce confusion, and will revise accordingly. This is a simple fix that can be lifted from L184-188, and that is aligned with standards and terminology from existing work.

---

re: "the obtained bound does not indicate the number of samples needed for exploration when the initial policy is not good"

First of all, the coverage coefficient is defined with respect to a distribution over initial/starting states. This is not an initial policy, or even a policy.

Moreover, a good coverage coefficient does not necessarily imply a good initial policy; the initial policy can still have poor performance, and we rely on gradient ascent to find a better policy through online interactions. This is what the PG literature means by “exploration”.

---

re: "Additionally, the proposed algorithm does not address the problem of exploration"

As argued above, the coverage coefficient in our online bounds characterizes the difficulty of exploration faced by policy gradient optimization algorithms, on par with the existing literature.

However, developing algorithms that actively explore to improve policy optimization is a valuable direction of future work, given that in practice the initial state distribution may not be very exploratory (aka have poor coverage over $d^{\pi^*}$ per Lemma 5).

---

Coverage coefficient in offline setting

Lastly, our results cover both online and offline PG, and the latter makes up half of the paper (pages 6-9). In the offline learning setting, the learner cannot interact with the environment and must learn from the given data, hence the problem of exploration simply does not exist there. As with all offline results, a coverage coefficient is expected in the guarantee, reflecting the quality of the offline data (since, if the data is poor, there is nothing one can do). In fact, the kind of coverage coefficient we use is already improved over previous offline PG results; previous results require all-policy coverage even just for estimating the gradient [LSAB19, NZJZW22], whereas our estimation guarantee only depends on single-policy coverage (Theorem 16).

---

References

[AKLM21] Alekh Agarwal, Sham M Kakade, Jason D Lee, and Gaurav Mahajan. “On the theory of policy gradient methods: Optimality, approximation, and distribution shift”. In: The Journal of Machine Learning Research 22.1 (2021), pp. 4431–4506.

[BR24] Jalaj Bhandari and Daniel Russo. “Global optimality guarantees for policy gradient methods”. In: Operations Research (2024)

[LSAB19] Yao Liu, Adith Swaminathan, Alekh Agarwal, and Emma Brunskill. “Off-policy policy gradient with state distribution correction”. In: arXiv preprint arXiv:1904.08473 (2019).

[NZJZW22] Chengzhuo Ni, Ruiqi Zhang, Xiang Ji, Xuezhou Zhang, and Mengdi Wang. “Optimal Estimation of Policy Gradient via Double Fitted Iteration”. In: International Conference on Machine Learning. PMLR. 2022, pp. 16724–16783.