离线强化学习(Offline RL)系列3: (算法篇) IQL(Implicit Q-learning)算法详解与实现

[更新记录]

论文信息：Ilya Kostrikov, Ashvin Nair, Sergey Levine: “Offline Reinforcement Learning with Implicit Q-Learning”, 2021; arXiv:2110.06169.

本篇论文由伯克利Sergey Levine团队的Ilya Kostrikov以第一作者提出，发表在ICLR2022顶会上,并被确定为Poster，接收意见是：“This paper proposes a new paradigm — called in-sample Q learning — to tackle offline reinforcement learning. Based on the novel idea of using expectile regression, the proposed algorithm enjoys stable performance by focusing on in-sample actions and avoiding querying the values of unseen actions. The empirical performance of the proposed algorithm is appealing, outperforming existing baselines on several tasks. The paper is also well written.”

摘要：目前基于策略约束和正则化的离线强化学习算法非常广泛，它们直接正面去解决OOD之外的动作分布，使得Learned policy能够很好的逼近行为策略，本文提出的 IQL(Implicit Q-learning) 直接没有去学习OOD之外的动作，而是用已知的state-action进行学习，通过使用SARSA style的方式重构策略和值函数（引入Expectile Regression） $L$ ，在策略的抽取方面采用了AWR(Advantage Weighted Regression)方式抽取，直接确定 $Q$ 值如何随着不同的动作而变化，并借助随机动态对未来结果进行平均，而不是确定 $Q$ 值如何随着不同的未来结果而变化.结果表明，该算法可以达到SOTA的效果。

1. 问题及背景简介

在前几篇博文[BCQ、BEAR、BRAC、TD3+BC、CQL、REM]中，这些方法都面临了OOD问题，为了解决这些问题，作者们都从station-action的轨迹优化着手，通过使用函数近似，行为克隆、支撑集、随机Q值函数等方面不断让学习策略（learned policy）和行为策略（behavior policy）距离变小或者处在一个范围内，不管从Policy Constraint还是regulization方面都达到了不错的效果。

于是，本文的作者就发出疑问：

这也是本篇论文要解决的一个核心问题。在开始前我们先说两个概念 Multi-step DP和 Single-step DP，以及前序工作。

1.1 Multi-step DP

最近提出的离线 RL 方法的很大一部分是基于约束或正则化的近似动态规划（例如，Q-learning 或 actor-critic 方法），constraint或regulizaion用于限制与行为策略的偏差。 我们将这些方法称为“多步动态规划(Multi-step DP)”算法，因为它们对多次迭代执行真正的动态规划，因此如果提供高覆盖率数据，原则上可以恢复最优策略。通常情况下Multi-step DP问题也可以分为：

• 显式密度模型(explicit density model)：BRAC，BCQ，BEAR等
• 隐式差异约束（implicit divergence constraints）：AWAC，CRR，AWR等

1.2 Single-step DP

与Multi-step DP相比，有一种方法依赖于单步策略迭代的方法，即对行为策略的价值函数或 $Q$ 函数进行拟合，然后提取相应的贪心策略，或者完全避免价值函数并利用行为克隆目标。 我们将这些统称为“单步（Single-step DP）”方法。 这些方法也避免了query看不见的station-action，因为它们要么根本不使用价值函数，要么学习行为策略的价值函数。

下图是一个关于one-step和multi-step的对比：

1.3 Offline RL 优化目标

我们知道对于离线强化学习来说，优化的TD目标为：

$$L_{TD}(\theta )={\mathbb{E}}_{\left(s,a,s^{\prime }\right)\sim \mathcal{D}}\left[{\left(r(s,a)+\gamma \underset{a^{\prime }}{\max }{Q}_{\hat{\theta }}\left(s^{\prime },a^{\prime }\right)-Q_{\theta }(s,a)\right)}^2\right]$$

其中公式中的 $Q_{\hat{\theta }}(\cdot )$ 表示target network的值函数，且 $\pi (s)=\arg {\max }_a{Q}_{\theta }(s,a)$ ，那么就可以发现对于处于OOD之外的动作对 $a^{{}^{\prime }}$ ，就会计算得到一个错误的 $Q_{\hat{\theta }}(s^{{}^{\prime }},a^{{}^{\prime }})$ , 从而通过 $\max$ 导致了Overestimate, 且这个误差会不断地增大而导致learned策略变废。

于是作者就根据 SARSA style重新构造了一个损失函数：
$$L(\theta )={\mathbb{E}}_{\left(s,a,s^{\prime },a^{\prime }\right)\sim \mathcal{D}}\left[{\left(r(s,a)+\gamma Q_{\hat{\theta }}\left(s^{\prime },a^{\prime }\right)-Q_{\theta }(s,a)\right)}^2\right]$$

对两个公式进行对比发现少了 $\max$ 操作，也就是说这里不在query 处于OOD之外的动作，只对分布内的数据学习（处于fixed Dataset $D$ ）。

那么如果假设 $D$ 的容量无限，且没有采样误差，则应满足最优参数（贝尔曼最优方程）

于是作者为了估计在支持数据分布动作上（support of the data distribution）的最大 $Q$ 值, 在不查询学习到的 $Q$ 函数的情况下通过使用期望回归(Expectile Regression)，目标定义如下：

说到这里，作者只是解释如何处理OOD之外的数据， 那么实际的策略、值函数等更新过程到底是啥样的？

1.4 期望回归(Expectile Regression)

【可跳过】(但建议看一看)

下面是一张关于Expectile Regression和Quantiles Regression之间的区别。

在论文中的定义如下（和图中的基本没差别）：
$$\underset{m_{\tau }}{\mathrm{argmin}}{\mathbb{E}}_{x\sim X}\left[L_2^{\tau }\left(x-m_{\tau }\right)\right]$$

$$L_2^{\tau }(u)=|\tau -1(u<0)|u^2$$
作者的目标是去预测条件期望回归，即找到最小函数下的 $m_{\tau }(x)$

$$\underset{m_{\tau }(x)}{\mathrm{argmin}}{\mathbb{E}}_{(x,y)\sim \mathcal{D}}\left[L_2^{\tau }\left(y-m_{\tau }(x)\right)\right]$$

通过实验，作者得出了不同的 $\tau$ 对应的函数，并在二维分布上实验得到图(right)

这里我在看文章时纳闷，为什么这个所谓的 Expectile Regression 要比直接的MSE以及quantiles好用，后来看了审稿意见，审稿人也提出了对应疑问

下面是作者给出的回复：总之就是我使用不同的 $\tau$ 效果就是好，有点Incremental的感觉。

具体的关于相关的分析如下：

2. IQL原理部分

作者在文章产生了IQL与普通算法的最大的区别在于这里：

译：我们的目标不是估计随机转换（stochastic transitions）产生的值的分布，而是估计状态值函数相对于随机动作的期望值。 这是一个非常不同的统计数据：我们的目标不是确定 $Q$ 值如何随着不同的未来结果而变化，而是确定 $Q$ 值如何随着不同的动作而变化，同时随机动态（stochastic dynamics.）的对未来结果进行平均。

2.1 值函数构造(Expectile Regression构造)

下面是作者根据1.4中将的Expectile Regression构造Loss函数
$$L(\theta )={\mathbb{E}}_{\left(s,a,s^{\prime },a^{\prime }\right)\sim \mathcal{D}}\left[L_2^{\tau }\left(r(s,a)+\gamma Q_{\hat{\theta }}\left(s^{\prime },a^{\prime }\right)-Q_{\theta }(s,a)\right)\right]$$

其中的 $L_2^{\tau }$ 表示为Expectile Regression函数，作者在不同的 $\tau$ 下进行对比实验，最终得出当 $\tau =0.9$ 时效果最佳，相比如MSE性能提升3倍多。

另外作者对值函数和状态值函数分别分开的进行了构造，其中separate value function和 Q-function 的计算过程如下：

$$L_V(\psi )={\mathbb{E}}_{(s,a)\sim \mathcal{D}}\left[L_2^{\tau }\left(Q_{\hat{\theta }}(s,a)-V_{\psi }(s)\right)\right]$$

$$L_Q(\theta )={\mathbb{E}}_{\left(s,a,s^{\prime }\right)\sim \mathcal{D}}\left[{\left(r(s,a)+\gamma V_{\psi }\left(s^{\prime }\right)-Q_{\theta }(s,a)\right)}^2\right]$$

备注：这里作者说一个大的target value不能代表一个好的action，

于是提出单独的网络来表示值函数，且该网络只近似于（函数逼近）对动作分布的期望（这里作者参考了Behavioral Modelling Priors for Offline Reinforcement Learning）中的 $V$ 的定义 $\text{ with }{\hat{V}}^{{\pi }_i}(s)={\mathbb{E}}_{a\sim {\pi }_i(\cdot \mid s)}\left[\hat{Q}\left(s,a;{\varphi }_{i-1}\right)\right]$

2.2 策略函数构造(Policy Extraction by AWR)

构造了值函数候，下一步就是构造策略函数，作者在这里提出了使用AWR中的策略抽取（policy extraction）的方法,数学表达如下：

$$L_{\pi }(\varphi )={\mathbb{E}}_{(s,a)\sim \mathcal{D}}\left[\exp \left(\beta \left(Q_{\hat{\theta }}(s,a)-V_{\psi }(s)\right)\right)\log {\pi }_{\varphi }(a\mid s)\right]$$

这里我找到了AWR的抽取方法，具体的伪代码如下：

从伪代码可以看出，作者构造的 $L_{\pi }(\varphi )$ 和AWR基本相似，接下来就是算法的执行过程。

2.3 伪代码

下图是本文的核心代码，

作者表明：可以简单地通过修改 SARSA style的 TD backup 损失函数来做到这一点，而无需在目标值 $Q$ 中使用样本外(OOD)操作。 一旦这个 $Q$ 函数收敛，就可以使用优势加权行为克隆提取相应的策略。 ,另外作者在文中强调：

3. 实验结果分析

4. 代码实现

在代码实现部分，作者使用了JAX实现ikostrikov implicit_q_learning，这个比较独特。

当然作者也比较贴心，提供了基于RLkit的pytorch实现版本： Github: pytorch based RLkit

"""
        QF Loss
        """
        q1_pred = self.qf1(obs, actions)
        q2_pred = self.qf2(obs, actions)
        target_vf_pred = self.vf(next_obs).detach()

        q_target = self.reward_scale * rewards + (1. - terminals) * self.discount * target_vf_pred
        q_target = q_target.detach()
        qf1_loss = self.qf_criterion(q1_pred, q_target)
        qf2_loss = self.qf_criterion(q2_pred, q_target)

        """
        VF Loss
        """
        q_pred = torch.min(
            self.target_qf1(obs, actions),
            self.target_qf2(obs, actions),
        ).detach()
        vf_pred = self.vf(obs)
        vf_err = vf_pred - q_pred
        vf_sign = (vf_err > 0).float()
        vf_weight = (1 - vf_sign) * self.quantile + vf_sign * (1 - self.quantile)
        vf_loss = (vf_weight * (vf_err ** 2)).mean()

        """
        Policy Loss
        """
        policy_logpp = dist.log_prob(actions)

        adv = q_pred - vf_pred
        exp_adv = torch.exp(adv / self.beta)
        if self.clip_score is not None:
            exp_adv = torch.clamp(exp_adv, max=self.clip_score)

        weights = exp_adv[:, 0].detach()
        policy_loss = (-policy_logpp * weights).mean()

        """
        Update networks
        """
        if self._n_train_steps_total % self.q_update_period == 0:
            self.qf1_optimizer.zero_grad()
            qf1_loss.backward()
            self.qf1_optimizer.step()

            self.qf2_optimizer.zero_grad()
            qf2_loss.backward()
            self.qf2_optimizer.step()

            self.vf_optimizer.zero_grad()
            vf_loss.backward()
            self.vf_optimizer.step()

        if self._n_train_steps_total % self.policy_update_period == 0:
            self.policy_optimizer.zero_grad()
            policy_loss.backward()
            self.policy_optimizer.step()

5. 关于OpenReview审稿意见

本篇论文的审稿意见非常的充实，通过作者和审稿人的讨论，最后审稿人给论文做了分数修改，同时作者也做了很多的修改与补充

查看全部意见与Rebuttal： https://openreview.net/forum?id=68n2s9ZJWF8

6. 项目部署过程

参考文献

[1]. Ilya Kostrikov, Ashvin Nair, Sergey Levine: “Offline Reinforcement Learning with Implicit Q-Learning”, 2021; arXiv:2110.06169.
[2]. Arthur Charpentier, Quantile and Expectile Regresions, Erasmus School of Economics, May 2017, Access
[3]. Hajo Holzmann, Bernhard Klar: “Expectile Asymptotics”, 2015; arXiv:1509.06866.
[4]. David Brandfonbrener, William F. Whitney, Rajesh Ranganath, Joan Bruna: “Offline RL Without Off-Policy Evaluation”, 2021; arXiv:2106.08909.
[5]. Noah Y. Siegel, Jost Tobias Springenberg, Felix Berkenkamp, Abbas Abdolmaleki, Michael Neunert, Thomas Lampe, Roland Hafner, Nicolas Heess, Martin Riedmiller: “Keep Doing What Worked: Behavioral Modelling Priors for Offline Reinforcement Learning”, 2020, ICLR 2020; arXiv:2002.08396.

---

OfflineRL推荐阅读

离线强化学习(Offline RL)系列3: (算法篇) CQL 算法详解与实现
离线强化学习(Offline RL)系列3: (算法篇) TD3+BC 算法详解与实现（经验篇）
离线强化学习(Offline RL)系列3: (算法篇) REM(Random Ensemble Mixture)算法详解与实现
离线强化学习(Offline RL)系列3: (算法篇)策略约束 - BRAC算法原理详解与实现(经验篇)
离线强化学习(Offline RL)系列3: (算法篇)策略约束 - BEAR算法原理详解与实现
离线强化学习(Offline RL)系列3: (算法篇)策略约束 - BCQ算法详解与实现
离线强化学习(Offline RL)系列2: (环境篇)D4RL数据集简介、安装及错误解决
离线强化学习(Offline RL)系列1：离线强化学习原理入门