ICML 2022 | 基于解耦梯度优化的可迁移模仿学习方法

2022 年 6 月 28 日 PaperWeekly

©作者 | 刘明桓

单位 | 上海交通大学

研究方向 | 强化学习，模仿学习

本文是一篇对自己工作的分享，算是一篇实验和方法上自己都还比较满意的工作。发表在 ICML 2022，文章提出了一种新的 policy structure 和优化方式，用一种解耦的方式学习 policy，使得 policy 中的模块可以被用于下游的迁移任务（不同的 action space 和 action dynamics）。

论文标题：

Plan Your Target And Learn Your Skills: Transferable State-Only Imitation Learning via Decoupled Policy Optimization

论文链接：

https://arxiv.org/abs/2203.02214

想法：模仿学习通常只是模仿在 state 上出 action 的 decision 能力，我们能不能模仿一个更加 genreral 的 plan 能力？

比如在自动驾驶的时候，虽然一个司机开不同的车的时候可能采取的油门刹车力度不一样，但是其实对道路规划的能力是同样共享的。有了这种 high-level 的 plan 能力，我们就可以 transfer 到不同的 action space 的情况中去。但是传统的 policy 建模，对于 action 是 ad-hoc 的，无法训练一次迁移多次。

通过观察我们可以发现如果把 policy 拆成两个模块，他们拼起来也是一个策略，但是我们先用一个规划模块从 s 去 plan 到 s'，再由一个 inverse dynamcis model（IDM）去输出此时应该采取的动作 a，这样连起来不就是一个 policy 了吗？通过 formally 的表示，我们也可以得到这样一个拆解的表示。

策略解耦

下面的图展示了一种这样的策略结构。

这里由于我们是从 imitation 出发的，所以也推了一下要 match 专家的状态分布也可以通过 match 一个 plan 模块来得到，但是前提是，这个 inverse dynamics 是给定而且提前学好的！

这时候如果我们可以训练得到这两个模块，他们拼在一起可以得到一个完整的 policy，而分开之后上层的 plan 模块其实可以直接用于一个新的 action space 不同的任务（只需要学一个底层的 IDM）。如果底层的 IDM 是一个给定的 ground truth，这样的 policy 也可以直接工作；如若不然，则需要通过学习得到。

监督学习

一种最直接的监督学习方式就是分开学习两个模块。但是这样做的话，上层模块只能在专家数据上学习，下层则可以在新收集的经验上学习。

一个很明显的问题，和 bc 一样，由于只在部分数据上学习，compounding error 会很严重。而且 planner 输出的 s' 也不一定能马上到，还要受到 inverse dynamics 的影响。下面我们将推导如何通过策略梯度来学习。

解耦策略梯度

首先推一个直接的形式。如果把这两部分看作是无约束无意义的 function 的话，可以直接梯度回传。这也是第一版投稿的时候的做法，但是当时这么做，始终无法得到很准的 state plan，也就是 plan 的 state 始终无法到达。

在重新推导之后，我们发现，我们想要求 state planner 是准的，必须是要假定有一个准的 IDM 啊！否则这样的学习是没有意义的，因为你学到的 s' 在和环境交互的时候也根本到不了。因此在学习的时候，底层的 IDM 应该是（local）准的，至少在学习的这部分数据上。这样的梯度才是朝向一个有意义的 plan state。

这个时候的解释就很明显了，我们希望 IDM 来告诉我们，如果要输出一个特定的 action，应该朝哪个方向去改 plan 的 state。大家可以发现我们的 formulation 虽然出发自 imitation，其实并不局限于 imitation，也适用于 RL。对于 imitation learning，只是我们要额外训练一个 Discriminator（像 GAIfO 一样）来获取 reward 信号。

但是如果这么做，还是有问题。因为我们的 IDM 一般是 NN 来拟合的，而 NN 的拟合无法做到一一映射，也就是 NN 在没见过的数据上是有不可知性的，意思是我们本来在环境里得到的训练数据是，两个相邻的 s,s' 预测一个 action a，这个是符合 MDP 的。但是如果训练了一个 NN，给一个 s 和一个很离谱的 s‘，这个 NN 也是可以输出这个 a！这个问题导致了，通过 IDM 回去的信号，不一定是一个正确的 correction，也就是 state的plan 可能是一个不相邻的 next state，但 IDM 仍可以输出一个合法的 action（因为 action space 就那么多，总会输出一个 action）来交互。但这不是我们想要的。

这个图里，这个问题更明显。在一个 grid world，专家是在黄色 grid 上的箭头，从左下到右上。智能体可以从除了阴影以外的任意区域出发。如果我们只做监督学习（b），可以发现 planner 的预测输出都是专家数据的 state，不会输出合法的相邻 plan。而只用上面的 DePG，也会输出不合法的 planing。而我们后面修正之后的方法，则是可以输出泛化并且合法的 plan。

修正策略梯度

怎么办呢？我们如果加一个额外的约束让前面的 planner 是输出合法的 next state 是不是就可以了？最直观的就是加一个 MLE，但是在我们的推导中我们推导了一个所谓的 CDePG 来更新，其实 PG 就是 weighted 的 MLE，用这样一个目标来更新也是更符合优化目标，这点在文中有相关推导和证明。但是在实验中，我们测试的 MLE 和 CDePG 的效果其实差不多。

但是如果只优化 CDePG，由于 s,a,s' 都是策略采集的，单纯这么做反而限制了 action 的探索。注意原本的 PG 的探索，通过 Q 值的高点不同，输出不同的 action 和环境交互再修正 Q 值。但是在我们这种方式下，action 不是直接得到的，而是先得到一个 plan，再由一个底层的 IDM 得到 action。因此，如果这里的 plan 都是见过的，底层的 action 也不会有什么变化。但是 DePG 没有对 s' 的约束，因此天然有着 explore 的能力（通过输出一个不合法的 s' 来得到一个想要的 a）。因此在实际中，我们是同时优化这两个目标（加上监督学习）。