Flow Models for Unbounded and Geometry-Aware Distributional Reinforcement Learning

TL;DR

This paper addresses the limitations of bounded-support and single-moment approaches in distributional RL by modeling full return densities with normalizing flows. It introduces NFDRL, a forward-flow architecture conditioned on state-action pairs that yields unbounded, flexible return distributions, and it trains using a geometry-aware Cramèr-distance surrogate with KDE-based target alignment. Key contributions include a CDF-based 1D flow design with an affine rescaling to support unbounded returns, a principled target distribution construction via composed flows, and a surrogate loss that provides unbiased gradients and contraction properties. Empirically, NFDRL demonstrates expressive, multimodal distributions in toy MDPS and Frozen Lake, and achieves competitive or superior performance on the Atari-5 benchmark while significantly reducing parameter count compared to C51, highlighting practical benefits in parameter efficiency and density-based learning for distributional RL.

Abstract

We introduce a new architecture for Distributional Reinforcement Learning (DistRL) that models return distributions using normalizing flows. This approach enables flexible, unbounded support for return distributions, in contrast to categorical approaches like C51 that rely on fixed or bounded representations. It also offers richer modeling capacity to capture multi-modality, skewness, and tail behavior than quantile based approaches. Our method is significantly more parameter-efficient than categorical approaches. Standard metrics used to train existing models like KL divergence or Wasserstein distance either are scale insensitive or have biased sample gradients, especially when return supports do not overlap. To address this, we propose a novel surrogate for the Cramèr distance, that is geometry-aware and computable directly from the return distribution's PDF, avoiding the costly CDF computation. We test our model on the ATARI-5 sub-benchmark and show that our approach outperforms PDF based models while remaining competitive with quantile based methods.

Figures (13)

• Figure 1: Architecture of the conditional flow model. A neural network $h_\theta$ maps each state $x$ to parameters $\{(w_i^j, \mu_i^j, \sigma_i^j)\}_{i=1}^n$ for each action $a_j\in\mathcal{A}$, defining a mixture CDF $F^{(a_j)}$. This CDF acts as a flow that transforms base noise samples $z_k\sim\mathcal{U}$ into return samples $y_k^j = F^{(a_j)}(z_k)$. The change of variable formula then makes use of the flows derivative to approximate the return distribution $\eta^\pi(y_k^j)$. To estimate the target distribution, the bootsrap function is implemented as a flow and takes the reward and next state as input. It finally outputs target distribution $T^\pi\eta$.
• Figure 2: Learnt return distributions for the state-action pair $(s_2, a_1)$ in $MDP_1$, under different values of the KDE bandwidth and final state's reward variance. The target reward is $0.8$. The top row shows distributions learned using the exact Cramèr loss, while the bottom row shows those obtained with our surrogate. Narrower distributions can be achieved with both losses, illustrating the method's flexibility.
• Figure 3: Learnt distributions for $(s_1,a_1)$ in $MDP_1$; $R_1=-0.8$, $R_2=0.3$; $\gamma=0.9$. Here we use a KDE bandwidth of $0.05$ and a final reward standard error of $0.01$.
• Figure 4: Learnt distributions for $(s_1,a_1)$ in $MDP_2$; $R_1=0.8$, $R_2=0.3$. We use a KDE bandwidth of $0.05$ and a final reward standard error of $0.1$
• Figure 5: Return distributions learnt for the final state of $MDP_3$ using IQN and NFDRL. Right: true quantile function (blue) and quantile function learnt using IQN, reproducing the results of expectile_DRL. Middle: True CDF and CDF approximated for the quantile function obtained from IQN. Left: True PDF, and approximations obtained using IQN and NFDRL. While IQN produces a noisy distribution blurring the models, our method outputs a smoother distribution that makes the modes apparent.