On the Sample Complexity of a Policy Gradient Algorithm with Occupancy Approximation for General Utility Reinforcement Learning

Question 1. We will provide further details as requested by the reviewer. Please see the above comments for an extended discussion. As for the choice of the ‘function class so that the true occupancy measure can be estimated well, i.e. Assumption 1(ii) is satisfied’:

• As we highlight in l. 242-246, state-occupancy measures are linear (or affine in the discounted setting) in density features in low-rank MDPs. We refer the reader to Appendix B for a proof of this statement. Therefore, in this case, it is natural to approximate occupancy measures via linear function approximation using some density features.

• As we mention in l. 323, this realizability assumption can be relaxed at the price of incurring a misspecification error. This assumption allows us to get rid of the approximation error stemming from the fact that the occupancy measure might not belong to our function class. This induced error would depend on the expressivity and richness of the function class. One may relax it by introducing an additional function approximation error and propagate it in our analysis. This would result in a constant approximation error $\epsilon_{\approx}$ (assuming this is the constant upper bound on that error, which is rather standard in the literature, see e.g. Agarwal et al. 2021 ‘On the Theory of Policy Gradient Methods: Optimality, Approximation, and Distribution Shift’) propagating in Proposition 1 and then in our theorems. The propagation can be seen in Theorem 1 where the last term corresponds to the occupancy measure estimation error. We will add a comment about this.

Questions 2 and 3. Please see the detailed discussion above.

Question 4: challenge in the proof of Theorem 1 and Theorem 2 compared to Zhang et al. (2021).

• Recall first that the analysis of Zhang et al. 2021 is restricted to the tabular setting and their algorithm uses non-scalable count-based estimates of the occupancy measure.

• Our analysis does of course use similar parts of the technical analysis of Zhang et al. 2021 concerning the optimization part (use of hidden convexity with the required policy parametrization assumption and leveraging smoothness of the objective precisely) and we do acknowledge it in the paper and the proofs. Our contributions from the technical viewpoint consist in (a) approximating the occupancy measure, provide statistical guarantees for the MLE approach we follow using prior work, introduce the right assumptions for our problem (see Assumption 1, proposition 1) and (b) propagate the estimation error into the optimization analysis to obtain Theorems 1 and 2 as well as setting up all the parameters (step sizes, number of samples, …) to derive sample complexities.

Thanks again for your review. Please let us know if you have any further comments or questions, we will be happy to address them.

2. Coverage and scalability problem. Now you may argue that it is enough to take an initial distribution (or just $\rho$ distribution for the expected loss if one assumes it is unrelated to the initial distribution) that just needs to cover the support of the occupancy measure we want to estimate. Note that the occupancy measure is unknown and we want to estimate it so we have a priori no clue about its support. One might then think about just taking the uniform distribution as an initial distribution to be sure to cover the whole state space equally. This choice is problematic for several reasons: (a) First, this introduces a bias: Why would we need to estimate the occupancy measure equally well in all the states if the occupancy measure is concentrated on a specific set of states which is not necessarily the entire state space? (b) Second and most importantly, if we make such a choice, we have now $\rho_min = 1/|S|$ (say in the discrete state space setting) and the first order stationarity bound in Barakat et al. 23 scales with $1/\rho_{\min}$, this makes the result not scalable to large state spaces. You might argue that we do not need the uniform distribution but just to take a distribution $\rho$ covering the entire state space (not necessarily equally well), i.e. which has a support equal to the entire state space. Then again, this introduces a bias as the loss minimization might focus on states which are irrelevant to the occupancy measure we want to estimate.

3. An additional point that we want to highlight here and that relates to the shortcomings above is that MSE might not be the best metric for comparing distributions because it focuses on pointwise differences (in our case states or state-action pairs). In our setting, how the weights of the MSE loss are chosen for fitting our probability distribution of interest is important.

There are a number of shortcomings of using MSE compared to MLE for fitting a probability distribution in general, we summarize them here:

1. Consistency and efficiency. MLE maximizes the likelihood function, ensuring the estimator is consistent (i.e. converges to the true parameter value as the sample size increases) and asymptotically efficient (i.e. achieves the lowest possible variance among unbiased estimators for large samples). In contrast, as MSE minimizes mean squared errors, it does not guarantee properties like consistency or efficiency unless under specific assumptions on the data such as normality of the errors in which case both coincide. Note that the approach in Barakat et al. 2023 does not fit a Gaussian distribution to the normalized occupancy measure or anything like that. We use favorable statistical properties of MLE (see our proposition 1).

2. Sensitivity to scaling. This is an important point regarding our discussion that also relates in a way to scaling to large state spaces. MLE operates on probabilities and likelihoods, these are normalized and scale-invariant. This makes MLE very suitable for probability distribution fitting in general. MSE is rather better used for point wise estimation in statistics (which is also indirectly used for distribution estimation via estimating parameters such as Gaussian means), MSE depends on the scale of the observations which might make it less robust in some settings. Please see also next point as a follow-up.

3. Robustness to outliers. As MLE models probabilities directly, it might be more robust to outliers depending on the distribution. MSE can be more sensitive to outliers and extreme values as it relies on squared errors which amplify such outliers. In our setting, this is also relevant as we are interested in estimating occupancy measures on large state spaces, this induces small probability values (even extremely small for some of them) and squaring differences makes it worse, this is what we have explained in the appendix with simple calculations showing how the state space size shows up.

4. Satisfying distribution constraints. MLE naturally adapts to the distribution's shape and constraints to satisfy them. In contrast, MSE can lead to estimates that violate distribution constraints such as probability normalization or predicting a negative variance for a normal distribution. Therefore, post-processing might be required to ensure these are satisfied.

We focused on the state space here, similar comments can be made for the choice of the uniform action distribution in the expected loss.

We hope that these clarifications have convinced the reviewer of the fundamental issues with such an MSE approach as well as advantages of ours. We believe that simulations are not even needed to illustrate this (some of these facts which are not necessarily specific to our setting are even general and well-known in statistics). We will add some of these points to our paper to strengthen our motivation for using our approach rather than theirs.

We thank the reviewer for their feedback and for their questions. We are glad that they found the paper ‘well-written and easy to read.’ We respond to each one of their comments and questions in the following to further clarify our contributions.

Preliminary points about novelty compared to Barakat et al. 2023. Let us first recall a few points regarding our contributions compared to that work along the lines of the detailed comparison we provided in the appendix (p. 15-16) and l. 265-267 in the main part.

1. Experiments. They do not provide any simulations testing their algorithm in section 5 for large state action spaces, Fig. 1 therein is only for the tabular setting for a different algorithm. We do provide simulations and compare to the tabular setting (with count-based approaches) to demonstrate the scalability issue.

2. Global convergence. In contrast to our work (see our theorem 2 and corollary 2), they only provide a first-order stationarity guarantee and they do not provide global convergence guarantees;

3. Technical analysis. From the technical viewpoint, our occupancy measure MLE estimation procedure combined with our PG algorithm requires a different theoretical analysis even for our first order stationarity guarantee. Please see appendix D for details.

Due to the length of the discussion, we have delayed it to the appendix while only keeping the salient points of the comparison (namely the advantage of MLE compared to their MSE approach) to avoid interrupting the flow of the paper. We will expand on this discussion in the main part to further clarify our contributions as suggested by the reviewer.

MSE in Barakat et al. 23 vs MLE (ours) for distribution estimation. While this might indeed be a simple algorithmic change (and simplicity of the algorithm should perhaps be seen positively), we argue that it is crucial given the main goal and motivation of scaling to larger state action spaces. Our discussion might have given the impression that the assumption that p^* is uniform is needed to showcase our point, it is not. Our discussion regarding this point might have confused the reviewer. We provide clarifications here regarding why the MSE formulation in Barakat et al. 2023 is problematic for occupancy measure approximation regarding scalability and other fundamental aspects. Recall from that paper (see Eq. 11 therein) that the loss for approximating an occupancy measure (induced here by policy $\pi_{\theta}$) via linear regression is $E_{s \sim \rho, a \sim \mathcal{U}(\mathcal{A})}[(\lambda^{\pi_{\theta}}(s,a) - \langle \phi(s,a), \omega \rangle)^2]$ where $\rho$ is the initial state distribution and $\mathcal{U}(\mathcal{A})$ the uniform action distribution. We now make a few points regarding why there are issues with this objective:

1. The expected loss is over the initial state distribution (defining the MDP). Take the extreme case where we initialize at a single state (note that this is also realistic, e.g. a robot starting at a given deterministic state). Then this means that the expected loss boils down to estimating the occupancy measure only at that state. However we need to estimate it as accurately as possible for all states and there is no reason why the occupancy measure should be supported by the same set of states as the initial distribution (which we should have freedom about). Note also that the occupancy measure itself depends on the initial distribution. In principle, the distribution used for defining the expected loss (over state action pairs) should be different from the initial state distribution defining MDP.

Thank you for your response. We want to clarify here that there is 'no rigorous analysis for the general case' to come up with here (as you suggest) as we are not considering a particular case in our discussion above (regarding the MSE approach in prior work) which starts from their general formulation of the problem. We have just clearly shown above that the MSE approach in prior work has fundamental issues even from the initial formulation of the problem of occupancy measure estimation and it has not been tested in practice in prior work.

We thank the reviewer for their positive feedback on our work. We respond to their comments and questions in the following.

Predictable Results: Although the introduction of MLE estimation and corresponding theoretical analysis is a novel and concrete contribution, the state space size independent performance bound is somewhat predictable given the assumptions. This result is not surprising and thus may not add substantial new insights. Additionally, the widespread use of MLE in practice diminishes the practical value of the proposal.

We thank the reviewer for acknowledging the novelty of our contributions. Leveraging existing statistical estimation error bounds for MLE, we achieve the first 1st order stationarity and global optimality results (with sample complexity guarantees) for solving RLGU in large state action spaces to the best of our knowledge. Estimating occupancy measures using function approximation is not common in the RL literature even if MLE is standard in statistics. We also stress the advantage of using MLE over prior work which proposed a MSE approach which is not well-suited for distribution estimation and which has a number of limitations as it is introduced in Barakat et al. 2023, please see our response to reviewer MseJ for a detailed discussion.

Assumption 3 (General utility smoothness):

• This assumption is quite standard in the RLGU literature: see e.g. Hazan et al. ICML 2019, Zhang et al. NeurIPS 2020, NeurIPS 2021, AAAI 2022, Barakat et al. ICML 2023, Ying et al. AAAI 2023, NeurIPS 2023.
• The assumption is satisfied for several functions of interest including the standard RL setting, quadratic objectives (recall that the utility variables are confined to the simplex) and more generally for any smooth objective. While the KL divergence is not strictly speaking smooth, the smoothed KL which only adds a small numerical constant inside the log term is actually smooth (as it allows to get rid of the singularity at the boundary of the simplex). The smoothed entropy defined as $H_{\sigma}(x) = - x log(x + \sigma)$ is $\frac{1}{2\sigma}$-smooth w.r.t. the infinity norm and has been used in the RLGU literature (and beyond), we refer the reader to the work of Hazan et al. 2019 (see e.g. Lemma 4.3 therein). Overall this assumption captures most of the problems of interest in RLGU including pure exploration (using the smoothed entropy), learning from demonstrations (using the smoothed KL) as well as standard linear RL and CMDPs. Other entropic measures or $l^2$ losses are also possible.

Assumption 5:

• This assumption is also commonly used in the RLGU literature, in all the aforementioned works that we mentioned for Assumption 3 (see also Ying et al. 2024 ‘Policy-based Primal-Dual Methods for Concave CMDP with Variance Reduction’).
• It is satisfied for instance for tabular policy parameterizations. However, as we mention in the paper, this assumption is difficult to verify in general and this is also acknowledged by prior work. We note that it is only a local assumption though. Trying to relax such an assumption is an interesting question, we leave this for future work.
• Note that this assumption is not needed for our first order stationarity results (Theorem 1 and Corollary 1). It is only required for Theorem 2 and Corollary 2.

Can you provide specific types of RLGU problems where these assumptions might hold or discuss how restrictive these assumptions are in practice?

Please see our responses above. All our assumptions are discussed in the paper and examples as well as references are provided, except for the smoothness one (Assumption 3). Concerning this standard assumption, we will add a brief discussion along the lines of our response above for completeness.

Thank you for your review and your positive assessment of our work. Please let us know if you have any further questions.

Question 4. (Dependence on the policy space size). We use Theorem 2 and Proposition 1 combined to obtain Corollary 2, i.e. we use Proposition 1 to choose $\epsilon_{MLE}$ appropriately by setting the upper bound in Prop. 1 to be equal to $\epsilon_{MLE}$. This gives our conditions on the number of samples to draw $n$ for occupancy measure estimation and the dependence on the dimension $d$ of the function approximation class. As we suppose realizability (Assumption 1 (ii)), $\epsilon_{MLE}$ only depends on the size of the parameter space $\Omega$ used in the function approximation class rather than the size of the policy space (i.e. dimension of the policy space parameter $\theta$). This dependence is clear via (a) the dimension $d$ of the function approximation space which appears in the final sample complexity and (b) an additional mild dependence captured by a constant inside the log terms in Proposition 1, namely the constant $B_{\omega}$, as the function approximation class ‘covers’ the set of occupancy measures induced by policies (see section D.2 for details).

Thanks again for your review. Please let us know if you have any further comments or questions. We will be happy to address them.

We thank the reviewer for their time and feedback on our work. We appreciate that the reviewer finds that our 'theoretical analysis is comprehensive' and that 'structure is logical, with ideas presented in a clear and coherent manner'. We reply to their comments and questions in what follows.

Comparison to Mutti et al. 2023. To complement our detailed comparison in appendix A p. 16 (along several dimensions including the problem formulation, the assumptions, the algorithm and the analysis) we provide a few additional points to answer the reviewer’s follow-up questions.

Advantages of our analysis/approach. Let us list first a few advantages/differences w.r.t. the aforementioned work:

1. Policy parameterization. We consider policy parameterization instead of tabular policies which results in a practical algorithm. We do require a strong assumption though (Assumption 5) to obtain our global optimality result.

2. No planning oracle access. We do not suppose access to a planning oracle able to solve any convex MDP efficiently. This is precisely the point of our optimization guarantees for our PG algorithm which updates policies incrementally. We discuss in particular how to set up the step size and other parameters to obtain our results. Nevertheless, we point out here that Mutti et al. 2023 address a slightly different (finite-trial) convex RL problem which is computationally intractable. Our problem coincides with their infinite trial variant of the problem.

3. No access required to feature vectors for function approximation. We learn occupancy measure approximations rather than supposing access to a set of features to approximate utilities (i.e. $F(\lambda)$ in our notations), i.e. we do not suppose access to a set of basis feature vectors for our approximation.

4. Model-free and no dependence on state action space size. We do not estimate the transition kernel, our algorithm is model-free as we only require access to sampled trajectories. In particular, we do not require to go through the entire state action space to estimate each entry of the transition kernel. More importantly, our performance bounds do not have dependence on the state action space sizes as their regret bound.

Advantages of our model-free PG algorithm vs their model-based algorithm. We inherit the usual advantages of model-free vs model-based algorithms. We list some of these:

1. Implementation. Model-free are often simpler to implement because they do not require learning or using a model of the environment. Our algorithm directly focuses on learning a policy (even if we do also approximate the occupancy measures but not the transitions themselves like in model-based approaches). We do not require to model dynamics explicitly and we require less assumptions on the environment’s dynamics.

2. Robustness. Inaccuracies in environment modeling propagate to policy optimization and can significantly degrade performance. Model-free methods directly learn policies from interaction with the environment.

3. Complex and high-dimensional environments. PG methods are particularly suitable for such settings. Estimating a model accurately in such settings can be difficult.

Of course, model-based methods also have advantages over model-free ones. For instance, model-free methods generally require more samples (note this is the case when comparing the result of Mutti et al. 2023 and ours) and model-based methods can also guide exploration.

Assumptions. Besides the points above, we provide a few comments regarding assumptions:

1. Some of our assumptions are quite similar. For instance, Mutti et al. 2023 assume linear realizability of the utility function $F$ with known feature vectors (Assumption 4, p. 17 therein) and then assume access to a regression problem solver with cross-entropy loss to approximate the utility function. We rather have a similar but different function approximation class regularity assumption (Assumption 1) and we suppose access to an optimizer which solves our log-likelihood loss maximization problem to approximate occupancy measures (see Eq. (8)). We both assume concavity of the utility function.
2. We require smoothness assumptions on the utility function (Assumption 3) whereas Mutti et al. only require Lipschitzness of the same function (Assumption 1 therein). This is because smoothness is important for deriving optimization guarantees as we make use of gradient information whereas Lipschitzness is enough for developing their statistical analysis.
3. Mutti et al. 2023 assume access to an optimal planner (as mentioned above), we do not need such a requirement as we provide optimization guarantees using our PG algorithm.
4. We need policy parametrization assumptions as previously discussed, Mutti et al. do not consider policy parametrization.

We hope this discussion further clarifies the comparison to Mutti et al. 2023 and our contributions.

Thank you for your review. As a first comment regarding the presentation that you find ‘poor’, we highlight that all other reviewers are unanimous on the fact that the presentation is good, the paper well-written and easy to read, we refer the reviewer to their numerous individual positive comments. Please let us know if you have any meaningful constructive comment to further improve our work and we will be happy to take it into account. We provide a detailed point-by-point response to your comments in what follows:

The proposed work learns state-action occupancy measures (Eq. 8). But it does not learn the occupancy measure defined in Eq. 1. In Eq 1, the occupancy measure is discounted by \gamma. However, Eq. 8 does not have a discount.

In contrast to what the reviewer is claiming, (Eq. 8) does allow us to learn the occupancy measure as we define it in (Eq. 1). Note that the samples used in (8) for the estimation are generated according to the discounted state action occupancy measure defined in (Eq. 1) which is indeed the discounted distribution. We invite the reviewer to Appendix D.1 (as we mentioned in l. 269 p. 5) to see our (standard) sampling procedure which allows us to generate such samples from the discounted distribution.

And we don't know if the learned occupancy measure satisfies the backward Bellman flow equation which is mentioned in Section 3.1.

• There is no reason why the approximated occupancy measure should satisfy the Bellman flow equation in general. For instance, to draw a parallel with value function approximation, the approximated value function when considering linear function approximation might only satisfy a projected Bellman equation when using TD learning with linear function approximation. There are also a number of different approaches in the literature to incorporate function approximation: minimize the squared Bellman error, the projected Bellman error, the TD error, … each one leading to an algorithm (see e.g. Sutton, Szepesvári and Maei (2009) and the thesis of the last aforementioned author). Similarly here, we do not necessarily require the approximation to satisfy the Bellman flow equation (although we use some of its characteristics for learning it of course, see our sampling procedure used for learning our maximum likelihood estimator).

• In addition, as we comment on in the paper, the occupancy measure does not satisfy a standard ‘forward’ Bellman equation which is immediately amenable to sampling as this is the case for standard value functions. To the best of our knowledge, there is no immediate way to address our online problem this way in a straightforward manner. This is one of the reasons why we follow a different approach.

• Even in the works that the reviewer mentions (for standard RL or CMDPs), there is no guarantee that occupancy measure approximations (in the function approximation setting) will satisfy the Bellman flow equations. While these works do incorporate them into the formulation to derive their algorithm in the tabular setting, considering neural network parametrization does not guarantee to preserve any such favorable property of the original problem formulation which becomes another problem formulation.

line 221: "Second and foremost, while prior work has used Monte Carlo estimates for this quantity, such count-based estimates are not tractable beyond small tabular settings." There are RL algorithms that focus on using stationary state-action distributions for learning policies that maximize expected returns with some constraint on the stationary distributions and these works do not use count-based estimates and works on continuous state-action space as well [1][2].

As a first preliminary comment, the references [1] and [2] that you mention do not address our general RLGU problem and do not provide any theoretical guarantee in terms of convergence nor sample complexity. [1] deals with a CMDP problem whereas [2] focuses on the standard expected return RL problem using diffusion models in offline RL settings. In both these cases, estimating occupancy measures is not needed (at least in the online setting which is our focus). There are of course a number of works in the literature to address standard RL and CMDP problems using PG methods for instance (in the online setting) with theoretical guarantees, including for larger state action spaces. We highlight that the main motivation of RLGU is to deal with objectives which are nonlinear function(al)s of the occupancy measure: standard RL and CMDP problems are essentially linear programs in the occupancy measures (if we do not think about complex practical policy parametrization) and widely studied while also being particular cases of RLGU. In contrast, RLGU calls for estimating occupancy measures (at least if a PG approach is followed like in our work) if one wants to tackle the RLGU problem without ad hoc approaches tailored to special instances of RLGU.

2. Our contributions are mainly: (a) proposing a simple model-free PG algorithm for RLGU replacing count-based occupancy measure estimators by scalable MLE which only use sampled trajectories, (b) showing that our approach is more suitable for the RLGU problem than the MSE approach proposed in prior work for the same problem (see our response to reviewer MseJ for clarifications) and (c) providing convergence and sample complexity guarantees for our algorithm which do not scale with the size of state action space. To the best of our knowledge, this is not covered by any of the works you mentioned (see our detailed discussion above) or any other work in the literature.

3. In general, function approximation errors in density estimation naturally depend on the class of probability distributions considered for the approximation. We mentioned Gaussian mixtures as an example in l. 254: ‘For continuous state spaces, practitioners can consider for instance Gaussian mixture models with means and covariance matrices encoded by trainable neural networks.’ For this specific example, one can try to tune the number of mixtures or even include this number in the learning parameters of the model. Of course, other parameterizations can be used. Note that the criticism you formulate can even be applied to policy parameterization. For continuous state action spaces, the most common policies used in RL in practice are Gaussians which are even unimodal and optimal policies need not be unimodal in general (note that Gaussian mixtures are at least multimodal).

About the papers you mention.

• Ma et al. 2023 consider transition occupancy matching, this is similar to the standard imitation learning problem and they consider a tabular problem formulation. We are rather proposing to approximate the occupancy measure itself using function approximation in the more general RLGU setting. The Bellman flow constraint you mentioned is the same in all the papers you mentioned and we have also discussed it in our paper.
• Yang et al., 2020 is about the same DICE approach that we have discussed in our response. The paper is focused on the policy evaluation problem in standard (expected return) RL (see their eq. (1)). They learn occupancy ratios. Similarly in Lee et al. 2021, for a similar problem, they learn occupancy ratios (a.k.a. stationary distribution corrections) as the title suggests. We have discussed this line of works and their limitations regarding our present motivation in our previous response.

Learning occupancy measure ratios (which is not our approach) has been used in some special instances of RLGU in the works you mention as we discussed it. This approach might not always be possible as the trick of considering the ratios and learning them only using offline data might not be possible because the problem reformulation using these ratios might fail when considering objectives beyond standard RL or learning from demonstrations (involving ratios). Furthermore, as we also mentioned, adopting such an approach also entails several additional difficulties to overcome from which we highlight again two main ones:

• One has to deal with an intractable number of constraints (even infinite uncountable in the continuous state action space setting) to incorporate Bellman flow constraints (supposing the Backward sampling issue can be handled), then using a Lagrangian approach introduces functional multipliers to handle this infinite dimensional problem which brings a number of technical difficulties. A possible approach would be to introduce function approximation at this stage but then an analysis would be required to see if any benefit can be concluded (how to introduce it? under which assumptions? …)
• Policy extraction from occupancy ratios in large and continuous state action settings as we previously discussed.

We propose a simple online approach to tackle our problem with theoretical guarantees. Of course, we do not exclude that there might be other relevant approaches to tackle our problem and improve over our results.

Dear AC,

Thank you for your email, and I apologize for the delay. We are going to post the rebuttal soon.

Regards,

Authors