拉格朗日神经网络

迈尔斯·克兰默、萨姆·格雷达努斯、施特凡·霍耶尔、彼得·巴塔利亚、大卫·斯珀格尔、雪莉·霍

• 论文
• 博客
• 自包含教程
• 论文示例笔记本：双摆
• 论文示例笔记本：狭义相对论
• 论文示例笔记本：波动方程

[!警告] 本项目是使用 JAX 0.1.55（2020年）开发的。为保持兼容性，请使用 pixi 安装原始环境（参见指南）。或者，您可以根据下方的核心运动方程片段为现代 JAX 进行适配。

摘要

在本项目中，我们提出了拉格朗日神经网络（LNNs），它能够利用神经网络参数化任意拉格朗日量。与哈密顿神经网络不同，这些模型不需要正则坐标，并且在广义动量难以计算的情况下表现良好（例如双摆系统）。这一点对于结合学习到的潜在表征尤为吸引人，而哈密顿神经网络在这种情况下往往难以处理。与先前关于学习拉格朗日量的工作相比，LNNs 具有完全通用性，并可扩展到非完整系统，如一维波动方程。

	神经网络	神经ODE	HNN	DLN (ICLR'19)	LNN（本工作）
学习动力学	✔️	✔️	✔️	✔️	✔️
学习连续时间动力学		✔️	✔️	✔️	✔️
学习精确守恒定律			✔️	✔️	✔️
从任意坐标学习	✔️	✔️		✔️	✔️
学习任意拉格朗日量					✔️

核心运动方程

LNNs 的关键创新在于能够自动从学习到的拉格朗日量推导出运动方程。核心运动方程是版本无关的，可以适配任何 JAX 版本：

import jax
import jax.numpy as jnp

def lagrangian_eom(lagrangian, state, t=None):
    """计算欧拉-拉格朗日运动方程。

    参数:
        lagrangian: 表示拉格朗日量的函数 L(q, q_dot)
        state: 位置 q 和速度 q_dot 的拼接向量

    返回:
        速度和加速度的拼接向量
    """
    q, q_dot = jnp.split(state, 2)

    # 欧拉-拉格朗日方程：d/dt(∂L/∂q̇) - ∂L/∂q = 0
    # 重新排列以求解加速度 q_ddot
    q_ddot = (jnp.linalg.pinv(jax.hessian(lagrangian, 1)(q, q_dot))
              @ (jax.grad(lagrangian, 0)(q, q_dot)
                 - jax.jacobian(jax.jacobian(lagrangian, 1), 0)(q, q_dot) @ q_dot))

    return jnp.concatenate([q_dot, q_ddot])


# 用于学习拉格朗日量的 MLP 自定义初始化
def custom_init(layer_sizes, seed=0):
    """初始化 MLP，使其更有效地学习拉格朗日量"""
    # 完整实现请参见 lnn/core.py
    # 使用每层优化后的尺度

该运动方程适用于任何 JAX 版本。代码库的其余部分则使用 JAX 0.1.55 维护，以确保与论文发表时的环境一致，从而实现结果的可重复性。

安装

为确保结果的可重复性，本项目需要特定版本的 JAX 及其他依赖项，这些依赖项均来自 2020 年初。我们使用 pixi 来管理可重复的环境。

快速入门

1. 安装 pixi（如果尚未安装）：

   curl -fsSL https://pixi.sh/install.sh | bash
   

   或通过 Homebrew：

   brew install pixi
   

2. 克隆本仓库并安装依赖项：

   git clone https://github.com/MilesCranmer/lagrangian_nns.git
   cd lagrangian_nns
   pixi install
   

3. 运行笔记本：

   pixi run jupyter notebook
   

依赖项

环境包含以下内容：

• Python 3.7
• JAX 0.1.55 和 jaxlib 0.1.37（2020年1月版本）
• NumPy 1.19.5
• Matplotlib 3.1.2（可视化工具）
• MoviePy 1.0.0（可视化工具）
• celluloid 0.2.0（可视化工具）

所有依赖项都被固定在 2020 年初的版本上，以确保代码的行为与论文发表时完全一致。

Lagrangian Neural Networks (LNN) 快速上手指南

Lagrangian Neural Networks (LNN) 是一种利用神经网络参数化任意拉格朗日量的模型。与哈密顿神经网络（HNN）不同，LNN 不需要正则坐标，特别适用于广义动量难以计算的场景（如双摆），且能处理非完整系统（如一维波动方程）。

环境准备

本项目基于 2020 年初的 JAX 版本（0.1.55）开发以确保论文结果的可复现性。为了兼容旧版依赖，官方推荐使用 pixi 进行环境管理。

• 操作系统: Linux 或 macOS (Windows 用户建议使用 WSL2)
• 前置工具: git, curl (或 brew)
• 核心依赖:
  • Python 3.7
  • JAX 0.1.55 & jaxlib 0.1.37
  • NumPy 1.19.5
  • Matplotlib, MoviePy, celluloid (用于可视化)

注意：由于依赖版本较老，不建议直接在现代 Python 环境中通过 pip 安装，务必使用下文提供的 pixi 方案以隔离环境。

安装步骤

1. 安装 Pixi

如果尚未安装 pixi，请通过以下命令安装：

Linux/macOS (通用脚本):

curl -fsSL https://pixi.sh/install.sh | bash

macOS (Homebrew):

brew install pixi

2. 克隆项目并安装依赖

获取源代码并使用 pixi 创建包含特定版本依赖的虚拟环境：

git clone https://github.com/MilesCranmer/lagrangian_nns.git
cd lagrangian_nns
pixi install

3. 启动开发环境

安装完成后，使用以下命令启动 Jupyter Notebook 以运行示例：

pixi run jupyter notebook

基本使用

LNN 的核心在于自动从学习到的拉格朗日量中推导运动方程。以下是最小化的核心代码示例，展示了如何定义欧拉 - 拉格朗日运动方程。此代码片段独立于 JAX 版本，可直接在现代环境中测试逻辑。

核心运动方程实现

import jax
import jax.numpy as jnp

def lagrangian_eom(lagrangian, state, t=None):
    """Compute Euler-Lagrange equation of motion.

    Args:
        lagrangian: A function L(q, q_dot) representing the Lagrangian
        state: Concatenated position q and velocity q_dot

    Returns:
        Concatenated velocity and acceleration
    """
    q, q_dot = jnp.split(state, 2)

    # Euler-Lagrange equation: d/dt(∂L/∂q̇) - ∂L/∂q = 0
    # Rearranged to solve for acceleration q_ddot
    q_ddot = (jnp.linalg.pinv(jax.hessian(lagrangian, 1)(q, q_dot))
              @ (jax.grad(lagrangian, 0)(q, q_dot)
                 - jax.jacobian(jax.jacobian(lagrangian, 1), 0)(q, q_dot) @ q_dot))

    return jnp.concatenate([q_dot, q_ddot])

运行示例

安装完成后，项目根目录下的 notebooks 文件夹包含了完整的教程和论文复现实例：

• LNN_Tutorial.ipynb: 自助式入门教程。
• DoublePendulum.ipynb: 双摆系统示例。
• SpecialRelativity.ipynb: 狭义相对论示例。
• WaveEquation.ipynb: 波动方程示例。

在启动的 Jupyter 界面中打开上述任意笔记本即可开始实验。