Uncertainty Estimation Using Riemannian Model Dynamics for Offline Reinforcement Learning
Authors: Guy Tennenholtz
Shie Mannor
Venue: NeurIPS 2022 Accept (2022)
Keywords: reinforcement learning
Riemannian geometry
uncertainty estimation

Abstract: Model-based offline reinforcement learning approaches generally rely on bounds of model error. Estimating these bounds is usually achieved through uncertainty estimation methods. In this work, we combine parametric and nonparametric methods for uncertainty estimation through a novel latent space based metric. 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. Our proposed metric measures both the quality of out-of-distribution samples as well as the discrepancy of examples in the data. We leverage our combined 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.

Uncertainty Estimation Using Riemannian Model Dynamics for Offline Reinforcement Learning
Guy Tennenholtz. Technion Institute of Technology
Shie Mannor
Technion Institute of Technology & Nvidia Research
Abstract
Model-based offline reinforcement learning approaches generally rely on bounds of model error. Estimating these bounds is usually achieved through uncertainty estimation methods. In this work, we combine parametric and nonparametric methods for uncertainty estimation through novel latent space based metric. In particular, we build upon recent advances in Riemannian geometry of generative models construct pullback metric of encoder-decoder based forward model. Our proposed metric measures both the quality of out-of-distribution samples well as the discrepancy of examples in the data. We leverage our method for uncertainty estimation in a pessimistic model-based framework, showing a significant improvement upon contemporary model-based offline approaches continuous control and autonomous driving benchmarks.
Introduction
Offline Reinforcement Learning (RL) Levine et al., 2020). a.k.a. batch-mode RL [Ernst et al. 2005 Riedmiller 2005 Fonteneau et al. 2013), involves learning policy from data sampled by potentally suboptimal policy. Offline RL seeks to surpass the average performance of the agents that generated the data. Traditional methodologies fall short in offline settings, causing overestimation of the return [Buckman et al.) 2020 Wang et al.) 2020 Zanette 2020).
One approach to overcome this in model-based settings is to penalize the return in out of distribution (OOD) regions, as depicted in Figure this manner, the agent is constrained to stay "near" areas of low model error, thereby limiting possible overestimation. However, reliable estimates of model error are key to the success such methods.
Estimating model eтror in OOD regions can be achieved through uncertainty estimation Yu et al. 2020). Methods of parametric uncertainty estimation such as bootstrap ensembles [Efron] 1982]. 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. In contrast, nonparametric methods such as k-nearest neighbors (k-NN) [Villa Medina et al. 2013
Fathabadi et al., 202T] are beneficial regions of limited data, yet require a proper metric be used.
We propose to combine parametric and nonparametric methods for uncertainty estimation. Particularly, define novel Riemmannian metric which captures the epistemic and aleatoric uncertainty of generative parametric forward model. This distance metric is then applied to measure the average geodesic distance the k k-nearest neighbors in the data. We derive analytical expressions for our metric and provide an efficient manner to estimate it. We then demonstrate the effectiveness of our metric for penalizing an offline RL agent compared to contemporary approaches on continuous control and autonomous driving henchmarks. 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-band offline RI.
"Correspondence to guyt enntennn@mmil... com
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

2 Preliminaries
2.1 Offline Reinforcement Learning
We consider the standard Markov Decision Process (MDP) framework (Puterman, 2014] defined by the tuple A. P. a). where S is the state space, A the action space, r : A [0, 1] the reward function, P x A x S + [0, the transition kemel,, and
(0, 1) is the discount factor.
In the online setting of reinforcement learning (RL), the environment initiates at some state 8p ~ po- At any time step the environment is in state S, agent takes an action E A and receives reward r(s, a) from environment as a result of this action. The environment transitions to state 8' according the transition function P(-ss, a). The goal of online RL is to find policy ([aas) that
maximizes the expected discounted return o" = Ex Σ arr(st,r 0z) 90 PO)
Unlike the online setting, the offline setup considers a dataset D, = {8i, 0s Tiss + of transitions generated by some unknown agents. The objective of offline RI. is find the best policy in the test environment (i.e., real MDP) given only access to the data generated by the unknown agents.
2.2 Riemannian Manifolds
We define the Riemannian pullback metric, a fundamental component of our proposed method. We refer the reader to Carmo [992] for further details on Riemannian geometry.
We are interested in studying smooth surface M with Riemannian metric g. Riemannian metric is smooth function that assigns a symmetric positive definite matrix to any point in M. At each point M tangent space T.M specifies the pointing direction of vectors "along" the surface.
Definition Let M be smooth manifold. A Riemannian metric g on M changes smoothly and defines a real scalar product on the tangent space T.M for any E M as
(z(2, y) (x. u) (z. G(2)u), 2, T.M,
where G(z) E Rddxxd xds is the corresponding metric tensor. (M,9) g) is called a Riemannian manifold.
The Riemannian metric enables us to easily de- fine geodesic curves. Consider some differentiable mapping 7 jo. 1] .M C Rd:, such that 7()) 2o: (1) = 21- The length of the curve measured on M is given by
model enror
R
r B.
I(7) S.
V V
(1(11) 0t
θy(1)) ōt /d. (1)
Empirical State-Action Manifold
G(T(())
3 Pa
The geodesic distance d(21, 22) between any two points 21, 22 E M then the infimum length over all curves 7 for which 7(0) 20. (1) 21. That is, d(21, 22) inf, L(7) s.t. 7(0) 70) = 21- The geodesic distance system of nonlinear ordinary differential equations (ODEs) defined in the intrinsic coordinates [Carmo
high model crror
Figure 1: Illustration of a data manifold corre- sponding model error. Curved lines represent trajectories different stages of training, Agent incurs a penalty in arcas of high model error.
be found by solving
1992].
Pullback Metric. Assume ambient (observation) space X' and its respective Riemannian manifold (Mx.9x) 9x). Learning 9x can be hard (c.g., learning the distance metric between images). Still, may be captured through low dimensional submanifold. As such, it is many times convenient to parameterize the surface Mxx by a latent space = Rdz and a smooth function f 2 + X, where 2 is low dimensional latent embedding space. As learning the manifold Mxx be hard, turn to learning the immersed low dimensional submanifold Mz (for which the chart maps are trivial, (af(((t)) Gx(f(r(())) aS((() 8
since 2 Rd22,, Given a curve y [0, 1] + Mz we have that (
(x(() #t (+((())))) (
(x( where the Jacobian matrix J/(2) = of Rddxxx

---

maps tangent vectors in TM2 to tangent vectors in T Mx. The induced metric is thus given by
(2)
The metric Gj is known the pullback metric, as it "pulls back" the metric Gx X' back to Gj
via f 2 + X. The pullback metric captures the intrinsic geometry of the immersed submanifold while taking into account the ambient space X. 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. Indeed, learning the latent space and observation function f can be achieved through a mecoder-decodee framework, such as a VAE Arvanitiiis et al. 2018).
3 Background: Penalty of Uncertainty for Offline Reinforcement Learning
A key element of model-based RL methods involves estimating model P(s'|s, a) to construct a pessimistic MDF This work builds upon MOPO, a recently proposed model-based offline RL framework Yuu et al., 2.20).. Particularly, we assume access to an approximate MDP (S, A,, P. o) (e.g., trained by maximizing the likelihood of the data), and define penalized MDP (S,A,F, a). such that for all E S.a A, F(s, a) = f(s, a) Ac(P(-Is,, a), P(-|s, a)), where penalizes the reward according to model error (e.g., the total variation distance) and À > O. The offline RL problem is then solved by executing
online algorithm in the reward- penallzed MDP. Unfortunately, as P(-|s, a) is unknown, and can only be estimated from the data, c( c(P(,)s, a), P(-Is, a)) cannot be calculated. Nevertheless, one can attempt to upper bound the distance, i.e., for some U S x A + R, ((P(s, a), P(-|s, a)) < U(s, a), Vs E S, E A. In this work we propose to use a naturally induced metric of variational forward model, which we show can introduce an effective penalty for offline RL. In Section 4 we define this metric, and finally, we leverage it in Section 4.4
4 Metrics of Uncertainty
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. Yu al. (2020) proposed achieve this through bootstrap ensembles, an out of distribution uncertainty estimation technique. Alternatively, we propose to employ well-known nonparametric approach for uncertainty estimation [Villa Medina et al. 2013 Fathabadi et al.) 2021], namely k-nearest neighbors (k-NN). Specifically, for any 8, E S x A. we estimate model error by
U(s,a)
Σ
a((s,a), oz));
K.
(Nt.a:)ENN, (s,a)
where d: : x R2 is a distance metric, and NN, (s, a) is the set of k-nearest neighbors of (s, a) in D, according to the distance metric d.
question arises: how to choose d? Using the Euclidean distance in ambient (state-action) space usually bad choice (e.g.. the la distance between natural images is not necessarily meaningful). Moreover, to correctly measure the error, model dynamics should be somehow taken into considera- tion. We therefore consider an alternative approach which leverages the latent space of variational forward model, as described next.
4.1 A Variational Latent Model of Dynamics
We begin by modeling P(s'(s, a) using generative latent model. Specifically, we consider a latent model which consists of an encoder E S x A + B(2) and decoder fo Z ++ B(S). where B(X') is set of probability measures on the Borel sets of X. While the encoder E learns latent representation of 8, a, the decoder fp estimates the next state s' according to P(- a). This model corresponds to the decomposition (l-|s, a) fp(- a)). Such a model can be trained by maximizing the Evidence Lower BOund (ELBO, Kingma and Welling (2013) over the data.
è Pessimism is a key element of offline RL algorithms Jin et al. 2020). limiting overestimation of a trained policy due the distribution shift between the data and the trained policy.
3

---

Algorithm GELATO: Geometrically Enriched LATent model for Offline reinforcement learning
1: Input: Offline dataset Dr. RL algorithm
2: Train variational latent forward model on dataset D. by maximizing ELBO.
3: Construct approximate MDP AS f, P,P,ov P.a)
4: Use distance dz induced by pullback metric Gyowefe (Theorem to penalize reward
fals, a) f(s, a) (s, a) where (s, a) = E (uNA)EENN, B(s,,,
5:
Train RL algorithm over penalized MDP is, P.a)
That is, given a prior P(z), we model Es, J'D as parametric functions and maximize the ELBO, m (x s.a) log (b.a(s((=) - Dx (Eg(s,s P())) We refer the reader to the appendix for exhaustive overview of training VAEs by maximum likelihood and the ELBO.
Recall that we wish to find good metric for estimating model error. Having learned a latent model for P(s'|s. a), its latent space Z can be used to define the metric in Equation 33. i.c., the Euclidean distance between latent representations of state-action pairs. Unfortunately, as was previously shown (Arvanitidis et al. 2018), latent codes in variational models contain sharp discontinuities, rendering Euclidean distances in latent space unreliable and inaccurate (as will also demonstrate in experiments). Instead, we propose to use the natural induced metric of our latent model, as described in the following subsection.
4.2 The Pullback Metric of Model Dynamics
In this part we define the metric d that will be used approximate model enror in Equation (3) Specifically, we consider a Riemannian submanifold defined by a latent space and observation function f. which induces minimum energy in the ambient space. We will later choose Z to be the
latent space our variational model (i.e., encoded state-action) and f to be the decoder function fo of next state transitions. We define the distance metric formally below.
Definition 2. We define Riemannian submanifold (M2,92) by differential functiooo 2 2: + and latent space Z such that
d2 d2(21, 22) inf / (f6())
dt 70) 21, 22.
Ot
A similar metric has been used in previous work on generative latent models [Chen et al. 2018 Arvanitidis et al.) 2018). By choosing f to be the decoder function f'D, latent codes that are close w.r.t. dz induce curves of minimal energy in the ambient observation space (i.e., next state). This metric is closely related to the pullback metric (see Section [.2]. as shown by the following proposition.
Proposition Let (Mz,9a) as defined above. Then =y
= (=)] for any = E 2.
Indeed, Proposition Isows us that G a pullback metric. Particularly Z and J, define the structure of the ambient observation space X (in our case, next state transitions).
By choosing f to be the decoder function fp. the metric Gjp becomes stochastic, complicating analysis. Instead, as proposed and analyzed in Arvanitidis et al. [2018). we use the expected pullback metric [G'z] as an approximation of the underlying stochastic metric. Similar to previous work on variational models, we use a normally distributed decoder to define the output. Using Proposition we have the following result (see Appendix for proof).
Theorem 1. Arvanitidis et al. [2018]] Assume fD(+)2) (()()). Then
E Eso(12) (Gro = Ga(z) + Go(z),
(4)
where G,(z)= J. and Jz)
Given an embedded latent space 2, the expected metric in Equation (4) gives us sense of the
topology of the latent space manifold induced by fD. The terms G,. J Jn and G. = JH J, are in fact the induced pullback metrics of and a,, respectively.

---

(a) fo)
(b) Latent Geodesic Distance (c) Latent Euclidean Distance
(a) The latent space (yellow markers) of the grid world environment and the geometric volume measure of the decoder-induced metric (background). (b, The geodesic distance of the decoder-induced submanifold and the Euclidean distance of latent states, viewed ambient space. All distances calculated w.r.t. the yellow marked state. Note: colors (a), which measure magnitude, are unrelated to colors (b,c), which measure distance to the yellow marked state.
Figure 2:
4.3 Capturing Epistemic and Aleatoric Uncertainty
The previously proposed encoder-decoder model induces a metric which captures the structure of the learned dynamics. However, the decoder variance, o=)), does not differentiate between aleatoric uncertainty (environment dynamics) and epistemic uncertainty (missing data).
We propose two methods to enrich the metric in Equation (4 in order to achieve better estimate of uncertainty. First, by using an ensemble of M decoder functions {_D,,)== trained using standard bootstrap techniques [Efron 1982]. we capture the traditional epistemic uncertainty of the decoder parameters. Second, to correctly distinguish epistemic and aleatoric uncertainty, we add latent forward function to our previously proposed variational model. Specifically, our latent model consists of an encoder E S A + B(Z). forward model fF Z B(t) and decoder functions fD. X' + B(S) such that P(-Is, a) fD, F.,((-ee (-1æz), and x ~ fF(-)E(s, a)). This structure enables us to capture the aleatoric uncertainty under the forward transition model fr, and the epistemic uncertainty using the decoders fD.; That is, for fr(-12) ~ (Nur((2) (z()))), the variance model r2(2 captures the stochasticity model dynamics. This decomposition also helpful whenever one wants to train agent in latent space (e.g., for planning Schrittwieser et al. [2020])
Next, we turn to analyze the pullback metric induced by the proposed forward transition model. As both fr and (/D.a are stochastic (capturing epistemic and aleatoric uncertainty), the result of Theorem cannot be directly applied to their composition. 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 the appendix).
fDa()/2)) ~ ~ Unif(l..... M). Then, the expected pullback metric of the composite function (Sn. o fr) given by
Theorem 2. Assume fF(112) ~ (u IF)
Er(oowefr) G)o - 1 (x) (2) + J (2) (z) (z),,,
where (z) =
1 1o
JT.
(a) (x) Jo (x) (z)oo(()) ub
Unlike the metric in Equation 4p. the composite metric distorts the decoder metric with Jacobian matrices of the forward model statistics. It takes into account both the alcatoric and epistemic uncertainty through the forward model as well as ensemble of decoders. As special case we note the metric for the case of deterministic model dynamics.
Corollary 1. Assume deterministic model dynamics, i.e., = fF(Z), and without loss of generality
assume fr = I. Then, the expected pullback metric of Theorem 2 is given by E by (f)
1 MN JT. (z)J (=) + JTH (z) (z).
M
Pb
5

---

la Distance
Geodesic Distance
time
80
2
3
6
Figure 3: Plots show t-SNE embeddings generated the intersection enviromment. Left plot depicts embed- dings using Euclidean distances. Right plot depicts embeddings using geodesics distances which induce curves of minimum energy ambient space. Colors correspond to the time trajectory (normalized w.r.t. longest trajectory). Numbering show mappings of specific trajectory's states onto the embedded space as controlled red car takes left turn at an intersection (trajectory visualization shown under plots).
4.4 GELATO: Geometrically Enriched LATent model for Offline reinforcement learning
Having defined our metric, we are now ready to leverage it in model based offline RL framework. Specifically, provided a dataset D, {(S;, az Pr, 1=) we train the variational latent forward model described in the previous section.
Algorithm presents GELATO, our proposed approach. In GELATO, estimate model error by measuring the distance of a new sample to the data manifold. We construct the reward-penalized MDP for which the егтоr acts as pessimistic regularizer. Finally, we train an RL agent over the pessimistic MDP with transition P(-|s, a) and reward r(s, a) \U a). By achieving an improved estimate for model eтror the model-based based pessimistic approach can significantly improve performance, as shown in the following section.
5 Experiments
This section is dedicated to quantitatively and qualitatively understand the benefits of our proposed metric and method. We validate two principal claims: (1) Our metric captures inherent character- isties of model dynamics. We demonstrate this by visualizing state-action geodesics of toy grid world problem and a multi-agent autonomous driving task. We show that curves of minimum energy in ambient space indeed capture intrinsic properties of the problem. (2) Our metric provides an improved OOD uncertainty estimate for offfine RL. We compare the traditionally used bootstrap ensemble method to approach, which leverages our pullback metric in a nonparametric nearest neighbors approach. We also compare our method to simple use of Euclidean distances in latent space. We run extensive experiments continuous control and autonomous driving benchmarks. show that our metric achieves significantly improved performance in tasks for which geodesics non-euclidean.
5.1 Metric Visualization
Four Rooms. To better understand the inherent structure our metric, we constructed grid-worid environment for visualizing our proposed latent representation and metric. The 15 x 15 environment, as depicted in Figure consists of four rooms, with impassable obstacles in their centers. The agent, residing at some position (z. y) E [-1, 1]2 in the environment can take one of four actions: up. down, steps (uniformly distributed) in that direction. We collected a dataset of 10000 samples, taking random actions at random initializations of the environment. The ambient state space was represented by the position of the agent a vector of dimension 2, normalized to values in [-1, 1]. Finally, we trained variational latent model with latent dimension dz 2.
left, or right moving the agent 1,

---

Table 1: Performance of GELATO and its variants in comparison contemporary model-based methods on D4RL datasets, Scores correspond the returm, averaged over 5 seeds (standard deviation removed to space constraints and is given in the appendix). Results for MOPO, MBPO, SAC, and imitation are taken from Yu et al. 2020 Mean score of dataset added for reference. Bold scores show an improved score w.r. other methods.
Hopper
Walker2d
Halfcheetah
Method
Rand
Med
Med-Expert
Rand
Med
Med-Expert
Rand
Med
Med-Expert
Data Score
299
1021
1849
498
1062
-303
3945
8059
GELATO (metric)
412
1515
2560
5168
7989
GELATO (22)
544
1320
815
388
312
1255
512
4096
7304
MOPO (houtstrap)
677
1202
1063
396
518
1296
4114
4974
7594
MBPO
370
222
3527
3228
SAC
664
325
1850
120
27
-2
3502
-33
-78
Imitation
615
1234
3907
47
193
329
-41
4201
4164
We used a standard encoder ~ (pe(s), (s)(s) and decoder s' (N(pp oz()) represented by neural networks trained end-to-end using the evidence lower bound. We refer the reader the appendix for an exhaustive description of the training procedure.
The latent space output of our model is depicted by yellow markers in Figure 2a. Indeed, the latent embedding consists of four distinctive clusters, structured a similar manner as our grid-world environment. Interestingly, the distortion of the latent space accurately depicts an intuitive notion of distance between states. As such, rooms are distinctively separated, with fair distance between each cluster. States of pathways between rooms clearly separate the room clusters, forming topology with four discernible bottlenecks.
In addition to the latent embedding, Figure 2 depicts the geometric volume measure v (Gt(G( of the trained pullback metric induced by JD. This quantity demonstrates the effective geodesic distances between states in the decoder-induced submanifold. 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. In contrast, geodesics near the data manifold would low values.
We visualize the decoder-induce geodesic distance and compare the latent Euclidean distance in Figures 2 and respectively. The plots depict the normalized distances of all states to the state marked by yellow square. Evidently, the geodesic distance captures a better notion of distance in the said environment, correctly exposing the "land distance" in ambient space. 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. In contrast, the latent Euclidean distance does not properly capture these geodesics, exhibiting a similar distribution of distances in other rooms. 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, 1J2. Particularly, the state coordinates (I. y) would wrongly assign short distances to states across impassible walls or obstacles, i.e., measuring the "air distance".
Intersection. We visualized our metric in the intersection environment proposed in Leurent 2018 Figure 3 compares the Euclidean and geodesic distances of partially trained agent. Unlike the previous toy example, to visualize the inherent manifolds we used t-SNE [Van der Maaten and Hinton 2008 projections computed with Fuclidean distance and compared them the projection computed with geodesic distance, i.e., curves of minimum energy in ambient space (Definition 2 Indeed, the geodesies captured the inherent structure the environment, whereas Euclidean distances only managed to capture general clusters. These suggest that Euclidean distance might not be representative for measuring distance in latent space, will also become evident by our experiments in the following subsections.
5.2
Datasets and Implementation Details
We used D4RL Fuu et al., 2020 and the autonomous vehicle environments highway-env [Leurent] 2018 as benchmarks for all of our experiments. We tested GELATO on three Mujoco (Toorov et al. environments (Hopper, Walker2d, Halfcheetah) on datasets generated by single policy and mixture of two policies. Specifically, we used datasets generated by a random agent (IM samples), partially trained agent, i.e, medium agent (IM samples), and mixture of partially trained and expert
2012
7

---

Highway
Highway (Fast)
Roundabout
S003
Intersection (Easy)
Intersection (Hard)
Racetrack
Jouns
Figure 4: Performance of GELATO with different uncertainty estimation methods the highway-eny benchmarks. Results suggest our induced semiparametric distance is an effective penalty for offline RL.
agents (2M samples). For autonomous driving, tested GELATO on four environments (Highway, Roundabout, Intersection, Racetrack), on datasets containing five episodes generated by a partially trained agent. We also tested faster (x 15 speedup) variant of the Highway environment, 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).
We trained our variational latent model two phases. First, the model was fully trained using a calibrated Gaussian decoder [Rybkin al. 2020]. Specifically, maximum-likeliben the variance was used a² MSE (i) E arg max, N (fju,, I² Then, in the second stage we fit the variance decoder networks. Hyperparameters for training are found in the appendix.
estimate of
In order to practically estimate the geodesic distance in Algorithm [] we defined parametric curve in latent space and used gradient descent to minimize the curve's 's energy. 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).
We used FAISS [Johnson et al. 2019] for efficient GPU-based k-nearest neighbors calculation. We set h = 5 neighbors for the penalized reward (Equation (3). Finally, we used a variant of Soft Learning, as proposed by Yu al. 2020 our RL algorithm for the continuous control benchmarks, and PPO [Schulman et al. [017]] for the autonomous driving tasks. All agents were trained for IM steps (for continuous control benchmarks) and 350K steps (for the driving benchmarks), using single GPU (RTX 2080), and averaged over seeds (see Appendix for more details).
5.3 Results
D4RL. Results for the continuous control domains are shown in Table [. We performed experiments to analyze GELATO on various continuous control datasets. We compared GELATO to contemporary model-based offline RI approaches; namely, MOPO YYu et al. 2020) and MBPO (Jnner et al.) 2019). as well as the standard baselines of SAC [Haarnoja et al.) 2018 and imitation (behavioral cloning,
Bain and Sammut 1995]. Ross et al. 12011)..
We found performance increase on most domains, and most significantly in the medium domain, i.e., partially trained agent. Since the medium dataset contained average behavior, combining the nonparametric nearest-neighbor uncertainty method with the bootstrap decoders benefited the agent's 's overall performance. In addition, GELATO with the latent & distance metric performed well on many of the benchmarks. We conjecture this is due to the inherent Euclidean nature of continuous control benchmarks. Unlike the embedding for the autonomous driving benchmarks
(Figure 3), we found the D4RI. data to project similarly when & and geodesic distances were used (we provide plots of these embeddings in the Appendix).

---

Highway-Env. Figure shows results for the autonomous driving benchmarks in highway-env. 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 Fuclidean distance. We credit this improvement to the non-euclidean nature of the environments, as previously described in Figure 3 While Huclidean distances were useful in the Mujoco environments, they performed distinctly worse in the autonomous driving environments.
Our results emphasize the importance of OOD uncertainty estimations methods in reinforcement learning on various types of datasets. While robotic control tasks provided useful insights, they did not capture the non-euclidean nature inherent in alternative tasks, such as autonomous driving.
Related Work
Offline Reinforcement Learning. The field of offline RL has recently received much attention as several algorithmic approaches were able to surpass standard off-policg algorithms. Value- based online algorithms do not perform well under highly off-policy batch data [Fujimoto et al. 2019 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- cistt--distriuttion (OOD) samples. These issues become more prominent the
offline setting, as new samples cannot be acquired.
Several works on offline RL have shown improved performance on standard continuous control benchmarks (Laroche et al. 20199] Kumar et al.] 2019 Fujimoto et al., 2019 Chen et al. 2020b Swazinna et al.] 2020, Kidambi et al. 20200 Yu et al. 2020]. This work focused on model- -bsed approaches Yu et al. 2020 Kidambi et al. 2020]. in which the agent is incentivized to remain close to areas of low uncertainty. Our work focused on controlling uncertainty estimation high dimensional environments. Our methodology utilized recent success the geometry of deep generative models Arvanitiiis et al. 2018 2020]. proposing an alternative approach to uncertainty estimation.
Representation Learning. Representation learning seeks to find an appropriate representation of data for performing a machine-learning task (Goodfellow al. 2016). Variational Auto Encoders Kingma and Welling 2013 Rezende et al. [014] have been a popular representation learning technique, particularly in unsupervised learning regimes [Chen et al. 2016 Van Den Oord al.
2017
Hsu et al. 2017 Serban et al. 2017 Engel et al. 2017 Booaanwski 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. [019]. Particularly, variational models have been shown
able to derive successful behaviors in high dimensional benchmarks [Hafner et al. 2020
Various representation techniques in reinforcement learning have also proposed to disentangle representation of both states (Engel and Mannor 20011] Littman and Sutton| 2002 Stooke al. 2020 Zhu et al. 2020), and actions (Tennenholtz and Mannor [019 Chandak et al. 2019]. 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]. Unlike previous work, GELATO is focused on a semiparametric approach for uncertainty estimation, enhancing offline reinforcement learning performance.
Discussion and Future Work
This work presented metric for model dynamics and its application to offline reinforcement learning. 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 's variance for uncertainty estimation. The apparent slowdown in performance geodesic distance. Improvement in this area may utilize techniques for efficient geodesic estimation.
mostly due to computation of the
We conclude by noting possible future applications of our work. In Section 5.1 we demonstrated the inherent geometry our model had captured, its corresponding metric, and geodesics. Still, in this work we focused specifically on metrics related to the decoded state. In fact, similar derivation to Theorem 2 could be applied to other modeled statistics, e.g., Q-values, rewards, future actions, and more. Each distinct statistic would induce its own unique metric w.r.t. its respective probability measure. Particularly, this concept may benefit a vast array of applications in continuous or large state and action spaces.

---

References
Rishabh Agarwal, Dale Schuurmans, and Mohammad Norouzi. An optimistic perspective on offline reinforcement learning. In International Conference Machine Learning, pages (14-114. PMLR, 2020.
Georgios Arvanitidis, Lars Kai Hansen, and Søren Hauberg. Latent space oddity: On the curvature of deep generative models. 6th International Conference on Learning Representations, ICLR 2018, 2018.
Georgios Arvanitidis, Søren Hauberg, and Bernhard Schölkopf. Geometrically enriched latent spaces. arXiv preprint UrXiv:2008.00560 2020.
Michael Bain and Claude Sammut. A framework for behavioural cloning. In Machine Intelligence 15, pages 103-129, 1995.
Piotr Bojanowski, Armand Joulin, David Lopez-Pas, and Arthur Szlam. Optimizing the latent space of generative networks. In International Conference on Machine Learning, pages 600-609, 2018.
Jacob Buckman, Carles Gelada, and Mare G Bellemare. The importance of pessimism fixed-dataset policy optimization. arXiv preprint urXiv:2009.06799 2020.
Manfredo Perdigao do Carmo. Riemannian geometry. Birkhäuser, 1992.
Yash Chandak, Georgios Theocharous, James Kostas, Scott Jordan, and Philip Thomas. Learning ac- tion representations for reinforcement learning. In International Conference Machine Learning, pages 941-950, 2019.
Nutan Chen, Alexej Klushyn, Richard Kurle, Xueyan Jiang, Justin Bayer, and Patrick Smagt. Metrics for deep generative models. In International Conference on Artificial Intelligence and Statistics, pages 1540-1550. PMLR, 2018.
Nutan Chen, Francesco Ferroni, Alexej Klushyn, Alexandros Paraschos, Justin Bayer, and Patrick van der Smagt. Fast approximate geodesics for deep generative models. In International Conference on Artificial Neural Networks, pages 554-566. Springer, 2019.
Nutan Chen, Alexej Klushyn, Francesco Ferroni, Justin Bayer, and Patrick van der Smagt. Learning flat latent manifolds with vaes. 2020a.
Xi Chen, Diederik P Kingma, Tim Salimans, Yan Duan, Prafulla Dhariwal, John Schulman, Ilya Sutskever, and Pieter Abbeel. Variational lossy autoencoder. arXiv preprint 1.2.711 2016.
Xinyue Chen, Zijian Zhou, Zheng Wang. Che Wang, Yanqiu Wu, and Keith Ross. Bail: Best- action imitation learning for batch deep reinforcement learning. Advances in Neural Information Processing Systems, 33, 2020b.
Zheng Ding, Yifan Xu, Weijian Xu, Gaurav Parmar, Yang Yang, Max Welling, and Zhuowen Tu. Guided variational autoencoder for disentanglement learning. In Proceedings of the IEEF/CVF Conference on Computer Vision and Puttern Recognition, pages 7920-7929, 2020.
Laurent Dinh, David Krueger, and Yoshua Bengio. Nice: Non-linear independent components
estimation. arXiv preprint wrXivv:1410.cc 16, 2014.
Bradley Efron. The jackknife, the bootstrap and other resampling plans. SIAM, 1982.
Jesse Engel, Matthew Hoffman, and Adam Roberts. Latent constraints: Learning to generate conditionally from unconditional generative models. arXiv preprint 17X1v:17 2017.
Yaakov Engel and Shie Mannor. Learning embedded maps of markov processes. In Proceedings of the Eighteenth International Conference on Machine Learning, pages 138-145, 2001.
Damien Ernst, Pierre Geurts, and Louis Wehenkel. Tree-based batch mode reinforcement learning. Journal of Machine Learning Research, 6(Apr):503-556, 2005.
10

---

Aboalhasan Fathabadi, Seyed Morteza Seyedian, and Arash Malekian. Comparison of bayesian, k- nearest neighbor and gaussian process regression methods for quantifying uncertainty of suspended sediment concentration prediction. Science of The Total Environment, page 151760, 2021.
William Fedus, Prajit Ramachandran, Rishabh Agarwal, Yoshua Bengio, Hugo Larochelle, Mark Rowland, and Will Dabney. Revisiting fundamentals of experience replay. In International Conference Machine Learning, pages 3061-3071. PMLR, 2020.
Raphael Fonteneau, Susan Murphy, Louis Wehenkel, and Damien Ernst. Batch mode reinforcement learning based on the synthesis of artificial trajectories. Annals of operations research, 208(1): 383-416, 2013.
Justin Fu, Aviral Kumar, Matthew Soh, and Sergey Levine. Diagnosing bottlenecks in deep q-learning algorithms. In International Conference on Machine Learning, pages 2021-2030, 2019.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep
data-driven reinforcement learning. arXiv preprint arXiv:2004.07219, 2020.
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. International Conference on Machine Learning, pages 0062-2062. PMLR, 2019.
Yarin Gal and Zoubin Ghahramani. Dropout as a bayesian approximation: Representing model uncertainty in deep learning. In international conference machine learning, pages 1050-1059, 2016.
Ian Goodfellow, Yoshua Bengio, Aaron Courville, and Yoshua Bengio. Deep learning, volume 1.
MIT press Cambridge, 2016.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International Conference on Machine Learning, pages 1861-1870. PMLR, 2018
Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference on Learning Representations, 2019.
Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. arXiv preprint arXiv:2010.02/93 2020.
Karol Hausman, Jost Tobias Springenberg, Ziyu Wang, Nicolas Heess, and Martin Riedmiller. Learning an embedding space for transferable robot skills. International Conference on Learning Representations, 2018.
Wei-Ning Hsu, Yu Zhang, and James Glass. Unsupervised learning of disentangled and interpretable representations from sequential data. In Advances in neural information processing systems, pages 1878-1889, 2017.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model-based policy optimization. arXiv preprint 1906v: 08255, 2019.
Ying Jin, Zhuoran Yang, and Zhaoran Wang. Is pessimism provably efficient for offline rI? arXiv preprint arXiv:2012. /505: 2020.
Jeff Johnson, Matthijs Douze, and Hervé Jégou. Billion-scale similarity search with gpus. IEEE Transactions on Big Data, 2019.
Dimitris Kalatzis, David Eklund, Georgios Arvanitidis, and Søren Hauberg. Variational autoencoders with riemannian brownian motion priors. In Proceedings of the 37th International Conference Machine Learning (ICML). PMLR, 2020.
Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model- based offline reinforcement learning. arXiv preprint arXiv:2005.05951, 2020.
Diederik P Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014.
11

---

Diederik P Kingma and Max Welling. arXiv:1312.6114, 2013.
Auto-encoding variational
bayes.
arXiv preprint
Alexej Klushyn, Nutan Chen, Richard Kurle, Botond Cseke, and Patrick van der Smagt. Learning hierarchical priors in vaes. In Advances in Neural Information Processing Systems, pages 2870-
2019.
Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off-policy q-learning via bootstrapping error reduction. In Advances Neural Information Processing Systems, pages 11784-11794, 2019.
Romain Laroche, Paul Trichelair, and Remi Tachet Des Combes. Safe policy improvement with baseline bootstrapping. International Conference on Machine Learning, pages 3652-3661. PMLR, 2019.
Edouard Leurent. An environment for autonomous driving decision-making. https://github com
2018.
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint (rxXi::2000.006443 2020.
Minne Li, Lisheng Wu, WANG Jun, and Haitham Bou Ammar. Multi-view reinforcement learning. In Advances in neural information processing systems, pages 1420-1433 14300-131,, 2019.
Bric Liang, Richard Liaw, Robert Nishihara, Philipp Moritz, Roy Fox, Ken Goldberg, Joseph Gonzalez, Michael Jordan, and Ion Stoica. Rllib: Abstractions for distributed reinforcement learning. International Conference on Machine Learning, pages 3053-3062. PMLR, 2018.
Michael L Littman and Richard S Sutton. Predictive representations of state. Advances in neural information processing systems, pages 1555-1561, 2002.
Ian Osband, John Aslanides, and Albin Cassirer. Randomized prior functions for deep reinforcement learning. Advances in Neural Information Processing Systems, 31, 2018.
Janith C Petangoda, Sergio Pascual- -iiz, Vincent Adam, Peter Vrancx, and Jordi Grau-Moya. Disentangled skill embeddings for reinforcement learning. arXiv preprint arXiv: 0ete:09009..233. 09223, 2019.
Martin L Puterman. Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 2014.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wienstra. Stochastic backpropagation and approximate inference in deep generative models. arXiv preprint 1401::/400.4822 2014.
Martin Riedmiller. Neural fitted q iheration-firs experiences with a data efficient neural reinforcement learning method. In European Conference on Machine Leurning, pages 317-328. Springer, 2005.
Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. A reduction of imitation learning and structured prediction no-regret online learning. In Proceedings of the fourteenth international conference artificial intelligence and statistics, pages 627-635. JMLR Workshop and Conference Proceedings, 2011.
Oleh Rybkin, Kostas Daniilidis, and Sergey Levine. Simple and effective vae training with calibrated decoders. arXiv preprint arXiv:2006. 13202, 2020.
Julian Schrittwiserr, Ioannis Antonoglou. Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with learned model. Nature, 588(7839):604 604-66909)6666 2020.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017.
Tulian Serban, Alessandro Sordoni, Ryan Lowe, Laurent Charlin, Joelle Pineau, Aaron Courville, and Yoshua Bengio. A hierarchical latent variable encoder-decoder model for generating dialogues. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 31, 2017.
12

---

Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin. Decoupling representation learning from reinforcement learning. arXiv preprint arXiv:2009.08319, 2020.
Phillip Swazinna, Steffen Udluft, and Thomas Runkler. Overcoming model bias for robust offline deep reinforcement learning. arXiv preprint 2XXi:20088.05333 2020.
Guy Tennenholtz and Shie Mannor. The natural language of actions. In International Conference Machine Learning, pages 6196-6205, 2019.
Emanuel Todorov, Tom Frez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, pages 5026-5033. IEEE, 2012.
Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. In Advances in Neural Information Processing Systems, pages 6306-6315, 2017.
Laurens Van der Maaten and Geoffrey Hinton. Visualizing data using t-sne. Journal of machine
learning research, 9(11), 2008.
Joe Luis Villa Medina al. Reliability of classification and prediction in k-nearest neighbours. PhD thesis, Universitat Rovira Virgili, 2013.
Ruosong Wang, Dean P Foster, and Sham M Kakade. What are the statistical limits of offline rl with linear function approximation? arXiv preprint arXiv:2010.11895, 2020.
Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-based offline policy optimization, arXiv preprint arXiv:2005./322 2020.
Andrea Zanette. Exponential lower bounds for batch reinforcement learning: Batch can be exponentially harder than online rl. arXiv preprint arXiv:2012.08005. 2020.
Jinhua Zhu, Yingee Xia, Lijun Wu, Jiajun Deng, Wengang Zhou, Tao Qin, and Houqiang Li. Masked contrastive representation learning for reinforcement learning. arXiv preprint arXiv:2010.07470, 2020.
13

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes]
(c) Did you discuss any potential negative societal impacts of your work? [Yes]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes]
(b) Did you include complete proofs of all theoretical results? [Yes]
3. If you experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either the supplemental material or as URL)? [Yes]
(b) Did you specify all the training details (e.g., data splits. hyperparameters, how they were chosen)? [Yes]
(c) Did you report error bars (e.g.. with respect the random seed after running experi- ments multiple times)? [Yes]
(d) Did you include the total amount of compute and the type of resources used (c.g., type of GPUs, internal cluster, or cloud provider)? [Yes]
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes]
(b) Did you mention the license of the assets? [N/A]
(c) Did you include any new assets either in the supplemental material or as a URL? [Yes]
(d) Did you discuss whether and how consent was obtained from people whose data you're using/curating? [N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, applicable? [N/A]
(b) Did you describe any potential participant risks, with links Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation? [N/A]
14


Reviews:
[{"review_id": "Haq7LVvG7mv", "reviewer": "Reviewer_TKTA", "summary": "Using Riemannian distance metrics instead of Euclidean metrics, the authors better estimate uncertainties for out-of-distribution (OOD) data and apply it to model-based Reinforcement Learning (RL) settings. The paper introduces a new pullback metric for such dynamic models, and verifies its performance on several RL benchmarks.", "questions": "1) In the pullback metric for model dynamics, how did the authors come up with this choice for metric? How is the time derivative computed? And what about the geodesic computation, are using gradient descent to solve an ODE?\n\n2) In the abstract, the authors mention how the \"proposed metric measures both the quality of OOD samples as well as the discrepancy of examples in the data\" (line 7). Don't they both in the end mean the same thing?\n\n3) In Figure 1, \"high model error\" is mentioned twice, but never do we see \"low model error\". This would be on the manifold, correct? Would it make sense to add it somewhere in the figure/caption?\n\n4) Line 121 mentions how the encoder and decoders are parametric functions, is there any details on this? If normal neural networks are used, would that not be nonparametric?\n\n5) The equation in Definition 2 defines the curve length form Eq. 1, correct? It could be helpful to link these two for better understanding if that's the case.\n\n6) The authors introduce an ensemble of decoders in Line 161, which I assume is then used for all experiments that follow? How about the forward model, do we continue with the assumption in Corollary 1 and practically neglect it? What is the \"meaning\" of this ambient space X? This could be made clearer. The need for this forward model in the first place is unclear, perhaps Corollary 1 could be entirely removed.\n\n7) The results are compared with a classical bootstrap ensemble method, but isn't that also used in the proposed method for the epistemic uncertainty? What is the difference there with the baseline?\n\n8) Is the four room experiment discrete or continuous? From the actions, it seems to be discrete, but since the latent space is continuous, I was wondering what happens if we sample more points in the latent space, and where they would correspond to in the rooms, i.e. how likely it is they will end up in the walls/obstacles. Also, would the input space be 3D? Isn't a 2D latent space a bit of an easy mapping in this case? \n\n9) Is \"Data Score\" in Table 1 the score that the collected data achieved? A quick explanation here would help.\n\n10) Why is t-SNE used for the intersection example latent space visualization? Is it because it is more than 2D? How many dimensions is the latent space here?\n\n", "limitations": "The limitations are properly addressed by the authors.", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "A very well written paper, where the novelty is clear. The results are quite nice, though showing the improvement in more scenarios would be great, since the authors noticed that the Riemannian manifold metric work better in some specific cases. However, it seems a bit limiting to define the metric for only dynamic systems. The authors could elaborate more on their choice for this specific setting for model-based RL uncertainty estimation as well, as opposed to general uncertainty measurement.\n\nWhile the whole paper is focused around offline RL, it would be great if online RL can be added into the introduction a bit more so the contrast between the two is clearer. It also isn't clear to the reader why parametric and nonparametric methods are combined for uncertainty estimation, this could be motivated better.\n\nLine 27 the \"metric\" used for kNN could be associated with how the distance is defined. That way the reader can understand why a \"proper metric\" is important. Similarly, line 29 could have just a quick explanation on what epistemic and aleatoric uncertainties are. Just model/data uncertainty would help the understanding for the reader. \n\nSection 2.1 is mainly about (online) RL, not sure if the title necessarily fits. Or the third paragraph on offline setting can be expanded and explained more in detail, since it's the focus of this paper. The equation in line 78-79 (which seems quite important) could be elaborated more, in what space the curve length is defined, and why we care about using a curve length for $f(\\gamma)$ instead of $\\gamma$. Also, it seems section 3 would fit into section 2, since it is also a preliminary/background. Is there a reason it is separate?\n\nThough the Algorithm 1 is presented quite nicely, it could be even further improved if it were a visual figure, I believe. It would make a very strong titlepage figure as well, outlining the novel framework the authors introduce in this paper.\n\nA major issue that can be easily fixed, is that many of the figures, tables, and algorithms are located pages before they are actually referenced in the text. \n\nThe whole idea of the metric capturing inherent characteristics of the model dynamics is still a bit vague to me. If it is, for instance, the discrete bottlenecks present in the four rooms example, more samples would be needed to see this clearly. If it is the clustering in the intersection example, more explanation is needed why the clusters for Euclidean distance is not enough. t-SNE visualization is only qualitative, and distances do not necessarily mean much, though maybe some conclusion can be drawn from the shape of the clusters.\n\nA lot of explanation on how the geodesics are found is missing, how the parametric curve is defined for instance. This could be made a much more important part of the paper. The computation time is also mentioned at some point, and this could be described more in detail, exact computation times would be appreciated.\n\nThe related works are divided quite nicely, but the last paragraph on disentangling latent space for RL feels a bit out of place and irrelevant. \n\nLastly, a lot of the discussion was in the results 5.3 already, so the section 7 title could just be conclusion. \n\nAll-in-all a well-written paper with plenty of preliminary information for readers of all backgrounds. The proposed framework could be shown using a titlepage figure describing the pipeline/block diagram for more impact. The overall contributions are clear, and the results agree with the conclusions from the authors. Although many of the core ideas follow from previous work, this work enables the application of geometry-aware learning methods in RL context, providing consistently better results than baselines.\n\n\nSmall details: \n- There is no appendix, even though it was referenced many times in the paper.\n- Missing figure reference in line 19\n- Typo line 29 \"Riemmannian\"\n- Line 231 \"would low values.\" misses some word", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "6baQMHgPFD_", "reviewer": "Reviewer_Jba5", "summary": "To address the uncertainty quantification (UQ) problem in offline model-based RL, the authors propose a Riemannian metric that captures both uncertainty in dynamics (aleatoric) and OOD data (epistemic). Their offline RL algorithm GELATO (cute name) learns an uncertainty-penalized reward function by learning the geodesic distance through a 'pullback' metric. The geodesic distance is combined with KNN evaluation and bootstrapping to generate improved UQ performance. In control and autonomous driving datasets GELATO achieves good performance.  ", "questions": "Also minor: I found section 4.3 to be a little bit hard to follow at first. $M$ is used to refer to the manifold in the paper except in this section.", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength:\n\n- Delightfully well written paper. The exposition of motivation, theory, and experiment results are clear and convincing. \n- Introduces a new representation for estimating distance of model dynamics in latent space, that is more geometrically salient then Euclidean distance.\n- Detailed analysis of success and failure cases in multiple domains in the experiment section.\n- Broad applicability of the contribution. As the authors mentioned, their method of capturing geodesics can be applied to many RL estimates such as Q values, rewards, etc. The authors provide a clear and effective way of doing so.\n\nminor details:\n- missing reference in line 19, typo of \"Riemannian\" in line 29.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "3JXdgqEgqoJ", "reviewer": "Reviewer_XJeV", "summary": "For reinforcement learning agents learned in an offline setting, it is especially important to be able to quantify uncertainty, as the agent may encounter out-of-distribution experiences when deployed in the online environment.  Learning to quantify uncertainty then enables the agent to avoid these regions, by operating under a penalized MDP.  In this work, uncertainty is estimated as a k-nearest neighbor between the current state action pair against the state actions pairs in the offline dataset, according to a distance metric.  The distance metric is formulated by the authors as the geodesic in a learned latent space, modeled as a VAE.  The latent encoder handles aleatoric uncertainty, whereas an ensemble of decoder functions is used to handle epistemic uncertainty.\n", "questions": "I am wondering how flexible the proposed method is to changes in the encoder and decoder beyond the standard Gaussian formulation.  Would the derived pullback metric work for more complex VAEs?  Could similar pullback metrics be derived for more general generative modeling techniques (my understanding of the work of Arvanitidis et al., 2018 is that it is a study on the geometry of generative models in general, not limited to VAEs)?\n", "limitations": "(Social Impact Limitations N/A)", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "To my knowledge, the paper is quite original.  It builds off of theoretical findings from past work on deep generative models and applies it well to the offline RL domain.\n\nThe quality of the paper is high - the idea is straightforward, expressed clearly, and demonstrated cleanly.  The clarity of the writing is excellent, and it was easy to follow what the authors proposed.\n\nI believe this paper has high significance - enabling RL agents trained on offline datasets to quantify their uncertainty is critical for ensuring good continued performance when deployed in an online environment.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}]

Author rebuttal:

Decision: Accept