Original: Uncertainty Estimation Using Riemannian Model Dynamics Anonymous Author(s) AffiliationAddress email
Revised: Uncertainty Estimation Using Riemannian Model Dynamics for Offline Reinforcement Learning

Original: Abstract
Revised: Abstract

Original: Model-based offline reinforcement learning approaches generally rely on boundsof model error.
Revised: Model-based offline reinforcement learning approaches generally rely on bounds of model error.

Original: Estimating these bounds is usually achieved through uncertaintyestimation methods.
Revised: Estimating these bounds is usually achieved through uncertainty estimation methods.

Original: In this work, we combine parametric and nonparametricmethods for uncertainty estimation through a novel latent space based metric.
Revised: In this work, we combine parametric and nonparametric methods for uncertainty estimation through a novel latent space based metric.

Original: Inparticular, we build upon recent advances in Riemannian geometry of generativemodels to construct a pullback metric of an encoder-decoder based forward model.
Revised: In particular, we build upon recent advances in Riemannian geometry of generative models to construct a pullback metric of an encoder-decoder based forward model.

Original: Our proposed metric measures both the quality of out-of-distribution samplesas well as the discrepancy of examples in the data.
Revised: Our proposed metric measures both the quality of out-of-distribution samples as well as the discrepancy of examples in the data.

Original: We leverage our methodfor uncertainty estimation in a pessimistic model-based framework, showing asignificant improvement upon contemporary model-based offline approaches oncontinuous control and autonomous driving benchmarks.
Revised: We leverage our method for uncertainty estimation in a pessimistic model-based framework, showing a significant improvement upon contemporary model-based offline approaches on continuous control and autonomous driving benchmarks.

Original:  Introduction
Revised:  Introduction

Original: Offline Reinforcement Learning (RL) [Levine et al., 2020], a.k.a.
Revised: Offline Reinforcement Learning (RL) [Levine et al., 2020], a.k.a.

Original: batch-mode RL [Ernst et al. ,2005, Riedmiller, 2005, Fonteneau et al., 2013], involves learning a policy from data sampled by apotentially suboptimal policy.
Revised: batch-mode RL [Ernst et al., 2005, Riedmiller, 2005, Fonteneau et al., 2013], involves learning a policy from data sampled by a potentially suboptimal policy.

Original: Offline RL seeks to surpass the average performance of the agents thatgenerated the data.
Revised: Offline RL seeks to surpass the average performance of the agents that generated the data.

Original: Traditional methodologies fall short in offline settings, causing overestimation ofthe return [Buckman et al., 2020, Wang et al., 2020, Zanette, 2020].
Revised: Traditional methodologies fall short in offline settings, causing overestimation of the return [Buckman et al., 2020, Wang et al., 2020, Zanette, 2020].

Original: One approach to overcome this in model-based settings is to penalize the return in out of distribution(OOD) regions, as depicted in ?? .
Revised: One approach to overcome this in model-based settings is to penalize the return in out of distribution (OOD) regions, as depicted in Figure 1.

Original: In this manner, the agent is constrained to stay “near" areas of lowmodel error, thereby limiting possible overestimation.
Revised: In this manner, the agent is constrained to stay “near" areas of low model error, thereby limiting possible overestimation.

Original: However, reliable estimates of model error arekey to the success of such methods.
Revised: However, reliable estimates of model error are key to the success of such methods.

Original: Estimating model error in OOD regions can be achieved through uncertainty estimation [Yu et al.
Revised: Estimating model error in OOD regions can be achieved through uncertainty estimation [Yu et al., 2020].

Original: Methods of parametric uncertainty estimation such as bootstrap ensembles [Efron, 1982],
Revised: Methods of parametric uncertainty estimation such as bootstrap ensembles [Efron, 1982],

Original: Monte Carlo Dropout [Gal and Ghahramani, 2016], and randomized priors [Osband et al., 2018], maybe susceptible to poor model specification and are most effective when dealing with large datasets.
Revised: Monte Carlo Dropout [Gal and Ghahramani, 2016], and randomized priors [Osband et al., 2018], may be susceptible to poor model specification and are most effective when dealing with large datasets.

Original: In contrast, nonparametric methods such as k-nearest neighbors (k-NN)
Revised: In contrast, nonparametric methods such as k-nearest neighbors (k-NN)

Original: [Villa Medina et al., 2013,
Revised: [Villa Medina et al., 2013,

Original: Fathabadi et al., 2021] are beneficial in regions of limited data, yet require a proper metric to be used.
Revised: Fathabadi et al., 2021] are beneficial in regions of limited data, yet require a proper metric to be used.

Original: We propose to combine parametric and nonparametric methods for uncertainty estimation.
Revised: We propose to combine parametric and nonparametric methods for uncertainty estimation.

Original: Particularly,we define a novel Riemmannian metric which captures the epistemic and aleatoric uncertainty of agenerative parametric forward model.
Revised: Particularly, we define a novel Riemmannian metric which captures the epistemic and aleatoric uncertainty of a generative parametric forward model.

Original: This distance metric is then applied to measure the averagegeodesic distance to the k -nearest neighbors in the data.
Revised: This distance metric is then applied to measure the average geodesic distance to the k -nearest neighbors in the data.

Original: We derive analytical expressions for ourmetric and provide an efficient manner to estimate it.
Revised: We derive analytical expressions for our metric and provide an efficient manner to estimate it.

Original: We then demonstrate the effectiveness ofour metric for penalizing an offline RL agent compared to contemporary approaches on continuouscontrol and autonomous driving benchmarks.
Revised: We then demonstrate the effectiveness of our metric for penalizing an offline RL agent compared to contemporary approaches on continuous control and autonomous driving benchmarks.

Original: As we empirically show, common approaches, including statistical bootstrap ensemble or Euclidean distances in latent space, do not necessarily capture theunderlying degree of error needed for model-based offline RL.
Revised: As we empirically show, common approaches, including statistical bootstrap ensemble or Euclidean distances in latent space, do not necessarily capture the underlying degree of error needed for model-based offline RL.

Original:  Preliminaries2.1 Offline Reinforcement Learning
Revised:  Preliminaries 2.1 Offline Reinforcement Learning

Original: We consider the standard Markov Decision Process (MDP) framework [Puterman, 2014] definedby the tuple ( S , A , r, P, α ) , where S is the state space, A the action space, r : S × A (cid:55)→ [0 , 1] thereward function, P : S×
Revised: We consider the standard Markov Decision Process (MDP) framework [Puterman, 2014] defined by the tuple ( S , A , r, P, α ) , where S is the state space, A the action space, r : S × A (cid:55)→ [0 , 1] the reward function, P : S

Revised: ×

Original: A × S (cid:55)→ [0 , 1] the transition kernel, and α ∈ (0 , 1) is the discount factor.
Revised: A × S (cid:55)→ [0 , 1] the transition kernel, and α ∈ (0 , 1) is the discount factor.

Original: In the online setting of reinforcement learning (RL), the environment initiates at some state s 0 ∼ ρ 0 .
Revised: In the online setting of reinforcement learning (RL), the environment initiates at some state s 0 ∼ ρ 0 .

Original: At any time step the environment is in a state s ∈ S , an agent takes an action a ∈ A and receives areward r ( s, a ) from the environment as a result of this action.
Revised: At any time step the environment is in a state s ∈ S , an agent takes an action a ∈ A and receives a reward r ( s, a ) from the environment as a result of this action.

Original: The environment transitions to state s ′according to the transition function P ( ·| s, a ) .
Revised: The environment transitions to state s ′ according to the transition function P ( ·| s, a ) .

Original: The goal of online RL is to find a policy π ( a | s ) thatmaximizes the expected discounted return v π = E π (cid:104)(cid:80) ∞ t =

Revised: The goal of online RL is to find a policy π ( a | s ) that maximizes the expected discounted return v π = E π (cid:104)(cid:80) ∞ t =

Original: α t r ( s t , a t ) | s 0 ∼ ρ 0 (cid:105) .
Revised: α t r ( s t , a t ) | s 0 ∼ ρ 0 (cid:105) .

Original: Unlike the online setting, the offline setup considers a dataset D n = { s i , a i , r i , s ′ i } ni =1 of transitionsgenerated by some unknown agents.
Revised: Unlike the online setting, the offline setup considers a dataset D n = { s i , a i , r i , s ′ i } ni =1 of transitions generated by some unknown agents.

Original: The objective of offline RL is to find the best policy in the testenvironment (i.e., real MDP) given only access to the data generated by the unknown agents.
Revised: The objective of offline RL is to find the best policy in the test environment (i.e., real MDP) given only access to the data generated by the unknown agents.

Original: 2.2 Riemannian Manifolds
Revised: 2.2 Riemannian Manifolds

Original: We define the Riemannian pullback metric, a fundamental component of our proposed method.
Revised: We define the Riemannian pullback metric, a fundamental component of our proposed method.

Original: Werefer the reader to Carmo [1992] for further details on Riemannian geometry.
Revised: We refer the reader to Carmo [1992] for further details on Riemannian geometry.

Original: We are interested in studying a smooth surface M with a Riemannian metric g .
Revised: We are interested in studying a smooth surface M with a Riemannian metric g .

Original: A Riemannian metricis a smooth function that assigns a symmetric positive definite matrix to any point in M .
Revised: A Riemannian metric is a smooth function that assigns a symmetric positive definite matrix to any point in M .

Original: At eachpoint z ∈ M a tangent space T z M specifies the pointing direction of vectors “along” the surface.
Revised: At each point z ∈ M a tangent space T z M specifies the pointing direction of vectors “along” the surface.

Original: Definition 1.
Revised: Definition 1.

Original: Let M be a smooth manifold.
Revised: Let M be a smooth manifold.

Original: A Riemannian metric g on M changes smoothly anddefines a real scalar product on the tangent space T z M for any z ∈ M as
Revised: A Riemannian metric g on M changes smoothly and defines a real scalar product on the tangent space T z M for any z ∈ M as

Original: [Equation]
Revised: [Equation]

Original: where G ( z ) ∈ R d z × d z is the corresponding metric tensor.
Revised: where G ( z ) ∈ R d z × d z is the corresponding metric tensor.

Original: ( M, g ) is called a Riemannian manifold.
Revised: ( M, g ) is called a Riemannian manifold.

Original: The Riemannian metric enables us to easily define geodesic curves.
Revised: The Riemannian metric enables us to easily define geodesic curves.

Original: Consider some differentiablemapping γ : [0 , 1] (cid:55)→ M ⊆ R d z , such that γ (0) =
Revised: Consider some differentiable mapping γ : [0 , 1] (cid:55)→ M ⊆ R d z , such that γ (0) =

Original: z 0 , γ (1) =
Revised: z 0 , γ (1) =

Original: z 1 .
Revised: z 1 .

Original: The length of the curve γ measuredon M is given by
Revised: The length of the curve γ measured on M is given by

Original: [Equation]
Revised: [Equation]

Original: The geodesic distance d ( z 1 , z 2 ) between any twopoints z 1 , z 2 ∈ M is then the infimum length over allcurves γ for which γ (0) =
Revised: The geodesic distance d ( z 1 , z 2 ) between any two points z 1 , z 2 ∈ M is then the infimum length over all curves γ for which γ (0) =

Original: z 0 , γ (1) =
Revised: z 0 , γ (1) =

Original: z 1 .
Revised: z 1 .

Original: That is,d ( z 1 , z 2 ) =
Revised: That is, [Equation]

Original: inf γ L ( γ ) s.t. γ (0)

Original: = z 0 , γ (1) = z 1 . The geodesic distance can be found by solving asystem of nonlinear ordinary differential equations(ODEs) defined in the intrinsic coordinates [Carmo,1992].
Revised: The geodesic distance can be found by solving a system of nonlinear ordinary differential equations (ODEs) defined in the intrinsic coordinates [Carmo, 1992].

Original: Pullback Metric.
Revised: Pullback Metric.

Original: Assume an ambient (observation) space X and its respective Riemannian manifold( M X , g X ) .
Revised: Assume an ambient (observation) space X and its respective Riemannian manifold ( M X , g X ) .

Original: Learning g X can be hard (e.g., learning the distance metric between images).
Revised: Learning g X can be hard (e.g., learning the distance metric between images).

Original: Still, itmay be captured through a low dimensional submanifold.
Revised: Still, it may be captured through a low dimensional submanifold.

Original: As such, it is many times convenient toparameterize the surface M X by a latent space 
Revised: As such, it is many times convenient to parameterize the surface M X by a latent space Z

Original: Z = R d Z and a smooth function f : Z (cid:55)→ X , where
Revised: = R d Z and a smooth function f : Z (cid:55)→ X , where

Original: Z is a low dimensional latent embedding space.
Revised: Z is a low dimensional latent embedding space.

Original: As learning the manifold M X can be hard, we turnto learning the immersed low dimensional submanifold M Z (for which the chart maps are trivial, since Z = R d Z ).
Revised: As learning the manifold M X can be hard, we turn to learning the immersed low dimensional submanifold M Z (for which the chart maps are trivial, since Z = R d Z ).

Original: Given a curve γ : [0 , 1] (cid:55)→ M Z we have that (cid:68) ∂f ( γ ( t )) ∂t , G X ( f ( γ ( t )))
Revised: Given a curve γ : [0 , 1] (cid:55)→ M Z we have that (cid:68) ∂f ( γ ( t )) ∂t , G X ( f ( γ ( t )))

Original: ∂f ( γ ( t )) [Equation]
Revised: ∂f ( γ ( t )) ∂t (cid:69) =

Original: (cid:69) =(cid:68) ∂γ ( t ) ∂t , J Tf ( γ ( t ))
Revised: (cid:68) ∂γ ( t ) ∂t , J Tf ( γ ( t ))

Original: G X ( f ( γ ( t )))
Revised: G X ( f ( γ ( t )))

Original: J f ( γ ( t )) ∂γ ( t ) ∂t (cid:69) , where the Jacobian matrix J f ( z )
Revised: J f ( γ ( t )) ∂γ ( t ) ∂t (cid:69) , where the Jacobian matrix J f ( z )

Original: = ∂f∂z ∈ R d X × d Zmaps tangent vectors in TM Z to tangent vectors in TM X .
Revised: = ∂f∂z ∈ R d X × d Z maps tangent vectors in TM Z to tangent vectors in TM X .

Original: The induced metric is thus given by
Revised: The induced metric is thus given by

Original: [Equation]
Revised: [Equation]

Original: [Equation] The metric G f is known as the pullback metric , as it “pulls back" the metric G X on X back to G fvia f : Z (cid:55)→ X .
Revised: The metric G f is known as the pullback metric , as it “pulls back" the metric G X on X back to G f via f : Z (cid:55)→ X .

Original: The pullback metric captures the intrinsic geometry of the immersed submanifoldwhile taking into account the ambient space X .
Revised: The pullback metric captures the intrinsic geometry of the immersed submanifold while taking into account the ambient space X .

Original: The geodesic distance in ambient space is capturedby geodesics in the latent space Z , reducing the problem to learning the latent embedding space Zand the observation function f .
Revised: The geodesic distance in ambient space is captured by geodesics in the latent space Z , reducing the problem to learning the latent embedding space Z and the observation function f .

Original: Indeed, learning the latent space and observation function f can beachieved through a encoder-decoder framework, such as a VAE [Arvanitidis et al., 2018].
Revised: Indeed, learning the latent space and observation function f can be achieved through a encoder-decoder framework, such as a VAE [Arvanitidis et al., 2018].

Original:  Background: Penalty of Uncertainty for Offline Reinforcement Learning
Revised:  Background: Penalty of Uncertainty for Offline Reinforcement Learning

Original: A key element of model-based RL methods involves estimating a model ˆ P ( s ′ | s, a ) to constructa pessimistic MDP 1 .
Revised: A key element of model-based RL methods involves estimating a model ˆ P ( s ′ | s, a ) to construct a pessimistic MDP 2 .

Original: This work builds upon MOPO, a recently proposed model-based offline RLframework [Yu et al., 2020]).
Revised: This work builds upon MOPO, a recently proposed model-based offline RL framework [Yu et al., 2020]).

Original: Particularly, we assume access to an approximate MDP ( S , A , ˆ r, ˆ P, α )
Revised: Particularly, we assume access to an approximate MDP ( S , A , ˆ r, ˆ P, α )

Original: (e.g., trained by maximizing the likelihood of the data), and define a penalized MDP ( S , A , ˜ r, ˆ P, α ) ,such that for all s ∈ S , a ∈ A ,
Revised: (e.g., trained by maximizing the likelihood of the data), and define a penalized MDP ( S , A , ˜ r, ˆ P, α ) , such that for all s ∈ S , a ∈ A ,

Original: ˜ r
Revised: ˜ r

Original: ( s, a )
Revised: ( s, a )

Original: =
Revised: =

Original: ˆ r
Revised: ˆ r

Original: ( s, a ) − λc ( P ( ·| s, a ) , ˆ P ( ·| s, a )) , where c penalizes thereward according to model error (e.g., the total variation distance) and λ > 0 .
Revised: ( s, a ) − λc ( P ( ·| s, a ) , ˆ P ( ·| s, a )) , where c penalizes the reward according to model error (e.g., the total variation distance) and λ > 0 .

Original: The offline RL problemis then solved by executing an online algorithm in the reward-penalized MDP.
Revised: The offline RL problem is then solved by executing an online algorithm in the reward-penalized MDP.

Original: Unfortunately, as
Revised: Unfortunately, as

Original: P ( ·| s, a ) is unknown, and can only be estimated from the data,
Revised: P ( ·| s, a ) is unknown, and can only be estimated from the data,

Original: c ( P ( ·| s, a ) , ˆ P ( ·| s, a )) cannot becalculated.
Revised: c ( P ( ·| s, a ) , ˆ P ( ·| s, a )) cannot be calculated.

Original: Nevertheless, one can attempt to upper bound the distance, i.e., for some U :
Revised: Nevertheless, one can attempt to upper bound the distance, i.e., for some U :

Original: S × A (cid:55)→ R ,c ( P ( ·| s, a ) , ˆ P ( ·| s, a )) ≤
Revised: S × A (cid:55)→ R , c ( P ( ·| s, a ) , ˆ P ( ·| s, a )) ≤

Original: U ( s, a ) , ∀ s ∈ S , a ∈ A .
Revised: U ( s, a ) , ∀ s ∈ S , a ∈ A .

Original: In this work we propose to use a naturally inducedmetric of a variational forward model, which we show can introduce an effective penalty for offline
Revised: In this work we propose to use a naturally induced metric of a variational forward model, which we show can introduce an effective penalty for offline

Original: RL.
Revised: RL.

Original: In Section 4 we define this metric, and finally, we leverage it in Section 4.4.
Revised: In Section 4 we define this metric, and finally, we leverage it in Section 4.4.

Original:  Metrics of Uncertainty
Revised:  Metrics of Uncertainty

Original: As described in the previous section, our goal is to estimate model error in order to penalize the agentin out of distribution (OOD) regions.
Revised: As described in the previous section, our goal is to estimate model error in order to penalize the agent in out of distribution (OOD) regions.

Original: Yu et al.
Revised: Yu et al.

Original: [2020] proposed to achieve this through bootstrapensembles, an out of distribution uncertainty estimation technique.
Revised: [2020] proposed to achieve this through bootstrap ensembles, an out of distribution uncertainty estimation technique.

Original: Alternatively, we propose toemploy a well-known nonparametric approach for uncertainty estimation [Villa Medina et al., 2013,
Revised: Alternatively, we propose to employ a well-known nonparametric approach for uncertainty estimation [Villa Medina et al., 2013,

Original: Fathabadi et al., 2021], namely k -nearest neighbors ( k -NN).
Revised: Fathabadi et al., 2021], namely k -nearest neighbors ( k -NN).

Original: Specifically, for any s, a ∈ S × A , weestimate model error by
Revised: Specifically, for any s, a ∈ S × A , we estimate model error by

Original: [Equation]
Revised: [Equation]

Original: where d : S ×
Revised: where d : S ×

Original: A × S × A (cid:55)→ R + is a distance metric, and NN k ( s, a ) is the set of k -nearest neighborsof ( s, a ) in D n according to the distance metric d .
Revised: A × S × A (cid:55)→ R + is a distance metric, and NN k ( s, a ) is the set of k -nearest neighbors of ( s, a ) in D n according to the distance metric d .

Original: A question arises: how to choose d ?
Revised: A question arises: how to choose d ?

Original: Using the Euclidean distance in ambient (state-action) spaceis usually a bad choice (e.g., the ℓ 2 distance between natural images is not necessarily meaningful).
Revised: Using the Euclidean distance in ambient (state-action) space is usually a bad choice (e.g., the ℓ 2 distance between natural images is not necessarily meaningful).

Original: Moreover, to correctly measure the error, model dynamics should be somehow taken into consideration.
Revised: Moreover, to correctly measure the error, model dynamics should be somehow taken into consideration.

Original: We therefore consider an alternative approach which leverages the latent space of a variationalforward model, as described next.
Revised: We therefore consider an alternative approach which leverages the latent space of a variational forward model, as described next.

Original: 4.1 A Variational Latent Model of Dynamics
Revised: 4.1 A Variational Latent Model of Dynamics

Original: We begin by modeling ˆ P ( s ′ | s, a ) using a generative latent model.
Revised: We begin by modeling ˆ P ( s ′ | s, a ) using a generative latent model.

Original: Specifically, we consider alatent model which consists of an encoder E : S × A (cid:55)→ B ( Z ) and a decoder f D : Z (cid:55)→ B ( S ) , where B ( X ) is set of probability measures on the Borel sets of X .While the encoder E learns a [Table] [Table]
Revised: Specifically, we consider a latent model which consists of an encoder E : S × A (cid:55)→ B ( Z ) and a decoder f D : Z (cid:55)→ B ( S ) , where B ( X ) is set of probability measures on the Borel sets of X .

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table]

Original: [Table] [Table] [Table] latent representation of s, a , the decoder f D estimates the next state s ′ according to P ( ·| s, a ) .
Revised: While the encoder E learns a latent representation of s, a , the decoder f D estimates the next state s ′ according to P ( ·| s, a ) .

Original: Thismodel corresponds to the decomposition P ( ·| s, a ) =
Revised: This model corresponds to the decomposition P ( ·| s, a ) =

Original: f D ( ·| E ( s, a )) .
Revised: f D ( ·| E ( s, a )) .

Original: Such a model can be trainedby maximizing the Evidence Lower BOund (ELBO, Kingma and Welling [2013]) over the data. 
Revised: Such a model can be trained by maximizing the Evidence Lower BOund (ELBO, Kingma and Welling [2013]) over the data. [Table]

Original: That is, given a prior P ( z ) , we model E ϕ , f D,θ as parametric functions and maximize the ELBO,
Revised: That is, given a prior P ( z ) , we model E ϕ , f D,θ as parametric functions and maximize the ELBO,

Original: [Equation]
Revised: [Equation]

Original: s
Revised: P

Original: [Equation]
Revised: [Equation]

Original: || P [Equation] z )) We refer the reader to the appendixfor an exhaustive overview of training VAEs by maximum likelihood and the ELBO.
Revised: We refer the reader to the appendix for an exhaustive overview of training VAEs by maximum likelihood and the ELBO.

Original: Recall that we wish to find a good metric for estimating model error.
Revised: Recall that we wish to find a good metric for estimating model error.

Original: Having learned a latent modelfor ˆ P ( s ′ | s, a ) , its latent space Z can be used to define the metric d in Equation (3), i.e., the Euclideandistance between latent representations of state-action pairs.
Revised: Having learned a latent model for ˆ P ( s ′ | s, a ) , its latent space Z can be used to define the metric d in Equation (3), i.e., the Euclidean distance between latent representations of state-action pairs.

Original: Unfortunately, as was previously shown[Arvanitidis et al., 2018], latent codes in variational models contain sharp discontinuities, rendering
Revised: Unfortunately, as was previously shown [Arvanitidis et al., 2018], latent codes in variational models contain sharp discontinuities, rendering

Original: Euclidean distances in latent space unreliable and inaccurate (as we will also demonstrate in ourexperiments).
Revised: Euclidean distances in latent space unreliable and inaccurate (as we will also demonstrate in our experiments).

Original: Instead, we propose to use the natural induced metric of our latent model, as describedin the following subsection.
Revised: Instead, we propose to use the natural induced metric of our latent model, as described in the following subsection.

Original: 4.2 The Pullback Metric of Model Dynamics
Revised: 4.2 The Pullback Metric of Model Dynamics

Original: In this part we define the metric d that will be used to approximate model error in Equation (3).
Revised: In this part we define the metric d that will be used to approximate model error in Equation (3).

Original: Specifically, we consider a Riemannian submanifold defined by a latent space Z and observationfunction f , which induces minimum energy in the ambient space.
Revised: Specifically, we consider a Riemannian submanifold defined by a latent space Z and observation function f , which induces minimum energy in the ambient space.

Original: We will later choose Z to be thelatent space of our variational model (i.e., encoded state-action) and f to be the decoder function f Dof next state transitions.
Revised: We will later choose Z to be the latent space of our variational model (i.e., encoded state-action) and f to be the decoder function f D of next state transitions.

Original: We define the distance metric formally below.
Revised: We define the distance metric formally below.

Original: Definition 2.
Revised: Definition 2.

Original: We define a Riemannian submanifold ( M Z , g Z ) by a differential function f : Z (cid:55)→ Sand latent space Z such that
Revised: We define a Riemannian submanifold ( M Z , g Z ) by a differential function f : Z (cid:55)→ S and latent space Z such that

Original: [Equation]
Revised: [Equation]

Original: A similar metric has been used in previous work on generative latent models [Chen et al., 2018,
Revised: A similar metric has been used in previous work on generative latent models [Chen et al., 2018,

Original: Arvanitidis et al., 2018].
Revised: Arvanitidis et al., 2018].

Original: By choosing f to be the decoder function f D , latent codes that are close w.r.t.d Z induce curves of minimal energy in the ambient observation space (i.e., next state).
Revised: By choosing f to be the decoder function f D , latent codes that are close w.r.t. d Z induce curves of minimal energy in the ambient observation space (i.e., next state).

Original: This metric isclosely related to the pullback metric (see Section 2.2), as shown by the following proposition.
Revised: This metric is closely related to the pullback metric (see Section 2.2), as shown by the following proposition.

Original: Proposition 1.
Revised: Proposition 1.

Original: Let ( M Z , g Z ) as defined above.
Revised: Let ( M Z , g Z ) as defined above.

Original: Then G f ( z ) =
Revised: Then G f ( z ) =

Original: J T f ( z ) J f ( z ) , for any z ∈ Z .
Revised: J Tf ( z ) J f ( z ) , for any z ∈ Z .

Original: Indeed, Proposition 1 shows us that G f is a pullback metric.
Revised: Indeed, Proposition 1 shows us that G f is a pullback metric.

Original: Particularly Z and J f define the structureof the ambient observation space X (in our case, next state transitions).
Revised: Particularly Z and J f define the structure of the ambient observation space X (in our case, next state transitions).

Original: By choosing f to be the decoder function f D , the metric G f D becomes stochastic, complicatinganalysis.
Revised: By choosing f to be the decoder function f D , the metric G f D becomes stochastic, complicating analysis.

Original: Instead, as proposed and analyzed in Arvanitidis et al.
Revised: Instead, as proposed and analyzed in Arvanitidis et al.

Original: [2018], we use the expected pullbackmetric E
Revised: [2018], we use the expected pullback metric E

Original: [ G Z ] as an approximation of the underlying stochastic metric.
Revised: [ G Z ] as an approximation of the underlying stochastic metric.

Original: Similar to previous work onvariational models, we use a normally distributed decoder to define the output.
Revised: Similar to previous work on variational models, we use a normally distributed decoder to define the output.

Original: Using Proposition 1,we have the following result (see Appendix for proof).
Revised: Using Proposition 1, we have the following result (see Appendix for proof).

Original: Theorem 1.
Revised: Theorem 1.

Original: [Arvanitidis et al.
Revised: [Arvanitidis et al.

Original: [2018]] Assume f D ( ·| z ) ∼ N ( µ ( z ) , σ ( z ) I ) .
Revised: [2018]] Assume f D ( ·| z ) ∼ N ( µ ( z ) , σ ( z ) I ) .

Original: Then
Revised: Then

Original: [Equation]
Revised: [Equation]

Original: where G µ ( z ) =
Revised: where G µ ( z ) =

Original: J Tµ ( z ) J µ ( z ) and G σ ( z ) =
Revised: J Tµ ( z ) J µ ( z ) and G σ ( z ) =

Original: J Tσ ( z ) J σ ( z ) .
Revised: J Tσ ( z ) J σ ( z ) .

Original: Given an embedded latent space Z , the expected metric in Equation (4) gives us a sense of thetopology of the latent space manifold induced by f D .
Revised: Given an embedded latent space Z , the expected metric in Equation (4) gives us a sense of the topology of the latent space manifold induced by f D .

Original: The terms G µ = J Tµ J µ and G σ = J Tσ J σ are infact the induced pullback metrics of µ and σ , respectively.
Revised: The terms G µ = J Tµ J µ and G σ = J Tσ J σ are in fact the induced pullback metrics of µ and σ , respectively.

Original: [Figure]
Revised: [Figure]

Original: 4.3 Capturing Epistemic and Aleatoric Uncertainty
Revised: 4.3 Capturing Epistemic and Aleatoric Uncertainty

Original: The previously proposed encoder-decoder model induces a metric which captures the structure ofthe learned dynamics.
Revised: The previously proposed encoder-decoder model induces a metric which captures the structure of the learned dynamics.

Original: However, the decoder variance, σ ( z ) , does not differentiate between aleatoricuncertainty (environment dynamics) and epistemic uncertainty (missing data).
Revised: However, the decoder variance, σ ( z ) , does not differentiate between aleatoric uncertainty (environment dynamics) and epistemic uncertainty (missing data).

Original: We propose two methods to enrich the metric in Equation (4) in order to achieve a better estimate ofuncertainty.
Revised: We propose two methods to enrich the metric in Equation (4) in order to achieve a better estimate of uncertainty.

Original: First, by using an ensemble of M decoder functions { f D,i } Mi =1 trained using standardbootstrap techniques [Efron, 1982], we capture the traditional epistemic uncertainty of the decoderparameters.
Revised: First, by using an ensemble of M decoder functions { f D,i } Mi =1 trained using standard bootstrap techniques [Efron, 1982], we capture the traditional epistemic uncertainty of the decoder parameters.

Original: Second, to correctly distinguish epistemic and aleatoric uncertainty, we add a latentforward function to our previously proposed variational model.
Revised: Second, to correctly distinguish epistemic and aleatoric uncertainty, we add a latent forward function to our previously proposed variational model.

Original: Specifically, our latent model consistsof an encoder E : S × A (cid:55)→ B ( Z ) , forward model f F : Z (cid:55)→ B ( X ) and decoder functionsf D,i : X (cid:55)→ B ( S ) such that P ( ·| s, a ) =
Revised: Specifically, our latent model consists of an encoder E : S × A (cid:55)→ B ( Z ) , forward model f F : Z (cid:55)→ B ( X ) and decoder functions f D,i : X (cid:55)→ B ( S ) such that P ( ·| s, a ) =

Original: f D,i ( ·| x ) , and x ∼ f F ( ·| E
Revised: f D,i ( ·| x ) , and x ∼ f F ( ·| E

Original: s, a )) .
Revised: s, a )) .

Original: This structure enablesus to capture the aleatoric uncertainty under the forward transition model f F , and the epistemicuncertainty using the decoders f D,i .
Revised: This structure enables us to capture the aleatoric uncertainty under the forward transition model f F , and the epistemic uncertainty using the decoders f D,i .

Original: That is, for f F ( ·| z ) ∼ N ( µ F ( z ) , σ F ( z ) I ) , the variance modelσ F ( z ) captures the stochasticity in model dynamics.
Revised: That is, for f F ( ·| z ) ∼ N ( µ F ( z ) , σ F ( z ) I ) , the variance model σ F ( z ) captures the stochasticity in model dynamics.

Original: This decomposition is also helpful wheneverone wants to train an agent in latent space (e.g., for planning Schrittwieser et al.
Revised: This decomposition is also helpful whenever one wants to train an agent in latent space (e.g., for planning Schrittwieser et al.

Original: Next, we turn to analyze the pullback metric induced by the proposed forward transition model.
Revised: Next, we turn to analyze the pullback metric induced by the proposed forward transition model.

Original: As both f F and { f D,i } mi =1 are stochastic (capturing epistemic and aleatoric uncertainty), the resultof Theorem 1 cannot be directly applied to their composition.
Revised: As both f F and { f D,i } mi =1 are stochastic (capturing epistemic and aleatoric uncertainty), the result of Theorem 1 cannot be directly applied to their composition.

Original: The following proposition providesan analytical expression for the expected pullback metric of a sampled next state and a uniformlysampled decoder (the proof is given in the appendix).
Revised: The following proposition provides an analytical expression for the expected pullback metric of a sampled next state and a uniformly sampled decoder (the proof is given in the appendix).

Original: Theorem 2.
Revised: Theorem 2.

Original: Assume
Revised: Assume

Original: [Equation]

Original: [Equation]

Original: [Equation]

Original: ∼ N

Original: [Equation]
Revised: [Equation]

Original: µ iD ( x ) , σ iD ( x ) I ) ,

Original: U ∼ Unif { 1 , . . .
Revised: U ∼ Unif { 1 , . . .

Original: , M } .
Revised: , M } .

Original: Then, the expected pullback metric of the composite function ( f D,U ◦ f F ) isgiven by
Revised: Then, the expected pullback metric of the composite function ( f D,U ◦ f F ) is given by

Original: [Equation]
Revised: [Equation]

Original: Unlike the metric in Equation (4), the composite metric distorts the decoder metric with Jacobianmatrices of the forward model statistics.
Revised: Unlike the metric in Equation (4), the composite metric distorts the decoder metric with Jacobian matrices of the forward model statistics.

Original: It takes into account both the aleatoric and epistemicuncertainty through the forward model as well as ensemble of decoders.
Revised: It takes into account both the aleatoric and epistemic uncertainty through the forward model as well as ensemble of decoders.

Original: As a special case we notethe metric for the case of deterministic model dynamics.
Revised: As a special case we note the metric for the case of deterministic model dynamics.

Original: Corollary 1.
Revised: Corollary 1.

Original: Assume deterministic model dynamics, i.e., x = f F ( z ) , and without loss of generalityassume f F ≡ I .
Revised: Assume deterministic model dynamics, i.e., x = f F ( z ) , and without loss of generality assume f F ≡ I .

Original: Then, the expected pullback metric of Theorem 2 is given by E (cid:104)
Revised: Then, the expected pullback metric of Theorem 2 is given by E (cid:104)

Original: G f D,U ◦ f F ( z) (cid:105) =
Revised: G f D U ◦ f F ( z (cid:105) =

Original: 1 M (cid:80) Mi

Original: =1 J Tµ iD ( z ) J µ iD ( z ) + J Tσ iD ( z ) J σ iD ( z ) .

Original: 4.4 GELATO: Geometrically Enriched LATent model for 
Revised: 4.4 GELATO: Geometrically Enriched LATent model for Offline reinforcement learning

Original: Having defined our metric, we are now ready to leverage it in a model based offline RL framework.
Revised: Having defined our metric, we are now ready to leverage it in a model based offline RL framework.

Original: Specifically, provided a dataset D n = { ( s i , a i , r i , s ′ i ) } ni =

Revised: Specifically, provided a dataset D n = { ( s i , a i , r i , s ′ i ) } ni =

Original: we train the variational latent forwardmodel described in the previous section.
Revised: we train the variational latent forward model described in the previous section.

Original: Algorithm 1 presents GELATO, our proposed approach.
Revised: Algorithm 1 presents GELATO, our proposed approach.

Original: In GELATO, we estimate model error bymeasuring the distance of a new sample to the data manifold.
Revised: In GELATO, we estimate model error by measuring the distance of a new sample to the data manifold.

Original: We construct the reward-penalized MDPfor which the error acts as a pessimistic regularizer.
Revised: We construct the reward-penalized MDP for which the error acts as a pessimistic regularizer.

Original: Finally, we train an RL agent over the pessimistic
Revised: Finally, we train an RL agent over the pessimistic

Original: MDP with transition ˆ P ( ·| s, a ) and reward r ( s, a ) − λU ( s, a ) .
Revised: MDP with transition ˆ P ( ·| s, a ) and reward r ( s, a ) − λU ( s, a ) .

Original: By achieving an improved estimate formodel error the model-based pessimistic approach can significantly improve performance, as shownin the following section.
Revised: By achieving an improved estimate for model error the model-based pessimistic approach can significantly improve performance, as shown in the following section.

Original:  Experiments
Revised:  Experiments

Original: This section is dedicated to quantitatively and qualitatively understand the benefits of our proposedmetric and method.
Revised: This section is dedicated to quantitatively and qualitatively understand the benefits of our proposed metric and method.

Original: We validate two principal claims: (1) Our metric captures inherent characteristics of model dynamics.
Revised: We validate two principal claims: (1) Our metric captures inherent characteristics of model dynamics.

Original: We demonstrate this by visualizing state-action geodesics of a toy gridworld problem and a multi-agent autonomous driving task.
Revised: We demonstrate this by visualizing state-action geodesics of a toy grid world problem and a multi-agent autonomous driving task.

Original: We show that curves of minimum energyin ambient space indeed capture intrinsic properties of the problem.
Revised: We show that curves of minimum energy in ambient space indeed capture intrinsic properties of the problem.

Original: (2) Our metric provides animproved OOD uncertainty estimate for offline RL.
Revised: (2) Our metric provides an improved OOD uncertainty estimate for offline RL.

Original: We compare the traditionally used bootstrapensemble method to our approach, which leverages our pullback metric in a nonparametric nearestneighbors approach.
Revised: We compare the traditionally used bootstrap ensemble method to our approach, which leverages our pullback metric in a nonparametric nearest neighbors approach.

Original: We also compare our method to simple use of Euclidean distances in latentspace.
Revised: We also compare our method to simple use of Euclidean distances in latent space.

Original: We run extensive experiments on continuous control and autonomous driving benchmarks.
Revised: We run extensive experiments on continuous control and autonomous driving benchmarks.

Original: We show that our metric achieves significantly improved performance in tasks for which geodesicsare non-euclidean.
Revised: We show that our metric achieves significantly improved performance in tasks for which geodesics are non-euclidean.

Original: 5.1 Metric Visualization
Revised: 5.1 Metric Visualization

Original: Four Rooms.
Revised: Four Rooms.

Original: To better understand the inherent structure of our metric, we constructed a grid-worldenvironment for visualizing our proposed latent representation and metric.
Revised: To better understand the inherent structure of our metric, we constructed a grid-world environment for visualizing our proposed latent representation and metric.

Original: The 15 × 15 environment,as depicted in Figure 2, consists of four rooms, with impassable obstacles in their centers.
Revised: The 15 × 15 environment, as depicted in Figure 2, consists of four rooms, with impassable obstacles in their centers.

Original: The agent,residing at some position ( x, y ) ∈ [ − 1 , 1] 2 in the environment can take one of four actions: up, down,left, or right – moving the agent 1 , 2 or 3 steps (uniformly distributed) in that direction.
Revised: The agent, residing at some position ( x, y ) ∈ [ − 1 , 1] 2 in the environment can take one of four actions: up, down, left, or right – moving the agent 1 , 2 or 3 steps (uniformly distributed) in that direction.

Original: We collecteda dataset of 10000 samples, taking random actions at random initializations of the environment.
Revised: We collected a dataset of 10000 samples, taking random actions at random initializations of the environment.

Original: Theambient state space was represented by the position of the agent – a vector of dimension 2, normalizedto values in [ − 1 , 1] .
Revised: The ambient state space was represented by the position of the agent – a vector of dimension 2, normalized to values in [ − 1 , 1] .

Original: Finally, we trained a variational latent model with latent dimension d Z = 2 .
Revised: Finally, we trained a variational latent model with latent dimension d Z = 2 .

Original: [Table]
Revised: [Table]

Original: We used a standard encoder z ∼ N ( µ θ ( s ) , σ θ ( s )) and decoder s ′ ∼ N ( µ ϕ ( z ) , σ ϕ ( z )) representedby neural networks trained end-to-end using the evidence lower bound.
Revised: We used a standard encoder z ∼ N ( µ θ ( s ) , σ θ ( s )) and decoder s ′ ∼ N ( µ ϕ ( z ) , σ ϕ ( z )) represented by neural networks trained end-to-end using the evidence lower bound.

Original: We refer the reader to theappendix for an exhaustive description of the training procedure.
Revised: We refer the reader to the appendix for an exhaustive description of the training procedure

Revised: [Table]

Original: The latent space output of our model is depicted by yellow markers in Figure 2a.
Revised: The latent space output of our model is depicted by yellow markers in Figure 2a.

Original: Indeed, the latentembedding consists of four distinctive clusters, structured in a similar manner as our grid-worldenvironment.
Revised: Indeed, the latent embedding consists of four distinctive clusters, structured in a similar manner as our grid-world environment.

Original: Interestingly, the distortion of the latent space accurately depicts an intuitive notion ofdistance between states.
Revised: Interestingly, the distortion of the latent space accurately depicts an intuitive notion of distance between states.

Original: As such, rooms are distinctively separated, with fair distance between eachcluster.
Revised: As such, rooms are distinctively separated, with fair distance between each cluster.

Original: States of pathways between rooms clearly separate the room clusters, forming a topologywith four discernible bottlenecks.
Revised: States of pathways between rooms clearly separate the room clusters, forming a topology with four discernible bottlenecks.

Original: In addition to the latent embedding, Figure 2a depicts the geometric volume measure (cid:112) det( G f D )of the trained pullback metric induced by f D .
Revised: In addition to the latent embedding, Figure 2a depicts the geometric volume measure (cid:112) det( G f D ) of the trained pullback metric induced by f D .

Original: This quantity demonstrates the effective geodesicdistances between states in the decoder-induced submanifold.
Revised: This quantity demonstrates the effective geodesic distances between states in the decoder-induced submanifold.

Original: Indeed geodesics between data pointsto points outside of the data manifold (i.e., outside of the red region), attain high values as integralsover areas of high magnitude.
Revised: Indeed geodesics between data points to points outside of the data manifold (i.e., outside of the red region), attain high values as integrals over areas of high magnitude.

Original: In contrast, geodesics near the data manifold would low values.
Revised: In contrast, geodesics near the data manifold would low values.

Original: We visualize the decoder-induced geodesic distance and compare it to the latent Euclidean distancein Figures 2b and 2c, respectively.
Revised: We visualize the decoder-induced geodesic distance and compare it to the latent Euclidean distance in Figures 2b and 2c, respectively.

Original: The plots depict the normalized distances of all states to the statemarked by a yellow square.
Revised: The plots depict the normalized distances of all states to the state marked by a yellow square.

Original: Evidently, the geodesic distance captures a better notion of distance inthe said environment, correctly exposing the “land distance" in ambient space.
Revised: Evidently, the geodesic distance captures a better notion of distance in the said environment, correctly exposing the “land distance" in ambient space.

Original: As expected, statesresiding in the bottom-right room are farthest from the source state, as reaching them comprisesof passing through at least two bottleneck states.
Revised: As expected, states residing in the bottom-right room are farthest from the source state, as reaching them comprises of passing through at least two bottleneck states.

Original: In contrast, the latent Euclidean distance doesnot properly capture these geodesics, exhibiting a similar distribution of distances in other rooms.
Revised: In contrast, the latent Euclidean distance does not properly capture these geodesics, exhibiting a similar distribution of distances in other rooms.

Original: Nevertheless, both geodesic and Euclidean distances reveal the intrinsic topological structure of theenvironment, that of which is not captured by the extrinsic coordinates ( x, y ) ∈ [ − 1 , 1] 2 .
Revised: Nevertheless, both geodesic and Euclidean distances reveal the intrinsic topological structure of the environment, that of which is not captured by the extrinsic coordinates ( x, y ) ∈ [ − 1 , 1] 2 .

Original: Particularly,the state coordinates ( x, y ) would wrongly assign short distances to states across impassible walls orobstacles, i.e., measuring the “air distance".
Revised: Particularly, the state coordinates ( x, y ) would wrongly assign short distances to states across impassible walls or obstacles, i.e., measuring the “air distance".

Original: Intersection.
Revised: Intersection.

Original: We visualized our metric in the intersection environment proposed in Leurent [2018].
Revised: We visualized our metric in the intersection environment proposed in Leurent [2018].

Original: Figure 3 compares the Euclidean and geodesic distances of a partially trained agent.
Revised: Figure 3 compares the Euclidean and geodesic distances of a partially trained agent.

Original: Unlike theprevious toy example, to visualize the inherent manifolds we used t-SNE [Van der Maaten and
Revised: Unlike the previous toy example, to visualize the inherent manifolds we used t-SNE [Van der Maaten and

Original: Hinton, 2008] projections computed with Euclidean distance and compared them to the projectioncomputed with geodesic distance, i.e., curves of minimum energy in ambient space (Definition 2).
Revised: Hinton, 2008] projections computed with Euclidean distance and compared them to the projection computed with geodesic distance, i.e., curves of minimum energy in ambient space (Definition 2).

Original: Indeed, the geodesics captured the inherent structure of the environment, whereas Euclidean distancesonly managed to capture general clusters.
Revised: Indeed, the geodesics captured the inherent structure of the environment, whereas Euclidean distances only managed to capture general clusters.

Original: These suggest that Euclidean distance might not berepresentative for measuring distance in latent space, as will also become evident by our experimentsin the following subsections.
Revised: These suggest that Euclidean distance might not be representative for measuring distance in latent space, as will also become evident by our experiments in the following subsections.

Original: 5.2 Datasets and Implementation Details
Revised: 5.2 Datasets and Implementation Details

Original: We used D4RL [Fu et al., 2020] and the autonomous vehicle environments highway-env [Leurent,2018] as benchmarks for all of our experiments.
Revised: We used D4RL [Fu et al., 2020] and the autonomous vehicle environments highway-env [Leurent, 2018] as benchmarks for all of our experiments.

Original: We tested GELATO on three Mujoco [Todorov et al. ,2012] environments (Hopper, Walker2d, Halfcheetah) on datasets generated by a single policy and amixture of two policies.
Revised: We tested GELATO on three Mujoco [Todorov et al., 2012] environments (Hopper, Walker2d, Halfcheetah) on datasets generated by a single policy and a mixture of two policies.

Original: Specifically, we used datasets generated by a random agent (1M samples), apartially trained agent, i.e, medium agent (1M samples), and a mixture of partially trained and expert agents (2M samples).
Revised: Specifically, we used datasets generated by a random agent (1M samples), a partially trained agent, i.e, medium agent (1M samples), and a mixture of partially trained and expert agents (2M samples).

Original: For autonomous driving, we tested GELATO on four environments (Highway,
Revised: For autonomous driving, we tested GELATO on four environments (Highway,

Original: Roundabout, Intersection, Racetrack), on datasets containing five episodes generated by a partiallytrained agent.
Revised: Roundabout, Intersection, Racetrack), on datasets containing five episodes generated by a partially trained agent.

Original: We also tested a faster ( × 15 speedup) variant of the Highway environment, as well asa harder instantiation of the Intersection environment in which the number of cars was tripled (furtherdetails can be found in the appendix).
Revised: We also tested a faster ( × 15 speedup) variant of the Highway environment, as well as a harder instantiation of the Intersection environment in which the number of cars was tripled (further details can be found in the appendix).

Original: We trained our variational latent model in two phases.
Revised: We trained our variational latent model in two phases.

Original: First, the model was fully trained using acalibrated Gaussian decoder [Rybkin et al., 2020].
Revised: First, the model was fully trained using a calibrated Gaussian decoder [Rybkin et al., 2020].

Original: Specifically, a maximum-likelihood estimate ofthe variance was used σ ∗ =
Revised: Specifically, a maximum-likelihood estimate of the variance was used σ ∗ =

Original: MSE ( µ, ˆ µ ) ∈ arg max σ N (ˆ µ | µ, σ 2 I ) .
Revised: MSE ( µ, ˆ µ ) ∈ arg max σ N (ˆ µ | µ, σ 2 I ) .

Original: Then, in the second stage we fitthe variance decoder networks.
Revised: Then, in the second stage we fit the variance decoder networks.

Original: Hyperparameters for training are found in the appendix.
Revised: Hyperparameters for training are found in the appendix.

Original: In order to practically estimate the geodesic distance in Algorithm 1, we defined a parametric curvein latent space and used gradient descent to minimize the curve’s energy.
Revised: In order to practically estimate the geodesic distance in Algorithm 1, we defined a parametric curve in latent space and used gradient descent to minimize the curve’s energy.

Original: The resulting curve andpullback metric were then used to calculate the geodesic distance by a numerical estimate of thecurve length (See Appendix for an exhaustive overview of the estimation method).
Revised: The resulting curve and pullback metric were then used to calculate the geodesic distance by a numerical estimate of the curve length (See Appendix for an exhaustive overview of the estimation method).

Original: We used FAISS [Johnson et al., 2019] for efficient GPU-based k -nearest neighbors calculation.
Revised: We used FAISS [Johnson et al., 2019] for efficient GPU-based k -nearest neighbors calculation.

Original: Weset k = 5 neighbors for the penalized reward (Equation (3)).
Revised: We set k = 5 neighbors for the penalized reward (Equation (3)).

Original: Finally, we used a variant of Soft
Revised: Finally, we used a variant of Soft

Original: Learning, as proposed by Yu et al.
Revised: Learning, as proposed by Yu et al.

Original: [2020] as our RL algorithm for the continuous control benchmarks,and PPO [Schulman et al., 2017] for the autonomous driving tasks.
Revised: [2020] as our RL algorithm for the continuous control benchmarks, and PPO [Schulman et al., 2017] for the autonomous driving tasks.

Original: All agents were trained for 1Msteps (for continuous control benchmarks) and 350K steps (for the driving benchmarks), using asingle GPU (RTX 2080), and averaged over 5 seeds (see Appendix for more details).
Revised: All agents were trained for 1M steps (for continuous control benchmarks) and 350K steps (for the driving benchmarks), using a single GPU (RTX 2080), and averaged over 5 seeds (see Appendix for more details).

Original: 5.3 Results
Revised: 5.3 Results

Original: D4RL.
Revised: D4RL.

Original: Results for the continuous control domains are shown in Table 1.
Revised: Results for the continuous control domains are shown in Table 1.

Original: We performed experimentsto analyze GELATO on various continuous control datasets.
Revised: We performed experiments to analyze GELATO on various continuous control datasets.

Original: We compared GELATO to contemporarymodel-based offline RL approaches; namely, MOPO [Yu et al., 2020] and MBPO [Janner et al., 2019],as well as the standard baselines of SAC [Haarnoja et al., 2018] and imitation (behavioral cloning,
Revised: We compared GELATO to contemporary model-based offline RL approaches; namely, MOPO [Yu et al., 2020] and MBPO [Janner et al., 2019], as well as the standard baselines of SAC [Haarnoja et al., 2018] and imitation (behavioral cloning,

Original: Bain and Sammut [1995], Ross et al.
Revised: Bain and Sammut [1995], Ross et al.

Original: We found performance increase on most domains, and most significantly in the medium domain,i.e., partially trained agent.
Revised: We found performance increase on most domains, and most significantly in the medium domain, i.e., partially trained agent.

Original: Since the medium dataset contained average behavior, combining thenonparametric nearest-neighbor uncertainty method with the bootstrap of decoders benefited theagent’s overall performance.
Revised: Since the medium dataset contained average behavior, combining the nonparametric nearest-neighbor uncertainty method with the bootstrap of decoders benefited the agent’s overall performance.

Original: In addition, GELATO with the latent ℓ 2 distance metric performedwell on many of the benchmarks.
Revised: In addition, GELATO with the latent ℓ 2 distance metric performed well on many of the benchmarks.

Original: We conjecture this is due to the inherent Euclidean nature of thecontinuous control benchmarks.
Revised: We conjecture this is due to the inherent Euclidean nature of the continuous control benchmarks.

Original: Unlike the embedding for the autonomous driving benchmarks(Figure 3), we found the D4RL data to project similarly when ℓ 2 and geodesic distances were used(we provide plots of these embeddings in the Appendix).
Revised: Unlike the embedding for the autonomous driving benchmarks (Figure 3), we found the D4RL data to project similarly when ℓ 2 and geodesic distances were used (we provide plots of these embeddings in the Appendix).

Original: Highway-Env.
Revised: Highway-Env.

Original: Figure 4 shows results for the autonomous driving benchmarks in highway-env.
Revised: Figure 4 shows results for the autonomous driving benchmarks in highway-env.

Original: Incontrast to the continuous control benchmarks, we found a significant improvement of our metricon the autonomous driving benchmarks compared to standard uncertainty estimation methods aswell as the latent Euclidean distance.
Revised: In contrast to the continuous control benchmarks, we found a significant improvement of our metric on the autonomous driving benchmarks compared to standard uncertainty estimation methods as well as the latent Euclidean distance.

Original: We credit this improvement to the non-euclidean nature of theenvironments, as previously described in Figure 3.
Revised: We credit this improvement to the non-euclidean nature of the environments, as previously described in Figure 3.

Original: While Euclidean distances were useful in the
Revised: While Euclidean distances were useful in the

Original: Mujoco environments, they performed distinctly worse in the autonomous driving environments.
Revised: Mujoco environments, they performed distinctly worse in the autonomous driving environments.

Original: Our results emphasize the importance of OOD uncertainty estimations methods in reinforcementlearning on various types of datasets.
Revised: Our results emphasize the importance of OOD uncertainty estimations methods in reinforcement learning on various types of datasets.

Original: While robotic control tasks provided useful insights, they didnot capture the non-euclidean nature inherent in alternative tasks, such as autonomous driving.
Revised: While robotic control tasks provided useful insights, they did not capture the non-euclidean nature inherent in alternative tasks, such as autonomous driving.

Original:  Related Work
Revised:  Related Work

Original: Offline Reinforcement Learning.
Revised: Offline Reinforcement Learning.

Original: The field of offline RL has recently received much attention asseveral algorithmic approaches were able to surpass standard off-policy algorithms.
Revised: The field of offline RL has recently received much attention as several algorithmic approaches were able to surpass standard off-policy algorithms.

Original: Value-basedonline algorithms do not perform well under highly off-policy batch data [Fujimoto et al., 2019,
Revised: Value-based online algorithms do not perform well under highly off-policy batch data [Fujimoto et al., 2019,

Original: Kumar et al., 2019, Fu et al., 2019, Fedus et al., 2020, Agarwal et al., 2020], much due to issues withbootstrapping from out-of-distribution (OOD) samples.
Revised: Kumar et al., 2019, Fu et al., 2019, Fedus et al., 2020, Agarwal et al., 2020], much due to issues with bootstrapping from out-of-distribution (OOD) samples.

Original: These issues become more prominent in theoffline setting, as new samples cannot be acquired.
Revised: These issues become more prominent in the offline setting, as new samples cannot be acquired.

Original: Several works on offline RL have shown improved performance on standard continuous controlbenchmarks [Laroche et al., 2019, Kumar et al., 2019, Fujimoto et al., 2019, Chen et al., 2020b,
Revised: Several works on offline RL have shown improved performance on standard continuous control benchmarks [Laroche et al., 2019, Kumar et al., 2019, Fujimoto et al., 2019, Chen et al., 2020b,

Original: Swazinna et al., 2020, Kidambi et al., 2020, Yu et al., 2020].
Revised: Swazinna et al., 2020, Kidambi et al., 2020, Yu et al., 2020].

Original: This work focused on model-basedapproaches [Yu et al., 2020, Kidambi et al., 2020], in which the agent is incentivized to remain close toareas of low uncertainty.
Revised: This work focused on model-based approaches [Yu et al., 2020, Kidambi et al., 2020], in which the agent is incentivized to remain close to areas of low uncertainty.

Original: Our work focused on controlling uncertainty estimation in high dimensionalenvironments.
Revised: Our work focused on controlling uncertainty estimation in high dimensional environments.

Original: Our methodology utilized recent success on the geometry of deep generative models[Arvanitidis et al., 2018, 2020], proposing an alternative approach to uncertainty estimation.
Revised: Our methodology utilized recent success on the geometry of deep generative models [Arvanitidis et al., 2018, 2020], proposing an alternative approach to uncertainty estimation.

Original: Representation Learning.
Revised: Representation Learning.

Original: Representation learning seeks to find an appropriate representation ofdata for performing a machine-learning task
Revised: Representation learning seeks to find an appropriate representation of data for performing a machine-learning task

Original: [Goodfellow et al., 2016].
Revised: [Goodfellow et al., 2016].

Original: Variational Auto Encoders[Kingma and Welling, 2013, Rezende et al., 2014] have been a popular representation learningtechnique, particularly in unsupervised learning regimes [Chen et al., 2016, Van Den Oord et al. ,2017, Hsu et al., 2017, Serban et al., 2017, Engel et al., 2017, Bojanowski et al., 2018, Ding et al. ,2020], though also in supervised learning and reinforcement learning [Hausman et al., 2018, Li et al. ,2019, Petangoda et al., 2019, Hafner et al., 2019].
Revised: Variational Auto Encoders [Kingma and Welling, 2013, Rezende et al., 2014] have been a popular representation learning technique, particularly in unsupervised learning regimes [Chen et al., 2016, Van Den Oord et al., 2017, Hsu et al., 2017, Serban et al., 2017, Engel et al., 2017, Bojanowski et al., 2018, Ding et al., 2020], though also in supervised learning and reinforcement learning [Hausman et al., 2018, Li et al., 2019, Petangoda et al., 2019, Hafner et al., 2019].

Original: Particularly, variational models have been shownable to derive successful behaviors in high dimensional benchmarks [Hafner et al., 2020].
Revised: Particularly, variational models have been shown able to derive successful behaviors in high dimensional benchmarks [Hafner et al., 2020].

Original: Various representation techniques in reinforcement learning have also proposed to disentanglerepresentation of both states [Engel and Mannor, 2001, Littman and Sutton, 2002, Stooke et al. ,2020, Zhu et al., 2020], and actions [Tennenholtz and Mannor, 2019, Chandak et al., 2019].
Revised: Various representation techniques in reinforcement learning have also proposed to disentangle representation of both states [Engel and Mannor, 2001, Littman and Sutton, 2002, Stooke et al., 2020, Zhu et al., 2020], and actions [Tennenholtz and Mannor, 2019, Chandak et al., 2019].

Original: Theseallow for the abstraction of states and actions to significantly decrease computation requirements, bye.g. , decreasing the effective dimensionality of the action space [Tennenholtz and Mannor, 2019].
Revised: These allow for the abstraction of states and actions to significantly decrease computation requirements, by e.g., decreasing the effective dimensionality of the action space [Tennenholtz and Mannor, 2019].

Original: Unlike previous work, GELATO is focused on a semiparametric approach for uncertainty estimation,enhancing offline reinforcement learning performance.
Revised: Unlike previous work, GELATO is focused on a semiparametric approach for uncertainty estimation, enhancing offline reinforcement learning performance.

Original:  Discussion and Future Work
Revised:  Discussion and Future Work

Original: This work presented a metric for model dynamics and its application to offline reinforcement learning.
Revised: This work presented a metric for model dynamics and its application to offline reinforcement learning.

Original: While our metric showed supportive evidence of improvement in model-based offline RL we notethat it was significantly slower – comparably, 5 times slower than using the decoder’s variance foruncertainty estimation.
Revised: While our metric showed supportive evidence of improvement in model-based offline RL we note that it was significantly slower – comparably, 5 times slower than using the decoder’s variance for uncertainty estimation.

Original: The apparent slowdown in performance was mostly due to computation of thegeodesic distance.
Revised: The apparent slowdown in performance was mostly due to computation of the geodesic distance.

Original: Improvement in this area may utilize techniques for efficient geodesic estimation.
Revised: Improvement in this area may utilize techniques for efficient geodesic estimation.

Original: We conclude by noting possible future applications of our work.
Revised: We conclude by noting possible future applications of our work.

Original: In Section 5.1 we demonstrated theinherent geometry our model had captured, its corresponding metric, and geodesics.
Revised: In Section 5.1 we demonstrated the inherent geometry our model had captured, its corresponding metric, and geodesics.

Original: Still, in thiswork we focused specifically on metrics related to the decoded state.
Revised: Still, in this work we focused specifically on metrics related to the decoded state.

Original: In fact, a similar derivationto Theorem 2 could be applied to other modeled statistics, e.g., Q-values, rewards, future actions,and more.
Revised: In fact, a similar derivation to Theorem 2 could be applied to other modeled statistics, e.g., Q-values, rewards, future actions, and more.

Original: Each distinct statistic would induce its own unique metric w.r.t.
Revised: Each distinct statistic would induce its own unique metric w.r.t.

Original: its respective probabilitymeasure.
Revised: its respective probability measure.

Original: Particularly, this concept may benefit a vast array of applications in continuous or largestate and action spaces.Checklist
Revised: Particularly, this concept may benefit a vast array of applications in continuous or large state and action spaces.

Original: For all authors...

Original: (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s

Original: contributions and scope?

Original: [Yes]

Original: (b) Did you describe the limitations of your work?

Original: [Yes]

Original: (c) Did you discuss any potential negative societal impacts of your work?

Original: [Yes]

Original: (d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?

Original: [Yes]

Original: If you are including theoretical results...

Original: (a) Did you state the full set of assumptions of all theoretical results?

Original: [Yes]

Original: (b) Did you include complete proofs of all theoretical results?

Original: [Yes]

Original: If you ran experiments...

Original: (a) Did you include the code, data, and instructions needed to reproduce the main experi-

Original: mental results (either in the supplemental material or as a URL)?

Original: [Yes]

Original: (b) Did you specify all the training details (e.g., data splits, hyperparameters,

Original: how they

Original: were chosen)?

Original: [Yes]

Original: (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?

Original: [Yes]

Original: (d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?

Original: [Yes]

Original: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...

Original: (a)

Original: If your work uses existing assets, did you cite the creators?

Original: [Yes]

Original: (b) Did you mention the license of the assets?

Original: [N/A]

Original: (c) Did you include any new assets either in the supplemental material or as a URL?

Original: [Yes]

Original: (d) Did you discuss whether and how consent was obtained from people whose data you’re

Original: using/curating?

Original: [N/A]

Original: (e) Did you discuss whether the data you are using/curating contains personally identifiableinformation or offensive content?

Original: [N/A]

Original: If you used crowdsourcing or conducted research with human subjects...

Original: (a) Did you include the full text of instructions given to participants and screenshots, if

Original: applicable?

Original: [N/A]

Original: (b) Did you describe any potential participant risks, with links to Institutional Review

Original: Board (IRB) approvals, if applicable?

Original: [N/A]

Original: (c) Did you include the estimated hourly wage paid to participants and the total amount

Original: spent on participant compensation?

Original: [N/A]