Understanding End-to-End Model-Based Reinforcement Learning Methods as Implicit Parameterization
Authors: Clement Gehring
Kenji Kawaguchi
Jiaoyang Huang
Leslie Pack Kaelbling
Venue: NeurIPS 2021 Poster (2021)
Keywords: model-based
reinforcement learning
end-to-end
implicit

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.

Understanding End-to-End Model-Based Reinforcement Learning Methods as Implicit Parameterization
Clement Gehring
Electrical Engineering and Computer Sciences Massachusetts Institute of Technology
cl clement@gebring.22
Kenji Kawaguchi
Center of Mathematical Sciences and Applications
Harvard University
harvard.edu
Jiaoyang Huang
Leslie Pack Kaelbling
Courant Institute of Mathematical Sciences New York University
Electrical Engineering and Computer Sciences Massachusetts Institute of Technology
jh4427@nyu edu edu
1pk@csail m.t. edu
Abstract
Estimating the per-state expected cumulative rewards is a critical aspect of rein- forcement learning approaches, however the experience is obtained, but standard deep neural-network dunctionapproxximatio methods are often inefficient in this setting. An alternative approach, exemplified by value iteration networks, is learn transition and reward models of a latent Markov decision process whose value predictions fit the data. This approach has been shown empirically converge faster to a more robust solution in many cases, but there has been little theoretical study of this phenomenon. 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 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.
1 Introduction
Whether done by averaging over trajectories or some intricate fitting procedure, estimating the per-state expected cumulative rewards is critical aspect of every reinforcement learning (RL) approach. Efficiently learning these values in large complex domains still remains a challenge despite advances in deep learning. One contributing factor comes from the difficulty of scaling model-based methods which model the transition probabilities and rewards directly. While known to be sample efficient, these methods have failed to fully leverage recent advances in deep learning, forcing the use of less efficient but more scalable model-frea methods which try to learn the values directly. Even when achieving low one-step prediction error, models learned with large highly expressive
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

estimators struggle to produce good policies or accurate value estimates in part due to mismatch in the objective [T]] or compounding errors [24 22 To address this, several recent methods propose to directly optimize predictions derived from the model and leverage modern automatic differentiation frameworks to propagate errors back to the model parameters [33 18 15 19 []. However, despite
promising results shown by these methods, often referred to as end-to-end model-based methods, there is little formal justification for this idea beyond intuitive arguments.
In many cases, this approach is motivated by the goal of varying the model parameters based on the current state or observation. Empirical results suggest that such end-to-ena model-based methods are capable of better generalization when observations can be used to formulate smaller subproblems, for example, by providing a map of the world [23], a description of goal [19], or the high-level
action being followed [15]. However, as we will show later, even without input dependent models, this indirect, or implicit, representation of values drastically alters the dynamics of gradient descent and, thus, the inductive bias, in interesting ways.
To better understand how these end-to-end model-based methods compare to more direct methods, we formulate these ideas as different parameterization of the same function class. Specifically, by formulating an end-to-end model-based approach as an implicit parameterization of a linear function, we can directly quantify differences between this implicit approach to the more common trivial explicit parameterization where linear weights are stored and learned directly.
Concurrent to our work was a recent paper [4 which focuses on implicit differentiation applied to end-to-end model based methods. This work compares this approach to model-based methods which maximize likelihood. Although we focus on the comparison between end-to-end methods and other direct methods, we view this related work highly relevant and complementary with ours. Both work and theirs aim to better understand the properties of end-to-end model-based methods.
Our contributions are primarily theoretical and aim to provide theoretical account of the performance of end-to-end model-based methods. To help in this matter, we also provide some empirical results in simple illustrative problems which serve to demonstrate properties derived from analysis. We start by showing that although the implicit formulation defines a non- linear and non-conves optimization problem, the implicit linear weights still converge global optimum. We continue by showing that stochastic gradient descent with the implicit parameterization has appealing properties when the residuals exhibit certain properties. Specifically, we derive the expected changes in the loss and show a surprisingly fast convergence rate along a specific component of the residual. We then use these results to discuss conditions under which end-to-end methods might be preferable. We conclude theoretical contributions by quantifying the variance-mean ratio updates under both parameterization and establishing conditions under which the implicit parameterization exhibits significantly better properties than its explicit counterpart. Finally, we conclude with some empirical results.
2 Framework
While our theoretical results consider supervised learning setting, the motivation for this work comes from the policy evaluation problem. In this setting, we assume there an MDP, M = (S, ra. T, }- with state space S whose dynamics for each action E A are described by the transition probabilities Ta generating rewards as described by ra- We are interested in learning the discounted value function for a given policy æ which is defined as the per-state expected cumulative discounted rewards. Formally, given a discount factor y [0. 1) and state s E S,
V V*(()) = AA
Σ 7t t=
Ray
So = = As (Ro + +V* So = =}]
(1)
where the second equality leverages the Markov assumption to define the values recursively.
2.1 Linear function approximation
We investigate the differences between parameterizations in the context of linear functions of the form ((s)- = (ss) 8, where o(s) E R is some fixed, possibly non- non-linear, encoding of the state s E S and 0 E R represents the linear coefficients. This type of function approximator is common in reinforcement learning and can be seen generalization of the discrete tabular setting, i.c., d(s)

---

corresponds to a one-hot encoding, and has yielded insight in the past, for example, by highlighting temporal difference learning convergence issues [ or by relating Bellman error to model error [T6] Additionally, in settings where domain knowledge can be used to generate useful features, these linear approximations can be significantly more efficient to learn and evaluate, making them not only theoretically interesting, but often practical as well.
To allow us to compare different parameterizations, we introduce a reparameterization function g which maps parameters vb E V to linear feature weights 0 E θ, i.e., 0 := g(()) and define linear estimators of the form:
() - +()) ( (s))
Note that while Vg is always linear in ((s) and 0, it is not necessarily linear with respect to the parameters v. For this reason, the dynamics of gradient descent can differ significantly across different parameterization despite all defining equally expressive function classes (if g spans R.).
2.2 Implicit parameterization
In order to explore the properties of end-to-end model-based methods, we consider a simple instance of such methods whose model corresponds to Markov reward chain, i.e., an MDP with a singleton action. We consider the case where the rewards are encoded as a vector w and the transition probabilities are kept as matrix of unnormalized bog-probabilities, F, such that o (F) corresponds to row-stochastic matrix, i.e., a applies the softmax function on each row. With the probabilities
encoded in this way, we can treat F as an unconstrained matrix which simplifies both the optimization
problem and our analysis. Finally, we define the linear weights 0 == g(ww, F)) to be the state values of the Markov reward chain defined by (w, F).
We note that this formulation closely resembles that of value iteration networks (VIN) 23 with few differences. Although VIN allows the model parameters to conditioned on the state, we wish to better understand the effect of representing values implicitly, as model parameters, without any conflating effect introduced by using state-dependent models. Note that any insight in this simpler framework is likely relevant to the state-dependent model setting since the gradients differ only by the Jacobian relating model parameters' differentials to learned parameters' differentials, i.e., the chain rule. Additionally, while VIN consider an internal action space similar to the action space of the true underlying MDP, we limit ourselves to the a singleton action in order to keep the analysis tractable. We provide a brief explanation for why insight in the singleton case is relevant to the general case and provide some supporting empirical results in the appendix.
We assume the linear feature weights, Ó, and the model parameters, th, satisfy some constraint h(0; (6) = 0 which can be used to implicitly define the function, g. which maps model parameters to feature weights, i.e., ((v)= = Ô. This implicit formulation allows us to analyse how & and relate without needing to consider the algorithm used to compute ô. For the Markov reward chain case previously described, we can use 4 to define the constraints as
h(,;w, F) w + (a(F) 0,
where 7) E [0, 1) is the internal discount factor for the modelled Markov reward chain. Note, that this discount factor 1 is different than the discount factor y which defines the true discounted values. We intentionally use different symbol to avoid any confusion later on. Also, for clarity, we allow ourselves to abuse the notation slightly by unpacking the parameters &b into a vector, w, and matrix, F, analogous to the rewards and the unnormalized transition log- pogo-abbbilliies, respectively.
2.3 Implicit differentiation
Under appropriate assumptions of continuity, an implicit parameterization can be differentiated just like any other explicit parameterization through the use of the implicit function theorem. Formally, for an implicit parameterization g, the derivatives can be defined solely through the derivatives of the constraint, h(@; w) = 0. That is, the constraint of 0 = h(0; Φb) = h(g ()); 1)) (where e = g(b))
implies that (
8h 8y + ah
(by the chain rule), which can be rearranged into
59 de
0g
0h
0h
30
O
3

---

Note that this results in differentiation rules that are independent of how the constraints were solved. In contrast, many methods like VIN will differentiate through the solver's computation, e-g., several iterations of value iteration. This approach is often only asymptotically equivalent to the implicit differentiation formulation and will otherwise alter the gradient dynamics.
3 Dynamics of gradient descent
We start by exploring the dynamics of gradient descent under the implicit parameterization induced by (2). We show that the gradient descent updates under the implicit parameterization are equivalent to preconditioned updates under the explicit parameterization. Furthermore, we show that despite no longer being a convex optimization problem, the gradient dynamics of the implicit parameterized objective function converges to a global optimum.
To keep the notation clear, we reserve 0 to represent parameters of the trivial explicit parameterization in which the reparameterization function the identity, Jemplicit (z) = x, and represent the linear weights under an implicit parameterization, e.g., Iimplicit ( (w, F) = 0 where is the solution to the constraint (2). For the remainder of this paper, we will refer to these two parameterizations as the explicit and implicit parameterization, respectively.
We assume we are given a set of n datapoints D = ((zi, : (1,..., m}} and loss L defined on the linear weights such that
12 L(0) E(o(x)) - y1)² 1 12 2 Σ (0)
(3)
where rf) is the residual of the predicted value of x; using the linear weights @ and @ is a fixed function mapping data points, Wi, feature vectors. Additionally, we assume fixed sequence of learning rates (ox is used. Let the sequence (0(z) represent the linear weights obtained by using gradient descent to optimize I. under the explicit parameterization with initial weights (00) and, similarly, let ( be the sequence of implicitly parameterized linear weights obtained by optimizing [[gampticits (W(e) F(o))) with corresponding parameters (wr()) and (F()) and initial parameters (w) F. Following previous works on gradient dynamics of deep learning models 17 5 8 we also consider the corresponding gradient dynamics d
at (()) and d dt (1) under
the gradient flow (where (00 is induced by the dynamics of w(()) and F(z)).
We start our analysis by presenting our results describing the gradient dynamics for both parameteri- zations. This will allow to compare the two and start building better understanding of the effects the implicit parameterization has on the gradient dynamics. First, we define A(e) = I (F)e)) and note that because or (F())) always forms a stochastic matrix, we can explicitly write out the implicitly parameterized linear weight, (()) = A (((e), enabling us to show the following theorem.
Theorem 1. Let P(y) = -A7) (17 D(o) + +) (AT))) where Duy = diag(d()) and [d()]x =
(o(Feeskss 03) then Amin (P()) 2 c for some time- time-independent constant 0, and Σ. ES (E, o(Pcolieg the following holds:
d = di
dt -
and
(4)
We provide the proofs for all our results in the supplementary material. From this result, we see that the updates to the linear weights under the implicit parameterization resemble that of preconditioned gradient descent under the explicit parameterization where P(t) acts as a preconditioner which transforms the gradient before updates. Ideally, the preconditioner would be chosen such that it accelerates convergence for some class of problems, for example, by efficiently approximating the inverse Hessian thus resulting in updates approximating Netwon's method. This potential to greatly accelerate convergence under the right circumstances is why our analysis will revolve around understanding the properties of P(1)-
We begin by noting that Ptty 's positive definiteness implies that the weights of the implicit parameter- ization converge to global minimum from any initialization.

---

a Corollary 1. For any initialization B/o-- the gradient dynamics a (1 converges to a global minimum.
is important to note that the loss landscape under the implicit parameterization is generally non- linear and non-convex, As consequence, convergence results applicable to linear regression cannot be trivially applied to our setting; additional properties of the optimization dynamics are required, such as those presented in Theorem [ convex optimization for reinforcement learning, whereas related works studied for deep learning 90, 26].
Thus, we proved non-trivial global convergence in non-
3.1 Properties of stochastic gradient descent
We now turn our focus towards stochastic gradient descent (SGD), e.g., gradient descent on sampled of training data, and present some additional properties of both parameterizations considered. In this section, we first define some additional notation and update rules to accommodate the SGD setting. We then quantify the expected improvement of the loss and show that both cases have comparable convergence rates.
We start by defining a sequence of random averaged loss functions L(o), I())..., corresponding to the loss L defined (3) evaluated on random subset of the data D. Formally,
(0) (1)
Σ
(b(x) -- v)
(2.))ED(O
where Dio) is random subset of the data D of size k drawn without replacement. In this setting, gradients are taken just as in gradient descent but using Lu) instead of the full loss.
Before presenting our results quantifying the dynamics of SGD with the explicit and implicit parameterization, we need the following assumptions on the feature matrix.
Assumption 1. Using the notation X S Y to mean X < eY for some constant C 0, we assume the features ((x)) satisfy
[0(x,)l12 2 f ss
Let X = ((((=1),((2)) ... ((xn)}, the sample covariance matrix XT X has bounded norm |X X||2 = O11, and ||X 1 Z n.
Theorem 2. Given a sequence of learning rates (a<()) t. for SGD with the explicit parameterization. the expected decrease of the loss function satisfies
E[(L(0++1) - L(( = a((t)
|r(0()))12 o(a2)) - a(t)) L((()) + (dta): n
where r(0) = r,
If a, 0(1,, in expectation, the convergence rate for the SGD with explicit parametrization is (11nn) which makes it slow, taking O(n) steps to converge. Note that this result isn't 't novel and follows from standard results of convex optimization 33. but is stated here to facilitate comparisons. For the implicit parameterization, we recall from Theorem (and (and Lemma ??) that the GD dynamics
given explicitly in terms of the matrix P(g) We state formally similar result for the expected SGD dynamics:
Theorem 3. Given a sequence of learning rate (at))). SGD with the implicit parameterization, the expected decrease of the loss function satisfies
B[(L((++1) ( I((,))))00)) -0(()
((t)
(Pe))((()) - -((ag))
n
where r(0) = r.))d())
3.2 Conditions favoring implicit parameterization
An important takeaway of our results so far is that the matrix Pto) is critical in understanding the effects of using such an implicit parameterization, and, more importantly, understanding when it would be expected outperform its explicit counterpart. For this reason, we first examine some
5

---

properties of P(t). To this end, we first must highlight some general properties of row-stochastic matrices.
Recall that, by construction, a (F(t) is an irreducible row-stochastic matrix for any finite ((e-- As a result, the Feiron-Frobenuus theorem tells us that the largest eigenvalue (in absolute value) of (F()) is 1 with all other eigenvalues strictly smaller and there exists an eigenvector V(t) such that
a (F(t)) To + M()- T = 1, M O,
Mo) =0;
where v(o) can be interpreted as the stationary distribution of Markov chain with transition proba- bilities (F This observation allows us to decompose A7e and identify its dominant terms:
17) (1 -1-100 1o (Fo())) I+ -1=
n M
77
1v +
Σ " Miyy
k21
#21
k20
Since all eigenvalues of Maa are strictly less than 1, see that for arbitrary vector
1
1 (vxx) + O(1),
A
(5)
as 7) approaches 1. This decomposition of ATo) allows us to see that the vectors 1 and V(x) also dominate the dynamics induced by P(t). Specifically, we see that
P(o = Aio) (n² + (47) = 1 (1-n)² 1v (n) 11 +
Recall that the expected improvement under the implicit parameterization is dominated by (P())(((()))(()))) This means that our results suggest that the component of the residuals, r(((x)) along the direction, i.e., 1 (r((())) plays a eritical role determining the performance under the implicit parameterization, We use this observation to better understand under which conditions we might expect the implicit parameterizations to perform better. We start by formally quantify the dynamics of the component of the vector of residual along this direction with the following theorem.
Theorem 4. Let r (0) be the component of the vector of residuals parallel to the 1 direction, i.e.. r(0) = (0)) + r-(0) (0), then
((0)))((0))) 110 (1)D() +1) v()
1
1
B|r (Ó.-1) - (1))
+ O(afi) + + "
(1 - mm)
(()))(((()))))) -1)
From this result we can conclude that this component of the residual converges very rapidly under the implicit parameterization. For instance, given a, = O(1,, the expected convergence rate along the 1 direction is (11((1 ?)) which is substantially faster than the more general rate, O(1/n), for the explicit and implicit case seen in Theorem 2 and respectively. Moreover, this suggest that
increasing 7 will accentuate this effect.
Intuitively, this suggests that the implicit parameterization is favored when the residuals have similar magnitudes and sign, such as when every state is being over or underestimated. In practice, this is not uncommon since the estimates are often initialized to similar values, e.g., ((s) R O. Additionally, the true state values are typically highly correlated, making this systemic error all the more likely.
To further understand how the properties of the two parameterizations differ, we continue our analysis by examining how the variance of the SGD updates compares to the magnitude of their expectation which we quantify with the ratio of the variance and the squared mean update. If the variance is large compared to the expected updates, then the optimization is likely to spend a significant amount of time bouncing around, leading to slower convergence 12 25). For this reason, it is desirable to have this ratio be as small as possible. We consider this pariance-meat ratio instead of just the variance itself to avoid penalizing cases where the variance is large due to large expected improvements. In the rest of this section, we compare the performance of SGD with the explicit and implicit parameterization through this ratio which we state formally in the following theorem.

---

Theorem 5. Given n data points and a batch size of k. the (variance V(mean square) ratio of SGD updates as follows.
1. For the explicit parametrization,
BiJl to
n(kk - 1) E k(n 1) k(n ||((
+ (n- k)n E, 1
2. For the implicit parametrization,
(n k)n [,1,, + O(1 -n). 1) k(n 1) (2,
n(k -1)
E A (c)
Once more, we see the 1 direction play an important role in determining the hehavior of the implicit Given the importance of the choice of features, o, it can be somewhat difficult to understand the implications of these results. To help build intuition, we conclude by examining the specials case when using tabular values where corresponds a one-hot encoding with the following corollaries.
Corollary 2. For the SGD with explicit parametrization and one-hot encoding the (variance)mmaan square) ratio of each noisy gradient update is given by
(8)
A
(n)
K
We see that under the explicit parametrization, the (variance)((mean square) ratio is given by n/k, and fixed throughout training. In contrast, CheoremmStlll that when using a one-hot encoding under the implicit parameterization, the dependency on the residual remains and will likely vary throughout the optimization. Similar before, can these observations help us better understand the conditions under which the implicit parameterization might exhibit preferable properties. The following corollary provides some insight to what some of these conditions might be.
Corollary 3. Assume the residuals are bounded such that 0 < Tmin S [r(0w))) I Mmaz for all 1 1 i n. then the (variance M/man square) ratio under the implicit parameterization is
B ( TE A 1
21
1 T
Tmaz
S1+
Tmin
This result tells that we should expect considerably better (variance)(mmean square) ratio under the implicit parameterization when the ratio, Tmaz/ Tman/Tmin Tmin. of the largest and the smallest residual among the data points is small, is the case when residuals are similar to each other. For instance, a "balanced" residuals such that (Tmaz/ monaz/mmmn )2 = O(1) means that the (variance /(mean square) ratio of SGD updates is also (1)., whereas it would be ((n) under the explicit parameterization. This tells us that we can expect notable benefit of using the implicit parameterization instead of the explicit one under balanced residuals. Finally, note that this observation mirrors well the results from Theorems] [ and which say that for residual with a dominant component, the convergence rate with the implicit parameterization is roughly times faster. These results strongly suggest that the implicit parameterization might be preferrable in situations where the residuals are expected to be highly correlated.
4 Empirical results
To illustrate our theoretical results, we evaluate the two parameterizations by fitting value functions from data collected offline. Specifically, we consider setting where value targets are generated by rolling out a fixed stochastic policy and are thus noisy and imperfect similar to what might be encountered applications of reinforcement learning. When evaluating, using the true values when dealing with a continuous state space, using unseen trajectories. All states in a trajectory are used with corresponding value targets, !/. computed from the empirical discounted return. Consequently, duplicate states might be present in the dataset and might be given
residual computed
7

---

Chaln MDP
Four Rooms
Mountain Car
1420
1200
90
1000
60
800
Component.
[|r-(0]]
4:00
20
0.0
0.2
0.0
0.2
0.4
0.6
0.0
0.2
Reward Offset
Reward Offict
Rewair Offset
Figure 1: The effect of the reward offset on the components of the residual
different targets. All results show here except for Figure 2 consider the residual after 300 training steps. Since the purpose of these experiments are to illustrate our theoretical results, we leave the temporal difference learning case to future work in which we hope to extend our results.
To ensure any observed differences are caused by differing dynamics and not by differing initialization strategy, we use the same 20 seeds for both parameterizations and initialize parameters such that identical seeds result in the same initial linear weights, and θ. The shaded areas in our figures visualize the sample standard deviation across the seeds. For all experiments, 25. We ran these experiments with several combinations of learning rates and internal discounts but only present few representative results here. Additional results and implementation details can be found the appendix. All other implementation details, data and code are publicly available]] We consider three simple, illustrative domains: chain MDP, the four rooms domain and the mountain car domain, which we describe below.
used a batch size
Chain MDP: two-action MDP which requires optimal agent to repeat the "correct" action until it has reached the end of the chain where the agent receives a large reward, +10, before being sent back to the start of the chain. At any point along the chain, the agent and collect a small reward, +1, before being prematurely sent back to the beginning. The chain MDP has 10 states along the good path and extra transient state when exccuting the wrong action. States were represented with one-hot encoding and 7 = 0.9 was used. Trajectories followed an e-greedy policy which picks random action with probability = 0.1 and an optimal action otherwise.
execute the "wrong" action
Four rooms [21 a gridworld navigation task with obstacles dividing the world into connected rooms. The agent can attempt to move in either of the four cardinal directions but will remain stationary with probability 1/3. The goal is a terminal state with a single -1 reward situated at (7, 1). As before, use a one-hot encoding, an e-greedy policy with = 0.1, and 0.9.
a domain which an agent must accumulate momentum by driving back and forth in order to get an underactuated car over a hill by applying three possible actions: idle, forward or backwards force. Each step results a reward except when transitioning to terminal state. We encode the continuous state space using uniformed 20x20 grid of normalized radial basis functions. Trajectories are sampled using an <-greedy energy pumping policy which picks random action with probability 0.1 and otherwise applies force in the direction of the car's velocity. Furthermore, in lieu of the true analytical values, we generate an additional test dataset and evaluate using similar but likely unseen states. In this domain, we use discount factor = 0.99.
Mountain car 13 20
We consider a set of experiments that observe the change in the performance of both parameterizations as the rewards are offset by constant. The purpose of these experiments is two fold: 1) to demonstrate the significance of the component of the residual in policy evaluation setting, and 2) to illustrate how increasing this component of the residual severely degrades the performance under the explicit case while having a lesser effect under the implicit parameterization as our theory would suggest.
The scale and skew of state values can have a significant impact on the learning performance. While these effects can be mitigated in a supervised learning setting, existing strategies aren't always applicable to RL setting where values might not be known beforehand such as when using temporal difference learning based approach. For this reason, it's not uncommon for values to be offiset away from zero in unison making the 1 component more significant.
http:://ghhh....om////////m////s
8

---

Chain MDP
Four Roomss
Mountain Car
50
400
Parameterizatio
exs
200
0.0
1.0
5.4
0.5
0.6
Reward Offset
Reward Offset
Reward Offiet
Figure 2: The effect of the reward offset on the norm of the final residual using 7) 0.95
Chain
Four Rooms
Mountain
5.
12
Reward Offiet
1.0
250
4.6
150
0.7
42
0
0.
1.0
0.8
0.0
0.95
2.875
0.8
0.9
0.95
0.975
0.8
0.9
0.95
0.975
Figure 3: The effect of 1) on the final residual under the implicit parameterization. The vertical dotted line represents the true discount factor used to generate values. Note that the x-axis uses a logit scale.
We illustrate this by adding constant offset to the reward while keeping policies and transitions unchanged. In the case mountain car, varying this offset from 0 +1 spans the two most common reward formulations for shortest-path like problems, the uniform step cost and the sparse reward approach where reaching the goal is the only non-zeri reward. Figure Siguressooss how the reward offset affects the residual by decomposing the residual at the start of training into two components: its projection on the 1 direction and the norm of the remainder. We see that the reward offset allows us to manipulate the relative magnitude of 1 component enabling to illustrate our theory discussed bellow.
We start by observing how changes in the 1 component affect each parameterizations. Figure compares the residual at the end of training and show that, unsurprisingly, increasing the norm of the initial residual decreases the final performance. However, we also see a significant difference the severity of this effect between the two cases. As established by Theorem 4 the implicit parameterization has significant advantage in convergence rate of the residual's 1 component when compared to the explicit case. Thus, our theory predicts that the implicit parameterization would be less sensitive to changes in this component which is exactly what empirical results show.
Finally, we conclude our empirical contributions by illustrating the role of r) reducing the sensitivity to the 1 component of the residual. Figure plots the final residual of each 7j and compares between the different offsets. We can compare this sensitivity under different r) 's by looking at the relative distance between the plots of given reward offset. Note that, as is the case for the Chain MDP, the least sensitive 1j isn't necessarily the best performing n. In addition, the best isn't necessarily equal to the true underlying discount factor, 7. This raises several interesting questions not covered by our theoretical results, such as how to select the "best" 7J, which we hope to explore in future work.
Conclusion
This work serves as first step towards better understanding the properties of and-to-end model-based methods. This was done by first formulating a simple instance of such methods as an implicit parameterization of a linear function and analysing the dynamics of gradient descent under this parameterization. We showed that despite resulting in non-linear and non-convex optimization problem, the implicitly parameterized linear weights still converges to a unique global optimal. Furthermore, quantified the convergence rate of the residual under the implicit parameterization

---

showing surprisingly fast convergence rate along the 1 direction. We then derive similar conclusions by deriving the ratio of the variance and squared mean SGD updates and identifying some conditions under which the implicit parameterization exhibits significantly faster convergence rate than its explicit counterpart. We conclude our contributions with some empirical results which illustrate some of the phenomena identified in our theoretical results.
Although our analysis is restricted likely to generalize to the full MDP setting since the subgradients of these two formulations have similar forms; both can seen backpropagating through Markov reward chain. This is because the subderivatives of the Bellman optimality constraints only depend on the transitions under the optimal policy, i.e. the Markov reward chain induced by the optimal policy. We hope extend our results to the more general MDP case in future work.
simple Markov reward chain model, we believe our results
6 Acknowledgements
We gratefully acknowledge support from NSF grant 1723381; from AFOSR grant FA9550-17-1- 0165; from ONR grant 1-0014-88-1-8444 from MIT-IBM Watson Lab; from the MIT Quest for Intelligence and from SUTD Temasek Laboratories.
References
[1] L. Baird. Residual algorithms: Reinforcement learning with function approximation. In Machine Learning Proceedings 1995, pages 30-37. Elsevier, 1995.
[2] S. Boyd and A. Mutapcic. Stochastic subgradient methods. Lecture Notes for EE364b, Stanford University, 2008
[3] S. Boyd and L. Vandenberghe. Convex optimization. Cambridge university press, 2004.
[4] Huang and H.-T. Yau. Dynamics of deep neural networks and neural tangent hierarchy, In International conference machine learning. 2020.
[5] Z. Ji and M. Telgarsky. Directional convergence and alignment in deep learning, Advances in Neural Information Processing Systems, 33, 2020.
[6] Karkus, X. Ma, D. Hsu, L. Kaelbling, W. S. Lee, and T. Lozano-Perez, Differentiable algorithm networks for composable robot learning. In Proceedings of Robotics: Science and Systems, Freiburg Breisgau, Germany, June 2019. doi: N..15607/RSS20119.X. XV.099.
[7] K. Kawaguchi. Deep learning without poor local minima. Advances in Neural Information Processing Systems, 29:586-594, 2016.
[8] K. Kawaguchi. On the theory of implicit deep learning: Global convergence with implicit layers. In International Conference Learning Representations, 2021.
[9] K. Kawaguchi and Y. Bengio. Depth with nonlinearity creates no bad local minima in resnets. Neural Networks, 118:167-174, 2019.
[10] K. Kawaguchi, L. Zhang, and Z. Deng. Understanding dynamics nonlinear representation learning and its application. arXiv preprint arXiv:2106.1448 14836, 2021.
[11] N. Lambert, Amos, O. Yadan, and R. Calandra. Objective mismatch in model-based reinforcement learning. In Learning for Dynamics and Control, pages 761-770. PMLR, 2020.
[12] A. Miller, Foti, A. D' Amour, and R. Adams. Reducing reparameterization gradient variance. In I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30, pages 18. Curran Associates, Inc., 2017.
[13] A. W. Moore. Efficient memory-based learning for robot control. Technical report, University of Cambridge Computer Laboratory, 1990.
10

---

[14] E. Nikishin, R. Abachi, R. Agarwal, and P.-L. Bacon. Control-oriented model- reinforce- ment learning with implicit differentiation. arXiv preprint arXiv:2106.03273 2021.
[15] Oh, S. Singh, and H. Lee. Value prediction network. I. Guyon, U. V. Luxburg, S. Ben- gio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 30. pages 61 6128. Curran Associates, Inc., 2017.
[16] R. Parr, L. Li, G. Taylor, C. Painter- Wakefield, and M. L. Littman. An analysis linear models, linear value-function approximation, and feature selection for reinforcement learning. In Proceedings of the 25th international conference on Machine learning, pages 752-759, 2008.
[17] A. M. Saxe, J. L. McClelland, and S. Ganguli. Exact solutions to the nonlinear dynamics of learn- ing deep linear neural networks. In International Conference on Learning Representations, 2014.
[18] D. Silver, H. Hasselt, M. Hessel, T. Schaul, A. Guez, T. Harley, G. Dulac-Arnold, D. Reichert, N. Rabinowitz, Barreto, International Conference Machine Learning, pages 3191-3199, PMLR, 2017.
The predictron: End-to-end learning and planning. In
[19] A. Srinivas, A. Jabri, P. Abbeel, S. Levine, and C. Finn. Universal planning networks: Learning generalizable representations for visuomotor control. In International Conference on Machine Learning, pages 4732-4744 PMLR, 2018.
[20] R. S. Sutton and A. G. Barto. Reinforcement learning: An introduction. MIT press, 2018.
[21] R. S. Sutton, Precup, and Singh. Between mdps and semi-mdps: framework for temporal abstraction in reinforcement learning. Artificial intelligence, (1111)):111-2 1-111, 1999.
[22] E. Talvitie. Self-correcting models for model-based reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, 2017.
[23] A. Tamar, Y. Wu, G. Thomas, S. Levine, and P. Abbeel. Value iteration networks. In Proceedings of the 30th International Conference Neural Information Processing Systems, pages 2154- 2162, 2016.
[24] A. Venkatraman, M. Hebert, and J. Bagnell. Improving multi-step prediction of learned time series models. In Proceedings of the AAAI Conference on Artificial Intelligence, 2015.
[25] C. Wang. X. Chen, A. J. Smola, and E. P. Xing. Variance reduction for stochastic gradient optimization. In C. J. C. Burges, L. Bottou, M. Welling, Z. Ghahramani, and K. Q. Weinberger, editors, Advances Neural Information Processing Systems, volume 26, pages 181-189. Curran Associates, Inc., 2013.
[26] K. Xu, M. Zhang, S. Jegelka, and K. Kawaguchi. Optimization of graph neural networks: Implicit acceleration by skip connections and more depth. In International Conference on Machine Learning, 2021.

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes] See Corollary D Theorem + and Section
(b) Did you describe the limitations of your work? (Yes] See Assumption and Section
and
(c) Did you discuss any potential negative societal impacts of your work? [No] We provide
theoretical results to better understand the strengths and weaknesses of existing methods
which might help practitioners apply them more appropriately.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes] See Assump- tion and
and Section 2
(b) Did you include complete proofs of all theoretical results? [Yes] Proofs are provided in the appendix due to limited space.
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either in the supplemental material or URL)? [Yes] The data from our experiments, and the code to produce and plot is provided.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Ys]] See [ Missing details are provided in the appendix, The exact seeds used are provided with the code.
(c) Did you report error bars (c.g.. with respect the random seed after running experi- ments multiple times)? [Yes] We report the empirical standard deviation over the 20 seeds used.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] We provide sufficient information estimate this in the appendix as well as the hardware used. Additionally, part of the data bundled with the code, we report the runtime of each run.
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) your work uses existing assets, did you cite the creators? [N/A]
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any new assets either in the supplemental material or as a URL? [N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you're 're using/curating? [N/A]
(e) Did you discuss whether the data you using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
12


Reviews:
[{"review_id": "uuk01L0WTE", "reviewer": "Reviewer_kLUp", "summary": "This paper presents a theoretical analysis of two different ways of parameterising a linear value function approximator. They consider an 'explicit' parameterisation, in which the weights connecting the state features to value estimate are parameterised directly, and an implicit parameterisation, which attempts to use knowledge of the properties of a Markov process to inform a more structured parameterisation. \n\nThe authors analyse the learning dynamics of these two parameterisations under stochastic gradient descent, and draw some conclusions about the properties of RL environments that might favour the implicit parameterisation. Finally, they carry out some empirical investigations on offline datasets of trajectories from several different environments.\n", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The authors attempt to motivate their work as a theoretical investigation into model-based methods that model transition probabilities and rewards. In particular, they attempt to draw a close parallel to \"value iteration networks\", which is a particular form of model-based method. \n\nDespite the undoubted importance of investigations into model-based planning, I am not convinced that this is really what this work achieves. In particular, their method is really a different parameterisation of a purely model-free method - as far as I can see there is no internal policy improvement or value iteration step based on a model, which is what a model-based method would be aiming to achieve. For example, value iteration networks use repeated applications of a convolutional network with max pooling as a form of differentiable \"end-to-end\" value iteration in order to inform an improved policy.\n\nAdditionally, their implicit parameterisation seems really most applicable to a discrete-state Markov process (relying on a vector of rewards and an explicit transition matrix as components of its parameterisation), which is not typically the case in more complex environments.  \n\nTheir theoretical analysis provides some interesting insights into how environments with reward offsets can impact learning in this setting. However, I am unconvinced that this would be an issue for modern algorithms that use temporal differences and baselines, in part precisely to help improve learning in such contexts. In the setting with no reward offset, it seems that both parameterisations perform similarly well in their empirical results.\n\nIn summary, despite an admirable intention to investigate an important area, I am not convinced this paper makes sufficient contributions to the community in its current form. \n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "There are no anticipated negative societal impacts from this work.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "k2dR5n84My", "reviewer": "Reviewer_CGbD", "summary": "This paper provides a theoretical analysis of learning value functions in a Markov reward chain by SGD. They propose to analyze these methods as indirectly parameterizing weights of a linear function (with fixed features) and compare to explicitly learning the linear weights. The authors prove that, under certain assumptions, the implicit parameterization converges to the global minimum. The authors note that the gradient updates differ because the implicit parameterization preconditions the gradient. The authors then characterize the conditions under which this implicit parameterization performs better (as measured  by the variance-to-mean-square ratio) by analyzing the dynamics of the primary component of this preconditioning matrix and find that this occurs when the sign of the components all point in the same direction. The authors corroborate this theory with a series of experiments on toy environments.", "questions": "", "limitations": "", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "This paper studies an interesting question regarding the dynamics of V-learning in RL.\n\nMy main concern of this paper is that it is unclear how relevant this analysis is to end-to-end RL algorithms. In particular, the premise that value-based learning methods learn linear weights of fixed features requires more justification. The restriction to Markov reward chain, rather than an MDP also greatly limits the scope of the analysis.\n\nFor example, the statement, “Note that any insight in this simpler framework is likely relevant to the state-dependent model setting since the gradients differ only by the Jacobian relating model parameters’ differentials to learned parameters’ differentials, i.e., the chain rule.” seems to sweep aside the fact that this Jacobian can be arbitrarily complex for deep neural networks. Similarly, it would be useful to explain why analysis on a Markov reward chain would translate to VIN, which internally takes a max over actions.\nAs another example, what are the restrictions or assumptions on the constraint h and why are those assumptions realistic?\n\nMoreover, does O(alpha^2) include terms that depend on F? If so, the abstract statement “We prove that, for a linear parametrization, gradient descent converges to global optima despite non-linearity and non-convexity introduced by the implicit representation” seems misleading as it misses an important condition. In particular, the paper proves that there is a global minimum, but I believe it requires the strong assumption that all second order terms are dominated by a quadratic. Given this strong assumption, the global convergence results are not as universal as the statement implies and are more limited in significance.\n\nOverall, although the paper performs interesting analysis, tying this analysis back to RL algorithms and justifying the simplifying assumptions made would greatly improve the impact of this paper.\n\ntypos:\nabstract: “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.”", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "N/A as this paper provides mainly theoretical analysis of V-learning in Markov reward chains, which is too generic to warrant discussing societal impact", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "1"}, "extra_text": {}}, {"review_id": "Ww9wDBOavOy", "reviewer": "Reviewer_g4M6", "summary": "This paper focuses on the area of theoretical reinforcement learning, where the authors consider the form of parameterization for the value function.  Specifically, the authors consider the case of a Markov decision process with a linear value function, and compare the explicit and implicit way of parameterizing the weights of the value function. The authors prove that using gradient descent, the implicit parameterization is equivalent to the explicit parameterization with a gradient preconditioner, which converges to the global minima. The authors then extend this result to stochastic gradient descent, and then compare the condition in which implicit parameterization gives better convergence speed. The authors provide empirical evidence in several discrete MDP environments.", "questions": "", "limitations": "", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Overall I think this paper presents an interesting study of the implicit parameterization of value functions for a Markov decision process. However, I believe that there are also some limitations that need to be addressed.\n\n\nPros:\n\nThe paper is well written. The structure of the paper is well organized and the authors present the theoretical results in a step-by-step way, which makes it easy for the reader to grasp the motivation behind the theoretical results.\n\nBeing a theoretical paper, the paper also includes empirical studies in several MDP environments, making the results more intuitive.\n\n\nCons:\n\nThe choice of the specific implicit parameterization (equation 2) requires more justification. We know that setting equation 2 to zero gives us the Bellman equation for the tabular Markov reward process with rewards $w$, discount factor $\\eta$ and transition probability $\\sigma(F)$. The authors suggest that the internal discount factor $\\eta$ is different from the true discount factor $\\gamma$. Under this assumption it is clear that the ground truth transition matrix and reward vector cannot be a good fit for parameters $\\sigma(F)$ and $w$. It is then unclear what the learned $\\sigma(F)$ and $w$ represent. Are we looking for a different MDP dynamics and rewards with discount $\\eta$ such that the value function happens to be the same as our ground truth MDP?\n\nMoreover, the authors claim that “this implicit formulation allows us to analyze how $\\psi$ and $\\hat{\\theta}$ relate without needing to consider the algorithm used to compute $\\hat{\\theta}$”. I believe this is not valid for this particular implicit parameterization, since it is well known that $\\hat{\\theta} = (I - \\eta \\sigma(F))^{-1} w$, which suggests that this implicit parameterization can be considered as just another form of explicit parameterization.\n\n\n\nOne important aspect that prevents implicit parameterization to be used in practice is the computation complexity of the implicit gradient, which the authors fail to consider in this paper. Computing the inverse Jacobian product in the implicit gradient has complexity $kd + d^2$, for each step where $d$ is the dimension of $\\hat{\\theta}$ and $k$ is the batch size. On the other hand normal stochastic gradient descent only has complexity $kd$. Since $d$ is usually larger than $k$, it is unclear that implicit SGD would have much advantage compared to explicit SGD given the same computation budget. Furthermore, under the one hot encoding assumption in the paper, it is not clear why SGD is even necessary if the data takes the form of empirical discounted sum of rewards, as one can obtain the exact minima of equation 3 by simply averaging all the data points of each state.\n\n\nGiven these concerns, it is hard to recommend the acceptance of this paper before they are addressed.\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "See main review for limitations, and N/A for social impact.", "review_text": "", "extra_scores": {"ethical_concerns": "N/A", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "TL5VdOV60h6", "reviewer": "Reviewer_AqM3", "summary": "This paper does some theoretical study to explain the fast convergence of end-to-end model-based RL in some domains. Specifically, the authors prove global optimal convergence for the linear approximation case and study the convergence rate under certain conditions that stochastic gradient descent (SGD) with implicit representation converges substantially faster than its explicit counterpart. The authors also provide empirical results in some simple environments to back up their theory.\n", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Originality, quality, clarity of the paper:\nThe paper studies an observation with little formal justification in RL literature: End-to-end model-based methods are heuristically shown to be capable of better generalization in some RL domains.\nThere are substantial theoretical results (but I have not checked the proof carefully, since I have no prior experience in working with gradient descent dynamics). Overall, the derivation is smooth and has good logic flow, with relatively enough explanation. The definitions and quantities introduced show the thoughtfulness of the authors. Thus I think this paper is of relatively good quality.\n\nSignificance of the paper:\nThere are 2 main theoretical contributions. The first is showing the implicitly parameterized linear weights in the model still converges to a unique global optimal. The second is finding conditions under which the implicit parameterization exhibits significantly better properties than explicit parameterization. Besides, the empirical results seem to support the second theoretical claim. Therefore I think this paper does provide some novel insight on the principle of end-to-end model-based RL methods. \n\n*post-rebuttal*:\n Thanks for your reply. I have read your response and my review keeps the same.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The empirical results are conducted on toy environments and have some disparity to the actual implementation of recent model-based RL algorithms. Meanwhile, if to supplement the empirical results, the authors could do similar experiments on more complex domains, such as the Atari or Mujoco environments, and with deep RL algorithms. Though this might introduce unknown factors that affect the final convergence, any implication in such a setting might be more helpful to algorithmic works in the community.\nOne last concern is that the authors could have spent more paragraphs discussing related literature (both empirical and theoretical) and summarize their implications. \n\nI don’t see any potential negative societal impact of this work.", "review_text": "", "extra_scores": {"ethical_concerns": "No ethical concerns", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "Ruoye7Rj5O", "reviewer": "Reviewer_LFGm", "summary": "The paper studies value estimation in end-to-end model-based RL methods.\nThe setting is simplified for analysis. They consider action-less MDP. The form of the value is assumed to be a linear combination of fixed state features \\phi(s) with constant coefficients \\theta. If \\phi is one-hot, \\theta would correspond to a vector of state values.\nThe implicit rep. for \\theta is a log Transition matrix + reward vector and \\theta is defined to satisfy the bellman condition.\nThe paper then studies gradient and gradient descent of regression for implicit and explicit representation.\nThe first main result is to show the similarity between explicit update and implicit update save for a ‘preconditioning’ matrix P (and higher order terms). They prove P positive definite and obtain global convergence from that.\nIn the next part they decompose the transition matrix and show how the preconditioner P behaves in the limit \\eta -> 1 and its consequences (preconditioner favors 1 component).\nFinally they present some experiments on toy tasks where they are showing that the 1 component of the residual is increased by a reward offset and that the final performance is less sensitive to this in the case of implicit compared to explicit.\n", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n1. Using this analytic implicit representation for analysis is a neat idea.\n2. Decomposing the transition matrix in terms of the stationary distribution is also very neat and I didn’t expect the arising form of the preconditioner.\n3. The writing presents the ideas quite clearly.\n\nIssues:\n1. I believe there might be an error in corollary 3. The sum of $r_i^2$ shouldn’t cancel with the denominator which is sum of squares: $||\\sum_i r_i \\phi(x_i)||_2^2$ (same as in implicit case). Actually I think in this case, implicit and explicit should have the same variance/mean ratio (except for $O(1-\\eta)$ term).\n2. Parallel residual convergence is a main result and is the only one that is tested experimentally. However, I’m not sure why the convergence rate for the sum of residuals is important? Can’t you get 0 sum of residual by overestimating some examples and underestimating others?\n3. Related to above, in Theorem 4 the authors claim faster convergence for the parallel residual by comparing the rates. This is an apples to oranges comparison though since in theorem 3 they are talking about convergence rate of loss and theorem 4 they are talking about convergence rate of sum of residual. Wouldn’t you want to compare to the rate for $\\sum_i (<r_i \\phi_i,1>^2$) instead?\n4. Experiments: Since the theoretical results mainly regard convergence rates, it would have been good to also show learning curves wrt # of iterations.\n5. The additional figures in the appendix are hard to read (why are there some learning rate runs missing for the implicit case?). There also seem to be graphical glitches on a few of them.\n6. Introducing implicit differentiation seems kind of superfluous here, since the actual “implicit” rep for \\theta can be easily expressed explicitly.\n7. There are a few syntactic typos in the proofs (f.e. on line 95 the brackets, on line 106 duplicate $\\sum_i$, line 108/109 there’s an extra 1/n out of nowhere that then disappears later)\n8. This might be obvious but I think you should still state if/when the implicit rep will be as expressive as the explicit one since this affects the minima.\n9. I believe you should state the $O(\\alpha^2)$ term in theorem 4 as well.\n10. In proof of theorem 3, near line 102, why is there no 1/n from the loss? Isn’t loss defined to have $1/|D_t|$?\n\nNovelty/Impact:\nI will first admit that optimization theory is not my area of expertise.\nI think studying end-to-end model-based RL is interesting. However, I think generally the benefit that people in the broader community are more interested in when using these methods is generalization and sample efficiency rather than optimization convergence rate.\nFurthermore, the simplifications used in this work are rather strong. I feel without TD learning, structural assumptions about the value targets and action selection the setting essentially reduces to supervised regression with a particular parameterization. \nThat being said, the paper does present some interesting results about the gradient of implicit param itself that could be built upon.\n\nRating:\nI think the work could be an alright theoretical analysis. However, I am giving a lower rating and also lower confidence primarily because I question some of the main contributions claimed: Specifically I question the validity of corollary 3 and don’t feel like theorem 4 warrants the interpretation it is given as well as a couple technical deficiencies and sloppy writing in the proofs. But I’d be willing to raise my score if these can be clarified.\n\nEdit:\nAfter rebuttal, authors clarified my doubts about corollary 3 and I increased my rating to 5.\nI still think the significance of the results are somewhat limited with the simplified setting.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "7"}, "extra_text": {}}, {"review_id": "I_UWrWSjwG", "reviewer": "Reviewer_EkYa", "summary": "This paper proposes a theoretical framework to analyze model-based end-to-end reinforcement learning. In particular, the authors focus on the MBRL with a linear approximation of the value function, and within the Markov reward chain setting (with singleton action). Theoretical results show that MBRL in an end-to-end way is actually doing implicit parameterization and convergence analysis.\n", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "### Originality\n\nThe main contribution of this paper is to provide a theoretical analysis of MBRL. With the implicit differentiation theorem, the authors link the implicit parameterization of the learned model and the linear feature weights, which looks interesting to me. The most interesting result of this paper, in my opinion, is that introducing reward and transition models is an implicit way of preconditioning in the gradient flow.\n\n### Quality\n\nAs claimed by the authors, all main results in this paper are derived under the Markov reward chain framework. However, what I can see in the main part, e.g., Theorem 1 and Corollary 1/2, they're all analyzed under finite-sum square-error and no direct link to the common reinforcement learning scenario. For example, how the value function is estimated and updated? How does the agent update/improve the policy? What's the end-to-end MBRL algorithm the agent is performing? Based on this, I cannot agree with the authors' conclusion that\n\n> ..., we believe our results are likely to generalize to the full MDP setting since ...\n\nI would like to explain this point with a few more details, which lead to my concerns regarding the quality of this paper. \n\n1. The constraint Eq. 2 appears quite confusing to me. I'd like to see more justification on why this constraint can hold. It seems to me that it can be derived from the Bellman equation, but it leads to another problem: the optimality of $\\hat{\\theta}$ is not identical to that of $\\psi$ when $g()$ is non-linear. For example, when the image of $g()$ is a subset of $\\hat{\\theta}$, the optimal $\\psi$ may not be obtained. For this reason, I'm not convinced by the proof of Corollary 1 (and 2). \n\n2. Another issue on implicit differentiation is about the assumption on $\\frac{\\partial g}{\\partial \\psi}$: it should exist, but may not hold in many cases, e.g., when linear piece-wise activations like ReLU are involved in a neural network.\n\n3. The analysis of finite-sum square error (Eq. 3) is not clear in MBRL. Does this equation apply at each timestep $t$, or at the end of each episode? For either case, how does the agent collect the datapoints (from the real environment or the model)? ... I think such ambiguity comes from the clarity issue of presentation. See the *Clarity* section for more details.\n    * I'm trying to understand this in two cases: a) when the gradient updates in Eq. (4) are performed once per timestep; b) when the steps in Eq. (4) are performed until converging in each step. \n    * In case a), the loss function $L$ and the datapoints should change along the time, which vanishes the assumption of the proof (a finite-sum result). \n    * In case b), we can still conclude that the value prediction converges to a global minimum at the step, but it relies on the fact that such value prediction can apply on all timesteps, which is not discussed in the paper. \n    * Note that this also holds for the Markov reward chain discussed in the paper: the datapoints collected at step $t$ could be different from that in other steps. \n\n4. I'm also confused about the explicit parameterization scheme in the paper. One reason might be from the above point of clarity issue, and the other one is: it seems to me that explicit parameterization is not using any models, so it's actually a model-free RL (MFRL)? Considering the fact that MBRL takes more calculations (planning steps) during training, it's not surprising to me that the implicit parameterization converges faster than explicit parameterization.\n\n### Clarity\n\nThis paper has some issues with writing and organization. For example:\n\n1. This paper doesn't have a clear formal definition of the Markov reward chain that can be usually found in the RL literature. As a result, readers would find it hard to understand why such a model can help for MBRL analysis.\n2. The authors don't provide a clear context on the gradient steps, especially when the steps are performed. As discussed in the *Quality* section, this would result in many confusions and weakens the conclusions of this paper.\n3. Given the results from the finite-sum minimization, the authors should provide clarifications on how this relates to MBRL.\n\n### Significance\n \nThe authors provide a convergence analysis to understand the learning of MBRL. However, as discussed above, the results are mostly derived from finite-sum with unknown sampling methods, it's hard for readers and other researchers to use the results for the MBRL study, either in theory or in practice.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "This paper only has results based on a simplified model, Markov reward chains (MRC), which is different from the MDP in the RL literature in the sense that, MRC doesn't involve the action selection. In the meanwhile, the obscureness of gradient updates in this paper makes the readers hard to understand the significance of the work, neither using nor applying the results in other researches. In particular, how to sample the datapoints is not discussed in the paper, therefore I'm convinced by the results can explain the dynamic nature of RL.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "20"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper aims at presenting a theoretical investigation of the gradient dynamics and the convergence properties of model-based RL methods. In order to make the analysis theoretically tractable, the authors analyze these methods by looking at a simplified setting: training a linear value function approximator using two formulations, an implicit one, i.e. parameterizing weights of a linear function with fixed features, and explicitly learning the linear weights.

The first contribution of the paper is providing proof that under some assumptions the implicit parameterization converges to the global minimum, and showing that the gradient updates differ in the two settings because the implicit parameterization induces a preconditioning on the gradients. A second contribution is in further analysing the properties of the preconditioning. The theoretical findings are finally verified experimentally on small scale problems using offline datasets, rendering the theoretical contributions easier to understand.

Whilst the paper attempts to address a very interesting and relevant problem in the field, and the submission indeed presents a stimulating study, all six reviewers have struggled to find a clear and strong connection between the insights proposed by the authors and "a theoretical account of the performance of end-to-end model-based methods" and lamented lack of substantial (practical) results. Therefore, I fear the paper in the current format will have limited impact at NeruIPS, and draw little follow up from this audience.

I hope that the authors will not give up on this work, and will consider either submitting to a venue with a more specialist audience, or acting on the specific recommendations some of the reviewers have provided, e.g. reframing the work to make it easier for the readers to draw connections between its results and the model-based methods discussed in the introduction.
Decision: Accept (Poster)