ADG: Ambient Diffusion-Guided Dataset Recovery for Corruption-Robust Offline Reinforcement Learning
Extracted source text · arXiv 2505.23871
\documentclass{article}
% if you need to pass options to natbib, use, e.g.:
% \PassOptionsToPackage{numbers, compress}{natbib}
% before loading neurips_2025
% ready for submission
\usepackage[preprint]{neurips_2025}
% to compile a preprint version, e.g., for submission to arXiv, add add the
% [preprint] option:
% \usepackage[preprint]{neurips_2025}
% to compile a camera-ready version, add the [final] option, e.g.:
% \usepackage[final]{neurips_2025}
% to avoid loading the natbib package, add option nonatbib:
% \usepackage[nonatbib]{neurips_2025}
\usepackage[utf8]{inputenc} % allow utf-8 input
\usepackage[T1]{fontenc} % use 8-bit T1 fonts
%\usepackage{hyperref} % hyperlinks
\usepackage{url} % simple URL typesetting
\usepackage{booktabs} % professional-quality tables
\usepackage{amsfonts} % blackboard math symbols
\usepackage{nicefrac} % compact symbols for 1/2, etc.
\usepackage{xcolor} % colors
\usepackage{graphicx}
\usepackage{subfigure}
\usepackage{booktabs}
\usepackage{colortbl}
\usepackage{setspace}
\usepackage{booktabs} % \toprule, \midrule, \bottomrule
\usepackage{multirow} % \multirow{}
\usepackage{colortbl} % \columncolor{}
\usepackage{graphicx} % 插入图形 (图像宽度设置用)
\usepackage{rotating} % \rotatebox{} 用于旋转文字
\usepackage{array} % >{\columncolor{}}c 之类需要这个
\usepackage{float} % 控制浮动体位置，例如 [H]
\usepackage{makecell} % 可选，改进单元格内容换行
\usepackage{caption} % 改进 caption 控制（可选）
%\usepackage{hyperref}
\usepackage{bm}
\usepackage{amsmath}
\usepackage{amssymb}
\usepackage{mathtools}
\usepackage{amsthm}
\usepackage{algorithm}
\usepackage{algorithmic}
\usepackage{wrapfig}
\usepackage[capitalize,noabbrev]{cleveref}
\newtheorem{theorem}{Theorem}[section]
\newtheorem{proposition}[theorem]{Proposition}
\newtheorem{lemma}[theorem]{Lemma}
\newtheorem{corollary}[theorem]{Corollary}
\newtheorem{definition}[theorem]{Definition}
\newtheorem{assumption}[theorem]{Assumption}
\newtheorem{remark}[theorem]{Remark}
\def\blue#1{{\color{blue}#1}}
\def\red#1{{\color{red}#1}}
\def\black#1{{\color{black}#1}}
\long\def\old#1{}
\title{ADG: Ambient Diffusion-Guided Dataset Recovery for Corruption-Robust
Offline Reinforcement Learning}
% The \author macro works with any number of authors. There are two commands
% used to separate the names and addresses of multiple authors: \And and \AND.
%
% Using \And between authors leaves it to LaTeX to determine where to break the
% lines. Using \AND forces a line break at that point. So, if LaTeX puts 3 of 4
% authors names on the first line, and the last on the second line, try using
% \AND instead of \And before the third author name.
% \author{%
% Zeyuan Liu\thanks{equal contribution} \\
% SIGS\\
% \And
% Zhihe Yang \\
% SIGS\\
% \And
% Jiawei Xu\\
% The Chinese University of Hong Kong, Shenzhen\\
% \And
% Rui Yang\\
% University of Illinois Urbana-Champaign\\
% \And
% Baoxiang Wang\\
% The Chinese University of Hong Kong, Shenzhen\\
% % examples of more authors
% % \And
% % Coauthor \\
% % Affiliation \\
% % Address \\
% % \texttt{email} \\
% % \AND
% % Coauthor \\
% % Affiliation \\
% % Address \\
% % \texttt{email} \\
% % \And
% % Coauthor \\
% % Affiliation \\
% % Address \\
% % \texttt{email} \\
% % \And
% % Coauthor \\
% % Affiliation \\
% % Address \\
% % \texttt{email} \\
% }
\author{%
Zeyuan Liu\textsuperscript{1}\thanks{Equal contribution},
Zhihe Yang\textsuperscript{2}\footnotemark[1],
Jiawei Xu\textsuperscript{3}\footnotemark[1],
Rui Yang\textsuperscript{4},
Jiafei Lyu\textsuperscript{1},
Baoxiang Wang\textsuperscript{3},\\
\textbf{Yunjian Xu}\textsuperscript{2},
\textbf{Xiu Li}\textsuperscript{1}
\\
\textsuperscript{1}Tsinghua Shenzhen International Graduate School, Tsinghua University,\\
\textsuperscript{2}The Chinese University of Hong Kong,\\
\textsuperscript{3}The Chinese University of Hong Kong, Shenzhen,
\textsuperscript{4}University of Illinois Urbana-Champaign
\\
}
\begin{document}
\maketitle
\input{section/0_abstract}
\input{section/1_introduction}
\input{section/2_related_works}
\input{section/2_preliminaries}
\input{section/3_methodologies}
\input{section/4_experiments}
\input{section/5_conclusion}
\bibliography{icml2025/reference}
\bibliographystyle{plain}
% \section*{References}
% References follow the acknowledgments in the camera-ready paper. Use unnumbered first-level heading for
% the references. Any choice of citation style is acceptable as long as you are
% consistent. It is permissible to reduce the font size to \verb+small+ (9 point)
% when listing the references.
% Note that the Reference section does not count towards the page limit.
% \medskip
% {
% \small
% [1] Alexander, J.A.\ \& Mozer, M.C.\ (1995) Template-based algorithms for
% connectionist rule extraction. In G.\ Tesauro, D.S.\ Touretzky and T.K.\ Leen
% (eds.), {\it Advances in Neural Information Processing Systems 7},
% pp.\ 609--616. Cambridge, MA: MIT Press.
% [2] Bower, J.M.\ \& Beeman, D.\ (1995) {\it The Book of GENESIS: Exploring
% Realistic Neural Models with the GEneral NEural SImulation System.} New York:
% TELOS/Springer--Verlag.
% [3] Hasselmo, M.E., Schnell, E.\ \& Barkai, E.\ (1995) Dynamics of learning and
% recall at excitatory recurrent synapses and cholinergic modulation in rat
% hippocampal region CA3. {\it Journal of Neuroscience} {\bf 15}(7):5249-5262.
% }
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\appendix
\newpage
\input{section/7_appendix}
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
\newpage
% \section*{NeurIPS Paper Checklist}
%%% BEGIN INSTRUCTIONS %%%
% The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: {\bf The papers not including the checklist will be desk rejected.} The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page
% limit.
% Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
% \begin{itemize}
% \item You should answer \answerYes{}, \answerNo{}, or \answerNA{}.
% \item \answerNA{} means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.
% \item Please provide a short (1–2 sentence) justification right after your answer (even for NA).
% % \item {\bf The papers not including the checklist will be desk rejected.}
% \end{itemize}
% {\bf The checklist answers are an integral part of your paper submission.} They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include it (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
% The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation. While "\answerYes{}" is generally preferable to "\answerNo{}", it is perfectly acceptable to answer "\answerNo{}" provided a proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "\answerNo{}" or "\answerNA{}" is not grounds for rejection. While the questions are phrased in a binary way, we acknowledge that the true answer is often more nuanced, so please just use your best judgment and write a justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer \answerYes{} to a question, in the justification please point to the section(s) where related material for the question can be found.
% IMPORTANT, please:
% \begin{itemize}
% \item {\bf Delete this instruction block, but keep the section heading ``NeurIPS Paper Checklist"},
% \item {\bf Keep the checklist subsection headings, questions/answers and guidelines below.}
% \item {\bf Do not modify the questions and only use the provided macros for your answers}.
% \end{itemize}
%%% END INSTRUCTIONS %%%
% \begin{enumerate}
% \item {\bf Claims}
% \item[] Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
% \item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification: There are three main claims laid out in the abstract. These concern ADG’s ability to recover near-original trajectories under heavy corruption, the tradeoff between corruption detection precision and policy learning stability, and improved performance from boundary-biased diffusion sampling. All of these points are addressed in the main text through theoretical analysis and empirical validation.
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the abstract and introduction do not include the claims made in the paper.
% \item The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
% \item The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
% \item It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
% \end{itemize}
% \item {\bf Limitations}
% \item[] Question: Does the paper discuss the limitations of the work performed by the authors?
% \item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification: We include the limitations of our work in~\cref{limitation}.
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
% \item The authors are encouraged to create a separate "Limitations" section in their paper.
% \item The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
% \item The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
% \item The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
% \item The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
% \item If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
% \item While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
% \end{itemize}
% \item {\bf Theory assumptions and proofs}
% \item[] Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
% \item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification: The proof of most statements are layed out in the supplemental, to save on space.
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the paper does not include theoretical results.
% \item All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced.
% \item All assumptions should be clearly stated or referenced in the statement of any theorems.
% \item The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
% \item Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
% \item Theorems and Lemmas that the proof relies upon should be properly referenced.
% \end{itemize}
% \item {\bf Experimental result reproducibility}
% \item[] Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
% \item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification: We explained our settings in~\cref{sec:experiments} and hyperparameters in~\cref{hyper-parameter}.
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the paper does not include experiments.
% \item If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
% \item If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
% \item Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
% \item While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
% \begin{enumerate}
% \item If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
% \item If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
% \item If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
% \item We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
% \end{enumerate}
% \end{itemize}
% \item {\bf Open access to data and code}
% \item[] Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?
% \item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification: We upload the codes and instructions to recover the results. Once the blind review period is finished, we’ll open-source all codes, instructions, and model checkpoints.
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that paper does not include experiments requiring code.
% \item Please see the NeurIPS code and data submission guidelines (\url{https://nips.cc/public/guides/CodeSubmissionPolicy}) for more details.
% \item While we encourage the release of code and data, we understand that this might not be possible, so “No” is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
% \item The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (\url{https://nips.cc/public/guides/CodeSubmissionPolicy}) for more details.
% \item The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
% \item The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.
% \item At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
% \item Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
% \end{itemize}
% \item {\bf Experimental setting/details}
% \item[] Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
% \item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification: Training and test details (including hyperparameters, optimizer settings, and data splits) are specified in~\cref{sec:experiments} and~\cref{hyper-parameter}.
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the paper does not include experiments.
% \item The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
% \item The full details can be provided either with the code, in appendix, or as supplemental material.
% \end{itemize}
% \item {\bf Experiment statistical significance}
% \item[] Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
% \item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification: We show the standard error in most training curves, with an average of over four random seeds.
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the paper does not include experiments.
% \item The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
% \item The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
% \item The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
% \item The assumptions made should be given (e.g., Normally distributed errors).
% \item It should be clear whether the error bar is the standard deviation or the standard error of the mean.
% \item It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96\% CI, if the hypothesis of Normality of errors is not verified.
% \item For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
% \item If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
% \end{itemize}
% \item {\bf Experiments compute resources}
% \item[] Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
% \item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification: We have a~\cref{app:computation} on this.
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the paper does not include experiments.
% \item The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
% \item The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
% \item The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
% \end{itemize}
% \item {\bf Code of ethics}
% \item[] Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics \url{https://neurips.cc/public/EthicsGuidelines}?
% \item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification: We do not believe that our work has any harmful consequences as layed out in the Code of Ethics.
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
% \item If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
% \item The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).
% \end{itemize}
% \item {\bf Broader impacts}
% \item[] Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
% \item[] Answer: \answerNA{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification: Our work involves small models and simulated datasets. It does not impact society at large, beyond improving our understanding of certain aspects of deep learning.
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that there is no societal impact of the work performed.
% \item If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
% \item Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
% \item The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
% \item The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
% \item If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
% \end{itemize}
% \item {\bf Safeguards}
% \item[] Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
% \item[] Answer: \answerNA{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification: No such models or datasets are involved.
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the paper poses no such risks.
% \item Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
% \item Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
% \item We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.
% \end{itemize}
% \item {\bf Licenses for existing assets}
% \item[] Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
% \item[] Answer: \answerYes{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification: Yes, we credited them in appropriate ways.
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the paper does not use existing assets.
% \item The authors should cite the original paper that produced the code package or dataset.
% \item The authors should state which version of the asset is used and, if possible, include a URL.
% \item The name of the license (e.g., CC-BY 4.0) should be included for each asset.
% \item For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
% \item If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, \url{paperswithcode.com/datasets} has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.
% \item For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
% \item If this information is not available online, the authors are encouraged to reach out to the asset's creators.
% \end{itemize}
% \item {\bf New assets}
% \item[] Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
% \item[] Answer: \answerNA{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification: The paper does not release new assets.
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the paper does not release new assets.
% \item Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
% \item The paper should discuss whether and how consent was obtained from people whose asset is used.
% \item At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
% \end{itemize}
% \item {\bf Crowdsourcing and research with human subjects}
% \item[] Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
% \item[] Answer: \answerNA{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification: The paper does not involve crowdsourcing nor research with human subjects.
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
% \item Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
% \item According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
% \end{itemize}
% \item {\bf Institutional review board (IRB) approvals or equivalent for research with human subjects}
% \item[] Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
% \item[] Answer: \answerNA{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification: We have no human participants in our study.
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
% \item Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
% \item We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
% \item For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
% \end{itemize}
% \item {\bf Declaration of LLM usage}
% \item[] Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
% %this research?
% \item[] Answer: \answerNA{} % Replace by \answerYes{}, \answerNo{}, or \answerNA{}.
% \item[] Justification: The LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research.
% \item[] Guidelines:
% \begin{itemize}
% \item The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
% \item Please refer to our LLM policy (\url{https://neurips.cc/Conferences/2025/LLM}) for what should or should not be described.
% \end{itemize}
% \end{enumerate}
\end{document}
\begin{abstract}\label{abstract}
Real-world datasets collected from sensors or human inputs are prone to noise and errors, posing significant challenges for applying offline reinforcement learning (RL). While existing methods have made progress in addressing corrupted actions and rewards, they remain insufficient for handling corruption in high-dimensional state spaces and for cases where multiple elements in the dataset are corrupted simultaneously. Diffusion models, known for their strong denoising capabilities, offer a promising direction for this problem—but their tendency to overfit noisy samples limits their direct applicability.
To overcome this, we propose \textbf{A}mbient \textbf{D}iffusion-\textbf{G}uided Dataset Recovery (\textbf{ADG}), a novel approach that pioneers the use of diffusion models to tackle data corruption in offline RL. First, we introduce Ambient Denoising Diffusion Probabilistic Models (DDPM) from approximated distributions, which enable learning on partially corrupted datasets with theoretical guarantees. Second, we use the noise-prediction property of Ambient DDPM to distinguish between clean and corrupted data, and then use the clean subset to train a standard DDPM. Third, we employ the trained standard DDPM to refine the previously identified corrupted data, enhancing data quality for subsequent offline RL training. A notable strength of ADG is its versatility—it can be seamlessly integrated with any offline RL algorithm. Experiments on a range of benchmarks, including MuJoCo, Kitchen, and Adroit, demonstrate that ADG effectively mitigates the impact of corrupted data and improves the robustness of offline RL under various noise settings, achieving state-of-the-art results.
% Offline reinforcement learning (RL) typically relies on real-world data collected through sensors or human interactions, which are often noisy and prone to errors.
% Although existing methods have made progress in addressing corrupted actions and rewards, they remain insufficient for handling corruption in high-dimensional states or scenarios involving multiple corrupted elements.
% Recent advancements in diffusion models have shown significant potential in tackling this challenge; however, their tendency to overfit noisy samples limits their direct applicability.
% In this work, we present \textbf{A}mbient \textbf{D}iffusion-\textbf{G}uided Dataset Recovery (\textbf{ADG}), a novel three-stage approach that pioneers the use of diffusion models to tackle corruption-robust offline RL.
% First, we introduce Ambient Denoising Diffusion Probabilistic Models (DDPM) from approximated distributions, which enable diffusion training on partially corrupted datasets with theoretical guarantees.
% Next, relying on its noise-prediction property, we identify the corrupted data within the dataset and take the remaining clean data to train a standard DDPM.
% Finally, we employ the standard DDPM to refine the previously identified corrupted data, resulting in a high-quality dataset for subsequent offline RL training.
% Our method is versatile and compatible with any offline RL approach and even their robust variants. Experiments on a range of benchmarks, including MuJoCo, Kitchen, and Adroit, demonstrate that ADG effectively reduces the impact of corrupted data and enhances the robustness of offline RL methods, achieving state-of-the-art performance.
\end{abstract}
%Diffusion models, known for their strong denoising capabilities, are promising for this problem but are prone to overfitting noisy samples. To address this, we propose \textbf{A}mbient \textbf{D}iffusion-\textbf{G}uided Dataset Recovery (\textbf{ADG}), a novel two-stage method that fully harness diffusion models.
%First, we introduce Ambient Denoising Diffusion Probabilistic Models (DDPM) to identify corrupted data using their noise-prediction properties. Then, we apply a standard diffusion model trained on the selected clean data to refine the previously identified corrupted data, improving the dataset quality for offline RL training.
% FOR MARKDOWN FORMAT
% Offline reinforcement learning (RL) typically relies on real-world data collected through sensors or human interactions, which are often noisy and prone to errors. Although existing methods have made progress in addressing corrupted actions and rewards, they remain insufficient for handling corruption in high-dimensional states or scenarios involving multiple corrupted elements. Recent advancements in diffusion models have shown significant potential in tackling this challenge; however, their tendency to overfit noisy samples limits their direct applicability. In this work, we present **A**mbient **D**iffusion-**G**uided Dataset Recovery (**ADG**), a novel three-stage approach that pioneers the use of diffusion models to tackle corruption-robust offline RL. First, we introduce Ambient Denoising Diffusion Probabilistic Models (DDPM) from approximated distributions, which enable diffusion training on partially corrupted datasets with theoretical guarantees. Next, relying on its noise-prediction property, we identify the corrupted data within the dataset and take the remaining clean data to train a standard DDPM. Finally, we employ the standard DDPM to refine the previously identified corrupted data, resulting in a high-quality dataset for subsequent offline RL training. Our method is versatile and compatible with any offline RL approach and even their robust variants. Experiments on a range of benchmarks, including MuJoCo, Kitchen, and Adroit, demonstrate that ADG effectively reduces the impact of corrupted data and enhances the robustness of offline RL methods, achieving state-of-the-art performance.
\section{Introduction}\label{intro}
Offline reinforcement learning (RL) has emerged as a prominent paradigm for learning decision-making policies from offline datasets~\citep{levine2020offlinereinforcementlearningtutorial,fujimoto2019off}. Existing approaches can be broadly categorized into MDP-based methods~\citep{fujimoto2019off, fujimoto2021minimalist, kostrikov2021offline, kumar2020conservative, an2021uncertainty, bai2022pessimistic, ghasemipour2022so} and non-MDP methods~\citep{chen2021decision, janner2021offline, shi2023unleashing}. However, due to the data-dependent nature of offline RL, it encounters significant challenges when dealing with offline data subjected to random noise or adversarial corruption~\citep{zhang2020robust, zhang2021robust, liang2024robust, yang2024regularizing,yang2023towards}. Such disturbances can cause substantial performance degradation or result in a pronounced deviation from the intended policy objectives. Therefore, ensuring robust policy learning is crucial for offline RL to function effectively in real-world scenarios.
Several previous studies have focused on the theoretical properties and certification of offline RL under corrupted data~\citep{zhang2022corruption, ye2024corruption, wu2022copa, chen2024exact, ye2024towards}. Empirical efforts have resulted in the development of uncertainty-weighted algorithms utilizing Q-ensembles~\citep{ye2024corruption} and robust value learning through Huber loss and quantile Q estimators~\cite{yang2023towards}. Additionally, sequence modeling techniques have been applied to mitigate the effects of data corruption while iteratively refining noisy actions and rewards in the dataset~\citep{xu2024robust}. However, as noted in~\citep{xu2024robust}, recovering observations remains challenging due to the high dimensionality.
\begin{figure}[t]
\centering
% 子图 1
\includegraphics[width=1.00\textwidth]{pics/ADG_algo_structure.pdf}
\caption{
Overview of the training processes for the detector and denoiser (left) and the dataset recovery process (right) in the proposed ADG method.
}\vspace{-2mm}
\label{fig:ADG_algo}
\end{figure}
In recent years, diffusion-based generative models~\citep{sohl2015deep, ho2020denoising, song2020score} have gained considerable attention for their ability to effectively model complex data distributions, making them increasingly important in offline RL~\citep{janner2022planning, ajay2023conditional, liang2023adaptdiffuser}. One promising application of these models is their usage in reducing or mitigating noise in data, thanks to their inherent denoising capabilities. For instance, DMBP~\citep{yang2024dmbp} proposed a diffusion-based framework aimed at minimizing noise in observations during the testing phase. Despite its success, this approach is specifically designed to work with clean training data and only addresses data perturbations during testing. It cannot manage perturbations in the training dataset, as current diffusion models typically assume the dataset is entirely clean or has a consistent noise distribution across all data points~\citep{aali2024ambient, daras24ambient}. Therefore, naive diffusion methods can encounter challenges, such as overfitting to noise data, when applied to partially corrupted data during training.
To gain deeper insight into which aspects of corrupted datasets degrade the performance of offline RL algorithms, we evaluate several offline RL methods on three types of datasets: a clean dataset, a partially corrupted dataset, and a filtered dataset, which is created by removing the noisy portions from the corrupted data, as illustrated in~\cref{fig:motivation} (a–c). Surprisingly, We find that current offline RL algorithms fail to fully restore their performance on the filtered dataset, particularly when the dataset size is limited. These results imply that the loss of critical sequential information plays a key role in performance degrading of decision-making. Additionally, while sequence modeling methods~\citep{chen2021decision,xu2024robust} exhibit robustness against data corruption, they fail to function when faced with incomplete trajectories. \textbf{This finding indicates that simply filtering out noisy samples is insufficient for handling corrupted datasets, underscoring the importance of recovering corrupted data.}
\begin{figure*}[h]
\centering
% 子图 1
\includegraphics[width=1.00\textwidth]{pics/motivation.pdf}
\caption{
The performance of baseline algorithms (CQL, IQL, RIQL, DT, and RDT) is evaluated under four dataset conditions: \textbf{Clean} (original dataset without corruption), \textbf{Noised} (dataset corrupted using Random State Attack as described in Appendix~\ref{noised_setting}), \textbf{Filtered} (Noised dataset with corrupted samples removed), and \textbf{Recovered} (dataset recovered using ADG). The comparison results show the average normalized scores across three MuJoCo tasks (\textit{halfcheetah}, \textit{hopper}, and \textit{walker2d}) using the ``medium-replay-v2'' datasets. We include the results for both the 100\% dataset and 10\% dataset settings. ADG effectively restores performance close to the results on clean datasets.
}\vspace{-0mm}
\label{fig:motivation}
\end{figure*}
%In this paper, we first analyze the performance of various offline RL algorithms on three types of datasets: a clean dataset, a partially corrupted dataset, and a filtered dataset, which is a subset of the corrupted dataset with the noised portion removed, as shown in~\cref{fig:motivation} (a–c). We find that state-of-the-art offline RL algorithms fail to fully restore their performance on the filtered dataset, particularly when the dataset size is limited. This is possibly due to the loss of critical information for decision-making caused by incomplete sequences. Additionally, while sequence modeling methods~\citep{chen2021decision, xu2024robust} exhibit robustness against data corruption, they fail to function when faced with incomplete trajectories. These results indicate that simply filtering out noisy samples is insufficient for handling corrupted datasets, underscoring the importance of recovering corrupted data.
To address this issue, \textit{we introduce the first diffusion-based denoising framework for handling corruption robust offline RL.} Our approach recovers clean data purely from the corrupted dataset, without requiring any external information or supervision. We name this novel three-stage diffusion-based method \textbf{A}mbient \textbf{D}iffusion-\textbf{G}uided Dataset Recovery (\textbf{ADG}).
% Our method leverages ambient diffusion model as powerful detector to distinguish between corrupted and clean data, significantly aiding the training of a naive diffusion model as a denoiser based on the selected training data.
The detailed diagram is presented in Figure~\ref{fig:ADG_algo}.
In the first stage, we introduce Ambient Denoising Diffusion Probabilistic Models (DDPM) from \textit{approximated distributions}, which enable diffusion training on partially corrupted datasets with theoretical guarantees.
In the second stage, leveraging the noise-prediction property of the well-trained ambient DDPM, we identify the corrupted data within the dataset. The remaining clean data is then used to train a denoiser within the framework of standard DDPM.
Finally, in the third stage, we apply the standard DDPM to refine the previously identified corrupted data. The corrected data is combined with the clean data to form a high-quality dataset, which is subsequently used for offline RL training.
% which provides a theoretical foundation for diffusion models to learn an approximation of the distribution in the presence of partial corruption within the dataset, as directly learning the distribution is infeasible.
% The detector then generates masks to mitigate the negative impact of noisy data, enabling the training of the denoiser using a naive loss function. A diagram of our method is shown in Figure~\ref{fig:ADG_algo}.
% To address this issue and fully harness the potential of diffusion models for recovering corrupted data, we investigate the feasibility of directly recovering noisy datasets using diffusion models.
%To address this issue and fully harness the potential of diffusion models for recovering corrupted data, we investigate the feasibility of directly recovering noisy datasets using diffusion models, relying solely on the information contained within the dataset itself without requiring any external information. We propose a novel two-stage diffusion-based approach called \textbf{A}mbient \textbf{D}iffusion-\textbf{G}uided Dataset Recovery (\textbf{ADG}). Our method leverages ambient diffusion model as powerful detector to distinguish between corrupted and clean data, significantly aiding the training of a naive diffusion model as a denoiser based on the selected training data. In the first stage, we introduce Ambient Denoising Diffusion Probabilistic Models from Approximated Distribution, which provides a theoretical foundation for diffusion models to learn an approximation of the distribution in the presence of partial corruption within the dataset, as directly learning the distribution is infeasible. The detector then generates masks to mitigate the negative impact of noisy data, enabling the training of the denoiser using a naive loss function. To the best of our knowledge, ADG is the first diffusion-based denoising framework for offline RL under data corruption during the training process. A diagram of our method is shown in Figure~\ref{fig:ADG_algo}.
We find that ADG is overall competitive with the ideal case of perfectly filtering out all noisy samples for MDP-based algorithms and significantly outperforms the filtering method for sequence modeling methods, as shown in Figure~\ref{fig:motivation}(d). Additionally, we provide a comprehensive analysis of ADG's performance under both random and adversarial data corruption scenarios, examining various levels of data availability, whether full or limited dataset size, as described in \cref{sec:experiments}.
Our empirical studies demonstrate that ADG consistently enhances the performance of all baseline algorithms, achieving a remarkable overall improvement. These findings indicate that ADG is versatile and compatible with any offline RL approach, including their robust variants. Notably, ADG exhibits consistent and robust performance across a variety of dataset qualities, corrupted scales, and corrupted ratios.
\section{Related Works}
\paragraph{Robust Offline RL.} Several works have focused on testing-time robustness against environment shifts~\citep{shi2024distributionally, yang2022rorl, panaganti2022robust, yang2024dmbp, xu2024robust}. For training-time robustness, Li et al.~\citep{li2024survival} investigate various reward attack strategies in offline RL and reveal that certain biases can unintentionally enhance robustness to reward corruption. Wu et al.~\citep{wu2022copa} introduce a certification framework to determine the tolerable number of poisoning trajectories based on different certification criteria. From a theoretical perspective, Zhang et al.~\citep{zhang2022corruption} propose a robust offline RL algorithm utilizing robust supervised learning oracles. Ye et al.~\citep{ye2024corruption} introduce uncertainty weighting to address reward and dynamics corruption, offering theoretical guarantees. Ackermann et al.~\citep{ackermann2024offline} develop a contrastive predictive coding-based approach to tackle non-stationarity in offline RL datasets. Yang et al.~\citep{yang2023towards} utilize the Huber loss to manage heavy-tailedness and adopt quantile estimators to balance penalization for corrupted data. Additionally, Xu et al.~\citep{xu2024robust} introduce a sequential modeling method to iteratively correct corrupted data for offline RL.
\paragraph{Diffusion Models in Offline RL.} Diffusion-based generative models \citep{sohl2015deep, ho2020denoising, song2020score} have been extensively utilized for synthesizing high-quality images from text descriptions \citep{rombach2022high}. More recently, they have gained significant attention in the RL community, serving as behavior policy replicators \citep{wang2022diffusion, hansen2023idql}, trajectory generators \citep{janner2022planning, ajay2023conditional, liang2023adaptdiffuser}, and state denoisers \citep{yang2024dmbp}.
\section{Preliminaries}
\label{sec:preliminaries}
\paragraph{RL and Offline RL.} Reinforcement Learning (RL) is typically formulated as a Markov Decision Process (MDP) defined by the tuple $(S, A, P, r, \gamma)$, where $S$ and $A$ represent the state and action spaces, $P$ denotes the transition function, $r$ is the reward function, and $\gamma \in [0, 1]$ is the discount factor. The objective of RL is to learn a policy $\pi(a|s)$ that maximizes the expected cumulative return. In offline RL, access to the online environment is restricted. Instead, the objective is to optimize the RL objective using a previously collected dataset, $\mathcal{D}=\left\{\left(s_t^i, a_t^i, r_t^i, s_{t+1}^i\right)\right\}^{N-1}_{i=0}$ which consists of $N$ transitions in total.
\paragraph{Corruption-robust Offline RL.} We adopt a unified trajectory-based storage approach, as proposed in prior works~\citep{xu2024robust}. An original trajectory is represented as
$\tau = (s_0, a_0, r_0, \ldots, s_{T-1}, a_{T-1}, r_{T-1})$, where each trajectory consists of three components: states, actions, and rewards. This trajectory can be reorganized into sequence data for DT~\citep{chen2021decision} and RDT~\citep{xu2024robust}, or split into transitions $\left(s_t, a_t, r_t, s_{t+1}\right)_{t=0}^{T-2}$ for Markov Decision Process (MDP)-based methods such as CQL~\citep{kumar2020conservative} and IQL~\citep{kostrikov2021offline}.
We investigate the impact of injecting random or adversarial noise into the dataset under two corruption scenarios. First, we examine \textit{state} corruption, where only the states in the trajectories are affected. Second, we introduce noise into state-action-reward triplets, which we refer to as \textit{``full-element''} in the following context.
Random corruption refers to the addition of noise drawn from a uniform distribution. For example, corrupting the state with uniform noise of corruption scale $\alpha$ can be written as $\hat{s}_0 = s_0 + \lambda \cdot \operatorname{std}(s), \quad \lambda \sim \text{Uniform}[-\alpha, \alpha]^{d_s}$, where $d_s$ is the dimensionality of the state, and $\operatorname{std}(s)$ represents the $d_s$-dimensional standard deviation of all states in the offline dataset. On the other hand, adversarial corruption employs a Projected Gradient Descent (PGD) attack~\citep{madry2017towards} with pretrained value functions. We build upon prior work~\citep{xu2024robust} by introducing learnable noise to the target elements and optimizing it through gradient descent to minimize the pretrained value functions. Further details refer to Appendix~\ref{noised_setting}.
\paragraph{Diffusion Models.}
Given any clean sample $\bm{x}$, the \emph{forward process} of diffusion models is a Markov chain that gradually adds Gaussian noise to data according to a variance schedule $\alpha_1,\dots, \alpha_K$:
\begin{equation}\label{forward_diffusion}
\begin{aligned}
&q(\bm{x}^{1:K} \mid \bm{x}^0):= \prod\limits_{k=1}^{K} q(\bm{x}^k \mid \bm{x}^{k-1}),
&q(\bm{x}^k \mid \bm{x}^{k-1}) := \mathcal{N}(\bm{x}^k; \sqrt{\alpha_k}\bm{x}^{k-1},(1-\alpha_k)\bm{I}).
\end{aligned}
\end{equation}
The \emph{reverse process} is likewise a Markov chain characterized by learned Gaussian transitions (parameterized by $\theta$), typically initiated at $p(x^K) = \mathcal{N}(x^K; 0, \bm{I})$:
\begin{equation}\label{reverse_diffusion}
\begin{aligned}
&p_{\theta}(\bm{x}^{0:K}) := p(\bm{x}^K) \prod\limits_{k=1}^{K}p_{\theta}(\bm{x}^{k-1} \mid \bm{x}^k),
&p_{\theta}(\bm{x}^{k-1} \mid \bm{x}^k) := \mathcal{N}(\bm{x}^{k-1}; \bm{\mu}_{\theta}(\bm{x}^k,k),\bm{\Sigma}_{\theta}(\bm{x}^k,k)).
\end{aligned}
\end{equation}
To ensure high-quality generation, the diffusion model learning process typically demands clean original data, i.e., we have direct access to $\bm{x}$. {Building on the foundation of initial Denoising Diffusion Probabilistic Models (DDPM) introduced by Ho et al.~\citep{ho2020denoising}, the most widely adopted training loss for diffusion models is formulated as}:
\begin{equation}\label{original_ddpm_loss}
\mathcal{L}_{\mathrm{DDPM}}(\theta) := \mathbb{E}_{k \sim [1,K],\bm{\epsilon} \sim \mathcal{N}(\bm{0},\bm{I})}[\Vert \bm{\epsilon}_{\theta}(\bm{x}^k,k) - \bm{\epsilon} \Vert^2],
\end{equation}
where $\bm{x}^k=\sqrt{\bar{\alpha}_i}\bm{x} + \sqrt{1 - \bar{\alpha}_i} \bm{\epsilon}$ with $\bar{\alpha}_k := \prod_i^k\alpha_i$, and $\bm{\epsilon}$ is a randomly sampled Gaussian noise.
Nevertheless, obtaining completely clean data poses a significant challenge in certain circumstances, making the loss in Eq.~\eqref{original_ddpm_loss} not directly applicable. Recently, several studies have proposed methods for handling corrupted datasets, while they assume a uniform noise scale across the dataset~\citep{aali2024ambient, daras24ambient}, leaving the problem of training diffusion models on partially corrupted datasets unresolved.
\section{Methodologies}
\subsection{Motivation}
A natural approach to addressing the partially corrupted dataset is to identify the corrupted samples within the dataset and leverage the uncorrupted samples to train a denoiser, which can subsequently be used to correct the corrupted samples.
{However, naive noise detectors trained with supervised methods relying solely on manually labeled clean and noisy samples may easily fail due to the inherent ambiguity of noise.}
As an alternative, diffusion models have garnered significant attention due to their demonstrated effectiveness in training on fully corrupted datasets and leveraging the trained models to extract uncorrupted samples~\cite{aali2024ambient, daras24ambient}.
Moreover, diffusion models have also been shown to be effective as state denoisers in the realm of offline RL~\cite{yang2024dmbp}.
Nevertheless, a key challenge is that existing diffusion models for corrupted samples typically assume a consistent noise distribution across the dataset, which does not hold for corruption-robust offline RL.
On the other hand, if the naive diffusion process is applied to corrupted samples during training, even small amounts of noise in the training set can cause the diffusion model to overfit to these perturbed points, leading to poor performance during subsequent sampling.
\textit{To address the discrepancy between clean and corrupted data within the same dataset, the development of a new diffusion training method is urgently needed.}
\subsection{Diffusion for Partially Corrupt data}
\paragraph{Extending Ambient Diffusion Models to DDPM.}
Daras et al.~\citep{daras24ambient} demonstrate that diffusion models can be trained on datasets corrupted by a consistent scale of noise in the context of score-based continuous diffusion models. We observe that their conclusions can be seamlessly extended to the discrete DDPM framework.
% By following the proof sketch of Theorem A.5 in \citep{daras24ambient}, we derive the following:
{Building upon Theorem A.5 in \citep{daras24ambient}, which addresses score-based variance preserving diffusion, we derive the following corollary that extends the result to the discrete DDPM.}
\begin{corollary}\label{corollary21}
(Ambient DDPM)
Let $k_a$ be a manually defined constant representing a noise-added diffusion timestep.
Suppose we are given samples $\bm{x}^{k_a}= \sqrt{\bar{\alpha}_{k_a}} \bm{x}^0 + \sqrt{1-\bar{\alpha}_{k_a}} \bm{\epsilon}$. Let $\bm{x}^{k}= \sqrt{\bar{\alpha}_{k}} \bm{x}^0 + \sqrt{1-\bar{\alpha}_{k}} \bm{\epsilon}$
be further en-noised samples with $1\leq k_a<k$.
Then, the unique minimizer of the objective
\begin{equation}\label{corollay_eq4}
\begin{aligned}
\mathbb{E}_{\bm{x}^{k_a}} & \mathbb{E}_{k\sim\mathcal{U}(k_a, K)}\mathbb{E}_{\bm{x}^k\vert \bm{x}^{k_a}}\bigg[ \Vert
\frac{\bar{\alpha}_{k_a}}{\sqrt{\bar{\alpha}_{k} \cdot \bar{\alpha}_{k_a}}} \bm{x}^k - \frac{(\bar{\alpha}_{k_a} - \bar{\alpha}_{k}) \bm{\epsilon}_{\theta}(\bm{x}^k, k)}{\sqrt{\bar{\alpha}_{k_a} \cdot \bar{\alpha}_{k} \cdot (1- \bar{\alpha}_{k})}}
- \bm{x}^{k_a}
\Vert^2\bigg]
\end{aligned}
\end{equation}
have $\bm{\epsilon}_{\theta^*}(\bm{x}^k, k) = \mathbb{E}[\bm{\epsilon}\vert\bm{x}^k], \forall k\geq k_a$
(cf. Appendix~\ref{proof_coll_41} for detailed proof).
\end{corollary}
Notably, Eq.~\eqref{corollay_eq4} cannot be directly applied in our settings because the noise-consistent data, $\bm{x}^{k_a}$, are not directly accessible. As discussed in Section~\ref{4_experiment_setting}, some portions of our dataset are corrupted while others remain clean.
\paragraph{Ambient DDPM from Approximated Distribution.}
{It is worth noting that all samples involved in the Ambient DDPM training process follow Gaussian distributions with parameters that are functions of $k$.
While direct access to noise-consistent data is not available, an ideal alternative is to approximate these distributions and train the diffusion model using the approximations.}
% Though we do not have direct access to noise-consistent data, as an ideal alternative, we can approximate these distributions and train the diffusion model using the approximated distributions.
Let $q (\bm{x}^k\vert\bm{x}^0)$ denote the ground-truth distribution of samples generated by the DDPM forward process at diffusion timestep $k$, as described in Eq.~\eqref{forward_diffusion}, and let $\varrho(\bm{x}^k\vert\bm{x}^0)$ represent an approximation of this distribution.
{To ensure effective learning under this approximation, we make the following.}
\begin{assumption}\label{assumption22}
There exists a positive constant $c$ such that, for any $k \geq k_a$, if the Kullback-Leibler (KL) divergence satisfies $\mathbb{D}_{KL}[q (\bm{x}^k\vert\bm{x}^0) \Vert \varrho(\bm{x}^k\vert\bm{x}^0)] < c$, then the ambient DDPM with $k \geq k_a$, as introduced in Corollary~\ref{corollary21}, can be effectively learned from samples drawn from the approximated distribution $\varrho(\bm{x}^k\vert\bm{x}^0)$.
\end{assumption}
% Basic idea: KL divergence between noised/clean <=c. for any c>0, we can always find ka corresponding to it.
{Following the standard setup of corruption-robust offline RL, we are provided with samples $\check{\bm{x}}$ that may or may not contain scaled Gaussian noise $\iota\cdot\bm{\epsilon}$. The distribution of such samples at diffusion timestep $k$ in the DDPM forward process is denoted as $q(\check{\bm{x}}^k\vert\check{\bm{x}}^0)$.
We have the following.}
\begin{theorem}\label{theorem23}
Let Assumption~\ref{assumption22} hold. For any bounded noise scale $\iota$, one can always find a diffusion timestep $k_a$ such that ambient DDPM with $k\geq k_a$, {which should have been learned from samples drawn from $q(\bm{x}^k\vert\bm{x}^0)$}, can instead be effectively learned from samples drawn from $q(\check{\bm{x}}^k\vert\check{\bm{x}}^0)$.
{That is, for any $c>0$, one can always find $k_a$ such that for any $k\geq k_a$, the following inequality holds:
$\mathbb{D}_{KL}[q (\bm{x}^k\vert\bm{x}^0) \Vert q (\check{\bm{x}}^k\vert\check{\bm{x}}^0)] < c$.}
\end{theorem}
We provide the detailed proof in Appendix~\ref{proof_therorm_43}. It is evident that a smaller $k_a$ retains more of the original information, thereby improving the accuracy of the noise predictor. However, a smaller $k_a$ also introduces a more relaxed threshold $c$, which may result in a larger discrepancy between the original distribution and the approximated distribution. The choice of $k_a$ necessitates a trade-off between these two factors. See Section~\ref{ablation_main} for detailed ablation studies.
\subsection{Corrupted Samples Detection}
Given a sample $\check{\bm{x}}$ that may or may not contain scaled Gaussian noise $\iota\cdot\bm{\epsilon}$, we propose using the squared Frobenius norm of the noise predictor, $\Vert \bm{\epsilon}_{\theta}(\check{\bm{x}}, k) \Vert_F^2$, to determine whether noise is present.
\begin{proposition}\label{prop34}
Assume the noise prediction error for $\bm{\epsilon}_{\theta}(\cdot, k)$ follows $\bm{\delta}_{\theta}^k \sim \mathcal{N}(\bm{0}, \sigma^2_k\bm{I})$.
Define the difference between the noisy and noise-free cases as: $\Delta=\mathbb{E}[\Vert \bm{\epsilon}_{\theta}(\check{\bm{x}}_{\mathrm{ns}}^k, k) \Vert_F^2] - \mathbb{E}[\Vert \bm{\epsilon}_{\theta}(\check{\bm{x}}_{\mathrm{nf}}^k, k) \Vert_F^2]$, where $\check{\bm{x}}_{\mathrm{ns}}^k$ and $\check{\bm{x}}_{\mathrm{nf}}^k$ denote noisy and noise-free samples with original information consistency operation, respectively.
The Signal-to-Noise Ratio (SNR) of the prediction is then expressed as:
\begin{equation}\label{eq_SNR}
\mathrm{SNR}(k) := \frac{\Delta}{\mathbb{E}[\Vert \bm{\epsilon}_{\theta}(\check{\bm{x}}_{\mathrm{nf}}, k) \Vert_F^2]}
= \frac{\iota^2 \cdot \bar{\alpha}_k}{(1- \bar{\alpha}_k) \cdot \sigma^2_k}.
\end{equation}
\end{proposition}
See Appendixa~\ref{proof_propsition_44} for detailed proof. Assume that the noise prediction error are the same across all diffusion timesteps, i.e., $\sigma_k=\sigma$ for any $k$, then $\mathrm{SNR}(k)$ achieves maximum value at $k=k_a$, {as $\bar{\alpha}_k$ is a strictly monotonically decreasing function of $k$.}
\subsection{Ambient Diffusion-Guided Dataset Recovery (ADG)}\label{method_recovery}
Having established the theoretical foundation, we now proceed to introduce our proposed ADG method for corruption-robust offline RL.
As there are two timesteps involved, we use superscripts $k$ to denote the diffusion timesteps and subscripts $t$ to denote the RL timesteps for clarity.
The ground truth trajectory matrix is defined as $\bm{\tau}_t := [\bm{z}_{t-H}, \dots, \bm{z}_{t+H}]\in \mathbb{R}^{M\times (2H+1)}$, where $\bm{z} \in \mathbb{R}^M$ represents the RL component (which can correspond to an observation $\bm{s}$ or a state-action-reward triplet $(\bm{s}, \bm{a}, r)$), $M$ denotes the dimensionality of $\bm{z}$, and $H$ specifies the temporal slice size.
Let $\check{\bm{z}}_t$ denote the RL component that may or may not contain scaled noise $\iota\cdot\bm{\epsilon}$. We only have access to the observed (partially corrupted) trajectory $\check{\bm{\tau}}_t = [\check{\bm{z}}_{t-H}, \dots, \check{\bm{z}}_{t+H}]\in \mathbb{R}^{M\times (2H+1)}$.
\paragraph{Ambient DDPM for Corrupted Samples Detection.}
Following Theorem~\ref{theorem23}, given a partially corrputed offline RL datastet, we pre-define $k_a$ and train the ambient DDPM through
\begin{equation}
\begin{aligned}
\mathbb{E}_{\check{\bm{\tau}}_t^{k_a}\sim\mathcal{D}, k\sim\mathcal{U}(k_a, K), \check{\bm{\tau}}_t^k\vert \check{\bm{\tau}}_t^{k_a}}\bigg[
\Vert \frac{\bar{\alpha}_{k_a}}{\sqrt{\bar{\alpha}_{k} \cdot \bar{\alpha}_{k_a}}} \check{\bm{\tau}}_t^k - \frac{(\bar{\alpha}_{k_a} - \bar{\alpha}_{k}) \bm{\epsilon}_{\theta}(\check{\bm{\tau}}_t^k, k)}{\sqrt{\bar{\alpha}_{k_a} \cdot \bar{\alpha}_{k} \cdot (1- \bar{\alpha}_{k})}}
- \check{\bm{\tau}}_t^{k_a}
\Vert^2\bigg] &,
\end{aligned}
\end{equation}
where $\check{\bm{\tau}}_t^{k}=\sqrt{\frac{\bar{\alpha}_k}{\bar{\alpha}_{k_a}}} \check{\bm{\tau}}_t^{k_a} + \sqrt{\frac{\bar{\alpha}_{k_a}-\bar{\alpha}_k}{\bar{\alpha}_{k_a}}} \bm{\epsilon}$. Once the training converges, we obtain the noise predictor $\bm{\epsilon}_{\theta}(\check{\bm{\tau}}_t, k)$, which achieves the largest SNR at $k=k_a$ as described in Proposition~\ref{prop34}.
It should be noted that for each sample $\check{\bm{\tau}}_t$, we focus solely on whether the RL component at the center position ($\check{\bm{z}}_t$) is corrupted, rather than evaluating the entire trajectory slice.
For this purpose, we further define $e_\theta(\check{\bm{z}}_t) = \Vert \bm{\epsilon}_{\theta}(\check{\bm{\tau}}_t, k_a)_{H+1} \Vert_F^2$, where $(\cdot)_{H+1}$ represents the (H+1)-th column of the matrix.
Subsequently, we utilize $e_\theta(\check{\bm{z}}_t)$ to evaluate every samples within the partially corrupted dataset, and rescale the prediction range to \([0, 1]\). With a manually defined threshold \(\zeta\), samples with \(e_\theta(\check{\bm{z}_t}) > \zeta\) are classified as noised samples ($\mathcal{D}_n$), while the remaining samples are considered clean ($\mathcal{D}_c$).
\paragraph{Naive DDPM for Corrupted Samples Recovery.}
Once the corrupted samples have been identified, the remaining uncorrupted samples can be utilized to train a denoiser (through naive DDPM), which can subsequently be applied to correct the corrupted samples.
To avoid overfitting of DDPM to misclassified noisy data, we reuse $\zeta$ to filter out training data for the naive DDPM.
We denote $\mathbb{I}_t$ as a binary indicator variable that specifies whether $\check{\bm{z}}_t$ is corrupted ($\mathbb{I}_t = 0$ for $e_\theta(\check{\bm{z}}_t)>\zeta$) or not ($\mathbb{I}_t = 1$ for $e_\theta(\check{\bm{z}}_t)\leq\zeta$). Given the mask defined as $\bm{m}_t:=[\mathbb{I}_{t-H}, \dots, \mathbb{I}_{t+H}]$, the training loss for naive DDPM follows
\begin{equation}\label{ddpm_loss}
\mathbb{E}_{\check{\bm{\tau}}_t\sim\mathcal{D}_c, k \sim [1,K],\bm{\epsilon} \sim \mathcal{N}(\bm{0},\bm{I})}[\Vert[\bm{\epsilon}_{\phi}(\check{\bm{\tau}}_t^k,k) - \bm{\epsilon})] \odot \bm{m}_t \Vert^2].
\end{equation}
where $\odot$ is the Hadamard product. We refer to this training process as \textit{selective training} in the following discussion. After the training coverges, we then conduct revese DDPM process $p_\phi(\check{\bm{\tau}}^{0:k_a}_t)$ as described in Eq.~\eqref{reverse_diffusion} to all samples within $\mathcal{D}_n$. Finally, we combine the denoised \(D_n\) with \(D_c\) to form the final dataset.
More implementation details of ADG can be found in~\cref{app:pseudocode,implementation_details}.
\section{Experiments}
\label{sec:experiments}
In this section, we conduct comprehensive experiments to empirically evaluate ADG by exploring three key questions: (1) How does ADG enhance the performance of both non-robust and robust offline RL methods across various data corruption scenarios? (2) What is the individual contribution of ambient loss and selective training to the overall effectiveness of ADG? (3) What are the advantages of ADG's structure, which incorporates two independent diffusion models?
\input{section/4_1_table}
\subsection{Experimental Setups}\label{4_experiment_setting}
We assess ADG on various widely used offline RL benchmarks~\cite{fu2020d4rl}, including MuJoCo, Kitchen, and Adroit. Since prior work~\cite{xu2024robust} has shown that the impact of data corruption becomes more severe when data is limited, we further evaluate the effectiveness of ADG across different dataset scales by conducting experiments on down-sampled tasks. Specifically, we down-sample 10\% and 1\% of the data from MuJoCo and Adroit tasks as the new testbed. We also include results under different down-sample ratio of 2\% and 5\% in Appendix~\ref{ablation_dataset_size}. For MuJoCo, we select ``medium-replay'' datasets. Additional results on ``medium'' and ``expert'' datasets are deferred to~\cref{ablation_different_quality_levels}.
For data corruption, we consider two data corruption scenarios: \textit{random} and \textit{adversarial} corruption, as introduced in \cref{sec:preliminaries}. These scenarios are applied either to states alone or to state-action-reward triplets (full-element). Following the settings of prior works~\cite{yang2023towards}, we set the corruption rate $\eta$ to 0.3 and the corruption scale $\alpha$ to 1.0. We provide the further implementation details on data corruption in Appendix~\ref{noised_setting}. We also investigate ADG under Gaussian noise corruption and under different noise ratio and scales in Appendix~\ref{ablation_gaussian}, \ref{ablation_corruption_rates_and_scales}.
We evaluate a diverse set of offline RL methods, including pessimistic value estimation method CQL~\cite{kumar2020conservative}, policy constraint methods along with their robust variants IQL~\cite{kostrikov2021offline} and RIQL~\cite{yang2023towards}, as well as sequence modeling methods such as DT~\cite{chen2021decision} and RDT~\cite{xu2024robust}. To ensure the robustness of the findings, each experiment is conducted using four distinct random seeds, with the standard deviation across these seeds also reported.
\subsection{Evaluation under Different Data Corruption}\label{exp:random_ad}
\textbf{Results under Random Corruption.}
We evaluate the improvement brought by ADG on various offline RL algorithms under random data corruption. The average scores presented in \cref{tab:random_experiment_results} show that ADG provides benefits to both non-robust and robust algorithms across all scenarios. Notably, ADG brings significant improvements to MDP-based algorithms (CQL, IQL, and RIQL), with an average performance boost of 69.1\%, demonstrating the effectiveness of sequence completion for missing information. ADG also provides substantial improvements to non-MDP algorithms (DT and RDT), with an average performance boost of 17.4\%.
When equipped with ADG, both IQL and DT outperform their \textit{Naive} robust variants in nearly all scenarios. This highlights ADG's core strength: by modifying only the dataset, it enables standard algorithms to surpass their robust variants explicitly designed for noise resistance. Moreover, this advantage can be further amplified using robust variants, achieving nearly a 60\% improvement on RIQL. Additionally, when equipped with ADG, most algorithms achieve similar performance under both state and full-element corruption, further demonstrating the scalability of ADG.
We provide visualizations of ADG's detection and denoising in \cref{ablation_visual_detection,ablation_visual_denoising} to clearly explain its effectiveness, and investigate the effect of using ADG solely for filtering in \cref{ablation_filted_vs_restored}.
\input{section/4_2_table}
\textbf{Results under Adversarial Corruption.} We further examine the robustness of ADG under adversarial data corruption. The results, summarized in Table~\ref{tab:adversarial_experiment_results}, show that ADG consistently improves baseline performance by an average of 24.41\%.
Notably, when equipped with ADG, the baselines IQL and DT outperform their \textit{Naive} robust variants RIQL and RDT across all scenarios. This further supports the conclusions drawn from the random corruption scenarios. These findings highlight ADG's ability to adapt to and mitigate adversarial data corruption.
\begin{figure}[hbt]
\centering
\includegraphics[width=0.49\textwidth]{pics/ablation_ambient_plot_big.pdf}
\includegraphics[width=0.49\textwidth]{pics/ablation_ambient_D4RL_performance_big.pdf}
\caption{
The FN rate during detector training (left), and the performance of IQL and DT using these detectors (right). ``Kit-c'' denotes Kitchen-complete, and ``Re'' denotes Relocate.
}
\label{fig:ablation_ambient}
\end{figure}
\subsection{Ablation Study}\label{ablation_main}
We conduct ablation studies to analyze the impact of each component on ADG's performance.
\textbf{Impact of Ambient Loss.} We assess the impact of ambient loss on detector performance by varying $t_n \in \{15, 30, 50\}$ on the ``kitchen-complete-v0'' and ``relocate-expert-v1'' datasets, selected for their complexity. Following Section~\ref{4_experiment_setting}, Random State Attacks are introduced. Samples detected as corrupted are labeled as positive and others as negative. The detector is trained for 5k steps. We detect the dataset with \( e_\theta(\check{\bm{z}}) \leq \zeta \) (as in~\cref{method_recovery}) at each step and plot false negatives (FNs), which represent the proportion of undetected corrupted samples. We also evaluate the D4RL scores of baseline algorithms on the datasets recovered using ADG with these trained detectors. As shown in Figure~\ref{fig:ablation_ambient}, the detector trained with the naive diffusion loss shows some detection capability initially, but quickly overfits to the corrupted portion of the training data. Ambient loss significantly improves it. Notably, $t_n = 30$, which is also used in the main experiments, achieves the lowest false negatives (FNs), implying that nearly all selected samples remain unaltered by attacks. This ensures that the denoiser is trained on nearly clean data, making the naive diffusion loss feasible. Moreover, the D4RL scores of baseline algorithms exhibit a clear correlation with detection performance.
\begin{wrapfigure}{r}{0.50\textwidth}
% \begin{figure}[htb]
\centering
% 子图 1
\includegraphics[width=0.50\textwidth]{pics/ablation_detecting_zeta.pdf}
\caption{
Results on ``walker2d-medium-replay-v2'' with Random State Attack (Appendix~\ref{noised_setting}): (a) MSE vs. ground truth, and (b) DT performance across $\zeta$. Best results at $\zeta = 0.20$.
}
\label{fig:ablation_zeta_detector}
% \end{figure}
\end{wrapfigure}
\textbf{Impact of Selective Training.} To evaluate the impact of selective training on denoiser, we vary $\zeta$ within the range $\{0.05, 0.10, 0.20, 0.50\}$ and measure the mean squared error (MSE) between the recovered dataset and the ground truth. The results are shown in Figure~\ref{fig:ablation_zeta_detector}. A very low $\zeta=0.05$ leads to poor denoiser performance and eventual overfitting, likely due to the insufficient information in a small dataset. Therefore, a very high $\zeta=0.50$ also degrades performance, possibly due to excessive corrupted data in the dataset. Moderate values of $\zeta=0.10 \ \text{or} \ 0.20$ yield similar, good performance, supporting the necessity of selective training and suggesting that $\zeta$ is somewhat robust, performing well within a certain range. The D4RL scores of the baseline algorithms exhibit a strong correlation with the performance of the denoiser, further supporting the necessity of using $\zeta$ for the selective training of the denoiser.
% \begin{figure}[htb]
\begin{wrapfigure}{r}{0.48\textwidth}
\centering
% 子图 1
\includegraphics[width=0.45\textwidth]{pics/ablation_single_diffusion_main.pdf}
\caption{
%Comparison of Naive, ADG (w/ SD: single diffusion), and ADG (w/ separate diffusion) for IQL (top) and DT (bottom).
Comparison results among Naive, ADG with a single diffusion model, and ADG with separate diffusion models for IQL (upper) and DT (lower). ADG (w/ SD) denotes ADG using a single diffusion model.
}
\label{fig:ablation_single_df}
% \end{figure}
\end{wrapfigure}
% Single Diffusion Ablation on Walker2d, Kitchen, Relocate
\textbf{Impact of Using Two Separate Diffusion Models for Detection and Denoising.}
As described in Section~\ref{method_recovery}, our approach employs a structure with two independent diffusion models: one serving as the detector and the other as the denoiser. However, it is technically possible to employ a single diffusion model for both detection and denoising tasks by training it concurrently with both Eq.~\ref{corollay_eq4} and Eq.~\ref{ddpm_loss}. We conduct ablation experiments on the ``walker2d-medium-replay-v2'' dataset under Random State Attack, using both full and limited dataset sizes to evaluate this configuration. The results are presented in Figure~\ref{fig:ablation_single_df}. Although ADG with single diffusion consistently improves the performance of all baselines, ADG with two independent diffusion models always outperforms ADG with a single model. This result is intuitive, as applying two different losses to a single diffusion model can easily cause interference, leading the network to converge to local minima. Using two independent models effectively addresses this issue by decoupling the mutual interference. For more implementation details and results, refer to Appendix~\ref{appendix_single_diffusion}.
Additionally, we also include the ablation study on the length of the slice window by varying the values of $H$ in Appendix~\ref{ablation_different_H}. The performance of ADG shows a positive correlation with the hyperparameter $H$ as it increases from 0, and becomes robust to further changes once $H$ reaches a certain range. This demonstrates the importance of incorporating sequential information and the robustness of ADG.
% \begin{table*}[ht]
% \centering
% \caption{Performance under random data corruption. \textit{Naive} refers to performance evaluated on the original corrupted dataset. \textit{AGD} refers to performance evaluated on the dataset recovered by our AGD.}
% \label{tab:random_experiment_results}
% \resizebox{\textwidth}{!}{
% \begin{tabular}{l|l|c>{\columncolor{gray!20}}cc>{\columncolor{gray!20}}cc>{\columncolor{gray!20}}cc>{\columncolor{gray!20}}cc>{\columncolor{gray!20}}c}
% %\hline
% \toprule
% \multirow{2}{*}{Attack} & \multirow{2}{*}{Task} & \multicolumn{2}{c}{CQL} & \multicolumn{2}{c}{IQL} & \multicolumn{2}{c}{RIQL} & \multicolumn{2}{c}{DT} & \multicolumn{2}{c}{RDT} \\
% & & Naive & AGD & Naive & AGD & Naive & AGD & Naive & AGD & Naive & AGD\\
% \hline
% \midrule
% \multirow{13}{*}{\rotatebox{90}{{\Large State}}}
% & halfcheetah & 23.7$\pm$4.4 & 41.5$\pm$1.7 & 20.2$\pm$2.6 & 28.9$\pm$3.8 & 19.9$\pm$2.1 & 36.2$\pm$4.2 & 27.5$\pm$2.5 & 39.8$\pm$0.4 & 30.8$\pm$1.8 & 34.2$\pm$1.4 \\
% & hopper & 34.4$\pm$22.8 & 78.8$\pm$11.1 & 56.0$\pm$18.4 & 56.8$\pm$21.8 & 34.0$\pm$13.4 & 62.2$\pm$4.9 & 51.3$\pm$14.0 & 79.1$\pm$6.7 & 56.6$\pm$2.9 & 65.2$\pm$7.5 \\
% & walker2d & 33.7$\pm$10.8 & 66.0$\pm$1.9 & 20.5$\pm$6.0 & 35.0$\pm$12.7 & 14.2$\pm$1.2 & 43.3$\pm$5.2 & 47.6$\pm$4.9 & 55.1$\pm$8.9 & 53.4$\pm$4.0 & 66.4$\pm$2.9 \\
% \cline{2-12}\rule{0pt}{2.5ex}
% & halfcheetah(10\%) & 9.7$\pm$0.8 & 13.8$\pm$9.6 & 2.8$\pm$0.2 & 4.6$\pm$0.7 & 4.4$\pm$0.9 & 13.1$\pm$1.7 & 6.3$\pm$0.4 & 26.4$\pm$3.0 & 8.3$\pm$1.5 & 22.4$\pm$0.8 \\
% & hopper(10\%) & 2.2$\pm$1.5 & 7.7$\pm$4.8 & 14.2$\pm$5.6 & 11.6$\pm$9.8 & 15.5$\pm$5.4 & 36.7$\pm$5.5 & 36.1$\pm$7.6 & 37.4$\pm$9.4 & 40.8$\pm$3.5 & 42.4$\pm$7.9 \\
% & walker2d(10\%) & 1.2$\pm$1.8 & 20.5$\pm$7.7 & 6.6$\pm$0.8 & 10.9$\pm$2.9 & 9.2$\pm$4.4 & 22.5$\pm$2.7 & 18.0$\pm$2.5 & 34.0$\pm$5.5 & 20.3$\pm$2.8 & 20.5$\pm$2.7 \\
% \cline{2-12}\rule{0pt}{2.5ex}
% & kitchen-complete & 12.9$\pm$4.3 & 8.8$\pm$5.4 & 33.6$\pm$7.3 & 50.0$\pm$3.1 & 37.5$\pm$6.4 & 52.5$\pm$5.9 & 37.0$\pm$6.2 & 61.9$\pm$12.4 & 52.8$\pm$1.8 & 58.1$\pm$10.1 \\
% & kitchen-partial & 0.0$\pm$0.0 & 7.4$\pm$1.8 & 13.5$\pm$3.4 & 14.4$\pm$9.4 & 25.9$\pm$3.4 & 48.1$\pm$11.1 & 31.0$\pm$8.1 & 43.8$\pm$1.3 & 36.8$\pm$5.8 & 49.4$\pm$11.1 \\
% & kitchen-mixed & 0.0$\pm$0.0 & 15.0$\pm$1.8 & 16.2$\pm$5.6 & 41.9$\pm$3.7 & 21.6$\pm$3.7 & 27.5$\pm$9.4 & 31.8$\pm$3.4 & 36.3$\pm$13.1 & 41.8$\pm$4.3 & 37.5$\pm$17.6 \\
% \cline{2-12}\rule{0pt}{2.5ex}
% & door(1\%) & -0.3$\pm$0.0 & 4.4$\pm$5.0 & 46.6$\pm$17.5 & 58.6$\pm$11.7 & 39.0$\pm$16.4 & 103.1$\pm$1.6 & 94.6$\pm$4.2 & 102.9$\pm$0.4 & 102.8$\pm$2.4 & 103.9$\pm$2.2 \\
% & hammer(1\%) & 0.2$\pm$0.0 & -0.4$\pm$0.5 & 64.6$\pm$17.3 & 67.2$\pm$9.7 & 70.0$\pm$12.6 & 100.0$\pm$18.9 & 97.8$\pm$12.3 & 115.7$\pm$0.7 & 113.8$\pm$1.6 & 126.6$\pm$26.9 \\
% & relocate(1\%) & -0.2$\pm$0.0 & 0.8$\pm$0.8 & 9.4$\pm$3.5 & 12.9$\pm$6.6 & 5.2$\pm$5.0 & 22.5$\pm$9.2 & 61.6$\pm$5.6 & 67.6$\pm$4.0 & 65.0$\pm$6.2 & 67.2$\pm$4.8 \\
% \cline{2-12}\rule{0pt}{2.5ex}
% & \textbf{Average Score} & 7.6 & \textbf{22.0} & 25.3 & \textbf{32.7} & 24.1 & \textbf{47.3} & 45.9 & \textbf{58.3} & 52.5 & \textbf{57.8} \\
% & \textbf{Improvement $\uparrow$} & & \textbf{189.47\%} & & \textbf{29.25\%} & & \textbf{96.27\%} & & \textbf{27.02\%} & & \textbf{10.10\%} \\
% \midrule
% \hline\rule{0pt}{2.5ex}
% \multirow{13}{*}{\rotatebox{90}{{\Large Full-Eelement}}}
% & halfcheetah & -0.7$\pm$2.2 & 15.5$\pm$1.4 & 21.8$\pm$2.9 & 32.1$\pm$3.8 & 22.5$\pm$2.3 & 33.2$\pm$2.8 & 28.0$\pm$5.0 & 39.2$\pm$0.9 & 20.7$\pm$2.9 & 38.1$\pm$1.2 \\
% & hopper & 34.0$\pm$4.2 & 54.3$\pm$16.2 & 35.2$\pm$19.9 & 38.8$\pm$20.0 & 16.4$\pm$3.2 & 64.5$\pm$5.2 & 53.0$\pm$14.2 & 55.1$\pm$11.9 & 58.6$\pm$11.1 & 51.2$\pm$20.1 \\
% & walker2d & 16.8$\pm$3.1 & 45.6$\pm$9.2 & 22.4$\pm$8.5 & 28.7$\pm$16.5 & 17.5$\pm$3.6 & 33.7$\pm$2.4 & 51.0$\pm$14.5 & 55.3$\pm$1.9 & 56.5$\pm$11.0 & 56.1$\pm$15.4 \\
% \cline{2-12}\rule{0pt}{2.5ex}
% & halfcheetah(10\%) & 7.6$\pm$3.7 & 22.7$\pm$6.6 & 1.3$\pm$0.4 & 9.0$\pm$4.6 & 2.0$\pm$0.5 & 10.2$\pm$1.8 & 10.0$\pm$4.5 & 28.3$\pm$2.6 & 16.8$\pm$5.1 & 25.2$\pm$2.1 \\
% & hopper(10\%) & 2.0$\pm$1.2 & 10.4$\pm$4.7 & 7.8$\pm$1.2 & 17.7$\pm$7.0 & 11.6$\pm$2.4 & 39.4$\pm$6.7 & 28.3$\pm$8.8 & 35.9$\pm$6.0 & 36.5$\pm$4.7 & 38.4$\pm$7.6 \\
% & walker2d(10\%) & 0.9$\pm$1.6 & 16.0$\pm$5.8 & 4.7$\pm$1.1 & 22.6$\pm$9.0 & 4.7$\pm$1.2 & 32.4$\pm$11.4 & 18.9$\pm$4.0 & 32.0$\pm$16.6 & 24.5$\pm$5.9 & 43.4$\pm$9.3 \\
% \cline{2-12}\rule{0pt}{2.5ex}
% & kitchen-complete & 4.4$\pm$3.2 & 5.6$\pm$3.7 & 23.8$\pm$6.0 & 47.5$\pm$7.1 & 26.$\pm$5.0 & 50.0$\pm$1.8 & 51.3$\pm$9.4 & 55.0$\pm$5.3 & 60.6$\pm$5.7 & 60.0$\pm$4.7 \\
% & kitchen-partial & 1.9$\pm$2.1 & 9.4$\pm$6.5 & 1.1$\pm$1.4 & 33.1$\pm$3.7 & 0.5$\pm$0.9 & 30.0$\pm$8.1 & 34.4$\pm$7.4 & 37.5$\pm$7.3 & 45.6$\pm$23.7 & 43.8$\pm$9.4 \\
% & kitchen-mixed & 0.0$\pm$0.0 & 9.4$\pm$4.1 & 0.8$\pm$0.8 & 35.0$\pm$2.5 & 8.4$\pm$1.8 & 27.5$\pm$8.3 & 21.9$\pm$13.3 & 23.8$\pm$22.4 & 39.4$\pm$17.1 & 35.6$\pm$9.7 \\
% \cline{2-12}\rule{0pt}{2.5ex}
% & door(1\%) & -0.3$\pm$0.0 & 4.5$\pm$7.8 & 47.1$\pm$12.7 & 54.6$\pm$23.0 & 64.3$\pm$11.1 & 68.0$\pm$11.9 & 33.9$\pm$24.3 & 45.1$\pm$26.3 & 92.1$\pm$17.5 & 105.6$\pm$2.0 \\
% & hammer(1\%) & 0.2$\pm$0.0 & 0.1$\pm$0.0 & 65.7$\pm$12.7 & 78.7$\pm$12.2 & 88.7$\pm$18.0 & 87.8$\pm$17.9 & 26.3$\pm$16.9 & 43.3$\pm$38.4 & 113.1$\pm$23.8 & 109.2$\pm$17.3 \\
% & relocate(1\%) & -0.3$\pm$0.0 & 0.3$\pm$0.3 & 7.2$\pm$2.0 & 10.5$\pm$6.4 & 12.5$\pm$1.5 & 24.4$\pm$2.4 & 35.9$\pm$18.7 & 42.0$\pm$21.1 & 56.7$\pm$13.2 & 61.8$\pm$3.4 \\
% \cline{2-12}\rule{0pt}{2.5ex}
% & \textbf{Average} & 5.5 & \textbf{16.2} & 19.9 & \textbf{34.0} & 22.9 & \textbf{41.8} & 32.7 & \textbf{50.8} & 53.0 & \textbf{54.8} \\
% & \textbf{Improvement $\uparrow$} & & \textbf{194.55\%} & & \textbf{70.85\%} & & \textbf{82.53\%} & & \textbf{55.35\%} & & \textbf{3.40\%} \\
% \bottomrule
% \end{tabular}
% }
% \end{table*}
\begin{table*}[ht]
\centering
\caption{Performance under random data corruption. Results are averaged over four random seeds.}
\label{tab:random_experiment_results}
\resizebox{\textwidth}{!}{
\begin{tabular}{l|l|c>{\columncolor{gray!20}}cc>{\columncolor{gray!20}}cc>{\columncolor{gray!20}}cc>{\columncolor{gray!20}}cc>{\columncolor{gray!20}}c}
%\hline
\toprule
Attack & \multirow{2}{*}{Task} & \multicolumn{2}{c}{CQL} & \multicolumn{2}{c}{IQL} & \multicolumn{2}{c}{RIQL} & \multicolumn{2}{c}{DT} & \multicolumn{2}{c}{RDT} \\
Element & & Naive & ADG & Naive & ADG & Naive & ADG & Naive & ADG & Naive & ADG\\
\hline
\midrule
\multirow{13}{*}{\rotatebox{90}{{\Large State}}}
& halfcheetah & 15.9$\pm$1.8 & 23.9$\pm$5.6 & 19.2$\pm$2.2 & 28.8$\pm$1.7 & 19.9$\pm$2.1 & 27.8$\pm$4.5 & 27.5$\pm$2.5 & 39.8$\pm$0.4 & 30.8$\pm$1.8 & 34.2$\pm$1.4 \\
& hopper & 55.4$\pm$6.4 & 78.8$\pm$11.1& 47.6$\pm$7.1 & 72.2$\pm$7.1 & 34.0$\pm$13.4 & 66.3$\pm$15.9& 51.3$\pm$14.0 & 79.1$\pm$6.7 & 56.6$\pm$2.9 & 65.2$\pm$7.5 \\
& walker2d & 39.9$\pm$5.2 & 46.0$\pm$9.7 & 17.5$\pm$6.8 & 27.9$\pm$3.7 & 14.2$\pm$1.2 & 39.5$\pm$4.3 & 47.6$\pm$4.9 & 55.1$\pm$8.9 & 53.4$\pm$4.0 & 66.4$\pm$2.9 \\
\cline{2-12}\rule{0pt}{2.5ex}
& halfcheetah(10\%) & 11.0$\pm$1.1 & 17.3$\pm$2.7& 6.1$\pm$1.3 & 12.0$\pm$0.9 & 4.4$\pm$0.9 & 8.3$\pm$2.2 & 6.3$\pm$0.4 & 26.4$\pm$3.0 & 8.3$\pm$1.5 & 22.4$\pm$0.8 \\
& hopper(10\%) & 1.8$\pm$0.7 & 3.1$\pm$0.6 & 13.3$\pm$3.3 & 19.6$\pm$6.3 & 15.5$\pm$5.4 & 15.7$\pm$6.1 & 36.1$\pm$7.6 & 37.4$\pm$9.4 & 40.8$\pm$3.5 & 42.4$\pm$7.9 \\
& walker2d(10\%) & -0.0$\pm$0.1 & 1.0$\pm$1.1 & 10.9$\pm$7.2 & 15.7$\pm$2.8 & 9.2$\pm$4.4 & 10.1$\pm$4.3 & 18.0$\pm$2.5 & 34.0$\pm$5.5 & 20.3$\pm$2.8 & 20.5$\pm$2.7 \\
\cline{2-12}\rule{0pt}{2.5ex}
& kitchen-complete & 3.8$\pm$2.8 & 15.0$\pm$6.8& 33.6$\pm$7.3 & 51.2$\pm$4.5 & 37.5$\pm$6.4 & 52.5$\pm$4.3 & 37.0$\pm$6.2 & 61.9$\pm$12.4 & 52.8$\pm$1.8 & 58.1$\pm$10.1 \\
& kitchen-partial & 0.0$\pm$0.0 & 0.6$\pm$1.1 & 13.5$\pm$3.4 & 23.8$\pm$2.2 & 25.9$\pm$3.4 & 26.9$\pm$7.6 & 31.0$\pm$8.1 & 43.8$\pm$1.3 & 36.8$\pm$5.8 & 49.4$\pm$11.1 \\
& kitchen-mixed & 0.0$\pm$0.0 & 0.0$\pm$0.0 & 16.2$\pm$5.6 & 41.2$\pm$4.1 & 21.6$\pm$3.7 & 41.9$\pm$3.2 & 31.8$\pm$3.4 & 36.3$\pm$13.1 & 41.8$\pm$4.3 & 37.5$\pm$17.6 \\
\cline{2-12}\rule{0pt}{2.5ex}
& door(1\%) & -0.3$\pm$0.0 &-0.4$\pm$0.0& 46.6$\pm$17.5 & 63.3$\pm$9.8 & 39.0$\pm$16.4 & 77.5$\pm$10.2& 94.6$\pm$4.2 & 102.9$\pm$0.4 & 102.8$\pm$2.4& 103.9$\pm$2.2 \\
& hammer(1\%) & 0.2$\pm$0.0 & 0.2$\pm$0.0& 64.6$\pm$17.3 & 78.2$\pm$17.3 & 70.0$\pm$12.6 & 88.4$\pm$5.2 & 97.8$\pm$12.3 & 115.7$\pm$0.7 & 113.8$\pm$1.6& 126.6$\pm$26.9 \\
& relocate(1\%) & -0.3$\pm$0.1 &-0.3$\pm$0.1& 9.4$\pm$3.5 & 14.4$\pm$4.9 & 5.2$\pm$5.0 & 27.8$\pm$4.8 & 61.6$\pm$5.6 & 67.6$\pm$4.0 & 65.0$\pm$6.2 & 67.2$\pm$4.8 \\
\cline{2-12}\rule{0pt}{2.5ex}
& \textbf{Average Score} & 10.6 & \textbf{15.4} & 24.9& \textbf{37.4} & 24.7& \textbf{40.2} & 45.0& \textbf{58.3} & 51.9& \textbf{57.8} \\
& \textbf{Improvement $\uparrow$} & & \textbf{45.28\%} & & \textbf{50.20\%} & & \textbf{62.75\%} & & \textbf{29.56\%} & & \textbf{11.37\%} \\
\midrule
\hline\rule{0pt}{2.5ex}
\multirow{13}{*}{\rotatebox{90}{{\Large Full-Eelement}}}
& halfcheetah & 0.3$\pm$0.2&33.9$\pm$2.6& 15.5$\pm$1.5 & 37.4$\pm$1.6 & 22.5$\pm$2.3 & 32.3$\pm$3.6 & 28.0$\pm$5.0 & 39.2$\pm$0.9 & 20.7$\pm$2.9 & 38.1$\pm$1.2 \\
& hopper &0.7$\pm$0.0 &20.3$\pm$7.3& 26.5$\pm$19.6& 36.1$\pm$20.1 & 16.4$\pm$3.2 & 49.4$\pm$19.7& 53.0$\pm$14.2 & 55.1$\pm$11.9 & 58.6$\pm$11.1 & 51.2$\pm$20.1\\
& walker2d & -0.1$\pm$0.0&22.2$\pm$4.2& 20.3$\pm$6.9 & 21.1$\pm$8.2 & 17.5$\pm$3.6 & 19.9$\pm$5.9 & 51.0$\pm$14.5 & 55.3$\pm$1.9 & 56.5$\pm$11.0& 56.1$\pm$15.4 \\
\cline{2-12}\rule{0pt}{2.5ex}
& halfcheetah(10\%)& 0.9$\pm$0.2 &20.3$\pm$4.0& 4.2$\pm$1.0 & 18.3$\pm$2.0 & 2.0$\pm$0.5 & 14.6$\pm$2.4 & 10.0$\pm$4.5 & 28.3$\pm$2.6 & 16.8$\pm$5.1 & 25.2$\pm$2.1 \\
& hopper(10\%) &2.0$\pm$1.2 &10.9$\pm$12.1& 9.9$\pm$0.3 &19.1$\pm$10.0 & 11.6$\pm$2.4 & 14.2$\pm$4.8 & 28.3$\pm$8.8 & 35.9$\pm$6.0 & 36.5$\pm$4.7 & 38.4$\pm$7.6 \\
& walker2d(10\%) & 0.9$\pm$1.6&1.2$\pm$2.0& 5.4$\pm$2.7 & 10.4$\pm$2.1 & 4.7$\pm$1.2 & 19.7$\pm$8.2 & 18.9$\pm$4.0 & 32.0$\pm$16.6 & 24.5$\pm$5.9 & 43.4$\pm$9.3 \\
\cline{2-12}\rule{0pt}{2.5ex}
& kitchen-complete &4.4$\pm$3.2&13.1$\pm$5.4& 23.8$\pm$6.0 & 48.8$\pm$3.8 & 26.$\pm$5.0 & 48.8$\pm$4.5 & 51.3$\pm$9.4 & 55.0$\pm$5.3 & 60.6$\pm$5.7 & 60.0$\pm$4.7 \\
& kitchen-partial &1.9$\pm$2.1&0.0$\pm$0.0& 1.1$\pm$1.4 & 27.5$\pm$6.8 & 0.5$\pm$0.9 & 15.6$\pm$8.7 & 34.4$\pm$7.4 & 37.5$\pm$7.3 & 45.6$\pm$23.7& 43.8$\pm$9.4 \\
& kitchen-mixed &0.0$\pm$0.0&1.2$\pm$2.2& 0.8$\pm$0.8 & 37.5$\pm$5.0 & 8.4$\pm$1.8 & 43.1$\pm$6.2 & 21.9$\pm$13.3 & 23.8$\pm$22.4 & 39.4$\pm$17.1& 35.6$\pm$9.7 \\
\cline{2-12}\rule{0pt}{2.5ex}
& door(1\%) &-0.3$\pm$0.0& -0.3$\pm$0.0& 47.1$\pm$12.7& 56.3$\pm$13.9 & 64.3$\pm$11.1& 71.1$\pm$0.8 & 33.9$\pm$24.3 & 45.1$\pm$26.3 & 92.1$\pm$17.5& 105.6$\pm$2.0 \\
& hammer(1\%) &0.2$\pm$0.0&0.2$\pm$0.0& 65.7$\pm$12.7& 84.5$\pm$12.2 & 88.7$\pm$18.0& 95.0$\pm$28.3& 26.3$\pm$16.9 & 43.3$\pm$38.4 &113.1$\pm$23.8 & 109.2$\pm$17.3 \\
& relocate(1\%) &-0.3$\pm$0.0&-0.3$\pm$0.0& 7.2$\pm$2.0 & 10.6$\pm$1.7 & 12.5$\pm$1.5 & 20.3$\pm$8.2 & 35.9$\pm$18.7 & 42.0$\pm$21.1 & 56.7$\pm$13.2 & 61.8$\pm$3.4 \\
\cline{2-12}\rule{0pt}{2.5ex}
& \textbf{Average Score} & 0.9 & \textbf{10.2} & 19.0 & \textbf{34.0} & 22.9 & \textbf{37.0} & 32.7 & \textbf{41.0} & 51.8 & \textbf{55.7} \\
& \textbf{Improvement $\uparrow$} & & \textbf{1033.33\%} & & \textbf{78.95\%} & & \textbf{61.57\%} & & \textbf{25.38\%} & & \textbf{7.52\%} \\
\bottomrule
\end{tabular}
}
\end{table*}
% \begin{table*}[ht]
% \centering
% \caption{Performance under adversarial data corruption.}
% \label{tab:adversarial_experiment_results}
% \resizebox{\textwidth}{!}{
% \begin{tabular}{l|l|c>{\columncolor{gray!20}}cc>{\columncolor{gray!20}}cc>{\columncolor{gray!20}}cc>{\columncolor{gray!20}}cc>{\columncolor{gray!20}}c}
% %\hline
% \toprule
% \multirow{2}{*}{Attack} & \multirow{2}{*}{Task} & \multicolumn{2}{c}{CQL} & \multicolumn{2}{c}{IQL} & \multicolumn{2}{c}{RIQL} & \multicolumn{2}{c}{DT} & \multicolumn{2}{c}{RDT} \\
% & & Naive & AGD & Naive & AGD & Naive & AGD & Naive & AGD & Naive & AGD\\
% \hline
% \midrule
% \multirow{13}{*}{\rotatebox{90}{{\Large State}}}
% & halfcheetah \\
% & hopper \\
% & walker2d \\
% \cline{2-12}\rule{0pt}{2.5ex}
% & halfcheetah(10\%) & 10.0$\pm$1.6 & 13.0$\pm$1.4 & 3.9$\pm$2.3 & 6.4$\pm$2.9 & 3.6$\pm$0.5 & 13.6$\pm$1.1 & 7.4$\pm$0.6 & 26.4$\pm$2.1 & 7.5$\pm$0.4 & 22.3$\pm$1.8 \\
% & hooper(10\%) & 1.8$\pm$0.8 & 6.5$\pm$3.9 & 19.1$\pm$5.2 & 13.9$\pm$2.5(?) & 19.1$\pm$6.9 & 39.3$\pm$6.7 & 38.6$\pm$4.7 & 49.4$\pm$6.6 & 39.3$\pm$5.1 & 41.8$\pm$8.8 \\
% & walker2d(10\%) & 0.7$\pm$1.5 & 16.4$\pm$2.8 & 7.4$\pm$3.9 & 11.4$\pm$1.8 & 10.4$\pm$1.9 & 25.6$\pm$6.6 & 22.3$\pm$2.4 & 48.4$\pm$13.4 & 21.1$\pm$2.6 & 41.5$\pm$4.8 \\
% \cline{2-12}\rule{0pt}{2.5ex}
% & kitchen-complete & 4.6$\pm$4.8 & 5.6$\pm$1.1 & 45.6$\pm$2.1 & 48.1$\pm$3.2 & 51.9$\pm$3.3 & 54.4$\pm$6.2 & 48.4$\pm$6.7 & 73.8$\pm$3.8 & 58.4$\pm$3.7 & 60.0$\pm$3.1 \\
% & kitchen-partial & 0.0$\pm$0.0 & 7.5$\pm$4.7 & 26.9$\pm$7.6 & 23.1$\pm$7.4(?) & 35.4$\pm$5.8 & 43.1$\pm$12.7 & 32.6$\pm$6.1 & 39.4$\pm$21.9 & 36.5$\pm$8.8 & 42.5$\pm$10.5 \\
% & kitchen-mixed & 2.0$\pm$3.5 & 8.1$\pm$4.5 & 41.2$\pm$4.5 & 45.6$\pm$7.2 & 33.9$\pm$11.2 & 50.0$\pm$3.1 & 28.2$\pm$9.9 & 38.8$\pm$9.6 & 30.0$\pm$5.5 & 45.6$\pm$10.1 \\
% \cline{2-12}\rule{0pt}{2.5ex}
% & door(1\%) & -0.3$\pm$0.0 & 12.0$\pm$10.6 & 49.3$\pm$11.4 & 58.1$\pm$10.2 & 47.3$\pm$24.8 & 64.8$\pm$7.9 & 99.0$\pm$0.9 & 103.3$\pm$2.0 & 104.7$\pm$0.5 & 105.5$\pm$1.3 \\
% & hammer(1\%) & 0.2$\pm$0.0 & 0.1$\pm$0.0 & 70.4$\pm$15.4 & 71.5$\pm$17.4 & 69.1$\pm$23.4 & 95.2$\pm$20.9 & 96.0$\pm$2.5 & 120.3$\pm$4.3 & 116.6$\pm$7.4 & 93.4$\pm$13.1(?) \\
% & relocate(1\%) & -0.3$\pm$0.0 & 0.0$\pm$0.0 & 7.0$\pm$1.4 & 14.6$\pm$2.8 & 17.1$\pm$9.5 & 31.7$\pm$7.7 & 76.2$\pm$5.0 & 80.8$\pm$5.9 & 69.0$\pm$4.4 & 70.9$\pm$8.8 \\
% \cline{2-12}\rule{0pt}{2.5ex}
% % & \textbf{Average Score} & 2.1$\pm$3.2 & 7.7$\pm$5.2 & 30.1$\pm$21.7 & 32.5$\pm$22.3 & 32.0$\pm$20.2 & 46.4$\pm$22.6 & 49.9$\pm$31.1 & 64.5$\pm$30.2 & 53.7$\pm$35.1 & 58.2$\pm$25.6 \\
% & \textbf{Average Score} & 2.1 & \textbf{7.7} & 30.1 & \textbf{32.5} & 32.0 & \textbf{46.4} & 49.9 & \textbf{64.5} & 53.7 & \textbf{58.2} \\
% & \textbf{Improvement $\uparrow$} & & \textbf{266.67\%} & & \textbf{7.97\%} & & \textbf{45.00\%} & & \textbf{29.26\%} & & \textbf{8.38\%} \\
% \bottomrule
% \end{tabular}
% }
% \end{table*}
\begin{table*}[ht]
\centering
\caption{Performance under adversarial data corruption. Results are averaged over four random seeds.}
\label{tab:adversarial_experiment_results}
\resizebox{\textwidth}{!}{
\begin{tabular}{l|l|c>{\columncolor{gray!20}}cc>{\columncolor{gray!20}}cc>{\columncolor{gray!20}}cc>{\columncolor{gray!20}}cc>{\columncolor{gray!20}}c}
%\hline
\toprule
\multirow{2}{*}{Attack} & \multirow{2}{*}{Task} & \multicolumn{2}{c}{CQL} & \multicolumn{2}{c}{IQL} & \multicolumn{2}{c}{RIQL} & \multicolumn{2}{c}{DT} & \multicolumn{2}{c}{RDT} \\
& & Naive & ADG & Naive & ADG & Naive & ADG & Naive & ADG & Naive & ADG\\
\hline
\midrule
\multirow{13}{*}{\rotatebox{90}{{\Large State}}}
& halfcheetah &31.2$\pm$1.4 &35.6$\pm$1.3& 23.1$\pm$6.7 & 29.4$\pm$2.2 & 26.5$\pm$5.3 & 27.4$\pm$1.9 & 34.8$\pm$0.6 & 35.6$\pm$2.1 & 34.5$\pm$1.9 & 35.2$\pm$1.3 \\
& hopper &55.3$\pm$11.3& 57.2$\pm$3.2 & 51.0$\pm$6.9 & 51.2$\pm$26.7 & 41.2$\pm$22.7 & 61.1$\pm$18.1 & 27.2$\pm$10.1& 55.8$\pm$21.2 & 55.2$\pm$28.3 & 70.5$\pm$7.8 \\
& walker2d &44.2$\pm$6.6 & 43.3$\pm$1.4 &36.2$\pm$14.0 & 42.4$\pm$7.4 & 16.7$\pm$3.7 & 32.4$\pm$15.4 & 44.1$\pm$8.1 & 62.2$\pm$3.2 & 42.4$\pm$6.4 & 62.2$\pm$5.7 \\
\cline{2-12}\rule{0pt}{2.5ex}
& halfcheetah(10\%) &10.0$\pm$2.0 &18.3$\pm$1.8 & 5.6$\pm$1.7 & 13.0$\pm$4.4 & 3.6$\pm$0.5 & 8.4$\pm$1.0 & 7.4$\pm$0.6 & 26.4$\pm$2.1 & 7.5$\pm$0.4 & 22.3$\pm$1.8 \\
& hopper(10\%) &2.0$\pm$1.1 &2.4$\pm$2.8& 20.0$\pm$3.8 & 19.9$\pm$2.3 & 19.1$\pm$6.9 & 22.6$\pm$12.3 & 38.6$\pm$4.7 & 40.3$\pm$5.9 & 39.3$\pm$5.1 & 42.2$\pm$6.9 \\
& walker2d(10\%) &-0.2$\pm$0.1 &0.2$\pm$0.5& 9.8$\pm$3.3 & 14.0$\pm$4.3 & 10.8$\pm$4.4 & 10.4$\pm$4.2 & 22.3$\pm$2.4 & 48.4$\pm$13.4 & 21.1$\pm$2.6 & 41.5$\pm$4.8 \\
\cline{2-12}\rule{0pt}{2.5ex}
& kitchen-complete &3.8$\pm$2.8 &7.5$\pm$7.3& 45.6$\pm$2.1 & 53.8$\pm$5.7 & 51.9$\pm$3.3 & 51.2$\pm$2.8 & 48.4$\pm$6.7 & 73.8$\pm$3.8 & 58.4$\pm$3.7 & 60.0$\pm$3.1 \\
& kitchen-partial &0.0$\pm$0.0 &2.5$\pm$4.3& 26.9$\pm$7.6 & 33.1$\pm$13.5& 35.4$\pm$5.8 & 44.4$\pm$8.7 & 32.6$\pm$6.1 & 39.4$\pm$21.9 & 36.5$\pm$8.8 & 42.5$\pm$10.5 \\
& kitchen-mixed &0.0$\pm$0.0 &1.9$\pm$3.2& 41.2$\pm$4.5 & 42.5$\pm$5.6 & 33.9$\pm$11.2 & 38.8$\pm$4.5 & 28.2$\pm$9.9 & 38.8$\pm$9.6 & 30.0$\pm$5.5 & 45.6$\pm$10.1 \\
\cline{2-12}\rule{0pt}{2.5ex}
& door(1\%) &-0.3$\pm$0.0 &-0.3$\pm$0.0& 49.3$\pm$11.4& 58.1$\pm$10.2 & 47.3$\pm$24.8 & 61.6$\pm$11.2 & 99.0$\pm$0.9 & 103.3$\pm$2.0 & 104.7$\pm$0.5 & 105.5$\pm$1.3 \\
& hammer(1\%) & 0.2$\pm$0.0 &0.2$\pm$0.0& 70.4$\pm$15.4& 78.4$\pm$8.9 & 69.1$\pm$23.4 & 94.1$\pm$11.7 & 96.0$\pm$2.5 & 120.3$\pm$4.3 & 116.6$\pm$7.4 & 93.4$\pm$13.1 \\
& relocate(1\%) &-0.2$\pm$0.0 &-0.2$\pm$0.1& 7.0$\pm$1.4 & 18.7$\pm$8.7 & 17.1$\pm$9.5 & 20.5$\pm$9.0 & 76.2$\pm$5.0 & 80.8$\pm$5.9 & 69.0$\pm$4.4 & 70.9$\pm$8.8 \\
\cline{2-12}\rule{0pt}{2.5ex}
& \textbf{Average Score} & 12.2 & \textbf{14.0} & 32.2 & \textbf{37.9} & 31.1 & \textbf{39.4} & 46.2 & \textbf{60.4} & 51.3 & \textbf{57.6} \\
& \textbf{Improvement $\uparrow$} & & \textbf{14.75\%} & & \textbf{17.70\%} & & \textbf{26.69\%} & & \textbf{30.74\%} & &\textbf{12.28\%} \\
\midrule
\hline\rule{0pt}{2.5ex}
\multirow{13}{*}{\rotatebox{90}{{\Large Full-Eelement}}}
& halfcheetah &0.2$\pm$0.7 &30.9$\pm$1.6 & 29.3$\pm$3.0 & 38.4$\pm$1.5 & 17.0$\pm$5.9 & 34.5$\pm$4.1 & 37.0$\pm$2.0 & 39.5$\pm$0.5 & 36.3$\pm$1.2 & 39.1$\pm$0.3 \\
& hopper &1.1$\pm$0.4 &12.3$\pm$1.3 & 51.4$\pm$24.1 & 68.8$\pm$7.7 & 27.8$\pm$12.5 & 40.0$\pm$0.8 & 56.4$\pm$11.3 & 65.9$\pm$18.5 & 62.3$\pm$4.6 & 64.7$\pm$10.5 \\
& walker2d &-0.3$\pm$0.1&37.6$\pm$13.0& 42.2$\pm$14.3 & 52.7$\pm$1.8 & 49.3$\pm$6.5 & 57.4$\pm$14.7 & 52.4$\pm$6.7 & 60.1$\pm$7.4 & 58.8$\pm$6.0 & 65.2$\pm$6.0 \\
\cline{2-12}\rule{0pt}{2.5ex}
& halfcheetah(10\%) &-1.1$\pm$0.6&13.1$\pm$1.9 & 4.6$\pm$1.2 & 17.1$\pm$3.4 & 4.8$\pm$2.1 & 17.3$\pm$3.7 & 12.7$\pm$2.4 & 24.4$\pm$1.7 & 14.2$\pm$1.5 & 26.7$\pm$1.5 \\
& hopper(10\%) &1.0$\pm$0.4 &3.8$\pm$2.4 & 21.3$\pm$5.4 & 20.5$\pm$2.7 & 30.3$\pm$5.3 & 35.1$\pm$2.4 & 40.0$\pm$12.1 & 40.2$\pm$3.8 & 39.7$\pm$3.8 & 49.4$\pm$13.2 \\
& walker2d(10\%) &-0.2$\pm$0.0&0.7$\pm$1.1 & 10.8$\pm$1.8 & 20.7$\pm$4.9 & 14.7$\pm$2.7 & 24.6$\pm$9.2 & 28.5$\pm$7.4 & 37.4$\pm$11.3 & 14.0$\pm$7.3 & 38.0$\pm$10.7 \\
\cline{2-12}\rule{0pt}{2.5ex}
& kitchen-complete &3.8$\pm$2.8 &7.5$\pm$4.7 & 46.2$\pm$6.2 & 48.1$\pm$6.5 & 50.0$\pm$9.8 & 56.2$\pm$9.4 & 55.0$\pm$6.1 & 65.6$\pm$8.0 & 59.4$\pm$6.2 & 58.1$\pm$13.8 \\
& kitchen-partial &5.0$\pm$4.0 &4.4$\pm$7.6 & 29.4$\pm$2.7 & 33.1$\pm$6.5 & 19.4$\pm$2.1 & 38.1$\pm$12.8 & 26.9$\pm$17.0 & 42.5$\pm$10.2 & 22.5$\pm$21.4 & 30.6$\pm$12.0 \\
& kitchen-mixed &0.0$\pm$0.0 &3.8$\pm$6.5 & 34.4$\pm$4.8 & 37.5$\pm$4.7 & 20.6$\pm$4.8 & 25.6$\pm$3.7 & 38.1$\pm$2.1 & 31.9$\pm$9.1 & 38.1$\pm$15.1 & 43.1$\pm$8.2 \\
\cline{2-12}\rule{0pt}{2.5ex}
& door(1\%) &-0.3$\pm$0.0&-0.4$\pm$0.0 & 66.9$\pm$15.0 & 68.1$\pm$11.1 & 37.7$\pm$3.2 & 55.3$\pm$5.7 & 96.5$\pm$11.3 & 98.9$\pm$3.5 & 87.9$\pm$20.5 & 94.1$\pm$7.0 \\
& hammer(1\%) & 0.1$\pm$0.1&0.2$\pm$0.0 & 61.5$\pm$10.1 & 92.6$\pm$13.5 & 50.9$\pm$20.5 & 80.7$\pm$13.2 & 75.1$\pm$20.6 & 75.2$\pm$25.5 & 62.5$\pm$26.5 & 75.4$\pm$35.3 \\
& relocate(1\%) &-0.2$\pm$0.0&-0.3$\pm$0.0 & 5.0$\pm$3.4 & 6.0$\pm$3.1 & 7.9$\pm$3.3 & 7.4$\pm$4.1 & 54.5$\pm$14.3 & 67.7$\pm$5.3 & 4.7$\pm$4.9 & 6.8$\pm$6.3 \\
\cline{2-12}\rule{0pt}{2.5ex}
& \textbf{Average Score} & 0.8 & \textbf{9.5} & 33.6 & \textbf{42.0} & 27.5 & \textbf{39.4} & 47.8 & \textbf{54.1} & 41.7 & \textbf{49.3} \\
& \textbf{Improvement $\uparrow$} & & \textbf{1087.50\%} & & \textbf{25.00\%} & & \textbf{43.27\%} & & \textbf{13.18\%} & &\textbf{18.23\%} \\
\bottomrule
\end{tabular}}
\end{table*}
\section{Conclusion}
% We propose Ambient Diffusion-Guided Dataset Recovery (ADG), the first diffusion-based denoising framework for offline RL under data corruption during the training process. The diffusion models in ADG are designed to be unconditional, eliminating the need for additional information. We introduce Ambient Denoising Diffusion Probabilistic Models (DDPM), enabling the diffusion model to distinguish between corrupted and clean samples, filtering the training data to facilitate the training of a naive diffusion model as a denoiser. Comprehensive empirical studies on D4RL benchmarks demonstrate that ADG consistently improves the performance of existing offline RL algorithms across various types, scales, and ratios of data corruption, and in most cases, allows baseline algorithms to outperform their robust variants.
We propose Ambient Diffusion-Guided Dataset Recovery (ADG), the first diffusion-based denoising framework for offline RL under data corruption during the training process. We introduce Ambient Denoising Diffusion Probabilistic Models (DDPM), which enable the diffusion model to distinguish between corrupted and clean samples. This mechanism effectively filters the training data, allowing training a naive diffusion model to serve as a denoiser. Comprehensive empirical studies on D4RL benchmarks demonstrate that ADG consistently improves the performance of existing offline RL algorithms across various types, scales, and ratios of data corruption, and in most cases, allows baseline algorithms to outperform their robust variants. We hope that this work establishes a new paradigm for more robust learning from noisy or corrupted data, ultimately benefiting the application of offine RL in real-world scenarios.
\section{Theoretical Interpretations}
\input{section/7_appendix_theory}
\section{Algorithm Pseudocode}
\label{app:pseudocode}
We provide the pseudocode of our proposed Ambient Diffusion-Guided Dataset Recovery (ADG) in Algorithm~\ref{alg:alg} for a comprehensive overview.
\input{section/7_pseduo}
\section{Implementation Details}\label{implementation_details}
\subsection{Data Corruption Details during Training Phase}\label{noised_setting}
We study two types of corruption: random noise and adversarial noise. For each type, we investigate two categories of elements to attack: (1) State corruption, where only the states in a portion of the samples are corrupted. According to Section~\ref{sec:preliminaries}, an original trajectory is defined as $\tau = (s_0, a_0, r_0, \ldots, s_{T-1}, a_{T-1}, r_{T-1})$. In MDP-based methods like IQL and CQL, corrupting a state $s_t$ affects two transitions: $(s_{t-1}, a_{t-1}, r_{t-1}, s_t)$ and $(s_t, a_t, r_t, s_{t+1})$. (2) Full-element corruption, where the states, actions, and rewards $(s_t, a_t, r_t)$ in a portion of the samples are corrupted, introducing a stronger challenge for robust offline RL algorithms.
We consider the tasks including MuJoCo, Kitchen and Adroit~\citep{fu2020d4rl}. We select the ``medium-replay-v2'' datasets in the MuJoCo tasks for our main experiments, using both the full datasets and down-sampled versions (reduced to 10\% of the original size). For Adroit tasks, we choose ``expert-v0'' datasets and down-sample them to 1\% of their original dataset. We use the full datasets for the tasks in the Kitchen, as their original dataset size is already limited. To control the overall level of corruption within the datasets, we introduce two parameters $\eta$ and $\alpha$. The parameter $\eta$ represents the proportion of corrupted data within a dataset, while $\alpha$ indicates the scale of corruption across each individual dimension. These settings are consistent with prior works~\citep{ye2024corruption, yang2023towards, xu2024robust}. We outline two types of random data corruption as follows:
\begin{itemize}
\item \textbf{Random State Attack:} We randomly sample $\eta \cdot N \cdot T$ states from all trajectories, where $N$ refers to the number of trajectories and $T$ represents the number of steps in a trajectory. The selected states are then modified as $\hat{s} = s + \lambda \cdot \text{std}(s)$, where $\lambda \sim \text{Uniform}[-\alpha, \alpha]^{d_s}$. Here, $d_s$ represents the dimension of states, and $\text{std}(s)$ is the $d_s$-dimensional standard deviation of all states in the offline dataset. The noise is scaled based on the standard deviation of each dimension and is independently added to each respective dimension.
\item \textbf{Random Full-element Attack:} We randomly sample $\eta \cdot N \cdot T$ state-action-reward triplets $(s_t, a_t, r_t)$ from all trajectories, and modify the action $\hat{a} = a + \lambda \cdot \text{std}(a)$, $\hat{r} \sim \text{Uniform}[-30 \cdot \alpha, 30 \cdot \alpha]^{d_s}$, where $\lambda \sim \text{Uniform}[-\alpha, \alpha]^{d_a}$, $d_a$ represents the dimension of actions and $\text{std}(a)$ is the $d_a$-dimensional standard deviation of all actions in the offline dataset. The corruption to rewards is multiplied by 30, following the setting in RDT~\citep{xu2024robust}, since offline RL algorithms tend to be resilient to small-scale random rewards corruption.
\end{itemize}
The two types of adversarial data corruption are detailed as follows:
\begin{itemize}
\item \textbf{Adversarial State Attack:} We first pretrain IQL agents with a Q-function $Q_p$ and policy function $\pi_p$ on clean datasets. Then, we randomly sample $\eta \cdot N \cdot T$ states and modify them as follows. Specifically, we perform the attack by solving for $\hat{s} = \min_{\hat{s} \in \mathbb{B}_d(s, \alpha)} Q_p(\hat{s}, a)$. Here, $\mathbb{B}_d(s, \epsilon)=\{\hat{s}| | \hat{s}-s \mid \leq \epsilon \cdot \operatorname{std}(s)\}$ regularizes the maximum difference for each state dimension. The optimization is implemented through Projected Gradient Descent, similar to prior works~\citep{madry2017towards, zhang2020robust, yang2023towards, xu2024robust}. In this approach, we first initialize a learnable vector $v \in [-\alpha, \alpha]^{d_s}$, and then conduct a 100-step gradient descent with a step size of 0.01 for $\hat{s} = s + v \cdot \text{std}(s)$. After each update, we clip each dimension of $z$ within the range $[-\alpha, \alpha]$.
\item \textbf{Adversarial Full-element Attack:} We use the pretrained IQL agent with a Q-function $Q_p$ and a policy function $\pi_p$. Then, we randomly sample $\eta \cdot N \cdot T$ state-action-reward triplets $(s_t, a_t, r_t)$, and modify $u = (s, a)$ to $\hat{u} = \min_{\hat{u} \in \mathbb{B}_d(u, \alpha)} Q_p(u)$. Here, $\mathbb{B}_d(u, \alpha)=\{\hat{u} \| \hat{u}-u \mid \leq \alpha \cdot \operatorname{std}(u)\}$ regularizes the maximum difference for each dimension of $u$. The optimization is implemented through Projected Gradient Descent, as discussed above. The rewards in these triplets are also modified to: $\hat{r}=-\alpha \cdot r$.
\end{itemize}
\subsection{ADG Network Structure}\label{ADG_network}
Both detector and denoiser in ADG is classical diffusion models~\citep{song2020score} with Unet Structure as shown in~\ref{fig:unet}. The diffusion models are unconditional, which do not need to introduce extra knowledge except the noised input, make ADG can be simply applied to new scenarios. For the noised prediction generation, we utilize a 3 Layer MLP with Mish activation. Further details including the hyper-parameters refer to Section~\ref{hyper-parameter}.
\begin{figure*}[htb]
\centering
\includegraphics[width=1.00\textwidth]{pics/Unet.pdf}
\caption{
Neural network structure of ADG.
}
\label{fig:unet}
\end{figure*}
\subsection{Hyperparameters}\label{hyper-parameter}
We present the hyperparameters and other details of ADG in Table~\ref{tab:general_hyper} and~\ref{tab:specific_hyper}.
\begin{table*}[ht]
\centering
\caption{Generic hyperparameters of ADG.}
\label{tab:general_hyper}
\begin{tabular}{l|cccccc}
\toprule
Hyper-parameter & Value \\
\midrule
Batch size & 256 \\
Total diffusion step~($K$) & 100 \\
Ambient Nature Timestep~($K_a$) & 30 \\
Threshold $\zeta$ & 0.20 \\
Temporal slice size $H$ & 5 \\
Learning Rate~(lr) & 5 \\
Noise prediction network & FC(256,256,256) with Mish activations \\
Dropout for predictor network & 0.1 \\
Variance schedule & Variance Preserving (VP) \\
Learning Rate & 1e-4 \\
\bottomrule
\end{tabular}
%}
\end{table*}
\begin{table*}[ht]
\centering
\caption{Hyperparameters of ADG for different benchmark environments and datasets.}
\label{tab:specific_hyper}
\begin{tabular}{l|cccc}
\toprule
Tasks & Embedded Dimension & Denoising Diffusion Steps & Total Training Steps \\
\midrule
MuJoCo & 128 & 100 & 20k \\
MuJoCo (10\%) & 128 & 100 & 20k \\
Kitchen & 512 & 10 & 40k \\
Adroit (1\%) & 512 & 10 & 40k \\
\bottomrule
\end{tabular}
\end{table*}
\subsection{Dataset Details}
Since we conduct the experiments on both the full and down-sampled datasets, we provide detailed information about the number of transitions and trajectories as shown in~\cref{tab:dataset_info}.
\begin{table}[htbp]
\centering
\begin{tabular}{l|ccc}
\toprule
Dataset & halfcheetah & hopper & walker2d\\ \hline
\# Transitions & 202000 & 402000 &302000 \\
\# Trajectories & 202 & 2041 & 1093 \\
\midrule
\midrule
Dataset & halfcheetah(10\%) & hopper(10\%) & walker2d(10\%) \\ \hline
\# Transitions & 20000 & 32921 & 27937 \\
\# Trajectories & 20 & 204 & 109 \\
\midrule
\midrule
Dataset & kitchen-complete & kitchen-partial & kitchen-mixed\\ \hline
\# Transitions & 3680 & 136950 & 136950 \\
\# Trajectories & 19 & 613 & 613 \\
\midrule
\midrule
Dataset & door & hammer & relocate \\ \hline
\# Transitions & 10000 & 10000 & 10000 \\
\# Trajectories & 50 & 50 & 50 \\
\bottomrule
\end{tabular}
\caption{Detailed information about the number of transitions and trajectories.}
\label{tab:dataset_info}
\end{table}
\subsection{Computation Overhead}\label{app:computation}
In Table~\ref{tab:training_time}, we compare the computational cost of ADG with baseline algorithms on a single GPU (P40). Each algorithm is trained on the ``walker2d-medium-replay-v2'' dataset, and we record the total training time. CQL requires significantly more time due to its reliance on a larger number of Q ensembles and other computationally intensive processes. IQL and DT exhibit similar computational costs to their robust counterparts, RIQL and RDT. Notably, ADG achieves substantially lower computational costs than all baselines, demonstrating its ability to deliver significant performance gains with minimal additional overhead. The total time and per-step time for training the detector and denoiser, as well as for dataset recovery, are also reported.
\begin{table*}[ht]
\centering
\caption{Training time of ADG and offline RL baseline algorithms.}
\label{tab:training_time}
\begin{tabular}{l|cccccc}
\toprule
Tasks & CQL & IQL & RIQL & DT & RDT & ADG \\
\midrule
Epoch Num & 1000 & 1000 & 1000 & 100 & 100 & 10 \\
Total Time (h) & 9.39 & 3.90 & 4.14 & 1.22 & 1.27 & 0.78 \\
Detector Training (h) & \multicolumn{6}{c}{0.37 (0.07s/step)} \\
Denoiser Training (h) & \multicolumn{6}{c}{0.22 (0.04s/step)}\\
Dataset Recovering (h) & \multicolumn{6}{c}{0.19} \\
\bottomrule
\end{tabular}
\end{table*}
% \clearpage
\section{Additional Experimental Results}\label{appendix_experiments}
We present additional experimental results in this section. The network architecture and hyperparameters remain the same as those specified in Table~\ref{tab:general_hyper}. All results are averaged over four seeds.
\begin{figure*}[htb]
\centering
\includegraphics[width=0.97\textwidth]{pics/ablation_box.pdf}
\caption{
Visualization of the distribution of $e(\check{\bm{x}})$ for samples in the datasets, including halfcheetah, hopper, walker2d, kitchen-complete, kitchen-mixed, kitchen-partial, door, hammer, and relocate, denoted as ``ha'', ``ho'', ``wa'', ``kit-c'', ``kit-m'', ``kit-p'', ``do'', ``hm'', and ``re'', respectively. We distinguish clean and corrupted samples using different colors.
}
\label{fig:visual_illustration_detection}
\end{figure*}
\subsection{Visualization of $e(\check{\bm{x}})$ Predicted by the Detector}\label{ablation_visual_detection}
To demonstrate the effectiveness of ADG's detection capability, we visualize the squared Frobenius norm $e(\check{\bm{x}})$ of each sample $\check{\bm{x}}$ across the MuJoCo, Kitchen, and Adroit datasets using trained detectors. Consistent with the experiments in Section~\ref{sec:experiments}, we use ``medium-replay'' datasets for MuJoCo and ``expert'' datasets for Adroit. The results are shown in Figure~\ref{fig:visual_illustration_detection} using box plots. From the results, we observe a clear difference between the clean samples and corrupted samples in the distribution of $e(\check{\bm{x}})$.
\subsection{Visualization of the Recovered Trajectories}\label{ablation_visual_denoising}
To demonstrate the effectiveness of our proposed approach, ADG, we visualize a partial trajectory of ``hopper'' in Figure~\ref{fig:visual_illustration}. From the results, we observe that ADG can effectively detect and recover the corrupted samples. The partially noised trajectory has its discontinuity significantly reduced after recovery. Additionally, ADG introduces nearly no further corruption to the clean samples during the recovery process.
\begin{figure*}[htb]
\centering
% 子图 1
\subfigure[]{
\includegraphics[width=0.47\textwidth]{pics/traj1_crop.pdf}
}
\hfill
% 子图 2
\subfigure[]{
\includegraphics[width=0.47\textwidth]{pics/traj2_crop.pdf}
}
\caption{
Visualization of the denoising effect of ADG on the ``hopper-medium-replay-v2'' dataset. In (a), we present the complete trajectories, while in (b), we provide a partial view of the trajectories along with their scatter plot representations to highlight the differences in trajectory coherence. In both subfigures, the first row depicts the clean trajectories from the offline dataset, the second row shows the trajectories perturbed with Random State Attack as described in Section~\ref{noised_setting}, and the third row illustrates the trajectories restored using ADG, which have largely corrected the corrupted samples.
}
\label{fig:visual_illustration}
\end{figure*}
\subsection{Performance Under Gaussian Noise}\label{ablation_gaussian}
As presented in Section~\ref{exp:random_ad}, we evaluated the performance of ADG under both random and adversarial noise. We further analyze the robustness of ADG against Gaussian noise in the state. The implementation simply changes the noise from Uniform $[-\alpha, \alpha]^{d_s}$ in~\ref{noised_setting} to Gaussian noise $\mathcal{N}(0, \alpha^2)$. We set $\eta = 0.3$ and $\alpha = 1.0$, as in the other experiments. The results are presented in Table~\ref{tab:gaussian}, where ADG consistently improves the performance of the baselines, including IQL and DT, further highlighting its efficiency.
%As presented in Section~\ref{exp:random_ad}, we evaluated the performance of ADG under both random and adversarial noise. We further analyze the robustness of ADG against Gaussian noise in the state. The implementation is simply change the noise Uniform $[-\alpha, \alpha]^{d_s}$ in~\ref{noised_setting} to Gaussain 帮我补全. We set $\eta = 0.3$ and $\alpha = 1.0$, as in the other experiments. The results are presented in Table~\ref{tab:gassain}, where ADG consistently improves the performance of the baselines, including IQL and DT, further highlighting its efficiency.
\begin{table*}[ht]
\centering
\caption{Results under Gaussian noise.}
\label{tab:gaussian}
\begin{tabular}{l|l|c>{\columncolor{gray!20}}cc>{\columncolor{gray!20}}c}
\toprule
\multirow{2}{*}{Attack} & \multirow{2}{*}{Task} & \multicolumn{2}{c}{IQL} & \multicolumn{2}{c}{DT} \\
& & Naive & ADG & Naive & ADG \\
\hline
\midrule
\multirow{8}{*}{\rotatebox{90}{{\Large State}}} & halfcheetah & 17.2$\pm$4.6 & 27.9$\pm$1.3 & 36.1$\pm$1.8 & 36.1$\pm$1.1 \\
& hopper & 33.3$\pm$9.1 & 54.9$\pm$2.6 & 41.8$\pm$12.6 & 60.8$\pm$13.7 \\
& walker2d & 18.0$\pm$5.6 & 40.3$\pm$16.0 & 50.9$\pm$5.2 & 58.2$\pm$10.1 \\
\cline{2-6}\rule{0pt}{2.5ex}
& halfcheetah(10\%) & 13.6$\pm$3.5 & 15.6$\pm$3.7 & 12.8$\pm$2.1 & 24.8$\pm$1.7 \\
& hopper(10\%) & 25.4$\pm$11.7 & 32.6$\pm$4.5 & 39.0$\pm$7.4 & 44.3$\pm$7.7 \\
& walker2d(10\%) & 15.3$\pm$5.3 & 6.7$\pm$3.1 & 27.1$\pm$8.2 & 33.7$\pm$4.0 \\
\cline{2-6}\rule{0pt}{2.5ex}
& \textbf{Average} & 20.5 & \textbf{29.7} & 34.6& \textbf{43.0} \\
& \textbf{Improvement $\uparrow$} & & \textbf{44.88\%} & & \textbf{24.28\%}\\
\midrule
\hline\rule{0pt}{2.5ex}
\multirow{8}{*}{\rotatebox{90}{{\Large Full-Element}}} & halfcheetah & 10.3$\pm$3.9 & 32.5$\pm$2.5 & 34.1$\pm$2.3 & 44.5$\pm$7.3 \\
& hopper & 28.0$\pm$9.9 & 31.9$\pm$2.6 & 38.1$\pm$9.2 & 49.7$\pm$9.4 \\
& walker2d & 8.9$\pm$2.7 & 38.7$\pm$13.4 & 59.3$\pm$7.2 & 60.2$\pm$3.1 \\
\cline{2-6}\rule{0pt}{2.5ex}
& halfcheetah(10\%) & 5.9$\pm$1.6 & 17.8$\pm$2.8 & 11.8$\pm$2.0 & 21.7$\pm$2.8 \\
& hopper(10\%) & 15.6$\pm$3.8 & 13.4$\pm$1.7 & 25.5$\pm$7.9 & 37.8$\pm$12.2 \\
& walker2d(10\%) & 11.7$\pm$5.7 & 15.8$\pm$3.4 & 24.2$\pm$6.3 & 27.6$\pm$15.8 \\
\cline{2-6}\rule{0pt}{2.5ex}
& \textbf{Average} & 13.4 & \textbf{25.0} & 32.2 & \textbf{40.2} \\
& \textbf{Improvement $\uparrow$} & & \textbf{86.57\%} & & \textbf{24.84\%}\\
\bottomrule
\end{tabular}
\end{table*}
\subsection{Performance Under MuJoCo Dataset with Different Quality Levels}\label{ablation_different_quality_levels}
We evaluate the robustness of ADG under varying dataset quality levels. Specifically, we select the ``medium-v2'' and ``expert-v2'' datasets from MuJoCo tasks and downsample each to 2\% of their original size, resulting in datasets containing 20k samples. For comparison, we include RIQL, DT, and RDT as baselines. We do not include the results of CQL and IQL as they perform poorly on these down-sampled datasets. The results, summarized in Table~\ref{tab:medium_and_expert}, demonstrate that ADG consistently enhances the performance of these baselines across different dataset quality levels.
\input{section/7_1_table}
\subsection{Filtered Datasets vs. Recovered Datasets in MDP-based Algorithms}\label{ablation_filted_vs_restored}
For MDP-based algorithms such as IQL and RIQL, which do not take sequences as input, filtering out corrupted samples from the dataset using a detector, without recovery, is also a viable method to mitigate the impact of corrupted data. Specifically, we construct the filtered dataset by excluding all $\check{\bm{x}}$ for which $e_\theta(\check{\bm{x}}) > \zeta = 0.20$ using ADG. To further investigate this approach, we also evaluate filtering with $\zeta = 0.10$ to assess whether a stricter threshold improves performance. Although ADG demonstrates strong detection capabilities, some corrupted data may remain in the filtered dataset. Therefore, we also evaluate the performance of the filtered dataset after removing the remaining corrupted data, which we refer to as the purified dataset. We then perform IQL and RIQL on these datasets. The results, shown in Figure~\ref{fig:ablation_filtered}, reveal the following insights: (1) The filtered dataset does not consistently improve the performance of the MDP-based algorithms. (2) The recovered dataset shows better robustness, even when compared to the purified dataset. (3) Using a lower value of $\zeta$ to filter the dataset more aggressively does not improve performance, possibly due to the reduced information content in the dataset. These findings highlight that dataset size is a crucial factor for RL performance, suggesting that filtering the dataset without recovery does not outperform recovery-based approaches. This emphasizes the importance of recovery in ADG.
\begin{figure*}[htb]
\centering
\includegraphics[width=0.97\textwidth]{pics/ablation_Filtered_Final.pdf}
\caption{
The performance of IQL (left) and RIQL (right) on the noised, filtered, purified, and restored datasets. Performance on the noised data is referred to as ``Naive''.
}
\label{fig:ablation_filtered}
\end{figure*}
\subsection{Evaluation of Performance Across Different Temporal Slice Sizes ($H$)}\label{ablation_different_H}
In this section, we evaluate ADG with varying temporal slice sizes by setting $H \in \{0, 1, 3, 5, 7\}$. As described in Section~\ref{method_recovery}, the temporal slice size is defined as $2 \cdot H + 1$. Notably, $H = 0$ means the diffusion model takes only a single step as input and cannot access sequential information. The results are presented in Figure~\ref{fig:ablation_H}. From the results, we observe that the performance of the detector and denoiser in ADG improves significantly as $H$ increases from $0$ to $3$, and remains robust for $H = 3, 5, 7$. The performance of IQL strongly correlates with that of the detector and denoiser, showing significant improvement for $H = 3, 5, 7$ compared to $H = 0, 1$. In contrast, DT is more resilient to variations in ADG's recovery performance across different values of $H$, likely due to its inherent ability to leverage temporal information and mitigate the effects of corruption. This observation aligns with the findings in RDT~\citep{xu2024robust}. These results highlight the critical role of incorporating temporal slices in enhancing the performance of the detector, denoiser, and overall RL performance on the denoised dataset.
\begin{figure*}[htb]
\centering
% 子图 1
\subfigure[]{
\includegraphics[width=0.45\textwidth]{pics/ablation_H_plot.pdf}
}
\hfill
% 子图 2
\subfigure[]{
\includegraphics[width=0.45\textwidth]{pics/ablation_H_D4RL_score.pdf}
}
\caption{
The performance of the ADG detector and denoiser under varying temporal slice sizes ($H$) is shown in (a), while (b) presents the corresponding performance of the baseline algorithms, IQL and DT, on the ``walker2d-medium-replay-v2'' dataset. FN represents the False Negative ratio (the proportion of corrupted samples incorrectly identified as clean), and MSE indicates the mean squared error between the restored dataset and the ground-truth dataset.
}
\label{fig:ablation_H}
\end{figure*}
\subsection{Evaluation under Various Dataset Scales}\label{ablation_dataset_size}
We further examine the robustness of ADG on MuJoCo tasks with dataset sizes of 20\% and 50\% of the ``medium-replay-v2'' datasets. We investigate corruption types including Random State Attack and Random Full-element Attack as described in Section~\ref{noised_setting}. The comparison results are shown in Figure~\ref{fig:ablation_data_size_ratio}. From the results, we observe that the performance of both the algorithms and their robust variants on both corrupted and restored datasets is positively correlated with dataset size. Moreover, ADG outperforms the baselines across most dataset sizes, validating the effectiveness and importance of dataset recovery.
\begin{figure*}[htb]
\centering
\subfigure[]{
\includegraphics[width=0.45\textwidth]{pics/ablation_dataset_ratio_IQL_random_state.pdf}
}
\hfill
\subfigure[]{
\includegraphics[width=0.45\textwidth]{pics/ablation_dataset_ratio_IQL_random_element.pdf}
}
\subfigure[]{
\includegraphics[width=0.45\textwidth]{pics/ablation_dataset_ratio_DT_random_state.pdf}
}
\hfill
\subfigure[]{
\includegraphics[width=0.45\textwidth]{pics/ablation_dataset_ratio_DT_random_element.pdf}
}
\subfigure[]{
\includegraphics[width=0.45\textwidth]{pics/ablation_dataset_ratio_RIQL_random_state.pdf}
}
\hfill
\subfigure[]{
\includegraphics[width=0.45\textwidth]{pics/ablation_dataset_ratio_RIQL_random_element.pdf}
}
\subfigure[]{
\includegraphics[width=0.45\textwidth]{pics/ablation_dataset_ratio_RDT_random_state.pdf}
}
\hfill
\subfigure[]{
\includegraphics[width=0.45\textwidth]{pics/ablation_dataset_ratio_RDT_random_element.pdf}
}
\caption{
Performance of ADG under random corruption across different dataset scales.
}
\label{fig:ablation_data_size_ratio}
\end{figure*}
\subsection{Varying Corruption Rates and Scales}\label{ablation_corruption_rates_and_scales}
We evaluate the robustness of ADG under various corruption rates \( \{0.0, 0.1, 0.3, 0.5\} \) and scales \( \{0.0, 1.0, 2.0\} \). We choose the ``walker2d-medium-replay-v2'' dataset downsampled to 10\% of its original size to make the results more sensitive to corruption rates and scales. As shown in Figure~\ref{fig:ablation_rate_and_scale}, increasing corruption rates and scales progressively degrade the performance of baseline algorithms like IQL and DT due to greater deviations between the corrupted and clean datasets. Nevertheless, ADG consistently enhances the overall performance of these baselines.
\begin{figure*}[htb]
\centering
% 子图 1
\subfigure[]{
\includegraphics[width=0.40\textwidth]{pics/Ablation_rate_walker2d.pdf} % 替换为实际图片路径
}
\hfill
% 子图 2
\subfigure[]{
\includegraphics[width=0.40\textwidth]{pics/Ablation_scale_walker2d.pdf} % 替换为实际图片路径
}
\caption{
Results under various corruption rates~(a) and scales~(b) on the ``walker2d-medium-replay-v2'' dataset, which is downsampled to 10\% of the original size.
}
\label{fig:ablation_rate_and_scale}
\end{figure*}
\subsection{Impact of $\zeta$ in Dataset Splitting}\label{zeta_in_data_spliting}
As outlined in Section~\ref{method_recovery}, we introduce a threshold hyperparameter $\zeta$ to differentiate between noisy and clean samples. To emphasize the importance of selective sampling during training, we evaluate ADG with a trained denoiser by varying $\zeta$ over the range $\{0.02, 0.05, 0.1, 0.2, 0.5, 1.0\}$. Note that when $\zeta = 1$, the dataset remains unchanged. The results are presented in Figure~\ref{fig:ablation_zeta_in_splitting}. A higher $\zeta$ incorporates more noisy samples into the restored dataset without applying restoration, increasing the risk of overfitting in the naive diffusion model. In contrast, a lower $\zeta$ results in more clean samples being misclassified as noisy, leading to unnecessary restoration and a loss of original dataset information. There is a trade-off in choosing the threshold $\zeta$, and $\zeta = 0.20$ yields the lowest MSE between the restored and clean datasets, as well as the best D4RL score across all baseline algorithms.
\begin{figure*}[htb]
\centering
% 子图 1
\subfigure[]{
\includegraphics[width=0.30\textwidth]{pics/ablation_c_plot.pdf} % 替换为实际图片路径
}
% \hfill
% 子图 2
\subfigure[]{
\includegraphics[width=0.49\textwidth]{pics/ablation_c_D4RL_performance.pdf} % 替换为实际图片路径
}
\caption{
Results of (a) the detection performance in the dataset recovery process and (b) D4RL score under various $\zeta$ on the ``walker2d-medium-replay-v2'' dataset. In (a), we label noisy samples as positive and clean samples as negative. The False Positive (FP) rate represents the ratio of clean samples incorrectly classified as noisy, while the False Negative (FN) rate corresponds to the ratio of noisy samples misclassified as clean. Both are expected to be low for a better dataset splitting result. Scaled MSE is the MSE scaled to \([0, 1]\), which directly reflects the deviation of the restored dataset from the clean dataset.
}
\label{fig:ablation_zeta_in_splitting}
\end{figure*}
\subsection{Additional Experiments Comparing the Use of Single vs. Double Diffusion Models}\label{appendix_single_diffusion}
In this section, we provide further details on the implementation of ADG using a single diffusion model. We reorganize Algorithm~\ref{alg:alg} to merge \textbf{Step 1} and \textbf{Step 2} into a single loop. In each step, the diffusion model is updated using both Eq.~\ref{corollay_eq4} and Eq.~\ref{ddpm_loss}. We continue to use $k_a$ to detect corrupted data in the partially corrupted dataset. The results for ``halfcheetah-medium-replay-v2'' and ``hopper-medium-replay-v2'' datasets are presented in Figure~\ref{fig:ablation_single_diffusion_mujoco} to complete the analysis. Employing a single diffusion model improves performance in 4 out of 6 tasks compared to the baselines. However, ADG with two independent models still achieves significantly higher performance. These findings align with the conclusions in Section~\ref{ablation_main}.
\begin{figure*}[htb]
\centering
\includegraphics[width=0.75\textwidth]{pics/ablation_single_diffusion_halfcheetah_hopper_appendix.pdf} % 替换为实际图片路径
\caption{
Comparison results among Naive, ADG with a single diffusion model, and ADG with separate diffusion models for IQL (upper) and DT (lower) under ``halfcheetah-medium-replay-v2'' and ``hopper-medium-replay-v2'' datasets. ADG (w/ SD) denotes ADG using a single diffusion model.
}
\label{fig:ablation_single_diffusion_mujoco}
\end{figure*}
\subsection{Missing data}
While our main focus is on additive corruption, ADG is theoretically applicable to a broader class of noise, including missing data such as dropped elements. We construct a missing-data variant of each dataset by randomly zeroing out 30\% of the actions. We evaluate baseline offline RL algorithms and ADG under this setting and report the results averaged over four seeds in Table~\ref{tab:missing_data}. Despite the substantial information loss, ADG consistently improves performance across all environments, recovering a significant portion of the original policy quality.
\begin{table}[ht]
\centering
\caption{Performance under missing data. Results are averaged over 4 random seeds. "Missing Data" refers to standard offline RL applied to the corrupted dataset. "Missing Data (w/ ADG)" denotes applying ADG for recovery.}
\label{tab:missing_data}
\begin{tabular}{lc>{\columncolor{gray!20}}cc}
\toprule
Environment & Missing Data & Missing Data (w/ ADG) & Clean \\
\midrule
halfcheetah-mr & 3.8$\pm$0.2 & 31.3$\pm$0.7 & 38.9$\pm$0.5 \\
hopper-mr & 14.0$\pm$1.0 & 45.5$\pm$16.4 & 81.8$\pm$6.9 \\
walker2d-mr & 7.2$\pm$0.3 & 32.7$\pm$4.1 & 59.9$\pm$2.7 \\
\bottomrule
\end{tabular}
\end{table}
These results indicate that ADG is not limited to additive noise settings and generalizes well to missing data, further demonstrating its robustness in practical offline RL scenarios where various forms of corruption may co-occur.
\subsection{Extending ADG to Recent Offline RL Methods}
To evaluate the broader applicability of ADG beyond standard baselines, we further apply it to two recent state-of-the-art offline RL algorithms: A2PR~\citep{liu2024adaptive} and NUNO~\citep{kopruluneural}. Specifically, we test their robustness under the Random State Attack corruption setting and assess the performance gains when combined with ADG. This allows us to examine whether ADG can consistently enhance policy quality across diverse algorithmic backbones.
As shown in Table~\ref{tab:sota_methods}, ADG significantly improves performance for both A2PR and NUNO in all environments, highlighting its general utility in mitigating data corruption across a range of policy learning strategies.
\begin{table}[ht]
\centering
\caption{Performance of A2PR and NUNO under Random State Attack, with and without ADG. Results are averaged over four random seeds.}
\label{tab:sota_methods}
\begin{tabular}{lcc>{\columncolor{gray!20}}c>{\columncolor{gray!20}}c}
\toprule
Environment & A2PR (noised) & NUNO (noised) & A2PR (w/ ADG) & NUNO (w/ ADG) \\
\midrule
halfcheetah-mr & 4.4$\pm$0.5 & 27.3$\pm$7.8 & 13.2$\pm$1.6 & \textbf{29.2}$\pm$7.5 \\
hopper-mr & 16.1$\pm$5.0 & 12.6$\pm$0.8 & 25.8$\pm$6.2 & \textbf{29.9}$\pm$0.6 \\
walker2d-mr & 4.6$\pm$2.3 & 11.0$\pm$5.5 & 9.3$\pm$0.8 & \textbf{14.6}$\pm$4.1 \\
\bottomrule
\end{tabular}
\end{table}
\section{Discussion and Limitation}\label{limitation}
While ADG demonstrates promising results in handling corrupted offline RL datasets, our work has several limitations that warrant discussion. The three-stage diffusion process (detection + denoising) introduces additional training time compared to standard offline RL methods. The performance of ADG depends on the choice of key hyperparameters such as the temporal slice size \( H \) and corruption threshold \( \zeta \). The trajectory-based recovery mechanism (Section 4.4) relies on temporal consistency. For environments with highly discontinuous dynamics or sparse rewards, the current slice-based approach may miss long-range dependencies.
\subsection{Proof for Corollary 4.1}\label{proof_coll_41}
% Firstly, we recap that Tweedie's Formula with:
% \begin{equation}
% \end{equation}
% $X = \gamma Y + \eta Z$
Firstly, we define $\bm{x}^{k_a}= \sqrt{\bar{\alpha}_{k_a}} \bm{x}^0 + \sqrt{1-\bar{\alpha}_{k_a}} \bm{\epsilon}_1$ and $\bm{x}^k= \sqrt{\bar{\alpha}_k} \bm{x}^0 + \sqrt{1-\bar{\alpha}_k} \bm{\epsilon}_2$, where $\bm{\epsilon}_1, \bm{\epsilon}_2 \sim \mathcal{N}(\bm{0}, \bm{I})$.
According to the forward process in DDPM, $\bm{x}^k$ can be expressed as a function of $\bm{x}^{k_a}$:
\begin{equation}\label{ddpm_forward}
\begin{aligned}
\bm{x}^k & = \sqrt{\bar{\alpha}_k} \bm{x}^0 + \sqrt{1-\bar{\alpha}_k} \bm{\epsilon}_2 \\
& = \sqrt{\bar{\alpha}_k} \frac{\bm{x}^{k_a} - \sqrt{1-\bar{\alpha}_{k_a}} \bm{\epsilon}_1}{\sqrt{\bar{\alpha}_{k_a}}} + \sqrt{1-\bar{\alpha}_k} \bm{\epsilon}_2 \\
& = \frac{\sqrt{\bar{\alpha}_k}}{\sqrt{\bar{\alpha}_{k_a}}} \bm{x}^{k_a} + \sqrt{\sqrt{1-\bar{\alpha}_{k_a}}^2 - {\frac{\sqrt{\bar{\alpha}_k} \sqrt{1-\bar{\alpha}_{k_a}}}{\sqrt{\bar{\alpha}_{k_a}}}}^2 }\bm{\epsilon} \\
& = \sqrt{\frac{\bar{\alpha}_k}{\bar{\alpha}_{k_a}}} \bm{x}^{k_a} + \sqrt{\frac{\bar{\alpha}_{k_a}-\bar{\alpha}_k}{\bar{\alpha}_{k_a}}} \bm{\epsilon}
\end{aligned}
\end{equation}
Following the proof sketch of Lemma A.4 in \citep{daras24ambient}, we apply Tweedie's formula to the pair of $\bm{x}^{k}$ and $\bm{x}^{0}$:
\begin{equation}\label{tweedies1}
\nabla \log p_k (\bm{x}^{k}) = \frac{ \sqrt{\bar{\alpha}_k} \mathbb{E}[\bm{x}^0\vert\bm{x}^k] - \bm{x}^k}{1-\bar{\alpha}_k}.
\end{equation}
Similarly, applying Tweedie's formula to the pair of $\bm{x}^{k}$ and $\bm{x}^{k_a}$, we derive:
\begin{equation}\label{tweedies2}
\nabla \log p_k (\bm{x}^{k})= \frac{ \sqrt{\frac{\bar{\alpha}_k}{\bar{\alpha}_{k_a}}} \mathbb{E}[\bm{x}^{k_a}\vert\bm{x}^k] - \bm{x}^k}{\frac{\bar{\alpha}_{k_a}-\bar{\alpha}_k}{\bar{\alpha}_{k_a}}}
\end{equation}
Equating Eq.~\eqref{tweedies1} and Eq.~\eqref{tweedies2}, we obtain:
\begin{equation}\label{exp_ka}
\mathbb{E}[\bm{x}^{k_a}\vert\bm{x}^k] =
\frac{\bar{\alpha}_{k_a} - \bar{\alpha}_k}{\sqrt{\bar{\alpha}_{k_a}} \cdot (1- \bar{\alpha}_k)} \mathbb{E}[\bm{x}^{0}\vert\bm{x}^k] +
\frac{\sqrt{\bar{\alpha}_{k_a}} \cdot (1 - \bar{\alpha}_{k_a}) }{\sqrt{\bar{\alpha}_{k}} \cdot (1- \bar{\alpha}_k)} \bm{x}^k
\end{equation}
Since we aim to predict the Gaussian noise rather than the original signal, the forward process of DDPM provides the relationship:
\begin{equation}\label{exp_noise_to_info}
\mathbb{E}[\bm{x}^0\vert\bm{x}^k] = \frac{\bm{x}^k - \sqrt{1 - \bar{\alpha}_k} \mathbb{E}[\bm{\epsilon}\vert\bm{x}^k]}{\sqrt{\bar{\alpha}_k}}
\end{equation}
Substituting Eq.~\eqref{exp_noise_to_info} into Eq.~\eqref{exp_ka}, we relate $\mathbb{E}[\bm{\epsilon}\vert\bm{x}^k]$ to $\mathbb{E}[\bm{x}^{k_a}\vert\bm{x}^k]$:
\begin{equation}\label{exp_noise_to_partial}
\mathbb{E}[\bm{x}^{k_a}\vert\bm{x}^k] = \frac{\bar{\alpha}_{k_a}}{\sqrt{\bar{\alpha}_{k} \cdot \bar{\alpha}_{k_a}}} \bm{x}^k -
\frac{(\bar{\alpha}_{k_a} - \bar{\alpha}_{k})\mathbb{E}[\bm{\epsilon}\vert\bm{x}^k]}{\sqrt{\bar{\alpha}_{k_a} \cdot \bar{\alpha}_{k} \cdot (1- \bar{\alpha}_{k})}}
\end{equation}
Building on Theorem A.3 in \citep{daras24ambient}, we minimize the objective:
\begin{equation}\label{original_objective}
\mathbb{E}_{\bm{x}^{k_a}}\mathbb{E}_{k\sim\mathcal{U}(k_a, K)}\mathbb{E}_{\bm{x}^k\vert \bm{x}^{k_0}}\left[ \Vert
\bm{g}_{\theta}(\bm{x}^k, k) - \bm{x}^{k_a}
\Vert^2\right],
\end{equation}
where the minimizer satisfies $\bm{g}_{\theta^*}(\bm{x}^k, k)=\mathbb{E}[\bm{x}^{k_a}\vert\bm{x}^k]$. Substituting Eq.~\eqref{exp_noise_to_partial} into Eq.~\eqref{original_objective}, we have the final form of the optimization objective:
\begin{equation}
\mathbb{E}_{\bm{x}^{k_a}} \mathbb{E}_{k\sim\mathcal{U}(k_a, K)}\mathbb{E}_{\bm{x}^k\vert \bm{x}^{k_0}}\bigg[ \Vert
\frac{\bar{\alpha}_{k_a}}{\sqrt{\bar{\alpha}_{k} \cdot \bar{\alpha}_{k_a}}} \bm{x}^k -
\frac{(\bar{\alpha}_{k_a} - \bar{\alpha}_{k}) \bm{\epsilon}_{\theta}(\bm{x}^k, k)}{\sqrt{\bar{\alpha}_{k_a} \cdot \bar{\alpha}_{k} \cdot (1- \bar{\alpha}_{k})}}
- \bm{x}^{k_a}
\Vert^2\bigg]
\end{equation}
with the minimizer $\bm{\epsilon}_{\theta^*}(\bm{x}^k, k) = \mathbb{E}[\bm{\epsilon}\vert\bm{x}^k], \forall k\geq k_a$. This completes the proof.
\subsection{Proof for Theorem 4.3}\label{proof_therorm_43}
Let $\check{\bm{x}}$ be the samples that may or may not contain scaled Gaussian noise $\iota\cdot\bm{\epsilon}$. It can be expressed in the form of:
\begin{equation}\label{def_partial}
\check{\bm{x}} = \bm{x}^0 + \iota\cdot \mathbb{I}_\mathrm{noise}\cdot\bm{\epsilon},
\end{equation}
where $\bm{x}^0 \in \mathbb{R}^{m\times n}$ is the noise-free matrix, $\iota$ is the noise scale, $\bm{\epsilon} \in \mathbb{R}^{m\times n}$ is the matrix of Gaussian noise with i.i.d. entries $\epsilon_{ij}\sim \mathcal{N}(0,1)$, and $\mathbb{I}_\mathrm{noise}$ is an unknown indicator variable that determines whether noise is present. Accordingly, we define the samples that contain consistent Gaussian noise $\iota\cdot\bm{\epsilon}$, which is expressed as:
\begin{equation}\label{def_all}
\tilde{\bm{x}} = \bm{x}^0 + \iota\cdot \bm{\epsilon}.
\end{equation}
\begin{lemma}\label{lemma_a1}
Given $\check{\bm{x}}$ and $\tilde{\bm{x}}$ as defined in Eqs.~\eqref{def_partial} and \eqref{def_all} respectively, let $q(\cdot^k\vert \cdot^0)$ denote the distribution of DDPM forward process at diffusion timestep $k$, the following inequality always holds:
\begin{equation}
\mathbb{D}_{KL}[q (\bm{x}^k\vert\bm{x}^0) \Vert q (\check{\bm{x}}^k\vert\check{\bm{x}}^0)] \leq
\mathbb{D}_{KL}[q (\bm{x}^k\vert\bm{x}^0) \Vert q (\tilde{\bm{x}}^k\vert\tilde{\bm{x}}^0)].
\end{equation}
\end{lemma}
\textit{Proof.}
We first denote the probability density function for distirbution $q (\bm{x}^k\vert\bm{x}^0)$, $q (\check{\bm{x}}^k\vert\check{\bm{x}}^0)$, and $q (\tilde{\bm{x}}^k\vert\tilde{\bm{x}}^0)$ as $\mathrm{Pr}(\cdot)$, $\check{\mathrm{Pr}}(\cdot)$, and $\tilde{\mathrm{Pr}}(\cdot)$ respectively. According to the definition of KL divergence, we have:
\begin{equation}\label{kl_diff_init}
\begin{aligned}
& \mathbb{D}_{KL}[q (\bm{x}^k\vert\bm{x}^0) \Vert q (\tilde{\bm{x}}^k\vert\tilde{\bm{x}}^0)] -
\mathbb{D}_{KL}[q (\bm{x}^k\vert\bm{x}^0) \Vert q (\check{\bm{x}}^k\vert\check{\bm{x}}^0)] \\
= & \int_{\bm{x}} \mathrm{Pr}(\bm{x}) \log\frac{\mathrm{Pr}(\bm{x})}{\mathrm{\tilde{Pr}}(\bm{x})} -
\int_{\bm{x}} \mathrm{Pr}(\bm{x}) \log\frac{\mathrm{Pr}(\bm{x})}{\mathrm{\check{Pr}}(\bm{x})} \\
= & \int_{\bm{x}} \mathrm{Pr}(\bm{x}) \log\frac{\mathrm{\check{Pr}}(\bm{x})}{\mathrm{\tilde{Pr}}(\bm{x})}.
\end{aligned}
\end{equation}
Note that the probability density function for $q (\check{\bm{x}}^k\vert\check{\bm{x}}^0)$ can be expressed as a mixture of two components: $\mathrm{\check{Pr}}(\bm{x})=\gamma \mathrm{\tilde{Pr}}(\bm{x}) + (1-\gamma) \mathrm{{Pr}}(\bm{x})$, where $\gamma \in [0,1]$ represents the probability that the indicator variable$\mathbb{I}_\mathrm{noise}=1$. Substituting this relation into Eq.~\eqref{kl_diff_init}, we have:
\begin{equation}
\begin{aligned}
& \mathbb{D}_{KL}[q (\bm{x}^k\vert\bm{x}^0) \Vert q (\tilde{\bm{x}}^k\vert\tilde{\bm{x}}^0)] -
\mathbb{D}_{KL}[q (\bm{x}^k\vert\bm{x}^0) \Vert q (\check{\bm{x}}^k\vert\check{\bm{x}}^0)] \\
= & \int_{\bm{x}} \mathrm{Pr}(\bm{x}) \log\frac{\gamma \mathrm{\tilde{Pr}}(\bm{x}) + (1-\gamma) \mathrm{{Pr}}(\bm{x})}{\mathrm{\tilde{Pr}}(\bm{x})} \\
= & - \int_{\bm{x}} \mathrm{Pr}(\bm{x}) \log\frac{\mathrm{\tilde{Pr}}(\bm{x})}{\gamma \mathrm{\tilde{Pr}}(\bm{x}) + (1-\gamma) \mathrm{{Pr}}(\bm{x})} \\
\stackrel{\romannumeral1}{\geq} & - \log \int_{\bm{x}} \mathrm{Pr}(\bm{x}) \frac{\mathrm{\tilde{Pr}}(\bm{x})}{\gamma \mathrm{\tilde{Pr}}(\bm{x}) + (1-\gamma) \mathrm{{Pr}}(\bm{x})} \\
\stackrel{\romannumeral2}{\geq} & - \log \int_{\bm{x}} \mathrm{Pr}(\bm{x}) \frac{\mathrm{\tilde{Pr}}(\bm{x})}{ {\mathrm{\tilde{Pr}(\bm{x})}^\gamma \cdot {\mathrm{{Pr}}(\bm{x})}^{(1-\gamma)}}} \\
= & - \log \int_{\bm{x}} {\mathrm{Pr}(\bm{x})}^{\gamma} \cdot {\mathrm{\tilde{Pr}}(\bm{x})}^{(1-\gamma)} \\
\stackrel{\romannumeral3}{\geq} & - \log \int_{\bm{x}} {\gamma} {\mathrm{Pr}(\bm{x})} + {(1-\gamma)} \cdot {\mathrm{\tilde{Pr}}(\bm{x})} \\
= & \log[\gamma + (1-\gamma)] = 0,
\end{aligned}
\end{equation}
where the inequality $(\romannumeral1)$ holds due to Jensen's inequality, and inequality $(\romannumeral1, \romannumeral3)$ follows from weighted Arithmetic-Geometric mean inequality. They all hold with equality if $\gamma=1$, i.e., the indicator takes $\mathbb{I}_\mathrm{noise}=1$ for all samples.
\begin{lemma}\label{lemma_a2}
For any $c>0$, with the bounded Gaussian noise scale $\iota$ and the samples $\tilde{\bm{x}}$ from Eq.~\eqref{def_all}, one can always find $k_a$ such that for any $k\geq k_a$, the KL divergence between $q (\bm{x}^k\vert\bm{x}^0)$ and $q (\tilde{\bm{x}}^k\vert\tilde{\bm{x}}^0)$ satisfying:
\begin{equation}\label{kl_divergence_threshold}
\mathbb{D}_{KL}[q (\bm{x}^k\vert\bm{x}^0) \Vert q (\tilde{\bm{x}}^k\vert\tilde{\bm{x}}^0)] < c.
\end{equation}
\end{lemma}
\textit{Proof.}
According to the definition of DDPM forward process as in Eq.~\eqref{ddpm_forward}, we have
\begin{equation}
\bm{x}^k\vert\bm{x}^0 = \sqrt{\bar{\alpha}_k} \bm{x}^0 + \sqrt{1-\bar{\alpha}_k} \bm{\epsilon} \qquad \qquad
\end{equation}
\begin{equation}
\begin{aligned}
\bm{\tilde{x}}^k\vert\bm{\tilde{x}}^0 & = \sqrt{\bar{\alpha}_k} \bm{\tilde{x}}^0 + \sqrt{1-\bar{\alpha}_k} \bm{\epsilon}_1 \\
& = \sqrt{\bar{\alpha}_k} (\bm{x}^0 + \iota\cdot \bm{\epsilon}) + \sqrt{1-\bar{\alpha}_k} \bm{\epsilon}_2 \\
& = \sqrt{\bar{\alpha}_k} \bm{x}^0 + \sqrt{1-\bar{\alpha}_k + \iota^2\cdot \bar{\alpha}_k} \bm{\epsilon}
\end{aligned}
\end{equation}
Note that $\bm{\epsilon} \in \mathbb{R}^{m\times n}$ is the matrix of Gaussian noise with i.i.d. entries $\epsilon_{ij}\sim \mathcal{N}(0,1)$.
By flattening the matrix into a vector, we can explicitly express the probability density function of the distribution $q (\bm{x}^k\vert\bm{x}^0)$ as follows:
\begin{equation}
\mathrm{Pr}(\bm{x}^k\vert\bm{x}^0) = \frac{1}{(2\pi)^{mn/2} |\bm{\Sigma}|^{1/2}} \cdot
\exp\left( -\frac{1}{2} (\bm{x}^k - \sqrt{\bar{\alpha}_k} \bm{{x}}^0)^\top \bm{\Sigma}^{-1} (\bm{x}^k - \sqrt{\bar{\alpha}_k} \bm{{x}}^0) \right)
\end{equation}
where $\bm{\Sigma} \in \mathbb{R}^{mn}$ is a diagonal matrix with all diagonal elements equal to $1-\bar{\alpha}_k$. Similarly, the probability density function of the distribution $q (\tilde{\bm{x}}^k\vert\tilde{\bm{x}}^0)$ follows:
\begin{equation}
\mathrm{\tilde{Pr}}(\bm{\tilde{x}}^k\vert\bm{x}^0) = \frac{1}{(2\pi)^{mn/2} |\bm{\tilde{\Sigma}}|^{1/2}} \cdot
\exp\left( -\frac{1}{2} (\bm{\tilde{x}}^k - \sqrt{\bar{\alpha}_k} \bm{{x}}^0)^\top \bm{\tilde{\Sigma}}^{-1} (\bm{\tilde{x}}^k - \sqrt{\bar{\alpha}_k} \bm{{x}}^0) \right)
\end{equation}
where $\bm{\tilde{\Sigma}}\in \mathbb{R}^{mn}$ is a diagonal matrix with all diagonal elements equal to $1-\bar{\alpha}_k + \iota^2\bar{\alpha}_k$. Then we can derive the KL divergence between $q (\bm{x}^k\vert\bm{x}^0)$ and $q (\tilde{\bm{x}}^k\vert\tilde{\bm{x}}^0)$:
\begin{equation}\label{kl_final_expression}
\begin{aligned}
& \mathbb{D}_{KL}[q (\bm{x}^k\vert\bm{x}^0) \Vert q (\tilde{\bm{x}}^k\vert\tilde{\bm{x}}^0)] \\
= & \frac{1}{2} \bigg[ \log\frac{|\bm{\tilde{\Sigma}}|}{|\bm{{\Sigma}}|} + tr\{ \bm{\tilde{\Sigma}}^{-1} \bm{{\Sigma}}\}- mn \bigg] \\
= & \frac{mn}{2} \bigg[ \log\frac{1-\bar{\alpha}_k + \iota^2\bar{\alpha}_k}{1-\bar{\alpha}_k} + \frac{1-\bar{\alpha}_k}{1-\bar{\alpha}_k + \iota^2\bar{\alpha}_k} - 1 \bigg]
\end{aligned}
\end{equation}
Let $f(k)=\frac{1-\bar{\alpha}_k + \iota^2\bar{\alpha}_k}{1-\bar{\alpha}_k}$, where $\bar{\alpha}_k$ is a monotonically decreasing function of $k$ with $\bar{\alpha}_k\in [0,1)$. It is straightforward to deduce that $f(k)$ is also a monotonically decreasing function of $k$, with $f(k)\in[1, \infty)$ if $\iota$ is bounded. Substituting this expression into Eq.~\eqref{kl_final_expression}, we obtain:
\begin{equation}\mathbb{D}_{KL}[q (\bm{x}^k\vert\bm{x}^0) \Vert q (\tilde{\bm{x}}^k\vert\tilde{\bm{x}}^0)]
= \frac{mn}{2} \bigg[ \log f(k) + \frac{1}{f(k)} - 1 \bigg],
\end{equation}
which is a monotonically increasing function of $f(k)$ within the range of $f(k)$. It attains its minimum value of 0 when $f(k)=1$. Therefore, for any $c>0$, we can always find $k_a$ such that for any $k\geq k_a$, the inequality in Eq.~\eqref{kl_divergence_threshold} holds.
By combining Lemma~\ref{lemma_a1} and Lemma~\ref{lemma_a2}, and under the validity of Assumption~\ref{assumption22}, the proof of Theorem~\ref{theorem23} is complete.
\subsection{Proof for Proposition 4.4}\label{proof_propsition_44}
We begin with the noise prediction at diffusion timestep $k$. If the noise is perfectly predicted, it should follow the form:
\begin{equation}\label{eps_pred}
\bm{\epsilon}_{\mathrm{pred}} = \frac{\bm{x}^k - \sqrt{\bar{\alpha}_k}\bm{x}^0}{\sqrt{1 - \bar{\alpha}_k}}
\end{equation}
Notebly, any noised/un-noised sample can be expressed as $\check{\bm{x}} = \bm{x}^0 + \iota\cdot \mathbb{I}_\mathrm{noise}\cdot\bm{\epsilon}$, where $\bm{x}^0 \in \mathbb{R}^{m\times n}$ is the noise-free matrix, $\iota$ is the noise scale, $\bm{\epsilon} \in \mathbb{R}^{m\times n}$ is the matrix of Gaussian noise with i.i.d. entries $\epsilon_{ij}\sim \mathcal{N}(0,1)$, and $\mathbb{I}_\mathrm{noise}$ is an unknown indicator variable that determines whether noise is present.
To predict the noise within $\check{\bm{x}}$ using diffusion timestep $k$, the first step is to perform the original information consistency operation:
\begin{equation}\label{check_x_k}
h_k(\check{\bm{x}}) = \sqrt{\bar{\alpha}_k} \cdot \check{\bm{x}} = \sqrt{\bar{\alpha}_k} \cdot (\bm{x}^0 + \iota\cdot \mathbb{I}_\mathrm{noise}\cdot\bm{\epsilon} ).
\end{equation}
Then substituting Eq.\eqref{check_x_k} into Eq.\eqref{eps_pred}, we obtain:
\begin{equation}\label{eps_pred2}
\bm{\epsilon}_{\mathrm{pred}} = \frac{\sqrt{\bar{\alpha}_k}\cdot \iota \cdot \mathbb{I}_\mathrm{noise}}{\sqrt{1 - \bar{\alpha}_k}} \cdot \bm{\epsilon}
\end{equation}
Assume we have a noise predictor $\bm{\epsilon}_{\theta}(\cdot, k)$ with prediction error $\bm{\delta}_{\theta}^k \sim \mathcal{N}(\bm{0}, \sigma^2_k\bm{I})$. Given a sample $\check{\bm{x}}$ that may or may not contain noise, we have the following two criteria:
% Firstly, we recap the definition of the forward process of diffusion models:
% \begin{equation}
% \bm{x}^k = \sqrt{\bar{\alpha}_k} \bm{x}^0 + \sqrt{1-\bar{\alpha}_k} \bm{\epsilon}
% \end{equation}
% where $\bm{x}^0 \in \mathbb{R}^{m\times n}$ is the noise-free matrix, $\bm{\epsilon} \in \mathbb{R}^{m\times n}$ is the matrix of Gaussian noise with i.i.d. entries $\epsilon_{ij}\sim \mathcal{N}(0,1)$, and $k$ is any diffusion timestep.
% Assume we have a noise predictor $\bm{\epsilon}_{\theta}(\cdot, k)$ with prediction error $\bm{\delta}_{\theta}^k \sim \mathcal{N}(\bm{0}, \sigma^2_k\bm{I})$. Given a sample $\check{\bm{x}}$ that may or may not contain noise, we have the following two critieria:
\paragraph{Case 1: Noise-free data ($\mathbb{I}_\mathrm{noise}=0$)}
Under such circumstance, the model's prediction is entirely determined by the prediction error:
\begin{equation}
\bm{\epsilon}_{\theta}(\check{\bm{x}}^k, k) = \bm{\epsilon}_{\mathrm{pred}} + \bm{\delta}_{\theta}^k = \bm{0} + \bm{\delta}_{\theta}^k,
\end{equation}
thereby the expected Frobenius norm squared of the prediction is:
\begin{equation}\label{noise_free_exp}
\begin{aligned}
\mathbb{E}[\Vert\bm{\epsilon}_{\theta}(\check{\bm{x}}^k, k)\Vert^2_F] &= \mathbb{E}[\Vert\bm{\delta}_{\theta}^k \Vert^2_F] \\
&= \mathbb{E}[\sum\nolimits_i\sum\nolimits_j (\delta_{\theta, ij}^k)^2] \\
& = m\cdot n \cdot \sigma_k^2
\end{aligned}
\end{equation}
\paragraph{Case 2: Noisy data ($\mathbb{I}_\mathrm{noise}=1$)} For noisy data, the model's prediction includes both the actual noise and the prediction error. Specifically, we have:
\begin{equation}
\bm{\epsilon}_{\theta}(\check{\bm{x}}^k, k) = \bm{\epsilon}_{\mathrm{pred}} + \bm{\delta}_{\theta}^k = \frac{\sqrt{\bar{\alpha}_k}\cdot \iota \cdot \mathbb{I}_\mathrm{noise}}{\sqrt{1 - \bar{\alpha}_k}} \cdot \bm{\epsilon} + \bm{\delta}_{\theta}^k,
\end{equation}
The expected squared Frobenius norm of the prediction is then given by:
\begin{equation}\label{noisy_exp}
\begin{aligned}
& \quad \mathbb{E}[\Vert\bm{\epsilon}_{\theta}(\check{\bm{x}}^k, k)\Vert^2_F] \\
& = \mathbb{E}[\Vert\frac{\sqrt{\bar{\alpha}_k}\cdot \iota \cdot \mathbb{I}_\mathrm{noise}}{\sqrt{1 - \bar{\alpha}_k}} \cdot \bm{\epsilon} + \bm{\delta}_{\theta}^k \Vert^2_F] \\
& \stackrel{i}{=} \mathbb{E}[\Vert\frac{\sqrt{\bar{\alpha}_k}\cdot \iota \cdot \mathbb{I}_\mathrm{noise}}{\sqrt{1 - \bar{\alpha}_k}} \cdot \bm{\epsilon} \Vert^2_F] + \mathbb{E}[\Vert\bm{\delta}_{\theta}^k \Vert^2_F] + \frac{2\sqrt{\bar{\alpha}_k}\cdot \iota \cdot \mathbb{I}_\mathrm{noise}}{\sqrt{1 - \bar{\alpha}_k}}\mathbb{E}[\mathrm{Tr({\bm{\epsilon}}^\top \bm{\delta}_{\theta}^k)}]\\
& = \frac{\iota^2 \cdot \bar{\alpha}_k}{(1- \bar{\alpha}_k)}\cdot \mathbb{E}[\Vert \bm{\epsilon} \Vert^2_F]+ m\cdot n \cdot \sigma_k^2 + 0 \\
& = \frac{\iota^2 \cdot \bar{\alpha}_k}{(1- \bar{\alpha}_k)}\cdot m\cdot n + m\cdot n \cdot \sigma_k^2
\end{aligned}
\end{equation}
The last term in equality $(i)$ vanishes because $\bm{\epsilon}$ and $\bm{\delta}_{\theta}^k$ are mutually independent.
Combining the results of Eq.~\eqref{noise_free_exp} and Eq.~\eqref{noisy_exp}, and substituting into the definition of the Signal-to-Noise Ratio (SNR) from Eq.~\eqref{eq_SNR}, we derive:
\begin{equation}
\begin{aligned}
\mathrm{SNR}(k) &:= \frac{\mathbb{E}[\Vert \bm{\epsilon}_{\theta}(\check{\bm{x}}_{\mathrm{ns}}^k, k) \Vert_F^2] - \mathbb{E}[\Vert \bm{\epsilon}_{\theta}(\check{\bm{x}}_{\mathrm{nf}}^k, k) \Vert_F^2]}{\mathbb{E}[\Vert \bm{\epsilon}_{\theta}(\check{\bm{x}}_{\mathrm{nf}}^k, k) \Vert_F^2]} \\
& = \frac{\frac{\iota^2 \cdot \bar{\alpha}_k}{(1- \bar{\alpha}_k)}\cdot m\cdot n + m\cdot n \cdot \sigma_k^2 - m\cdot n \cdot \sigma_k^2}{m\cdot n \cdot \sigma_k^2} \\
& = \frac{\iota^2 \cdot \bar{\alpha}_k}{(1- \bar{\alpha}_k) \cdot \sigma^2_k},
\end{aligned}
\end{equation}
which complete the proof.
\begin{algorithm}[htb]
\caption{Ambient Diffusion-Guided Dataset Recovery~(ADG)}
\label{alg:alg}
\begin{algorithmic}
\REQUIRE Offline partially corrupted dataset $D$, initialized noise predictors detector $\epsilon_\theta$ and the denoiser $\epsilon_\phi$.
\STATE \textbf{Step 1: Update the detector $\epsilon_\theta$ using Ambient Loss}
\FOR{each iteration}
\STATE Sample a trajectory mini-batch $B = \{(\check{\bm{z}}_{t-H}, \dots, \check{\bm{z}}_{t+H})\} \sim D$, where $\check{\bm{z}}$ represents either the state $s$ or the concatenation of $(s, a, r)$ may or may not contain noise
\STATE Sample uniformly distributed diffusion timestep $k \sim \{k_a, \dots, K\}$
\STATE Sample random Gaussian noise ${\epsilon}^k_t \sim \mathcal{N}(\mathbf{0}, \boldsymbol{I})$
\STATE Produce noised element through $\tilde{\bm{z}}_t^k=\sqrt{\bar{\alpha}_k} \check{\bm{z}}_t+\sqrt{1-\bar{\alpha}_k} \epsilon_t^k$
\STATE Get trajectory $\check{\bm{\tau}}_t = [\check{\bm{z}}_{t-H}, \dots, \check{\bm{z}}_{t+H}]\in \mathbb{R}^{M\times (2H+1)}$
\STATE Update the $\epsilon_\theta$ through Eq.~\ref{corollay_eq4}.
\ENDFOR
\STATE \textbf{Step 2: Update the denoiser $\epsilon_\phi$ using Naive loss}
\FOR{each iteration}
\STATE Sample a trajectory mini-batch $B = \{(\check{\bm{z}}_{t-H}, \dots, \check{\bm{z}}_{t+H})\} \sim D$, where $\check{\bm{z}}$ represents either the state $s$ or the concatenation of $(s, a, r)$ may or may not contain noise
\STATE Sample uniformly distributed diffusion timestep $k \sim \{1, \dots, K\}$
\STATE Sample random Gaussian noise ${\epsilon}^k_t \sim \mathcal{N}(\mathbf{0}, \boldsymbol{I})$
\STATE Produce noised element through $\tilde{\bm{z}}_t^k=\sqrt{\bar{\alpha}_k} \check{\bm{z}}_t+\sqrt{1-\bar{\alpha}_k} \epsilon_t^k$
\STATE Get trajectory $\check{\bm{\tau}}_t = [\check{\bm{z}}_{t-H}, \dots, \check{\bm{z}}_{t+H}]\in \mathbb{R}^{M\times (2H+1)}$
\STATE Get $e_\theta(\check{\bm{z}}_t)$ using detector $\epsilon_\theta(\check{\bm{\tau}}_t, k)$ with $k=k_a$ \STATE Get mask $\bm{m}_t:=[\mathbb{I}_{t-H}, \dots, \mathbb{I}_{t+H}]$, where $\mathbb{I}_t = 0$ for $e_\theta(\check{\bm{z}})>\zeta$ and $\mathbb{I}_t = 1$ for $e_\theta(\check{\bm{z}})\leq\zeta$
\STATE Update the $\epsilon_\phi$ through Eq.~\ref{ddpm_loss}.
\ENDFOR
\STATE \textbf{Step 3: Detect and Recover the noised dataset}
\FOR{each $\check{\bm{z}}_t$ in the noised dataset with $e_\theta(\check{\bm{z}})>\zeta$}
\STATE Get the trajectory $\check{\bm{\tau}}_t := [\check{\bm{z}}_{t-H}, \dots, \check{\bm{z}}_{t+H}]$
\STATE Recover the trajectory $\breve{\bm{\tau}}_t=\frac{1}{\sqrt{\bar{\alpha}_k}}\left[\check{\bm{\tau}}_t^k-\sqrt{1-\bar{\alpha}_k} \epsilon_\theta\left(\check{\bm{\tau}}_t^k, k\right)\right]$
\STATE Replace $\check{\bm{z_t}}$ with $(\breve{\bm{\tau}}_{t})_{H+1}$, which represents the (H+1)-th column of $\breve{\bm{\tau}}$
\ENDFOR
\end{algorithmic}
\end{algorithm}
\begin{table*}[ht]
\centering
\caption{Results on ``medium-v2'' and ``expert-v2'' datasets under Random State Attack.}
\label{tab:medium_and_expert}
\resizebox{\textwidth}{!}{
\begin{tabular}{l|l|c>{\columncolor{gray!20}}cc>{\columncolor{gray!20}}cc>{\columncolor{gray!20}}c}
\toprule
\multirow{2}{*}{Dataset} & \multirow{2}{*}{Task} & \multicolumn{2}{c}{RIQL} & \multicolumn{2}{c}{DT} & \multicolumn{2}{c}{RDT} \\
& & Naive& ADG & Naive& ADG & Naive & ADG \\
\hline
\midrule
\multirow{5}{*}{medium} & halfcheetah(2\%) & 18.0$\pm$1.5 & 24.8$\pm$2.1 & 15.7$\pm$1.3 & 31.1$\pm$1.4 & 22.3$\pm$1.1 & 22.0$\pm$3.0 \\
& hopper(2\%) & 47.5$\pm$7.3 & 44.1$\pm$2.7 & 48.6$\pm$3.2 & 51.4$\pm$4.0 & 52.2$\pm$5.7 & 54.1$\pm$4.9 \\
& walker2d(2\%) & 25.4$\pm$5.0 & 26.4$\pm$4.0 & 20.5$\pm$6.0 & 50.0$\pm$4.7 & 28.0$\pm$8.0 & 46.8$\pm$3.6 \\
\cline{2-8}\rule{0pt}{2.5ex}
& \textbf{Average} & 30.3 & \textbf{31.8} & 28.3 & \textbf{44.2} & 34.2 & \textbf{41.0}\\
& \textbf{Improvement $\uparrow$} & & \textbf{4.95\%} & & \textbf{56.18\%} & & \textbf{19.88\%}\\
\midrule
\hline\rule{0pt}{2.5ex}
\multirow{5}{*}{expert} & halfcheetah(2\%) & 0.5$\pm$1.2 & 1.2$\pm$0.7 & 2.9$\pm$0.8 & 4.2$\pm$1.0 & 4.4$\pm$0.2 & 2.02$\pm$0.6 \\
& hopper(2\%) & 32.0$\pm$4.4 & 39.7$\pm$12.9 & 38.5$\pm$7.4 & 41.3$\pm$4.3 & 48.6$\pm$7.0 & 42.6$\pm$7.0 \\
& walker2d(2\%) & 21.7$\pm$4.6 & 46.6$\pm$2.7 & 40.7$\pm$4.8 & 53.2$\pm$3.2 & 41.6$\pm$4.2 & 55.3$\pm$6.1 \\
\cline{2-8}\rule{0pt}{2.5ex}
& \textbf{Average} & 18.1 & \textbf{29.2} & 27.4 & \textbf{32.9} & 31.5 & \textbf{33.3} \\
& \textbf{Improvement $\uparrow$} & & \textbf{61.33\%} & & \textbf{20.07\%} & & \textbf{5.71\%}\\
\bottomrule
\end{tabular}
}
\end{table*}
\begin{thebibliography}{10}
\bibitem{aali2024ambient}
Asad Aali, Giannis Daras, Brett Levac, Sidharth Kumar, Alexandros~G Dimakis, and Jonathan~I Tamir.
\newblock Ambient diffusion posterior sampling: Solving inverse problems with diffusion models trained on corrupted data.
\newblock {\em arXiv preprint arXiv:2403.08728}, 2024.
\bibitem{ackermann2024offline}
Johannes Ackermann, Takayuki Osa, and Masashi Sugiyama.
\newblock Offline reinforcement learning from datasets with structured non-stationarity.
\newblock {\em arXiv preprint arXiv:2405.14114}, 2024.
\bibitem{ajay2023conditional}
Anurag Ajay, Yilun Du, Abhi Gupta, Joshua~B Tenenbaum, Tommi~S Jaakkola, and Pulkit Agrawal.
\newblock Is conditional generative modeling all you need for decision making?
\newblock In {\em The Eleventh International Conference on Learning Representations}, 2023.
\bibitem{an2021uncertainty}
Gaon An, Seungyong Moon, Jang-Hyun Kim, and Hyun~Oh Song.
\newblock Uncertainty-based offline reinforcement learning with diversified q-ensemble.
\newblock {\em Advances in neural information processing systems}, 34:7436--7447, 2021.
\bibitem{bai2022pessimistic}
Chenjia Bai, Lingxiao Wang, Zhuoran Yang, Zhihong Deng, Animesh Garg, Peng Liu, and Zhaoran Wang.
\newblock Pessimistic bootstrapping for uncertainty-driven offline reinforcement learning.
\newblock {\em arXiv preprint arXiv:2202.11566}, 2022.
\bibitem{chen2021decision}
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch.
\newblock Decision transformer: Reinforcement learning via sequence modeling.
\newblock {\em Advances in neural information processing systems}, 34:15084--15097, 2021.
\bibitem{chen2024exact}
Yiding Chen, Xuezhou Zhang, Qiaomin Xie, and Xiaojin Zhu.
\newblock Exact policy recovery in offline rl with both heavy-tailed rewards and data corruption.
\newblock In {\em Proceedings of the AAAI Conference on Artificial Intelligence}, volume~38, pages 11416--11424, 2024.
\bibitem{daras24ambient}
Giannis Daras, Alex Dimakis, and Constantinos~Costis Daskalakis.
\newblock Consistent diffusion meets tweedie: Training exact ambient diffusion models with noisy data.
\newblock In {\em Forty-first International Conference on Machine Learning}, 2024.
\bibitem{fu2020d4rl}
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine.
\newblock D4rl: Datasets for deep data-driven reinforcement learning.
\newblock {\em arXiv preprint arXiv:2004.07219}, 2020.
\bibitem{fujimoto2021minimalist}
Scott Fujimoto and Shixiang~Shane Gu.
\newblock A minimalist approach to offline reinforcement learning.
\newblock {\em Advances in neural information processing systems}, 34:20132--20145, 2021.
\bibitem{fujimoto2019off}
Scott Fujimoto, David Meger, and Doina Precup.
\newblock Off-policy deep reinforcement learning without exploration.
\newblock In {\em International conference on machine learning}, pages 2052--2062. PMLR, 2019.
\bibitem{ghasemipour2022so}
Kamyar Ghasemipour, Shixiang~Shane Gu, and Ofir Nachum.
\newblock Why so pessimistic? estimating uncertainties for offline rl through ensembles, and why their independence matters.
\newblock {\em Advances in Neural Information Processing Systems}, 35:18267--18281, 2022.
\bibitem{hansen2023idql}
Philippe Hansen-Estruch, Ilya Kostrikov, Michael Janner, Jakub~Grudzien Kuba, and Sergey Levine.
\newblock {IDQL}: Implicit {Q}-learning as an actor-critic method with diffusion policies.
\newblock {\em arXiv preprint arXiv:2304.10573}, 2023.
\bibitem{ho2020denoising}
Jonathan Ho, Ajay Jain, and Pieter Abbeel.
\newblock {Denoising diffusion probabilistic models}.
\newblock {\em Advances in Neural Information Processing Systems}, 33:6840--6851, 2020.
\bibitem{janner2022planning}
Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine.
\newblock Planning with diffusion for flexible behavior synthesis.
\newblock In {\em International Conference on Machine Learning}, pages 9902--9915. PMLR, 2022.
\bibitem{janner2021offline}
Michael Janner, Qiyang Li, and Sergey Levine.
\newblock Offline reinforcement learning as one big sequence modeling problem.
\newblock {\em Advances in neural information processing systems}, 34:1273--1286, 2021.
\bibitem{kopruluneural}
Cevahir Koprulu, Franck Djeumou, et~al.
\newblock Neural stochastic differential equations for uncertainty-aware offline rl.
\newblock In {\em The Thirteenth International Conference on Learning Representations}.
\bibitem{kostrikov2021offline}
Ilya Kostrikov, Ashvin Nair, and Sergey Levine.
\newblock Offline reinforcement learning with implicit q-learning.
\newblock {\em arXiv preprint arXiv:2110.06169}, 2021.
\bibitem{kumar2020conservative}
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine.
\newblock Conservative q-learning for offline reinforcement learning.
\newblock {\em Advances in Neural Information Processing Systems}, 33:1179--1191, 2020.
\bibitem{levine2020offlinereinforcementlearningtutorial}
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu.
\newblock Offline reinforcement learning: Tutorial, review, and perspectives on open problems, 2020.
\bibitem{li2024survival}
Anqi Li, Dipendra Misra, Andrey Kolobov, and Ching-An Cheng.
\newblock Survival instinct in offline reinforcement learning.
\newblock {\em Advances in neural information processing systems}, 36, 2024.
\bibitem{liang2024robust}
Xize Liang, Chao Chen, Jie Wang, Yue Wu, Zhihang Fu, Zhihao Shi, Feng Wu, and Jieping Ye.
\newblock Robust preference optimization with provable noise tolerance for llms.
\newblock {\em arXiv preprint arXiv:2404.04102}, 2024.
\bibitem{liang2023adaptdiffuser}
Zhixuan Liang, Yao Mu, Mingyu Ding, Fei Ni, Masayoshi Tomizuka, and Ping Luo.
\newblock Adaptdiffuser: Diffusion models as adaptive self-evolving planners.
\newblock In {\em International Conference on Machine Learning}, pages 20725--20745. PMLR, 2023.
\bibitem{liu2024adaptive}
Tenglong Liu, Yang Li, Yixing Lan, Hao Gao, Wei Pan, and Xin Xu.
\newblock Adaptive advantage-guided policy regularization for offline reinforcement learning.
\newblock {\em arXiv preprint arXiv:2405.19909}, 2024.
\bibitem{madry2017towards}
Aleksander Madry.
\newblock Towards deep learning models resistant to adversarial attacks.
\newblock {\em arXiv preprint arXiv:1706.06083}, 2017.
\bibitem{panaganti2022robust}
Kishan Panaganti, Zaiyan Xu, Dileep Kalathil, and Mohammad Ghavamzadeh.
\newblock Robust reinforcement learning using offline data.
\newblock {\em Advances in neural information processing systems}, 35:32211--32224, 2022.
\bibitem{rombach2022high}
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bj{\"o}rn Ommer.
\newblock High-resolution image synthesis with latent diffusion models.
\newblock In {\em Proceedings of the IEEE/CVF conference on computer vision and pattern recognition}, pages 10684--10695, 2022.
\bibitem{shi2024distributionally}
Laixi Shi and Yuejie Chi.
\newblock Distributionally robust model-based offline reinforcement learning with near-optimal sample complexity.
\newblock {\em Journal of Machine Learning Research}, 25(200):1--91, 2024.
\bibitem{shi2023unleashing}
Ruizhe Shi, Yuyao Liu, Yanjie Ze, Simon~S Du, and Huazhe Xu.
\newblock Unleashing the power of pre-trained language models for offline reinforcement learning.
\newblock {\em arXiv preprint arXiv:2310.20587}, 2023.
\bibitem{sohl2015deep}
Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli.
\newblock {Deep unsupervised learning using nonequilibrium thermodynamics}.
\newblock In {\em International Conference on Machine Learning}, pages 2256--2265. PMLR, 2015.
\bibitem{song2020score}
Yang Song, Jascha Sohl-Dickstein, Diederik~P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole.
\newblock Score-based generative modeling through stochastic differential equations.
\newblock In {\em International Conference on Learning Representations}, 2020.
\bibitem{wang2022diffusion}
Zhendong Wang, Jonathan~J Hunt, and Mingyuan Zhou.
\newblock Diffusion policies as an expressive policy class for offline reinforcement learning.
\newblock In {\em The Eleventh International Conference on Learning Representations}, 2022.
\bibitem{wu2022copa}
Fan Wu, Linyi Li, Chejian Xu, Huan Zhang, Bhavya Kailkhura, Krishnaram Kenthapadi, Ding Zhao, and Bo~Li.
\newblock Copa: Certifying robust policies for offline reinforcement learning against poisoning attacks.
\newblock {\em arXiv preprint arXiv:2203.08398}, 2022.
\bibitem{xu2024robust}
Jiawei Xu, Rui Yang, Shuang Qiu, Feng Luo, Meng Fang, Baoxiang Wang, and Lei Han.
\newblock Tackling data corruption in offline reinforcement learning via sequence modeling.
\newblock In {\em The Thirteenth International Conference on Learning Representations}, 2025.
\bibitem{yang2022rorl}
Rui Yang, Chenjia Bai, Xiaoteng Ma, Zhaoran Wang, Chongjie Zhang, and Lei Han.
\newblock Rorl: Robust offline reinforcement learning via conservative smoothing.
\newblock {\em Advances in neural information processing systems}, 35:23851--23866, 2022.
\bibitem{yang2024regularizing}
Rui Yang, Ruomeng Ding, Yong Lin, Huan Zhang, and Tong Zhang.
\newblock Regularizing hidden states enables learning generalizable reward model for llms.
\newblock {\em arXiv preprint arXiv:2406.10216}, 2024.
\bibitem{yang2023towards}
Rui Yang, Han Zhong, Jiawei Xu, Amy Zhang, Chongjie Zhang, Lei Han, and Tong Zhang.
\newblock Towards robust offline reinforcement learning under diverse data corruption.
\newblock In {\em The Twelfth International Conference on Learning Representations}, 2024.
\bibitem{yang2024dmbp}
Zhihe Yang and Yunjian Xu.
\newblock Dmbp: Diffusion model-based predictor for robust offline reinforcement learning against state observation perturbations.
\newblock In {\em The Twelfth International Conference on Learning Representations}, 2024.
\bibitem{ye2024towards}
Chenlu Ye, Jiafan He, Quanquan Gu, and Tong Zhang.
\newblock Towards robust model-based reinforcement learning against adversarial corruption.
\newblock {\em arXiv preprint arXiv:2402.08991}, 2024.
\bibitem{ye2024corruption}
Chenlu Ye, Rui Yang, Quanquan Gu, and Tong Zhang.
\newblock Corruption-robust offline reinforcement learning with general function approximation.
\newblock {\em Advances in Neural Information Processing Systems}, 36, 2024.
\bibitem{zhang2021robust}
Huan Zhang, Hongge Chen, Duane Boning, and Cho-Jui Hsieh.
\newblock Robust reinforcement learning on state observations with learned optimal adversary.
\newblock {\em arXiv preprint arXiv:2101.08452}, 2021.
\bibitem{zhang2020robust}
Huan Zhang, Hongge Chen, Chaowei Xiao, Bo~Li, Mingyan Liu, Duane Boning, and Cho-Jui Hsieh.
\newblock Robust deep reinforcement learning against adversarial perturbations on state observations.
\newblock {\em Advances in Neural Information Processing Systems}, 33:21024--21037, 2020.
\bibitem{zhang2022corruption}
Xuezhou Zhang, Yiding Chen, Xiaojin Zhu, and Wen Sun.
\newblock Corruption-robust offline reinforcement learning.
\newblock In {\em International Conference on Artificial Intelligence and Statistics}, pages 5757--5773. PMLR, 2022.
\end{thebibliography}