博客园链接

世界模型阅读笔记


文献(按阅读顺序编排):

《Recurrent World Models Facilitate Policy Evolution》——世界模型
《Mastering Diverse Domains through World Models》——DreamerV3


世界模型:

本段的世界模型理解仅供个人读论文,不会按正规写法。
可供理解《Recurrent World Models Facilitate Policy Evolution》

简要描述:世界模型模仿人类基于现有感官对世界形成的心里模型,人类通常会接受世界输入的信号,在脑中分析当前情况并作出预测与判断,来决定自己的行为。世界模型大致上在模拟这一过程。

结构:

世界模型主干是一个大型的RNN(循环神经网络),其周围还存在其它组件。
Encoder(VAE,个人习惯称为Encoder):核心组件,将外部世界的输入信号(通常是图像帧)进行压缩,输出高维潜在向量 Z Z Z(layer通常不止一层)
Decoder:与Encoder作用相反,通常被用来可视化。
Model(即MDN-RNN,这里统称为Model):核心组件,在世界模型领域内Model通常为RNN,RNN接受一个记忆信号 h t − 1 h_{t-1} ht1,一个当前观测值 Z t − 1 Z_{t-1} Zt1和一个动作 a t − 1 a_{t-1} at1,输出预测值 Z t Z_t Zt(通常是下一帧图像)同时继续迭代隐藏记忆 h h h
(!!!注意:这里的预测值是一个分布 p ( Z t + 1 ) p(Z_{t+1}) p(Zt+1),即建模 P ( Z t + 1   ∣   a t , z t , h t ) P(Z_{t+1} \ | \ a_t,z_t,h_t) P(Zt+1  at,zt,ht),并不是确定的,因为复杂环境通常会带有随机性,这被称为混合密度网络,在这里与RNN结合,即MDN-RNN)
Controller:核心组件,用来对Model提供的 Z t Z_t Zt h t h_t ht预测下一步动作 a t a_t at,具体就是在当前状况和历史状况的结合下,通过学习和经验,考虑下一步动作的可能分布。

(还有其它组件类似Reward回报,即计算当前操作产生的价值,但并不是核心组件且在本阶段与训练无关,所以暂不说明)

注意在本文中,Controller将用C代替,Model为M,Encoder用V

具体构建:

参考下段代码,论文中的原文

def rollout(controller):
    # 假设 env, rnn, vae 是全局变量(外部定义)
    obs = env.reset()                # 重置环境,获取初始观测
    h = rnn.initial_state()          # 获取RNN的初始隐藏状态(用于处理时序信息)
    done = False
    cumulative_reward = 0

    while not done:
        z = vae.encode(obs)          # 将观测压缩为潜在向量(特征提取)
        a = controller.action([z, h])# 控制器根据当前特征和RNN状态选择动作
        obs, reward, done = env.step(a)  # 执行动作,获得下一状态、奖励、终止标志
        cumulative_reward += reward
        h = rnn.forward([a, z, h])   # 更新RNN隐藏状态(输入动作、特征、旧状态)

    return cumulative_reward

多次循环,每次按照先用V读取输入帧,将输入与记忆通过决策C产生动作,然后计算回报,预测下一帧的画面,然后进行RNN forward,重复操作。
注意到里面有一个done,通常由模型M预测的结果输出,表示当前游戏(或其它)是否结束。
这里来了一个关键点
为何模型要预测下一帧的画面,为何不直接从现实中读取。
在世界模型的研究中,直接读取现实通常耗费大量算力和时间,拖慢训练。而通过模型根据前画面和动作来预测下一画面,即可无需从现实中读取,除了初始,所有的画面都由模型内部自己“想象”,这直接在训练中省去了V和Decoder的工作。最后直接将次迁移到现实世界进行操作。(当然训练时还是需要现实数据进行训练,这种方法真正起作用的是C的训练和时间资源成本的控制)
大量事实表面这种方式的迁移泛化仅会丝丝影响模型现实操作的能力,于是研究者又设置了一个超参数 τ τ τ,让模型生成更加不确定的更复杂的画面,让模型在更难的空间中训练,回到现实时可以游刃有余一点。
同时这也可以一定程度上缓解模型利用机制钻空子,即M和C通过训练发现一种不符合操作规则的高回报路径,而产生路径依赖。

训练:

研究者将C,V,M分开单独训练,以获取最高效率。
其中对于C(Controller),它仅吃 Z t , h t Z_t,h_t Ztht,输出 a t a_t at,为一个简单的单层模型: a t = W [ Z t   h t ] + b c a_t=W[Z_t \ h_t] + b_c at=W[Zt ht]+bc,这种情况允许我们使用非反向传播算法来训练C,因为C的参数量相比之下是在太小。本文使用的是进化策略ES(类似于自然界优胜劣汰,同一时间生成多种情况,保留优质情况,然后根据优质情况继续生成子代,重复操作)
V,M都用传统的深度学习反向传播训练。
值得一提,进化策略ES的适用范围在世界模型种十分有趣,感兴趣可以找论文研究。

实验:

本文的模型基于一个2D赛车控制游戏进行实验,模型需操控赛车在随机乘生成的跑道上通过4种操作(左打方向盘,右打方向盘,加速,刹车),在最短时间内通过尽可能多的图块。

结合训练来看,大致流程如下:
1:从随机策略收集10000个轨迹(给模型的训练数据)。
2:训练V将帧编码。
3:训练MDN-RNN建模 P ( Z t + 1   ∣   Z t , h t , a t ) P(Z_{t + 1} \ | \ Z_t,h_t,a_t) P(Zt+1  Zt,ht,at)
4:训练C以最大化所有轨迹的期望得分。

具体细节:
先让一个Agent在这个环境里随机跑10000次,并记录每次的行动 a t a_t at和产生的环境观测。先通过Encoder-Decoder最小化差异训练出V,得到 Z t Z_t Zt。然后将所有数据 Z t , a t Z_t,a_t Zt,at丢进MDN-RNN中进行训练,以高斯混合模型的形式构建 P ( Z t + 1   ∣   a t , h t , Z t ) P(Z_{t+1} \ | \ a_t,h_t,Z_t) P(Zt+1  at,ht,Zt)。然后由ES策略训练C。由于V,M观测不到动作回报,所以认为它们是“无知的”,即你不可通过提前知道未来的回报而去改变生成模型的概率分布。

实验结果:

研究者先仅使用V和C构成的模型来测试,这意味着C只能读到当前帧向量,无法通过记忆系统。
实验发现这种模型仍然可以导航,但在急弯道处会左右摇晃冲出赛道。获得了632±251的分数。(此外,研究者在C中添加了一层隐藏层,可将分数提高到788±141)

方法平均得分
DQN343 ± 18
A3C (连续)591 ± 45
A3C (离散)652 ± 10
Gym 榜首838 ± 11
V 模型632 ± 251
带隐藏层的 V 模型788 ± 141
完整世界模型906 ± 21

实验展示了本文完整世界模型在这个任务上的良好表现,这主要依赖于M和C的良好表示。
(值得注意的是,通过之前的对比实验,我们发现运行模型同时访问 h t , z t h_t,z_t ht,zt极大提高了驾驶能力,这种M输出包含了对未来预测的概率分布,可以直接被C调用,这与现实中赛车手面对情况是通过经验和肌肉记忆进行本能反应的情况类似,C可以本能利用MDN-RNN中的未来预测指导决策。)

实验2:VizDoom

这里引出了在由模型自主生成的空间中学习,并迁移到现实空间的情况。
显然,如果模型的预测训练的越好,越接近真实环境,那么这是可行的。
这里用到了另一个游戏:DoomTakeCover。智能体必须学会躲避房间另一侧怪物射出的火球,这些怪物的唯一目的就是杀死智能体。最终得分为智能体在一次回合中存活的时间步。
与赛车实验的区别:模型预测要增加一个 d o n e t done_t donet,表示下一帧是否死亡。
本实验直接在虚拟环境中训练,使用的都是V编码的向量,不接触任何真实环境。(除了V(encoder)训练并编码第一帧给模型)最后在测试时,将模型迁移到真实环境进行观测结果。
这里,研究人员提出可以向虚拟环境增加更多的不确定性,具体由温度 τ τ τ来控制,使得模型的游戏过程变得更加困难。(甚至出现单纯因为运气而死亡)
此外,这里的RNN网络稍有不同。额外引入了 c c c记忆细胞,负责记录长时间的信息。它和 h h h的区别就是它们虽然都是记忆,但 h h h记忆是工作状态, c c c是信息。
实验结果:在较高温度下,模型在虚拟环境中得到了918的分数,并在真实环境中得到了1092的分数,这高于在虚拟环境中获得的分数。

温度 τ τ τ虚拟得分实际得分
0.102086 ± 140193 ± 58
0.502060 ± 277196 ± 50
1.001145 ± 690868 ± 511
1.15918 ± 5461092 ± 556
1.30732 ± 269753 ± 139
随机策略N/A210 ± 108
Gym 榜首N/A820 ± 58

尽管V无法捕捉每一帧的细节(例如怪物的数量),C仍然能够有良好的学习情况和表现。
增加 τ τ τ确实会对智能体的训练和迁移有一定好处,这将在下一段进行讨论。

欺骗世界模型(超参数 τ τ τ和高斯混合模型的作用)

像现实中,对于一款游戏,人们总会有卡BUG的手段以获取高分。在世界模型中也存在这种情况。
具体来说,对于初步世界模型,C可能会发现一种对抗性策略,使得以这种策略移动,模型预测的虚拟环境中不会有怪物发射火球。但是迁移到现实环境中这种策略会马上失效。
由于M生成虚拟环境,C也能访问M中的所有隐藏状态。这本质上是在给予智能体访问游戏内部引擎的权限,并不仅仅是表面的画面观测。这在世界模型中是一个经典问题,智能体能较容易地找到一种在动力学模型(预测模型)上表现良好,而在现实中失败的策略。这通常是因为C访问了模型错误的,偏离的预测。

高斯混合模型和 τ τ τ
本文的方法是使用高斯混合模型作为模型的预测结果,并通过超参数 τ τ τ进行随机量采样。高斯混合模型并不会表示一个确定的概率,而是一个概率分布。在通过 τ τ τ对分布进行采样,得到一个具有一定随机化的结果。
在这种随机化的加持下的虚拟环境中,模型无法通过一种确定的操作方式进行作弊,因为每次预测的结果具有随机性,都可能不是作弊策略所对应的。这基本解决了这个问题。
(但 τ τ τ也不能太高,这会导致虚拟环境过难,模型无法学习任何东西就死亡。上表展示了不同 τ τ τ的结果)

讨论:

1-使用RNN:

文中讨论的是使用RNN相对于FNN的优势,但这里想着重探讨一下为何当前主流世界模型基本使用RNN或其变体而不是Transformer这种也可以处理序列问题甚至有注意力机制的方法。
可能的原因:
① 考虑人类的智能,在原始情况下,人们通常会按照时间顺序进行注意力分配(即我们更容易记住刚刚发生的事情而不是一段时间前发生的事情,在通常情况下),RNN这种可在时间顺序层面进行递归的网络似乎更能处理这种情况。而Transformer在不给序列编码的情况下,注意力分配是均匀的。
② 由于RNN这种基于时间顺序的机制和其独特的设计,使得它在世界模型这种需要频繁预测未来的工作中能通过更多更好的方法来预测未来。而Transformer对于未来预测的工作似乎难以找到优秀的方法。(也可能是我还没接触到)
③ 考虑复杂度,RNN这种几乎可以说是O(1)的网络可以考虑更多的时间步。而Transformer的复杂度在这方面有所限制,可能无法大规模调用记忆信息。
④ 训练和更新,RNN可以根据当前状态直接连续训练,直接考虑当前真实情况输入将其加入自己的记忆系统。Transformer通常会将记忆截断(健忘)。
⑤ 世界模型被用于世界中行动,世界中的一些基本准则在Transformer机制下可能会出现因注意力分散和Text限制而淡化或丢失的情况。而RNN可以通过双RNN一个生成确定性预测(基本准则)一个生成混合模型分布(随机事件)进行平衡。

但是:其中有些问题Transformer并非不能解决(例如⑤),但在表现性上更优
2-使用ES(进化算法)

现在DL是主流算法,DL靠堆积参数量和训练量可以做到优秀。但仍然使用ES,肯定不是因为算力和资源的问题。其实ES这种优胜劣汰的训练方式更像人的学习过程。人在没有指导和经验的情况下去尝试多种不同路径,然后根据效果产生经验进行学习。而DL这种对学习率这种超参数较敏感的方法从这个方面考虑似乎有点太过死板。

改进:

关于V

文中提到V,M,C分开训练。但存在一个问题是V这种无监督训练不知道什么是重要的,什么是无关紧要的。文中就提到V将游戏操作决策过程中无关紧要的元素(比如墙壁上的瓷砖)给压缩,却没有压缩车道线。
一种思考是人的大脑本来就是一个整体,比如眼睛看到事物传给大脑,然后大脑做决策并产生反馈,这种反馈会通过某种方式作用回我们的眼部注意力。比如对于一个现实场景,我们更会去注意那些重要的事物(比如钱,美食),而忽略其它的事物(比如一张用过的纸巾),虽然眼睛捕捉到的画面并没有对这两者产生权重,但这种注意力也是一种隐性的学习训练过程。
所以考虑此,让V通过一种潜在的奖励机制来学习训练可能对于这种单一任务更有效果,但从单一任务迁移到多任务则需要更多的参数和训练量,需要进行权衡。

更加复杂的任务

目前(在那时)对于这种简单的任务,模型不需要更复杂的模块去引导学习过程,就能通过在有限且小范围的空间中通过随机训练而完成任务。对于更加复杂的任务,这种随机策略显然不太适用,这需要引入好奇心机制,去引导模型对真正感兴趣(潜在价值更高)的方向进行探究学习。

模型记忆

模型的容量始终有限度,如果接受的信息超过了这个限度,模型可能产生轻微或严重遗忘,这不相似于人对某些事物能记忆长时间而不会遗忘。(也有可能是人反复训练对某一事物的记忆而导致的,但这不妨碍我们的模型也可以对此类比出相似的算法)
文中提到的可能做法是,给模型增加外部记忆模块(类似语言模型的RAG检索)。但这产生的问题也不计其数。——有兴趣可以关注Memory-Augmented Neural Networks

动作堆叠

对于这种小型世界模型,C必须考虑每一帧的每一次动作。这相当于我们在走路时考虑每一次肌肉细胞的收缩(夸张了)。这种将微小(即不需要任何思考)的动作堆叠在时间线上的线性决策内显然非常无脑。
文中的办法是让C只负责大脑的决策层,M中加入多模块函数负责将更小的动作堆叠成一个大型动作,C可以调用这些函数进行操作。这好比一个总的C和M负责推理思考,其下还有多个次C和M进行动作,甚至更其下的C和M,根据任务的复杂程度定义深度。
研究者在文中还提出了一种更激进的策略,PowerPlay,通过强制性迫使模型学会一种大操作。具体的,模型先找到一个不会的任务,然后强制性学会这个任务(类似于学生不会的题目,老师直接给你讲完整的方法,然后学生就会了),其中学习的这个过程直接将所有这个新任务需要的小任务操作合并(这保证模型不会在学会新任务后忘记小任务)。这种类似于搜索新任务,然后合并其它小任务为一整个新任务的模块的方法,在哲学上其实类似于人的学习过程。

补充

参数:

实验一模型的参数:

ModelParameter Count
VAE4,348,547
MDN-RNN422,368
CONTROLLER867

实验二模型的参数:

ModelParameter Count
VAE4,446,915
MDN-RNN1,678,785
CONTROLLER1,088
工作图(方便不懂的理解)

V的工作图
在这里插入图片描述

训练细节在这里就不细讲,原文中描述的很清楚。

RNN工作图
在这里插入图片描述

总结

到这里这篇论文的核心工作思想大致讲完,仅供更好理解论文内容用,如果出现某些很SB的错位,希望包容。(本人还在初步机器学习研究中)
接下来会更新当下优秀的世界模型DreamerV3的论文,还在阅读中。

DreamerV3

DreamerV3作为当今世界最先进的世界模型之一,阅读其论文对于目前的我来说实在有障碍,我也是查了很多资料才能大致理解整个模型。所以本文有错误请多包涵。

前言

DreamerV3是一个非常强大的模型。它仅需单一配置就可在150多项不同任务上超越专用模型,它真正做到了模仿人从零对于环境的学习。而且当时这是唯一的在无任何引导和数据的情况下从零开始在Minecraft中收集钻石的算法。在本文有许多创新点,需要非常多前置知识才能完全理解。

目前许多算法虽然能在交互任务上超越人,但始终只能作用于单一领域,迁移则需要大量新资源和修改。

所以DreamerV3被创造了。
(DreamerV3的基基基础可以参考上文世界模型的解释)

  ~  

结构:

基础大框架结构由三个神经网络组成:
世界模型用来预测潜在动作(当前+动作分布)的结果。
评论家Critic用来评判每个结果的价值。
操纵者Actor选择操作以达到最佳价值的结果。
它们在与环境交互时就能进行并发训练。

  ~  

世界模型

世界模型吃三个东西,一个 z t − 1 z_{t-1} zt1(由Encoder压缩编码),一个循环状态 h t − 1 h_{t-1} ht1(记忆),一个动作 a t − 1 a_{t-1} at1,吐出来的东西作为预测,这里将一一解释。

结构:

  ~  

  • 序列模型: h t = f ϕ ( h t − 1 , z t − 1 , a t − 1 ) h_t=f_ϕ(h_{t-1},z_{t-1},a_{t-1}) ht=fϕ(ht1,zt1,at1)
    RNN的必要,描述记忆状态。

  ~  

  • 编码器: z t z_t zt ~ q ϕ ( z t   ∣   h t , x t ) q_ϕ(z_t \ | \ h_t,x_t) qϕ(zt  ht,xt)
    使用的是卷积神经网络(图像处理的通常做法),对当前环境输入 x t x_t xt进行编码,考虑 h t h_t ht的情况下效果更好(对原来世界模型的提升)
    如果是输入向量用MLP

  ~  

  • 动力学预测器: z ^ t \hat{z}_t z^t ~ p ϕ ( z ^ t   ∣   h t ) p_ϕ( \hat{z}_t \ | \ h_t) pϕ(z^t  ht)
    模型内部的根据记忆对 z t z_t zt进行的预测,使用多层感知机(MLP)

  ~  

  • 奖励预测器: r ^ t \hat{r}_t r^t ~ p ϕ ( r ^ t   ∣   h t , z t ) p_ϕ( \hat{r}_t \ | \ h_t,z_t) pϕ(r^t  ht,zt)
    MLP,模型接受当前状态和记忆状态后预测的奖励(通常是对后续总奖励进行预测,而不是单单这一时刻这一步的奖励)

  ~  

  • 继续预测器: c ^ t \hat{c}_t c^t ~ p ϕ ( r ^ t   ∣   h t , z t ) p_ϕ( \hat{r}_t \ | \ h_t,z_t) pϕ(r^t  ht,zt)
    MLP,预测模型是否继续任务

  ~  

  • 解码器: x ^ t \hat{x}_t x^t ~ p ϕ ( x ^ t   ∣   h t , z t ) p_ϕ( \hat{x}_t \ | \ h_t,z_t) pϕ(x^t  ht,zt)
    解码器,顾名思义,使用的同样是是卷积神经网络(或MLP)。

  ~  
在这里,~代表从分布中随机采样。而采样使用了softmax函数,但这个函数只能前向传递(不可微),所以无法进行反向传播更新梯度,所以在反向时直接跳过这个采样过程,将 1.包含了所有预测信息的,2.各自的采样结果 传过去。

损失函数:

其中,计算loss的公式:
L ( ϕ ) ≜ E ϕ [ ∑ t = 1 T ( β pred L pred ( ϕ ) + β dyn L dyn ( ϕ ) + β rep L rep ( ϕ ) ) ] \mathcal{L}(\phi) \triangleq \mathbb{E}_{\phi} \left[ \sum_{t=1}^{T} (\beta_{\text{pred}} \mathcal{L}_{\text{pred}}(\phi) + \beta_{\text{dyn}} \mathcal{L}_{\text{dyn}}(\phi) + \beta_{\text{rep}} \mathcal{L}_{\text{rep}}(\phi)) \right] L(ϕ)Eϕ[t=1T(βpredLpred(ϕ)+βdynLdyn(ϕ)+βrepLrep(ϕ))]
其中:
  ~  

  • L pred ( ϕ ) \mathcal{L}_{\text{pred}}(\phi) Lpred(ϕ):预测损失,预测项包括图像预测Decoder,奖励预测 r r r
    ,继续预测 c c c,通常用均方误差或者交叉熵

  • L dyn ( ϕ ) \mathcal{L}_{\text{dyn}}(\phi) Ldyn(ϕ):动力学损失,预测项即为动力学预测 z ^ \hat{z} z^,直接用真 z z z z ^ \hat{z} z^的KL散度

  • L rep ( ϕ ) \mathcal{L}_{\text{rep}}(\phi) Lrep(ϕ):表征损失,预测项为编码器Encoder,它在逻辑上与动力学损失有很强的关系。

为了防止表征崩塌(即编码器不做任何事情,仅仅靠将编码方式表示为一个确定的策略,让动力学模型的预测变得无限容易,从而导致损失没有任何意义),而设计这个损失。让编码器和动力学预测器一开始的分布带有位置随机性,然后通过计算双方之间的距离(即损失,所以它们都用KL散度来计算,且公式相差不大),不断对齐彼此分布。这一关系将在下文中详细讲解。

然后它们前面的为损失权重,即 β pred = 1 \beta_{\text{pred}}=1 βpred=1 β dyn = 1 \beta_{\text{dyn}}=1 βdyn=1 β rep = 0.1 \beta_{\text{rep}}=0.1 βrep=0.1

(另外原文中提到,这是对一个序列批次进行计算loss(即公式中的T和t),这有利于模型基于时间进行训练预测,对任务进行展开)

关于动力学损失与表征损失:

对于上文提到的问题,具体做法,研究者使用了停止梯度算子sg和free bits。

  • sg的作用是停止梯度,即切断反向传播,防止在计算动力学损失的时候通过反向传播传到编码器中,同时也在预测表征损失的时候防止传到动力学预测器中。

  • free bits,将表征损失截断在1 nat ≈ 1.44 bits 以下,即小于这个损失时,就不去更新编码器的参数。这带来的好处是,当编码器试图通过特定策略编码时,这个机制会导致编码器的预测不能无限靠近动力学预测,而动力学预测会继续学习并更新,从而使表征损失增加,编码器必须采取另一种策略重新对齐动力学预测,降低了表征崩塌的可能性。

三种loss的计算公式:
L pred ( ϕ ) ≐ − log ⁡ p ϕ ( x t ∣ z t , h t ) − log ⁡ p ϕ ( r t ∣ z t , h t ) − log ⁡ p ϕ ( c t ∣ z t , h t ) \mathcal{L}_{\text{pred}}(\phi) \doteq -\log p_{\phi}(x_t|z_t, h_t) - \log p_{\phi}(r_t|z_t, h_t) - \log p_{\phi}(c_t|z_t, h_t) Lpred(ϕ)logpϕ(xtzt,ht)logpϕ(rtzt,ht)logpϕ(ctzt,ht)

L dyn ( ϕ ) ≐ max ⁡ ( 1 , KL [ sg ( q ϕ ( z t ∣ h t , x t ) ) ∥ p ϕ ( z t ∣ h t ) ] ) \mathcal{L}_{\text{dyn}}(\phi) \doteq \max(1, \text{KL}[\text{sg}(q_{\phi}(z_t|h_t, x_t)) \| p_{\phi}(z_t|h_t)]) Ldyn(ϕ)max(1,KL[sg(qϕ(ztht,xt))pϕ(ztht)])

L rep ( ϕ ) ≐ max ⁡ ( 1 , KL [ q ϕ ( z t ∣ h t , x t ) ∥ sg ( p ϕ ( z t ∣ h t ) ) ] ) \mathcal{L}_{\text{rep}}(\phi) \doteq \max(1, \text{KL}[q_{\phi}(z_t|h_t, x_t) \| \text{sg}(p_{\phi}(z_t|h_t))]) Lrep(ϕ)max(1,KL[qϕ(ztht,xt)sg(pϕ(ztht))])

其中文中还提到对于复杂场景和简单图形的平衡(跨多领域任务),使用 free bits 和 超参数 β rep = 0.1 \beta_{\text{rep}}=0.1 βrep=0.1 的结合可以有效地平衡。

此外,论文提到了KL散度出现的尖峰问题(分布退化,原本是混合分布,退化成了确定性分布,根本原因来源于梯度下降法不断堆砌正确类的分布。这会导致模型探索能力下降,训练崩溃)。文中的解决方法是99%的网络预测+1%均匀分布。

至此,世界模型学习部分告一段落。

  ~  

评论家(Critic)

Critic和下文的Actor都在虚拟环境中学习(世界模型的预测),且与真实环境交互时不会使用前瞻来选择动作(即在部署到实际环节中时,每一次动作采样都是基于当下状态,不会预测后续未来的多步。充分利用世界模型的 h t , z t h_t,z_t ht,zt

Actor的目的是通过学习,选择动作,在每个模型状态下都能采样出最大回报的动作分布。具体回报计算: R t ≜ ∑ τ = 0 ∞ γ τ r t + τ R_t \triangleq \sum_{\tau=0}^{\infty} \gamma^\tau r_{t+\tau} Rtτ=0γτrt+τ,其中折扣因子 γ = 0.997 \gamma=0.997 γ=0.997,这代表Actor最大化回报不能只考虑当前产生的奖励,还必须考虑未来的奖励。

具体的定义: Actor: a t ∼ π θ ( a t ∣ s t ) a_t \sim \pi_\theta(a_t | s_t) atπθ(atst),Critic: v ψ ( R t ∣ s t ) v_\psi(R_t | s_t) vψ(Rtst)
这里 π \pi π为一个动作概率分布函数, v v v为价值概率分布函数,但最终的输出还是转化为数值期望: v t ≜ E [ v ϕ ( ⋅ ∣ s t ) ] v_t \triangleq \mathbb{E}[v_\phi(\cdot|s_t)] vtE[vϕ(st)] s t s_t st h t h_t ht z t z_t zt拼接。

损失函数:

Critic损失函数:

这里loss使用最大似然损失:
L ( ψ ) ≜ − ∑ t = 1 T ln ⁡ p ψ ( R t λ ∣ s t ) \mathcal{L}(\psi) \triangleq -\sum_{t=1}^{T} \ln p_{\psi}(R_t^{\lambda}|s_t) L(ψ)t=1Tlnpψ(Rtλst)

其中, R t λ R_t^{\lambda} Rtλ回报公式
R t λ ≜ r t + γ c t ( ( 1 − λ ) v t + λ R t + 1 λ ) R_t^\lambda \triangleq r_t + \gamma c_t((1 - \lambda)v_t + \lambda R_{t+1}^\lambda) Rtλrt+γct((1λ)vt+λRt+1λ)
(注:这里要分清它不是前面Actor的 R t R_t Rt回报)

解释:当前奖励+未来回报, c t c_t ct直接表示了是否产生未来回报

λ \lambda λ表示了偏差与方差的平衡,通常为0.95。如果越大,表示更依赖于未来回报,偏差会小(轨迹足够长足够准),方差会大(随机性强)。如果越小,表示更依赖于当下估计,偏差大(估计不准),方差小(足够稳定)。

R t λ R_t^\lambda Rtλ递归终止锚点:
R T λ ≜ v T R^{\lambda}_{T} \triangleq v_T RTλvT

对于上式所有的 T T T都等于16。(预测未来16步,这是质量和时间算力成本的权衡)

  ~  

Critic输出:

在这里,考虑到不同任务间回报的量级差异可能很大,且回报分布的输出不一定只有单个模态(类似分布有多个峰这样的),因此这里提出一种将Critic的分布输出变为类别分布,即一个向量,其中每一个数值代表有多少概率落在这个类别上。然后类别通过指数间隔来定义(类似于1 ~ 2,2 ~ 4,4 ~ 8……这种)。

  • 训练时,将 R t λ R^{\lambda}_{t} Rtλ通过前文的计算公式计算后(这里注意,不要把训练时用于计算的 v t v_t vt(当前网络生成)和推理要预测的 v t v_t vt搞混),落在指数分布的哪个类别。然后通过one-hot编码输出一个分布,与Critic输出的分布计算交叉熵来得到loss。

  • 推理时,将Critic的输出与对应类别的中位数进行加权和得到最终回报 v t v_t vt

其中, 为了让Critic有更多从模型的“想象”中学习,又不脱离现实环境,Critic的最终损失计算使用加权loss: L total = β v a l ⋅ L val + β r e p v a l ⋅ L repval \mathcal{L}_{\text{total}} = \beta_{val} \cdot \mathcal{L}_{\text{val}} + \beta_{repval} \cdot \mathcal{L}_{\text{repval}} Ltotal=βvalLval+βrepvalLrepval

在这里, β v a l \beta_{val} βval被设为1.0, β r e p v a l \beta_{repval} βrepval被设为0.3,代表以想象轨迹为主,以回放轨迹为辅。

  • 想象轨迹,即模型在虚拟环境中预测推理,将连续状态组合成一条状态轨迹线。正常计算loss。

  • 回放轨迹,模型从真实交互经验中,提取出初始状态 s 0 s_0 s0,连续交互轨迹中的真实奖励 r t r_t rt和继续标志 c t c_t ct。然后以当前模型对 s 0 s_0 s0初始状态进行想象,得到一条虚拟轨迹,以现实中提取的 r t r_t rt c t c_t ct和当前评论家预测的 v t v_t vt,用前文计算公式计算出这条轨迹的 R t λ R^{\lambda}_{t} Rtλ。最后用这条轨迹的 R t λ R^{\lambda}_t Rtλ,再去计算loss。

由于每个 R t λ R^{\lambda}_{t} Rtλ取取决于Critic预测的 v t v_t vt,加上预测的 v t v_t vt通过训练不断逼近 R t λ R^{\lambda}_{t} Rtλ,一旦 v t v_t vt在某步中预测过大或过小,很容易通过正反馈循环不断放大预测误差,导致一系列的训练问题。(震荡或发散)
于是这里引入了正则化:
使用一个参数缓慢更新的网络 ψ e m a \psi_{ema} ψema ψ e m a ← τ ⋅ ψ e m a + ( 1 − τ ) ⋅ ψ \psi_{ema} \leftarrow \tau \cdot \psi_{ema} + (1 - \tau) \cdot \psi ψemaτψema+(1τ)ψ(其中 τ \tau τ是一个接近1的超参数)
训练时,为了训练Critic的当前网络,让一个轨迹中的 R 0 λ R^{\lambda}_0 R0λ v 0 v_0 v0使用当前网络的(为了更新),轨迹中的 v t v_t vt使用目标网络 ψ e m a \psi_{ema} ψema的(为了稳定)。

对于初始化Critic预测和奖励预测,如果随机化,那么可能导致训练过程不稳定(比如预测很大,梯度很大),可能需要大量时间去消除这种随机带来的影响。所以研究者把权重初始化为0,从一个平缓、无知的状态开始。

  ~  

Actor(演员)

感觉叫操纵者有点奇怪,用一个更贴切幽默的名称 演员吧。

Actor学习能够最大化回报的动作。但为了避免单一策略,加入一个正则化项——熵正则化器(探索程度)。
但不同环境的奖励的量级,频率都会有差异。如果奖励稀疏,则要模型能增加探索能力(增大熵);如果奖励密集则反之。如果奖励很大,则为了防止模型直接产生依赖,也要适当增加探索能力。
同时,这个探索程度(熵)的值不应受不同环境中奖励缩放的影响。(如果另一个环境中奖励比其它的环境中都大十倍,那本质上它们的探索程度也不应受这个缩放的影响)于是还需要一个合适的归一化器。

  ~  

损失函数:

L ( θ ) ≜ − ∑ t = 1 T sg ( ( R t λ − v t ) / max ⁡ ( 1 , S ) ) log ⁡ π θ ( a t ∣ s t ) + η H [ π θ ( a t ∣ s t ) ] \mathcal{L}(\theta) \triangleq -\sum_{t=1}^{T} \text{sg}((R_t^\lambda - v_t)/\max(1, S)) \log \pi_\theta(a_t|s_t) + \eta H[\pi_\theta(a_t|s_t)] L(θ)t=1Tsg((Rtλvt)/max(1,S))logπθ(atst)+ηH[πθ(atst)]
其中
S ≜ EMA ( Per ( R t λ , 95 ) − Per ( R t λ , 5 ) , 0.99 ) S \triangleq \text{EMA}(\text{Per}(R_t^\lambda, 95) - \text{Per}(R_t^\lambda, 5), 0.99) SEMA(Per(Rtλ,95)Per(Rtλ,5),0.99)

  ~  

  • R t λ − v t R^{\lambda}_{t}-v_t Rtλvt
    R t λ R^{\lambda}_{t} Rtλ是Actor的目标值, v t v_t vt是Critic的预测值。这个式子表示的是优势(如果Actor的预测大于Critic的,就增加这个预测的概率。否则反之)

  ~  

  • S S S m a x ( 1 , S ) max(1,S) max(1,S)
    通常归一化时除以的是标准差。但在现在情况下,由于奖励可能稀疏也可能密集,导致这种方法会放大噪声。(比如对于稀疏奖励,基本上所有的 R t λ − v t R^{\lambda}_{t}-v_t Rtλvt都接近0,于是标准差也接近0。这时如果除以标准差,这些接近0的奖励都被放大了,但这却不是真正的模型要的奖励(噪声)。)
    同时,如果使用 最大回报–最小回报 作为分母,如果回报是多模态的(多峰),则可能导致出了最大回报那附件的,其它的都被大幅度缩小,过分低估了其它动作的回报,从而导致学习率下降。(次优收敛)
    所以这里用了两个方法:

    • ①百分位数范围:用回报分布的95%位数 – 5%位数,排除两段的极端值或出现异常奖励,既能找到一个稳定的分母,又能反映有效奖励的分布的情况。

    • ②指数移动平均:使用衰减参数为0.99的前向网络对S进行处理,使得S更加平滑稳定。

    • ③阈值最小限制:通过 m a x ( 1 , S ) max(1,S) max(1,S),当奖励稀疏而S过小时,则使用1作为分母(即不缩放),防止放大噪声。

    m a x ( 1 , σ ) max(1,\sigma) max(1,σ) m a x ( 1 , S ) max(1,S) max(1,S)的区别:
    如果奖励稀疏(偶尔出现一个大奖励),则 σ \sigma σ会因那个大奖励被放大,从而超过1的阈值,导致某些正常的比较 小的奖励被过度缩小,从而导致失效。而 S S S则用了百分位数进行计算(极巧妙的设计),95%–5%能够反映奖励的波动(大的话,那整个奖励就较为稀疏或存在大奖励;小的话,就较为密集且奖励之间相差不大),排除异常,分母稳定。在 加上阈值的限制,这种方法确实能起到很好的作用。

  ~  

  • s g ( ) sg() sg()
    停止梯度算子。在原式中,这个函数将 ( R t λ − v t ) / max ⁡ ( 1 , S ) (R_t^\lambda - v_t)/\max(1, S) (Rtλvt)/max(1,S)设计为了权重,即反映优势的大小,真正需要反向传播的是后面的 π θ ( a t ∣ s t ) \pi_{\theta}(a_t | s_t) πθ(atst)(Actor的预测)。

  ~  

  • η H [ π θ ( a t ∣ s t ) ] \eta H[\pi_\theta(a_t|s_t)] ηH[πθ(atst)]
    即熵正则化器, η = 3 × 10 − 4 \eta=3 \times 10^{-4} η=3×104,可以说上面所有设计都是为了能够在不同环境中使用同一正则项系数而服务的。

  ~  

至此,这种特别设计的loss函数不仅能够在单一任务上保持高效预测和探索,而且对于不同领域内任务也能保持良好的能力。

  ~  
  ~  

鲁棒预测(Robust predictions)

总体的loss公式讲完,接下来讲loss内部计算。
因为是在多环境中,预测量级差异过大,传统的损失计算方法存在一系列问题:直接均方误差(MSE)会过于放大梯度,L1或Huber损失对小值学习率过低,运行统计化(按均值和标准差归一化)则不稳定。

symlogMSE计算:

这是个对确定性的预测使用的方法

在跨环境训练时,即使是确定性目标(如图像像素或物理状态),其数值范围也可能差异巨大:

  • 有的环境像素值在0~255之间,有的则归一化到[-1, 1];
  • 有的环境速度值在m级别,有的则在km级别。

如果直接对这些目标使用 MSE:

  1. 大值目标会产生巨大的梯度,导致训练初期不稳定甚至梯度爆炸;
  2. 小值目标的梯度则过于微小,网络几乎学不到信号。

所以这里引入了 symlog \text{symlog} symlog函数

具体地,一个输入为 x x x,参数为 θ \theta θ的神经网络 f ( x , θ ) f(x,\theta) f(x,θ),学习预测其目标 y y y的变换版本—— symlog(y) \text{symlog(y)} symlog(y),即:
L ( θ ) ≜ 1 2 ( f ( x , θ ) − symlog ( y ) ) 2 \mathcal{L}(\theta) \triangleq \frac{1}{2} (f(x, \theta) - \text{symlog}(y))^2 L(θ)21(f(x,θ)symlog(y))2
当网络学会了如何预测变换版本的目标 s y m l o g ( y ) symlog(y) symlog(y),最后只需要通过逆变换 s y m e x p ( f ( x , θ ) ) symexp(f(x,\theta)) symexp(f(x,θ))得到预测的目标 y ˙ \dot{y} y˙
y ˙ ≜ symexp ( f ( x , θ ) ) \quad \dot{y} \triangleq \text{symexp}(f(x, \theta)) y˙symexp(f(x,θ))
其中:
symlog ( x ) ≜ sign ( x ) log ⁡ ( ∣ x ∣ + 1 ) \text{symlog}(x) \triangleq \text{sign}(x) \log(|x| + 1) symlog(x)sign(x)log(x+1)
symexp ( x ) ≜ sign ( x ) ( exp ⁡ ( ∣ x ∣ ) − 1 ) \text{symexp}(x) \triangleq \text{sign}(x)(\exp(|x|) - 1) symexp(x)sign(x)(exp(x)1)
(这两个函数互为逆函数)

这个函数使得:

  • 大值被压缩:将较大的值压缩到合理范围内,网络可以稳定且快速地学习。
  • 小值几乎不变:函数在 x → 0 x \rightarrow 0 x0时斜率趋近于 1,更好保留数据信息。
  • 正负对称:对大正数和大负数压缩对称,避免了不对称变换在不同环境中表现不稳定的问题。

  ~  

two-hot交叉熵计算:

这是对含有噪声(奖励,价值)和跟需要随机性预测使用的方法

有时候奖励的量级也会存在差异,并且其中包含的噪声有可能被不合适的计算策略给放大。

所以这里使用了 two-hot + s y m e x p \text{two-hot}+symexp two-hot+symexp 计算交叉熵。

  • 首先网络得到输入 x x x与目标 y y y,然后先通过 symlog ( y ) \text{symlog}(y) symlog(y)变换得到 y t r a i n y_{train} ytrain,作为网络要预测逼近的值。

  • 然后网络通过 f ( x , θ ) f(x,\theta) f(x,θ)计算出一个分布向量 logits \text{logits} logits,然后使用交叉熵计算最终损失。(内部softmax)
    由于要将 y t r a i n y_{train} ytrain也变成向量形式以计算交叉熵,这里使用了 symexp two-hot \text{symexp two-hot} symexp two-hot 编码。

具体的,使用一个间隔向量 b = [ − 20 , − 19 , … , 19 , 20 ] b=[-20,-19,…,19,20] b=[20,19,,19,20]和指数桶向量 B = symexp ( b ) B=\text{symexp}(b) B=symexp(b),然后网络最终的训练目标就是通过 softmax \text{softmax} softmax网络输出的 logits \text{logits} logits B B B进行加权平均: y ^   = ˙  softmax ( f ( x , θ ) ) T B \hat{y} ~ \dot{ =} ~ \text{softmax}(f(x,\theta))^TB y^ =˙ softmax(f(x,θ))TB,让 y ^ \hat{y} y^去逼近 y y y

但要使用交叉熵来计算loss,于是将这种变换方式换一个角度,让 y t r a i n y_{train} ytrain去生成一个编码向量,与网络生成的 logits \text{logits} logits就可以计算交叉熵。

这个编码过程使用了 two-hot \text{two-hot} two-hot编码,将 y t r a i n y_{train} ytrain通过 B B B指数间隔桶向量上寻找位置: b i ≤ y t r a i n ≤ b i + 1 b_i \leq y_{train} \leq b_{i+1} biytrainbi+1,然后通过 w = ( y t r a n s − b i ) / ( b i + 1 − b i ) w = (y_{trans} - b_i) / (b_{i+1} - b_i) w=(ytransbi)/(bi+1bi) e i = 1 − w e_i=1-w ei=1w e i + 1 = w e_{i+1}=w ei+1=w,其它位置 e i = 0 e_i=0 ei=0,得到编码向量 e e e。然后就可以去和 logits \text{logits} logits计算交叉熵了。

这里的写法可能有点绕,可能原因是弄不清这里如何实际计算loss,和如何理解网络要预测的值,这两个事情混在一起。总而言之,就是计算loss时,是要通过交叉熵来计算的,所以要用网络输出logits(未softmax向量)和 y t r a i n y_{train} ytrain通过编码得到的向量。而如何理解,则通过将网络的logits通过softmax和symexp操作映射回实际目标y的范围中去逼近它。

这个计算方法:

  • 将回归问题转化为分类问题,分类问题更适合对噪声较多,更需要稳定的计算。

  • 网络更适合预测多分布,更加具有探索和随机能力。

  • 指数桶的合理压缩与信息保留,让小的奖励能不丢失,也能让大的奖励得到合理压缩。

  ~  

消融实验

文中通过消融实验测试每个算法部分,对模型的整体效果的影响,其中:

  • 鲁棒性技术:(对模型性能而言)

    • KL平衡 + free bit
    • 归一化回报
    • symexp + two-hot编码
  • 学习信号:

    • 任务相关信号
    • 任务无关信号

总体来看,对于鲁棒性技术,KL平衡 + free bit 的效果最好,且对于学习模式,模型更倾向于无监督学习。(这在某种程度上意味着模型真正地在理解感知世界)

模型评估

  • 文中具体列出了在8个不同领域(超过150个任务)中的通用性,并与其中最佳方法进行比较。

  • 进一步研究者将这模型与PPO(一种高鲁棒性算法)的高质量实现进行比较,且在ProGen上复现了其强劲结果。

  • 此外,还将此模型用于Minecraft进行测试,展现了惊人的能力。

(而且这个模型是在A100上训练的,不用大型算力资源)

原文中直白提到了具体在8个领域内测试的结果,并且附赠了图标。这里不再复述,可以知道的是它的能力超越了几乎所有优秀算法。

(先搁这里,基本大致的模型讲完了,后续的论文部分如果有时间和精力的话应该会写)

更多推荐