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

[更新记录]

论文信息：Conservative Q-Learning for Offline Reinforcement Learning]
[Code]

本文由UC Berkeley的Sergey Levine团队(一作是Aviral Kumar)于2020年提出，并发表在NIPS2020会议上。论文的主要思想是在$Q$值基础上增加一个regularizer，学习一个保守的Q函数，作者从理论上证明了CQL可以产生一个当前策略的真实值下界，并且是可以进行策略评估和策略提升的过程。从代码的角度上来说，本文的regularizer只需要20行代码即可实现，大幅提升了实验结果。同时作者也全部opensource了代码，非常推荐研究。

摘要：在CQL算法出来之前，离线强化学习中对于分布偏移问题的解决思路是将待优化策略的动作选择限制在离线数据集的动作分布上，从而避免分布外的动作出现Q值的过估计问题，进而减少了未知的动作在策略训练学习过程中的影响，这种方法被称为策略约束（Policy constraint），比如离线强化学习中的BCQ和BEAR算法。CQL尝试通过修改值函数的back up方式，在$Q$值的基础上添加一个regularizer，得到真实动作值函数的下界估计。实验表明，CQL的表现非常好，特别是在学习复杂和多模态数据分布的时候。

1、预备知识

1.1 sample error

离线数据集$\mathcal{D}$是通过使用行为策略${\pi }_{\beta }(\mathbf{a}\mid \mathbf{s})$采样得到的，$d^{{\pi }_{\beta }}(\mathbf{s})$是折扣的边缘状态分布，$\mathcal{D}\sim d^{{\pi }_{\beta }}(\mathbf{s}){\pi }_{\beta }(\mathbf{a}\mid \mathbf{s})$，采样的过程会因为状态动作对的采样不充分产生sample error。

1.2 Operator

关于对Bellman算子的理解和策略迭代过程可以参考这篇文章，通过Bellman算子理解动态规划

1.2.1 Bellman operator

$${\mathcal{B}}^{\pi }Q=r+\gamma P^{\pi }Q$$
$$P^{\pi }Q(\mathbf{s},\mathbf{a})={\mathbb{E}}_{{\mathbf{s}}^{\prime }\sim T\left({\mathbf{s}}^{\prime }\mid \mathbf{s},\mathbf{a}\right),{\mathbf{a}}^{\prime }\sim \pi \left({\mathbf{a}}^{\prime }\mid {\mathbf{s}}^{\prime }\right)}\left[Q\left({\mathbf{s}}^{\prime },{\mathbf{a}}^{\prime }\right)\right]$$

1.2.2 Empirical Bellman operator

离线数据集无法包含所有动作的转移数据，因此只能用$\mathcal{D}$中包含的数据进行back up，用${\hat{\mathcal{B}}}^{\pi }$表示。

1.2.3 Optimal Bellman operator

$${\mathcal{B}}^{\ast }Q(\mathbf{s},\mathbf{a})=r(\mathbf{s},\mathbf{a})+\gamma {\mathbb{E}}_{{\mathbf{s}}^{\prime }\sim P\left({\mathbf{s}}^{\prime }\mid \mathbf{s},\mathbf{a}\right)}\left[\underset{{\mathbf{a}}^{\prime }}{\max }Q\left({\mathbf{s}}^{\prime },{\mathbf{a}}^{\prime }\right)\right]$$

1.3 策略迭代

1.3.1 策略评估

当前我们在优化这个策略的过程中，会得到对应策略的值函数，根据值函数估计策略的价值。
$${\hat{Q}}^{k+1}\leftarrow \arg \underset{Q}{\min }{\mathbb{E}}_{\mathbf{s},\mathbf{a},{\mathbf{s}}^{\prime }\sim \mathcal{D}}\left[{\left(\left(r(\mathbf{s},\mathbf{a})+\gamma {\mathbb{E}}_{{\mathbf{a}}^{\prime }\sim {\hat{\pi }}^k\left({\mathbf{a}}^{\prime }\mid {\mathbf{s}}^{\prime }\right)}\left[{\hat{Q}}^k\left({\mathbf{s}}^{\prime },{\mathbf{a}}^{\prime }\right)\right]\right)-Q(\mathbf{s},\mathbf{a})\right)}^2\right]\text{ (policy evaluation) }$$

1.3.2 策略提升

通过在$Q$函数上取极大化，然后在这个$Q$函数上面做一个贪心的搜索来进一步改进它的策略。
$${\hat{\pi }}^{k+1}\leftarrow \arg \underset{\pi }{\max }{\mathbb{E}}_{\mathbf{s}\sim \mathcal{D},\mathbf{a}\sim {\pi }^k(\mathbf{a}\mid \mathbf{s})}\left[{\hat{Q}}^{k+1}(\mathbf{s},\mathbf{a})\right]\text{ (policy improvement) }$$

2、算法框架

离线强化学习算法的关键在于避免因为分布偏移导致的 $Q$ 值过估计问题，CQL算法直接从值函数出发，旨在找到原本 $\mathrm{Q}$ 值函数的下界估计，进而使用其去优化具有更加保守的policy value的策略。在离线数据集中采样状态 $s$ ，并且 $\mu$ 的分布要能够和 $d^{\pi \beta }(s)$ 匹配： $\mu (s,a)=d^{\pi \beta }(s)\mu (a\mid s)$ 。基于这个思路，给Q函数添加正则项使得Q的估计值变的保守。

2.1 Q_1

Q函数的更新公式：
$${\hat{Q}}^{k+1}\leftarrow \arg \underset{Q}{\min }\alpha {\mathbb{E}}_{\mathbf{s}\sim \mathcal{D},\mathbf{a}\sim \mu (\mathbf{a}\mid \mathbf{s})}[Q(\mathbf{s},\mathbf{a})]+\frac{1}{2}{\mathbb{E}}_{\mathbf{s},\mathbf{a}\sim \mathcal{D}}\left[{\left(Q(\mathbf{s},\mathbf{a})-{\hat{\mathcal{B}}}^{\pi }{\hat{Q}}^k(\mathbf{s},\mathbf{a})\right)}^2\right]$$
对于任意的策略$\mu (a\mid s)$，其中$\mu \in {\hat{\pi }}_{\beta }$。

$$\forall \mathbf{s}\in \mathcal{D},\mathbf{a},{\hat{Q}}^{\pi }(s,a)\le Q^{\pi }(\mathbf{s},\mathbf{a})-\alpha \left[{\left(I-\gamma P^{\pi }\right)}^{-1}\frac{\mu }{{\hat{\pi }}_{\beta }}\right](\mathbf{s},\mathbf{a})+\left[{\left(I-\gamma P^{\pi }\right)}^{-1}\frac{C_{r,T,\delta }{R}_{\max }}{(1-\gamma )\sqrt{|\mathcal{D}|}}\right](\mathbf{s},\mathbf{a})$$

当$\alpha$足够大时，${\hat{Q}}^{\pi }(\mathbf{s},\mathbf{a})\le Q^{\pi }(\mathbf{s},\mathbf{a})$，如果${\hat{\mathcal{B}}}^{\pi }={\mathcal{B}}^{\pi }$, 也就没有sample error，$\alpha \ge 0$，${\hat{Q}}^{\pi }$ 为 $Q^{\pi }$ 的逐点下界。

2.2 Q_2

Q函数的更新公式：
$$\begin{array}{rl}{\hat{Q}}^{k+1}\leftarrow \arg \underset{Q}{\min }\alpha \cdot ({\mathbb{E}}_{\mathbf{s}\sim \mathcal{D},\mathbf{a}\sim \mu (\mathbf{a}\mid \mathbf{s})}[Q(\mathbf{s},\mathbf{a})]& -{\mathbb{E}}_{\mathbf{s}\sim \mathcal{D},\mathbf{a}\sim {\hat{\pi }}_{\beta }(\mathbf{a}\mid \mathbf{s})}[Q(\mathbf{s},\mathbf{a})])\\ & +\frac{1}{2}{\mathbb{E}}_{\mathbf{s},\mathbf{a},{\mathbf{s}}^{\prime }\sim \mathcal{D}}\left[{\left(Q(\mathbf{s},\mathbf{a})-{\hat{\mathcal{B}}}^{\pi }{\hat{Q}}^k(\mathbf{s},\mathbf{a})\right)}^2\right]\end{array}$$
令 $\mu (a\mid s)=\pi (a\mid s)$
$$\forall \mathbf{s}\in \mathcal{D},{\hat{V}}^{\pi }(\mathbf{s})\le V^{\pi }(\mathbf{s})-\alpha \left[{\left(I-\gamma P^{\pi }\right)}^{-1}{\mathbb{E}}_{\pi }\left[\frac{\pi }{{\hat{\pi }}_{\beta }}-1\right]\right](\mathbf{s})+\left[{\left(I-\gamma P^{\pi }\right)}^{-1}\frac{C_{r,T,\delta }{R}_{\max }}{(1-\gamma )\sqrt{|\mathcal{D}|}}\right])$$
此时得到的估计 ${\hat{Q}}^{\pi }$ 不一定是真实Q值的下界，但是此时策略 $\pi$ 的值是真实值函数的下界 ${\mathbb{E}}_{\pi (a\mid s)}[\hat{Q}(s,a)]\le V^{\pi }(s)$。

2.3 CQL

在$Q\_2$中，$\mu =\pi$，也就是说$\mu$需要与当前待提升策略相同，此时策略 $\pi$ 的值是真实值函数的下界，相仿online中策略迭代的过程，直接将$\mu$定义为能够最大化Q值的策略。
$$\begin{array}{rl}\underset{Q}{\min }\underset{\mu }{\max }\alpha ({\mathbb{E}}_{\mathbf{s}\sim \mathcal{D},\mathbf{a}\sim \mu (\mathbf{a}\mid \mathbf{s})}& [Q(\mathbf{s},\mathbf{a})]-{\mathbb{E}}_{\mathbf{s}\sim \mathcal{D},\mathbf{a}\sim {\hat{\pi }}_{\beta }(\mathbf{a}\mid \mathbf{s})}[Q(\mathbf{s},\mathbf{a})])\\ +& \frac{1}{2}{\mathbb{E}}_{\mathbf{s},\mathbf{a},{\mathbf{s}}^{\prime }\sim \mathcal{D}}\left[{\left(Q(\mathbf{s},\mathbf{a})-{\hat{\mathcal{B}}}^{{\pi }_k}{\hat{Q}}^k(\mathbf{s},\mathbf{a})\right)}^2\right]+\mathcal{R}(\mu )(\mathrm{CQL}(\mathcal{R})).\end{array}$$
添加了正则化项$\mathcal{R}(\mu )$，如果选择$\mathcal{R}(\mu )$是当前策略和先验策略的KL散度，那么原式就是：
$$\underset{\mu }{\max }{\mathbb{E}}_{\mathbf{x}\sim \mu (\mathbf{x})}[f(\mathbf{x})]+D_{\mathrm{K}\mathrm{L}}(\mu \Vert \rho )\text{ s.t. }\sum _{\mathbf{x}}\mu (\mathbf{x})=1,\mu (\mathbf{x})\ge 0\forall \mathbf{x}.$$
最优解决方法就是${\mu }^{\ast }(\mathbf{x})=\frac{1}{Z}\rho (\mathbf{x})\exp (f(\mathbf{x}))$，其中$Z$是正则化因子，将${\mu }^{\ast }$带回原式就得到：
$$\underset{Q}{\min }\alpha {\mathbb{E}}_{\mathbf{s}\sim d^{{\pi }_{\beta }}(\mathbf{s})}\left[{\mathbb{E}}_{\mathbf{a}\sim \rho (\mathbf{a}\mid \mathbf{s})}\left[Q(\mathbf{s},\mathbf{a})\frac{\exp (Q(\mathbf{s},\mathbf{a}))}{Z}\right]-{\mathbb{E}}_{\mathbf{a}\sim {\pi }_{\beta }(\mathbf{a}\mid \mathbf{s})}[Q(\mathbf{s},\mathbf{a})]\right]+\frac{1}{2}{\mathbb{E}}_{\mathbf{s},\mathbf{a},{\mathbf{s}}^{\prime }\sim \mathcal{D}}\left[{\left(Q-{\mathcal{B}}^{{\pi }_k}{\hat{Q}}^k\right)}^2\right]$$
令$\rho =Unif(a)$，就得到了$CQL(H)$
$$\underset{Q}{\min }\alpha {\mathbb{E}}_{\mathbf{s}\sim \mathcal{D}}\left[\log \sum _{\mathbf{a}}\exp (Q(\mathbf{s},\mathbf{a}))-{\mathbb{E}}_{\mathbf{a}\sim {\hat{\pi }}_{\beta }(\mathbf{a}\mid \mathbf{s})}[Q(\mathbf{s},\mathbf{a})]\right]+\frac{1}{2}{\mathbb{E}}_{\mathbf{s},\mathbf{a},{\mathbf{s}}^{\prime }\sim \mathcal{D}}\left[{\left(Q-{\hat{\mathcal{B}}}^{{\pi }_k}{\hat{Q}}^k\right)}^2\right]$$

2.4 策略提升

使用推导出的保守Q进行策略评估，能够使得提升后的策略值依然是保守的。$\hat{Q}(s,a)$ 保证了${\mathbb{E}}_{\pi (a\mid s)}[\hat{Q}(s,a)]\le V^{\pi }(s)$，修改$\hat{Q}$使得${\pi }_k$到${\pi }_{k+1}$变化尽可能小，这样引起的策略批偏移也就更小。
学习到的Q值下界是：
$${\mathbb{E}}_{{\pi }_{{\hat{Q}}^k}(\mathbf{a}\mid \mathbf{s})}\left[\frac{{\pi }_{{\hat{Q}}^k}(\mathbf{a}\mid \mathbf{s})}{{\hat{\pi }}_{\beta }(\mathbf{a}\mid \mathbf{s})}-1\right]\ge \underset{\mathbf{a}\text{ s.t. }{\hat{\pi }}_{\beta }(\mathbf{a}\mid \mathbf{s})>0}{\max }\left(\frac{{\pi }_{{\hat{Q}}^k}(\mathbf{a}\mid \mathbf{s})}{{\hat{\pi }}_{\beta }(\mathbf{a}\mid \mathbf{s})}\right)\cdot \epsilon$$
增加了Q估计值和真实值之间的距离，使得Q更加保守。

2.5 伪代码

• 如果是Q-learning模式：$\mu$可以作为最终的策略
• 如果是Actor-Critic模式：需要使用SAC的训练方式额外训练actor

3、结果

Gym结果

D4RL结果

Atari结果

4、代码实现

Github Github

参考文献

[1]. Aviral Kumar, Aurick Zhou, George Tucker, Sergey Levine:“Conservative Q-Learning for Offline Reinforcement Learning”,2020;arXiv:2006.04779.
[2]. Conservative Q-Learning

---

OfflineRL推荐阅读

离线强化学习(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：离线强化学习原理入门