[论文解读] Reinforcement Learning with Neural Radiance Fields

Danny Driess, Ingmar Schubert+3 more|arXiv (Cornell University)|Jun 3, 2022

Adversarial Robustness in Machine Learning被引用 15

一句话总结

该论文提出NeRF-RL，一种利用神经辐射场（NeRFs）作为自监督信号来训练强化学习（RL）潜在状态表征的方法。通过训练编码器将多视角RGB图像映射到一个潜在空间，该空间可通过潜在条件化的NeRF重建场景，该方法在机器人操作任务中实现了卓越的样本效率——优于2D CNN、对比学习，甚至优于人工提供的关键点——这是由于NeRF强大的3D归纳偏置。

ABSTRACT

It is a long-standing problem to find effective representations for training reinforcement learning (RL) agents. This paper demonstrates that learning state representations with supervision from Neural Radiance Fields (NeRFs) can improve the performance of RL compared to other learned representations or even low-dimensional, hand-engineered state information. Specifically, we propose to train an encoder that maps multiple image observations to a latent space describing the objects in the scene. The decoder built from a latent-conditioned NeRF serves as the supervision signal to learn the latent space. An RL algorithm then operates on the learned latent space as its state representation. We call this NeRF-RL. Our experiments indicate that NeRF as supervision leads to a latent space better suited for the downstream RL tasks involving robotic object manipulations like hanging mugs on hooks, pushing objects, or opening doors. Video: https://dannydriess.github.io/nerf-rl

研究动机与目标

• 通过从高维图像观测中学习更好的状态表征，提升视觉强化学习的样本效率。
• 探究神经辐射场（NeRFs）的3D归纳偏置是否能带来比2D CNN或对比学习更有效的潜在表征。
• 在不依赖人工设计状态信息的前提下，实现在多样化物体形状下的泛化能力。
• 证明基于NeRF的监督可超越学习得到的和人工提供的状态表征，在复杂3D操作环境中表现更优。

提出的方法

• 训练一个自编码器，其中编码器将多视角RGB图像映射到潜在向量，解码器为基于潜在条件的NeRF，可从新视角重建场景。
• 使用多视角图像重建损失作为自监督信号，联合训练编码器与解码器，其中NeRF解码器提供强大的3D归纳偏置。
• 实现全局与组合式NeRF解码器：前者编码完整场景，后者使用实例分割掩码编码单个物体。
• 在随机动作数据集上离线预训练编码器，随后在强化学习循环中使用固定编码器微调策略。
• 利用相机矩阵和多视角图像提升3D场景理解能力，以解决遮挡问题。
• 在预训练阶段应用NeRF重建损失，引导潜在空间捕捉与任务相关的几何与空间结构。

实验结果

研究问题

• RQ1 基于NeRF的自监督是否能带来比2D CNN或对比学习更有效的视觉强化学习状态表征？
• RQ2 NeRF的3D归纳偏置是否能提升涉及可变物体形状的机器人操作任务中的样本效率与泛化能力？
• RQ3 在下游强化学习性能方面，NeRF-RL与人工提供的关键点监督相比如何？
• RQ4 使用建模单个物体的组合式NeRF解码器是否能优于全局NeRF解码器？
• RQ5 使用多视角图像是否能提升所学习潜在表征的质量？

主要发现

• NeRF-RL在机器人操作任务中的样本效率和最终性能上显著优于基于2D CNN的自编码器和对比学习基线方法。
• 组合式NeRF-RL变体实现了最高的样本效率，表明使用NeRF建模单个物体比全局场景建模更有效。
• NeRF-RL的性能甚至超过了人工提供的关键点监督，证明了NeRF解码器中3D归纳偏置的优越性。
• 该方法在同分布内的不同起始状态和物体形状下均表现出良好的泛化能力，显示出鲁棒性与泛化潜力。
• 即使不依赖物体掩码，全局NeRF-RL变体仍优于所有基线方法，表明3D归纳偏置在无实例级监督下依然有效。
• 尽管预训练成本较高（最高达2天），但推理时间仍保持在较低水平（RTX 3090上约7 ms），使该方法适用于在线强化学习部署。