SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \left(\sum_ {k = 0} ^ {K - 1} \beta_ {k}\right) ^ {- 1} \left(\frac {1}{1 - \Gamma} \sum_ {k = 0} ^ {K - 1} \beta_ {k} \left(\rho^ {\pi_ {k}} - \rho^ {*}\right) + \beta_ {K - 1} \left(\mu + \frac {1}{\eta_ {K - 1}}\right) \mathbb {E} \left[ D \left(\pi_ {K}, \pi^ {*}\right) \right]\right) \\ \leq \left(\sum_ {k = 0} ^ {K - 1} \beta_ {k}\right) ^ {- 1} \left(\frac {\beta_ {0}}{\eta_ {0}} D \left(\pi_ {0}, \pi^ {*}\right) + \sum_ {k = 0} ^ {K - 1} \beta_ {k} \eta_ {k} \left[ \left(2 M _ {\bar {\mathcal {Q}}} + M _ {\bar {\mathcal {Q}}} + M _ {h}\right) ^ {2} / 2 + \left(\sigma^ {\omega}\right) ^ {2} \right]\right) \\ + \varsigma^ {\bar {Q}} + \varsigma^ {\omega} \bar {D} _ {\mathcal {A}}, \tag {25} \\ \end{array}
$$
where  $\bar{D}_{\mathcal{A}}\coloneqq \max_{a_1,a_2\in \mathcal{A}}D(a_1,a_2)$
Proof can be found in A.4. Equation 25 encapsulates a convergence overview without specifying explicit step size choices. To be specific,  $\alpha_{k},\beta_{k}$  are both step sizes, and the central term  $\sum \rho_{k} - \rho^{*}$  (distance to the optimal value) plus  $D(\pi_k,\pi^*)$  (the distance of the current policy to the optimal policy) will shrink as  $k$  grows. The sum is bounded by some combination of  $\alpha_{k}$  and  $\beta_{k}$  plus an irreducible function approximation error  $\varsigma$ . The left-hand side is a weighted average of function value errors plus an average error from step 1 to  $K$ . The following result summarizes the convergence rate using particular step size choices. We obtain the first part of Corollary 3.5 by fixing the number of iteration  $K$  and optimizing the right-hand side of Eq. 25. We get the rest by straightforward computation.
Corollary 3.5. a) If  $\mu = 0$ , and  $\eta_{k} = \sqrt{\frac{\mathcal{D}(\pi_{1},\pi^{*})}{K\left[\left(2M_{\bar{Q}} + M_{\bar{Q}} + M_{h}\right)^{2} + (\sigma^{\omega})^{2}\right]}}$  and  $\beta_{k} = 1, k = 1,2,3\dots ,K,$  then the average reward will converge to the global optimal  $\rho^{*}$
$$
\frac {1}{K} \sum_ {k = 0} ^ {K - 1} \left(\rho^ {\pi_ {k}} - \rho^ {*}\right) / (1 - \Gamma) = \mathcal {O} \left(K ^ {- 1 / 2}\right) + \varsigma^ {\bar {Q}} + \varsigma^ {\omega} \bar {D} _ {\mathcal {A}}. \tag {26}
$$
b) If  $\mu > 0$ , and  $\eta_k = \frac{1}{\mu k}$  and  $\beta_k = 1$ , then the average reward will converge to the global optimal
$$
\frac {1}{K} \sum_ {k = 0} ^ {K - 1} \left(\rho^ {\pi_ {k}} - \rho^ {*}\right) / (1 - \Gamma) + \mu \mathbb {E} \left[ D \left(\pi_ {K}, \pi^ {*}\right) \right] = \tilde {\mathcal {O}} \left(K ^ {- 1}\right) + \varsigma^ {\bar {Q}} + \varsigma^ {\omega} \bar {D} _ {\mathcal {A}}. \tag {27}
$$
c) If  $\mu > 0$ ,  $\eta_k = \frac{2}{\mu k}$  and  $\beta_k = k + 1$ , then the average reward will converge to the global optimal
$$
\frac {2}{K (K + 1)} \sum_ {k = 0} ^ {K - 1} \frac {k + 1}{1 - \Gamma} \left(\rho^ {\pi_ {k}} - \rho^ {*}\right) + \mu \mathbb {E} \left[ D \left(\pi_ {K}, \pi^ {*}\right) \right] = \mathcal {O} \left(K ^ {- 2}\right) + \mathcal {O} \left(K ^ {- 1}\right) + \varsigma^ {\bar {Q}} + \varsigma^ {\omega} \bar {D} _ {\mathcal {A}}. \tag {28}
$$
Remark The above result shows that if we carefully choose the function approximation  $\bar{\mathcal{Q}}$  and the regularizer  $h$  so that  $\mu > 0$ , we need  $\mathcal{O}(\varepsilon^{-1})$  number of SPMD iterations to obtain  $\varepsilon$  precision solution, i.e.,  $\frac{1}{K}\sum_{k=0}^{K-1}\rho_k^\pi - \rho^* + \mathbb{E}[D(\pi_K, \pi^*)] \leq \varepsilon$ . The complexity deteriorates to  $\mathcal{O}(\varepsilon^{-2})$  if  $\mu = 0$ . Notice that  $\varsigma^Q$  and  $\varsigma^\omega$  are irreducible solely by policy optimization, as they arise due to function approximation and stochastic estimation errors.
When  $\mu < 0$ , global optimal solutions are unobtainable. The best we can do is to find a stationary point, that is when  $\psi$  and  $D(\pi_k,\pi_{k + 1})$  tend to 0. In this case, we only need to assume that the estimation error terms  $\delta_k^{\omega ,det}(s)$  and  $\delta_k^{\omega ,sto}(s)$  are bounded, i.e.,
Assumption 3.6. When  $\mu < 0$ , the critic step has bounded errors for any  $s \in S$ , i.e., there exist some constants  $\bar{\varsigma}^{\omega}$ ,  $\bar{\sigma}^{\omega}$  such that
$$
\left\| \delta_ {k} ^ {\omega , d e t} (s) \right\| _ {*} \leq \bar {\varsigma} ^ {\omega} \text {a n d} \left\| \delta_ {k} ^ {\omega , s t o} (s) \right\| _ {*} \leq \bar {\sigma} ^ {\omega}. \tag {29}
$$
Theorem 3.7. Let the number of iterations  $K$  be fixed. Suppose that  $\boxed{2.1}$  and  $\boxed{3.2}$  hold. Also assume that  $\mu < 0$  and that  $\eta_{k} = \eta = \min \{| \mu | / 2, 1 / \sqrt{K} \}$ ,  $k = 0, \dots, K - 1$ . Then for any  $s \in S$ , there exist iteration indices  $k(s)$  found by running  $K$  iterations of the SPMD method such that both the generalized advantage function and the distance between the iterates will converge to 0, i.e.,
$$
\psi^ {\pi_ {k (s)}} \left(s, \pi_ {k (s) + 1} (s)\right) = \mathcal {O} \left(K ^ {- 1}\right) + \mathcal {O} \left(K ^ {- 1 / 2}\right) + \Gamma \bar {\zeta} ^ {\omega} \bar {D} _ {\mathcal {A}} / (1 - \Gamma) \tag {30}
$$
$$
\begin{array}{l} \frac {1}{2 \eta} D (\pi_ {k (s)}, \pi_ {k (s) + 1} (s)) + (\mu + \frac {1}{\eta}) D (\pi_ {k (s) + 1} (s), \pi_ {k (s)} (s)) \\ = \mathcal {O} (K ^ {- 1}) + \mathcal {O} (K ^ {- 1 / 2}) + \bar {\varsigma} ^ {\omega} \bar {D} _ {\mathcal {A}} / (1 - \Gamma) \tag {31} \\ \end{array}
$$
Remark This is the most general case, as we don't assume any restrictions on the function approximation class. As expected, the complexity bound becomes worse compared to previous results. To reach a  $\varepsilon$ -precision stationary point, we require at least  $\mathcal{O}(\varepsilon^{-2})$  iterations. Note that  $\psi \approx 0$  indicates that  $\bar{Q}^{\pi_k}(s,\pi_{k + 1}(s))$  has virtually no difference than  $\bar{V}^{\pi_k}(s)$ . This has a similar implication of  $D(\pi_{k + 1}(s),\pi_k(s))$  approaching 0 which implies that we are reaching a stationary point.
3.1 Practical algorithm for Entropy Regularized AMDP
The above SPMD algorithm computes  $\pi(s)$  for every state. In each iteration, we need to solve a subproblem Eq. [15] for every state (or for any state encountered). In practice, we prefer performing each SPMD iteration in a mini-batch style, i.e., approximately solving Eq. [15] from some trajectories collected. If the regularizer  $h$  and  $\omega$  are the negative entropy, we can approximately solve the SPMD actor step Eq. [15] by
$$
\pi_ {k + 1} = \underset {\pi \in \Pi} {\arg \min } K L \left(\pi (s) \left\| \frac {1}{Z \left(\phi_ {k}\right)} \exp \left(\tilde {\mathcal {Q}} ^ {\pi_ {k}} \left(s, a; \phi_ {k}\right)\right)\right), \right. \tag {32}
$$
where  $Z(\phi_k)$  is some normalization constant and  $\tilde{\mathcal{Q}}^{\pi_k}(s,a;\phi_k) := -\frac{\eta_k}{1 + \eta_k}\bar{\mathcal{Q}}^{\pi_k}(s,a;\phi_k) - \frac{1}{1 + \eta_k}\log \pi_k(s)$ . In practice, we apply the so-called reparameterization trick [13] to obtain an unbiased gradient estimator. Let the policy be parameterized by
$$
\pi (a | s) = \pi \left(f _ {\xi} (\epsilon ; s) | s\right), \forall (s, a) \in \mathcal {S} \times \mathcal {A}, \tag {33}
$$
where  $\epsilon$  is an input noise sampled from a fixed distribution;  $\xi$  represents the parameters of the policy. Then we can approximately perform the actor step by solving the following problem:
$$
\min  _ {\xi} \mathbb {E} _ {(s, a) \sim \zeta , \epsilon} \left[ \log \left(\pi \left(f _ {\xi} (\epsilon ; s) | s\right)\right) - \tilde {\mathcal {Q}} ^ {\pi_ {k}} (s, f _ {\xi} (\epsilon ; s)) \right]. \tag {34}
$$
For large action spaces, the reparameterization trick efficiently approximates the solution of Eq. [15]. As for the critic step, we can approximate  $\bar{Q}$  by minimizing the Temporal Difference (TD) error, i.e., solving the following problem:
$$
\min  _ {\phi} \mathbb {E} _ {(s, a, c, s ^ {\prime}, a ^ {\prime}) \sim \zeta} \left(c (s, a) + h ^ {a} (s) - \tilde {\rho} ^ {\pi} + \bar {\mathcal {Q}} ^ {\pi , \zeta} \left(s ^ {\prime}, a ^ {\prime}; \phi\right) - \bar {\mathcal {Q}} ^ {\pi , \zeta} (s, a; \phi)\right) ^ {2}, \tag {35}
$$
where  $\tilde{\rho}^{\pi}$  is the estimated average cost, e.g., by taking  $\tilde{\rho}^{\pi} = \frac{1}{|\zeta|}\sum c(s,a) + h^{a}(s)$ .
4 Inverse Policy Mirror Descent for IRL
Equipped with an efficient solver for AMDPs presented in section 3, we can now solve the IRL problem in the form of (Dual IRL). We suppose  $\theta$  parameterizes the reward function with arbitrary methods, e.g., neural networks. To solve the dual problem, we update the parameter  $\theta$  by performing a gradient descent step  $\theta_{k + 1} = \theta_k - \alpha_k\nabla_\theta L(\theta)$  where  $\alpha_{k}$  is some predefined step size. The gradient computation of the dual objective function is
$$
\nabla L (\theta) = \mathbb {E} _ {(s, a) \sim d ^ {\varepsilon}} \left[ \nabla_ {\theta} c (s, a; \theta) \right] - \mathbb {E} _ {(s, a) \sim d ^ {\pi}} \left[ \nabla_ {\theta} c (s, a; \theta) \right]. \tag {36}
$$
Algorithm 2: The Inverse Policy Mirror Descent (IPMD) algorithm
1: Input: Initialize random policy  $\pi_0$  and step size sequence  $\{\alpha_k\}$
2: for  $k = 0,1,\dots ,K$  do
3: Collect samples  $\{(s_t, a_t)\}_{t \geq 1}$  and compute  $c(s_t, a_t; \theta_k)$  based on current reward estimation.
4: Critic step: Implement a policy evaluation algorithm to evaluation  $\bar{Q}^{\pi_k}$
$$
\bar {Q} ^ {\pi_ {k}} (s, a) \approx \bar {Q} ^ {\pi , \zeta_ {k}} (s, a). \tag {37}
$$
5: Actor step: Update the policy by solving the following for every  $s \in S$
$$
\pi_ {k + 1} (s) = \underset {\pi \in \Pi} {\arg \min } K L \left(\pi (s) \| \frac {1}{Z} \exp \left(\tilde {\mathcal {Q}} ^ {\pi_ {k}} (s, a)\right)\right). \tag {38}
$$
6: Dual Update step: Perform the stochastic update
$$
\theta_ {k + 1} = \theta_ {k} - \alpha_ {k} g _ {k}. \tag {39}
$$
Without accessing the transition kernel  $\mathbb{P}$ , the true gradient is unavailable. So we use a stochastic approximation instead, denoted as  $g_{k} \coloneqq g(\theta_{k};\zeta_{k}^{E}) - g(\theta_{k};\zeta_{k}^{\pi})$  where  $g(\theta ;\zeta)\coloneqq \frac{1}{N}\sum_{t = 1}^{N}\nabla c(s_{t},a_{t};\theta)$  is the stochastic estimator of the gradient of the average reward. We describe the proposed Inverse Policy Mirror Descent (IPMD) algorithm in Algorithm 2. In this section, we provide our analysis that captures the algorithm behaviors across iterations since, in each iteration,  $\bar{Q}$  is only an approximation of the true differential  $Q$ -function, which itself alters due to the change of the reward estimation. The idea of the proof is based on the Lipschitz continuity of the iterates, as it controls the difference of these functions across iterations. We make the following formal assumptions for such a purpose.
Assumption 4.1. For any  $s \in S, a \in \mathcal{A}$ , the gradient of the reward function is bounded and Lipschitz continuous, i.e., there exist some constant real numbers  $L_{r}, L_{q}$  so that the following holds:
$$
\| \nabla c (s, a; \theta) \| _ {2} \leq L _ {r}, \| \nabla c (s, a; \theta_ {1}) - \nabla c (s, a; \theta_ {2}) \| _ {2} \leq L _ {g} \| \theta_ {1} - \theta_ {2} \| _ {2}.
$$
Note that  $\bar{Q}$  is also a function of  $\theta$ . As shown in section 3,  $\bar{Q}$  can be parameterized but we only denote the estimator as  $\bar{Q}^{\pi}(s,a;\theta_k)$  since its parameters do not contribute to the analysis.
Assumption 4.2. Suppose that at least one of  $S, \mathcal{A}$  is continuous, we assume that
$$
\max  _ {\theta} \| \nabla_ {\theta} \bar {\mathcal {Q}} ^ {\pi} (s, a; \theta) \| _ {2} \leq L _ {q}, \tag {40}
$$
where  $L_{q}$  is some positive constant. For convenience we denote  $\bar{\mathcal{Q}}^{\pi}(s,a;\theta_{k})$  also as  $\bar{\mathcal{Q}}_{\theta_k}^\pi (s,a)$ .
We further assume that the estimation error from using  $\bar{Q}_{\theta}^{\pi}$  is bounded.
Assumption 4.3. For any  $s \in S, a \in \mathcal{A}$ , there exist some constants  $\varsigma_k, \varsigma, \nu_k, \nu$  such that the following inequalities hold:
$$
\left\| \mathbb {E} _ {\zeta_ {k}} \left[ \bar {Q} _ {\theta_ {k}} ^ {\pi_ {k}, \zeta_ {k}} \right] - \bar {Q} _ {\theta_ {k}} ^ {\pi_ {k}} \right\| _ {s p, \infty} \leq \varsigma_ {k} \leq \varsigma , \quad \mathbb {E} _ {\zeta_ {k}} \left[ \left\| \bar {Q} _ {\theta_ {k}} ^ {\pi_ {k}, \zeta_ {k}} - \bar {Q} _ {\theta_ {k}} ^ {\pi_ {k}} \right\| _ {s p, \infty} ^ {2} \right] \leq \nu_ {k} ^ {2} \leq \nu^ {2}. \tag {41}
$$
Notice that this error bound subsumes both the estimation errors and approximation errors. In general, the Bellman operator is nonexpansive in the average-reward setting, so analysis based on the contraction property of the Bellman operator with the infinity norm, as what is used in [37], fails in this case. To this end, we assume that the operator is span contractive, i.e.,
Assumption 4.4. In the critic step, there's a way to construct a span contractive Bellman operator  $\mathcal{T}$  such that there exists  $0 < \gamma < 1$  for any differential  $Q$ -functions  $\bar{Q}_1, \bar{Q}_2$ ,
$$
\left\| \mathcal {T} \bar {Q} _ {1} - \mathcal {T} \bar {Q} _ {2} \right\| _ {s p, \infty} \leq \gamma \left\| \bar {Q} _ {1} - \bar {Q} _ {2} \right\| _ {s p, \infty}. \tag {42}
$$
In [A.6], we provide conditions when the Bellman operator has a J-step span contraction. Now we can discuss the convergence results in the following theorems. The convergence of the differential  $Q$ -function is characterized in the following result.
Theorem 4.5. Suppose that assumptions  $\boxed{2.1}$ ,  $\boxed{4.1}$ ,  $\boxed{4.4}$  hold. If  $\alpha = \frac{\alpha_0}{\sqrt{K}}$ , the differential  $Q$ -function will converge to the optimal solution for a given reward function parameterized by  $\theta_k$ , i.e.,
$$
\frac {1}{K} \sum_ {k = 0} ^ {K - 1} \left\| \mathbb {E} _ {\zeta_ {k}} \left[ \bar {Q} _ {\theta_ {k}} ^ {\pi_ {k}, \zeta_ {k}} \right] - \bar {Q} _ {\theta_ {k}} ^ {\pi_ {\theta_ {k}}} \right\| _ {s p, \infty} = \mathcal {O} (K ^ {- 1}) + \mathcal {O} (K ^ {- 1 / 2}) + \frac {\varsigma}{1 - \gamma}, \tag {43}
$$
where  $\alpha_0 > 0$  is some step size chosen, and  $\gamma$  is some constant defined in Assumption 4.4.
Proof can be found in A.7. The above theorem shows that the differential  $Q$ -function will approach the optimal differential  $Q$ -function  $\bar{Q}_{\theta_k}^{\pi_{\theta_k}}$  with respect to the current reward estimation  $\theta_k$ . If the reward estimation is accurate, so is the differential  $Q$ -function. Finally, in the next theorem, we show that the policy converges to the optimal policy for a given reward function, and the reward function approximation converges to a stationary point. Proof can be found in A.8.
Theorem 4.6. Suppose that assumptions  $\boxed{2.1, 4.1, 4.4}$  hold. If  $\alpha = \frac{\alpha_0}{\sqrt{K}}$  and run the proposed IPMD algorithm  $K$  iterations, the algorithm will produce near stationary solutions for the reward function and the optimal policy for such reward function. Specifically, the following holds,
$$
\frac {1}{K} \sum_ {k = 0} ^ {K - 1} \| \mathbb {E} _ {\zeta_ {k}} [ \log \pi_ {k + 1} ] - \log \pi_ {\theta_ {k}} \| _ {\infty} = \mathcal {O} (K ^ {- 1}) + \mathcal {O} (K ^ {- 1 / 2}) + \frac {\varsigma}{1 - \gamma}, \tag {44}
$$
$$
\frac {1}{K} \mathbb {E} [ \| \nabla L (\theta_ {k}) \| _ {2} ^ {2} ] = \mathcal {O} (K ^ {- 1}) + \mathcal {O} (K ^ {- 1 / 2}) + \frac {1}{1 - \gamma} 2 \nu L _ {c} L _ {r} C _ {d} \sqrt {| \mathcal {S} | \cdot | \mathcal {A} |}, \tag {45}
$$
where  $\alpha_0, \nu, L_c, L_r, C_d$  are some constants and  $|\mathcal{S}|, |\mathcal{A}|$  some measure of the state and action space.
Remark First, the above theorem shows that the policy will converge to the optimal policy given the reward parameterization, although each iteration's reward parameter differs. Second, the reward parameter  $\theta$  will converge to a stationary point of (Dual IRL), as the objective can be highly nonconvex under general function approximation. Third, for general AMDPs without 1-step span contraction, we can use a  $J$ -step Bellman operator (see Appendix A.7) for policy evaluation to maintain a  $\mathcal{O}(\varepsilon^{-2})$  complexity for the entire algorithm.
In this section, we showcase the performance of the proposed SPMD and IPMD algorithms. Our code can be found at https://anonymous.4open.science/r/IPMD-9D60. See more detail about all our implementation in the Appendix A.9.
MuJoCo Robotics Manipulation Tasks for RL This experiment tests our RL agent's performance on robotics manipulation tasks. Our SPMD algorithm is based on the stable-baselines3 [28]. We compare the performance of our algorithm with Soft Acrot-Critic (SAC) [13] implemented in [28]. The policy network employs two fully connected hidden layers of dimension 256 each, taking actions as input and outputting a distribution. Both the  $Q$  network and reward function share the same architecture, with ReLU activation used in hidden layers. A double  $Q$ -learning technique is used to minimize overestimation [13]. During training, we found that setting the entropy coefficient term to 0.01 makes training stable and efficient. The learning rate is  $3e^{-4}$ . Each step of the algorithm samples 512 state-action sample pairs. Table I reports the numerical results of each model. Our proposed SPMD achieves on-par performance with SAC and exceptionally better performance in the Humanoid environment.
MuJoCo benchmark for IRL In this experiment, we compare the proposed IPMD method with IQ-Learn [11] and  $f$ -IRL [23]. The authors of ML-IRL have not released its implementation at the time we experimented. Nevertheless, the performance of ML-IRL is comparable to  $f$ -IRL. It is observed in the literature that imitation learning algorithms have inferior performance [37]. Hence we omit these methods. Table 2 reports the numerical results of each model. Note that we cannot record a competitive result for IQ-Learn with Humanoid as claimed in the original paper, which we highlight using *. The result shows that IPMD outperforms in a variety of environments. One possible reason Ant is falling behind is that Ant has more ground contact since it has more legs. This will impact the mixing time of the MDP and our assumption on the 1-step contraction. Compared to Half-Cheetah, or Humanoid, it is harder for the Ant to transition from an arbitrary state to another arbitrary state as it involves multiple legs working together. One remedy is that we construct a  $J$ -step contractive operator for policy evaluation, as done in [7]. The success of humanoid is possibly due to a slightly different policy evaluation scheme compared to the discounted setting, where the entropy
Table 1. MuJoCo Results for RL. The average performance of RL agents across five runs.
Task	SAC	SPMD (Ours)
Hopper	3876 ± 78	3619 ± 5
Half-Cheetah	13300 ± 46	13025 ± 37
Walker	6115 ± 41	4454 ± 26
Ant	5118 ± 57	5230 ± 118
Humanoid	6923 ± 11	10249 ± 7
Table 2. MuJoCo Results for IRL. The average performance of IRL agents across five runs.
IQL	f-IRL	IPMD (ours)	Expert
1909	3083	3564	4115
9132	11259	12634	15467
5155	5378	5423	5323
3486	5460	4053	5783
661*	3580	7379	7137
Figure 1. Reward recovery performance of IPMD and IQ-Learn. The lower the bar, the better the performance.
term from the policy no longer plays a part, as the term  $c - \rho$  cancels out the additional regularization and entropy of the policy. We suspect this brings more stable training and thus higher performance.
Reward Recovery Finally, we compare the proposed IPMD method against IQ-Learn on recovering the expert's reward function in two environments: Pendulum and LunarLanderContinuous, both implemented in [4]. The result of the experiment is shown in Figure 1. Using state-action pairs from 4 different episodes, we compare the span seminorm of the predicted reward and the true reward (the reward the expert uses). The experts are trained with a discount factor 0.99. IQ-Learn's discount factor is set to 0.99. The result shows IPMD's superiority in reward recovery.
In this paper, we formulate the Maximum Entropy IRL problem with the average-reward criterion and propose efficient algorithms to solve it. We devise the SPMD method for solving RL problems in general state and action spaces beyond the scope of linear function approximation with strong theoretical guarantees. We integrate this method to develop the IPMD method for solving the entire IRL problem and provide convergence analysis. To reach a  $\varepsilon$  precision stationary point, IPMD requires  $\mathcal{O}(\varepsilon^{-2})$  policy optimization steps.
We also notice there are possible improvements from this work. For example, we impose various assumptions both the function approximation classes and strong reliance on the accuracy of the critic step. It is natural to seek more robust methods when facing inaccurate estimations. Additionally, we impose continuity assumption in various Assumptions (e.g.  $\mathrm{Assump}\boxed{4.2}$ ). Such assumptions can be violated if neural networks are used as function approximations. In practice, it is possible to alleviate the issue by "clipping" the gradient norm when performing gradient descent, which proves effective in related works [30]. Furthermore, we only consider uniform ergodic chains. It is interesting to see how our method would hold under the unichain or multi-chain setting. Lastly, the analysis based on span contraction might be improved for general MDPs. This requires nontrivial work and thus we leave it for future development.
[1] Anas Barakat, Ilyas Fatkhullin, and Niao He. Reinforcement learning with general utilities: Simpler variance reduction and large state-action space. arXiv preprint arXiv:2306.01854, 2023.
[39] Junyu Zhang, Alec Koppel, Amrit Singh Bedi, Csaba Szepesvari, and Mengdi Wang. Variational policy gradient method for reinforcement learning with general utilities. Advances in Neural Information Processing Systems, 33:4572-4583, 2020.
[40] Brian D Ziebart, Andrew L Maas, J Andrew Bagnell, Anind K Dey, et al. Maximum entropy inverse reinforcement learning. In AAAI Conference on Artificial Intelligence, 2008.
[41] Brian D Ziebart, J Andrew Bagnell, and Anind K Dey. Modeling interaction via the principle of maximum causal entropy. In ICML, 2010.
[42] Konrad Zolna, Scott Reed, Alexander Novikov, Sergio Gomez Colmenarejo, David Budden, Serkan Cabi, Misha Denil, Nando de Freitas, and Ziyu Wang. Task-relevant adversarial imitation learning. In Conference on Robot Learning, pages 247-263. PMLR, 2021.