Control-Oriented Model-Based Reinforcement Learning with Implicit Differentiation
Abstract
The shortcomings of maximum likelihood estimation in the context of model-based reinforcement learning have been highlighted by an increasing number of papers.
When the model class is misspeciﬁed or has a limited representational capacity, model parameters with high likelihood might not necessarily result in high performance of the agent on a downstream control task.
To alleviate this problem, we propose an end-to-end approach for model learning which directly optimizes the expected returns using implicit differentiation.
We treat a value function that satisﬁes the Bellman optimality operator induced by the model as an implicit function of model parameters and show how to differentiate the function.
We provide theoretical and empirical evidence highlighting the beneﬁts of our approach in the model misspeciﬁcation regime compared to likelihoodbased methods.
1. Introduction
The conceptual separation between model learning and policy optimization is the basis for much of the work on modelbased reinforcement learning (MBRL) (Grefenstette et al., 1990; Sutton, 1991; Lin, 1992; Boots et al., 2011; Chua et al., 2018; Hafner et al., 2019; Janner et al., 2019; Kaiser et al., 2019).
A standard MBRL agent ﬁrst estimates the transition parameters and the reward function of a Markov
Decision Process and then uses the approximate model for planning (Theil, 1957; Kurano, 1972; Mandl, 1974; Georgin, 1978; Borkar & Varaiya, 1979; Hern´andez-Lerma & Marcus, 1985; Manfred, 1987;
Sato et al., 1988).
If the estimated model perfectly captures the actual system, the resulting Reinforcement Learning Workshop at the 37 th International Conference on Machine Learning , 2021. Copyright 2021 by the author(s).
Model
Q*
Planning
IFT
Loss
Acting
Backprop policies are not affected by the model approximation error. However, if the model is imperfect, the inaccuracies can lead to nuanced effects on the policy performance (Abbad & Filar, 1992; Manfred, 1987).
Several works (Skelton, 1989; Joseph et al., 2013; Lambert et al., 2020) have pointed out on the objective mismatch in MBRL and demonstrated that optimization of model likelihood might be unrelated to optimization of the returns achieved by the agent that uses the model.
For example, accurately predicting individual pixels of the next state (Kaiser et al., 2019) might be neither easy nor necessary for decision making.
Motivated by these observations, our paper studies control-oriented model learning that takes into account how the model is used by the agent.
While much of the work in control-oriented model learning has focused on robust or uncertainty-based methods (Ludwig & Walters, 1982; Nilim & Ghaoui, 2005; Iyengar, 2005;
Xu & Mannor, 2010; Yu et al., 2020), we propose an algorithm for learning a model that directly optimizes the expected return using implicit differentiation (Christianson, 1994; Griewank & Walther, 2008).
Speciﬁcally, we assume that there exists an implicit function that takes the model as input and outputs a value function that is a ﬁxed point of the
Bellman optimality operator (Denardo, 1967) induced by the model.
We then calculate the derivatives of the optimal value function with respect to the model parameters using the implicit function theorem (IFT), allowing us to form a differentiable computational graph from model parameters to the sum of rewards.
In reference to (Rust, 1988; Sorg et al., 2010; Bacon et al., 2019), we call our control-oriented method optimal model design (OMD).
Our contributions can be summarized as follows:
• We propose OMD, an end-to-end MBRL method that optimizes expected returns directly.
• We characterize the set of OMD models in the tabular case and derive an approximation bound on the optimal Q-function that is tighter than a likelihood-based one.
• We propose a series of approximations to scale our approach to non-tabular environments.
• We demonstrate that OMD outperforms likelihoodbased MBRL agents under the model misspeciﬁcation in both tabular and non-tabular settings.
This ﬁnding suggests that our method should be preferred when we cannot approximate the true model accurately.
• We empirically demonstrate that models obtained by OMD can have lower likelihood than a random model yet generate useful targets for updating the value function.
This ﬁnding suggests that likelihood optimization might be an unnecessary step for MBRL.
2. Related work
Learning control-oriented models.
Earlier work in optimal control and econometrics (Skelton, 1989; Rust, 1988) studied the relation between the model approximation error and the control performance and noted that true parameter identiﬁcation could be suboptimal when the model class is limited.
Joseph et al.
(2013) were one of the ﬁrst to address the objective mismatch (Lambert et al., 2020) and proposed an algorithm for training models that maximize the expected returns using zero-order optimization.
Several papers have proposed model learning approaches that optimize other return-aware objectives.
Farahmand et al.
(2017) train a model to minimize the difference between values of the real next states and the next states predicted by the dynamics.
Abachi et al.
(2020) use the norm of the difference between policy gradients as the model objective.
D’Oro et al.
(2020) use a weighted maximum likelihood objective where the weights are chosen to minimize the difference between the true policy gradient and the policy gradient in the MDP induced by the model.
Schrittwieser et al.
(2019) use tree search and train models for image-based states by encoding them into a latent space and predicting rewards, a policy, and values without reconstruction.
The idea of differentiable planning has also been investigated.
Amos et al.
(2018) learn a model via differentiating the KKT conditions in the LQR setting (Dorato et al., 1994).
Tamar et al.
(2016) uses a differentiable approximation of the value iteration algorithm to learn a planner.
Amos &
Yarats (2020) optimize the parameters of a sampling distribution in Cross-Entropy Method (Rubinstein, 1997) using a differentiable approximation of Top-K operation.
Several works have theoretically studied the control-oriented model learning.
Ayoub et al.
(2020) derive regret bounds for models used to predict values.
Grimm et al.
(2020) introduce the principle of value equivalence for MBRL deﬁning two models to be equivalent if they induce the same Bellman operator.
Our work is closely related to the above papers but proposes to learn models by directly optimizing the sum of rewards in an end-to-end manner via gradient-based methods.
Implicit function theorem.
Implicit differentiation has been applied for a variety of bi-level optimization problems.
Lorraine et al.
(2020) treat weights of a neural network as an implicit function of hyperparameters and use IFT to optimize the hyperparameters.
Rajeswaran et al.
(2019) study meta-learning and apply IFT to compute the outer loop gradient without the need to differentiate through the inner loop iterations.
Instead of treating a neural network as a sequence of layers that transform an input, Bai et al.
(2019) propose an implicit layer that corresponds to an inﬁnite depth neural network and ﬁnd a ﬁxed point of the layer via IFT.
Our method also solves a bi-level problem: in the inner loop, we train an action-value function compatible with the model, while in the outer loop we update the model parameters towards maximizing the expected returns.
3. Preliminaries
Reinforcement Learning (RL) (Sutton & Barto, 2018) methods follow the Markov Decision Process (MDP) formalism.
An MDP is deﬁned as M = ( S , A , γ, p, r, ρ 0 ) , where
S is a state space, A is an action space, p ( s (cid:48) | s, a ) is a transition probability distribution (often called dynamics ), r ( s, a ) is a reward function, γ ∈ [0 , 1) is a discount factor, and ρ 0 ( s ) is an initial state distribution.
The pair ( p , r ) is jointly called the true model .
The goal of an agent is to learn a policy π ( a | s ) that maximizes the expected discounted sum of rewards J ( π ) =
{+E π [ (cid:80) ∞ t =
+}
γ t r ( s t , a t )] .
The performance of the agent following the policy π can also be quantiﬁed using the action value function Q
[Equation]
Model-based RL algorithms typically train a model ( p θ , r θ ) and use it for policy or value learning.
Traditional methods based on Dyna (Sutton, 1991) rely on maximum likelihood estimation (MLE) of model parameters θ .
For example, if the true model is assumed to be Gaussian with a parameterized mean and a ﬁxed variance, maximizing the likelihood is equivalent to minimizing the mean squared error of the prediction, namely, to solving
[Equation]
4. Optimal Model Design for Tabular MDPs
Consider a modiﬁcation of the original RL problem statement, which was ﬁrst proposed by Rust (1988) and revisited by Bacon et al.
In addition to maximizing the expected returns J , we introduce a constraint forcing the action value function Q to satisfy the Bellman equation induced by the model.
The optimization problem becomes
[Equation]
where π
[Equation]
B θ here is the soft Bellman optimality operator with respect to the model parameters θ :
[Equation]
We choose the soft
Bellman operator with log (cid:80) a (cid:48) exp Q ( s (cid:48) , a (cid:48) ) over the “hard” version with max a (cid:48) Q ( s (cid:48) , a (cid:48) ) because of the differentiability of logsum-exp.
We also use a temperature α in softmax and log-sum-exp but omit it from the expressions for simplicity.
Note that ﬁnding a ﬁxed point of the soft Bellman optimality operator corresponds to solving the MaxEnt RL formulation (Ziebart et al., 2008), but for a sufﬁciently small value of α , the difference is negligible.
Suppose there exists an implicit function ϕ ( θ ) =
Q ∗ that takes as input a model and outputs a Q-function that satisﬁes the constraint in (2).
The sequence of transformations from the model parameters to the agent’s performance can be described then using the following graph:
[Equation]
In Section 4.1, we show how ∂ϕ ( θ ) ∂θ can be calculated using the implicit function theorem (IFT).
Since ∂J ( π ) ∂π can be calculated using the policy gradient theorem (Sutton et al., 1999), we can apply automatic differentiation to calculate the gradient with respect to θ :
[Equation]
Given the expression for the gradient of J with respect to θ , we use an appropriate optimization method to train the model.
We call the approach optimal model design (OMD).
Note that Dyna-based methods also train the Q-function to satisfy the constraint in (2) while using the likelihood as the objective for model parameters θ (Rajeswaran et al., 2020).
In contrast, we train θ to directly optimize the returns.
The optimization problem (2) suggests that OMD is a policybased method (Sutton et al., 1999).
However, we can turn it into a value-based approach (Watkins & Dayan, 1992) by replacing the objective J ( π Q ) with the Bellman error:
[Equation]
where B , similarly to B θ , is the soft Bellman operator but induced by the true reward r and dynamics p .
We discuss the relation between the models obtained by solving problems (2) and (6) in Section 5.1.
While the constraint Q ( s, a ) =
B θ Q ( s, a ) has to be satisﬁed for all state-action pairs limiting the approach to tabular
MDPs, we show an extention to the function approximation case in Section 6.
4.1. Implicit Differentiation
In this subsection, we state the implicit function theorem used to calculate ∂ϕ ( θ ) ∂θ .
Theorem 1.
(Cauchy,
Implicit
Function)
Let f : Θ × W
→ W be a continuously differentiable function and (˜ θ, ˜ w ) be a point satisfying f (˜ θ, ˜ w ) = 0 .
If the Jacobian ∂f (˜ θ, ˜ w ) ∂w is invertible, then there exists an open set U ⊆
Θ containing ˜ θ and a unique continuously differentiable function ϕ such that ϕ
(˜ θ )
=
˜ w and f ( θ, ϕ ( θ ))
=for all θ ∈ U .
Moreover,
[Equation]
We provide a proof in Appendix A. Note that (7) requires only a ﬁnal point w ∗ satisfying the constraint and does not require knowledge about ϕ itself.
Hence, ϕ can be any black-box function outputting w ∗ .
The gradient of the scalar objective J or L true is calculated using (7).
To use backpropagation, we only need to deﬁne the product of a vector and ∂ϕ ( θ ) ∂θ .
We provide an implementation of a custom vector-Jacobian product for the implicit function ϕ in Appendix B allowing to use ϕ as a block in a differentiable computational graph.
We provide an implementation of a custom vector-jacobian product for the implicit function ϕ in Appendix B allowing to use ϕ as a block in a differentiable computational graph.
4.2. Beneﬁts under Model Misspeciﬁcation
In the previous subsection, we showed how to use implicit differentiation for training a model that aims to maximize the expected returns.
In this subsection, we demonstrate that such a control-oriented model is preferable over a likelihoodbased in the setting where the true model is not representable by a chosen parametric class.
Let r θ ∈ R |S|×|A| and p θ ( s (cid:48) | s, a ) ∈ R |S|×|S|×|A| be a parametric model, where parameters in p θ denote the corresponding logits and each parameter in r θ is a reward for a stateaction pair.
We consider a set of parameters { θ : (cid:107) θ (cid:107) ≤ κ } with the bounded norm and use κ as a measure of the model misspeciﬁcation.
By decreasing the bound of the norm of θ , we get a more misspeciﬁed model class.
To isolate the model learning aspect, we consider the exact RL setting without sampling.
We take a 2 state, 2 action MDP shown in Figure 3 with a discount factor γ = 0 .
9 and a uniform initial distribution ρ 0 .
For every θ , a function ϕ outputs the corresponding Q ∗ via performing the ﬁxed point iteration until convergence.
Q ∗ is transformed into the policy π Q ∗ via softmax with the temperature α = 0 .
We then calculate
J given π Q ∗ in a closed form (Sutton & Barto, 2018).
For OMD, we obtain the gradient of J with respect to θ using the expression (5).
We then apply the projected gradient ascent where after each step we make a projection on a space of bounded parameters via
[Equation]
Finding an MLE solution corresponds to minimizing the average KL divergence
D KL ( p || p θ ) =
[Equation]
S| · |A| (cid:88) s,a,s (cid:48) p ( s [Equation] [Equation]
[Equation]
a ) p
[Equation]
s
[Equation]
a ) for optimizing the dynamics p θ and minimizing the squared error for the reward r θ .
We similarly perform the projected gradient descent and call the agent MLE (even though we do not use the samples for estimation).
The resulting J as a function of the norm bound κ is shown in Figure 2.
When the true model is not representable by a chosen class, OMD learns a model that uses its representational capacity for helping the agent to maximize the expected returns, while the MLE agent tries to predict the next states and rewards accurately while discarding the true objective function the agent seeks to optimize.
In MDPs with high-dimensional state spaces (Bellemare et al., 2013; Beattie et al., 2016) where the underlying dynamics are complex, having a model that will accurately predict the next observation might be expensive and unnecessary for decision making.
Figure 2 reﬂects the problems
[Figure]
an MLE-based model will face for such environments and provides evidence for using control-oriented models that leverage the available capacity of the model effectively.
5. Theoretical Analysis
In the previous section, we have empirically demonstrated that OMD outperforms Dyna-style (Sutton, 1991)
MBRL agents when the model capacity is limited.
This section characterizes the set of optimal solutions of OMD and compares the Q ∗ approximation bounds for OMD and MLE agents.
5.1. Optimal Solutions for OMD
We use the principle of value equivalence for
MBRL (Grimm et al., 2020) and argue that value equivalent models are optimal solutions to (2) and (6).
Deﬁnition 1 (Optimal value equivalence) .
Let Q ∗ be an optimal action-value function for the unconstrained RL problem.
The models with parameters θ and θ (cid:48) are Q ∗ - equivalent if
[Equation]
The deﬁnition is a slight modiﬁcation of the value equivalence used in (Grimm et al., 2020): instead of requiring the Bellman operators to be equal for a set of value functions and policies, we require the equality for a chosen Q ∗ only.
The subset of models that are Q ∗ -equivalent forms an equivalence class Θ Q ∗ .
Proposition 1.
If we let the soft Bellman operator (3) temperature α → 0 and let θ be any model parameters from
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure]
[Figure] [Figure] [Figure] the equivalence class Θ Q ∗ , then ( Q ∗ , θ ) is a solution for (2) and (6) .
This property holds by construction.
The optimal Q-function maximizes the objective in the true MDP.
As the log-sumexp temperature in (3) approaches 0, we recover the “hard” target in the Bellman optimality operator:
[Equation]
Thus, if we set θ to the true model, Q ∗ will satisfy the
Bellman equation Q ∗ ( s, a ) =
B θ Q ∗ ( s, a ) .
But even though the true model belongs to the equivalence class Θ Q ∗ , it is not identiﬁable : all models from Θ Q ∗ are going to be indistinguishable for OMD.
Seemingly undesirable at ﬁrst glance, it allows OMD choosing any model that induces the same Bellman operator, which is beneﬁcial under the model misspeciﬁcation as shown in Section 4.2.
We provide an example of a model that is Q ∗ -equivalent with the true model in Figure 3.
The model differs signiﬁcantly, demonstrating that the equivalence class Θ Q ∗ consists of multiple elements.
Moreover, the dynamics learned by OMD are deterministic, suggesting that OMD can choose a simpler model that will have the same Q ∗ as the true model.
Drawing the connection to the prior work on state abstractions (Li et al., 2006), the fact that MDPs have the same optimal action values indicates that the learned models can be seen as Q ∗ -irrelevant with respect to a state abstraction over S .
5.2. Approximation Bound
Our next result relates approximation errors for the optimal
Q-functions under the OMD and MLE models.
For sim-
[Figure]
plicity, we analyze the setting with α → 0 and the Bellman error (6) as the objective.
Theorem 2.
( Q ∗ approximation error)
Let Q ∗ be the optimal action-value function for the true MDP.
Let ˆ Q OMD and ˆ
Q MLE be the ﬁxed
• If the MLE dynamics ˆ p and reward ˆ r have the bounded errors
[Equation]
and max s,a | r ( s, a ) − ˆ r ( s, a ) | = (cid:15) r , and the reward function is bounded r
( s, a ) ∈ [0 , r max ] ∀ s, a , we have
[Equation]
• If the Bellman optimality operator induced by the OMD model ˆ θ has the bounded error
[Equation]
, we have
[Equation]
We prove the bounds in Appendix H using similar arguments as the proof of the simulation lemma (Kearns &
Singh, 2002).
The MLE bound has
a
[Equation] 1 − γ [Equation] 2 term making the bound loose compared to the OMD bound with only the 11 − γ term.
The bound suggests that OMD approximation error translates into a lower Q ∗ approximation error than for the MLE model.
Figure 4 compares empirically the errors and the tightness of the bounds for a tabular MDP where Q ∗ can be computed exactly.
The result provides evidence that OMD indeed achieves a lower Q ∗ approximation error compared to an agent that seeks to estimate p and r accurately.
Motivated by the theoretical ﬁndings, the next section discusses a practical version of OMD.
6. OMD with Function Approximation
Section 4 describes optimal model design, a non-likelihoodbased method for learning models in tabular MDPs.
In this section, we propose several approximations to make OMD practically applicable.
We analyze the effect of the approximations and perform an ablation study in Appendix F.
Q-network.
We use a neural network with parameters w to approximate the Q-values.
The network is trained to minimize the Bellman error induced by the model θ :
L ( θ, w ) (cid:44) E s,a [ Q w ( s, a ) − B θ Q ¯ w ( s, a )] 2 → min w , (11) where ¯ w is a target copy of parameters w updated using exponential moving average, a standard practice to increase the stability of deep Q-learning (Mnih et al., 2015).
We also use double Q-learning (Hasselt, 2010; Fujimoto et al., 2018) but omit it from the equations for simplicity.
To estimate the expectation, we use a replay buffer (Mnih et al., 2015).
Constraint.
The constraint in (2) and (6) should be satisﬁed for all state-action pairs making it impractical for non-tabular MDPs.
We introduce an alternative but similar constraint, the ﬁrst-order optimality condition for minimizing the Bellman error (11):
[Equation]
We note that the 0 is vector-valued and has the same dimensionality as w .
Implicit differentiation.
The process of training θ is bilevel: in the inner loop, we optimize the Q-function parameters to get optimal w ∗ corresponding to a ﬁxed model θ ; in the outer loop, we make a gradient update of θ .
We make K steps of an optimization method to approximate w ∗ =
ϕ ( θ ) where K is a hyperparameter and reuse the weights from the previous outer loop iterations.
We follow Rajeswaran et al.
and approximate the inverse Jacobian term in ∂ϕ ( θ )
∂θ with the identity matrix.
Surprisingly, we did not observe beneﬁts when using the inverse Jacobian term.
We investigate the phenomenon deeper and discuss possible explanations in Appendix C.
Objective.
We consider the problem (6) and use the Bellman error as the outer loop objective:
[Equation]
where B , again, is a soft Bellman operator induced by the true reward r and dynamics p :
[Equation]
Note that the objective (13) is used for estimating the gradient with respect to θ only and w is trained to optimize Algorithm
Input: Initial parameters w , θ , empty replay buffer D .
repeat
Set s to be the current state.
Sample an action a using softmax over Q w ( s, a ) .
Apply a to get r = r
( s, a ) , s (cid:48)
∼ p ( s (cid:48) | s, a ) .
Append ( s, a, s (cid:48) , r ) to buffer D .
for i = 1 to K do
Sample ( s, a ) from buffer D .
Apply θ to get r = r θ ( s, a ) , s (cid:48) ∼ p θ ( s (cid:48) | s, a ) .
Update Q w parameters w to minimize L ( θ, w ) .
end for
Update model parameters θ according to (15).
until the maximum number of interactions is reached
L ( θ, w ) .
While both L true and J objectives could be used for training θ , we found that the latter requires more samples to converge.
Note that optimizing the L true still corresponds to maximizing the (entropy-regularized) expected returns.
Resulting gradient.
The changes above in the objective function and the constraint yield the following optimization problem:
[Equation]
The Q function and IFT approximations and result in the following gradient with respect to the model parameters:
[Equation]
The OMD algorithm is summarised in Algorithm 1.
The only difference between Dyna-based approaches and OMD (highlighted in blue) is given by the gradient used to train model parameters θ .
7. Experiments with Function Approximation
This section aims to test the following hypotheses:
• The OMD agent with approximations from Sectionachieves near-optimal returns.
• The performance of OMD is better compared to MLE under the model misspeciﬁcation.
• The parameters θ of the OMD model have low likelihood, yet the agent that acts using the Q-function trained with the model achieves near-optimal returns in the true MDP.
[Figure]
Setup .
We provide full details about the experimental setup and hyperparameters in Appendix D. We choose CartPole (Barto et al., 1983) to have controllable experiments but also include results on MuJoCo
HalfCheetah (Todorov et al., 2012) with similar ﬁndings in Appendix G further supporting our conclusions.
Since OMD learns one of the
Q ∗ -equivalent models as shown in Section 5.1, a close nonMLE baseline would be the algorithm used in the value equivalence principle (VEP) paper (Grimm et al., 2020).
The VEP model minimizes the difference between the Bellman operators:
[Equation]
where B θπ V ( s
[Equation]
= E [Equation]
B π is the real model counterpart estimated from samples, and Π and V are predeﬁned sets of policies and value functions.
Performance under model misspeciﬁcation .
We design two experiments that allow measuring the misspeciﬁcation in isolation.
First, we limit the model class representational capacity by controlling the number of units in hidden layers of the model.
Next, we add distracting states by sampling noise from a standard gaussian and vary the number of distractors.
Figure 5 and Figure 6 show the returns achieved by the agents after training in the two regimes.
Note that the
Q-function is updated using only the next states and rewards produced by the model, and even when the hidden dimensionality of the model is 1, the OMD model encodes useful information for taking optimal actions.
Returns achieved by
OMD are also more robust to the distractors indicating that the MLE focuses on predicting the parts of a state that might not be relevant for decision making.
The relatively poor performance of VEP suggests that learning a model to predict
[Figure]
values for a ﬁxed set of policies and value functions is not as effective, especially if some states are non-informative.
The experiments reﬂect the challenges an MBRL agent will face in complex domains such as (Bellemare et al., 2013; Beattie et al., 2016;
Kalashnikov et al., 2021): accurately predicting the next observations can be infeasible because the underlying dynamics can be too involved and there might be few components that are important for taking action.
Figuresand 6 provide evidence that using control-oriented methods would allow using the model capacity more effectively.
Likelihood of OMD model.
We show the mean squared error (MSE) of OMD and MLE dynamics predictions in
Figure 7.
Quantitatively, the MSE of OMD models is higher than the MSE of a randomly initialized model , while OMD achieves higher returns than MLE.
This ﬁnding suggests that the dynamics might not need to produce predictions close to the true states to be useful for planning.
Qualitatively, we visualize individual coordinates of the predictions under the model misspeciﬁcation in Figure 8.
We predict the immediate next states given a sequence of states and actions from an optimal trajectory.
We use one of the runs of the OMD agent with the hidden size 1 that achieves optimal returns of 500 and observe that Cart Velocity predictions are nearly constant.
On the other hand, an MLE model with the hidden size 1 spends its limited capacity to predict the ﬂuctuations in Cart Velocity leading to a signiﬁcant deterioration in obtained returns.
Overall, these ﬁndings suggest that the OMD agent achieves near-optimal returns, performs better than the MLE-based
MBRL agent as well as VEP under model misspeciﬁcation, and learns a model that is useful for control despite having low likelihood.
[Figure]
8. Discussion and Future Work
An exciting direction for future work is the extension of OMD to environments with image-based observations where model misspeciﬁcation naturally arises.
We expect that for complex visual domains, learning a control-oriented model should be more effective compared to model-based methods that rely on reconstruction (Kaiser et al., 2019;
Hafner et al., 2020).
Based on Figure 8, we also conjecture that OMD can learn an abstract model that can ignore parts of the original state space that are irrelevant for control.
This might allow applying OMD in a zero-shot manner for transfer learning tasks where the underlying dynamics remain unchanged.
Implicit differentiation is not the only way to solve the described constrained optimization problems.
Other alternatives include using Lagrangian methods as proposed for the tabular case in (Bacon et al., 2019).
Since extending the approach to non-tabular MDPs would require an additional approximator for Lagrange multipliers, we conjecture that ﬁnding a saddle point is going to be less stable than using the IFT.
Finally, it is worth theoretically studying the sensitivity of the IFT to the approximations to the inverse Jacobian term and the inner loop solutions.
Our ﬁndings, as well as ﬁndings of (Rajeswaran et al., 2019; 2020; Lorraine et al., 2020) suggest that there is a gap between the assumptions of the IFT and its applicability in practice.
9. Conclusion
The paper proposes optimal model design (OMD), a method for learning control-oriented models that addresses the short-
[Figure]
comings of likelihood-based MBRL approaches.
OMD optimizes the expected returns in an end-to-end manner and alleviates the objective mismatch of standard MBRL methods that train models using a proxy of the true RL objective.
Theoretically, we characterize the set of optimal solutions to
OMD and illustrate the efﬁcacy of OMD over MLE agents for approximating optimal value functions.
Empirically, we introduce approximations to apply OMD to non-tabular environments and demonstrate the improved performance of OMD in settings with limited model capacity.
Perhaps surprisingly, we ﬁnd that the OMD model can have low likelihood, yet the model is useful for maximizing returns.
Overall, OMD sheds light on the potential of control-oriented models for model-based reinforcement learning. A CKNOWLEDGMENTS
EN thanks Iurii Kemaev and Clement Gehring for invaluable help with JAX; Tristan Deleu, Gauthier Gidel, Amy Zhang, Aravind Rajeswaran, Ilya Kostrikov, Brandon Amos, and Aaron Courville for insightful discussions; Pierluca D’Oro, David Brandfonbrener, Valentin Thomas, and Timur Garipov for providing useful suggestions on the early draft of the paper; Compute Canada for providing computational resources.
We acknowledge the Python community (Van Rossum &
Drake Jr, 1995; Oliphant, 2007) for developing the core set of tools that enabled this work, including JAX (Bradbury et al., 2018), Jupyter (Kluyver et al., 2016), Matplotlib (Hunter, 2007), numpy (Oliphant, 2006; Van Der Walt et al., 2011), pandas (McKinney, 2012), and SciPy (Jones et al., 2014).
A. Proof of IFT
The implicit function theorem is a well known result discussed in, for example, (Krantz & Parks, 2012).
Theorem 1.
(Cauchy,
Implicit
Function)
Let f : Θ × W
→ W be a continuously differentiable function and (˜ θ, ˜ w ) be a point satisfying f (˜ θ, ˜ w ) = 0 .
If the Jacobian ∂f (˜ θ, ˜ w ) ∂w is invertible, then there exists an open set U ⊆
Θ containing ˜ θ and a unique continuously differentiable function ϕ such that ϕ
(˜ θ )
=
˜ w and f ( θ, ϕ ( θ ))
=for all θ ∈ U .
Moreover, ∂ϕ ( θ )
[Equation]
Proof.
By the assumption, we have
[Equation]
Taking the total derivative of f with respect to θ , we have
[Equation]
Rearranging the terms and using the invertibility of the
Jacobian, we get
[Equation]
[Equation]
[Equation]
The implicit function theorem (IFT) allows differentiating outputs of black-box functions
ϕ ( θ ) that do not have a closed (differentiable) form.
For example, in Section 4.2, function ϕ takes as input the parameters of the model and applies ﬁxed-point iteration until convergence to ﬁnd the optimal value function for the given parameters.
This contrasts implicit differentiation to other meta-learning algorithms like MAML (Finn et al., 2017) that differentiate through the iterations of the inner loop procedure.
B. Implicit Differentiation in JAX
We provide an implementation of implicit differentiation in
JAX (Babuschkin et al., 2020; Heek et al., 2020) by adapting the code from the library for ﬁnding ﬁxed points (Gehring et al., 2019).
The implementation requires a solver that takes parameters θ and an initial w 0 as input and outputs ϕ ( θ ) =
w ∗ such that f ( θ, w ∗ )
= 0 .
Then, we deﬁne a custom vector-Jacobian product that allows using ϕ ( θ ) as a block in a differentiable computational graph.
We highlight the importance of having an implementation without explicitly forming the matrices in (7) which is crucial for large-scale applications.
The implementation allows using both the version with the inverse Jacobian term as well as with the identity approximation as discussed in Section 6.
) vdp = dp_vjp(-vdsoli[0])[0] return jnp.zeros_like(sol), vdp root_solve.defvjp(fwd, rev) sol =
root_solve(f, w0, p, solver)
# solver returns sol: f(p, sol) =
# sol is differentiable w.r.t.
p
C. Sensitivity to IFT Approximations
The IFT provides a way to calculate the derivatives of a black-box function ϕ .
The expression (7) is valid under the assumptions that
the inner loop solution w ∗ satisﬁes the equation f
( θ, w ∗ )
=
exactly; 2.
the Jacobian term ∂f ( θ,w ∗ ) ∂w is inverted accurately.
Ensuring both of the conditions can be challenging for largescale applications.
This appendix analyzes the sensitivity to the conditions in a controlled manner for the 2-state MDP from Figure 3, while Appendix F studies the sensitivity for the function approximation case.
At every outer loop iteration, we calculate the exact w ∗ , add gaussian noise with standard deviation σ , and observe the effect on the expected returns J after the convergence of θ .
Figure 9 demonstrates the results for the exact outer loop gradient ∂ϕ ( θ )
∂θ as well as the gradient using the identity approximation of the inverse
Jacobian term.
Surprisingly, we did not observe signiﬁcant beneﬁts of using the Jacobian ∂f ( θ,w ∗ ) ∂w .
We conjecture that the inverse Jacobian acts like a preconditioner (Boyd et al., 2004) on ∂ϕ ( θ )
∂θ and the preconditioner is useful in our setting only near the exact inner loop solutions w ∗ .
We leave the theoretical investigation of the IFT sensitivity as future work and refer the reader to Lorraine et al. (2020) for a discussion about approximations of the Jacobian term.
[Figure] Std 0.
E x p e c t e d R e t u r n
With Inverse Jacobian
With Identity Approximation
D. Experimental Details
In Section 7, we use CartPole (Barto et al., 1983), an environment with 2 actions, 4-dimensional continuous state space, and optimal returns of 500.
We train the agents for 200000 environment steps.
The temperature α is 0.01.
We sample from the replay buffer with a mini-batch size of 256.
The discount factor γ is 0.99.
At each time step during training, the agent chooses a random action with a probability of 0.1 for exploration.
We have a separate copy of the environment where we evaluate the agent and take the average over 10 runs to estimate the returns.
We run each experiment using 10 random seeds.
We set the number of Q-function updates K equal to 1.
For the results with K = 3 and K = 10 , see Figure 11.
We highlight that after each outer loop step, weights w are warm-started using the last iterate of the previous inner loop (instead of randomly initializing w and training from scratch).
We use Adam optimizer with the learning rate 10 − 3 for updating θ and perform a hyperparameter sweep over the learning rate for w in { 3 · 10 − 4 , 10 − 3 , 3 · 10 − 3 } .
We make a sweep over the moving average coefﬁcient τ for the target network ¯ w in { 0 . 005 , 0 . 01 } .
Both of the parameters control how fast the Q-network parameters are updated relatively to the model parameters.
Since the CartPole environment is non-stochastic, we use a deterministic dynamics model.
All networks have two hidden layers and ReLU activations (Nair & Hinton, 2010).
For both hidden layers in all networks, we set the dimensionality to 32.
In the experiment with the limited model class capacity, we vary the hidden dimensionality in { 1 , 2 , 3 , 4 , 6 , 12 } for the dynamics and [Figure] reward networks to measure how the limitation affects the agent’s performance.
In the experiment with the distractors, we vary the number of gaussians in { 2 4 , 2 5 , 2 6 , 2 7 , 2 8 , 2 9 } to measure how the uninformative state components affect the returns.
For the value equivalence principle (VEP) baseline, we have followed the experimental setup from the original paper (Grimm et al., 2020).
Perhaps surprisingly, the authors ﬁnd that it sufﬁces to use a set of all deterministic stateindependent policies as Π and 5 random value functions as V for CartPole (see Appendix A.2.3 in (Grimm et al., 2020)).
We have used CPU-only nodes of the internal cluster.
Each experiment requires 10 seeds × 3 algorithms × 2 τ ’s ×hidden sizes / 6 numbers of distractors × 3 LRs resulting in 2160 total jobs.
E. Returns under Model Misspeciﬁcation
In Section 7, we demonstrate that OMD models have a higher mean squared error (MSE) than predictions of a randomly initialized model (Figure 7).
We provide the corresponding returns of the OMD and MLE agents in Figure 10.
While the MLE model has more accurate predictions, OMD achieves higher returns using the next states signiﬁcantly deviating from the real next states.
The result suggests that likelihood optimization may be an unnecessary step for
MBRL algorithms.
[Figure]
F. Ablation Study
Section 6 introduces a series of approximations to scale the
OMD algorithm to non-tabular environments.
We analyze the effect of the approximations by varying the number of inner loop steps K , using the inverse Jacobian term, and using a single Q-function estimator (without double
Q-learning).
Figure 11 shows the (undiscounted) returns for K varying in { 1 , 3 , 10 } .
We did not observe signiﬁcant changes in performance for different values of K for CartPole.
The result suggests that as long as the Q-network update speed (which is also controlled by the target update coefﬁcient τ and learning rates) stays aligned with the model update speed, adding more inner loop steps is not necessary.
Figure 12 demonstrates the returns without using double Q learning and without using the identity approximation of the inverse jacobian term (cid:16) ∂f ( θ,w ∗ ) ∂w (cid:17) − 1 .
We observed that Qfunctions trained with OMD can be prone to overestimation of Q-values showing that double Q-learning is important for
OMD.
We conjecture that training a model that maximizes the returns can amplify the overestimation bias (Hasselt, 2010) caused by using (soft) maximized sampled targets.
Surprisingly, we did not observe signiﬁcant beneﬁts from using the inverse Jacobian (cid:16) ∂f ( θ,w ∗ ) ∂w (cid:17) − 1 .
We conjecture that there are two reasons explaining the phenomenon.
First, the modiﬁed constraint in (14) forces the gradient of L ( θ, w ) to be zero, implying that the Jacobian is in fact the Hessian matrix ∂ 2 L ( θ,w ∗ ) ∂w.
Dauphin et al.
; Sagun et al.
(2017) observed that Hessians of neural networks tend to be sin- [Figure]
With Inv [Figure] gular. Since a system of linear equations with a singular matrix has multiple solutions, it is up to a linear algebra solver to choose the solution.
One of the alternatives would be a min-norm solution corresponding to the Moore-Penrose pseudoinverse and the solution might not be providing a useful inductive bias for the learning process of θ .
Second, the Jacobian term could be useful only in proximity to the exact inner loop solution w ∗ .
Since the practical algorithm performs only K inner loop steps and does not reach the exact w ∗ , the curvature information provided by the Jacobian might not be beneﬁcial for training θ .
Finally, we tried to use the gradient constraint in (14) in the tabular setting.
We got similar results as with the constraint on Q-values in (2) suggesting that the two constraints have similar effects on the learning process.
Overall, the ablation study provides evidence that OMD is robust to the choice of the number of inner loop steps and the IFT approximations, while double Q-learning is the only important algorithmic modiﬁcation.
G. Results on HalfCheetah
We provide an additional comparison of OMD and MLE agents under the model misspeciﬁcation on MuJoCo
HalfCheetah (Todorov et al., 2012).
For both agents, the inner optimizer is Soft Actor-Critic (SAC) (Haarnoja et al., 2018a;b) with the default conﬁguration.
OMD trains the model using 15.
The MLE agent trains the model with MSE effectively becoming the MBPO algorithm (Janner et al., 2019) without having an ensemble of models and learning [Figure] the variance of the predictions.
We perform a hyperparameter sweep over the model learning rate in { 10 − 4 , 3 · 10 − 4 } , over the SAC networks learning rate in { 10 − 4 , 3 · 10 − 4 } , and over K in { 1 , 3 } .
Model hidden size is 64 for the experiment with distractors.
Figures 13 and 14 summarize the results.
Similarly to the observations on the tabular and CartPole environments, the experiments provide evidence that OMD should be preferred over the likelihood-based agent in the model misspeciﬁcation setup.
H. Proof of OMD Bounds
Section 5.2 discusses the bounds on Q ∗ approximation error obtained by the MLE and OMD agents.
We ﬁrst prove a lemma relating the error of the model approximation and the
Bellman operator approximation.
We then prove a theorem giving a bound on Q ∗ error.
For simplicity, we focus on the case with the Bellman error in the true MDP (6) as the objective function and “hard” versions of the Bellman optimality operators which are obtained by taking the limit of the log-sum-exp temperature α → 0 .
Note that results for MLE hold for any agent that approximates the reward and dynamics functions, but we call the agent MLE since it is a common choice for model parameters estimation.
Notation.
We denote p ( ·| s, a ) and Q ( · , a ) as vectors of transition probabilities and Q-values for all states in S .
The MLE model is given by (ˆ p, ˆ r ) and the corresponding
Bellman optimality operator is denoted as ˆ BQ .
To have a distinction between OMD and MLE, we denote OMD parameters as ˆ θ and the corresponding operator as B ˆ θ Q . [Figure] [Figure] [Figure]
[Figure]
f R
[Figure]
(cid:107)
[Equation]
sup
[Equation]
| f
[Equation]
x ) | is the inﬁnity norm of a function f .is a vector of an appropriate size with ones as entries.
Lemma 1.
(Bellman operator error bound)
Let Q be an action-value function.
If the dynamics ˆ p and the reward ˆ r have the bounded errors max s,a (cid:107) p ( ·| s, a ) − ˆ p ( ·| s, a ) (cid:107)
=
(cid:15) p and max s,a | r ( s, a ) − ˆ r ( s, a ) |
= (cid:15) r , and the reward function is bounded r ( s, a ) ∈ [0 , r max ] ∀ s, a , we have
[Equation]
Proof.
Using the derivations similar to the proof of the simulation lemma (Jiang, 2018), we obtain for any ( s, a ) [Equation] a [Equation]
a
[Equation]
) (cid:33) − (cid:32) ˆ
[Equation] a [Equation] s (cid:48)
[Equation]
a [Equation] a (cid:48) [Equation]
a
[Equation]
) (cid:33)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12) ≤ [Equation] a
[Equation]
= (cid:15) r + ∵ p and ˆ p are distributions γ (cid:12)(cid:12)(cid:12)(cid:12)(cid:12)(cid:88) s (cid:48) ( p ( s (cid:48) [Equation] a ) − ˆ p (
[Equation]
|
[Equation]
a ))
[Equation]
max a (cid:48) Q ( s (cid:48) , a (cid:48) )
− r max 2(1 − γ ) (cid:19)(cid:12)(cid:12)(cid:12)(cid:12)(cid:12)
[Equation]
inequality
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
Since the inequalities hold for all state-action pairs, we can take the maximum over ( s, a ) and obtain
[Equation]
Theorem 2.
( Q ∗ approximation error)
Let Q ∗ be the optimal action-value function for the true MDP.
Let ˆ Q OMD and ˆ Q MLE be the ﬁxed points of the Bellman optimality operators for approximate OMD and MLE models respectively.
• If the MLE dynamics ˆ p and reward ˆ r have the bounded errors max s,a (cid:107) p {+( ·| s, a ) − ˆ p ( ·| s, a ) (cid:107) 
+}
=
(cid:15) p and max s,a | r ( s, a ) − ˆ r ( s, a ) |
= (cid:15) r , and the reward function is bounded r ( s, a ) ∈ [0 , r max ] ∀ s, a , we have
[Equation]
• If the Bellman optimality operator induced by the OMD model ˆ θ has the bounded error
[Equation]
, we have
[Equation]
Proof.
(OMD)
For all state-action pairs ( s, a ) we get (cid:12) (cid:12)(cid:12) Q ∗ ( s, a )
[Equation]
ˆ Q ( s, a ) (cid:12) (cid:12)(cid:12) = (cid:12)(cid:12)(cid:12)
[Equation]
∵ Fixed point = [Equation] ≤ (cid:12)(cid:12)(cid:12) [Equation]
[Table] [Equation] [Table] [Equation]
[Table]
[Equation] [Table] [Equation] [Table]
[Equation] [Table] [Equation] [Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation] [Table] [Equation] [Table]
[Equation] [Table] [Equation] [Table]
[Equation]
[Table]
[Equation]
Taking the maximum over ( s
[Equation]
a )
[Equation]
we get the recursion:
[Equation]
(MLE)
The proof for MLE can be obtained using the same derivations and additionally using the result of the lemma bounding the difference between the Bellman operators:
[Equation]
The last inequality demonstrates that the OMD bound is tighter.
OMD model directly optimizes (cid:12)(cid:12)(cid:12)
[Equation]
, while
MLE minimizes (cid:15) r and (cid:15) p that only upper bound (cid:12)(cid:12)(cid:12)
[Equation]
as suggested by the lemma.
Hence, given the same budget of representational capacity, OMD will learn a model that is more helpful for approximating the optimal Q-function.
Finally, Figure 4 empirically supports our theory showing that both the Q ∗ approximation error and the error-bound gap are smaller for OMD.