O2O : OLLIE: Imitation Learning from Offline Pretraining to Online Finetuning

ICML 2024
paper
code
解决离线到在线过程中，判别器对齐问题。

Intro

传统采用离线模仿学习结合基于GAIL的在线模仿学习，容易因为在线阶段初始化的判别器表现具有随机性，与离线获得的策略不一致。因此，本文提出的OLLIE，便是利用混合质量的数据，实现判别器与策略之间的对齐，从何防止O2O的performance drop。

Method

GAIL

GAIL是一种传统在线模仿学习算法，其目标是对抗的学习一个判别器，用于区分专家数据与在线数据
$$\underset{\pi }{\min }\underset{D}{\max }{\mathbb{E}}_{{\rho }^{\pi }}[\log D(s,a)]+{\mathbb{E}}_{{\bar{\rho }}^e}[\log (1-D(s,a))].(2)$$.
最优判别器输出表示为$D^{\ast }(s,a)=\frac{{\rho }^{\pi }(s,a)}{{\rho }^{\pi }(s,a)+{\tilde{\rho }}^e(s,a)}.$

Offline IL

假设专家数据$D_e$以及混合数据${\mathcal{D}}_o\doteq {\mathcal{D}}_e\cup {\mathcal{D}}_s$的状态动作分布分别为${\tilde{\rho }}^e>0$以及${\tilde{\rho }}^o>0$。离线学习的目标可以看作状态动作分布匹配问题，通过一个逆KL散度表达
$$\underset{\pi }{\min }{D}_{\mathrm{KL}}({\rho }^{\pi }\Vert {\tilde{\rho }}^e)={\mathbb{E}}_{(s,a)\sim {\rho }^{\pi }}\left[\log \frac{{\rho }^{\pi }(s,a)}{{\tilde{\rho }}^e(s,a)}\right]$$
为了让混合数据参与上式的处理过程，在对数项的分子分母同时添加${\tilde{\rho }}^o$，原问题改造为
$$\underset{\pi }{\max }{\mathbb{E}}_{(s,a)\sim {\rho }^{\pi }}[\tilde{R}(s,a)]-D_{\mathrm{KL}}({\rho }^{\pi }\Vert {\tilde{\rho }}^o)（5）$$
其中$\tilde{R}(s,a)\doteq \log \frac{{\bar{\rho }}^e(s,a)}{{\bar{\rho }}^o(s,a)}$。对于低维表格环境可以通过统计计算状态动作分布，而对于高维环境则可以通过训练一个判别器进行区分
$$\underset{d}{\max }{\mathbb{E}}_{{\tilde{\rho }}^e}[\log d(s,a)]+{\mathbb{E}}_{{\tilde{\rho }}^o}[\log (1-d(s,a))](6)$$
进而得到$\tilde{R}(s,a)=\log \frac{{\tilde{\rho }}^e(s,a)}{{\tilde{\rho }}^o(s,a)}=\log \frac{d^{\ast }(s,a)}{1-d^{\ast }(s,a)}.(7)$。
而对于KL正则，在Offline设定下无法保证策略性能。因此，算法提出状态动作的分布需要满足贝尔曼流约束(Bellman flow constraints),即
Z ≐ { ρ : ρ ( s , a ) ≥ 0 , f s ( ρ ) = 0 , ∀ s ∈ S , a ∈ A } ( 8 ) w h e r e     f s ( ρ ) ≐ ( 1 − γ ) μ ( s ) + γ ∑ a , s ′ T ( s ∣ s ′ , a ) ρ ( s ′ , a ) − ∑ a ρ ( s , a ) . \mathcal{Z}\doteq\left\{\rho:\rho(s,a)\geq0,f_{s}(\rho)=0,\forall s\in\mathcal{S},a\in\mathcal{A}\right\}\quad(8)\\\mathrm{where~~~}f_s(\rho)\doteq(1-\gamma)\mu(s)+\gamma\sum_{a,s^{\prime}}T(s|s^{\prime},a)\rho(s^{\prime},a)-\sum_a\rho(s,a). Z≐{ρ:ρ(s,a)≥0,fs​(ρ)=0,∀s∈S,a∈A}(8)where   fs​(ρ)≐(1−γ)μ(s)+γa,s′∑​T(s∣s′,a)ρ(s′,a)−a∑​ρ(s,a).
上述条件下，策略与其平稳状态-动作分布之间存在一一对应关系：如果 $\rho \in \mathcal{Z}$，则$\rho$是策略 ${\pi }_{\rho }(a|s)\doteq \rho (s,a)/{\sum }_{a^{\prime }}\rho (s,a^{\prime })$下的状态动作平稳分布，而${\pi }_{\rho }$是唯一$\rho$的平稳策略。因此，问题 (5) 可以等价地写成以下形式：
$$\begin{gathered} \underset{\rho \ge 0}{\max }{\mathbb{E}}_{(s,a)\sim \rho }[\tilde{R}(s,a)]-D_{\mathrm{KL}}(\rho \Vert {\tilde{\rho }}^o)\text{(9)} \\ s.t.f_s(\rho )=0,\forall s\in \mathcal{S}.\text{(10)} \end{gathered}$$
由于目标和约束分别在$\rho$上是凹的和仿射的，问题 (9)-(10) 是一个凸优化问题。考虑上述问题的拉格朗日：
$$L(\rho ,\nu )\doteq {\mathbb{E}}_{s,a\sim \rho }[\tilde{R}(s,a)]-D_{\mathrm{KL}}(\rho \Vert {\tilde{\rho }}^o)+\sum _s\nu (s)f_s(\rho )$$
其中$D_{\mathrm{KL}}(\rho \Vert {\tilde{\rho }}^o)={\sum }_{s,a}\rho (s,a)\log (\rho (s,a)/{\tilde{\rho }}^o(s,a))$,进一步将上述等式$L(\rho ,v)$转化为
$$L(\rho ,\nu )=\sum _{s,a}\rho (s,a)({\delta }_{\nu }(s,a)-\log \frac{\rho (s,a)}{{\tilde{\rho }}^o(s,a)})+(1-\gamma )\sum _s\nu (s)\mu (s)（11）$$
其中${\delta }_{\nu }(s,a)\doteq \tilde{R}(s,a)+\gamma {\sum }_{s^{\prime }}\nu (s^{\prime })T(s^{\prime }|s,a)-\nu (s)$
对其求最优$\rho$, 令$\frac{\partial L}{\partial \rho (s,a)}=0$得到$$\rho (s,a)={\tilde{\rho }}^o(s,a)\exp \left({\delta }_{\nu }(s,a)-1\right).(13)$$
将等式（13）带入等式（11）得到关于v的优化问题
$$\begin{array}{rlr}\underset{\nu }{\min }L(\nu )& \doteq {\mathbb{E}}_{(s,a)\sim {\tilde{\rho }}^o}\left[\exp \left({\delta }_{\nu }(s,a)-1\right)\right]+(1-\gamma ){\mathbb{E}}_{s\sim \mu }\left[\nu (s)\right].& \text{(14)}\end{array}$$
直接优化问题 (14) 是有问题的，因为${\delta }_{\nu }$中的期望由于双采样导致随机梯度存在偏差，且等式 (14) 中的指数项在实践中很容易导致数值不稳定。因此，文章引入凸共轭（convex conjugate）解决：

将x用${\delta }_{\nu }$代替得到$$\begin{array}{rl}\exp \left({\delta }_{\nu }(s,a)-1\right)& =\underset{y(s,a)}{\max }{\delta }_{\nu }(s,a)y(s,a)-y(s,a)\log y(s,a)(16)\end{array}$$
等式（16）带入（14)得到一个min-max问题
$$\underset{\nu }{\min }\underset{y}{\max }F(\nu ,y)\doteq {\mathbb{E}}_{(s,a)\sim {\tilde{\rho }}^o}[{\delta }_{\nu }(s,a)y(s,a)-y(s,a)\cdot \log y(s,a)]+(1-\gamma ){\mathbb{E}}_{s\sim \mu }[\nu (s)].(17)$$
由于 F (·, y) 在固定 y 下是凸的，并且 F (ν, ·) 在固定 ν 的凹的，极大极小定理成立 ，问题 (17) 实际上是一个凸凹随机鞍点 (SSP) 问题。对于固定转换 $(s,a,s^{\prime })$，将 ${\delta }_{\nu }(s,a,s^{\prime })$表示为
$${\tilde{\delta }}_{\nu }(s,a,s^{\prime })\doteq \tilde{R}(s,a)+\gamma \nu (s^{\prime })-\nu (s).(18)$$.
那（17）的无偏对应为：
$$\underset{\nu }{\min }\underset{y}{\max }\tilde{F}(\nu ,y)\doteq {\mathbb{E}}_{(s,a,s^{\prime })\sim {\mathcal{D}}_o}[{\tilde{\delta }}_{\nu }(s,a,s^{\prime })y(s,a)-y(s,a)\log y(s,a)]+(1-\gamma ){\mathbb{E}}_{s\sim {\mathcal{D}}_o(s_0)}[\nu (s)].(19)$$

Offline Policy Extraction

根据等式（13）可以获得最优策略与状态动作分布见存在如下关系
$${\rho }^{\ast }(s,a)={\tilde{\rho }}^o(s,a)\exp \left({\delta }_{{\nu }^{\ast }}(s,a)-1\right).(20)$$
$${\pi }^{\ast }(a|s)=\frac{{\rho }^{\ast }(s,a)}{{\sum }_{a^{\prime }}{\rho }^{\ast }(s,a^{\prime })}\propto {\tilde{\rho }}^o(s,a)\exp \left({\delta }_{{\nu }^{\ast }}(s,a)-1\right).(21)$$
当$\frac{\partial F}{\partial y(s,a)}=0$,可以得到$$y^{\ast }(s,a)=\exp ({\delta }_{{\nu }^{\ast }}(s,a)-1)(22)$$
因此${\pi }^{\ast }(a|s)=\frac{{\tilde{\rho }}^o(s,a)y^{\ast }(s,a)}{z(s)}(23)$

那么策略优化可以采用两种形式

1. Reverse KL-divergence
   类似SAC中的重参数策略优化${\min }_{\pi }J(\pi )={\mathbb{E}}_{s\sim {\mathcal{D}}_o}[D_{\text{KL}}\left(\pi (\cdot |s)\Vert \frac{q(s,\cdot )}{z(s)})\right](25)$。其中公式（7）带入公式（23）得到$q(s,a)\doteq {\tilde{\rho }}^e(s,a)y^{\ast }(s,a)\left(\frac{1}{d^{\ast }(s,a)}-1\right)$。
2. Forward KL-divergence
   $$\begin{array}{rlr}& {\mathbb{E}}_{s\sim {\rho }^{\ast }}\left[D_{\mathrm{KL}}({\pi }^{\ast }(\cdot |s)\Vert \pi (\cdot |s))\right]& \\ & ={\mathbb{E}}_{s\sim {\rho }^{\ast }}\left[{\mathbb{E}}_{a\sim {\pi }^{\ast }(\cdot |s)}\left[\log {\pi }^{\ast }(a|s)-\log \pi (a|s)\right]\right]& \\ & \iff {\mathbb{E}}_{(s,a)\sim {\rho }^{\ast }}\left[\begin{array}{c}-\log \pi (a|s)\end{array}\right]& \\ & ={\mathbb{E}}_{(s,a)\sim {\tilde{\rho }}^o}\left[-\frac{{\rho }^{\ast }(s,a)}{{\tilde{\rho }}^o(s,a)}\log \pi (a|s)\right]& \text{(26)}\end{array}$$
   由于$\frac{{\rho }^{\ast }(s,a)}{{\tilde{\rho }}^o(s,a)}=\exp \left({\delta }_{{\nu }^{\ast }}(s,a)-1\right)=y^{\ast }(s,a).(27)$，问题转化为加权模仿学习的优化问题
   $$\underset{\pi }{\max }J(\pi )={\mathbb{E}}_{(s,a)\sim {\mathcal{D}}_o}\left[\begin{array}{c}{y}^{\ast }(s,a)\log \pi (a|s)\end{array}\right](28)$$

Aligned Discriminator

离线获得最优策略${\pi }^{\ast }$后，便是解决在线判别器对齐问题。对于GAIL，
$$D_0(s,a)\doteq \frac{{\rho }^{\ast }(s,a)}{{\rho }^{\ast }(s,a)+{\tilde{\rho }}^e(s,a)}={\left(1+\frac{{\tilde{\rho }}^e(s,a)}{{\rho }^{\ast }(s,a)}\right)}^{-1}(29)$$
由等式（20）（22）对上式推到得到
$$\begin{array}{rl}{D}_0(s,a)& ={\left(1+\frac{{\tilde{\rho }}^e(s,a)}{{\tilde{\rho }}^o(s,a)}\cdot \frac{{\tilde{\rho }}^o(s,a)}{{\rho }^{\ast }(s,a)}\right)}^{-1}\\ & ={\left(1+\frac{d^{\ast }(s,a)}{1-d^{\ast }(s,a)}\cdot \frac{1}{\exp \left({\delta }_{{\nu }^{\ast }}(s,a)-1\right)}\right)}^{-1}\\ & ={\left(1+\frac{d^{\ast }(s,a)}{1-d^{\ast }(s,a)}\cdot \frac{1}{y^{\ast }(s,a)}\right)}^{-1}.(30)\end{array}$$
判别器的输出天然与离线下学习的$d$以及$y$挂钩，自然而然实现对齐。

Implementation with Function Approximation

在实际操作中，使用参数化模型对$\varphi ,d,\nu$建模。使用前向策略优化(Forward KL-divergence)。在在线微调过程中，通过离线学习 得到的由参数${\varphi }_d$ 和 ${\varphi }_y$ 构造的$d$与$y$来初始化判别器：
$$D_{{\varphi }_y,{\varphi }_d}(s,a)={\left(1+\frac{{\varphi }_d(s,a)}{1-{\varphi }_d(s,a)}\cdot \frac{1}{{\varphi }_y(s,a)}\right)}^{-1}.\text{(31)}$$

伪代码

结果

Offline to Online

不同数量专家数据下进行在线微调


后续还有AntMaze、FrankaKitchen.以及vision-based输入下的微调结果。效果都非常好。