Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Vector-Valued Distributional Reinforcement Learning Policy Evaluation: A Hilbert Space Embedding Approach

Mehrdad Mohammadi, Qi Zheng, Ruoqing Zhu

TL;DR

The paper tackles offline evaluation of policies in high-dimensional, multi-dimensional reinforcement learning settings by representing return distributions through kernel mean embeddings in a reproducing kernel Hilbert space. By employing a kernel IPM based on Ma térn kernels, KE-DRL achieves a contraction-like property for the distributional Bellman operator and derives finite-sample convergence guarantees for the conditional mean embeddings. The method estimates off-policy quantities with nonparametric importance ratios in RKHS, uses a finite return grid to approximate embeddings, and provides a practical algorithm (KE-DRL) with stability-driven optimization via AdamW and fixed-point/mass penalties. Empirically, KE-DRL yields accurate mean embeddings and enables recovery of various distributional statistics in offline multi-dimensional RL tasks, demonstrating robust performance and offering a scalable framework for risk-aware decision making in complex settings.

Abstract

We propose an (offline) multi-dimensional distributional reinforcement learning framework (KE-DRL) that leverages Hilbert space mappings to estimate the kernel mean embedding of the multi-dimensional value distribution under a proposed target policy. In our setting, the state-action variables are multi-dimensional and continuous. By mapping probability measures into a reproducing kernel Hilbert space via kernel mean embeddings, our method replaces Wasserstein metrics with an integral probability metric. This enables efficient estimation in multi-dimensional state-action spaces and reward settings, where direct computation of Wasserstein distances is computationally challenging. Theoretically, we establish contraction properties of the distributional Bellman operator under our proposed metric involving the Matern family of kernels and provide uniform convergence guarantees. Simulations and empirical results demonstrate robust off-policy evaluation and recovery of the kernel mean embedding under mild assumptions, namely, Lipschitz continuity and boundedness of the kernels, highlighting the potential of embedding-based approaches in complex real-world decision-making scenarios and risk evaluation.

Vector-Valued Distributional Reinforcement Learning Policy Evaluation: A Hilbert Space Embedding Approach

TL;DR

The paper tackles offline evaluation of policies in high-dimensional, multi-dimensional reinforcement learning settings by representing return distributions through kernel mean embeddings in a reproducing kernel Hilbert space. By employing a kernel IPM based on Ma térn kernels, KE-DRL achieves a contraction-like property for the distributional Bellman operator and derives finite-sample convergence guarantees for the conditional mean embeddings. The method estimates off-policy quantities with nonparametric importance ratios in RKHS, uses a finite return grid to approximate embeddings, and provides a practical algorithm (KE-DRL) with stability-driven optimization via AdamW and fixed-point/mass penalties. Empirically, KE-DRL yields accurate mean embeddings and enables recovery of various distributional statistics in offline multi-dimensional RL tasks, demonstrating robust performance and offering a scalable framework for risk-aware decision making in complex settings.

Abstract

We propose an (offline) multi-dimensional distributional reinforcement learning framework (KE-DRL) that leverages Hilbert space mappings to estimate the kernel mean embedding of the multi-dimensional value distribution under a proposed target policy. In our setting, the state-action variables are multi-dimensional and continuous. By mapping probability measures into a reproducing kernel Hilbert space via kernel mean embeddings, our method replaces Wasserstein metrics with an integral probability metric. This enables efficient estimation in multi-dimensional state-action spaces and reward settings, where direct computation of Wasserstein distances is computationally challenging. Theoretically, we establish contraction properties of the distributional Bellman operator under our proposed metric involving the Matern family of kernels and provide uniform convergence guarantees. Simulations and empirical results demonstrate robust off-policy evaluation and recovery of the kernel mean embedding under mild assumptions, namely, Lipschitz continuity and boundedness of the kernels, highlighting the potential of embedding-based approaches in complex real-world decision-making scenarios and risk evaluation.
Paper Structure (34 sections, 13 theorems, 118 equations, 8 figures, 5 tables, 2 algorithms)

This paper contains 34 sections, 13 theorems, 118 equations, 8 figures, 5 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Background
  3. Methodology
  4. Reproducing Kernel Hilbert Space for Return Distributions
  5. Kernel IPM for Return Distributions
  6. Off-Policy Evalutaion: Non-Parametric Importance Ratio in RKHS
  7. Theoretical Properties and Bounds
  8. Fixed Point Property
  9. Error Bounds
  10. Evaluation Metric in Offline-RL
  11. Statistics Recovery from the Estimated Embedding
  12. Algorithm and Empirical Implementation
  13. KE-DRL Algorithm
  14. Optimization
  15. Numerical Examples
  16. ...and 19 more sections

Key Result

Theorem 3.2

Assuming (i)${\mathbf R}|{\bf s},{\bf a}$ has an integrable Lebesgue density $f_{{\mathbf R}}(\cdot|{\bf s},{\bf a})$. (ii) For $p^\pi({{\bf s}',{\bf a}'} |{\bf s},{\bf a})$-a.e. , $({{\bf s}',{\bf a}'}), ({\mathbf Z}| {{\bf s}',{\bf a}'})$ has a Lebesgue density $f_{\mathbf Z}(\cdot|{\bf s}',{\bf a

Figures (8)

  • Figure 1: The $\hbox{Mat'{e}rn}$ kernel as a function of distance $d$: (a) Varying smoothness parameter $\nu$ with fixed length scale $\ell = 1.0$. (b) Varying length scale $\ell$ with fixed $\nu = 1.5$. (c) Quotient $\sqrt{2(\sigma^2-k_\nu(d))}/d$ near $d = 0$ for fixed $\ell = 1.0$, whose supremum defines the Lipschitz constant $L_k$ in the kernel embedding. The finite, bounded values confirm kernel Lipschitz continuity (required for $\nu > 1$), which is essential for establishing the Hölder contraction property in Theorem \ref{['thm:fixed_point']}.
  • Figure 2: Behavioral Policy:Uniform-Target Policy:Gaussian, $\hbox{Mat'{e}rn}$ Kernel Parameters $(\nu,\ell,\sigma)=(6.5,2,0.6)$, dims$(S,R,A)=(5,3,1)$, $\lambda_{Reg}=\text{5e-4}$, Penalty $\lambda_{FP}=\text{100}$, Policy evaluated at $({\bf s},{\bf a})=([-1.294,-0.917,0.219,0.283,1.466],[0.434])$
  • Figure 3: Behavioral Policy:Logistic-Target Policy:Gaussian, $\hbox{Mat'{e}rn}$ Kernel Parameters $(\nu,\ell,\sigma)=(7.5,3,0.6)$,dims$(S,R,A)=(5,3,1)$, $\lambda_{Reg}=\text{2e-5}$, Penalty $\lambda_{FP}=\text{200}$, Policy evaluated at $({\bf s},{\bf a})=([-1.638,-0.234,-0.264,-0.671,-0.205],[0.465])$
  • Figure 4: Behavioral Policy:Gaussian - Target Policy: Uniform, $\hbox{Mat'{e}rn}$ Kernel Parameters $(\nu,\ell,\sigma)=(7.5,2,0.6)$,dims$(S,R,A)=(5,3,1)$, $\lambda_{Reg}=\text{1e-4}$, Penalty $\lambda_{FP}=\text{100}$, $Z$-grid expansion factor $= 1.1$, Policy evaluated at $({\bf s},{\bf a})=([1.447, 0.075, -1.163, -0.182, 0.924],[0.81])$
  • Figure 5: Behavioral Policy:Gaussian - Target Policy: Logistic, $\hbox{Mat'{e}rn}$ Kernel Parameters $(\nu,\ell,\sigma)=(6.5,1.5,0.8)$,dims$(S,R,A)=(5,3,1)$, $\lambda_{Reg}=\text{5e-4}$, Penalty $\lambda_{FP}=\text{100}$, $Z$-grid expansion factor $= 1.0$, Policy evaluated at $({\bf s},{\bf a})=([0.51, 0.308, -0.201, 0.734, 0.807],[0.81])$
  • ...and 3 more figures

Theorems & Definitions (26)

  • Remark 3.1
  • Theorem 3.2: Distributional Bellman Operator PDF
  • Corollary 3.3
  • Remark 3.4
  • Theorem 4.1: Equivalence of $\raisebox{.12\baselineskip}{\large$\boldsymbol{\gamma}$}_k$ and $W_1$
  • Theorem 4.2: Fixed Point Property
  • Corollary 4.3: Iteration Complexity
  • Definition 4.1
  • Theorem 4.4: Statistical Pointwise Error Bound
  • Corollary 4.5: Weak Consistency
  • ...and 16 more