Enhancing Inverse Reinforcement Learning through Encoding Dynamic Information in Reward Shaping
Extracted source text · arXiv 2410.03847
\documentclass{article} % For LaTeX2e
% \usepackage{iclr2025_conference,times}
\usepackage[preprint]{neurips_2023}
% Optional math commands from https://github.com/goodfeli/dlbook_notation.
\input{math_commands.tex}
\input{header.tex}
% \newcommand{\chao}[1]{{{\color{red} (chao: #1)}}}
% \title{Model-Enhanced Adversarial Inverse Reinforcement Learning with Model Estimation Reward Shaping in Stochastic Environments}
\title{Model-Based Reward Shaping for Adversarial Inverse Reinforcement Learning in Stochastic Environments}
% Authors must not appear in the submitted version. They should be hidden
% as long as the \iclrfinalcopy macro remains commented out below.
% Non-anonymous submissions will be rejected without review.
\author{Simon Sinong Zhan\\
% Electrical and Computer Engineering\\
Northwestern University\\
Evanston, IL 60208, USA \\
% \texttt{SinongZhan2028@u.northwestern.edu} \\
\And
Qingyuan Wu \\
University of Southampton \\
Southampton, SO17 1BJ, UK\\
% \texttt{\{robot,net\}@wits.ac.za} \\
\And
Philip Wang \\
Northwestern University \\
% Address \\
% \texttt{email}
\And
Yixuan Wang \\
Northwestern University \\
\And
Ruochen Jiao \\
Northwestern University \\
\And
Chao Huang \\
University of Southampton\\
\And
Qi Zhu \\
Northwestern University \\
}
% \author{Antiquus S.~Hippocampus, Natalia Cerebro \& Amelie P. Amygdale \thanks{ Use footnote for providing further information
% about author (webpage, alternative address)---\emph{not} for acknowledging
% funding agencies. Funding acknowledgements go at the end of the paper.} \\
% Department of Computer Science\\
% Cranberry-Lemon University\\
% Pittsburgh, PA 15213, USA \\
% \texttt{\{hippo,brain,jen\}@cs.cranberry-lemon.edu} \\
% \And
% Ji Q. Ren \& Yevgeny LeNet \\
% Department of Computational Neuroscience \\
% University of the Witwatersrand \\
% Joburg, South Africa \\
% \texttt{\{robot,net\}@wits.ac.za} \\
% \AND
% Coauthor \\
% Affiliation \\
% Address \\
% \texttt{email}
% }
% 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 linebreak 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.
\newcommand{\fix}{\marginpar{FIX}}
\newcommand{\new}{\marginpar{NEW}}
%\iclrfinalcopy % Uncomment for camera-ready version, but NOT for submission.
\begin{document}
\maketitle
\begin{abstract}
In this paper, we aim to tackle the limitation of the Adversarial Inverse Reinforcement Learning (AIRL) method in stochastic environments where theoretical results cannot hold and performance is degraded. To address this issue, we propose a novel method which infuses the dynamics information into the reward shaping with the theoretical guarantee for the induced optimal policy in the stochastic environments. Incorporating our novel model-enhanced rewards, we present a novel Model-Enhanced AIRL framework, which integrates transition model estimation directly into reward shaping. Furthermore, we provide a comprehensive theoretical analysis of the reward error bound and performance difference bound for our method. The experimental results in MuJoCo benchmarks show that our method can achieve superior performance in stochastic environments and competitive performance in deterministic environments, with significant improvement in sample efficiency, compared to existing baselines.
% In this paper, we tackle stochastic limitations of Adversarial Imitation Learning (AIL) by using model-based reinforcement learning (RL) techniques and model estimation reward shaping. While adversarial methods demonstrate theoretical elegance and promising performance in deterministic setting, their theoretical results and performance often degrade in stochastic environments. To address these challenges, we propose a novel model-based framework for Adversarial Inverse Reinforcement Learning (IRL). By seamlessly integrating transition model estimation directly into the reward shaping, our method establishes theoretical optimal guarantees for state-action shaped reward functions, improving performance in stochastic environments. Additionally, our approach significantly enhances sample efficiency while maintaining comparable performance in deterministic settings by generating synthetic trajectories for updating reward and policy. We further provide both theoretical and empirical analyses of the optimality gap under the transition model error, demonstrating the effectiveness of our approach.
\end{abstract}
%===============================================================================
\input{section/introduction}
\input{section/related_work}
\input{section/preliminary}
\input{section/reward_shaping}
\input{section/algorithm}
\input{section/experiment}
% \input{section/discussion}
\input{section/conclusion}
%===============================================================================
\section{Reproducible Statement}
% The full source code for our model-based adversarial Inverse Reinforcement Learning framework, along with all necessary dependencies and configurations, will be available in a public repository once cleaned up.
This work uses the open-source MuJoCo~\citep{mujoco} as the benchmark.
The practical implementation of our method is built on the \texttt{CleanRL} repository~\citep{huang2022cleanrl}.
All hyperparameters to reproduce our experimental results, including learning rates and transition model settings, are explicitly listed in Appendix~\ref{sec::implementation}. For every reported result, we averaged the performance over three random seeds, and the seed initialization is included for exact reproducibility.
% \subsubsection*{Acknowledgments}
% Use unnumbered third level headings for the acknowledgments. All
% acknowledgments, including those to funding agencies, go at the end of the paper.
\clearpage
\bibliography{iclr2021_conference}
\bibliographystyle{iclr2025_conference}
%===============================================================================
\input{section/appendix}
\end{document}
%%%%% NEW MATH DEFINITIONS %%%%%
\usepackage{amsmath,amsfonts,bm}
% Mark sections of captions for referring to divisions of figures
\newcommand{\figleft}{{\em (Left)}}
\newcommand{\figcenter}{{\em (Center)}}
\newcommand{\figright}{{\em (Right)}}
\newcommand{\figtop}{{\em (Top)}}
\newcommand{\figbottom}{{\em (Bottom)}}
\newcommand{\captiona}{{\em (a)}}
\newcommand{\captionb}{{\em (b)}}
\newcommand{\captionc}{{\em (c)}}
\newcommand{\captiond}{{\em (d)}}
% Highlight a newly defined term
\newcommand{\newterm}[1]{{\bf #1}}
% Figure reference, lower-case.
\def\figref#1{figure~\ref{#1}}
% Figure reference, capital. For start of sentence
\def\Figref#1{Figure~\ref{#1}}
\def\twofigref#1#2{figures \ref{#1} and \ref{#2}}
\def\quadfigref#1#2#3#4{figures \ref{#1}, \ref{#2}, \ref{#3} and \ref{#4}}
% Section reference, lower-case.
\def\secref#1{section~\ref{#1}}
% Section reference, capital.
\def\Secref#1{Section~\ref{#1}}
% Reference to two sections.
\def\twosecrefs#1#2{sections \ref{#1} and \ref{#2}}
% Reference to three sections.
\def\secrefs#1#2#3{sections \ref{#1}, \ref{#2} and \ref{#3}}
% Reference to an equation, lower-case.
\def\eqref#1{equation~\ref{#1}}
% Reference to an equation, upper case
\def\Eqref#1{Equation~\ref{#1}}
% A raw reference to an equation---avoid using if possible
\def\plaineqref#1{\ref{#1}}
% Reference to a chapter, lower-case.
\def\chapref#1{chapter~\ref{#1}}
% Reference to an equation, upper case.
\def\Chapref#1{Chapter~\ref{#1}}
% Reference to a range of chapters
\def\rangechapref#1#2{chapters\ref{#1}--\ref{#2}}
% Reference to an algorithm, lower-case.
\def\algref#1{algorithm~\ref{#1}}
% Reference to an algorithm, upper case.
\def\Algref#1{Algorithm~\ref{#1}}
\def\twoalgref#1#2{algorithms \ref{#1} and \ref{#2}}
\def\Twoalgref#1#2{Algorithms \ref{#1} and \ref{#2}}
% Reference to a part, lower case
\def\partref#1{part~\ref{#1}}
% Reference to a part, upper case
\def\Partref#1{Part~\ref{#1}}
\def\twopartref#1#2{parts \ref{#1} and \ref{#2}}
\def\ceil#1{\lceil #1 \rceil}
\def\floor#1{\lfloor #1 \rfloor}
\def\1{\bm{1}}
\newcommand{\train}{\mathcal{D}}
\newcommand{\valid}{\mathcal{D_{\mathrm{valid}}}}
\newcommand{\test}{\mathcal{D_{\mathrm{test}}}}
\def\eps{{\epsilon}}
% Random variables
\def\reta{{\textnormal{$\eta$}}}
\def\ra{{\textnormal{a}}}
\def\rb{{\textnormal{b}}}
\def\rc{{\textnormal{c}}}
\def\rd{{\textnormal{d}}}
\def\re{{\textnormal{e}}}
\def\rf{{\textnormal{f}}}
\def\rg{{\textnormal{g}}}
\def\rh{{\textnormal{h}}}
\def\ri{{\textnormal{i}}}
\def\rj{{\textnormal{j}}}
\def\rk{{\textnormal{k}}}
\def\rl{{\textnormal{l}}}
% rm is already a command, just don't name any random variables m
\def\rn{{\textnormal{n}}}
\def\ro{{\textnormal{o}}}
\def\rp{{\textnormal{p}}}
\def\rq{{\textnormal{q}}}
\def\rr{{\textnormal{r}}}
\def\rs{{\textnormal{s}}}
\def\rt{{\textnormal{t}}}
\def\ru{{\textnormal{u}}}
\def\rv{{\textnormal{v}}}
\def\rw{{\textnormal{w}}}
\def\rx{{\textnormal{x}}}
\def\ry{{\textnormal{y}}}
\def\rz{{\textnormal{z}}}
% Random vectors
\def\rvepsilon{{\mathbf{\epsilon}}}
\def\rvtheta{{\mathbf{\theta}}}
\def\rva{{\mathbf{a}}}
\def\rvb{{\mathbf{b}}}
\def\rvc{{\mathbf{c}}}
\def\rvd{{\mathbf{d}}}
\def\rve{{\mathbf{e}}}
\def\rvf{{\mathbf{f}}}
\def\rvg{{\mathbf{g}}}
\def\rvh{{\mathbf{h}}}
\def\rvu{{\mathbf{i}}}
\def\rvj{{\mathbf{j}}}
\def\rvk{{\mathbf{k}}}
\def\rvl{{\mathbf{l}}}
\def\rvm{{\mathbf{m}}}
\def\rvn{{\mathbf{n}}}
\def\rvo{{\mathbf{o}}}
\def\rvp{{\mathbf{p}}}
\def\rvq{{\mathbf{q}}}
\def\rvr{{\mathbf{r}}}
\def\rvs{{\mathbf{s}}}
\def\rvt{{\mathbf{t}}}
\def\rvu{{\mathbf{u}}}
\def\rvv{{\mathbf{v}}}
\def\rvw{{\mathbf{w}}}
\def\rvx{{\mathbf{x}}}
\def\rvy{{\mathbf{y}}}
\def\rvz{{\mathbf{z}}}
% Elements of random vectors
\def\erva{{\textnormal{a}}}
\def\ervb{{\textnormal{b}}}
\def\ervc{{\textnormal{c}}}
\def\ervd{{\textnormal{d}}}
\def\erve{{\textnormal{e}}}
\def\ervf{{\textnormal{f}}}
\def\ervg{{\textnormal{g}}}
\def\ervh{{\textnormal{h}}}
\def\ervi{{\textnormal{i}}}
\def\ervj{{\textnormal{j}}}
\def\ervk{{\textnormal{k}}}
\def\ervl{{\textnormal{l}}}
\def\ervm{{\textnormal{m}}}
\def\ervn{{\textnormal{n}}}
\def\ervo{{\textnormal{o}}}
\def\ervp{{\textnormal{p}}}
\def\ervq{{\textnormal{q}}}
\def\ervr{{\textnormal{r}}}
\def\ervs{{\textnormal{s}}}
\def\ervt{{\textnormal{t}}}
\def\ervu{{\textnormal{u}}}
\def\ervv{{\textnormal{v}}}
\def\ervw{{\textnormal{w}}}
\def\ervx{{\textnormal{x}}}
\def\ervy{{\textnormal{y}}}
\def\ervz{{\textnormal{z}}}
% Random matrices
\def\rmA{{\mathbf{A}}}
\def\rmB{{\mathbf{B}}}
\def\rmC{{\mathbf{C}}}
\def\rmD{{\mathbf{D}}}
\def\rmE{{\mathbf{E}}}
\def\rmF{{\mathbf{F}}}
\def\rmG{{\mathbf{G}}}
\def\rmH{{\mathbf{H}}}
\def\rmI{{\mathbf{I}}}
\def\rmJ{{\mathbf{J}}}
\def\rmK{{\mathbf{K}}}
\def\rmL{{\mathbf{L}}}
\def\rmM{{\mathbf{M}}}
\def\rmN{{\mathbf{N}}}
\def\rmO{{\mathbf{O}}}
\def\rmP{{\mathbf{P}}}
\def\rmQ{{\mathbf{Q}}}
\def\rmR{{\mathbf{R}}}
\def\rmS{{\mathbf{S}}}
\def\rmT{{\mathbf{T}}}
\def\rmU{{\mathbf{U}}}
\def\rmV{{\mathbf{V}}}
\def\rmW{{\mathbf{W}}}
\def\rmX{{\mathbf{X}}}
\def\rmY{{\mathbf{Y}}}
\def\rmZ{{\mathbf{Z}}}
% Elements of random matrices
\def\ermA{{\textnormal{A}}}
\def\ermB{{\textnormal{B}}}
\def\ermC{{\textnormal{C}}}
\def\ermD{{\textnormal{D}}}
\def\ermE{{\textnormal{E}}}
\def\ermF{{\textnormal{F}}}
\def\ermG{{\textnormal{G}}}
\def\ermH{{\textnormal{H}}}
\def\ermI{{\textnormal{I}}}
\def\ermJ{{\textnormal{J}}}
\def\ermK{{\textnormal{K}}}
\def\ermL{{\textnormal{L}}}
\def\ermM{{\textnormal{M}}}
\def\ermN{{\textnormal{N}}}
\def\ermO{{\textnormal{O}}}
\def\ermP{{\textnormal{P}}}
\def\ermQ{{\textnormal{Q}}}
\def\ermR{{\textnormal{R}}}
\def\ermS{{\textnormal{S}}}
\def\ermT{{\textnormal{T}}}
\def\ermU{{\textnormal{U}}}
\def\ermV{{\textnormal{V}}}
\def\ermW{{\textnormal{W}}}
\def\ermX{{\textnormal{X}}}
\def\ermY{{\textnormal{Y}}}
\def\ermZ{{\textnormal{Z}}}
% Vectors
\def\vzero{{\bm{0}}}
\def\vone{{\bm{1}}}
\def\vmu{{\bm{\mu}}}
\def\vtheta{{\bm{\theta}}}
\def\va{{\bm{a}}}
\def\vb{{\bm{b}}}
\def\vc{{\bm{c}}}
\def\vd{{\bm{d}}}
\def\ve{{\bm{e}}}
\def\vf{{\bm{f}}}
\def\vg{{\bm{g}}}
\def\vh{{\bm{h}}}
\def\vi{{\bm{i}}}
\def\vj{{\bm{j}}}
\def\vk{{\bm{k}}}
\def\vl{{\bm{l}}}
\def\vm{{\bm{m}}}
\def\vn{{\bm{n}}}
\def\vo{{\bm{o}}}
\def\vp{{\bm{p}}}
\def\vq{{\bm{q}}}
\def\vr{{\bm{r}}}
\def\vs{{\bm{s}}}
\def\vt{{\bm{t}}}
\def\vu{{\bm{u}}}
\def\vv{{\bm{v}}}
\def\vw{{\bm{w}}}
\def\vx{{\bm{x}}}
\def\vy{{\bm{y}}}
\def\vz{{\bm{z}}}
% Elements of vectors
\def\evalpha{{\alpha}}
\def\evbeta{{\beta}}
\def\evepsilon{{\epsilon}}
\def\evlambda{{\lambda}}
\def\evomega{{\omega}}
\def\evmu{{\mu}}
\def\evpsi{{\psi}}
\def\evsigma{{\sigma}}
\def\evtheta{{\theta}}
\def\eva{{a}}
\def\evb{{b}}
\def\evc{{c}}
\def\evd{{d}}
\def\eve{{e}}
\def\evf{{f}}
\def\evg{{g}}
\def\evh{{h}}
\def\evi{{i}}
\def\evj{{j}}
\def\evk{{k}}
\def\evl{{l}}
\def\evm{{m}}
\def\evn{{n}}
\def\evo{{o}}
\def\evp{{p}}
\def\evq{{q}}
\def\evr{{r}}
\def\evs{{s}}
\def\evt{{t}}
\def\evu{{u}}
\def\evv{{v}}
\def\evw{{w}}
\def\evx{{x}}
\def\evy{{y}}
\def\evz{{z}}
% Matrix
\def\mA{{\bm{A}}}
\def\mB{{\bm{B}}}
\def\mC{{\bm{C}}}
\def\mD{{\bm{D}}}
\def\mE{{\bm{E}}}
\def\mF{{\bm{F}}}
\def\mG{{\bm{G}}}
\def\mH{{\bm{H}}}
\def\mI{{\bm{I}}}
\def\mJ{{\bm{J}}}
\def\mK{{\bm{K}}}
\def\mL{{\bm{L}}}
\def\mM{{\bm{M}}}
\def\mN{{\bm{N}}}
\def\mO{{\bm{O}}}
\def\mP{{\bm{P}}}
\def\mQ{{\bm{Q}}}
\def\mR{{\bm{R}}}
\def\mS{{\bm{S}}}
\def\mT{{\bm{T}}}
\def\mU{{\bm{U}}}
\def\mV{{\bm{V}}}
\def\mW{{\bm{W}}}
\def\mX{{\bm{X}}}
\def\mY{{\bm{Y}}}
\def\mZ{{\bm{Z}}}
\def\mBeta{{\bm{\beta}}}
\def\mPhi{{\bm{\Phi}}}
\def\mLambda{{\bm{\Lambda}}}
\def\mSigma{{\bm{\Sigma}}}
% Tensor
\DeclareMathAlphabet{\mathsfit}{\encodingdefault}{\sfdefault}{m}{sl}
\SetMathAlphabet{\mathsfit}{bold}{\encodingdefault}{\sfdefault}{bx}{n}
\newcommand{\tens}[1]{\bm{\mathsfit{#1}}}
\def\tA{{\tens{A}}}
\def\tB{{\tens{B}}}
\def\tC{{\tens{C}}}
\def\tD{{\tens{D}}}
\def\tE{{\tens{E}}}
\def\tF{{\tens{F}}}
\def\tG{{\tens{G}}}
\def\tH{{\tens{H}}}
\def\tI{{\tens{I}}}
\def\tJ{{\tens{J}}}
\def\tK{{\tens{K}}}
\def\tL{{\tens{L}}}
\def\tM{{\tens{M}}}
\def\tN{{\tens{N}}}
\def\tO{{\tens{O}}}
\def\tP{{\tens{P}}}
\def\tQ{{\tens{Q}}}
\def\tR{{\tens{R}}}
\def\tS{{\tens{S}}}
\def\tT{{\tens{T}}}
\def\tU{{\tens{U}}}
\def\tV{{\tens{V}}}
\def\tW{{\tens{W}}}
\def\tX{{\tens{X}}}
\def\tY{{\tens{Y}}}
\def\tZ{{\tens{Z}}}
% Graph
\def\gA{{\mathcal{A}}}
\def\gB{{\mathcal{B}}}
\def\gC{{\mathcal{C}}}
\def\gD{{\mathcal{D}}}
\def\gE{{\mathcal{E}}}
\def\gF{{\mathcal{F}}}
\def\gG{{\mathcal{G}}}
\def\gH{{\mathcal{H}}}
\def\gI{{\mathcal{I}}}
\def\gJ{{\mathcal{J}}}
\def\gK{{\mathcal{K}}}
\def\gL{{\mathcal{L}}}
\def\gM{{\mathcal{M}}}
\def\gN{{\mathcal{N}}}
\def\gO{{\mathcal{O}}}
\def\gP{{\mathcal{P}}}
\def\gQ{{\mathcal{Q}}}
\def\gR{{\mathcal{R}}}
\def\gS{{\mathcal{S}}}
\def\gT{{\mathcal{T}}}
\def\gU{{\mathcal{U}}}
\def\gV{{\mathcal{V}}}
\def\gW{{\mathcal{W}}}
\def\gX{{\mathcal{X}}}
\def\gY{{\mathcal{Y}}}
\def\gZ{{\mathcal{Z}}}
% Sets
\def\sA{{\mathbb{A}}}
\def\sB{{\mathbb{B}}}
\def\sC{{\mathbb{C}}}
\def\sD{{\mathbb{D}}}
% Don't use a set called E, because this would be the same as our symbol
% for expectation.
\def\sF{{\mathbb{F}}}
\def\sG{{\mathbb{G}}}
\def\sH{{\mathbb{H}}}
\def\sI{{\mathbb{I}}}
\def\sJ{{\mathbb{J}}}
\def\sK{{\mathbb{K}}}
\def\sL{{\mathbb{L}}}
\def\sM{{\mathbb{M}}}
\def\sN{{\mathbb{N}}}
\def\sO{{\mathbb{O}}}
\def\sP{{\mathbb{P}}}
\def\sQ{{\mathbb{Q}}}
\def\sR{{\mathbb{R}}}
\def\sS{{\mathbb{S}}}
\def\sT{{\mathbb{T}}}
\def\sU{{\mathbb{U}}}
\def\sV{{\mathbb{V}}}
\def\sW{{\mathbb{W}}}
\def\sX{{\mathbb{X}}}
\def\sY{{\mathbb{Y}}}
\def\sZ{{\mathbb{Z}}}
% Entries of a matrix
\def\emLambda{{\Lambda}}
\def\emA{{A}}
\def\emB{{B}}
\def\emC{{C}}
\def\emD{{D}}
\def\emE{{E}}
\def\emF{{F}}
\def\emG{{G}}
\def\emH{{H}}
\def\emI{{I}}
\def\emJ{{J}}
\def\emK{{K}}
\def\emL{{L}}
\def\emM{{M}}
\def\emN{{N}}
\def\emO{{O}}
\def\emP{{P}}
\def\emQ{{Q}}
\def\emR{{R}}
\def\emS{{S}}
\def\emT{{T}}
\def\emU{{U}}
\def\emV{{V}}
\def\emW{{W}}
\def\emX{{X}}
\def\emY{{Y}}
\def\emZ{{Z}}
\def\emSigma{{\Sigma}}
% entries of a tensor
% Same font as tensor, without \bm wrapper
\newcommand{\etens}[1]{\mathsfit{#1}}
\def\etLambda{{\etens{\Lambda}}}
\def\etA{{\etens{A}}}
\def\etB{{\etens{B}}}
\def\etC{{\etens{C}}}
\def\etD{{\etens{D}}}
\def\etE{{\etens{E}}}
\def\etF{{\etens{F}}}
\def\etG{{\etens{G}}}
\def\etH{{\etens{H}}}
\def\etI{{\etens{I}}}
\def\etJ{{\etens{J}}}
\def\etK{{\etens{K}}}
\def\etL{{\etens{L}}}
\def\etM{{\etens{M}}}
\def\etN{{\etens{N}}}
\def\etO{{\etens{O}}}
\def\etP{{\etens{P}}}
\def\etQ{{\etens{Q}}}
\def\etR{{\etens{R}}}
\def\etS{{\etens{S}}}
\def\etT{{\etens{T}}}
\def\etU{{\etens{U}}}
\def\etV{{\etens{V}}}
\def\etW{{\etens{W}}}
\def\etX{{\etens{X}}}
\def\etY{{\etens{Y}}}
\def\etZ{{\etens{Z}}}
% The true underlying data generating distribution
\newcommand{\pdata}{p_{\rm{data}}}
% The empirical distribution defined by the training set
\newcommand{\ptrain}{\hat{p}_{\rm{data}}}
\newcommand{\Ptrain}{\hat{P}_{\rm{data}}}
% The model distribution
\newcommand{\pmodel}{p_{\rm{model}}}
\newcommand{\Pmodel}{P_{\rm{model}}}
\newcommand{\ptildemodel}{\tilde{p}_{\rm{model}}}
% Stochastic autoencoder distributions
\newcommand{\pencode}{p_{\rm{encoder}}}
\newcommand{\pdecode}{p_{\rm{decoder}}}
\newcommand{\precons}{p_{\rm{reconstruct}}}
\newcommand{\laplace}{\mathrm{Laplace}} % Laplace distribution
\newcommand{\E}{\mathbb{E}}
\newcommand{\Ls}{\mathcal{L}}
\newcommand{\R}{\mathbb{R}}
\newcommand{\emp}{\tilde{p}}
\newcommand{\lr}{\alpha}
\newcommand{\reg}{\lambda}
\newcommand{\rect}{\mathrm{rectifier}}
\newcommand{\softmax}{\mathrm{softmax}}
\newcommand{\sigmoid}{\sigma}
\newcommand{\softplus}{\zeta}
\newcommand{\KL}{D_{\mathrm{KL}}}
\newcommand{\Var}{\mathrm{Var}}
\newcommand{\standarderror}{\mathrm{SE}}
\newcommand{\Cov}{\mathrm{Cov}}
% Wolfram Mathworld says $L^2$ is for function spaces and $\ell^2$ is for vectors
% But then they seem to use $L^2$ for vectors throughout the site, and so does
% wikipedia.
\newcommand{\normlzero}{L^0}
\newcommand{\normlone}{L^1}
\newcommand{\normltwo}{L^2}
\newcommand{\normlp}{L^p}
\newcommand{\normmax}{L^\infty}
\newcommand{\parents}{Pa} % See usage in notation.tex. Chosen to match Daphne's book.
\DeclareMathOperator*{\argmax}{arg\,max}
\DeclareMathOperator*{\argmin}{arg\,min}
\DeclareMathOperator{\sign}{sign}
\DeclareMathOperator{\Tr}{Tr}
\let\ab\allowbreak
\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{microtype} % microtypography
\usepackage{xcolor} % colors
\usepackage{multirow}
\usepackage{graphicx}
\usepackage{amsmath}
\usepackage{amssymb,amsthm,mathtools}
\usepackage{nicefrac}
\usepackage{mathrsfs}
%\usepackage{amsthm}
\usepackage{xspace}
\usepackage[inline]{enumitem}
\usepackage{comment}
\usepackage{thm-restate}
\usepackage{graphicx}
% \usepackage[colorlinks=true,linkcolor=violet,citecolor=blue]{hyperref}
\usepackage{hyperref}
\hypersetup{colorlinks=true,linkcolor=violet,citecolor=blue}
\usepackage[capitalize,nameinlink]{cleveref}
\usepackage{adjustbox}
\usepackage{bm}
\usepackage{tikz}
\usepackage{ifthen}
%\usetikzlibrary{patterns}
\usetikzlibrary {arrows.meta}
\usetikzlibrary{backgrounds}
\usetikzlibrary{arrows,shapes}
\usetikzlibrary{tikzmark}
\usetikzlibrary{calc}
%\usepackage{fontawesome5}%include some useful icons
\usepackage{utfsym}
\usepackage{natbib}
%\usepackage[prologue,dvipsnames,table,usenames]{xcolor}
\usepackage{bussproofs}
\usepackage{latexsym} % More elegant symbol of eventually and always
\usepackage{tcolorbox}
\usepackage{algorithm}
\usepackage{algorithmic}
% \usepackage{algpseudocode}
\usepackage{booktabs}
\usepackage{caption}
\usepackage{wrapfig}
\usepackage{makecell}
\usepackage{multirow}
\usepackage{bbding}
\usepackage{subfigure}
\usepackage{subcaption}
\usepackage[normalem]{ulem}
% \usepackage[dvipsnames]{xcolor}
% \usepackage{hyperref}
\usepackage{url}
\usepackage{algorithm}
\usepackage{algorithmic}
\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}
% \newtheorem{proof}{Proof}
% \newtheorem{proof}[proof]{proof}
\newtheorem{subproblem}{Problem~1.\!\!}{\itshape}{}
\Crefname{problem}{Problem}{Problem}
\Crefname{subproblem}{Problem~1.\!\!}{Problem~1.\!\!}
%\Crefname{align}{Eq.}{Eq.}
\crefname{lemma}{Lem.}{Lem.}
\crefname{example}{Exmp.}{Exmp.}
\crefname{section}{Sec.}{Sec.}
\Crefname{appendix}{Appendix}{Appendix}
%\crefname{appendix}{App.}{App.}
\crefname{definition}{Def.}{Def.}
\crefname{theorem}{Thm.}{Thm.}
\crefname{corollary}{Cor.}{Cor.}
\crefname{algorithm}{Alg.}{Alg.}
\crefname{proof}{Prf.}{Prf.}
\newcommand{\todo}[1]{{{\color{red} (todo: #1)}}}
\newcommand{\chao}[1]{{{\color{brown} (Chao: #1)}}}
\newcommand{\qingyuan}[1]{{{\color{green} (qy: #1)}}}
\newcommand{\simon}[1]{{{\color{cyan} (simon: #1)}}}
\newcommand{\yixuan}[1]{{{\color{orange} (yixuan: #1)}}}
\newcommand\boldblue[1]{\textcolor{blue}{\textbf{#1}}}
\section{Introduction}
Reinforcement learning (RL) has achieved considerable success across various domains, including board game~\citep{schrittwieser2020mastering}, MOBA game~\citep{berner2019dota}, time-delayed system~\citep{wuboosting,wu2024variational}, and cyber-physical systems~\citep{wang2023empowering,wang2023joint,wang2023enforcing,zhan2023state}. Despite these advances, RL highly depends on the quality of reward function design which demands expertise, intensive labour, and a great amount of time~\citep{russell1998learning}. To address this, imitation learning (IL) methods, such as Behavior Cloning (BC)~\citep{torabi2018behavioral} and Inverse Reinforcement Learning (IRL)~\citep{arora2021survey}, leverage human or expert demonstrations to bypass the need for explicit reward functions. These methods aim to learn from the demonstrations to eventually match the distribution of expert behavior, and have shown great promise in applications like autonomous driving \citep{codevilla2018end,sun2018fast}, legged locomotion \citep{peng2020learning,ratliff2009learning}, and planning tasks \citep{choudhury2018data,yin2022planning}.
The notable approaches within IRL are Adversarial Imitation Learning~(AIL) methods that build upon maximum entropy framework~\citep{ziebart2008maximum}. These adversarial methods frame imitation learning as a maximum likelihood estimation problem on trajectory distributions, converting the distribution into a Boltzmann distribution parameterized by rewards under \textbf{deterministic} environment settings~\citep{wu2024variational}. This closely mirrors the distribution approximation found in generative models~\citep{finn2016connection,swamy2021moments}. Thus, model-free AIL approaches often follow generative model structures, such as GANs \citep{ho2016generative,fu2017learning} or diffusion models \citep{reuss2023goal}, and require extensive sampling for distribution matching gradient updates \citep{orsini2021matters}~conducting. Model-based IRL and AIL frameworks have also emerged, where model-based framework is designed to provide estimation for gradient and planning, leading to innovative combinations such as gradient-based IRL with model predictive control (MPC) \citep{das2021model} and end-to-end differentiable IRL frameworks for complex robotics tasks \citep{baram2016model, baram2017end, sun2021adversarial, rafailov2021visual}. However, these approaches primarily address deterministic settings and struggle when applied to stochastic environments.
The only learning "deterministic" reward techniques among the existing AIL methods face significant performance degeneration in stochastic environments, leading to risk-seeking behavior and increased data requirements \citep{ziebart2010modeling}. For example, an agent trained under the deterministic Markov Decision Process (MDP) might aim to imitate expert behavior by seeking high rewards, yet fail to account for the low probability of some transitions in stochastic MDP settings. This happens because, in stochastic environments, the assumption that trajectory distributions are aligned with a Boltzmann distribution solely parameterized by deterministic rewards no longer holds. Instead, the dynamics information must also be incorporated into the formulation. There are two possible solutions. One is massive sampling to cover all possible outcomes, which is computationally expensive in large state action spaces~\citep{devlin2011theoretical,gupta2022unpacking}. The other is estimating the dynamics information and integrating it into the reward design, making the reward "stochastic". Traditional reward design is usually based on state only $R(s_t)$~\citep{torabi2018generative}, state-action pair $R(s_t, a_{t})$~\citep{blonde2019sample}, or transition tuple $R(s_t, a_t, s_{t+1})$~\citep{fu2017learning}, where the information inputted can be thought as a deterministic sample piece under the stochastic setting. The challenge in stochastic environments calls for a different perspective of rewards -- stochastic rewards absorbing the transition information.
Inspired by this idea, we propose a novel model-enhanced AIRL framework with a specifically tailored model-enhanced reward design approach to elevate performance in stochastic environments while remaining competitive in deterministic settings. In contrast to existing methods, our approach leverages the predictive power of the estimated transition model to shape rewards, represented as $\hat{R}(s_t, a_t, \hat{\mathcal{T}})$. This also enables us to generate synthetic trajectories to help guide policy optimization and reduce dependency on costly real-world interactions. As part of our analysis, we provide a theoretical guarantee on the optimal behavior for policies induced by our reward shaping and derive a bound on the performance gap with respect to the transition model errors. Empirically, we demonstrate that this integration significantly enhances sample complexity and policy performance in both settings, providing a comprehensive solution to the limitations of existing AIL methods in uncertain environments.
\textbf{Contributions of this work} include:
\begin{itemize}
\item A novel reward shaping method with model estimation under the stochastic MDP setting, which provides the optimal policy invariance guarantee.
\item A novel model-based adversarial IRL framework that seamlessly incorporates transition model training, adversarial reward learning with model estimation and forward model-based RL process, enhancing performance in stochastic environments, and sample efficiency.
\item Theoretical analysis on reward learning with model estimation under the adversarial framework and performance difference under transition model learning errors.
\item Empirical validation that demonstrates our approach's performance improvements in stochastic environments and comparable performance in deterministic environments as well as significant improvement in sample efficiency.
\end{itemize}
In \cref{section::related_work}, we introduce related works in AIL and reward shaping. In \cref{section::preliminaries}, we provide the necessary preliminaries for MDP and IRL. In \cref{section::reward_shaping}, we present our model-enhanced reward shaping method and corresponding theoretical guarantee. In \cref{section::algorithm}, we present our Model-Enhanced AIRL framework design together with theoretical analysis on reward error bound and performance difference bound. In \cref{section::experiment}, we show the experimental results in \texttt{Mujoco} for various benchmarks. \cref{section::conclusion} concludes the paper.
% \todo{Content of each section}
% This proposed framework not only addresses the sample inefficiency of adversarial IRL but also extends its applicability to more complex, uncertain environments, paving the way for more robust and efficient RL solutions in real-world applications.
\section{Related Works}
\label{section::related_work}
% Behavior Cloning (BC)\citep{torabi2018behavioral} and Inverse Reinforcement Learning (IRL)\citep{arora2021survey} have been two main directions within the research area of imitation learning\citep{hussein2017imitation}, which has already proved efficacy in real world applications such as driving\citep{codevilla2018end, sun2018fast}, legged locomotion\citep{peng2020learning, ratliff2009learning}, and various planning problems\citep{choudhury2018data, yin2022planning}. IRL, aiming to recover reward function from demonstrations, is inherently more generalizable comparing to BC-like approaches which directly observe and imitate the policy\citep{russell1998learning}. However, due to the ill-posed nature of the IRL problem, this kind of problem is general hard to solve. There are several core approaches for IRL problem. First, margin optimization approaches\citep{ng2000algorithms, abbeel2004apprenticeship, ratliff2006maximum} try to learn the reward function that better explain the expert demonstration than other policies by a margin. Secondly, the Bayesian update methods are introduced with different prior assumptions on the reward functions, such as Boltzmann distribution\citep{ramachandran2007bayesian, choi2011map, chan2021scalable}, Gaussian Process\citep{levine2011nonlinear}, etc. There are also some classical statistical learning approaches such as multi-class classification\citep{klein2012inverse,brown2019extrapolating}, regression tree\citep{levine2010feature}, etc.
\paragraph{Adversarial Imitation Learning.}
Margin optimization based IRL methods \citep{ng2000algorithms, abbeel2004apprenticeship, ratliff2006maximum} aim to learn reward functions that explain expert behavior better than other policies by a margin. Bayesian approaches were introduced with different prior assumptions on reward distributions, such as Boltzmann distributions \citep{ramachandran2007bayesian, choi2011map, chan2021scalable} or Gaussian Processes \citep{levine2011nonlinear}. Other statistical learning methods include multi-class classification \citep{klein2012inverse, brown2019extrapolating} and regression trees \citep{levine2010feature}. The entropy optimization approach has seen significant development. To avoid biases from maximum margin methods, the maximum entropy principle \citep{shore1980axiomatic} is used to infer distributions over trajectories parameterized by reward weights. \cite{ziebart2008maximum,ziebart2010modeling} proposed a Lagrangian dual framework to cast the reward learning into a maximum likelihood problem with linear-weighted feature-based reward representation. \citet{wulfmeier2015maximum} extended the framework to nonlinear reward representations, and \citet{finn2016guided} combines importance sampling techniques to enable model-free estimation. Inspired by GANs, adversarial methods were introduced for policy and reward learning in IRL~\citep{ho2016generative, fu2017learning, torabi2018generative}. However, these methods typically work with Maximum Entropy (ME) formulation yet suffer from sample inefficiency. Although there have been efforts to combine adversarial methods with off-policy RL agents to improve sample efficiency \citep{kostrikov2018discriminator, blonde2019sample, blonde2022lipschitzness}, few extend it to the model-based setting which might further the improvement, and none of these approaches addresses the rewards learning in stochastic MDP settings.
% Among them, entropy optimization approach has been well-developed. To avoid bias introduced by maximum margin approaches, maximum entropy principle\citep{shore1980axiomatic} can be used to infer the distribution over trajectories, which is parameterized by the reward function weights. \cite{ziebart2008maximum,ziebart2010modeling} proposed a Lagrangian dual framework to cast the reward learning into a maximum likelihood problem with linear-weighted feature-based reward representation, and \cite{wulfmeier2015maximum,finn2016guided} extended the framework to nonlinear reward representations with importance sampling techniques. \cite{finn2016connection} pointed out the connection between Maximum Entropy IRL(MEIRL) and generative model training. Inspired by GANs, \cite{ho2016generative} and \cite{fu2017learning} proposed the model free adversarial structure for policy learning and reward learning respectively. However, both methods are based on deterministic MDP assumptions and suffer from sample inefficiency. There are several attempts to integrate Adversarial methods with off-policy RL to mitigate sample inefficiency\citep{kostrikov2018discriminator}. However, still numerous interactions with environment are needed to learn the reward function.
\paragraph{Rewards Shaping.}
Reward shaping~\citep{dorigo1994robot,randlov1998learning} is a technique that enhances the original reward signal by adding additional domain information, making it easier for the agent to learn optimal behavior. This can be defines as $\hat{R}=R+F$, where $F$ is the shaping function and $\hat{R}$ is the shaped reward function. Potential-based reward shaping (PBRS)~\citep{ng2000algorithms} builds the potential function on states, $F(s,a,s')=\phi(s')-\phi(s)$, while ensuring the policy invariance property, which refers to inducing the same optimal behavior under different rewards $R$ and $\hat{R}$. Nonetheless, there also exist other variants on the inputs of the potential functions such as state-action~\citep{wiewiora2003principled}, state-time~\citep{devlin2012dynamic}, and value function~\citep{harutyunyan2015expressing} as potential function input. There are also some latest attempts of reward shaping without utilization of domain knowledge potential function to solve exploration under sparse rewards~\citep{hu2020learning,devidze2022exploration,gupta2022unpacking,skalse2023invariance}.
% \todo{state-only, state-state, state-action-state reward formulation discussion missing}~\citet{ng2000algorithms} has proved hard optimal policy reward shaping with the form $\hat{R}(s_t,a_t,s_{t+1})=\lambda\left(R(s_t,a_t,s_{t+1})+\gamma\phi(s_{t+1})-\phi(s_t)\right)$
% As a promising candidate for solving real world complex control and planning problems, model-based RL has also proved its additional advantages in sample efficiency and generalization\citep{janner2019trust,yu2020mopo}. Besides, there is an extensive attempts to combine receding planning module with learned dynamic\citep{hansen2022temporal, sikchi2022learning}. Introduction of world model seamlessly combine the transition dynamic modelling and RL techniques, and also enables vision-based control/planning tasks\citep{hafner2019dream, zhan2023state}. There has been some successful attempts to combine IRL techniques with model-based approaches. \cite{das2021model} presents a gradient based IRL approach with MPC for linear-weighted feature rewards, where dynamic model is used for approximated/exact gradient update and provide prediction for planning. Similarly, \cite{baram2016model, baram2017end, sun2021adversarial, rafailov2021visual} introduces an end-to-end differentiable AIL framework for nonlinear rewards, where dynamic model serves a similar role. However, none of the aforementioned methods try to improve the reward learning under the stochastic setting and provide corresponding theoretical analysis.
\paragraph{MBIRL.}
Model-Based RL (MBRL) has emerged as a promising direction for improving sample efficiency and generalization \citep{janner2019trust, yu2020mopo}. MBRL combines various learned dynamics neural network structures with planning \citep{hansen2022temporal, sikchi2022learning}.This framework has been successfully extended to vision-based control tasks \citep{hafner2019dream, zhan2023state}. Integrating IRL with MBRL has also shown success. For example, \citet{das2021model} and \citet{herman2016inverse} presented a gradient-based IRL approach using different policy optimization methods with dynamic models for linear-weighted features reward learning. In \citet{das2021model}, the dynamic model is used to pass forward/backward the gradient in order to update the IRL and policy optimization modules. Similarly, end-to-end differentiable adversarial IRL frameworks to various state spaces have also been explored \citep{baram2016model, baram2017end, sun2021adversarial, rafailov2021visual}, where dynamic model serves a similar role. Despite these advancements, existing methods rarely address the specific challenges posed by stochastic environments, which limit reward learning performance. To our knowledge, this is the first study that provide a theoretical analysis on the performance difference with learned dynamic model for the adversarial IRL problem under stochastic MDP.
\section{Preliminaries}
\label{section::preliminaries}
% \textbf{MDP}
\paragraph{MDP.}
RL is usually formulated as a Markov Decision Process (MDP) $\mathcal{M}$~\citep{puterman2014markov} denoted as a tuple $\langle\mathcal{S},\mathcal{A},\mathcal{T},\gamma, R, \rho_0\rangle$. $\rho_0$ is the initial distribution of the state. $s\in\mathcal{S},a\in\mathcal{A}$ stands for the state and action space respectively. $\mathcal{T}$ stands for the transition dynamic such that $\mathcal{T}:\mathcal{S}\times\mathcal{A}\times\mathcal{S}\rightarrow[0,1]$. $\gamma\in(0,1)$ is the discounted factor, $R$ stands for reward function such that $R:\mathcal{S}\times\mathcal{A}\rightarrow\mathbb{R}$ and $\|R\|_{\infty}\leq R_{\max}$. The discounted visitation distribution of trajectory $\tau$ with policy $\pi$ is given by:
\begin{equation}
p(\tau) = \rho_0\prod_{t=0}^{T-1}\gamma^t\mathcal{T}(s_{t+1}|s_t,a_t)\pi(a_t|s_t).
\end{equation}
The objective function of RL is $\max\mathbb{E}_{\tau\sim p(\tau)}\left[\sum_{t=0}^{T}\gamma^t R(\tau)-H(\pi)\right]$, where $H$ is the log likelihood of the policy. We introduce Soft Value Iteration for bellmen update~\citep{haarnoja2018soft}, where $Q^{soft}$ and $V^{soft}$ denotes the soft Q function and Value function respectively:
\begin{align}\label{equation::soft_VI}
V^{soft}(s_t) &= log\sum_{a_t\in\mathcal{A}}\exp{Q^{soft}(s_t, a_t)}da_t,\\
Q^{soft}(s_t,a_t) &= R(s_t, a_t) + \gamma\mathbb{E}_{\mathcal{T}}\left[V^{soft}(s_{t+1})|s_t,a_t\right],\\
\pi(a_t|s_t) &= \exp{(Q^{soft}(s_t,a_t) - V^{soft}(s_t))},
\end{align}
where the soft Advantage function is defined as $A^{soft}(s_t,a_t)=Q^{soft}(s_t,a_t)-V^{soft}(s_t)$.
\paragraph{Inverse RL.}
In IRL setting, we usually consider the MDP without reward as $\mathcal{M}'$ where $R$ is also unknown. We denote the data buffer $\mathcal{D}_{exp}$ which collects trajectories from an expert policy $\pi^E$. We consider a reward function $R_\theta:\mathcal{S}\times\mathcal{A}\rightarrow\mathbb{R}$, where $\theta$ is the reward parameter. An IRL problem can be defined as a pair $\mathcal{B}=(\mathcal{M}',\pi^E)$. A reward function $R_\theta$ is feasible for $\mathcal{B}$ if $\pi^E$ is an optimal policy for the MDP $\mathcal{M}'\cup R_\theta$, and we denote the set of feasible rewards as $\mathcal{R}_{\mathcal{B}}$.
Using maximize likelihood estimation framework, we can formulate the IRL as the following maximum causal entropy problem:
\begin{equation}
\arg\max_{\theta} \mathbb{E}_{\tau\sim\mathcal{D}_{exp}}\log p_\theta(\tau),
\end{equation}
where $Q_{R_\theta}^{soft}$ and $V^{soft}_{R_\theta}$ are based on $R_\theta$ and $p_\theta(\tau)\propto\rho_0\prod_{t=0}^{T-1}\mathcal{T}(s_{t+1}|s_t,a_t)\exp(Q^{soft}_{\hat{R}_\theta}(s_t,a_t)-V^{soft}_{R_\theta}(s_t))$~\citep{ziebart2010modeling}. Under deterministic MDP, the above problem can be simplified as maximum entropy, where $p_\theta(\tau)\propto\frac{1}{Z_\theta}\exp{\sum_{t=0}^{T-1} R_\theta(s_t,a_t)}$ and $Z_\theta$ is the temperature factor of the Boltzmann Distribution~\citep{ziebart2008maximum}.
% Under derteminism assumption of $\mathcal{T}$, the gradient of loss function of above optimization is as follow:
% \begin{equation}\label{equaiton::maxent_loss}
% \nabla_{\theta}\mathcal{L} = \mathbb{E}_{\tau\sim\mathcal{D}_{exp}}\sum_{t=1}^{T-1}\left[\nabla_{\theta} R_{\theta}(s_t, a_t)\right] - \mathbb{E}_{\pi_\theta}\sum_{t=1}^{T-1}\left[\nabla_{\theta} R_{\theta}(s_t,a_t)\right]
% \end{equation}
% The first part of the above loss function is easy to compute, which is the gradient of learned reward function on expert trajectory with respect to the parameter $\theta$. However, the latter part is hard to compute, which requires to solve either a inner loop MDP problem which requires calculation of forward/backward messages~\citep{levine2018reinforcement} or conduct the Value Iteration, which easily takes hours to converge, such as DQN~\citep{van2016deep}.
% Recent methods~\citep{fu2017learning,ho2016generative} proposing a smart reformulation of above problem by framing above optimization into the GAN formulation~\citep{goodfellow2014generative}. Specifically, we need to optimize the discriminator $D$ and generator $G$.
% $$
% D_G ^*(x)=\frac{p_{data}(x)}{p_{data}(x)+p_g(x)}
% $$
% Under the IRL setting, the generator is the policy, and the generator has the following format.
% $$
% D_{\phi,\varphi}(s,a,s')=\frac{\exp\{f_{\phi,\varphi}(s,a,s')\}}{\exp\{f_{\phi,\varphi}(s,a,s')\}+\pi_\theta(a|s)}
% $$
% where $f_{\phi,\varphi}(s,a,s')=g_\phi(s,a)+\gamma h_\varphi(s')-h_\varphi(s)$, $g_\phi$ is a reward approximator, and $h_\varphi$ is a shaping term. The reward function~\citep{fu2017learning} is defined as
% $$
% R_{\phi,\varphi}(s,a,s')=\log D_{\phi,\varphi}(s,a,s')-\log(1-D_{\phi,\varphi}(s,a,s'))
% $$
% The core idea is to train the discriminator $D_{\phi,\varphi}$ to distinguish expert policy data from trained policy samples and to train the policy to fool the trained discriminator. However, the above methods are settled for the deterministic environments. Under stochastic setting, these Maximum Entropy based method might lead to risk-seeking behavior~\citep{gleave2022primer,levine2018reinforcement}.
\section{Model Estimation in reward shaping}
\label{section::reward_shaping}
% % \subsection{Reward Shaping for Stochastic MDP}
% $$
% D_{\theta}(s,a)=\frac{\exp\{f_{\theta}(s,a)\}}{\exp\{f_{\theta}(s,a)\}+\pi(a|s)}
% $$
% where $f_{\theta,\phi}$ is restricted to a reward approximator $g_\theta$ and a shaping term $h_\phi$ as
% \begin{equation}
% f_{\theta,\phi}(s,a) = g_\theta(s,a)+\gamma\mathbb{E}_{\mathcal{T}}\left[h_\phi(s')|s,a\right] - h_\phi(s)
% \end{equation}
% We can interpret $f^*(s,a)$ as the advantage under stochastic dynamics
% $$
% f^*(s,a) = \underbrace{R^*(s,a)+\gamma\mathbb{E}_{\mathcal{T}} \left[V^*(s')|s,a\right]}_{Q^*(s,a)} - V^*(s) = A^*(s,a)
% $$
% Then, we can learn a model $\hat{\mathcal{T}}$ here to substitute the original provided transition model.
% where $f_{\theta}(s,a)=g_\theta(s,a)+\gamma \mathbb{E}_{\hat{s'\sim\mathcal{T}}}[h_\theta(s')|s,a]-h_\theta(s)$
% state-only reward ignores influence of action, leading to
% state-action pair reward does not consider the outcome
% transition-tuple based reward works needs large data, thus not efficient.
% \todo{draft revise afterwards, math symbols not align yet, double check later}
% We first discuss motivations of defining rewards on state-action pair and involving transition dynamic in reward shaping to better represent the reward in stochastic MDP settings.
\begin{table*}[h]
\caption{We summarize the different reward formulations and their dynamic properties in this table. Components refer to the input pair that the reward functions take. Reward Shaping indicates whether there is the extra physical potential information involved where X means no reward shaping used. Dynamics information shows whether transitions are involved in the reward function.}
\scalebox{0.8}{
\begin{tabular}{c|ccc}
\hline
Methods
& Components
& Reward Shaping
% & Reward for policy update
& Dynamics Information
\\ \hline
AIRL~\citep{fu2017learning}
& $s_t,a_t,s_{t+1}$
& $R(s_t,a_t)+\gamma\phi(s_{t+1})-\phi(s_t)$
% & $\log(D)-\log(1-D)$
& single sample
\\
AIRL(State Only)
& $s_t$
& $R(s_t)+\textrm{constant}$
% & $\log(D)-\log(1-D)$
& X
\\
DAC~\citep{kostrikov2018discriminator}
& $s_t, a_t$
& X
% & $\log(D)-\log(1-D)$
& X
\\
SAM~\citep{blonde2019sample}
& $s_t,a_t$
& X
% & $-\log(1-D)$
& X
\\
SQIL~\citep{reddy2019sqil}
& $s_t,a_t$
& binary %Take 1 or 0 as reward
% & 1 or 0
& X
\\
GAIfO~\citep{torabi2018generative}
& $s_t, s_{t+1}$
& X
% & $-\log(D)$
& single sample
\\
% Ant-v3
% & $3370.4_{\pm 868.1}$
% &
% &
% &
% &
% \\
Ours
& $s_t, a_t, \mathcal{T}$
& $R(s_t,a_t)+\gamma\mathbb{E}_{\mathcal{T}}[\phi(s_{t+1})|s_t,a_t]-\phi(s_t)$
% & $\log(D)-\log(1-D)$
& transition model
\\
\hline
\end{tabular}
}
\label{reward_shaping_table}
\end{table*}
In this section, we illustrate the advantages of involving transition dynamics into the reward shaping, especially in stochastic MDP settings. Most of literature work has various formulations and definitions~(Table~\ref{reward_shaping_table}), but few considers transition dynamic information in the reward shaping. Defining rewards solely based on states, $R^{s}(s_t)$, offers limited utility in environments where actions are critical. Even though the state-action pair-based rewards $R^{sa}(s_t, a_t)$ can capture the missing information of the taken action, it fails to consider any future information, the successive state $s_{t+1}$.
Transition tuple-based rewards $R^{tuple}(s_t, a_t, s_{t+1})$ incorporate the dynamics information in a sampling-based way, which requires abundant data to learn the underlying relationship of two consecutive states, potentially raising the sample efficiency issue in the stochastic environment with the huge state space. To address this issue, we propose dynamics-based rewards shaping $\hat{R}(s_t,a_t,\mathcal{T})$, which explicitly infuse the dynamics information $\mathcal{T}$ on the potential function, thus significantly improving the sample-efficiency. Specifically, our rewards shaping is defined as
\begin{equation}
\label{def::transition_reward}
\hat{R}(s_t,a_t,\mathcal{T})=R(s_t,a_t)+\gamma\mathbb{E}_{\mathcal{T}}\left[\phi(s_{t+1})\vert s_t,a_t\right]-\phi(s_t),
\end{equation}
where $\phi$ is a state-only potential function, $\mathcal{T}$ is the dynamics. Another insight of the above reward shaping is to resemble the advantage function with the soft value function as the potential function, which we will elaborate in \cref{section::adversarial_reward}. With the given reward shaping $\hat{R}$, it is crucial to show that it induces the same optimal behaviour as the ground-true reward $R$. We formally define this policy invariance property as follows. \begin{definition}\citep{memarian2021self}
\label{def::soft_optimal_policy}
Let $R$ and $\hat{R}$ be two reward functions. We say they induce the same soft optimal policy under transition dynamics $\mathcal{T}$ if, for all states $s\in\mathcal{S}$ and actions $a\in\mathcal{A}$:
\begin{equation}
A_{R}^{soft}(s_t,a_t) = A_{\hat{R}}^{soft}(s_t,a_t).
\end{equation}
\end{definition}
With the above definition, we can transfer the proof of policy invariant property of our designed reward shaping (\cref{def::transition_reward}) to showing the equivalence of soft advantage functions, which is proved in the following theorem. The detailed proof can be found in \ref{proof::soft_optimal_policy}.
\begin{theorem}[Reward Invariance]
\label{thm::soft_optimal_policy}
Let $R$ and $\hat{R}$ be two reward functions. $R$ and $\hat{R}$ induce the same soft optimal policy under all transition dynamics $\mathcal{T}$ if $\hat{R}(s_t,a_t,\mathcal{T})=R(s_t,a_t)+\gamma\mathbb{E}_{\mathcal{T}}[\phi(s_{t+1})\vert s_t,a_t] - \phi(s_t)$ for some potential-shaping function $\phi:\mathcal{S}\rightarrow\mathbb{R}$.
\end{theorem}
% \qingyuan{a sentence illustrate the invariant property, e.g., the reward shaping method can induce the same optimal policy as the original method...}
\cref{thm::soft_optimal_policy} implies that the optimal policy induced from our model-enhanced rewards shaping $\hat{R}$ (\cref{def::transition_reward}) is equivalent to the optimal policy trained by the ground-truth reward function $R$ under the soft Value Iteration fashion.
\section{Model Enhanced Adversarial IRL}
\label{section::algorithm}
In this section, we first elaborate on the adversarial formulation of our reward shaping (\cref{equation::discriminator}) and present the theoretical insight (\cref{prop::gradient_align}) of the equivalence between cross-entropy training loss of adversarial reward shaping formulation and maximum log-likelihood loss of original maximum causal entropy IRL problem.
Then, in the \cref{section::algorithm_framework}, we showcase our practical algorithm framework with trajectory generation and transition model learning in the loop, as shown in ~\cref{fig::framework}.
% Together, we introduce implementation details to mitigate distribution shift and to ensure stability.
Furthermore, we theoretically investigate the reward function bound (\cref{thm:reward_func_Error_Bound}) and performance difference bound (\cref{thm:value_error_bound}) under the transition model learning error.
\begin{figure}[h]
\centering
\includegraphics[width=1.05\textwidth]{figures/MEAIRL_overall_figuare.pdf}
% \caption{Framework overview of Model-Enhanced Adversarial IRL. Different color arrows stand for different sample flows. \textcolor{Purple}{Purple} stands for real environmental interaction samples, \textcolor{Rhodamine}{pink} stands for synthetic samples generated from learned transition model, and \textcolor{Cyan}{blue} stands for mixed of both.}
\caption{Framework overview of Model-Enhanced Adversarial IRL. Different color arrows stand for different sample flows. Purple stands for real environmental interaction samples, pink stands for synthetic samples generated from learned transition model, and blue stands for mixed of both.}
\label{fig::framework}
\end{figure}
\subsection{Adversarial Formulation of Reward Shaping}
\label{section::adversarial_reward}
In order to learn a robust reward function, we apply the AIRL framework~\citep{fu2017learning}. Inspired by GANs~\citep{goodfellow2014generative}, the idea behind AIRL is to train a binary discriminator $D(s_t,a_t,s_{t+1})$ to distinguish state-action-transition samples from an expert and those generated by imitator policy. However, as mentioned above, we only take in state-action pair and transition function to define our reward function which also extends to our discriminator as follows:
\begin{equation}
\label{equation::discriminator}
D_\theta(s_t,a_t,\mathcal{T})=\frac{\exp\{f_{\theta}(s_t,a_t,\mathcal{T})\}}{\exp\{f_{\theta}(s_t,a_t,\mathcal{T})\}+\pi(a_t\vert s_t)},
\end{equation}
where $f_\theta(s_t,a_t,\mathcal{T})=R_\theta(s_t,a_t)+\gamma\mathbb{E}_{\mathcal{T}}\left[\phi_\theta(s_{t+1})|s_t,a_t\right]-\phi_\theta(s_t)$ resembles the reward shaping defined above. The loss function for the training discriminator is defined below.
\begin{equation}
\label{equation::disc_loss}
\mathcal{L}_{disc}=-\mathbb{E}_{\mathcal{D}_{exp}}\left[\log D_{\theta}(s,a,\mathcal{T})\right]-\mathbb{E}_{\pi}\left[\log(1-D_\theta(s,a,\mathcal{T}))\right].
\end{equation}
% \todo{Connect the adversarial reward formulation with the gradient match-up in the latter part , may be we need a thm? Or a insight might be more appropriate?}
With the model-enhanced stochastic reward shaping, we can bridge this adversarial formulation with the original maximum causal entropy IRL problem. In the following proposition, we give a sketch of proof to show the link between the loss of discriminator and loss of MCE IRL, which sheds the light on how this adversarial formulation can recover the rewards of the task. Proof details can be found in Appendix \ref{proof::gradient alignment}.
\begin{proposition}
\label{prop::gradient_align}
Consider an undiscounted MDP. Suppose $f_\theta$ and $\pi$ at the current iteration are the soft-optimal advantage function and policy for reward function $R_\theta$. Minimising the cross-entropy loss of the discriminator under generator $\pi$ is equivalent to maximising the log-likelihood under MCE IRL.
\end{proposition}
With above given proposition, we can construct a direct intuition that $f^*_\theta$ should be equal to $\hat{R}_\theta$ the reward shaping we introduced early and resemble the soft advantage function. To extract rewards to represent reward used for policy optimization, we use $\log(D_\theta(s,a,\mathcal{T}))-\log(1-D_\theta(s,a,\mathcal{T}))$, which resembles the entropy-regularized reward shaping $f_\theta(s,a,\mathcal{T})-\log\pi(a|s)$. Given this entropy-regularized reward, it is straightforward to see why the optimal policy can satisfy the RL objectives, $\max\mathbb{E}_{\tau\sim p(\tau)}\left[\sum_{t=0}^{T-1}\gamma^t \hat{R}_\theta(\tau)-H(\pi)\right]$.
\subsection{Algorithm Framework}
\label{section::algorithm_framework}
\begin{algorithm*}[tb]
\caption{Model Enhanced Adversarial IRL}
\label{alg:meairl}
\begin{algorithmic}[1]
\STATE Obtain expert buffer $\mathcal{D}_{exp}$.
\STATE Initialize policy $\pi$, discriminator $D_{\theta}$, buffers $\mathcal{D}_{env}, \mathcal{D}_{gen}$, and transition model $\hat{\mathcal{T}}$.
\FOR{step $t$ in \{1, \dots, N\}}
\STATE Interact with real environments and add state-action pair to $\mathcal{D}_{env}$.
\IF{$t<\text{STARTING\_STEP}$}
\STATE Pretrain transition model $\hat{\mathcal{T}}$.
\ELSE
\STATE Sample state-action batch from $\mathcal{D}_{exp},\mathcal{D}_{env}$ respectively.
% \STATE Update discriminator with loss function.
\STATE Train $D_{\theta}$ via cross entropy loss ~\cref{equation::disc_loss} to classify expert data from samples.
\STATE Update dynamic model $\hat{\mathcal{T}}$ with MLE loss and generate $H$-steps trajectories to $\mathcal{D}_{gen}$.
\STATE Sample state-action batches from $\mathcal{D}_{env},\textrm{and }\mathcal{D}_{gen}$ with varying ratio.
\STATE Update $\pi$ with respect to $\hat{R}_{\theta}$ using Soft Actor Critic policy optimization.
\ENDIF
\ENDFOR
\end{algorithmic}
\end{algorithm*}
In this section, we present the overall framework of Model-Enhanced Adversarial IRL and illustrate how transition model training is incorporated into the learning loop. We assume the estimated transition distribution $\hat{\mathcal{T}}(\cdot|s,a)$ follows a Gaussian distribution with mean and standard deviation parameterized by the MLP, and the model is updated with standard maximum likelihood loss. At each iteration, the updated transition model is utilized for reward calculation and synthetic data generation, which is stored in the synthetic trajectory replay buffer. Unlike AIRL and GAIL, our framework operates in an off-policy manner, where samples used for both discriminator and policy update are drawn from a combination of the environmental replay buffer and the synthetic replay buffer. An overview of our framework is shown in ~\cref{fig::framework}, and detailed algorithmic steps are provided in ~\cref{alg:meairl} and Appendix~\ref{sec::implementation}.
\textbf{Sample Efficiency:} To improve sample efficiency, we leverage the estimated transition model to generate $H$-steps synthetic trajectories data alongside real interaction data, facilitating policy optimization. Given that the estimated transition model is inaccurate at the beginning, we employ a dynamic ratio between real and synthetic data to prevent the model from being misled by unlikely synthetic transitions~\citep{janner2019trust,zhan2023state}. Specifically, early-stage generated trajectories are not stored persistently, unlike real interactions which are fully stored in the off-policy environmental replay buffer. To maintain training stability, we use a synthetic replay buffer with a size that gradually increases as training progresses, ensuring a balanced inclusion of synthetic data over time. The growth rates of the data ratio and buffer size are adjusted based on the complexity of the transition model learning process and can be fine-tuned via hyperparameters.
\textbf{Distribution Shift:} To mitigate distribution shift~\citep{lee2020addressing, lin2020model} during training, we employ a strategy involving the learned transition model. Typically, during interaction, the real state $s_t$ is used as input to the actor, and the resulting action $a_t$ is applied in the environment. To incorporate the transition model, we predict a synthetic state $\hat{s}_t$ from previous $s_{t-1}$ and $a_{t-1}$. This generated $\hat{s}_t$ is then fed into the actor to produce action $\hat{a}_t$. The actions $a_t$ and $\hat{a}_t$ are mixed and applied to the environment with a certain ratio, and the resulting pairs $(s_t, a_t)$ or $(s_t, \hat{a}_t)$ are stored in the environmental replay buffer. This approach helps balance the exploration of real and model-predicted dynamics, reducing the impact of distributional discrepancies.
% \todo{some reason why this trick works, references}
%===============================================================================
\subsection{Performance Analysis}
\label{section::performance_gap}
In this section, we analyze the optimal performance bound in the presence of transition model learning errors. Our results show that as the transition model error approaches zero, the performance gap at the optimal point vanishes at the same time. The learned transition model $\hat{\mathcal{T}}$ persists in some errors compared with the ground-true transition dynamic. In this section, we investigate how this error will affect performance of our method. As a reminder, we define an IRL problem as $\mathfrak{B}=(\mathcal{M}',\pi^E)$, where $\mathcal{M}'$ is a MDP without $R$ and $\pi^E$ is an optimal expert policy. We denote $\mathcal{R}_{\mathfrak{B}}$ as the set of feasible rewards set for $\mathfrak{B}$. Since under our case $\mathcal{T}$ is approximated by $\hat{\mathcal{T}}$, we have another IRL problem defined as $\hat{\mathfrak{B}}=(\hat{\mathcal{M}'},\pi^E)$ where $\hat{\mathcal{M}}'$ has the same state, action space, discount factor, and initial distribution but an estimated transition model $\hat{\mathcal{T}}$. For notation, we use \( D_{\text{TV}} \) to denote the total variation distance, \( \Vert \cdot \Vert \) to represent the infinity norm (with \( \infty \) omitted for simplicity), \( \vert \mathcal{S} \vert \) to denote the cardinality of the state space, and \( V_{\mathcal{M}' \cup R}^{\pi^*} \) to represent the value function of policy \( \pi^* \) under the MDP \( \mathcal{M}' \) with reward \( R \), and vice versa.
\begin{assumption}[Transition Model Error]
\label{assump:transition_error}
Since transition model is trained through a supervised fashion, we can use a PAC generalization bound~\citep{shalev2014understanding} for sample error. Therefore, we assume that the total variation distance between $\mathcal{T}$ and $\hat{\mathcal{T}}$ is bounded by $\epsilon_\mathcal{T}$ through $[0,T]$:
\begin{equation}
\max_t\mathbb{E}_{s\sim\pi_{D,t}}\left[D_{TV}(\mathcal{T}(s'|s,a)\vert\hat{\mathcal{T}}(s'|s,a))\right]\leq\epsilon_\mathcal{T},
\end{equation}
which is a common assumption that adopted in literature~\citep{janner2019trust,sikchi2022learning}.
\end{assumption}
Next, with the assumed total visitation bound on transition models (~\cref{assump:transition_error}), we can bound the error on rewards learning through model-enhanced reward shaping.
\begin{theorem}[Reward Function Error Bound]
\label{thm:reward_func_Error_Bound}
Let $\mathfrak{B}=(\mathcal{M}',\pi^*)$ and $\hat{\mathfrak{B}}=(\hat{\mathcal{M}'},\pi^*)$ be two IRL problems, then for any $R^E\in\mathcal{R}_{\mathfrak{B}}$ there is a corresponding $\hat{R}^E\in\mathcal{R}_{\hat{\mathfrak{B}}}$ such that
\begin{equation}
\Vert R^E-\hat{R}^E\Vert\leq\frac{\gamma}{1-\gamma}|\mathcal{S}|\epsilon_{\mathcal{T}}R_{\max}.
\end{equation}
\end{theorem}
Proof of ~\cref{thm:reward_func_Error_Bound} can be found in Appendix \ref{proof:reward_error_bound}. With rewards bound above, we can extend the bound to the value function, which represents the performance difference brought up by estimated transition model error under RL setting.
\begin{theorem}[Performance Difference Bound]
\label{thm:value_error_bound}
% $\Vert V_{\mathcal{M}'\cup R}^{\pi^*}(s) - V_{\hat{\mathcal{M}}'\cup \hat{R}}^{\hat{\pi}^*}(s)\Vert$:
The performance difference between the optimal policies ($\pi^*$ and $\hat{\pi}^*$) in corresponding MDPs ($\mathcal{M}'\cup R$ and $\hat{\mathcal{M}}'\cup \hat{R}$) can be bounded as follows:
\begin{equation}
\Vert V_{\mathcal{M}'\cup R}^{\pi^*} - V_{\hat{\mathcal{M}}'\cup \hat{R}}^{\hat{\pi}^*} \Vert\leq
\epsilon_\mathcal{T}\left[ \frac{\gamma}{(1-\gamma)^2} R_{\max}+\frac{1+\gamma}{(1-\gamma)^2}R_{\max} |\mathcal{S}|\right].
\end{equation}
\end{theorem}
The detailed proof of ~\cref{thm:value_error_bound} is presented in Appendix \ref{proof:value_error_bound}. The above theorem highlights the relationship between the performance gap and the transition model error, also implying that a perfectly-learned transition model ($\epsilon_\mathcal{T} \rightarrow 0$) could make the performance difference negligible.
\section{Experiments}
\label{section::experiment}
% In this section, we investigate the performance and sample efficiency of our Model Enhanced Adversarial IRL framework. We try to prove our method can maintain competitive performance under deterministic setting with much better sample efficiency. Furthermore, we also demonstrate our method's superiority under the stochastic setting comparing with other methods, where our method can achieve even surpass expert performance with faster rate. All the experiments are conducted in the Mujoco Environment Benchmarks. To mimic stochastic behavior, we add a Gaussian noise with $0$ mean and $0.5$ stander deviation to the interaction steps. We mainly compare with Adversarial Imitation Learning (AIL) methods including on-policy algorithm GAIL~\citep{ho2016generative} and AIRL~\citep{fu2017learning} and off-policy method Discriminator Actor Critic (DAC)~\citep{kostrikov2018discriminator}. We use PPO~\citep{schulman2017proximal} as the policy optimization method for both GAIL and AIRL. We use Soft Actor Critic~\citep{haarnoja2018soft} as the policy optimization methods for both DAC and our method. Both PPO and SAC are referenced from \texttt{Clean RL}.
In this section, we evaluate the performance and sample efficiency of our Model-Enhanced Adversarial IRL framework.
We aim to demonstrate the superiority of our method in stochastic environments, achieving better performance and sample efficiency compared to existing approaches.
Additionally, in deterministic settings, our method not only maintains competitive performance, but also achieves better sample efficiency with baselines.
% We aim to demonstrate that, in deterministic settings, our method not only maintains competitive performance, but also achieves better sample efficiency compared to existing approaches.
% Additionally, we show the superiority of our method in stochastic environments.
All experiments are conducted on the MuJoCo benchmarks~\citep{mujoco}.
To simulate stochastic dynamics in MuJoCo, we introduce the agent-unknown Gaussian noise with a mean of 0 and a standard deviation of 0.5 to the environmental interaction steps. All the expert data are generated by an expert trained with standard SAC under deterministic or stochastic MuJoCo environments. Our experiments are designed to highlight the key advantages of our framework:
\begin{itemize}
\item \textbf{Superiority in Stochastic Environments:} In stochastic settings, our method significantly outperforms other approaches, consistently surpassing expert-level performance more rapidly. This enhanced ability to learn under uncertainty is attributed to our framework’s effectiveness in leveraging model-based synthetic data.
\item \textbf{Performance in Deterministic Environments:} We demonstrate that our method is either better or competitive with existing AIL methods' performances in deterministic settings.
\item \textbf{Sample Efficiency} For both deterministic and stochastic settings, our method can reach expert performance with fewer training steps than all the other baselines.
\end{itemize}
We primarily compare our approach with other Adversarial Imitation Learning (AIL) methods, including the on-policy algorithms GAIL \citep{ho2016generative} and AIRL \citep{fu2017learning}, and the off-policy method Discriminator Actor-Critic (DAC) \citep{kostrikov2018discriminator}. For policy optimization, we use Proximal Policy Optimization (PPO) \citep{schulman2017proximal} for both GAIL and AIRL, and Soft Actor-Critic (SAC) \citep{haarnoja2018soft} for both DAC and our proposed method.
All implementations of PPO and SAC are referenced from the \texttt{Clean RL} library~\citep{huang2022cleanrl}.
Each algorithm is trained with 100k environmental steps and evaluated each 1k steps across 5 different seeds for tasks including \texttt{InvertedPendulum-v4} and \texttt{InvertedDoublePendulum-v4}.
For \texttt{Hopper-v3}, AIRL and GAIL are trained with 10M steps and evaluated each 100k steps across 5 different seeds, but DAC and our algorithm are trained with 1M environmental steps and evaluated each 10k steps across 5 different seeds. We conduct the experiments under 3 different random seeds to seek more general trend and conclude systematic results. All the experiments are run on the Desktop equipped with RTX 4090 and Core-i9 13900K. The learning curves of all methods are provided in \cref{appendix:learning_curves}.
\paragraph{Performance in Stochastic MuJoCo.}
In ~\cref{performance_table_stoch}, we present the performance under the stochastic MuJoCo environments.
\textbf{Our method can achieve the best performance compared to all of the baselines in stochastic environments.}
% For InvertedPendulum-v4, all four methods are able to reach expert-level performance.
% However, as shown in Fig~\ref{Fig::invertedpendulum} and Table~\ref{performance_table_reach}, our method reaches expert performance in the fewest training steps.
Specifically, for \texttt{InvertedDoublePendulum-v4}, introducing stochasticity into the dynamics significantly degrades the performance. Our method still can still achieve performance better than all baselines.
% By the end of training, all methods except DAC manage to exceed expert performance.
Furthermore, in the stochastic settings, the performance of DAC is decreased significantly, as DAC's ineffective reward formulation on state-action pairs under uncertain conditions, also resulting a trend of learning instability (\cref{appendix:learning_curves}).
For \texttt{Hopper-v4}, our method remarkably surpasses the all baselines with significantly large margins.
These results indicate that our approach can effectively learn the better reward function from demonstrations in stochastic environments, resulting in significant performance improvement.
% In the higher-dimensional Walker2d task, our method also closely matches the expert’s performance in relatively few steps whereas both GAIL and AIRL fail to recover expert-level performance within the allotted training steps.
\begin{table*}[h]
\centering
\caption{Performance of expert and all algorithms in stochastic MuJoCo Environments.}
\scalebox{0.8}{
\begin{tabular}{c|ccccc}
\hline
Environment
& Expert
& GAIL
& AIRL
& DAC
& Ours
\\ \hline
InvertedPendulum-v4(Stoch)
& $1000.0_{\pm 0.0}$
& $1000.0_{\pm 80.8}$
& $1000.0_{\pm 25.4}$
& $1000.0_{\pm 69.2}$
& $1000.0_{\pm 31.2}$
\\
InvertedDoublePendulum-v4(Stoch)
& $210.1_{\pm 98.6}$
& $276.9_{\pm 25.1}$
& $234.9_{\pm 38.1}$
& $142.3_{\pm 31.2}$
& \boldblue{$286.7_{\pm 39.2}$}
\\
% Walker2d-v3(Stoch)
% & $3330.2_{\pm 672.5}$
% & $1855.0_{\pm 198.1}$
% & $2301.0_{\pm 341.2}$
% &
% & \boldblue{$3557.0_{\pm 204.9}$}
% \\
% Pusher-v4(Stoch)
% & $3330.2_{\pm 672.5}$
% &
% &
% &
% &
% \\
Hopper-v3(Stoch)
& $2257.8_{\pm 317.4}$
& $2832.0_{\pm 105.5}$
& $2834.0_{\pm 333.2}$
& $2408.0_{\pm 78.6}$
& \boldblue{$3139.0_{\pm 109.9}$}
\\
% Ant-v3(Stoch)
% & $1534.4_{\pm 89.3}$
% &
% &
% &
% &
% \\
% HalfCheetah-v3(Stoch)
% & $1534.4_{\pm 89.3}$
% &
% &
% &
% &
% \\
\hline
\end{tabular}
}
\label{performance_table_stoch}
\end{table*}
\paragraph{Performance in Deterministic MuJoCo.}
The performance of deterministic MuJoCo environments can be found in ~\cref{performance_table_det}.
For the tasks with deterministic dynamics, our method can achieve the performance aligning with all of baselines and the expert in \texttt{InvertedDoublePendulum-v4}, \texttt{InvertedDoublePendulum-v4} and \texttt{Hopper-v3}.
Additionally, our method shows superiority in learning stability (\cref{appendix:learning_curves}) and sample efficiency (\cref{performance_table_reach}).
Therefore, \textbf{our method shows competitive performance with the baselines in the deterministic environments with better sample efficiency and learning stability.}
% However, in terms of sample efficiency, as reported in Table~\ref{performance_table_reach}, our approach shows a significant advantage over the on-policy methods, namely GAIL and AIRL. When compared to DAC, our advantage is less prominent in low-dimensional tasks with simple transitions such as InvertedPendulum. This is because in deterministic settings, our model-enhanced reward shaping does not provide any immediate benefits for learning the reward function. Additionally, at the early stages of training, the transition model is not yet accurate, so the generated samples do not contribute effectively to policy updates. However, for tasks with more complex transitions, as training progresses, the synthetic samples through the learned transition become more and more realistic. Meanwhile, the ratio of the synthetic samples also increases as these samples accumulate in the growing buffer. Consequently, in tasks such as InvertedDoublePendulum, and Hopper, our method demonstrates a clear advantage in sample efficiency compared to the other methods.
\begin{table*}[h]
\centering
\caption{Performance of expert and all algorithms in deterministic MuJoCo Environments.}
\scalebox{0.90}{
\begin{tabular}{c|ccccc}
\hline
Environment
& Expert
& GAIL
& AIRL
& DAC
& Ours
\\ \hline
InvertedPendulum-v4
& $1000.0_{\pm 0.0}$
& $1000.0_{\pm 0.0}$
& $1000.0_{\pm 0.0}$
& $1000.0_{\pm 0.0}$
& $1000.0_{\pm 0.0}$
\\
InvertedDoublePendulum-v4
& $9320.0_{\pm 578.8}$
& $9322.0_{\pm 39.6}$
& $9359.0_{\pm 88.3}$
& $9360.0_{\pm 12.3}$
& \boldblue{$9361.0_{\pm 0.2}$}
\\
% Walker2d-v3
% & $4493.8_{\pm 782.8}$
% & $4302.9_{\pm 104.4}$
% & $3994.3_{\pm 12.3}$
% & \boldblue{$4865.7_{\pm 45.3}$}
% & $4602.1_{\pm 307.6}$
% \\
% Pusher-v4
% & $-42.5_{\pm 4.6}$
% &
% &
% &
% &
% \\
Hopper-v3
& $3329.3_{\pm 306.6}$
& $3485.0_{\pm 5.6}$
& $3472.0_{\pm 3.6}$
& \boldblue{$3495.0_{\pm 4.5}$}
& $3463.0_{\pm 2.7}$
\\
% Ant-v3
% & $3370.4_{\pm 868.1}$
% &
% &
% &
% &
% \\
% HalfCheetah-v3
% & $12988.0_{\pm 819.4}$
% &
% &
% &
% &
% \\
\hline
\end{tabular}
}
\label{performance_table_det}
\end{table*}
\paragraph{Sample Efficiency.}
In ~\cref{performance_table_reach}, we report the number of training steps required to achieve expert-level performance.
\textbf{Our method shows significant superiority in sample efficiency across all of the benchmarks under both deterministic and stochastic settings.}
Specifically, for \texttt{InvertedPendulum-v4}, all four methods are able to reach expert-level performance. However, our method reaches expert performance in the fewest training steps, $9.7k$ and $11.7k$ for deterministic and stochastic settings respectively.
In the case of \texttt{InvertedDoublePendulum-v4}, introducing stochasticity into the dynamics makes it challenging for all algorithms to reach expert-level performance from noisy demonstrations. Especially, the strong baseline, DAC, failed to reach the expert's performance, while our method can arrive at the expert-level performance within $27.3k$ steps.
For \texttt{Hopper-v3}, all four methods achieve expert-level performance, but our method reaches it with the fewest samples whether in deterministic ($195.3k$) or stochastic ($170.0k$) settings.
\begin{table*}[h]
\centering
\caption{Averaged steps used for all algorithms to reach experts' performances in deterministic and stochastic Environments. $X$ stands for that the algorithm does not reach the expert performance within given training steps.}
\scalebox{0.95}{
\begin{tabular}{c|cccc}
\hline
Environment
& GAIL
& AIRL
& DAC
& Ours
\\ \hline
InvertedPendulum-v4
& $37.7k$
& $38.0k$
& $15.3k$
& \boldblue{$9.7k$}
\\
InvertedPendulum-v4(Stoch)
& $54.0k$
& $58.0k$
& $11.7k$
& \boldblue{$11.7k$}
\\
InvertedDoublePendulum-v4
& $87.0k$
& $91.5k$
& $28.3k$
& \boldblue{$11.3k$}
\\
InvertedDoublePendulum-v4(Stoch)
& $72.3k$
& $75.3k$
& X
& \boldblue{$27.3k$}
\\
Hopper-v3
& $5400.0k$
& $6400.0k$
& $243.5k$
& \boldblue{$195.3k$}
\\
Hopper-v3(Stoch)
& $4170.0k$
& $3200.0k$
& $786.7k$
& \boldblue{$170.0k$}
\\
% HalfCheetah-v3
% & X
% & X
% &
% &
% \\
% HalfCheetah-v3(Stoch)
% & X
% & X
% &
% &
% \\
\hline
\end{tabular}
}
\label{performance_table_reach}
\end{table*}
\section{Conclusion}
\label{section::conclusion}
In this paper, we presented a novel model-enhanced adversarial inverse reinforcement learning framework by incorporating model-based techniques with reward shaping, specifically designed to enhance performance in stochastic environments and maintain competitive performance in deterministic setting while significantly improve performance under the traditional AIL approaches. The theoretical analysis provides guarantees on the optimal policy invariance under the transition model involved reward shaping and highlight the relationship between performance gap and transition model error, showing that the gaps becomes negligible with a well-learned model. Empirical evaluations on Mujoco benchmark environments validate the effectiveness of our method, showcasing its superior performance and sample efficiency across different tasks.
Future works will focus on further refining the model estimation process to handle more complex and dynamic environments and exploring extensions of the framework to multi-agent and hierarchical reinforcement learning scenarios. Additionally, it would be valuable to investigate the generalization ability of our framework in a transfer learning tasks. Overall, our approach offers a promising direction for advancing model-based adversarial IRL, with the potential to scale to a broader range of real-world applications.
\newpage
\appendix
\section{Reward Shaping Soft Optimal Policy}
\begin{theorem}
\label{proof::soft_optimal_policy}
Let $R$ and $\hat{R}$ be two reward functions. $R$ and $\hat{R}$ induce the same soft optimal policy under all transition dynamics $\mathcal{T}$ if $\hat{R}(s_t,a_t)=R(s_t,a_t)+\gamma\mathbb{E}_{\mathcal{T}}[\phi(s_{t+1})\vert s_t,a_t] - \phi(s_t)$ for some potential-shaping function $\phi:\mathcal{S}\rightarrow\mathbb{R}$.
\end{theorem}
\begin{proof}
According to Soft VI (~\cref{equation::soft_VI}), we can expand the representation of $Q^{soft}_{\hat{R}}(s_t,a_t)$ as follows.
\begin{align*}
Q^{soft}_{\hat{R}}(s_t,a_t) &= R(s_t,a_t)+\gamma\mathbb{E}_{\mathcal{T}}\left[\phi(s_{t+1})\vert s_t, a_t\right]-\phi(s_t)+\gamma\mathbb{E}_{\mathcal{T}}\left[V_{\hat{R}}^{soft}(s_{t+1})\vert s_t, a_t\right],\\
Q^{soft}_{\hat{R}}(s_t,a_t) + \phi(s_t) &= R(s_t,a_t) + \gamma\mathbb{E}_{\mathcal{T}}\left[V_{\hat{R}}^{soft}(s_{t+1})+\phi(s_{t+1})\vert s_t, a_t\right],\\
Q^{soft}_{\hat{R}}(s_t,a_t) + \phi(s_t) &= R(s_t,a_t) + \gamma\mathbb{E}_{\mathcal{T}}\left[\log\sum_{a\in\mathcal{A}}\exp\left(Q_{\hat{R}}^{soft}(s_{t+1},a)+\phi(s_{t+1})\right)\vert s_t, a_t\right].
\end{align*}
From above induction, we can tell that $Q^{soft}_{\hat{R}}(s_t,a_t) + \phi(s_t)$ satisfy the soft bellmen update with original $R$. Thus, with simple induction, we can arrive that $Q_{R}^{soft}(s_t,a_t) = Q_{\hat{R}}^{soft}(s_t,a_t)+\phi(s_t)$.
Then, we can derive the advantage function
\begin{align*}
A_{\hat{R}}^{soft}(s_t,a_t) &= Q_{\hat{R}}^{soft}(s,a) - V_{\hat{R}}^{soft}(s_t)\\
&= Q_{\hat{R}}^{soft}(s_t,a_t)-\log\sum_{a\in\mathcal{A}}\exp\left(Q_{\hat{R}}^{soft}(s_t,a_t)\right)\\
&= Q_{\hat{R}}^{soft}(s_t,a_t) + \phi(s_t) -\log\sum_{a\in\mathcal{A}}\exp\left(Q_{\hat{R}}^{soft}(s_t,a_t) + \phi(s_t)\right)\\
&= Q_{R}^{soft}(s_t,a_t)-\log\sum_{a\in\mathcal{A}}\exp\left(Q_{R}^{soft}(s_t,a_t)\right)\\
&= A_{R}^{soft}(s_t,a_t).
\end{align*}
\end{proof}
\section{Adversarial Reward Learning}
\begin{proposition}
\label{proof::gradient alignment}
Consider an undiscounted MDP. Suppose $f_\theta$ and $\pi$ at current iteration are the soft-optimal advantage function and policy for reward function $R_\theta$. Minimising the cross-entropy loss of the discriminator under generator $\pi$ is equivalent to maximising the log-likelihood under MCE IRL.
\end{proposition}
\begin{proof}
\begin{align*}
\mathcal{L}_{IRL}(\mathcal{D}_{exp},\theta) &= \mathbb{E}_{\mathcal{D}_{exp}}[\log p_\theta(\tau)]\\
&= \mathbb{E}_{\mathcal{D}_{exp}}\left[\sum_{t=0}^{T-1}\log\pi(a_t|s_t)+\log\rho_0+\sum_{t=1}^T\mathcal{T}(s_{t+1}|s_t,a_t)\right]\\
&= \mathbb{E}_{\mathcal{D}_{exp}}\left[\sum_{t=0}^{T-1}\left(Q^{soft}_\theta(s_t,a_t)-V^{soft}_\theta(s_t)\right)\right] + \text{constant}.
\end{align*}
Breaking down above equations with soft VI (\cref{equation::soft_VI}), we can arrive the following.
\begin{equation}
\mathbb{E}_{\mathcal{D}_{exp}}\left[\sum_{t=0}^{T-1} R_\theta(s_t,a_t)\right]+\mathbb{E}_{\mathcal{D}_{exp}}\left[\sum_{t=0}^{T-2}\mathbb{E}_{\mathcal{T}}\left[V_\theta^{soft}(s_{t+1})|s_t,a_t\right]\right]-\mathbb{E}_{\mathcal{D}_{exp}}\left[\sum_{t=0}^{T-1} V_\theta^{soft}(s_t)\right].
\end{equation}
Next we will derive the gradient of the loss.
\begin{multline}
\nabla_{\theta}\mathcal{L}(\mathcal{D}_{exp},\theta) = \underbrace{\nabla_{\theta}\mathbb{E}_{\mathcal{D}_{exp}}\left[\sum_{t=0}^{T-1} R_{\theta}(s_t,a_t)\right]}_{A} + \\\underbrace{\nabla_{\theta}\mathbb{E}_{\mathcal{D}_{exp}}\left[\sum_{t=0}^{T-2}\left(\mathbb{E}_{\mathcal{T}}\left[V_{\theta}^{soft}(s_{t+1})\vert s_t,a_t\right]\right) -V_{\theta}^{soft}(s_{t+1})\right]}_{B} - \underbrace{\nabla_{\theta}V_{\theta}^{soft}(s_0)}_C.
\end{multline}
Let's get explicit expression of each part.
\begin{align*}
A &= \mathbb{E}_{\mathcal{D}_{exp}}\left[\sum_{t=0}^{T-1}\nabla_{\theta}R_{\theta}(s_t,a_t)\right]\\
C &=\nabla_\theta\log\sum_{a_t\in\mathcal{A}}\exp{Q_\theta^{soft}(s_t,a_t)}\\
&=\sum_{a_t\in\mathcal{A}}\pi(a_t\vert s_t)\nabla_\theta Q_{\theta}^{soft}(s_t,a_t)\\
& = \mathbb{E}_{\pi}\left[\sum_{t=0}^{T-1}\nabla_\theta R_{\theta}(s_t,a_t)\right].
\end{align*}
The transitions in the demonstration $\mathcal{D}_{exp}$ will follow $\mathcal{T}$. Thus we can drop the $\mathbb{E}_{\mathcal{T}}$, and $B$ term will cancel out, ending up to $0$. To summarize, the gradient of log MLE loss of MCE IRL is the following.
\begin{equation}
\nabla_\theta\mathcal{L}_{IRL}(\mathcal{D}_{exp},\theta)=\mathbb{E}_{\mathcal{D}_{exp}}\left[\sum_{t=0}^{T-1}\nabla_{\theta}R_{\theta}(s_t,a_t)\right] -\mathbb{E}_{\pi}\left[\sum_{t=0}^{T-1}\nabla_\theta R_{\theta}(s_t,a_t)\right].
\end{equation}
Next, we will start to derive the gradient of cross-entropy discriminator training loss. Remember the discriminator loss is defined in Eq~\ref{equation::disc_loss}.
\begin{align*}
\log D_\theta(s,a) &= f_\theta(s,a)-\log(\exp\{f_\theta(s,a)\}+\pi(a|s)),\\
\log(1-D_\theta(s,a)) &= \log\pi(a|s)-\log(\exp\{f_\theta(s,a)\}+\pi(a|s)).
\end{align*}
Then, the gradient of each term is as follow:
\begin{align*}
\nabla_\theta\log D_\theta(s,a) &= \nabla_\theta f_\theta(s,a)-\frac{\exp\{f_\theta(s,a)\}\nabla_\theta f_\theta(s,a)}{\exp\{f_\theta(s,a)\}+\pi(a|s)},\\
\nabla_\theta\log(1-D_\theta(s,a)) &= -\frac{\exp\{f_\theta(s,a)\}\nabla_\theta f_\theta(s,a)}{\exp\{f_\theta(s,a)\}+\pi(a|s)}.
\end{align*}
Since $\pi$ is trained by using $f_\theta$ as shaped reward, from soft VI we can derive that $\pi^*_{f_\theta}(a|s)=\exp A^{soft}_{f_\theta}(s,a)$. By assumption, we assume that $f_\theta$ is the advantage function of $R_\theta$, $f_\theta(s,a)=A^{soft}_{R_\theta}(s,a)$. From Thm~\ref{thm::soft_optimal_policy}, we know that $A^{soft}_{R_\theta}(s,a)=A^{soft}_{f_\theta}(s,a)$, which also implies that $\pi^*_{f_\theta}=\pi^*_{R_\theta}$. Then, we can deduce the gradient of the loss of discriminator.
\begin{align*}
-\nabla_\theta\mathcal{L}_{disc} &= \mathbb{E}_{\mathcal{D}_{exp}}\left[\nabla_\theta\log D_\theta(s,a)\right]+\mathbb{E}_{\pi}\left[\nabla_\theta\log(1-D_\theta(s,a))\right]\\
&= \mathbb{E}_{\mathcal{D}_{exp}}\left[\frac{1}{2}\nabla_\theta f_\theta(s,a)\right]-\mathbb{E}_{\pi}\left[\frac{1}{2}\nabla_\theta f_\theta(s,a)\right],\\
-2\nabla_\theta\mathcal{L}_{disc} &= \mathbb{E}_{\mathcal{D}_{exp}}\left[\nabla _\theta f_\theta(s,a)\right] - \mathbb{E}_{\pi}\left[\nabla_\theta f_\theta(s,a)\right].
\end{align*}
\end{proof}
\section{Performance Gap Analysis}
\begin{lemma}[Implicit Feasible Reward Set~\citep{ng2000algorithms}]
Let $\mathfrak{B}=(\mathcal{M}',\pi^*)$ be an IRL problem. Then $R\in\mathcal{R}_\mathfrak{B}$ if and only if for all $(s,a)\in\mathcal{S}\times\mathcal{A}$ the following holds:
\begin{align*}
Q_{\mathcal{M}'\cup R}^{\pi^*}(s,a)-V_{\mathcal{M}'\cup R}^{\pi^*}(s) &= 0\quad \textrm{if}\; \pi^*(a|s)>0,\\
Q_{\mathcal{M}'\cup R}^{\pi^*}(s,a)-V_{\mathcal{M}'\cup R}^{\pi^*}(s) &\leq 0\quad \textrm{if}\; \pi^*(a|s)=0.
\end{align*}
\end{lemma}
Combined with the traditional Value Iteration of RL, we can write out the explicit form of the reward function $R$.
\begin{lemma}[Explicit Feasible Reward Function~\citep{metelli2021provably}]
\label{lemma:feasi_reward}
With the above lemma conditions, $R\in\mathcal{R}_\mathfrak{B}$ if and only if there exist $\xi\in\mathbb{R}_{\geq 0}^{\mathcal{S}\times\mathcal{A}}$ and value function $V\in\mathbb{R}^\mathcal{S}$ such that:
\begin{equation}\label{equation::explicit_reward}
R(s,a) = V(s) - \gamma \sum_{s'\in\mathcal{S}} \mathcal{T}(s'|s,a)V(s') - \xi(s,a)\mathbb{I}\{\pi^*(a|s)=0\}.
\end{equation}
\end{lemma}
With ~\cref{equation::explicit_reward}, we can derive the following error bound between $R\in \mathcal{R}^E_{\mathfrak{B}}$ and $\hat{R}^E\in\mathcal{R}_{\hat{\mathfrak{B}}}$.
% \begin{theorem}[Reward Function Learning Error With Estimated MDP]
% Let $\mathfrak{B}=(\mathcal{M}',\pi^*)$ and $\hat{\mathfrak{B}}=(\hat{\mathcal{M}'},\pi^*)$ be two IRL problems, then for any $R^E\in\mathcal{R}_{\mathfrak{B}}$ there is a corresponding $\hat{R}^E\in\mathcal{R}_{\hat{\mathfrak{B}}}$ such that
% \begin{equation}
% \Vert R^E-\hat{R}^E\Vert\leq\frac{\gamma}{1-\gamma}|\mathcal{S}|\epsilon_{\mathcal{T}}R_{\max}.
% \end{equation}
% \end{theorem}
\begin{theorem}[Reward Function Error Bound]
\label{proof:reward_error_bound}
Let $\mathfrak{B}=(\mathcal{M}',\pi^*)$ and $\hat{\mathfrak{B}}=(\hat{\mathcal{M}'},\pi^*)$ be two IRL problems, then for any $R^E\in\mathcal{R}_{\mathfrak{B}}$ there is a corresponding $\hat{R}^E\in\mathcal{R}_{\hat{\mathfrak{B}}}$ such that
\begin{equation}
\Vert R^E-\hat{R}^E\Vert\leq\frac{\gamma}{1-\gamma}|\mathcal{S}|\epsilon_{\mathcal{T}}R_{\max}.
\end{equation}
\end{theorem}
\begin{proof}
From ~\cref{lemma:feasi_reward}, we can derive the following representations of $R$ and $\hat{R}$ with the same set of $V$ and $\xi$:
\begin{align*}
R^E(s,a) &= V(s) - \gamma \sum_{s' \in \mathcal{S}} \mathcal{T}(s'|s,a)V(s') - \xi(s,a)\mathbb{I}\{\pi^*(a|s) = 0\}, \\
\hat{R}^E(s,a) &= V(s) - \gamma \sum_{s' \in \mathcal{S}} \hat{\mathcal{T}}(s'|s,a)V(s') - \xi(s,a)\mathbb{I}\{\pi^*(a|s) = 0\}.
\end{align*}
The difference between $R^E$ and $\hat{R}^E$ can be bounded as follows:
\begin{align*}
\Vert R^E - \hat{R}^E \Vert &\leq \gamma \sum_{s' \in \mathcal{S}} D_{TV}\left(\mathcal{T}(s'|s,a)\vert\hat{\mathcal{T}}(s'|s,a)\right) \cdot \Vert V(s')\Vert.
\end{align*}
Given that the total variation distance between the two dynamics is bounded by $\epsilon_{\mathcal{T}}$, and the reward function is bounded by $R_{\max}$, together with the definition of the value function, we have $\Vert V \Vert_{\infty} \leq \frac{R_{\max}}{1 - \gamma}$. Substituting these bounds, we derive the following inequality:
\begin{equation*}
\Vert R^E - \hat{R}^E\Vert \leq \frac{\gamma}{1 - \gamma} |\mathcal{S}| \epsilon_{\mathcal{T}} R_{\max}.
\end{equation*}
\end{proof}
Next, we will propagate this bound to the value functions of optimal policy regarding different reward functions $R$ and $\hat{R}$. From the traditional Value iteration, we can write out the value function.
\begin{equation}
V_{\mathcal{M}'\cup R}^\pi(s) =\sum_{a\in\mathcal{A}}\pi(a|s)\sum_{s'\in\mathcal{S}}\mathcal{T}(s'|s,a)\left[R(s,a)+\gamma V_{\mathcal{M}'\cup R}^\pi(s')\right].
\end{equation}
\begin{lemma}[Value Function Error under same policy and different rewards and MDP]
$||V_{\mathcal{M}'\cup R}^\pi(s) - V_{\hat{\mathcal{M}}'\cup \hat{R}}^\pi(s)||$: the performance difference of the same policy in different MDPs.
\begin{equation}
||V_{\mathcal{M}'\cup R}^\pi(s) - V_{\hat{\mathcal{M}}'\cup \hat{R}}^\pi(s)||\leq
\epsilon_\mathcal{T} \frac{1+\gamma}{(1-\gamma)^2}R_{\max} |\mathcal{S}|.
\end{equation}
\end{lemma}
\begin{proof}
\begin{equation}
\begin{aligned}
&||V_{\mathcal{M}'\cup R}^\pi(s) - V_{\hat{\mathcal{M}}'\cup \hat{R}}^\pi(s)||\\
&\leq\sum_{a\in\mathcal{A}}\pi(a|s)\sum_{s'\in\mathcal{S}}||\mathcal{T}(s'|s,a)\left[R(s,a)+\gamma V_{\mathcal{M}'\cup R}^\pi(s')\right]-\hat{\mathcal{T}}(s'|s,a)\left[R(s,a)+\gamma V_{\mathcal{M}'\cup R}^\pi(s')\right]\\
&+\hat{\mathcal{T}}(s'|s,a)\left[R(s,a)+\gamma V_{\mathcal{M}'\cup R}^\pi(s')\right]-\hat{\mathcal{T}}(s'|s,a)\left[\hat{R}(s,a)+\gamma V_{\hat{\mathcal{M}}'\cup \hat{R}}^\pi(s')\right]||\\
&\leq \sum_{a\in\mathcal{A}}\pi(a|s)\sum_{s'\in\mathcal{S}} (\epsilon_\mathcal{T} \frac{R_{\max}}{1-\gamma}+\hat{\mathcal{T}}(s'|s,a)\epsilon_{\mathcal{T}} (\frac{\gamma R_{\max}}{1-\gamma}|\mathcal{S}| + \gamma ||V_{\mathcal{M}'\cup R}^\pi(s') - V_{\hat{\mathcal{M}}'\cup \hat{R}}^\pi(s')||))
\\
&=\sum_{a\in\mathcal{A}}\pi(a|s)(\epsilon_\mathcal{T} \frac{R_{\max}}{1-\gamma}|\mathcal{S}|+\epsilon_{\mathcal{T}} \frac{\gamma R_{\max}}{1-\gamma}|\mathcal{S}|+ \gamma ||V_{\mathcal{M}'\cup R}^\pi(s') - V_{\hat{\mathcal{M}}'\cup \hat{R}}^\pi(s')||)\\
&\leq\epsilon_\mathcal{T} \frac{1+\gamma}{1-\gamma}R_{\max} |\mathcal{S}|
+ \gamma ||V_{\mathcal{M}'\cup R}^\pi(s') - V_{\hat{\mathcal{M}}'\cup \hat{R}}^\pi(s')||\\
&\leq\epsilon_\mathcal{T} \frac{1+\gamma}{(1-\gamma)^2}R_{\max} |\mathcal{S}|.
\end{aligned}
\end{equation}
\end{proof}
\begin{lemma}
\label{lemma:different_policy}
Let $\Vert V_{\hat{\mathcal{M}'}\cup\hat{R}}^{\pi_{1}}(s) - V_{\hat{\mathcal{M}'}\cup\hat{R}}^{\pi_{2}}(s)\Vert$ denote the performance difference between different policies $\pi_1$ and $\pi_2$ in the same learned MDP \citep{viano2021robust, zhang2020learning}. The following inequality holds:
\begin{equation*}
\Vert V_{\hat{\mathcal{M}'}\cup\hat{R}}^{\pi_{1}}(s) - V_{\hat{\mathcal{M}'}\cup\hat{R}}^{\pi_{2}}(s)\Vert \leq \frac{\gamma}{(1 - \gamma)^2} \epsilon_{\mathcal{T}} R_{\max}.
\end{equation*}
\end{lemma}
% \begin{theorem}[Value Function Error under between expert and imitator]
% The value function difference bound between the expert and the imitator is as follow.
% \begin{equation}
% ||V_{\mathcal{M}'\cup R}^{\pi^*}(s) - V_{\hat{\mathcal{M}}'\cup \hat{R}}^{\hat{\pi}^*}(s)||\leq
% \epsilon_\mathcal{T}\left[\frac{\gamma}{(1-\gamma)^2}R_{\max}+\frac{1+\gamma}{(1-\gamma)^2}R_{\max}\vert\mathcal{S}\vert\right].
% \end{equation}
% \end{theorem}
\begin{theorem}[Performance Difference Bound]
\label{proof:value_error_bound}
% $\Vert V_{\mathcal{M}'\cup R}^{\pi^*}(s) - V_{\hat{\mathcal{M}}'\cup \hat{R}}^{\hat{\pi}^*}(s)\Vert$:
The performance difference between the optimal policies ($\pi^*$ and $\hat{\pi}^*$) in corresponding MDPs ($\mathcal{M}'\cup R$ and $\hat{\mathcal{M}}'\cup \hat{R}$) can be bounded as follows:
\begin{equation}
\Vert V_{\mathcal{M}'\cup R}^{\pi^*} - V_{\hat{\mathcal{M}}'\cup \hat{R}}^{\hat{\pi}^*} \Vert\leq
\epsilon_\mathcal{T}\left[ \frac{\gamma}{(1-\gamma)^2} R_{\max}+\frac{1+\gamma}{(1-\gamma)^2}R_{\max} |\mathcal{S}|\right].
\end{equation}
\end{theorem}
\begin{proof}
\begin{align*}
&||V_{\mathcal{M}'\cup R}^{\pi^*}(s) - V_{\hat{\mathcal{M}}'\cup \hat{R}}^{\hat{\pi}^*}(s)||\\
&\leq||V_{\hat{\mathcal{M}}'\cup \hat{R}}^{\pi^*}(s) - V_{\hat{\mathcal{M}'}\cup \hat{R}}^{\hat{\pi}^*}(s)|| + ||V_{\mathcal{M}'\cup R}^{\hat{\pi}^*}(s) - V_{\hat{\mathcal{M}}'\cup \hat{R}}^{\hat{\pi}^*}(s)|| \\
&=\epsilon_\mathcal{T} \frac{\gamma}{(1-\gamma)^2} R_{\max}+ \epsilon_\mathcal{T} \frac{1+\gamma}{(1-\gamma)^2}R_{\max} |\mathcal{S}|\\
&=\epsilon_\mathcal{T}\left[\frac{\gamma}{(1-\gamma)^2}R_{\max}+\frac{1+\gamma}{(1-\gamma)^2}R_{\max}\vert\mathcal{S}\vert\right].
\end{align*}
\end{proof}
% \begin{theorem}
% \label{thm:mbpo_monotonic}
% \cite{janner2019trust} Let the expected TV-distance
% between two transition distributions be bounded at each timestep by $\epsilon_m$ and the policy divergence be bounded by $\epsilon_\pi$. Then the true returns and model returns of the policy are bounded as:
% \begin{equation}
% \eta[\pi] \ge \hat{\eta}[\pi] -
% \underbrace{\left[\frac{2\gamma r_\mathrm{max} (\epsilon_m+2\epsilon_\pi)}{(1-\gamma)^2} + \frac{4r_\mathrm{max}\epsilon_\pi}{(1-\gamma)} \right]}_{C(\epsilon_m, \epsilon_\pi)}
% \end{equation}
% \end{theorem}
% This bound implies that as long as we improve the returns under the model $\hat{\eta}[\pi]$ by more than $C(\epsilon_m, \epsilon_\pi)$, we can guarantee improvement under the true returns.
% The detailed model-based H-step look-ahead planning is listed in the below section. We define a few terms for the necessary analysis and proof. We define a finite H-length MDP $\mathcal{H}$ as $\langle\mathcal{S},\mathcal{A},R_{mix},\gamma,\mathcal{T}\rangle$. Correspondingly, there is a MDP $\hat{\mathcal{H}}$ as $\langle\mathcal{S},\mathcal{A},\hat{R}_{mix},\gamma,\hat{\mathcal{T}}\rangle$, $R_{mix}(s_t,a_t)=R^E(s_t,a_t)$\;if $t<H$ and $R_{mix}(s_H,a_H)=V(s_H)$\;if $t=H$. Given a value function $V$, we define the policy $\pi_{H,V}$ obtained by maximizing the H-step lookahead objective.
% \begin{equation}
% \pi_{H,V}(s_0)=\arg\max_{a_0}\;\max_{a_1,a_2,\cdots,a_{H-1}}\mathbb{E}_{\hat{\mathcal{T}}}\left[\sum_{t=0}^{H-1}\gamma^t \hat{R}^E(s_t,a_t)+\gamma^H \hat{V}_{\mathcal{M}\cup\hat{R}^E}^{\pi_{\hat{\mathcal{H}}}}(s_H)\right]
% \end{equation}
% We have the total variation between $\mathcal{T}$ and $\hat{\mathcal{T}}$ bounded by $\epsilon_\mathcal{T}$. From Lemma B.2~\citep{janner2019trust}, we can have the state distribution bound after t step propagation under $\mathcal{H}$ and $\hat{\mathcal{H}}$ as $|\rho^t(s,a)-\hat{\rho}^t(s,a)|\leq t\epsilon_{\mathcal{T}}$, where $\rho^t(s,a)$ is the state-action visitation induced by $\pi$ on $\mathcal{H}$ till t time step, vice versa for $\hat{\rho}^t(s,a)$ on $\hat{\mathcal{H}}$. To start the proof, we need to define a few trajectories. $\tau$ is the H-step trajectory of running $\pi^*_{\mathcal{H}}$ in the MDP $\mathcal{M}\cup R^E$; $\tau^*$ is trajectory of running $\pi^*_{\mathcal{M}}$ in $\mathcal{M}\cup R^E$; $\hat{\tau}$ is the trajectory of running $\pi_{\hat{\mathcal{H}}}^*$ in $\mathcal{M}\cup\hat{R}^E$.
% \begin{theorem}
% Given $\mathcal{M}, \hat{\mathcal{M}}, \mathcal{H}, \hat{\mathcal{H}}$, and bound between $R^E$ and $\hat{R}^E$ above. Assumed the suboptimal bound caused in solving equation 12 as $\epsilon_p$ and bound in value function approximation as $\epsilon_v$. The optimality gap between traditional model-based IRL with Value Iteration and model-based IRL with unknown model and H-step lookahead planning is as follow.
% \begin{equation}
% J^{\pi^*} - J^{\pi_{\hat{\mathcal{H}}}} \leq 2C\epsilon_\mathcal{T}+\epsilon_p + \gamma^H\epsilon_v
% \end{equation}
% where $$C=\left[\frac{\gamma^H R_\max}{1-\gamma}\left(H - \frac{1+\gamma}{1-\gamma}|\mathcal{S}|\right)+\sum_{t=0}^{H-1}\gamma^t R_{\max}\left(t+\frac{\gamma|\mathcal{S}|}{1-\gamma}\right)\right]$$
% \end{theorem}
% \begin{proof}
% \begin{align*}
% J^{\pi^*} - J^{\pi_{\hat{\mathcal{H}}}} &= V_{\mathcal{M}'\cup R^E}^{\pi^*}(s_0) - \hat{V}_{\mathcal{M}'\cup\hat{R}^E}^{\pi_{\hat{\mathcal{H}}}}(s_0)\\
% &= \mathbb{E}_{\tau^*\sim p_{\tau^*}}\left[\sum_{t=0}^{H-1}\gamma^tR^E(s,a)+\gamma^HV^*_{\mathcal{M}'\cup R}(s_H)\right] - \mathbb{E}_{\hat{\tau}\sim p_{\hat{\tau}}}\left[\sum_{t=0}^{H-1}\gamma^t\hat{R}^E(s,a)+\gamma^H \hat{V}_{\mathcal{M}'\cup\hat{R}}^{\pi_{\hat{\mathcal{H}}}} (s_H)\right]\\
% &= \mathbb{E}_{\tau^*\sim p_{\tau^*}}\left[\sum_{t=0}^{H-1}\gamma^tR^E(s,a)+\gamma^HV^*_{\mathcal{M}'\cup R}(s_H)\right] -
% \mathbb{E}_{\hat{\tau}\sim p_{\hat{\tau}}}\left[\sum_{t=0}^{H-1}\gamma^t\hat{R}^E(s,a)+\gamma^H V_{\mathcal{M}'\cup\hat{R}}^{*} (s_H)\right]\\
% &+ \mathbb{E}_{\hat{\tau}\sim p_{\hat{\tau}}}\left[\sum_{t=0}^{H-1}\gamma^t\hat{R}^E(s,a)+\gamma^H V_{\mathcal{M}'\cup\hat{R}}^{*} (s_H)\right] -
% \mathbb{E}_{\hat{\tau}\sim p_{\hat{\tau}}}\left[\sum_{t=0}^{H-1}\gamma^t\hat{R}^E(s,a)+\gamma^H \hat{V}_{\mathcal{M}'\cup\hat{R}}^{\pi_{\hat{\mathcal{H}}}} (s_H)\right]\\
% &= \mathbb{E}_{\tau^*\sim p_{\tau^*}}\left[\sum_{t=0}^{H-1}\gamma^tR^E(s,a)+\gamma^HV^*_{\mathcal{M}'\cup R}(s_H)\right] -
% \mathbb{E}_{\hat{\tau}\sim p_{\hat{\tau}}}\left[\sum_{t=0}^{H-1}\gamma^t\hat{R}^E(s,a)+\gamma^H V_{\mathcal{M}'\cup\hat{R}}^{*} (s_H)\right]\\
% &+ \gamma^H\mathbb{E}_{\hat{\tau}\sim p_{\hat{\tau}}}\left[V^*_{\mathcal{M}'\cup\hat{R}}(s_H) - \hat{V}^{\pi_{\hat{\mathcal{H}}}}_{\mathcal{M}'\cup\hat{R}}(s_H)\right]
% \end{align*}
% We assume that the value function learning error is bounded by $\epsilon_v$ such that $\max_s|\hat{V}-V^*|\leq\epsilon_v$. Then we can further confine the above optimality gap.
% \begin{align*}
% &\leq\mathbb{E}_{\tau^*\sim p_{\tau^*}}\left[\sum_{t=0}^{H-1}\gamma^tR^E(s,a)+\gamma^HV^*_{\mathcal{M}'\cup R}(s_H)\right] -
% \mathbb{E}_{\tau\sim p_{\tau}}\left[\sum_{t=0}^{H-1}\gamma^tR^E(s,a)+\gamma^HV^*_{\mathcal{M}'\cup R}(s_H)\right]\\
% &+ \mathbb{E}_{\tau\sim p_{\tau}}\left[\sum_{t=0}^{H-1}\gamma^tR^E(s,a)+\gamma^HV^*_{\mathcal{M}'\cup R}(s_H)\right] -
% \mathbb{E}_{\hat{\tau}\sim p_{\hat{\tau}}}\left[\sum_{t=0}^{H-1}\gamma^t\hat{R}^E(s,a)+\gamma^H V_{\mathcal{M}'\cup\hat{R}}^{*} (s_H)\right] + \gamma^H\epsilon_v\\
% &\leq \mathbb{E}_{\tau\sim p_{\tau}}\left[\sum_{t=0}^{H-1}\gamma^tR^E(s,a)+\gamma^HV^*_{\mathcal{M}'\cup R}(s_H)\right] -
% \mathbb{E}_{\hat{\tau}\sim p_{\hat{\tau}}}\left[\sum_{t=0}^{H-1}\gamma^t\hat{R}^E(s,a)+\gamma^H V_{\mathcal{M}'\cup\hat{R}}^{*} (s_H)\right]+\gamma^H\epsilon_v
% \end{align*}
% From the definition of trajectories $\tau^*$ and $\tau$, $\mathbb{E}_{\tau^*\sim p_{\tau^*}}\left[\sum_{t=0}^{H-1}\gamma^tR^E(s,a)+\gamma^HV^*_{\mathcal{M}'\cup R}(s_H)\right] - \mathbb{E}_{\tau\sim p_{\tau}}\left[\sum_{t=0}^{H-1}\gamma^tR^E(s,a)+\gamma^HV^*_{\mathcal{M}'\cup R}(s_H)\right]\leq 0$ must hold. Thus, we can get rid the first two terms of the above equation and still make the inequality sound. Next, we transfer the original gap into optimal return gap between two optimal policies in $\pi^*_\mathcal{H}$ and $\pi^*_\hat{\mathcal{H}}$ in $\mathcal{H}$.
% \begin{align*}
% \mathbb{E}_{\tau\sim p_{\tau}}&\left[\sum_{t=0}^{H-1}\gamma^tR^E(s,a)+\gamma^HV^*_{\mathcal{M}'\cup R}(s_H)\right] -
% \mathbb{E}_{\hat{\tau}\sim p_{\hat{\tau}}}\left[\sum_{t=0}^{H-1}\gamma^t\hat{R}^E(s,a)+\gamma^H V_{\mathcal{M}'\cup\hat{R}}^{*} (s_H)\right]+\gamma^H\epsilon_v\\
% &=J^{\pi^*_{\mathcal{H}}}_{\mathcal{H}} - J^{\pi^*_{\hat{\mathcal{H}}}}_{\mathcal{H}} + \gamma^H\epsilon_v \\
% &= J^{\pi^*_{\mathcal{H}}}_{\mathcal{H}} - J^{\pi^*_{\mathcal{H}}}_{\hat{\mathcal{H}}}+J^{\pi^*_{\mathcal{H}}}_{\hat{\mathcal{H}}}-J^{\pi^*_{\hat{\mathcal{H}}}}_{\hat{\mathcal{H}}}+J^{\pi^*_{\hat{\mathcal{H}}}}_{\hat{\mathcal{H}}}-J^{\pi^*_{\hat{\mathcal{H}}}}_{\mathcal{H}}+\gamma^H\epsilon_v\\
% &\leq2\max_{\pi\in\{\pi^*_{\hat{\mathcal{H}}}, \pi^*_{\mathcal{H}}\}} |J^{\pi}_{\mathcal{H}}-J^{\pi}_{\hat{\mathcal{H}}}|+\epsilon_p+\gamma^H\epsilon_v
% \end{align*}
% where $\epsilon_p = |J^{\pi^*_{\mathcal{H}}}_{\hat{\mathcal{H}}}-J^{\pi^*_{\hat{\mathcal{H}}}}_{\hat{\mathcal{H}}}|$ is the suboptimality bound might cause from the the optimization of the H-step lookahead algorithm of equation 12. We can write our the specific discounted return in the following format $J_\mathcal{H}^\pi=\sum_{s,a}\sum_{t=0}^H \gamma^t\rho_1^t(s,a)R_{mix}(s,a)$, then we can derive the bound of $|J^{\pi}_{\mathcal{H}}-J^{\pi}_{\hat{\mathcal{H}}}|$.
% \begin{align*}
% |J^{\pi}_{\mathcal{H}}-J^{\pi}_{\hat{\mathcal{H}}}| &\leq \sum_{s,a}\sum_{t=0}^H |\gamma^t\left(\rho^t(s,a)R_{mix}(s,a)-\hat{\rho}^t(s,a)\hat{R}_{mix}(s,a)\right)| \\
% &= \sum_{s,a}\sum_{t=0}^H |\gamma^t\left(\rho^t(s,a)R_{mix}(s,a)-\hat{\rho}^t(s,a)R_{mix}(s,a)+\hat{\rho}^t(s,a)R_{mix}(s,a)-\hat{\rho}^t(s,a)\hat{R}_{mix}(s,a)\right)| \\
% &= \sum_{s,a}\sum_{t=0}^H |\gamma^t\left(\left(\rho^t(s,a)-\hat{\rho}^t(s,a)\right)R_{mix}(s,a)+\hat{\rho}^t(s,a)\left(R_{mix}(s,a)-\hat{R}_{mix}(s,a)\right)\right)|\\
% &\leq \gamma^H H\epsilon_{\mathcal{T}}\frac{R_\max}{1-\gamma}+\gamma^H\underbrace{|V^\pi_{\mathcal{M}'\cup R}(s_H)-V^\pi_{\hat{\mathcal{M}'}\cup\hat{R}}(s_H)|}_\text{already derived}+\sum_{t=0}^{H-1}\left[\gamma^t t\epsilon_{\mathcal{T}}R_{\max}+\gamma^t \underbrace{|R^E-\hat{R}^E|}_\text{already derived}\right]\\
% &\leq \frac{\gamma^H\epsilon_\mathcal{T}R_\max}{1-\gamma}\left(H - \frac{1+\gamma}{1-\gamma}|\mathcal{S}|\right)+\sum_{t=0}^{H-1}\gamma^t\epsilon_\mathcal{T}R_{\max}\left(t+\frac{\gamma|\mathcal{S}|}{1-\gamma}\right)
% \end{align*}
% By putting everything together, we can derive a final performance optimallity gap which correlate with $\epsilon_\mathcal{T}, \epsilon_p,$ and $\epsilon_v$.
% $$2\left[\frac{\gamma^H R_\max}{1-\gamma}\left(H - \frac{1+\gamma}{1-\gamma}|\mathcal{S}|\right)+\sum_{t=0}^{H-1}\gamma^t R_{\max}\left(t+\frac{\gamma|\mathcal{S}|}{1-\gamma}\right)\right]\epsilon_\mathcal{T}+\epsilon_p + \gamma^H\epsilon_v $$
% \end{proof}
% \begin{lemma}
% Given $\hat{V}$ is an approximated learned value function such that $\max_s|V^*(s)-\hat{V}(s)|\leq\epsilon_v$. Then, we can derive the following optimality gap.
% \begin{equation*}
% J^{\pi^*}-J^{\pi_{H,\hat{V}}}\leq\frac{2}{1-\gamma^H}\left[R_{\max}\left[\sum_{t=0}^{H-1}\gamma^t t\epsilon_{\mathcal{T}}+\frac{\gamma^H H\epsilon_\mathcal{T}}{1-\gamma}\right]+\frac{\epsilon_p}{2}+\gamma^H\epsilon_v\right]
% \end{equation*}
% where $\epsilon_p$ is the suboptimality incurred in H-step lookahead optimization such that $J^*-\hat{J}\leq\epsilon_p$ ($J^*$ is the optimal return for the H-step optimization and $\hat{J}$ is the result of the suboptimal H-step optimization.)
% \end{lemma}
\newpage
\section{Implementation Details}
\label{sec::implementation}
We used two identical 2-layer MLP with 100 hidden units and ReLU activiations for the reward term $r$ and the shaping term $\phi$. The neural nets are trained using Adam optimizer~\citep{kingma2014adam} with learning rate being $0.0003$. For the learned dynamic model, we also used a 2-layer MLP with ReLU activations and 128 hidden units. Gradient clipping is added on the model network with the value of 10. The model network is trained by the Adam optimizer as well with learning rate $0.0003$. For SAC and PPO, we follow the implementation from \texttt{CleanRL} repository~\citep{huang2022cleanrl}.
% Similar to DAC~\citep{kostrikov2018discriminator}, we added gradient penalties~\citep{gulrajani2017improved} as the regularization term to help the model to be more stable by preventing overfitting of the training data.
\newpage
\section{Graph Results}
\label{appendix:learning_curves}
Below is the testing return diagrams from both deterministic and stochastic Mujoco Environments
\begin{figure}[h]
\centering
\includegraphics[width=\linewidth]{figures/InvertedPendulum.png}
\caption{Training return diagram of benchmarks and our algorithm in the InvertedPendulum-v4. We only show the training return until reaching the experts performance of all the other benchmarks, but the thorough training return of our algorithm.}
\label{Fig::invertedpendulum}
\end{figure}
\begin{figure}[h]
\centering
\includegraphics[width=\linewidth]{figures/InvDblPen.png}
\caption{Training return diagram of benchmarks and our algorithm in the InvertedDoublePendulum-v4.}
\label{fig:inverteddoublependulum}
\end{figure}
\begin{figure}[h]
\centering
\includegraphics[width=\linewidth]{figures/Hopper.png}
\caption{Training return diagram of benchmarks and our algorithm in the Hopper-v3.}
\label{fig:hopper}
\end{figure}
% \begin{figure}
% \centering
% \includegraphics[width=\linewidth]{figures/Walker2d.png}
% \caption{Training return diagram of benchmarks and our algorithm in the Walker2d-v3.}
% \label{fig:walker2d}
% \end{figure}
\begin{thebibliography}{78}
\providecommand{\natexlab}[1]{#1}
\providecommand{\url}[1]{\texttt{#1}}
\expandafter\ifx\csname urlstyle\endcsname\relax
\providecommand{\doi}[1]{doi: #1}\else
\providecommand{\doi}{doi: \begingroup \urlstyle{rm}\Url}\fi
\bibitem[Abbeel \& Ng(2004)Abbeel and Ng]{abbeel2004apprenticeship}
Pieter Abbeel and Andrew~Y Ng.
\newblock Apprenticeship learning via inverse reinforcement learning.
\newblock In \emph{Proceedings of the twenty-first international conference on Machine learning}, pp.\ ~1, 2004.
\bibitem[Arora \& Doshi(2021)Arora and Doshi]{arora2021survey}
Saurabh Arora and Prashant Doshi.
\newblock A survey of inverse reinforcement learning: Challenges, methods and progress.
\newblock \emph{Artificial Intelligence}, 297:\penalty0 103500, 2021.
\bibitem[Baram et~al.(2016)Baram, Anschel, and Mannor]{baram2016model}
Nir Baram, Oron Anschel, and Shie Mannor.
\newblock Model-based adversarial imitation learning.
\newblock \emph{arXiv preprint arXiv:1612.02179}, 2016.
\bibitem[Baram et~al.(2017)Baram, Anschel, Caspi, and Mannor]{baram2017end}
Nir Baram, Oron Anschel, Itai Caspi, and Shie Mannor.
\newblock End-to-end differentiable adversarial imitation learning.
\newblock In \emph{International Conference on Machine Learning}, pp.\ 390--399. PMLR, 2017.
\bibitem[Berner et~al.(2019)Berner, Brockman, Chan, Cheung, D{\k{e}}biak, Dennison, Farhi, Fischer, Hashme, Hesse, et~al.]{berner2019dota}
Christopher Berner, Greg Brockman, Brooke Chan, Vicki Cheung, Przemys{\l}aw D{\k{e}}biak, Christy Dennison, David Farhi, Quirin Fischer, Shariq Hashme, Chris Hesse, et~al.
\newblock Dota 2 with large scale deep reinforcement learning.
\newblock \emph{arXiv preprint arXiv:1912.06680}, 2019.
\bibitem[Blond{\'e} \& Kalousis(2019)Blond{\'e} and Kalousis]{blonde2019sample}
Lionel Blond{\'e} and Alexandros Kalousis.
\newblock Sample-efficient imitation learning via generative adversarial nets.
\newblock In \emph{The 22nd International Conference on Artificial Intelligence and Statistics}, pp.\ 3138--3148. PMLR, 2019.
\bibitem[Blond{\'e} et~al.(2022)Blond{\'e}, Strasser, and Kalousis]{blonde2022lipschitzness}
Lionel Blond{\'e}, Pablo Strasser, and Alexandros Kalousis.
\newblock Lipschitzness is all you need to tame off-policy generative adversarial imitation learning.
\newblock \emph{Machine Learning}, 111\penalty0 (4):\penalty0 1431--1521, 2022.
\bibitem[Brown et~al.(2019)Brown, Goo, Nagarajan, and Niekum]{brown2019extrapolating}
Daniel Brown, Wonjoon Goo, Prabhat Nagarajan, and Scott Niekum.
\newblock Extrapolating beyond suboptimal demonstrations via inverse reinforcement learning from observations.
\newblock In \emph{International conference on machine learning}, pp.\ 783--792. PMLR, 2019.
\bibitem[Chan \& van~der Schaar(2021)Chan and van~der Schaar]{chan2021scalable}
Alex~J Chan and Mihaela van~der Schaar.
\newblock Scalable bayesian inverse reinforcement learning.
\newblock \emph{arXiv preprint arXiv:2102.06483}, 2021.
\bibitem[Choi \& Kim(2011)Choi and Kim]{choi2011map}
Jaedeug Choi and Kee-Eung Kim.
\newblock Map inference for bayesian inverse reinforcement learning.
\newblock \emph{Advances in neural information processing systems}, 24, 2011.
\bibitem[Choudhury et~al.(2018)Choudhury, Bhardwaj, Arora, Kapoor, Ranade, Scherer, and Dey]{choudhury2018data}
Sanjiban Choudhury, Mohak Bhardwaj, Sankalp Arora, Ashish Kapoor, Gireeja Ranade, Sebastian Scherer, and Debadeepta Dey.
\newblock Data-driven planning via imitation learning.
\newblock \emph{The International Journal of Robotics Research}, 37\penalty0 (13-14):\penalty0 1632--1672, 2018.
\bibitem[Codevilla et~al.(2018)Codevilla, M{\"u}ller, L{\'o}pez, Koltun, and Dosovitskiy]{codevilla2018end}
Felipe Codevilla, Matthias M{\"u}ller, Antonio L{\'o}pez, Vladlen Koltun, and Alexey Dosovitskiy.
\newblock End-to-end driving via conditional imitation learning.
\newblock In \emph{2018 IEEE international conference on robotics and automation (ICRA)}, pp.\ 4693--4700. IEEE, 2018.
\bibitem[Das et~al.(2021)Das, Bechtle, Davchev, Jayaraman, Rai, and Meier]{das2021model}
Neha Das, Sarah Bechtle, Todor Davchev, Dinesh Jayaraman, Akshara Rai, and Franziska Meier.
\newblock Model-based inverse reinforcement learning from visual demonstrations.
\newblock In \emph{Conference on Robot Learning}, pp.\ 1930--1942. PMLR, 2021.
\bibitem[Devidze et~al.(2022)Devidze, Kamalaruban, and Singla]{devidze2022exploration}
Rati Devidze, Parameswaran Kamalaruban, and Adish Singla.
\newblock Exploration-guided reward shaping for reinforcement learning under sparse rewards.
\newblock \emph{Advances in Neural Information Processing Systems}, 35:\penalty0 5829--5842, 2022.
\bibitem[Devlin \& Kudenko(2011)Devlin and Kudenko]{devlin2011theoretical}
Sam Devlin and Daniel Kudenko.
\newblock Theoretical considerations of potential-based reward shaping for multi-agent systems.
\newblock In \emph{Tenth International Conference on Autonomous Agents and Multi-Agent Systems}, pp.\ 225--232. ACM, 2011.
\bibitem[Devlin \& Kudenko(2012)Devlin and Kudenko]{devlin2012dynamic}
Sam~Michael Devlin and Daniel Kudenko.
\newblock Dynamic potential-based reward shaping.
\newblock In \emph{11th International Conference on Autonomous Agents and Multiagent Systems (AAMAS 2012)}, pp.\ 433--440. IFAAMAS, 2012.
\bibitem[Dorigo \& Colombetti(1994)Dorigo and Colombetti]{dorigo1994robot}
Marco Dorigo and Marco Colombetti.
\newblock Robot shaping: Developing autonomous agents through learning.
\newblock \emph{Artificial intelligence}, 71\penalty0 (2):\penalty0 321--370, 1994.
\bibitem[Finn et~al.(2016{\natexlab{a}})Finn, Christiano, Abbeel, and Levine]{finn2016connection}
Chelsea Finn, Paul Christiano, Pieter Abbeel, and Sergey Levine.
\newblock A connection between generative adversarial networks, inverse reinforcement learning, and energy-based models.
\newblock \emph{arXiv preprint arXiv:1611.03852}, 2016{\natexlab{a}}.
\bibitem[Finn et~al.(2016{\natexlab{b}})Finn, Levine, and Abbeel]{finn2016guided}
Chelsea Finn, Sergey Levine, and Pieter Abbeel.
\newblock Guided cost learning: Deep inverse optimal control via policy optimization.
\newblock In \emph{International conference on machine learning}, pp.\ 49--58. PMLR, 2016{\natexlab{b}}.
\bibitem[Fu et~al.(2017)Fu, Luo, and Levine]{fu2017learning}
Justin Fu, Katie Luo, and Sergey Levine.
\newblock Learning robust rewards with adversarial inverse reinforcement learning.
\newblock \emph{arXiv preprint arXiv:1710.11248}, 2017.
\bibitem[Goodfellow et~al.(2014)Goodfellow, Pouget-Abadie, Mirza, Xu, Warde-Farley, Ozair, Courville, and Bengio]{goodfellow2014generative}
Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio.
\newblock Generative adversarial nets.
\newblock \emph{Advances in neural information processing systems}, 27, 2014.
\bibitem[Gupta et~al.(2022)Gupta, Pacchiano, Zhai, Kakade, and Levine]{gupta2022unpacking}
Abhishek Gupta, Aldo Pacchiano, Yuexiang Zhai, Sham Kakade, and Sergey Levine.
\newblock Unpacking reward shaping: Understanding the benefits of reward engineering on sample complexity.
\newblock \emph{Advances in Neural Information Processing Systems}, 35:\penalty0 15281--15295, 2022.
\bibitem[Haarnoja et~al.(2018)Haarnoja, Zhou, Abbeel, and Levine]{haarnoja2018soft}
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine.
\newblock Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor.
\newblock In \emph{International conference on machine learning}, pp.\ 1861--1870. PMLR, 2018.
\bibitem[Hafner et~al.(2019)Hafner, Lillicrap, Ba, and Norouzi]{hafner2019dream}
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi.
\newblock Dream to control: Learning behaviors by latent imagination.
\newblock \emph{arXiv preprint arXiv:1912.01603}, 2019.
\bibitem[Hansen et~al.(2022)Hansen, Wang, and Su]{hansen2022temporal}
Nicklas Hansen, Xiaolong Wang, and Hao Su.
\newblock Temporal difference learning for model predictive control.
\newblock \emph{arXiv preprint arXiv:2203.04955}, 2022.
\bibitem[Harutyunyan et~al.(2015)Harutyunyan, Devlin, Vrancx, and Now{\'e}]{harutyunyan2015expressing}
Anna Harutyunyan, Sam Devlin, Peter Vrancx, and Ann Now{\'e}.
\newblock Expressing arbitrary reward functions as potential-based advice.
\newblock In \emph{Proceedings of the AAAI conference on artificial intelligence}, volume~29, 2015.
\bibitem[Herman et~al.(2016)Herman, Gindele, Wagner, Schmitt, and Burgard]{herman2016inverse}
Michael Herman, Tobias Gindele, J{\"o}rg Wagner, Felix Schmitt, and Wolfram Burgard.
\newblock Inverse reinforcement learning with simultaneous estimation of rewards and dynamics.
\newblock In \emph{Artificial intelligence and statistics}, pp.\ 102--110. PMLR, 2016.
\bibitem[Ho \& Ermon(2016)Ho and Ermon]{ho2016generative}
Jonathan Ho and Stefano Ermon.
\newblock Generative adversarial imitation learning.
\newblock \emph{Advances in neural information processing systems}, 29, 2016.
\bibitem[Hu et~al.(2020)Hu, Wang, Jia, Wang, Chen, Hao, Wu, and Fan]{hu2020learning}
Yujing Hu, Weixun Wang, Hangtian Jia, Yixiang Wang, Yingfeng Chen, Jianye Hao, Feng Wu, and Changjie Fan.
\newblock Learning to utilize shaping rewards: A new approach of reward shaping.
\newblock \emph{Advances in Neural Information Processing Systems}, 33:\penalty0 15931--15941, 2020.
\bibitem[Huang et~al.(2022)Huang, Dossa, Ye, Braga, Chakraborty, Mehta, and Ara{\~A}{\v{s}}jo]{huang2022cleanrl}
Shengyi Huang, Rousslan Fernand~Julien Dossa, Chang Ye, Jeff Braga, Dipam Chakraborty, Kinal Mehta, and Jo{\~A}{\c{G}}o~GM Ara{\~A}{\v{s}}jo.
\newblock Cleanrl: High-quality single-file implementations of deep reinforcement learning algorithms.
\newblock \emph{Journal of Machine Learning Research}, 23\penalty0 (274):\penalty0 1--18, 2022.
\bibitem[Janner et~al.(2019)Janner, Fu, Zhang, and Levine]{janner2019trust}
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine.
\newblock When to trust your model: Model-based policy optimization.
\newblock \emph{Advances in neural information processing systems}, 32, 2019.
\bibitem[Kingma(2014)]{kingma2014adam}
Diederik~P Kingma.
\newblock Adam: A method for stochastic optimization.
\newblock \emph{arXiv preprint arXiv:1412.6980}, 2014.
\bibitem[Klein et~al.(2012)Klein, Geist, Piot, and Pietquin]{klein2012inverse}
Edouard Klein, Matthieu Geist, Bilal Piot, and Olivier Pietquin.
\newblock Inverse reinforcement learning through structured classification.
\newblock \emph{Advances in neural information processing systems}, 25, 2012.
\bibitem[Kostrikov et~al.(2018)Kostrikov, Agrawal, Dwibedi, Levine, and Tompson]{kostrikov2018discriminator}
Ilya Kostrikov, Kumar~Krishna Agrawal, Debidatta Dwibedi, Sergey Levine, and Jonathan Tompson.
\newblock Discriminator-actor-critic: Addressing sample inefficiency and reward bias in adversarial imitation learning.
\newblock \emph{arXiv preprint arXiv:1809.02925}, 2018.
\bibitem[Lee et~al.(2020)Lee, Seo, Lee, Abbeel, and Shin]{lee2020addressing}
Seunghyun Lee, Younggyo Seo, Kimin Lee, Pieter Abbeel, and Jinwoo Shin.
\newblock Addressing distribution shift in online reinforcement learning with offline datasets.
\newblock 2020.
\bibitem[Levine et~al.(2010)Levine, Popovic, and Koltun]{levine2010feature}
Sergey Levine, Zoran Popovic, and Vladlen Koltun.
\newblock Feature construction for inverse reinforcement learning.
\newblock \emph{Advances in neural information processing systems}, 23, 2010.
\bibitem[Levine et~al.(2011)Levine, Popovic, and Koltun]{levine2011nonlinear}
Sergey Levine, Zoran Popovic, and Vladlen Koltun.
\newblock Nonlinear inverse reinforcement learning with gaussian processes.
\newblock \emph{Advances in neural information processing systems}, 24, 2011.
\bibitem[Lin et~al.(2020)Lin, Thomas, Yang, and Ma]{lin2020model}
Zichuan Lin, Garrett Thomas, Guangwen Yang, and Tengyu Ma.
\newblock Model-based adversarial meta-reinforcement learning.
\newblock \emph{Advances in Neural Information Processing Systems}, 33:\penalty0 10161--10173, 2020.
\bibitem[Memarian et~al.(2021)Memarian, Goo, Lioutikov, Niekum, and Topcu]{memarian2021self}
Farzan Memarian, Wonjoon Goo, Rudolf Lioutikov, Scott Niekum, and Ufuk Topcu.
\newblock Self-supervised online reward shaping in sparse-reward environments.
\newblock In \emph{2021 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)}, pp.\ 2369--2375. IEEE, 2021.
\bibitem[Metelli et~al.(2021)Metelli, Ramponi, Concetti, and Restelli]{metelli2021provably}
Alberto~Maria Metelli, Giorgia Ramponi, Alessandro Concetti, and Marcello Restelli.
\newblock Provably efficient learning of transferable rewards.
\newblock In \emph{International Conference on Machine Learning}, pp.\ 7665--7676. PMLR, 2021.
\bibitem[Ng et~al.(2000)Ng, Russell, et~al.]{ng2000algorithms}
Andrew~Y Ng, Stuart Russell, et~al.
\newblock Algorithms for inverse reinforcement learning.
\newblock In \emph{Icml}, volume~1, pp.\ ~2, 2000.
\bibitem[Orsini et~al.(2021)Orsini, Raichuk, Hussenot, Vincent, Dadashi, Girgin, Geist, Bachem, Pietquin, and Andrychowicz]{orsini2021matters}
Manu Orsini, Anton Raichuk, L{\'e}onard Hussenot, Damien Vincent, Robert Dadashi, Sertan Girgin, Matthieu Geist, Olivier Bachem, Olivier Pietquin, and Marcin Andrychowicz.
\newblock What matters for adversarial imitation learning?
\newblock \emph{Advances in Neural Information Processing Systems}, 34:\penalty0 14656--14668, 2021.
\bibitem[Peng et~al.(2020)Peng, Coumans, Zhang, Lee, Tan, and Levine]{peng2020learning}
Xue~Bin Peng, Erwin Coumans, Tingnan Zhang, Tsang-Wei Lee, Jie Tan, and Sergey Levine.
\newblock Learning agile robotic locomotion skills by imitating animals.
\newblock \emph{arXiv preprint arXiv:2004.00784}, 2020.
\bibitem[Puterman(2014)]{puterman2014markov}
Martin~L Puterman.
\newblock \emph{Markov decision processes: discrete stochastic dynamic programming}.
\newblock John Wiley \& Sons, 2014.
\bibitem[Rafailov et~al.(2021)Rafailov, Yu, Rajeswaran, and Finn]{rafailov2021visual}
Rafael Rafailov, Tianhe Yu, Aravind Rajeswaran, and Chelsea Finn.
\newblock Visual adversarial imitation learning using variational models.
\newblock \emph{Advances in Neural Information Processing Systems}, 34:\penalty0 3016--3028, 2021.
\bibitem[Ramachandran \& Amir(2007)Ramachandran and Amir]{ramachandran2007bayesian}
Deepak Ramachandran and Eyal Amir.
\newblock Bayesian inverse reinforcement learning.
\newblock In \emph{IJCAI}, volume~7, pp.\ 2586--2591, 2007.
\bibitem[Randl{\o}v \& Alstr{\o}m(1998)Randl{\o}v and Alstr{\o}m]{randlov1998learning}
Jette Randl{\o}v and Preben Alstr{\o}m.
\newblock Learning to drive a bicycle using reinforcement learning and shaping.
\newblock In \emph{ICML}, volume~98, pp.\ 463--471, 1998.
\bibitem[Ratliff et~al.(2006)Ratliff, Bagnell, and Zinkevich]{ratliff2006maximum}
Nathan~D Ratliff, J~Andrew Bagnell, and Martin~A Zinkevich.
\newblock Maximum margin planning.
\newblock In \emph{Proceedings of the 23rd international conference on Machine learning}, pp.\ 729--736, 2006.
\bibitem[Ratliff et~al.(2009)Ratliff, Silver, and Bagnell]{ratliff2009learning}
Nathan~D Ratliff, David Silver, and J~Andrew Bagnell.
\newblock Learning to search: Functional gradient techniques for imitation learning.
\newblock \emph{Autonomous Robots}, 27:\penalty0 25--53, 2009.
\bibitem[Reddy et~al.(2019)Reddy, Dragan, and Levine]{reddy2019sqil}
Siddharth Reddy, Anca~D Dragan, and Sergey Levine.
\newblock Sqil: Imitation learning via reinforcement learning with sparse rewards.
\newblock \emph{arXiv preprint arXiv:1905.11108}, 2019.
\bibitem[Reuss et~al.(2023)Reuss, Li, Jia, and Lioutikov]{reuss2023goal}
Moritz Reuss, Maximilian Li, Xiaogang Jia, and Rudolf Lioutikov.
\newblock Goal-conditioned imitation learning using score-based diffusion policies.
\newblock \emph{arXiv preprint arXiv:2304.02532}, 2023.
\bibitem[Russell(1998)]{russell1998learning}
Stuart Russell.
\newblock Learning agents for uncertain environments.
\newblock In \emph{Proceedings of the eleventh annual conference on Computational learning theory}, pp.\ 101--103, 1998.
\bibitem[Schrittwieser et~al.(2020)Schrittwieser, Antonoglou, Hubert, Simonyan, Sifre, Schmitt, Guez, Lockhart, Hassabis, Graepel, et~al.]{schrittwieser2020mastering}
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et~al.
\newblock Mastering atari, go, chess and shogi by planning with a learned model.
\newblock \emph{Nature}, 588\penalty0 (7839):\penalty0 604--609, 2020.
\bibitem[Schulman et~al.(2017)Schulman, Wolski, Dhariwal, Radford, and Klimov]{schulman2017proximal}
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov.
\newblock Proximal policy optimization algorithms.
\newblock \emph{arXiv preprint arXiv:1707.06347}, 2017.
\bibitem[Shalev-Shwartz \& Ben-David(2014)Shalev-Shwartz and Ben-David]{shalev2014understanding}
Shai Shalev-Shwartz and Shai Ben-David.
\newblock \emph{Understanding machine learning: From theory to algorithms}.
\newblock Cambridge university press, 2014.
\bibitem[Shore \& Johnson(1980)Shore and Johnson]{shore1980axiomatic}
John Shore and Rodney Johnson.
\newblock Axiomatic derivation of the principle of maximum entropy and the principle of minimum cross-entropy.
\newblock \emph{IEEE Transactions on information theory}, 26\penalty0 (1):\penalty0 26--37, 1980.
\bibitem[Sikchi et~al.(2022)Sikchi, Zhou, and Held]{sikchi2022learning}
Harshit Sikchi, Wenxuan Zhou, and David Held.
\newblock Learning off-policy with online planning.
\newblock In \emph{Conference on Robot Learning}, pp.\ 1622--1633. PMLR, 2022.
\bibitem[Skalse et~al.(2023)Skalse, Farrugia-Roberts, Russell, Abate, and Gleave]{skalse2023invariance}
Joar Max~Viktor Skalse, Matthew Farrugia-Roberts, Stuart Russell, Alessandro Abate, and Adam Gleave.
\newblock Invariance in policy optimisation and partial identifiability in reward learning.
\newblock In \emph{International Conference on Machine Learning}, pp.\ 32033--32058. PMLR, 2023.
\bibitem[Sun et~al.(2021)Sun, Yu, Dong, Lu, and Zhou]{sun2021adversarial}
Jiankai Sun, Lantao Yu, Pinqian Dong, Bo~Lu, and Bolei Zhou.
\newblock Adversarial inverse reinforcement learning with self-attention dynamics model.
\newblock \emph{IEEE Robotics and Automation Letters}, 6\penalty0 (2):\penalty0 1880--1886, 2021.
\bibitem[Sun et~al.(2018)Sun, Peng, Zhan, and Tomizuka]{sun2018fast}
Liting Sun, Cheng Peng, Wei Zhan, and Masayoshi Tomizuka.
\newblock A fast integrated planning and control framework for autonomous driving via imitation learning.
\newblock In \emph{Dynamic Systems and Control Conference}, volume 51913, pp.\ V003T37A012. American Society of Mechanical Engineers, 2018.
\bibitem[Swamy et~al.(2021)Swamy, Choudhury, Bagnell, and Wu]{swamy2021moments}
Gokul Swamy, Sanjiban Choudhury, J~Andrew Bagnell, and Steven Wu.
\newblock Of moments and matching: A game-theoretic framework for closing the imitation gap.
\newblock In \emph{International Conference on Machine Learning}, pp.\ 10022--10032. PMLR, 2021.
\bibitem[Todorov et~al.(2012)Todorov, Erez, and Tassa]{mujoco}
Emanuel Todorov, Tom Erez, and Yuval Tassa.
\newblock Mujoco: A physics engine for model-based control.
\newblock In \emph{2012 IEEE/RSJ international conference on intelligent robots and systems}, pp.\ 5026--5033. IEEE, 2012.
\bibitem[Torabi et~al.(2018{\natexlab{a}})Torabi, Warnell, and Stone]{torabi2018behavioral}
Faraz Torabi, Garrett Warnell, and Peter Stone.
\newblock Behavioral cloning from observation.
\newblock \emph{arXiv preprint arXiv:1805.01954}, 2018{\natexlab{a}}.
\bibitem[Torabi et~al.(2018{\natexlab{b}})Torabi, Warnell, and Stone]{torabi2018generative}
Faraz Torabi, Garrett Warnell, and Peter Stone.
\newblock Generative adversarial imitation from observation.
\newblock \emph{arXiv preprint arXiv:1807.06158}, 2018{\natexlab{b}}.
\bibitem[Viano et~al.(2021)Viano, Huang, Kamalaruban, Weller, and Cevher]{viano2021robust}
Luca Viano, Yu-Ting Huang, Parameswaran Kamalaruban, Adrian Weller, and Volkan Cevher.
\newblock Robust inverse reinforcement learning under transition dynamics mismatch.
\newblock \emph{Advances in Neural Information Processing Systems}, 34:\penalty0 25917--25931, 2021.
\bibitem[Wang et~al.(2023{\natexlab{a}})Wang, Jiao, Lang, Zhan, Huang, Wang, Yang, and Zhu]{wang2023empowering}
Yixuan Wang, Ruochen Jiao, Chengtian Lang, Sinong~Simon Zhan, Chao Huang, Zhaoran Wang, Zhuoran Yang, and Qi~Zhu.
\newblock Empowering autonomous driving with large language models: A safety perspective.
\newblock \emph{arXiv preprint arXiv:2312.00812}, 2023{\natexlab{a}}.
\bibitem[Wang et~al.(2023{\natexlab{b}})Wang, Zhan, Wang, Huang, Wang, Yang, and Zhu]{wang2023joint}
Yixuan Wang, Simon Zhan, Zhilu Wang, Chao Huang, Zhaoran Wang, Zhuoran Yang, and Qi~Zhu.
\newblock Joint differentiable optimization and verification for certified reinforcement learning.
\newblock In \emph{Proceedings of the ACM/IEEE 14th International Conference on Cyber-Physical Systems (with CPS-IoT Week 2023)}, pp.\ 132--141, 2023{\natexlab{b}}.
\bibitem[Wang et~al.(2023{\natexlab{c}})Wang, Zhan, Jiao, Wang, Jin, Yang, Wang, Huang, and Zhu]{wang2023enforcing}
Yixuan Wang, Simon~Sinong Zhan, Ruochen Jiao, Zhilu Wang, Wanxin Jin, Zhuoran Yang, Zhaoran Wang, Chao Huang, and Qi~Zhu.
\newblock Enforcing hard constraints with soft barriers: Safe reinforcement learning in unknown stochastic environments.
\newblock In \emph{International Conference on Machine Learning}, pp.\ 36593--36604. PMLR, 2023{\natexlab{c}}.
\bibitem[Wiewiora et~al.(2003)Wiewiora, Cottrell, and Elkan]{wiewiora2003principled}
Eric Wiewiora, Garrison~W Cottrell, and Charles Elkan.
\newblock Principled methods for advising reinforcement learning agents.
\newblock In \emph{Proceedings of the 20th international conference on machine learning (ICML-03)}, pp.\ 792--799, 2003.
\bibitem[Wu et~al.()Wu, Zhan, Wang, Wang, Lin, Lv, Zhu, Schmidhuber, and Huang]{wuboosting}
Qingyuan Wu, Simon~Sinong Zhan, Yixuan Wang, Yuhui Wang, Chung-Wei Lin, Chen Lv, Qi~Zhu, J{\"u}rgen Schmidhuber, and Chao Huang.
\newblock Boosting reinforcement learning with strongly delayed feedback through auxiliary short delays.
\newblock In \emph{Forty-first International Conference on Machine Learning}.
\bibitem[Wu et~al.(2024)Wu, Zhan, Wang, Wang, Lin, Lv, Zhu, and Huang]{wu2024variational}
Qingyuan Wu, Simon~Sinong Zhan, Yixuan Wang, Yuhui Wang, Chung-Wei Lin, Chen Lv, Qi~Zhu, and Chao Huang.
\newblock Variational delayed policy optimization.
\newblock \emph{arXiv preprint arXiv:2405.14226}, 2024.
\bibitem[Wulfmeier et~al.(2015)Wulfmeier, Ondruska, and Posner]{wulfmeier2015maximum}
Markus Wulfmeier, Peter Ondruska, and Ingmar Posner.
\newblock Maximum entropy deep inverse reinforcement learning.
\newblock \emph{arXiv preprint arXiv:1507.04888}, 2015.
\bibitem[Yin et~al.(2022)Yin, Ye, Chen, and Gao]{yin2022planning}
Zhao-Heng Yin, Weirui Ye, Qifeng Chen, and Yang Gao.
\newblock Planning for sample efficient imitation learning.
\newblock \emph{Advances in Neural Information Processing Systems}, 35:\penalty0 2577--2589, 2022.
\bibitem[Yu et~al.(2020)Yu, Thomas, Yu, Ermon, Zou, Levine, Finn, and Ma]{yu2020mopo}
Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James~Y Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma.
\newblock Mopo: Model-based offline policy optimization.
\newblock \emph{Advances in Neural Information Processing Systems}, 33:\penalty0 14129--14142, 2020.
\bibitem[Zhan et~al.(2023)Zhan, Wang, Wu, Jiao, Huang, and Zhu]{zhan2023state}
Simon~Sinong Zhan, Yixuan Wang, Qingyuan Wu, Ruochen Jiao, Chao Huang, and Qi~Zhu.
\newblock State-wise safe reinforcement learning with pixel observations.
\newblock \emph{arXiv preprint arXiv:2311.02227}, 2023.
\bibitem[Zhang et~al.(2020)Zhang, Sodhani, Khetarpal, and Pineau]{zhang2020learning}
Amy Zhang, Shagun Sodhani, Khimya Khetarpal, and Joelle Pineau.
\newblock Learning robust state abstractions for hidden-parameter block mdps.
\newblock \emph{arXiv preprint arXiv:2007.07206}, 2020.
\bibitem[Ziebart et~al.(2008)Ziebart, Maas, Bagnell, Dey, et~al.]{ziebart2008maximum}
Brian~D Ziebart, Andrew~L Maas, J~Andrew Bagnell, Anind~K Dey, et~al.
\newblock Maximum entropy inverse reinforcement learning.
\newblock In \emph{Aaai}, volume~8, pp.\ 1433--1438. Chicago, IL, USA, 2008.
\bibitem[Ziebart et~al.(2010)Ziebart, Bagnell, and Dey]{ziebart2010modeling}
Brian~D Ziebart, J~Andrew Bagnell, and Anind~K Dey.
\newblock Modeling interaction via the principle of maximum causal entropy.
\newblock 2010.
\end{thebibliography}