REVIEW 4 major objections 4 minor 34 references
Variational OOD State Correction for Offline Reinforcement Learning
T0 review · 4 major / 4 minor · reviewed 2026-08-16 · deepseek-v4-flash
Pith's one-line read A single density-aware regularizer in the actor loss can correct out-of-distribution states in offline reinforcement learning, the paper argues.
desk verdict A simple, empirically strong density-aware regularizer for offline RL whose variational lower-bound proof does not hold as written; worth reviewing with major revisions. 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 DASP variational lower bound in Theorem 1. It replaces the intractable expectation $\mathbb{E}_{s'\sim P(\cdot|s,a)}\log d(s')$ with an objective over a latent code $z$: reconstruction of $s'$ from $q_\psi(z|s')$, a prior penalty on $q_\theta(z|s,a)$, and a divergence term pulling the two encoders together. This machinery matters because it lets one module supply both the forward model $\hat{P}(s'|\hat{s},a)$ and the density score $f_\tau(L_{\text{dasp}})$, making the safety regularizer a compact component rather than a separately estimated distribution.
What would settle it
Train the DASP module exactly as described and build a held-out set of transitions. Compare the regularizer's scores with a separately fitted density estimate of next states: if the variational score fails to order safe versus unsafe transitions, or if the lower-bound inequality in Theorem 1 is violated numerically, the mechanism's stated guarantee collapses.
Extended reading notes
Core claim
The central discovery is that the one-step forward density, meaning how likely a state's successor is under the offline dataset, can be turned directly into a policy training signal. DASP is a variational module with two encoders, $q_\psi(z|s')$ and $q_\theta(z|s,a)$, and a decoder $P_\phi(s'|z)$; trained with reconstruction, prior, and divergence terms, it doubles as a forward dynamics model and a density indicator. The actor is updated to maximize $Q(s,a)$ plus $\alpha R(\hat{s},a)$, where $R$ is the clipped DASP score of the consequence of the perturbed state-action pair. The paper reports that this single regularizer benchmarks above prior state-correction and conservative methods on MuJoCo, with an average score of 89.0, and on AntMaze, with an average of 68.6, and that it yields smaller performance drops when states are perturbed out of sample.
Load-bearing premise
Everything rests on the identity that the distribution of latent codes encoded from real next states equals the distribution encoded from the current state-action pair; the paper enforces this by matching the two encoders on single data samples, and the appendix concedes that the two objectives are not obviously equivalent.
Editorial extensions
If this is right
- On standard D4RL MuJoCo tasks, the paper reports an average score of 89.0, ahead of the compared baselines in Table 1.
- On AntMaze, DASP reports an average of 68.6, above the compared state-correction and conservative methods across the six sparse-reward navigation tasks.
- In the OOSMuJoCo perturbation benchmarks, DASP maintains higher absolute scores than the compared methods under slight, moderate, and large force perturbations.
- On sub-optimal and downsampled datasets, DASP degrades more slowly than the compared baselines as the ratio of random data or the data shortage increases.
- The regularizer can be added to a standard actor-critic loop with only a pretrained variational module, avoiding a separately trained dynamics model.
- If the regularizer truly captures one-step forward density, the same DASP module can be applied in an online or model-based setting to keep rollouts inside the data manifold.
Reading between the lines
- If one-step forward density is a sufficient safety signal, the same regularizer could be layered onto model-based planners or value functions, not just actor updates, to suppress low-density rollouts before they accumulate.
- The paper's equivalence argument suggests a sharper diagnostic: on transitions where the two encoders disagree strongly, the density score is least trustworthy, so the divergence between them could serve as an uncertainty estimate for when correction will fail.
- An implicit consequence is that DASP collapses the action-suppression and state-recovery stories: raising one-step forward density at every step should, over multiple steps, reduce cumulative state deviation without explicitly aligning transition distributions.
- A direct extension would compare R(hat(s),a) against ground-truth next-state density on held-out transitions; a mismatch would point to the single-sample KL approximation rather than the safety objective itself.
Signed reviews
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper proposes Density-Aware Safety Perception (DASP), a regularizer for offline reinforcement learning that encourages the policy to prefer actions whose one-step consequences land in high-density regions of the offline dataset. The authors formulate the one-step forward density as a variational lower bound (Theorem 1), approximate the required distributions with two encoders and a decoder trained by an ELBO-style loss, and use the resulting score R(ŝ,a) to augment the actor objective. Empirical results are reported on D4RL MuJoCo and AntMaze, on an out-of-sample MuJoCo benchmark, and on suboptimal and data-inefficient settings, with DASP reported to outperform several recent baselines.
Significance. If the theoretical claim were sound, the idea of encoding one-step forward density into an actor regularizer would be a useful and modular contribution to the OOD-state-correction literature. The paper has clear strengths: the method is conceptually simple, it requires only a lightweight addition to an off-the-shelf offline RL algorithm, and the experimental evaluation covers multiple suites and robustness settings. The validation study in Table 3, however, is circular in that the same DASP model is used to produce and to score the safety measure, and the central theoretical result currently rests on an unenforced equality between a variational marginal and an encoder. The paper also leaves the clipping function in Eq. (11) undefined. These issues limit the current significance of the work.
major comments (4)
- [Appendix A, Eq. (7) and Eq. (10)] The proof of Theorem 1 relies on the identity ∫ P(s'|s,a) q1(z|s') ds' = q2(z|s,a), but this identity is never enforced. The training loss in Eq. (10) minimizes only the per-sample KL(qψ(z|s')∥qθ(z|s,a)); the appendix's N=1 argument equates the marginal over the dynamics with a single observed next state, which is not valid. By joint convexity, KL(E_P q1∥q2) ≤ E_P KL(q1∥q2), so minimizing the sample KL does not drive q2 toward the true marginal. Consequently Eq. (7) is not established as a lower bound on E_{s'∼P(·|s,a)} log d(s'), and R(ŝ,a) is not shown to measure one-step forward data density. The paper should either enforce the marginalization identity explicitly or substantially weaken the theoretical claim.
- [Eq. (11), Tables 5-6] The clip function fτ is never specified. The text only says it is a 'clip function with threshold τ,' but the actual mapping—whether it clips values to an interval, clamps gradients, or thresholds and zeroes—is absent. Since fτ determines the sign and scale of R(ŝ,a), and since the threshold τ is never reported, the method is underspecified for reproduction. This is a load-bearing implementation detail, not merely a typo.
- [Eq. (10)-(12), Table 3] Even if Theorem 1 were valid, the quantity L_dasp in Eq. (10) is a composite of a reconstruction error (negative MSE), a KL prior penalty, and a cross-encoder KL term. The regularizer R(ŝ,a) = E_{ŝ'∼P̂} fτ(L_dasp) therefore mixes transition-model likelihood with variational regularization penalties; it is not directly a log-density of d_{πβ}(ŝ'). The paper does not show that maximizing this composite drives the agent toward higher data density, and the validation in Table 3 uses exp(R(s,a)) computed by the same DASP model that is being validated. An independent density estimate is needed to substantiate the claim that R scores in-distribution versus OOD consequences.
- [Section 5.1, Appendix B, Table 4] The hyperparameters α and σ are tuned per task on the very benchmarks used for the final comparison (Table 4 reports per-task α selection; Table 5/6 give per-task σ values). Baselines are mostly taken from previous papers rather than re-run under the same protocol, and baseline standard deviations are not reported in Table 1 or Table 2. The claimed state-of-the-art margins may therefore be inflated by this selection protocol. The authors should report results with a fixed hyperparameter setting, or provide nested model selection, and should include baseline variances to support the comparison.
minor comments (4)
- [Section 4.3, Eq. (8) vs. Eq. (12)] The perturbation radius is denoted B_σ(s) in Eq. (8) but B_ε(s) in Eq. (12); the notation should be made consistent.
- [Eq. (7) and surrounding text] The phrase 'poster distribution' appears several times and should be 'posterior distribution'; this is a persistent typo.
- [Table 1] Several AntMaze entries are marked with a dash ('-') for methods that were not evaluated; the reader should be told whether these results are unavailable or simply not run.
- [Table 2 and Figure 3] Baseline scores in Table 2 and the curves in Figure 3 are reported without error bars; adding standard deviations (or at least stating the seed protocol for baselines) would make the robustness claims more credible.
Circularity Check
No significant circularity: DASP's headline results are benchmarked against external baselines, and the Appendix A identity gap is a proof gap rather than a circular reduction.
full rationale
The derivation is not circular. The DASP regularizer R(ŝ,a) in Eq. (11) is a learned function trained by Eq. (10) on the offline dataset, and the paper's headline results (Table 1) are evaluated on D4RL MuJoCo/AntMaze against external baselines (CQL, PBRL, SPOT, SVR, EDAC, RORL, SDC, OSR-10, SCAS), so the central empirical claim is not equivalent to the training loss by construction. The only self-citations—[Jiang et al. 2023] for OOSMuJoCo and state-shift context and [Zhang and Tan 2024] as general offline-RL background—are not load-bearing: no uniqueness theorem or unverified premise from these works forces Eq. (12). One caveat belongs in correctness, not circularity: Appendix A's proof of Theorem 1 requires the identity ∫P(s'|s,a)q1(z|s')ds' = q2(z|s,a), and the appendix itself concedes 'the relationship between the two formulas is not obvious' while substituting an N=1 per-sample KL; this leaves the lower-bound interpretation of R unproven but does not make the empirical comparison circular, since the benchmark scores come from external environments and baselines.
Assumptions & free parameters
free parameters (3)
- alpha (DASP balance coefficient) =
0.1 / 0.1 / 0.05 for Halfcheetah / Hopper / Walker2d; 0.1 for OOS benchmarks
- sigma (OOD perturbation noise) =
0.001 / 0.005 / 0.01 for standard MuJoCo; 0.05 / 0.005 / 0.07 for OOS MuJoCo
- tau (clip threshold in f_tau)
assumptions (5)
- standard math Bellman optimality equations define Q-learning in offline RL
- domain assumption The stationary state distribution d_{pi_beta} of the behavior policy exists and can be approximated by maximum likelihood on the dataset
- domain assumption Gaussian perturbation ball B_sigma(s) generates useful OOD states for training correction
- ad hoc to paper The variational encoders satisfy the marginalization identity ∫ P(s'|s,a) q1(z|s') ds' = q2(z|s,a)
- domain assumption The learned decoder P_phi(s'|z) combined with q_theta(z|s,a) yields a low-bias dynamics model
Cite this review
Pith. "Pith review of Variational OOD State Correction for Offline Reinforcement Learning." pith.science (2026). https://pith.science/paper/AVZW4IFF
@misc{pith2026250500503,
author = {Pith},
title = {Pith review of: Variational OOD State Correction for Offline Reinforcement Learning},
year = {2026},
howpublished = {\url{https://pith.science/paper/AVZW4IFF}},
note = {Machine review of arXiv:2505.00503}
}
read the original abstract
The performance of Offline reinforcement learning is significantly impacted by the issue of state distributional shift, and out-of-distribution (OOD) state correction is a popular approach to address this problem. In this paper, we propose a novel method named Density-Aware Safety Perception (DASP) for OOD state correction. Specifically, our method encourages the agent to prioritize actions that lead to outcomes with higher data density, thereby promoting its operation within or the return to in-distribution (safe) regions. To achieve this, we optimize the objective within a variational framework that concurrently considers both the potential outcomes of decision-making and their density, thus providing crucial contextual information for safe decision-making. Finally, we validate the effectiveness and feasibility of our proposed method through extensive experimental evaluations on the offline MuJoCo and AntMaze suites.
Figures
Figures from the paper (2 more)
Reference graph
Works this paper leans on
-
[1]
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774 , 2023
arXiv 2023
-
[2]
Learning markov state abstractions for deep reinforcement learning
Cameron Allen, Neev Parikh, Omer Gottesman, and George Konidaris. Learning markov state abstractions for deep reinforcement learning. In Marc'Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Syste...
work page 2021
-
[3]
Uncertainty-based offline reinforcement learning with diversified q-ensemble
Gaon An, Seungyong Moon, Jang - Hyun Kim, and Hyun Oh Song. Uncertainty-based offline reinforcement learning with diversified q-ensemble. In Marc'Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing S...
work page 2021
-
[4]
Pessimistic bootstrapping for uncertainty-driven offline reinforcement learning
Chenjia Bai, Lingxiao Wang, Zhuoran Yang, Zhi - Hong Deng, Animesh Garg, Peng Liu, and Zhaoran Wang. Pessimistic bootstrapping for uncertainty-driven offline reinforcement learning. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022 . OpenReview.net, 2022
work page 2022
-
[5]
Label-noise robust logistic regression and its applications
Jakramate Bootkrajang and Ata Kab \'a n. Label-noise robust logistic regression and its applications. In Joint European conference on machine learning and knowledge discovery in databases , pages 143--158. Springer, 2012
work page 2012
-
[6]
Importance weighted autoencoders
Yuri Burda, Roger Grosse, and Ruslan Salakhutdinov. Importance weighted autoencoders. arXiv preprint arXiv:1509.00519 , 2015
arXiv 2015
-
[7]
Tutorial on variational autoencoders
Carl Doersch. Tutorial on variational autoencoders. arXiv preprint arXiv:1606.05908 , 2016
arXiv 2016
-
[8]
D4RL: datasets for deep data-driven reinforcement learning
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4RL: datasets for deep data-driven reinforcement learning. CoRR , abs/2004.07219, 2020
arXiv 2004
Show all 34 references
-
[9]
Off-policy deep reinforcement learning without exploration
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, Cali...
2019
-
[10]
A comprehensive survey on safe reinforcement learning
Javier Garc a and Fernando Fern \'a ndez. A comprehensive survey on safe reinforcement learning. Journal of Machine Learning Research , 16(1):1437--1480, 2015
2015
-
[11]
Estimation of non-normalized statistical models by score matching
Aapo Hyv \"a rinen and Peter Dayan. Estimation of non-normalized statistical models by score matching. Journal of Machine Learning Research , 6(4), 2005
2005
-
[12]
Planning with diffusion for flexible behavior synthesis
Michael Janner, Yilun Du, Joshua B Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. arXiv preprint arXiv:2205.09991 , 2022
2022 arXiv
-
[13]
Recovering from out-of-sample states via inverse dynamics in offline reinforcement learning
Ke Jiang, Jia-Yu Yao, and Xiaoyang Tan. Recovering from out-of-sample states via inverse dynamics in offline reinforcement learning. In Thirty-seventh Conference on Neural Information Processing Systems , 2023
2023
-
[14]
Ying Jin, Zhuoran Yang, and Zhaoran Wang. Is pessimism provably efficient for offline rl? In Marina Meila and Tong Zhang, editors, Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event , volume 139 of Proceedings of Mac...
2021
-
[15]
Scalable deep reinforcement learning for vision-based robotic manipulation
Dmitry Kalashnikov, Alex Irpan, Peter Pastor, Julian Ibarz, Alexander Herzog, Eric Jang, Deirdre Quillen, Ethan Holly, Mrinal Kalakrishnan, Vincent Vanhoucke, and Sergey Levine. Scalable deep reinforcement learning for vision-based robotic manipulation. In 2nd Annual Conferenc...
2018
-
[16]
Lyapunov density models: Constraining distribution shift in learning-based control
Katie Kang, Paula Gradu, Jason J Choi, Michael Janner, Claire Tomlin, and Sergey Levine. Lyapunov density models: Constraining distribution shift in learning-based control. In International Conference on Machine Learning , pages 10708--10733. PMLR, 2022
2022
-
[17]
Kingma and Max Welling
Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. In Yoshua Bengio and Yann LeCun, editors, 2nd International Conference on Learning Representations, ICLR 2014, Banff, AB, Canada, April 14-16, 2014, Conference Track Proceedings , 2014
2014
-
[18]
Conservative q-learning for offline reinforcement learning
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. In Hugo Larochelle, Marc'Aurelio Ranzato, Raia Hadsell, Maria - Florina Balcan, and Hsuan - Tien Lin, editors, Advances in Neural Information Processing Sys...
2020
-
[19]
Batch reinforcement learning
Sascha Lange, Thomas Gabel, and Martin Riedmiller. Batch reinforcement learning. In Reinforcement learning: State-of-the-art , pages 45--73. Springer, 2012
2012
-
[20]
Supported value regularization for offline reinforcement learning
Yixiu Mao, Hongchang Zhang, Chen Chen, Yi Xu, and Xiangyang Ji. Supported value regularization for offline reinforcement learning. In Thirty-seventh Conference on Neural Information Processing Systems , 2023
2023
-
[21]
Offline reinforcement learning with OOD state correction and OOD action suppression
Yixiu Mao, Qi Wang, Chen Chen, Yun Qu, and Xiangyang Ji. Offline reinforcement learning with OOD state correction and OOD action suppression. CoRR , abs/2410.19400, 2024
2024 arXiv
-
[22]
Human-level control through deep reinforcement learning
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature , 518(7540):529--533, 2015
2015
-
[23]
Deeploco: Dynamic locomotion skills using hierarchical deep reinforcement learning
Xue Bin Peng, Glen Berseth, KangKang Yin, and Michiel Van De Panne. Deeploco: Dynamic locomotion skills using hierarchical deep reinforcement learning. Acm transactions on graphics (tog) , 36(4):1--13, 2017
2017
-
[24]
Mastering the game of go without human knowledge
David Silver, Julian Schrittwieser, Karen Simonyan, Ioannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, et al. Mastering the game of go without human knowledge. nature , 550(7676):354--359, 2017
2017
-
[25]
Llama 2: Open foundation and fine-tuned chat models
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288 , 2023
2023 arXiv
-
[26]
Robust distance metric learning in the presence of label noise
Dong Wang and Xiaoyang Tan. Robust distance metric learning in the presence of label noise. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 28, 2014
2014
-
[27]
Christopher J. C. H. Watkins and Peter Dayan. Technical note q-learning. Mach. Learn. , 8:279--292, 1992
1992
-
[28]
Behavior regularized offline reinforcement learning
Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint arXiv:1911.11361 , 2019
1911 arXiv
-
[29]
Supported policy optimization for offline reinforcement learning
Jialong Wu, Haixu Wu, Zihan Qiu, Jianmin Wang, and Mingsheng Long. Supported policy optimization for offline reinforcement learning. In Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems 35: An...
2022
-
[30]
RORL: robust offline reinforcement learning via conservative smoothing
Rui Yang, Chenjia Bai, Xiaoteng Ma, Zhaoran Wang, Chongjie Zhang, and Lei Han. RORL: robust offline reinforcement learning via conservative smoothing. CoRR , abs/2206.02829, 2022
2022 arXiv
-
[31]
An implicit trust region approach to behavior regularized offline reinforcement learning
Zhe Zhang and Xiaoyang Tan. An implicit trust region approach to behavior regularized offline reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 38, pages 16944--16952, 2024
2024
-
[32]
State deviation correction for offline reinforcement learning
Hongchang Zhang, Jianzhun Shao, Yuhang Jiang, Shuncheng He, Guanwen Zhang, and Xiangyang Ji. State deviation correction for offline reinforcement learning. In Thirty-Sixth AAAI Conference on Artificial Intelligence, AAAI 2022 , pages 9022--9030. AAAI Press, 2022
2022
-
[33]
Constrained policy optimization with explicit behavior density for offline reinforcement learning
Jing Zhang, Chi Zhang, Wenjia Wang, and Bingyi Jing. Constrained policy optimization with explicit behavior density for offline reinforcement learning. In Alice Oh, Tristan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine, editors, Advances in Neural Infor...
2023
-
[34]
write newline
" write newline "" before.all 'output.state := FUNCTION fin.entry add.period write newline FUNCTION new.block output.state before.all = 'skip after.block 'output.state := if FUNCTION new.sentence output.state after.block = 'skip output.state before.all = 'skip after.sentence '...
Reviewed August 16, 2026 · model on record in the stance chip above.
Discussion (0). Continue with ORCID to comment.