Year: 2021

Abstract:
Model-free reinforcement learning algorithms can compute policy gradients given sampled environment transitions, but require large amounts of data. In contrast, model-based methods can use the learned model to generate new data, but model errors and bias can render learning unstable or sub-optimal. In this paper, we present a novel method that combines real world data and a learned model in order to get the best of both worlds. The core idea is to exploit the real world data for onpolicy predictions and use the learned model only to generalize to different actions. Specifically, we use the data as time-dependent on-policy correction terms on top of a learned model, to retain the ability to generate data without accumulating errors over long prediction horizons. We motivate this method theoretically and show that it counteracts an error term for model-based policy improvement. Experiments on MuJoCoand PyBullet-benchmarks show that our method can drastically improve existing model-based approaches without introducing additional tuning parameters.

heading: 1 INTRODUCTION
text: Model-free reinforcement learning (RL) has made great advancements in diverse domains such as single- and multi-agent game playing (Mnih et al., 2015; Silver et al., 2016; Vinyals et al., 2019), robotics (Kalashnikov et al., 2018) and neural architecture search (Zoph & Le, 2016). All of these model-free approaches rely on large numbers of interactions with the environment to ensure successful learning. While this issue is less severe for environments that can easily be simulated, it limits the applicability of model-free RL to (real-world) domains where data is scarce.
Model-based RL (MBRL) aims to reduce the amount of data required for policy optimization by learning a model that approximates the true environment, which we can use to generate cheap, simulated state transitions (Sutton, 1990; Racanière et al., 2017; Moerland et al., 2020). While early approaches on low-dimensional tasks by Schneider (1997); Deisenroth & Rasmussen (2011) used probabilistic models with closed-form posteriors, recent methods rely on neural networks to scale to more complex tasks on discrete (Kaiser et al., 2019) and continuous (Chua et al., 2018; Kurutach et al., 2018) action spaces. However, the learned representation of the true environment always remains imperfect and as such introduces an error source to the RL problem (Atkeson & Santamaria, 1997; Abbeel et al., 2006). Hence, a key challenge in MBRL is model-bias; small errors in the learned models that can compound over multi-step predictions and lead to lower asymptotic performance in contrast to model-free approaches.
To address these challenges with both model-free and model-based RL, Levine & Koltun (2013); Chebotar et al. (2017) propose to combine the merits of both. While there are multiple possibilities to combine the two methodologies, in this work we focus on improving the model’s predictive state distribution such that it more closely resembles the data distribution of the true environment.
Contributions The main contribution of this paper is on-policy corrections (OPC), a novel hyperparameter-free methodology that uses on-policy transition data on top of a separately learned model to enable accurate long-term predictions for MBRL. A key strength of our approach is that it does not introduce any new parameters that need to be hand-tuned for specific tasks. We theoretically motivate our approach by means of a policy improvement bound and show that we can recover the true state distribution when generating trajectories on-policy with the model. We illustrate how OPC improves the quality of policy gradient estimates in a simple toy example and evaluate it on various continuous control tasks from the MuJoCo control suite and their PyBullet variants. There, we
demonstrate that OPC improves current state-of-the-art MBRL algorithms in terms of data efficiency, especially for the more difficult PyBullet environments.
Related Work To counteract model-bias, several approaches combine ideas from model-free and model-based RL. For example, Levine & Koltun (2013) guide a model-free algorithm via modelbased planning towards promising regions in the state space, Kalweit & Boedecker (2017) augment the training data by an adaptive ratio of simulated transitions, Talvitie (2017) use ‘hallucinated’ transition tuples from simulated to observed states to self-correct the model, and Feinberg et al. (2018); Buckman et al. (2018) use a learned model to improve the value function estimate. Janner et al. (2019) mitigate the issue of compounding errors for long-term predictions by simulating short trajectores start from real states. Cheng et al. (2019) extend first-order model-free algorithms via adversarial online learning to leverage prediction models in a regret-optimal manner. Clavera et al. (2020) employ a model to augment an actor-critic objective, which they optimize via back-propagation through time, adapting the planning horizon interpolates between a purely model-based and a modelfree approach. Morgan et al. (2021) combine actor-critic methods with model-predictive rollouts to guarantee near-optimal simulated data and retain exploration on the real environment. A downside of most existing approaches is that they introduce additional hyperparameters that are critical to the learning performance (Zhang et al., 2021).
In addition to empirical performance, recent work builds on the theoretical guarantees for model-free approaches by Kakade & Langford (2002); Schulman et al. (2015) to provide guarantees for MBRL. Luo et al. (2018) provide a general framework to show monotonic improvement towards a local optimum of the value function, while Janner et al. (2019) present a lower-bound on performance for different rollout schemes and horizon lengths. Yu et al. (2020) show guaranteed improvement in the offline MBRL setting by augmenting the reward with an uncertainty penalty, while Clavera et al. (2020) present improvement guarantees in terms of the model’s and value function’s gradient errors.
A related field to MBRL that also combines models with data is iterative learning control (ILC) (Bristow et al., 2006). The differences between the two methods are the following: RL typically focuses on finding parametric feedback policies for general reward functions, whereas ILC seeks an open-loop sequence of actions with fixed length to improve state tracking performance. Further, the model in ILC is often derived from first principles and then kept fixed, whereas in MBRL the model is continuously improved upon observing new data. The method most closely related to RL and our approach is optimization-based ILC (Owens & Hätönen, 2005; Schöllig & D’Andrea, 2009), in which a linear dynamics model is used to guide the search for optimal actions. Recently, Baumgärtner & Diehl (2020) extended the ILC setting to nonlinear dynamics and more general reward signals. Little work is available that draws connections between RL and ILC (Zhang et al., 2019) with one notable exception: Abbeel et al. (2006) use the observed data from the last rollout to account for a mismatch in the dynamics model. The limitations of this approach are that deterministic dynamics are assumed, the policy optimization itself requires a line search procedure with rollouts on the true environment and that it was not combined with model learning. We build on this idea and extend it to the stochastic setting of MBRL by making use of recent advances in RL and model learning.
Lastly, Harutyunyan et al. (2016) propose a similar correction term as the one introduced in this paper in the context of off-policy policy evaluation and correct the state-action value function instead of the transition dynamics. Similarly, Fonteneau et al. (2013) consider the problem of off-policy policy evaluation but in the batch RL setting and propose to generate ‘artificial’ trajectories from observed transitions instead of using an explicit model for the dynamics.

heading: 2 PROBLEM STATEMENT AND BACKGROUND
text: We consider the Markov decision process (MDP) (S,A, p, r, γ, ρ) where S ⊆ RdS and A ⊆ RdA are the continuous state and action spaces, respectively. The unknown environment dynamics are described by the transition probability p(st+1 | st,at), an initial state distribution ρ(s0) and the reward signal r(s,a). The goal in RL is to find a policy πθ(at | st) parameterized by θ that maximizes the expected return, discounted by γ ∈ [0, 1] over episodes of length T ,
η = E s0:T ,a0:T [∑T t=0 γtr(st,at) ] , st+1 ∼ p(st+1 | st,at), s0 ∼ ρ, at ∼ πθ(at | st). (1)
Algorithm 1 General Model-based Reinforcement Learning 1: for n = 1, . . . do 2: for b = 1, . . . , B do 3: Rollout policy πθn on environment and store transitions Dbn = {(ŝn,bt , ân,bt , ŝn,bt+1)}T−1t=0 4: p̃n(st+1 | st,at): Learn a global dynamics model given all data D1:n = ⋃n i=1 ⋃B b=1Dbi
5: θn+1 = θn + α∇η̃n: Optimize the policy based on the model p̃ with any RL algorithm
The expectation is taken with respect to the trajectory under the stochastic policy πθ starting from some stochastic initial state. Direct maximization of Eq. (1) is challenging, since we do not know the environment’s transition probabilities p. In MBRL, we learn a model for the transitions and reward function from data, p̃(st+1 | st,at) ≈ p(st+1 | st,at) and r̃ respectively, and subsequently maximize
η̃ = E s0:T ,a0:T [∑T t=0 γtr̃(st,at) ] , st+1 ∼ p̃(st+1 | st,at), s0 ∼ ρ̃, at ∼ πθ(at | st). (2)
as a surrogate problem for Eq. (1). The choice of initial state distribution ρ̃(s0) depends on the optimization method: For model predictive control, Chua et al. (2018) use the current state to optimize the next action based on a rollout, while Janner et al. (2019) use the empirical state distribution observed from the true environment. In the following, we focus on the latter and for ease of exposition assume a known reward function r̃ = r, even though we learn it jointly with p̃ in our experiments.
We let ηn denote the policy’s performance under the policy πn = πθn at iteration n and use ŝ and â for states and actions that are observed on the true environment. Algorithm 1 summarizes the overall procedure for MBRL: At each iteration n we store B on-policy trajectories Dbn = {(ŝn,bt , ân,bt , ŝn,bt+1)}T−1t=0 obtained by rolling out the current policy πn on the real environment in Line 3. Lastly, we approximate the environment with a learned model p̃n(st+1 | st,at) based on the data in Line 4, and optimize the policy based on the proxy objective in Eq. (2) in Line 5. Note that the policy optimization algorithm can be off-policy and employ its own, separate replay buffer.
Model Choices The choice of model p̃ plays a key role, since it defines the surrogate problem Eq. (2). Inside the optimization loop, the model is typically used to predict sequences of transitions τ with length H . One option is to use a purely data-based model that replays past transitions from a recent on-policy trajectory b starting at time step t,
p̃datan (τt:t+H | t, b) = t+H∏ t′=t p̃datan (st′ | t′, b) with p̃datan (st | t, b) = δ(st − ŝn,bt ), (3)
where δ(·) denotes a Dirac-delta distribution and 1 ≤ b ≤ B. This model choice effectively reduces model-based algorithms back to the model-free setting. An alternative is to summarize all the past dataD1:n = ∪ni=1∪Bb=1Dbi in a parametric model coming from a distribution family P by maximizing the marginal log-likelihood L,
p̃modeln (st+1 | st,at) = arg max p̃∈P ∑ (ŝt,ât,ŝt+1)∈D1:n L(ŝt+1; ŝt, ât), (4)
which for each state-action pair (st,at) models a distribution over the next state st+1, such that p̃modeln (τt:t+H | t, b) = δ(st−ŝn,bt ) ∏t+H−1 t′=t p̃ model n (st′+1 | st′ ,at′) with at′ ∼ πθ. For both models, we can sample H-step predictions by first sampling a trajectory index b ∼ Categorical({0, . . . , B}) and then rolling out following the respective state sequence distributions p̃(τt:t+H | t, b).

heading: 3 ON-POLICY CORRECTIONS
text: In this section, we analyze how the choice of model impacts policy improvement, develop OPC as a model that can eliminate one term in the improvement bound, and analyze its properties. In the following, we drop the n sub- and superscript when the iteration is clear from context.

heading: 3.1 POLICY IMPROVEMENT
text: Independent of whether we use a the data directly Eq. (3) or summarize it in a world model Eq. (4), our goal is to find an optimal policy that maximizes Eq. (1) via the corresponding model-based proxy
objective Eq. (2). To this end, we would like to know how a policy improvement η̃n+1 − η̃n ≥ 0 based on the model p̃, which is what we optimize in Eq. (2), relates to the true gain in performance ηn+1 − ηn on the environment with unknown transitions p. While the two are equal without model errors, p̃ = p, in general the larger the model error, the worse we can expect our proxy objective to be. More concretely, we show in Appendix A.1 that the policy improvement can be decomposed as
ηn+1 − ηn︸ ︷︷ ︸ True policy improvement ≥ η̃n+1 − η̃n︸ ︷︷ ︸ Model policy improvement − |ηn+1 − η̃n+1|︸ ︷︷ ︸ Off-policy model error − |ηn − η̃n|︸ ︷︷ ︸ On-policy model error , (5)
where a performance improvement on our model-based objective in Eq. (2) only translates to a gain in Eq. (1) if two error terms are sufficiently small. These terms depend on how well the performance estimate based on our model, η̃, matches the true performance, η. If the reward function is known, this term only depends on the model quality of p̃ relative to p. Note that in contrast to the result by Janner et al. (2019), Eq. (5) is a bound on the policy improvement instead of a lower bound on ηn+1.
The first error term compares ηn+1 and η̃n+1, the performance estimation gap under the optimized policy πn+1 that we obtain in Line 5 of Algorithm 1. Since at this point we have only collected data with πn in Line 3, this term depends on the generalization properties of our model to new data; what we call the off-policy model error. For our data-based model p̃data that just replays data under πn independently of the action, this term can be bounded for stochastic policies. For example, Schulman et al. (2015) bound it by the average KL-divergence between πn and πn+1. For learned models p̃model, it depends on the generalization properties of the model (Luo et al., 2018; Yu et al., 2020). While understanding model generalization better is an interesting research direction, we will assume that our learned model is able to generalize to new actions in the following sections.
While the first term hinges on model-generalization, the second term is the on-policy model error, i.e., the deviation between ηn and η̃n under the current policy πn. This error term goes to zero for p̃data as we use more on-policy data B → ∞, since the transition data are sampled from the true environment and the mean converges, c.f., Appendix A.2. While the learned model is also trained with on-policy data, small errors in our model compound as we iteratively predict ahead in time. Consequently, for prediction horizons shorter than the effective horizon encoded in the discount factor, H < γ/(1 − γ), the on-policy error term grows as O(min(H/(1 − γ), H2)), c.f., (Janner et al., 2019) and Appendix A.3. The root cause is the same as in behavior cloning (Ross & Bagnell, 2010, Theorem 2.1), where small errors aggregate over time.

heading: 3.2 COMBINING LEARNED MODELS AND REPLAY BUFFER
text: The key insight of this paper is that the replay buffer in Eq. (3) and the parametric model in Eq. (4) have opposing strengths: The replay buffer has low error on-policy, but high error off-policy since it replays transitions from past data, i.e., they are independent of the actions chosen under the new policy. In contrast, the learned model can generalize to new actions by extrapolating from the data and thus has lower error off-policy, but errors compound over the prediction horizon H .
An ideal model would combine the model-free and model-based approaches in a way such that it retains the unbiasedness of on-policy generated data, but also generalizes to new policies via the model. To this end, we propose to adjust the observed transitions by an offset that accounts for the model predictions given a new state-action pair. In particular, we use the model’s mean prediction f̃(st,at) = Est+1 [p̃modeln (st+1 | st,at)] to construct the joint model
p̃opcn (st+1 | st,at, b) = δ ( st+1 − ŝn,bt+1 )︸ ︷︷ ︸ p̃datan (st+1 | t+ 1, b) ∗ δ ( st+1 − [ f̃(st,at)− f̃(ŝn,bt , ân,bt ) ]) ,︸ ︷︷ ︸ Model mean correction to generalize to st, at
(6)
where ∗ denotes the convolution of the two distributions and b refers to a specific rollout stored in the replay buffer that was observed in the true environment. Given a batch-index b, p̃opc in Eq. (6) transitions deterministically according to st+1 = ŝ n,b t+1 + f̃(st,at) − f̃(ŝn,bt , ân,bt ), resembling the equations in ILC (c.f., Baumgärtner & Diehl (2020) and Appendix D). If we rollout p̃opc along a trajectory, starting from a state ŝn,bt and apply the recorded actions from the replay buffer, â n,b t , the correction term on the right of Eq. (6) cancels out and we have p̃opc(st+1 | ŝn,bt , ân,bt , b) = p̃data(st+1 | t + 1, b) = δ(st+1 − ŝn,bt+1). Thus OPC retrieves the true on-policy data distribution independent of the prediction quality of the model, which is why we refer to this method as on-policy
corrections (OPC). This behaviour is illustrated in Fig. 1a, where even though the model (blue) is biased on-policy, OPC corrects the model’s prediction to match the true data. In Fig. 1b, we show how this affects predicted rollouts on a simple stochastic double-integrator environment. Even though small on-policy errors in p̃model (blue) compound over time, the corresponding p̃opc matches the ground-truth environment data closely. Note that even though the model in Eq. (6) is deterministic, we retain the environment’s stochasticity from the data in the transitions to ŝt+1, so that we recover the on-policy aleatoric uncertainty (noise) from sampling different reference trajectories via indexes b.
When our actions at are different from â n,b t , p̃ opc still uses the data from the environment’s transitions, but the correction term in Eq. (6) uses the learned model to predict how the next state changes in expectation relative to the prediction under ân,bt . That is, in Fig. 1a for a new at the model predicts the state distribution shown in red. Correspondingly, we shift the static prediction ŝn,bt+1 by the difference in means (gray arrow) between the two predictions; i.e., the change in trajectory by changing from ân,bt to at. Since we shift the model predictions by a time-dependent but constant offset, this does not recover the true state distribution unless the model has zero error. However, regardless of this it can help with long-term predictions in Fig. 1b by shifting the model off-policy predictions (red) to the OPC predictions (green), which are closer to the environment’s state distribution under the new policy.

heading: 3.3 THEORETICAL ANALYSIS
text: As mentioned in the previous sections, OPC is motivated to reduce the on-policy model error from Eq. (5) to tighten the improvement bound. In this section, we analyze the on-policy performance gap from a theoretical perspective and show that with OPC this error can be reduced independently from the learned model’s error. To this end, we derive a more general form of the predictive model for OPC where we consider ŝn,bt+1 as a random variable following the true environment’s transition dynamics instead of an observed quantity. While impossible to implement in practice, this formulation helps in the theoretical analysis. The generalized OPC-model is defined as
p̃opc? (st+1 | st,at, b) = p(ŝt+1 | ŝbt , âbt)︸ ︷︷ ︸ On-policy transition
∗ δ ( st+1 − [ f̃(st,at)− f̃(ŝbt , âbt) ]) ︸ ︷︷ ︸
Mean correction term
, (7)
which highlights that it transitions according to the true on-policy dynamics conditioned on data from the replay buffer, combined with a correction term. We provide a detailed derivation for the generalized model in Appendix A (Lemma 4. With Eq. (7), we have the following result: Theorem 1. Let η̃opc? and η be the expected return under the generalized OPC-model Eq. (7) and the true environment, respectively. Assume that the learned model’s mean transition function f̃(st,at) = E[p̃model(st+1 | st,at)] is Lf -Lipschitz and the reward r(st,at) is Lr-Lipschitz. Further,
if the policy π(at | st) is Lπ-Lipschitz with respect to st under the Wasserstein distance and its (co-)variance Var[π(at | st)] = Σπ(st) ∈ SdA+ is finite over the complete state space, i.e., maxst∈S trace{Σπ(st)} ≤ σ̄2π , then with C1 = √ 2(1 + L2π)LfLr and C2 = √ L2f + L 2 π
|η − η̃opc? | ≤ C1 1− γ dA 1 4CT2 √ T σ̄π. (8)
We provide a proof in Appendix A.4. From Theorem 1, we can observe the key property of OPC: for deterministic policies, the on-policy model error from Eq. (5) is zero and independent of the learned models’ predictive distribution p̃model, so that η = η̃opc? . For policies with non-zero variance, the bound scales exponentially with T , highlighting the problem of compounding errors. In this case, as in the off-policy case, the model quality determines how well we can generalize to different actions. To alleviate this issue, one could extend Theorem 1 similarly to the results in Janner et al. (2019), such that the rollouts are only of length H T . In practice, p̃opc? can never be realized as it requires the true (unknown) state distribution. However, we can approximate it by p̃opc Eq. (6) using an ever-increasing size of reference trajectories such that the on-policy error converges to zero in probability.
Lemma 1. LetM be a MDP with dynamics p(st+1 | st,at) and reward r < rmax. Let M̃ be another MDP with dynamics p̃model 6= p. Assume a deterministic policy π : S 7→ A and a set of trajectories D = ⋃Bb=1{(ŝbt , âbt , ŝbt+1)}T−1t=0 collected from M under π. If we use OPC Eq. (6) with data D, then
lim B→∞
Pr (|η − η̃opc| > ε) = 0 ∀ε > 0 with convergence rate O(1/B). (9)
We provide a proof in Appendix A.4. Lemma 1 states that given sufficient reference on-policy data, the performance gap due to model errors becomes arbitrarily small for any model p̃model when using OPC. While the assumption for infinite on-policy data in Lemma 1 is unrealistic for practical applications, we found empirically that OPC drastically reduces the on-policy model error even when the assumptions are violated. In our implementation, we use stochastic policies as well as trajectories from previous policies, i.e., off-policy data, for the corrections in OPC (see also Section 4.2).

heading: 3.4 DISCUSSION
text: Epistemic uncertainty So far we have only considered aleatoric uncertainty (noise) in our transition models. In practice, modern methods additionally distinguish epistemic uncertainty that arises from having seen limited data (Deisenroth & Rasmussen, 2011; Chua et al., 2018). This leads to a distribution (or an ensemble) of models, where each sample could explain the data. In this setting, we apply OPC by correcting each sample individually. This allows us to retain epistemic uncertainty estimates after applying OPC, while the epistemic uncertainty is zero on-policy.
Limitations Since OPC uses on-policy corrections, it is inherently limited to local policy optimization where the policy changes slowly over time. As a consequence, it is not suitable for global exploration schemes like the one proposed by Curi et al. (2020). Similarly, since OPC uses the observed data and corrects it only with the expected learned model, p̃opc always uses the on-policy transition noise (aleatoric uncertainty) from the data, even if the model has learned to represent it. While not having to learn a representation for aleatoric uncertainty can be a strength, it limits our approach to environments where the aleatoric uncertainty does not vary significantly with state/actions. It is possible to extend the method to the heteroscedastic noise setting under additional assumptions that enable distinguishing model error from transition noise (Schöllig & D’Andrea, 2009). Lastly, our method applies directly only to MDPs, since we rely on state-observations. Extending the ideas to partially observed environments is an interesting direction for future research.

heading: 4 EXPERIMENTAL RESULTS
text: We begin the experimental section with a motivating example on a toy problem to highlight the impact of OPC on the policy gradient estimates in the presence of model errors. The remainder of the section focuses on comparative evaluations and ablation studies on complex continuous control tasks.

heading: 4.1 ILLUSTRATIVE EXAMPLE
text: In Section 3.3, we investigate the influence of the model error directly on the expected return from a theoretical perspective. From a practical standpoint, another relevant question is how the policy optimization and the respective policy gradients are influenced by model errors. For general environments and reward signals, this question is difficult to answer, due to the typically highdimensional state/action spaces and large number of parameters governing the dynamics model as well as the policy. To shed light on this open question, we resort to a simple low-dimensional example and investigate how OPC improves the gradient estimates under a misspecified dynamics model.
In particular, we assume a one-dimensional deterministic environment with linear transitions and a linear policy. The benefits of this example are that we can 1) compute the true policy gradient based on a single rollout (determinism), 2) determine the environment’s closed-loop stability under the policy (linearity), and 3) visualize the gradient error as a function of all relevant parameters (low dimensionality). The dynamics and initial state distribution are specified by p(st+1 | st,at) = δ(Ast + Bat | st,at) with ρ(s0) = δ(s0) where A,B ∈ R and δ(·) denotes the Dirac-delta distribution. We define a deterministic linear policy πθ(at | st) = δ(θst | st) that is parameterized by the scalar θ ∈ R. The objective is to drive the state to zero, which we encode with an exponential reward r(st,at) = exp { −(st/σr)2 } . Further, we assume an approximate dynamics model
p̃(st+1 | st,at) = δ((A + ∆A)st + (B + ∆B)at | st,at), (10) where ∆A,∆B quantify the mismatch between the approximate model and the true environment. In practice, the mismatch can arise due to noise-corrupted observations of the true state or, in the case of stochastic environments, due to a finite amount of training data.
With the setting outlined above, we investigate how model errors influence the estimation of policy gradients. To this end, we roll out different policies under models with varying degrees of error ∆B. For each policy/model combination, we compute the model-based policy gradient and compare it to the true gradient. The results are summarized in Fig. 2, where the background’s opacity depicts the gradient error’s magnitude and its color indicates whether the respective signs of the gradients are the same (blue, ≥ 0) or differ (red, < 0). For policy optimization, the sign of the gradient estimate is paramount. However, we see in the left-hand image that even small model errors can lead to the wrong sign of the gradient. OPC significantly reduces the magnitude of the gradient error and increases the robustness towards model errors. See also Appendix B for a more in-depth analysis.

heading: 4.2 EVALUATION ON CONTINUOUS CONTROL TASKS
text: In the following section, we investigate the impact of OPC on a range of continuous control tasks and compare it to the current state-of-the-art model-based and model-free RL algorithms MBPO (Janner et al., 2019) and SAC (Haarnoja et al., 2018). Further, we aim to answer the important question about how data diversity affects MBRL, and OPC in particular. While having a model that can generate (theoretically) an infinite amount of data is intriguing, the the benefit of having more data is limited
by the model quality in terms of being representative of the true environment. For OPC, the following questions arise from this consideration: Do longer rollouts help to generate better data? And is there a limit to the value of simulated transition data, i.e., is more always better?
For the dynamics model p̃model, we follow Chua et al. (2018); Janner et al. (2019) and use a probabilistic ensemble of neural networks, where each head predicts a Gaussian distribution over the next state and reward. For policy optimization, we employ the soft actor critic (SAC) algorithm by Haarnoja et al. (2018). All learning curves are presented in terms of the median (lines) and interquartile range (shaded region) across ten independent experiments, where we smooth the evaluation return with a moving average filter to accentuate the results of particularly noisy environments. Apart from small variations in the hyperparameters, the only difference between OPC and MBPO is that our method uses p̃opc, while MBPO uses p̃model. Generally, we found that OPC was more robust to the choice of hyperparameters. The rollout horizon to generate training data is set to H = 10 for all experiments. Note that when using p̃data to generate data, we retain the standard (model-free) SAC algorithm.
Our implementation is based upon the code from Janner et al. (2019). We made the following changes to the original implementation: 1) The policy is only updated at the end of an epoch, not during rollouts on the true environment. 2) The replay buffer retains data for a fixed number of episodes, to clearly distinguish on- and off-policy data. 3) For policy optimization, MBPO uses a small number of environment transitions in addition to those from the model. We found that this design choice did not consistently improve performance and added an additional level of complexity. Therefore, we refrain from mixing environment and model transitions and only use simulated data for policy optimization. While we stay true to the key ideas of MBPO under these changes, we denote our variant as MBPO(?) to avoid ambiguity. See Appendices C.6 and C.7 for a comparison to the original MBPO algorithm.
Comparative Evaluation We begin our analysis with a comparison of our method to MBPO(?) and SAC on four continuous control benchmark tasks from the MuJoCo control suite (Todorov et al., 2012) and their respective PyBullet variants (Ellenberger, 2018–2019). The results are presented in Fig. 3. We see that the difference in performance between both methods is only marginal when evaluated on the MuJoCo environments (Fig. 3, top row). Notably, the situation changes drastically for the PyBullet environments (Fig. 3, bottom row). Here, MBPO(?) exhibits little to no learning progress, whereas OPC succeeds at learning a good policy with few interactions in the environment. One of the main differences between the environments (apart from the physics engine itself) is that the PyBullet variants have initial state distributions with significantly larger variance.
Influence of State Representation In general, the success of RL algorithms should be agnostic to the way an environment represents its state. In robotics, joint angles ϑ are often re-parameterized by a sine/cosine transformation, ϑ 7→ [sin(ϑ), cos(ϑ)]. We show that even for the simple CartPole environment, the parameterization of the pole’s angle has a large influence on the performance of MBRL.
of simulated transitions N and vary the rollout lengths H = {1( ), 5( ), 10( )}. The results indicate that multi-step rollouts generate more diverse data, which is required for efficient learning.
In particular, we compare OPC and MBPO(?) on the RoboSchool and PyBullet variants of the CartPole environment, which represent the pole’s angle with and without the sine/cosine transformation. The results are shown in Fig. 4. To rule out other effects than the angle’s representation, we repeat the experiment for each implementation but transform the state to the other representation, respectively. Notably, OPC successfully learns a policy irrespective of the state’s representation, whereas MBPO(?) fails if the angle of the pole is represented by the sine/cosine transformation.
Influence of Data Diversity Here, we investigate whether multi-step predictions are in fact beneficial compared to single-step predictions during the data generation process. To this end, we keep the total number of simulated transitions N for training constant, but choose different horizon lengths H = {1, 5, 10}. The corresponding numbers of simulated rollouts are then given by nrollout = N/H . The results for N = {20, 40, 100, 200} × 103 on the HalfCheetah environment are shown in Fig. 5. First, note that more data leads to a higher asymptotic return, but after a certain point more data only leads to diminishing returns. Further, the results indicate that one-step predictions are not enough to generate sufficiently diverse data. Note that this result contradicts the findings by Janner et al. (2019) that one-step predictions are often sufficient for MBPO.

heading: 5 CONCLUSION
text: In this paper, we have introduced on-policy corrections (OPC), a novel method that combines observed transition data with model-based predictions to mitigate model-bias in MBRL. In particular, we extend a replay buffer with a learned model to account for state-action pairs that have not been observed on the real environment. This approach enables the generation of more realistic transition data to more closely match the true state distribution, which was further motivated theoretically by a tightened improvement bound on the expected return. Empirically, we demonstrated superior performance on high-dimensional continuous control tasks as well as robustness towards state representations.

heading: SUPPLEMENTARY MATERIAL
text: In the appendix we provide additional details on our method, ablation studies, and the detailed hyperparameter configurations used in the paper. An overview is shown below.

heading: Table of Contents

heading: A Theoretical Analysis of On-Policy Corrections 14
text: A.1 General Policy Improvement Bound . . . . . . . . . . . . . . . . . . . . . . . . 14 A.2 Properties of the Replay Buffer . . . . . . . . . . . . . . . . . . . . . . . . . . 14 A.3 Properties of the Learned Model . . . . . . . . . . . . . . . . . . . . . . . . . . 15 A.4 Properties of OPC . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 15 A.5 Model Errors in OPC . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 23

heading: B Motivating Example – In-depth Analysis 24
text: B.1 Setup . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 24 B.2 Reward Landscapes . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 25 B.3 Influence of Model Error . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 25 B.4 Influence of Off-Policy Error . . . . . . . . . . . . . . . . . . . . . . . . . . . 26 B.5 Additional Information . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 27

heading: C Additional Experimental Results 29
text: C.1 Comparison with Other Baseline Algorithms . . . . . . . . . . . . . . . . . . . 29 C.2 Ablation - Retain Epochs . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 30 C.3 Influence of State Representation: In-Depth Analysis . . . . . . . . . . . . . . . 30 C.4 Improvement Bound: Empirical Investigation . . . . . . . . . . . . . . . . . . . 32 C.5 Off-Policy Analysis . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 33 C.6 Implementation Changes to Original MBPO . . . . . . . . . . . . . . . . . . . . 34 C.7 Results for Original MBPO Implementation . . . . . . . . . . . . . . . . . . . . 35

heading: D Connection between MBRL and ILC 38
text: D.1 Norm-optimal ILC . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 38 D.2 Model-based RL to Norm-optimal ILC . . . . . . . . . . . . . . . . . . . . . . 39 D.3 Extensions . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 40

heading: E Hyperparameter Settings 40
text: F Implementation and Computational Resources 41

heading: A THEORETICAL ANALYSIS OF ON-POLICY CORRECTIONS
text: In this section, we analyze OPC from a theoretical perspective and how it affects policy improvement.
Notation In the following we drop the n superscript for states and actions for ease of exposition. That is, ŝn,b = ŝb and ân,b = âb.

heading: A.1 GENERAL POLICY IMPROVEMENT BOUND
text: We begin by deriving inequality Eq. (5), which serves as motivation for OPC and is the foundation for the theoretical analysis. Our goal is to bound the difference in expected return for the policies before and after the policy optimization step, i.e., ηn+1 − ηn. Since we are considering the MBRL setting, it is natural to express the improvement bound in terms of the expected return under the model Eq. (2) and thus obtain the following
ηn+1 − ηn = ηn+1 − ηn + η̃n+1 − η̃n+1 + η̃n − η̃n = η̃n+1 − η̃n + ηn+1 − η̃n+1 + η̃n − ηn ηn+1 − ηn︸ ︷︷ ︸ True policy
improvement ≥ η̃n+1 − η̃n︸ ︷︷ ︸ Model policy improvement − |ηn+1 − η̃n+1|︸ ︷︷ ︸ Off-policy model error − |η̃n − ηn|︸ ︷︷ ︸ On-policy model error .
According to this bound, the improvement of the policy under the true environment is governed by the three terms on the LHS:
• Model policy improvement: This term is what we are directly optimizing in MBRL Eq. (2) offset by the return of the previous iteration η̃n, which is constant given the current policy πn. Assuming that we are not at an optimum, standard policy optimization algorithms guarantee that this term is non-negative.
• Off-policy model error: The last term is the difference in return for the true environment and model under the improved policy πn+1. This depends largely on the generalization properties of our model, since it is not trained on data under πn+1.
• On-policy model error: This term compares the on-policy return under πn between the true environment and the model and it is zero for any model p̃ = p. Since we have access to transitions from the true environment under the πn, the replay buffer Eq. (3) fulfills this condition under certain circumstances and the on-policy model error vanishes, see Lemma 2.
Note that the learned model Eq. (4) is able to generalize to unseen state-action pairs better than the replay buffer Eq. (3) and accordingly will achieve a lower off-policy model error. The motivation behind OPC is therefore to combine the best of the learned model and the replay buffer to reduce both on- and off-policy model errors.

heading: A.2 PROPERTIES OF THE REPLAY BUFFER
text: The benefit of the replay buffer Eq. (3) is that it can never introduce any model-bias such that any trajectory sampled from this model is guaranteed to come from the true state distribution. Accordingly, if we have collected sufficient data under the same policy, the on-policy model error vanishes. Lemma 2. Let M be the true MDP with (stochastic) dynamics p, bounded reward r < rmax and let p̃data be the transition model for the replay buffer Eq. (3). Further, consider a set of trajectories D = ⋃Bb=1{(ŝbt , âbt , ŝbt+1)}T−1t=0 collected from M under policy π. If we collect more and more on-policy training data under the same policy, then
lim B→∞
Pr ( |η − η̃replay| > ε ) = 0 ∀ε > 0
where η̃replay is the expected return under the replay buffer using the collected trajectories D.
Proof. First, note that the corresponding expected return for the replay-buffer model is given by
η̃replay = 1
B B∑ b=1 T−1∑ t=0 r(ŝbt , â b t),
which is a sample-based approximation of the true reward η. By the weak law of large numbers (see e.g., Blitzstein & Hwang (2019, Theorem 10.2.2)), the Lemma then holds.

heading: A.3 PROPERTIES OF THE LEARNED MODEL
text: Following Janner et al. (2019, Lemma B.3), a general bound on the performance gap between two MDPs with different dynamics can be given by
|η1 − η2| ≤ 2rmax m H∑ t=1 tγt. (11)
where m ≥ maxt Es∼pt1(s) KL(p1(s ′,a)||p2(s′,a)) bounds the mismatch between the respective transition models. Now, the final form of Eq. (11) depends on the horizon length H . For H →∞, we obtain the original result form Janner et al. (2019) with ∑ t≥1 tγ
t = γ/(1−γ)2. For the finite horizon case and H being smaller than the effective horizon encoded in the discount factor H < γ/(1− γ), the following bound holds:
H∑ t=1 tγt ≤ min {H(H − 1) 2 , H 1− γ } , (12)
which can be easily verified by upper bounding t ≤ H to obtain H/(1− γ) or by bounding γ ≤ 1 to obtain O(H2). Note that this bound is vacuous for deterministic policies, since the KL divergence between two distributions with non-overlapping support is infinite. In the following we focus on the Wasserstein metric under the euclidean distance.

heading: A.4 PROPERTIES OF OPC
text: In this section, we analyze the properties of OPC relative to the true, unknown environment’s transition distribution p and a learned representation p̃. In general, the OPC-model mixes observed transitions from the environment with the learned model. The resulting transitions are then a combination of the mean transitions from the learned model, the aleatoric noise from the data (the environment), and the mean-error between our learned model and the environment.

heading: A.4.1 PROOF OF LEMMA 1
text: In this section we proof Lemma 1 by showing that the OPC-model coincides with the replay-buffer Eq. (3) in the case of a deterministic policy and thus lead to the same expected return η̃.
Lemma 3. Let M be the true MDP with (stochastic) dynamics p and let M̃ be a MDP with the same reward function r and initial state distribution ρ0, but different dynamics p̃model, respectively. Further, assume a deterministic policy π : S 7→ A and a set of trajectories D = ⋃Bb=1{(ŝbt , âbt , ŝbt+1)}T−1t=0 collected from M under π. If we extend the approximate dynamics p̃model by OPC with data D, then
η̃replay = ηopc,
where η̃replay and η̃opc are the model-based returns following models Eqs. (3) and (6), respectively.
Proof. For the proof, it suffices to show that the resulting state distributions of the two transition models p̃data and p̃opc under the deterministic policy π are the same for all b with 1 ≤ b ≤ B. We show this by induction:
For t = 0 the states are sampled from the initial state distribution, thus we have s0 = ŝb0 by definition. Assume that st = ŝbt as an induction hypothesis. Then
p̃opc(st+1 | st, π(st), b) = δ ( st+1 − ŝbt+1 ) ∗ δ ( st+1 − [ f(st, π(st))− f(ŝbt , âbt) ]) = δ ( st+1 − [ ŝbt+1 + f(st, π(st))− f(ŝbt , âbt)
]) = δ ( st+1 − [ ŝbt+1 + f(ŝ b t , â b t)− f(ŝbt , âbt)
]) = δ ( st+1 − ŝbt+1
) = p̃data(st+1 | t+ 1, b)
where the second step follows by the induction hypothesis and due to the deterministic policy. Thus, for any index b we have τopcb = τb and the result follows.
Now, combining Lemmas 2 and 3 proofs the result in Lemma 1.

heading: A.4.2 PROOF OF THEOREM 1
text: Remark on Notation In the main paper, we unify the notation for state sequence probabilities of the different models Eqs. (3), (4) and (6) as p̃x(τt:t+H | t, b) with x ∈ {replay,model, opc}. This allows for a consistent description of the respective rollouts independent of the model being a learned representation, a replay buffer or the OPC-model. For that notation, the index b denotes the sampled trajectory from the collected data on the real environment. Implicitly, we therefore condition the state sequence probability on the observed transition tuples [(ŝbt+1, ŝ b t , â b t), . . . , (ŝ b t+H , ŝ b t+H−1, â b t+H−1)], i.e.,
p̃x(τt:t+H | t, b) = p̃x(τt:t+H | (ŝbt+1, ŝbt , âbt), . . . , (ŝbt+H , ŝbt+H−1, âbt+H−1)), (13) where we omit the explicit conditioning for the sake of brevity in the main paper. Similarly, we can write the one-step model Eq. (6) for OPC in an explicit form as
p̃opc(st+1|st,at, b) = p̃opc(st+1|st,at, ŝbt+1, ŝbt , âbt). (14) Note that with the explicit notation, the relation between the OPC-model Eq. (6) and generalized OPCmodel Eq. (7) becomes clear:
p̃opc? (st+1 | st,at, b) = p̃opc(st+1 | st,at, ŝbt , âbt) = ∫ p̃opc(st+1 | st,at, ŝbt+1ŝbt , âbt) dŝbt+1. (15)
For the following proofs, we will stay with the explicit notation for sake of clarity and instead omit the conditioning on b.
Generalized OPC-Model In this section, we have a closer look at the generalized OPC-model Eq. (7). The main difference between Eqs. (6) and (7) is that the former is in fact transitioning deterministically (the stochasticity arises from the environment’s aleatoric uncertainty which manifests itself in the observed transitions). The two models can be related via marginalization of ŝbt+1, see Eq. (15). The resulting generalized OPC-model can then be related to the true transition distribution according to the following Lemma.
Lemma 4. For all ŝbt , âbt ∈ S ×A with ŝbt+1 ∼ p(ŝt+1 | ŝbt , âbt) it holds that
p̃opc(st+1 | st,at, ŝbt , âbt) = p ( st+1 − [ f̃(st,at)− f̃(ŝbt , âbt) ]︸ ︷︷ ︸ Mean correction | st,at, ŝbt , âbt ) , (16)
where f̃(st,at) = E[p̃model(st+1 | st,at)] is the mean transition of the learned model Eq. (4) and p̃opc(st+1 | st,at, ŝbt , âbt) denotes the OPC-model if we marginalize over the distribution for ŝbt+1 instead of using its observed value.
Proof. Using the explicit notation Eq. (14), the OPC-model from Eq. (6) is defined as
p̃opc(st+1 | st,at, b) = p̃opc(st+1 | st,at, ŝbt , âbt , ŝbt+1) (17) = δ ( st+1 − ŝbt+1 ) ∗ δ ( st+1 − [ f̃(st,at)− f̃(ŝbt , âbt) ] ) , (18)
= δ ( st+1 − [ ŝbt+1 + f̃(st,at)− f̃(ŝbt , âbt) ] ) . (19)
With ŝbt+1 ∼ p(ŝt+1 | ŝbt , âbt), marginalizing ŝt+1 yields (note that we use ŝt+1 instead of st+1 to denote the random variable for the next state in order to distinguish it from the random variable for p̃opc(st+1 | st,at, b) under the integral)
p̃opc(st+1 | st,at, ŝbt , âbt) = ∫ δ ( st+1 − [ ŝt+1 + f̃(st,at)− f̃(ŝbt , âbt) ] ) p(ŝt+1 | ŝbt , âbt) dŝt+1,
= p ( st+1 − [ f̃(st,at)− f̃(ŝbt , âbt) ] | ŝbt , âbt ) .
Remark 1. An alternative way of writing the general OPC-model is the following,
p̃opc(st+1 | st,at, ŝbt , âbt) = p(ŝt+1 | ŝbt , âbt)︸ ︷︷ ︸ On-policy transition
∗ δ ( st+1 − [ f̃(st,at)− f̃(ŝbt , âbt) ]) ︸ ︷︷ ︸
Mean correction term
, (20)
which highlights that the p̃opc transitions according to the true on-policy dynamics conditioned on data from the replay buffer, combined with a correction term. We can further explicitly see why an implementation of this model wouldn’t be possible due to its dependency on the true transition probabilities. Thus, in practice, we’re limited to the sample-based approximation shown in the paper.
The fundamental idea for the proof of Theorem 1 lies in the following Lemma 5, which is the foundation for bounding the on-policy error. The Wasserstein distance naturally arises in bounding this type of error model as it depends on the expected return under two different distributions. The final result is then summarized in Theorem 1.
Lemma 5. Let p̃opc be the generalized OPC-model (cf. Lemma 4) and η̃opc be its corresponding expected return. Assume that the return r(st,at) is Lr-Lipschitz and the policy π(at | st) is Lπ-Lipschitz with respect to st under the Wasserstein distance, then
|η − η̃opc| ≤ Lr √ 1 + L2π ∑ t≥0 γtW2(p(st), p̃opc(st)) (21)

heading: Proof.
text: |η − η̃opc| = ∣∣∣ E τ∼p [∑ t≥0 γtr(ŝt, ât) ] − E τ∼p̃opc [∑ t≥0 γtr(st,at) ]∣∣∣ (22)
= ∣∣∣∑ t≥0 γt ( E ŝt,ât∼p(ŝt,ât) [ r(ŝt, ât) ] − E st,at∼p̃opc(st,at) [ r(st,at) ])∣∣∣ (23) ≤ ∑ t≥0 γt ∣∣∣ E ŝt,ât∼p(ŝt,ât) [ r(ŝt, ât) ] − E st,at∼p̃opc(st,at) [ r(st,at)
]∣∣∣ (24) Applying Lemma 8:
≤ ∑ t≥0 γtLrW2(p(ŝt, ât), p̃opc(st,at)) (25)
Writing the joint distributions for state/action in terms of their conditional (i.e., policy) and marginal distributions p(st,at) = π(at | st)p(st):
= ∑ t≥0 γtLrW2(π(ât | ŝt)p(ŝt), π(at | st)p̃opc(st)) (26)
Under the assumption that the policies π are Lπ-Lipschitz under the Wasserstein distance, application of Lemma 10 concludes the proof:
≤ ∑ t≥0 γtLr √ 1 + L2πW2(p(ŝt), p̃opc(st)). (27)
Lemma 6 (Wasserstein Distance between Marginal State Distributions). Let p̃opc(st) and p(ŝt) be the marginal state distributions at time t when rolling out from the same initial state ŝ0 under the same policy with the OPC-model and the true environment, respectively. Assume that the underlying learned dynamics model is Lf -Lipschitz continuous with respect to both its arguments and the policy π(a | s) is Lπ-Lipschitz with respect to s under the Wasserstein distance. If it further holds that the policy’s (co-)variance Var[π(a | s)] = Σπ(s) ∈ SdA+ is finite over the complete state space, i.e., maxs∈S trace{Σπ(s)} ≤ σ̄2π, then the discrepancy between the marginal state distributions of the two models is bounded
W22 (p̃opc(st), p(ŝt)) ≤ 2 √ dAσ̄ 2 πL 2 f t−1∑ t′=0 (L2f + L 2 π) t′ (28)
Proof. We proof the Lemma by induction.
Base Case: t = 1 For the base case, we need to show that starting from the same initial state ŝ0 the following condition holds:
W22 (p̃opc(s1), p(ŝ1)) ≤ 2 √ dAσ̄ 2 πL 2 f
For ease of readability, we define z = (s,a) and use the notation dp(x,y) = p(x,y) dx dy whenever no explicit assumptions are made about the distributions.
W2(p̃opc(s1), p(ŝ1)) (29) =W2 (∫∫ p̃opc(s1 | ẑ0, z0) dp(ẑ0, z0), ∫ p(ŝ1 | ẑ0) dp(ẑ0) ) (30)
Recall that we can write both p̃opc and p as convolution between p and a Dirac delta (see Lemma 4). Together with the identity Eq. (55), the Wasserstein distance for sums of random variables Eq. (54) and notingWp(p ( ), p ( )) = 0:
≤ W2 (∫∫ δ(s1 − [f(ẑ0) + f̃(z0)− f̃(ẑ0)]) dp(ẑ0, z0), ∫ δ(ŝ1 − f(ẑ0)) dp(ẑ0) ) (31)
Squaring and using Lemma 9: W22 (∫∫ δ(s1 − [f(ẑ0) + f̃(z0)− f̃(ẑ0)]) dp(ẑ0)p(z0), ∫ δ(ŝ1 − f(ẑ0)) dp(ẑ0) ) ≤ ∫∫ ‖f(ẑ0) + f̃(z0)− f̃(ẑ0)− f(ẑ0)‖2 dp(ẑ0, z0) (32)
= ∫∫ ‖f̃(z0)− f̃(ẑ0)‖2 dp(ẑ0, z0) (33)
≤ L2f ∫ ‖s0 − ŝ0‖2 + ‖a0 − â0‖2 dp(ŝ0, â0, s0,a0) (34)
We are assuming that the initial states of the trajectory rollouts coincide. The joint state/action distribution can then be written as p(ŝ0, â0, s0,a0) = p(ŝ0)π(â0 | ŝ0)δ(s0 − ŝ0)π(a0 | s0). Integrating with respect to s0 leads to:
= L2f ∫ ‖a0 − â0‖2p(ŝ0)π(â0 | ŝ0)π(a0 | ŝ0) dŝ0 dâ0 da0 (35)
This term describes the mean squared distance between two random actions. Since we condition π on the same state ŝ0, the policy distributions coincide. Define ∆a = a0 − â0,
= L2f Ê s0 [ E ∆a [‖∆a‖2] ]
(36)
= L2f Ê s0
[ trace{Var[∆a] ] } (37)
Now Var[∆a] = Var[π(â0 | ŝ0)] + Var[π(a0 | ŝ0)] = 2 Var[π(a0 | ŝ0)] = 2L2f Ê
s0
[ trace{Var[π(a0 | ŝ0)]} ] (38)
This term is in fact less than Eq. (28) for t = 0, thus proofing the base case. ≤ 2 √ dAσ̄ 2 πL 2 f (39)
Inductive Step We will show that if the hypothesis holds for t then it holds for t + 1 as well. We explicitly write the following intermediate bound such that its application in the proof is more apparent, i.e.,
W22 (p(ŝt), p̃opc(st)) ≤ ∫ ‖f̃(zt−1)− f̃(ẑt−1)‖2 dp(ẑt−1, zt−1) (40)
≤ 2 √ dAσ̄ 2 πL 2 f t−1∑ t′=0 (L2f + L 2 π) t′ , (41)
where the first inequality immediately follows from the same reasoning as in the base case Eq. (29)– Eq. (33).
W22 (p(ŝt+1, p(st+1)) (42) ≤ ∫ ‖f̃(zt)− f̃(ẑt)‖2 dp(ẑt, zt) (43)
≤ L2f ∫ ‖st − ŝt‖2 dp(ẑt, zt) + L2f ∫ ‖at − ât‖2 dp(ẑt, zt) (44)
Applying Lemma 11 to the second integral ≤ (L2f + L2π) ∫ ‖st − ŝt‖2 dp(ẑt, zt) + 2 √ dAL 2 f σ̄ 2 π (45)
We predict along a consistent trajectory, i.e., st = ŝt + f̃(st−1,at−1)− f̃(ŝt−1, ât−1) ≤ (L2f + L2π) ∫ ‖f̃(zt−1)− f̃(ẑt−1)‖2 dp(ẑt−1, zt−1) + 2 √ dAL 2 f σ̄ 2 π (46)
Assume that the hypothesis Eq. (40) holds for t
≤ (L2f + L2π)× 2 √ dAσ̄ 2 πL 2 f t−1∑ t′=0 (L2f + L 2 π) t′ + 2 √ dAL 2 f σ̄ 2 π (47)
= 2 √ dAσ̄ 2 πL 2 f [ 1 + t−1∑ t′=0 (L2f + L 2 π) t′ ]
(48)
= 2 √ dAσ̄ 2 πL 2 f t∑ t′=0 (L2f + L 2 π) t′ (49)
Theorem 1. Let η̃opc? and η be the expected return under the generalized OPC-model Eq. (7) and the true environment, respectively. Assume that the learned model’s mean transition function f̃(st,at) = E[p̃model(st+1 | st,at)] is Lf -Lipschitz and the reward r(st,at) is Lr-Lipschitz. Further, if the policy π(at | st) is Lπ-Lipschitz with respect to st under the Wasserstein distance and its (co-)variance Var[π(at | st)] = Σπ(st) ∈ SdA+ is finite over the complete state space, i.e., maxst∈S trace{Σπ(st)} ≤ σ̄2π , then with C1 = √ 2(1 + L2π)LfLr and C2 = √ L2f + L 2 π
|η − η̃opc? | ≤ C1 1− γ dA 1 4CT2 √ T σ̄π. (8)
Proof. From combining Lemmas 5 and 6 it follows that
|η − η̃opc? | ≤ √ 2 √ dA(1 + L2π)σ̄πLfLr ∑ t≥0 γt √√√√ t∑ t′=0 (L2f + L 2 π) t′ (50)
with the shorthand notations C1 = √ 2(1 + L2π)LfLr and C2 = L 2 f + L 2 π
= C1|A| 1 4 σ̄π T∑ t=0 γt √√√√ t∑ t′=0 Ct ′ 2 (51)
Since t ≤ T , we have that∑tt′=0 Ct′2 ≤ TCT2 ≤ C1|A| 1 4 σ̄πC T 2
2
√ T T∑ t=0 γt (52)
Since T ≤ ∞, we have with the geometric series∑Tt=0 γt ≤ 1/(1− γ) ≤ C1
1− γ |A| 1 4C
T 2 2 √ T σ̄π (53)

heading: A.4.3 DEFINITIONS, HELPFUL IDENTITIES AND SUPPORTING LEMMAS
text: Here we briefly summarize some basic definitions and properties that will be used throughout the following.
• The Wasserstein distance fulfills the properties of a metric: Wp(p1, p3) ≤ Wp(p1, p2) + Wp(p2, p3).
• Wasserstein distance of sums of random variables (see, e.g., Mariucci et al. (2018, Corollary 1) for a proof):
Wp(p1 ∗ · · · ∗ pn, q1 ∗ · · · ∗ qn) ≤ n∑ i=1 Wp(pi, qi) (54)
• For any function g(z, ẑ) we have∫∫ p(st+1 − g(z, ẑ))ν(z, ẑ) dzdẑ = p ∗ ∫∫ δ(st+1 − g(z, ẑ))ν(z, ẑ) dz dẑ (55)
• For any multivariate random variable z1 and z2 with probability distributions p(z1) = p1(x)q(y) and p(z2) = p2(x)q(y), respectively, we have that (Panaretos & Zemel, 2019)
W22 (p1(x)q(y), p2(x)q(y)) =W22 (p1(x), p2(x)). (56)
Further, the following Lemmas are helpful for the proof of Theorem 1.
Lemma 7 (Kantorovich-Rubinstein (cf. Mariucci et al. (2018) Proposition 1.3)). Let X and Y be integrable real random variables. Denote by µ and µ their laws [. . . ]. Then the following characterization of the Wasserstein distance of order 1 holds:
W1(ν, µ) = sup ‖φ‖Lip≤1 E x∼ν(·) [φ(x)]− E y∼µ(·) [φ(y)], (57)
where the supremum is being taken over all φ satisfying the Lipschitz condition |φ(x)−φ(y)| ≤ |x−y|, for all x, y ∈ R. Lemma 8. Let f be Lf -Lipschitz with respect to a metric d. Then
| E x∼ν(·) [f(x)]− E y∼µ(·) [f(y)]| ≤ LfW1(ν, µ) ≤ LfW2(ν, µ) (58)
Proof. The first inequality is a direct consequence of Lemma 7 and the second inequality comes from the well-known fact that if 1 ≤ p ≤ q, thenWp(µ, ν) ≤ Wq(µ, ν) (cf. Mariucci et al. (2018, Lemma 1.2)).
Lemma 9. For any two functions f(s) and g(s) and probability density p(s) that govern the distributions defined by
p1(x1) = ∫ δ(x1 − f(s))p(s) ds and p2(x2) = ∫ δ(x2 − g(s))p(s) ds, (59)
it holds for any q ≥ 1 that Wqq (p1, p2) ≤ ∫ ‖f(s)− g(s)‖qp(s) ds. (60)
Proof. We have
Wqq (p1(x1), p2(x2)) = inf γ∈Γ(p1,p2)
∫∫ ‖ξ1 − ξ2‖qγ(ξ1, ξ2) dξ1 dξ2 (61)
Enforcing the following structure on γ(ξ1, ξ2) reduces the space of possible distributions: γ(ξ1, ξ2) =∫ δ(ξ1 − f(s))δ(ξ2 − g(s))p(s) ds, so that γ(ξ1, ξ2) ∈ Γ(p1, p2) and thus
≤ ∫ ‖ξ1 − ξ2‖q ∫ δ(ξ1 − f(s))δ(ξ2 − g(s))p(s) ds dξ1 dξ2 (62)
Integrating over ξ1 and ξ2 yields
= ∫ ‖f(s)− g(s)‖qp(s) ds (63)
Lemma 10. If the policy π(at | st) is Lπ-Lipschitz with respect to st under the Wasserstein distance, then with p(ŝ, â) = π(â | ŝ)p(ŝ) and p̃(s,a) = π(a | s)p̃(s),
W22 (p(ŝ, â), p̃(s,a)) ≤ (1 + L2π)W22 (p(ŝ), p̃(s)) (64)

heading: Proof.
text: W22 (p(ŝ, â), p̃(s,a)) (65)
= inf γ∈Γ(p(ŝ,â),p̃(s,a))
∫ [ ‖â− a‖2 + ‖ŝ− s‖2 ] γ(â,a, ŝ, s) dâ da dŝ ds (66)
Enforcing the following structure on γ reduces the space of possible distributions: γ(ŝ, s, â,a) = γ(ŝ, s)γ(â,a | ŝ, s) with γ(ŝ, s) ∈ Γ(p(ŝ), p̃(s)) and γ(â,a|ŝ, s) ∈ Γ(π(â | ŝ), π(a | s)).
≤ inf γ(ŝ,s)∈Γ(p(ŝ),p̃(s))
γ(â,a|ŝ,s)∈Γ(π(â|ŝ),π(a|s))
∫ {∫ [ ‖â− a‖2γ(â,a | ŝ, s) dâ da } + ‖ŝ− s‖2 ] γ(ŝ, s) dŝ ds
(67)
Interchange infimum and the integral: Rockafellar (1976, Theorem 3A)
= inf γ(ŝ,s)
∫ { inf
γ(â,a|ŝ,s)
∫ ‖â− a‖2γ(â,a | ŝ, s) dâda + ‖ŝ− s‖2 } γ(ŝ, s) dŝ ds (68)
= inf γ(ŝ,s)∈Γ(p(ŝ),p̃(s))
∫ { W22 (π(â | ŝ), π(a | s)) + ‖ŝ− s‖2 } γ(ŝ, s) dŝds (69)
Using the assumption that the action distribution is Lπ-Lipschitz continuous under the Wasserstein metric with respect to the state s:
≤ inf γ(ŝ,s)∈Γ(p(ŝ),p̃(s))
∫ (1 + L2π)‖ŝ− s‖2]γ(ŝ, s) dŝ ds (70)
= (1 + L2π)W22 (p(ŝ), p̃(s)) (71)
Lemma 11 (Average Squared Euclidean Distance Between Actions). If the policy π(a | s) is LπLipschitz with respect to s under the Wasserstein distance and the policy’s (co-)variance Var[π(a | s)] = Σπ(s) ∈ SdA+ is finite over the complete state space, i.e., maxs∈S trace{Σπ(s)} ≤ σ̄2π , then
E ât∼π(ŝt) at∼π(st)
[ ‖ât − at‖2 ] ≤ L2π‖ŝt − st‖2 + 2 √ dAσ̄ 2 π (72)
Proof. Straightforward application of Corollary 1 and Lemma 13
Lemma 12 (Average Squared Euclidean Distance). Consider two random variables x,y with distributions px, py, mean vectors µx, µy ∈ Rm and covariance matrices Σx,Σy ∈ Sm+ , respectively. Then the average squared Euclidean distance between the two is
E x,y
[ ‖x− y‖2 ] = ‖µx − µy‖2 + trace { Σx + Σy } . (73)
Proof. Define z = x− y with mean µz = µx − µy and variance Σz = Σx + Σy.
E [ ‖z‖2 ] = E [∑ i z2i ]
= ∑ i E [ z2i ]
= ∑ i E[zi]2 + Var[zi]
= µ>z µz + trace { Σz }
= ‖µx − µy‖2 + trace { Σx + Σy } .
Theorem 2 (Gelbrich Bound (from Kuhn et al. (2019))). If ‖ · ‖ is the Euclidean norm, and the distributions px and py have mean vectors µx, µy ∈ Rm and covariance matrices Σx,Σy ∈ Sm+ , respectively, then
W2(px, py) ≥ √ ‖µx − µy‖2 + trace { Σx + Σy − 2 ( Σ 1/2 x ΣyΣ 1/2 x )1/2} . (74)
The bound is exact if px and py are elliptical distributions with the same density generator. Corollary 1. Consider the same setting as in Lemma 12, then the average squared Euclidean distance is bounded by
E x,y
[ ‖x− y‖2 ] ≤ W22 (px, py) + 2 trace {( Σ1/2x ΣyΣ 1/2 x )1/2} . (75)
Proof. Straightforward application of the results from Lemma 12 and Theorem 2.
Lemma 13. If the policy’s (co-)variance Var[π(a | s)] = Σπ(s) ∈ SdA+ is finite over the complete state space, i.e., maxs∈S trace{Σπ(s)} ≤ σ̄2π , then
trace {( Σπ(ŝ) 1/2Σπ(s)Σπ(ŝ) 1/2 )1/2} ≤√dAσ̄2π (76)
Proof.
trace {( Σπ(ŝ) 1/2Σπ(s)Σπ(ŝ) 1/2 )1/2}
(77)
The trace of a matrix is the same as the sum of its eigenvalues, and the square root of a matrix has eigenvalues that are square root of its eigenvalues. From Jensen’s inequality we know
that ∑dA i=1 √ λi ≤ √ dA ∑dA i=1 λi and consequently it holds for a matrix M ∈ RdA×dA that
trace{M1/2} ≤ √ dA trace{M} , so that
≤ √ dA trace{Σπ(ŝ)1/2Σπ(s)Σπ(ŝ)1/2} (78)
The trace is invariant under cyclic permutation = √ dA trace{Σπ(ŝ)Σπ(s)} (79)
Since both matrices are positive semi-definite, it follows from the Cauchy-Schwartz inequality that ≤ √ dA trace{Σπ(ŝ)} trace{Σπ(s)} (80)
By assumption, the covariance matrices’ traces are bounded ≤ √ dAσ̄ 2 π (81)

heading: A.5 MODEL ERRORS IN OPC
text: While Theorem 1 highlights that OPC counteracts the on-policy error in predicted performance, for stochastic policies we use the Lipschitz continuity of the model to upper-bound errors. In this section, we look at the impact of model errors in combination with OPC. Specifically we focus on the one-step prediction case from a known initial state ŝ0. There, while for the model p̃model without OPC the prediction error only depends on the quality of the model, with OPC it is instead the minimum of the model error and the policy variance. This is advantageous, since typical environments tend to have more states that actions, so that the trace of the policy variance can be significantly smaller than the full-state model error. Lemma 14. Under the assumptions of Theorem 1, starting from an initial state ŝ0 the following condition holds:
W22 (p̃opc? (s1), p(ŝ1)) ≤ min ( O ( trace{Var[π(· | ŝ0)]} ) ,O (∫ ‖f(ŝ0,a)− f̃(ŝ0,a)‖2 dπ(a | ŝ0)︸ ︷︷ ︸
One-step model error
))
Proof. The first term in the minimum follows directly from the base-case of Lemma 6. For the second term, follow the same derivation, but note that under the distribution p(ŝ0, â0, s0,a0) = p(ŝ0)π(â0 | ŝ0)δ(s0 − ŝ0)π(a0 | s0) we have ∫ p(ŝ1 | ẑ0) dp(ẑ0) = ∫ p(s1 | z0) dp(z0). Inserting this into the r.h.s. of Eq. (29) and following the same steps we obtain
W22 (p̃opc? (s1), p(ŝ1)) ≤ W22 (∫∫ δ(s1 − [f(ẑ0) + f̃(z0)− f̃(ẑ0)]) dp(ẑ0)p(z0), ∫ δ(s1 − f(z0)) dp(z0) ) (82)
≤ ∫∫ ‖f(ẑ0) + f̃(z0)− f̃(ẑ0)− f(z0)‖2 dp(ẑ0, z0) (83)
= ∫∫ ‖f(ẑ0)− f̃(ẑ0)‖2 + ‖f(z0)− f̃(z0)‖2 dp(ẑ0, z0) (84)
= 2 ∫ ‖f(ẑ0)− f̃(ẑ0)‖2 dp(z0) (85)
= 2 ∫ ‖f(ŝ0,a)− f̃(ŝ0,a)‖2 dp(a | ŝ0) (86)
Note the additional factor of two in front of the upper bound on the model error, which comes from using the model ‘twice’: once with ẑ and once with z. In practice we do not see any adverse effects of this error, presumably because either the variance of the policy is sufficiently small, or due to the upper bound being lose in practice.

heading: B MOTIVATING EXAMPLE – IN-DEPTH ANALYSIS
text: In this section, we re-visit the motivating example presented in Section 4.1 of the main paper. For completeness, we re-state all assumptions that lead to the simplified system at hand. We continue with an analysis of the reward landscape and how OPC influences its shape. Next, we investigate how an increasing mismatch of the dynamics model impacts the gradient error. In addition to the result presented in the main paper, we here show the influence of different model errors, i.e., ∆A as well as ∆B. While OPC is motivated for the use case of on-policy RL algorithms, we further show that the resulting gradients are robust with respect to differences in data-generating and evaluation policy, i.e., the off-policy setting. Lastly, we state the the signed gradient distance that we use for evaluation of the gradient errors, state the relevant theorem for determining the closed-loop stability of linear systems, as well as all numerical values used for the motivating example.

heading: B.1 SETUP
text: Here, we assume a linear system with deterministic dynamics
p(st+1 | st,at) = δ(Ast + Bat | st,at), ρ0(s0) = δ(s0) (87) with A,B ∈ R and δ(·) denoting the Dirac-delta distribution. The linear policy and bell-shaped reward are given by the following equations
πθ(at | st) = δ(θst | st) with θ ∈ R and r(st,at) = exp { − ( st σr )2} . (88)
Further, we assume to have access to an approximate dynamics model p̃
p̃(st+1 | st,at) = δ((A + ∆A)st + (B + ∆B)at | st,at), (89) where ∆A,∆B quantify the mismatch between the approximate model and the true system. For completeness, the (deterministic) policy gradient is defined as
∇θ 1
T T−1∑ t=0 r(st,at), (90)
where the state/action pairs are obtained by simulating any of the two above models for T time-steps and following policy πnθ resulting in the trajectory τ̃
n = {(snt ,ant )}T−1t=0 . Because both the model and policy are deterministic, we can compute the analytical policy gradient from only one rollout.

heading: B.2 REWARD LANDSCAPES
text: In a first step, we will look at the cumulative rewards as a function of the policy parameter for the different systems at hand: 1) the true system, 2) the approximate model without OPC and 3) the approximate model with OPC. Further, let’s assume that the model mismatch is fixed to some arbitrary value. The resulting reward landscapes are shown in Fig. 7. We would like to emphasize several key aspects in the plots: First, as one would expect the model mismatch leads to different optimal policies as well as misleading policy gradients for large parts of the policy parameter space. Second, the reward landscape for the model with OPC depends on the respective reference policy πn that was used to generate the data for the corrections. Consequently, the correct reward is recovered at θ = θn. More importantly, the OPC reshape the reward landscape such that the policy gradients point towards the correct optimum (left plot). Lastly, even when using OPC the policy gradient’s sign is not guaranteed to have the correct sign (right plot). The extent of this effect strongly depends on the model mismatch, which we will investigate in the next section.
B.3 INFLUENCE OF MODEL ERROR
As shown in the previous section, the estimated policy gradient depends on the current policy as well as the mismatch between the true system and the approximate model. Fig. 2 depicts the (signed) differences between the true policy gradient as well as the approximated gradient as a function of model mismatch and the reference policy. Here, the opacity of the background denotes the magnitude of the error and the color denotes if the true and estimated gradient have the same (blue) or oppposite (red) sign. In the context of policy learning, the sign of the gradient is more relevant than the actual magnitude due to internal re-scaling of the gradients in modern implementations of stochastic optimizers such as Adam (Kingma & Ba, 2015). In our example, even for negligible model errors (either in ∆A or ∆B), the model-based approach can lead to gradient estimates with the opposite sign, indicated by the large red areas for the left figures in Fig. 2. On the other hand, applying OPC to the model, we gradient estimates are significantly more robust with respect to errors in the dynamics.x
B.4 INFLUENCE OF OFF-POLICY ERROR
Until now we have considered the case in which the reference trajectory used for OPC is generated with the same policy as the one used for gradient estimation, i.e., the on-policy setting. In this case, we have observed that the true return could be recovered (see Fig. 7) when using OPC and that the gradient estimates are less sensitive to model errors (see Fig. 8). The off-policy case corresponds to the policy gains in Fig. 7 that are different from the reference policy πn indicated by the dashed line. Fig. 9 summarizes the results for the off-policy setting. Here, we varied the policy error and the reference policy itself for varying model errors. Note that for the correct model, we always recover the true gradient. But also for inaccurate models, the gradient estimates retain a good quality in most cases, with the exception for some model/policy combinations that are close to unstable.

heading: B.5 ADDITIONAL INFORMATION
text: Next, we provide some additional information about how we compute gradient distances, properties of linear systems, and exact numerical values used.
B.5.1 COMPUTING THE SIGNED GRADIENT DISTANCE
In order to compare two (1-dimensional) gradients in terms of sign and magnitude, we use the following formula
d(g1, g2) = 1
π  sign(g2) ·∆g, if g1 = 0 sign(g1) ·∆g, if g2 = 0 sign(g1 · g2) ·∆g, otherwise , with ∆g = |arctan g1 − arctan g2| . (91)
The magnitude of this quantity depends on the normalized difference between the tangent’s angles ∆g and is positive for gradients with the same sign and vice versa it is negative for gradients with opposing signs. See also Fig. 10 for a sketch.

heading: B.5.2 DETERMINING THE CLOSED-LOOP STABILITY FOR LINEAR SYSTEMS
text: For linear and deterministic systems, we can easily check if the system is (asymptotically) stable for a particular linear policy using the following standard result from linear system theory: Theorem 3 (Exponential stability for linear time-invariant systems (Callier & Desoer, 1991)). The solution of xt+1 = Fxt is exponentially stable if and only if σ(F) ⊂ D(0, 1), i.e., every eigenvalue of F has magnitude strictly less than one.
In our setting, this means that the closed-loop systems fulfilling the following are unstable,
|A + ∆A + (B + ∆B)θ| > 1, (92) i.e., the state and input grow exponentially. We therefore refrain from including unstable system in the results to avoid numerical issues for the gradients’ computation. The respective areas in the plots are not colored, see e.g., bottom left corner in Fig. 8b.

heading: B.5.3 NUMERICAL VALUES
text: The numerical values for all parameters used in the motivating example are given as follows:
• True system dynamics: A = 1.0,B = 1.0 • Initial condition: s0 = 1.0 • Reward width parameter: σr = 0.05 • Optimal policy gain: θ∗ = −1.0 • Rollout horizon: T = 60

heading: C ADDITIONAL EXPERIMENTAL RESULTS
text: In this section, we provide additional experimental results that did not fit into the main body of the paper.

heading: C.1 COMPARISON WITH OTHER BASELINE ALGORITHMS
text: Fig. 11 shows our method compared to a range of baseline algorithms. The results for all baselines were obtained from Janner et al. (2019) via personal communication. Note that all results are presented in terms of mean and standard deviation. The comparison includes the following methods:
• MBPO (Janner et al., 2019), • PETS (Chua et al., 2018), • SAC (Haarnoja et al., 2018), • PPO (Schulman et al., 2017), • STEVE (Buckman et al., 2018), • SLBO (Luo et al., 2018).

heading: C.2 ABLATION - RETAIN EPOCHS
text: One of the hyperparameters that we found is critical to both OPC and MBPO(?), is retain epochs, i.e., the number of epochs that are kept in the data buffer for the simulated data generated with the p̃x with x = {OPC,model}. The results for a comparison are shown in Fig. 12. For MBPO(?), we found that for some environments (HalfCheetah, AntTruncatedObs) smaller values for retain epochs are helpful, i.e., simulated data is almost only on-policy, and for other environments larger values are beneficial (Hopper, Walker2d). For OPC on the other hand, we found that retain epochs = 50 almost always leads to better results.
C.3 INFLUENCE OF STATE REPRESENTATION: IN-DEPTH ANALYSIS
In the following, we will have a closer look at the surprising result from Fig. 4 (left). In there, the results indicate that MBPO(?) is not able to learn a stabilizing policy within the first 7’500 steps on the RoboSchool variant of the CartPole environment. We hypothesize that the failure of MBPO(?) is due to a mismatch of the simulated data and the true state distribution. As a result, the policy that is optimized with the simulated data cannot stabilize the pole in the true environment.
To validate our hypothesis, we perform the following experiment: First, we train a policy π∗ that we know performs well on the true environment, leading to the maximum evaluation return of 1000. With this policy, we rollout a reference trajectory on the true environment τ ref = {(ŝt, ât)}Tt=0. To perform branched rollouts with the respective methods, OPC and MBPO(?), we use the learned transition models after 20 epochs (5’000 time steps) that were logged during a full learning process for each method. We then perform 100 branched rollouts of length H = 20 starting from randomly sampled initial states of the reference trajectories.
Fig. 13 shows the difference between the true and predicted state trajectories (median and 95th
percentiles) for each state in [x, cos(ϑ), sin(ϑ), ẋ, ϑ̇]. Since we start each branched rollout from a state on the real environment, the initial error is always zero. Across all states, the errors are drastically reduced when using OPC. Fig. 14 shows the predicted trajectories for the cosine of the pole’s angle, cos(ϑ). For MBPO(?), we observe that the trajectories often diverge and attain values that are clearly out of distribution, i.e., cos(ϑ) > 1.
With OPC ( ), the simulated rollouts follows the true state trajectories much closer, whereas with MBPO(?) ( ) the prediction errors quickly accumulate over time.
n ( ). For the return on the true environment ηn (top),
we additionally show the sample distribution of the rollouts’ returns. This nicely demonstrates how the policy smoothly transitions from failing consistently (n ≤ 8) to successfully stabilizing the pole (n ≥ 12). Additionally, note that the on-policy model error is almost always smaller for OPC compared to MBPO(?), which supports the theoretical motivation that our method is build upon.
C.4 IMPROVEMENT BOUND: EMPIRICAL INVESTIGATION
In this section, we empirically investigate to what extent OPC is able to tighten the policy improvement bound Eq. (5) compared to pure model-based approaches. As mentioned in the main paper, the motivation behind OPC is to reduce the on-policy error and we assume that the off-policy error is not affected too badly by the corrected transitions. Generally speaking, it is difficult to quantify a-priori how OPC compares to a purely model-based approach as this depends on the generalization capabilities of the learned dynamics model, the environment itself and the reward signal.
Here, we analyze the CartPole environment (PyBullet implementation) and estimate the respective error terms that appear in the policy improvement bound Eq. (5). To this end, we roll out a sequence of policies πn on the true environment (to estimate ηn and ηn+1) and on the learned model with and without OPC (to estimate η̃n and η̃n+1). The sequence of policies was obtained during a full policy optimization run (the policy was logged after each update) and we roll out each policy 100 times on the respective environment/models. The corresponding learned transition models were similarly logged during a full policy optimization. For OPC, we estimate the off-policy return η̃n+1 using the learned model from iteration n and reference trajectories from the true environment that were collected under πn, but then rollout the model with πn+1. The results are shown in Fig. 15.

heading: C.5 OFF-POLICY ANALYSIS
text: In this section, we investigate the robustness of OPC towards ‘off-policy-ness’ of the reference trajectories that are simulated under the data-generating policy πn. To this end, we manually increase the stochasticity of the behaviour policy πrollout by a factor β such that
V[πrollout(· | s)] = β2V[πn(· | s)], (93) and we keep the mean the same for both policies. Fig. 16 shows the distributions of the returns from 100 rollouts with varying degree of ‘off-policy-ness’ on the CartPole environment (PyBullet). Note that the data-generating policy consistently leads to the maximum return of 1000. For β close to one, both OPC and MBPO(?) lead to almost ideal behaviour and correctly predict the return in more than 95% of the cases. As we increase the policy’s stochasticity (from left to right), the respective performance of the policy decreases until all rollouts terminate prematurely (β ≥ 2.3). Notably, the extent of this effect is almost identical between OPC and MBPO(?). We conclude that OPC is robust towards ‘off-policy-ness’ of the reference trajectories. Otherwise, we should observe a more pronounced degradation of the policy’s performance with increased stochasticity of the behaviour policy since this leads to observing more off-policy state/actions.
Algorithm 2 Original MBPO algorithm 1: Initialize policy πφ, predictive model pθ, environment dataset Denv, model dataset Dmodel 2: for N epochs do 3: Train model pθ on Denv via maximum likelihood 4: for E steps do 5: Take action in environment according to πφ; add to Denv 6: for M model rollouts do 7: Sample st uniformly from Denv 8: Perform k-step model rollout starting from st using policy πφ; add to Dmodel 9: for G gradient updates do 10: Update policy parameters on model data: φ← φ− λπ∇̂φJπ(φ,Dmodel)
Algorithm 3 Our version of MBPO algorithm, denoted as MBPO(?) 1: Initialize policy πφ, predictive model pθ, environment dataset Denv, model dataset Dmodel 2: for N epochs do 3: Train model pθ on Denv via maximum likelihood 4: for E steps do 5: Take action in environment according to πφ; add to Denv 6: for M model rollouts do 7: Sample st uniformly from Denv 8: Perform k-step model rollout starting from st using policy πφ; add to Dmodel 9: for G gradient updates do 10: Update policy parameters on model data: φ← φ− λπ∇̂φJπ(φ,Dmodel)
C.6 IMPLEMENTATION CHANGES TO ORIGINAL MBPO
Here, we provide details to the changes we made to the original implementation of MBPO. We denote our variant as MBPO(?).
Episodic Setting MBPO updates the policy during rollouts on the real environment. Algorithms 2 and 3 show the original and our version of MBPO, respectively. While the original version might be more sample efficient because it allows for more policy gradient updates based on more recent data, it is not realistic to update the policy during a rollout on the real environment.
Mix-in of Real Data As mentioned in the main paper, one of the key problems with model-based RL (MBRL) is that the generated data might exhibit so-called model-bias. Biased data can be problematic for policy optimization as the transition tuples possibly do not come from the true state distribution of the real environment, thus misleading the RL algorithm. In the original implementation of MBPO, the authors use a mix of simulated data from the model as well as observed data from the true environment for policy optimization (https://github.com/JannerM/mbpo/blob/ master/mbpo/algorithms/mbpo.py#L430) to alleviate the issue of model bias. While this design choice might help in practice, it adds an additional hyperparameter and the exact influence of this parameter is difficult to interpret. We therefore refrain from mixing in data from the true environment but instead only use simulated transition tuples – staying true to the spirit of MBRL.
Fix Replay Buffer The replay buffer that stores the simulated data Dmodel for policy optimization is allocated to a fixed size of rollout length×rollout batch size×retain epochs, meaning that per epoch rollout length×rollout batch size transition tuples are simulated and only the data from the last retain epochs are kept in the buffer (https://github. com/JannerM/mbpo/blob/master/mbpo/algorithms/mbpo.py#L351). As the buffer is implemented as a FIFO queue and rollouts might terminate early such that the actual number of simulated transitions is less than rollout length, data from older episodes as specified by retain epochs are retained in the buffer. We assume that this is not the intended behaviour and correct for that in our implementation.

heading: C.7 RESULTS FOR ORIGINAL MBPO IMPLEMENTATION
text: The following results were obtained with the original MBPO implementation without the changes described in Appendix C.6. Consequently, the results for OPC in this section also do not include these changes.

heading: Comparative Evaluation
text: We evaluate the original implementation on three continuous control benchmark tasks from the MuJoCo control suite (Todorov et al., 2012). The results for OPC, two variants of MBPO and SAC are presented in Fig. 3. Both OPC and MBPO use a rollout horizon ofH = 10 to generate the training data. The difference between the two MBPO variants lies in the mix-in ratio β of real off-policy transitions to the simulated training data. Especially for highly stochastic environments such as the Hopper, this mix-in ratio is a critical hyperparameter that requires careful tuning (see also Fig. 18). Fig. 3 indicates that OPC is on par with MBPO on both the InvertedPendulum and HalfCheetah environments that exhibit little stochasticity. On the Hopper environment, OPC outperforms both MBPO variants. Note that the mix-in ratio for MBPO is critical for successful learning (the original implementation uses β = 5%). OPC on the other hand, does not require any mixed-in real data. SAC learns slower on the more complex Hopper and HalfCheetah environments, re-iterating that model-based approaches are significantly more data-efficient than model-free methods.

heading: Large Ablation Study – Hopper
text: The full study investigates the influence of the following hyperparameters and design choices:
• Rollout length H and total number of simulated transitions N .
• Mix-in ratio of real transitions into the training data β.
• Deterministic or stochastic rollouts (for MBPO): Current state-of-the-art methods in MBRL rely on probabilistic dynamic models to capture both aleatoric and epistemic uncertainty. Accordingly, when rolling out the model, these two sources of uncertainty are accounted for. However, we show that in terms of evaluation return, the stochastic rollouts do not always lead to the best outcome.
• Re-setting the buffer of simulated data after a policy optimization step: We found that re-setting the replay buffer for simulated data after each iteration of Algorithm 2 can have a large influence. In particular, the replay buffer is implemented as a FIFO queue with a fixed size. Hence, if the buffer is not emptied after each iteration, it still contains (simulated) off-policy transitions.
Fig. 18 presents the results of the ablation study. We want to highlight a few core insights:
1. When choosing the best setting for each method, OPC improves MBPO by a large margin (bottom row right). Generally, for long rollouts H = 20 (bottom row), OPC improves MBPO.
2. Across all settings, OPC performs well and is more robust with respect to the choices of hyperparameters (e.g., bottom row left and center, third row left). Only for few exceptions, re-setting the buffer can have detrimental effects (e.g., top right, third row right).
3. Mixing in real transition data can be highly beneficial for MBPO (second row left and center) but it can also have the opposite effect (bottom row).
4. Using deterministic rollouts can be beneficial (third row right and left), detrimental (third row center) or have no influence (bottom row left) for MBPO.
5. It is not clear, if re-setting the buffer after each iteration should overall be recommended or not. This remains an open question left for future research.

heading: D CONNECTION BETWEEN MBRL AND ILC
text: In this section we compare optimization-based or so-called norm-optimal ILC (NO-ILC) with MBRL. In particular, we show that under certain assumptions we can reduce the MBRL setting to NO-ILC. This comparison is structured as follows: First, we review the basic assumptions and notations for NO-ILC. While there are many variations on NO-ILC, we will only consider the very basic setting, i.e., linear dynamics, fully observed state and deterministic state evolution. Then, based on the lifted state representation of the problem, we derive the solution to the optimization problem that leads to the input sequence of the next iteration / rollout. Next, we will state the general MBRL problem and pose the simplifications that we need to make in order to be equivalent to the NO-ILC problem. Last, we show that the solution to the reduced MBRL problem is equivalent to the one of NO-ILC.

heading: D.1 NORM-OPTIMAL ILC
text: The goal of NO-ILC is to find a sequence of inputs a = [a>0 , . . . ,a > T−1] > with length T + 1 such that the outputs y = [y>0 , . . . ,y > T ] > follow a desired output trajectory ŷ = [ŷ>0 , . . . , ŷ > T ] >. In the simplest setting we assume that the system evolves according to the following linear dynamics
st+1 = Ast + Bat + dt (94) yt = Cst, (95)
with state st ∈ RdS , action at ∈ RdA , output yt ∈ Rp and disturbance dt ∈ RdS . One of the major assumption in ILC is that the disturbances dt are repetitive, meaning that the sequence d = [d>0 , . . . ,d > T ] > does not change (or only varies slightly) across multiple rollouts. While these disturbances can be considered to come from some exogenous error source, one can also interpret these as unmodeled effects of the dynamics, e.g., nonlinearities stemming from friction, aerodynamic effects, etc.
For the derivation, let’s assume C = I so that we’re operating on an MDP. Consequently, the goal of ILC is to track a sequence of desired states ŝ instead of outputs ŷ. In order to find an optimal input sequence, we minimize the squared 2-norm of a state’s deviation from the reference for each time-step t of the sequence, i.e., et = ŝt − st. As is common in the ILC literature, we make use of the so-called lifted system formulation such that we can conveniently write the state evolution for one rollout using a single matrix/vector multiplication. Assuming zero initial state (which can always be done in the linear setting by just shifting the state by a constant offset), we obtain the following formulation
s = Fa + d, with F =  0 · · · B 0 AB B 0 A2B AB B . . .
... . . . . . .  ∈ RdS(T+1)×dAT . (96) Thus, we can write the state-error at the j-th iteration of ILC in the lifted representation as (recall that the disturbance d does not change across iterations)
e(i) = ŝ− s(i) = ŝ− Fa(i) − d, (97) e(i+1) = e(i) − F(a(i+1) − a(i)). (98)
Now, the resulting optimization problem then becomes
a (i+1) ∗ = arg min a(i+1) J ( a(i+1) ) with J ( a(i+1) ) = 1 2 ‖e(i+1)‖2.
In order to be less sensitive to noise and the inherent stochasticity of real-world problems, one typically also adds a regularizing term that penalizes the changes in the input sequence. While this additional penalization term can slow down the learning process it makes it more robust by avoiding overcompensation to the disturbances. The full objective then becomes
J ( a(i+1) ) = 1
2 ‖e(i+1)‖2M +
1 2 ‖a(i+1) − a(i)‖2W , (99)
where we additionally added positive semi-definite cost matrices M,W for the respective norms in order to facilitate tuning of the corresponding terms in the objective. Given the regularized cost function we obtain the optimal sequence of inputs at the next iteration as
a (i+1) ∗ = a (i) ∗ + ( F>MF +W )−1 F>Me(i). (100)

heading: D.2 MODEL-BASED RL TO NORM-OPTIMAL ILC
text: Assumptions In order to show the equivalence of MBRL and NO-ILC we need to make some assumptions to the above stated optimization problem.
• No aleatoric uncertainty in the model, i.e., no transition noise ωt = 0 ∀t. • Typically in RL we assume the policy to be stationary, however, in this setting we will allow
for non-stationary policies that are indexed by time t (the result will not necessarily depend on the state such that we essentially just obtain a feedforward control sequence. To include feedback one can always just combine the feedforward signal with a local controller that tracks the desired state/action trajectory).
• The reward is given as the negative quadratic error of the state w.r.t. a desired state trajectory, r(st,at) = − 12‖ŝt − st‖2 • Typically in RL we have constraints on the policy such that it does not change too much after every iteration, see e.g. TRPO, REPS, etc. While in the mentioned approaches, the policy are often constrained in terms of their parameterization vectors, we constrain it as ‖π(i+1)t − π(i)t ‖2 ≤ π ∀t.
• Assume that the model is given by f̃t(s,a) = As + Ba + dt, where A and B are fixed system matrices and dt is a time-dependent offset that we learn.
The learned dynamics model Given state/input pairs of the i-th trajectory τ (i) = {(s(i)t ,a(i)t )}Tt=0 from the true system, we can now improve our dynamics model. In particular, if we minimize the prediction error over τ (i), we obtain
d (i) t = s (i) t+1 −
( As
(i) t + Ba (i) t
) (101)
The resulting dynamics for the optimal control problem in Eq. (1) become
f̃ (i) t (s,a) = s (i) t+1 + A(s− s (i) t ) + B(a− a(i)t ) (102)
which are the error dynamics around the trajectory. Now in the noisy case, taking the last trajectory is not necessarily the best thing one can do. E.g., (Schöllig & D’Andrea, 2009) instead integrate the information of all past trajectories via Kalman-filtering. In the fully observed case, one way to think of this is as low-pass filtering dt in order to account for the transition noise ω.
The resulting MBRL problem Now, let’s plug in all assumptions into the MBRL problem and have a look at how to solve it.
π (i+1) ∗ = arg min
π={π0,...,πT } T∑ t=0 1 2 ‖ŝt − st‖2
st+1 = s (i) t+1 + A(s− s (i) t ) + B(a− a(i)t )
at = πt
‖π(i)t − πt‖2 ≤ π ∀t
(103)
where we have flipped the reward’s sign to transform it into a minimization problem. In the presented form, this optimization problem has a well-defined unique solution, however, it is a-priori not clear if the trust-region constraint is active for some time-steps. To facilitate an analytical solution to equation 103, we incorporate the constrain on the policy’s stepsize by a soft-constraint such that
π (i+1) ∗ = arg min
π={π0,...,πT } T∑ t=0 1 2 ‖ŝt − st‖2 + Cπt 2 ‖π(i)t − πt‖2
st+1 = s (i) t+1 + A(s− s (i) t ) + B(a− a(i)t )
at = πt,
(104)
with Cπt being constants that weight the respective trust-region terms.
Generally, the MBRL problem Eq. (2) cannot be solved analytically due to the (possibly highly nonlinear, non-differentiable, etc.) reward signal and the need for propagating the (uncertain) dynamics model forwards in time. However, using the simplifying assumptions above, the reduced MBRL problem equation 104 can in fact be solved analytically. We can circumvent the equality constraints by predicting the state trajectory using the error corrected dynamics such that we obtain a lifted dynamics formulation similar to the analysis for ILC,
s(i) = Fa(i) + d(i), with d(i) = s(i) − Fa(i), (105) with s,a denoting the stacked states and actions of the recorded trajectory. Using this notation, the optimization problem reduces to
π (i+1) ∗ = arg min
π
1 2 ‖ŝ− s(i)‖2 + 1 2 ‖π(i) − π‖2Cπ , (106)
with Cπ = diag[Cπ0 , . . . , C π H ]. By inserting equation 105 into equation 106 we obtain the closedform solution as
π (i+1) ∗ = ( F>F + Cπ )−1 F> ( ŝ− d(i) ) = π (i) ∗ + ( F>F + Cπ )−1 F> ( ŝ− s(i) ) , (107)
which, clearly, is equivalent to equation 100 for M = I andW = Cπ .

heading: D.3 EXTENSIONS
text: In the previous section, we analyzed the most basic setting for NO-ILC, i.e., linear dynamics, no transition noise in the dynamics and no state/input constraints. Some of these assumptions can easily be liftd to genearlized the presented framework.
Nonlinear Dynamics Instead of the system being defined by fixed matrices A,B, we can just linearize a non-linear dynamics model f(s,a) around the last state/input trajectory such that
A (i) t =
∂f
∂s ∣∣∣∣ st=s (i) t ,at=a (i) t , B (i) t = ∂f ∂a ∣∣∣∣ st=s (i) t ,at=a (i) t
(108)
and the corresponding dynamics matrix for the lifted representation becomes (dropping the superscript notation indicating the iteration for clarity)
F =  0 · · · B0 0 A0B0 B1 0 A0A1B0 A0B1 B2 . . .
... . . . . . .
 ∈ RdS(T+1)×dAT . (109)
State/Input Constraints Recall that we could solve the quadratic problems equation 99 and equation 104 analytically because we assumed no explicit inequality constraints on the state and inputs. In practice, however, both states and actions are limited by physical or safety constraints typically given by polytopes. While a closed-form solution is not readily available for this case, one can easily employ any numerical solver to deal with such constraints. See, e.g., https://en.wikipedia.org/w/index.php?title=Quadratic_programming for a comprehensive list of available solvers.
Stochastic Dynamics Schöllig & D’Andrea (2009) generalize the ILC setting by considering two separate sources of noise: 1) transition noise in the dynamics, i.e., st+1 = f(st,at) + ω f t as well as 2) varying disturbances, i.e., d(i+1) = d(i) + ωd. Based on this model, a Kalman-Filter in the iteration-domain is developed that estimates the respective random variables and the ILC scheme is adapted to account for the estimated quantities.

heading: E HYPERPARAMETER SETTINGS
text: Below we list the most important hyperparameter settings.
F IMPLEMENTATION AND COMPUTATIONAL RESOURCES
Our implementation is based on the code from MBPO (Janner et al., 2019), which is open-sourced under the MIT license. All experiments were run on a HPC cluster, where each individual experiment used one nVidia V100 GPU and four Intel Xeon CPUs. All experiments (including early debugging and evaluations) amounted to a total of 84‘713 hours, which corresponds to roughly 9.7 years if the jobs ran sequentially. Most of this compute was required to ensure reproducability (ten random seeds per job and ablation studies over the effects of parameters).
Since the cluster is part of an overall carbon-neutral infrastructure, these experiments did not contribute to climate change.

Review summaries:
The paper considers model-based reinforcement learning (MBRL) of the MBPO flavour, where model-free RL methods are accelerated using rollouts of a learned model.  THis paper proposes to alleviate (off-policy) model bias per rollout using the (on-policy) data and calls this approach on-policy correction (OPC).  With this perspective, a comparison is made between MBRL and iterative learning control (ILC). This method is compared to MBPO on standard mujoco tasks and pybullet versions.  

This paper addresses on-policy errors in model-based Reinforcement Learning. The authors decompose the policy improvement bound in terms that depend on the off-policy and on-policy model errors. They first present two major techniques, one purely based on the replay buffer, which asymptotically obtains zero on-policy error, and one based on a learned model, which leverages off-policy data. Then they propose On-Policy Corrections (OPC), a combination of the two techniques which can use a model and still maintain low on-policy error using the replay buffer. This method has asymptotically (with infinite data) zero on-policy error when using deterministic policies. If the policy is stochastic, the error grows exponentially with the trajectory length T. The authors implement their algorithm on top of MBPO and show experimentally that OPC-MBPO has superior performance in some environments.

The paper proposes a way to combine observed data with a learned model to mitigate the cost of each choice separately - the observed data is over-fitted to a specific policy while the learned model has approximation errors. The authors mix the two options using probabilistic corrections based on the learned model applied to observed transitions.  The authors also provide convergence results for their algorithm,