Let Offline RL Flow: Training Conservative Agents in the Latent Space of Normalizing Flow
Authors: Dmitry Akimov
Vladislav Kurenkov
Alexander Nikulin
Denis Tarasov
Sergey Kolesnikov
Venue: Submitted to ICLR 2023 (2023)
Keywords: Offline Reinforcement Learning
Normalizing Flows

Abstract: Offline reinforcement learning aims to train a policy on a pre-recorded and fixed dataset without any additional environment interactions. There are two major challenges in this setting: (1) extrapolation error caused by approximating the value of state-action pairs not well-covered by the training data and (2) distributional shift between behavior and inference policies. One way to tackle these problems is to induce conservatism - i.e., keeping the learned policies closer to the behavioral ones. To achieve this, we build upon recent works on learning policies in latent action spaces and use a special form of normalizing flow for constructing a generative model, which we use as a conservative action encoder. This normalizing flow action encoder is pre-trained in a supervised manner on the offline dataset, and then an additional policy model - controller in the latent space - is trained via reinforcement learning. This approach avoids querying actions outside of the training dataset and therefore does not require additional regularization for out-of-dataset actions. We evaluate our method on various locomotion and navigation tasks, demonstrating that our approach outperforms recently proposed algorithms with generative action models on a large portion of datasets.

Under review as a conference paper at ICLR 2023
LET OFFLINE RL FLOW: TRAINING CONSERVATIVE AGENTS IN THE LATENT SPACE OF NORMALIZING FLOWS
Anonymous authors Paper under double-bling brind review
ABSTRACT
Offline reinforcement learning aims to train policy on a pre-recorded and fixed dataset without any additional environment interactions. There are two major challenges in this setting: (I) extrapolation error caused by approximating the value of actton pairs not well-covered by the training data and (2) distri- butional shift between behavior and inference policies. One way to tackle these problems is to induce conservatism - -i..,, keeping the learned policies closer to the behavioral ones. To achieve this, we build upon recent works on learning policies in latent action spaces and use a special form of Normalizing Flows for construct- ing a generative model, which we use as a conservative action encoder. This Normalizing Flows action encoder is pre-trained in supervised manner on the offline dataset, and then an additional policy model - controller in the latent space is trained via reinforcement learning. This approach avoids querying actions out- side of the training dataset and therefore does not require additional regularization for out-of-dataset actions. We evaluate our method on various locomotion and navigation tasks, demonstrating that our approach outperforms recently proposed algorithms with generative action models on a large portion of datasets.
INTRODUCTION
Offline Reinforcement Learning (ORL.) addresses the problem of training new decision-making pol- icy from static and pre-recorded dataset collected by some other policies without any additional data collection (Lange et al.) 2012 Levine et al., 2020). One of the main challenges in this set- ting is the extrapolation error (Fujimoto et al. 2019) i.c. inability to properly estimate values of state-action pairs not well-supported by the training data, which in turn leads to overestimation bias. This problem is typically resolved with various forms of conservatism, for example, Implicit Q-Learning (Kostrikov et al.] (2021) completely avoids estimates of out-of-sampp actions, Conser- vative Q-Learning (Kumar et al. 2020) penalizes q-values for out-of-distribution actions and others (Fujimoto & Gu. 2021
Kumar al. 2019) put explicit constraints stay closer to the behavioral
policies.
An alternative approach constraint-trained policies was introduced in PLAS (Zhou et al. 2020) where authors proposed to construct a latent space that maps to the actions well-supported by the training data. To achieve this, Zhou et al. (2020) use Variational Autoencoder (VAE) (Kingma et al. 2019) to learn latent action space and then train a controller within it. However, as was demonstrated in Chen al., (2022), their specific use of VAE leads necessity for clipping the latent space. Otherwise, the training process becomes unstable, and the optimized controller can exploit the newly constructed action space, arriving at the regions resulting out-of-distributioo actions in the original space. While the described clipping procedure found to be effective, this solution is rather ad-hoc and discards some of the in-dataset actions which could potentially limit the performance of the trained policies.
In this work, inspired by the recent success of Normalizing Flows (NFs) (Singh et al. 2020) in the online reinforcement learning setup, propose Flows (CNF) for constructing a latent action space useful for offline RL. First, we describe why a naive approach for constructing latent action spaces with NFs is also prone extrapolation error, and
method called Conservative Normalizing

---

Under review as a conference paper at ICLR 2023
then outline straightforward architectural modification that allows avoiding this without need for manual post-hoe clipping. Our method is schematically presented in Figure [ where we highlight key differences between method and the previous approach. We benchmark our method against other competitors based on generative models and show that it performs favorably on a large portion of the D4RL (Fu et al. 2020) locomotion and maze2d datasets.
2 PRELIMINARIES
Offline RL The goal of offline RL is to find policy that maximizes the expected discounted return given a static and pre-recorded dataset D consisting of state-action-reward tuples. Normally, the underlying decision-making
problem is formulated via Markov Decision Process (MDP) that is defined as a 4- 4-elmments tuple, consisting of state space S, action space A, state transition probability
S x S x A + |0, so), which represents probability density of the next state s' E S given the current state a and action a E A; bounded reward function : S x A x - [Tmin, Tmaz] and scalar discount factor y E (0, 1). We denote the reward r(s;, a, S(+1)
Tt. The discounted return defined as R₂ = Ex= T Also, the notion of the advantage function A(s, a) is introduced difference between state-action value Q(ss, a) and state value ((s) functions:
Q (=(s at) = E
retal
k=0
(I)
V*(s) = Ean 10" a)] (s, a) Q" (s. a) (*)
Advantage Weighted Actor Critic One way to learn a policy in an offline RL setting is by following the gradient of the expected discounted return estimated via importance sampling (Levine et al. 2020), however, methods employing estimation of the Q-function were found to be more empirically successful (Kumar al. 20200 Nair al. 2020 Wang et al., 2020). Here, we describe Advantage Weighted Actor Critic (Nair et al., 2020) where the policy is trained by optimization of log- probabilities of the actions from the data buffer re- re-weighted by the exponentiated advantage.
In practice, there are two trained models: policy Ty with parameters @ and critic Qp with parameters wb. The training process consists of two alternating phases: policy evaluation and policy improve- ment. During the policy evaluation phase, the critic Q* (s, a) estimates the action-value function for the current policy, and during the policy improvement phase, the actor æ is updated based on the current estimation of advantage. Combining all together, two following losses are minimized using the gradient descent:
(+) (ss.a)=Dl- log No (als) exp(AA(s, a)/
LTD(W) ((s.a.ns)>>((1> + 2 a' m-(sss)) Q+(s,
Where As(s, a) is computed according to Equation using critic Q+ and À is temperature hyper- parameter.
Normalizing Flows Given a dataset D = 1, with points (00) from unknown distribution with density px the goal of a Normalizing Flow model (Dinh et al. 2016 Kingma & Dhariwal 2018 is to train an invertible mapping z = (x) with parameters Φ to a simpler base distribution with density pz, typically spherical Gaussian:
10, This mapping is required vertible by design to sample new points from data distribution by applying the inverse mapping to samples from the base distribution: 'z) full flow model is a composition of K invertible functions f, and the relationship between Z and x can be written as:
be in-
- h la h
{x
Log-likelihood of data point x is obtained by using the change of variable formula and can be written as:
log (x) = log (z) + log (ez//d/
k
(4)
= log Pzz (z)) + log /dh;
2

---

Under review as a conference paper at ICLR 2023
Normalizing Flows models are optimized to directly maximize ogglikelihood of data points using Equation
METHOD
Supervisedly pre-tralned models
Action space
space
PLAS
state Re
latent vector
VAE decader D3
action
Environment
state W0
latent vector
Flow model
action
CNF
U 11)"
Figure Schematic visualization and comparison PLAS and CNF (ours) approaches. Both methods an action encoder-decoder model trained in supervised manner and a controller model to select actions from the latent space of the encoder. PLAS algorithm uses VAE with normal latent distribution with unbounded support (represented restricts latent policy outputs to only a part the latent space (represented as the black borders inside of latent space). Our algorithm uses Normalizing Flow instead of VAE and bound base distribution itself, allowing the latent policy to use the whole latent space.
offline dataset
the blue circle) and
In this section, we describe our approach named Conservative Normalizing Flows (CNF) in detail (Figure I... We start by delineating the proposed architecture and then describe both the pre-training and policy optimization phases.
3.1 CONSERVATIVE NORMALIZING FLOWS
Normalizing Flows models are trainable bijective mappings, composed of fully invertible layers, that transform the original data space to the base distribution space. The latter is often referred to as the latent space. Typically, the base distribution is modeled as normal since it provides a tractable density and is easy to sample from (Kingma & Dhariwal 2018). However, similar to the PLAS (Zhou et al. 2020), this distribution would result in unbounded support. Therefore, when learning a policy in the latent space maximize q-values it may exploit the regions that lead the out-of- distribution actions in the original space exaggerating the extrapolation error.
To illustrate this problem better, consider toy 2-d task with modeling moons dataset which consists of points of two interleaving half circles (Buitinck et al. 2013) using normalizing flows. We train two normalizing flow models, first with the normal latent distribution and second with the uniform latent distribution. We additionally squeeze data points to lie in (1) 2-d region via linear transformation in this experiment, and add the inverse of tanh function as the first layer of NFs models, so that tanh function applied the last transformation during sampling from the model. Both trained NFs model the underlying training distribution well (Figure 2). We then select the NF model with normal latent space and gradually increase the amplitude of the latent samples drawn from a- N(0, I) during sampling new data points from the model, this process is presented in Figure 3 And, expected, higher amplitude values result in more out- ant-of-distriuution data points.
One solution to avoid exploiting regions leading to oft-of--istribuue actions is to restrict policy output amplitude by some value. This approach was proposed in /houu et al. (2020); Chen et al.
3

---

Under review as a conference paper at ICLR 2023
10
10
0.5
0.5
-0.5
-0.
-1.0
-1.0
-3.5
0.0
0.5
1.0
-1.0
1.0
-1.
-055
0.0
0.5
-1.0
-0.5
10
(a) NF-Normal, Density (b) NF-Normal, Samples (c) NF-Uniform, Density (d) NF-Unformm, Samples
Figure 2: A toy example demonstrate that both NFs either with Normal or Uniform latent dis- tributions can recover the training data. However, as we demonstrated above, a potential controller trained the latent space normal-based NFs is still able sample actions outside of the training dataset.
2022), where the output of the deterministic policy in the latent space was modeled as = Emar (s)) and Zmaz was set to 2.
However, it was shown that the optimal clipping value Smar different for every locomotion dataset in the D4RL benchmark (Zhou al. (2020), Appendix C). Since Zmer value essentially part of the latent space policy model, it is necessary to evaluate multiple values online after training to select the best one, which may be feasible in some tasks.
To tackle this problem we would like to avoid this post-hoc clipping and construct latent action space model that would prohibit the exploitation of out-of-distribution actions by design. To do so, we make use of NFs versatility, and add invertible tanh activation after the last layer of the Normaliz- ing Flows model it makes NFs outputs lie bounded n dimenssioo ininerraall 1-1,+11", which in turn allows to substitute normal base distribution with n nn-dimesional uniform With these changes, the potential actor model should not be capable generating actions outside of the training distribution.
1.0
10
10
-1.0
-10
-1.0
-0.5
0.0
0.5
10
-1.
-0.5
0.0
0.5
100
-10
-0.5
0.0
0.5
1.0
-10
-0.5
0.0
0.5
10
(a) Amplitude =2
(b) Amplitude 4
(c) Amplitude = 10 (d) Amplitude
Figure 3: A toy example demonstrates that the normal- based NF can be manipulated by controller to produce fuitt--distribution training points. We model this by increasing the amplitude of the latent space samples, tries to maximize q-values. Note that the model with uniform latent space does not suffer from this problem because it already uses the whole latent distribution support during the training and sampling processes.
found this to happen in the preliminary experiments when the controller
3.2 LATENT ACTION SPACE PRE-TRAINING
First, start by pre-training conditional Normalizing Flow model f(als) with parameters á on the actions and states from the offline dataset, this is essentially a supervised learning problem. We use the same Normalizing Flow conditioning scheme as in the SSnngh et al. (2020) (for more details, Figure 8 in the original paper): for each NF layer, we add additional input for the state vector. Given
is the action space dimensionality.
4

---

Under review as a conference paper at ICLR 2023
conditional NFs f with parameters o, use it to compute log-likelihood log Pş(a(s) of actions conditioned on states s from the offline dataset D. minimizing the following loss:
Le(Ø)= B(s,a))oll (ogpo(als))
(5)
Algorithm CNF, Pre-training
1: Input: offline data replay buffer !
2: Initialize flow action encoder parameters ø
3: repeat until convergence
4: Sample a mini-batch B = {(s, } from D
5: Compute flow loss using Equation compute loss gradient and update flow model:
V
| loggpo(aass)
(x.a)EB
The final result of this optimization process is mapping A Z, where Z is the latent action space bounded by the (-1,1) interval. As this mapping is invertible, we can further transform latent vectors into the original action space for policy optimization.
3.3 POLICY OPTIMIZATION
The whole procedure is outlined in Algorithm 2 Here, we describe this phase in detail as follows. Given a pre-trained latent action space model, we freeze its parameters and add it as an action encoder during RI training. We modify policy optimization loss from Advantage Weighted Actor Critic (Nair et al. 2020) to use it with our flow model. Specifically, we train stochastic latent policy model [= with parameters 0, which predicts μ and o² for tanh(N (n, (u,o²)) distribution, and two critic models (s, a). Q2(s, a) with parameters t and thb to mitigate positive bias in the policy improvement step that is known degrade the performance of value-based methods (Hasselt) 2010 (ujimoto et al.] 2018). Note that our policy model operates in the latent space, not in the original space. By combining latent policy and a pre-trained normalizing flow model, obtain the following loss function for policy optimization:
(0) = (s..)) (s, a) la - f- ~ TT(-ss)|s)]
(6)
where weights comes from exponentiation of the Advantage function:
w a, s) = exp ((Q(s, flssff (T-(s)ss))) /A)
= exp (A(s, a)/A)
and Q-function is the minimum between two trained models: Q(s, a) = min (s a), as proposed by Fujimoto et al. (2018). Here, À E (0, oc) is temperature hyperparameter: for higher values training objective behaves similarly to behavioral cloning and for lower values aims to maximize advantage. Overall, this loss function train a policy that maximizes the Q-values subject distribution constraint (Nair et 2020).
Together with the latent policy model, we optimize two critics with the standard Q-learning loss:
Lo 122) = E(s.a..s) (QQ - y)² (Q2(s, a) - 1
y = r + ((()))))))))))))))))))))))) [min
[min Qi(s', a')) 1=1,2
(8)
Since NFs are differentiable, we can compute the gradient of the given loss functions for the policy model weights using the chain rule and reparametrization trick (Kingma al. 2019). Note, that one could bypass the differentiation through the NFs model, and optimize policy and crities the latent space directly. However, show in the Appendix (Section A.4), that this results in worse policies.

---

Under review as a conference paper at ICLR 2023
Algorithm 2 Offline RL with CNF training
1: Input: Pre-trained flow action encoder model f(a]s) with parameters o, offline data replay
buffer !
2: Initialize actor with parameters 0 and two crities Qı and Q2 with parameters ta and 4/2
3: repeat for given number of train-ops
4: Sample mini-batch {(s. a, r; s')} from D
5: Sample next-state actions a' using policy π (-|ss) and flow f(-ls") models, compute Q-target:
a' (= n (-1s)) for all s' E B
> min Qi(s', 1=1.2
6: Compute critics loss using Equation
8 compute loss gradient and update models:
Σ
(s,r y)² + (Q2(s, a) -
y)
B
(x.0,r.a")BB
7: Sample actions a using policy
(-|ss) and flow [-[(8) models, compute advantage weights
à = f (z (T||s) for all E B
A ( a) = Q ((s,
w = 6 compute loss gradient and update policy model:
A)
8: Compute policy loss using Equation
V 0
axp(A(s,a)/A) -å
B
(s
EXPERIMENTS
To show how the proposed method works, we benchmark on various locomotion and navigation tasks from the popular D4RL benchmark (Fu et al. 22220) comparing it to the other methods based on generative models PLAS (Zhou et al.) 2020) and LAPO (Chen et al. (022). We also include AWAC (Nair et al. 2020) algorithm in our comparisons because build policy optimization method on top of it, and IQL (Kostrikov et al.) 2021) algorithm because of its competitive performance across non-ensemble methods.
4.1 D4RL BENCHMARK
Locomotion We focus on three locomotion environments from the D4RL dataset: Walker2d-v2, Hopper-v2, and HalfCheetah-v2. For the Normalizing Flows model's pre-training phase, divide the training dataset into two parts by separating 10% portion of randomly selected data for validation. We run 50 experiments with the random search of hyperparameters from Table 3 and then select the best model according to the log-likelihood on the validation dataset for RL training. We train all models, including Normalizing Flows, latent policies, and critics, using Adam optimizer KinimiaerKinnmma & Ba 2014). During the RL phase, we run million training steps for actor and critic models on all environments except HalfCheetah-v2. where we use only 200.000 training steps as it is enough for convergence. We evaluate the agent by running 10 episodes and averaging the scores over them once per 5000 (100 for HalfCheetah-v2)) training steps. Hyperparameters for RI. training are listed in Table 2
For comparison, we implemented AWAC and IQL algorithms in our code base and used the publicly available PLAS implementation. We run AWAC and IQL algorithms with hyperparameters found in the papers Kostrikov et al. (2021) : Nair et al. p Zhou et al. (2020) and PLAS with hyperparam- reported in the paper. The final
eters recommended in Zhou al. (2020). For LAPO, scores of benchmarked algorithms are summarized in Table Missing environments are labeled as t_" in the table.

---

Under review as a conference paper at ICLR 2023
To highlight the performance of our method, we include training curves in Figure 4 It can be that the proposed algorithm exhibits preferable performance on all 9 locomotion datasets, especially on the HalfCheetah--2 environment.
Normllized Score
wellxer2d-edduumz2
malleer2-medium-reelay..
walker med mdd
expertv2
100
100
50
600
54
05
04
10
0.
32
3.4
0,6
0.
2.5
20
22
644
86
20
1.0
Training Steps
Training Steps
Training Steps
Normalzee malazedd Score
hopper-medium-w2
ropper-medumm-relayy-22
hopper-medium-experrs
50
50
PJ6
Training Steps
166
Training Steps
1e6
Training Steps
166
Nomalized Score
halfchectah- mediuum--
maaftchetah-mmeiiumae
ha halfcheetah--meee medium medium-expert.v2
25
19392 Training Steps
50000
Training Steps
Training Steps
Figure 4: Average normalized performance D4RL locomotion tasks. The x-axis denotes the training steps. Each curve is averaged over 3 random seeds. Shaded area represents one standard deviation.
Maze2d For our next experiment, we choose three maze2d datasets with increasing layout com- plexity: umaze-v1, medium-vl, and large-vl. We use the same flow pre-training scheme as before: we run 20 experiments with random hyperparameters from Table 3 and select the model with the highest log- bigllikel on the validation dataset. We list the hyperparameters for RL training used in this experiment Table 2 For comparison, we use our implementation of IQL, rely on a publicly available PLAS implementation, and report scores from the LAPO paper. As can be seen in Table [ our method outperforms baselines out of 3 environments.
4.2 ABLATIONS
To study the importance of each major component in the proposed CNF method, we conduct addi- tional ablation experiments. We begin by comparing training policies in uniform and normal latent spaces. Then, we examine additional clipping for the latent policy to see if it can operate in normal latent space. And finally, integrate the VAE model into our approach and compare performance between Normalizing Flows and VAES for training policies in latent spaces.
Normalizing Flows with normal latent distribution, no clipping In the first experiment, we com- pare the performance of the proposed method but with normal latent distribution in Normalizing Flow and latent policy models. We use conventional Normalizing Flows encoder without tanh activation after the last layer and with the normal latent distribution. We pre-train action encoders in the same way as before, using only the best hyperparameters from previous experiments (note that these hyperparameters result in very similar models in terms of the training and validation losses, as depicted Figure g). We also remove tanh activation after the last layer from the latent policy model, letting it operate over the whole latent space of the pre-trained Normalizing Flow action encoder, which we make to predict and o² for N N(u,o3 distribution. We again select the HalfCheetah-Ve v2 environment and run 3 experiments with different random seeds per dataset. We plot the results in Figure 5) One can see that the performance degraded substantially, indicating that latent policy without clipping could not be trained to produce competitive performance.

---

Under review as a conference paper at ICLR 2023
Table 1: Normalized performance on D4RL benchmark. The scores are averaged over 10 final evaluations and 3 random seeds. The results are reproduced for all of the algorithms except LAPO, for which we take the values stated in the original paper. CNF outperforms other methods on of 9 locomotion datasets, and on two out of three maze2d datasets.
Environment
AWAC
IQL
PLAS
LAPO
CNF
walker22-mecium--22
78.20
78.80
71.2
80.75
83.60 3.01
walker2d--mediu--rrpla.
76.76
74.49
2.74
81.96 + 1.98
walker2d-mcciim-expprr
112.97
111.96
108.13
112.32 = 0.21
hopper-medium-v2
62.59
63.7
52.93
51.63
69.32 + 1.04
hopper-mcdium--eplayy-2
73.12
87.72
3.17
89.04 + 10.39
hopper-meddum--expertttV
109.64
109.05
106.5
108.6 + 5.45
halfcheetah-mmchiu--22
43.15
47.4
43.78
45,97
50.55 + 0.53
hallcheetai-mmcduu--rrpa------
42.00
43.2
44.8
45.84 0.31
halicheetad-mcdiuu--eerttttzz
87.40
78.95
86.63
96.23 + 0.20
locomotion-v2 average
76.20
77.25
57.76
81.94
maze2d-umaze-vl
37.69
53.9
118.9
62.9 = 10.36
maze2d-mediun-vl
35.45
66.4
142.8
155.89 10.49
maze2d-large-vl
49.64
107.2
200.6
212.8 + 2.23
Normalized Score
hallcheetah-mediuum22
hellcheetah-medduu--eaaayyyy
59392
100028
300020
Training Steps
Trainingg Steps
Training Steps
Figure 5: Comparison of the proposed method with uniform (CNF) and normal (NF) latent spaces. Policy performance is significantly worse when the latent space is normal.
Normalizing Flows with normal latent distribution, manual clipping As it was shown in the previous ablation, latent policies without clipping in normal space perform poorly. To test how different clipping values affect the agent's performance, we did the following experiment, which is similar to the previous one, except manually clip latent policy output by some value. To do so, we return (anh activation at end the policy model and multiply by a, treated as a hyperparameter. Latent policy output is modeled as 2 distribution with parameters predicted by the policy model. This experiment is similar to Zhou et al. Normalizing Flows extract latent policies using
(anh(N(mo((
where N is normal
(2020) Ablation C, but instead of VAEs, Also, we use datasets from the Walker2d-v2 environment. We examine several values for parameter a and compare them with the proposed method. The results are averaged over three random seeds and are shown in Figure 6 One can see that optimal clipping values are different for each dataset: for the medium dataset, there is no disparity in performance, but the clipping value of 3 produces slightly better results and almost matches the performance of the uniform latent space; for the medium-replas it equals 2 and for the medium-expert it equals 1. On the other hand, CNF, parameterized with the uniform latent distribution, does not add extra clipping value as hyperparameter and performs better on each dataset.
Algorithm [
Action encoder model: Normalizing Flow and VAE In this experiment, we compare the perfor- mance of latent policies obtained by the CNF with different action encoders, namely, Normalizing Flow and VAE. We adopt the best VAE architecture and training parameters from Zhou et al. 2020)) and then integrate pre-trained models into the latent policy training as Algorithm 2 For this com- parison, rely the HalfChectah-r2 environment and plot training curves for PLAS, CNF, and CNF- VAE. Results are presented in Figure?7 One can see that the CNF performs best with the Nor- malizing Flow action encoder. On the other hand, the of the VAE encoder shows performance
8

---

Under review as a conference paper at ICLR 2023
walker2d- 2ilker22-mediumm222
wallee22d-eediummrepleee..
Walker22dmeen www.mmedduummeee eexpeteeeeeee-2222222
75
50
50
25
unifomm
0
uniform
0
uniform
Figure 6: Normalized average performance of the proposed method (uniform on the X-axis) and latent policies with clipped normal latent distribution. Number on the X-axis the clipping value. Optimal clipping value is different for each dataset and final performance is better for the proposed method.
similar PLAS. On the wwwllchcetah-meduum--ree formance to CNF, but exhibit marginal improvement during the training process. This experiment indicates the importance of the whole pipeline with the use Normalizing Flows as opposed the VAEs.
dataset, it starts from a commensurate per-
Nomalizzd Score
halfcheetah-medlum-22
hallcheetah-meddu---paaa---
ppllchhetth-meelum-ennnee
33292
300800 10000 Training Steps
200000
30000 300000 100000 Training Steps
200c00
30000
100000 130030 Training Steps
200000
Figure 7: Comparison of CNF, CNF with VAE (CNF-VAE) and PLAS. CNF with Normalizing Flow and uniform distribution performs above all. For normal distribution optimal clipping value is different for each dataset.
CONCLUSION
In this work, we presented new deep offline RL method called Conservative Normalizing Flow (CNF). It constructs a latent action space with the use of the NFs model and then runs policy optimization within
This approach makes trained policies capable of fully utilizing the latent space opposed to the post-hoc manual clipping procedures in PLAS PLAS et al. (2020) and LAPO Chen et al. (2022). We benchmarked our method against other competitors based on generative models and showed that it performs favorably large portion of D4RL Fu et al. (2020) locomotion and maze2d datasets.
REFERENCES
Lars Buitinck, Gilles Louppe, Mathieu Blondel, Fabian Pedregosa, Andreas Mueller, Olivier Grisel, Vlad Niculae, Peter Prettenhofer, Alexandre Gramfort, Jaques Grobler, Robert Layton, Jake Van- derPlas, Arnaud Joly, Brian Holt, and Gaël Varoquaux. API design for machine learning software: experiences from the seikit-learn project. In ECML PKDD Workshop: Languages for Data Min- ing and Machine Learning, pp. 108-122. 2013.
Xi Chen, Ali Ghadirzadeh, Tianhe Yu. Yuan Gao, Jianhao Wang. Wenzhe Li, Bin Liang, Chelsea Finn, and Chongjie Zhang. Latent-variable advantage-weighted policy optimization for offline rl. arXiv preprint wwx.i::203.08844. 2022.
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. Density estimation using real nvp. arXiv preprint arXiv:1605.08800 2016.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprini arXiv:2004.07219, 2020.

---

Under review as a conference paper at ICLR 2023
Scott Fujimoto and Shixiang Shane Gu. A minimalist approach to offline reinforcement learning. Advances in neural information processing systems, 1420132--21155 2021.
Scott Fujimoto, Herke van Hoof, and David Meger. Addressing function approximation error in actor-critic methods, 2018. htthtts::://aat httttps://axxx..orrraaa//////// 09477
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without PMLR, 2019.
exploration. In International conference on machine learning, pp.
Tuomas Haarnoja, Kristian Hartikainen, Pieter Abbeel, and Sergey Levine. Latent space policies for hicrarchical reinforcement learning. In International Conference on Machine Learning, pp- 1851-1860. PMLR, 2018.
Hado Hasselt. Double q-learning. Advances neural information processing systems, 23, 2010.
Natasha Jaques, Judy Hanwen Shen, Asma Ghandeharioun, Craig Ferguson, Agata Lapedriza, Noah Jones, Shixiang Shane Gu, and Rosalind Picard. Guman-centrie dialog training via offline rein- forcement learning. arXiv preprint arXiv:2010.05848, 2020.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint 141.6904,, 2014.
Diederik Kingma, Max Welling, et al. An introduction variational autoencoders. Foundations and Trends@ in Machine Learning, 3074)33 2019.
Durk P Kingma and Prafulla Dhariwal. Glow:
Generative flow with invertible 1x1 convolutions. Advances in neural information processing systems, 31, 2018.
Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q- learning. arXiv preprint arXiv:2110.06169, 2021.
Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off- policy q- learning via bootstrapping error reduction. Advances in Neural Information Processing Systems, 32, 2019.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems, 33:1179 19111999119111 1191, 2020.
Sascha Lange, Thomas Gabel, and Martin Riedmiller. Batch reinforcement learning. In Reinforce- ment learning, pp. 45-73. Springer, 2012.
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tuto- rial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643, 2020.
Ashvin Nair, Abhishek Gupta, Murtaza Dalal, and Sergey Levine. Awac: Accelerating online rein- forcement learning with offline datasets. arXiv preprint arXiv:2006.095: 2020.
Avi Singh, Huihan Liu, Gaoyue Zhou, Albert Yu, Nicholas Rhinehart, and Sergey Levine. Parrot: Data-driven behavioral priors for reinforeement learning. arXiv preprint arXiv:2011.1000 2020.
Ziyu Wang, Alexander Novikov, Konrad Zolna, Josh S Merel, Jost Tobias Springenberg, Scott E Reed, Bobak Shahriari, Noah Siegel, Caglar Gulcehre, Nicolas Heess, et al. Critic regularized regression. Advances in Neural Information Processing Systems, 33:7768-7778. 2020.
Wenxuan Zhou, Sujay Bajracharya, and David Held. Plas: Latent action space for offline reinforce- ment learning. arXiv preprint 2rXiv:20/1.072/5 2020.
10

---

Under review as a conference paper at ICLR 2023
A
APPENDIX
A.1 RELATED WORK
Normalizing Flows in RL Several prior works applied Normalizing Flows models for reinforce- ment learning tasks. In the work Haarnoja et al. (2018) Normalizing Flows were used for training hierarchical policies using the Soft Actor-Critic RL algorithm. They choose Normalizing Flows be- cause it provides expression for exact likelihood computation and has an intuitive way to stack sequence of models to construct one hierarchical model. In contrast to our work, they did not modify NF's latent space, they did not use offline pre- ree-triniig of flow models and train models in the online RL framework.
The PARROT work (Singh et al. 2020) proposes to use the Normalizing Flows action encoder dur- ing the behavioral cloning pre-traning phase before running reinforcement learning on the target task. After pre-training, they freeze Normalizing Flow and run reinforcement learning and unseen task. They aimed to develop a method leveraging near-optimal demonstrations dur- ing Normalizing Flows pre-training and speed up the convergence of RL by better and meaningful exploration, while we use Normalizing Flows in purely offline RL setting and to make agent conser- vative.
new
Offline RL A large portion of recently proposed deep offline RL algorithms focuses on addressing the extrapolation issue, trying impose certain degree of conservatism limiting the deviation of the final policy from the behavioral one. Researchers approached this problem from multiple angles. For example. Kumar al. (2020) proposed directly penalize out-of-distribution actions, et al. (2021) avoids estimating values for out-of- actions completely. Others Fujimoto Gui (2021): Kumar et al. (2019); Jaques et al.] (2020) put explicit constraints to stay closer to the behavioral policy. Here, we took different approach by constructing latent action space that allows us to bypass the need for explicit regularizations.
Offinee RL with generative models In the method named PLAS (Zhou et al. 2020), the authors proposed to pre-train conditional variational autoencoder on actions from an offline dataset. This idea resembles ours, but make agent conservative authors restrict policy outputs in the latent space to a fixed range. In our work, we aimed to make a better action encoder model by switching from VAEs Normalizing Flows. This allows encoder, avoid manual clipping, and we experimentally demonstrate that our approach leads to better performance the popular offline RL benchmark.
utilize the whole latent space of the action
One recent approach called LAPO (Chnn et al. 2022) proposes to train an action encoder together with a reinforcement learning agent. They motivated this by observing that the action distribution does not match the return distribution in the training data set, and therefore actions that lead to higher returns are more important for action encoder training. In our work, we examine an orthogonal direction, studying a different generative model for action encoding.

---

Under review as a conference paper at ICLR 2023
A.2 HYPERPARAMETERS
Table 2: Hyperparameters for latent policy and critics models for RL training phase on locomo- tion and maze2 tasks. We grid-search for the values written in square brackets. Rest of the parameter were fixed for all of the datasets.
Hyperparameter
value
number of training steps number of training steps (HallChcetah-v22)
1000000
200000
number of layers (locomotion)
[3, 4]
number of layers (maze2d)
4
hidden size learning rate batch size (locomotion) batch size (maze2d)
256
3e-4
[256, 512, 1024]
10240
X-temperature (locomotion)
1/3
X-temperature (maze2d)
1/10
Table NFs hyperparameters for the supervised pre-training phase on locomotion and maze2 tasks. Medium, medium-replay, and medium-exper datasets are marked as m, m-r, and m-è correspond- ingly. We sample learning rate and weight decay from the continuous uniform distribution.
Hyperparameter
value
number of training steps
100000
number of layers (locomotion m m-e datasets)
number of layers (locomotion m-r and maze2d datasets)
4
hidden size (locomotion m and m-c datasets) hidden size (locomotion m-r and maze2d datasets)
256
64
learning rate
min = le-5, max 3e-3
weight decay
min 0.0, max 1e-2
batch size
[512, 1024, 2048)
A.3 NORMALIZING FLOWS TRAINING CURVES
Neggiive togs leeitoood
halicheetal-mmedimm222
ha hnlicheraat-meedu--egiii
holicheetal-mediimme exerrar--teeuuueeeeeeeee
Uniform trii
Unifrrm
Unifom valid
Uniform volid
homal train
Normal valid
Uniform volid
Unitorm trai-
Normal vold
Normall train
-0
-5.0
Normel tnain
-7.5
20000
40000 60000
50000 100000
20000
40000 60000 80000 200000 Training Steps
20000 40000 G0000
80000 100000
Training
Training
Figure 8: Training and validation loss for Normalizing Flows pre-training with Uniform and Normal latent spaces. Both training and validation curves almost identical between models with different latent spaces, which means that models have similar encoding and reconstruction quality. For each latent policy training, we select NF model with the lowest validation metric.
12

---

Under review as a conference paper at ICLR 2023
A.4 TRAINING CONTROLLER LATENT SPACE
Ideologically, policy optimization can be carried out directly in the latent space. This can be done by simply substituting the original actions with their latent counterparts. After this substitution, a myriad of offiin RL algorithms can be used to extract a new policy. To test if this is a viable approach, we train actor and critic models in the latent space of the Normalizing Flows without utilizing gradients from the action encoder during the policy optimization phase. First, we convert actions from the original environment's s action space the latent space by encoding them with the pre-trained action encoder. After that, we run the AWAC algorithm with no additional changes to the dataset with converted actions. We conduct this experiment on HalfCheetah-v2 datasets: medium, medium-replay, and medium-expert. Results are presented in Figure ? It can be seen that this approach shows promising results, but convergence speed and the final score are slightly lower all datasets. We conjecture that the performance is lower because metric in the latent space induced from the original action space by the flow model, and optimization of the distance between latent vectors corresponds optimization of the distance between actions only implicitly. When training policy using metrics from the original action space, as suggested by Equation [6] the Normalizing Flows model provides gradients to the policy model that guide minimize the actual distance between actions.
Nomalized Score
hallcheetah-medlum-22
hallchhetah-mmdduu---eeaa--x
33227
350828
150000
200000
30028
300028
120000
200000
30000 100000 330000 Training Steps
200000
Training Steps
Training Steps
Figure 9: Comparison of CNF trained in original (blue) and latent (red) action spaces. Final perfor- mance is slightly lower for the training directly in latent space.
13


Reviews:
[{"review_id": "oGZikrrPHB-", "reviewer": "Reviewer_17pi", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper introduces Conservative Normalizing Flow (CNF) which encourages the final offline RL policy to remain withing the data distribution while being able to run policy optimization over an unconstrained latent space. They achieve this by pre-training the latent space to be a uniform distribution using normalizing flows that learns bijective mappings using fully invertible layers. The key novelty here is the use of bounded functions such as tanh at the last layer of normalizing flow such that the latent distribution may be learned as a uniform distribution between -1 and +1. This allows us to model/sample the behavioral action distribution using the entire latent space. (no clipping required.). Experimental results show that they are competitive with other SOTA offline RL methods in many domains.", "summary_of_the_review": "Overall the paper provides a novel approach on sampling well covered actions for the downstream task of offline RL, However the experimental results show some anomalies that I would like to be been elaborated further by the authors. I would be willing to change my score with some explanations regarding the anomalies in the scores and some additional experiments.", "strength_and_weaknesses": "Overall the paper is written well and is intuitive. The use of normalizing flows to construct a latent space that leads to conservative agents by design is unique yet neat. The ability to learn better action-spaces that allows to model the state-action coverage of the dataset by default can be very useful for a many downstream and RL learning problems.\n\nIt would be wonderful to see some experimental analysis on the qualitative or quantitative differences between the action sampling portion of the current model and other methods in terms of being able to sample state action pairs that are well covered in the data set along with if and when these models are prone to over-fitting as discussed in the paper. Some empirical evaluation of different action generating functions in real wold datasets would be really helpful to push the community towards using Normalizing Flows with uniform latent distribution as a principled way of sampling actions in offline RL settings.\n\nRegarding the experimental results, while it is competitive with the SOTA methods in many cases, it is not trivial why it fails to perform well when it does not. For example it is surprising that the method does not work as well for domains like maze2d umaze and medium expert tasks. (It is expecially surprising that it is performing worse in hopper medium expert as compared to hopper medium replay. This is an anomaly as generally  most offline RL algorithms tend to perform better with expert dataset. I would love to know more about what authors think about this result. Moreover it would be wonderful to have results on more involved domains such as antmaze to further bolsterd the claims made in the paper.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "dPtKvBv5Ex", "reviewer": "Reviewer_iyr1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "Offline RL is challenging when the learned RL policy drifts too far from the support of the dataset. Thus, many offline RL methods use some form of constrained or conservative policy update to ensure that the RL policy remains close to the behavior policy of the dataset. One method for doing this is to train a generative model (e.g. a VAE) on the data, and learn an RL policy that picks embedding vectors which are decoded into actions by the VAE decoder. This paper proposes a slight tweak on that prior approach: rather than learning a VAE on the dataset, what if we learned a model using normalizing flows? Using this Conservative Normalizing Flows approach, the authors transform the prior into a Uniform distribution rather than a Normal distribution, which ensures that the RL policy can choose samples anywhere within the support of the Uniform distribution without generated OOD samples, unlike with the VAE. The results show some improvement over existing methods on a sample of tasks from the D4RL offline RL benchmark. ", "summary_of_the_review": "In summary, the paper clearly and precisely investigates a specific hypothesis related to using normalizing flows rather than a VAE for a particular type of Offline RL approach. It is benchmarked against relevant Offline RL baselines and shows an improvement in 67% of environments, and a strong detriment in others. Relevant ablations are conducted to give intuition as to why the proposed approach is important. The novelty above prior work is somewhat limited. ", "strength_and_weaknesses": "A strength of the paper is that it precisely targets a clear hypothesis: whether NF can improve the ability for an offline RL agent to learn to control a pretrained model vs. a VAE. The motivation for the paper is clear, as is the distinction from prior work.\n\nA weakness of the paper is its significance, or impact. It makes a relatively minor tweak to the approach proposed by Zhou et al. 2020, and while the results demonstrate gains on 67% of environments sampled for the paper (8/12), it also leads to dramatically worse performance in some environments. Thus, the novelty and effectiveness of the approach is limited. \n\nOne possible avenue for future work that the authors could consider in order to increase the significance of their results, is to think about applying this approach to learning to control large language models with RL. Currently, it's difficult to train LLMs with RL due to the complexity of LLM infrastructure and training with RL on such a large scale (i.e. fine-tuning all the parameters could be prohibitively expensive). One approach would be to impose a bottleneck embedding layer within an LLM (see e.g. https://arxiv.org/abs/2209.06792), and then train an RL agent to pick embedding vectors to optimize some language objective. The issue is that this can be quite difficult to get working. The linked vec2text paper suggests that using a VAE doesn't work well in this context. So perhaps normalizing flows could be useful there as well. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "KlVNIglq1x5", "reviewer": "Reviewer_gFEy", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper presents a method using normalizing flows for offline reinforcement learning (ORL).\nTo alleviate the difference between training and testing datasets, the proposed method maps uniform distributions to action spaces using the flow model.\nExperimental results present that the proposed method outperforms the state-of-the-art methods in many benchmark datasets.\n", "summary_of_the_review": "The proposed method is interesting and the paper is well-written.\nHowever, some ablation studies and clarifications of the performance improvements are missing.", "strength_and_weaknesses": "This paper is generally well-written and easy to follow.\nIt is reasonable that the uniform distribution is better than the normal distribution for the latent space in ORL.\n\nThe reviewer has two concerns as follows.\nFirst, the necessity of the flow model is not well supported.\nFigure 6 presents that the proposed model performs poor scores when using the normal distribution as the latent space.\nPlease report an ablation study for the pre-trained models (VAE decoder vs. flow model).\n\nSecond, the performance improvement by the latent space clipping is missing.\nThe latent space with the uniform distribution has clipped values in [-1,1], while the latent space with the normal distribution is unbounded.\nNormal distribution with value clipping in [-1,1] will provide fair comparisons.\nIn this regard, the reviewer wants to see the results of NF-uniform under the experiment setting in Figure 3.\nIf the results are the same as in Figure 2(b) regardless of the value of amplitude, the proposed method does not address \"extrapolation\".\n\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "Not applicable"}, "extra_text": {}}, {"review_id": "ATs5t9xD1C", "reviewer": "Reviewer_DZCL", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper utilized *Normalizing Flows* models, which are trainable bijective mappings, composed of fully invertible layers. The core idea is to obtain a bounded support in the latent space by design via shifting the base distribution space (that of the latent variables) from a normal distribution to a bounded uniform one. The proposed method seems to be based on PLAS (Zhuo et al.), which constructs a latent space using VAE that maps to the actions which are well-supported by the training data. It then learns optimal policy in this space. ", "summary_of_the_review": "I had a hard time following the paper. It contains too many mistakes, inconsistency, and unclear/confusing parts. I would strongly recommend the authors should start fresh and completely rewrite the draft.\n\nHere are my main comments; hopefully they help:\n\n- The choice of $f(x) = tanh(x)$ as the activation needs to be discussed. In particular, $tanh$ has a strong saturation for when $x$ passes one. In other words, its behaviour is *almost* linear for when the advantage function is in $[-1,1]$ and somehow clipping $A(s,a)$ when falling outside unit interval (or $[-2,2]$).  Is it a limitation? Discussion needed. \n\n- In equation 2, $\\textbf{f}$ applies on the advantage function, while later $\\textbf{f}$ is defined over $(s,a)$. If by these you mean the same function, fix your definitions. If not, then use different letters. \n\n- Section 3 is quite unclear. Specially the paragraph before Fig. 3. Further, the range of data (from -1 to 1 on each axis) seems to be arbitrary with no explanation provided. \n\n- Page 4, end of first paragraph: “one is still able to find latent vectors that would result in out-of-distribution points.”  -> it is not inferred from the figures. Perhaps consider an example in the figures for this statement. \n\n- The idea of mapping the value function and training the Q network in the mapped space has also been studied in the recent literature. Though it might not be directly related, it could be beneficial to borrow some ideas, or proofs. See for example: https://arxiv.org/abs/2203.07171\n\n- Equation 6 is very confusing: sampling from $\\pi(\\cdot |s)$ gives rise to actual actions (not the z space).\n\n- Consider discussing equations 6-9 in much more details, as they are the core of your paper. \n\n\nOther comments:\n\n- Figure 1 seems to never have been referred to in the paper. Also, Figure 2 is referred after Figure 3. \n\n- When citing a paper indirectly (many times throughout the manuscript), use “\\citep” rather than “\\cite” to include parentheses outside both the names and the date. \n\n- Your definition of reward is indeed the expected reward (over next state). \n\n- You allow for $\\gamma = 1$ and at the same time you define return on an infinite horizon, which may become unbounded. Either remove the “no discount” scenario, or rather make your MDP episodic by introducing terminal states and necessitating termination. \n\n- Equation 1 should be the definition of advantage function (it is misplaced).\n\n- Your notation is not consistent. From what I see, you use CAP-BOLDFACE for sets, but then later you also use it for random variables X and Z. Similarly for the lower-case letters. Please consider a full overhaul of your notations. \n\n- Section \"Normalizing Flow\": X and Z are random variables, NOT distributions. \n\n- In section “Critic Regularized Regression” you used script-D for the dataset, but in the next part you used D. Again, I strongly recommend consistency in your notation to help your readers following the concepts better. \n\n- Figure 1 -> in the CNF graph, $z\\sim U(-1,1)^n$. Therefore, the color should be a solid blue instead of radial gradient.\n", "strength_and_weaknesses": "Strength:\n\n- The core idea of using NF is quite interesting and might be a good replacement for PLAS.\n\nWeakness:\n\n- The paper is too immature with numerous mistakes and quite poor presentation. In the current form it is far from acceptable. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject