在2020年之前,Offline RL的解决思路通常是将待优化策略的动作选择限制在离线数据集的分布/support区域上,从而避免分布外动作的选择,进而规避因为认知不确定性给动作值函数学习带来的过高估计等影响。Conservative Q-Learning(CQL) 发表于NeurIPS 2020,和之前进行policy constraint的算法(BCQ,BEAR)不同,CQL尝试通过修改动作值函数的back up方式得到真实动作值函数$Q$的一个下界估计$\hat{Q}$。在此之后,同一个组在NeurIPS 2021上发表了基于CQL改进的算法COMBO,二者的想法和思路基本类似,使用的推导方法也大致相同。鉴于这两篇文章在proof部分的证明挺有启发性,其中的一些concentration在其他的工作中也看到过,因此就借着Conservative Q Learning这个主题做点简单的总结。

CQL: Conservative Q-Learning for Offline Reinforcement Learning

Preliminaries

文章中对offline rl中的一些概念做了补充定义,比如empirical policy/bellman op等:

  • 离线数据集$ \mathcal{D} $是通过使用行为策略$\pi_\beta(a|s)$采样得到的:$ \mathcal{D}\sim d^{\pi_\beta(s)}\pi_\beta(a|s) $,因此在这sampling的过程中会引入sampling error,也就是因为离线数据集采样不充分导致使用$\pi_\beta$进行error估计时额外引入的误差。
  • Empirical policy:$\hat{\pi}_{\beta}(a \mid s):=\frac{\sum_{s, a \in \mathcal{D}} 1[s=s, a=a]}{\sum_{s \in \mathcal{D}} 1[s=s]}$,即离线数据集$ \mathcal{D} $对应的行为策略。
  • 各种算子

    • Bellman Operator:$\mathcal{B}^{\pi} Q=r+\gamma P^{\pi} Q$,其中$P^{\pi} Q(s, a)=\mathbb{E}_{s^{\prime} \sim T\left(s^{\prime} \mid s, a\right), a^{\prime} \sim \pi\left(a^{\prime} \mid s^{\prime}\right)}\left[Q\left(s^{\prime}, a^{\prime}\right)\right]$。
    • Empirical Bellman Operator:由于离线数据集中无法做到包含所有动作的转移数据,因此我们只能用$ \mathcal{D} $中包含的数据进行back up,这就是Emperical Bellman Operator $\hat{\mathcal{B}}$。
    • Optimal Bellman Operator:也就是最优贝尔曼算子,不多解释了,用$ \mathcal{B}^* $表示。

Concervative Off-Policy Evaluation

离线强化学习算法的关键在于规避因为认知不确定性(或者分布偏移)等因素所导致的Q值高估问题,因此CQL旨在找到原本Q值函数的一个下届估计$\hat{Q}$,进而使用这个$\hat{Q}$优化具有更加保守的policy value的策略。

具体而言,CQL提出要在原本Bellman Backup的基础上,额外minimize某个分布$\mu(s, a)$下的Q值。由于我们实际上是在离线数据集中sample状态$s$,因此这里$\mu$的marginal distribution要能够和$d^{\pi_\beta}(s)$匹配:$\mu(s, a) = d^{\pi_\beta}(s)\mu(a|s)$。基于这个思路,CQL提出了两种迭代方法:

  • CQL-1:

$$ \begin{equation} \hat{Q}^{k+1} \leftarrow \arg \min _{Q} \alpha \mathbb{E}_{s \sim \mathcal{D}, a \sim \mu(a \mid s)}[Q(s, a)]+\frac{1}{2} \mathbb{E}_{s, a \sim \mathcal{D}}\left[\left(Q(s, a)-\hat{\mathcal{B}}^{\pi} \hat{Q}^{k}(s, a)\right)^{2}\right] \end{equation} $$

  • CQL-2:

$$ \begin{equation} \begin{aligned} \hat{Q}^{k+1} \leftarrow \arg \min _{Q} \alpha \cdot\left(\mathbb{E}_{s \sim \mathcal{D}, a \sim \mu(a \mid s)}[Q(s, a)]\right.&\left.-\mathbb{E}_{s \sim \mathcal{D}, a \sim \hat{\pi}_{\beta}(a \mid s)}[Q(s, a)]\right) \\ &+\frac{1}{2} \mathbb{E}_{s, a, s^{\prime} \sim \mathcal{D}}\left[\left(Q(s, a)-\hat{\mathcal{B}}^{\pi} \hat{Q}^{k}(s, a)\right)^{2}\right] \end{aligned} \end{equation} $$

使用上面两种迭代方法可以分别得到具备一下性质的$\hat{Q}^\pi$:

  • 使用CQL-1,当$\alpha$足够大时,$\hat{Q}^\pi$为真实Q值$Q^\pi$的逐点下界。
  • 使用CQL-2,我们令$\mu(a|s)=\pi(a|s)$,此时得到的估计$\hat{Q}^\pi$不一定是真实Q值的下界,但是此时策略$\pi$的值是真实值函数的下界$ \mathbb{E}_{\pi(a|s)}[\hat{Q}(s, a)] \leq V^\pi(s) $。由于我们更加关注的其实是策略的值,因此使用这个迭代方法能够得到更为tight,但是也能避免over-estimation的下界估计。

Proof Sketch:
两条定理的证明都可以在Appendix C找到,这里简单提几点证明中比较重要的部分,以CQL-1为例

  • 直接取eq(1) RHS的derivative,得到$\hat{Q}^{k+1}(s, a)=\hat{\mathcal{B}^\pi}\hat{Q}^k(s,a)-\alpha\frac{\mu(a|s)}{\hat{\pi_\beta}(a|s)}$
  • 考虑因为sample error带来的$\hat{\mathcal{B}}$与$ \mathcal{B} $之间的差异,这里使用了之前文献中提到的关于reward和transition的concentration:$\forall Q, s, a \in \mathcal{D}, \quad\left|\hat{\mathcal{B}}^{\pi} Q(s, a)-\mathcal{B}^{\pi} Q(s, a)\right| \leq \frac{C_{r, T, \delta} R_{\max }}{(1-\gamma) \sqrt{|\mathcal{D}(s, a)|}}$,得到

$$ \hat{Q}^{k+1}(s, a) \leq \mathcal{B}^{\pi} \hat{Q}^{k}(s, a)-\alpha \frac{\mu(a \mid s)}{\hat{\pi}_{\beta}(a \mid s)}+\frac{C_{r, T, \delta} R_{\max }}{(1-\gamma) \sqrt{|\mathcal{D}(s, a)|}} $$

  • 左右两端取不动点,则有

$$ \begin{aligned} \hat{Q}^{\pi} &\leq \mathcal{B}^{\pi} \hat{Q}^{\pi}-\alpha \frac{\mu(a \mid s)}{\hat{\pi}_{\beta}(a \mid s)}+\frac{C_{r, T, \delta} R_{\max }}{(1-\gamma) \sqrt{|\mathcal{D}(s, a)|}}\\ &\leq(I-\gamma P^\pi)^{-1}\left[R - \alpha \frac{\mu}{\hat{\pi}_\beta} + \frac{C_{r,T,\delta}R_{\text{max}}}{(1-\gamma)\sqrt{|\mathcal{D}|}}\right] \\ &=Q^{\pi}(s, a)-\alpha\left[\left(I-\gamma P^{\pi}\right)^{-1}\left[\frac{\mu}{\hat{\pi}_{\beta}}\right]\right](s, a)+\left[\left(I-\gamma P^{\pi}\right)^{-1} \frac{C_{r, T, \delta} R_{\max }}{(1-\gamma) \sqrt{|\mathcal{D}|}}\right](s, a) \end{aligned} $$

  • 因此,只需要$\alpha$能够满足一定的条件保证第二项大于第三项即可。特别的,如果$ \hat{\mathcal{B}}=\mathcal{B} $,也就是没有sample error,那么只需要$\alpha>0$即可。本质上这个定理想要阐述的就是如何选取$\alpha$来克服sample error以保证得到一个Q值的下界估计。

对于CQL-2,其实证明方法是类似的,具体而言就是要考察(假设不考虑sample error)

$$ \hat{V}^{k+1}(s) = \mathbb{E}_{a\sim \pi(a|s)}[\hat{Q}^{k+1}(s, a)]=\mathcal{B}^{\pi} \hat{V}^{k}(s)-\alpha \mathbb{E}_{a \sim \pi(a \mid s)}\left[\frac{\mu(a \mid s)}{\pi_{\beta}(a \mid s)}-1\right] $$

其中后一项在$\mu(a|s)=\pi(a|s)$时永远是positive的,并且当且仅当$\pi=\pi_\beta$时等于0;我们把这个值叫做underestimation value。

Implementation for $\mu$

正如上文所述,如果我们希望得到的估值$\hat{Q}$能有CQL-2带来的良好性质,那么条件是$\mu=\pi$,也就是$\mu$需要与当前待提升的策略(如果有的话)相同。但在一些online算法例如SAC当中,$\pi$实际上是根据$Q$导出的,因此我们可以借鉴这一做法,直接将$\mu$定义为能够maximize Q值的策略

$$ \begin{equation} \begin{aligned} \min _{Q} \max _{\mu} \alpha\left(\mathbb{E}_{s \sim \mathcal{D}, a \sim \mu(a \mid s)}[Q(s, a)]-\mathbb{E}_{s \sim \mathcal{D}, a \sim \hat{\pi}_{\beta}(a \mid s)}[Q(s, a)]\right) &\\ +\frac{1}{2} \mathbb{E}_{s, a, s^{\prime} \sim \mathcal{D}}\left[\left(Q(s, a)-\hat{\mathcal{B}}^{\pi_{k}} \hat{Q}^{k}(s, a)\right)^{2}\right]+\mathcal{R}(\mu) \quad(\operatorname{CQL}(\mathcal{R})) \end{aligned} \end{equation} $$

其中$\mathcal{R}(\mu)$是对策略的某种正则化,可以是:

  • 熵正则,则得到CQL($\mathcal{H}$)评估算法

$$ \min _{Q} \alpha \mathbb{E}_{s \sim \mathcal{D}}\left[\log \sum_{a} \exp (Q(s, a))-\mathbb{E}_{a \sim \hat{\pi}_{\beta}(a \mid s)}[Q(s, a)]\right]+\frac{1}{2} \mathbb{E}_{s, a, s^{\prime} \sim \mathcal{D}}\left[\left(Q-\hat{\mathcal{B}}^{\pi_{k}} \hat{Q}^{k}\right)^{2}\right] $$

这是因为我们清楚的知道在熵正则下,maximize Q的最优策略$\mu$满足$\mu(s, a)\sim Q(s, a)$,因此可以直接代入,这样我们的算法中就不存在显式的$\mu$了。

Policy Improvement

最后我们证明,使用上面的保守Q估计进行策略提升(Actor-Critic算法有显式的策略,而Q-Learning算法也可看作由Q导出的隐式策略),能够使得提升后的策略值依然是保守的。这是因为,尽管$\hat{Q}$已经能够确保$ \mathbb{E}_{a\sim\pi_k}[\hat{Q}^k(s, a)]\leq V^k(s) $,在使用上面的方法把$\pi_k$更新到$\pi_{k+1}$之后是无法保证依然满足concervative性质的。但直观上,由于我们也可以不断修改$\hat{Q}$,因此可以让$\pi_k$到$\pi_{k+1}$的变化步长尽可能小,小到因为策略的偏移带来的策略值的上升不足以抵消$\hat{Q}^k$带来的underestimation。这也就是Theorem 3.3:

Practical Implementation for CQL

实际上CQL的算法的实现相当简单,这里主要分为两种,一种是Q-Learning Style,一种是Actor-Critic Style。伪代码如下:

  • 如果是Q-Learning Style,那么实际上中间的$\mu$本身就可以作为最终的策略
  • 如果是Actor-Critci Style,还需要用SAC的训练方式额外训练actor。
  • 两点细节:

    • 关于trade-off参数$\alpha$,使用的是转化为拉格朗日问题自动调参,其中$\tau$是某种目标阈值

$$ \min _{Q} \max _{\alpha \geq 0} \alpha\left(\mathbb{E}_{s \sim d^{\pi} \beta(s)}\left[\log \sum_{a} \exp (Q(s, a))-\mathbb{E}_{a \sim \pi_{\beta}(a \mid s)}[Q(s, a)]\right]-\tau\right)+\frac{1}{2} \mathbb{E}_{s, a, s^{\prime} \sim \mathcal{D}}\left[\left(Q-\mathcal{B}^{\pi_{k}} \hat{Q}^{k}\right)^{2}\right] $$

  • 如果问题是连续控制问题,中间的$\log\sum_a\exp$采用的是Soft Q-Learning中的importance sampling方法来进行估计。

COMBO: Conservative Offline Model-Based Policy Optimization

与CQL算法同期发表的offline rl算法还有一类Model-Based算法,比如MOPO。这一类算法往往借助uncertainty oracle来给出智能体对于某一片区域的认知不确定性的度量,从而修改MDP来限制策略的轨迹。但问题在于,这类算法在实验中使用的uncertainty oracle往往并不可靠,不能准确地反映出model error和ensemble uncertainty之间的关系。与之相反,Tianhe Yu和Aviral Lumar等人在前作CQL的基础上,以一种不同的方法将Model融入到了保守Q函数的学习过程当中,这就是COMBO算法。

我们还是按照Notation -> Evaluation -> Policy Learning的顺序展开

Notation

因为引入了Model,所以会额外多一个MDP:

Conservative Policy Evaluation

具体而言,COMBO计算$\hat{Q}$使用的迭代公式如下

$$ \begin{equation} \begin{aligned} \hat{Q}^{k+1} \leftarrow \arg \min _{Q} \beta\left(\mathbb{E}_{s, a \sim \rho(s, a)}[Q(s, a)]-\mathbb{E}_{s, a \sim \mathcal{D}}[Q(s, a)]\right)+\frac{1}{2} \mathbb{E}_{s, a, s^{\prime} \sim d_{f}}\left[\left(Q(s, a)-\widehat{\mathcal{B}}^{\pi} \hat{Q}^{k}(s, a)\right)^{2}\right] \end{aligned} \end{equation} $$

其中有一些notation需要说明

  • $\rho(s,a)$和$d_f$是我们用于采样的两种分布。其中,和CQL一样,我们要求$\rho(s, a)$的状态边际分布能够匹配模型在使用某个策略$\pi$rollout时给出的状态分布:$\rho(s, a)=d^\pi_{\widehat{M}}(s)\rho(a|s)$。
  • 分布$d_f$则是offline data分布和模型使用策略$\mu$进行rollout的数据分布的混合:$d_f^\mu(s, a)=f\ d(s, a)+(1-f)d^\mu_{\widehat{M}}(s, a)$。简记为$d_f(s, a)$。
  • 这里分别说一下具体实现当中$\rho, f, \mu$的形式。$\rho(s, a)$分成两部分$\rho(s), \rho(a|s)$,前者从$\{d_{\widehat{M}^\pi}, d_f\}$中选择,后者为log-sum-exp也就是soft-max形式。$f$的可选集合为$\{0.5, 0.8\}$,$\mu$为$\{Unif(a), \pi(a|s)\}$.

利用相同的方法,我们可以得到

$$ \begin{equation} \hat{Q}^{k+1}(s, a) = \widehat{B}^\pi \hat{Q}^k(s, a)-\beta \frac{\rho(s, a)-d(s, a)}{d_f(s, a)} \end{equation} $$

我们下面来考虑使用这个迭代目标得到的$\hat{Q}$有什么性质。

  • 首先给出Lemma A.1,讨论了在任意分布$\rho(s, a)$下expected penalty的性质
  • 使用和CQL中相类似的方法,我们可以得到当$\beta$足够大时,$\hat{Q}^\pi$是真实Q值的逐点下界
  • 同样,我们期望关注的是在某个初始状态$\mu_0(s)$的分布下值函数$V(s)^\pi$的下界,因此作者进一步给出

一直到这里,证明过程都与CQL基本一致,区别只是在于引入了$d_f$作Bellman Backup。但是Remark1中提到,COMBO并不像CQL-2一样会在每一个状态$s$上都产生下界的值估计$\hat{V}(s)$,而是只会在模型的初始状态分布上给出下界估计。这本质上是因为现在的$\rho(s, a)$匹配的是模型的初始状态分布$\rho(s, a)=\mu_0(s, a)\rho(a|s)$,这个初始状态分布和离线数据集的初始状态分布$d(s)$是有区别的。$d(s)$越大、$\rho(s, a)$越小的数据,上述下界估计越容易失效。

  • 最后,需要着重强调的是,在离线数据集的状态分布上,COMBO算法相比CQL算法在一定条件下可以给出更为tight的下界,这也就是Proposition 4.2。这也是做的claim的比CQL更好的原因。

Practical Implementation for COMBO

最后的伪代码实现并不复杂

  • 其中,rollout distribution实际上是soft-max策略。也就是说在评估步我们依然是类似CQL($ \mathcal{H} $)那样利用采样计算log-sum-exp
  • 如果是Actor-Critic算法,那么策略提升步就会相应提升策略;如果是Q-Learning类型的value-based算法,这一步跳过即可。

写在后面

个人认为CQL是非常solid的工作。COMBO在CQL的基础上融合了model,从model-based offline算法的视角来看它移除了其他model-based算法对uncertainty度量的依赖,从CQL的角度来看它在一定条件下取得了更紧的下界度量。但是,伴随COMBO而来的是额外的超参数和采样策略抉择,从附录就可以看出实现过程中的复杂度。除此以外,COMBO算法在附录部分的证明个人认为还有一些存在问题的地方,但是作者并没有开源自己的代码供大家讨论,因此我本人对COMBO算法持保留态度。

最后修改:2022 年 02 月 21 日
ヽ(✿゚▽゚)ノ