SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
L \left(\theta^ {*}\right) - L (\hat {\theta}) \leq \varepsilon . \tag {11}
$$
4 Algorithm design
In the following sections, we will design a computationally efficient algorithm to optimize  $\widehat{L}(\cdot)$ , and obtain its corresponding optimal reward parameter  $\hat{\theta}$ .
From the definition (7), it is clear that  $\widehat{L} (\cdot)$  depends on the optimal soft value function  $V_{\theta}(\cdot)$  in (3a), which in turn depends on the optimal policy  $\pi_{\theta}$  as defined in (2b). Therefore, the surrogate objective maximization problem can be formulated as a bi-level optimization problem, expressed below, where the upper-level problem optimizes  $\widehat{L} (\cdot)$  to search for a good reward estimate  $\theta$  , while the lower-level problem solves the optimal policy  $\pi_{\theta}$  in a conservative MDP under the current reward estimate:
$$
\max  _ {\theta} \widehat {L} (\theta), \quad \text {s . t .} \quad \pi_ {\theta} := \arg \max  _ {\pi} \mathbb {E} _ {\tau \mathrm {A} \sim (\eta , \pi , \widehat {P})} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} \left(r \left(s _ {t}, a _ {t}; \theta\right) + U \left(s _ {t}, a _ {t}\right) + \mathcal {H} \left(\pi \left(\cdot \mid s _ {t}\right)\right)\right) \right]. \tag {12}
$$
In order to avoid the computational burden from repeatedly solving the optimal policy  $\pi_{\theta}$  under each reward estimate  $\theta$ , we aim to design an algorithm which alternates between a policy optimization step and a reward update step. That is, at each iteration  $k$ , based on the current policy estimate  $\pi_{k}$  and the reward parameter  $\theta_{k}$ , two steps will be performed consecutively: (1) the algorithm generates an updated policy  $\pi_{k+1}$  through performing a conservative policy improvement step under the estimated world model  $\widehat{P}$ , and (2) it obtains an updated reward parameter  $\theta_{k+1}$  through taking a reward update step. Next, we describe the proposed algorithm in detail.
Policy Improvement Step. Under the reward parameter  $\theta_{k}$ , we consider generating a new policy  $\pi_{k+1}$  towards approaching the optimal policy  $\pi_{\theta_k}$  as defined in (12). Similar to the definitions of  $V_{\theta}$  and  $Q_{\theta}$  in (3a) - (3b), under the current policy estimate  $\pi_{k}$ , the reward estimate  $r(\cdot, \cdot; \theta_k)$  and the estimated world model  $\widehat{P}$ , we define the corresponding soft value function as  $V_{k}(\cdot)$  and the soft Q-function as  $Q_{k}(\cdot, \cdot)$ . Please see (28a) - (28b) in Appendix for the precise definitions.
In order to perform a policy improvement step, we first approximate the soft Q-function by using an estimate  $\widehat{Q}_k(s,a)$ , which satisfies the following: (where  $\epsilon_{\mathrm{app}} > 0$  is an approximation error)
$$
\left\| \widehat {Q} _ {k} - Q _ {k} \right\| _ {\infty} := \max  _ {s \in \mathcal {S}, a \in \mathcal {A}} \left| \widehat {Q} _ {k} (s, a) - Q _ {k} (s, a) \right| \leq \epsilon_ {\text {a p p}}. \tag {13}
$$
With the approximator  $\widehat{Q}_k$ , an updated policy  $\pi_{k + 1}$  can be generated by a soft policy iteration:
$$
\pi_ {k + 1} (a | s) \propto \exp (\widehat {Q} _ {k} (s, a)), \quad \forall s \in \mathcal {S}, a \in \mathcal {A}. \tag {14}
$$
As a remark, in practice, one can follow the popular reinforcement learning algorithms such as soft Q-learning [38] and soft Actor-Critic (SAC) [39] to obtain accurately approximated soft Q-function with low approximation error  $\epsilon_{\mathrm{app}}$  (as outlined in (13)), so to achieve stable updates for the soft
Algorithm 1 A Model-based Approach for Offline Maximum Likelihood IRL (Offline ML-IRL)
Input: Initialize reward parameter  $\theta_0$  and policy  $\pi_0$ . Set the reward parameter's stepsize as  $\alpha$ . Train the world model  $\widehat{P}$  on the transition dataset  $\mathcal{D}$ . Specify the penalty function  $U(\cdot, \cdot)$  based on  $\widehat{P}$ .  
for  $k = 0,1,\ldots,K-1$  do  
    Policy Evaluation: Approximate the soft Q-function  $Q_k(\cdot, \cdot)$  by  $\widehat{Q}_k(\cdot, \cdot)$   
    Policy Improvement:  $\pi_{k+1}(\cdot|s) \propto \exp(\widehat{Q}_k(s, \cdot)), \forall s \in S$   
    Data Sampling I: Sample an expert trajectory  $\tau_k^{\mathrm{E}} := \{s_t, a_t\}_{t \geq 0}$   
    Data Sampling II: Sample  $\tau_k^{\mathrm{A}} := \{s_t, a_t\}_{t \geq 0}$  from  $\pi_{k+1}$  and  $\widehat{P}$   
    Estimating Gradient:  $g_k := h(\theta_k; \tau_k^{\mathrm{E}}) - h(\theta_k; \tau_k^{\mathrm{A}})$  where  $h(\theta; \tau) := \sum_{t \geq 0} \gamma^t \nabla_\theta r(s_t, a_t; \theta)$   
    Reward Parameter Update:  $\theta_{k+1} := \theta_k + \alpha g_k$   
end for
policy iteration in (14). In the literature of the model-based offline reinforcement learning, the state-of-the-art methods [25, 37, 47] also build their framework upon the implementation of SAC.
Reward Optimization Step. At each iteration  $k$ , given the current reward parameter  $\theta_{k}$  and the updated policy  $\pi_{k+1}$ , we can update the reward parameter to  $\theta_{k+1}$ . First, let us compute the gradient of the surrogate objective  $\nabla \widehat{L}(\theta_{k})$ . Please see Appendix F for the detailed proof.
Lemma 3. The gradient of the surrogate objective  $\widehat{L} (\theta)$  defined in (7), can be expressed as:
$$
\nabla \widehat {L} (\theta) = \mathbb {E} _ {\tau^ {\mathrm {E}} \sim (\eta , \pi^ {\mathrm {E}}, P)} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} \nabla_ {\theta} r \left(s _ {t}, a _ {t}; \theta\right) \right] - \mathbb {E} _ {\tau^ {\mathrm {A}} \sim (\eta , \pi_ {\theta}, \widehat {P})} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} \nabla_ {\theta} r \left(s _ {t}, a _ {t}; \theta\right) \right]. \tag {15}
$$
In practice, we do not have access to the optimal policy  $\pi_{\theta}$ . This is due to the fact that repeatedly solving the underlying offline policy optimization problem under each reward parameter is computationally intractable. Therefore, at each iteration  $k$ , we construct an estimator of the exact gradient based on the current policy estimate  $\pi_{k+1}$ .
To be more specific, we take two approximation steps to develop a stochastic gradient estimator of  $\nabla \widehat{L}(\theta)$ : 1) choose one observed expert trajectory  $\tau_k^{\mathrm{E}}$ ; 2) sample a trajectory  $\tau_k^{\mathrm{A}}$  from the current policy estimate  $\pi_{k+1}$  in the estimated world model  $\widehat{P}$ . Following these two approximation steps, the stochastic estimator  $g_k$  which approximates the exact gradient  $\nabla \widehat{L}(\theta_k)$  in (15) is defined as follows:
$$
g _ {k} := h \left(\theta_ {k}; \tau_ {k} ^ {\mathrm {E}}\right) - h \left(\theta_ {k}; \tau_ {k} ^ {\mathrm {A}}\right), \tag {16}
$$
where  $h(\theta ;\tau)\coloneqq \sum_{t = 0}^{\infty}\gamma^{t}\nabla_{\theta}r(s_{t},a_{t};\theta)$  denotes the cumulative reward gradient under a trajectory.
Then we can update reward parameter according to the following update rule:
$$
\theta_ {k + 1} = \theta_ {k} + \alpha g _ {k} \tag {17}
$$
In Alg. 1, we summarize the proposed algorithm (named the Offline ML-IRL).
5 Convergence analysis
In this section, we present a theoretical analysis to show the finite-time convergence of Alg.1.
Before starting the analysis, let us point out the key challenges in analyzing the Alg. 1. Note that the algorithm relies on the updated policy  $\pi_{k + 1}$  to approximate the optimal policy  $\pi_{\theta_k}$  at each iteration  $k$ . This coarse approximation can potentially lead to the distribution mismatch between the gradient estimator  $g_{k}$  in (16) and the exact gradient  $\nabla \widehat{L} (\theta_k)$  in (15). To maintain the stability of the proposed algorithm, we can use a relatively small stepsize  $\alpha$  to ensure the policy estimates are updated in a faster time-scale compared with the reward parameter  $\theta$ . This allows the policy estimates  $\{\pi_{k + 1}\}_{k\geq 0}$  to closely track the optimal solutions  $\{\pi_{\theta_k}\}_{k\geq 0}$  in the long run.
To proceed, let us introduce a few assumptions.
Assumption 2 (Ergodic Dynamics). Given any policy  $\pi$ , the Markov chain under the estimated world model  $\widehat{P}$  is irreducible and aperiodic. There exist constants  $\kappa > 0$  and  $\rho \in (0,1)$  to ensure:
$$
\max _ {s \in \mathcal {S}} \| \widehat {P} (s _ {t} \in \cdot | s _ {0} = s, \pi) - \mu_ {\widehat {P}} ^ {\pi} (\cdot) \| _ {\mathrm {T V}} \leq \kappa \rho^ {t}, \quad \forall t \geq 0
$$
where  $\| \cdot \|_{\mathrm{TV}}$  denotes the total variation (TV) norm;  $\mu_{\widehat{P}}^{\pi}$  is the stationary distribution of visited states under the policy  $\pi$  and the world model  $\widehat{P}$ .
The assumption about the ergodic dynamics is common in the literature of reinforcement learning [48, 49, 21, 31, 50], which ensures the Markov chain mixes at a geometric rate.
Assumption 3 (Lipschitz Reward). Under any reward parameter  $\theta$ , the following conditions hold for any  $s \in S$  and  $a \in \mathcal{A}$ :
$$
\left\| \nabla_ {\theta} r (s, a; \theta) \right\| \leq L _ {r}, \quad \left\| \nabla_ {\theta} r (s, a; \theta_ {1}) - \nabla_ {\theta} r (s, a; \theta_ {2}) \right\| \leq L _ {g} \| \theta_ {1} - \theta_ {2} \|, \tag {18}
$$
where  $L_{r}$  and  $L_{g}$  are positive constants.
According to Assumption 3, the parameterized reward has bounded gradient and is Lipschitz smooth. This assumption is common for min-max / bi-level optimization problems [41, 51, 52, 31, 50].
Based on Assumptions 2 - 3, we show that certain Lipschitz properties of the optimal soft Q-function and the surrogate objective hold. Please see the detailed proof in Appendix G.
Lemma 4. Suppose Assumptions 2 - 3 hold. Under any reward parameter  $\theta_{1}$  and  $\theta_{2}$ , the optimal soft  $Q$ -function and the surrogate objective satisfy the following Lipschitz properties:
$$
\left| Q _ {\theta_ {1}} (s, a) - Q _ {\theta_ {2}} (s, a) \right| \leq L _ {q} \| \theta_ {1} - \theta_ {2} \|, \forall s \in \mathcal {S}, a \in \mathcal {A} \tag {19a}
$$
$$
\left\| \nabla \widehat {L} \left(\theta_ {1}\right) - \nabla \widehat {L} \left(\theta_ {2}\right) \right\| \leq L _ {c} \| \theta_ {1} - \theta_ {2} \| \tag {19b}
$$
where  $L_{q}$  and  $L_{c}$  are positive constants.
Our main convergence result is summarized in the following theorem, which characterizes the convergence speed of the estimates  $\{\pi_{k + 1}\}_{k\geq 0}$  and  $\{\theta_k\}_{k\geq 0}$ . The detailed proof is in Appendix J.
Theorem 2 (Convergence Analysis). Suppose Assumptions 2 - 3 hold. Let  $K$  denote the total number of iterations to be run in Alg. 1. Setting the stepsize as  $\alpha = \alpha_0 \cdot K^{-\frac{1}{2}}$  where  $\alpha_0 > 0$ , we obtain the following convergence results:
$$
\frac {1}{K} \sum_ {k = 0} ^ {K - 1} \mathbb {E} \left[ \left\| \log \pi_ {k + 1} - \log \pi_ {\theta_ {k}} \right\| _ {\infty} \right] = \mathcal {O} \left(K ^ {- \frac {1}{2}}\right) + \mathcal {O} (\epsilon_ {\mathrm {a p p}}) \tag {20a}
$$
$$
\frac {1}{K} \sum_ {k = 0} ^ {K - 1} \mathbb {E} \left[ \| \nabla \widehat {L} \left(\theta_ {k}\right) \| ^ {2} \right] = \mathcal {O} \left(K ^ {- \frac {1}{2}}\right) + \mathcal {O} \left(\epsilon_ {\mathrm {a p p}}\right) \tag {20b}
$$
where we have defined  $\| \log \pi_{k + 1} - \log \pi_{\theta_k}\|_\infty \coloneqq \max_{s\in \mathcal{S},a\in \mathcal{A}}\big|\log \pi_{k + 1}(a|s) - \log \pi_{\theta_k}(a|s)\big|$  and  $\epsilon_{\mathrm{app}}$  is the approximation error defined in (13).
In Theorem 2, we demonstrated that Alg.1 identifies the approximate stationary solution of the surrogate problem (12) in finite time. Next, we show that when the reward is parameterized linearly, an improved result can be obtained, where the stationary solutions of problem (12) can be connected to the optimal solutions of the offline-IRL problem (12). See Appendix K for detailed proof.
Theorem 3 (Optimality Guarantee). Assume that the reward function is linearly parameterized, i.e.,  $r(s,a;\theta)\coloneqq \phi (s,a)^{\top}\theta$  where  $\phi (s,a)$  is the feature vector of the state-action pair  $(s,a)$ . Then any stationary point of the surrogate problem (12) is a global optimum. Furthermore, for any  $\varepsilon \in (0,\frac{4\gamma C_v}{1 - \gamma})$ , suppose there are more than  $N$  data points on each state-action pair  $(s,a)\in \Omega$  and the number of transition dataset  $\mathcal{D}$  satisfies:
where  $c$  is a constant dependent on  $\delta$ . With probability greater than  $1 - \delta$ , any stationary point  $\tilde{\theta}$  of the surrogate objective  $\widehat{L}(\cdot)$  is an epsilon-optimal solution to the maximum likelihood problem (2):
$$
L \left(\theta^ {*}\right) - L (\tilde {\theta}) \leq \varepsilon \tag {21}
$$
where  $\theta^{*}$  is defined as the optimal reward parameter of the log-likelihood objective  $L(\cdot)$ .
As a remark, the epsilon-optimal solution on the MLE problem implies:
$$
L (\theta^ {*}) - L (\tilde {\theta}) = \frac {1}{1 - \gamma} \mathbb {E} _ {s \sim d ^ {\mathrm {E}} (\cdot), a \sim \pi^ {\mathrm {E}} (\cdot | s)} \big [ \log \big (\frac {\pi_ {\theta^ {*}} (a | s)}{\pi_ {\tilde {\theta}} (a | s)} \big) \big ] \leq \varepsilon .
$$
When the expert trajectories are consistent with the optimal policy under a ground truth reward parameter  $\theta^{*}$ , we have  $\pi^{\mathrm{E}} = \pi_{\theta^{*}}$ . Due to this property, we can obtain that
$$
L (\theta^ {*}) - L (\tilde {\theta}) = \frac {1}{1 - \gamma} \mathbb {E} _ {s \sim d ^ {\mathrm {E}} (\cdot), a \sim \pi^ {\mathrm {E}} (\cdot | s)} \left[ \log \left(\frac {\pi^ {\mathrm {E}} (a | s)}{\pi_ {\tilde {\theta}} (a | s)}\right) \right] = \frac {1}{1 - \gamma} \mathbb {E} _ {s \sim d ^ {\mathrm {E}} (\cdot)} \left[ D _ {K L} \left(\pi^ {\mathrm {E}} (\cdot | s) | | \pi_ {\tilde {\theta}} (\cdot | s)\right) \right] \leq \varepsilon .
$$
Hence, Theorem 3 also provides a formal guarantee that the recovered policy  $\pi_{\tilde{\theta}}$  is  $\epsilon$ -close to the expert policy  $\pi^{\mathrm{E}}$  measured by the KL divergence.
We remark that even under the linear parameterization assumption, showing the optimality of the stationary solution for the surrogate problem (12) is still non-trivial, since the problem is still not a concave problem with respect to  $\theta$ . In our analysis, we translate this problem to a certain saddle point problem. Under linear reward parameterization, we show that any stationary solution  $\tilde{\theta}$  of the surrogate problem (12) together with the corresponding optimal policy  $\pi_{\tilde{\theta}}$  consists of a saddle point to the saddle point problem. By further leveraging the property of the saddle point, we show the optimality of the stationary solution for the surrogate problem. Finally, by utilizing the statistical guarantee in Theorem 1, we obtain the performance guarantee for any stationary point  $\tilde{\theta}$  in (21).
In this section, we present numerical results for the proposed algorithm. More specifically, we intend to address the following questions: 1) How does the proposed algorithm compare with other state-of-the-art methods? 2) Whether offline ML-IRL can recover a high-quality reward estimator of the expert, in the sense that the reward can be used across different datasets or environment?
We compare the proposed method with several benchmarks. Two classes of the existing algorithms are considered as baselines: 1) state-of-the-art offline IRL algorithms which are designed to recover the ground-truth reward and expert policy from demonstrations, including a model-based approach CLARE [15] and a model-free approach IQ-Learn [14]; 2) imitation learning algorithms which only learn a policy to mimic the expert behaviors, including behavior cloning (BC) and ValueDICE [28].
We test the performance of the proposed Offline ML-IRL on a diverse collection of robotics training tasks in MuJoCo simulator [53], as well as datasets in D4RL benchmark [54], which include three environments (halfcheetah, hopper and walker2d) and three dataset types (medium-replay, medium, and medium-expert). In each experiment set, both environment interactions and the ground-truth reward are not accessible. Moreover, we train the algorithm until convergence and record the average reward of the episodes over 6 random seeds.
In Offline ML-IRL, we estimate the dynamics model by neural networks which model the location of the next state by Gaussian distributions. Here, we independently train an ensemble of  $N$  estimated world model  $\{\widehat{P}_{\phi,\varphi}^i(s_{t+1}|s_t,a_t) = \mathcal{N}(\mu_\phi^i(s_t,a_t),\Sigma_\varphi^i(s_t,a_t))\}_{i=1}^N$  via likelihood maximization over transition samples. Then we can quantify the model uncertainty and construct the penalty function. For example, in [25], the aleatoric uncertainty is considered and the penalty function is constructed as  $U(s,a) = -\max_{i=1,\dots,N}\|\Sigma_\varphi^i(s,a)\|_F$ . For the offline RL subroutine in (13) - (14), we follow the implementation of MOPO [25]. We follow the same setup in [37] to select the key hyperparameters for MOPO. We parameterize the reward function by a three-layer neural network. The reward parameter and the policy are updated alternatingly, where the former is updated according to (17), while the latter is optimized according to MOPO. More experiment details are in Appendix A.
In Figure 2 and Table 1, we show the performance comparison between the proposed algorithm and benchmarks. The results of IQ-Learn is not presented since it suffers from unstable performance. To obtain better performance in BC and ValueDICE, we only use the expert demonstrations to train their agents. From the results, it is clear that the proposed Offline ML-IRL outperforms the benchmark algorithms by a large margin in most cases. In the experiment sets where Offline ML-IRL does not obtain the best performance, we notice the performance gap is small compared with the leading benchmark algorithm. Moreover, among the three dataset types (medium-replay, medium and medium-expert), medium-expert dataset has the most complete coverage on the expert-visited
Figure 2: The performance of Offline ML-IRL given 5,000 expert demonstrations.
Dataset type	Environment	Offline ML-IRL	BC	ValueDICE	CLARE	Expert Performance
medium	hopper	2453.25 ± 717.30	2801.19 ± 330.88	3073.16 ± 538.67	3015.37 ± 474.38	3512.09 ± 21.65
medium	halfcheetah	7640.73 ± 195.00	4471.72 ± 2835.55	1125.17 ± 959.45	841.46 ± 344.06	12174.61 ± 91.45
medium	walker2d	3989.20 ± 487.82	2328.75 ± 906.96	3191.47 ± 1887.90	237.49 ± 160.62	5383.98 ± 52.15
medium-replay	hopper	3046.36 ± 429.25	2801.19 ± 330.88	3073.16 ± 538.67	2888.04 ± 844.48	3512.09 ± 21.65
medium-replay	halfcheetah	9236.84 ± 309.10	4471.72 ± 2835.55	1125.17 ± 959.45	437.18 ± 182.99	12174.61 ± 91.45
medium-replay	walker2d	3995.32 ± 487.82	2328.75 ± 906.96	3191.47 ± 1887.90	291.71 ± 75.66	5383.98 ± 52.15
medium-exp	hopper	3347.11 ± 238.18	2801.19 ± 330.88	3073.16 ± 538.67	3350.47 ± 245.78	3512.09 ± 21.65
medium-exp	halfcheetah	11231.40 ± 585.21	4471.72 ± 2835.55	1125.17 ± 959.45	622.79 ± 56.46	12174.61 ± 91.45
medium-exp	walker2d	4201.40 ± 637.99	2328.75 ± 906.96	3191.47 ± 1887.90	959.50 ± 470.64	5383.98 ± 52.15
Table 1: MuJoCo Results. The performance versus different datasets and 5,000 expert demonstrations. The bolded numbers are the best ones for each data set, among Offline ML-IRL, BC, ValueDICE, and CLARE.
state-action space. As a result, Offline ML-IRL can nearly match the expert performance when the world model is trained from the medium-expert transition dataset. This numerical result matches our theoretical understanding in Theorem 1. As a remark, in most cases, we notice that Offline ML-IRL outperforms CLARE by a large margin which is also a model-based offline IRL algorithm. As we discussed in the related work (Appendix B), this is due to the fact that when the expert demonstrations are limited and diverse transition samples take a large portion of all collected data, CLARE learns a policy to match the joint visitation measure of all collected data, thus fails to imitate expert behaviors.
In Appendix A, we further present an experiment showcasing the quality of the reward recovered by the Offline ML-IRL algorithm. Towards this end, we use the recovered reward from the medium-expert dataset to estimate the reward value for all transition samples  $(s,a,s^{\prime})$  in the medium-replay dataset. With the recovered reward function  $r(\cdot ,\cdot ;\theta)$  and given those transition samples with estimated reward labels  $\{(s,a,r(s,a;\theta),s^{\prime})\}$ , we run MOPO to solve Offline RL tasks given these transition samples with estimated rewards. In Fig. 5, we can infer that the Offline ML-IRL can recover high-quality reward functions, since the recovered reward can label transition samples for solving offline RL tasks. Comparing with the numerical results of directly doing offline-IRL on the respective datasets in Table 1, we show that solving Offline RL by using transferred reward function and unlabelled transition datasets can achieve similar performance in Fig. 5.
In this paper, we model the offline Inverse Reinforcement Learning (IRL) problem from a maximum likelihood estimation perspective. Through constructing a generative world model to estimate the environment dynamics, we solve the offline IRL problem through a model-based approach. We develop a computationally-efficient algorithm that effectively recovers the underlying reward function and its associated optimal policy. We have also established statistical and computational guarantees for the performance of the recovered reward estimator. Through extensive experiments, we demonstrate that our algorithm outperforms existing benchmarks for offline IRL and imitation Learning, especially on high-dimensional robotics control tasks. One limitation of our method is that we focus solely on aligning with expert demonstrations during the reward learning process. In an ideal scenario, reward learning should incorporate diverse metrics and data sources, such as expert demonstrations and preferences gathered through human feedback. One direction for future work is to broaden our algorithm framework and theoretical analysis for a wider scope in reward learning.
M. Hong and S. Zeng are supported by NSF grant CIF-1910385.