AlignIQL: Policy Alignment in Implicit Q-Learning through Constrained Optimization

## Abstract

Abstract Implicit Q-learning (IQL) serves as a strong baseline for offline RL, which learns the value function using only dataset actions through quantile regression. However, it is unclear how to recover the implicit policy from the learned implicit Q-function and why IQL can utilize weighted regression for policy extraction. IDQL reinterprets IQL as an actor-critic method and gets weights of implicit policy, however, this weight only holds for the optimal value function. In this work, we introduce a different way to solve the implicit policy-finding problem (IPF) by formulating this problem as an optimization problem. Based on this optimization problem, we further propose two practical algorithms AlignIQL and AlignIQL-hard, which inherit the advantages of decoupling actor from critic in IQL and provide insights into why IQL can use weighted regression for policy extraction. Compared with IQL and IDQL, we find our method keeps the simplicity of IQL and solves the implicit policy-finding problem. Experimental results on D4RL datasets show that our method achieves competitive or superior results compared with other SOTA offline RL methods. Especially in complex sparse reward tasks like Antmaze and Adroit, our method outperforms IQL and IDQL by a significant margin. Code is available at https://github.com/felix-thu/AlignIQL.git

## 1 Introduction

Offline Reinforcement Learning (RL), or Batch RL aims to seek an optimal policy without environmental interactions [, ]. This is compelling for having the potential to transform large-scale datasets into powerful decision-making tools and avoids costly and risky online environmental interactions, which offers significant application prospects in fields such as healthcare [, ] and autopilot [, ]. Notwithstanding its promise, applying off-policy RL algorithms [, , , ] directly into the offline context presents challenges due to out-of-distribution actions that arise when evaluating the learned policy.[, ].

Despite a variety of methods based on constraint and conservative Q-learning have been proposed to address this, IQL [] stands out among them since IQL avoids query out-of-distribution (OOD) actions entirely and decouples the critic from the actor, which contributes to stability and hyperparameter robust. However, IQL extracts policy through advantage-weighted regression (AWR) [, , ], resulting in the inconsistency between the extracted policy and the implicit policy in the learned Q-function. To solve this problem, IDQL [] reinterprets IQL as
an actor-critic method and derives the implicit optimal policy weights. Nevertheless, this optimal weight hinges on the assumption that the optimal value function can be learned through parametrization (e.g. neural networks) and certain critic loss functions. It remains unclear whether extracting implicit policy with AWR is feasible and what is the policy implied by an arbitrary value function. Moreover, many recent offline RL [] and safe RL methods [] use IQL to learn the Q-function. Addressing these issues can lead to a better understanding of the bottlenecks of IQL-style methods, thereby promoting the development of offline RL.

In this paper, we solve this problem by formulating the policy-finding problem as an optimization problem, where the objective function is a generalized form of behavior regularizers and the constraint is policy alignment, which ensures the learned policy is the policy implied in the value function. By solving this optimization problem, we can get a closed-form solution, which can be expressed by imposing weight on the behavior policy. The weight can be computed by a value function, an action-value function, and multipliers. This weight shows that using AWR for policy extraction in IQL is feasible when the certain multiplier $<0$. Furthermore, our work explains how the implicit policy in IQL-style methods addresses OOD actions from the perspective of behavior regularizers.

Based on the optimization problem, we further propose two practical algorithms, AlignIQL-hard and AlignIQL. Both inherit the characteristics of IQL, namely the decoupling of actor and critic training and implicitly trained value function while achieving policy alignment. AlignIQL-hard can theoretically achieve a globally optimal solution, but it is more vulnerable to hyperparameter choices than AlignIQL. AlignIQL is simpler to implement and performs better in complex tasks like sparse rewards tasks, but it does not guarantee convergence to the global optimum.

Since behavior policy is often complex and potentially multi-modal, the unimodal Gaussian policy used with IQL is unlikely to accurately approximate the complex behavior policy [, , , ], which in turn affects the implicit policy extraction, as the general form of implicit policy is $\pi({\bm{s}}|{\bm{a}})\propto\mu({\bm{a}}|{\bm{s}})w({\bm{s}},{\bm{a}})$, where $\mu({\bm{a}}|{\bm{s}})$ is the behavior policy.
We, therefore, use the diffusion model [, , ] to model the behavior policy.
For policy extraction, we select actions from a diffusion-parameterized behavior model with weights derived from our policy-finding problem, which enjoy both superior expressivity and policy alignment. We verify the effectiveness of AlignIQL on D4RL datasets and demonstrate the state-of-the-art performance, especially on sparse reward datasets in which the critic learning is hard and unstable. We also show that, compared to IDQL, AlignIQL is more robust to hyperparameters and achieves more stable training.

To summarize, our main contributions are as follows:

- •
We propose the policy-finding problem, where the policy alignment term is added as a constraint. By solving this problem, we provide insights into why and when IQL can use weighted regression for policy extraction, and in turn, make it better to understand the bottlenecks of the IQL-style algorithms.
- •
We demonstrate that there is no price to achieving policy alignment in IQL-style methods, all we need is to modify the importance weights of the policy extraction based on the sign of multiplier or hyperparameter $\eta$. These results can be generalized to any generalized value loss function, which greatly extends the theoretical results of IDQL.
- •
We introduce two practical IQL-style algorithms: AlignIQL-hard and AlignIQL, which obtain competitive performance on D4RL datasets and show robustness on sparse reward tasks. Extensive experiments, together with the theoretical analysis, validate the effectiveness of our proposed approaches.

## 2 Background

Offline RL. Consider a Markov decision process (MDP): $M=\{{\mathcal{S}},{\mathcal{A}},P,R,\gamma,d_{0}\}$, with state space $S$, action space ${\mathcal{A}}$, environment dynamics ${\mathcal{P}}({\bm{s}}^{\prime}|{\bm{s}},{\bm{a}}):S\times S\times{\mathcal{A}}\rightarrow[0,1]$, reward function $R:S\times{\mathcal{A}}\rightarrow{\mathbb{R}}$, discount factor $\gamma\in[0,1)$, policy $\pi({\bm{a}}|{\bm{s}}):{\mathcal{S}}\times{\mathcal{A}}\rightarrow[0,1]$, and initial state distribution $d_{0}$. The action-value or Q-value of policy $\pi$ is defined as $Q^{\pi}({\bm{s}}_{t},{\bm{a}}_{t})=\mathbb{E}_{{\bm{a}}_{t+1},{\bm{a}}_{t+2},...\sim\pi}\left[\sum_{j=0}^{\infty}\gamma^{j}r({\bm{s}}_{t+j},{\bm{a}}_{t+j})\right]$. The goal of RL is to get a policy to maximize the cumulative discounted reward
$J(\theta)=\int_{\mathcal{S}}d_{0}({\bm{s}})Q^{\pi}({\bm{s}},{\bm{a}})d{\bm{s}}$. $d^{\pi}({\bm{s}})=\sum_{t=0}^{\infty}\gamma^{t}p_{\pi}({\bm{s}}_{t}={\bm{s}})$ is the state visitation distribution induced by policy $\pi$ [, ], and $p_{\pi}({\bm{s}}_{t}={\bm{s}})$ is the likelihood of the policy being in state ${\bm{s}}$ after following $\pi$ for $t$ timesteps. In offline setting [], environmental interaction is not allowed, and a static dataset ${\mathcal{D}}\triangleq\left\{({\mathcal{S}},{\mathcal{A}},R,{\mathcal{S}}^{\prime},\text{done})\right\}$ is used to learn a policy.

Implicit Q-learning (IQL). To avoid OOD actions in offline RL, IQL [] uses the state conditional upper expectile of action-value function $Q({\bm{s}},{\bm{a}})$ to estimate the value function $V({\bm{s}})$, which avoid directly querying a Q-function with unseen action. For a parameterized critic $Q_{\theta}(s,a)$, target critic $Q_{\hat{\theta}}({\bm{s}},{\bm{a}})$, and value network $V_{\psi}({\bm{s}})$ the value objective is learned by

$$ $\displaystyle{\mathcal{L}}_{V}(\psi)=\mathbb{E}_{({\bm{s}},{\bm{a}})~{}\sim{\mathcal{D}}}[L_{2}^{\tau}(Q_{\hat{\theta}}({\bm{s}},{\bm{a}})-V_{\psi}({\bm{s}}))]$ (1) $\displaystyle\text{where}\quad L_{2}^{\tau}(u)=|\tau-\mathbbm{1}(u<0)|u^{2},$ $$

where $\mathbbm{1}$ is the indicator function.
Then, the Q-function is learned by minimizing the MSE loss

$$ $\displaystyle\mathcal{L}_{Q}(\theta)$ $\displaystyle=\mathbb{E}_{({\bm{s}},{\bm{a}},{\bm{s}}^{\prime})~{}\sim\mathcal{D}}[(r({\bm{s}},{\bm{a}})+\gamma V_{\psi}({\bm{s}}^{\prime})-Q_{\theta}({\bm{s}},{\bm{a}}))^{2}].$ (2) $$

Note that, in IQL, the corresponding policy is not explicitly represented, it is implicit in the learned value function. For policy extraction, IQL uses AWR [, , ]

$$ $\displaystyle\mathcal{L}_{\pi}(\phi)$ $\displaystyle=\mathbb{E}_{({\bm{s}},{\bm{a}})\sim\mathcal{D}}[\exp(\alpha(Q_{\hat{\theta}}({\bm{s}},{\bm{a}})-V_{\psi}({\bm{s}})))\log\pi_{\phi}({\bm{a}}|{\bm{s}})].$ (3) $$

Implicit Diffusion Q-learning (IDQL). To find the implicit policy in the learned value function, IDQL [] generalizes the value loss in Equation LABEL:eqn:fit_v_expectiles with an arbitrary convex loss $U$ on the difference $Q-V$.

$$ $V^{*}({\bm{s}})=\operatorname*{arg\,min}_{V({\bm{s}})}\mathbb{E}_{{\bm{a}}\sim\mu({\bm{a}}|{\bm{s}})}[U(Q({\bm{s}},{\bm{a}})-V({\bm{s}}))]=\operatorname*{arg\,min}_{V({\bm{s}})}\mathcal{L}_{V}^{U}(V({\bm{s}})).$ (4) $$

Under some assumptions about $U$, IDQL derives the implicit policy in optimal $V$ defined in Equation [4](#S2.E4)

$$ $w({\bm{s}},{\bm{a}})=\frac{|U^{\prime}(Q({\bm{s}},{\bm{a}})-V^{*}({\bm{s}}))|}{|Q({\bm{s}},{\bm{a}})-V^{*}({\bm{s}})|},$ (5) $$

which yields an expression for the implicit actor as $\pi_{\text{imp}}({\bm{a}}|{\bm{s}})\propto\mu({\bm{a}}|{\bm{s}})w({\bm{s}},{\bm{a}})$.

However, it remains unclear whether it is feasible to extract policy with AWR. Answering this question can help us better understand the bottlenecks of IQL-style methods. Before drawing our conclusion, we first introduce the form of our Implicit Policy-finding Problem.

## 3 Implicit Policy-finding Problem

Before introducing the implicit policy-finding problem, we first introduce Definition [3.1](#S3.Thmtheorem1), which defines what a policy implied by a value function $Q({\bm{s}},{\bm{a}})$ is.

###### Definition 3.1 .

We refer to a policy as one implied by the value function $Q({\bm{s}},{\bm{a}}),V({\bm{s}})$, when

$$ $Q({\bm{s}},{\bm{a}})-r({\bm{s}},{\bm{a}})-\gamma\mathbb{E}_{{\bm{s}}^{\prime}\sim p({\bm{s}}^{\prime}|{\bm{s}},{\bm{a}}),{\bm{a}}^{\prime}\sim\pi({\bm{a}}|{\bm{s}})}\left[Q({\bm{s}}^{\prime},{\bm{a}}^{\prime})\right]=0.$ (6) $$

$$ $\mathbb{E}_{{\bm{a}}\sim\pi({\bm{a}}|{\bm{s}})}\left[Q({\bm{s}},{\bm{a}})\right]=V({\bm{s}}),$ (7) $$

Definition [3.1](#S3.Thmtheorem1) is derived from IDQL [] and the conventional definition of the value function in actor-critic methods.
Note that in IQL, the $Q$-function is updated by minimizing Equation [2](#S2.E2), which implies if we can ensure Equation [7](#S3.E7), Equation [6](#S3.E6) can be derived by substituting Equation [7](#S3.E7) back in Equation [2](#S2.E2). So in IQL, we can define the implicit policy when it satisfies Equation [7](#S3.E7).

It is known that the offline RL problem can be solved by constrained policy search (CPS) problem (aka AWR) [, , ], where a policy is sought to maximize cumulative rewards under the constraint of policy divergence from the behavior policy. (See more details in Appendix [C](#A3)) Inspired by CPS, we formulate the implicit policy-finding problem (IPF) as a constrained optimization, where a policy is sought to minimize policy divergence from the behavior policy under policy alignment

$$ $\displaystyle\min_{\pi}\quad\mathbb{E}_{{\bm{s}}\sim d^{\pi}({\bm{s}}),{\bm{a}}\sim\pi({\bm{a}}|{\bm{s}})}\left[f\left(\frac{\pi({\bm{a}}|{\bm{s}})}{\mu({\bm{a}}|{\bm{s}})}\right)\right]$ (IPF) $\displaystyle s.t.\quad\pi({\bm{a}}|{\bm{s}})\geq 0,\quad\forall{\bm{s}},\forall{\bm{a}}$ $\displaystyle\quad\int_{\bm{a}}\pi({\bm{a}}|{\bm{s}})d{\bm{a}}=1,\quad\forall{\bm{s}}$ $\displaystyle\quad\mathbb{E}_{{\bm{a}}\sim\pi({\bm{a}}|{\bm{s}})}\left[Q({\bm{s}},{\bm{a}})\right]-V({\bm{s}})=0,\quad\forall{\bm{s}},$ $$

where $V({\bm{s}}),Q({\bm{s}},{\bm{a}})$ is the learned value function, which does not have to be the optimal value function. $f(\cdot)$ is a regularization function which aims to avoid out-of-distribution actions. The third constraint ensures that the learned policy is the policy implied in $Q,V$.

Here we briefly describe the characteristics of the solution to problem [IPF](#S3.Ex1). In Problem [IPF](#S3.Ex1), when the feasible set includes multiple policies ($i.e.$
multiple implicit policies satisfy Definition [3.1](#S3.Thmtheorem1)), problem [IPF](#S3.Ex1) aims to find an optimal implicit policy that deviates least from the behavior policy while satisfying the requirements of policy alignment. In other cases, when the feasible set has a unique policy, problem [IPF](#S3.Ex1) will return the unique policy as the optimal implicit policy. The above analysis shows that we can model the implicit policy-finding problem in IQL as problem [IPF](#S3.Ex1).

###### Assumption 3.2 .

Assume $\pi({\bm{a}}|{\bm{s}})>0\Longrightarrow\mu({\bm{a}}|{\bm{s}})>0$ so that $\frac{\pi({\bm{a}}|{\bm{s}})}{\mu({\bm{a}}|{\bm{s}})}$ is well-defined. []

###### Assumption 3.3 .

Assume that $f(x)$ is differentiable on $(0,\infty)$ and that $h_{f}(x)=xf(x)$ is strictly convex and $f(1)=0$. []

###### Remark 3.4 .

Under the above assumptions, problem [IPF](#S3.Ex1) is a convex optimization problem and assumption [3.3](#S3.Thmtheorem3) makes the regularization term positive due to Jensen’s inequality as $\mathbb{E}_{\mu}[\frac{\pi}{\mu}f(\frac{\pi}{\mu})]\geq 1f(1)=0$ []. Slater’s conditions hold since constraint 1 and constraint 2 define a probability simplex, and constraint 3 defines a hyperplane in the tabular setting. The intersection of these convex sets is nonempty if the optimal policy exists, $i.e.$ the optimal policy is not a uniform distribution. The analysis described above shows that this convex optimization problem is feasible and Slater’s conditions are satisfied.

###### Definition 3.1 .

###### Assumption 3.2 .

###### Assumption 3.3 .

###### Remark 3.4 .

## 4 Optimization

In this section, we introduce two methods AlignIQL-hard and AlignIQL for solving problem [IPF](#S3.Ex1). Theoretically, AlignIQL-hard is more rigorous as it strictly ensures policy alignment and provides insights into why IQL can use AWR for policy extraction, but it suffers from complex training. AlignIQL avoids the training complexity of AlignIQL-hard while also guaranteeing local convergence to the optimal solution of problem [IPF](#S3.Ex1) through soft constraints. All proof can be found in Appendix [A](#A1)

### 4.1 AlignIQL-hard

We first consider directly solving [IPF](#S3.Ex1) with KKT conditions and get the following theorems.

###### Theorem 4.1 .

For problem [IPF](#S3.Ex1), the optimal policy $\pi^{*}$ and its optimal Lagrange multipliers satisfy the following optimality condition for all states and actions:

$$ $\pi^{\star}({\bm{a}}|{\bm{s}})=\mu({\bm{a}}|{\bm{s}})\max\left\{g_{f}\left(-\alpha^{*}({\bm{s}})-\beta^{*}({\bm{s}})Q({\bm{s}},{\bm{a}})\right),0\right\}.$ (8) $$

$$ $\mathbb{E}_{{\bm{a}}\sim\mu}\left[\max\left\{g_{f}(-\alpha^{*}({\bm{s}})-\beta^{*}({\bm{s}})Q({\bm{s}},{\bm{a}})),0\right\}\right]=1,$ (9) $$

$$ $\mathbb{E}_{{\bm{a}}\sim\mu({\bm{a}}|{\bm{s}})}\left[Q({\bm{s}},{\bm{a}})\max\left\{g_{f}(-\alpha^{*}({\bm{s}})-\beta^{*}({\bm{s}})Q({\bm{s}},{\bm{a}})),0\right\}-V({\bm{s}})\right]=0,$ (10) $$

where $\alpha^{*},\beta^{*}$ is the Lagrange multiplier, $g_{f}$ is the inverse function of h${}^{\prime}_{f}(x)$.

###### Remark 4.2 .

Note that $\alpha^{*}$ is a normalization term, it does not affect the action generated by the policy.
Let $f(x)=\log x$, then $g_{f}(x)=\exp{(x-1))}>0$, we can get $\pi^{*}({\bm{a}}|{\bm{s}})\propto\mu({\bm{a}}|{\bm{s}})\exp{(-\beta^{*}Q({\bm{s}},{\bm{a}}))}$ In most environments (especially MuJoCo tasks), $\beta^{*}$ we learned through the neural network is negative. Because only the positive and negative of $\beta^{*}$ affect the action generated by the policy, we can approximate $\beta^{*}$ with a fixed $\beta\in(0,\infty]$, $i.e.$ $\pi^{*}({\bm{a}}|{\bm{s}})\propto\mu({\bm{a}}|{\bm{s}})\exp{(\beta Q({\bm{s}},{\bm{a}}))}$, which is exactly what optimal policy obtained by AWR. This explains why IQL can learn implicit policy with weighted regression and shows implicit policy avoids the OOD actions through the regularization function $f$, which gives a deeper understanding of how IQL-style methods handle the distribution shift. This also addresses the issue in IDQL, that is, they find that simply taking the action with the highest Q-value tends to yield better performance at evaluation time.

Although implicit policy can be extracted by Equation [8](#S4.E8), how to handle weights is a key issue for IQL, especially when we use the diffusion model to approximate the behavior policy. Because we use the importance sampling weights to select actions from learned behavior policy. The weight used in prior works is an increasing function of $Q({\bm{s}},{\bm{a}})$, however according to Theorem [4.1](#S4.Thmtheorem1), when $\beta^{*}({\bm{s}})\geq 0$, we need to be more conservative, that is, we should choose actions with lower $Q({\bm{s}},{\bm{a}})$. To calculate the weights, we need to solve the closed-form solution of Equation [9](#S4.E9), Equation [10](#S4.E10), which is usually intractable. However, we can resort to the parameterized neural network to approximate solve it.

###### Lemma 4.3 .

Following EQL [], let $f(x)=\log x$, then $g_{f}(x)=\exp{(x-1))}>0$. We can approximate $\alpha^{*}({\bm{s}})$, $\beta^{*}({\bm{s}})$ through neural network with the following loss function:

$$ $\max_{\alpha,\beta}{\mathcal{L}}_{M}=-\mathbb{E}_{{\bm{a}}\sim\mu}\left[\exp{\left(-\alpha({\bm{s}})-\beta({\bm{s}})Q({\bm{s}},{\bm{a}})-1\right)}\right]-\alpha({\bm{s}})-\beta({\bm{s}})V({\bm{s}}),$ (11) $$

###### Proof.

Then Lemma [4.3](#S4.Thmtheorem3) can be get through setting the gradient of Equation [11](#S4.E11) to 0 0 with respect to $\alpha,\beta$, which is Equation [9](#S4.E9), Equation [10](#S4.E10) respectively.
∎

###### Remark 4.4 .

Now we can obtain $\alpha^{*},\beta^{*}$ by iteratively updating $\alpha,\beta$ following Equation [11](#S4.E11).

Based on Theorem [A.1](#A1.SS1) and Lemma [4.3](#S4.Thmtheorem3), we can get a practical algorithm AlignIQL-hard, where hard means we rigidly constrain the policy to satisfy policy alignment. AlignIQL-hard shows when multiplier $\beta({\bm{s}})<0$, we can use AWR for extracting the implicit policy in IQL. However, for strict policy alignment, AlignIQL-hard needs to train an additional two multiplier networks, which increases the training costs and compound errors. Moreover, the exponential term in Equation [11](#S4.E11) makes the unstable training. In the remainder of this section, we introduce a more simple and effective method AlignIQL to solve problem [IPF](#S3.Ex1).

Figure: Algorithm 1 AlignIQL Training

###### Theorem 4.1 .

###### Remark 4.2 .

###### Lemma 4.3 .

###### Proof.

###### Remark 4.4 .

### 4.2 AlignIQL

In this section, we introduce AlignIQL to solve the alignment problem of IQL. Firstly, we introduce the soft constraint form of problem [IPF](#S3.Ex1). Given $\eta\in\mathbb{R}\setminus{0}$, IPF-Soft is defined as

$$ $\displaystyle\min_{\pi,V({\bm{s}})}\quad\mathbb{E}_{{\bm{s}}\sim d^{\pi}({\bm{s}}),{\bm{a}}\sim\pi({\bm{a}}|{\bm{s}})}\left[f\left(\frac{\pi({\bm{a}}|{\bm{s}})}{\mu({\bm{a}}|{\bm{s}})}\right)+\eta\left(Q({\bm{s}},{\bm{a}})-V({\bm{s}})\right)^{2}\right]$ (IPF-Soft) $\displaystyle s.t.\quad\quad\pi({\bm{a}}|{\bm{s}})\geq 0,\quad\forall{\bm{s}},\forall{\bm{a}}$ $\displaystyle\quad\quad\int_{\bm{a}}\pi({\bm{a}}|{\bm{s}})d{\bm{a}}=1,\quad\forall{\bm{s}}.$ $$

###### Remark 4.5 .

We refer to the above problem as problem [IPF-Soft](#S4.Ex2), since it does not rigidly keep the policy alignment. But we will show problem [IPF-Soft](#S4.Ex2) approximately solves problem [IPF](#S3.Ex1) through Proposition [4.8](#S4.Thmtheorem8). More importantly, in practice, AlignIQL does not require training multiplier networks, making the training more stable, efficient, and resulting in better performance.

###### Theorem 4.6 .

Suppose that $f(x)=\log x$, then the optimal policy of problem [IPF-Soft](#S4.Ex2) satisfies

$$ $\pi^{\star}({\bm{a}}|{\bm{s}})\propto\exp{\left\{-\eta\left(Q({\bm{s}},{\bm{a}})-V({\bm{s}})\right)^{2}\right\}}.$ (12) $$

###### Remark 4.7 .

Equation [12](#S4.E12) and Equation [8](#S4.E8) are very similar, but the difference is that Equation [12](#S4.E12) extract policy based on $\eta$ rather than multiplier $\beta({\bm{s}})$, which avoids training the multiplier network. For $\eta>0$, Equation [12](#S4.E12) prefers actions that minimize $(Q({\bm{s}},{\bm{a}})-V({\bm{s}}))^{2}$. For $\eta<0$, Equation [12](#S4.E12) prefers actions that maximize the absolute value of the advantage function $|A({\bm{s}},{\bm{a}})|=|Q({\bm{s}},{\bm{a}})-V({\bm{s}})|$. However, the absolute function can not distinguish between positive and negative values. To solve this problem, we suppose that the probabilities of $A({\bm{s}},{\bm{s}})>0$ and $A({\bm{s}},{\bm{s}})<0$
for the sampled actions are equal and choose actions with $A({\bm{s}},{\bm{a}})>0$ since $A({\bm{s}},{\bm{a}})\geq 0$ means better actions, that is, for $\eta<0$, we use the $w({\bm{s}},{\bm{a}})\propto\exp{\left\{|\eta|(Q({\bm{s}},{\bm{a}})-V({\bm{s}}))\right\}}$ as the weight. Note that in this case, $w({\bm{s}},{\bm{a}})$ is actually the same as the weight used by IQL, SfBC [] and it is very similar to the $w_{\text{exp}}$ of IDQL. This demonstrates the connection between IQL-style algorithms and Weighted Regression.

Next, we show the relationship between problem [IPF](#S3.Ex1) and problem [IPF-Soft](#S4.Ex2).

###### Proposition 4.8 .

Suppose that $\pi^{*}({\bm{a}}|{\bm{s}})$ is a global solution to the convex optimization problem [IPF](#S3.Ex1),
with its corresponding value function (denoted as $V^{*}({\bm{s}})$). Then $\pi^{*},V^{*}({\bm{s}})$ is a strict local minimizer of problem [IPF-Soft](#S4.Ex2) for all

$$ $\eta(h^{*}-k^{*})<0,$ $$

where
$h^{*}=\mathbb{E}_{{\bm{a}}\sim\pi^{*}}\left[\left(Q({\bm{s}},{\bm{a}})-V^{*}({\bm{s}})\right)^{2}\right],k^{*}=\inf_{\pi\in{\mathcal{T}}}\mathbb{E}_{{\bm{a}}\sim\pi}\left[\left(Q({\bm{s}},{\bm{a}})-V({\bm{s}})\right)^{2}\right],$

$$ ${\mathcal{T}}=\left\{\pi|\pi({\bm{a}}|{\bm{s}})\geq 0,\int_{\bm{a}}\pi({\bm{a}}|{\bm{s}})d{\bm{a}}=1,V({\bm{s}})=\mathbb{E}_{{\bm{a}}\sim\pi({\bm{a}}|{\bm{s}})}\left[Q({\bm{s}},{\bm{a}})\right],\pi\in\mathring{U}(\pi^{*},\sigma)\right\}.$ $$

###### Remark 4.9 .

Note that problem [IPF-Soft](#S4.Ex2) may not be a convex optimization problem, as the objective function may not be a convex function. Through Proposition [4.8](#S4.Thmtheorem8), solution of problem [IPF](#S3.Ex1) can be found by approximately solving the problem [IPF-Soft](#S4.Ex2).

Note that in both AlignIQL-hard and AlignIQL, we do not impose a limit on the loss function of the $Q-V$, which means that our conclusion can be generalized to the arbitrary critic loss function and the arbitrary value function. To summarize, both the AlignIQL-hard and AlignIQL are IQL-style algorithms, which means the training of actor and critic are decoupled and the critic is learned by expectile regression. The difference between AlignIQL-hard and AlignIQL lies in the calculation of weights and the necessity of training multiplier networks. We summarize the procedure of AlignIQL-hard and AlignIQL in Algorithm [1](#alg1) and Algorithm [2](#alg2).

###### Remark 4.5 .

###### Theorem 4.6 .

###### Remark 4.7 .

###### Proposition 4.8 .

###### Remark 4.9 .

## 5 Experiments

In this section, we conduct extensive experiments and specifically answer the following questions

- •
Q1: Can AlignIQL match the performance with other SOTA offline RL baselines?
- •
Q2: Can AlignIQL derive any benefits from the alignment?
- •
Q3: Both IDQL and AlignIQL derive the weights needed for policy alignment. Are our weights better than IDQL?

### 5.1 Results on D4RL Tasks (Q1)

We conduct extensive experiments on D4RL datasets [] to verify the performance of AlignIQL. For the selection of baselines, we select Conservative Q-learning (CQL) [], DiffusionQL [], Implicit-Q Learning (IQL) [], SQL [], SfBC [] because of their strong performance in the offline RL setting. We also select the DD [] and Diffuser [] as baselines since they represent sequential-modeling-based RL algorithms. Aggregated results can be found in Table [1](#S5.T1). In MuJoCo tasks, where the performance is already saturated, AlignIQL shows slightly worse results compared to the other methods. Since the performance in MuJoCo tasks is already saturated, the effect of policy alignment is not evident. However, in more challenging AntMaze tasks, AlignIQL outperforms other methods by a large margin. Above all, results in Table [1](#S5.T1) show that AlignIQL can match the SOTA performance in MuJoCo tasks and achieve SOTA results in AntMaze tasks. More importantly, we find that compared to IDQL, AlignIQL is less sensitive to hyperparameter $N$. We will elaborate on this in the following sections. We include learning curves and experiment details in Appendix E to ensure the reproducibility of our results in Table [1](#S5.T1).

**Table 1: The performance of our method and other SOTA baselines on D4RL tasks. The performance of AlignIQL and IDQL is obtained by evaluating the trained policy over $10$ different random seeds. We report the performance of baseline methods using the best results reported from their papers except for IDQL. We rerun the official code of IDQL and report the results on the same hardware (RTX 4090 with 24GB memory) for a fair comparison. We report the result of AlignIQL and IDQL by taking the average of the last evaluation over $10$ seeds. Results within $3\%$ of the maximum in every D4RL task and the best average result are highlighted in SkyBlue.**
| Dataset | Environment | CQL | Diffusion-QL | SfBC | SQL | DD | Diffuser | IDQL | IQL | AlignIQL (ours) |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Medium-Expert | HalfCheetah | $62.4$ | $96.8$ | $92.6$ | $94.0$ | $90.6$ | $79.8$ | $89.2$ | $86.7$ | $89.1$ |
| Medium-Expert | Hopper | $98.7$ | $111.1$ | $108.6$ | $111.8$ | $111.8$ | $107.2$ | $108.2$ | $91.5$ | $107.1$ |
| Medium-Expert | Walker2d | $111.0$ | $110.1$ | $109.8$ | $110.0$ | $108.8$ | $108.4$ | $111.7$ | $109.6$ | $111.9$ |
| Medium | HalfCheetah | $44.4$ | $51.1$ | $45.9$ | $48.3$ | $49.1$ | $44.2$ | $46.0$ | $47.4$ | $46.0$ |
| Medium | Hopper | $58.0$ | $90.5$ | $57.1$ | $75.5$ | $79.3$ | $58.5$ | $56.3$ | $66.3$ | $56.1$ |
| Medium | Walker2d | $79.2$ | $87.0$ | $77.9$ | $84.2$ | $82.5$ | $79.7$ | $77.6$ | $78.3$ | $78.5$ |
| Medium-Replay | HalfCheetah | $46.2$ | $47.8$ | $37.1$ | $44.8$ | $39.3$ | $42.2$ | $41.1$ | $44.2$ | $41.1$ |
| Medium-Replay | Hopper | $48.6$ | $101.3$ | $86.2$ | $99.7$ | $100.0$ | $96.8$ | $86.2$ | $94.7$ | $74.8$ |
| Medium-Replay | Walker2d | $26.7$ | $95.5$ | $65.1$ | $81.2$ | $75.0$ | $61.2$ | $85.1$ | $73.9$ | $76.5$ |
| Average (Locomotion) | $63.9$ | $87.9$ | $75.6$ | $83.3$ | $81.8$ | $75.3$ | $78.0$ | $76.9$ | $75.7$ |  |
| Default | AntMaze-umaze | $74.0$ | $93.4$ | $92.0$ | $92.2$ | - | - | $93.4$ | $87.5$ | $94.8$ |
| Diverse | AntMaze-umaze | $84.0$ | $66.2$ | $85.3$ | $74.0$ | - | - | $75.2$ | $62.2$ | $82.4$ |
| Play | AntMaze-medium | $61.2$ | $76.6$ | $81.3$ | $80.2$ | - | - | $85$ | $71.2$ | $80.5$ |
| Diverse | AntMaze-medium | $53.7$ | $78.6$ | $82.0$ | $79.1$ | - | - | $74.4$ | $70.0$ | $85.5$ |
| Play | AntMaze-large | $15.8$ | $46.4$ | $59.3$ | $53.2$ | - | - | $60.0$ | $39.6$ | $65.2$ |
| Diverse | AntMaze-large | $14.9$ | $57.3$ | $45.5$ | $52.3$ | - | - | $58.4$ | $47.5$ | $66.4$ |
| Average (AntMaze) | $50.6$ | $69.8$ | $74.2$ | $-$ | - | - | $74.4$ | $63.0$ | $79.1$ |  |
| # Diffusion steps | - | $5$ | $15$ | $-$ | $100$ | $100$ | $5$ | - | $5$ |  |

### 5.2 The Role of Policy Alignment (Q2)

In this section, we conduct a detailed comparison of the impact of the sign of $\eta$ in AlignIQL. According to Theorem [A.2](#A1.SS2), the sign of $\eta$ determines what policy the learned value function is actually evaluating. If $\eta=1$ represents the aligned policy, then $\eta=-1$ corresponds to the misaligned policy and vice versa. Note that, when $\eta=-1$, the weight is equivalent to the weight of AWR, which selects actions based on the Q-function and is widely used in offline RL [, , ].
We choose the AntMaze tasks since challenging tasks can better show the effect of alignment. In the experiment shown in Figure [1](#S5.F1), we use a diffusion-based behavior policy and guarantee that all the hyperparameters and network structures are the same except for $\eta$, which is used for policy alignment.

Figure [1](#S5.F1) shows that under the same $N$ ( where $N$ represents the samples per state),
The correctly aligned policy ($\eta=1$ in Figure [1](#S5.F1)) converges faster and does not experience a significant performance drop at higher $N=256$. Generally, the greedy selection of actions based on Q-values [, ], which usually yields better results, performed poorly here. This is due to policy misalignment, leading to AWR-based weights assigning high values to out-of-distribution actions potentially generated by the behavior policy. In MuJoCo tasks, due to the relative simplicity of the tasks, the Q-function can achieve good generalization performance through training, so this issue is not pronounced.

Above all, the results in Figure [1](#S5.F1) show the benefits of policy alignment and demonstrate the correctness of our theory through the performance difference of different $\eta$.

### 5.3 Comparison with IDQL (Q3)

Since both IDQL and AlignIQL provide weights under policy alignment, in this part we evaluate which weights are better through D4RL sparse reward tasks Adroit and ablation study on $N$. In the experiments in this section, both IDQL and AlignIQL use expectile regression to learn the critic and employ greedy extraction ${\bm{a}}_{\text{taken}}=\operatorname*{arg\,max}_{\bm{a}}w({\bm{s}},{\bm{a}})$ during evaluation. We use expectile regression to train critic since the expectile objective performs the strongest with greedy extraction as mentioned in IDQL. Of course, our AlignIQL framework can also be extended to other generalized critic loss functions mentioned in IDQL.

Figure: Figure 1: Performance of different $\eta$ in AlignIQL on AntMaze tasks with different training steps. The horizontal axis represents time (s) on a logarithmic scale. Results are averaged over $10$ random seeds.
Refer to caption: /html/2405.18187/assets/x1.png

Results on D4RL Sparse Reward Tasks. We conduct the experiment on Adroit tasks. Compared with MuJoCo tasks, the Adroit tasks are high dimensional and feature sparse rewards. For human and expert datasets, the data is collected in from human demonstrators.
Results can be found in Table [2](#S5.T2). Table [2](#S5.T2) shows that compared to IDQL, AlignIQL can achieve competitive results and significantly outperform IDQL on some challenging tasks like relocate-human.

Ablation Study. Since the main limitation of diffusion-based methods is the running speed, we
compare the performance of IDQL and AlignIQL at different training times in Figure [2](#S5.F2). Figure [2](#S5.F2) shows that as training time increases, the performance of AlignIQL with different $N$ convergences to the same value, which shows that AlignIQL is insensitive to $N$. We also
observe that AlignIQL converges faster compared to IDQL.

Overall, compared to IDQL, the weights computed by our method not only have better theoretical properties (applicable to any Q-loss, without requiring optimal $V$) but also perform better in practice.

**Table 2: Normalized scores of AlignIQL against other baselines on D4RL sparse-reward tasks. We bold the mean values that $\geq 0.99$ $*$ highest value. “-A“ refers to we only sweep over $\eta\in\left\{-1,1\right\}$ while keeping other parameters the same as the IDQL. AlignIQL refers to we fix $\eta=1$.**
| Task | BC | BCQ | CQL | IQL | AlignIQL-A | Algae-<br>DICE | IDQL | AlignIQL |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| pen-human | 63.9 | 68.9 | 37.5 | 71.5 | 76.0±4.8 | -3.3 | 70.7±8.4 | 76.0±4.8 |
| hammer-human | 1.2 | 0.5 | 4.4 | 1.4 | 2.7±1.2 | 0.3 | 2.8±0.8 | 2.0±0.7 |
| door-human | 2.0 | 0.0 | 9.9 | 4.3 | 7.3±4.9 | 0.0 | 5.2±1.3 | 6.0±3.6 |
| relocate-human | 0.1 | -0.1 | 0.2 | 0.1 | 0.3±0.1 | -0.1 | 0.09±0.02 | 0.28±0.34 |
| pen-expert | 85.1 | 114.9 | 107.0 | 111.7 | 139.2±6.5 | -3.5 | 132.9±4.5 | 116.0±4.3 |
| hammer-expert | 125.6 | 107.2 | 86.7 | 116.3 | 124.7±1.9 1 | 0.3 | 126.5±0.7 | 124.7±1.9 |
| door-expert | 34.9 | 99.0 | 101.5 | 103.8 | 105.2±0.2 | 0.0 | 105.0±0.1 | 104.6±0.5 |
| relocate-expert | 101.3 | 41.6 | 95.0 | 102.7 | 107.7±2.7 | -0.1 | 108.3±1.3 | 106.0±1.5 |

Figure: Figure 2: Performance of AlignIQL and IDQL on AntMaze tasks with different training steps. The different sizes, from smallest to largest, represent $N=16,N=64,N=256$, respectively.
Refer to caption: /html/2405.18187/assets/x2.png

## 6 Related Work

Offline RL. Offline RL algorithms need to avoid OOD actions. Previous methods to mitigate this issue under the model-free offline RL setting generally fall into three categories: 1) value function-based approaches, which implement pessimistic value estimation by assigning low values to out-of-distribution actions [, ], or implicit TD backups [, ] to avoid the use of out-of-distribution actions 2) sequential modeling approaches, which casts offline RL as a sequence generation task with return guidance [, , , ], and 3) constrained policy search (CPS) approaches, which regularizes the discrepancy between the learned policy and behavior policy [, , ].

Implicit Q-learning. Recently, implicit Q-learning [] has attracted interest due to its stable training and simplicity. Many offline RL methods [, , ] use IQL-style expectile regression to learn Q-function and realize the advantage of decoupling the training of actor and critic. While IQL achieves superior performance, several issues remain unsolved. SQL [] reinterprets IQL in the Implicit Value Regularization (IVR) framework and provides insights about why in practice a large $\tau$ may give a worse result in IQL. However, there is another important open question about IQL, that is, what policy the learned value function is evaluating. IDQL [] solves this by reinterpreting the IQL as an actor-critic method and getting the corresponding implicit policy for the (generalized) IQL loss function. However, the corresponding implicit policy in IDQL only holds for optimal value function under certain critic loss functions.

Diffusion Model in Offline RL. Due to our method using the diffusion model for modeling behavior policy, we review works that incorporate the Diffusion model in offline RL. There exist several works that introduce the diffusion model to RL. Diffuser [] uses the diffusion model to directly generate trajectory guided with gradient guidance or reward. DiffusionQL [] uses the diffusion model as an actor and optimizes it through the TD3+BC-style objective with a coefficient $\eta$ to balance the two terms. AdaptDiffuser [] uses a diffusion model to generate extra trajectories and a discriminator to select desired data to add to the training set to enhance the adaptability of the diffusion model. DD [] uses a conditional diffusion model to generate trajectory and compose skills. Unlike Diffuser, DD diffuses only states and trains inverse dynamics to predict actions. QGPO [] uses the energy function to guide the sampling process and proves that the proposed CEP training method can get an unbiased estimation of the gradient of the energy function under unlimited model capacity and data samples. IDQL [] reinterpret IQL as an Actor-Critic method and extract the policy through sampling from a diffusion-parameterized behavior policy with weights computed from the IQL-style critic. EDP [] focuses on boosting sampling speed through approximated actions. SRPO [] uses a Gaussian policy in which the gradient is regularized by a pretrained diffusion model to recover the IQL-style policy. Our method is distinct from these methods because we aim to
align the implied policy in the value function.

The closest work to ours is IDQL [], which derives the implicit policy for optimal value function under different critic loss functions. Our method is related, but features with AlignIQL apply to arbitrary value functions and arbitrary critic loss functions. More importantly, our method explains when and why IQL can use AWR for policy extraction while providing theoretical insights for IQL and other RL paradigms that use Q-values to guide sampling.

## 7 Discussion and Future Work

In our work, we define the implicit policy-finding problem in IQL and propose two practical algorithms AlignIQL-hard and AlignIQL to solve it. The optimal policy (Theorem [A.1](#A1.SS1)) in AlignIQL-hard shows that it is feasible to extract policy with AWR in certain cases, which builds the bridge between the Implicit Q-learning and Weighted Regression. Our theoretical findings also extend the policy alignment of IDQL to arbitrary critic loss and value functions. Besides the theoretical findings, we also verify the effectiveness of our algorithm on D4RL datasets. Experimental results show that compared to other IQL-style algorithms, our algorithm achieves SOTA performance and is more stable, especially in sparse reward tasks. One future work is to explore better methods for training multiplier networks and explore the impact of different regularization functions of problem [IPF](#S3.Ex1). Another future work is to extend our approach to fields of safe RL and offline-to-online (O2O) learning. In safe RL, prior works [, ] have used IQL to learn the Q-function. Investigating how to ensure policy alignment while satisfying safety constraints is an interesting research direction.

Limitation. Our work focus on regularization function $f(x)=\log x$ due to its simplicity. However, different regularization functions can result in different type policy. For example, SQL [] shows using $f(x)=x-1$ can introduce sparsity in learning the state-value function. In future work, we can explore other regularization functions to obtain better policy.

Boarder Impact. Our method will promote the development of offline reinforcement learning, thus facilitating the implementation of offline reinforcement learning in practical scenarios, such as robotic control. Since our work primarily focuses on RL theory, it will not raise ethical issues.

## Appendix A Proof

### A.1 Proof of Theorem 3.4

###### Proof.

The Lagrange function of Equation [IPF](#S3.Ex1) is written as follows

$$ $\displaystyle L(\pi,\alpha({\bm{s}}),\beta({\bm{s}}),\lambda)$ $\displaystyle=\mathbb{E}_{{\bm{s}}\sim d^{\pi}({\bm{s}}),{\bm{a}}\sim\pi({\bm{a}}|{\bm{s}})}\left[f(\frac{\pi({\bm{a}}|{\bm{s}})}{\mu({\bm{a}}|{\bm{s}})})\right]-\mathbb{E}_{{\bm{s}}\sim d^{\pi}({\bm{s}}),{\bm{a}}\sim\pi({\bm{a}}|{\bm{s}})}\left[\lambda({\bm{a}}|{\bm{s}})\pi({\bm{a}}|{\bm{s}})\right]$ (13) $\displaystyle+\mathbb{E}_{{\bm{s}}\sim d^{\pi}({\bm{s}})}\left[\alpha({\bm{s}})\left(\int_{\bm{a}}\pi({\bm{a}}|{\bm{s}})d{\bm{a}}-1\right)\right]+$ $\displaystyle\mathbb{E}_{{\bm{s}}\sim d^{\pi}({\bm{s}})}\left[\beta({\bm{s}})\left(\mathbb{E}_{{\bm{a}}\sim\pi({\bm{a}}|{\bm{s}})}\left[Q({\bm{s}},{\bm{a}})\right]-V({\bm{s}})\right)\right],$ $$

where $d^{\pi}({\bm{s}})$ represents the state distribution induced by policy $\pi$, $\alpha({\bm{s}})$, $\beta({\bm{s}})$, and $\lambda$ are Lagrangian multipliers for the equality and inequality constraints respectively.

Let $h_{f}(x)=xf(x)$. Then for all states and actions, the KKT conditions can be written as follows

$$ $\displaystyle\pi({\bm{a}}|{\bm{s}})\geq 0$ (14) $\displaystyle\int_{\bm{a}}\pi({\bm{a}}|{\bm{s}})d{\bm{a}}=1$ (15) $\displaystyle\mathbb{E}_{{\bm{a}}\sim\pi({\bm{a}}|{\bm{s}})}\left[Q({\bm{s}},{\bm{a}})-V({\bm{s}})\right]=0$ (16) $\displaystyle\lambda({\bm{a}}|{\bm{s}})\geq 0$ (17) $\displaystyle\lambda({\bm{a}}|{\bm{s}})\pi({\bm{a}}|{\bm{s}})=0$ (18) $\displaystyle h_{f}^{\prime}(\frac{\pi({\bm{a}}|{\bm{s}})}{\mu({\bm{a}}|{\bm{s}})})+\alpha({\bm{s}})+\beta({\bm{s}})Q({\bm{s}},{\bm{a}})-\lambda({\bm{a}}|{\bm{s}})=0$ (19) $$

We eliminate $d^{\pi}({\bm{s}})$ due to irreducible Markov chain assumption. Note that in our derivation, we assume that $V({\bm{s}})$ and $Q({\bm{s}},{\bm{a}})$ are known.

Since $h_{f}^{\prime}$ is a strictly increasing function, its inverse function exists and is also a strictly increasing function. Let $g_{f}=(h_{f}^{\prime})^{-1}(x)$ be its inverse function.
From Equation [19](#A1.E19), we can get

$$ $\pi({\bm{a}}|{\bm{s}})=\mu({\bm{a}}|{\bm{s}})g_{f}\left(\lambda({\bm{a}}|{\bm{s}})-\alpha({\bm{s}})-\beta({\bm{s}})Q({\bm{s}},{\bm{a}})\right)$ (20) $$

Given a state ${\bm{s}}$, we can get $\lambda({\bm{a}}|{\bm{s}})=h_{f}^{\prime}(\frac{\pi}{\mu})+\alpha({\bm{s}})+\beta({\bm{s}})Q({\bm{s}},{\bm{a}})$ from Equation [19](#A1.E19), then

- (a)
If $\lambda({\bm{a}}|{\bm{s}})=h_{f}^{\prime}(\frac{\pi}{\mu})+\alpha({\bm{s}})+\beta({\bm{s}})Q({\bm{s}},{\bm{a}})>0$, then $\pi({\bm{a}}|{\bm{s}})$ is zero due to complementary slackness. Note that $\pi({\bm{a}}|{\bm{s}})=0$, thus $h_{f}^{\prime}(0)+\alpha({\bm{s}})+\beta({\bm{s}})Q({\bm{s}},{\bm{a}})>0$ and we can get $g_{f}(-\alpha({\bm{s}})-\beta({\bm{s}})Q({\bm{s}},{\bm{a}}))<g_{f}(h_{f}^{\prime}(0))=0$.
- (b)
If $\lambda({\bm{a}}|{\bm{s}})=0$, then $h_{f}^{\prime}(\frac{\pi}{\mu})+\alpha({\bm{s}})+\beta({\bm{s}})Q({\bm{s}},{\bm{a}})$ is zero and $\pi({\bm{a}}|{\bm{s}})=\mu({\bm{a}}|{\bm{s}})g_{f}\left(-\alpha({\bm{s}})-\beta({\bm{s}})Q({\bm{s}},{\bm{a}})\right)\geq 0$. Note that $\pi({\bm{a}}|{\bm{s}})\geq 0$, thus $h_{f}^{\prime}(0)+\alpha({\bm{s}})+\beta({\bm{s}})Q({\bm{s}},{\bm{a}})\leq 0$ and we can get $g_{f}(-\alpha({\bm{s}})-\beta({\bm{s}})Q({\bm{s}},{\bm{a}}))\geq g_{f}(h_{f}^{\prime}(0))=0$.

Through analysis (a) and (b), we can resolve optimal policy $\pi^{*}({\bm{a}}|{\bm{s}})$ as

$$ $\pi^{\star}({\bm{a}}|{\bm{s}})=\mu({\bm{a}}|{\bm{s}})\max\left\{g_{f}\left(-\alpha({\bm{s}})-\beta({\bm{s}})Q({\bm{s}},{\bm{a}})\right),0\right\}.$ (21) $$

Substituting back in Equation [15](#A1.E15) and Equation [16](#A1.E16) with Equation [8](#S4.E8), we can get

$$ $\mathbb{E}_{{\bm{a}}\sim\mu}\left[\max\left\{g_{f}(-\alpha^{*}({\bm{s}})-\beta^{*}({\bm{s}})Q({\bm{s}},{\bm{a}})),0\right\}\right]=1,$ (22) $$

$$ $\mathbb{E}_{{\bm{a}}\sim\mu({\bm{a}}|{\bm{s}})}\left[Q({\bm{s}},{\bm{a}})\max\left\{g_{f}(-\alpha^{*}({\bm{s}})-\beta^{*}({\bm{s}})Q({\bm{s}},{\bm{a}})),0\right\}-V({\bm{s}})\right]=0,$ (23) $$

∎

###### Proof.

### A.2 Proof of Theorem 3.8

###### Proof.

The Lagrange function of Equation [IPF-Soft](#S4.Ex2) is written as follows

$$ $\displaystyle L(\pi,\alpha({\bm{s}}),\beta({\bm{s}}),\lambda)$ $\displaystyle=\mathbb{E}_{{\bm{s}}\sim d^{\pi}({\bm{s}}),{\bm{a}}\sim\pi({\bm{a}}|{\bm{s}})}\left[f(\frac{\pi({\bm{a}}|{\bm{s}})}{\mu({\bm{a}}|{\bm{s}})})+\eta\left(Q({\bm{s}},{\bm{a}})-V({\bm{s}})\right)^{2}\right]$ (24) $\displaystyle-\mathbb{E}_{{\bm{s}}\sim d^{\pi}({\bm{s}}),{\bm{a}}\sim\pi({\bm{a}}|{\bm{s}})}\left[\lambda({\bm{a}}|{\bm{s}})\pi({\bm{a}}|{\bm{s}})\right]$ $\displaystyle+\mathbb{E}_{{\bm{s}}\sim d^{\pi}({\bm{s}})}\left[\alpha({\bm{s}})\left(\int_{\bm{a}}\pi({\bm{a}}|{\bm{s}})d{\bm{a}}-1\right)\right].$ $$

Let $h_{f}(x)=xf(x)$. Then for all states and actions, the KKT conditions can be written as follows

$$ $\displaystyle\pi({\bm{a}}|{\bm{s}})\geq 0$ (25) $\displaystyle\int_{\bm{a}}\pi({\bm{a}}|{\bm{s}})d{\bm{a}}=1$ (26) $\displaystyle\mathbb{E}_{{\bm{a}}\sim\pi({\bm{a}}|{\bm{s}})}\left[Q({\bm{s}},{\bm{a}})-V({\bm{s}})\right]=0$ (27) $\displaystyle\lambda({\bm{a}}|{\bm{s}})\geq 0$ (28) $\displaystyle\lambda({\bm{a}}|{\bm{s}})\pi({\bm{a}}|{\bm{s}})=0$ (29) $\displaystyle h_{f}^{\prime}(\frac{\pi({\bm{a}}|{\bm{s}})}{\mu({\bm{a}}|{\bm{s}})})+\alpha({\bm{s}})+\eta\left(Q({\bm{s}},{\bm{a}})-V({\bm{s}})\right)^{2}-\lambda({\bm{a}}|{\bm{s}})=0$ (30) $$

Since $h_{f}^{\prime}$ is a strictly increasing function, its inverse function exists and is also a strictly increasing function. Let $g_{f}=(h_{f}^{\prime})^{-1}(x)$ be its inverse function.
From Equation [30](#A1.E30), we can get

$$ $\pi({\bm{a}}|{\bm{s}})=\mu({\bm{a}}|{\bm{s}})g_{f}\left(\lambda({\bm{a}}|{\bm{s}})-\alpha({\bm{s}})-\eta\left(Q({\bm{s}},{\bm{a}})-V({\bm{s}})\right)^{2}\right)$ (31) $$

Given a state ${\bm{s}}$, we can get $\lambda({\bm{a}}|{\bm{s}})=h_{f}^{\prime}(\frac{\pi}{\mu})+\alpha({\bm{s}})+\eta\left(Q({\bm{s}},{\bm{a}})-V({\bm{s}})\right)^{2}$ from Equation [30](#A1.E30), then

- (a)
If $\lambda({\bm{a}}|{\bm{s}})=h_{f}^{\prime}(\frac{\pi}{\mu})+\alpha({\bm{s}})+\eta\left(Q({\bm{s}},{\bm{a}})-V({\bm{s}})\right)^{2}>0$, then $\pi({\bm{a}}|{\bm{s}})$ is zero due to complementary slackness. Note that $\pi({\bm{a}}|{\bm{s}})=0$, thus $h_{f}^{\prime}(0)+\alpha({\bm{s}})+\eta\left(Q({\bm{s}},{\bm{a}})-V({\bm{s}})\right)^{2}>0$ and we can get $g_{f}(-\alpha({\bm{s}})-\eta\left(Q({\bm{s}},{\bm{a}})-V({\bm{s}})\right)^{2})<g_{f}(h_{f}^{\prime}(0))=0$.
- (b)
If $\lambda({\bm{a}}|{\bm{s}})=0$, then $h_{f}^{\prime}(\frac{\pi}{\mu})+\alpha({\bm{s}})+\eta\left(Q({\bm{s}},{\bm{a}})-V({\bm{s}})\right)^{2}$ is zero and $\pi({\bm{a}}|{\bm{s}})=\mu({\bm{a}}|{\bm{s}})g_{f}\left(-\alpha({\bm{s}})-\eta\left(Q({\bm{s}},{\bm{a}})-V({\bm{s}})\right)^{2}\right)\geq 0$. Note that $\pi({\bm{a}}|{\bm{s}})\geq 0$, thus $h_{f}^{\prime}(0)+\alpha({\bm{s}})+\eta\left(Q({\bm{s}},{\bm{a}})-V({\bm{s}})\right)^{2}\leq 0$ and we can get $g_{f}(-\alpha({\bm{s}})-\eta\left(Q({\bm{s}},{\bm{a}})-V({\bm{s}})\right)^{2})\geq g_{f}(h_{f}^{\prime}(0))=0$.

Through analysis (a) and (b), we can resolve optimal policy $\pi^{*}({\bm{a}}|{\bm{s}})$ as

$$ $\pi^{\star}({\bm{a}}|{\bm{s}})=\mu({\bm{a}}|{\bm{s}})\max\left\{g_{f}\left(-\alpha({\bm{s}})-\eta\left(Q({\bm{s}},{\bm{a}})-V({\bm{s}})\right)^{2}\right),0\right\}.$ (32) $$

let $f(x)=\log x$, then $g_{f}(x)=\exp{(x-1))}>0$. Substituting back in Equation [32](#A1.E32) with $g_{f}(x)=\exp{(x-1))}$, we can get Equation [12](#S4.E12).
∎

###### Proof.

### A.3 Proof of Proposition 3.10

###### Proof.

Let ${\mathcal{U}}=\left\{\pi({\bm{a}}|{\bm{s}})|\pi({\bm{a}}|{\bm{s}})\geq 0,\int_{\bm{a}}\pi({\bm{a}}|{\bm{s}})d{\bm{a}}=1\right\}$.
Since $\inf_{x,y}g(x,y)=\inf_{x}\inf_{y}g(x,y)$ and the constraints about $\pi$ and $V$ in Problem [IPF-Soft](#S4.Ex2) are independent, we can reformulate Problem [IPF-Soft](#S4.Ex2) as

$$ $\displaystyle\min_{\begin{subarray}{c}\pi,V\\ s.t.\pi\in{\mathcal{U}}\end{subarray}}\mathbb{E}_{{\bm{s}}\sim d^{\pi}({\bm{s}}),{\bm{a}}\sim\pi({\bm{a}}|{\bm{s}})}\left[f\left(\frac{\pi({\bm{a}}|{\bm{s}})}{\mu({\bm{a}}|{\bm{s}})}\right)+\eta\left(Q({\bm{s}},{\bm{a}})-V({\bm{s}})\right)^{2}\right]$ (33) $\displaystyle=\min_{\begin{subarray}{c}\pi\\ s.t.\pi\in{\mathcal{U}}\end{subarray}}\min_{V}\mathbb{E}_{{\bm{s}}\sim d^{\pi}({\bm{s}}),{\bm{a}}\sim\pi({\bm{a}}|{\bm{s}})}\left[f\left(\frac{\pi({\bm{a}}|{\bm{s}})}{\mu({\bm{a}}|{\bm{s}})}\right)+\eta\left(Q({\bm{s}},{\bm{a}})-V({\bm{s}})\right)^{2}\right].$ $$

For $V$, this is an unconstrained problem. Setting the gradient with respect to $V$ to 0 0, we obtain that

$$ $V({\bm{s}})=\mathbb{E}_{{\bm{a}}\sim\pi({\bm{a}}|{\bm{s}})}\left[Q({\bm{s}},{\bm{a}})\right].$ (34) $$

Substituting back in Equation [33](#A1.E33), we can get

$$ $\displaystyle\min_{\begin{subarray}{c}\pi\\ s.t.\pi\in{\mathcal{U}}\end{subarray}}\min_{V}\mathbb{E}_{{\bm{s}}\sim d^{\pi}({\bm{s}}),{\bm{a}}\sim\pi({\bm{a}}|{\bm{s}})}\left[f\left(\frac{\pi({\bm{a}}|{\bm{s}})}{\mu({\bm{a}}|{\bm{s}})}\right)+\eta\left(Q({\bm{s}},{\bm{a}})-V({\bm{s}})\right)^{2}\right]$ (35) $\displaystyle=\min_{\begin{subarray}{c}\pi\\ s.t.\pi\in{\mathcal{V}}\end{subarray}}\mathbb{E}_{{\bm{s}}\sim d^{\pi}({\bm{s}}),{\bm{a}}\sim\pi({\bm{a}}|{\bm{s}})}\left[f\left(\frac{\pi({\bm{a}}|{\bm{s}})}{\mu({\bm{a}}|{\bm{s}})}\right)+\eta\left(Q({\bm{s}},{\bm{a}})-V({\bm{s}})\right)^{2}\right],$ $$

where ${\mathcal{V}}=\left\{\pi({\bm{a}}|{\bm{s}})|\pi({\bm{a}}|{\bm{s}})\in{\mathcal{U}},V({\bm{s}})=\mathbb{E}_{{\bm{a}}\sim\pi({\bm{a}}|{\bm{s}})}\left[Q({\bm{s}},{\bm{a}})\right]\right\}$. Note that ${\mathcal{V}}$ is the feasible set of Problem [IPF](#S3.Ex1).

Since $\inf\left\{f(x)+g(x)\right\}\geq\inf f(x)+\inf g(x)$, we can get

$$ $\displaystyle\min_{\begin{subarray}{c}\pi\\ s.t.\pi\in{\mathcal{V}}\end{subarray}}\mathbb{E}_{\begin{subarray}{c}{\bm{s}}\sim d^{\pi}({\bm{s}})\\ {\bm{a}}\sim\pi({\bm{a}}|{\bm{s}})\end{subarray}}\left[f\left(\frac{\pi({\bm{a}}|{\bm{s}})}{\mu({\bm{a}}|{\bm{s}})}\right)+\eta\left(Q({\bm{s}},{\bm{a}})-V({\bm{s}})\right)^{2}\right]$ (36) $\displaystyle\geq\min_{\begin{subarray}{c}\pi\\ s.t.\pi\in{\mathcal{V}}\end{subarray}}\mathbb{E}_{\begin{subarray}{c}{\bm{s}}\sim d^{\pi}({\bm{s}})\\ {\bm{a}}\sim\pi({\bm{a}}|{\bm{s}})\end{subarray}}\left[f\left(\frac{\pi({\bm{a}}|{\bm{s}})}{\mu({\bm{a}}|{\bm{s}})}\right)\right]+\min_{\begin{subarray}{c}\pi\\ s.t.\pi\in{\mathcal{V}}\end{subarray}}\mathbb{E}_{\begin{subarray}{c}{\bm{s}}\sim d^{\pi}({\bm{s}})\\ {\bm{a}}\sim\pi({\bm{a}}|{\bm{s}})\end{subarray}}\left[\eta\left(Q({\bm{s}},{\bm{a}})-V({\bm{s}})\right)^{2}\right].$ $$

The LHS is exactly Problem [IPF](#S3.Ex1), which has the global optimal value at $\pi^{*},V^{*}$, so that

$$ $\displaystyle\min_{\begin{subarray}{c}\pi\\ s.t.\pi\in{\mathcal{V}}\end{subarray}}\mathbb{E}_{\begin{subarray}{c}{\bm{s}}\sim d^{\pi}({\bm{s}})\\ {\bm{a}}\sim\pi({\bm{a}}|{\bm{s}})\end{subarray}}\left[f\left(\frac{\pi({\bm{a}}|{\bm{s}})}{\mu({\bm{a}}|{\bm{s}})}\right)\right]+\min_{\begin{subarray}{c}\pi\\ s.t.\pi\in{\mathcal{V}}\end{subarray}}\mathbb{E}_{\begin{subarray}{c}{\bm{s}}\sim d^{\pi}({\bm{s}})\\ {\bm{a}}\sim\pi({\bm{a}}|{\bm{s}})\end{subarray}}\left[\eta\left(Q({\bm{s}},{\bm{a}})-V({\bm{s}})\right)^{2}\right]$ (37) $\displaystyle\geq p^{*}+\min_{\begin{subarray}{c}\pi\\ s.t.\pi\in{\mathcal{V}}\end{subarray}}\mathbb{E}_{\begin{subarray}{c}{\bm{s}}\sim d^{\pi}({\bm{s}})\\ {\bm{a}}\sim\pi({\bm{a}}|{\bm{s}})\end{subarray}}\left[\eta\left(Q({\bm{s}},{\bm{a}})-V({\bm{s}})\right)^{2}\right].$ $$

Let ${\mathcal{T}}=\left\{\pi|\pi\in{\mathcal{V}},\pi\in\mathring{U}(\pi^{*},\sigma)\right\}$, we can get ${\mathcal{T}}\notin\emptyset$, since ${\mathcal{V}}$ is a convex set and $\pi^{*}\in{\mathcal{V}}$.

Let $k^{*}=\inf\left\{\mathbb{E}_{\begin{subarray}{c}{\bm{s}}\sim d^{\pi}({\bm{s}})\\
{\bm{a}}\sim\pi({\bm{a}}|{\bm{s}})\end{subarray}}\left[\left(Q({\bm{s}},{\bm{a}})-V({\bm{s}})\right)^{2}\right]|\pi\in{\mathcal{T}}\right\}$, we can get $p^{*}+\eta k^{*}$ is the minimum value for $\pi\in{\mathcal{T}},\forall V$.

Therefore, if the value of Problem [IPF](#S3.Ex1) at $\pi^{*},V^{*}$ is less than $p^{*}+k^{*}$, then $\pi^{*},V^{*}$ is a strict local minimizer of Problem [IPF](#S3.Ex1). Let $h^{*}=\mathbb{E}_{{\bm{a}}\sim\pi^{*}}\left[\left(Q({\bm{s}},{\bm{a}})-V({\bm{s}})\right)^{2}\right]$, we can get

$$ $p^{*}+\eta h^{*}<p^{*}+\eta k^{*}.$ (38) $$

So if $\eta(h^{*}-k^{*})<0$, we can get $\pi^{*},V^{*}$ is a strict local minimizer of Problem [IPF](#S3.Ex1).
∎

###### Proof.

## Appendix B Diffusion model

Diffusion Probabilistic Model (DPM).
Diffusion models [, , ] are composed of two processes: the forward diffusion process and the reverse process. In the forward diffusion process, we gradually add Gaussian noise to the data ${\bm{x}}_{0}\sim q({\bm{x}}_{0})$ in $T$ steps. The step sizes are controlled by a variance schedule $\beta_{i}$:

$$ $\displaystyle q({\bm{x}}_{1:T}\,|\,{\bm{x}}_{0}):=\textstyle\prod_{i=1}^{T}q({\bm{x}}_{i}\,|\,{\bm{x}}_{i-1}),$ (39) $\displaystyle q({\bm{x}}_{i}\,|\,{\bm{x}}_{i-1}):={\mathcal{N}}({\bm{x}}_{i};\sqrt{1-\beta_{i}}{\bm{x}}_{i-1},\beta_{i}{\bm{I}}).$ $$

In the reverse process, we can recreate the true sample ${\bm{x}}_{0}$ through $p({\bm{x}}^{i-1}|{\bm{x}}^{i})$:

$$ $\displaystyle p({\bm{x}})$ $\displaystyle=\int p({\bm{x}}^{0:T})d{\bm{x}}^{1:T}$ (40) $\displaystyle=\int{\mathcal{N}}({\bm{x}}^{T};\mathbf{0},{\bm{I}})\prod_{i=1}^{T}p({\bm{x}}^{i-1}|{\bm{x}}^{i})d{\bm{x}}^{1:T}.$ $$

The training objective is to maximize the ELBO of $\mathbb{E}_{{\bm{q}}_{x_{0}}}\left[\log p({\bm{x}}_{0})\right]$. Following DDPM [], we use the simplified surrogate loss

$$ ${\mathcal{L}}_{d}(\theta)=\mathbb{E}_{i\sim\left[1,T\right],{\mathbf{\epsilon}}\sim{\mathcal{N}}(\mathbf{0},{\bm{I}}),{\bm{x}}_{0}\sim q}\left[||{\mathbf{\epsilon}}-{\mathbf{\epsilon}}_{\theta}({\bm{x}}_{i},i)||^{2}\right]$ (41) $$

to approximate the ELBO. After training, sampling from the diffusion model is equivalent to running the reverse process.

Conditional DPM.
There are two kinds of conditioning methods: classifier-guided [] and classifier-free []. The former requires training a classifier on noisy data ${\bm{x}}_{i}$ and using gradients $\nabla_{\bm{x}}\log f_{\phi}({\bm{y}}|{\bm{x}}_{i})$ to guide the diffusion sample toward the conditioning information ${\bm{y}}$. The latter does not train an independent $f_{\phi}$ but combines a conditional noise model $\epsilon_{\theta}({\bm{x}}_{i},i,{\bm{s}})$ and an unconditional model $\epsilon_{\theta}({\bm{x}}_{i},i)$ for the noise. The perturbed noise $w\epsilon_{\theta}({\bm{x}}_{i},i)+(w+1)\epsilon_{\theta}({\bm{x}}_{i},i,{\bm{s}})$ is used to later generate samples. However [] shows this combination will degrade the policy performance in offline RL. Following [, ] we solely employ a conditional noise model $\epsilon_{\theta}({\bm{x}}_{i},i,{\bm{s}})$ to construct our noise model ($w=0$).

## Appendix C Advantage Weighted Regression (AWR)

Prior works [, ] formulate offline RL as a constrained policy search (CPS) problem with the following form:

$$ $\displaystyle\mu^{*}=\operatorname*{arg\,max}_{\mu}\$ $\displaystyle J(\mu)=\operatorname*{arg\,max}_{\mu}\int_{\mathcal{S}}d_{0}({\bm{s}})\int_{\mathcal{A}}Q^{\mu}({\bm{s}},{\bm{a}})d{\bm{a}}d{\bm{s}}$ $\displaystyle s.t.\quad$ $\displaystyle D_{\mathrm{KL}}(\pi_{b}(\cdot|{\bm{s}})\|\mu(\cdot|{\bm{s}}))\leq{\mathbf{\epsilon}},\quad\forall{\bm{s}}$ (42) $\displaystyle\int_{\bm{a}}\mu({\bm{a}}|{\bm{s}})d{\bm{a}}=1,\quad\forall{\bm{s}},$ $$

Previous works [, , ] solve Equation [C](#A3.Ex5) through KKT conditions and get the optimal policy $\pi^{*}$ as:

$$ $\displaystyle\pi^{*}({\bm{a}}|{\bm{s}})$ $\displaystyle=\frac{1}{Z({\bm{s}})}\ \pi_{b}({\bm{a}}|{\bm{s}})\ \mathrm{exp}\left(\alpha Q_{\phi}({\bm{s}},{\bm{a}})\right),$ (43) $$

where $Z({\bm{s}})$ is the partition function. Intuitively we can use Equation [43](#A3.E43) to optimize policy $\pi$. However, the behavior policy may be very diverse and hard to model. To avoid modeling the behavior policy, prior works [, , ] optimize $\pi^{*}$ through a parameterized policy $\pi_{\theta}$, known as AWR:

$$ $\displaystyle\mathop{\mathrm{arg\ min}}_{\theta}\mathbb{E}_{{\bm{s}}\sim\mathcal{D}^{\mu}}\left[D_{\mathrm{KL}}\left(\pi^{*}(\cdot|{\bm{s}})\middle|\middle|\pi_{\theta}(\cdot|{\bm{s}})\right)\right]$ (44) $\displaystyle=$ $\displaystyle\mathop{\mathrm{arg\ max}}_{\theta}\mathbb{E}_{({\bm{s}},{\bm{a}})\sim\mathcal{D}^{\mu}}\left[\frac{1}{Z({\bm{s}})}\mathrm{log}\ \pi_{\theta}({\bm{a}}|{\bm{s}})\ \mathrm{exp}\left(\alpha Q_{\phi}({\bm{s}},{\bm{a}})\right)\right].$ $$

where $\mathrm{exp}(\alpha Q_{\phi}({\bm{s}},{\bm{a}}))$ being the regression weights. However, AWR requires the exact probability density of policy, which restricts the use of generative models like diffusion models. In this paper, we directly utilize the diffusion-based policy to address Eq. ([C](#A3.Ex5)). Therefore, our method not only avoids the need for explicit probability densities in AWR but also solves the limited policy expressivity problem.

**Table 3: Hyperparameters for reproducing our results on Table [1](#S5.T1).**
| Dataset | Environment | Learning Rate | $\bf{\eta}$ | Reward Tune | Dropout [] | Expectiles $\tau$ | Samples $N$ |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Medium-Expert | HalfCheetah | 3e-4 | $-1$ | Normalization | $0.1$ | $0.7$ | $16$ |
| Medium-Expert | Hopper | 3e-4 | $-1$ | Normalization | $0.1$ | $0.7$ | $16$ |
| Medium-Expert | Walker2d | 3e-4 | $-1$ | Normalization | $0.1$ | $0.7$ | $256$ |
| Medium | HalfCheetah | 3e-4 | $-1$ | Normalization | $0.1$ | $0.7$ | $256$ |
| Medium | Hopper | 3e-4 | $-1$ | Normalization | $0.1$ | $0.7$ | $64$ |
| Medium | Walker2d | 3e-4 | $-1$ | Normalization | $0.1$ | $0.7$ | $256$ |
| Medium-Replay | HalfCheetah | 3e-4 | $-1$ | Normalization | $0.1$ | $0.7$ | $64$ |
| Medium-Replay | Hopper | 3e-4 | $-1$ | Normalization | $0.1$ | $0.7$ | $256$ |
| Medium-Replay | Walker2d | 3e-4 | $-1$ | Normalization | $0.1$ | $0.7$ | $256$ |
| Default | AntMaze-umaze | 3e-4 | $1$ | IQL | $0.1$ | $0.9$ | $16$ |
| Diverse | AntMaze-umaze | 3e-4 | $1$ | IQL | $0.1$ | $0.9$ | $16$ |
| Play | AntMaze-medium | 3e-4 | $-1$ | IQL | $0.1$ | $0.9$ | $64$ |
| Diverse | AntMaze-medium | 3e-4 | $-1$ | IQL | $0.1$ | $0.9$ | $16$ |
| Play | AntMaze-large | 3e-4 | $1$ | IQL | $0.1$ | $0.9$ | $64$ |
| Diverse | AntMaze-large | 3e-4 | $1$ | IQL | $0.1$ | $0.9$ | $64$ |

## Appendix D Experimental Details

Our implementation is based on IDQL [] and jaxrl repo which uses the JAX framework to implement RL algorithms. All networks are optimized through the Adam []. We clip the multiplier network gradient to prevent gradient explosion due to the exponential term.

### D.1 Main Results on D4RL

For the D4RL experiments in Table [1](#S5.T1), we only sweep $\eta$ while keeping other hyperparameters consistent with IDQL to explore the role of alignment. We sweep $\eta$ since the alignment of policy depends on the sign of $\eta$ in different environments. We provide the main hyperparameters in Table [3](#A3.T3) to reproduce our results in Table [1](#S5.T1). Following IDQL, we use normalization to adjust the rewards, which means $r=r/(r_{\text{max}}-r_{\text{min}})$. For AntMaze tasks, $r=r-1$.

Networks We follow the [] to build our policy with an LN_Resnet architecture based conditional diffusion model. To fairly compare our algorithm with IDQL, we use a three-layer network with 256 hidden dim size $256$, consistent with IDQL. The critic and value network consistent with IDQL, is a $2$ layer MLP with hidden size $256$ and ReLU activation function.

We train for $300000$ epochs for D4RL tasks with batch size $512$, consistent with IDQL. The training curve can be found in Figure [3](#A5.F3). We can see due to the policy alignment the training is more stable, even in hard sparse reward tasks like AntMaze large.

| LR (For all networks except for multiplier) | 3e-4 |
| --- | --- |
| Critic Batch Size | 512 |
| Actor Batch Size | 512 |
| $\tau$ Expectiles | $0.7$ (locomotion), $0.9$ (AntMaze) |
| LR | 3e-5 for multiplier networks |
| Grad norm for multiplier on MuJoCo | $1.0$ ($\alpha$), $0.5$ ($\beta$) |
| Grad norm for multiplier on AntMaze | $1.0$ ($\alpha$), $1$ ($\beta$) |
| Critic Grad Steps | 3e6 |
| Actor Grad Steps | 3e6 |
| Target Critic EMA | 0.005 |
| T | 5 |
| Beta schedule | Variance Preserving [] |
| Actor Dropout Rate | 0.1 for actor on all tasks |
| Critic Dropout Rate | 0.1 for AntMaze Tasks in AlignIQL-hard |
| Number Residual Blocks | 3 |
| Actor Cosine Decay [] | Number of Actor Grad Steps |
| Optimizer | Adam [] |

### D.2 Sparse Rewards Tasks

For sparse reward tasks, we use $\eta=1>0$ and $\tau=0.7$ and keep other hyperparameters the same as Table [3](#A3.T3). The weight of IDQL is $w_{2}^{\tau}(s,a)=|\tau-\mathbbm{1}(Q(s,a)<V_{\tau}^{2}(s))|.$ since we use expectile regression to train the critic. We also follow the IDQL’s advice the take the max probability action at evaluation time.

### D.3 Results of AlignIQL-hard

In this chapter, we report the results of AlignIQL-hard. Table [D.1](#A4.SS1) reports the hyperparameters we used for AlignIQL-hard.

**Table 4: Average Results of AlignIQL-hard on D4RL tasks.**
|  | AlignIQL-hard | AlignIQL |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| D4RL Tasks | $N=16$ | $N=64$ | $N=256$ | $N=16$ | $N=64$ | $N=256$ |
| Locomotion | $69.7$ | $71.0$ | $71.5$ | $73.2$ | $74.0$ | $72.3$ |
| AntMaze | $54.2$ | $57.9$ | $56.7$ | $65.8$ | $70.2$ | $70.7$ |

We also report the performance of AlignIQL under different $N$. For the MuJoCo tasks, we report the results for $\eta=-1$, and for AntMaze, we report the results for $\eta=1$. Therefore, the results in Table [4](#A4.T4) are slightly lower than those in Table [1](#S5.T1). For AlignIQL and AlignIQL-hard, especially for AlignIQL-hard, we perform minimal hyperparameter tuning. In most cases, we use the default parameters of IDQL. Therefore, the performance of our algorithm can be further improved with additional tuning.

## Appendix E Training Curves D4RL Tasks

Figure: Figure 3: Training curves of AlignIQL and IDQL, where AlignIQL-P means we use hyperparameter $\eta>0$ and AlignIQL-P means $\eta<0$. The normalized score is calculated by averaging the scores for the three different $N$ ($N=16,64,256$), where $N$ represents the number of actions generated by the diffusion-based behavior policy.
Refer to caption: /html/2405.18187/assets/x3.png

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
questions: see above
soundness: 3
strengths: 1. The proposed method is derived rigorously.
2. The experiment shows that the proposed method has good empirical performance compared with other baselines on standard benchmarks.
summary: The paper considers policy extraction problem, where sometimes in offline RL, existing algorithms only learn value function, and policy extraction problem is to find a policy that coorespond to the policy and does not perform OOD actions. The paper considers distilling policies from value functions learned with IQL algorithm, and propose the implicit policy-finding problem. The solution of the IPF problem leads to the proposed AlignIQL algorithm, from a careful derivation of the IPF formulation. In the experiment, the proposed method is compared with several baselines with competitive performance.
weaknesses: 1. The formulation aims to use a general regularization function $f$, which is a good attempt. However, the remaining results seems to rely on the case that $f(x) = \log(x)$. Does the result generalize to any other regularization function?
2. Remark 5.7 seems very hand-wavy. How does the algorithm ensure that the action with the positive advantage is chosen? It does not seem to be reflected in the loss function. 
3. While the result in table 1 looks impressive, I am not sure if this can serve a strong evidence that the proposed method is better than AWR. The proposed method is equipped with diffusion policies, but the IQL (AWR) baseline seem to only use MLP so it might not be a fair comparison. 
4. The result in table 1 is missing standard deviation. 
5. The goal of section 6.2 is unclear. What is the baseline that is compared against in this section? 
6. Some minor issues: a) in eq. 1, is a $\pi(a \mid s)$ missing? b) in eq. 2, where is the $Q_{\theta}$ from? It does not appear eq. 1. c) in eq. 7, the notation $a$ is overloaded in $a' \sim \pi(a \mid s)$.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: See Weaknesses
soundness: 2
strengths: - This paper introduces a new approach to tackle the implicit policy-finding problem, combining theoretical rigor with practical effectiveness in offline RL.
- The proposed algorithm, AlignIQL, performs well across varied tasks, demonstrating versatility and effectiveness across different offline RL benchmarks.
summary: This paper proposes AlignIQL (in two versions) to address the implicit policy-finding problem. The authors formulate it as a constrained optimization problem and derive a closed-form solution. The performance of AlignIQL is competitive compared to the baselines.
weaknesses: - While AlignIQL is rigorous, it adds complexity to training by requiring additional multiplier networks and diffusion models, which may increase computational costs and sensitivity to hyperparameters. The scalability of the method is also a concern; can it be extended to image-based tasks?
- The authors do not explain the use of diffusion modeling in the methods section.
- The performance of AlignIQL raises some concerns:
    - The authors argue that MuJoCo tasks are already saturated for offline RL, which I agree with. However, AlignIQL's performance is also considerably worse than Diffusion QL and even worse than IQL in 4 out of 9 tasks. Given that AlignIQL consumes more computational resources, this discrepancy is problematic.
    - There is a significant performance difference between the authors' version and the original IDQL paper, which further leaves the reader uncertain about the supposed improvements in AlignIQL's performance.
    - The results were obtained using inconsistent hyperparameters, yet the authors under-analyze the ablation study and hyperparameter sensitivity.
    - The authors state, “Figure 2 shows that as training time increases, the performance of AlignIQL with different N converges to the same value, which shows that AlignIQL is insensitive to N.” This conclusion is not obvious from Figure 2. A clearer approach would be to report the mean and standard deviation of these scores.
    - The optimization techniques may risk overfitting in AntMaze environments with sparse rewards, potentially reducing generalization to new scenarios. More testing on sparse reward tasks would benefit this submission.
- In this submission, "policy alignment" is defined differently from its use in language models. A more formal definition of “policy alignment” should be provided, or the authors could consider renaming it.
- A minor issue: multiple duplicate references appear in the bibliography (e.g., lines 568-573). Additionally, lines 916-917 may contain editing errors.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
questions: - How does AlignIQL perform in real-world environments with noisy or incomplete datasets? The paper evaluates performance on D4RL benchmarks, but it would be interesting to see how the method handles imperfect data.
- What are the key factors that affect the alignment between the Q-values and the learned policy in AlignIQL? Understanding the sensitivity of the method to different alignment parameters could help clarify its robustness.
- How does the computational complexity of AlignIQL compare to other state-of-the-art offline reinforcement learning methods in terms of training time and resource usage? This would help evaluate the method’s scalability for larger or more complex tasks.
- Is the approach compatible with more advanced neural network architectures, such as transformers, for offline reinforcement learning? Could integrating more modern architectures improve its performance?
- What are the potential limitations of applying AlignIQL to tasks outside of continuous control, such as discrete action spaces or hierarchical tasks? This could provide insight into the method’s broader applicability.
soundness: 3
strengths: - The introduction of AlignIQL as a constrained optimization approach represents a significant advancement in offline reinforcement learning, providing a fresh perspective on implicit policy extraction.
- The empirical results demonstrate that AlignIQL and its variant achieve competitive performance across a variety of D4RL benchmarks, particularly in challenging tasks with sparse rewards, indicating the effectiveness of the proposed methods.
- Theoretical Insights: The paper offers valuable theoretical analysis regarding the use of weighted regression for policy extraction, enhancing the understanding of the underlying mechanisms that contribute to the success of IQL methods.
- By incorporating policy alignment constraints, the approach ensures that the extracted policies are both effective and representative of the learned Q-function, leading to improved stability and reliability in offline settings.
- AlignIQL shows increased robustness to variations in hyperparameters compared to existing methods, which is crucial for practical applications where tuning can be challenging.
summary: The paper introduces AlignIQL, a novel approach to extracting implicit policies in offline reinforcement learning by formulating the implicit policy-finding problem as a constrained optimization problem. AlignIQL and its variant AlignIQL-hard leverage policy alignment constraints to ensure that the extracted policy reflects the learned Q-function while maintaining the advantages of decoupling the actor and critic in Implicit Q-Learning (IQL). The authors demonstrate that their method achieves competitive or superior performance on D4RL datasets, particularly excelling in complex sparse reward tasks like AntMaze and Adroit, while also being more robust to hyperparameter variations than existing methods. Additionally, the study provides theoretical insights into the conditions under which weighted regression can be effectively utilized for policy extraction in IQL. Overall, the proposed algorithms contribute to a better understanding of the bottlenecks in IQL-style methods and offer a more effective means for implicit policy extraction in offline reinforcement learning.
weaknesses: - While the experiments demonstrate competitive performance on specific D4RL benchmarks, the applicability of AlignIQL to other domains or more diverse environments may not be fully established, limiting its generalizability.
- The proposed framework may introduce additional complexity in implementation compared to existing methods, which could deter practitioners who seek simpler solutions for offline reinforcement learning.
- Although the paper includes comparisons with several baseline methods, it may benefit from a more comprehensive analysis against a broader range of state-of-the-art algorithms to fully contextualize its contributions.
- The performance improvements may be contingent on the quality of the dataset used, raising concerns about the approach’s robustness in real-world scenarios where data can be noisy or incomplete.
- The computational requirements for training AlignIQL could be higher than those of simpler methods, potentially limiting its scalability for larger-scale applications or real-time scenarios.

Meta-review: The paper studied how to distill policies from value functions learned with the IQL algorithm and proposed the implicit policy-finding problem. The author formulated this as a constrained optimization problem and derived two versions of alignIQL. Experiment results show that AlignIQL is compatible with diffusion model based policies and achieved competitive performance under standard offline RL benchmarks. The weaknesses of the paper are in the experiment,

Decision: Reject