学习离散世界模型用于启发式搜索

https://openreview.net/pdf?id=k0ob5K8Wqe

Learning Discrete World Models for HeuristicSearch

学习离散世界模型用于启发式搜索

摘要

对于许多顺序决策问题，规划通常是找到解决方案的必要手段。然而，对于诸如机器人等领域，状态转移函数（也称为世界模型）通常是未知的。虽然基于模型的强化学习方法可以学习世界模型并用于规划，但这种方法受到模型在多个时间步长应用时累积的误差以及无法在规划期间重新识别状态的限制。为了解决这些问题，我们提出了DeepCubeAI算法，该算法学习一个在离散潜在空间中表示状态的世界模型，使用强化学习通过学习模型学习一个在起始状态和目标状态上泛化的启发式函数，并将学习到的模型和启发式函数与启发式搜索结合以解决问题。由于潜在空间是离散的，我们可以通过舍入防止小误差的累积，并通过简单地比较两个二进制向量来重新识别状态。在Rubik's Cube、Sokoban、IceSlider和DigitJump的像素表示实验中，我们发现DeepCubeAI能够在数千步内应用模型而不会累积任何误差。此外，DeepCubeAI在所有领域中解决了超过99%的测试实例，能够在目标状态上泛化，并显著优于不使用学习世界模型进行规划的贪婪策略。

1 引言

规划需要一个状态转移函数（也称为世界模型），它能够准确地将状态和动作映射到下一个状态。虽然对于具有符号表示的环境，手动构建世界模型通常很方便，但对于具有子符号表示（例如像素）的环境，这种方法变得不切实际。另一方面，使用机器学习技术从观察到的转移中学习模型提供了一种领域无关的模型构建方法。然后，强化学习可以与这些学习到的模型一起使用，以学习策略或启发式函数，而无需收集任何额外的真实世界数据。此外，在测试时，学习到的模型可以与搜索结合以提高性能。

然而，这种方法有两个主要障碍：1）许多学习到的模型存在模型退化问题，从而使其在长期规划中无效；2）许多学习到的模型无法在搜索期间重新识别状态，导致搜索树中的循环，从而降低规划效率。

模型退化发生在模型的预测误差随时间步长累积时，导致长期预测的可靠性下降。模型退化已在Atari学习环境（Oh等，2015）、Sokoban（Racanière等，2017）和机器人操作任务（Finn等，2016）等领域中观察到。由于这将学习模型的使用限制在短期任务中，如果使用这种学习模型来学习启发式函数，代理将仅限于探索接近真实世界中观察到的状态，这可能导致泛化能力差。虽然可以通过更多的真实世界探索来弥补这一点，但真实世界探索通常比使用模拟真实世界的学习模型耗时得多。

当使用退化的模型进行规划时，只能考虑相对接近起始状态的状态。这可能导致规划效果差，并需要频繁重新规划（Tian等，2021）。状态重新识别是知道两个潜在嵌入是否表示相同状态的能力。这对于规划至关重要，因为如果没有状态重新识别，搜索过程中会多次访问相同的状态。在最坏的情况下，这会导致计算时间和内存随着搜索树深度的增加呈指数增长。

为了解决这些问题，我们将学习从状态到离散潜在空间的映射，并学习一个在该离散潜在空间中捕捉状态转移的模型。这将使我们能够对抗模型退化，因为小于0.5的误差可以通过舍入轻松修复。这将使模型能够在数千个时间步长内使用而不会累积任何误差。此外，这种离散表示使得状态重新识别成为两个二进制向量之间的简单比较。一旦模型被学习，将使用深度神经网络（DNN）（Schmidhuber，2015），即深度Q网络（DQN）（Mnih等，2015），通过Q学习（Watkins & Dayan，1992；Sutton & Barto，2018）学习启发式函数。由于在测试时指定的目标不假定为事先已知，启发式函数将使用一种受后见经验回放（Andrychowicz等，2017）启发的方法进行训练，以使其能够在目标上泛化。这导致了一种领域无关的算法，用于训练能够在问题实例上泛化的领域特定启发式函数。然后，该启发式函数将与Q*搜索（Agostinelli等，2024b）（一种适用于DQN的A*搜索变体）结合使用以解决问题。由于该方法基于DeepCubeA算法（Agostinelli等，2019），该算法结合深度强化学习和搜索来解决经典规划问题，因此我们将我们的方法称为DeepCubeA-Imagination（DeepCubeAI），其中“想象”指的是使用学习到的模型“想象”未来场景的能力（Racanière等，2017）。

2 相关工作

基于模型的强化学习（RL）方法试图利用学习到的模型来减少学习策略或价值函数所需的真实世界训练数据量，并在测试时进行策略改进。最早的实例之一是Dyna架构（Sutton，1991）。Dyna架构方法与当今的许多方法类似，即使用观察到的转移来训练模型，随后用于学习和规划。尽管在表格设置中展示了强大的结果，但在无法用表格表示的大状态空间中，可靠的结果仍未实现，至今仍然难以捉摸。

现代基于模型的RL方法的一个例子是**基于模型的离线学习距离强化学习（MBOLD）**（Tian等，2021）。MBOLD提出了一种使用离线数据训练模型以预测下一状态像素的方法。它使用这些离线数据和模型来训练启发式函数以估计“成本到目标”。然而，该模型在连续潜在空间中运行并累积误差。因此，它在生成启发式函数的训练数据方面受到限制，无法在达到目标之前进行规划，并且无法重新识别状态。

Bagatella等人（2021）的工作提出了一种名为**通过图搜索从像素规划（PPGS）**的方法，该方法学习在连续潜在空间中表示状态。状态重新识别通过比较两个向量之间的距离并设置重新识别阈值来完成。通过利用状态重新识别，他们创建了一个潜在图并部署图搜索算法来解决经典规划问题。该架构包括一个编码器、一个前向模型和一个逆模型，后者用于确保潜在状态包含模型所需的相关信息。随后，他们引入了两个新环境——IceSlider和DigitJump，这些环境具有潜在的组合结构，并在其中验证了PPGS相对于无模型方法（如PPO（Schulman等，2017））的优越性能。然而，学习到的模型会累积误差，并且在预测的潜在状态与观察到的状态不匹配时需要重新规划。此外，PPGS不学习启发式函数，因此依赖于广度优先搜索来解决问题，这无法扩展到更复杂的问题。

DreamerV3（Hafner等，2023）使用**递归状态空间模型（RSSM）（Hafner等，2019）在离散潜在空间中建模状态。他们使用该模型和演员-评论家方法训练策略函数。DreamerV3能够从零开始在Minecraft中收集钻石，而无需人类数据。然而，DreamerV3仅在训练时使用学习到的模型，而不是在测试时用于规划；因此，它尚未展示出在达到目标之前进行规划或重新识别状态的能力。

与学习在人类不易理解的潜在空间中运行的黑箱模型不同，已有研究致力于学习可以显式表示为**规划域定义语言（PDDL）**格式的模型（Asai & Fukunaga，2018；Asai等，2022）。给定这种表示，可以使用领域无关的规划器来解决问题。然而，这些领域无关的规划器在解决诸如魔方（Muppasani等，2023；Agostinelli等，2024a）等问题时常常失败。此外，由于从数据中学习PDDL模型并不总是可行的，因此可能必须使用学习到的黑箱模型和适用于黑箱模型的领域无关启发式方法（例如目标计数启发式）。

3 预备知识

在本工作中，我们设计了一种能够在确定性、完全可观察的领域中学习离散世界模型的算法。一个领域D可以表示为一个确定性无折扣的马尔可夫决策过程（MDP）（Puterman，2014），它是一个元组 ，其中S是所有状态的集合，A是所有动作的集合，T是将状态和动作映射到下一个状态的状态转移函数，G是将状态、动作和下一个状态映射到转移成本的转移成本函数。它也可以表示为一个加权有向图（Pohl，1970），其中节点表示状态，边表示状态之间的转移，边权重表示转移成本。目标对应于一组被视为目标状态的状态。给定一个起始状态，目标是找到一个动作序列，将起始状态转换为目标状态，同时尝试最小化路径成本，即转移成本的总和。状态转移函数和转移成本函数共同构成了世界模型。当转移成本是均匀的时，学习模型简化为仅学习状态转移函数。

4 方法

4.1 学习离散世界模型

我们旨在学习一个模型m，该模型在某个潜在空间中表示状态转移函数T。在此设置中，我们假设所有转移成本均为1。与Tian等人（2021）类似，我们将从随机探索收集的离线数据中学习模型。该数据集将包含一组状态、动作和下一状态的元组（s, a, s'）。编码器将被训练以将给定状态投影到潜在空间中的状态。编码器将在其输出层使用逻辑函数，该函数将被舍入为0或1。在梯度下降期间，将使用直通梯度估计器（Bengio等，2013）来解决舍入函数相对于其输入的导数为零的问题。解码器将潜在空间映射回状态空间。将使用均方误差作为重建误差，以使解码器的输出尽可能接近编码器的输入。这确保了编码能够捕捉状态中的内容。重建误差如公式1所示，其中N是批量大小，是解码器的输出，θ是自动编码器和模型的参数。

同时，将训练一个模型以将潜在状态和动作映射到下一个潜在状态。将使用损失函数来鼓励模型的输出和编码器的输出尽可能接近。在我们的实验中，我们发现将模型与自动编码器一起训练的最佳方式是鼓励模型的输出与编码器的输出匹配，同时鼓励编码器的输出与模型的输出匹配。然而，我们仅在鼓励编码器的输出与模型的输出匹配时对模型的输出进行舍入，而编码器的输出始终被舍入。如下面的公式2所示，其中r()是使用直通估计器进行梯度下降的舍入函数，.detach()将张量从计算图中移除。

在我们的实验中，我们观察到先训练自动编码器再训练模型会导致模型不完美，这意味着它无法以100%的准确率预测下一个潜在状态。因此，我们发现需要同时训练自动编码器和模型，以确保自动编码器的参数能够学习到模型也能学习的表示。然而，公式1和公式2中的损失函数可能会相互冲突，因为Lr没有考虑模型预测潜在状态的能力，而Lm没有考虑重建误差。因此，我们使用权重ω首先将Lr损失加权高于Lm，并逐渐调整ω为0.5以使它们权重相等。损失如公式3所示。

训练过程如图1所示。训练后，每次应用模型时，都会对其输出进行舍入操作以纠正错误并防止误差累积。

4.2 学习启发式函数

给定训练好的模型和离线数据，可以生成由起始状态和目标状态对组成的训练数据，以训练一个在起始状态和目标状态上都能泛化的启发式函数。对于每个训练示例，从离线数据中采样一个真实世界状态。然后使用编码器获得相应的潜在状态。通过从采样的潜在状态开始并使用模型在潜在空间中随机采取ts步来获得起始状态，其中ts在0到Ts之间均匀随机分布。然后通过从起始状态开始并采取tg步来获得目标状态，其中tg在0到Tg之间随机分布。利用这些数据，通过强化学习训练一个深度Q网络（DQN），将起始状态和目标状态映射到每个动作的“成本到目标”。

DQN是一种将状态映射到大小为|A|的向量的神经网络，其中索引a处的每个元素表示从给定状态开始并采取动作a时的预期“成本到目标”，记为Q(s, a)。在无折扣的确定性设置中，Q(s, a)的估计值迭代更新为G(s, a, s') + mina' Q(s', a')。然而，由于Q是由参数ϕ表示的DQN，qϕ，从自身引导会导致非平稳目标问题。为了解决这个问题，根据之前的工作（Mnih等，2015），维护一个具有参数ϕ−的目标网络，并在训练期间定期更新为ϕ。使用的损失函数为L(ϕ) = (G(s, a, s') + mina' qϕ− (s', a', sg) − qϕ(s, a, sg))，其中是目标状态。

为了选择要更新Q学习的动作，我们优先选择更有希望的动作而不是不太有希望的动作，因为在许多环境中，给定状态下的大多数动作都不在最短路径上，从而导致偏差。因此，我们根据玻尔兹曼分布选择动作，其中每个动作a根据公式4中的概率被选择，其中τ是温度。

4.3 使用学习模型和学习启发式函数进行规划

给定学习到的模型和启发式函数，规划可以通过状态空间搜索的形式完成。虽然可以通过将启发式函数h(s, sg)设置为mina′ qϕ(s, a′, sg)将DQN与A*搜索结合使用，但A*搜索要求每次迭代使用模型|A|次。鉴于模型是一个计算成本高昂的DNN，我们希望尽量减少使用它的次数。因此，我们改用Q*搜索（Agostinelli等，2024b），这是A*搜索的一种改进，利用了Q*搜索可以通过一次DQN计算所有下一个状态的启发式值的特点。实际上，Q*搜索已被证明与A*搜索性能相似，但速度快了几个数量级且内存效率更高。为了利用GPU并行性并加速搜索，我们还使用了DeepCubeA在A*搜索中使用的批量加权（Pohl，1970）版本的Q*搜索。

5 实验

我们在魔方、Sokoban、IceSlider和DigitJump上测试了我们的方法。对于魔方，状态由两个32x32的RGB图像表示，每个图像显示魔方的三个面。对于Sokoban，状态由一个40x40的RGB图像表示，显示代理、墙壁和箱子。IceSlider和DigitJump（由Bagatella等人（2021）引入）由一个64x64的RGB图像表示，代表一个8x8的二维网格。在IceSlider中，代理必须在冰面上滑动，只有岩石或环境边界才能阻止其移动，以到达给定的单元格。在我们的工作中，我们使用代理作为指示器来表示目标单元格，而不是像之前的工作中使用紫色方块。在DigitJump中，代理从左上角开始，目标是到达右下角。代理在给定单元格中跳跃的单元格数量由该单元格中的MNIST（LeCun等，1998）数字表示，数字范围为1到6。状态示例如图3所示。

在我们的实验中，我们通过观察跨剧集的转移生成离线数据集，其中每个剧集中代理采取随机动作1。对于魔方和Sokoban，我们在10,000个剧集中生成300,000个转移，每个剧集中采取30个随机动作。对于IceSlider和DigitJump，我们以与Bagatella等人（2021）类似的方式生成离线数据。具体来说，我们在20,000个剧集中生成400,000个转移，每个剧集中采取20个随机动作。对于魔方，每个剧集的起始状态通过随机打乱目标状态100到200次获得。对于Sokoban，每个剧集的起始状态从Guez等人（2018）提供的训练示例中随机采样。对于IceSlider和DigitJump，每个剧集的起始状态从Bagatella等人（2021）使用的相同1,000个关卡中获得。对于魔方和Sokoban，生成的数据中90%用于训练模型，10%用于验证。对于IceSlider和DigitJump，验证数据通过在5,000个剧集中重复该过程生成，使用另一组1,000个不同的关卡，每个剧集中采取20个随机动作，总共生成100,000个转移。在训练和搜索期间，如果潜在状态中的100%位相等，则认为两个潜在状态相等。

对于魔方，自动编码器架构是一个全连接神经网络，编码器和解码器各有一个隐藏层，编码维度为400。编码器使用逻辑激活函数，而解码器使用线性激活函数。尽管RGB值在0到1之间，但我们发现解码器最后一层的线性层表现更好。离散世界模型是一个具有四层（大小为500、500、500和400）的全连接神经网络。它在所有层中使用批量归一化，除了最后一层。此外，除了最后一层使用逻辑激活函数外，其他层都使用修正线性单元（ReLU）（Glorot等，2011）。模型使用独热编码表示动作，并将其与潜在状态连接。

对于Sokoban，自动编码器架构使用卷积编码器和解码器，两者都有两层，每层16个通道，核大小为2，步幅为2，第一层使用批量归一化。解码器使用一个额外的卷积层，核大小为1，激活函数为线性。离散世界模型是一个卷积神经网络，具有三层，通道大小分别为32、32和16，核大小均为3，步幅为1，前两层使用批量归一化，前两层使用修正线性单元，最后一层使用逻辑激活函数。模型使用独热编码表示动作，并将其扩展为一个张量，其长度和宽度与潜在表示的大小相同，通道数等于动作数。然后将其与潜在状态沿通道维度连接。

对于IceSlider，自动编码器架构与Sokoban类似，使用两层卷积编码器，第一层32个通道，最后一层3个通道。解码器使用转置卷积层，32个通道。卷积层的核大小分别为4和2，步幅分别为4和2。批量归一化和激活函数与Sokoban相同。解码器还包括一个额外的线性激活函数层。在离散世界模型中，我们遵循Sokoban的连接过程。首先，一个核大小为1、步幅为1的卷积层处理输入。然后将其传递给四个残差块，保持与输入相同的通道数。最后一个残差块的输出传递给一个核大小为1、步幅为1的卷积层。最后，一个核大小为1的额外层作为输出层。除了第一层使用层归一化，最后一层使用无归一化的逻辑激活函数外，其他层都应用批量归一化和修正线性单元。DigitJump与IceSlider共享相同的架构层，编码器输出12个通道，残差块使用47个通道。

所有模型均使用ADAM优化器（Kingma & Ba，2014）进行梯度下降训练，学习率为0.001，衰减率为0.9999993，批量大小为100。ω初始化为0.0001，并在120,000次迭代时逐渐调整为0.5。然后在180,000次迭代时训练神经网络，并每20,000次迭代将学习率降低10倍。

然后使用Q学习来训练启发式函数。为了生成起始状态和目标状态对，对于魔方和Sokoban，Ts和Tg均设置为30，对于IceSlider和DigitJump，Ts和Tg均设置为20。启发式函数使用ADAM优化器进行Q学习训练，学习率为0.001，衰减率为0.9999993，批量大小为10,000，训练1百万次迭代。根据公式4选择动作，τ设置为3.0。为了在学习过程中更好地探索状态空间，还通过贪婪行为生成新状态，对于魔方和Sokoban最多30步，对于IceSlider和DigitJump最多20步。每5,000次迭代使用贪婪策略测试DQN，如果贪婪策略有所改进，则更新目标网络的参数。

对于Q*搜索，魔方的批量大小和路径成本权重分别为10,000和0.6，Sokoban为100和0.8，IceSlider和DigitJump为1和0.7。为了指定目标状态，给出目标图像，编码到潜在空间，并传递给启发式函数。对于Sokoban，目标由箱子应放置位置的图像指定，代理随机选择放置在箱子旁边。我们将在未来工作部分讨论更鲁棒的目标指定方法。我们将DeepCubeAI与DeepCubeA以及利用群论知识的领域特定PDB（Rokicki，2016；Rokicki等，2014）进行比较。

5.1模型性能

为了评估模型的性能，我们在10个10,000步的序列上测试它，其中动作是随机选择的。我们为每一步获取真实图像，并在潜在空间中采取步骤，从解码器获取重建图像。此外，我们与一个具有相同架构和训练过程的连续模型进行比较，但没有离散化。比较结果如图2所示。图中显示，虽然连续模型在Sokoban、IceSlider和DigitJump上没有累积误差，但在魔方上累积了显著的误差。图3a展示了魔方的一个示例，其中连续模型产生了显著误差，而离散模型没有。图3b、3c和3d展示了Sokoban、IceSlider和DigitJump的示例，其中连续模型和离散模型都没有产生显著误差。这可能归因于这些环境在多个时间步长上更容易重建。例如，在Sokoban中，箱子很快被推到墙边，因此在此之后变得不可移动，只有代理的位置在转移之间发生变化。同样，在IceSlider和DigitJump中，动作仅影响代理的位置，并且存在代理无法移动的状态，如图3d所示。

5.2 问题解决性能

我们在1,000个魔方和Sokoban的测试实例上评估DeepCubeAI，这些实例来自DeepCubeA存储库（Agostinelli等，2020），并在用于评估PPGS的100个IceSlider和DigitJump测试实例上进行评估（Bagatella等，2021）。为了确定规划在解决这些测试实例中的重要性，我们还使用它们来评估通过贪婪行为训练的DQN在100步内获得的贪婪策略。为了测试DeepCubeAI对新目标状态的泛化能力，我们包括一个测试集，其中魔方的起始状态和目标状态被反转。因此，每个测试实例都有不同的目标状态。我们注意到，DeepCubeAI在训练期间并未被告知测试目标状态。

表1详细比较了DeepCubeAI与DeepCubeA、PDB和贪婪策略的表现。结果显示，DeepCubeAI解决了100%的魔方（使用规范目标状态）、Sokoban、IceSlider和DigitJump的测试实例。对于反转的起始状态和目标状态，DeepCubeAI解决了99.9%的测试实例（仅遗漏1个），并且表现与规范目标测试实例相似。我们注意到，DeepCubeA和PDB不能直接应用于此测试集，因为它们是针对规范目标状态设计的。要将DeepCubeA应用于不同的目标状态，我们必须为1,000个目标状态中的每一个训练一个新的DNN。贪婪策略未能解决任何魔方问题实例，并分别解决了Sokoban、IceSlider和DigitJump的41.9%、46.0%和90.0%的问题实例。这表明，使用学习到的世界模型进行规划对于解决这些问题至关重要。DeepCubeAI的路径成本比DeepCubeA更长。这可能是因为DeepCubeAI学习了一个更复杂的启发式函数，能够在目标上泛化，而DeepCubeA是针对预定义目标训练的。然而，对于魔方，尽管由于学习模型比手工编码模型计算成本更高，每秒处理的节点更少，但DeepCubeAI在找到解决方案时生成的节点更少，花费的时间也更短。这可能部分归功于Q*搜索提供的加速。然而，对于Sokoban，我们发现DeepCubeAI在执行Q*搜索时需要100的批量大小，而DeepCubeA使用1的批量大小进行A*搜索，因此DeepCubeAI生成的节点数仍然比DeepCubeA多。

6 未来工作

在DeepCubeAI未能找到路径的个别案例中，我们发现它无法正确识别潜在目标状态。这可能是因为模型在搜索过程中产生了大于0.5的误差，意味着舍入无法纠正它。未来的工作可以通过训练一个DNN来纠正轻微损坏的潜在状态来解决这些罕见错误。

与基于模型的强化学习研究类似（Tian等，2021），我们使用目标图像指定目标。虽然这对于某些环境可能是可行的，但在难以生成目标图像的环境中，这种方法变得不切实际。此外，如果只知道目标的高级信息而不了解低级细节，则无法生成目标图像。为了解决这个问题，已有研究使用形式逻辑来指定目标，其中目标可以是一组状态（Agostinelli等，2024a）。这种方法可以扩展到学习模型，并允许在不生成任何目标图像的情况下指定目标。

对于某些机器人操作任务，如果有足够的传感器和环境经验，可以将领域视为确定性和完全可观察的。然而，由于领域的固有特性或对环境动态的缺乏了解，许多机器人任务是随机的，并且由于传感有限而部分可观察。已有研究通过在随机环境中训练DNN来采样可能的下一个状态（Kaiser等，2020；Hafner等，2021）。序列模型，如递归神经网络（Hochreiter & Schmidhuber，1997），已被用于学习嵌入信念状态（Hausknecht & Stone，2015；Cassandra等，1994），我们可以在这些状态上进行规划。离散模型的优势也可以扩展到这些领域，允许模型在长时间范围内应用，以改进训练探索并在搜索期间获得更多的前瞻性。

7 结论

我们介绍了DeepCubeAI，这是一种领域无关的方法，用于学习在离散潜在状态上操作的模型。然后使用该学习模型来学习在问题实例上泛化的启发式函数。学习模型和学习启发式函数与搜索结合以解决问题。在魔方的情况下，结果表明，拥有离散模型对于防止误差累积至关重要。在魔方、Sokoban、IceSlider和DigitJump的所有情况下，结果表明DeepCubeAI解决了超过99%的测试案例，并有效地在目标状态上泛化。

原文链接：https://openreview.net/pdf?id=k0ob5K8Wqe