Normalizing Flows are Capable Models for Continuous Control

We thank the reviewer for their insightful feedback and constructive comments. The reviewer’s main feedback is:

(1) regarding the conceptual overlap and comparisons with prior work.

(2) a deeper investigation of why NFs are particularly suited for RL.

(3) analysis of hyperparameter sensitivity and broader implementation complexity.

For (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.

For (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.

For (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.

Do these answers, revisions and experiments address all the reviewer's concerns?

Comparison with the architecture of [1], on conditional IL and offline RL

Success rates (higher is better):

SAC-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.

conceptual overlap with [2–6]

We 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:

[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).

[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]).

[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.

[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.

[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.

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.

deeper investigation why NFs are particularly suited for RL

The 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:

Log-likelihoods (higher is better):

NF-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].

We have also added new experiments to show the contribution of the components we use:

Success rates (higher is better):

These 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.

illustrate the benefits of NFs over other continuous generative models (e.g., VAEs, EBMs, and diffusion)

In 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]):

Goal reaching distribution:

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, 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.

deeper analysis of simplicity

We 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.

NF-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).

We also present ablations demonstrating the robustness of NF-BC to key hyperparameters:

Success rates (higher is better):

The above results show that NF-BC is robust to many values.

underperforms on a number of benchmarks

Our 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.

comparing NF-GCBC and GCIQL

Due to space, please refer to our Table 7 and OGbench (pg. 41) for GCIQL comparisons.

NFs are continuous generative models

In 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.

Focus on online reward-based RL

Due 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.

simulation free training

Thank you for pointing this out, we have updated our paper to remove this.

We thank the reviewer for engaging and providing thoughtful feedback.

The current title appears to overstate the generality of the method

We will update our title. More suitable titles we were thinking of are -

1. Plug-and-Play Normalizing Flows for Continuous Control.
2. Normalizing Flows are Capable Models for Continuous Control.

Do one of these titles better reflect the scope of our experiments?

how NFs influence the dynamics of the agent?

Currently, we have concrete evidence that expressiveness is an important property provided by NFs (see rebuttal above). We include below, some new results which hint towards surprising ways in which NFs affect the dynamics of RL. We believe the new results provide useful starting points to analyze how NFs affect learning dynamics.

NF policies are able to efficiently maximize complex action-value optimization landscapes

In the puzzle task, none of the baselines achieve better than 30% success rate (see Fig 6). Because solving the puzzle requires hitting specific buttons, we hypothesize that the Q function has narrow modes of good actions and a large number of bad actions.

We compare with NF-AWR which optimizes an NF policy with a sampling based optimization method AWR [4]. This does not work well despite using the same expressive NFs.

NF policies are able to maximize Q-values faster than gaussian policies in online RL

1. Validation mean value of $Q_{SAC-NF} - Q_{SAC}$ : $-0.35, -0.01, 0.20, 0.04, 0.01, 0.24, 0.32, 0.07, 0.19, 0.36, 0.26$
2. Validation mean value of $(\sum_{i=t}^{H} \gamma^{i} r_i - Q_{SAC-NF}) / (\sum_{i=t}^{H} \gamma^{i} r_i)$ : $0.09$

The first point shows that Q values learnt by SAC-NF are higher, indicating that the NF policy is better able to optimize the Q-function. Importantly this is not caused by overestimation of the true returns. As indicated by the second point, the Q values of SAC-NF slightly underestimate the Q values on average (by 9%).

The above two results show that apart from expressiveness, the ability of NFs to perform direct gradient based optimization is crucial in many tasks.

NF-GCBC's stitching performance improves with training-steps

Unlike TD methods, GCBC does not have an explicit mechanism for stitching. But surprisingly, we find the NF-GCBC's performance scales with training steps on tasks designed to explicitly test for stitching. A deeper investigation of this phenomena is an exciting future direction.

NF for VI and MLE is well established & the idea that IL and CGRL can be modeled using NF is relatively straightforward

Even if certain ideas may appear straightforward in hindsight, there is substantial value in rigorously demonstrating their effectiveness in practice. While our experiments are based on simple ideas—which we consider a strength—the fact remains that NFs are still underutilized in RL. In contrast, diffusion models / transformers are widely used in IL [1–3], and contrastive / MLP models remain the norm for GCRL. By actually showing that NFs can often outperform these models and providing open-sourced implementations (in jax and torch), we are offering value for the community to build on.

deeper connection between RL and NF

One conclusion we would like the readers of this paper to make (see line 72-74) is about rethinking how one should make RL algorithms better. Let's look at NLP. For a long time, there was lots of discussion about the right objectives in NLP: some objectives were good for sentiment classification, others for analogy making, or learning representations. However, after transformers, the field has shifted somewhat to stick with one objective (next token prediction) and instead focus on optimizing the architecture (transformers variants, mamba, etc). Our paper aims to take a first step towards doing this for RL. Of course, this paper is just a start: we're just looking at 4 problem settings. We don't claim that NFs are the best possible architecture -- it seems rather plausible that some yet-to-be-invented architecture will work even better! But the paper will hopefully help shift perceptions and show that there's this other important avenue for improving performance.

We have included all the new results / text updates during the rebuttal in our paper, and we thank the reviewer for helping us do that! We hope that they address reviewer's concerns and assessment of the paper

[1] π0: A Vision-Language-Action Flow Model for General Robot Control

[2] Behavior Generation with Latent Actions

[3] Q-Transformer: Scalable Offline Reinforcement Learning via Autoregressive Q-Functions

[4] Advantage-Weighted Regression: Simple and Scalable Off-Policy Reinforcement Learning

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.

To 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.

To 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.

Do these answers, revisions and experiments address all the reviewer's concerns?

motivating design decisions

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 [1] an old NF architecture and TarFlow [2] a SOTA NF architecture in computer vision.

Success rates (higher is better):

Log-likelihoods (higher is better):

From the above results, it is evident that

1. 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.
2. 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.

We also add new experiments ablating the inclusion of LayerNorm, a generalized permutation layer, and the sampling trick:

Success rates (higher is better):

These 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.

I would find a more faithful plot to be one that simply reports that number directly, rather than relabelling it "simplicity".

We 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.

NF-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).

Additionally, 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:

Success rates (higher is better):

Success rates (higher is better):

Success rates (higher is better):

The 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.

Related work section

We have added a detailed related work section in the appendix of the paper.

Typos / Writing Suggestions

We appreciate all of these suggestions and have updated the paper to incorporate them.

[1] Estimating the Dimension of a Model

I thank the authors for their extensive and thoughtful reply, both to my review and the other reviewers. I have read the other reviews, and maintain my positive assessment of the work.

In particular I believe reviewer oKbR raised several excellent questions I would like to comment on -- first, regarding whether RL can be well-thought of as primarily MLE and VI. I agree with this question completely, and also believe that the authors' response does a good job of addressing this point. All of RL is not reducible to these two problems, so making it very clear what the scope of the proposed use of NFs in the context of RL is will sharpen the work considerably. The proposed changed paragraph accounts for this (though adding slight changes to the text throughout the paper can also help strengthen the paper).

Second, clarifying why exactly NFs work well in RL. The rebuttal indicates the reasoning is due to expressivity ("The main reason why NFs perform so well is their expressivity."). While this is an important part of the story, it remains unconvincing that it is the sole reason NFs perform well. That being said, my take is that this question does not need to be fully addressed in the present paper; it is sufficient to show that NFs can act as generally capable models for a variety of sub-problems within RL, which I believe the paper does, and to leave some of this explanation as room for future development (even if expressivity does yield a partial answer). Giving a completely comprehensive explanation to such a question is likely a longer-term research program. I am happy to discuss this point further.

Lastly, regarding my point about motivating the design decisions more thoroughly: the added details to the plots (especially the new plot taking over for the simplicity plot, described at "We now compare other details of the algorithms") offer compelling richness to the design space. I suggest including additional text in the paper to reflect this design space, and to leave readers with a sense of why certain design choices are made (as suggested in the rebuttal).

Overall I am happy with the proposed changes, and I believe the paper should be accepted, barring any other major concerns from other reviewers that I may have missed. I am happy to discuss with the other reviewers about these points.

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.

To address (1), we have run new ablation experiments comparing the performance, speed and stability with two other popular NF architectures.

For (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.

For (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.

Do these answers, revisions and experiments address all the reviewer's concerns?

analysis of the motivations for choosing coupling networks as reversible transformations

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 [1] (an early NF architecture) and TarFlow [2] (a SOTA NF architecture in computer vision).

Success rates (higher is better):

Log-likelihoods (higher is better):

From the above results, it is evident that

1. 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.
2. 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.

What are the differences between the NF architecture chosen in this paper and those of NICE and RealNVP?

Our 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).

Why was the coupling network chosen?

Like 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).

steps of invertible transformations are also an important hyperparameter in NFs

Indeed, 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.

Success rates (higher is better):

Success rates (higher is better):

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

We 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.

Measuring the model's simplicity solely by the number of hyper-parameters

We 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.

NF-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).

Additionally, 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:

Success rates (higher is better):

Success rates (higher is better):

Success rates (higher is better):

The 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.

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.

Because 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.

Can both the policy and Q-value learning in these scenarios be modeled using NFs?

In 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.

[1] Density estimation using Real NVP

[2] Normalizing Flows are Capable Generative Models

[3] Estimating the Dimension of a Model

We thank the reviewer for engaging with us and providing thoughtful feedback. It seems the remaining feedback is regarding (1) writing clarity (introduction of NFs and efficiency) and (2) plots showing convergence gradient steps vs performance.

To address (1), we have added new text to clarify the NF introduction and made it inline with how prior works introduce NFs. We also clarify what we mean by efficiency. To address (2), we have presented the desired plots in the rebuttal.

Do these answers and revisions address the reviewer's remaining concerns?

then Eq.(3) and Eq.(5) appear to represent one-step transformations

We agree with the reviewer that the way NFs are currently introduced can cause confusion. We would like to note that Eq. (3) and Eq. (5) are indeed supposed to represent the entire multi-step NF transformations. We have added new text to resolve this confusion :

It is standard practice to stack multiple single step invertible mappings ($f^i_\theta$) to make the entire NF ($f_\theta$), i.e. $f_\theta = f^T_\theta \circ \dots \circ f^2_\theta \circ f^1_\theta$. Each single step mapping is called an NF block. Because each $f^i_\theta$ are invertible, the entire NF $f_\theta$, which is their composition, is also invertible. The log-det of the entire NF can be estimated by summing up the log-dets of all its blocks (chain rule of derivatives):

$\log |\frac{df_\theta(x)}{dx}| = \sum_{i=1}^T \log |\frac{df^i_\theta(x)}{ df^{i-1}_\theta(x) }|,$

where $f^{0}_\theta(x) = x$.

We hope the changed paragraph adds clarity. We would like to note that this is in line with how NFs are introduced in prior works [3-5].

clarify "efficient" sampling

Thank you for pointing this out. By efficiency, we mean computational efficiency, which is the computational speed of sampling and training a normalizing flow. We have clarified this in the paper and have also referenced the speed comparison shared in the previous rebuttal.

provide the corresponding data at specific convergence steps in the rebuttal

Below we provide the total gradient steps and the convergence gradient steps of all experiments. The convergence gradient step is the first gradient when the performance exceeds 95% of the reported performance in the paper.

Post note - For all reported success rates in the main paper, we always set the number of gradient steps as done for baselines in prior work. For example, all results in Fig 5 and Fig 6 show the success rates after exactly 1 million gradient steps. This ensure a fairer protocol for comparisons with baselines.

BC experiments (Figure 3)

The batch size for all experiments is 256 and the dataset sizes. For more details of these datasets see Appendix A.1, page 13 of [1].

GCBC experiments (Figure 5)

The batch size for all experiments is 256. For more details of these datasets see section 7 of [2].

Offline RL experiments (Figure 6)

The batch size for all experiments is 256. For more details of these datasets see section 7 of [2].

sample efficiency does not play a role with a fixed dataset

We agree with the reviewer that sample efficiency matters, especially if data is scarce. We apologize for this remark in the previous rebuttal. We would like to point that the paper does not claim NFs are sample-efficient, nor does it make any claims about the importance of sample efficiency. (see also clarification about efficiency^).

We have included all the new results / text updates suggested during the rebuttal in our paper, and we thank the reviewer for helping us do that! We hope that they address reviewer's concerns and assessment of the paper

[1] Behavior Generation with Latent Actions

[2] OGBENCH: BENCHMARKING OFFLINE GOAL-CONDITIONED RL

[3] Density estimation using Real NVP

[4] Normalizing Flows are Capable Generative Models

[5] Glow: Generative Flow with Invertible 1x1 Convolutions

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.

To 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.

Regarding (2), we have changed our claim to be more specific, and have provided counter-examples (algorithms that don’t fall into this category).

Do these answers, revisions and experiments address all the reviewer's concerns?

a fundamental understanding of why NFs perform well in RL. What are the specific properties of NFs for RL performance contribution?

The 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:

Log-likelihoods (higher is better)

NF-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].

In 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:

Goal coverage (more uniform is better):

Clearly, 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.

The above two experiments suggests that expressiveness and the ability to do VI are the most important property of NFs for RL performance contribution.

How is the performance of other NF architectures evaluated? distinguish the contribution of NF properties and architecture properties

In 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).

Success rates (higher is better):

Log-likelihoods (higher is better):

From the above results, it is evident that

1. 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
2. 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.

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.

The paper's central claim: "RL algorithms are MLE and VI" (Section 4) is too strong"

We 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:

1. BC / GCBC with SVMs or other non probabilistic losses like 0–1 loss, hinge loss.
2. temporal difference learning.
3. State / Action abstraction learning algorithms.
4. Count based exploration algorithms, Planning algorithms like MCTS, algorithms based on game theory (GANs, minmax regret). These 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.

detailed computational cost analysis, such as training/inference time, memory, compared to baselines?

specific RL scenarios where NFs fail?

We have added a list of detailed failure cases of NFs in RL in the appendix. We present an outline below:

1. Without LayerNorm and linear permutation flows, NFs are unstable to train and sometimes result in exploding losses.
2. Without adding noise to the actions (noise_std hyperparameter in table 9), NFs are unstable to train and sometimes result in exploding losses.
3. Training with larger values of $\lambda$ (in equation 10) leads to lower performance.

To 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:

The above two tables suggest that NFs fail without permutation flows and LayerNorm, and performance decreases at very low and very high noise_std values.

[1] Behavior Transformers: Cloning-K modes with one stone

[2] Diffusion Policy: Visuomotor Policy Learning via Action Diffusion

[3] Density estimation using Real NVP

[4] Normalizing Flows are Capable Generative Models

[5] Improving Variational Inference with Inverse Autoregressive Flow

[6] Reinforcement Learning and Control as Probabilistic Inference: Tutorial and Review

[7] What type of inference is planning?

[8] General duality between optimal control and estimation

[9] Robot trajectory optimization using approximate inference

[10] Modeling purposeful adaptive behavior with the principle of maximum causal entropy

[11] Maximum Entropy Reinforcement Learning via Energy-Based Normalizing Flow

I am really thankful for the detailed response, comprehensive explanation, and hard work of the authors. My concerns are almost resolved. But I needed some time for careful consideration about the why question. While authors' comprehensive experiments give a significant explanation, they still lack essential (theoretical analysis or well-designed toy/counter experiment, etc.) understanding about 'why NFs are good?' It's not about expansion (of experimental task, etc.). I think it's about narrowing down/restricting to get a fundamental understanding.

That's my last concern, but I think it's too harsh or too high bar for a newly-discovered sub-area. I believe authors will find and release the answers in camera-ready or future work.

I raise my score to 4. Thanks for the great work.

We thank the reviewer for their appreciation of our work and for helping us improve our paper.