1 Introduction↩︎

Offline reinforcement learning (RL) aims to derive an optimal decision-making policy from a previously collected dataset without further environment interaction [1]. Circumventing risky, costly, and inefficient online interactions, offline RL enables models to fully exploit prior data, and thus gained significant traction in safety-critical applications such as autonomous driving and robotic manipulation [2]. However, learning exclusively from previously collected data can be more than challenging. Conventional reinforcement learning algorithms typically suffer from distributional shift as they evaluate actions outside the support of the behavior policy, where value estimations are unreliable [3]. Moreover, as the datasets are growing larger and more diverse, the behavioral distribution is also more complex and multi-modal, which necessitates a more expressive policy class to represent the complex policies [4].

To model these complex and potentially multimodal policy distributions, Gaussian mixture model (GMM) [5], [6] and variational auto-encoders (VAE) [7] are frequently employed for policy representation, as their tractable sampling and efficient optimization properties allow them to capture the underlying probability distributions of expert behaviors. More advanced frameworks, such as Langevin dynamics-based approaches [8] and probability flow methods [9] offer superior expressivity, and thus they have been extensively integrated into both imitation learning and offline RL scenarios. However, the slow sampling speed for these generative models remains unacceptable for various computationally intensive online tasks. Consequently, there is an imperative need for more computationally efficient sampling methods.

To address this problem, we introduced Moment Matching Q-Learning (MoMa QL), which involves a stable and theoretically rigorous policy learning procedure based on actor-critic style algorithm. MoMa QL enables actor’s policy to operate on the time-dependent marginal distributions of stochastic interpolates [10], which connects two arbitrary probability density functions. By learning a function mapping from any marginal distribution at time \(t\) to any marginal at time \(s < t\), this framework facilitates a seamless transition across the probability flow trajectory, MoMa QL naturally supports both single- or multi-step action sampling. With this accelerated action sampling, MoMa QL enables a simultaneous critic update through it’s policy training procedure. In summary, our contribution can be summarized as follows:

• We introduce MoMa QL, which motivates from minimizing the Maximum Mean Discrepancy (MMD) of conditional distributions derived from two intermediate diffusion steps \(r < s\) to enable the diffusion process to “jump” from \(r\) to \(s\) via one-step induction.

• Theoretically, we prove that the effectiveness and convergence of the MoMa QL with the fact that the consistency models can be viewed as a special case of our results.

• Empirically, we demonstrated that MoMa QL exhibit superior performance comparing with a series of offline RL methods. Specially, thanks to the accelerated action sampling process, we show that MoMa QL can be efficiently finetuned with online rollouts and outperforms existing offline-to-online RL methods.

2 Preliminaries↩︎

2.1 Offline Reinforcement Learning↩︎

We consider the offline reinforcement learning setting in this paper. The environment in RL can be formulated as a Markov decision process \(\mathcal{M} := \left(\mathcal{S}, \mathcal{A}, \mathcal{P}, R, \rho, \gamma\right )\), with state space \(\mathcal{S}\), action space \(\mathcal{A}\), environment dynamics \(\mathcal{P}(\mathbf{s}' | \mathbf{s}, \mathbf{a}) : \mathcal{S} \times \mathcal{S} \times \mathcal{A} \rightarrow I\), reward function \(R(\mathbf{s}, \mathbf{a}) : \mathcal{S} \times \mathcal{A} \rightarrow \mathbb{R}\), initial distribution \(\rho \in \Delta(S)\) and discounted factor \(\gamma \in [0, 1)\) [11], where we denote the set \([0, 1]\) as \(I\) and the set of probability distributions over a space \(\mathcal{X}\) as \(\Delta(\mathcal{X})\).

The goal of Offline RL is to learn a policy \(\pi_\theta(\mathbf{a} | \mathbf{s}) : \mathcal{S} \rightarrow \Delta(\mathcal{A})\), a conditional distribution parametrized by \(\theta\), which maximizes the cumulative discounted reward \(\mathbb{E}\big[\sum_{t = 0}^H \gamma^t r(\mathbf{s}_t, \mathbf{a}_t) \big]\) over horizon \(H\). This is equivalent to a constrained optimization problem with respect to \(\theta\).

In the offline setting, the dataset \(\mathcal{D} := \left \{ \tau^{n} \right \}_{n=1}^N\) is static and doesn’t involve environment interaction, where each \(\tau^{i} := \left ( \mathbf{s}_0, \mathbf{a}_0, \dots, \mathbf{s}_H, \mathbf{a}_H\right )\) denotes one single trajectory.

In addition to the purely offline setting, we also consider the offline-to-online RL regime, in which we aim to fine-tune a pre-trained policy via online interactions with the environment after pretraining using offline dataset.

2.2 Continuous-time Probabilistic Generative Models↩︎

For a given data distribution \(q(\mathbf{x}) \in \Delta(\mathbb{R}^d)\), diffusion models [12], [13] and flow matching [14], [15] construct time-dependent variable \(x_t\) which can be viewed as an interpolation between sampled data \(\mathbf{x} \sim q(\mathbf{x})\) and noise \(\boldsymbol{\varepsilon} \sim \mathcal{N}(0, I)\) in the form as \(\mathbf{x}_t = \alpha_t\mathbf{x} + \sigma_t \boldsymbol{\varepsilon}\) with \(\alpha_0 = \sigma_1 = 1, \alpha_1 = \sigma_0 = 0\). Essentially, the problem can be formulated as solving a Schrödinger bridge [16], [17] between \(\boldsymbol{\varepsilon}\) and \(q(x)\). Within this perspective, Variance Preserved diffusion conventionally chooses \(\alpha_t = \cos\!\left(\frac{\pi}{2} t\right), \sigma_t = \sin\!\left(\frac{\pi}{2} t\right)\) and flow matching chooses \(\alpha_t = 1 - t, \sigma_t = t\). Both diffusion models and flow matching aim to learn a time-dependent vector field \(\mathbf{v}_t(\mathbf{x}) := \frac{\mathrm{d}\mathbf{x}_t}{\mathrm{d}t}\) parameterized by a neural network. For score based diffusion model with deterministic sampling, the vector field can be represented as: \[\mathbf{v}_t(\mathbf{x}) = \alpha_t - \sigma_t^2 \mathbf{s}_\theta(\mathbf{x}, t) / 2, \label{eq2}\tag{1}\] where \(\mathbf{s}_\theta(\mathbf{x}_t, t) := \nabla_{\mathbf{x}} \log p_t(\mathbf{x})\) denotes the score function. The model is trained with the objective \[\mathcal{L}_{\mathrm{DSM}}(\theta) = \mathbb{E}_{t, \mathbf{x}_0, \boldsymbol{\varepsilon}} \left[ \left\| \mathbf{s}_\theta(\mathbf{x}_t, t) + {\boldsymbol{\varepsilon}}/{\sigma_t} \right\|^2 \right], \label{eq3}\tag{2}\] where \(\boldsymbol{\varepsilon} \sim \mathcal{N}(0, I)\). In contrast, flow matching directly parameterizes the velocity field \(\mathbf{v}_t(\mathbf{x})\) and trains the model by regressing toward a target conditional velocity. Given a stochastic interpolation path \(\mathbf{x}_t\) between samples from the base and target distributions, the optimal velocity field is defined as \(\mathbf{v}_t^\star(\mathbf{x}_t) = \mathbb{E}\!\left[ \frac{\mathrm{d}\mathbf{x}_t}{\mathrm{d}t} | \mathbf{x}_t \right]\), and the corresponding flow matching loss function is \[\mathcal{L}_{\mathrm{FM}}(\theta) = \mathbb{E}_{t, \mathbf{x}_0, \mathbf{x}_1} \left[ \left\| G_\theta(\mathbf{x}_t, t) - \mathbf{v}_t^\star(\mathbf{x}_t) \right\|^2 \right], \label{eq5}\tag{3}\] where \(G_\theta(\mathbf{x}_t, t)\) is a parameterized neural network.

Stochastic Interpolants [10] provides a unified and flexible perspective that encapsulates both SDE-based models like diffusion and ODE-based models like flow matching. It constructs a conditional interpolation \(q_t(\mathbf{x}_t | \mathbf{x}, \boldsymbol{\varepsilon}) = \mathcal{N}(I_t(\mathbf{x}, \boldsymbol{\varepsilon}), \gamma_t^2 I)\) with constraints \(I_0(\mathbf{x}, \boldsymbol{\varepsilon}) = \mathbf{x}\), \(I_1(\mathbf{x}, \boldsymbol{\varepsilon}) = \boldsymbol{\varepsilon}\) and \(\gamma_0 = \gamma_1 = 0\). The conditional interpolant velocity can thus be represented as \(\mathbf{v}_t = \partial_t I_t(\mathbf{x}, \boldsymbol{\epsilon}) + \dot{\gamma}_t \mathbf{z}\) where \(\mathbf{z} \sim \mathcal{N}(0, I)\). Similar to flow matching, we can learn a deterministic sampler via \(G_\theta(\mathbf{x}_t, t) \approx \mathbb{E}_{\mathbf{x}, \boldsymbol{\varepsilon}, \mathbf{z}}\left [\mathbf{v}_t | \mathbf{x}_t\right ]\). Given linear interpolation condition \(I_t(\mathbf{x}, \boldsymbol{\epsilon}) = \alpha_t \mathbf{x} + \sigma_t \boldsymbol{\epsilon}\), it’s easy to show that stochastic interpolants reduces to flow matching when \(\gamma_t \equiv 0\) and reduces to diffusion when \(\boldsymbol{\varepsilon} \sim \mathcal{N}(0, I)\).

2.3 Policy Regularization↩︎

Policy regularization adds constraints or penalties to an agent’s learning process to prevent it from overfitting and promote stable and robust behavior. Actor-Critic (AC) plays an important role in policy regularization, which casts an explicit constraint on the policy and thus enhances both the stability and the asymptotic performance of the agent. Following conventional behavior-regularized actor-critic framework (BRAC) [18], [19], we minimize the empirical risk for both actor and critic by: \[\begin{align} \mathcal{L}_Q(\phi) &= \mathbb{E}_{\substack{\mathbf{s},\mathbf{a},r,\mathbf{s}' {\sim}\mathcal{D}, \\ \mathbf{a}' \sim \pi_\theta}} [ ( Q_\phi(\mathbf{s},\mathbf{a}) {-} r {-}\gamma Q_{{\phi}}^\perp(\mathbf{s}',\mathbf{a}') )^2 ],\tag{4} \\ \mathcal{L}_\pi(\theta) &= \mathbb{E}_{\mathbf{s},\mathbf{a} \sim \mathcal{D}, \mathbf{a}^\pi \sim \pi_\theta} [ \underbrace{- \eta Q_\phi(\mathbf{s},\mathbf{a}^\pi)}_{\text{Q loss}} - \underbrace{ \log \pi(\mathbf{a}|\mathbf{s})}_{\text{BC loss}} ], \tag{5} \end{align}\] where \(Q_\phi(\mathbf{s}, \mathbf{a}) : \mathcal{S} \times \mathcal{A} \rightarrow \mathbb{R}\) denotes the state-action value function parameterized by \(\phi\), \(Q_{{\phi}}^\perp(\mathbf{s}, \mathbf{a})\) is a target network which stops gradient and delays the update [20] and \(\alpha\) controls the strength of the BC regularizer. The critic loss \(\mathcal{L}_Q(\phi)\) minimizes the standard Bellman error, while the actor loss \(\mathcal{L}_\pi(\theta)\) maximizes values with reparameterized gradients through \(\mathbf{a}^\pi\). For the actor, the BC loss is additionally applied to prevent the policy from deviating too much from the behavioral policy’s distribution. Despite its simplicity, BRAC is one of the most performant frameworks on standard D4RL tasks, and thus we try to build our algorithm on a variant of the BRAC framework.

2.4 Maximum Mean Discrepancy↩︎

Consider a reproducing kernel Hilbert space (RKHS) \(\mathcal{H}\) of functions \(\mathbb{R}^d \rightarrow \mathbb{R}\) associated to a bounded positive definite kernel \(k\) on \(R^d\), the Maximum Mean Discrepancy (MMD) between two probability measures \(\mu\) and \(\nu\) on \(\mathbb{R}^d\) is: \[\mathsf{MMD}(\mu, \nu) = \sup_{f \in \mathcal{H}, \lVert f \rVert_{\mathcal{H}} \le 1} \left| \scalebox{1.4}{\int} f \, \mathrm{d}\mu - \scalebox{1.4}{\int} f \, \mathrm{d}\nu \right|. \label{eq8}\tag{6}\] \(\mathsf{MMD}\) is a metric if and only if the kernel mean embedding \(\mu \rightarrow \int k(x, \cdot) \mu(\mathrm{d}x)\) is injective. Kernels which satisfies such a property are called characteristic, and one of the most important example is the RBF kernel \(k(x, y) = e^{-\frac{||x - y||^2}{2\sigma^2}}\) [21]. RBF kernel also implies an inner product of infinite-dimensional feature maps consisting of all orders of moments of \(\mu\) and \(\nu\), which can be respectively denoted and \(\mathbb{E}_\mu(\mathbf{x}^j)\) and \(\mathbb{E}_\nu(\mathbf{x}^j)\) for any \(j \geq 1\) [22].

3 Related Works↩︎

3.1 Offline Reinforcement Learning↩︎

Offline RL is the problem of policy optimization with previously collected data only, which is well-known for suffering from the value overestimation problem for out-of-distribution states and actions. Previous methods for solving this issue include explicit behavioral regularization [18], [19], [23], pessimistic value estimation [24]–[26], out-of-distribution detection [27]–[29] and dual RL [30], [31].

3.2 Offline-to-Online RL↩︎

Besides the advancements only using offline data, recent works starts to allow the offline-pretrained RL agents to be finetuned through online interactions, referred to as offline-to-online RL. In this regime, we fine-tune the policy with additional online rollouts, and this process usually suffers from a catastrophic degraded performance at initial online training stage due to the distribution shift of training samples. Prior research has studied online fine-tuning with offline data or pre-trained policies, including Off2OnRL [32], Hybrid Q learning [33], Cal-QL [34], RLPD [35] and ACA [36]. Our method, MoMa QL, mainly focuses on offline RL setting, but we will also show it can also be directly fine-tuned with online rollouts.

3.3 Score-based and Flow-based Decision Making↩︎

Previous works have applied various iterative score-based models to a series of RL tasks. Diffusion QL [3] initiates the usage of DDPM [12] as policy representation in a Q-Learning+BC framework. Implicit Diffusion Q-learning (IDQL) [37] improves IQL via integrating diffusion policy. Diffusion policies [8] applies DDIM model for both state-based and vision-based imitation learning tasks in Robotics domain. Consistency-AC [38] utilizes CM to accelerate the action sampling. Methods such as QGPO [39], EDP [40] and QIPO [41] model the offline RL objective as an energy-guided diffusion process, which can be viewed as a variant of AWR [42], [43]. Flow matching is also widely used in offline RL. GFlower [9] apply flow matching models to policy optimization. FlowPolicy [44] employed Consistency Flow Matching for robot manipulation tasks. FQL [45] introduced a distilled neural network for a flow-based policy, which accelarates the inference. In our experiments, we compare our method against a representative subset of these models to demonstrate its superior empirical performance and inference efficiency.

4 Moment Matching Q-Learning↩︎

We now introduce our method, Moment Maching Q-Learning (MoMa QL), for effective decision-making tasks. The cornerstone of the whole algorithm is to learn an implicit action sampler which can harness the expressive power of score and flow based models and expedite the training and inference process.

4.1 General Objective↩︎

For the generative policies, the model is targeted to learn an implicit conditional sampler for \(p^\theta(\mathbf{x} | \boldsymbol{\varepsilon}, \mathbf{s})\), where \(\mathbf{x} \sim q(\mathbf{x})\), \(\boldsymbol{\varepsilon} \sim p(\boldsymbol{\varepsilon})\) and \(\mathbf{s}\) denotes the state of the agent which serves as the guidance for the reversed process. In the BRAC framework, the loss function in Eq. 5 can be rewritten as: \[\mathcal{L}_\pi(\theta)=\mathbb{E}_{\mathbf{s},\mathbf{a} \sim \mathcal{D}, \mathbf{a}^\pi \sim \pi_\theta} [ \underbrace{-\eta Q_\phi(\mathbf{s},\mathbf{a}^\pi)}_{\text{Q loss}} ] + \underbrace{\mathcal{L}_{D}(\theta)}_{\text{BC loss}}, \label{eq9}\tag{7}\] where \(D\) is usually designed as a sample-based metric to calculate the divergence between the marginal \(p\) and \(q\) measure. We utilize MMD loss due to its great optimization stability and superior statistical property, and thus match higher order moment of the distribution. Details related to MMD will be further introduced in Appendix 7.1. For the simplicity, we will ignore the guidance \(\mathbf{s}\) in further derivation.

4.2 Marginal Interpolants↩︎

Assuming a time-augmented interpolation \(\mathbf{x}_t \sim q_t(\mathbf{x}_t | \mathbf{x}, \boldsymbol{\varepsilon})\) between original data distribution \(q(\mathbf{x})\) and prior distribution \(p(\boldsymbol{\varepsilon})\), we can formulate the marginal interpolating distribution as: \[q_t(\mathbf{x}_t) = \scalebox{1.4}{\iint} q_t(\mathbf{x}_t | \mathbf{x}, \boldsymbol{\varepsilon}) q(\mathbf{x}) p(\boldsymbol{\varepsilon}) \mathrm{d}\mathbf{x} \mathrm{d}\boldsymbol{\varepsilon}. \label{eq10}\tag{8}\] Suppose the conditional distribution of \(\mathbf{x}_s | \mathbf{x}_t, \mathbf{x}\) follows as: \[q_{s|t}(\mathbf{x}_s | \mathbf{x}, \mathbf{x}_t) = \mathcal{N} \left( I_{s|t}(\mathbf{x}, \mathbf{x}_t), \gamma_{s|t}^2 I \right) \label{eq11}\tag{9}\] with constraints \(I_{t|t}(\mathbf{x}, \mathbf{x}_t) = \mathbf{x}_t, I_{0|t}(\mathbf{x}, \mathbf{x}_t) = \mathbf{x}, \gamma_{t|t} = \gamma_{0|t} = 0\). This is a natural extrapolation of conditional interpolation \(q_t(\mathbf{x}_t | \mathbf{x}, \boldsymbol{\varepsilon})\). Conceptually, it extends the traditional framework by dynamically shifting the prior from endpoint 1 to an arbitrary time \(t\). We also adopt the notion of marginal-preserving [46] as follows:

This constraint can provide strong theoretical guarantee for the existence and convergence of the minimizer, as further elucidated in Appendix 7.2. Similarly, for \(\forall t \in [0, 1]\) and \(\forall s \in [0, t]\), we can thus define the marginal model on \(p\) measure as: \[p_{s|t}^{\theta}(\mathbf{x}_s) = \scalebox{1.4}{\iint} q_{s|t}(\mathbf{x}_s|\mathbf{x}, \mathbf{x}_t) p_{s|t}^{\theta}(\mathbf{x}|\mathbf{x}_t) q_t(\mathbf{x}_t) \mathrm{d}\mathbf{x}_t \mathrm{d}\mathbf{x} \label{eq14}\tag{12}\] where the interpolate is marginal preserving and \(p_{s|t}^{\theta}(\mathbf{x}|\mathbf{x}_t)\) serves as the implicit one-step sampler we intend to attain. So, our main target is to build a model which can help us to minimize the gap between \(p_{s|t}^{\theta}(\mathbf{x}_s)\) and \(q_s(\mathbf{x}_s)\).

To produce a clean sample \(\mathbf{x}\) given \(\mathbf{x}_t \sim q(\mathbf{x}_t)\) via an intermediate \(s\), we first sample a denoised \(\tilde{\mathbf{x}} \sim p_{s|t}^{\theta}(\mathbf{x}|\mathbf{x}_t)\) and then sample \(\tilde{\mathbf{x}}_s \sim q_{s|t}(\mathbf{x}_s|\mathbf{x}, \mathbf{x}_t)\) and then (3) sample \(\mathbf{x} \sim p^\theta_{0|s} (\mathbf{x} | \tilde{\mathbf{x}}_s)\). Similar to DDIM song2020denoising?, this process can be inductive, and thus we can also obtain a multi-step sampler for \(p^\theta\) from \(t = 1\) to \(t = 0\).

4.3 Actor Loss Formulation↩︎

Directly matching the MMD loss between \(q_s(\mathbf{x}_s)\) and \(p_{s|t}^{\theta}(\mathbf{x}|\mathbf{x}_t)\) is not a good idea: there can be great discrepancy between \(q_t(\mathbf{x}_t)\) and \(q_s(\mathbf{x}_s)\) when \(t\) is far from \(s\), so \(p^\theta_{s|t} (\mathbf{x}_s)\) can be difficult to estimate. We adopt a training paradigm inspired by consistency training [47] to address this challenge. Specifically, we introduce an intermediate timestep \(r\) such that \(s < r < t\), and enforce a self-consistency constraint by aligning the denoising distributions \(p^\theta_{s|r}(\mathbf{x}_s)\) and \(p^\theta_{s|t}(\mathbf{x}_s)\): \[\mathcal{L}_D(\theta) = \mathbb{E}_{s,r,t} \left[\mathsf{MMD}^2 \left( p_{s|r}^{\theta^-}(\mathbf{x}_s), p_{s|t}^{\theta}(\mathbf{x}_s) \right) \right] \label{eq15}\tag{13}\] We can utilize a neural network \(G_\theta(\mathbf{x}, s, t)\) to attain the denoised sample \(\tilde{\mathbf{x}}\) and thus define \(p^\theta_{s|t} (\mathbf{x} | \mathbf{x}_t) = \delta(\mathbf{x} - G_\theta(\mathbf{x}_t, s, t))\). With the denoised sample \(\tilde{\mathbf{x}}\), we choose to utilize DDIM interpolants, denoted as \(f^\theta_{s, t}(\mathbf{x}_t)\) and \(f^\theta_{s, r}(\mathbf{x}_r)\), to obtain a sample from \(p^\theta_{s|t}(\mathbf{x}_s)\) and \(p^\theta_{s|r}(\mathbf{x}_s)\). Empirically, we choose an unbiased estimator of squared MMD Loss [48] to compute \(\mathcal{L}_D\). Further details and theoretical derivation related to actor loss function can be found in Appendix 7.3.

4.4 Algorithm Implementation↩︎

We present a brief MoMa QL algorithm in Algorithm 1 and defer the full version in Appendix 8. For parameterized \(Q_\phi(\mathbf{s},\mathbf{a})\), we utilized the double Q-Learning loss [19] and with batched data \(\mathcal{B} \subseteq \mathcal{D}\), which mitigates Q-value overestimation in Eq. 4 : \[\begin{align} \mathcal{L}(\phi) = &\mathbb{E}_{(\mathbf{s}, \mathbf{a}, \mathbf{s}') \sim \mathcal{B}, \mathbf{a}' \sim \pi_{\theta^{\mathsf{T}}}(\cdot | \mathbf{s}')} \Big[ \big( ( r(\mathbf{s}, \mathbf{a}) + \\ &\gamma \min_{i \in \{1, 2\}} Q_{\phi_i^{\mathsf{T}}}(\mathbf{s}', \mathbf{a}') ) - Q_{\phi_i}(\mathbf{s}, \mathbf{a}) \big)^2 \Big] \end{align} \label{eq16}\tag{14}\] The regularized policy \(\pi_\theta\) is learned via minimizing the objective function in Eq. 7 via backpropagation. This loss comprises two components: (i) a Q-loss, where actions \(\mathbf{a}\) are sampled inductively according to the procedure in Algorithm 2 and subsequently fed into the critic network \(Q_{\phi_i}\); and (ii) a behavior cloning (BC) loss, as formulated in Eq. 13 . We update the actor and critic networks using an exponential moving average (EMA) with a smoothing coefficient \(\alpha\).

5 Experiments↩︎

This section empirically evaluates our proposed RL algorithm, MoMa QL. We conduct comprehensive experiments on standard offline RL and offline-to-online RL benchmarks. We also conducted experiments to demonstrate the effectiveness and stability of our method.

5.1 Experiment Setup↩︎

5.1.1 Benchmarks↩︎

We primarily evaluate the expressiveness and computational efficiency of our proposed algorithm and the responding algorithms on three task suites (Gym, Adroit and Kitchen) in D4RL benchmarks [49]. D4RL benchmarks include a diverse set of offline RL tasks spanning locomotion, manipulation, and Gym Mujoco environments. Tasks in D4RL are formulated with fixed offline datasets and standardized evaluation protocols.

5.1.2 Baseline Methods↩︎

Our experiments evaluate three distinct learning paradigms: behavior cloning (BC), offline reinforcement learning (RL), and offline-to-online RL. We select a range of established and state-of-the-art algorithms as baselines, drawing from both classical approaches and recent advances in generative modeling for RL.

Behavior Cloning (BC). We compare our method with vanilla Behavior Cloning (BC) parameterized by Gaussian policy, Consistency-BC [38], and Diffusion-BC [3]. Additionally, we include classic BC baselines reported in previous works: Diffuser [50], MoRel [28], Onestep RL [23], TD3+BC [19] and Decision Transformer(DT)[51].

Offline RL. For offline RL, we compare methods across different algorithmic families. In Gym locomotion tasks, baselines include Conservative Q-Learning (CQL) [24], Implicit Q-Learning (IQL) [26], Extreme Q-learning (\(\mathcal{X}\)-QL)[52], Actor-Restricted Q-learning (ARQ)[53], IDQL-A[37], Diffusion-QL [3], and Consistency-AC (CAC) [38]. For Adroit manipulation tasks, we consider BC, IQL[26], ReBRAC [54], Implicit Diffusion Q-Learning (IDQL) [37], SRPO [55], CAC[38], Flow advantage-weighted actor-critic (FAWAC) [43], Flow behavior-regularized actor-critic (FBRAC)[3], Implicit Flow Q-Learning (IFQL)[45], and Flow Q-Learning (FQL) [45]. For Kitchen tasks, baselines include CQL, IQL, \(\mathcal{X}\)-QL, ARQ, Diffusion-QL, and Consistency-AC.

Offline-to-Online RL. In the offline-to-online setting, we initialize the policy using offline data and fine-tune it with online interactions. We compare MoMa QL against IQL, Cal-QL [34], RLPD [35], Diffusion-QL, and Consistency-AC. We also include comparisons with Soft Actor-Critic (SAC) [56], Advantage-Weighted Actor-Critic (AWAC) [43], and Actor-Critic Alignment (ACA) [36] in detailed evaluations.

For each method, We adopted the same results from the previous works[38], [45]. We do not use the same parameters for all the methods.

5.2 Main Results For Behavior Cloning↩︎

We evaluate our method, MoMa QL, against various baseline algorithms in the behavior cloning (BC) setting. Table 1 summarizes the performance comparisons on the D4RL Gym tasks. The hyperparameters for training can be found in Appendix 9.1.

The quantitative results indicate that our method, MoMa QL, achieves the highest overall average score of 89.8, surpassing all baseline methods. Compared to vanilla BC (51.9) and Diffusion-BC (76.3), our approach achieves relative improvements of \(1.73\times\) and \(1.18\times\), respectively. This suggests that the moment-matching objective effectively captures the underlying distribution of the dataset, providing a more robust policy representation than standard supervised learning or diffusion-based counterparts.

Notably, MoMa QL excels in tasks with medium and medium-replay datasets. For instance, on walker2d-mr and halfcheetah-m, our method surpasses the strongest baselines by factors of \(1.26\times\) and \(1.44\times\), respectively. This highlights the effectiveness of our method in handling suboptimal and highly stochastic data. While the performance on hopper-me is lower than some baselines, likely due to the specific characteristics of that expert dataset, the consistent superiority across the majority of tasks confirms the strong generalization capability of our proposed framework.

5.3 Results For Offline RL↩︎

We evaluate MoMa QL on standard D4RL benchmarks across three task suites: Gym locomotion tasks, Adroit manipulation tasks, and Kitchen tasks. Following key baselines, we compare MoMa QL against a variety of methods representing distinct policy classes. Table 4 presents the aggregated results across dataset variants. Detailed per-dataset results can be found in Table 9 in Appendix 9.2.

Gym Locomotion Tasks. In the standard D4RL Gym locomotion suite, MoMa QL achieves a dominant average score of 95.5, surpassing the strongest baseline, Diffusion-QL (87.9), by a factor of approximately \(1.09\times\). Specifically, on HalfCheetah and Walker2D, our method outperforms the nearest competitors by margins of \(14\%\) and \(8.5\%\), respectively. Although slightly trailing Diffusion-QL on Hopper (\(99.9\) vs. \(101.0\)), MoMa QL demonstrates superior consistency across tasks, effectively balancing distribution matching with reward maximization.

Adroit Manipulation Tasks. For the challenging high-dimensional Adroit tasks, MoMa QL achieves an average score of 56.7, which is highly competitive with the state-of-the-art method ReBRAC (55.4). Notably, on the door task, our method attains the highest score of 38.5, outperforming all baselines. While ReBRAC shows advantages in expert datasets, MoMa QL demonstrates robustness across diverse dataset qualities, particularly excelling in scenarios where capturing multi-modal distributions is critical.

Kitchen Tasks. In the long-horizon Kitchen environments, MoMa QL achieves a strong average score of 73.1, outperforming Diffusion-QL (69.0) and \(\mathcal{X}\)-QL (72.9). This highlights the advantage of our sample-efficient few-step generation for sequential control tasks requiring long-term planning.

Overall Performance. Across all benchmarks, MoMa QL demonstrates consistent state-of-the-art or competitive performance. The significant improvements in Gym tasks (\(1.09\times\) over Diffusion-QL) and robust results in Adroit and Kitchen domains validate the effectiveness of our MMD-based moment matching framework.

5.4 Offline-to-Online Fine-tuning↩︎

We further investigate the capability of MoMa QL in the offline-to-online setting. In this regime, we initialize the policy with the pre-trained offline model (trained for 100K steps) and then fine-tune it with online interactions for an additional 400K steps. We use the same hyperparameters as in the offline setting, which are detailed in Appendix 9.1.

Table 5 highlights the performance improvement of MoMa QL after online fine-tuning. Our method achieves significant gains, particularly on the challenging medium-replay datasets, where the offline initialization provides a strong starting point for further effective learning. For instance, on halfcheetah-medium-replay, performance improves from \(63.3\) to \(80.9\) (\(\approx 28\%\) increase).

Figure 3 visualized the improvement of MoMa QL after online finetuning. More detailed comparisons with baselines are provided in Appendix 9.3.

5.5 Ablation Studies↩︎

We conduct comprehensive ablation studies to demonstrate that MoMa QL is highly robust to hyperparameter variations. We examine six parameters: \(\eta\), \(N\), \(a\), \(b\), \(p_{\text{mean}}\), and \(p_{\text{std}}\). Detailed results are provided in Appendix 9.4.

First, we analyze the impact of the Q-learning weight \(\eta\) and sampling steps \(N\). As shown in Figure 4 (a), performance remains stable and high for \(\eta \in [0.5, 1.0]\). This indicates that the method is not sensitive to precise tuning of the balance parameter, effectively learning across a broad range of weights. Regarding the number of sampling steps, Figure 4 (b) demonstrates that performance saturates at \(N=2\). MoMa QL is thus computationally efficient, working well with as few as 2-4 steps without performance degradation at higher steps.

Furthermore, MoMa QL performs consistently well across variations in MMD kernel parameters (\(a, b\)) and noise schedule parameters (\(p_{\text{mean}}, p_{\text{std}}\)). Figure 5 (a, b) shows that the method is insensitive to kernel parameters \(a, b \in [1, 8]\), requiring no specific tuning. Similarly, Figure 5 (c, d) indicates that performance is largely invariant to the noise schedule parameters. This overall robustness ensures that MoMa QL is easy to deploy without complex hyperparameter search.

5.6 Computational Cost Analysis↩︎

We evaluate the computational efficiency of MoMa QL by measuring the training time per 1,000 gradient steps across different tasks. We ensure a fair comparison by using identical batch sizes across all methods. Figure 6 compares the training efficiency of our method against diffusion-based (Diffusion-BC) and consistency-based (Consistency-BC) baselines. Baseline results are normalized from literature values, assuming 1 epoch = 500 steps vs.our 1,000 steps.

The results demonstrate that MoMa QL achieves superior computational efficiency. On average, our method is approximately \(\mathbf{6\times}\) faster than Diffusion-BC and \(\mathbf{1.5\times}\) faster than Consistency-AC on complex tasks like Adroit and Kitchen. Notably, MoMa QL maintains a constant training cost (\(\approx 21-23\) seconds per 1k steps) regardless of task dimensionality, whereas baselines exhibit significant slowdowns on high-dimensional tasks.

Impact of Sampling Steps. We further analyze how the number of sampling steps affects training efficiency. Figure 7 shows the relationship between the number of sampling steps (1, 2, 4, 8, 16) and training time per 1,000 gradient steps across Gym, Adroit, and Kitchen domains.

The results reveal several key insights: (1) Training time scales approximately linearly with the number of sampling steps, increasing from \(\sim\)​17-18s for single-step sampling to \(\sim\)​71-73s for 16-step sampling across all domains. (2) MoMa QL maintains consistent efficiency regardless of task complexity—the training cost remains nearly identical across Gym locomotion (\(17.70 \rightarrow 72.56\)s), Adroit manipulation (\(17.64 \rightarrow 73.12\)s), and Kitchen (\(17.71 \rightarrow 70.94\)s) tasks, demonstrating excellent scalability. (3) Even with 16 sampling steps, MoMa QL remains faster than Diffusion-BC (\(\sim\)​132-189s) and Consistency-BC (\(\sim\)​83-130s) baselines. This linear scaling behavior enables practitioners to flexibly trade off between sample quality and computational cost based on task requirements. Detailed per-step timing results are provided in Appendix 9.5.2.

Detailed per-task training and inference time comparisons are provided in Appendix 9.5.

6 Conclusion↩︎

In this paper, we presented Moment Matching Q-Learning (MoMa QL), an efficient RL framework that integrates MMD regularization with stochastic interpolants. MoMa QL strikes a balance between modeling expressivity and computational efficiency, achieving state-of-the-art results. Empirically, it surpasses Diffusion-QL by \(1.09\times\) on Gym tasks and maintains competitive performance on complex Adroit and Kitchen domains. Crucially, our method is approximately and \(1.5\times\) faster than Consistency-AC, eliminating the inference bottleneck of generative policies. Beyond efficiency, MoMa QL exhibits strong robustness to hyperparameters and effective offline-to-online fine-tuning capabilities. We hope this simple yet rigorous approach spurs future research into scalable generative policies.

Acknowledgements↩︎

We thank the anonymous reviewers for their helpful comments. This research was supported by WZ’s startup funding provided by the School of Data Science and Society at UNC Chapel Hill. This research supported by the NVIDIA Academic Grant Program with the cloud GPU credits and infrastructure.

Impact Statement↩︎

This paper presents work whose goal is to advance the field of Reinforcement Learning and Robotics. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.

langley00?

7 Theorems and Derivations↩︎

7.1 Maximum Mean Discrepancy↩︎

To facilitate a further discussion of MMD, it is necessary to first introduce several key mathematical definitions.

For instance, 1-Wasserstein distance \(W_1\) is an IPM equipped with function class \(\mathcal{F} = \mathrm{Lip}_1 := \{ f: \mathcal{X} \rightarrow \mathbb{R} \mid |f(x) - f(y)| \leq d(x, y), \forall x, y \in \mathcal{X} \}\), which is exactly the loss function in WGAN [57]. We mentioned Reproducing Kernel Hilbert Space (RKHS) in section 2.4, and now we give the formal definition of RKHS as:

Since \(L_x\) is continuous and linear, we can apply Riesz Representation Theorem presented as follows, and thus prove the existence of a kernel function \(K(\cdot, \cdot), K_x := K(x, \cdot) \in \mathcal{H}\) which satisfies the \(f(x) = \langle f(\cdot), K(\cdot, x) \rangle_{\mathcal{H}}\) (referred as "reproducing property").

We can also prove the finiteness of MMD with this useful proposition.

Furthermore, the MMD is bounded if the kernel \(k\) is bounded: \[\begin{align} \text{MMD}(\mu, \nu) &\leq \int \|k(x, \cdot)\|_{\mathcal{H}} \mu(dx) + \int \|k(x, \cdot)\|_{\mathcal{H}} \nu(dx) \\ &\leq 2 \sup_{x \in \mathbb{R}^d} \sqrt{k(x, x)} < \infty, \end{align}\] where we utilize the fact that \(\|k(x, \cdot)\|_{\mathcal{H}} = \sqrt{\langle k(x, \cdot), k(x, \cdot) \rangle_{\mathcal{H}}} = \sqrt{k(x, x)}\). Unlike Wasserstein distances, MMD does not suffer from the curse of dimensionality, and thus is a powerful metric for matching the gap between high-dimensional distributions.

7.2 Existence of Divergence Minimizer↩︎

In this section, we intend to prove the existence of the optimal \(p^{\theta^*}_{s|t}\), and we can further prove that this optimal solution can be learnt inductively in the form of Eq. 13 .

In general, the minimizer \(q_t(\mathbf{x}|\mathbf{x}_t)\) exists, although it may not be unique. We also should discuss more about the mapping function \(0 \leq s \leq r := r(s, t) \leq t \leq 1\). We set \(r(s, t) = \mathrm{max} (s, t - \Delta_t)\) where \(\Delta_t \geq \epsilon > 0\). In this case, we have can prove the convergence.

7.3 Empirical Loss Function↩︎

For ODE based diffusion model, DDIM Interpolants can be formally defined as: \[\mathsf{DDIM}(\mathbf{x}_t, \mathbf{x}, s, t) = \left( \alpha_s - \frac{\sigma_s}{\sigma_t} \alpha_t \right) \mathbf{x} + \frac{\sigma_s}{\sigma_t} \mathbf{x}_t \label{eq31}\tag{22}\] and sample \(\mathbf{x}_s = \mathsf{DDIM}(\mathbf{x}_t, \mathbb{E}_\mathbf{x}[\mathbf{x}|\mathbf{x_t}], s, t)\) can be drawn when \(\mathbb{E}_\mathbf{x}[\mathbf{x}|\mathbf{x_t}]\) is approximated by a network. DDIM can be viewed as an interpolant with \(\gamma_{s|t} \equiv 0\) and \(I_{s|t} = \mathsf{DDIM}(\mathbf{x}_t, \mathbf{x}, s, t)\), and it is a marginal-preserving interpolants, which means that there is a deterministic sampler which satisfies [thma7].

To further justify the marginal-preserving property used in Theorem 3, we provide a proof sketch showing that DDIM interpolants intrinsically satisfy the marginal-preserving equality.

Consider a DDIM interpolant defined as \[\mathbf{x}_t = I(\mathbf{x}, \boldsymbol{\epsilon}, t) = a_t \mathbf{x} + b_t \boldsymbol{\epsilon},\] where \(\mathbf{x} \sim q(\mathbf{x})\) and \(\boldsymbol{\epsilon} \sim p(\boldsymbol{\epsilon})\). The marginal distribution at time \(t\) is \[q_t(\mathbf{x}_t) = \iint \delta(\mathbf{x}_t - I(\mathbf{x}, \boldsymbol{\epsilon}, t)) q(\mathbf{x}) p(\boldsymbol{\epsilon}) \,\mathrm{d}\mathbf{x} \,\mathrm{d}\boldsymbol{\epsilon}.\]

Under the DDIM formulation, the trajectory is deterministic conditioned on \(\mathbf{x}\). Therefore, for arbitrary \(s,t \in [0,1]\), the state \(\mathbf{x}_s\) is uniquely determined by \(\mathbf{x}\) and \(\mathbf{x}_t\), and the conditional distribution can be written as \[q_{s|t}(\mathbf{x}_s|\mathbf{x}, \mathbf{x}_t) = \delta \left( \mathbf{x}_s - \left[ a_s \mathbf{x} + b_s \frac{\mathbf{x}_t - a_t \mathbf{x}}{b_t} \right] \right).\]

Substituting the DDIM conditional distribution into the right-hand side of the marginal-preserving equality yields \[\begin{align} &\iint q_{s|t}(\mathbf{x}_s|\mathbf{x}, \mathbf{x}_t) q_t(\mathbf{x}, \mathbf{x}_t) \,\mathrm{d}\mathbf{x}_t \,\mathrm{d}\mathbf{x} \nonumber\\ = &\iiint \delta \left( \mathbf{x}_s - \left[ a_s \mathbf{x} + \frac{b_s}{b_t} (\mathbf{x}_t - a_t \mathbf{x}) \right] \right) \delta(\mathbf{x}_t - (a_t \mathbf{x} + b_t \boldsymbol{\epsilon})) q(\mathbf{x}) p(\boldsymbol{\epsilon}) \,\mathrm{d}\boldsymbol{\epsilon} \,\mathrm{d}\mathbf{x}_t \,\mathrm{d}\mathbf{x}. \end{align}\]

Using the DDIM construction, \[a_s \mathbf{x} + \frac{b_s}{b_t} \left( (a_t \mathbf{x} + b_t \boldsymbol{\epsilon}) - a_t \mathbf{x} \right) = a_s \mathbf{x} + b_s \boldsymbol{\epsilon} = I(\mathbf{x}, \boldsymbol{\epsilon}, s),\] we obtain \[\begin{align} &\iint q_{s|t}(\mathbf{x}_s|\mathbf{x}, \mathbf{x}_t) q_t(\mathbf{x}, \mathbf{x}_t) \,\mathrm{d}\mathbf{x}_t \,\mathrm{d}\mathbf{x} \nonumber\\ = &\iint \delta(\mathbf{x}_s - I(\mathbf{x}, \boldsymbol{\epsilon}, s)) q(\mathbf{x}) p(\boldsymbol{\epsilon}) \,\mathrm{d}\boldsymbol{\epsilon} \,\mathrm{d}\mathbf{x} = q_s(\mathbf{x}_s). \end{align}\]

Therefore, DDIM interpolants intrinsically satisfy the marginal-preserving equality.

Since DDIM interpolants satisfy the marginal-preserving equality, we can reuse the sampled state \(\mathbf{x}_t\) to construct \(\mathbf{x}_r\) through the conditional distribution \[\mathbf{x}_r \sim q_{r|t}(\mathbf{x}_r \mid \mathbf{x}, \mathbf{x}_t),\] instead of independently sampling \(\mathbf{x}_r\) from the forward flow \[\mathbf{x}_r = \alpha_r \mathbf{x} + \sigma_r \boldsymbol{\epsilon}.\] In particular, we use the deterministic DDIM mapping \[\mathbf{x}_r = \mathsf{DDIM}(\mathbf{x}_t, \mathbf{x}, r, t),\] which reduces the variance of empirical estimation.

More generally, for any marginal-preserving interpolant, the marginal distribution of \(\mathbf{x}_r\) remains unchanged under conditional resampling: \[\begin{align} q_r(\mathbf{x}_r) = q_{r|t}(\mathbf{x}_r) &= \iint q_{r|t}(\mathbf{x}_r \mid \mathbf{x}, \mathbf{x}_t) q_t(\mathbf{x} \mid \mathbf{x}_t) q_t(\mathbf{x}_t) \,\mathrm{d}\mathbf{x} \,\mathrm{d}\mathbf{x}_t \\ &= \iiint q_{r|t}(\mathbf{x}_r \mid \mathbf{x}, \mathbf{x}_t) q_t(\mathbf{x}_t \mid \mathbf{x}, \boldsymbol{\epsilon}) q(\mathbf{x}) p(\boldsymbol{\epsilon}) \,\mathrm{d}\boldsymbol{\epsilon} \,\mathrm{d}\mathbf{x} \,\mathrm{d}\mathbf{x}_t. \end{align}\] Therefore, sampling \((\mathbf{x}, \mathbf{x}_t)\) first and then drawing \[\mathbf{x}_r \sim q_{r|t}(\mathbf{x}_r \mid \mathbf{x}, \mathbf{x}_t)\] still preserves the target marginal distribution \(q_r(\mathbf{x}_r)\).

Next, we try to simplify the objective function 13 . Given MMD loss function we defined in 6 , we have: \[\begin{align} \mathcal{L}_{D}(\theta) &= \mathbb{E}_{s,t} \left[\left\| \mathbb{E}_{\mathbf{x}_t} [k(f_{s,t}^\theta(\mathbf{x}_t), \cdot)] - \mathbb{E}_{\mathbf{x}_r} [k(f_{s,r}^{\theta^-}(\mathbf{x}_r), \cdot)] \right\|_{\mathcal{H}}^2 \right] \\ &= \mathbb{E}_{s,t} \left[ \left\| \mathbb{E}_{\mathbf{x}_t, \mathbf{x}_r} [k(f_{s,t}^\theta(\mathbf{x}_t), \cdot) - k(f_{s,r}^{\theta^-}(\mathbf{x}_r), \cdot)] \right\|_{\mathcal{H}}^2 \right] \quad (\text{reuse the sample \mathbf{x}_t})\\ &= \mathbb{E}_{s,t} \left[ \left\langle \mathbb{E}_{\mathbf{x}_t, \mathbf{x}_r} [k(f_{s,t}^\theta(\mathbf{x}_t), \cdot) - k(f_{s,r}^{\theta^-}(\mathbf{x}_r), \cdot)], \mathbb{E}_{\mathbf{x}_t', \mathbf{x}_r'} [k(f_{s,t}^\theta(\mathbf{x}_t'), \cdot) - k(f_{s,r}^{\theta^-}(\mathbf{x}_r'), \cdot)] \right\rangle_\mathcal{H} \right] \\ &= \mathbb{E}_{s,t} \left[ \mathbb{E}_{\mathbf{x}_t, \mathbf{x}_r, \mathbf{x}_t', \mathbf{x}_r'} \left[ \left \langle k(f_{s,t}^\theta(\mathbf{x}_t), \cdot), k(f_{s,t}^\theta(\mathbf{x}_t'), \cdot) \right \rangle_\mathcal{H} + \left \langle k(f_{s,r}^{\theta^-}(\mathbf{x}_r), \cdot), k(f_{s,r}^{\theta^-}(\mathbf{x}_r'), \cdot) \right \rangle_\mathcal{H} \right. \right. \\ &\qquad\qquad \left. \left. - \left \langle k(f_{s,t}^\theta(\mathbf{x}_t), \cdot), k(f_{s,r}^{\theta^-}(\mathbf{x}_r'), \cdot) \right \rangle_\mathcal{H} - \left \langle k(f_{s,t}^\theta(\mathbf{x}_t'), \cdot), k(f_{s,r}^{\theta^-}(\mathbf{x}_r), \cdot) \right \rangle_\mathcal{H} \right] \right] \\ &= \mathbb{E}_{\mathbf{x}_t, \mathbf{x}_r, \mathbf{x}_t', \mathbf{x}_r', s, t} \left[\left[ k(f_{s,t}^\theta(\mathbf{x}_t), f_{s,t}^\theta(\mathbf{x}_t')) + k(f_{s,r}^{\theta^-}(\mathbf{x}_r), f_{s,r}^{\theta^-}(\mathbf{x}_r')) \right. \right. \\ &\qquad\qquad \left. \left. - k(f_{s,t}^\theta(\mathbf{x}_t), f_{s,r}^{\theta^-}(\mathbf{x}_r')) - k(f_{s,t}^\theta(\mathbf{x}_t'), f_{s,r}^{\theta^-}(\mathbf{x}_r)) \right] \right] \quad (\text{reproducing property}) \end{align}\]

If we select kernel function \(k(x, y) = -||x - y||_2^2\), \(\mathbf{x}_t = \mathbf{x}_t', \mathbf{x}rt = \mathbf{x}_r'\), and \(s = \epsilon\), we can see that the MMD loss function is exactly the same as the consistency training loss in the form of \(\mathcal{L}_{CT} = \mathbb{E}_{\mathbf{x}_t, \mathbf{x}, t} \left[ \left\| f_\theta(\mathbf{x}_t, t) - f_{\theta^-}(\mathbf{x}_r, r) \right\|^2 \right]\). From a moment-matching perspective, \(\mathcal{L}_{CT}\) significantly deviates from a proper divergence between distributions, as (i) \(\mathcal{L}_{CT}\) only conduct multi-particle estimation, which may pause mode collapse and training instability and (ii) this kernel is not positive definite kernel and can only match the first order (Expectation) and second order (Variance) moment, which can hardly match two high-dimensional distributions.

As proposed in [48], MMD loss can be estimated with V-statistics by instantiating a matrix of size \(M \times M\) such that a batch of \(B\) samples, \(\{x^{(i)}\}_{i=1}^B\), is separated into groups of \(M\) (assume \(B\) is divisible by \(M\)) particles \(\{x^{(i,j)}\}_{i=1,j=1}^{B/M,M}\) where each group share a \((s^i, r^i, t^i)\) sample. We also add the weight function \(w(s, t) = \frac{1}{\alpha_t^2 + \sigma_t^2}\) to improve the performance. The estimation of the MMD loss can be represented as: \[\begin{align} \hat{\mathcal{L}}_{D}(\theta) = \frac{1}{B/M} \sum_{i=1}^{B/M} w(s^i, t^i) \frac{1}{M^2} \sum_{j=1}^M \sum_{k=1}^M \Big[ &k(f_{s^i,t^i}^\theta(\mathbf{x}_{t^i}^{(i,j)}), f_{s^i,t^i}^\theta(\mathbf{x}_{t^i}^{(i,k)})) + k(f_{s^i,r^i}^{\theta^-}(\mathbf{x}_{r^i}^{(i,j)}), f_{s^i,r^i}^{\theta^-}(\mathbf{x}_{r^i}^{(i,k)})) \nonumber \\ &- 2k(f_{s^i,t^i}^\theta(\mathbf{x}_{t^i}^{(i,j)}), f_{s^i,r^i}^{\theta^-}(\mathbf{x}_{r^i}^{(i,k)})) \Big] \label{eq32} \end{align}\tag{23}\]

This is the final loss firm that we utilize to train the model. Compared to the loss function in consistency models, this empirical MMD loss can capture the high-order moment information of target distribution via a kernel function, and thus can match the distribution more accurately.

8 Full Algorithm↩︎

A more detailed training algorithm is presented in Algorithm 8. We utilized RBF kernel as default kernel function \(k(\cdot, \cdot)\), other kernels such as Laplacian kernel are also applicable.

9 Experiments↩︎

9.1 Experimental Hyperparameters↩︎

We provide the detailed hyperparameters used for training MoMa QL in Table 6. We use a consistent set of hyperparameters across most D4RL Gym locomotion tasks to demonstrate the robustness of our method. For the Kitchen and Adroit environments, minor adjustments (e.g., learning rate or epoch count) were made to accommodate the different task complexities, as noted.

9.2 Extended Offline RL Results and Analysis↩︎

This appendix provides a detailed breakdown of MoMa QL’s performance across all individual D4RL datasets. Table 9 presents the full comparison with standard deviations included.

9.2.1 Detailed Analysis↩︎

Gym Locomotion Tasks. MoMa QL demonstrates remarkable efficacy on suboptimal and diverse datasets. On halfcheetah-m, our method achieves a score of \(\mathbf{72.6 \pm 1.1}\), surpassing the best baseline (Consistency-AC, \(69.1\)) by approximately \(5\%\). Similarly, on the challenging walker2d-mr dataset, MoMa QL attains \(\mathbf{104.3 \pm 3.3}\), outperforming Diffusion-QL (\(95.5\)) by a factor of \(1.09\). These results underscore the robustness of our moment-matching objective in extracting optimal policies from noisy data. While Diffusion-QL holds a slight edge on hopper-me (\(111.1\) vs.our \(88.9 \pm 9.6\)), MoMa QL consistently dominates on medium and medium-replay datasets.

Adroit Manipulation Tasks. The detailed results reveal specific strengths of our method. On door-expert-v1, MoMa QL achieves a near-perfect score of \(\mathbf{108 \pm 1}\), narrowly edging out ReBRAC (\(106\)) and IQL (\(107\)). In the pen-cloned-v1 task, our method reaches \(\mathbf{104 \pm 8}\), matching the performance of ReBRAC. However, on human datasets like hammer-human, all methods struggle (scores \(<5\)), indicating a general difficulty in learning from these demonstrations, though MoMa QL remains competitive with a score of \(\mathbf{4 \pm 0}\).

Kitchen Tasks. Our method shines particularly on the kitchen-complete task, achieving a score of \(\mathbf{91.3 \pm 3.6}\), which is an improvement of roughly \(1.09\times\) over Diffusion-QL (\(84.0\)) and \(1.11\times\) over \(\mathcal{X}\)-QL (\(82.4\)). This suggests that for tasks requiring the completion of full sequential goals, MoMa QL’s generative policy is highly effective. On partial and mixed datasets, performance is comparable to leading baselines, maintaining stable results around the \(60.0\) mark.

9.3 Extended Offline-to-Online RL Results↩︎

This appendix provides comprehensive analysis of MoMa QL’s performance in the offline-to-online setting, including detailed per-dataset results.

Table 10 presents the complete offline-to-online results across all datasets with standard deviations. For each task, we compare against strong baselines including Diffusion-QL and Consistency-AC.

Analysis. MoMa QL achieves an average score of 101.4 across all tasks, demonstrating superior final performance.

• On Medium-Replay datasets, our method excels, achieving 113.5 on Hopper-mr and 117.9 on Walker2d-mr, significantly outperforming Diffusion-QL (68.4 and 95.7, respectively).

• On Medium-Expert datasets, MoMa QL remains competitive, with particularly strong results on Walker2d-me (118.0).

• The transition from offline to online is stable, with no catastrophic forgetting observed, a common issue in fine-tuning generative policies.

9.4 Detailed Ablation Studies↩︎

In this section, we provide a comprehensive analysis of the hyperparameters used in MoMa QL. We systematically vary one parameter at a time while keeping others fixed.

9.4.1 Hyperparameter Definitions↩︎

• \(a\): Controls the power of \(\alpha_t\) in the kernel weighting function \(\alpha_t^a / (\alpha_t^2 + \sigma_t^2)\). A larger \(a\) emphasizes cleaner samples (small noise), allowing the model to capture fine details.

• \(b\): Appears in the sigmoid weighting term \((b - \log \text{SNR}_t)\). It shifts the center of the weighting function. A higher \(b\) focuses on noisier samples, while a lower \(b\) focuses on cleaner samples.

• \(\eta\): The regularizing weight for the Q-value loss. It controls how much we trust the Q-function versus the behavioral distribution.

• Number of Steps: The number of denoising steps used during action sampling.

• \(P_{\text{mean}}\): The mean of the log-normal distribution for sampling time steps \(t\) during training. a negative value biases sampling towards \(t=0\) (clean data).

• \(P_{\text{std}}\): The standard deviation of the time step distribution. A larger value covers a wider range of noise levels.

9.4.2 Additional Results↩︎

We present the full ablation results for all parameters across HalfCheetah, Hopper, and Walker2d environments.

9.5 Detailed Computational Cost Results↩︎

9.5.1 Baseline Comparison↩︎

Table 11 reports the detailed computational cost for MoMa QL compared with Diffusion-BC and Consistency-BC. We report both Training Time (seconds per 1,000 steps) and Inference Time (seconds per episode) for MoMa QL with 2 sampling steps. Baseline training times are converted from reported values (seconds per 500-step epoch \(\times 2\)) for fair comparison.

MoMa QL maintains a consistent training cost (\(\approx 21-23\)s per 1k steps) across all tasks, demonstrating superior scalability compared to baselines which slow down significantly on higher-dimensional tasks (Adroit, Kitchen). All methods use comparable batch sizes and hardware configurations.

9.5.2 Impact of Sampling Steps↩︎

Tables 12, 13, and 14 present the detailed training time (seconds per 1,000 steps) as a function of the number of sampling steps for MoMa QL across all D4RL tasks. The results demonstrate how computational cost scales with sampling steps while maintaining efficiency.

Key observations: The training time exhibits nearly linear scaling with respect to the number of sampling steps across all domains. For single-step sampling (\(n=1\)), the average training time is 17.70s (Gym), 17.64s (Adroit), and 17.71s (Kitchen). This increases to 72.56s, 73.12s, and 70.94s respectively for 16-step sampling, representing approximately a \(4\times\) increase for a \(16\times\) increase in steps. This sub-linear scaling demonstrates the computational efficiency of our MMD-based approach. Notably, task complexity has minimal impact on training time—high-dimensional Kitchen and Adroit tasks require similar computational resources as lower-dimensional Gym tasks, highlighting MoMa QL’s superior scalability compared to diffusion-based baselines.

References↩︎