基于潜世界模型的强化规划

# 基于潜在世界模型的强化规划  
来源：https://arxiv.org/html/2608.18669  

## 基于潜在世界模型的强化规划  
Armin Sommer，Jannik Schilling  
所属机构：Pantheon Industries  

###### 摘要  
人类通过构建计划并在内部世界模型中模拟结果来解决复杂问题。机器学习已产出类似的世界模型，能预测动作序列的结果，但对候选计划的改进仍未完全实现自动化。当前的规划器要么是手工设计的，要么是从手工设计的优化器中蒸馏而来，或者仅用于指导分摊策略，而非修订计划本身。我们提出**强化规划**方法，其核心思想是通过将优良搜索规则强化到神经规划器中来学习搜索过程。我们的实现 RP1 既能通过评论家学习如何评估想象的结果，也能通过一个完全离线训练的优化器学习如何改进多步骤计划，该优化器基于想象的世界模型展开序列进行训练。据我们所知，RP1 是首个完全学习如何改进多步骤计划的方法。此外，它可以独立于任何预训练的潜在世界模型进行训练，并能与之结合。在视觉导航、机械臂伸展和机器人操作这三个领域，并基于两种世界模型骨干架构的实验表明，RP1 显著优于手工设计的搜索算法，在多个设置中达到近乎完美的成功率，同时使用的世界模型展开序列比最强竞争对手方法少 1000×，在并行推理下速度提升最高达 67×。  

## 1 引言  
人们通常认为，人类通过想象可能的未来并评估其后果来解决复杂问题。海马体活动可以在动作发生前表征预期轨迹 [20, 37]，这支持了大脑使用已学习的认知地图进行内部模拟的观点 [49]。在计算层面，这将规划分为两个组成部分：**世界模型**预测假设动作的结果，而**规划器**决定如何生成、评估和改进候选动作序列。机器学习在第一部分已取得重大进展。潜在世界模型现在支持高维视觉预测 [7, 9, 10]、自监督预测性表征 [23, 3]，以及使用预训练的、无奖励模型进行规划 [54, 44, 28]。然而，运行在这些模型之上的规划器通常仍是手工设计的。给定一个候选动作序列，世界模型可以预测其结果，但并未指明应如何修改该序列以产生更好的序列。因此，现有系统依赖固定的搜索规则，往往每次决策需要数千次世界模型评估，并且需要为不同的任务和模型单独选择配置 [9, 12, 54, 44, 38]。虽然规划已以多种形式被探索，但在基于模型的多步骤规划中学习更新规则至今尚未实现：规划规则要么是固定的，要么继承自传统优化器，要么通过在线交互学习，或者仅应用于下一个动作而非整个计划（第 2 节）。  

我们介绍**强化规划**方法及其首个实现 RP1，它既能学习如何评估想象的结果，也能学习如何改进多步骤计划。RP1 使用时序差分学习从离线轨迹中学习一个目标条件下的拟度量评论家 [26, 52]，然后训练一个神经规划器，在推理时通过强化良好的规划规则到网络权重中，反复改进动作计划。在每个优化步骤中，RP1 接收当前动作计划，并通过世界模型展开评估其结果。此更新过程中不执行传统优化器，也不使用其作为训练目标。据我们所知，RP1 是首个完全学习针对多步骤动作计划更新规则的基于模型的规划器（第 2 节）。  

我们使用两个预训练的世界模型骨干——LeWorldModel 和 PLDM——在视觉导航（TwoRoom）、连续控制伸展（Reacher）和接触丰富的操作（OGBench Cube）上评估 RP1。在这三个领域中，RP1 超越了现有的最强规划器，且每次决策仅使用 99 次世界模型展开，而最强的竞争对手方法需要 9,000 次，并且在多个控制循环共享一个 GPU 时，规划延迟最多降低 67×。  

## 2 背景  
#### 固定或继承的规划规则。  
大多数基于模型的规划器使用手工设计的更新规则：如 PlaNet 和 DINO-WM 中的 CEM，TD-MPC 系列中的 MPPI，或通过可微世界模型展开进行梯度下降 [9, 12, 54, 44, 38]。通用规划网络优化多步骤动作序列，但将梯度下降固定为计划优化器 [45]。DMPO 保留 MPPI 更新和移位操作，并从在线任务回报中学习对它们的修改 [40]。L2O-MPC 学习运行时更新，但仅通过模仿一个更高预算的 MPPI 专家来实现，该专家必须在训练期间运行 [39]。在所有这些方法中，规划规则要么是（部分）手工设计的，要么是从手工设计的优化器中学习的。  

#### 分摊的基于模型的控制。  
Dreamer 系列方法使用其世界模型来训练分摊策略，但在推理时并不通过世界模型进行规划 [8]。Diffuser 学习一个状态-动作轨迹的生成模型，该模型同时捕捉动力学和规划，直接通过去噪来改进轨迹，而不是通过一个单独的世界模型显式地展开和评估连续的候选动作计划 [19]。  

#### 为分摊策略提供信息的规划。  
基于想象的规划器（IBP）和 Thinker 方法学习构建或检查哪些想象轨迹，但并不迭代地改进候选计划。相反，通过想象收集到的信息为智能体的分摊动作策略提供条件 [36, 5]。因此，它们学习到的规划行为是为了改进固定的动作分布，而不是改进候选计划本身，并且需要在线学习。  

#### 迭代式下一步动作优化。  
迭代分摊策略优化（IAPO）则学习一个针对当前状态动作分布 π(aₜ|sₜ) 的迭代优化器 [30]。即使在其基于模型的变体中，世界模型展开中的未来项仍然是分摊策略输出，而非联合优化的决策变量。因此，学习到的优化仍然是单步改进，而不是学习针对多步骤计划的更新规则。  

#### 想象计划的目标。  
基于 JEPA 的世界模型文献倾向于通过想象结果在潜在空间中与目标的欧几里得距离来评分 [28, 44, 54]。这种选择在网格细胞表征中具有生物学类比，网格细胞被认为为基于向量的导航提供了空间度量 [11, 4]。然而，潜在接近性不一定反映时间可达性（定理 1），最近的研究表明，学习可达性目标可以优于潜在距离 [24]。使用学习的价值函数评估结果也与涉及眶额皮层和腹内侧前额皮层在前瞻性价值评估中的证据一致 [33, 53, 42]，并且在基于模型的控制中已得到充分确立 [13, 12]。因此，我们使用目标条件下的拟度量式评论家 [21, 41, 1, 14, 26, 52, 51] 从经验中估计时间上的剩余代价。  

## 3 预备知识  
我们考虑一个目标条件下的马尔可夫决策过程 (S, A, T, g, ρ₀)，其状态空间 S，动作空间 A ⊂ ℝ^|a|，目标状态 g ⊆ S，转移函数 T: S × A → S，以及初始状态分布 ρ₀。智能体不直接观察 s，而是接收形式为视觉图像的观测 o。  

#### 世界模型。  
世界模型预测在给定当前状态和某个候选动作下的下一个状态。具体来说，神经网络通过编码器将观测 oₜ 映射到其潜在表示 zₜ = Eφ(oₜ) ∈ Z。世界模型随后通过预测映射 hφ(z, a) = ẑ 预测下一个潜在状态，其中 ẑ 是预测的（或想象的）下一个状态的潜在表示。我们定义世界模型的 N 步展开算子为：  
Hφ(a, ẑ₀) := hφ( hφ( ⋯ hφ(ẑ₀, a₀) ⋯, aₙ₋₂), aₙ₋₁) = ẑₙ,  (1)  
即世界模型 N 次前向展开的组合。  

#### 基于模型的规划。  
给定起始潜在状态 zₜ 和编码的目标 zᵍ，一个计划通过世界模型 Hφ 预测的最终潜在状态上的终端代价 C 来评分。该计划旨在优化目标：  
J(a; zₜ, zᵍ) := C( Hφ(a, zₜ), zᵍ ).  (2)  
基于模型的规划器是一个针对动作序列的搜索过程：它持有候选计划，并查询世界模型 Hφ 以在 J 下评估它们，并应用更新规则：  
F: aₖ ↦ aₖ₊₁  (3)  
进行 K 轮。基于模型的规划器仅在其实例化 F 上有所不同。相比之下，策略 π(a | zₜ, zᵍ) 则将目标 (2) 分摊为直接的状态到动作映射，不执行任何规划。我们的方法遵循演员-评论家架构：一个**评论家**根据目标对潜在状态进行评分，而一个**学习的规划器**（此处作为演员）根据评论家对最终状态的估计优化动作计划。  

#### 评论家。  
评论家是一个目标条件下的价值函数 Vψ(zₜ, zᵍ)，它估计从潜在状态 zₜ 到编码的目标状态 zᵍ 的剩余代价。此处，较低的值对应到达目标所需的更少步骤，因此表示占据更好的状态。我们通过离线时序差分学习来学习这个评论家，尽管理论上它可以是任何代价函数。在规划过程中，评论家仅评估由展开算子产生的终端状态 ẑₙ。  

#### 规划器。  
规划器是一个学习的算子：  
ℱθ: ( aₖ （当前计划）, vₖ （终端价值）, gₖ （价值梯度） ) ⟼ aₖ₊₁ （改进计划）,  (4)  
由 θ 参数化，它根据当前计划、评论家在计划终端状态的价值以及计划的价值梯度，输出改进的计划。从初始计划 a₀ 开始，规划在每一步执行三件事：  
*展开：* ẑₙ⁽ᵏ⁾ = Hφ(aₖ, zₜ),  (5)  
*评估：* vₖ = Vψ(ẑₙ⁽ᵏ⁾, zᵍ), gₖ = ∇ₐₖ Vψ(ẑₙ⁽ᵏ⁾, zᵍ),  (6)  
*改进：* aₖ₊₁ = ℱθ(aₖ, vₖ, gₖ).  (7)  
优化后的计划是最终迭代 a⋆ = aₖ。虽然规划器可以访问价值和梯度，但它并不受限于遵循计划的梯度 -gₖ，并且可以学会何时信任和何时不信任它。任务特定信息仅通过 vₖ 和 gₖ 传递给规划器，迫使其学习计划更新规则，而不是直接的状态-目标到动作映射。  

#### 强化优良的规划规则。  
应用 ℱθ 产生一条想象的优化轨迹：  
a₀ →ℱθ a₁ →ℱθ ⋯ →ℱθ aₖ,  (8)  
其中每个新计划都是通过将相同的学习更新规则应用于前一个计划而获得的。规划器被训练以最小化由冻结的世界模型 Hφ 和价值函数 Vψ 预测的终端剩余代价：  
θ⋆ = arg minθ E₍z₀, zᵍ₎~D [ Vψ( Hφ(aₖ, zₜ), zᵍ ) ] + C,  (9)  
其中 C 是对中间计划价值的某种正则化项。产生更低代价想象计划的更新会减少优化目标，并在 ℱθ 的共享参数中得到强化，而更新