Transformers 权威指南——强化学习 Transformer

0 阅读6分钟

强化学习(RL)是一种基于“通过环境反馈来做序列决策”的范式。你仔细想想就会发现,通过交互来学习其实并不是什么新概念。事实上,你这一生都一直在接触 RL 的基本原则。因为很多时候,你就是通过和周围环境互动来学习的,而不需要一个显式的老师。比如小时候,你可能会摸各种东西,或者把物体放进嘴里,看看环境会如何反应。根据结果——比如被家里的猫抓了之后哭起来——你就学会了下一次该怎么做。通过这样的经历,你逐渐形成了未来行动的“策略(policy)”,比如以后不再去拽猫尾巴。

这个学习过程,其实和一个 RL agent 的工作方式非常相似。因此,本章会从你这种天然的“通过环境学习”的直觉出发,进一步引入强化学习中的关键概念。具体来说,我会先给出传统 RL 算法的入门概览,包括在线与离线学习的区别、基于模型与无模型方法的区别,以及 on-policy 与 off-policy RL 的区别。

当你对这些基础概念有了基本理解之后,我们会进一步探索近些年 RL Transformer 中的一些代表性架构,例如 Decision Transformer、Online Decision Transformer,以及高效的基于随机 Transformer 的世界模型。

强化学习入门

通常来说,agent 与环境之间的交互,会被放进马尔可夫决策过程(MDP)的框架中来理解。MDP 是一种数学框架,用来描述“在结果既部分随机、又部分由决策者控制的环境中如何做决策”,其核心组成包括状态、动作、转移概率和奖励。它的运行方式如下:agent 会根据环境当前状态 (s \in S),选择一个动作 (a \in A)。之后,agent 会收到环境反馈,也就是奖励 (r \in R)。接着,它所处的状态会更新到一个新状态 ((s' \in S))。图 7-1 展示了这种交互。

image.png

图 7-1. agent 与环境交互的流程图。

从数学上说,可以写成:

Rssa=E{rt+1st=s,;at=a,;st+1=s}R_{ss'}^a = E\{r_{t+1} \mid s_t = s,; a_t = a,; s_{t+1} = s'\}

马尔可夫过程

如果一个过程满足:未来状态只依赖当前状态和当前动作,而不依赖更早之前发生的事件序列,那么这个过程就称为马尔可夫的。这一性质被称为马尔可夫性质(Markov property)。

策略 (π\pi) 会把每个状态 (sSs \in S) 和动作 (aAa \in A) 映射到一个概率 (π(s,a)\pi(s,a)),这个概率表示:当 agent 处于某个状态 ss 时,它采取动作 aa 的概率。策略本质上是一条规则,用来指导 agent 在给定状态下如何选择动作。也就是说,它是 agent 根据当前情境来决定下一步行动的策略。

大多数 RL 算法都建立在价值函数(value function)的思想上。价值函数用于评估:对于 agent 来说,处在某个状态、或者在某个状态下采取某个动作,到底有多大潜在收益。借助价值函数,agent 可以量化未来可能获得的奖励,更准确地说,是量化期望回报(expected return)。当然,agent 未来能获得什么奖励,取决于它会选择什么动作。因此,价值函数的计算总是和 agent 所遵循的具体策略绑定在一起的。

更技术化地说,我们通常有两个价值函数:Vπ(s)V^\pi(s)(某个策略下状态的价值)和 Qπ(s,a)Q^\pi(s,a)(某个策略下在某状态采取某动作的价值)。前者叫 value function,后者叫 action-value function。策略、状态和动作之间的这些关系,是 agent 学习过程的基础。

价值函数是其 Bellman 方程的唯一解

价值函数有一个关键特性:它们满足特定的递归关系。Bellman 方程正是用来说明这一点的。它通过“未来奖励的期望值”来定义一个策略下的价值函数。本质上,它把当前状态的价值,和未来可能到达状态的价值连接了起来,像是在当前状态向未来“看一步”。Bellman 方程会考虑所有可能结果,并根据它们发生的概率进行加权,因此能大幅简化状态价值或状态—动作价值的计算。

强化学习的目标,是找到最优策略 π\pi^*,使得在时间步序列 tTt \in T 中,对于每个状态 ss,期望累计奖励在整个时间范围 TT 内达到最大。换句话说,在 RL 中,你要解决两个不同的学习目标:一个是学习如何表示状态,另一个是学习如何采取动作。学习状态表示,意味着要理解并编码环境当前的情境,使其适合决策;而学习如何采取动作,则是指基于当前状态表示,决定应该执行什么动作,通常目标是最大化长期累计奖励。下一节里,你会看到如何做到这一点。

强化学习中的基础概念

在这一节里,你将看到不同的方法:它们分别如何教会 RL 算法去编码环境,以及在此基础上决定该采取哪些动作。

在线与离线强化学习

强化学习通常意味着:agent 反复与某个具体环境进行交互,使用最近更新的策略收集经验,再直接利用这些经验来改进策略。这种方式被称为在线学习(online learning)。但对某些应用——例如医疗或自动驾驶——来说,这种方法并不现实。

在离线强化学习(也叫 batch RL)中,agent 会依赖一个静态数据集来训练,而在训练过程中不再和环境交互。这意味着,训练期间数据集本身不会变化。等策略完全在这个数据集上训练好之后,再部署出去。它的问题在于:agent 无法对新信息做适应。图 7-2 展示了这两种方式的差异。

image.png

图 7-2. 在线强化学习(右)与离线强化学习(左)的区别。

基于模型与无模型方法

在基于模型的 RL 中,agent 会根据和环境交互得到的转移样本 (s,a,r,ss,a,r,s') 去学习一个转移函数。这个函数描述的是:在给定当前状态 ss 和动作 aa 的条件下,下一个状态 ss' 以及奖励 rr 的概率分布 p(s,rs,a)p(s', r \mid s, a)。然后 agent 会基于这个环境模型进行规划,并选择能够最大化期望回报的动作。虽然这种方法通常计算开销更高,但相比其他方法,它往往能用更少的环境交互获得不错的表现。

在无模型 RL 中,agent 不会显式去建模状态转移动力学,而是直接通过与环境交互学习最优动作。这会让学习速度相对更慢;不过,无模型 RL 在适应环境变化方面通常更强,因此在复杂或有噪声的环境中更鲁棒。它的计算需求也更低,因为不需要去学习环境模型。所谓环境模型,就是对环境动态状态及其如何通向奖励的机制进行建模。

On-policy 与 Off-policy 强化学习

为了保证 agent 能持续地、无限次地尝试所有动作,通常有两种策略,这就是所谓的 on-policy 和 off-policy 方法。

On-policy 关注的是:评估或改进 agent 当前正在使用的那个策略。也就是说,agent 一边学习,一边跟随它正在尝试优化的那个策略。

而 off-policy 则是在评估或改进一个“不同于数据收集策略”的目标策略。agent 会把价值函数学习过程与其实际执行的策略分离开来,因此它可以利用由另一种策略生成的数据。这种灵活性意味着:agent 可以从自己以前的经验、甚至其他 agent 的经验中学习。此外,agent 还可以在执行探索性或随机动作的同时,逐步学到最优策略。

例如,在 SARSA 算法中——它是一个 on-policy 方法——采集的转移数据是五元组 (s,a,r,s,a)(s, a, r, s', a'),它包含当前状态—动作对 (s,a)(s,a)、即时奖励 rr,以及下一个状态—动作对 (s,a)(s',a')。这也是这个算法叫 SARSA 的原因。这个五元组被用于估计当前状态—动作对和下一状态—动作对的回报。

尽管 on-policy 方法通常比较容易实现,但它们也有一些明显缺点。首先,它们往往样本效率较低,也就是需要大量环境交互才能学到一个不错的策略。其次,它们容易出现策略振荡(policy oscillation,也叫 policy chattering),即策略反复改变方向,却无法收敛到最优策略;也可能出现 alternating policies,也就是策略在不同方案之间来回切换。所谓不稳定性,指的是策略无法持续改进或稳定下来。这些问题会妨碍探索,并拖慢学习过程,最终往往导致次优策略。在像金融市场或复杂仿真这样的动态环境中,这会带来很大挑战。

Off-policy 学习则使用一个目标策略,但它的评估依据是由另一个行为策略(behavior policy)收集到的经验,而不是直接按目标策略本身采样。Q-learning 就是一个 off-policy 方法,也是 RL 历史上的一个重要突破。这里,学到的动作价值函数 QQ 会直接逼近最优动作价值函数,而不依赖于当前实际执行的是哪一个策略。不过,实际执行的策略仍然有影响,因为它决定了哪些状态—动作对会被访问并更新。

因此,off-policy 学习中会涉及两种策略:一种是行为策略,它负责数据收集,并允许探索更多状态和动作;另一种是目标策略,通常是基于当前动作价值函数估计得到的确定性贪心策略。正是这种分工,使 off-policy 方法能够在“并没有真正执行目标策略”的情况下,仍然评估某些动作和状态的价值。

Off-policy 方法的主要优点,是它能够从多样化经验中学习,因此有可能加速学习过程。不过,它也通常比 on-policy 更复杂;而且 agent 还可能会遇到数据分布带来的问题,例如过拟合到某些特殊经验。

时序差分学习(Temporal Difference Learning)

时序差分学习(TD learning)是 RL 中一个非常重要的方法,它允许 agent 直接从原始经验中学习,而不需要知道环境动力学模型。TD 学习融合了 Monte Carlo 方法和动态规划(DP)的思想。与必须等到整个 episode 结束后才能更新价值估计的 Monte Carlo 方法不同,TD 学习可以在每一步都做增量更新。而与需要完整环境模型的动态规划不同,TD 学习依赖 bootstrapping——也就是基于已有估计去更新新的估计,而不需要完整环境知识。

TD 学习之所以重要,是因为它允许在复杂环境中持续更新和持续学习。它是很多实用 RL 算法的基础,比如 SARSA 和 Q-learning,这些算法都通过 TD 更新来不断改进自己的策略。

强化学习中的世界模型

作为人类,我们会根据感官收集到的信息,在脑中构建一个关于世界的内部模型。我们的决策和行动,正是由这种内部表示所引导的。不过在日常生活中,我们接收到的信息太多了,所以我们会自动对其中的关键信息做抽象。这一点和 RL 算法学习世界模型的方式很像。也就是说,在 RL 里,重点不只是“具备感知知识”,而是如何利用这种感知知识来做决策。

RL 中的世界模型,是指 agent 学到的一种内部环境表示。它会被用来根据当前状态和动作,预测未来状态和奖励。这个概念在基于模型的 RL 中尤其重要,因为 agent 可以借助这个世界模型,在不持续和真实环境交互的情况下做规划与决策。

世界模型之所以有价值,是因为它允许 agent 在内部模拟各种场景,在真正执行动作之前,先评估潜在后果。这种模拟会带来更高效的学习,因为 agent 可以在不承担真实世界交互成本或风险的情况下,探索大量可能性。它可以自己生成假设情境,并基于这些内部模拟来优化策略。世界模型已经在多个应用中获得成功,包括机器人和自动驾驶,在这些场景中,它能帮助完成规划与决策。

如果你想更深入理解经典强化学习

如果你想系统深入地理解强化学习,Richard S. Sutton 和 Andrew G. Barto 所著的 Reinforcement Learning: An Introduction(MIT Press)是公认的经典资源。它对强化学习的理论与实践都提供了非常扎实的基础。

强化学习中的 Transformer

Transformer 在 NLP 等领域取得巨大成功之后,被引入 RL 的速度相对较慢,这主要是因为 RL 本身有一些特殊挑战。最开始,很多基于 Transformer 的 RL 工作主要集中在状态表示学习和记忆增强上。但这些方法依然会面临传统 RL 框架中的一些问题,比如 TD 学习和策略优化中的困难。另一个显著挑战是训练过程中的非平稳性(non-stationarity),也就是环境的动态会随时间改变,从而导致一致学习和稳定优化都变得困难。

不过,Transformer 在 RL 里的一个关键优势,是它擅长建模时间依赖关系。在很多 RL 场景中,agent 需要结合一长段过去的观测或动作,才能在当前做出正确决策。借助自注意力机制,Transformer 在处理这类长程依赖时表现得非常强。

Transformer 在多智能体强化学习中也很有帮助。在这类场景下,理解多个 agent 的动作以及它们如何共同影响环境非常关键。Transformer 可以同时处理并上下文化所有 agent 的动作和状态,从而形成一个更完整的环境表示。

这种能力也同样适用于层次化强化学习(hierarchical RL):Transformer 可以帮助 agent 把策略分解为高层战略与低层具体动作,因此能在不同抽象层次上进行学习。

在基于模型的强化学习中,Transformer 还可以帮助 agent 学习环境模型,从而提升规划能力。它们可以把一串过去的状态和动作作为输入,预测未来状态。此外,Transformer 在 imitation learning 和 inverse RL 这类“从示范中学习”的设置下也很有用,因为它们可以学习把状态和动作序列映射到专家策略上。

Transformer 能处理多种类型数据,这也使它们适合作为 AI 系统中的“通才型 agent”,可以在极少任务特定改动的情况下执行多种任务。不过,尽管前景很诱人,Transformer 在实时 RL 场景中仍面临挑战,主要原因是它们参数量大、计算开销高,而这对实时应用来说并不理想。

然而,当 Transformer 开始结合离线 RL 方法时,一个关键转折点就出现了。通过在预先收集好的静态数据上训练 Transformer,研究者开始获得最先进的成果。接下来我们要讲的 Decision Transformer,就是这一方向中的先驱之一。它把 RL 建模成一个自回归的序列生成任务,通过预测轨迹中的下一个元素来生成目标轨迹。

Decision Transformer

Decision Transformer(DT)把 RL 重新解释成一个条件序列建模任务。它使用自回归模型,将目标奖励与过去的状态、动作结合起来进行建模,从而摆脱了像 TD learning 这样的传统 RL 方法。这个模型依赖的是一个固定且有限的数据集,其中包含由不同策略 rollout 出来的轨迹。正如你前面学到的,这种方式符合离线 RL 和无模型 RL 的原则。也就是说,这种策略性的转变利用已有数据,绕开了传统 RL 中的一些常见难题,从而使学习过程更稳定、更有效。其中一个典型问题就是 bootstrapping error,也就是 agent 基于其他估计值而不是真实完整结果来更新策略估计时产生的误差。TD 方法本质上就是在做 bootstrapping,而不是像 Monte Carlo 那样依赖完整回报。

DT 使用的是一种自回归(从左到右)的方式,类似 GPT 架构。它建立在概率推断和序列建模的原则之上,也就是说,每个状态通向一个动作,而每个动作又通向一个后续状态与观测。此外,DT 会按顺序处理整条轨迹,在给定过去元素的条件下预测后续元素。正是这种自回归建模方式,使得该过程不再是马尔可夫的。这也意味着,模型可以利用更丰富的历史信息来做决策,而不再局限于只看当前状态这一马尔可夫假设。

轨迹表示 τ\tau由状态 ss、动作 aa 和 returns-to-go

R^t=t=tTrt\hat{R}_t = \sum_{t' = t}^{T} r_{t'}

组成,可以写成:

τ=(R^1,s1,a1,R^2,s2,a2,,R^T,sT,aT)\tau = (\hat{R}_1, s_1, a_1, \hat{R}_2, s_2, a_2, \dots, \hat{R}_T, s_T, a_T)

这里的 returns-to-go 表示的是未来期望获得的回报,而不是过去已经得到的奖励。模型的输出则是一串预测动作:

a^1,a^2,,a^T\hat{a}_1, \hat{a}_2, \cdots, \hat{a}_T

正是这种严格的自回归序列建模方式,让模型能够在决策时使用更加丰富的历史信息。

由于 return、state 和 action 都会先经过各自特定模态的线性嵌入,再叠加一个位置时间步编码,因此你最终会得到三种不同类型的 token。这和只有文本 token 的 LLM 形成鲜明对比。示例 7-1 展示了如何构造这些 embedding。

示例 7-1. Decision Transformer 的嵌入

embed_timestep = nn.Embedding(max_ep_len, hidden_size) 
embed_return = torch.nn.Linear(1, hidden_size)
embed_state = torch.nn.Linear(self.state_dim, hidden_size)
embed_action = torch.nn.Linear(self.act_dim, hidden_size)

state_embeddings = embed_state(states) 
action_embeddings = embed_action(actions)
returns_embeddings = embed_return(returns_to_go)
time_embeddings = embed_timestep(timesteps)

state_embeddings = state_embeddings + time_embeddings 
action_embeddings = action_embeddings + time_embeddings
returns_embeddings = returns_embeddings + time_embeddings
  • 每个 token 都有自己的 embedding。
  • 每一种模态都由不同的 head 来嵌入。
  • 时间 embedding 的处理方式类似于位置 embedding。

图 7-3 展示了 Decision Transformer 的整体架构。

image.png

图 7-3. Return、state 和 action 会经过各自特定模态的线性 embedding,并额外叠加一个位置时间步编码。之后,这些 token 被送入一个 GPT 架构,并通过 masked self-attention 以自回归方式预测动作。图片改编自 Lili Chen 等人(2021)。

为什么在 Decision Transformer 里位置嵌入是“加”上去,而不是“拼接”上去?

你可能会好奇:为什么位置嵌入是和 token embedding 相加,而不是拼接?答案在于:相加能够把位置信息直接融合进原始 embedding 中,同时保持向量维度不变,从而不破坏模型整体结构。这样一来,位置编码就能以一种后续网络层可以自然处理的方式嵌入进去。在本书 GitHub 仓库中,ch_03_add_vs_cat_embeddings.ipynb 这个 notebook 也通过代码对这一点做了进一步演示。

为了亲自试一试 DT,你将使用 Gymnasium 中两个很经典的 RL 环境:HalfCheetah 和 Walker2D。Gymnasium 是一个提供各种单智能体强化学习环境 API 的项目。HalfCheetah(图 7-4 左)是一个二维机器人,它有 9 个身体部件和 8 个关节,其中包括两只爪子。目标是对关节施加扭矩,让这只“猎豹”尽可能快地向前奔跑:向前运动会得到正奖励,向后运动则得到负奖励。可以施加扭矩的地方是大腿、小腿和脚上的 6 个关节,而躯干和头部是固定的。

Walker2D 则是在 Hopper 环境基础上再多加了一条腿,从而让机器人不再是“跳着走”,而是真正地“走路”,任务因此更复杂。Hopper 是一个二维单腿角色,有四个主要身体部件:躯干、大腿、小腿和一只脚。它的目标是通过对连接这些部件的三个关节施加扭矩,向前跳跃。而 Walker 则是一个二维双腿角色,拥有七个主要身体部件:一个躯干、两条大腿、两条小腿和两只脚。它的目标是通过对六个关节施加扭矩,让它向前行走。图 7-4 展示了这几个环境。

image.png

图 7-4. 左边是 HalfCheetah,中间是 Walker2D,右边是 Hopper 环境。

接下来,为了在这些环境里使用 Decision Transformer,你需要先计算动作。示例 7-2 展示了这一过程。

示例 7-2. 计算动作

def get_action(model, states, actions, rewards, returns_to_go, timesteps):
    max_length = model.config.max_length
    state_dim = model.config.state_dim
    action_dim = model.config.act_dim

    states = states.reshape(1, -1, state_dim)[:, -max_length:]
    actions = actions.reshape(1, -1, action_dim)[:, -max_length:]
    returns_to_go = returns_to_go.reshape(1, -1, 1)[:, -max_length:]
    timesteps = timesteps.reshape(1, -1)[:, -max_length:]
    padding = max_length - states.shape[1]

    attention_mask = torch.cat([torch.zeros(padding),
                    torch.ones(states.shape[1])],
                    dim=0).to(dtype=torch.long).reshape(1, -1) 
    pad_tensor = lambda x, dim: torch.cat([torch.zeros((1,
                 padding, dim)), x], dim=1).float()

    states = pad_tensor(states, state_dim)
    actions = pad_tensor(actions, action_dim)
    returns_to_go = pad_tensor(returns_to_go, 1)
    timesteps = torch.cat([torch.zeros((1, padding),
                            dtype=torch.long), timesteps], dim=1)

    state_preds, action_preds, return_preds = model(
        states=states,
        actions=actions,
        rewards=rewards,
        returns_to_go=returns_to_go,
        timesteps=timesteps,
        attention_mask=attention_mask,
        return_dict=False,
    )

    return action_preds[0, -1]

把所有 token pad 到统一序列长度。

然后,你还需要一个和环境交互的函数,示例 7-3 给出了实现。

示例 7-3. 与环境交互

def interact_with_environment(env, model, state_mean, state_std, max_ep_len,
                            target_return_value, state_dim, act_dim, scale, device):
    episode_return, episode_length = 0, 0
    state = env.reset()
    target_return = torch.tensor([target_return_value], device=device,
                    dtype=torch.float32).reshape(1, 1)
    states = torch.from_numpy(state).reshape(1, state_dim).to(
                device=device, dtype=torch.float32)
    actions = torch.zeros((0, act_dim), device=device, dtype=torch.float32)
    rewards = torch.zeros(0, device=device, dtype=torch.float32)
    timesteps = torch.tensor([[0]], device=device, dtype=torch.long)

    for t in range(max_ep_len):
        actions = torch.cat([actions,
                    torch.zeros((1, act_dim), device=device)], dim=0)
        rewards = torch.cat([rewards, torch.zeros(1, device=device)], dim=0)

        normalized_states = (states - state_mean) / state_std
        action = get_action(model, normalized_states, actions,
                            rewards, target_return, timesteps)
        actions[-1] = action

        state, reward, done, _ = env.step(action.detach().cpu().numpy())
        cur_state = torch.from_numpy(state).reshape(1, state_dim).to(device)
        states = torch.cat([states, cur_state.to(device)], dim=0)
        rewards[-1] = reward

        pred_return = target_return[0, -1] - (reward / scale)
        target_return = torch.cat([target_return,
                        pred_return.reshape(1, 1).to(device)], dim=1)
        timesteps = torch.cat([timesteps, torch.tensor([[t + 1]],
                    device=device, dtype=torch.long)], dim=1)

        episode_return += reward
        episode_length += 1

        if done:
            break

    return episode_return, episode_length

为了把你选择的环境传进上面的函数,你只需要像示例 7-4 那样,在配套 notebook 里运行一个选择函数。

示例 7-4. 选择环境

user_input = input("""Enter the environment name
                    (e.g., 1 for 'HalfCheetah-v3' or
                    2 for 'Walker2d-v3'): """)

随后,示例 7-5 会根据你的选择,在 Gymnasium 中创建环境。

示例 7-5. 创建环境

if user_input in env_options:
    env_name = env_options[user_input]
else:
    print("Invalid input. Please enter 1 or 2.")


env = gym.make(env_name)
env = Recorder(env, directory, fps=30)

env, model, state_mean, state_std, state_dim, act_dim = setup_environment(env_name)

接下来,你就可以像示例 7-6 那样,调用环境并运行 play 函数。

示例 7-6. 从环境交互结果生成视频

episode_return, episode_length = interact_with_environment(env,
                                model, state_mean, state_std,
                                max_ep_len, TARGET_RETURN,
                                state_dim, act_dim, scale, device)

env.play()

播放交互后的环境过程。

这样你就会看到一段短视频,展示 Decision Transformer 在你所选环境中的表现。

离线 RL——就像 DT 这种方式——允许策略只利用预收集数据来训练,而不需要在真实环境中探索。这能显著降低训练过程中的事故风险或不安全行为,对于机器人、自动驾驶这类真实世界应用尤为重要。不过,这种方法也会把 agent 的训练上限限制在离线数据集本身的质量上。下一节里,你会看到如何通过在线交互,让 agent 针对具体任务进一步微调。

走向在线:Online Decision Transformer

Online Decision Transformer(ODT)把离线策略训练与在线微调结合到了一起。

在某些场景里,在线强化学习的数据采集既昂贵又危险;但另一方面,在线 RL 对于适应环境变化、处理离线数据中没有覆盖的意外情境又非常关键。ODT 对 Decision Transformer 的训练方式做了增强,这有点像语言和视觉领域里常见的做法:只用极少量在线数据,也能显著提升模型表现。

为此,ODT 对 DT 的架构做了扩展:它把原来的确定性策略改成了随机策略,也就是动作会基于最近 KK 个状态和 returns-to-go 来进行采样,而不是直接输出唯一动作。除此之外,它还引入了最大熵序列建模(max-entropy sequence modeling),也就是试图最大化预测序列的熵,从而促使生成结果具备更多探索性和多样性。在这种设置下,策略熵依赖于轨迹分布 τ\tau:在离线预训练阶段它是静态的,而在微调阶段,由于会不断引入在线探索得到的新数据,因此它会变成动态的。ODT 和经典最大熵 RL 算法之间的一个关键区别在于:后者通常关注通过“每一步动作的期望熵”来最大化回报,而 ODT 完全依赖于对动作序列的监督学习。同时,它的熵定义是在“序列层面”而不是“单步转移层面”上的。

它的训练管线还引入了 replay buffer,用来存储过去的转移轨迹,并在训练中随机采样复用。这能提升学习效率和稳定性,让 agent 能从更多样的数据中学习,也有助于提高策略的泛化能力。

在 ODT 中,replay buffer 存储的不是单步 transition,而是一整条 trajectory。在离线预训练结束后,replay buffer Treplay\mathcal{T} _{replay} 会用离线数据集 Toffline\mathcal{T}_{offline} 中回报最高的前 NN 条轨迹来初始化。然后,算法会迭代执行 RR 轮,或者直到收敛为止。在每一轮中,会使用当前策略 πθ\pi_\theta 和探索性 returns-to-go 收集新轨迹 τ\tau。随后,replay buffer 中最旧的一条轨迹会被新的 τ\tau 替换掉,也就是用先进先出(FIFO)的方式更新。之后,策略 πθ\pi_\theta 会在更新后的 replay buffer Treplay\mathcal{T}_{replay} 上执行 II 次 ODT 训练。示例 7-7 展示了论文中的原始实现。

示例 7-7. Online Decision Transformer 中的 replay buffer

class ReplayBuffer(object):
    def __init__(self, capacity, trajectories=[]):
        self.capacity = capacity
        if len(trajectories) <= self.capacity:
            self.trajectories = trajectories
        else:
            returns = [traj["rewards"].sum() for traj in trajectories]
            sorted_inds = np.argsort(returns)  # lowest to highest
            self.trajectories = [
                trajectories[ii] for ii in sorted_inds[-self.capacity :]
            ]

        self.start_idx = 0

    def __len__(self):
        return len(self.trajectories)

    def add_new_trajs(self, new_trajs):
        if len(self.trajectories) < self.capacity:
            self.trajectories.extend(new_trajs)
            self.trajectories = self.trajectories[-self.capacity :]
        else:
            self.trajectories[
                self.start_idx : self.start_idx + len(new_trajs)
            ] = new_trajs
            self.start_idx = (self.start_idx + len(new_trajs)) % self.capacity

        assert len(self.trajectories) <= self.capacity

为了更直观理解 replay buffer 的工作机制,你可以运行示例 7-8 中的代码。

示例 7-8. Replay buffer 的内部工作方式

trajectories = [    {"rewards": np.array([1, 2, 3]),
    "states": np.array([0, 1, 2]), "actions": np.array([0, 1, 2])},
    {"rewards": np.array([2, 3, 4]),
    "states": np.array([1, 2, 3]), "actions": np.array([1, 2, 3])},
    {"rewards": np.array([3, 4, 5]),
    "states": np.array([2, 3, 4]), "actions": np.array([2, 3, 4])},
    {"rewards": np.array([4, 5, 6]),
    "states": np.array([3, 4, 5]), "actions": np.array([3, 4, 5])}
] 

replay_buffer = ReplayBuffer(capacity=3, trajectories=trajectories[:2]) 

print("Initial replay buffer:")
for traj in replay_buffer.trajectories:
    print(traj)

new_trajectories = [
    {"rewards": np.array([5, 6, 7]),
    "states": np.array([4, 5, 6]), "actions": np.array([4, 5, 6])},
    {"rewards": np.array([6, 7, 8]),
    "states": np.array([5, 6, 7]), "actions": np.array([5, 6, 7])}
] 

replay_buffer.add_new_trajs(new_trajectories) 

print("\nReplay buffer after adding new trajectories:")
for traj in replay_buffer.trajectories:
    print(traj)
  • 创建示例轨迹数据。
  • 初始化容量为 3 的 replay buffer。
  • 构造新的轨迹。
  • 把新轨迹加入 buffer。

输出如下:

Initial replay buffer:
{'rewards': array([1, 2, 3]), 'states': array([0, 1, 2]), 'actions': array([0, 1, 2])}
{'rewards': array([2, 3, 4]), 'states': array([1, 2, 3]), 'actions': array([1, 2, 3])}

Replay buffer after adding new trajectories:
{'rewards': array([2, 3, 4]), 'states': array([1, 2, 3]), 'actions': array([1, 2, 3])}
{'rewards': array([5, 6, 7]), 'states': array([4, 5, 6]), 'actions': array([4, 5, 6])}
{'rewards': array([6, 7, 8]), 'states': array([5, 6, 7]), 'actions': array([5, 6, 7])}

ODT 训练管线中另一个重要补充,是 hindsight return relabeling(HER)。这是一种用来提升目标条件 agent 样本效率的技术。它的核心思想是:把 agent 轨迹中的“目标”改标成它真正达成的目标,而不是最初设定的目标。对 ODT 来说,策略学习是基于初始 returns-to-go 的;但在 rollout 中实际获得的回报,以及由此诱导出的 returns-to-go,可能与原来设想的 returns-to-go 不一致。受 HER 启发,ODT 会把 rollout 轨迹中的 returns-to-go token 重标成实际 achieved return,从而确保最后一个时间步上的 returns-to-go token 与 agent 最终获得的奖励一致。这种 return relabeling 策略既适用于稀疏奖励环境,也适用于稠密奖励环境。

一个全新的世界:基于随机 Transformer 的世界模型

Efficient stochastic transformer-based world models for reinforcement learning(STORM)是一种基于模型的强化学习算法。也就是说,它内部显式包含了一个世界模型。不过,与其他基于模型的 RL 算法相比,它在实时交互上更高效,并且刷新了最先进方法的基准表现。

在世界模型中做“想象(imagining)”本质上是一个自回归过程,因此误差会随着时间不断积累。原因在于,被想象出来的轨迹和真实轨迹之间可能存在偏差。这样一来,agent 就可能无意中去追逐“虚拟目标”,最终导致在真实环境中表现很差。为了解决这个问题,STORM 引入了变分自编码器(VAE)。VAE 能自动学习高维数据的低维 latent 表示,并在 latent 空间中引入随机噪声。图 7-5 展示了 VAE 的概览。

image.png

图 7-5. 变分自编码器的架构概览。

这种在世界模型中加入噪声的方式,可以帮助 agent 更好地泛化,并且对环境变化更鲁棒,从而提升它在真实世界场景中的表现。

过去大多数基于模型的深度强化学习(DRL)方法都采用循环神经网络(RNN)作为序列模型。DRL 把深度学习和强化学习结合起来,用于处理高维输入。但正如你已经知道的,RNN 的递归结构会阻碍并行计算,因此训练速度较慢。而且,RNN 也一直难以处理长程依赖。因此,使用 Transformer 来充当序列模型,就是一种很自然的选择:它正好能同时解决这两个问题。

STORM 的世界模型从环境中生成图像观测 oto_t。但如果直接在原始图像上建模环境动力学,会带来巨大的计算成本,并且容易出错。为了解决这个问题,VAE 会把图像转换成 latent stochastic categorical distributions ZtZ_t。这个分布由 32 个类别组成,而每个类别又有 32 个类。编码器 qϕq_\phi 和解码器 pϕp_\phi 都是卷积神经网络(CNN)。之后,会从 ZtZ_t 中采样一个 latent variable ztz_t 来表示 oto_t。为了在采样过程中保留梯度,STORM 使用了 straight-through gradients trick,也就是在反向传播时把离散变量当成连续变量处理,从而使涉及随机采样的模型仍能高效训练。

接下来,图像编码器

ztqϕ(ztot)=Ztz_t \sim q_\phi(z_t \mid o_t)=\mathcal{Z}_t

和解码器

o^t=pϕ(zt)\hat{o}_t = p_\phi(z_t)

会处理观测。在进入序列模型之前,latent sample 和 action 会先通过多层感知机(MLP)并进行拼接,从而构造成一个单独 token。这个过程如示例 7-9 所示。

示例 7-9. 构造单个 token

stem = nn.Sequential(
    nn.Linear(stoch_dim + action_dim, feat_dim, bias=False), 
    nn.LayerNorm(feat_dim),
    nn.ReLU(inplace=True),
    nn.Linear(feat_dim, feat_dim, bias=False),
    nn.LayerNorm(feat_dim)
)

把图像表示和动作 embedding 组合起来。

序列模型会接收这串 token 作为输入,并输出隐藏状态。它采用 GPT 风格的 Transformer 结构,其中 self-attention block 会使用 mask,从而保证模型只能沿序列方向做因果注意力。然后,它还会通过 MLP 去预测当前 reward、continuation flag,以及下一个 latent 分布。

整个模型以自监督方式训练。总损失函数包含以下几部分:重建损失、奖励预测损失、continuation flag 预测损失、动力学损失和表示损失:

L(ϕ)=1BTn=1Bt=1T[Ltrec(ϕ)+Ltrew(ϕ)+Ltcon(ϕ)+β1Ltdyn(ϕ)+β2Ltrep(ϕ)]\mathcal{L}(\phi)= \frac{1}{BT}\sum_{n=1}^{B}\sum_{t=1}^{T} \left[ \mathcal{L}^{rec}_t(\phi)+ \mathcal{L}^{rew}_t(\phi)+ \mathcal{L}^{con}_t(\phi)+ \beta_1 \mathcal{L}^{dyn}_t(\phi)+ \beta_2 \mathcal{L}^{rep}_t(\phi) \right]

其中,超参数 β1\beta_1β2\beta_2 固定为 0.5 和 0.1。BB 是 batch size,TT 是 batch length,ϕ\phi 表示模型各组件的参数,这些参数共同决定各项损失的计算方式。

重建损失 Ltrec(ϕ)\mathcal{L}^{rec}_t(\phi) 是原始图像与重建图像之间的误差,形式为:

Ltrec(ϕ)=o^tot2\mathcal{L}^{rec}_t(\phi)=|\hat{o}_t-o_t|_2

奖励预测损失

Ltrew(ϕ)=Lsym(r^t,rt)\mathcal{L}^{rew}_t(\phi)=\mathcal{L}^ {sym}(\hat{r}_t,r_t)

其中用到 symlog two-hot loss Lsym\mathcal{L}^ {sym},它会把原本的回归问题转成分类问题,从而让损失在不同环境之间都保持一致尺度。而 continuation flag 的预测损失写作:

Ltcon(ϕ)=ctlogc^t+(1ct)log(1c^t)\mathcal{L}^{con}_t(\phi)= c_t \log \hat{c}_t + (1-c_t)\log(1-\hat{c}_t)

这里,(gϕCg^C_\phi) 是一个 MLP,它会接收序列模型输出的隐藏状态 hth_t,并预测 continuation flag c^t\hat{c}_t

动力学损失用 KL 散度表示,但在反向传播和加权方式上做了特别处理,其目标是引导序列模型去预测下一个分布:

Ltdyn(ϕ)=max(1,KL[sg(qϕ(zt+1ot+1))gϕD(z^t+1ht)])\mathcal{L}^{dyn}_t(\phi)= \max\Big( 1, KL\big[ sg(q_\phi(z_{t+1}\mid o_{t+1})) || g^D_\phi(\hat{z}_{t+1}\mid h_t) \big] \Big)

其中 sg()sg(\cdot) 表示 stop-gradients。

表示损失同样也用 KL 散度表示,但反向传播与权重处理不同,它让编码器的输出受到序列模型预测结果的弱影响,从而保证分布动力学的学习不会变得过难。其形式是:

Ltrep(ϕ)=max(1,KL[qϕ(zt+1ot+1)sg(gϕD(z^t+1ht))])\mathcal{L}^{rep}_t(\phi)= \max\Big( 1, KL\big[ q_\phi(z_{t+1}\mid o_{t+1}) || sg(g^D_\phi(\hat{z}_{t+1}\mid h_t)) \big] \Big)

这里 sg()sg(\cdot) 同样表示 stop-gradients。

在 imagination 过程中,会从 replay buffer 中随机选取一段短上下文轨迹,并计算初始后验分布。而在推理阶段,则从先验分布中采样。agent 的状态由 ztz_thth_t 拼接而成,记作:

st=[zt,ht]s_t = [z_t, h_t]

critic 负责逼近状态价值 Vψ(st)V_\psi(s_t),actor 则根据策略

atπθ(atst)a_t \sim \pi_\theta(a_t \mid s_t)

来采样动作。图 7-6 展示了 STORM 的整体架构。

image.png

图 7-6. STORM 的结构与 imagination 过程。图中的 Transformer blocks 表示序列模型;而显示为神经网络的 agent block 则表示策略函数,它负责在给定状态下进行动作选择。图片改编自 Weipu Zhang 等人(2023)。

actor 的学习设置借鉴自 DreamerV3——这是一种使用基于 RNN 的世界模型来预测未来状态与奖励的强化学习算法。完整的 actor-critic 损失函数中包含奖励预测项和 continuation flag 项。actor 的损失依赖 λ\lambda-return,并会在 batch 维度上做归一化。critic 负责通过估计价值函数来评估动作,它会对自身参数 ψ\psi 使用指数滑动平均(EMA)进行正则化。这样可以稳定训练并避免过拟合。

用 Road Runner 测试 STORM

现在你已经理解了 STORM 的细节,是时候真正测试它了。这里你将使用 Arcade Learning Environment(ALE)中的 Road Runner 环境。ALE 是一个面向研究者和开发者的框架,用来构建 Atari 2600 游戏 AI agent。它建立在 Atari 2600 模拟器 Stella 之上,把底层模拟细节都封装起来,让用户能专注于 agent 设计。Road Runner 则是 ALE 中非常经典的一个 benchmark 游戏。

我按照论文仓库中的说明训练了 STORM,并保存了训练好的 checkpoint。为了在 Google Colab 中运行模型评估,你需要先写一个 evaluation script。示例 7-10 展示了其中一部分代码,其余部分你可以在本书仓库中找到。

示例 7-10. 评估脚本

def build_single_env(env_name, image_size, seed):
    env = gymnasium.make(env_name, full_action_space=False,
    render_mode="rgb_array", frameskip=1) 
    env = env_wrapper.SeedEnvWrapper(env, seed=seed) 
    env = env_wrapper.MaxLast2FrameSkipWrapper(env, skip=4) 
    env = gymnasium.wrappers.ResizeObservation(env, shape=image_size)
    env = env_wrapper.LifeLossInfo(env) 
    return env

def evaluate(env, world_model, agent, video_path):
    obs, _ = env.reset()
    total_reward = 0
    done = False
    frames = []

    while not done:
        frames.append(env.render())
        action = agent.select_action(world_model.encode_obs(obs)) 
        obs, reward, done, _, _ = env.step(action)
        total_reward += reward

    env.close()
    imageio.mimsave(video_path, frames, fps=30)
    return total_reward
  • 使用 Gymnasium 按指定设置创建环境。
  • 通过 seed wrapper 保证结果可复现。
  • 使用 frame skip wrapper 提升环境运行速度。
  • 记录 life loss 信息。
  • 基于世界模型编码后的观测,由 agent 选择动作。

如果你运行 ch03_STORM.ipynb notebook 里的代码,就会得到平均奖励。我在 STORM 上得到的输出是:

Mean reward: 16960.0

而一个随机 agent 在这个 benchmark 上只能得到大约 200–300 分,DreamerV3 在同样 benchmark 上大约是 15,565 分。人类在 Atari 100k benchmark 上的最高得分则是 7,845。这里的 “100k” 指的是:强化学习 agent 在训练过程中,最多只能和环境交互 100,000 次(帧数)。如果一个算法能在 Atari 100k benchmark 上取得很高表现,就说明它能在非常有限的数据下快速而有效地学习。

TWISTER:带对比预测编码的 Transformer 世界模型

虽然 STORM 已经证明了基于 Transformer 的世界模型在强化学习中的有效性,但很多过往方法——包括 STORM——都有一个固有局限:它们主要依赖“预测下一个状态”的目标。这个目标虽然听起来直观,但对于 Transformer 这种强表征模型来说,尤其是在 Atari 这种相邻视频帧差异很微妙的环境里,它往往不足以真正挖掘 Transformer 的表示能力。如果相邻 latent 状态之间的余弦相似度很高(如图 7-7 所示),那么“只预测下一帧”并不会迫使模型真正理解更深层的过去上下文,而这与神经语言建模那种明显依赖上下文的场景不同。

image.png

图 7-7. TWISTER 中 latent state ztz_t 与未来状态 zt+kz_{t+k} 的余弦相似度,在 Atari 100k benchmark 的 26 个游戏上聚合后的结果。数值是基于 5 个不同随机种子求平均得到的。图片改编自 Maxime Burchi 等人。

TWISTER(Transformer-based World model wIth contraSTivE Representations)正是为了解决这个问题。它通过把世界模型的预测目标扩展到更长时间范围,并引入动作条件的对比预测编码(action-conditioned contrastive predictive coding, AC-CPC)目标,来增强世界模型的表示能力。这种设计让世界模型能够学习更高层次的时序特征表示,从而在不依赖前瞻式搜索的情况下提升 agent 表现。

和你前面看到的其他基于模型的 RL 算法类似,TWISTER 主要由三个神经网络组成:世界模型、actor 网络和 critic 网络。它们会并行训练,并共同利用 replay buffer 中积累的历史经验。TWISTER 的核心创新在于它的世界模型设计,以及 AC-CPC 这个目标。和 STORM 类似,TWISTER 也采用基于 Transformer 的世界模型,更具体地说,是 transformer state-space model(TSSM)。图 7-8 展示了这个“带对比表示的 Transformer 世界模型”。

image.png

图 7-8. 这个基于 Transformer 的世界模型通过最大化当前模型状态 sts_t 与未来随机状态 zt:t+Kz_{t:t+K'}(从增强后的图像观测中得到)之间的互信息,来捕捉时序特征。图片改编自 Maxime Burchi 等人(2025)。

TSSM 使用带相对位置编码的 masked self-attention 来预测未来随机状态。它还有一个关键特性:会跨自回归片段缓存隐藏状态,这在功能上很像自回归 LLM 中的 KV cache,因此 TWISTER 可以在不反复重算全部上下文注意力的情况下,高效模拟长时程轨迹。和 STORM 一样,世界模型工作在 latent space 中:它会用一个带离散分类 latent 的卷积 VAE,把高维图像观测 oto_t 编码成离散随机状态 ztz_t,然后再通过解码器从这些随机状态重建图像。Transformer 网络会输出隐藏状态 hth_t,之后它会和随机状态 ztz_t 拼接起来,得到模型状态

st=(zt,ht)s_t = (z_t, h_t)

这和 STORM 类似。这些模型状态对于预测环境奖励、episode continuation flag,以及 AC-CPC 特征都很关键。

示例 7-11 展示了 TSSM 类中 forward_img 函数的实现(取自原始论文代码),它通过 masked attention 来编码时间上下文,并预测下一个 latent。

示例 7-11. 编码时间上下文

def forward_img(self, prev_states, prev_actions, mask,
                return_att_w=False, return_blocks_deter=False): 

        if self.action_clip > 0.0:
            prev_actions = prev_actions * (self.action_clip / torch.clip(
                            torch.abs(prev_actions),
                            min=self.action_clip)).detach() 

        if self.discrete:
            stoch = prev_states["stoch"].flatten(start_dim=-2, end_dim=-1) 
        else:
            stoch = prev_states["stoch"]

        x = self.action_mixer(torch.concat([stoch, prev_actions], dim=-1)) 

        outputs = self.transformer(x, hidden=prev_states["hidden"], mask=mask,
                                    return_hidden=True, return_att_w=return_att_w,
                                    return_blocks_x=return_blocks_deter) 
        deter, hidden = outputs.x, outputs.hidden

        add_out_dict = {}
        if return_att_w:
            add_out_dict["att_w"] = outputs.att_w
        if return_blocks_deter:
            add_out_dict["blocks_deter"] = outputs.blocks_x

        logits = self.dynamics_predictor(deter).reshape(
                deter.shape[:-1] + (self.stoch_size, self.discrete)) 
        dist_params = {'logits': logits}

        stoch = self.get_dist(dist_params).rsample() 

        return {"stoch": stoch, "deter": deter, "hidden": hidden,
                **dist_params, **add_out_dict}
  • 定义 latent space 中的主转移函数,用来预测下一个状态。
  • 裁剪动作幅度,让训练更稳定。
  • 把离散 latent 变量按类别维度展平,供 Transformer 输入。
  • 通过一个 MLP 把随机 latent 和动作组合成 Transformer 输入 embedding。
  • 使用带因果掩码的 Transformer 做时间建模。
  • 输出下一步 latent 分布的 logits(离散分类)。
  • 从预测分布中采样出下一个 latent state。
  • 返回更新后的随机状态、确定性状态和隐藏状态。

这个函数构成了训练时 rollout 和 agent 在 planning 过程中的 imagination 基础,使 TWISTER 能够完全在 latent space 中模拟序列。

TWISTER 最核心的创新,是把动作条件对比预测编码(AC-CPC)整合进世界模型的训练目标中。对比预测编码(CPC)最初被用于语音和图像等领域,它是一种自监督学习方法,目标是通过最大化“自回归模型输出特征”和“未来编码表示”之间的互信息,来学到更好的表示。在 TWISTER 中,这个思想被改造到了 model-based RL 里。不同于过去很多 Transformer 世界模型只关注预测下一个 latent 状态,TWISTER 的 AC-CPC 目标会把预测范围扩展到未来 (K=10) 个随机状态上。这迫使模型学习更抽象、更长程的时间表示。

借助这个目标,世界模型会学习区分两类样本:一类是正样本,也就是真实未来随机状态,它们与当前模型状态以及未来动作序列相对应;另一类是负样本,也就是 batch 中其他增强后的随机状态。AC-CPC 会利用 InfoNCE loss 训练世界模型,让它能够区分真实未来状态与负样本。每一个正负对都由两个部分组成:一个是 projected latent state,另一个是由当前模型状态和未来动作条件化得到的 prediction。它们之间的相似度通过点积计算。CPC 的损失目标,就是让真实配对的相似度变高,同时把负样本推远。

这些预测会以动作序列 at:t+ka_{t:t+k}作为条件,从而降低未来不确定性。对于每一个预测步 kk,模型会学习两个 MLP:

  • qϕkq_\phi^k:把离散随机未来状态 zt+kz'_{t+k} 投影到对比空间中。
  • pϕkp_\phi^k:基于当前模型状态 sts_t 和未来动作,去预测同一个对比空间表示。

InfoNCE 损失会在所有 k1,,Kk \in {1,\dots,K} 上计算,形式为:

Lcpc(ϕ)=1Kk=0K1logexp(sim(zt+k,st))zjZexp(sim(zj,st))L_{cpc}(\phi)= -\frac{1}{K}\sum_{k=0}^{K-1} \log \frac{ \exp(sim(z' _{t+k}, s_t)) }{ \sum_{z'_j \in Z'} \exp(sim(z'_j, s_t)) }

世界模型需要在 batch 中所有增强样本中,预测未来 K=10K=10 个随机状态。这里的相似度函数 simsim 定义为:

sim(zj,st)=qϕk(zj)Tpϕk(st,at:t+k)sim(z'_j, s_t)= q_\phi^k(z'_j)^T p_\phi^k(s_t, a_{t:t+k})

也就是说,针对每个步长 kk,都要学习一对 MLP:qϕkq_\phi^kpϕkp_\phi^k。这一点与原始 CPC 论文不同,原始 CPC 处理的是连续特征状态,而 TWISTER 的世界模型则使用离散 latent state。示例 7-12 展示了原始代码中这两个 MLP 的实现,它们负责生成 InfoNCE loss 中做相似度比较的特征对。

示例 7-12. 用于对比计算的两个 MLP

class ContrastiveNetwork(nn.Module): 

    def __init__(
        self,
        hidden_size=512,
        out_size=512,
        feat_size=32*32+512,
        embed_size=4*4*8*32,
        act_fun="ELU",
        num_layers=2
    ):
        super(ContrastiveNetwork, self).__init__()

        self.mlp_feats = modules.MultiLayerPerceptron( 
            dim_input=feat_size,
            dim_layers=[hidden_size for _ in range(num_layers-1)] + [out_size],
            act_fun=act_fun
        )
        self.mlp_embed = modules.MultiLayerPerceptron( 
            dim_input=embed_size,
            dim_layers=[hidden_size for _ in range(num_layers-1)] + [out_size],
            act_fun=act_fun
        )

    def forward(self, feats, embed): 
        x = self.mlp_feats(feats) 
        y = self.mlp_embed(embed) 

        return x, y
  • 定义 ContrastiveNetwork 类,把 projection head 和 prediction head 封装在一起。
  • 构造 projector MLP (q_\phi^k),用于把未来 latent state 映射到对比空间。
  • 构造 predictor MLP (p_\phi^k),用于根据当前模型状态和动作序列进行预测。
  • forward 方法接收未来 latent 和当前模型—动作 embedding。
  • 把 latent state 投影到对比空间。
  • 根据模型状态和动作上下文预测对比特征。
  • 返回这对特征,用于 InfoNCE loss 中的相似度计算。

在 TWISTER 中,actor 和 critic 网络与 DreamerV3 类似,都是基于学习到的世界模型生成的 imagination 轨迹来训练的。学习过程完全发生在 latent space 中,因此可以使用更大的 batch,也能获得更好的泛化。世界模型训练阶段缓存下来的 self-attention key/value,在行为学习时也会被复用,以便在 imagination 过程中保留上下文。actor 学习选择那些能够最大化预测回报的动作,而 critic 则评估状态的价值。

与其他领先的 model-based agent 相比,TWISTER 在样本效率和整体性能上都实现了提升。它在 Atari 100k benchmark 上取得了 162% 的 human-normalized mean score,刷新了这一类别 SOTA 方法的纪录。

总结

这一章向你介绍了强化学习中的关键概念,包括在线与离线学习、基于模型与无模型方法,以及 on-policy 与 off-policy 方法。你了解了传统 RL 算法及其基本原则,这为理解 RL Transformer 的最新进展打下了扎实基础。

你还学习了 Decision Transformer,它把 RL 重新定义为一个条件序列建模任务:通过一个自回归模型,把目标奖励与过去的状态和动作整合起来。这种方式符合离线 RL 和无模型 RL 的原则,因此可以提供更稳定、更有效的学习环境。

Online Decision Transformer 则把离线预训练和在线微调结合起来,从而增强了在动态环境中的适应能力。它通过 replay buffer 和 hindsight return relabeling 来进一步优化学习过程。

之后你又了解了 STORM,这是一种带有随机性的 Transformer 世界模型,它通过引入变分自编码器,增强了在有噪声、真实世界环境中的鲁棒性。STORM 在 Road Runner Atari benchmark 上的成功,也展示了它的性能,超过了随机基线以及更早的 model-based 方法。

最后,你接触到了 TWISTER,这是一种使用对比预测编码训练的 Transformer 世界模型。TWISTER 通过鼓励更深层的时间抽象,以及带动作条件的长时程规划,推动了这一方向的发展。它在 Atari 100k benchmark 上达到了最先进结果,也展示了对比式目标如何改善强化学习中的表示学习。

这一章中获得的知识,在下一章里会特别有价值,因为我们将从“基于奖励的训练”继续迈向“用于规划、推理和编程的 Transformer”。你刚刚学到的一些 RL 基础,会帮助你理解:Transformer 是如何为那些能够规划、推理并在多样任务中写代码的 agent 提供动力的。