Understanding End-to-End Model-Based Reinforcement Learning Methods as Implicit Parameterization

Abstract: Estimating the per-state expected cumulative rewards is a critical aspect of reinforcement learning approaches, however the experience is obtained, but standard deep neural-network function-approximation methods are often inefficient in this setting. An alternative approach, exemplified by value iteration networks, is to learn transition and reward models of a latent Markov decision process whose value predictions fit the data. This approach has been shown empirically to converge faster to a more robust solution in many cases, but there has been little theoretical study of this phenomenon. In this paper, we explore such implicit representations of value functions via theory and focused experimentation. We prove that, for a linear parametrization, gradient descent converges to global optima despite non-linearity and non-convexity introduced by the implicit representation. Furthermore, we derive convergence rates for both cases which allow us to identify conditions under which stochastic gradient descent (SGD) with this implicit representation converges substantially faster than its explicit counterpart. Finally, we provide empirical results in some simple domains that illustrate the theoretical findings.


Reviews and author responses:

 Thank you for the concise and specific response to the concerns raised in the review. I have updated my score on the basis of this response. I would still strongly urge the authors to devote some space to a discussion of the possible drawbacks of the simplifying assumptions (as per the rebuttal), and clarify the specific variant of model-based RL that is being considered in the paper.

The paper presents a theoretical analysis of implicit parametrization of value functions in Markov reward chains. Results analyzing the gradient dynamics, the convergence rates of stochastic gradient descent, and the relative trade-offs of implicit and explicit parameterization are presented with suitable assumptions. Experimental results are presented to illustrate how ideas from the theoretical results might be used, even in domains that do not satisfy all the required assumptions.  It is difficult to understand the significance and novelty of the theoretical analysis presented in the paper. Some areas of concern are discussed below.

Singleton Action. The analysis is only presented for the case of an MDP with a singleton action. As the authors point out, this means the theoretical analysis is only valid for Markov reward chains, not Markov Decision Processes. It is stated in the Conclusion that the results might generalize to full MDPs, but there does not seem to be any evidence for this assertion. In the absence of this generalization, it is confusing to frame the analysis in terms of understanding RL methods. Being restricted to a singleton action is far too stringent an assumption for the study of RL methods.

Supervised Learning setting. The authors should clarify how the assumption of working in the supervised learning setting impacts the theoretical analysis, and what results one might expect in the full reinforcement learning case. In particular, it seems that convergence to the global optimum would not hold in the absence of a fixed set of datapoints (since convergence on each set individually is not sufficient for convergence globally). 

Model-Based RL. The analysis is presented in a formulation that closely resembles that of value iteration networks (VIN) (Line 88). However, as the VIN paper points out “[This] approach is different from model-based RL, which requires system identification to map the observations to a dynamics model, which is then solved for a policy.” Considering this difference, it would be helpful if the authors would clearly comment on how this formulation is related and differentiated from typical model-based RL. Presenting the analysis for end-to-end model based RL, while only considering the VIN formulation is likely to create confusion. 

Novelty. It would help if the authors would clearly delineate the novel technical innovations in the theoretical analysis. The authors seem to state that results for the explicit parametrization case are not novel (Line 189). Is this the case for all the results regarding explicit parametrization? It would also be helpful to highlight the key steps and theorems that do not simply follow from standard results. The paper would greatly benefit from a more detailed discussion of the fundamental limitations of the proposed method. In particular, the two main limitations that should be discussed are the assumption of a singleton action, and the dependence on the supervised learning setting.

 It would help us improve our paper if the reviewer would let us know if we have not fully addressed their concerns. If most or some of the major concerns were addressed, we would also greatly appreciate if the reviewer adjusted their review to reflect this.

 Thank you for thoroughly answering my questions and comments! Having read the replies from the authors and other reviewers, I'm inclined to maintain my score.

This paper proposes a theoretical analysis of implicit-model value learning in Markov Reward Processes. The paper compares two linear parametrizations of the value function, the usual explicit one and a generic implicit parametrization. The latter is expressed as any differentiable function of a reward vector and transition matrix. 
The central theorem of this work shows that this parameterization induces a preconditoning of SGD in the supervised/Monte Carlo return regression setting. The usefulness of this preconditioning is then analyzed, and is strongly linked to properties of the residuals, for example their correlation. In the right setting, this suggests a much faster convergence.
The theoretical claims are then backed by empirical results on simple exemplar domains.
  While this analysis is fairly disconnected from existing empirical implicit-model methods, what it suggests has very interesting implications. The disconnect to me is the MDP->MRP simplification, and the control->evaluation->supervised simplification; both are fair simplifications for theoretical analysis, but they still reduce its predictive power. 
The interesting implications I see is that this intuitively translates very well to gradient dynamics in non-linear (deep) models: if the partial derivatives between the implicit model predictions and the residual are well conditioned, we should expect the same effect regardless of how complicated that implicit model is.

I am unfortunately not comfortable enough with the theory here to spot non-obvious mistakes, but I did not spot any superficial mistake. I am also only superficially familiar with the implicit-model literature, but as far as I know analysing particular RL methods as preconditioners has mostly only been done in the options literature; this seems like a new take.

The paper was fairly easy to read, and the ideas here are well communicated. The empirical results link the theoretical claims to empirical claims, although very basic. I think that, without starting to solve Atari, more could have been done to illustrate the implications of the theorem. For example, one could repeat these experiments in the control case, or in the temporal difference case (while we shouldn't expect the theory to exactly hold, we may still observe effects which, in the long run, we would like to see in complicated environments).

I usually write a long list of comments and nitpicking in my reviews, but this paper was fairly straightforward and I don't have much to add here.

Rebuttal update: I've read the other reviews and rebuttals and I'm inclined to keep my current score. I think this is a good contribution. The authors correctly address the fundamental limitations of their work.

There are no mentions of potential negative societal impacts, but this is fairly abstract work. As an exercise, the authors may want to situate the impact of RL on society.

This work introduces learning models under a different parameterization that the authors call implicit parameterization, where instead of training value estimation parameters directly, it is learned implicitly through reward and dynamics model estimation. But then the evolution of the implied value estimation parameters is analyzed and compared with that of the direct parameterization. With theoretical and empirical analysis, cases were established where implicit parameterization is more beneficial.
  The paper is well written, concepts are carefully and clearly exposed. My main issue is that more theoretical results could be empirically verified such as the superiority of the variance to expectation ratio for updates. Was it tried and not found conclusively?

Theorem 1 is a bit unclear. The two updates in (4) that are claimed to hold, among them the first one is by choice as mentioned earlier in line 132 and it is not clear what the optimization procedure involved is for the implicit parameterization. The only knowledge provided is in line 133 “implicitly parameterized linear weights obtained by optimizing L.” But optimizing using what? SGD? Least squares? Depending on the choice of the optimizer, the evolution may follow different paths. Also, the theorem statement should explicitly contain this information instead of relying on text in the prior paragraph.

Line 87: it will be better to replace (w, F) with psi = (w, F).

**** post rebuttal ****

Thanks for the clarifications. N/A

 Thank you for the update. I think considering the setting you address, it's good to have these experiments, and more generally I think, to make sure the conceptual gap (between the simple settings explored in the paper and large scale MDPs/model-based agents) is bridged clearly. In my understanding, as I've mentioned in my review, this _conceptually_ applies easily to complicated models, but perhaps being even more explicit about why would be good for readers.

 Thank you for the thoughtful feedback and corrections. We individually address each point bellow.

---
> A limitation of the results proved in the paper is that the questions it answers are all concerned with the optimization of the ERM loss, so it tells us essentially nothing about the generalization properties of the solutions found by the two parameterizations, even though the generalization abilities of end-to-end methods are the stated motivation for studying them.

The primary motivation for this work is the lack of theoretical understanding of end-to-end MBRL methods. We mention generalization when giving an overview of the kind of settings for which these types of methods have been proposed. By considering a simplified setting, our work identifies properties of the gradient dynamics of these methods inherent in the parameterization they induce. Although it most certainly would have been interesting to see results for some of these more complex settings, we believe our results serve as a useful first step towards better understanding these types of methods.

> If my understanding is correct, the paper’s result about (S)GD finding global optima tells us that the asymptotic solutions of the two parameterizations should be exactly the same if the global minimum of L(theta) is unique (which it will if the feature matrix X  has full rank). The corresponding (w, F) may not be unique, but any optimal solution would still have to induce the same theta.

Correct! This is an accurate summary of some of our results.

> In experiments, it is interesting that the optimal eta can be different from the true gamma that generates the labels. However, no insight regarding how to set eta is provided.

Our results provide some insight into the effect of eta on the gradient dynamics and overall behavior, but it is correct that it would be non-trivial to use those results to understand how best to set eta. It is not clear that this question would have a general answer using only the assumptions we’ve made. We agree that this is an interesting observation. Prior work has made similar (theoretical) observations under a more conventional model-based setting where similar “mismatched” discount factors are analyzed [1]. However, it is unclear how applicable that line of work is given it primarily focuses on planning.

[1] Jiang, Nan, et al. "The dependence of effective planning horizon on model accuracy." Proceedings of the 2015 International Conference on Autonomous Agents and Multiagent Systems. 2015.

> In the experiments, do you stop running SGD before convergence? As mentioned above, my understanding is that the two parameterizations should converge to the same value, even if one converges more quickly than the other. However in the results there is a clear difference between the two.

Your understanding about convergence is correct, and we are sorry not to have been clear about the fact that we run each method for a fixed number of updates.  The difference in performance in a fixed amount of time reflects, we believe, our theoretical results concerning the convergence rates of gradient descent under the different parameterization.

> In (48) and (49), what is U? It seems that it should be A?

Yes, this should be A. We have fixed it in the revised manuscript. 

> In places you write “if O(alpha^2_t) <= alpha^2_t C” for some C>= 0. Is this not always true by the definition of the O notation?

We use the usual definition of the O notation. For example, in (4), as usual, the $\mathcal{O}$ notation means that the absolute value of the error $\hat \theta_{(t+1)} - (\hat \theta_{(t)} -\alpha_{(t)} P_{(t)}\nabla L( \hat \theta_{(t)}))$ is at most some constant times $|\alpha_{(t)}^2|$ when $\alpha_{(t)}$ is close enough to zero. To clarify this point, we revised the paper by replacing “if O(alpha^2_t) <= alpha^2_t C” with the following: Let the $\mathcal{O}(\alpha_{(t)}^2)$ term be upper bounded by $C \alpha_{(t)}^2$ for some $C\ge0$. So, this means that $\hat \theta_{(t+1)} - (\hat \theta_{(t)} -\alpha_{(t)} P_{(t)}\nabla L( \hat \theta_{(t)})) \le C \alpha^2_t$.

> The statement of Corollary 1 says that the learning rates alpha_t are given, but go from (7) to (8) you specify a value for alpha_t. Please justify.

By this, we did not mean that the learning rates alpha_t are given by external environments. What we meant by this is the following: given that the learning rates alpha_t are sufficiently small, the statement holds. Or equivalently, provided that the learning rates alpha_t are sufficiently small, the statement holds. In the proof, we are providing examples for the specific values of sufficiently small alphas that work for the convergence. To clarify this point, we revised the paper with a more concrete statement by replacing “given a sufficiently small learning rate $\alpha_{(t)}>0$” by “there exists $\bar C > 0$ such that with $\alpha_{(t)}=\bar C\|\nabla L(\hat{\theta}_{(t)})\|_2^2$”, which corresponds to the example constructed in the proof. 

> Equation (87) seems to have a typo (extra phi).

Yes, that’s a typo, which we have fixed by deleting the extra phi

> In (106), the index i is reused.

This is a typo. We fixed this by deleting the extra summation over the index i.

> In the second-to-last step in the proof of Corollary 4, r_max should be squared since it is bounding r_i^2. (This seems to be just a typo, as the final result has r_max^2).

Correct, this is a typo. We have fixed this by replacing the r_max by the r_max^2 there. 


 We would like to first thank the reviewers for their thoughtful feedback, comments, and corrections. We've opted to answer each individually but wanted to highlight here some changes to come in the final version of the paper.

* We will add to the appendix similar figures to the existing ones showing results for the case when using a 4 action MDP (instead of a Markov reward chain). The results look qualitatively similar and suggest that our results might generalize to the MDP case.
* We will add figures showing the mean-variance ratio of the updates in the appendix. Currently we've ran these for the chain MDP and four room domain. We found that they agree with our theory but the effect is weak due to the reward offset having only a minor effect on the ratio $(r_{max}/r_{min})^2$. We will add figures for all 3 domains in the appendix of the final paper.

 > While this analysis is fairly disconnected from existing empirical implicit-model methods, what it suggests has very interesting implications. The disconnect to me is the MDP->MRP simplification, and the control->evaluation->supervised simplification; both are fair simplifications for theoretical analysis, but they still reduce its predictive power. The interesting implications I see is that this intuitively translates very well to gradient dynamics in non-linear (deep) models: if the partial derivatives between the implicit model predictions and the residual are well conditioned, we should expect the same effect regardless of how complicated that implicit model is.

We think this is a fair assessment of the value and limitations of our contributions. In order to address concerns from other reviewers, we’ve run our experiments using 4 action MDPs. The results are qualitatively similar which suggests that our results hold in the more complex setting, at least in the domains we considered. We will add these results to the appendix in the final version of the paper.


 Thank you for your questions and comments. We will answer them individually below.

---

> My main issue is that more theoretical results could be empirically verified such as the superiority of the variance to expectation ratio for updates. Was it tried and not found conclusively?

We didn’t think to try this since the experiments were designed to manipulate the components of the residual along the [1 1 … 1]^T direction and not necessarily the ratio $(r_{max}/r_{min})^2$. However, to address the reviewer’s concern, we reran our experiments and examined this ratio and the mean-var update ratio on the chain MDP and four rooms. We found that the reward offset has a weak effect on the ratio $(r_{max}/r_{min})^2$. As expected, the mean-var ratio under the implicit parameterization appears correlated with the ratio of residuals while under the explicit parameterization the mean-var ratio is unchanged, i.e., a flat curve.

These new results require computing many more gradients and are notably slower to run as currently implemented which is why the mountain car domain was not included (we ran into memory issues). We will fix the memory issues and add these figures, including mountain car, in the appendix of the final paper.

> Theorem 1 is a bit unclear. The two updates in (4) that are claimed to hold, among them the first one is by choice as mentioned earlier in line 132 and it is not clear what the optimization procedure involved is for the implicit parameterization. The only knowledge provided is in line 133 “implicitly parameterized linear weights obtained by optimizing L.” But optimizing using what? SGD? Least squares? Depending on the choice of the optimizer, the evolution may follow different paths. Also, the theorem statement should explicitly contain this information instead of relying on text in the prior paragraph.

Both parameterizations are using (S)GD on the loss defined in (3) for the GD case and the loss defined above Corollary 2 for the SGD case. We will make sure this is stated more clearly.


 Thank you for the thoughtful feedback. We will attempt to address each concern individually below.

---

> Singleton Action. The analysis is only presented for the case of an MDP with a singleton action. As the authors point out, this means the theoretical analysis is only valid for Markov reward chains, not Markov Decision Processes. It is stated in the Conclusion that the results might generalize to full MDPs, but there does not seem to be any evidence for this assertion. In the absence of this generalization, it is confusing to frame the analysis in terms of understanding RL methods. Being restricted to a singleton action is far too stringent an assumption for the study of RL methods.

We want to emphasize that when considering a singleton action for the implicit parameterization, the resulting estimator can still be used to learn value functions on MDPs. As for generalizing our results to parameterizations using regular MDPs, most of our derivations involving a single step apply to the MDP case since gradients will correspond to the gradient of the Markov reward chain (MRC) under the optimal policy (assuming no ties and an argmax policy) where the transition matrix of the induced MRC will simply be a selection of the rows of the transition for each corresponding optimal state action pair. For this reason, we expect many of the properties of the dynamics of the gradient that are present in the MRC case will also be present in the MDP case.

However, the full MDP case makes it harder to make general statements about convergence and adds significant complexity to everything else. Since, as far as we know, there doesn’t exist any published theoretical analysis of end-to-end model-based RL methods, we opted to focus on a more tractable setting in order to first lay the groundwork. That being said, for the reason we’ve mentioned above, we are hopeful we will be able to generalize these results in the future but as the reviewer points out, this isn’t necessarily trivial.

In an effort to address this concern, we also ran our experiments with 4 actions as a sanity check and are happy to report that the results look qualitatively very similar to the singleton action case. We will add these results in the appendix.

> Model-Based RL. The analysis is presented in a formulation that closely resembles that of value iteration networks (VIN) (Line 88). However, as the VIN paper points out “[This] approach is different from model-based RL, which requires system identification to map the observations to a dynamics model, which is then solved for a policy.” Considering this difference, it would be helpful if the authors would clearly comment on how this formulation is related and differentiated from typical model-based RL. Presenting the analysis for end-to-end model based RL, while only considering the VIN formulation is likely to create confusion.

This statement from the VIN paper seems to refer to “classical” model-based RL. In the classical setting, model parameters are explicitly fitted to match observed transitions. We (and the VIN authors) consider methods which fit models based on the value predictions generated. We propose to examine these types of methods as an implicit parameterization to allow for the analysis we present . Our theoretical contributions are strictly for this end-to-end model-based setting and do not apply to the classical setting (which would have a completely different objective function.)

> Novelty. It would help if the authors would clearly delineate the novel technical innovations in the theoretical analysis. The authors seem to state that results for the explicit parametrization case are not novel (Line 189). Is this the case for all the results regarding explicit parametrization? It would also be helpful to highlight the key steps and theorems that do not simply follow from standard results.

We don’t consider any of the results for the explicit case to be noteworthy contributions. Although some might be technically novel given the specific form and setting we consider, we consider them important only as comparisons. All results for the implicit parameterizations are novel, as far as we know, and represent our main contributions. We will emphasize this in the next draft.


This paper studies a simplified abstract model of end-to-end model-based reinforcement learning (RL). An implicit model is defined as the solution $\hat\theta$ to $\hat\theta = w + \eta \sigma(F)\hat\theta$, where $w$ models the rewards for all states, $\eta$ is a discount factor, and $\sigma(F)$ is the transition matrix of a Markov chain ($\sigma$ is softmax). The paper studies the optimization of least-squares linear regression where the linear weights are $\theta$, computed as an implicit function of $(w, F)$. The results are contrasted with the corresponding explicit parameterization, where the linear weights are themselves parameters.

The paper shows that the optimization landscape when using an implicit parameterization is still tractable, and can even be beneficial in some scenarios.
More specifically:
* Despite non-convexity in $(w, F)$, gradient descent with respect to these variables converges to a global minimum under mild assumptions. Similarly, SGD converges to a global minimum in expectation.
* Expected convergence rates are derived for SGD under both parameterizations
* Expected converges rates for a particular direction of the residual are derived for SGD under the implicit parameterization. This is used to analyze when the implicit parameterization is beneficial. In particular, convergence is fast along the [1, …, 1] component of the residual, meaning that the end-to-end method will quickly correct if universally under- or over-estimating.
* The variance/mean ratio of the parameter changes is derived for both parameterizations. These are used to analyze when the implicit parameterization is beneficial. In particular, the end-to-end method exhibits low variance when the residuals are “balanced” ($r_{max}$ is not much bigger than $r_{min}$).

There are also some experiments designed to corroborate some of the theoretical results. For each task, data is generated using an $\epsilon$-greedy policy, and then both parameterizations are fit to the empirical returns. The conclusions can be summarized as follows:
* Shifting the reward by a constant affects the implicit parameterization less than the explicit parameterization.
* The best value for $\eta$ (the discount factor in the implicit MDP) is not necessarily the same as the true discount factor of the MDP.  Most of the theoretical results are novel to my knowledge (some were already known but are included for comparison), and I am glad to see some progress being made on the theoretical analysis of end-to-end methods, which has so far been limited to empirical study. The results in this paper provide some insight as to why one parameterization may be easier to optimize than the other, which is useful to practitioners.

A limitation of the results proved in the paper is that the questions it answers are all concerned with the optimization of the ERM loss, so it tells us essentially nothing about the generalization properties of the solutions found by the two parameterizations, even though the generalization abilities of end-to-end methods are the stated motivation for studying them. 

If my understanding is correct, the paper’s result about (S)GD finding global optima tells us that the asymptotic solutions of the two parameterizations should be exactly the same if the global minimum of $L(\theta)$ is unique (which it will if the feature matrix $X$ has full rank). The corresponding $(w, F)$ may not be unique, but any optimal solution would still have to induce the same $\hat\theta$.

That said, some of the analysis performed in this paper (e.g., Theorem 1) may be useful for analyzing the generalization as well, even though this paper does not itself present results regarding generalization.

In experiments, it is interesting that the optimal $\eta$ can be different from the true $\gamma$ that generates the labels. However, no insight regarding how to set $\eta$ is provided.

The writing of the paper is clear.

I have a few miscellaneous comments/questions:
* In the experiments, do you stop running SGD before convergence? As mentioned above, my understanding is that the two parameterizations should converge to the same value, even if one converges more quickly than the other. However in the results there is a clear difference between the two.
* In (48) and (49), what is $U$? It seems that it should be $A$?
* In places you write “If $\mathcal{O}(\alpha_{(t)}^2) \leq \alpha_{(t)}^2 C$” for some $C \ge 0$. Is this not always true by the definition of the $\mathcal{O}(\cdot)$ notation? (Perhaps for clarify you can write exactly what is meant by $\mathcal{O}(\cdot)$.)
* The statement of Corollary 1 says that the learning rates $\alpha_{(t)}$ are given, but to go from (7) to (8) you specify a value for $\alpha_{(t)}$. Please justify.
* Equation (87) seems to have a typo (extra $\phi$)
* In (106), the index $i$ is reused.
* In the second-to-last step in the proof of Corollary 4, $r_{max}$ should be squared since it is bounding $r_i^2$. (This seems to be just a typo, as the final result has $r_{max}^2$.) Mostly; I have included additional comments about the limitations above

Meta-review: Initially some of the reviewers were concerned with the limiting assumptions used. However, after the author response, all of the reviewers agree that the theoretical results while limited in some ways are interesting and significant enough for acceptance. As reviewer 6XLT points out, adding a balanced discussion of the additional complexities of the full MDP case and the additional experiments will strengthen the paper significantly.

Decision: accept