专栏名称: PaperWeekly

PaperWeekly是一个分享知识和交流学问的学术组织，关注的领域是自然语言处理的各个方向。我们热爱知识，分享知识，希望通过我们大家的努力为自然语言处理的发展做出一点点贡献。我们每周会分享一期特定话题的论文笔记和本周值得读的相关论文。

用强化学习解决现实问题：Stochasticity、Scale、GAE与Curriculum Learning

PaperWeekly · 公众号 · 科研 · 2024-10-17 17:56

正文

©PaperWeekly 原创 · 作者 | 白昊

单位 | UIUC

研究方向 | 强化学习、表征学习

强化学习（RL）在游戏领域上已经取得了巨大成功，比如下围棋；但是目前在现实问题上的应用还有很多困难。俞老师有一个很有名的回答：强化学习领域目前遇到的瓶颈是什么？

据我的观察，现实任务相较于游戏，最大的挑战就是 stochasticity（随机性）。这篇文章我就用我们最近的工作作为例子来讲一下用 RL 解决现实问题的一些经验，包括怎么去建模现实问题、什么是随机性、解决的算法应该有怎样的性质、base model 怎么选等等。

这里有一个选择：RL 可以用来做 pre-training 和 post-training，但是 post-training 更亲民一点，作为科研乞丐我也只能玩的起微型的 post-training，比如这份工作从头到尾用的 GPU 甚至没超过 13G，一张 T4 就够了。

但是模型小并不代表性能差，关键还是怎么去建模这个问题然后选择合适的算法。例如我们把这个小模型的性能从 18%（跟 GPT-4 水平差不多）直接飙到了 82%，翻了好几倍。现在 LLM 无脑堆数据和模型，我感觉还是有不少别的东西可以做的。

如果大家关心工作的内容，也可以先看看我们的 website 有很多 demo 和可以玩的东西：

https://digirl-agent.github.io

现实环境中的RL困境

1.1 Stochasticity

RL 目前打游戏很强，而游戏往往是符合某种可预测的规则的。比如下围棋，规则是放在那里的，你下 5 行 4 列的黑子，这个黑子就会出现在 5 行 4 列；玩 atari 也是，你按一下往右的按键，人物就一定会往右一格。

但是 real-world 的 task 不一定，会有几乎无限的扰动，比如如果你要买一个东西，输入物品名字后淘宝每次会给你推不同的东西。现在很多的 RL 工作做机械臂的一些模拟，都是在同一个工作台上做实验，如果换一个桌面环境，就秒挂了。

但是现实任务下，几乎没有办法避免环境中出现随机性，因此应该去显式地建模这种随机性（将解决随机性作为任务的一部分）。

为了体现随机性，我们选了一个用手机完成查资料、购物这些现实任务的方向（之前没有人在这个上面做过 RL），这个任务的所有环境都是在手机模拟器里面运行的，所以用户自己玩会出现的各种随机性，环境都会建模进去。一些比较典型的随机性有这些：

比如网页过一段时间主页会变，这个会导致每次点击网页同一个地方都可能看到不同的结果，比如前两张图，有时候网页上面的广告只有两行，有时候有三行，这样会导致整个网页的布局变化；再比如后面有几张图，有时候网页会有弹窗或者广告，必须先解决掉才能继续；又比如商品每次搜索的顺序都可能不同，是根据服务器的推荐算法排序的。

这些随机性会导致直接 behavior cloning 比较困难，直接用 monte carlo 去做 RL 也很困难，因为之前 monte carlo 用的 data 很可能没法 generalize 到 stochastic 的情况上。

1.2 Scale

要解决 stochasticity，一般来说需要巨大的数据量。可以这么理解：每一个可能存在 stochasticity 的自由度都需要一批 cover 到尽可能多情况的数据进行拟合。比如上面提到的 google 弹窗，理想情况下我们应该尽可能多地收集在 google 页面可能弹出的弹窗，以及成功的 trajectory 怎么解决这些弹窗的。

当类似这样的 stochasticity 的量多起来时，所需要的数据也是成倍增长的。要解决这个问题，如果用 offline RL，显然需要大量数据；如果是 online RL，那每次 training 间隔时收集的数据也是越多越好：因为如果有偏地建立对 stochasticity 的表示，一定会降低 collected trajectory 的质量。

例如我们的工作中，为了解决这个问题，就专门写了一个多机 distributed parallel 的脚本，支持多机大规模收集 trajectory。例如一台机子能开 8 个安卓模拟器的话，可以轻松 scale 到四机甚至十六机，收集速度是线性增长的，而且不需要任何 GPU。

我们对比了只在一个更多 CPU 的机子上开更多的模拟器的方案，这个单机方案的速度是 logarithm 增长的，比我们的方法慢很多，在 repo 里面有详细的分析：

https://github.com/DigiRL-agent/digirl/blob/master/multimachine/README.md

工作早期四天的训练，有了多机支持，后面只需要两天甚至一天就搞定了。

1.3 Reward: Step-level or Trajectory-level?

之前在 Jiayi 的工作里面大家尝试了用 step-level reward 去解这个问题，但是发现 evaluator 很难正确地给出 reward。比如如果我们要查询淘宝上某一个商品的价格，那么从淘宝跳到百度应该是一个负的 reward。

但是百度搜索也是可以找到淘宝的商品的，只要给一个“淘宝”的关键字就可以了（有时候甚至不需要关键字搜出来的也都是淘宝的结果）。这使得这个 reward 具体的数值很难确定（有时候甚至正负也不好确定）。

因此我们认为应该将 reward 做成 trajectory-level 的，因为判断“一个任务是否成功”比“这一步是否更靠近结果”简单很多，现在的 LLM 例如 Gemini 都可以完美解决了；而对于 step-level 上是否离目标更近，LLM 还不太行，所以应该转而交由一个专用算法解决，例如训练一个很靠谱的 critic。

小模型能否逆袭？

2.1 “小”不一定“差”

大家会好奇我们用的是什么模型，好像想要解这个问题，需要一个很 fancy 的模型架构。其实我们用的模型非常小而且朴实，参数加起来才 1.5B，一张 T4 就能跑，无 lora 的 training 才用了 13G 的 vram。

模型的架构是一个 freeze 的 CLIP 配一个要 tune 的 T5，之前一个大哥 pre-train 的，模型名字叫做 AutoUI。模型的输入是一张当前和上一个 state 的手机的屏幕截图 + 上一步的 action + 最终的 goal。

模型会输出接下来要做什么，例如文字或者坐标。注意输出的是坐标，比如 [528, 129]，屏幕上是没有任何提示的。之前的工作，例如 GPT，都是屏幕上有很多 tag，让 GPT 去选一个出来，比如 GPT 的屏幕被弄成了这样来让他好做一点：

我们的是没有这些 tag 的，直接输出一个 text 形式的坐标。这个难度要大得多，也更 realistic 一点，但是我们一个 1.5B 的模型直接把 GPT 干趴了，领先三倍多（详见后面的结果部分）。大模型固然好，但也不是人人都玩得起的，这个结果充分说明了一个专精的模型不需要很大。当然，如果要 scale 到 general tasks，可以想像 1.5B 的容量还是有点小了。

2.2 选择 base 模型的一些通用建议

这里我有一些选 base 模型的建议，首当其冲的就是 base 模型的性能不能太差。如果 policy initialization 太差，那么采集到的 trajectory 大概率是错的，这个时候 RL 必死无疑，不管你用什么算法都救不回来，没有 date 拿什么学。

那怎么获得一个还可以的 policy initialization？往往需要先从 human/ 很厉害的模型（比如 GPT-4）那边用监督学习学到一点知识，然后再上 RL。比如我们的工作中，大概的流程是这样的：

务必先拿一个稍微靠谱的 pre-trained 模型（性能有 10% 以上），然后 RL。如果性能太差（比如只有 1%），还是早点换一个 base model 比较好，不然不知道要学到什么时候。

有同学会问为什么自己上课时作业里遇到的 RL 可以用 random initialization：这是因为游戏的 action space 很小，盲猜也能对的。在 real-world problem 上一抓几千个 potential action，给 random policy 直接干玉玉了。

第二个是 base 模型的大小不能太大。一个原因是模型大，就需要更多的 data 去 fit，简称训不动。特别是 online RL，所有 data 都是边训边收集的，如果 model需要成倍的数据，那么 collect 也会需要成倍的数据，这需要很多机器一起跑 collection。

如果想偷懒少收集一点，那得调 learning rate，可能更痛苦。试想你跑一轮实验需要一个星期，还要不要毕业了。另一个原因是穷。这个就不太需要解释了。作为一名合格的科研乞丐，我深刻认识到穷是不需要任何解释的。

算法

这一部分会稍微烧脑一点，可能需要反复读几次。我会把用到的知识都做成 link 放在旁边。

3.1 Filtered AWR

环境的基础搭好以后，就可以着手思考什么样的算法适合这个问题了。一个能想到的最简单的解法就是 filtered behavior cloning（filtered BC），就是让模型自己玩，玩对了就拿对的整条 trajectory 让模型去学习自己是怎么做对的，玩错了就扔掉这个数据。

train 的方法也很简单，直接用 language modeling 的老方法去训练即可。但是 filtered BC 有两个问题：第一是只要 trajectory 成功了，那么其中的每一步，不管好坏，都会被模仿到；第二是没法高效地建模 stochasticity，因为没有 critic 去专门 fit 这种随机性。

因此我们想到了，可以不使用最终的 reward，转而使用 advantage去filter data。这种算法会产出一个模型，这个模型在 infer 时的架构与原来完全相同，不需要为了 RL 往模型里面加新的 module；而在 train 时，我们可以加入 critic 来估计这个 advantage，从而用这个 advantage 来 filter data；然后用这个 filtered data 来用和 BC 一样的 loss 训练模型。

这是一个非常简单的想法，灵感来源于AWR（没听说过的同学可以略看这篇：AWR（Advantage-Weighted Regression）算法详解与实现 [1]），但是我们把 AWR 硬化成了 hard filtering。

3.2 GAE的通俗理解

现在我们就要思考如何去设计这个 advantage 了。前面提到了 stochasticity 的问题，其实有关随机性，RL 的前人已经有过非常多的讨论，其中 John Schulman 的 GAE（Generalized Advantage Estimation）很大程度上缓解了随机性的问题。

对 GAE 还不了解的同学强推这篇：强化学习技术—GAE [2]。其实 GAE 本质上就是下面这一串 advantage estimate 的加权平均：

其中，称为 -step advantage。当 时，就是 1-step TD；如果我们把 拉到无穷大，那其实就是 时刻的 monte carlo estimate 减掉 baseline 。

这体现了GAE “generalized” 的原因：将 advantage 考虑成 TD 和 MC 的某种融合。因为 MC 是低 bias 的，TD 是低 variance 的，所以可以根据实际情况调这个比重。

GAE 被当成各种 baseline（包括 PPO 也用 GAE）是有原因的 - 在实验中，我们发现 GAE 用来做 real-world task 效果比单纯的 TD/MC 的 estimate 要好不少，因为 TD 很大程度上看 value 的准确度，而因为现实任务下的训练数据非常 noisy，这个 value critic 是极难学的。

这个时候就需要 MC 过来帮忙。但是单纯的 MC 又是不行的，因为现实任务有随机性，之前拿到的 data 没法 generalize 到随机出现的情况。

3.3 GAE的简化

我们的 state-level advantage 就是受到 GAE 的启发，只取了 GAE 的第一项和无穷项。

对于每一个 state action pair，我们取它的 MC return advantage

（这个是因为我们是 trajectory-level reward，只有最后一步 reward 可能为 1）和 one-step TD advantage

，两者 weighted sum 可得：

我们发现这个简化版的 GAE 能够比较 robust 地估计每个 step 的 advantage。从行为的角度看，MC 使得如果 trajectory 是成功的，那么其中的每一个 step 都会有比较大的 advantage（在实践中我们发现，这样导致不成功的 trajectory 基本没有机会，这是件好事）。

TD 使得 critic 认为 value gain 比较大的 transition 会获得更大的 advantage，从而相对更不会被 filter 掉。下面是一个例子：

左图中要买一个笔记本电脑，agent 从搜索框跳到商品展示页面，我们发现算法给了一个正的 transition difference；右图中也是买一个笔记本电脑，但是 agent 点错了，从 costco 的主页跳到了地址栏，这个时候我们发现算法给了一个负的 transition difference。这些都是比较符合我们设计的结果。由于这些 filter 都是 step-level 的，我们将这里的 critic 称为 state-level value critic。

3.4 Automatic Curriculum

在实践中，我们发现了另外一个问题：由于我们训练时用的 task 是从 task set 里 randomly sample 出来的，所以当一个 task 已经被完全掌握了以后，重新 sample 到不会带来很大的提升，也就是说正信息量比较小。

因此我们的 filter 同时需要过滤掉那些经常成功的 trajectory。但是，如果一个 task 没做对，它带来的信息量大概率是负的。要同时考虑这两种情况，就可以设计一个简单的 instruct-level advantage：

其中这个 instruct-level value critic 直接去回归 task 的成功率即可。这个算法的一个妙处在于，一个简单的 task 会在早期被解决，这个时候它应该被学到；一个稍微难一点的 task 在早期没法被解决，但是由于早期的任务解决后，模型通过 automatic curriculum 有更多资源去学习更难的 task，从而加速中期 task 的解决，但这个时候依然无法解决最难的 task。

等中期的也被解决的差不多后，训练资源又会向最难的方向倾斜。一个训练到中期的 automatic curriculum 的功能大致如下：

如图，在中期阶段，最简单的 task 已经学会了，因此这个 instruction（也就是task）的 value 会比较大，导致 instruction-level advantage 小；较难的 task 在这个阶段做对了以后会有更小的 instruction value，因此反而可能不被 filter 掉。

在实操中有一个细节，如果是 offline-to-online 的学习，那么 offline 阶段不要用 instruction-level critic，因为模型一开始还没有学会最简单的 task，而 pre-collected data 里面有很多成功的简单的 task，这会导致 filtered buffer 损失大批简单 task 的 trajectory，模型被 critic 坑了，完全学不到。

在 fully online 的学习下这倒不是个问题，因为 critic 的学习和模型一样，也是渐进的，一开始大家的 value 都小，那成功了也就都会进 buffer。

其实 Automatic Curriculum 是从更大的尺度上在解决 stochasticity 的问题，因为解决 stochasticity 需要 scale，而大家的计算资源都是有限的。如果一个简单的 task 已经被解决了，那么它一定在可观察的范围内已经解决了 stochasticity 的问题，那么更多的计算资源应该被分配到更难的 task 上，因为更难的 task 往往难在更强的 stochasticity。