Diffusing States and Matching Scores: A New Framework for Imitation Learning

## Abstract

Abstract Adversarial Imitation Learning is traditionally framed as a two-player zero-sum game between a learner and an adversarially chosen cost function, and can therefore be thought of as the sequential generalization of a Generative Adversarial Network (GAN).
However, in recent years, diffusion models have emerged as a non-adversarial alternative to GANs that merely require training a score function via regression, yet produce generations of a higher quality. In response, we investigate how to lift insights from diffusion modeling to the sequential setting. We propose diffusing states and performing score-matching along diffused states to measure the discrepancy between the expert’s and learner’s states. Thus, our approach only requires training score functions to predict noises via standard regression, making it significantly easier and more stable to train than adversarial methods.
Theoretically, we prove first- and second-order instance-dependent bounds with linear scaling in the horizon, proving that our approach avoids the compounding errors that stymie offline approaches to imitation learning. Empirically, we show our approach outperforms both GAN-style imitation learning baselines and discriminator-free imitation learning baselines across various continuous control problems, including complex tasks like controlling humanoids to walk, sit, crawl, and navigate through obstacles.

## 1 Introduction

Fundamentally, in imitation learning (IL, Osa et al. ()), we want to match the sequential behavior of an expert demonstrator. Different notions of what matching should mean for IL have been proposed in the literature, from $f$-divergences (Ho & Ermon, ; Ke et al., ) to Integral Probability Metrics (IPMs, Müller (); Sun et al. (); Kidambi et al. (); Swamy et al. (); Chang et al. (); Song et al. ()). To compute the chosen notion of divergence from the expert demonstrations so that the learner can then optimize it, it is common to train a discriminator (i.e. a classifier) between expert and learner data. This discriminator is then used as a reward function for a policy update, an approach known as inverse reinforcement learning (IRL, Abbeel & Ng (); Ziebart et al. ()). Various losses for training discriminators have been proposed (Ho & Ermon, ; Kostrikov et al., ; Fu et al., ; Ke et al., ; Swamy et al., ; Chang et al., ), and IRL has been applied in real life in domains like routing (Barnes et al., ) and LLM fine-tuning (Wulfmeier et al., ).

As observed by Finn et al. (), inverse RL can be seen as the sequential generalization of a Generative Adversarial Network (GAN) (Goodfellow et al., ). However, even in the single-shot setting, GANs are known to suffer from issues like unstable training dynamics and mode collapse  (Miyato et al., ; Arjovsky & Bottou, ). In contrast, Score-Based Diffusion Models (Ho et al., ; Song & Ermon, ; Song et al., ) are known to be more stable to train and produce higher quality samples in domains like audio and video (Rombach et al., ; Ramesh et al., ; Kong et al., ). Intuitively, diffusion models corrupt samples from a target distribution with noise and train a generative model to reverse this corruption. Critically, these models are trained via score-matching: a non-adversarial, purely regression-based procedure in which a network is trained to match the score (i.e., gradient of the log probability) of the target distribution.

In fact, prior work has shown that diffusion models are powerful policy classes for IL due to their ability to model multi-modal behavior (Chi et al., ). However, these approaches are purely offline behavioral cloning (Pomerleau, ), and therefore can suffer from the well-known compounding error issue (Ross et al., ) that affects all offline approaches to IL (Swamy et al., ).

Taken together, the preceding points beg the question:

Can we lift the insights of diffusion models to inverse reinforcement learning?

Our answer to this question is SMILING: a new IRL framework that abandons unstable discriminator training in favor of a simpler, score-matching based objective. We first fit a score function to the expert’s state distribution. Then, at each iteration of our algorithm, we fit a score function to the state distribution of the mixture of the preceding policies via standard regression-based score matching, before using the combination of these score functions to define a cost function for the policy search step. Our framework treats diffusion model training (score-matching) as a black box which allows us to transfer any advancements in diffusion model training (e.g., better noise scheduling or better training techniques) to IRL.
In theory, rather than optimizing either an $f$-divergence or an IPM, this corresponds to minimizing a novel sequential generalization of the Fisher divergence (Johnson, ) we term the Diffusion Score Divergence (DS Divergence).

We demonstrate the advantages of our framework in both theory and practice. In theory, we show that our approach can achieve first- and second-order instance-dependent regret bounds, as well as a linear scaling in the horizon to model-misspecification errors arising from expert not being realizable by the learner’s policy class or potential optimization error in score-matching and policy search. Thus, we establish that

By lifting score-matching to IRL, SMILING provably avoids compounding errors and achieves instance-dependent bounds.

Intuitively, the second-order bounds mean that the performance gap between our learned policy and the expert automatically shrinks when either the expert or the learned policy has low variance in terms of their performance under the ground-truth reward (e.g. when the expert and dynamics are relatively deterministic). This is because when we perform score-matching, we are actually minimizing the *squared Hellinger distance* between the learner and the expert’s state distributions, which has shown to play an important role in achieving second-order regret bounds in the Reinforcement Learning setting (Wang et al., ). The ability to achieve instance-dependent bounds demonstrates the theoretical benefit of Diffusion Score divergence over other metrics such as IPMs and $f$-divergences.
In practice, we show that under in the IL from observation only setting (Torabi et al., ; Sun et al., ), SMILING outperforms adversarial GAN-based IL baselines, discriminator-free IL baselines, and Behavioral Cloning(^1^11This is despite the fact that BC requires the expert actions, while we do not.) on complex tasks such as controlling humanoids to walk, sit, crawl, and navigate through poles (see Figure 1). This makes SMILING the first IRL method to solve multiple tasks on the recently released HumanoidBench benchmark (Sferrazza et al., ) using only the state information of the expert demonstrations. We release the code for all experiments
at [https://github.com/ziqian2000/SMILING](https://github.com/ziqian2000/SMILING).

## 2 Related Works

Inverse Reinforcement Learning. Starting with the seminal work of Abbeel & Ng (), various authors have proposed solving the problem of imitation via inverse RL (Ziebart et al., ; Ratliff et al., ; Sun et al., ; Kidambi et al., ; Chang et al., ). As argued by Swamy et al. (), these approaches can be thought of as minimizing an integral probability metric (IPM) between learner and expert behavior via an adversarial training procedure – we refer interested readers to their paper for a full list. Many other IRL approaches can instead be framed as minimizing an $f$-divergence via adversarial training (Ke et al., ), including GAIL (Ho & Ermon, ), DAC (Kostrikov et al., ), AIRL (Fu et al., ), FAIRL (Ghasemipour et al., ), and f-GAIL (Zhang et al., ). However, all of the preceding approaches involve training a discriminator, while our technique only requires fitting score functions and does not seek to approximate either an $f$-divergence or an IPM. Wang et al. () also explored the use of diffusion models for imitation learning. They insert the score matching loss into the $f$-divergence objective of the discriminator. Thus, their method also belongs to the $f$-divergence framework. Huang et al. () also utilize diffusion models for imitation learning, proposing to train a conditional diffusion model as a discriminator for single-step state transitions between the expert and the current policy. Their approach remains within the classic $f$-divergence framework as well. Additionally, various authors have proposed using techniques to either stabilize this training procedure (e.g. via boosting, Chang et al. ()) or reducing the amount of interaction required during the RL step (Swamy et al., ; Ren et al., ; Sapora et al., ) – as we focus on improving the reward function used in inverse RL, our approach is orthogonal to and could be naturally applied on top of these techniques.

Discriminator-Free Inverse Reinforcement Learning. SQIL (Reddy et al., ) replaces training a discriminator with a fixed reward function ($+1$ for expert data, 0 0 for learner data). Unfortunately, this means that a performant learner is dis-incentivized to perform expert-like behavior, leading to dramatic drops in performance (Barde et al., ). The AdRIL algorithm of Swamy et al. () uses techniques from functional gradients to address this issue, but requires the use off an off-policy RL algorithm to implement, while our framework makes no such assumptions. ASAF (Barde et al., ) proposes using the prior policy to compute the optimal $f$-divergence discriminator in closed form, while we focus on a different class of divergences. IQ-Learn (Garg et al., ) proposes to perform IRL in the space $Q$ functions, but can therefore suffer from poor performance on problems with stochastic dynamics (Ren et al., ) or when the $Q$-function is more complex than the reward function (e.g. navigating to a goal in a maze). Jain et al. () propose measuring differences between learner and expert behavior in terms of successor features, but are only able to optimize over deterministic policies. DRIL (Brantley et al., ) relies on ensemble methods and uses the disagreement among models in the ensemble as a surrogate cost function, making it challenging to prove similar theoretical guarantees to our technique.

In contrast, our method instead trains score functions on the state distribution of the policy via the standard score matching objective. We show that our approach can require less expressive function approximators than discriminator-based methods (see Section 5.1). Additionally, our approach achieves solid first- and second-order sample complexity bound and also delivers strong empirical performance on challenging humanoid control tasks.

Diffusion Models for Decision Making.
Diffusion models have been widely studied in the context robot learning, where diffusion models have demonstrated a strong ability to capture complex and multi-modal data.
Recent works have explored using diffusion models to perform behavioral cloning   Chi et al. (); Pearce et al. (); Block et al. (); Team et al. () as well as using diffusion models to capture trajectory-level dynamics  (Janner et al., ; Ajay et al., ). These prior work exclusively rely on offline data for training.
We instead focus on lifting insights from diffusion modeling and score matching to IRL, which uses both offline expert data and also online interaction with the environment. Unlike prior diffusion BC work, we use diffusion models for both the expert’s and the learner’s state distributions, not just the expert’s action-conditional distributions. This means we do not require expert action labels, which can be challenging to acquire in practice.

## 3 Preliminaries

### 3.1 Markov Decision Processes

We consider finite-horizon Markov decision processes (MDPs) with state space $\mathcal{S}\subseteq\mathbb{R}^{d}$, action space $\mathcal{A}$, transition dynamics $P$, (unknown) cost function $c^{\star}:\mathcal{S}\to\mathbb{R}$,(^2^22We assume the cost depends only on states. The extension to state-action costs is straightforward.) and horizon $H$.
The goal is to find a policy $\pi:\mathcal{S}\to\Delta(\mathcal{A})$ that minimizes the expected cost $H\cdot\mathop{\mathbb{E}}_{s\sim d^{\pi}}c^{\star}(s)$ where $d^{\pi}$ denotes the average state distribution induced by the policy $\pi$, i.e., $d^{\pi}(s)=\sum_{h=1}^{H}\operatorname{Pr}(s_{h}=s)/H$ where $s_{h}$ denotes the state at time $h$. The state value function of a policy $\pi$ is defined as $V^{\pi}(s)=\mathop{\mathbb{E}}_{\tau\sim\pi}\sum_{h=1}^{H}c^{\star}(s_{h})$ where $\tau=(s_{1},\dots,s_{H})$ denotes the trajectory induced by the policy $\pi$.

### 3.2 Imitation Learning

In imitation learning, we are usually given expert demonstration in the form of state-action-next-state tuples, and the goal is to learn a policy that mimics the expert behavior. We do not have access to the ground-truth cost function $c^{\star}$.
In this paper, we consider a harder setting where only states are given.
Specifically, given a dataset $\mathcal{D}^{e}=\{s^{(i)}\}_{i=1}^{N}$ of expert demonstrations, the goal is to learn a policy $\pi$ that minimizes the discrepancy between the state distributions of policy $\pi$ and the expert policy $\pi^{e}$. In the inverse reinforcement learning setting, in addition to the expert demonstrations $\mathcal{D}^{e}$, we can interact with the environment to collect more (reward-free) learner trajectories.

### 3.3 Diffusion Models

Generally speaking, the forward process of a diffusion model adds noise to a sample from the data distribution $p_{0}$. We can formalize this via a stochastic differential equation (SDE). For simplicity in this paper, we consider the Ornstein-Uhlenbeck (OU) process:

$$ $\mathop{}\mathrm{d}\mathbf{x}_{t}=-\mathbf{x}_{t}\mathop{}\mathrm{d}t+\sqrt{2} \mathop{}\mathrm{d}B_{t},\quad\mathbf{x}_{0}\sim p_{0},$ (1) $$

where $B_{t}$ is the standard Brownian motion in $\mathbb{R}^{d}$. It is known that the reverse process $(\mathbf{y}_{t})_{t\in[0,T]}$ satisfies the following reverse-time SDE (Anderson, ):

$$ $\mathop{}\mathrm{d}\mathbf{y}_{t}=(\mathbf{y}_{t}+2\nabla\log p_{T-t}(\mathbf{ y}_{t}))\mathop{}\mathrm{d}t+\sqrt{2}\mathop{}\mathrm{d}B_{t},\quad\mathbf{y}_ {0}\sim p_{T}$ $$

where $U(T)$ denotes the uniform distribution on $[0,T]$, and $B_{t}$ now denotes reversed Brownian motion. The gradient of the log probability $\nabla\log p_{t}$ is called the score function of the distribution $p_{t}$. When the score function is known, we can use it to generate samples from $p_{0}$ by simulating the reverse process. Estimating the score function $\nabla\log p_{t}$ is called score matching, which is typically done via minimizing the following regression-based loss:

$$ $\min_{g}\mathop{\mathbb{E}}_{\mathbf{x}\sim p_{0}}\mathop{\mathbb{E}}_{t\sim U (T)}\mathop{\mathbb{E}}_{\mathbf{x}_{t}\sim q_{t}(\cdot\,|\,\mathbf{x})}\Big{ \|}g(\mathbf{x}_{t},t)-\nabla_{\mathbf{x}_{t}}\log q_{t}(\mathbf{x}_{t}\,|\, \mathbf{x})\Big{\|}_{2}^{2}.$ $$

We use $q_{t}(\mathbf{x}_{t}\,|\,\mathbf{x})$ to denote the conditional distribution at time $t$ of the forward process conditioned on the initial state $\mathbf{x}$, which has closed form $q_{t}(\mathbf{x}_{t}\,|\,\mathbf{x})=\mathcal{N}(\mathbf{x}e^{-t},(1-e^{-2t})I)$.

Applying Diffusion to State Distributions.
For a policy $\pi$, we will use $p_{t}^{\pi}$ to denote the marginal distribution at time $t$ obtained by applying the forward diffusion process to $d^{\pi}$, i.e., initial samples are drawn from $p_{0}^{\pi}\coloneqq d^{\pi}$. When presenting asymptotic results, we will use $O(\cdot)$ to hide constants and $\widetilde{O}(\cdot)$ to hide constants and logarithmic factors.

## 4 Algorithm

Figure: Algorithm 1 SMILING (Score-Matching Imitation LearnING)

We begin by introducing a novel discrepancy measure between two distributions that will be leveraged by algorithm:

###### Definition 1 (Diffusion Score Divergence) .

For two distributions $P$ and $Q$, we define the *Diffusion Score Divergence (DS Divergence)* as

$$ $\displaystyle D_{\rm{DS}}(P,Q):=\mathop{\mathbb{E}}_{s\sim P}\mathop{\mathbb{E }}_{t\sim U(T)}\mathop{\mathbb{E}}_{s_{t}\sim q_{t}(\cdot\,|\,s)}\Big{\|} \nabla\log P_{t}(s_{t})-\nabla\log Q_{t}(s_{t})\Big{\|}_{2}^{2}.$ $$

Here $q_{t}(\cdot\,|\,s)$ represents the conditional probability of the forward diffusion process at time $t$ conditioned on the initial state $s$; $P_{t}$ and $Q_{t}$ denote the marginal distributions at time $t$ obtained by applying the forward diffusion process to $P$ and $Q$, respectively. We call $\nabla\log P_{t}$ and $\nabla\log Q_{t}$ the *Diffusion Score Function* of $P$ and $Q$.

Figure 2(a) illustrates the DS divergence.
It measures the difference between two distributions by comparing their diffusion score functions within one diffusion process.
It is analogous to Fisher divergence but differs by incorporating an expectation over the diffusion process. DS divergence is a strong divergence in the sense that, whenever the DS divergence between the two distributions is small,
the KL divergence, Hellinger distance, and total variation distance are all small (see Lemma 2 and also Chen et al. (); Oko et al. ()).

In our algorithmic framework Score-Matching Imitation LearnING — SMILING,
we propose to frame imitation learning as the minimization of the DS divergence between the expert’s and the (history of) learner state distributions. We begin by discussing how to do so before discussing the theoretical benefits of doing so.

The first step is to *pre-train* a score function estimator $g^{e}(s,t)$ on expert’s state distribution $\nabla_{s}\log p^{\pi^{e}}_{t}(s)$. This can be done via standard least-squares regression-based score-matching, as shown in Line 1. Here $\mathcal{G}$ denotes the function class for score estimators.
Then, we seek to find a policy $\pi$ that minimizes the DS divergence between learner and expert state distributions:

Figure: (a) Illustration of DS Divergence (Definition 1).
Refer to caption: x2.png

$$ $\displaystyle\ell(\pi):=\mathop{\mathbb{E}}_{s\sim d^{\pi},t\sim U(T),s_{t} \sim q_{t}(\cdot|s)}\left\|g^{e}(s_{t},t)-\nabla_{s_{t}}\log p^{\pi}_{t}(s_{t} )\right\|_{2}^{2},$ (2) $$

where we have approximated the expert’s diffusion score function by $g^{e}(s_{t},t)$.
However, $\ell(\pi)$ is not directly computable since we do not know the learner’s score function $\nabla_{s_{t}}\log p^{\pi}_{t}(s_{t})$. A naive approach would be to directly learn an estimator for $\nabla_{s_{t}}\log p^{\pi}_{t}(s_{t})$ via score matching and substitute it into the equation. However, as we prove in Appendix A, this can introduce significant errors due to the unboundedness of a difference of score functions and variance-related concerns. We now derive a method that does not suffer from this issue.

Recall that we defined $q_{t}(s_{t}\,|\,s)$ as the distribution of $s_{t}$ given the initial sample $s$, which is a simple Gaussian distribution with an appropriately scaled variance.
To develop our method, we first note that, luckily, $\nabla_{s_{t}}\log q_{t}(s_{t}|s)$ is an unbiased estimator of the score $\nabla_{s_{t}}\log p^{\pi}_{t}(s_{t})$, because $\mathop{\mathbb{E}}_{s|s_{t}}[\nabla_{s_{t}}\log q_{t}(s_{t}|s)]=\nabla_{s_{t}
}\log p^{\pi}_{t}(s_{t})$
(Song et al., ). Given this, perhaps the most immediate strategy would be to simply replace $\nabla_{s_{t}}\log p^{\pi}_{t}(s_{t})$ by its unbiased estimator, $\nabla_{s_{t}}\log q_{t}(s_{t}|s)$, in Eq. 2. However, in contrast to linear objectives like an IPM, to approximate a squared objective like the DS Divergence accurately, we need to get both the expectation (i.e. have an unbiased estimator) as well as the variance correct. If we don’t, then $\mathop{\mathbb{E}}_{s\sim d^{\pi},t\sim U(T),s_{t}\sim q_{t}(\cdot|s)}\|g^{e}
(s_{t},t)-\nabla_{s_{t}}\log q_{t}(s_{t}|s)\|_{2}^{2}$ will differ from $\ell(\pi)$ by a term related to *variance* of the estimator $\nabla_{s_{t}}\log q_{t}(s_{t}|s)$, i.e.

$$ $\displaystyle\mathop{\mathbb{E}}_{s\sim d^{\pi},t\sim U(T),s_{t}\sim q_{t}( \cdot|s)}\left\|\nabla_{s_{t}}\log p^{\pi}_{t}(s_{t})-\nabla_{s_{t}}\log q_{t} (s_{t}|s)\right\|_{2}^{2}.$ (3) $$

To faithfully approximate $\ell(\pi)$, we need to estimate the variance term and subtract it from $\mathop{\mathbb{E}}_{s\sim d^{\pi},t\sim U(T),s_{t}\sim q_{t}(\cdot|s)}\|g^{e}
(s_{t},t)-\nabla_{s_{t}}\log q_{t}(s_{t}|s)\|_{2}^{2}$. We can do this by adding a term to our objective that is minimized at the variance and using another function / network to optimize it. This leads us to the following estimator for $\ell(\pi)$:

$$ $\displaystyle\hat{\ell}(\pi):=\mathop{\mathbb{E}}_{s\sim d^{\pi},t\sim U(T),s_ {t}\sim q_{t}(\cdot|s)}\left\|g^{e}(s_{t},t)-\nabla_{s_{t}}\log q_{t}(s_{t}|s) \right\|_{2}^{2}$ $\displaystyle\qquad\qquad-\underbrace{{\color[rgb]{1,.5,0}\min_{g\in\mathcal{G }}\mathop{\mathbb{E}}_{s\sim d^{\pi},t\sim U(T),s_{t}\sim q_{t}(\cdot|s)}\left \|g(s_{t},t)-\nabla_{s_{t}}\log q_{t}(s_{t}|s)\right\|_{2}^{2}}}_{\text{ Minimized at the variance of Eq.~{}\ref{eq:variance}}}$ $$

where $\mathcal{G}$ is the function class for score estimators. Crucially, the highlighted term (orange) in the above expression
estimates the variance in Eq. 3 since one of its minimizers will be the Bayes optimal $\nabla_{s_{t}}\log p^{\pi}_{t}(s_{t})$, i.e. the conditional expectation $\mathop{\mathbb{E}}_{s|s_{t}}\left[\nabla_{s_{t}}\log q_{t}(s_{t}|s)\right]$.(^3^33Note that similar ideas have been used in the offline RL literature for designing min-max based algorithms for estimating value functions (e.g., Chen & Jiang (); Uehara et al. ()).) Observe that minimizing this orange term just requires a simple regression-based score matching objective. So, while naive score matching causes issues due to the variance of the estimator, a more clever application of score matching can be used to fix this concern.

With $\hat{\ell}(\pi)$ now serving as a valid approximation of $\ell(\pi)$, the IL problem reduces to searching for a policy to minimize $\hat{\ell}(\pi)$, i.e., $\min_{\pi}\hat{\ell}(\pi)$. To facilitate this, we define payoff $\mathcal{L}(\pi,g)$ as:

$$ $\displaystyle\mathcal{L}(\pi,g)$ $\displaystyle:=\mathop{\mathbb{E}}_{s\sim d^{\pi},t\sim U(T),s_{t}\sim q_{t}( \cdot|s)}\left\|g^{e}(s_{t},t)-\nabla_{s_{t}}\log q_{t}(s_{t}|s)\right\|_{2}^{2}$ $\displaystyle\qquad-\mathop{\mathbb{E}}_{s\sim d^{\pi},t\sim U(T),s_{t}\sim q_ {t}(\cdot|s)}\left\|g(s_{t},t)-\nabla_{s_{t}}\log q_{t}(s_{t}|s)\right\|_{2}^{ 2}.$ $$

Then, minimizing $\hat{\ell}(\pi)$ is equivalent to solving the two-player zero sum game $\min_{\pi}\max_{g}\mathcal{L}(\pi,g).$
To solve this game, we propose following a no-regret strategy over $g$ against a best response over $\pi$ (i.e. a dual algorithm (Swamy et al., )). Specifically,
Algorithm 1 applies Follow-the-Leader (Shalev-Shwartz et al., )) to optimize $g$ in Line 3,(^4^44Note that $\mathcal{L}(\pi,g)$ is a square-loss functional with respect to $g$. Since square loss is strongly convex, FTL is no-regret. In contrast, when optimizing an IPM, one has to use Follow-the-Regularized-Leader (FTRL) to have the no-regret property (e.g., Sun et al. (); Swamy et al. ()), making implementation more complicated.) and performs best response computation over $\pi$ via RL (Line 4) under cost function $c^{k}(s)$:

$$ $\displaystyle c^{(k)}(s):=\mathop{\mathbb{E}}_{t\sim U(T)}\mathop{\mathbb{E}}_ {s_{t}\sim q_{t}(\cdot\,|\,s)}\left[\Big{\|}g^{e}(s_{t},t)-\nabla_{s_{t}}\log q _{t}(s_{t}\,|\,s)\Big{\|}_{2}^{2}-\Big{\|}g^{(k)}(s_{t},t)-\nabla_{s_{t}}\log q _{t}(s_{t}\,|\,s)\Big{\|}_{2}^{2}\right].$ (4) $$

###### Remark 1 (Noise-prediction Form of the Cost Function) .

The cost function in Eq. 4 involves the score function $\nabla\log q_{t}(s_{t}\,|\,s)$, which has a closed-form expression for most modern diffusion models including DDPM. Specifically, when the diffusion follows the OU process (Eq. 1), the score function takes the form $\nabla\log q_{t}(s_{t}\,|\,s)=(1-e^{-2t})^{-1}(se^{-t}-s_{t})$, which is exactly the noise added to the original sample $s$ during the diffusion process (recalling that for the OU process, $q_{t}(s_{t}\,|\,s)=\mathcal{N}(se^{-t},(1-e^{-2t})I)$). We denote this noise by $\epsilon$. Then, the cost function (Eq. 4) is equivalent to

$$ $\displaystyle c^{(k)}(s):=\mathop{\mathbb{E}}_{t\sim U(T)}\mathop{\mathbb{E}}_ {\epsilon\sim\mathcal{N}(0,I)}\left[\Big{\|}g^{e}(se^{-t}+\epsilon,t)-\epsilon \Big{\|}_{2}^{2}-\Big{\|}g^{(k)}(se^{-t}+\epsilon,t)-\epsilon\Big{\|}_{2}^{2} \right].$ $$

This closely resembles the noise prediction in DDPM.

To implement FTL, in Line 3, we use the classic idea of Data Aggregation (DAgger) (Ross et al., ), which corresponds to aggregating all states collected from prior learned policies $\pi^{(1)},\dots,\pi^{(k-1)}$, and perform a score-matching least square regression on the aggregated dataset. The RL procedure can take advantage of any modern RL optimizers. In our experiments, we use SAC (Haarnoja et al., ) and DreamerV3 (Hafner et al., ), which serve as representative model-free and model-based RL algorithms.

It is worth to pause for a moment and contrast the procedure we derived above with more traditional inverse RL / adversarial imitation learning procedures. At the highest level, we do not explicitly train a discriminator and instead simply perform score matching on the expert and learner state distributions. Doing so only requires the standard regression-based score-matching optimization, which has shown to be more stable than adversarial training in certain domains. On a more theoretical level, the DS-Divergence we’re optimizing in our procedure is an upper-bound on the worst-case IPM (the total variation distance) one often considers in IRL (Abbeel & Ng, ; Ziebart et al., ; Swamy et al., ). On the positive side, minimizing this upper bound merely requires score matching rather than explicitly training a discriminator.
Score matching gives one stronger theoretical guarantees w.r.t. second-order bounds than IPM-based methods, as we discuss in the following section. Intuitively, the benefit of matching the second moments and not just the first are performance bounds that scale with the variance of the data-generating process, which can be significantly tighter on problems where stochasticity is restricted to a subset of states (e.g. driving on highways, where merges are significantly more variable than actually being on the highway). On the negative side, for problems where we can a-priori specify a class of moments (e.g. autonomous driving, where it is relatively simple to write down the set of features a person could care about when driving), minimizing an upper bound can be needlessly complex. Thus, we believe that as the difficulty of specifying the class of reward functions grows, our method’s benefits over more traditional IPM-based IRL techniques grows as well.

###### Definition 1 (Diffusion Score Divergence) .

###### Remark 1 (Noise-prediction Form of the Cost Function) .

## 5 Theoretical Results

Our goal in this section is to demonstrate that SMILING achieves instance-dependent regret bounds while at the same time avoid compounding errors when model misspecification, optimization error, and statistical error exist. In practice, we can only estimate the diffusion score function up to some statistical error or optimization error due to finite samples. For similar reasons, the RL procedure (4) can only find a near-optimal policy especially when the policy class is not rich enough to capture the expert’s policy. To demonstrate that our algorithm can tolerate these errors, we explicitly study these errors in our theoretical analysis, particually how our regret bound scales with respect to these errors. We provide the following assumptions to formalize these errors.

###### Assumption 1 .

We have the following error bounds, corresponding respectively to 1, 3 and 4 of Algorithm 1:

- (a)
There exists $\epsilon_{\rm{score}}>0$ such that the estimator $g^{e}$ is accurate up to $\epsilon_{\rm{score}}$:
$\displaystyle\mathop{\mathbb{E}}_{s\sim d^{\pi^{e}}}\mathop{\mathbb{E}}_{t\sim
U
(T)}\mathop{\mathbb{E}}_{s_{t}\sim q_{t}(\cdot\,|\,s)}\big{[}\big{\|}g^{e}(s_{
t},t)-\nabla_{s_{t}}\log p^{\pi^{e}}_{t}(s_{t})\big{\|}_{2}^{2}\big{]}\leq
\epsilon_{\rm{score}}^{2};$
- (b)
There exists a function ${\rm{Regret}}(K)$ sublinear in $K$ such that the sequence of score function estimators $\{g^{(k)}\}_{k=1}^{K}$ has regret bounded by ${\rm{Regret}}(K)$:
$\displaystyle\sum_{k=1}^{K}\mathop{\mathbb{E}}_{s\sim d^{\pi^{(k)}}}\mathop{
\mathbb{E}}_{t\sim U(T)}\mathop{\mathbb{E}}_{s_{t}\sim q_{t}(\cdot\,|\,s)}
\left[\left\|g^{(k)}(s_{t},t)-\nabla_{s_{t}}\log q_{t}(s_{t}\,|\,s)\right\|_{2
}^{2}\right]$
$\displaystyle\leq\min_{g}\sum_{k=1}^{K}\mathop{\mathbb{E}}_{s\sim d^{\pi^{(k)}
}}\mathop{\mathbb{E}}_{t\sim U(T)}\mathop{\mathbb{E}}_{s_{t}\sim q_{t}(\cdot\,
|\,s)}\left[\left\|g(s_{t},t)-\nabla_{s_{t}}\log q_{t}(s_{t}\,|\,s)\right\|_{2
}^{2}\right]+{\rm{Regret}}(K);$
- (c)
There exists $\epsilon_{\rm{RL}}>0$ such that, for all $k=1,\dots,K$, the RL procedure finds an $\epsilon_{\rm{RL}}$-optimal policy within some function class $\Pi$:
$\mathop{\mathbb{E}}_{s\sim d^{\pi^{(k)}}}c^{(k)}(s)-\min_{\pi\in\Pi}\mathop{
\mathbb{E}}_{s\sim d^{\pi}}c^{(k)}(s)\leq\epsilon_{\rm{RL}}.$
Note that we do not assume $\pi^{e}\in\Pi$.

Item (a) is a standard assumption in diffusion process and score matching (Chen et al., ; ; Lee et al., ). When the function class $\mathcal{G}$ (hypothesis space) is finite, $\epsilon^{2}_{\rm{score}}$ typically scales at a rate of $\widetilde{O}\left(\ln(|\mathcal{G}|)/N\right)$ in terms of the number of samples $N$. For an infinite function class, more advanced bounds can be established (e.g., see Theorem 4.3 in Oko et al. ()). We emphasize that the discretization error arising from approximating the diffusion SDE using Markov chains is included in $\epsilon_{\rm{score}}$ as well.
Item (b) is similar to Item (a) but for the sequence of score function estimators. It can be satisfied by applying Follow-the-Leader (FTL) since the loss functional for $g$ is a square loss. Other no-regret algorithms, such as Follow-the-regularized-leader (FTRL) or online gradient descent, can also be used.
Under similar conditions as in Item (a), one typically achieves ${\rm{Regret}}(K)=\widetilde{O}(\sqrt{K})$. Item (c) can be satisfied as long as an efficient RL algorithm is applied at each iteration in 4. The rate of $\epsilon_{\rm{RL}}$ has been well-studied in the RL theory literature for various of MDPs. For instance, in tabular MDPs, one can achieve $\epsilon_{\rm{RL}}=\widetilde{O}(\sqrt{SA/M})$ where $S$ is the number of states, $A$ is the number of actions, and $M$ is the number of RL rollout samples. We emphasize that $\epsilon_{\rm{RL}}$ in this case does not scale with $H$ as it is defined as the upper bound on the *averaged* return instead of the cumulative return over $H$ steps.

In the theoretical analysis, we will not assume any misspecification of score functions. Specifically, for any policy $\pi$, we assume that its score function lies in the function class $\mathcal{G}$ (i.e., $\nabla\log p_{t}^{\pi}\in\mathcal{G}$). However, we allow for misspecification in the RL procedure (i.e., $\pi^{e}$ may not be in $\Pi$). This enables us to focus more on the misspecification of the RL procedure under our defined framework.

Next, we introduce the regularity assumptions on the diffusion process, which are standard in the literature on diffusion models

###### Assumption 2 .

For all $\pi$, the following two conditions hold:
(1) the diffusion score function $\nabla\log p_{t}^{\pi}$ is Lipschitz continuous with a finite constant for all $t\in[0,T]$, and
(2) the second moment of $d^{\pi}$ is bounded: $\mathop{\mathbb{E}}_{s\sim d^{\pi}}\big{[}\|s\|_{2}^{2}\big{]}\leq m$ for some $m>0$.

The first condition ensures that the score function behaves well so we can transfer the DS divergence bound to a KL divergence bound.
The second condition is needed for the exponential convergence of the forward diffusion process (i.e., the convergence to the standard Gaussian in terms of KL divergence).
We note that the second condition is readily satisfied in certain simple cases such as when the state space is bounded.

Notably, if we disregard discretization errors, the Lipschitz constant will not appear in our theoretical results as long as it is finite, even if it is arbitrarily large. Similarly, the constant $m$ only appears within logarithmic factors so it can be exponentially large. However, when discretization errors are considered, both the Lipschitz constant and $m$ will appear in $\epsilon_{\rm{score}}$ and ${\rm{Regret}}(T)$ (see, e.g., Chen et al. (; ); Oko et al. ()).

Now we are ready to present our main theoretical result.
Let $\pi^{(1:K)}$ denote the mixture policy of $\{\pi^{(1)},\dots,\pi^{(K)}\}$. Particularly, $\pi^{(1:K)}$ is executed by first choosing a policy $\pi^{(k)}$ uniformly at random and then executing the chosen policy. For any policy $\pi$, we define the variance of its return as ${\rm{Var}}^{\pi}\coloneqq{\rm{Variance}}(\sum_{h=1}^{H}c^{\star}(s_{h}))$. Below is our main theoretical result with proof provided in Appendix B.

###### Theorem 1 .

Under Assumptions 1 and 2, Algorithm 1 achieves the following instance-dependent bounds:

$$ (Second-order) $\displaystyle V^{\pi^{(1:K)}}-V^{\pi^{e}}=\widetilde{O}\left(\sqrt{\min\big{(} {\rm{Var}}^{\pi^{e}},{\rm{Var}}^{\pi^{(1:K)}}\big{)}\cdot\epsilon}+\epsilon H \right);$ (First-order) $\displaystyle V^{\pi^{(1:K)}}-V^{\pi^{e}}=\widetilde{O}\left(\sqrt{\min\big{(} V^{\pi^{e}},V^{\pi^{(1:K)}}\big{)}\cdot\epsilon H}+\epsilon H\right)$ $$

where we define total error

$$ $\epsilon\coloneqq\begin{cases}\epsilon^{2}_{\rm{score}}+\epsilon_{\rm{RL}}+{ \rm{Regret}}(K)/K&\text{if }\pi^{e}\in\Pi,\\ \epsilon^{2}_{\rm{score}}+\epsilon_{\rm{RL}}+{\rm{Regret}}(K)/K+\epsilon_{\rm{ mis}}&\text{if }\pi^{e}\not\in\Pi,\end{cases}$ $$

and the misspecification
$\epsilon_{\rm{mis}}\coloneqq\min_{\pi\in\Pi}\ell(\pi)$ where we recall that $\ell(\cdot)$ is defined in Eq. 2.

Here the misspecification error $\epsilon_{\rm{mis}}$ measures the minimum possible DS divergence to the pre-trained expert score $g^{e}$. We note that $\epsilon_{\rm{mis}}$ is algorithmic-path independent once given $g^{e}$, i.e., it does not depend on what the algorithm does — it is a quantity that is pre-determined when the IL problem is formalized and the expert score $g^{e}$ is pre-trained. If we increase the expressiveness of $\Pi$, $\epsilon_{\rm{mis}}$ will decrease. In contrast, interactive IL algorithms DAgger  (Ross et al., ) and AggreVate(D) (Ross & Bagnell, ; Sun et al., ), which claim to have no compounding errors, actually have algorithmic-path dependent misspecification errors. This means their misspecification error does not necessarily decrease and in fact can increase when one increases the capacity of $\Pi$ (since it affects the algorithm’s behavior).

It is known that the second-order bound is tighter and subsumes the first-order bound (Wang et al., ). Importantly, previous results on second-order bounds for sequential decision-making have relied exclusively on the Maximum Likelihood Estimator (MLE) (Wang et al., ; ; Foster et al., ; Wang et al., ; ),
which in general is not computationally tractable even for simple exponential family distributions (Pabbaraju et al., ).
In this work, we show for the first time that a computationally tractable alternative — diffusion score matching, can achieve these instance-dependent bounds in the context of IRL.
In particular, our second-order bounds (or the first-order bound) scale with the *minimum* of ${\rm{Var}}^{\pi^{e}}$ and ${\rm{Var}}^{\pi^{(1:K)}}$ (or minimum of $V^{\pi^{e}}$ and $V^{\pi^{(1:K)}}$). Hence, our bounds are sharper whenever one of the two is small. In contrast, bounds in these prior RL work do not scale with the min of variances associated with the learned policies and the comparator policy $\pi^{e}$. This subtlety is likely due to the difference between the IRL setting and the general RL setting instead of techniques in the analysis.

###### Assumption 1 .

###### Assumption 2 .

###### Theorem 1 .

#### Improved sample efficiency with respect to the size of expert data.

Let’s now focus on sample complexity with respect to expert data, which is typically the most expensive data to collect in IL, to highlight the benefits of our method over previous approaches. For conciseness, we ignore all errors that are purely related to computation. Specifically, we set $\epsilon_{\text{RL}}\rightarrow 0$, as it only depends on the time spent to run the RL procedure, let the number of iterations $K\rightarrow\infty$ so $\text{Regret}(K)/K\rightarrow 0$, and assume $\pi^{e}\in\Pi$.
Under these conditions, our bound solely depends on $\epsilon_{\text{score}}$. For simplicity, we further assume a finite score function class $\mathcal{G}$, neglect misspecification errors (i.e., $\nabla\log p^{\pi}_{t}(x)\in\mathcal{G}$ for all $\pi$),
and assume the magnitude of the score loss is $O(1)$. Now standard regression analysis (Agarwal et al., ; Oko et al., ) yields $\epsilon_{\text{score}}=O(\sqrt{\log(|\mathcal{G}|)/N})$ where $N$ is the number of expert samples.
Plugging this back into our second-order bound, we obtain

$$ $\displaystyle\frac{1}{K}\sum_{k=1}^{K}V^{\pi^{(k)}}-V^{\pi^{e}}=O\left(\sqrt{ \min\big{(}{\rm{Var}}^{\pi^{e}},{\rm{Var}}^{\pi^{(1:K)}}\big{)}\cdot\frac{\log (|\mathcal{G}|)}{N}}+\frac{H\cdot\log(|\mathcal{G}|)}{N}\right).$ $$

When the expert’s cost or the learned policy’s cost has low variance, i.e., ${\rm{Var}}^{\pi^{e}}\to 0$ or ${\rm{Var}}^{\pi^{(1:K)}}\to 0$, our bound simplifies to $O(H\cdot\log(|\mathcal{G}|)/N)$.
This shows a significant improvement over prior sample complexity bounds for IRL with general function approximation, which are typically $O(H\sqrt{\log(|\mathcal{G}|)/N})$ (e.g., IPM-based methods (Kidambi et al., ; Chang et al., )). This demonstrates the benefit of using diffusion processes and score matching over IPMs.

### 5.1 Comparison to Discriminator-based (f-divergence and IPM) Methods

In this section, we compare DS divergence with $f$-divergence, particularly its discriminator-based approximation (the key metric used by IL algorithms such as GAIL and DAC):

$$ $\displaystyle\max_{f\in\mathcal{F}}\;\;\mathop{\mathbb{E}}_{s\sim p}\big{[} \log f(s)\big{]}+\mathop{\mathbb{E}}_{s\sim q}\big{[}\log(1-f(s))\big{]}$ $$

where $\mathcal{F}\subseteq(0,1)^{\mathcal{S}}$ is a restricted function class. It is known that the above is a lower bound of the Jensen-Shannon (JS) divergence between $p$ and $q$. In particular, when $p/(p+q)\in\mathcal{F}$, it is exactly equal to the JS divergence. However, representing $p/(p+q)$ may be challenging.
Below we show why score-matching is more preferred than discriminator-based $f$-divergence via the example of exponential family distributions,

#### On the expressiveness of score functions and discriminators.

Consider the distributions $p$ and $q$ to be from the exponential family, $p(s)=\exp(w^{T}\phi(s))/Z_{p}$ (and similar for $q$), where $\phi$ is a quadratic feature map, and $Z_{p}$ is the partition function. In this case, the score function $\nabla_{s}\log p(s)=\nabla_{s}\phi(s)$ is simply linear in $s$.
In contrast, the optimal discriminator for JS divergence of $p$ and $q$ has the form of $p/(p+q)$, which is inherently nonlinear and cannot be captured by a linear function on $s$. As a result, if we use linear functions to model discriminators, the discriminator-based objective cannot even faithfully serve as a tight lower bound for the JS divergence. Thus minimizing a loose lower bound of the JS-divergence does not imply minimizing the JS-divergence itself.

Now consider integral probability metric (IPM), i.e., $\max_{f\in\mathcal{F}}\left(\mathop{\mathbb{E}}_{s\sim q}f(s)-\mathop{\mathbb{
E}}_{s\sim p}f(s)\right)$. When $p$ and $q$ are from the exponential family, to perfectly distinguish $q$ and $p$, one needs to design a discriminator in the form of $\theta^{\top}\phi(x)$ since $\phi$ is the sufficient statistics of the exponential family distribution. In our example, $\phi$ is a quadratic feature mapping, meaning that we need a non-linear discriminator for IPM as well.

Hence, score matching only needs relatively weaker score function class to represent complicated distributions to control certain powerful divergence including KL divergence and Hellinger distance. However, for $f$-divergence-based or IPM-based methods, we need a more expressive discriminator function class to make them serve as a tight lower bound of the divergences. This perhaps explains why in practice for generative models, score matching based approach in general outperforms more traditional GAN based approach.
In our experiment, we conduct an ablation study in Section 6.3 to compare our approach to a GAN-based baseline under linear function approximation for score functions and discriminators, and we show that SMILING can still learn well even when the score functions are linear with respect to the state $s$.

#### Can f 𝑓 f italic_f -divergence or IPM-based approaches achieve second-order bounds?

Wang et al. () has demonstrated that simply matching expectations (first moment) is insufficient to achieve second-order bounds in a simpler supervised learning setting. Mapping this to the IL setting, if the true cost is linear in feature $\phi(s)$, i.e., $c^{\star}(s)=\theta^{\star}\cdot\phi(s)$, and we design discriminator $g(s)=\theta^{\top}\phi(s)$, then IL algorithm that learns a policy $\pi$ to only match policy’s expected feature $\mathop{\mathbb{E}}_{s\sim\pi}\phi(s)$ to the expert’s $\mathop{\mathbb{E}}_{s\sim\pi^{e}}\phi(s)$ cannot achieve a second-order regret bound.
For f-divergence such as JS-divergence, unless the discriminator class is rich enough to include the optimal discriminator, the discriminator-based objective is only a loose lower bound of the true $f$-divergence. On the other hand, making the discriminator class rich enough to capture the optimal discriminator can increase both statistical and computational complexity for optimizing the discriminator. Taking the exponential family distribution as an example again, there the optimal discriminator relies on the partition functions $Z_{q}$ and $Z_{p}$ which can make optimization intractable (Pabbaraju et al., ).

## 6 Experiments

We evaluate our method, SMILING, on a set of continuous control tasks from the Deepmind Control Suite (Tassa et al., ) and HumanoidBench (Sferrazza et al., ). These tasks vary in difficulty, as summarized in Table 1. We elaborate on the experimental setup below.
While we will compare to multiple baselines including non-adversarial baselines, we focus on testing the our hypothesis that our score-matching based approach should outperform f-divergence based adversarial IL approach. Additional details can be found in Appendix C.

**Table 1: Spread of evaluated environments. DMC denotes DeepMind Control Suite and HB denotes HumanoidBench.**
| Task | Difficulty | Benchmark |
| --- | --- | --- |
| Ball in Cup Catch | Easy | DMC |
| Cheetah Run | Medium | DMC |
| Humanoid Walk | Hard | DMC |
| Humanoid Sit | Hard | HB |
| Humanoid Crawl | Hard | HB |
| Humanoid Pole | Hard | HB |

Baselines.
We include Discriminator Actor-Critic (DAC), IQ-Learn, and Behavior Cloning (BC) as baselines. Our DAC implementation is based on the original design by Kostrikov et al. () that uses JS-divergence based GAN style objective.
The implementation of IQ-Learn directly uses the official version. We tried several configurations provided in their implementation for the MuJoCo tasks and reported the best one.
The BC performance is reported as the maximum average episode reward obtained across all checkpoints during training. We experimented with both squared loss and MLE for BC and found MLE is better. Consequently, all BC results below are using MLE.

Diffusion Models.
We employ the Denoising Diffusion Probabilistic Model (DDPM) as the diffusion model. The score function is modeled via an MLP with one hidden layer of 256 units, and we use the same architecture for the discriminator in DAC. The diffusion process is discretized into 5K steps and we use a learnable time embedding of 16 dimensions.
We approximate the cost function (Eq. 4) via empirical mean using 500 samples. Since the scale of the score matching loss (and consequently our cost function) is sensitive to the data scale (in this context, the scale of state vector), we normalize the cost of each batch to have zero mean and a standard deviation of 0.1.

RL.
For DMC tasks, we use Soft Actor-Critic (SAC) (Haarnoja et al., ) as the RL algorithm for both DAC and our method, based on the implementation provided by Yarats & Kostrikov (). For HumanoidBench tasks, we adopt DreamerV3 (Hafner et al., ) using the implementation from the HumanoidBench repository (Sferrazza et al., ). Additionally, we maintain a separate “state buffer” that operates the same as the replay buffer but only stores states. This buffer is used to sample states to update the score function in our method and the discriminator in DAC.

Overall, we keep the implementation of SMILING and the baseline DAC as close as possible (e.g., using the same RL solver) so the only difference is different objective functions (i.e., SMILING learns score functions while DAC learns discriminators to approximate JS divergence). This allows us to demonstrate the benefits of our score-matching approach to the discriminator-based GAN-style approaches. For IQ-Learn, we simply use the official implementation.

Expert Policy and Demonstrations.
We use the aforementioned RL algorithms (SAC and DreamerV3) to train expert policies. Expert demonstrations are collected by executing the expert policy in the environment for five trajectories. We set the horizon of all tasks to be 1K without early termination, and thus the resulting dataset consists of 5K steps per task.

### 6.1 Learning from State-only Demonstrations

Figure: Figure 3: Learning curves for learning from state-only data across five random seeds. The x-axis corresponds to the number of environment steps (also the number of policy updates). The y-axis is normalized such that the expert performance is one and the random policy is zero. Our method clearly outperforms all baselines in five tasks out of six, with a substantial performance gap against baselines in four of them.
Refer to caption: x4.png

Figure 3 illustrates the results across five control tasks when the expert demonstrations are state-only. We trained for 0.4M steps on Ball-in-cup-catch and for 3M steps on all other tasks. Our method matches expert performance and outperforms both DAC and IQ-Learn on five of them. Specifically, our method outperforms DAC and IQ-Learn by a large margin on four tasks: ball-in-cup-catch, humanoid-walk, humanoid-crawl, and humanoid-pole. For humanoid-sit, while none achieves expert performance, ours converges faster. For cheetah-run, although DAC initially shows faster convergence, it becomes unstable soon and oscillates significantly after 2M steps. In contrast, our algorithm remains much more stable and surpasses DAC in 3M steps.
We provide video demos of SMILING and DAC on the humanoid-crawl and humanoid-pole tasks in Figure 1 to give a clear visual comparison of their performance differences.

We include BC in the plot as a reference, though it is not directly comparable since it uses state-action demonstrations. It serves as a performance upper bound for Behavioral Cloning from Observation (BCO) (Torabi et al., ), which infers the expert’s actions from state-only data before training. Nevertheless, we observe that both our method and DAC consistently outperform BC (and thus, BCO).

### 6.2 Learning from State-action Demonstrations

Our method can be easily extended to learning from state-action demonstrations by appending the action vector to the state vector for both training and computing rewards. Hence, we also explore its performance in such a setting. The experimental setup is identical to that of learning from state-only data. The training curves are presented in Figure 4. The results are highly consistent with those from the state-only setting, indicating that the performance gap is stable whether training on state data alone or state-action pairs. This also demonstrates that SMILING is robust to different data types.

Figure: Figure 4: Learning curves for learning from state-action data across five random seeds. The x-axis corresponds to the number of environment steps (also the number of policy updates). The y-axis is normalized such that the expert performance is one and the random policy is zero. The results are consistent with those from the state-only setting.
Refer to caption: x5.png

### 6.3 Expressiveness of Diffusion Models

In Section 5.1, we argued that score matching is more expressive than discriminator-based methods, which may explain why our approach generally outperforms DAC in previous experiments. To further support this argument, we conducted an ablation study on the cheetah-run task where our method previously showed a slower convergence compared to DAC.
Now we remove the activation functions in both the discriminator of DAC and the score function of our method, making them purely linear, and then re-run the experiments.
The results are in Figure 5. Different from the previous experiments where DAC outperformed our method, in this case, DAC’s performance degrades notably by showing slower convergence and significantly higher variance. In comparison, our method remains more effective and outperforms DAC clearly.

We attribute the performance degradation of DAC to the limited expressiveness of the discriminator, which results in mode collapse—a well-known issue in discriminator-based methods, originally identified in GANs. In such methods, if the discriminator’s performance falls significantly behind the generator (e.g., due to insufficient training or a lack of discriminator expressiveness, which is the case here), the generator will exploit this gap and produce single-mode data

Figure: Figure 5: Learning curves with linear functions. Compared to using MLPs, DAC shows a significant decline in performance of slower convergence and higher variance, while ours remains more effective.
Refer to caption: x6.png

that can easily fool the discriminator. In this case, the generator is falling into a local minima.
As the discriminator gradually catches up and learns to identify this data, the generator will discover new data to fool the updated discriminator, falling into another local minimum. This perpetual cycle of chasing leads to continuous oscillations in the training process. We conjecture this is the exact reason for DAC’s instability in this experiments—With the discriminator reduced to a purely linear function, its expressiveness is significantly weakened, allowing the RL policy to exploit it easily, and thus a perpetual chasing emerges. However, our method can remain more effective even with linear score functions, which aligns with our hypothesis in Section 5.1.

## 7 Conclusion

In this paper, we propose a new imitation learning framework, Score-Matching Imitation Learning (SMILING), that leverages the diffusion score function to learn from expert demonstrations. Different from previous methods, our approach is not formulated from any $f$-divergence or IPM perspective. Instead, we directly match the score function of the learned policy to that of the expert policy.
Theoretically, we show that our score function is more expressive than learning a discriminator on $f$-divergence in some sense.
Furthermore, we show that our method can achieve first- and second-order bounds in sample complexity, which are instance-dependent and tighter than previous results. In practice, we demonstrate that our method outperforms DAC and IQ-Learn on a set of continuous control tasks, especially on harder tasks.

## Appendix A Error Analysis of Objective

In this section, we show the following two things: (1) why directly replacing the diffusion score function of $\pi$ with an estimator $g^{\pi}$ in $\ell(\pi)$ (Eq. 2) incurs large error; (2) why the error is small in our derived objective $\widehat{\ell}(\pi)$

On the one hand, replacing the score function of $\pi$ with $g^{\pi}$ incurs the following error:

$$ $\displaystyle\mathop{\mathbb{E}}_{s\sim d^{\pi}}\mathop{\mathbb{E}}_{t\sim U(T )}\mathop{\mathbb{E}}_{s_{t}\sim q_{t}(\cdot\,|\,s)}\left[\left\|g^{e}(s_{t},t )-g^{\pi}(s_{t})\right\|_{2}^{2}-\left\|g^{e}(s_{t},t)-\nabla_{s_{t}}\log p_{t }^{\pi}(s_{t})\right\|_{2}^{2}\right]$ $\displaystyle=\mathop{\mathbb{E}}_{s\sim d^{\pi}}\mathop{\mathbb{E}}_{t\sim U( T)}\mathop{\mathbb{E}}_{s_{t}\sim q_{t}(\cdot\,|\,s)}\Big{[}\underbrace{\left( \nabla_{s_{t}}\log p_{t}^{\pi}(s_{t})-g^{\pi}(s_{t})\right)}_{(a)}\underbrace{ \left(2g^{e}(s_{t},t)-g^{\pi}(s_{t})-\nabla_{s_{t}}\log p_{t}^{\pi}(s_{t}) \right)}_{(b)}\Big{]}.$ $$

Here the scale of term (a) is bounded by the statistical error from score matching. However, the scale of term (b) is unbounded since $g^{e}$ can arbitrarily deviate from both $g^{\pi}$ and $\nabla_{s_{t}}\log p_{t}^{\pi}(s_{t})$. Hence, the error is potentially unbounded.

On the other hand, our derived objective is bounded. To see this, we note that our objective incur the following error:

$$ $\displaystyle\mathop{\mathbb{E}}_{s\sim d^{\pi}}\mathop{\mathbb{E}}_{t\sim U(T )}\mathop{\mathbb{E}}_{s_{t}\sim q_{t}(\cdot\,|\,s)}\left[\left\|\nabla_{s_{t} }\log p_{t}^{\pi}(s_{t})-\nabla_{s_{t}}\log q_{t}(s_{t}\,|\,s)\right\|_{2}^{2} -\left\|g^{\pi}(s_{t})-\nabla_{s_{t}}\log q_{t}(s_{t}\,|\,s)\right\|_{2}^{2}\right]$ $\displaystyle=\mathop{\mathbb{E}}_{s\sim d^{\pi}}\mathop{\mathbb{E}}_{t\sim U( T)}\mathop{\mathbb{E}}_{s_{t}\sim q_{t}(\cdot\,|\,s)}\left[\left\|\nabla_{s_{t }}\log p_{t}^{\pi}(s_{t})-g^{\pi}(s_{t})\right\|_{2}^{2}\right].$ $$

Here the equality is by Lemma 4. We observe that it is exactly bounded by the statistical error from score matching.

## Appendix B Proof of Theorem 1

### B.1 Supporting Lemmas

The following lemma is from Wang et al. () and the fact that triangular discrimination is equivalent to the squared Hellinger distance up to a multiplicative constant: $2D_{H}^{2}\leq D_{\triangle}\leq 4D_{H}^{2}$(Topsoe, ).

###### Lemma 1 .

(Wang et al., , Lemma 4.3)
For two distributions $p,q\in[0,1]$ over some random variable, denote their respective means by $\overline{p}$ and $\overline{q}$. Then, it holds that

$$ $\displaystyle\left|\overline{p}-\overline{q}\right|\leq 8\sqrt{{\rm{Var}}(p)D^ {2}_{H}(p,q)}+20D^{2}_{H}(p,q)$ $$

where ${\rm{Var}}(\cdot)$ denotes the variance, and $D^{2}_{H}(p,q)\coloneqq\int(\sqrt{p(x)}-\sqrt{q(x)})^{2}\mathop{}\mathrm{d}x/2$ is the squared Hellinger distance.

The following lemma is adapted from some known results from diffusion processes and score matching. It shows that as long as the DS divergence (Definition 1) is small, the Hellinger distance is also small under some mild conditions. Similar results can be found in some theoretical papers on diffusion models such as Chen et al. (); Oko et al. ().

###### Lemma 2 .

Let $P\in\Delta(\mathbb{R}^{d})$ and $g\in\mathbb{R}^{d}\times[0,T]\rightarrow\mathbb{R}^{d}$. Define $Q\in\Delta(\mathbb{R}^{d})$ as the distribution obtained through the reverse process of diffusion starting from $\mathcal{N}(0,I)$ by treating $g$ as the “score function” of the process. Specifically, $Q$ is the distribution of $\mathbf{z}_{T}$ of the following SDE:

$$ $\displaystyle\mathop{}\mathrm{d}\mathbf{z}_{t}=$ $\displaystyle(\mathbf{z}_{t}+2g(\mathbf{z}_{t},T-t))\mathop{}\mathrm{d}t+\sqrt {2}\mathop{}\mathrm{d}B_{t},\,\mathbf{z}_{0}\sim\mathcal{N}(0,I).$ $$

We assume the following:

- (a)
For all $t\geq 0$, the score $\nabla\log P_{t}$ is Lipschitz continuous with finite constant;
- (b)
The second moment of $P$ is upper bounded: $\mathop{\mathbb{E}}_{\mathbf{x}\sim P}\big{[}\left\|\mathbf{x}\right\|_{2}^{2}
\big{]}\leq m$ for some $m>0$.

Then, if

$$ $\displaystyle\mathop{\mathbb{E}}_{\mathbf{x}\sim P}\mathop{\mathbb{E}}_{t\sim U (T)}\mathop{\mathbb{E}}_{\mathbf{x}_{t}\sim q_{t}(\cdot\,|\,\mathbf{x})}\left[ \left\|g(\mathbf{x}_{t},t)-\nabla_{\mathbf{x}_{t}}\log P_{t}(\mathbf{x}_{t}) \right\|_{2}^{2}\right]\leq\epsilon^{2},$ $$

for some $\epsilon>0$, we have

$$ $\displaystyle D_{H}^{2}(P,Q)=O\left(T\epsilon^{2}+(m+d)\exp(-T)\right).$ $$

###### Proof of Lemma 2 .

We consider the following three stochastic processes specified by SDEs:

$$ $\displaystyle\mathop{}\mathrm{d}\mathbf{x}_{t}=$ $\displaystyle(\mathbf{x}_{t}+2\nabla\log P_{T-t}(\mathbf{x}_{t}))\mathop{} \mathrm{d}t+\sqrt{2}\mathop{}\mathrm{d}B_{t},\,\mathbf{x}_{0}\sim P_{T};$ $\displaystyle\mathop{}\mathrm{d}\mathbf{y}_{t}=$ $\displaystyle(\mathbf{y}_{t}+2g(\mathbf{y}_{t},T-t))\mathop{}\mathrm{d}t+\sqrt {2}\mathop{}\mathrm{d}B_{t},\,\mathbf{y}_{0}\sim P_{T};$ $\displaystyle\mathop{}\mathrm{d}\mathbf{z}_{t}=$ $\displaystyle(\mathbf{z}_{t}+2g(\mathbf{z}_{t},T-t))\mathop{}\mathrm{d}t+\sqrt {2}\mathop{}\mathrm{d}B_{t},\,\mathbf{z}_{0}\sim\mathcal{N}(0,I).$ $$

To clarify the notation: the variables $\mathbf{x}_{t},\mathbf{y}_{t},\mathbf{z}_{t}$ above are defined in the reverse time order, whereas in the lemma statement, $\mathbf{x}$ follow forward time order.

By triangle inequality for Hellinger distance:

$$ $\displaystyle D_{H}^{2}(P,Q)$ $\displaystyle=D_{H}^{2}(\mathbf{x}_{T},\mathbf{z}_{T})\leq\big{(}D_{H}(\mathbf {x}_{T},\mathbf{y}_{T})+D_{H}(\mathbf{y}_{T},\mathbf{z}_{T})\big{)}^{2}$ $\displaystyle\leq 2D_{H}^{2}(\mathbf{x}_{T},\mathbf{y}_{T})+2D_{H}^{2}(\mathbf {y}_{T},\mathbf{z}_{T}).$ $$

We will bound the two terms separately. Since the squared Hellinger distance is upper bounded by KL divergence (i.e., $D_{H}^{2}\leq D_{\rm{KL}}$), we seek to establish the KL divergence bounds instead.

First, by Item (a) and Girsanov’s Theorem (Karatzas & Shreve, ) (also see, e.g., Chen et al. (); Oko et al. ()), we have

$$ $\displaystyle D_{H}^{2}(\mathbf{x}_{T},\mathbf{y}_{T})\leq D_{\rm{KL}}(\mathbf {x}_{T}\,\|\,\mathbf{y}_{T})=O\left(T\cdot\mathop{\mathbb{E}}_{t\sim U(T)}\Big {\|}\nabla\log P_{T-t}(\mathbf{x}_{t})-g(\mathbf{x}_{t},T-t)\Big{\|}_{2}^{2} \right)=O\left(T\epsilon^{2}\right).$ $$

Next, by Item (b) and the convergence of OU process (e.g., Lemma 9 in Chen et al. ()), we have:

$$ $\displaystyle D_{\rm{KL}}(P_{T}\,\|\,\mathcal{N}(0,I))$ $\displaystyle=O\big{(}(m+d)\cdot\exp(-T)\big{)}.$ $$

By data-processing inequality for $f$-divergence, we have

$$ $\displaystyle D_{H}^{2}(\mathbf{y}_{T},\mathbf{z}_{T})$ $\displaystyle\leq D_{H}^{2}(P_{T},\mathcal{N}(0,I))$ $\displaystyle\leq D_{\rm{KL}}(P_{T}\,\|\,\mathcal{N}(0,I))$ $\displaystyle=O\big{(}(m+d)\cdot\exp(-T)\big{)}.$ $$

Plugging them back, we complete the proof.
∎

The following lemmas are standard results from diffusion processes and score matching. We show them here for completeness.

###### Lemma 3 .

Given any $g$, we have

$$ $\displaystyle\mathop{\mathbb{E}}_{\mathbf{x}_{t}\sim p_{t}}\Big{\langle}g( \mathbf{x}_{t},t),\nabla_{\mathbf{x}_{t}}\log q_{t}(\mathbf{x}_{t})\Big{ \rangle}=\mathop{\mathbb{E}}_{\mathbf{x}_{0}\sim p_{0}}\mathop{\mathbb{E}}_{ \mathbf{x}_{t}\sim p_{t}(\mathbf{x}_{t}\,|\,\mathbf{x}_{0})}\Big{\langle}g( \mathbf{x}_{t},t),\nabla_{\mathbf{x}_{t}}\log q_{t}(\mathbf{x}_{t}\,|\,\mathbf {x}_{0})\Big{\rangle}$ $$

where $\langle\cdot,\cdot\rangle$ denotes the inner product.

###### Proof of Lemma 3 .

It basically follows from integration by parts. We start from the left-hand side (LHS):

$$ LHS $\displaystyle=\int_{\mathbf{x}_{t}}g(\mathbf{x}_{t},t)\cdot\nabla_{\mathbf{x}_ {t}}\log p_{t}(\mathbf{x}_{t})\cdot p_{t}(x_{t})\mathop{}\mathrm{d}\mathbf{x}_ {t}$ $\displaystyle=\int_{\mathbf{x}_{t}}g(\mathbf{x}_{t},t)\cdot\nabla_{\mathbf{x}_ {t}}p_{t}(\mathbf{x}_{t})\mathop{}\mathrm{d}\mathbf{x}_{t}$ $\displaystyle=0-\int_{\mathbf{x}_{t}}\nabla_{\mathbf{x}_{t}}\cdot g(\mathbf{x} _{t},t)\cdot p_{t}(\mathbf{x}_{t})\mathop{}\mathrm{d}\mathbf{x}_{t}$ (integration by parts) $\displaystyle=-\int_{\mathbf{x}_{0}}\int_{\mathbf{x}_{t}}\nabla_{\mathbf{x}_{t }}\cdot g(\mathbf{x}_{t},t)\cdot p_{t}(\mathbf{x}_{t}\,|\,\mathbf{x}_{0})\cdot p _{0}(\mathbf{x}_{0})\mathop{}\mathrm{d}\mathbf{x}_{t}\mathop{}\mathrm{d} \mathbf{x}_{0}$ $\displaystyle=0+\int_{\mathbf{x}_{0}}\int_{\mathbf{x}_{t}}g(\mathbf{x}_{t},t) \cdot\nabla_{\mathbf{x}_{t}}p_{t}(\mathbf{x}_{t}\,|\,\mathbf{x}_{0})\cdot p_{0 }(\mathbf{x}_{0})\mathop{}\mathrm{d}\mathbf{x}_{t}\mathop{}\mathrm{d}\mathbf{x }_{0}$ (integration by parts again) $\displaystyle=\int_{\mathbf{x}_{0}}\int_{\mathbf{x}_{t}}g(\mathbf{x}_{t},t) \cdot\nabla_{\mathbf{x}_{t}}\log p_{t}(\mathbf{x}_{t}\,|\,\mathbf{x}_{0})\cdot p _{t}(\mathbf{x}_{t}\,|\,\mathbf{x}_{0})\cdot p_{0}(\mathbf{x}_{0})\mathop{} \mathrm{d}\mathbf{x}_{t}\mathop{}\mathrm{d}\mathbf{x}_{0}$ $\displaystyle=\text{RHS}.$ $$

Hence, the lemma is proved.
∎

###### Lemma 4 .

Given any $g$, we have

$$ $\displaystyle\mathop{\mathbb{E}}_{\mathbf{x}_{t}\sim p_{t}}\Big{\|}g(\mathbf{x }_{t},t)-\nabla_{\mathbf{x}_{t}}\log p_{t}(\mathbf{x}_{t})\Big{\|}^{2}_{2}- \mathop{\mathbb{E}}_{\mathbf{x}_{0}\sim p_{0},\mathbf{x}_{t}\sim p_{t}(\mathbf {x}_{t}\,|\,\mathbf{x}_{0})}\Big{\|}g(\mathbf{x}_{t},t)-\nabla_{\mathbf{x}_{t} }\log p_{t}(\mathbf{x}_{t}\,|\,\mathbf{x}_{0})\Big{\|}^{2}_{2}$ $\displaystyle\qquad=\mathop{\mathbb{E}}_{\mathbf{x}_{t}}\Big{\|}\nabla_{ \mathbf{x}_{t}}\log p_{t}(\mathbf{x}_{t})\Big{\|}^{2}_{2}-\mathop{\mathbb{E}}_ {\mathbf{x}_{0},\mathbf{x}_{t}}\Big{\|}\nabla_{\mathbf{x}_{t}}\log p_{t}( \mathbf{x}_{t}\,|\,\mathbf{x}_{0})\Big{\|}^{2}_{2}$ $$

where we note that the right-hand side of the equation is independent of $g$.

###### Proof of Lemma 4 .

We prove the lemma by expanding the first term on the left-hand side:

$$ $\displaystyle\mathop{\mathbb{E}}_{\mathbf{x}_{t}\sim p_{t}}\Big{\|}g(\mathbf{x }_{t},t)-\nabla_{\mathbf{x}_{t}}\log p_{t}(\mathbf{x}_{t})\Big{\|}^{2}_{2}$ $\displaystyle=\mathop{\mathbb{E}}_{\mathbf{x}_{t}\sim p_{t}}\Big{\|}g(\mathbf{ x}_{t},t)\Big{\|}^{2}_{2}-2\mathop{\mathbb{E}}_{\mathbf{x}_{t}\sim p_{t}}\Big{ \langle}g(\mathbf{x}_{t},t),\nabla_{\mathbf{x}_{t}}\log p_{t}(\mathbf{x}_{t}) \Big{\rangle}+\mathop{\mathbb{E}}_{\mathbf{x}_{t}\sim p_{t}}\Big{\|}\nabla_{ \mathbf{x}_{t}}\log p_{t}(\mathbf{x}_{t})\Big{\|}^{2}_{2}$ $\displaystyle=\mathop{\mathbb{E}}_{\mathbf{x}_{0},\mathbf{x}_{t}}\Big{\|}g( \mathbf{x}_{t},t)\Big{\|}^{2}_{2}-2\mathop{\mathbb{E}}_{\mathbf{x}_{0},\mathbf {x}_{t}}\Big{\langle}g(\mathbf{x}_{t},t),\nabla_{\mathbf{x}_{t}}\log p_{t}( \mathbf{x}_{t}\,|\,\mathbf{x}_{0})\Big{\rangle}+\mathop{\mathbb{E}}_{\mathbf{x }_{t}}\Big{\|}\nabla_{\mathbf{x}_{t}}\log p_{t}(\mathbf{x}_{t})\Big{\|}^{2}_{2}$ ( Lemma 3 ) $\displaystyle=\mathop{\mathbb{E}}_{\mathbf{x}_{0},\mathbf{x}_{t}}\Big{\|}g( \mathbf{x}_{t},t)\Big{\|}^{2}_{2}-2\mathop{\mathbb{E}}_{\mathbf{x}_{0},\mathbf {x}_{t}}\Big{\langle}g(\mathbf{x}_{t},t),\nabla_{\mathbf{x}_{t}}\log p_{t}( \mathbf{x}_{t}\,|\,\mathbf{x}_{0})\Big{\rangle}+\mathop{\mathbb{E}}_{\mathbf{x }_{0},\mathbf{x}_{t}}\Big{\|}\nabla_{\mathbf{x}_{t}}\log p_{t}(\mathbf{x}_{t} \,|\,\mathbf{x}_{0})\Big{\|}^{2}_{2}$ $\displaystyle\qquad-\mathop{\mathbb{E}}_{\mathbf{x}_{0},\mathbf{x}_{t}}\Big{\| }\nabla_{\mathbf{x}_{t}}\log p_{t}(\mathbf{x}_{t}\,|\,\mathbf{x}_{0})\Big{\|}^ {2}_{2}+\mathop{\mathbb{E}}_{\mathbf{x}_{t}}\Big{\|}\nabla_{\mathbf{x}_{t}} \log p_{t}(\mathbf{x}_{t})\Big{\|}^{2}_{2}$ $\displaystyle=\mathop{\mathbb{E}}_{\mathbf{x}_{0},\mathbf{x}_{t}}\Big{\|}g( \mathbf{x}_{t},t)-\nabla_{\mathbf{x}_{t}}\log p_{t}(\mathbf{x}_{t}\,|\,\mathbf {x}_{0})\Big{\|}^{2}_{2}-\mathop{\mathbb{E}}_{\mathbf{x}_{0},\mathbf{x}_{t}} \Big{\|}\nabla_{\mathbf{x}_{t}}\log p_{t}(\mathbf{x}_{t}\,|\,\mathbf{x}_{0}) \Big{\|}^{2}_{2}+\mathop{\mathbb{E}}_{\mathbf{x}_{t}}\Big{\|}\nabla_{\mathbf{x }_{t}}\log p_{t}(\mathbf{x}_{t})\Big{\|}^{2}_{2}.$ $$

This completes the proof.
∎

###### Lemma 1 .

###### Lemma 2 .

###### Proof of Lemma 2 .

###### Lemma 3 .

###### Proof of Lemma 3 .

###### Lemma 4 .

###### Proof of Lemma 4 .

### B.2 Statistical results

For the ease of presentation, we will leverage the following quantities:

$$ $\displaystyle\mathcal{L}(\pi,g)=\mathop{\mathbb{E}}_{s\sim d^{\pi}}\mathop{ \mathbb{E}}_{t\sim U(T)}\mathop{\mathbb{E}}_{s_{t}\sim q_{t}(\cdot\,|\,s)} \left[\left\|g^{e}(s_{t},t)-\nabla_{s_{t}}\log q_{t}(s_{t}\,|\,s)\right\|_{2}^ {2}-\left\|g(s_{t},t)-\nabla_{s_{t}}\log q_{t}(s_{t}\,|\,s)\right\|_{2}^{2}\right]$ $$

By Lemma 4, the above is equivalent to the following by replacing the conditional score functions with the marginal ones for both terms in the expectation:

$$ $\displaystyle\mathcal{L}(\pi,g)=\mathop{\mathbb{E}}_{s\sim d^{\pi}}\mathop{ \mathbb{E}}_{t\sim U(T)}\mathop{\mathbb{E}}_{s_{t}\sim q_{t}(\cdot\,|\,s)} \left[\left\|g^{e}(s_{t},t)-\nabla_{s_{t}}\log p_{t}^{\pi}(s_{t})\right\|_{2}^ {2}-\left\|g(s_{t},t)-\nabla_{s_{t}}\log p_{t}^{\pi}(s_{t})\right\|_{2}^{2} \right].$ (5) $$

One can show that, fixing a policy $\pi$, we have

$$ $\max_{g}\mathcal{L}(\pi,g)=\mathop{\mathbb{E}}_{s\sim d^{\pi}}\mathop{\mathbb{ E}}_{t\sim U(T)}\mathop{\mathbb{E}}_{s_{t}\sim q_{t}(\cdot\,|\,s)}\left[\left \|g^{e}(s_{t},t)-\nabla_{s_{t}}\log p_{t}^{\pi}(s_{t})\right\|_{2}^{2}\right]= \ell(\pi)$ (6) $$

where the first equality is by observing that the second term in Eq. (5) can be minimized to zero by choosing $g=\nabla_{s_{t}}\log p_{t}^{\pi}(s_{t})$ (recalling that we do not assume misspecification for score functions), and the second equality is by definition.

Bounding the difference between $\pi^{(1:K)}$ and $\pi^{e}$ is what we aim to do in the following lemma.

###### Lemma 5 .

Under Assumptions 1 and 2, assuming misspecification of RL (i.e., $\pi^{e}$ may not be in $\Pi$), we have

$$ $\displaystyle D_{H}^{2}(d^{\pi^{(1:K)}},d^{\pi^{e}})=O\bigg{(}T\cdot\epsilon_{ \rm{score}}^{2}+T\cdot\epsilon_{\rm{mis}}+T\cdot\epsilon_{\rm{RL}}+\frac{T \cdot{\rm{Regret}}(K)}{K}+(m+d)\exp(-T)\bigg{)}.$ $$

Setting $T=\log(\frac{m+d}{\epsilon_{\rm{score}}^{2}+\epsilon_{\rm{RL}}+\epsilon_{\rm{
mis}}})$ yields the following bound:

$$ $\displaystyle D_{H}^{2}(d^{\pi^{(1:K)}},d^{\pi^{e}})=\widetilde{O}\bigg{(} \epsilon_{\rm{score}}^{2}+\epsilon_{\rm{mis}}+\epsilon_{\rm{RL}}+\frac{{\rm{ Regret}}(K)}{K}\bigg{)}$ $$

where $\widetilde{O}$ hides logarithmic factors.

###### Proof of Lemma 5 .

Items (b) and (c) in Assumption 1 implies the following:

$$ $\max_{g}\sum_{k=1}^{K}\mathcal{L}(\pi^{(k)},g)\leq\sum_{k=1}^{K}\mathcal{L}( \pi^{(k)},g^{(k)})+{\rm{Regret}}(K);$ (7) $$

$$ $\mathcal{L}(\pi^{(k)},g^{(k)})-\min_{\pi\in\Pi}\mathcal{L}(\pi,g^{(k)})\leq \epsilon_{\rm{RL}}.$ (8) $$

Combining Equation 7 and Equation 8, we have

$$ $\displaystyle\max_{g}\sum_{k=1}^{K}\mathcal{L}(\pi^{(k)},g)\leq\sum_{k=1}^{K} \min_{\pi\in\Pi}\mathcal{L}(\pi,g^{(k)})+K\cdot\epsilon_{\rm{RL}}+{\rm{Regret} }(K).$ $$

Now the right-hand side can be upper bounded by replacing all $g^{(k)}$ with maximization over $g$:

$$ $\displaystyle\max_{g}\sum_{k=1}^{K}\mathcal{L}(\pi^{(k)},g)$ $\displaystyle\leq\sum_{k=1}^{K}\min_{\pi\in\Pi}\max_{g}\mathcal{L}(\pi,g)+K \cdot\epsilon_{\rm{RL}}+{\rm{Regret}}(K)$ $\displaystyle=\sum_{k=1}^{K}\min_{\pi\in\Pi}\ell(\pi)+K\cdot\epsilon_{\rm{RL}} +{\rm{Regret}}(K)$ (by Eq. 6 ) $\displaystyle=K\cdot\epsilon_{\rm{mis}}+K\cdot\epsilon_{\rm{RL}}+{\rm{Regret}} (K)$ (9) $$

Note that the left-hand side is equivalent to the following

$$ $\displaystyle\max_{g}\sum_{k=1}^{K}\mathcal{L}(\pi^{(k)},g)=K\cdot\max_{g} \mathcal{L}(\pi^{(1:K)},g)=K\cdot\ell(\pi^{(1:K)}).$ $$

where the first equality is by the definition of $\mathcal{L}$.
Inserting it back to the previous inequality yields an upper bound on $\ell(\pi^{(1:K)})$:

$$ $\displaystyle\ell(\pi^{(1:K)})\leq\epsilon_{\rm{mis}}+\epsilon_{\rm{RL}}+\frac {{\rm{Regret}}(K)}{K}$ $$

Now we define a new distribution to facilitate the analysis: let $\widehat{d}^{\pi^{e}}$ denote the distribution induced by the reverse diffusion process starting from $\mathcal{N}(0,I)$ by treating $g^{e}$ as the “score function”. Specifically, $\widehat{d}^{\pi^{e}}$ is the distribution of $\overline{s}_{T}$ of the following SDE:

$$ $\displaystyle\mathop{}\mathrm{d}\overline{s}_{t}=$ $\displaystyle(\overline{s}_{t}+2g^{e}(\overline{s}_{t},T-t))\mathop{}\mathrm{d }t+\sqrt{2}\mathop{}\mathrm{d}B_{t},\,\overline{s}_{0}\sim\mathcal{N}(0,I)$ $$

where we recall that $B_{t}$ is the standard Brownian motion. We can see that $\widehat{d}^{\pi^{e}}$ is approximating $d^{\pi^{e}}$. Then, we invoke Lemma 2 and get

$$ $\displaystyle D_{H}^{2}(d^{\pi^{(1:K)}},\widehat{d}^{\pi^{e}})=O\bigg{(}T\cdot \epsilon_{\rm{mis}}+T\cdot\epsilon_{\rm{RL}}+\frac{T\cdot{\rm{Regret}}(K)}{K}+ (m+d)\exp(-T)\bigg{)}.$ $$

Now it remains to bound $D_{H}^{2}(d^{\pi^{e}},\widehat{d}^{\pi^{e}})$ so that we can apply the triangle inequality to get an upper bound on $D_{H}^{2}(d^{\pi^{(1:K)}},d^{\pi^{e}})$.
Since $\ell(\pi^{e})\leq\epsilon_{\rm{score}}^{2}$ (Item (a) in Assumption 1), we invoke Lemma 2 again and immediately get

$$ $\displaystyle D_{H}^{2}(d^{\pi^{e}},\widehat{d}^{\pi^{e}})\leq O\bigg{(}T\cdot \epsilon_{\rm{score}}^{2}+(m+d)\exp(-T)\bigg{)}.$ $$

We conclude the proof by putting them together via triangle inequality:

$$ $\displaystyle D_{H}^{2}(d^{\pi^{(1:K)}},d^{\pi^{e}})$ $\displaystyle\leq 2D_{H}^{2}(d^{\pi^{(1:K)}},\widehat{d}^{\pi^{e}})+2D_{H}^{2} (\widehat{d}^{\pi^{e}},d^{\pi^{e}})$ $\displaystyle=O\bigg{(}T\cdot\epsilon_{\rm{mis}}+T\cdot\epsilon_{\rm{score}}^{ 2}+T\cdot\epsilon_{\rm{RL}}+\frac{T\cdot{\rm{Regret}}(K)}{K}+(m+d)\exp(-T) \bigg{)}.$ $$

∎

###### Lemma 6 .

Under the same condition as Lemma 5 but assuming no misspecification (i.e., $\pi^{e}\in\Pi$), we have

$$ $\displaystyle D_{H}^{2}(d^{\pi^{(1:K)}},d^{\pi^{e}})=\widetilde{O}\bigg{(} \epsilon_{\rm{score}}^{2}+\epsilon_{\rm{RL}}+\frac{{\rm{Regret}}(K)}{K}\bigg{)}$ $$

###### Proof of Lemma 6 .

The proof is almost identical to that of Lemma 5 except that, when there is no misspecification, $\min_{\pi\in\Pi}\ell(\pi)$ in Eq. 9 can be simply upper bounded by $\ell(\pi^{e})$, which is bounded by $\epsilon_{\rm{score}}^{2}$ (Item (a) in Assumption 1). The rest of the proof is the same.
∎

###### Lemma 5 .

###### Proof of Lemma 5 .

###### Lemma 6 .

###### Proof of Lemma 6 .

### B.3 Main proof

In this section, we provide the main proof of Theorem 1. First, the following is by definition of value functions:

$$ $\displaystyle V^{\pi^{(1:K)}}-V^{\pi^{e}}$ $\displaystyle=H\cdot\left(\mathop{\mathbb{E}}_{s,a\sim d^{\pi^{(1:K)}}}c^{ \star}(s,a)-\mathop{\mathbb{E}}_{s,a\sim d^{\pi^{e}}}c^{\star}(s,a)\right)$ $$

By Lemma 1, it is bounded by

$$ $\displaystyle V^{\pi^{(1:K)}}-V^{\pi^{e}}$ (10) $\displaystyle\leq H\cdot\left(8\sqrt{{\rm{Var}}\left(\mathop{\mathbb{E}}_{s,a \sim d^{\pi^{e}}}c^{\star}(s,a)\right)\cdot D^{2}_{H}(c^{\star}\sim d^{\pi^{(1 :K)}},c^{\star}\sim d^{\pi^{e}})}+20D^{2}_{H}(c^{\star}\sim d^{\pi^{(1:K)}},c^ {\star}\sim d^{\pi^{e}})\right)$ $\displaystyle=8\sqrt{{\rm{Var}}^{\pi^{e}}\cdot D^{2}_{H}(c^{\star}\sim d^{\pi^ {(1:K)}},c^{\star}\sim d^{\pi^{e}})}+20HD^{2}_{H}(c^{\star}\sim d^{\pi^{(1:K)} },c^{\star}\sim d^{\pi^{e}})$ $$

where $c^{\star}\sim d^{\pi^{(1:K)}}$ (and $c^{\star}\sim d^{\pi^{e}}$) denotes the distribution of cost rather than the distribution of $d^{\pi^{(1:K)}}$ (and $d^{\pi^{e}}$) itself. Now we apply the data processing inequality and get

$$ $\displaystyle V^{\pi^{(1:K)}}-V^{\pi^{e}}$ $\displaystyle\leq 8\sqrt{{\rm{Var}}^{\pi^{e}}\cdot D_{H}^{2}(d^{\pi^{(1:K)}},d ^{\pi^{e}})}+20HD_{H}^{2}(d^{\pi^{(1:K)}},d^{\pi^{e}}).$ $$

Plugging in the bound from Lemma 5 if there is misspecification error or Lemma 6 otherwise, we get

$$ $\displaystyle V^{\pi^{(1:K)}}-V^{\pi^{e}}$ $\displaystyle=\widetilde{O}\bigg{(}\sqrt{{\rm{Var}}^{\pi^{e}}\cdot\epsilon}+ \epsilon H\bigg{)}.$ $$

Similarly, we can derive the second half of the second-order bound (with $V^{\pi^{(1:K)}}$ replaced with $V^{\pi^{e}}$ inside the square root) by applying Lemma 1 to Eq. 10 in the other direction. These conclude the proof of the second-order bound.

The first-order bound is a direct consequence of the second-order bound as shown in Theorem 2.1 in Wang et al. ().

## Appendix C Details of Experiments

Implementation of RL Algorithms.
We use the SAC implementation from Yarats & Kostrikov () and the DreamerV3 implementation from Sferrazza et al. () without modifying their hyperparameters. The only change we made is substituting their ground-truth cost (reward) functions with our cost function. For example, whenever the SAC agent samples a batch of transitions from the replay buffer, we replace the corresponding cost values with those defined in our algorithm. We do the same for DreamerV3.

Implementation of Our Method and DAC.
We made efforts to align the implementation details of both algorithms to ensure a fair comparison. A comprehensive list of hyperparameters is provided in Table 2. We perform one update of the score function in our method and the discriminator in DAC per 1K RL steps. The expert diffusion model is trained for 4K epochs using 5K expert states (and actions, if learning from state-action pairs).

Implementation of IQ-Learn.
We used the official implementation from Garg et al. (). We note that, unlike SMILING and DAC that can be built on any RL algorithm, IQ-Learn is exclusively tied to SAC. This makes it unclear how to implement it on top of DreamerV3, as we did for SMILING and DAC. Therefore, we adhered to their SAC implementation to run all experiments including the HumanoidBench tasks.

**Table 2: Hyperparameters of our method and DAC. Shared columns indicate the same hyperparameters.**
| Hyperparameter | Ours | DAC |
| --- | --- | --- |
| Neural Net | MLP (1 hidden layer w/ 256 units) |  |
| Learning rate | $5\times 10^{-3}$ |  |
| Samples used per update | 100,000 |  |
| Batch size | 1,024 |  |
| Discretization steps | 5,000 | N/A |
| Time embedding size | 16 | N/A |

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: Q1. Is the theoretical result specific to state-matching only? or can it also be applied to (state-action)-pair score matching?

Q2. LobsDICE [G] highlights a limitation of the state-matching objective, specifically that $D_f (d_E(s) | d_\pi (s))$ fails when it ignores the directionality of expert trajectories, and they suggest using $D_f (d_E(s,s’) | d_\pi (s,s’))$ instead (see the appendix of [G]). 
Can the score-matching objective $D_{DS}$ circumvent these issues?

**Reference** 

[G] Kim et al., “LobsDICE: Offline Learning from Observation via Stationary Distribution Correction Estimation”, NeurIPS 2022.
soundness: 3
strengths: This work presents a straightforward yet powerful objective: to match the score functions of expert and imitator state distributions. Despite its simplicity, this objective is strongly supported by theoretical results grounded in plausible assumptions. It is expected that the proposed method, SMILING, will show enhanced data efficiency relative to traditional imitation learning (IL) methods where the expert's (cumulative) cost is nearly deterministic as supported by the analysis in line 368.
summary: This work focuses on imitation learning from observations through the introduction of SMILING (Score-Matching Imitation LearnING). SMILING proposes to match the state distributions of experts and imitators by leveraging a diffusion model. 
This study presents theoretical results related to SMILING, specifically regarding first and second-order approximation bounds of suboptimality. Additionally, the authors theoretically demonstrate that score-matching methods can achieve superior data efficiency compared to traditional f-divergence or Integral Probability Metric (IPM) matching methods, particularly under deterministic transition scenarios.
weaknesses: I have main concerns in this study as follows:

**1. Insufficient comparative analysis with other diffusion-based IL methods:**

While this work introduces new theoretical contributions to the score-matching IL objective, the method itself lacks novelty. 
Previous studies [A, B, C, D] have already explored integrating diffusion models with imitation learning. 
However, the manuscript does not provide a detailed comparison between SMILING and these existing approaches, both conceptually and empirically. 
For instance, what are the key distinctions or what is the novelty compared to these methods? Under what conditions does the proposed method offer advantages over similar methods? Such discussions should be provided to clarify the novelty of this work.

This work appears to attempt to distinguish itself from other approaches by addressing the more specific problem of LfO, rather than standard IL. However, the manuscript does not show specific considerations for LfO; the stated objective could readily apply to state-action matching, as demonstrated in experiments (Figure 4).

**2. Weak Experimental Evaluations:**

The experimental evaluation presented also raises concerns. In the manuscript, SMILING is compared only against two obsolete baselines: BC and DAC. It is recommended to include comparisons with recent state-of-the-art IL methods such as ValueDICE [E], IQlearn [F], etc.
Moreover, the absence of comparisons with similar diffusion-based IL methods (e.g. DiffAIL [A], Diffusion-KDE [B], DBC [C], DRAIL [D] ...) further weakens the empirical validation. 
Such comparisons are insufficient to establish SMILING as a competitive method. 

In addition, the lack of experimental validation on theoretical results makes the paper less compelling. 
Given the theoretical emphasis on data efficiency in Section 5, experiments quantifying the relationship between the number of expert trajectories and the converged normalized score are essential for substantiating these claims. 
Additionally, providing experimental outcomes in both stochastic and deterministic cost scenarios would enhance the reliability of the theoretical results as well.

**References**

[A] Wang et al., “DiffAIL: Diffusion Adversarial Imitation Learning”, AAAI 2024.

[B] Pearce et al., “Imitating Human Behaviour with Diffusion Models”, ICLR 2024.

[C] Chen et al., “Diffusion Model-Augmented Behavioral Cloning”, ICML 2024.

[D] Lat et al., "Diffusion-Reward Adversarial Imitation Learning", NeurIPS 2024.

[E] Kostrikov et al., "Imitation Learning via Off-Policy Distribution Matching", ICLR 2020.

[F] Garg et al., "IQlearn: Inverse soft-Q Learning for Imitation", NeurIPS 2021.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
questions: As noted in the weaknesses section, I have the following questions:

- Are there any results that compare the proposed method to recent adversarial imitation learning baselines, such as DiffAIL?
- Are there any results that compare the proposed method to offline imitation methods that utilize diffusion models?
- Could the authors provide the performance and learning details of both the expert policy and the random policy?
soundness: 3
strengths: **Motivation and intuition**
- The motivation behind the proposed Diffusion Score Divergence is convincing, as diffusion models have demonstrated a strong capability for distribution modeling.

**Theoretical contribution**
- This paper presents solid theoretical contributions, including well-defined assumptions, lemmas, and derivations. Additionally, it discusses related topics such as sample efficiency, second-order bounds and comparisons between score functions and discriminators.

**Clarity**
- The writing is clear overall, providing both theoretical and intuitive explanations. The notations, formulations, and theorems are well-explained and easy to follow.
summary: This paper addresses the problem of online imitation learning. To this end, the paper proposes an algorithm that matches the learner's distributions to the expert's using a cost function computed by a diffusion model. Theoretical derivations demonstrate the advantages of minimizing the proposed Diffusion Score Divergence. However, the experiment results are currently insufficient to fully support the effectiveness of the proposed approach.
weaknesses: **Experiment setup**

The experiments are not sufficient. While the claims look promising, the proposed method only compares two baselines, behavior cloning (BC) and discriminator actor-critic (DAC), which is not sufficiently convincing.

First, recent adversarial imitation learning baselines, such as DiffAIL [1], should be considered. 
Although these methods do not explicitly learn reward functions but instead train discriminators, I am not convinced that this work is orthogonal to these previous works. They are feasible in the experimental setup in this paper, which includes limited expert demonstrations and online learner demonstrations.

Second, prior works [2,3] demonstrate that diffusion models are effective policies for offline imitation learning. Although offline methods can suffer from compounding error issues, it is still necessary to compare to one of the above methods to show that the proposed score-matching approach is beneficial when online learner demonstrations are accessible. 

**Experiment details**
- The details of the expert policy are missing. Since all experiments are normalized by the performance of the expert and random policies, the authors should report the specific performance values of these policies as reference points; otherwise, the reported results remain unclear.
- The number of environment steps used to train the expert policy is also missing, which makes it difficult to draw conclusions about sample efficiency from the results in Figures 3, 4, and 5. Specifically, online imitation learning aims to leverage expert demonstrations to improve sample efficiency. If the number of training steps for the expert policy is on a similar scale (e.g., 10^6), directly applying reinforcement learning to learn the expert policy could be a simpler solution.

[1] Wang, Bingzheng, et al. "DiffAIL: Diffusion Adversarial Imitation Learning." Proceedings of the AAAI Conference on Artificial Intelligence (AAAI). 2024

[2] Chi, Cheng, et al. "Diffusion policy: Visuomotor policy learning via action diffusion." Robotics: Science and Systems (RSS). 2023

[3] Pearce, Tim, et al. "Imitating human behaviour with diffusion models."  International Conference on Learning Representations (ICLR). 2023

---

confidence: 2
contribution: 3
initial_rating: 5
presentation: 2
questions: See weaknesses
soundness: 3
strengths: - The paper's contributions seem fair. They apply previously established score matching using diffusion models to IL. There have been many methods that used diffusion models in imitation learning (Section 2) and thus I rate the novelty as low. However, the authors provide new theoretical bounds and comparisons to discriminator-based methods. 

- The paper is fairly well-written and straightforward. I did not check for the correctness of the proofs and thus cannot comment.
summary: This paper proposes a new imitation learning method using score-based diffusion models. Specifically, the authors perform score-matching to minimize the difference between the expert’s and learners’  state distributions. In this way, they can perform IL using states only and also state-action pairs. They provide theoretical bounds on sample complexity when score functions are used, and show that they provide tighter bounds than when IPMs or discriminators are used. The empirical results show that their method SMILING performs favorably compared to other baselines.
weaknesses: - The framing of the method in terms of inverse reinforcement learning confused me. As far as I know, inverse reinforcement learning infers a reward function from the state-action-next state tuples and uses that to learn a policy. However, this method seems to be pure imitation learning where actions are learned directly from states or state-action pairs. Section 2 seems to imply that SMILING does IRL but I  see this method as an imitation learning method (without a discriminator). Can the authors clarify if a reward function is inferred or clarify on how SMILING is indeed an IRL method?
- Some parts of the paper are disconnected with the rest of the paper and might possibly be left over from previous drafts  (e.g. connection to IRL, SAC/DreamerV3, etc.).
   - Lines 267-269: Were SAC and DreamerV3 used in the experiments? The experiments compared against DAC and BC. 
- It's unclear how Section 6.3 helps to show that score functions are more expressive than discriminators. The results just seem to exploit a known weakness with adversarial training (mode collapse). Does the expressivity here mean that score matching does not suffer from mode collapse?

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: The paragraph at line 081 could be clarified for readability and comprehension.

Could the authors elaborate on the sentence: "We first fit a score function to the expert demonstrations. Then, at each iteration of our algorithm, we fit a score function to the state distribution of the mixture of the preceding policies, before using the combination of these score functions to define a reward function for the policy search step." Specifically, what does "mixture of the preceding policies" mean?

Your explanation of the algorithm here is helpful, but this part could be more precise. For example, does "mixture of preceding policies" refer to averaging the state distributions from previous iterations, or does it imply some other form of combining past policies? It would also improve clarity if you could provide a concrete example or a diagram illustrating how this mixture is constructed and applied in the algorithm. These additions would address potential confusion and make the process more accessible to readers.

I think there can be additional clarifications in the method: Algorithm 1, would be good to define $c^{(k)}$ explicitly either in the algorithm pseudocode or in the accompanying text. Additionally, adding how $c^{(k)}$ relates to the score functions mentioned earlier and if it's connected to the cost function in Eq. 3 would provide a more comprehensive understanding of the algorithm.
soundness: 4
strengths: Impressive experimental results, clearly demonstrating the benefit from not having a GAN approach for IRL. 
Clear motivation and formulation and writing is of good quality.
summary: Traditionally in Inverse Reinforcement Learning (IRL), when expert demonstrations are available, a generative and adversarial setup is often used. In this setup, the discriminator’s goal is to learn the distribution of the desired behavior from expert demonstrations, while the generator’s goal is to produce behaviors that fit the distribution implicitly defined by the discriminator. Methods such as GAIL have proven powerful but come with a drawback: the optimization process. In these settings with a generator and discriminator, the optimal solution is a saddle point; however, it is often difficult to achieve due to issues like mode collapse and instability in training dynamics. In this work, the authors propose using diffusion models as an alternative. They achieve this by learning the score of the state distribution from expert demonstrations and then applying the DS divergence, which measures the difference in scores throughout the diffusion process. This metric is then used as the reward function to train the policy. The authors demonstrate impressive results on challenging humanoid tasks, particularly the Humanoid Walk task.
weaknesses: I think there are prior work with a similar motivation.  Can you elaborate on how your is novel compared to the following:

https://openreview.net/pdf?id=NN60HKTur2
https://arxiv.org/pdf/2004.09395
https://arxiv.org/abs/2407.00626

---

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: The same with the weaknesses I mentioned above.
soundness: 3
strengths: I appreciate the numerous contributions and the excellent presentation of this paper:

- The idea is an innovative application of score matching techniques to model the state distributions of the policy. Unlike previous approaches that rely on training a discriminator between expert and learner data, the theoretical results demonstrate that the score-based method is more effective and sample-efficient, avoiding issues such as unstable training and mode collapse.
- The presentation is clear and highly informative. I particularly enjoyed the writing in Section 5.1, which intuitively explains why the score-based approach is more expressive than discriminators, especially in the context of linear models. The analysis in Section 5.1 is further validated by the experiments in Section 6.3, which is a great reinforcement of the theoretical insights.
- Across various benchmarks, the method provides a significant improvement over baseline techniques in the experiments, showing strong empirical results.

Overall, the paper's theoretical and experimental results convincingly demonstrate the novelty and effectiveness of the proposed method. This approach could potentially inspire other fields that rely on training discriminators to explore score-based techniques.
summary: This paper introduces an Inverse Reinforcement Learning (IRL) algorithm called SMILING, which replaces the commonly used f-divergence with a novel approach called Diffusion Score Divergence. This new method measures the discrepancy between the expert's and the learner's state distributions along the same forward diffusion process. The proposed algorithm achieves significantly tighter first- and second-order instance-dependent bounds, supported by thorough error analysis throughout the entire algorithm. Experimental results demonstrate the effectiveness of SMILING across various standard control benchmarks.
weaknesses: Although this paper presents promising results, there are three primary weaknesses that should be addressed to some extent. I will consider any responses to these concerns when recalibrating my final score.

- **Time Complexity of Evaluating the Cost Function $c^k(s)$ in Eq.3**:
  - The main concern is the time complexity associated with evaluating the cost function. In the experimental sections, the authors employ a DDPM framework with over 5,000 steps. This implies that evaluating the cost function for a given state requires 5,000 evaluations of the model, which could be prohibitive for more complex scenarios, particularly with higher-dimensional state spaces.
  - If the authors are using a single randomly sampled time step to estimate Eq.3, the variance of this estimation across time should be incorporated into the final theoretical results. However, this issue is not addressed in the current analysis, raising questions about the method's scalability to high-dimensional state spaces.
- **Choice of Data Aggregation**:
  - In Line 3 of the algorithm, the score matching target aims to force $g(s_t, t)$ to model the marginal score of the state distributions from a mixture of policies $\pi^{(1:k-1)}$. However, the paper lacks a clear justification for using this mixture-based approach instead of the more intuitive choice of selecting the prior learned policy $\pi^{k-1}$. While the theoretical results focus on the mixture of policies, it would be valuable to demonstrate the necessity of this choice.
  - Providing more empirical or theoretical results to verify the effectiveness of this data aggregation technique during training would help clarify its benefits. For instance, comparing the proposed approach with a baseline that only uses $\pi^{k-1}$ could offer insights into the significance of data aggregation.
- **Assumption (c) and the Rate of $\epsilon_{RL}$ in Continuous State-Space Settings**:
  - The rate of $\epsilon_{RL}$ should be discussed in more detail, especially considering that score matching is typically applied in continuous spaces. Although some recent algorithms extend score matching to discrete spaces [1, 2], incorporating them into the current framework may not be straightforward.
  - It would be beneficial to discuss how the proposed framework might be adapted or fit into established theoretical results in discrete spaces. This could include exploring potential modifications or providing insights into the limitations of the current approach when dealing with discrete settings.

Addressing these concerns would significantly strengthen the paper by enhancing its scalability, clarifying the choice of data aggregation, and broadening the discussion on the applicability of the proposed framework in various state-space settings.

[1] Meng, Chenlin, et al. "Concrete score matching: Generalized score matching for discrete data." *Advances in Neural Information Processing Systems* 35 (2022): 34532-34545.

[2] Lou, Aaron, Chenlin Meng, and Stefano  Ermon. "Discrete diffusion language modeling by estimating the ratios of the data distribution." *arXiv preprint arXiv:2310.16834* (2023).

Meta-review: All the reviewers unanimously recommend accepting this paper, given its clarity and theoretical and empirical results. Also, the author's rebuttal has sufficiently addressed some concerns and questions raised by the reviewers, such as comparisons to recent works, the importance of learning from observations, and the unclear experimental setup and notations. Consequently, I recommend accepting the paper.

Decision: Accept (Poster)