REVIEW 3 major objections 6 minor 10 references
Offline Reinforcement Learning with Wasserstein Regularization via Optimal Transport Maps
T0 review · 3 major / 6 minor · reviewed 2026-08-06 · deepseek-v4-flash
Pith's one-line read By modeling the policy as the gradient of an input-convex network, Q-DOT regularizes offline RL with the 2-Wasserstein distance and no adversarial critic.
desk verdict A clean OT-map idea for offline RL that deserves peer review, but the empirical headline is undercut by test-set hyperparameter selection. read the letter →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
What carries the argument
The load-bearing object is the gradient map $\nabla_a\psi_\omega(s,a)$ of an input-convex neural network (ICNN), which is convex in the action $a$. The optimal transport theorem says that for an absolutely continuous source distribution, the unique optimal transport map to any target is the gradient of a convex potential; the paper uses this to define the learned policy as the push-forward $\pi_\psi = \nabla_a\psi_\omega \sharp \pi_\mathcal{D}$ and to evaluate the squared 2-Wasserstein distance as the expectation in Eq. (5). The same potential then appears in the policy objective Eq. (6), an explicit stochastic policy $\pi_\rho$ is trained from transported actions by advantage-weighted regression, and the value function is learned in-sample by implicit Q-learning.
What would settle it
Compute the right-hand side of Eq. (5) on a dataset with a clearly discontinuous or multi-modal behavior policy, such as actions drawn from two well-separated Gaussians per state, and compare it with a numerical ground-truth squared 2-Wasserstein distance obtained by linear programming or Sinkhorn; if the two disagree substantially, or if Q-DOT's returns degrade sharply against a variant using the ground-truth distance, the central claim fails.
Extended reading notes
Core claim
The central claim is that the objective $J_\psi(\omega) = \mathbb{E}_{(s,a)\sim\mathcal{D}}[\hat{Q}_\theta(s, \nabla_a\psi_\omega(s,a)) - \alpha\|a - \nabla_a\psi_\omega(s,a)\|_2^2]$ is a principled Wasserstein-regularized offline RL objective, because when the dataset behavior policy is absolutely continuous, the classical optimal transport theorem identifies $\nabla_a\psi_\omega$ as the unique optimal transport map and Eq. (5) equals the squared 2-Wasserstein distance $W_2^2(\pi_\psi, \pi_\mathcal{D})$. The authors combine this input-convex network transport map with implicit Q-learning for in-sample value estimation and advantage-weighted regression for an explicit sampling policy, yielding Q-DOT. They argue that this discriminator-free construction avoids the instability of dual-form Wasserstein regularization, where a Lipschitz-constrained discriminator must be trained adversarially; their reproduced adversarial baseline (AdvW) failed on many tasks, while Q-DOT achieved the best or comparable total score in every D4RL domain. The paper also reports that the learned transport displaces actions from low-reward trajectories more than from high-reward ones, indicating the regularizer preferentially modifies poor behavior.
Load-bearing premise
The method's central identity assumes the dataset behavior policy is a smooth, absolutely continuous distribution; with a finite dataset it is an empirical distribution, so Eq. (5) is not exactly the 2-Wasserstein distance.
Editorial extensions
If this is right
- Q-DOT turns Wasserstein policy regularization into a least-squares regression problem on dataset actions, so no discriminator or Lipschitz constraint is needed.
- The regularizer reaches scores comparable to or better than TD3+BC, CQL, IQL, and DT on the D4RL benchmark, with the largest margins on hopper-medium-v2 (+10.4 over the best baseline) and kitchen-partial-v0 (+21.5).
- Because the transport map starts from the identity and the hyperparameter $\alpha$ controls deviation, the method interpolates between behavior cloning and return maximization, letting practitioners tune conservatism continuously.
- The transport-distance analysis indicates the ICNN mapping mainly moves actions from low-return trajectories, which supports the interpretation that the regularizer repairs poor behavior rather than perturbing good behavior.
- Adversarial Wasserstein regularization, even with in-sample value learning and large regularization weights, scores far below Q-DOT on expert datasets, which the paper reads as evidence that discriminator-free modeling is necessary for Wasserstein regularization in offline RL.
Reading between the lines
- The paper's exactness claim for Eq. (5) relies on the behavior policy being absolutely continuous; with a finite dataset the empirical behavior distribution is discrete, so a practical implementation computes an approximate Wasserstein term, and smoothing the empirical measure or using entropic optimal transport could make the claim rigorously testable.
- Because transported actions are drawn only from the dataset's action support, Q-DOT inherits a conservatism that may limit exploration; combining the ICNN transport with explicit out-of-distribution sampling or model-based rollouts could widen the method's reach.
- The mechanism is not tied to implicit Q-learning: any in-sample value learner that supplies an advantage signal could be paired with the ICNN transport regularizer, so the architecture could serve as a generic plug-in policy regularizer for offline RL.
- The observed correlation between low trajectory reward and large transport distance suggests an implicit data-quality weighting; an ablation with intentionally mislabeled rewards could test whether the effect is causal and whether it drives the benchmark gains.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper proposes Q-DOT, an offline reinforcement learning algorithm that regularizes the learned policy with the squared 2-Wasserstein distance between the policy and the dataset behavior policy. Instead of using a discriminator, Q-DOT models the optimal transport map as the gradient of an input-convex neural network (ICNN), justified by Brenier's theorem. The value function is learned via Implicit Q-Learning (IQL), and the policy is trained with Advantage Weighted Regression (AWR) on actions transported by the ICNN gradient. The method is evaluated on D4RL locomotion, antmaze, and kitchen domains, reporting total returns that are often competitive with or better than existing baselines, and an adversarial Wasserstein variant (AdvW) is shown to be less stable. The paper also provides an analysis of transport distance versus trajectory quality.
Significance. If the empirical claims hold, the paper makes a useful contribution by demonstrating that Wasserstein regularization can be implemented in offline RL without adversarial training, leveraging ICNN-based optimal transport maps. The theoretical derivation is largely correct in the idealized absolutely-continuous setting, and the connection between Brenier's theorem and policy regularization is novel. The paper also provides a clear comparison with an adversarial baseline and a qualitative analysis of transport behavior. However, the empirical support is weakened by the hyperparameter selection protocol and the lack of error bars for baselines, which limits the strength of the central 'comparable or superior performance' claim.
major comments (3)
- [Section 4.1 and Supplementary Section 7] Hyperparameters for Q-DOT are selected on the same D4RL benchmark used for reporting: α is chosen from (1, 5, 10, 20, 100, 400) and β from (0.5, 3, 10, 20), with per-domain choices (20, 3) for locomotion, (20, 20) for antmaze, and (400, 0.5) for kitchen. Figure 2 shows that α substantially affects returns, so selecting α on the test tasks can inflate the reported scores. Baselines, by contrast, are taken from Kostrikov et al. (2022) with their own fixed hyperparameters, making the comparison uncontrolled. To support the central empirical claim, please either fix hyperparameters a priori, use a validation split, or report performance across the hyperparameter sweep to demonstrate robustness.
- [Section 3.1, Eq. (5)] The statement that the expectation E_{a~πD}[||a - ∇aψω(s,a)||^2] equals the exact squared 2-Wasserstein distance W2^2(πψ, πD) relies on πD being absolutely continuous with respect to Lebesgue measure. In practice, πD is the empirical distribution of a finite dataset, which is discrete, so Brenier's theorem does not directly apply and Eq. (5) is only an approximation. The Introduction and Section 3.1 claim that 'the exact Wasserstein distance is consistently computed,' which is not justified. Please either relax this claim to 'approximate' and discuss the finite-sample discrepancy, or provide a formal justification under appropriate assumptions.
- [Table 1] Baselines are reported as point estimates without error bars or number of seeds, while Q-DOT is reported with standard error over six seeds. This asymmetry makes it difficult to assess whether the differences are meaningful. For example, halfcheetah-medium-v2 shows Q-DOT 47.9±0.1 vs IQL 47.4, and hopper-medium-expert-v2 shows Q-DOT 93.1±13.0 vs CQL 105.4. Please provide error bars for the baselines (e.g., re-run them under the same protocol or cite published error bars) and, if possible, include a statistical comparison to support the 'comparable or superior' claim.
minor comments (6)
- [Section 2.2, Eq. (1)] There is an extra closing parenthesis in the term D(π(·|s)∥πD(·|s))); please remove it.
- [Section 4.1] The text says 'refer to the scores reported in Kostrikov et al. (2022)' but does not specify which of the seven baselines were taken from that source and whether all were; please clarify the provenance of each baseline.
- [Figure 2] The x-axis of Figure 2 ranges only up to 20, but the hyperparameter sweep in Supplementary Section 7 includes α values of 100 and 400; please either show the full range or explain why those values are omitted from the figure.
- [Section 4.3 / Figure 1] The three panels in Figure 1 are all Hopper variants (medium-v2, medium-expert-v2, medium-replay-v2), so describing them as 'three tasks' is misleading; please say 'three Hopper datasets' or include a more diverse set of environments.
- [References] In the Kingma & Ba (2015) reference, 'San Diega' should be corrected to 'San Diego'.
- [Section 3.1, last paragraph] The limitation 'cannot sample actions for state-action pairs that are not present in the dataset' is imprecise; since actions are generated from dataset actions conditioned on state, the limitation is for states not present in the dataset, not state-action pairs. Please rephrase for clarity.
Circularity Check
No significant circularity: the Wasserstein penalty is backed by Brenier's external theorem, and the method's components come from independent prior work.
full rationale
The paper's derivation chain is not circular. The central identity, Eq. (5), is not fitted from data; it follows from Brenier's theorem, an external mathematical result: for an absolutely continuous dataset policy pi_D and a convex psi with pi_psi = grad_psi # pi_D, one has W2^2(pi_psi, pi_D) = E_{a ~ pi_D} ||a - grad_psi(a)||^2. The paper explicitly states this assumption ('assuming that pi_D is absolutely continuous'). Eq. (6) is the direct plug-in objective used to train the transport map. The final evaluation policy pi_rho is then trained separately via AWR (Eq. (7)) from samples of pi_psi, which is a distillation step rather than a renaming of the objective. The value function is learned with IQL from Kostrikov et al. (2022), an external method, and the ICNN transport-map construction is cited to Makkuva et al. (2020) and Korotin et al., also external prior work. No load-bearing claim rests on a self-citation, and no fitted parameter is later relabeled as a prediction. Two concerns are real but are not circularity: (i) in the finite-dataset implementation, pi_D is empirical rather than absolutely continuous, so Eq. (5) is an approximation, not an exact Wasserstein distance; and (ii) the supplementary material describes per-domain selection of alpha and beta on the D4RL benchmark without a validation split, which may inflate the reported scores. These are approximation and evaluation-protocol issues, not reductions of the derivation to its own inputs, so the appropriate finding is no significant circularity.
Assumptions & free parameters
free parameters (2)
- alpha =
20 (MuJoCo locomotion), 20 (AntMaze), 202 (Kitchen)
- beta =
3 (locomotion), 20 (AntMaze), 0.5 (Kitchen)
assumptions (4)
- standard math Brenier's theorem: for P absolutely continuous, the optimal transport map is the gradient of a convex function
- domain assumption The behavior policy pi_D is absolutely continuous with respect to Lebesgue measure
- domain assumption ICNN parameterization preserves convexity of psi
- domain assumption Empirical expectations over D approximate true expectations
Cite this review
Pith. "Pith review of Offline Reinforcement Learning with Wasserstein Regularization via Optimal Transport Maps." pith.science (2026). https://pith.science/paper/CY2767GP
@misc{pith2026250710843,
author = {Pith},
title = {Pith review of: Offline Reinforcement Learning with Wasserstein Regularization via Optimal Transport Maps},
year = {2026},
howpublished = {\url{https://pith.science/paper/CY2767GP}},
note = {Machine review of arXiv:2507.10843}
}
abstract
Offline reinforcement learning (RL) aims to learn an optimal policy from a static dataset, making it particularly valuable in scenarios where data collection is costly, such as robotics. A major challenge in offline RL is distributional shift, where the learned policy deviates from the dataset distribution, potentially leading to unreliable out-of-distribution actions. To mitigate this issue, regularization techniques have been employed. While many existing methods utilize density ratio-based measures, such as the $f$-divergence, for regularization, we propose an approach that utilizes the Wasserstein distance, which is robust to out-of-distribution data and captures the similarity between actions. Our method employs input-convex neural networks (ICNNs) to model optimal transport maps, enabling the computation of the Wasserstein distance in a discriminator-free manner, thereby avoiding adversarial training and ensuring stable learning. Our approach demonstrates comparable or superior performance to widely used existing methods on the D4RL benchmark dataset. The code is available at https://github.com/motokiomura/Q-DOT .
Figures
Reference graph
Works this paper leans on
-
[2]
Other implementation details follow Kostrikov et al. (2022). 2.5 5.0 7.5 10.0 12.5 15.0 17.5 20.0 50 55 60 65 70 75 80 85Return medium-v2 2.5 5.0 7.5 10.0 12.5 15.0 17.5 20.0 40 50 60 70 80 90 100Return medium-replay-v2 2.5 5.0 7.5 10.0 12.5 15.0 17.5 20.0 0 20 40 60 80 100Return medium-expert-v2 halfcheetah hopper walker2d Figure 2: The average returns f...
work page 2022
-
[3]
In contrast, this tendency was not as clearly observed in the HalfCheetah environment
In the Walker2d environment, similar to the Hopper environment, the transport distance was larger for lower-quality trajectories. In contrast, this tendency was not as clearly observed in the HalfCheetah environment. A smaller transport distance indicates that the transport that increases the advantage is not being identified by the value function. Thus, ...
work page 2000
-
[6]
Yicheng Luo, zhengyao jiang, Samuel Cohen, Edward Grefenstette, and Marc Peter Deisen- roth
URL https://arxiv.org/abs/2005.01643. Yicheng Luo, zhengyao jiang, Samuel Cohen, Edward Grefenstette, and Marc Peter Deisen- roth. Optimal transport for offline imitation learning. In The Eleventh International Confer- ence on Learning Representations ,
arXiv 2005
-
[8]
Ashvin Nair, Abhishek Gupta, Murtaza Dalal, and Sergey Levine
URL https://arxiv.org/abs/2001.01866. Ashvin Nair, Abhishek Gupta, Murtaza Dalal, and Sergey Levine. Awac: Accelerating online reinforcement learning with offline datasets. arXiv preprint arXiv:2006.09359 ,
arXiv 2001
-
[9]
Harshit Sikchi, Qinqing Zheng, Amy Zhang, and Scott Niekum
URL https://arxiv.org/abs/2006.09359. Harshit Sikchi, Qinqing Zheng, Amy Zhang, and Scott Niekum. Dual RL: Unification and new methods for reinforcement and imitation learning. In The Twelfth International Confer- ence on Learning Representations ,
arXiv 2006
-
[2015]
Wasserstein-2 generative networks
Alexander Korotin, Vage Egiazarian, Arip Asadulaev, Alexander Safin, and Evgeny Burnaev. Wasserstein-2 generative networks. In International Conference on Learning Representations , 2021a. URL https://openreview.net/forum?id=bEoxzW_EXsa. Reinforcement Learning Journal 2025 Alexander Korotin, Lingxiao Li, Justin Solomon, and Evgeny Burnaev. Continuous wass...
work page 2025
-
[2019]
URL https://arxiv.org/abs/1911.11361. Haoran Xu, Li Jiang, Jianxiong Li, Zhuoran Yang, Zhaoran Wang, Victor Wai Kin Chan, and Xi- anyuan Zhan. Offline RL with no OOD actions: In-sample learning via implicit value regular- ization. In The Eleventh International Conference on Learning Representations ,
arXiv 1911
-
[2020]
Scott Fujimoto and Shixiang Gu
URL https: //arxiv.org/abs/2004.07219. Scott Fujimoto and Shixiang Gu. A minimalist approach to offline reinforcement learning. In A. Beygelzimer, Y . Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural In- formation Processing Systems,
arXiv 2004
Show all 10 references
-
[2021]
Optimal transport tools (ott): A jax toolbox for all things wasserstein
Marco Cuturi, Laetitia Meng-Papaxanthos, Yingtao Tian, Charlotte Bunne, Geoff Davis, and Olivier Teboul. Optimal transport tools (ott): A jax toolbox for all things wasserstein. arXiv preprint arXiv:2201.12324,
-
[2024]
7 Experimental Details In AdvW and Q-DOT, the actor, critic, discriminator (for AdvW), and ICNN (for Q-DOT) are all two-layer MLPs with ReLU activations and 256 hidden units
Reinforcement Learning Journal 2025 Supplementary Materials The following content was not necessarily subject to peer review. 7 Experimental Details In AdvW and Q-DOT, the actor, critic, discriminator (for AdvW), and ICNN (for Q-DOT) are all two-layer MLPs with ReLU activation...
2019
Reviewed August 6, 2026 · model on record in the stance chip above.
Discussion (0). Continue with ORCID to comment.