Normalizing Flows are Capable Models for Continuous Control
Authors: Raj Ghugare
Benjamin Eysenbach
Venue: NeurIPS 2025 poster ()
Keywords: Reinforcement Learning
Normalizing Flows
Control as Inference
Offline RL
Imitation Learning
Unsupervised Goal Conditioned RL
Primary area: reinforcement_learning
TL;DR: We take a probabilistic view of decision making tasks, and show that normalizing flows offer many properties which make them highly performant models in imitation learning, offline, goal conditioned and unsupervised RL.

Abstract: Modern reinforcement learning (RL) algorithms have found success by using powerful probabilistic models, such as transformers, energy-based models, and diffusion/flow-based models. To this end, RL researchers often choose to pay the price of accommodating these models into their algorithms -- diffusion models are expressive, but are computationally intensive due to their reliance on solving differential equations, while autoregressive transformer models are scalable but typically require learning discrete representations. Normalizing flows (NFs), by contrast, seem to provide an appealing alternative, as they enable likelihoods and sampling without solving differential equations or autoregressive architectures. However, their potential in RL has received limited attention, partly due to the prevailing belief that normalizing flows lack sufficient expressivity. We show that this is not the case. Building on recent work in NFs, we propose a single NF architecture which integrates seamlessly into RL algorithms, serving as a policy, Q-function, and occupancy measure. Our approach leads to much simpler algorithms, and achieves higher performance in imitation learning, offline, goal conditioned RL and unsupervised RL.

Normalizing Flows are Capable Models for Continuous Control
Raj Ghugare
Benjamin Eysenbach
Department of Computer Science
Princeton University rg9360@princeton edu
Abstract
Modern reinforcement learning (RL) algorithms have found success by using prob- abilistic models, such transformers, energy-based models, and diffusion/flow- based models. To this end. researchers often choose to pay the price of accommo- dating these models into their algorithms diffusson models are expressive, but are computationally intensive due to their reliance solving differential equations, while autoregressive transformer models are scalable but typically require learning discrete representations. Normalizing flows (NFs), by contrast, seem to provide an appealing alternative, as they enable likelihoods and sampling without solving differential equations or autoregressive architectures. However, their potential in RL has received limited attention, partly due to the prevailing belief that normal- izing flows lack sufficient expressivity. We show that this is not the case. Building on recent work in NFs, we propose a single NF architecture which integrates seam- policy, Q-function, and occupancy measure. Our approach leads to much simpler algorithms, and achieves higher performance in imitation learning, offline, goal conditioned RL and unsupervised RL.
lessly into RL algorithms, serving
Exact Likelihoods
Sampling
Variational Inference
GCRL Value Func. Behavior Reg.
RL Policies
Offline RL Exact Max-Ent. RI. Control as Inference
Behavior Cloning
Unsupervised RL
World Models
(-
Ppost
P.
R(-)
Po(+)
What's the likelihood of a image?
Sample times from the fitted distribution.
What is the posterior if my likelihood is R(-)?
Probabilistic models and RL. Three capabilities of probabilistic models are important for RL: efficient likelihood computation, sampling, and compatibility with variational inference (VT). These capabilities are the core of most RL algorithms (See Section 4 for details). For example, behavior cloning requires maximizing likelihoods probabilities [23], and MaxEnt RL and RL finetuning requires variational inference [54]. However, models like diffusion and autoregressive transformers used in RL today support only subset of these capabilities, making their application complex. this paper, show how normalizing flows, which posses these properties, can be a generally capable model family for RL.
Figure 1:
sampling [4], GCRL value functions require estimating future discounted state
'The code for all experiments can be found here https://gthha..mmm//ccccotoo cewtrrccmeet-leeannnn
flows-4-
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Table 1: Capabilities of probabilistic models. NFs have capabilities that make them applicable a large number of RL problems (Section 4), including the ability compute exact likelihoods, sample from complex distributions, and be trained via both maximum likelihood and variational inference (Section 3). Despite these capabilities, research on NFs in RL has been limited compared to other state of the art models. Is this because NFs lack expressivity? Our experiments (Section 6) provide compelling evidence that this is not the case.
Model family
Exact Likelihoods
Sampling
Variational Inference"
Variational Autoencoders
Lower bound No
Yes
No
Generative Adversarial networks
Yes
No
Energy Based models
Unnormalized Yes, via ODEs Yes
MCMC
No
Diffusion models
Unrolling SDEs
No
Normal formalizing Flows
Yes
Yes
Although directly performing VI (posterior sampling and estimation) is generally intractible for the first four model families we have still marked them with yellow because additional algorithms for VI do exist [8].
1 Introduction
Reinforcement learning (RL) algorithms often require probabilistic modeling of high dimensional and complex distributions. For example, offline RI. or behavior cloning (BC) from large scale robotics datasets requires learning multi-modal policy [4, 16, 7] and value function estimation in goal-conditioned RI. (GCRL) requires estimating probabilities over future goals [23]. Different algorithms require different capabilities, including likelihood estimation, variational inference (VI) and efficient sampling (see Figure 1). A model trained via BC must maximize action likelihoods during training and support fast action sampling for real-time control. Most offline or finetuning RL algorithms require policy models to be able to backpropagate gradients of the value function through sampled actions. For MaxEnt RL with continuous actions, models should be optimizable using variational inference. However, obtaining these abilities in current models often comes at the cost of compute and complexity (Table 1). For example, diffusion models [78] are highly expressive yet computationally demanding: training, sampling, and evaluating log likelihoods all require solving differential equation. Autoregressive transformer models [86] are scalable and expressive, but typically limited to discrete settings; using them for RL in continuous spaces requires learning discrete representations of actions [76] or states [60]. Gaussian models offer computational efficiency but lack expressivity. Can we design expressive probabilistic models with all the capabilities in Fig. I?
Normalizing Flows (NFs) are among the most flexible probabilistic models: they (1) support both likelihood maximization and variational inference training, and (2) enable efficient sampling and exact likelihood computation (see Table 1). Yet, NFs have received far less attention from the RL community. perhaps due to the (mis)conceptiing that they have restricted architectures or training insta- bilities [65]. This paper demonstrates that neither of these notions hold NFs back in practice: NFs can match the expressivity of diffusion and autoregressive models while being significantly simpler. More- over, unlike other models, NFs do not require additional machinery to be integrated with existing RL frameworks, enabling simpler algorithms that are competitive with if not better than their counterparts.
Our primary contribution is to shed light on the effectiveness and NFs in RL. Building prior work [46, 20], we develop a simple NF architecture and evaluate it across 82 tasks spanning five different RL settings, including offline imitation learning, offline RL, conditinn RL, and unsupervised RL:
On offline imitation learning, NFs are competitive with baselines using state of the gener- ative models such as diffusion and transformers, while using significantly less compute (6x fewer parameters than diffusion policies).
On offline conditional imitation learning, NF's outperform BC baselines using models like flow matching, and offline RL algorithms which learn a value function, across 45 tasks.
On offline RL problems, NFs are competitive with strong baselines using diffusion flow matching models on most tasks, outperforming all of them on 15/30 tasks.
On unsupervised GCRL,, NFs can rival contrastive learning based algorithms, despite making fewer assumptions.
2

---

2 Related Work
Normalizing flows. This work builds on the rich literature of normalizing flow architec- tures [19, 74,4,, 46, 47, 40, 66, 91]. NICE [19] introduce the idea of stacking neural function blocks that perform shifting, making density estimation efficient. RealNVP [20] extends NICE to non volume preserving layers that also perform a scaling operation. GLOW [46] further propose applying an in- vertible generalized permutation and a special normalization layer (ActNorm) to improve expressivity.
Probabilistic models in RL. Strides in probabilistic modeling research consistently push the boundaries in deep RL. Variational autoencoders gave rise family of model-based [35, 36, 93] and representation learning [52, 90] algorithms. Contrastive learning has been used to learn value functions [59, 75, 25, 62] as well representations [51, 81] for RL. Advances in sequence modeling led to improvements in BC [14, 11, 53, 76], RL [63, 32] and planning [42]. More recently, diffusion models [80, 38, 78] and their variants [79, 57, 5, 6] have become the popular models for large scale BC pre-training [15, 56, 89, 7], RL finetuning [73, 37], offline RL 12, 68, 87] and planning [41, 92, 58].
NFs in RL. In RL, NFs have been used to learn policies from uncurated data [77] and model exact MaxEnt policies [12. 33, 12]. But despite their flexibility (see Fig. I and Table 1), they have not been go-to models in RL; research exploring the application of NFs in RL remains limited. We believe this is because historically, NFs were thought to have limited expressivity [91, 65]. In our experiments, we show that this is not the case (Section 6). We hope that our work brings greater attention towards the potential of NFs in RL, and the development of scalable probabilistic models with similar desirable properties (Fig. 1).
3 Preliminaries
This section describes two major approaches for training probabilistic models and how NFs can be trained using both of them. In the next section, we show that this flexibility allows NFs to be directly plugged into various RL algorithms, improving their expressivity and performance.
Maximum likelihood estimation (MLE). Let
12d be a random variable with an unknown density p(x). Given dataset of samples from this density D = { p(r), the goal of MLE is to estimate the parameters of probabilistic model (z( that maximize the likelihood of the dataset:
MMLE arg max EEND [o)])
If (1)) belongs to a universal family of distributions, then the MLE estimate po (x) converges towards the true density, as dataset size (N) increases [43].
Variational inference (VI). In many problems, instead of samples from the true density p(x:), one has access to its unnormalized part p(z) p(x) x Z. The goal of variational inference is to approximate p(x) using a distribution from some variational family of distributions @:
q" (I) arg min Dxt.qq |I p) arg max Exng(a) [log p(x) log q(x)],
9ES
where 7. can be ignored from the last equation. If Q is universal, the global optima of Eq. will be the true density g" (z) = p(x). To optimize Eq. using gradient descent, one needs to choose a family Q. from which it is easy generate samples (the expectation requires sampling from q), evaluate likelihoods (second term), and backpropagate through both these operations. Due to these constraints, one typically has to restrict Q to be simpler family [44] (e.g., exponential family), resort to other inference algorithms [31] that require additional complexity (more ppper--rrmeeters or models).
Normalizing flows (NFs). NFs learn an invertible mapping So Rd - Rdd, from true density p(x) to the a simple prior density (2))). This function is parameterized such that the inverse fo as well as the determinant of its jacobian matrix () (dfa())) |) easily calculable. The learned density Pe (z) can then be estimated using the change of variable formula:
po (x)) = (o)() dfo(x))
(3)
dx
3

---

To sample from an NE, one first samples from the prior, and return the inverse function:
~ po(2,, and return a - To (x)
Because Eq. (3) and Eq. (4) are passes through a single neural network, paopaprog gradients through them straightforward. Hence, the NF can be trained efficiently using both MLE (Eq. (1)) and VI (Eq. (2)) [83, 74]:
B arg max E an (o(Po((()))) log(1 dfo(x) [)] dr
(5)
0
0v1= arg max E po(p(=) [ogg(z>> - log)o(=) + log( dfe (x)
()].
(6)
dx
r5 (z)
Equation (5) follows from substituting the change of variable formula in Eq. (1). Equa- tion (6) is derived by the same substitution, followed by noting that fo( and
log( de (x) log( dfll)) I). is standard practice stack multiple such mappings together fo fg fg f Each mapping referred as a block in the entire NF.
4 Many RL Algorithms do MLE and VI.
In this section, we first introduce multiple RL problems settings. We will then show how popular representative algorithms in each setting are instantiations of MLE and VI. This will (1) underscore the importance of flexible models that can be directly optimized using MLE and VI, and (2) motivate our use of NF' as a generally capable family models for RL. We will then describe our architecture in the next section (Section 5). In all the following problems, we assume underlying controlled markov decision process (MDP, M), with states s S and actions a E A. The dynamics are p(s" [sslss,0 a), the initial state distribution is po()o) and H is the episode horizon.
4.1 Offline Imitation Learning (IL)
Given dataset of state-action trajectories D {56 0g. .g....,, & y aim in IL is to match the data generating behavior without any environment interactions. collected by expert(s), the
Behavior Cloning (BC) [4, 71] is an algorithm that models the conditional distribution of actions given states directly using MLE:
arg max E,, D log Ta (az ss)].
MLE
Goal Conditioned Behavior Cloning (GCBC) [22, 18, 30] is simple variant of BC which uses goal S as additional conditioning information. It is based the insight that trajectory leading to any goal, can be treated as an optimal trajectory to reach that goal. Given state-action pair, 8t, a.) is treated as the goal. Formally, GCBC maximizes the state sampled from the future Pt+ (8t+ following objective:
arg max E
log Ra (ar | St, St44)
(8)
MLE
Following prior work [67, 24], we set the future time ++ be truncated geometric distribution (from + 1 H).
4.2 Offline RI.
Offline RL assumes MDP M with reward function r(s,a). Given dataset of trajectories D = {sh. ab. ...,, s rh IN the goal is to find a policy sa s), which maximizes the expected discounted return Ex, 1 (" (rrs a:)] [55, 49}.

---

MaxEnt RL + BC [28, 54, 34]. We write out minimal objective for an offline actor-criti algorithm [28]:
- (oo(s)) -1 - 2700 ac+1)) )
21
arg min ES Ee Qo (+
(9)
arg max E De mol 22) Q0((), Alog To (a [t)] + a BeBeaDD log (al [8)] (10)
ML.
where Qa(s, a) is the action-value function and Qz(s, a) denotes the target network [61]. A and a control the strength of the entropy and behavior regularization respectively. For À 0. this algorithm exactly matches TD3+BC [28]. In all our experiments, we use À = 0 for minimality". The policy has to sample actions that have high values while staying close to the (potentially multi-modall behavior
policy. NFs ability to be trained using both MLE and allows the use of expressive policies, without any changes to the simplest offline RI. algorithms.
4.3 Goal conditioned RL (GCRL)
GCRL [45] is a multi-task RL problem where tasks correspond reaching goals states g S sampled from a test goal distribution (g). We define the discounted state occupancy distribution as P 8, a) (1 7) Er 7 (p(5 Se4 $0; 0gg 8, a).
Q-function estimation. Prior work has shown that the Q-function for GCRL is equal to the discounted state occupancy distribution Q*T(s. a, (se+ 8,a) [25]. Hence estimating the Q-function can be framed as a maximum likelihood problem:
arg max E
log po St, a:)].
(11)
81 Se Ge Peepoeelle [e,0e)
MLE
algorithms do not assume that the test time goal dis- tribution (gs)(g) is known a priori. The main aim UGS is to maximize exploration by commanding goals that are neither too easy nor too difficult for the current agent. This inclines the agent to explore goals on the edge of its goal coverage. But to recognize which goals are on the edge, one needs an estimate of the goal coverage density (g). Many methods [72, 70, 26] primarily fit density model and then choose goals that have low density. natural way of estimating densities is MLE, though some prior works use secondary objectives [64, 82, 26, 72] or non-parametric models [70]:
Unsupervised goal sampling (UGS) [26, 70]
arg max
Br ge4 pepepp (8e4 [8e.0e)
[log pa (9++)]
(12)
MLE
Using an estimate of this density, the canonical UGS algorithm is to sample large batch of goals from the replay buffer and command the agent to reach the goal with the minimum density under Po (o()))) [70].
5 A Simple and Efficient NF Architecture for RL
This section will introduce the NF architecture we use and Section 6 will show how the this ar- chitecture fits seamlessly into many RL algorithms (those from Section 4). Many design choices affect model expressivity, the speed of computing the forward (f)) and inverse (fo (f)'(z)), (x)). and the efficiency of computing the determinant of the Jacobian [65]. Our desiderata is to build simple and fast architecture that is expressive enough to compete and even surpass the most popular approaches across a range of RL problems. To achieve this, we build each NF block S by leveraging two com- ponents from prior work: (/) fully connected "coupling" network [20], and (2) linear flow [46]. In all of our experiments, the model learns conditional distribution pa (x g), so each NF block fg Rd x R ddd takes in y E IR* as additional input.
2This a special case variational inference. bayesian inference this corresponds to approximating a maximum a-posteriori estimate.
We will work with in infinite horizon case, with H cc.
5

---

Coupling network. We use an expressive and simple transformation proposed in Dinh et al. [19, 20]. Let Ty be the input for the tt block and y be the conditioning information. For the tt block, the coupling network f {ag s |R14/21 I2* + E 4/21} does the following computation:
(1) # 2 split(z'), (2) 2 (2 ag (+1 v)) x exp( (1(1.))) (3) i' concat (2) x22)
Note that the log-det det Jacobian of fg is -um (s/g (a/, y)). Further, given ä' and y, the inverse fb can be calculated as easily as the forward map (see Appendix B for details).
Linear flow. We adapt the generalized permutation proposed by Kingma and Dhariwal [46] to one dimensional inputs, to linearly transform the output of the coupling layer i². The intuition behind adding this linear flow that each dimension of je can be transformed by other dimensions in an end to end manner:
- (=")", where W PLaUUs and log [W] sum log((diag(U))).
Here, P is fixed permutation matrix and Lg lower triangular with ones on the diagonal and Ug is upper triangular. Hence, inverting them incurs similar costs as matrix multiplication [65] and calculating the log-det Jacobian requires only the diagonal elements of U.
While previous approaches used special variants of normalization layers ([20), [46]), we find that just using LayerNorm [3] in the coupling network's weights was sufficient of training deep NFs.
Normalization.
Overall, architecture can be thought of as RealNVP [20] with generalized permutation, or linear version of GLOW [46] without ActNorm. We have released a simple code for this architecture both in Jax and PyTorch.
6 Experiments
The main aim of our experiments is to use NFs and question some of the standard assumptions across several RL settings. (1) In II., methods
B
like diffusion policies [15] popular because they provide expressivity but at a high computa- tional cost. Our experiments with NF' show that this extra compute might not be necessary (Sec-
tion 6. (2) In conditional IL, prior work ar- gues that learning a value function is necessary
for strong performance [50, 13, 29, 1]. Our ex- periments suggest that sufficiently expressive policy, such as an NF, can often suffice (Section 6.2). (3) In offline RL, auxiliary techniques are typically used to support expressive policies. We show that NF policies can be directly integrated with existing offline RI. algorithms to achieve high perfor- mance (Section 6.3). (4) In UGS, we show that single NF can estimate both the value function and goal density, enabling better exploration compared to baseline that make stronger assumptions.
Figure 2: Visualization of tasks
In all experiments, we use the architecture described in Section 5 to implement all algorithms described in Section 4. Implementation details for all algorithms are provided in Appendix A. Across all problems (II, offline RL, online GCRL), we present results on a total of 82 tasks (see Fig. 2 for visualization of tasks) with 5 seeds each, which show that NFs are generally capable models that can be effectively applied to a diverse set of RL problems. Unless stated otherwise, all eιror bars represent + standard deviation across seeds. Finally, in Appendix D, we provide ablation experiments for various design choices and hyper-paramen values for some of the algorithms presented in this paper.
6.1 Imitation Learning (IL)
In this section, we use our NF architecture with BC (Section 4.1, Eq. (7)) for imitation learning (NF- BC). Please implementation details of training and sampling with IL algorithm in Appendix A. Tasks. PushT requires fine-grain control to maneuver T-block in a fixed location. Performance measured by the final normalized area of intersection between the block and the goal. Multimodal

---

BC
BCT
DifFolcy'C
DitiPolcy-T
Vg BeT
NF BC (ours)
24
2.4
2eforma
254
(a) Performance comparison. Performance is measured using normalized rewards.
Method
Time (sec) inference; training
Params (M)
Model types
Objectives
NF-GCBC
0.127 0.023
10 M
Fully connected NFs Transformer + VQ-VAES
1: end end
VQ-BET
0.273 0.021
4.23 M
2: req. pre-train
Diffusion Policy
1.28 0.019
65.78 M
Convolutional network
1; end end
(b) Conceptual comparison provide more holistic
RL practitioners.
Figure 3: Imitation Learning. Results multi-modal behavior datasets which require fine grained control expressivity. NF-BC faster in inference, easier to implement (end to end optimization of single loss function, contains 2-3 times less sypee--arameters and achieves competitive performance across all tasks (best success rate on two tasks and in top methods on the other two tasks). In Fig. 3a, we only report standard errors for method baseline results were taken from prior work which did not include error bars [53]
Ant, UR3 BlockPush, and Kitchen contain multiple behaviors (4, and respectively). Performance is measured by the normalized number of distinct behaviors the agent is able to imitate successfully over independent trials. We also report the simplicity of algorithms, as it often decisive factor in determining whether the algorithm can be practically employed. Simplicity is measured by the negating the number of hyperparameters used. We list and briefly describe the hyperparameters of all the methods used to make Fig. 3a in Appendix E. Baselines. We compare algorithms with 4 baselines. Diffusion policies [15] use high performing diffusion model architectures. We use both the convolutional and transformer (DiffPolicy-CC and DiftPolicy-T) variant of the diffusion policy. VQ-BeT [53] uses VQ-S [85] for discretizing actions and learns a transformer model chunked representations. BeT [76] is similar to VQ-BeT but uses K-means for action chunking.
Results. In Fig. 3a we see that NF-BC outperforms all baselines on PushT and Kitchen and performs competitively in other two tasks. Of all the baselines, NF-BC is fastest in inference and easier to implement (Fig. 3b). For example, DiffPolicy-T uses 3 times more hyperparameters than NF-BC. These results suggest that NF policies do not lack in expressivity, compared to both diffusion and auto-regressive policies despite introducing fewer additional hyper-parameterr In fact, NF-BC can achieve these results using just fully connected network and fraction of parameters compared to diffusion policies (see rightmost plot Fig. 3a). Unlike BeT [76] and VQ-BeT [53], NF-BC does not autoregressive models VQ- VAEs (see Fig. 3b for comparison of architectures). Reducing the number of hyperparameters is especially important for practitioners because it makes implementation casier and decreases the need for hyperparameter tuning.
6.2 Conditional Imitation Learning
In this section, we use our NF architecture with GCBC (Section 4.1, Eq. (8) for conditional imitation learning (NF-GCBC).
Tasks. We use total of 45 tasks from OGBench [67] meant test diverse capabilities like learning from suboptimal trajectories with high dimensional actions, trajectory stitching and long- horizon reasoning. Performance is measured using the average success rate reaching goal in 50 independent trials. Baselines. We compare with different baselines. GCBC is the GCBC 4.1 algorithm with a gaussian policy, and FM-GCBC is the GCBC 4.1 algorithm with a flow-matching based policy [57]. The velocity field for FM-GCBC is fully connected network trained using flow

---

matching, FM-GCBC is the closest baseline to NF-GCBC as it uses the same underlying algorithm, uses fully connected network like NF-GCBC, and uses an expressive family of policy models. GCIQL [48] is an offline RI algorithm that learns the optimal value function corresponding to the best actions the dataset. QRI. [88] uses a quasimetric architecture to learn the optimal value function and CRL [25] uses contrastive learning to estimate behavior policy's value function.
Improvement over GCBC
0.00
Aggregated perfomance
0:45
0.00
0.00
NF-0CBC FM-CCaC
Gcac
aCiOL
QRL
cii.
Offline GCRL benchmark (45 tasks)
Offline GCRL benchmark (45 tasks)
Figure 4: Conditional Imitation Learning, (Left) NE-GCBC exceeds its gaussian counterpart GCBC on 40/45 tasks, demonstrates that using expressive models like NFs can significantly improve performance. (Right) Despite being BC-based algorithm, NF-GCBC outperforms all offiine RL basclines across 45 tasks. NF-GCBC performs 77% better than FM-GCBC which the closest bascline that uses a flow matching policy with GCBC.
that NE-GCBC exceeds the performance (often significantly) of GCBC on 40/ 45 tasks. This underscores the importance of using an expressive policy with the GCBC algorithm. In Fig. 4 (right), NF-GCBC, despite being BC based algorithm, outperforms all baselines, including dedicated offline RL algorithms. We have added separate task wise plots in Appendix C. NF-GCBC outperforms FM-GCBC. which is the closest baseline, by 77%. While this does not necessarily imply that NFs are more expressive than flow- matching models, we hypothesize that all else being equal (e.g., fully connected architectures and identical underlying algorithms), NFs are better able to imitate the data using simple architectures as they directly optimize for log likelihoods. Flow matching models the other hand estimate probabilities indirectly (by learning SDE/ ODE), and might require more complex architectures (transformers or convolutional networks) work well.
Results. In Fig. (left),
6.3 Offline RL
NF-RLBC (ouns)
FOL
IFQL
IDOL
ReBRAC
BC
0.-0
MUUCCHN rARA
3:0216-3x3-glay
scene-glay
antmaze-lurge-rav..
cube-single-play
antsaccer-amnnaanaa
Figure 5: Offline RI. NE-RLEC outperforms all baselines 15/30 tasks, and achieves top 2 results 10/30 tasks. Notably, NF-RLBC achieves much higher results than baselines on puzzle-3x3-pla scene-play, which require the most long horizon and sequential reasoning. For example, on puzzle-3x3-play NF-RLBC performs 230% better than the next best baseline.
While NE-GCBC performs better than offline RL algorithms on average, there are certain tasks like puzzle-3x3-ple where only TD based methods succeed (Appendix C). This task, and many others as we will see, require learning a value function and accurately extracting behaviors from Fortunately, NFs can be trained using both VI and MLE (3). This allows to train MaxEnt RL+BC (Section 4.2, Eq. (9)) using our NF architecture as a policy (NF-RLBC).
Tasks. We use total of 30 tasks from prior work [68] testing offline RL with expressive policies. Each task has goal which requires completing K subtasks. At each step, the agent gets a reward equal to the negative of the number of subtasks left to be completed. For our experiments, we choose tasks that cover a diverse set of challenges and robots. Tasks like humanoidmaze-meddi navigate and codmmaze-medium-naae actions. Antsoccer-arena-naan
contain diverse suboptimal trajectories and high dimensional requires controlling quadruped agent to dribble ball goal location. Scene-plat play requires manipulating multiple objects and muzzle-3x3-ple requires

---

CRL-minmax
CRL-unffomm
271-6-17922
ant-ig-mazz aze
anshaard-azee
agoregated scoris
DUCCOS rale
unvionnmntt Steps
environment steps
ervironment steps
Figure 6: GCRL. (first 3 from left) NF-UGS consistently outperforms other unsupervised baselines. (rightmost) NF-UGS achieves higher asymptotic success rates than the CRL-oracle averaged across all tasks. Note that CRI.-oracle makes extra assumptions about the availability of test-time goals.
combinatorial generalization and long-horizen reasoning. We hypothesize that the Q function for the some of these tasks (esp manipulation tasks like puzzle and scene which require long horizon reasoning) can have narrow modes of good actions and a large number of bad actions. Hence using direct gradient based optimization can be crucial search for good actions. Baselines. We choose 5 representative baselines. BC performs BC with a gaussian policy, ReBRAC [84] is an offline RL algorithm with gaussian policy that achieves impressive results on prior benchmarks
[27]. IDQL [37] uses importance resampling to improve over diffusion policy. IFQL is the flow counterpart of IDQL [68]. FQL [68] strong offline RL algorithm that distills a flow-matching policy into a one-step policy, while jointly maximizing the Q-function. Notably, all offline RL algorithms (IDQL, IFQL, FQL) have incorporate additional components such distillation or resampling to leverage policy expressiveness.
Figure 5, shows that on average NE-RLBC outperforms all baselines 15/30 tasks, and achieves top 2 results on 10/30 tasks. Notably, NE-RLBC achieves much higher results than all baselines on puzzle-3x3-pla and scene-play, which require most long horizon and sequential reasoning (on puzzle-3x3-pp
Results.
NF-RLBC performs 230% better than the next best baseline). NF- RLBC is able to search for good actions because it directly uses gradient based optimization to backpropagate action gradients through the policy. All other offline RL algorithms with expressive policy classes (IDQL, IFQL, FQL) have to use extra intermediate steps.
6.4 Goal Conditioned RL
In this section, we use our NF architecture to estimate both the Q-function (Eq. (11) and the coverage density (Eq. (12)). Both quantities model the distribution over the same random variable (future goals), with the Q-function conditioned on a state-action pair. We use single neural network to estimate both by employing masking scheme indicate the presence or absence of conditioning [39]. We then train a goal-conditioned policy to maximize the Q-function in a standard actor-critic setup, and sample training goals using the canonical UGS algorithm (Section 4.3), resulting in our NF-UGS method.
Tasks. We use three problems from the UGS literature [70, ant-u-maze, ant-big-maze and ant-hard-maze. These environments are difficult because the goal test distribution is unknown, and success requires thorough exploration of the entire maze. Baselines. We compare against three baselines. CRL-oracle [25, 9] uses contrastive learning to estimate the Q-function and serves as an oracle bascline as it uses test goals Piest (g) to collect data during training. CRL-uniform is a variant that samples goals uniformly from the replay buffer, while CRL-minmax selects goals with the lowest Q-value from the initial state of the MDP intuitively targeting the goals the agent believes are hardest to reach. The aim of these comparisons isn't I to propose a state of the art method, but rather to check whether NFs can efficiently estimate the coverage densities of the ota-stationary RL data for unsupervised exploration. Results. These experiments aim to evaluate whether NFs can efficiently estimate both the Q-function and the coverage density. As shown in Fig. 6, NF-UGS outperforms all baselines, and notably the oracle algorithm as well (aggregated score is the average of the final success rates across all three tasks). Unlike the minmax strategy, which plateaus over time, NE-UGS continues to improve by sampling goals to maximize the entropy of the coverage distribution [70].

---

7 Conclusion
In this paper, we view RL problems from a probabilistic perspective, asking: which probabilistic models provide the capabilities necessary for efficient RL? We argue that NFs offer a compelling answer. NFs outperform strong baselines across diverse RL settings, often with fewer parameters and lower compute, leading to simpler algorithms. By making density estimation simpler and effective, we hope our work opens new avenues for areas such as distributional RL, bayesian RL, and RL safety,
A primary limitation of NFs is that they have restrictive architectures, in that they have to be invertible. Nevertheless, there is rich line of work characterizing which NF architectures are universal [21, 17, 65]. Another limitation of our work is that we do not explore new NF architectures; the design we use is mainly adapted from prior work (Section 5).
Acknowledgments and Disclosure of Funding
The authors are pleased to acknowledge that the work reported on in this paper substantially performed using the Princeton Research Computing resources at Princeton University which is consortium of groups led by the Princeton Institute for Computational Science and Engineering (PIC- SciE) and Office of Information Technology's 's Research Computing. The authors thank Emmanuel Bengio, Liv d' Aliberti and Eva Yi Xie for their detailed reviews thank Seohong Park for providing code for some baselines and Catherine for helping find a bug in code. The authors also thank Shuangfei Zhai and the members of Princeton RL lab for helpful discussions.
initial drafts. The authors
References
[1] Agarwal, R., Schuurmans, D., and Norouzi, M. (2020). An optimistic perspective offline reinforcement learning. In III, H. D. and Singh, A., editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 Proceedings of Machine Learning Research, pages 104 114. PMLR.
[2] Ajay, A., Du, Y., Gupta, A., Tenenbaum, J., Jaakkola, T., and Agrawal, P. (2023). Is conditional generative modeling all you need for decision-making
[3] Ba, J. L., Kiros, R., and Hinton, G. (2016). Layer normalization.
[4] Bain, M. and Sammut, C. (1995). framework for behavioural cloning. Machine Intelligence 15.
[S] Bengio, E., Jain, M., Korablyov, M., Precup, D., and Bengio, Y. (2021). Flow network based generative models for non iterative diverse candidate generation.
[6] Bengio, Y., Lahlou, S., Deleu, T., Hu, E. J., Tiwari, M., and Bengio, E. (2023). Gflownet foundations.
[7] Black, K., Brown, N., Driess, D., Esmail, A., Equi, M., Finn, C., Fusai, N., Groom, L., Hausman, K., Ichter, B., Jakubezak, S., Jones, T., Ke, L., Levine, S., Li-Bell, A.. Mothukuri, M., Nair. S., Pertsch, K.,, Shi, L. X.. Tanner, J., Vuong, Q.. Walling, A., Wang, H., and Zhilinsky. U. (2024). Not A vision- laiion-lnggagge-cct flow model for general robot control.
[8] Blei, D. M., Kucukelbir, A., and and, J. D. M. (2017). Variational inference: A review for statisticians. Journal of the American Statistical Association, 1 11((518)8899-877
[9] Bortkiewicz, M., Palucki, W. Myers, V., Dziarmaga. T., Arczewski, T., Kuciáski, L., and Eysenbach, B. (2025). Accelerating goal-conditioned reinforcement learning algorithms and research. In The Thirteenth International Conference Learning Representations.
[10] Bradbury, J., Frostig, R., Hawkins, Johnson, M. J., Leary, C., Maclaurin, D., Necula, G., Paszke, A., VanderPlas, J.. Wanderman Minne,, S., and Zhang, Q. (2018). JAX: composable transformations of Pythou+NumPy programs. GitHub repository.
[11] Brohan, A., Brown, N., Carbajal, J., Chebotar, Y., Dabis, J., Finn, C., Gopalakrishnan, K., Hausman, K., Herzog, A., Hsu, J., et al. (2023). Rt-1: Robotics transformer for real-world control at scale.
[12] Chao, C.HH., Feng, C., Sun, W.-F. Lee, C.-K., See, S., and Lee, C.-Y. (2024). Maximum entropy reinforcement learning via ensrdy-based based normalizing flow. In The Thirt--eihtth Annual Conference on Neural Information Processing Systems.
10

---

[13] Chebotar, Y., Vuong, Q. Hausman, K., Xia, F., Lu, Y., Irpan, A., Kumar, A., Yu. T., Herzog, A., Pertsch, K.. Gopalakrishnan, K., Ibarz, J., Nachum, O., Sontakke, S. A., Salazar, G., Tran, H1. T., Peralta, J., Tan, C.. Manjunath, D., Singh, J., Zitkovich, B., Jackson, T., Rao, K., Finn, C., and Levine, (2023). Q-transformer: Scalable offiine reinforeement learning via autoregressive q-functions. In Tan, J. Toussaint, M., and Darvish, K., editors, Proceedings of The 7th Conference on Robot Learning, volume 229 of Proceedings of Machine PMLR.
Learning Research, pages
[14] Chen, L., Lu, K., Rajeswaran, A., Lee, K., Grover, A., Laskin, M., Abbeel, P., Srinivas, A., and Mordatch, I. (2021). Decision transformer: Reinforcement learning via sequence modeling. In Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W. editors, Advances Neural Information Processing Systems.
[15] Chi, C., Xu, Z. Feng, S., Cousineau, E., Du, Y., Burchficl, B., Tedrake, R., and Song, S. (2024), Diffusion policy: Visuomotor policy learning via action diffusion.
[16] Collaboration, E., O'Neill, A., Rehman, A., Cupta, A., Maddukuri, A.. Gupta, A., Padalkar, A., Lee, A.. Pooley, A., Gupta, A., Mandlekar, A., et al. (2024). Open x-embodiment: Robotic learning datasets and rt-x models.
[17] Cornish, R., Caterini, A., Deligiannidis, G., and Doucet, A. (2020). Relaxing bijectivity constraints with continuously indexed normalising flows. In III, H. D. and Singh, A., editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 2133 2143. PMLR.
[18] Ding, Y.. Florensa, C., Abbeel. P., and Phielipp. M. (2019). Goal-conditioned imitation learning. In Wallach, H., Larochelle, H., Beygelzimer, A., d' d'Alché-Buee F., Fox, E., and Garnett, R., editors, Advances Neural Information Processing Systems, volume 32. Curran Associates, Inc.
[19] Dinh, L.,, Krueger, D., and Bengio, Y. (2014), Nice: Non-linear independent components estimation. arXiv: Learning.
[20] Dinh, L., Sohl-Dickstein. J. and Bengio, S. (2017). Density estimation using real NVP. In International Conference Learning Representations.
[21] Draxler, F., Wahl, S., Schnörr, C., and Köthe, U. (2024). On the Universality of Volume-Preserving and Coupling- Normalizing Flows. arXiv preprint arXiv:2402.06577
[22] Emmons, S., Eysenbach, B., Kostrikov, L,, and Levine, S. (2022). Rvs: What is essential for offline RL. via supervised learning? In International Conference on Learning Representations.
[23] Eysenbach, B., Salakhutdinov, R., and Levine, S. (2021). C-lcaming: Learning to achieve goals via recursive classification. In International Conference on Learning Representations
[24] Eysenbach, B., Udatha, S., Salakhutdinov, R. R., and Levine, S. (2022a). Imitating past successes can he very suboptimal. Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A., editors, Advances in Neural Information Processing Systems, volume 35, pages 6047-6059. Curran Associates, Inc.
[25] Eysenbach, B., Zhang, T., Levine, S., and Salakhutdinov, R. (2022b). Contrastive learning as goal- conditioned reinforeement lcarning. In Oh, A. H., Agarwal, A., Belgrave, D., and Cho, K., editors, Advances in Neural Information P Processing Systems.
[26] Florensa, C., Held, D., Geng, X., and Abheel, P. (2018). Automatic goal generation for reinforcement learning agents. In Dy, J. and Krause, A., editors, Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 1515-1528, PMI..
[27] Fu, J., Kumar, A., Nachum, O., Tucker, G., and Levine, S. (2021). D4rl: Datasets for deep data-driven reinforcement learning.
[28] Fujimoto, S. and Gu, S. S. (2021). A minimalist approach to offfine reinforcement learning. In Ranzato, M., Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W., editors, Advonces in Neural Information Processing Systems, volume 34, pages 20132-20145. Curran Associates, Inc.
[29] Fujimoto, S., Meger, D., and Precup, D. (2019). Off-policy deep reinforcement learning without exploration. In Chaudhuri, K. and Salakhutdinov, R., editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages
PMLR.
[30] Ghosh, D.,. Gupta, A., Reddy, A., Fu, J. Devin, C., Eysenbach, B., and Levine, S. (2019). Learning to reach goals via iterated supervised learning. arXiy preprint arXiv:1912.06088
11

---

[31] Graikos, A.. Malkin, N., Jojic, N., and Samaras, D. (2022). Diffusion models plug-and- ande-ald priors. In Oh, A. H. Agarwal, A., Belgrave, D., and Cho, K., editors, Advances in Neural Information Processing Systems.
[32] Grigsby, J., Fan, L., and Zhu, Y. (2024). AMAGO: Scalable in-context reinforcement learning for adaptive agents. The Twelfth International Conference on Learning Representations.
[33] Haarnoja, T., Hartikainen, K., Abbeel, P., and Levine, S. (2018a). Latent space policies for hierarchical reinforcement learning. In Dy, J. and Krause, A., editors, Proceedingy of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 1851-1860. PMLR.
[34] Haarnoja. T., Zhou, A., Abbeel, P., and Levine, S. (2018b). Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In Dy, J. and Krause, A., editors, Proceedings of the 35th International Conference Machine Learming, volume 80 of Proceedings of Machine Learning Research, pages 18700 PMLR.
[35] Hafner, D., Lillicrap. T., Fischer, I., Villegas, R., Ha, D., Lee, H., and Davidson, J. (2019). Learning latent dynamics for planning from pixels. In Chaudhuri, K. and Salakhutdinov, R., editors, Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 2555 2565. PMLR.
[36] Hafner, D., Pasukonis, J., Ba, and Lillicrap, T. (2024). Mastering diverse domains through world models.
[37] Hansen-Estruch, P., Kostrikov, I, Janner, M., Kuba, J. G., and Levine, S. (2023). Idql: Implicit q-learning as actor-critic methad with diffusion policies.
[38] Ho, J., Jain, A., and Abbeel, P. (2020). Denoising diffusion probabilistic models.
[39] Ho, J. and Salimans, T. (2022). Classifier-fre diffusion guidance.
[40] Huang, C.-W., Krueger, D., Lacoste, A., and Courville, A. (2018). Neural autoregressive flows. In Dy. J. and Krause, A., editors. Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Muchine Learning Research, pages 2078-2087, PMLR.
[41] Janner, M., Du, Y., Tenenbaum, J., and Levine. S. (2022). Planning with diffusion for flexible behavior synthesis. In Chaudhuri, K., Jegelka, S., Song. L., Szepesvari, C., Niu, G., and Sabato. S., editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 9902-9915. PMLR.
[42] Janner, M., Li, Q., and Levine, S. (2021). Offiin reinforcement learning one big sequence modeling problem. In Ranzato, M., Beygelzimer, A., Dauphin, Y., Liang. P., and Vaughan. J. W., editors, Advances Neural Information Processing Systems, volume 34, pages 1273 1286. Curran Associates, Inc.
[43] Jordan, M. I. and Bishop, C. M. (2001). An introduction to graphical models. Tutorial material.
[44] Jordan, M. I., Ghahramani, Z., Jaakkola, S., and Saul, L. K. (1999), An introduction to variational methods for graphical models, page 105-161. MIT Press, Cambridge, MA, USA.
[45] Kaelbling. L. P. (1993). Learning to achieve goals. International Joint Conference on Artificial Intelligence.
[46] Kingma, D. P. and Dhariwal, P. (2018). Glow: Generative flow with invertible Ix1 convolutions. In Bengio, N., and Gamett, R., editors, Advances in Neural
S., Wallach, H., Larochelle, H., Grauman, K., Information Processing Systems, volume 31. Curran Associates, Inc.
[47] Kingma, D. P., Salimans, T., Jozefowicz, R., Chen, X., Sutskever, I., and Welling, M. (2016). Improved variational inference with inverse autoregressive flow. In Lee. D., Sugiyama. M., Luxburg, U.. Guyon, I, and Garnett, R., editors, Advances in Neural Information Processing Systems, volume 29, Curran Associates, Inc.
[48] Kostrikov, I., Nair, A., and Levine, S. (2021). Offline reinforcement learning with implicit g-learning.
[49] Kumar, A., Fu, Soh, M., Tucker, G., and Levine, (2019). Stabilizing off-policy q-lcarning via bootstrapping error reduction. Wallach, H. Larochelle, H., Beygelzimer, A., d'Alché-Buu Bue, F., Fox, E., and Garnett, R., editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc.
[50] Kumar, A., Zhou, A., Tucker, G., and Levine, S. (2020). Conservative leaamnngg for offline reinforcement learning. In Larochelle, H., Ranzato, M.. Hadsell, R., Balcan, M.. and Lin, H., editors, Advances in Neural Information Processing Systems, volume 33, pages 1179 1191. Curran Associates, Inc
12

---

[SI] Laskin, M., Srinivas, A., and Abbeel. P. (2020). CURL: Contrastive unsupervised representations for reinforcement learning. In III, H. D. and Singh, A., editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 Proceedings of Muchine Learning Research, pages 5639-5650. PMLR.
[52] Lee, A. X., Nagahandi, A., Abbeel, P., and Levine, S. (2020). Stochastic latent actor-critic: Deep reinforcement learning with a latent variable model. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H., editors, Advances in Neural Information Processing Systems, volume 33, pages 741-752. Curran Associates, Inc.
[53] Lee, S., Wang, Y., Etukuru, H., Kim, H. J.. Shafiullah, N. M. M., and Pinto. L. (2024). Behavior generation with latent actions. In Salakhutdinov, R., Kolter, Z., Heller, K., Weller, A., Oliver, N., Scarleu, J., and Berkenkamp, F., editors, Proceedings of the 41 International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pages 26991 177008. PMLR.
[54] Levine, S. (2018). Reinforcement learning and control as probabilistic inference: Tutorial and review.
[55] Levine, S., Kumar, A., Tucker, G., and Fu, J. (2020). Offline reinforcement learning: Tutorial, review, and perspectives on open problems.
[56] Lin, F., Hu, Y., Sheng, Wen, C., You, and Gao, Y. (2025). Data scaling laws in imitation learning for robotic manipulation.
[57] Lipman, Y., Chen, R. T. Q. Ben-Hamu, H., Nickel, M., and Le, M. (2023). Flow matching for generative modeling.
[58] Lu, C., Ball, J., Teh, Y.W. and Parker-Holder, J. (2023). Synthetic experience replay. In Thirty-seventh Conference Neural Information Processing Systems.
[59] Ma, Y. J., Sodhani, S., Jayaraman, D., Bastani, O., Kumar, V., and Zhang, A. (2023). VIP: Towards universal visual reward and representation via value-implicit pre-training. In The Eleventh International Conference on Learning Representations.
[60] Micheli, V., Alonso, E.,, and Fleuret, F. (2023). Transformers are sample-efficient world models. In The Eleventh International Conference on Learning Representations.
[61] Mnih, V., Kavukcuoglu, K.. Silver, D., Graves, A., Antonogloo. L. Wierstra, D., and Riedmiller, M. (2013). Playing atan with deep reinforcement learning. arXiv preprint 3/2.5623..
[62] Myers, V., Zheng, C., Dragan, A. D., Levine, S., and Eysenbach, (2024). Learning temporal distances: Contrastive successor features can provide a metric structure for decision becision-making. In ICML.
[63] Ni, T., Eysenbach, B., and Salakhutdinov, R. (2022). Recurrent model-free RL can be a strong baseline for many POMDPs. In Chaudhuri, K. Jegelka, S., Song, L., Szepesvari, C., Niu, G. and Sabato. S., editors. Proceedings of the 39th International Conference Machine Learning. volume 162 of Proceedings of Machine Learning Research, pages 16691-16723, PMLR.
[64] OpenAL, O., Plappert, M., Sampedro, R., Xu, T., Akkaya, I., Kosaraju, V., Welinder, P., D'Sa, R., Petron, A., d. O. Pinto, H. P., Paino, A., Noh, H., Weng, L., Yuan, Q., Chu, C., and Zaremba, W. (2021). Asymmetric pelf-play for automatic goal discovery in robotic manipulation.
[65] Papamakarios, G., Nalisnick, F.,, Rezende, D. J., Mohamed, S., and Lakshninarayanan, B. (2021). Normalizing flows for probabilistic modeling and inference. Journal Machine Learning Research, 22(57):1-64.
[66] Papamakarios, G., Pavlakou, T., and Murray, (2017). Masked autoregressive flow for density estimation. In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Gamett, R., editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc.
[67] Park, S., Frans, K., Eysenhach, B., and Levine, S. (2025a). OGBench: Benchmarking offiie goal- conditioned RL. In The Thirteenth International Conference on Learning Representations.
[68] Park, S., Li, Q., and Levine, S. (2025b). Flow q-learning. ArXiv.
[69] Paszke, A., Gross, S., Massa, F, Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., Desmaison, A., Köpf, A., Yang, E., Ve Z., Raison, M., Tejani, A., Chilamkurthy, S., Steiner, B., Fang, L., Bai, J., and Chintala, S. (2019). Pytorch: An imperative style, perforrrforence deep learning library.
13

---

[70] Pitis, S., Chan, H., Zhao, S., Stadie, B., and Ba, J. (2020). Maximum entropy gain exploration for long horizon multi-goal reinforcement learning. In III, H. D. and Singh. A., editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 Proceedings of Machine Learning Research, pages 7750-7761, PMLR.
[71] Pomerleau, D. A. (1988). Alvinn: An autonomous land vehicle in a neural network. In Touretzky, D., editor, Advances in Neural Information Processing Systems, volume 1. Kragan-aafimann.
[72] Pong, V. H., Dalal, M., Lin, S., Nair, A., Bahl, S., and Levine, S. (2020). Skew-fit: Statc-covering self-supervised reinforcement learning.
[73] Ren, A. Z., Lidard. J., Ankile, L. L.. Simeonov, A., Agrawal, P., Majumdar, A., Burchfiel, B., Dai, H., and Simchowitz, M. (2024). Diffusion policy policy optimization.
[74] Rezende. D. and Mohamed, S. (2015). Variational inference with normalizing flows. In Bach, F, and Blei, D., editors, Proceedings of the 32nd International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pages 1530-1533 Lille, France. PMLR.
[75] Sermanet, P., Lynch, C., Chebotar, Y., Hsu, J., Jang, E., Schaal, S., Levine, S., and Brain, G. (2018). Time- contrasive networks: Supe-super learning from video. 2018 IEEE International Conference on Robotics and Automation (ICRA), pages
[76] Shafiullah, N. M. M., Cui, Z. J., Altanzaya, A., and Pinto, L. (2022). Behavior transformers: Cloning Sk$ modes with one stone. In Oh, A. H., Agarwal, A., Belgrave, D., and Cho, K., editors, Advances in Neurul Information Processing Systems.
[77] Singh, A., Liu, H., Zhou, G., Yu, A., Rhinehart, N., and Levine, S. (2021). Parrot: Data-driven behavioral priors for reinforcement learning. In International Conference on Learning Representations.
[78] Sohl-Dickstein, J., Weiss, F. A., Maheswaranathan, N., and Ganguli, S. (2015). Deep unsupervised learning using nonequilibrium thermodynamies.
[79] Song, Y., Dhariwal, P., Chen, M., and Sutskever, I. (2023), Consistency modcls.
[80] Song, Y.. Sohl-Dickstein, J., Kingma. D. P.. Kumar, A., Ermon, S., and Poole, B. (2021). Score-based generative modeling through stochastic differential equations.
[81] Stooke, A., Lee, K., Abbeel, P., and Laskin, M. (2021). Decoupling representation leamning from reinforce- ment learning. Meila, M. and Zhang. T., editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 9870-9879 PMLR.
[82] Sukhhaatar, S., Lin, Z., Kostrikov, I,, Synnaeve, G., Szlam, A., and Fergus, R. (2018). Intrinsic motivation and automatic curricula via asymmetric self-play.
[83] Tabak, E. and Fijnden E. (2010). Density estimation by dual ascent of the log-likelihood. Communications in Mathematical Sciences, 8(1):217-233.
[84] Tarasov, D.,, Kurenkov, V., Nikulin, A., and Kolesnikov, S. (2023). Revisiting the minimalist approach to offline reinforcement learning.
[85] van den Oord, A., Vinyals, O., and kavukcuoglu, k. (2017). Neural discrete representation learning. In Guyon, I., Luxburg, U. V., Bengio, S., Wallach, H., Fergus, R., Vishwanathan, S., and Garnett, R., editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc.
[86] Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L. u., and Polosukhin. (2017). Attention is all you need. In Guyon, L., Luxburg, U. V., Bengio, S.. Wallach, H., Fergus, R., Vishwanathan, S., and Gamett, R., editors, Advances in Neural Information Processing Systems, volume 30. Curran Associates, Inc.
[87] Venkatraman, S., Khaitan. S., Akella, R. T., Dolan, J., Schneider. J., and Berseth, G. (2024). Reasoning with latent diffusion in offline reinforcement learning. In The Twelfth International Conference Learning Representations.
[88] Wang, T., Torralba, A., Isola, P., and Zhang, A. (2023). Optimal goal- reaching reinforcement learning via quasimetric learning. In Krause, A.. Brunskill, E., Cho, K., Engelhardt, B., Sabato. S., and Searlett, J., editors, Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 14413-36430. PMLR.
14

---

[89] Weng, T., Bajracharya, S. M.. Wang. Y., Agrawal, K., and Held, D. (2022). Fabricflownet: Bimanual cloth manipulation with flow-based policy. In Faust, A., Hsu, D., and Neumann, G., editors. Proceedings of the Conference on Robot Learning, volume of Proceedings of Machine Learning Research, pages 192-202. PMLR.
[90] Yarats, D., Zhang, A., Kostrikov, I,, Amos, B., Pineau, and Fergus, R. (2021). Improving sample efficiency in model-free reinforcement learning from images. Proceedings of the AAAI Conference on Artificial Intelligence, 35(12):10674110681
[91] Zhai, S., Zhang. R., Nakkiran, P., Berthelot, D., Gu, J., Zheng, H., Chen, T., Bautista, M. A., Jaitly, N., and Susskind, J. (2024). Normalizing flows are capable generative models. arXiv preprint arXiv:2412.06322
[92] Zheng, Q., Le, M.. Shaul, N., Lipman, Y., Grover, A.. and Chen, R. T. (2023). Guided flows for generative modeling and decision making. arXiv preprint arXiv:2311.1344.
[93] Lukasz Kaiser, Babaeizadeh, M., Milos, P., Osiński, B., Campbell, R. H., Czechowski, K., Erhan, D., Finn,
C., Kozakowski, P., Levine, S., Mohiuddin, A., Sepassi. R., Tucker, G., and Michalewski, H. (2020). Model based reinforcement learning for atari. In International Conference on Learning Representations.
15

---

A Implementation details.
In this section we provide all the implementation details as well as hyper-parameter used for all the algorithms in our experiments NF-BC, NF-GCBC, NF-RLBC, and NF-UGS. Unless specified otherwise, we chose common hyper-parameters used in prior work without tuning.
A.1 Imitation Learning: NF-BC
For all imitation learning, we use a policy with 12 NF blocks. The fully connected networks of the coupling network for each NF block consist of two hidden layers with 512 activations each. We use fully connected network for the state-goal encoder with 4 hidden layers having 512 activations each.
It is common in the NF literature to add some noise to the training inputs before training an NF using MLE [74, 91]. For NF-BC, we add gaussian noise the actions before training. Hence, NF-BC models the noisy behavior policy instead n nenoise (a s). To reduce noise from the obtained samples, Zhai et al. [91] propose a denoising trick:
2 ~ po: y 1
pog)o())
We use this trick all our NF-BC experiments. Table contains the values of hyperparameters used by NF-BC. All experiments were done on RTX 3090s, A5000s or A6000s and did not require more than 20 hours. Experiment speed is mostly bottlenecked by intermediate evaluation of BC policies because the environments are slow and the number of rollouts per evaluation is 50. For example, the experiments on the multimodal-am environment took 20 hours, but on pusht they were finished in less than 90 minutes.
Table 2: NF-BC: Hyper-parameters and their values.
Hyper-parametis
Description
channels noise std
512
0.1
blocks
12
encoder num layers rep dims
512
A.2 Conditional Imitation Learning: NE-GCBC
For all conditional imitation learning experiments, we use policy with 6 NF blocks. The fully connected networks of the coupling network for each NF block consist of two hidden layers with 256 activations each. We use a fully connected network for the state-goal encoder with hidden layers having 512 activations each. The output representation dimensions
64. We also employ the demoising trick proposed by [91] and mentioned Appendix A.1. Table 3 contains the values of hyperparamters used by NF-GCBC All experiments done on RTX 3090s, A5000s A6000s and did not require more than 3 hours each.
Table 3: NF-GCBC: Hyper-parametors and their values.
Hyper-parameter
Description
channels
512
noise std
0.1
blocks
encoder num layers
4
rep dims
512
Future goal sampling discount
Table
16

---

Table NF-GCBC:
Environment values for future goal sampling discount values. For all GCBC algorithmss including the GCBC baselines, we tuned this value in the set (0.97. 0.99)
Environment Name
Future goals sampling discount value 7
puzzle 3 play
0.97
scene play
0.97
cube single play
0.97
cube double play
0.99
humanoidmaze medium navigation
0,97
antmaze large navigation
0.99
antsoccer arena navigation
0,97
antmaze medium stitch
0.97
antsoccer arena stitch
0.97
A.3
Offline RL: NF-RLBC
For all offline RI. experiments, we use a policy with 6 NF blocks. The fully connected networks the coupling network for each NF block consist of hidden layers with 256 activations each. We use fully connected network for the state-goal encoder with 4 hidden layers having 512 activations each. The output representation dimensions are 64. Table 5 contains the values of main hyperparameters used by NF-RLBC. All experiments were done on RTX 3090s, A5000s or A6000s and did not require more than 3 hours each.
Table 5: NF-RLBC: tlyer-parameee and their values.
Hyper-parameter
Description
channels
512
noise std
0.1
actor entropy coefficient (A in Eq. (9))
0.0
actor BC loss coefficient (a in Eq. (9))
Table 6
blocks
6
encoder num layers rep dims
4
512
Table 6: NF-RLBC: Environment wise actor BC loss coefficient. We the same hyperparameter tuning procedure as the baselines [68]. See App E.2, of Park et al. [68]
Environment Name
Actor BC loss coefficient
puzzle 3 3 play
10.0
scene play antmaze large navigation
10.0
1.0
cube single play
10.0
humanoidmaze medium navigation antsoccer arena navigation
1.0
1.0
A.4 Unsupervised RL: NF-UGS
For all unsupervised RL experiments, we use goal conditioned value function with 6 NF blocks. The fully connected networks of the coupling network for each NF block consist of two hidden layers with 256 activations each. We use fully connected network for the state-action encoder with 4 hidden layers having 1024 activations each. We train both the conditional Pe(g s, s,a) and the unconditional pe (g) using same network. We apply a mask on conditional variable inputs with probability 0.1. The discount factor we use is 0.99. We add gaussian noise with standard deviation 0.05. Because we do not sample from the model in NE-UGS, we do need to employ the denoising trick. Instead, 1024 random goals from the buffer are sampled and the one with the lowest
17

---

marginal density Pe (g) is selected for data collection. All experiments were done on RTX 3090s, A5000s A6000s and did not require more than 20 hours each.
B More details about NF block.
1+1
++1
2/2
(+ W, where W PLU
D2
22
MLP
Layerfoom
y
C2
Figure NF block.
It requires the same cost the coupling of the forward block to calculate the coupling of the inverse block f3
$ - = split(x'),
o = at.
x 2y x exp(>b(af v))
x' concat(af xb))
See Fig. for visualization of each NF block.
Table 7: Imitation Learning: Performance with varying numbers of NF blocks on the Push-T task.
number of NE blocks
Push-T
0.77
14
0.79
24
0.827
28
0.859
32
0.85
40
0.844
48
0.87
0.85
Table 8: Imitation Learning: Performance for different NF coupling layer widths on the Push-T
task.
width of NF's coupling layers
Push-T
128
0.79
256
0.82
1024
0.823
18

---

Table 9: Imitation Learning: Performance with varying noise levels for the denoising trick (Ap-
pendix A.I, [91]) on the Push-T task.
noise_st _td
Push-T
0.5
0.77
0.3
0.79
0.1
0.827
0.05
0.84
0.03
0.873
0.01
0.83
0.0
NaN
Table 10: Conditional Imitation Learning: Ablation study for the importance various components of NF-GCBC. These experiments show that the architectural modifications propose are crucial for performance NFs in the context of RL. We hope that the above experiments serve as sufficient motivation for design choices.
Method
antmaze-large-navigat
scene-play
cube-single-play
NF-GCBC
0.32
0.42
0.795
without permutation flows
0.26
0.04
0.0
without LayerNorm
0.29
0.3
0.0
without sampling trick
0.3
0.42
0.69
C
Additional results.
Table 11 and Table 12 contain additional experimental results.
Ablation experiments.
In Table 8, Table 7, Table 9 ablate the values of the width of the coupling networks, the number of NF blocs, and the levels of noise for the denoising trick Appendix A. for NF-BC. These ablation experiments aim show which hyper-parametest NF-BC is robust to.
important tune and which hyper-parameters
In Table 10, present experiments ablating the inclusion of LayerNorm, a generalized permutation layer, and the sampling trick in NF-GCBC.
Table 11: Task wise performance for Fig. 4.
Environment
Score
puzzle-3x3-play-v0
0.045=0.000
cube-single-play-v00
0.795=0.100
scene-play-vO
0.420=0.000
cube-double playy-ll--ayyvv
0.351 00.087
mutmaze-mediuu--tiieee
0.566=0.00m
0.31 0.:14=20.0
entmaze-large-navigeee
0.320
untsocoer-arcma-navigae
0.621 =0. 007
aalaoccer-rrena-stiichhh...
0.643=0.016
19

---

Table 12: Task wise performance for Fig. Scores are mean +
Environment
Score
0.764410.073
0.724:10.088
0.024
0.616+0.311
0.608
antsoccer-arent
0.600:10:122
antsoccer -rennacee--re---aaa
0.556:10.000
0.420=#0.077
0.480
0.336:+0.022
0.4400+0.057
0.640=0.0087
cubb-singg--llayyiinnlla-------
0.768:20.060
cub--singleppaayssslllllssssskka
0.600=1.082
0.616*a.067
0.036+a.045
0.188.0.097
humanoidmaze-meddim navigate-singlotask
1..40700.072
humanoidmaze-mcciimm navigate-singlant bingletask-tskk4-a.
0.016:100015
0.208+
puzzle-3x3 3x3 1-a33--ll--saae wii-s-innnst-aaskee -vo
0.996:10.008
puzzle-3x3 1eae--33---llyyylle play-singletask-tee -a0k220
-ingletass--tss22-00
0.996:10.00s
puzzle-3x3-play-sseeee singgggsiilttttttaskk.
0.976+10.029
0.996=0.008
puzzlee-kk-llaay-iinllttaa---k
0.996:0.000s
sese-play-sinnlettskktaaass
1.000++..00
hete-plly-ssiglltts-taassk
0. 4.800+9.075
scene play-singletask-taskk
0.988
scene-play-stee siaaa-singlltsss-tsssk---- task4-v0
0.85610.056
scene-play gally--inglleass-taae
0.000:1..00
E Hyperparameters used by behavior cloning algorithms based different probabilistic models.
Table 13: Hyper- parameesssetters used NF-BC.
Hyper-parameter
Description
channels
width of each layer of the coupling MLP network. added noise to labels for training the NF using MLE,
noise std
blocks
number of NF blocks.
encoder num layers
number of layers in the observation encoder MLP network. output dimension of observation encoder MLP network.
rep dims
20

---

Table 14:
used by BET.
Hapeee-arameter
Description
embedding dims num layers
dimension length for attention vectors.
number of attention blocks.
num heads
number of heads in attention.
act scale
action scaling parameter.
num bins
number of centroids for k-means. whether to predict offsets for actions or not. action offset prediction loss coefficient. focal loss parameter.
offset prediction
offset loss scale
gamma
Table 15: Hyper--prameee used by DiffPolicy-C.
Description
time-step embedding dims down dims
embedding dimension of diffusion timestep.
downsampling dimensions for Unet architecture.
kernel size
kernel size for CNNs in Unet.
groups
number of groups for group normalization. Length of the forward diffusion chain while training.
scheduler train timesteps
num inference steps
Length of the denoising steps while sampling.
beta schedule type
type of scheduler for noising.
beta start
first-step noise for diffusion scheduler.
beta end
last-step noise for diffusion scheduler.
ema inverse gamma
parameter for ema rate schedule.
ema power
parameter for ema rate schedule. minimum value ema decay rate. maximum value of ema decay rate.
ema min
ema max
Table 16: Hyper-parameters used by DiffPolicy-T.
Description
time-step embedding dims embedding dims num layers
embedding dimension of diffusion timestep.
dimension length for attention vectors.
number of attention blocks.
num heads
number of heads in attention.
scheduler train timesteps num inference steps beta schedule type
Length of the forward diffusion chain while training.
Length of the denoising steps while sampling.
type of scheduler for noising.
beta start
first-step noise for diffusion scheduler.
beta end
last-step noise for diffusion scheduler.
ema inverse gamma
parameter for ema rate schedule.
ema power
parameter for ema rate schedule. minimum value ema decay rate. maximum value of ema decay rate.
ema min
ema max
21

---

Table 17: Sypperparaaeeters used by VQ-BET.
Hyper-parameter
Description
embedding dims
dimension length for attention vectors.
num layers
number of attention blocks.
heads
number of heads in attention.
action seq len vqvae num embed
sequence length used to train the behavior transformer.
cookbook size of vqvae. embedding dimensions of vqvae codes. number of groups in residual vqvae. encoder loss coefficient. action offset prediction loss coefficient. secondary code prediction loss coefficient. focal loss parameter.
vqvac embedding dims vqvae groups
encoder loss multiplier
offset loss multiplier
secondary code multiplier gamma
F
Societal Impacts
Positive Impacts. Our work directly improves various RL algorithms in various RL problem settings, which can benefit domains like roboties, healthcare, and energy systems by enabling efficient autonomous decision-making.
Negative Impacts. As with many advances in reinforcement learning, improvements in policy expressiveness can be misused domains such as surveillance, autonomous weapons, or manipulation in recommendation systems. Further, increased training complexity may lead higher computational costs, raising concerns about energy use and environmental impact.
NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The main claims made in the abstract and introduction are that our experiments show that NFs are expressive, flexible and generally capable models for many RL problems like imitation learning, offline, goal conditioned RL and unsupervised RL. Towards the end our introduction, we highlight the success results of NF based algorithms in all of these problem settings. Section 6 contains detailed discussions of these experiments and results which concretely justify these claims.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Yes, the second paragraph of Section discussion the main limitations of our work.
Guidelines:
22

---

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't It acknowledged judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proof'
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
Our paper is primarily an empirical paper. While Section does provide a unifying proba-
bilistic framework to think about many RL algorithms, the specific algorithm-we results are known and have been cited accordingly.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Yes, the paper provide all information needed to reproduce the all the experimental results. Appendix contains the implementation details of all the algorithms,, including the hyper- parameters used and their values. We have also provided the code for all our experiments organized in readable manner in supplementary experiments. We have also provided the code for our architecture in both Jax [10] and Torch [69].
23

---

Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-sonnee open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Like mentioned in the previous answer, we have provided the code for all our experiments organized in readable manner in supplementary experiments. We have also provided the code for our architecture in both Jax [10] and PyTorch [69].
Guidelines:
The answer NA means that paper does not include experiments requiring code. Please see the NeurIPS code and data submission guidelines (https://nips.cm cc/ (eelli//guide//CCe htt////oeeeiiiioiioooooooo
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
/ /nips. wttpcccccppcccceeeeoee.ec.ccccct guwdes/ddeeeeeeeo /Cooeeubmiiieeeeeeeee
hotoliicggiiees/oduuuuiiiiii cetimisssion awssonooliigoo for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
24

---

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Appendix contains the implementation details of all the algorithms, including the hyper- parameters used, their values, and details of how they were chosen. Every subsection of Section contains detailed description of the tasks and baselines for all RL problems (offline imitation learning, conditional offline imitation learning, offline RL, unsupervised goal conditioned RL).
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: Yes, all experiments are conducted across five seeds and error bars are shown for all bar / graph plots.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2- 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
25

---

Justification: In Appendix A, we provide the type of GPUs used and the time taken for all our experiments.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't oodn't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: /tte:://neeurpeeeeeeeeee httppp////////iicccccuiiieeeeee
Answer: [Yes]
Justification: have read the guidelines and our research does indeed conform, in every respect, with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: In Appendix F we have added brief discussion of both positive and potentially negative societal impacts of our work.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (c.g., deploymeent of technologics that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
26

---

11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
We do not use any of data or models that have a high risk for misuse.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: Wherever necessary, we have eited the original source of the datasets and benchmarks used in our paper Section 6.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: Yes, our code well documented. Details about architecture, experiments and results have been provided in the paper Section 6, Appendix A.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/codeemmodel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
27

---

The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification:
There are no crowdsoureing experiments in our paper.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/teview based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: the paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMSs if it is important, original, or non-standard component of the core methods this research? Note that the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: The core method development this research does not involve LLMS as any important, original, or non-standard components.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer LLM policy (https: http://epss.cc/oooeeeeeeeeeeeeee for what should or should not be described.
28


Reviews:
[{"review_id": "u1scLdYfN4", "reviewer": "Reviewer_oKbR", "summary": "This paper implements Normalizing Flows(NFs) to RL problems.\nThe main claim is that NFs are generally capable models for RL.\nThe authors propose an NF architecture and algorithms using coupling networks and linear flows.\nBy implementing and evaluating NF to offline IL, conditional IL, offline RL, GCRL, and URL across diverse tasks, it shows a significant performance gain.", "questions": "A more precise context of my questions is described in the weakness section. I briefly summarize my questions below.\n\n1. How is the performance of other NF architectures evaluated? \n    \n2. Can you provide an analysis of why NFs work well? What are the specific properties of NFs for RL performance contribution?\n    \n3. What is the detailed computational cost analysis, such as training/inference time, memory, compared to baselines?\n    \n4. What are the specific RL scenarios where NFs fail?", "limitations": "The authors acknowledge limitations such as invertible architectures and new NF architectures, but the limitation section should be more comprehensive. My feedbacks are given in the weakness sections.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "### Strengths\n\n**Novel perspective for RL architecture**: \nThis paper opens and expands overlooked research directions on NFs for RL.\n\n**Comprehensive study and experiment**:\nThis paper provides comprehensive experiments across 82 tasks in 5 RL settings.\nThe results show competitive performance compared to previous sota baselines with fewer parameters. \nAlso, this work provides practical insights about NF architecture for RL, which is noteworthy for RL practitioners to consider NF beyond diffusion and transformers. \n\n\n### Weaknesses\n**Overstated theoretical framework and claims**: \nThe paper's central claim: \"RL algorithms are MLE and VI\" (Section 4) is too strong. Several specific issues:\n\n1. Claiming Q-function estimation equals occupancy distribution modeling requires very specific assumptions about reward structure (as given in references)\n\n2. Claiming behavior cloning as MLE couldn't contain other loss function variants. (MSE loss ≠ MLE when action distribution is not Gaussian)\n\n3. Claiming Q-learning and entropy-regularized policy optimization as variational inference is too general for what we choose for the posterior.\n\n4. By this logic, all ML problems could be reframed as MLE+VI, which is feasible to NFs, making the argument meaningless\n\nI hope these issues will be addressed in the rebuttal phase or further experiments.\n\n**Lack of fundamental understanding**: \n\nThe paper lacks a fundamental understanding of why NFs perform well in RL.\nBeyond \"NF is good for RL\", we need to find the underlying reason, such as \"It's because general expressivity/specific architectural or algorithmic properties/inductive biases\"\nAlso, failure case analysis is needed for fundamental understanding and practical implementation.\nIn addition, beyond the proposed architecture (RealNVP + linear flow + LayerNorm), the authors should compare other NF variants.\nIt is crucial to distinguish the contribution of NF properties and architecture properties.\nBefore that, the general claim about NFs remains unclear.", "was_revised": true, "final_justification": "My concerns were addressed during the rebuttal. I increased my score to a 4(borderline accept).", "rebuttal": "We thank the reviewer for their insightful feedback and constructive comments. It seems that the reviewer’s main feedback is regarding (1) explaining the importance of design choices for NFs and (2) the classification of RL algorithms as VI and MLE. \n\nTo address (1), we have run new experiments comparing with two other NF architectures, evaluated expressivity and the ability to do VI as the primary reason behind NF's performance, added ablations for all design choices (normalization, scaling depth, sampling trick, linear flow), computational costs and identified failure cases.\n\nRegarding (2), we have changed our claim to be more specific, and have provided counter-examples (algorithms that don’t fall into this category). \n\n**Do these answers, revisions and experiments address all the reviewer's concerns?**\n\n>  a fundamental understanding of why NFs perform well in RL. What are the specific properties of NFs for RL performance contribution?\n\nThe main reason why NFs perform so well is their expressivity. Almost all datasets we use have multi-modal, noisy and high dimensional action spaces. We compare the validation likelihoods of the policies learnt by GCBC and NF-GCBC on three tasks: \n\n**Log-likelihoods (higher is better)**    \n| Method  | antmaze-large-navigate | scene-play | cube-single-play |\n|---------|----------|----------|----------| \n| NF -GCBC     | -0.673   | 2.975 | 2.756 |\n| GCBC | -7.662 | -4.65 |  -4.661 |\n\nNF-GCBC is able to model the behavior policies more accurately. It is well established that higher-likelihoods are correlated with performance and generalisation in imitation learning [1,2].\n\nIn offline RL experiments, the ability to do VI is important for efficient policy improvement. We add a new didactic 4 goal environment (sec 4.1 in [11]). This is an online RL experiment which consists of 4 distinct goals. We report the number of times a trained policy of any method reaches a goal, out of 20 independent rollouts:\n\n**Goal coverage (more uniform is better):**\n| Method  | goal 1 | goal 2 | goal 3 | goal 4 |\n|---------|----------|----------|----------|----------| \n| SAC     | 0   | 0 | 0 | 20 |\n| SAC-VAE | 0 | 0 |  0 | 20 | \n| SAC-NF | 8 | 2 |  7 | 3 |\n\nClearly, SAC-NF is able to reach all 4 goals, whereas SAC-VAE, which uses a VAE state encoder and a gaussian action decoder still shows mode-seeking behavior which is a problem for reverse KL optimization. This shows that NFs which directly allow VI optimization are able to efficiently cover all the modes of the Q-function. \n\nThe above two experiments suggests that expressiveness and the ability to do VI are the most important property of NFs for RL performance contribution.\n\n> How is the performance of other NF architectures evaluated? distinguish the contribution of NF properties and architecture properties\n\nIn NFs, the architecture is the primary choice that affects modelling and inference. The architectural choices we make are optimized for performance, speed, and stability with MLE and VI. Below we perform new experiments to compare our method with RealNVP [3] (an old NF architecture) and TarFlow [4] (a SOTA NF architecture in computer vision).\n\n**Success rates (higher is better):**\n| Method  | scene-play | cube-single-play |\n|---------|---------|----------| \n| NF-GCBC      | 0.42 | 0.795 |\n| TarFlow | 0.456 |  0.594 |\n| RealNVP | 0.02 |  0.04 |\n\n**Log-likelihoods (higher is better):**\n| Method  | scene-play | cube-single-play |\n|---------|---------|----------| \n| NF-GCBC | 2.756 | 2.975 |\n| TarFlow | 2.902 |  3.01 |\n| RealNVP | 1.22 |  NAN |\n\n| Method  | Inference Speed (lower is better) |\n|---------|---------|\n| NF-GCBC      | 0.127|\n| TarFlow | 0.38 |\n| RealNVP | 0.120 | \n\nFrom the above results, it is evident that \n1) RealNVP is fast but unstable (likelihoods explode). Hence the combination we propose – which interleaves linear permutation flows and layernormed coupling networks – is important for high performance\n2) TarFlow is able to model the distributions well and perform at par with NF-GCBC on one task, but it is three times slower than NF-GCBC. This is because in auto-regressive architectures like Tarflow (and IAF [5]), the map from action to noise can be parallelized, but the inverse map cannot. The inverse map needs to auto-regressively generate every index of the action outputs. This makes sampling and VI very slow. For tasks like offline RL or unsupervised goal sampling, where all three are required (sampling, VI, likelihoods), training as well as inference are slowed down. Because coupling networks are fast for both sampling and likelihoods, we chose them over auto-regressive models. \n\n**We hope that this clarifies both experimentally and conceptually, the explanation behind why our design choices work well, and the properties / failure cases of different NF architectures.**\n\n> The paper's central claim: \"RL algorithms are MLE and VI\" (Section 4) is too strong\"\n\nWe agree with the reviewer that all RL algorithms are not versions of MLE and VI. We have reduced our claim and made it more specific -- *Some versions of BC, GCBC, MaxEnt RL, and GCRL are instantiations of MLE or VI*. This claim is grounded and clearly explained in equations [7,8,10,11,12]. We take the reviewer's note and have added a paragraph about counterexamples of RL algorithms that cannot be traditionally viewed as versions of MLE / VI, for example:\n\n1) BC / GCBC with SVMs or other non probabilistic losses like 0–1 loss, hinge loss.\n2) temporal difference learning.\n3) State / Action abstraction learning algorithms.\n4) Count based exploration algorithms, Planning algorithms like MCTS, algorithms based on game theory (GANs, minmax regret).\nThese counterexamples show that the framework we propose in the paper does not apply to all algorithms trivially. Moreover, a rich history of work viewing RL algorithms from a probabilistic perspective [6-10], backs that the claims we make are useful for algorithmic development.\n\n> detailed computational cost analysis, such as training/inference time, memory, compared to baselines?\n\n| Method  | Inference Speed | Training Speed (seconds; per batch) | Number of parameters (millions; re: memory) |\n|---------|---------|---------|---------|\n| NF-GCBC      | 0.127 | 0.023 | 10 M |\n| VQ-BET | 0.273 | 0.021 | 4.23 M |\n| Diffusion Policy | 1.28 | 0.019 | 65.78 M |\n| TarFlow | 0.273 | 0.042 | 15.M |\n\n> specific RL scenarios where NFs fail?\n\nWe have added a list of detailed failure cases of NFs in RL in the appendix. We present an outline below:\n1) Without LayerNorm and linear permutation flows, NFs are unstable to train and sometimes result in exploding losses.\n2) Without adding noise to the actions (noise_std hyperparameter in table 9), NFs are unstable to train and sometimes result in exploding losses.\n3) Training with larger values of $\\lambda$ (in equation 10) leads to lower performance. \n\nTo further clarify the success / failure of NFs with various hyper-parameters & design choices, we run ablation experiments controlling our main design choices and hyper-parameters:\n\n| Method  | antmaze-large-navigate | scene-play | cube-single-play |\n|---------|----------|----------|----------| \n| NF -GCBC     | 0.32   | 0.42 | 0.795 |\n| without permutation flows | 0.26 | 0.04 |  0.0 |\n| without LayerNorm | 0.29 | 0.3 |  0.0 |\n| without sampling trick | 0.3 | 0.42 |  0.69 |\n\n| noise_std  | Push-T |\n|---------|----------|\n| 0.5     | 0.77   |\n| 0.3 | 0.79 |\n| 0.1 | 0.827|\n| 0.05 | 0.84 |\n| 0.03 | 0.873 |\n| 0.01 | 0.83 |\n| 0.0 | NAN |\n\nThe above two tables suggest that NFs fail without permutation flows and LayerNorm, and performance decreases at very low and very high noise_std values.\n\n[1] Behavior Transformers: Cloning-K modes with one stone\n\n[2] Diffusion Policy: Visuomotor Policy Learning via Action Diffusion\n\n[3] Density estimation using Real NVP\n\n[4] Normalizing Flows are Capable Generative Models \n\n[5] Improving Variational Inference with Inverse Autoregressive Flow\n\n[6] Reinforcement Learning and Control as Probabilistic Inference: Tutorial and Review\n\n[7] What type of inference is planning?\n\n[8] General duality between optimal control and estimation\n\n[9] Robot trajectory optimization using approximate inference\n\n[10] Modeling purposeful adaptive behavior with the principle of maximum causal\nentropy\n\n[11] Maximum Entropy Reinforcement Learning via Energy-Based Normalizing Flow"}, {"review_id": "tucCT55KTx", "reviewer": "Reviewer_CtEo", "summary": "This paper examines reinforcement learning (RL) problems from a probabilistic perspective, framing RL settings as exact likelihoods, sampling, and variational inference. The authors analyze the capabilities of Normalizing Flows (NFs) and other probabilistic methods, arguing that NFs provide the capabilities necessary for efficient RL. Methodologically, they utilize existing coupling networks and linear flows to construct NF blocks, which are integrated into RL algorithms to model the policy, Q-function, and occupancy measure under different RL settings. The paper conducts extensive experiments in imitation learning, offline RL, and goal-conditioned RL, demonstrating the effectiveness of NFs in these settings.", "questions": "1.What is the number of steps for the invertible transformations in the experiments? Is it related to the \"blocks\" hyperparameter mentioned in the appendix?\n\n2.According to Section 4, NFs are used to model the policy network, Q-value, and occupancy measure in imitation learning, offline RL, and goal-conditioned RL, respectively. Can both the policy and Q-value learning in these scenarios be modeled using NFs?\n\n3.What are the differences between the NF architecture chosen in this paper and those of NICE and RealNVP? Why was the coupling network chosen?\n\nHope the authors provide relevant explanations for mentioned weaknesses above.", "limitations": "No, refer to the second and third point in weakness.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strength:\n\n(1) This paper presents a unique perspective on the classification of existing RL problems and offers a thorough discussion of the properties of various generative models, characterized by clear logic and fluent language.\n\n(2) The experiments in the paper are extensive and comprehensive, considering different reinforcement learning scenarios such as imitation learning, offline RL, and goal-conditioned RL, which demonstrate the effectiveness of the proposed method.\n\nWeaknesses:\n\n(1) This paper primarily utilizes existing NF architectures to construct its NF block. However, it lacks an analysis of the motivations for choosing coupling networks as reversible transformations.\n\n(2) Similar to the diffusion steps in diffusion models, the steps of invertible transformations are also an important hyperparameter in NFs, affecting both model expressiveness and complexity. However, there is a lack of relevant discussion on this topic. Additionally, based on Eq. (3) and (5), it appears that this paper only performs one-step transformation, making it challenging to ensure the model's expressiveness.\n\n(3) Measuring the model's simplicity solely by the number of hyperparameters, as shown in Figure 3, is unreasonable and insufficient. I suggest that the authors incorporate model complexity into their evaluation metrics. Furthermore, most experiments in this paper do not present training curves or complexity comparisons, making it unconvincing to claim that NFs are sample-efficient for RL.", "was_revised": true, "final_justification": "During the rebuttal, the authors conducted additional experiments to validate and analyze the motivation behind their specific NF design choices. Moreover, they provided further analysis on how the performance of their proposed method varies with different numbers of NF blocks. They also introduced an additional metric to measure model simplicity.\n\nMy only concern is that they did not provide training curves for most settings; instead, they only reported the convergence step for their own method. This makes it difficult to assess the convergence behavior of other baselines.\n\nTaking all of the above into consideration, I recommend a Borderline Accept for this paper.", "rebuttal": "We thank the reviewer for their insightful feedback and constructive comments. It seems that the reviewer’s main feedback is regarding (1) more motivation behind the particular NF design choices, (2) a better measure of simplicity, and (3) more discussion of the number of NF blocks and expressiveness. \n\nTo address (1), we have run new ablation experiments comparing the performance, speed and stability with two other popular NF architectures. \n\nFor (2), we now additionally compare the number of model parameters (complexity), number of training objectives (end to end vs step wise) with baselines, and provide plots for hyper-parameters sensitivity of our method to give a holistic view of implementation simplicity for practitioners. \n\nFor (3), we add new experiments showing how performance scales with the number of NF blocks. Additionally, we apologize for the reviewer's confusion and point that (see line 104) all NF models in the paper contained multi-step transformations. \n\n**Do these answers, revisions and experiments address all the reviewer's concerns?**\n\n> analysis of the motivations for choosing coupling networks as reversible transformations\n\nThe architectural choices we make are optimized for performance, speed, and stability with MLE and VI. Below we perform new experiments to compare our method with RealNVP [1] (an early NF architecture) and TarFlow [2] (a SOTA NF architecture in computer vision).\n\n**Success rates (higher is better):**\n| Method  | scene-play | cube-single-play |\n|---------|---------|----------|\n| NF-GCBC      | 0.42 | 0.795 |\n| TarFlow | 0.456 |  0.594|\n| RealNVP | 0.02 |  0.04 |\n\n**Log-likelihoods (higher is better):**\n| Method  | scene-play | cube-single-play |\n|---------|---------|----------|\n| NF-GCBC | 2.756 | 2.975 |\n| TarFlow | 2.902 |  3.01 |\n| RealNVP | 1.22 |  NAN |\n\n| Method  | Inference Speed (lower is better) |\n|---------|---------|\n| NF-GCBC      | 0.127|\n| TarFlow | 0.38 |\n| RealNVP | 0.120 | \n\nFrom the above results, it is evident that \n1) RealNVP is fast but unstable (likelihoods explode). Hence the combination we propose, of interleaving linear permutation flows and layernormed coupling networks is necessary for high performance.\n2) TarFlow is able to model the distributions well and perform at par with NF-GCBC on one task, but it is three times slower than NF-GCBC. This is because in auto-regressive architectures like Tarflow (and IAF [5]), the map from action to noise can be parallelized, but the inverse map cannot. The inverse map needs to auto-regressively generate every index of the action outputs. This makes sampling and VI very slow. For tasks like offline RL or unsupervised goal sampling, where all three are required, training as well as inference are slowed down. Because coupling networks are fast for both sampling and likelihoods, we chose them over auto-regressive models. \n\n> What are the differences between the NF architecture chosen in this paper and those of NICE and RealNVP? \n\nOur architecture is based on RealNVP. While the RealNVP on vision tasks, our focus was on boosting performance in RL tasks. The main difference which facilitates this (see table above for new comparisons) is using a generalized permutation flow layer after every coupling network, and LayerNorm (instead of the original normalization proposed in RealNVP). \n\n> Why was the coupling network chosen?\n\nLike mentioned above (after the experimental results presented above), we chose RealNVP's coupling layers for their high speed in both sampling and likelihood estimation. This is crucial for RL tasks like offline RL or unsupervised goal sampling (where all three are required) or robot imitation learning (where controllers with low latency are preferred). \n\n> steps of invertible transformations are also an important hyperparameter in NFs\n\nIndeed, the number of steps (or number of NF blocks as we call them) are an important hyper-parameter. Below we provide how the performance of NF-BC and NF-GCBC changes with different numbers of NF blocks.\n\n**Success rates (higher is better):**\n| number of NF blocks  | Push-T |\n|---------|----------|\n| 6   | 0.77   |\n| 14 | 0.79 |\n| 24 | 0.827 |\n| 28 | 0.859 |\n| 32 | 0.85 |\n| 40  | 0.844 |\n| 48  | 0.87 |\n| 54  | 0.85 |\n\n**Success rates (higher is better):**\n| number of NF blocks  | antmaze-large-navigate | scene-play | cube-single-play |\n|---------|----------|----------|----------| \n| 6     | 0.32   | 0.42 | 0.795 |\n| 12 | 0.34 | 0.36 |  0.86 |\n| 24 | 0.36 | 0.5 |  1 |\n\n> Additionally, based on Eq. (3) and (5), it appears that this paper only performs one-step transformation, making it challenging to ensure the model's expressiveness\n\nWe would like to point out that (see line 104) all NF models in the paper contained multi-step transformations. The blocks parameter is the number of NF blocks (invertible non linear transformations) stacked together to create the NF. It is mentioned in Tables 2,3,5 and A.4 for all our methods.\n\n> Measuring the model's simplicity solely by the number of hyper-parameters\n\nWe agree that there exist other factors which also determine simplicity of implementing an algorithm (from a practitioner's perspective). We have removed this plot from our revised version of the paper. We now compare other details of the algorithms like inference and training speed, number of parameters (related to bayesian information criterion [3]), and number of training objectives. We believe that these new comparisons provide a more holistic view for RL practitioners. \n\n| Method  | Inference Speed (seconds) | Training Speed (seconds; per batch) | Number of parameters (millions) | Number of objectives |\n|---------|---------|---------|---------|---------|\n| NF-GCBC      | 0.127 | 0.023 | 10 M | 1; end to end |\n| VQ-BET | 0.273 | 0.021 | 4.23 M | 2; requires pre-training a hierarchical VQ-VAE |\n| Diffusion Policy | 1.28 | 0.019 | 65.78 M | 1; end to end |\n| TarFlow | 0.273 | 0.042 | 15.M | 1; end to end |\n\nNF-BC, while being competitive on all tasks, is 2 and 10 times faster than VQ-BET and Diffusion policy. Unlike VQ-BET, which requires pretraining a VQ-VAE, NF-GCBC is an end to end algorithm which directly optimizes for the desired likelihood. Our updated claim is that NF-BC is faster in inference, easier to implement (end to end optimization of a single loss function, contains 2-3 times less hyper-parameters ) and achieves competitive performance across all tasks (best success rate on two tasks and in top 2 methods on the other two tasks).\n\nAdditionally, algorithms with fewer hyper-parameters can still be difficult to tune if their performance is sensitive to the values of those hyper-parameters. Hence, we present ablation experiments to show which hyper-parameter are important to tune and which hyper-parameters is NF-BC is robust to:\n\n**Success rates (higher is better):**\n| noise_std  | Push-T |\n|---------|----------|\n| 0.5     | 0.77   |\n| 0.3 | 0.79 |\n| 0.1 | 0.827|\n| 0.05 | 0.84 |\n| 0.03 | 0.873 |\n| 0.01 | 0.83 |\n| 0.0 | NAN |\n\n**Success rates (higher is better):**\n| width of NF's coupling layers  | Push-T |\n|---------|----------|\n| 128     | 0.79   |\n| 256 | 0.82 |\n| 1024 | 0.823 |\n\n**Success rates (higher is better):**\n| number of NF blocks  | Push-T |\n|---------|----------|\n| 6   | 0.77   |\n| 14 | 0.79 |\n| 24 | 0.827 |\n| 28 | 0.859 |\n| 32 | 0.85 |\n| 40  | 0.844 |\n| 48  | 0.87 |\n| 54  | 0.85 |\n\nThe above results show that NF-BC is generally robust to many values. It is important to have a nonzero and small noise_std value, but many values between 0.01 - 0.1 work well.\n\n> most experiments in this paper do not present training curves or complexity comparisons, making it unconvincing to claim that NFs are sample-efficient for RL.\n\nBecause most experiments have a fixed dataset, sample efficiency does not play a role. We do provide complete training curves for all online experiments (see Figure 7). Importantly, in any of the experiments, we do not make the claim that NFs are sample-efficient for RL. Nonetheless, we have added the gradient steps vs performance curves in the appendix of the revised paper.\n\n> Can both the policy and Q-value learning in these scenarios be modeled using NFs?\n\nIn GCRL, NFs can indeed be used to approximate policies and Q-functions. In other offline RL they can only be used for policies. We have added a sentence in section 4 to emphasize this.\n\n[1] Density estimation using Real NVP\n\n[2] Normalizing Flows are Capable Generative Models \n\n[3] Estimating the Dimension of a Model"}, {"review_id": "fjucFGhXmR", "reviewer": "Reviewer_8s3P", "summary": "This paper advocates for Normalizing Flows (NFs) as a useful building block for many RL algorithms. The core of the argument comes in two pieces: First, the work showcases that a wide variety of subproblems within RL involve maximum likelihood estimation and variational inference (Section 4) such as Imitation Learning, Offline RL, and Goal-Conditioned RL; Second, the paper presents a simple NF architecture that can be used as a modular tool across many different RL algorithms and problem settings. The design goal of the new NF architecture is characterized by the desiderata to simple, fast, and expressive. The resultant architecture combines a dense coupling network and a linear flow from prior work, along with standard layer normalization. The paper then presents a comprehensive experimental study involving a diversity of problem settings and domains, again including Imitation Learning, Offline RL, and Goal-Conditioned RL. The experimental data support the conclusion that the proposed NF architecture can act as a simple, expressive algorithmic building block across these many settings.", "questions": "I have two primary suggestions, along with a variety of low-level writing suggestions and comments which I include below. My two higher level suggestions are as follows:\n- \"Our desiderata is to build a simple and fast architecture that is expressive enough to compete and even surpass the most popular approaches across a range of RL problems.\": I like this, and believe it could be emphasized even more. Additionally, you can probably weave discussion of this desiderata even more into the rest of Section 5---what alternatives did you consider? Why did you not use them? At the moment the design choices feel somewhat arbitrary (even if they are simple), so it could be useful to motivate them slightly more. similarly, the three steps around line 166-167 are quite important: it's worth spending a sentence or two expanding on them more (before the comment about the log-det).\n- As mentioned in the weaknesses, I believe a slightly longer related work section will add more context to the contribution. Is it possible to provide more high-level information about the subfields discussed?\n\n__Typos / Writing Suggestions__\n\nLastly, here are some low-level writing suggestions and comments:\n\nIntroduction:\n- Typo: \" probabilistic models which all the\" --> \" probabilistic models with all the\"\n\nRelated work:\n- The three paragraphs at present would each benefit from some synthesizing text that summarises the space; what are the open questions, what are the current limitations, and so on.\n\nPreliminaries:\n- \"as dataset size N increases\" --> \"as the size of the dataset ($N$) increases\"\n- Line 85: Can you briefly mention what $Z$ is here: \"... its unnormalized part $\\tilde{p}(x) = p(x) \\times Z$, and later state why it can be ignored?\n- \"jacobian matrix\" --> \"Jacobian matrix\"\n- The notation in the subscript of the Expectation in Equation 6 is slightly non-standard: $x$ is used in $z \\sim p_0(x)$ and again $x = f_\\theta^{-1}(z)$; the latter seems unecessary.\n- Typo: \" multiple RL problems settings\" --> \" multiple RL problem settings\"\n- \" underlying controlled markov process\" --> \"underlying controlled Markov process\"\n- Line 126: Personal preference, but I would not index $s$ and $a$ by $t$ when mentioning the reward function here: \" a reward function r(s_t, a_t)\" --> \" a reward function $r(s, a)$\" (but I would use the time indexing on line 128 when the states and actions are inside of the sum)\n- $\\mathcal{D}$ is a set, but the notation $\\sim \\mathcal{D}$ is used---of course I assume this means there is some implicit probability mass function on $\\mathcal{D}$, but this could be stated explicitly somewhere the make the notation more precise.\n- Line 139: \"We will define the\" --> \"We define the\"\n- Line 145: \"apriori\" --> \"a priori\"\n- 'connected “coupling” network': Why the quotes?\n- In the Figure 3 caption, the \"(a − d)\" reads as just slightly confusing, in part because the \"-\" is rendered as a subtraction symbol (you can replace with a-d instead of $a-d$). Or alternatively you might replace with (Subfigures a-d).\n- \"Simplicity is measured by the inverse of the number of hyper-parameters\": This is reasonable, but could use with slightly more justification (it's the same justification as BIC and AIC---Bayesian and Akaikei Information Criterion---which provide different accounts of simplicity).\n- I believe the acronym for MDP was introduced without being defined. You introduce controlled Markov processes but not MDPs.", "limitations": "Yes, I believe as needed for this kind of work the authors disclose potential limitations of NFs.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "__Strengths__\n- [Sharp Narrative] This is a very focused paper: the whole narrative is framed around a single idea, and as such the communication of the idea is clear from start to finish.\n- [Code] Implementations are provided in both Jax and PyTorch.\n- [Simple Architecture] The proposed architecture is quite simple, and I imagine others in the community can quickly make use of it.\n- [Broad Experimental Study] The described experimental study is diverse and comprehensive. Many baselines are used, results are thoughtfully communicated, and the strength of the claims made matches the strength of the evidence.\n\n__Weaknesses__\n- [Justification] A few design decisions could be motivated in more detail. For instance, simplicity of an algorithm is measured according to the inverse of the number of its hyperparameters. This is one choice---I would find a more faithful plot to be one that simply reports that number directly, rather than relabelling it \"simplicity\".\n- [Clarity] There are a few areas where some added elucidation of the central ideas will help readers new to some of the ideas. For instance, MDPs are never defined, and the introduction of the coupling networks is quite quick with little added context.\n- [Related Work] This is a minor one, but the discussion around the related work would benefit from a deeper discourse on the areas discussed.", "was_revised": true, "final_justification": "This is a strong paper, with a sharp focus and a compelling experimental study: The paper effectively demonstrates that Normalizing Flows can be an important building block for many RL components.\n\nSeveral poignant questions were raised during the initial reviews and rebuttals. Most notably, my two concerns throughout the discussion were regarding the motivation of various design choices made, and the scope of the claims made regarding RL, MLE, and Variational Inference. To me, these concerns have been completely addressed in the rebuttal, and I maintain a strong outlook on the work.", "rebuttal": "We thank the reviewer for their insightful feedback and constructive comments. It seems that the reviewer’s main feedback are regarding (1) motivations of the chosen design choices, (2) elucidation of the central ideas like coupling networks and MDPs and a more detailed related work section. \n\nTo address (1), we have added new ablation experiments for all our design choices, underscoring their importance. We have also added new comparisons with two other NF architectures, showing why our proposed architecture is best suited for RL algorithms. \n\nTo address (2), we have modified our paper to introduce MDPs clearly, and added a deeper discussion of coupling networks. We have also added a detailed related work section in the appendix. \n\n**Do these answers, revisions and experiments address all the reviewer's concerns?**\n\n> motivating design decisions\n\nThe architectural choices we make are optimized for performance, speed, and stability with MLE and VI. Below we perform new experiments to compare our method with RealNVP [1] an old NF architecture and TarFlow [2] a SOTA NF architecture in computer vision.\n\n**Success rates (higher is better):**\n| Method  | scene-play | cube-single-play |\n|---------|---------|----------|\n| NF-GCBC      | 0.42 | 0.795 |\n| TarFlow | 0.456 |  0.594|\n| RealNVP | 0.02 |  0.04 |\n\n**Log-likelihoods (higher is better):**\n| Method  | scene-play | cube-single-play |\n|---------|---------|----------|\n| NF-GCBC | 2.756 | 2.975 |\n| TarFlow | 2.902 |  3.01 |\n| RealNVP | 1.22 |  NAN |\n\n| Method  | Inference Speed (seconds, lower is better) |\n|---------|---------|\n| NF-GCBC      | 0.127|\n| TarFlow | 0.38 |\n| RealNVP | 0.120 | \n\nFrom the above results, it is evident that \n1) RealNVP is fast but unstable (likelihoods explode). Hence the combination we propose, of interleaving linear permutation flows and layernormed coupling networks is necessary for high performance.\n2) TarFlow is able to model the distributions well and perform at par with NF-GCBC on one task, but it is three times slower than NF-GCBC. This is because in auto-regressive architectures like Tarflow (and IAF [5]), the map from action to noise can be parallelized, but the inverse map cannot. The inverse map needs to auto-regressively generate every index of the action outputs. This makes sampling and VI very slow. For tasks like offline RL or unsupervised goal sampling, where all three are required, training as well as inference are slowed down. Because coupling networks are fast for both sampling and likelihoods, we chose them over auto-regressive models. \n\nWe also add new experiments ablating the inclusion of LayerNorm, a generalized permutation layer, and the sampling trick:\n\n**Success rates (higher is better):**\n| Method  | antmaze-large-navigate | scene-play | cube-single-play |\n|---------|----------|----------|----------| \n| NF -GCBC     | 0.32   | 0.42 | 0.795 |\n| without permutation flows | 0.26 | 0.04 |  0.0 |\n| without LayerNorm | 0.29 | 0.3 |  0.0 |\n| without sampling trick | 0.3 | 0.42 |  0.69 |\n\nThese experiments show that above architectural modifications are crucial for performance of NFs in the context of RL. We hope that the above experiments serve as sufficient motivation for our design choices.\n\n> I would find a more faithful plot to be one that simply reports that number directly, rather than relabelling it \"simplicity\".\n\nWe agree with the reviewer and have removed that plot from our revised paper. We now compare other details of the algorithms like inference and training speed, number of parameters (related to bayesian information criterion [1]), number of training objectives, to provide a more holistic view for RL practitioners. \n\n| Method  | Inference Speed (seconds) | Training Speed (seconds; per batch) | Number of parameters (millions) | Number of objectives |\n|---------|---------|---------|---------|---------|\n| NF-GCBC      | 0.127 | 0.023 | 10 M | 1; end to end |\n| VQ-BET | 0.273 | 0.021 | 4.23 M | 2; requires pre-training a hierarchical VQ-VAE |\n| Diffusion Policy | 1.28 | 0.019 | 65.78 M | 1; end to end |\n| TarFlow | 0.273 | 0.042 | 15.M | 1; end to end |\n\nNF-BC, while being competitive on all tasks, is 2 and 10 times faster than VQ-BET and Diffusion policy. Unlike VQ-BET, which requires pretraining a VQ-VAE, NF-GCBC is an end to end algorithm which directly optimizes for the desired likelihood. Our updated claim is that NF-BC is faster in inference, easier to implement (end to end optimization of a single loss function, contains 2-3 times less hyper-parameters ) and achieves competitive performance across all tasks (best success rate on two tasks and in top 2 methods on the other two tasks).\n\nAdditionally, algorithms with fewer hyper-parameters can still be difficult to tune if their performance is sensitive to the values of those hyper-parameters. Hence, we present ablation experiments to show which hyper-parameters are important to tune and which hyper-parameters NF-BC is robust to:\n\n**Success rates (higher is better):**\n| noise_std  | Push-T |\n|---------|----------|\n| 0.5     | 0.77   |\n| 0.3 | 0.79 |\n| 0.1 | 0.827|\n| 0.05 | 0.84 |\n| 0.03 | 0.873 |\n| 0.01 | 0.83 |\n| 0.0 | NAN |\n\n**Success rates (higher is better):**\n| width of NF's coupling layers  | Push-T |\n|---------|----------|\n| 128     | 0.79   |\n| 256 | 0.82 |\n| 1024 | 0.823 |\n\n**Success rates (higher is better):**\n| number of NF blocks  | Push-T |\n|---------|----------|\n| 6   | 0.77   |\n| 14 | 0.79 |\n| 24 | 0.827 |\n| 28 | 0.859 |\n| 32 | 0.85 |\n| 40  | 0.844 |\n| 48  | 0.87 |\n| 54  | 0.85 |\n\nThe above results show that NF-BC is generally robust to many values. It is important to have a non zero and small noise_std value, but many values between 0.01 - 0.1 work well.\n\n> Related work section\n\nWe have added a detailed related work section in the appendix of the paper.\n\n> Typos / Writing Suggestions\n\nWe appreciate all of these suggestions and have updated the paper to incorporate them.\n\n[1] Estimating the Dimension of a Model"}, {"review_id": "NJoPtxTlXW", "reviewer": "Reviewer_aQtu", "summary": "This paper proposes using normalizing flows (NF) to model policy distributions (i.e., action distributions) in offline reinforcement learning (RL) and state occupancy distributions in goal-conditioned RL (GCRL). The authors demonstrate that, due to NF’s ability to support both maximum likelihood estimation (MLE) and variational inference (VI), the objectives of offline RL and GCRL can be computed efficiently. Experimental results show that the proposed approach improves performance on certain tasks in imitation learning (IL), offline RL, and GCRL.", "questions": "1. Given the conceptual overlap with prior work [1–6], could the authors provide experimental comparisons with these baselines? Also, including recent GCRL approaches such as [7] would improve the completeness of the evaluation.\n2. Could the authors include an example to illustrate the benefits of NFs over other continuous generative models (e.g., VAEs, EBMs, and diffusion)? A minimal setup, similar to Section 5.2 in [8] or Section 4.1 in [3], could provide helpful intuition for why NFs may be preferable in offline RL and GCRL.\n3. The simplicity metric in Fig. 3 seems oversimplified and biased toward counting hyperparameters. This overlooks their sensitivity and broader implementation complexity. Could the authors elaborate on the rationale behind this simplicity metric and why the sensitivity is not considered?\n4. Can the authors provide performance comparisons between NF-GCBC and GCIQL in a manner similar to Fig. 5 (Left)?\n5. Since NFs are inherently continuous generative models, they are not directly applicable to environments with discrete action (Sections 4.1 and 4.2) or state (Section 4.3) spaces. In these cases, the claim that “Normalizing Flows are Capable Models for RL” may be overstated. Could the authors clarify how their approach can generalize to discrete settings and provide empirical evidence?\n6. The scope of this work is limited to IL, offline RL, and GCRL. Is there a specific reason online RL was not considered? Could the authors comment on the applicability of NFs to online RL?\n\n\\\n**References**\n\n[1] Haarnoja et al. Latent space policies for hierarchical reinforcement learning. ICML 2018. \\\n[2] Mazoure et al. Leveraging Exploration in Off-policy Algorithms via Normalizing Flows. CoRL 2019. \\\n[3] Chao et al. Maximum entropy reinforcement learning via energy-based normalizing flow. NeurIPS 2024. \\\n[4] Akimov et al. Let Offline RL Flow: Training Conservative Agents in the Latent Space of Normalizing Flows. NeurIPS Workshop 2022. \\\n[5] Liu et al. Uncertainty-Aware Reinforcement Learning for Risk-Sensitive Player Evaluation in Sports Game. NeurIPS 2022. \\\n[6] Yang et al. Flow to Control: Offine Reinforcement Learning with Lossless Primitive Discovery. AAAI 2023. \\\n[7] Park et al. HIQL: Offline Goal-Conditioned RL with Latent States as Actions. NeurIPS 2023. \\\n[8] Haarnoja et al. Reinforcement Learning with Deep Energy-Based Policies. ICML 2017.", "limitations": "Yes. The authors adequately addressed the limitations and potential negative societal impact.", "rating": 4, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**Strengths**\n1. The proposed method is simple, with writing that is easy to follow.\n2. The paper provides a thorough discussion of the loss functions and presents a straightforward implementation based on Glow and RealNVP architectures.\n3. Experimental results demonstrate performance improvements when applying NF to some of the selected RL tasks.\n\n**Weaknesses**\n1. The use of NFs in both online RL [1–3] and offline RL [4–6] has been well-explored in prior work. This paper largely repurposes NFs in the context of offline RL and GCRL without introducing fundamentally new insights. (See Question 1)\n2. While the paper spans a broad range of topics, it does not deeply investigate why NFs are particularly suited for offline RL and GCRL. Many generative models (e.g., VAEs, EBMs, and diffusion models) also support MLE and VI, as shown in Table 1. A more thorough justification—beyond ease of loss computation or sampling—is needed. (See Question 2)\n3. The proposed approach underperforms on a number of benchmarks. For instance, in IL (Fig. 3), NF-BC is inferior to VQ-BeT on the Kitchen and Multi-modal Ant tasks, while also requiring twice as many learnable parameters. In the offline RL experiments (Fig. 6), NF-RLBC is outperformed by FQL on tasks such as `antsoccer-arena-nav`, `humanoidmaze-medium-nav`, and `cube-single-play`.\n4. Minor technical inaccuracy: Line 29 claims that training diffusion models requires solving an ODE. However, standard diffusion models are trained using denoising score matching, which is simulation-free and computationally efficient.", "was_revised": true, "final_justification": "My questions were addressed. I increased my score to a borderline accept.", "rebuttal": "We thank the reviewer for their insightful feedback and constructive comments. The reviewer’s main feedback is:\n\n(1) regarding the conceptual overlap and comparisons with prior work.\n\n(2) a deeper investigation of why NFs are particularly suited for RL.\n\n(3) analysis of hyperparameter sensitivity and broader implementation complexity.\n\nFor (1), we agree that there are prior papers that apply NFs to RL. Our paper is conceptually different in ways we highlight below. We have revised the related work section and introduction to better explain and clarify these important conceptual differences. Moreover, we have added new experiments comparing the architecture of SAC-LSP [1] with the architecture proposed in our paper.\n\nFor (2), we have added new ablation experiments, in benchmark tasks and a didactic task. These experiments show that expressivity and unbiased VI are the primary properties of NFs which make them particularly suited for RL algorithms.\n\nFor (3), we add comparisons for the number of model parameters (complexity), number of training objectives (end to end vs step wise) with baselines, and provide plots for hyperparameters sensitivity of our method to give a holistic view of implementation simplicity for practitioners.\n\n**Do these answers, revisions and experiments address all the reviewer's concerns?**\n\n> Comparison with the architecture of [1], on conditional IL and offline RL\n\n**Success rates (higher is better):**\n| Method  | scene-play | cube-single-play |\n|---------|---------|----------|\n| NF-GCBC      | 0.42 | 0.795 |\n| GCBC-LSP | 0.02 |  0.04 |\n\n| Method  | antmaze-large-navigate | antsoccer-arena-navigate | cube-single-play | puzzle-3x3-play |  scene-play | humanoidmaze-medium-navigate |\n|---------|---------|----------|---------|----------|----------| ----------|\n| SAC-NF      | 0.76 | 0.6 | 0.440 | 0.996 | 1.0 | 0.03 |\n| SAC-LSP | 0.86 | 0.43 | 0.03 | 0.26 | 0.49 | 0.0 |\n\nSAC-LSP uses the RealNVP architecture, which we find to be sometimes unstable to train (exploding likelihoods). **These results show our architecture is more stable and performant.**\n\n> conceptual overlap with [2–6]\n\nWe agree that these papers use different versions of NFs in the context of RL, but other than that there are distinct conceptual / methodological differences. We first emphasize the main argument of our paper -- because NFs do MLE and VI efficiently, using them as plug and play models in various RL algorithms can reap significant performance gains, rivaling and even surpassing SOTA / complex RL algorithms. We will now note how our argument is conceptually different than the papers pointed out by the reviewer:\n\n[2] use a high level policy ( $p(z \\mid s)$ ) and an NF low level policy ( $p(a \\mid z)$ ). Hence they can only obtain a lower bound on the true log-likelihood of the policy $E_{p(z \\mid s)} \\log [ \\pi(a \\mid z) ] \\leq \\log E_{p(z \\mid s)} [ \\pi(a \\mid z) ] = \\log \\pi(a \\mid z) $ (see 5th line of sec 6 or sec 6.1 where they allude to this).\n\n[3] Use a special kind of NF called Energy-Based Normalizing Flows, this puts constraints on the NF (state independent part of jacobian det should be equal to the value function, and for the state independent part be equal to the Q function). These constraints don't satisfy in general and hence require new techniques like LRS or SCDQ (see sec 3.2 of [3]).\n\n[4] Like [2], they also train a high level policy followed by an NF low level policy. Moreover, they train the NF using weighted regression and introduce new techniques for conservatism (see sec 3.1). This is in contrast to our approach of doing VI directly using NFs.\n\n[5] Does not use NFs directly for RL, but only uses them for uncertainty estimation (sec 4.2). Both the architectures introduced and the tasks considered in our papers are very different.\n\n[6] The main reason they use an NF is because they require an invertible network. Moreover their algorithm is a skill learning algorithm, again optimizing an evidence lower bound. This is in contrast to our argument of doing VI directly using NFs. Additionally, a comparison between a hierarchical algorithm with our flat algorithm wouldn't be fair.\n\n**Our central insight is that NFs unify MLE and VI objectives in RL while being expressive, stable, and fast to train. This allows them to serve as plug-and-play models without needing additional priors or multi-stage training. This is in contrast to prior works that use NFs in nested or constrained forms ([2-4,6]). By showing this across diverse RL settings, we believe our work repositions NFs as a general tool in the RL toolbox, not just a modeling novelty.**\n\n> deeper investigation why NFs are particularly suited for RL\n\nThe main reasons why NFs perform so well is their expressivity and unbiased VI (see the next 4 goal experiment). Almost all datasets we use have multi-modal and noisy distributions. We compare the validation likelihoods of the policies learnt by GCBC and NF-GCBC on three tasks:\n\n**Log-likelihoods (higher is better):**\n| Method | antmaze-large-navigate | scene-play | cube-single-play |\n|---------|----------|----------|----------| \n| NF -GCBC     | -0.673   | 2.975 | 2.756 |\n| GCBC | -7.662 | -4.65 |  -4.661 |\n\nNF-GCBC is able to model the behavior policies more accurately. It is well established that higher-likelihoods are correlated with performance and generalisation in imitation learning [1,2].\n\nWe have also added new experiments to show the contribution of the components we use:\n\n**Success rates (higher is better):**\n| Method | antmaze-large-navigate | scene-play | cube-single-play |\n|---------|----------|----------|----------|\n| NF -GCBC     | 0.32   | 0.42 | 0.795 |\n| without permutation flows | 0.26 | 0.04 |  0.0 |\n| without LayerNorm | 0.29 | 0.3 |  0.0 |\n| without sampling trick | 0.3 | 0.42 |  0.69 |\n\nThese experiments show that our architectural modifications are important for performance of NFs in the context of RL. We hope that the above experiments serve as sufficient motivation for our design choices.\n\n> illustrate the benefits of NFs over other continuous generative models (e.g., VAEs, EBMs, and diffusion)\n\nIn our main paper, we compare and show benefits over SOTA algorithms using VAEs, EBMs and diffusion. Fig 3 compares with diffusion policy / VQ-BET, Fig 5 with FM-GCBC (flow matching), Fig 6 with FQL, IFQL, IDQL, all use flow / diffusion models, Fig 6 compares with CRL (contrastive energy based models). **Of all the 85 tasks, our proposed algorithms exceed or match in 65 and is amongst the top two in seven tasks**. As requested by the reviewer, we have further added a didactic 4 goal environment (sec 4.1 in [3]):\n\n**Goal reaching distribution:**\n| Method | goal 1 | goal 2 | goal 3 | goal 4 |\n|---------|----------|----------|----------|----------| \n| SAC     | 0   | 0 | 0 | 20 |\n| SAC-VAE | 0 | 0 |  0 | 20 | \n| SAC-NF | 8 | 2 |  7 | 3 |\n\nSAC-NF is able to reach all 4 goals, whereas SAC-VAE, which uses a VAE state encoder and a gaussian action decoder still shows mode-seeking behavior, a problem faced by reverse KL optimization. This defect of SAC-VAE can also be caused because only a lower bound on the policy entropy is estimated.\n\n> deeper analysis of simplicity\n\nWe agree that there are other factors which determine the simplicity of implementing an algorithm. We removed this plot from our revised paper and now also compare other details like inference and training speed, number of parameters and training objectives, to provide a more holistic view for RL practitioners.\n\n| Method | Inference Speed | Training Speed (per batch) | Number of parameters (millions) | Number of objectives |\n|---------|---------|---------|---------|---------|\n| NF-GCBC      | 0.127 | 0.023 | 10 M | 1; end to end |\n| VQ-BET | 0.273 | 0.021 | 4.2 M | 2; pretrain a hierarchical VQ-VAE |\n| Diffusion Policy | 1.28 | 0.019 | 65.7 M | 1; end to end |\n| TarFlow | 0.273 | 0.042 | 15M | 1; end to end |\n\nNF-BC, while being competitive on all tasks, is 2, 10 times faster than VQ-BET and Diffusion policy. Unlike VQ-BET, which requires pretraining a VQ-VAE, NF-GCBC is an end to end algorithm directly optimizing for the desired likelihood. Our updated claim is that NF-BC is faster in inference, easier to implement (end to end optimization, contains 2-3 times fewer hyperparameters) and achieves competitive performance across all tasks (best success rate on two tasks and in top 2 methods on the other two tasks).\n\nWe also present ablations demonstrating the robustness of NF-BC to key hyperparameters:\n\n**Success rates (higher is better):**\n| noise_std  | Push-T |\n|---------|----------|\n| 0.5     | 0.77   |\n| 0.3 | 0.79 |\n| 0.1 | 0.827|\n| 0.05 | 0.84 |\n| 0.03 | 0.873 |\n| 0.01 | 0.83 |\n| 0.0 | NAN |\n\n| number of NF blocks  | Push-T |\n|---------|----------|\n| 6   | 0.77   |\n| 14 | 0.79 |\n| 28 | 0.859 |\n| 32 | 0.85 |\n| 48  | 0.87 |\n| 54  | 0.85 |\n\nThe above results show that NF-BC is robust to many values.\n\n> underperforms on a number of benchmarks\n\nOur methods exceed or match *all other baselines* on 76% (65 / 85) of tasks and is amongst the top two in seven tasks. We revised the paper to clarify that there are some small fraction of tasks where our method does not outperform the best prior methods.\n\n> comparing NF-GCBC and GCIQL\n\nDue to space, please refer to our Table 7 and OGbench (pg. 41) for GCIQL comparisons. \n\n> NFs are continuous generative models\n\nIn our paper, we only focus on continuous spaces. But using NFs for large discrete spaces in RL is an exciting direction (Discrete Flows; Tran et al.). We have added this in the limitation and future work section.\n\n> Focus on online reward-based RL\n\nDue to orthogonal challenges like exploration, we chose mostly offline experiments. The toy goal reaching task described above is an online RL task, where NFs indeed perform well. Moreover, we also have experiments on online GCRL.\n\n> simulation free training\n\nThank you for pointing this out, we have updated our paper to remove this."}]

Author rebuttal:

Meta-review: After a robust discussion, the reviewers agree that the paper makes a valuable and novel contribution, both by providing clear conceptual frameworks for categorizing RL algorithms and by exploring a rarely used model class for reinforcement learning agents. The authors should attend carefully to points that tripped up the reviewers, particularly overemphasis on the centrality of value iteration in RL, so that these questions can be anticipated and addressed for other readers.
Decision: Accept (poster)