ICLR2022：清华、腾讯AI Lab共同提出等变图力学网络，实现多刚体物理系统模拟...

来源：机器之心
本文约3000字，建议阅读10分钟
本文介绍了图力学网络（GMN）的构造与理论分析。

近年来，AI for Science 利用人工智能方法与物理、化学、生物等自然科学进行交叉融合，在一些重要的科学问题上（如蛋白质结构预测）取得了瞩目的进展。鉴于物理学科的基础性与重要性，AI+Physics 无疑是 AI for Science 不可缺失的一环。

为此，清华 AIR、计算机系与腾讯 AI Lab 合作共同发表论文《Equivariant Graph Mechanics Networks with Constraints》, 提出等变图力学网络，实现了理论力学中的一类重要任务—多刚体系统模拟。论文已被 ICLR2022 接收。这一项目也收到了来自腾讯 AI Lab 犀牛鸟专项研究计划的资助。 

论文地址：https://openreview.net/forum?id=SHbhHHfePhP

开源代码：https://github.com/hanjq17/GMN

一、引言

多体问题（N-body problem）是物理学的基础问题，广泛存在于诸多任务中，例如电子绕原子核运动、分子动力学模拟、机械系统（如机械臂）的动力控制、宇宙天体的轨迹预测等，都可以表示成多体问题。多体问题通常不存在解析解，其求解需要冗长的数值计算。近年来，机器学习研究人员利用图神经网络（Graph Neural Network, GNN）从数据中挖掘多体系统的演化规律，为多体问题的求解提供了一个全新的思路。该类方法把物体建模成节点，把物体之间的交互建模成连边，训练好的 GNN 直接预测系统的未来状态。

图 1 从多体问题到多刚体系统模拟

尽管如此，现有方法仍然没有很好实现对多刚体系统（Multi-rigid system）的模拟。多体问题一般假设每个物体是独立粒子，而在实际问题中，物体之间很可能存在几何约束，形成多刚体系统。实际上，多体问题是多刚体系统的一种特殊形式。如图 1（b）所示，粒子和粒子之间通过连杆连接（连杆两端的粒子只能相对转动），不同连杆可以通过铰链连接（铰链两端的连杆只能相对转动）。多刚体系统是理论力学模型中的最常见形式，在机械系统动力学控制中尤为普遍。

为了更好地刻画多刚体系统的动力学规律，本文提出一种新颖的图神经网络——图力学网络（Graph Mechanics Network, GMN）。GMN 具备几个优点:

1.GMN 满足刚体的几何约束，输出不会改变输入的几何属性（如棍子的长度、铰链的连接等）。

2.GMN 是欧氏等变（Equivariant）的，即对输入做任何平移、旋转、翻转，输出都相应地改变。这利用了物理规律关于空间是均匀对称的这一特点，作为模型的归纳偏置，很好地提高了模型的泛化性和解释性。

3. 理论上，GMN 具有良好的表达能力，能很好地逼近任何等变的函数。

我们构造了一个由一定数量的球、棍子和铰链组成的虚拟物理系统，GMN 比其他方法更准确地预测这些系统演变，并满足上述所说的性质。此外，在真实的应用场景包括分子动力学模拟和人体骨架轨迹预测等，GMN 也表现出最好的预测效果。

二、方法简介

图 2 刚体系统的两种表示方法：左子图利用笛卡尔坐标描述系统状态，不同坐标存在几何约束；右子图利用广义坐标描述系统状态，不同坐标完全解耦

我们的任务是给定多刚体系统的初始状态，利用提出的图力学网络预测其未来状态。在理论力学中，为了满足刚体系统的几何约束，需要引入广义坐标来表示系统的状态。如图 2 所示，我们以棍子为例，在左子图中，棍子两端球 1、球 2 的笛卡尔坐标分别是。

因为存在棍子长度约束，所以两个球的运动过程中需要满足两者相对距离不变；换句话说，利用笛卡尔坐标系表示系统的状态，需要引入额外的约束条件，不便于采用神经网络来端到端地表征系统演化规律。为了克服这个困难，我们引入广义坐标系。如图 2 右子图所示，可以通过棍子的中心 0 点的笛卡尔坐标，以及棍子的相对转角，来表示棍子的状态。同时，可以独立变化，始终不改变棍子的长度，特别利于使用神经网络进行建模。

具体而言，为了刻画多刚体系统的变化过程，我们的方法主要包含 4 部分：

1. 在笛卡尔坐标空间计算每个球受到其他刚体的作用力；

2. 利用逆向动力学把每个刚体的受力转化成广义坐标的加速度变化；

3. 对广义坐标进行更新；

4. 利用前向运动学，将更新后的广义坐标映射到笛卡尔坐标空间。

值得指出的是，受力的计算、逆向动力学、前向运动学都是通过等变神经网络来实现，极好地满足了物理学的对称性。如果我们把整个系统的初始状态旋转 90 度，我们的模型对未来状态的预测也会精准地旋转 90 度。通过引入这样的物理先验，我们能提高模型的泛化性和可解释性。

下面简要介绍每一部分过程。

1. 相互作用力的计算。在本文中，我们认为不同刚体中不同的球会发生相互作用，即所有的球组成全连接图。对于每个球i，它的受力的计算为：，即，每两个球的受力与其相对距离有关。

2. 逆向动力学。基于每个刚体上所有受力及其分布，我们计算广义坐标的加速度为：，是相对位置，是相对速度。直观上，这个公式是基于刚体上每个球的受力及其在刚体局部坐标下（也就是广义坐标系）的状态推算出广义坐标系的加速度。

3. 广义坐标更新。有了加速度，我们就可以更新广义坐标的速度和位置：。

4. 前向运动学。得到了更新后的广义位置和速度，我们就可以通过前向运动学对每个刚体内每个球的笛卡尔坐标进行更新：。需要指出的是，FK 可以是手工设置（Hard version），保证笛卡尔坐标更新始终满足约束；FK 也可以是自动学习的（Soft version），从数据中学习约束，但是不能精确保证约束条件。

为了简便，以上过程没有考虑每个节点的隐层变量，和广义坐标角度。读者可以通过查阅原论文得到更多这方面的细节。

三、理论分析

上文提到的函数均为欧氏等变函数。不失一般性，我们记正交等变函数。我们定义，

其中，为参数是的任意多层感知机（MLP）。很容易可以证明上述定义的等变性，即对于任意正交矩阵。实际上，我们这里的定义是 EGNN[1]的更一般形式，从向量形式扩展为矩阵形式。

我们还可以证明，这样定义的正交等变函数具有很好的表达能力。

【定理 1】如果，行满秩，即，那么对于任意连续的正交等变函数，一定存在某个使得充分逼近，即,任意小。

三、实验结果

1. 为了验证 GMN 的有效性，我们构造了一个仿真数据集 Constrained N-body。该数据集的每个数据（也就是每个系统）是由一定比例的孤立球、棍子、铰链组成。下图，GMN 利用连续的初始状态预测未来多帧状态。通过对比 GMN 预测效果（蓝色）与真实状态（红色），可以看到，GMN 能非常准确地拟合真实的动态过程。

图 3 多刚体系统仿真结果（蓝色为预测结果，红色为真实状态）

以下表格提供了在不同刚体组合的系统中，GMN 与当前方法的预测精度对比，可以看出 GMN 输出更小的预测误差，显著优于其他方法。 

2. 除了仿真数据，我们还在两个真实数据集：人体运动捕捉 CMU Motion Capture 和分子动力学轨迹 MD17，验证 GMN 在实际应用中的潜力。

2.1 CMU Motion Capture 是一个记录人类动作捕捉轨迹的数据集。在这个数据中，人体骨架可以看成是一种刚性约束。在 CMU Motion Capture 上，GMN 的预测轨迹几乎与真实轨迹重合。

图 4 人体骨架运动捕捉（蓝色为预测结果，红色为真实状态）

以下是不同方法的预测误差：

2.2 MD17 是分子动力学模拟的经典数据集，其包含 8 种分子的动力学运动轨迹，以下展示了 8 个分子的的分子结构：

在这个数据集上的任务是，根据当前状态分子构象预测间隔 t 时间后分子的状态。下表是 GMN 和各个方法的对比结果。

四、结论

本文介绍了图力学网络（GMN）的构造与理论分析。借助广义坐标，GMN 能有效刻画几何约束；借助等变神经网络，GMN 能满足物理对称性。在多刚体仿真系统 Constrained N-body、人体骨架预测 CMU Motion Capture、分子动力学模拟 MD-17 等任务上都验证了 GMN 的有效性。未来，GMN 可以用于更复杂、更大体系的物理系统动力学模拟。

[1] Satorras V G, Hoogeboom E, Welling M. E (n) equivariant graph neural networks[C]//International Conference on Machine Learning. PMLR, 2021: 9323-9332.

编辑：黄继彦