SNeRL:用于强化学习的语义感知神经辐射场

SNeRL: Semantic-aware Neural Radiance Fields for Reinforcement Learning

SNeRL:用于强化学习的语义感知神经辐射场

https://proceedings.mlr.press/v202/shim23a/shim23a.pdf

摘要

由于先前的强化学习表示无法有效地结合人类对 3D 环境的直观理解，它们通常会表现出次优的性能。在本文中，我们提出了用于强化学习的语义感知神经辐射场 (SNeRL)，它联合优化了语义感知神经辐射场 (NeRF) 和卷积编码器，以从多视图图像中学习 3D 感知的神经隐式表示。我们在 NeRF 中引入了与 RGB 辐射场并行的 3D 语义和蒸馏特征场，以学习强化学习的语义和以对象为中心的表示。SNeRL 不仅在模型无关和基于模型的强化学习中优于先前的基于像素的表示，还优于最近的 3D 感知表示。

1. 引言

开发能够直接从图像输入实现复杂控制任务的代理一直是强化学习 (RL) 中的一个长期问题。过去几年中，相关工作在数据效率方面取得了显著进展，尤其是在学习视觉控制问题方面。解决视觉控制任务最具挑战性的部分是从高维观察中获取低维潜在表示。为此，他们以各种方式预训练编码器，例如通过使用离线数据集进行图像重建的无监督表示学习 (Finn 等人，2016；Kulkarni 等人，2019；Islam 等人，2022)、对比学习 (Zhan 等人，2022)、重建任务信息 (Yang & Nachum，2021；Yamada 等人，2022) 以及训练多视图一致性 (Dwibedi 等人，2018)。

其他方法利用辅助无监督任务的联合学习（Laskin 等人，2020b；Schwarzer 等人，2020）以及数据增强的强化学习（Laskin 等人，2020a；Yarats 等人，2021）。

尽管已经提出了许多工作来提高视觉控制问题中的数据效率，但通过这些方法训练的大多数编码器在获取 3D 结构信息方面能力有限，并且缺乏对 3D 变换的等变性。这些限制源于忽略了 3D 结构信息，并从单视图观察中学习视觉表示。

最近，通过神经辐射场（NeRF）（Mildenhall 等人，2020）学习隐式空间表示，在机器人控制和操作中考虑环境 3D 信息的尝试有所增加（Li 等人，2022；Driess 等人，2022）。他们通过自编码器结构将场景的像素级多视图观察映射到潜在向量，其中 NeRF 解码器通过每个视图的 RGB 自监督提供 3D 结构化的神经场景表示。

尽管上述先驱者在性能上优于使用单视图观察的先前 RL 算法，但他们仍然没有充分利用 3D 感知表示学习。这是因为这些方法仅利用 RGB 监督来训练 NeRF，这使得编码器难以学习面向对象或语义表示以用于 RL 下游任务。尽管 NeRF-RL（Driess 等人，2022）提出了组合 NeRF 以缓解这些限制，但 NeRF-RL 中的 RL 代理在训练和部署期间需要对象个体掩码来利用语义表示，这在现实中非常不切实际。

在这项工作中，我们提出了用于强化学习的语义感知神经辐射场（SNeRL），它为 3D 环境中的 RL 代理学习 3D 感知的语义和几何表示。首先，我们提出的方法通过预测带有地面实况标签的 3D 语义场来学习 3D 感知的语义表示。因此，SNeRL 使下游视觉控制任务无需对象个体掩码，并解决了先前工作（Driess 等人，2022）的限制。此外，为了捕捉语义场中无法充分表达的更细粒度特征，并利用数据驱动方法的优势，我们采用现成的特征描述符（Caron 等人，2021）作为教师网络，并通过蒸馏方法（如 Kobayashi 等人，2022）学习预测特征场。

我们还引入了最近的自预测表示学习（Chen & He，2021）的多视图适应作为辅助任务，进一步提高了 SNeRL 的性能。在提出的辅助任务中，SNeRL 通过利用同一时间步中不同相机视图的观察来计算目标表示，以学习空间一致的表示。

我们提出的 SNeRL 不仅在 RL 的先前单视图表示学习算法（Laskin 等人，2020b；Yarats 等人，2021）中表现出色，而且在具有多视图观察的最先进方法（Driess 等人，2022）中，在四个不同的视觉控制任务中也表现优异。

总结起来，我们的贡献可以概括如下：

- 我们提出了 SNeRL，这是一个利用 NeRF 与语义和蒸馏特征场来学习强化学习的 3D 感知语义表示的框架。

- 我们验证了 SNeRL 在模型无关和基于模型的方法中的有效性。据我们所知，SNeRL 是首个在 RL 下游任务中利用语义感知表示而不需要对象掩码的工作。此外，这是首次利用 3D 感知表示来建模基于模型的 RL 的研究。

- 提出的 SNeRL 在 Meta-world 中的四个不同 3D 环境中优于先前的单视图和多视图基于图像的 RL 算法。此外，提出的多视图观察辅助自预测表示学习可以实现空间一致表示，从而进一步改进。

2. 相关工作

2.1. 3D 场景表示学习

为了从单视图图像中学习 3D 感知表示，先前的方法利用标准卷积自编码器架构，该架构由相机姿态条件化，生成任意视图的场景，使用确定性（Tatarchenko 等人，2016；Worrall 等人，2017）或随机（Eslami 等人，2018）潜在向量。最近，神经辐射场（NeRF）在理解 3D 场景和合成新视图方面取得了显著进展。随后，一些方法提出了潜在条件化的 NeRF（Martin-Brualla 等人，2021；Yu 等人，2021；Wang 等人，2021），但上述方法的主要目标是提高合成图像的质量，而不是从多视图输入中提取具有 3D 动态场景理解的时变潜在向量。在本文中，我们利用带有卷积编码器和 NeRF 风格解码器的自编码器（Li 等人，2022；Driess 等人，2022），以便编码器可以从多视图输入中提取 3D 感知表示，用于 RL 下游任务。

2.2. RL 的表示学习

带有图像输入的 RL 框架通常有一个编码器，将高维观察映射到低维潜在向量。RL 代理在潜在状态空间上进行训练，以最大化其目标函数，例如每个情节的总折扣奖励。尽管许多工作取得了显著进展，但这仍然是一个具有挑战性的开放问题。

为了解决基于图像的 RL 的样本效率问题，先前的工作采用了各种数据增强技术（Laskin 等人，2020a；Yarats 等人，2021），带有数据增强的对比学习（Laskin 等人，2020b；Schwarzer 等人，2020；Stooke 等人，2021；Liu & Abbeel，2021；Zhan 等人，2022），从图像重建中进行表示学习（Islam 等人，2022；Kulkarni 等人，2019），或任务信息重建（Yang & Nachum，2021；Yamada 等人，2022）。其他方法提出捕捉多视图数据之间的关系（Dwibedi 等人，2018；Kinose 等人，2022；Sermanet 等人，2018）或关键点（Manuelli 等人，2020）。还有一些方法利用过渡序列数据（Hansen 等人，2020；You 等人，2022），或使用离线基于图像的 RL 进行预训练（Wang 等人，2022）。不幸的是，这些工作在学习 3D 结构信息方面能力有限，无法获得人类对 3D 环境的直观理解，这是由于 2D 卷积神经网络的 2D 偏差所致。

最近，有一些尝试学习现实世界的 3D 结构（Li 等人，2022；Driess 等人，2022）。Li 等人（2022）首先提出了带有卷积编码器和 NeRF（Mildenhall 等人，2020）解码器的自编码器，以通过学习到的动力学模型和模型预测控制（MPC）来控制视觉运动。随后，NeRF-RL（Driess 等人，2022）扩展了先前的研究，并将基于 NeRF 的架构首次引入到通用的模型无关 RL 框架中。然而，由于简单的 NeRF 的有限 RGB 监督，他们无法学习语义特征。为了仅通过 RGB 监督学习以对象为中心的表示，NeRF-RL 提出了带有对象个体掩码的组合 NeRF，但在 RL 代理部署期间需要掩码似乎是一个强假设。在本文中，我们提出了 SNeRL，它在推理阶段无需任何对象掩码的情况下，通过 RGB、语义和蒸馏特征监督为 RL 下游任务学习几何和语义信息。

3. 预备知识

3.1. 神经辐射场

3.2 强化学习

4. 方法

在本节中，我们将详细介绍 SNeRL，它由一个多视图卷积图像编码器和一个潜在条件化的 NeRF 解码器组成，用于学习 3D 感知表示。与先前的方法（Driess 等人，2022）相比，该方法也为 RL 提出了 NeRF 监督，SNeRL 能够在部署期间无需任何对象个体掩码的情况下提取以对象为中心或语义表示。预训练的图像编码器被用作下游 RL 任务的特征提取器，SNeRL 框架的概述如图 2 所示。

4.1 多视图编码器

4.2 语义感知NeRF解码器

通过估计三种不同的辐射场（语义、特征和 RGB），潜在向量 z 被联合优化以学习 3D 环境的几何和语义表示。与依赖于位置 x 和观察方向 d 的 RGB 值 c 不同，我们将语义标签和蒸馏特征设计为对观察方向 d 不变，因为场景或对象的固有属性不会根据相机光线的方向而改变。

由于 SNeRL 通过添加字段分支来预测三种不同的场（RGB、语义和蒸馏特征），它们共享神经渲染函数 fθ，直到估计密度 σ。这表明三个辐射场在沿光线 r = o + td 的深度 t ∈ [tn, tf ] 处具有相同的累积透射率 T(t)，如下所示：

为了捕捉语义场中无法充分表达的更细粒度特征，SNeRL 还合成了蒸馏特征场（Kobayashi 等人，2022），以知识蒸馏的方式预测预训练特征描述符的输出（Hinton 等人）。根据先前的文献（Caron 等人，2021），以自监督方式训练的视觉变换器（ViT）（Dosovitskiy 等人，2020），例如 DINO（Caron 等人，2021），可以作为一个优秀的特征描述符，明确表示场景布局，如对象边界。由于 ViT 特征描述符的输出包含高维信息，所有像素中的不同值取决于几何关系或语义意义，预训练的 ViT 成为一个具有语义标签优势的良好特征描述符。

因此，我们利用这些优势到基于 NeRF 的解码器中，使得潜在向量 z 通过从 ViT 教师网络中蒸馏知识来学习高级信息，而这些信息无法通过地面实况语义监督学习。蒸馏特征场可以渲染为：

4.3 多视图自预测表示

5. 实验

在本节中，我们在 3D 环境中进行了几个实验，以探索 SNeRL 与现有最先进的 RL 算法在模型无关和基于模型设置中的有效性。我们固定了下游 RL 算法，并在模型无关设置中采用 Soft Actor-Critic（Haarnoja 等人，2018），在基于模型设置中采用 Dreamer（Hafner 等人，2019），用于 SNeRL 和所有基线进行公平评估。

环境 我们在基于 MuJoCo（Todorov 等人，2012）的四个视觉控制环境中评估了 SNeRL，包括一些需要巧妙利用对象间交互以获得高奖励的复杂控制任务。所有任务均由模拟的 Sawyer 机器人执行，该机器人具有单臂和手上的夹持器（4-DoF）。Sawyer 机器人的动作空间包括末端执行器的位置（x,y,z）和夹持器控制（打开/关闭）。代理从三个不同相机视图获取 128x128 图像作为像素级输入，并从 Meta-world（Yu 等人，2020）提供的环境中接收密集奖励。

- **Window-open-v2**：此环境涉及 Sawyer 机器人打开带有把手的滑动窗户。机器人手的初始状态为 [0, 0.4, 0.2]，机器人通过推动把手并打开位于 [-0.1, 0.785, 0.16] 的窗户来获得奖励。

- Hammer-v2：Sawyer机器人应该抓住随机位置生成的锤子的手柄，并击打钉子的头部以将其钉入。机器人手的初始状态在\(\{(x, y, z)| -0.5 \leq x \leq 0.5, 0.4 \leq y \leq 1, 0.05 \leq z \leq 0.5\}\)中随机生成。机器人在捡起锤子并将钉子插入木头时获得奖励。

- Drawer-open-v2：Sawyer机器人应该通过抓住抽屉的手柄并拉动它来打开抽屉。机器人手的初始状态与Hammer-v2相同。机器人在打开抽屉时获得奖励。

- Soccer-v2：在这个任务中，Sawyer机器人试图通过推动一个随机位置生成的足球来得分。机器人手的初始状态与Hammer-v2相同。机器人在触碰足球并将其放入球门时获得奖励。

我们参考Meta-world（Yu等人，2020年）以获取更多细节，包括奖励函数和随机位置的范围。

基线 我们将 SNeRL 与几种最先进的视觉 RL 方法和一种 3D 感知 RL 方法进行了比较，简要描述如下。DrQ-v2（Yarats 等人，2021）是 DrQ（Yarats 等人，2020）的改进版本，通过数据增强和计划探索噪声解决视觉控制任务。CURL（Laskin 等人，2020b）通过辅助对比损失训练 RL 代理，确保数据增强版本的观察的嵌入匹配。CURL-multiview 是 CURL 的多视图适应版本，利用 3 个不同相机视图，并具有与 SNeRL 相同结构的 CNN 编码器。CNN-AE 使用标准 CNN 自编码器（而不是 NeRF 解码器），使用 Finn 等人（2016）提出的重建损失预训练编码器。NeRF-RL（Driess 等人，2022）使用卷积编码器和简单的 NeRF 风格解码器预训练自编码器，没有语义和特征监督。

我们注意到，CNN-AE 和 NeRF-RL 中的下游 RL 任务学习与 SNeRL 使用相同的方法，并且它们使用由随机动作和 Meta-world 提供的策略（一半一半混合）收集的相同离线数据集。我们建议读者参阅附录 B.2 了解其他数据集的实验。此外，所有多视图方法（CURL-multiview、CNN-AE、NeRF-RL、SNeRL）接收相同的观察，并且不从环境中接收每个对象的掩码。对于其余在单视图上操作的基线，我们选择一个可以从该位置清晰观察每个对象状态的单相机位置。

5.1. 实验结果

图 3 显示了 SNeRL 和基线在 4 个不同视觉控制任务中的情节回报。得益于通过语义和蒸馏特征监督学习的以对象为中心的表示，SNeRL 在数据效率和性能方面始终优于最先进的视觉 RL 方法和先前的 3D 感知 RL 方法（NeRF-RL）。

具体来说，对比基线（CURL、CURL-multiview）和 DrQ-v2 在困难环境（足球和锤子）中无法实现高回报，尽管它们在相对简单的环境（窗户）中成功。结果还显示，通过离线数据使用简单的重建损失（CNN-AE）预训练 CNN 在所有环境中都无法成功。这表明从多视图观察中提取 3D 感知几何信息以及以对象为中心和语义信息对 RL 性能至关重要。

有趣的是，我们观察到仅通过 RGB 监督预训练基于 NeRF 的自编码器（NeRF-RL）不足以学习 RL 下游任务的特征，并且它无法超越带有对比损失的视觉 RL 方法的多视图适应版本（CURL-multiview）。这与先前工作（Driess 等人，2022）报告的结果相反，我们分析如下：我们在这项工作中采用的环境比 Driess 等人（2022）的环境更具挑战性，后者由简单形状和原色的对象组成。因此，仅使用 RGB 监督相对难以获得语义信息。因此，利用语义感知的 NeRF 解码器是提取特征以在 3D 感知 RL 的实际应用中在 RL 下游任务中获得更好性能的必要条件，这与我们的分析一致。

5.2. 消融研究

**语义和蒸馏特征场**。为了验证 SNeRL 中每个语义感知辐射场（语义和蒸馏特征场）对下游 RL 性能的贡献，我们分别评估了没有语义和特征监督的两个消融变体。如图 4(a) 所示，利用地面实况标签和基于 ViT 的特征描述符的语义和特征监督的 SNeRL 在所有消融模型中实现了最佳性能。我们观察到 SNeRL 与消融模型之间的性能差距取决于环境，因为语义标签在相对简单的环境（窗户）中足以学习语义信息，但在复杂环境（足球）中需要两者监督。

**多视图表示学习**。我们还引入了适合多视图观察的自预测损失的辅助表示学习。为了证明其在下游 RL 任务中的有效性，我们评估了两种不同模型，即带有和不带有多视图自预测表示学习的 SNeRL。通过强制潜在向量对观察视点不变，我们报告称，所提出的表示学习在一些环境中提高了 RL 代理的性能，如图 4(b) 所示。

5.3. 通过神经渲染的图像重建

尽管我们仅利用卷积编码器进行下游 RL 任务，但我们比较了 SNeRL 和 NeRF-RL 中基于 NeRF 的解码器的图像渲染性能，以探索合成图像质量与 RL 性能之间的关系。由于 NeRF 最初旨在从静态场景中合成任意相机视图的图像，仅通过 RGB 监督训练体积场的 NeRF-RL 无法在没有语义信息的情况下重建输入图像中的动态对象，例如机器人手臂。另一方面，利用 ViT 教师网络的语义标签和特征输出作为额外监督信号的 SNeRL 不仅实现了更好的 RL 性能，而且很好地表示了动态场景，并产生了高保真渲染输出，如图 5 所示。

5.4. SNeRL 用于基于模型的 RL

在本节中，我们评估通过 SNeRL 学习的表示是否也可以应用于现成的基于模型的强化学习算法，这些算法训练一个世界模型来表征环境并在学习到的模型上进行规划。我们采用 Dreamer（Hafner 等人，2019）作为下游基于模型的 RL 代理，并用我们预训练的编码器替换表示模型的编码器。有关额外的实现细节和架构概述，请参阅附录 A。

我们的结果如图 4(c) 所示。我们观察到，使用 SNeRL 预训练编码器学习基于模型的 RL 优于先前 3D 感知 RL 方法（NeRF-RL）和简单 CNN 自编码器的预训练权重。这一实证证据与第 5.1 节中模型无关 RL 的情况一致，表明所提出的方法允许编码器学习对通用现成 RL 代理重要的表示。

6. 结论

在本文中，我们提出了 SNeRL，一种用于 RL 的语义感知辐射场，它在四个不同 3D 环境中优于现有的 RL 算法表示学习方法。SNeRL 利用潜在条件 NeRF 自编码器以及 RGB 监督的语义和蒸馏特征监督，使图像编码器能够在下游 RL 任务中表达 3D 感知的几何和语义表示。我们还提出了多视图自预测损失作为辅助表示学习，以强制潜在向量对视点不变。最后，我们验证了 SNeRL 在模型无关和基于模型的 RL 算法中都是有效的。

局限性 尽管有这些改进，SNeRL 继承了先前 3D 感知 RL 方法的局限性。首先，SNeRL 需要多视图离线数据，并且在一些复杂控制任务中收集覆盖状态空间的离线数据集可能具有挑战性。此外，我们的方法使用消耗更多计算预算的 NeRF 解码器，因此在扩展我们的方法以在线设置时可能存在限制，该设置同时训练编码器和 RL 代理。

https://proceedings.mlr.press/v202/shim23a/shim23a.pdf