adaptive_prediction 自适应动力学预测实战：基于 JAX 与 PyBullet 的刚体运动在线预测指南-平芜编程栈

人工智能
深度学习
NLP
计算机视觉
强化学习

【免费下载链接】google-research

Google Research

项目地址：https://gitcode.com/gh_mirrors/go/google-research

本篇技术指南以 google-research 仓库中的 adaptive_prediction 包为核心，系统讲解"自适应动力学预测"（Adaptive Dynamics Prediction）这一在线学习方法：预测器在滚动预测轨迹的同时，利用测量反馈实时更新模型参数。文章从数学框架、核心类实现、超参数配置到两类刚体动力学实验（PyBullet 抛掷小球、自实现旋转动力学）逐层展开，读者学完后可复现 Notebook 中的完整预测-适应闭环，并具备将其迁移到其他动力学系统（含镜面下降势函数扩展）的动手能力。

包概览：README 与自包含 Notebook

adaptive_prediction包的结构非常精简，全部功能集中在三个文件中：

文件	作用
README.md	包概览与使用说明
rigid_body_adaptive_prediction.ipynb	自包含的完整实现与演示（交互式 Notebook）
requirements.txt	运行依赖清单

README 明确指出："该包包含通过自适应动力学预测对动力系统进行在线学习的代码，代码自包含在rigid_body_adaptive_prediction.ipynb中，并以刚体旋转动力学作为完整的演练示例（worked-through example）。"

Notebook 自身的"Notebook overview"小节进一步说明了内容组织：它提供自适应动力学预测的交互式演示，数据来自 PyBullet 对抛掷物体（thrown objects）的仿真输出；先描述自适应预测器的通用结构，再从数学上为刚体动力学构造一个自适应预测器，最后给出具体实现。因此，本文的正文脉络严格遵循 Notebook 的这一三段式结构。

运行环境与快速启动

根据 requirements.txt，运行该 Notebook 需要以下依赖：

jupyter>=6.4.0 matplotlib>=3.4 numpy>=1.20.3 jax>=0.2 pybullet>=3.05 seaborn>=0.11.1

启动方式即 README 给出的唯一使用指令：

jupyter notebook rigid_body_adaptive_prediction.ipynb

Notebook 的第一批代码单元会先通过 pip 安装 PyBullet 相关的仿真依赖（pybullet与bullet），随后引入全部库。值得注意的工程细节是：Notebook 在导入 JAX 后立即执行config.update('jax_enable_x64', True)开启 64 位浮点运算，这是数值精度敏感的动力学仿真所必需的（默认的 float32 会在长时间积分中积累不可忽视的误差）。

import pybullet as p import bullet as b import pybullet_data import numpy as onp import jax from jax import numpy as jnp from matplotlib import pyplot as plt import matplotlib as mpl import seaborn as sns import functools from jax import config config.update('jax_enable_x64', True)

从源码结构看，Notebook 全程使用jax.jit对热路径函数做即时编译加速，并使用jax.lax.scan实现批量 rollout，这与 JAX 的惯用写法保持一致。

自适应预测的通用数学框架

Notebook 中"Adaptive prediction: general formulation"一节给出了自适应预测器的通用连续时间形式（这也是理解整个包的理论入口）：

$$\dot{\hat{\mathbf{x}}} = \mathbf{f}(\hat{\mathbf{x}}, t) + \mathbf{Y}(\hat{\mathbf{x}}, t){\text{phys}}\hat{\boldsymbol{\alpha}}{\text{phys}} + \mathbf{Y}(\hat{\mathbf{x}}, t){\text{rf}}\hat{\boldsymbol{\alpha}}{\text{rf}} + \mathbf{g}(\mathbf{x}, \hat{\mathbf{x}})$$

各组成部分的含义如下：

$\mathbf{f}(\hat{\mathbf{x}}, t)$：动力学中已知的部分。例如对刚体动力学，我们已知姿态四元数（orientation quaternion）的动力学，以及速度中两个分量的动力学；
$\mathbf{g}(\mathbf{x}, \hat{\mathbf{x}})$：测量函数，用于在有测量可用时将测量信息注入预测器（本文采用简单的线性反馈 $k(\mathbf{x} - \hat{\mathbf{x}})$）；
$\mathbf{Y}(\hat{\mathbf{x}}, t){\text{phys}}\hat{\boldsymbol{\alpha}}{\text{phys}}$：物理近似项，$\mathbf{Y}{\text{phys}}$ 为物理基函数，$\hat{\boldsymbol{\alpha}}{\text{phys}}$ 为对应参数（例如惯性张量 $\mathbf{J}$ 的过参数化线性展开）；
$\mathbf{Y}(\hat{\mathbf{x}}, t){\text{rf}}\hat{\boldsymbol{\alpha}}{\text{rf}}$：无结构近似项，$\mathbf{Y}{\text{rf}}$ 为非结构化基函数（此处选用随机特征Random Features），$\hat{\boldsymbol{\alpha}}{\text{rf}}$ 为对应参数。

两组参数分别按照如下连续时间更新律（镜像下降形式的自适应律）演化：

$$ \frac{d}{dt}\nabla\psi_{\text{phys}}\left(\hat{\boldsymbol{\alpha}}{\text{phys}}\right) = - \eta_p \mathbf{Y}\left(\hat{\mathbf{x}}, t\right){\text{phys}}^{\mathsf{T}}\mathbf{e}{\text{phys}}(\mathbf{x}, \hat{\mathbf{x}}), \qquad \frac{d}{dt}\nabla\psi{\text{rf}}\left(\hat{\boldsymbol{\alpha}}{\text{rf}}\right) = - \eta{\text{rf}} \mathbf{Y}\left(\hat{\mathbf{x}}, t\right){\text{rf}}^{\mathsf{T}}\mathbf{e}{\text{rf}}(\mathbf{x}, \hat{\mathbf{x}}) $$

其中：

$\psi_i$：镜像下降（mirror descent）势函数，用于正则化两组预测器动力学，或根据基函数的几何结构改进自适应行为。Notebook 实现了三类势函数：欧几里得（euclid）、p 范数（p）、超熵（hyp，hypentropy）；
$\mathbf{e}_i(\cdot, \cdot)$：驱动自适应的误差信号；
$\eta_i$：两组参数向量各自的学习率。

面向刚体动力学的数学实例化

在"Adaptive prediction for rigid-body dynamics"一节中，上述通用框架被实例化到刚体动力学。刚体的旋转与平动运动在连续时间下可写为：

$$ \dot{\mathbf{q}} = \frac{1}{2}\boldsymbol{\omega}\circ\boldsymbol{q},\qquad \dot{\mathbf{p}} = \mathbf{v},\qquad \dot{\mathbf{v}} = -\mathbf{g},\qquad \dot{\boldsymbol\omega} = - \mathbf{J}^{-1}\boldsymbol{\omega}\times\mathbf{J}\boldsymbol{\omega} $$

其中 $\circ$ 表示四元数乘积，$\times$ 表示叉积，$\mathbf{J}$ 为惯性张量矩阵。Notebook 的两个关键坐标系选择值得强调：

平动动力学写在世界坐标系，因此与旋转动力学解耦；
旋转动力学写在刚体坐标系，此时惯性张量 $\mathbf{J}$ 为常量，便于参数学习。

由于唯一的未知物理参数是 $\mathbf{J}$ 的分量，所以 $\mathbf{q}$、$\mathbf{p}$、$\mathbf{v}$ 的动力学是已知的。虽然 $- \mathbf{J}^{-1}\boldsymbol{\omega}\times\mathbf{J}\boldsymbol{\omega}$ 对未知惯性矩阵是非线性的，但可以写成更大维数参数集上的线性展开（这正是代码中inertial_basis构造 18 维过参数化基函数的原因）。

对应的自适应预测器方程为：

$$ \dot{\hat{\mathbf{q}}} = \frac{1}{2}\hat{\boldsymbol{\omega}}\circ\hat{\boldsymbol{q}} + k(\mathbf{q} - \hat{\mathbf{q}}),\qquad \dot{\hat{\mathbf{p}}} = \hat{\mathbf{v}} + k(\mathbf{p} - \hat{\mathbf{p}}),\qquad \dot{\hat{\mathbf{v}}} = -\mathbf{g} + k\left(\mathbf{v} - \hat{\mathbf{v}}\right),\qquad \dot{\hat{\boldsymbol\omega}} = \mathbf{Y}(\hat{\boldsymbol\omega})\hat{\boldsymbol\alpha}_{\text{phys}} + k(\boldsymbol\omega - \hat{\boldsymbol\omega}) $$

物理参数的更新律为：

由于参数估计 $\hat{\boldsymbol{\alpha}}{\text{phys}}$ 只出现在 $\dot{\hat{\boldsymbol\omega}}$ 动力学中，因此物理误差信号取 $\mathbf{e}{\text{phys}}(\mathbf{x}, \hat{\mathbf{x}}) = \hat{\boldsymbol\omega} - \boldsymbol{\omega}$。若再向整个预测器添加随机特征，则其参数更新为：

$$\frac{d}{dt}\nabla\psi_{\text{rf}}\left(\hat{\boldsymbol{\alpha}}{\text{rf}}\right) = - \eta{\text{rf}} \mathbf{Y}\left(\hat{\mathbf{x}}, t\right)_{\text{rf}}^{\mathsf{T}}\left(\hat{\mathbf{x}} - \mathbf{x}\right)$$

此时误差信号取全状态误差$\hat{\mathbf{x}} - \mathbf{x}$；Notebook 同时指出，若只计划估计少数分量的扰动，随机特征的自适应可仅由这些分量的误差驱动，而不必是完整预测器状态的函数——代码中rf_error_signal正是只用了线速度分量的误差。

AdaptivePredictor：核心类与源码解读

Notebook 以AdaptivePredictor类（位于"Class for adaptive dynamics prediction"代码单元）承载全部实现逻辑。该类将"连续时间动力学定义"与"离散化更新"分离，是整个包架构的核心。

类属性与状态

从类定义可见其维护的状态包括：

系统信息：sys_dim（系统状态维度）、unknown_dim；
时间信息：dt（预测器步长）、nsteps（两次测量之间的预测步数）、nframes（帧数，即观测次数）、meas_count（测量计数）、start_ind、prediction_index；
自适应参数信息：params（当前参数估计）、n_features（随机特征数）、n_params（总参数数）；
镜像映射：grad_psi_phys/inverse_grad_psi_phys、grad_psi_rf/inverse_grad_psi_rf（物理与随机特征两套正反镜像映射）；
连续时间形式：dynamics、dynamics_with_measurement；
离散化更新：predict_state、predict_state_with_measurement、norm_state；
轨迹存储：state_prediction（整个 rollout 的状态预测）。

构造函数参数（默认值见源码）

构造函数签名及其语义（均以源码 docstring 为准）如下表所示：

参数	默认值	含义
`predictor_continuous`	必填	预测器的连续时间定义
`measurement_function`	必填	反馈测量函数 $g(\mathbf{x}, \hat{\mathbf{x}})$
`adapt_continuous`	必填	物理基函数与随机特征的连续自适应动力学
`norm_state`	必填	每步之后的可选状态归一化（如四元数投影）
`ob`	必填	系统状态的初始观测
`init_params`	必填	自适应参数初始化（物理在前，随机特征在后）
`n_features`	0	随机特征数量
`dt`	1e-3	预测器时间步长
`nsteps`	25	两次测量之间的预测器步数
`nframes`	40	帧数；加上初始化，预测器总积分步数为`1 + nframes*nsteps`
`potentials`	`{'phys': 'euclid', 'rf': 'euclid'}`	物理与随机特征自适应的势函数，可取值`'euclid'`、`'p'`（p 范数）、`'hyp'`（超熵），键为`'phys'`与`'rf'`
`potential_params`	`{'phys': 0.0, 'rf': 0.0}`	势函数参数（p 范数的 p 值、超熵的 beta 值）
`prediction_index`	40	超过该点后，将状态钳制到观测值以获得最准确的预测

关键方法与离散化

rollout_prediction（模块级函数）：给定一次测量后更新参数并滚动预测轨迹，核心是jax.lax.scan迭代执行scan_fn；@functools.partial(jax.jit, static_argnums=(0, 1))将norm_state与predict_state声明为静态参数。类方法rollout_prediction(t0)采用"每次始终 rollout 固定nsteps*nframes步、再截取所需数据点"的策略，以规避 JAX 因时间视野变化导致的重复重编译（源码注释明确说明了这一设计动机）；
euler：一阶欧拉离散化x + dt*dynamics(x, u, t)；
rk4：经典四阶 Runge-Kutta 离散化；
rk4_adaptive：自适应系统的离散化，等价于RK4 镜像下降——由于 RK4 要在中间时刻求值右端项，因此必须在中间时刻重新施加逆镜像映射（inverse_mirror）以恢复参数；
setup_dynamics_with_measurement：把已知动力学、测量函数、自适应动力学拼接为dynamics_with_measurement；
setup_mirrors：依据potentials配置选择势函数（p 范数对应grad_p/inv_grad_p，超熵对应grad_hyp/inv_grad_hyp，否则为恒等映射）；
apply_mirrors/invert_mirrors：对物理参数与随机特征参数分别施加/逆推镜像映射，并将状态部分原样透传；
update_state：测量到达时的核心入口——先对当前状态与镜像参数拼接，调用predict_state_with_measurement前进一步并更新参数；随后依据meas_count >= prediction_index判断：若已进入预测阶段则把估计钳制到观测，否则让预测继续推进，并滚动填充后续轨迹。

辅助数学函数

Notebook 还提供一组@jax.jit编译的四元数工具函数：skew（构造反对称矩阵）、to_axis_angle（四元数转轴角）、normalize（归一化）、multiply（四元数乘法）、vector_multiply（向量与四元数相乘）、to_rotation_matrix（由四元数经轴角与 Rodrigues 公式构造旋转矩阵）。状态归一化norm_state正是利用四元数归一化将预测状态重新投影到单位四元数流形上。

关键超参数配置详解

Notebook 中"Adaptive predictor parameters"代码单元给出了完整的运行配置，这是复现实验的关键（以下数值均来自源码，可直接复制运行）：

# 镜像下降势函数 psi_i 及其参数 potentials = {'phys': 'euclid', 'rf': 'euclid'} potential_params = {'phys': 1.0, 'rf': 1.0} # 帧：两次测量之间的一连串步。nsteps 为每帧内预测步数， # nframes 为总帧数（轨迹过程中接收到的观测次数） nsteps = 25 nframes = 40 # 学习率 eta_p = 10.0 * nsteps eta_rf = 20.0 * nsteps # 测量反馈增益，用于 g(x, x_hat) k = 25 * nsteps # 随机特征信息 n_features = 50 n_avg = 10 input_dim = 6 approx_dim = 3 # 初始化参数：18 个物理基参数 + n_features 个随机特征参数 nparams = 18 + n_features init_params = onp.zeros(nparams) # 添加到观测的噪声强度 noise_str = 0.0 # 刚体动力学状态维度：四元数(4) + 位置(3) + 线速度(3) + 角速度(3) sys_dim = 13

要点解析：

状态布局为(quaternion, position, linear velocity, angular velocity)，共 13 维，与sys_dim = 13对应；predictor_continuous的 docstring 明确给出了这一约定；
物理基函数inertial_basis对 3 维角速度构造 3×18 的过参数化矩阵（每行 18 个二次项组合），对应 18 个物理参数；
随机特征init_random_features构造权重w ~ N(0,1)、偏置b ~ U(0,1)的余弦特征：sum(cos(w@x + b), axis=2)/sqrt(n_avg)，其中n_avg=10表示对 10 组随机特征取平均以降低方差，input_dim=6对应（线速度, 角速度）拼接的 6 维输入，approx_dim=3对应 3 个输出分量；随机特征被加在线速度动力学上，用于估计马格努斯效应（Magnus effect）等与平动/角速度相关的未建模气动力；
自适应信号：phys_error_signal返回角速度估计差 $\hat{\omega} - \omega$；rf_error_signal返回线速度估计差 $\hat{v} - v$。

实验一：PyBullet 抛掷小球的在线预测

Notebook 的第一个完整实验以 PyBullet 仿真"小球被抛掷"的场景验证上述实现，流程如下。

仿真环境搭建

dt = 1e-3 p.connect(p.DIRECT) p.resetSimulation() p.setGravity(0, 0, -9.8) p.setTimeStep(dt) useFixedBase = True flags = p.URDF_INITIALIZE_SAT_FEATURES plane_pos = [0,0,-0.625] p.setAdditionalSearchPath(pybullet_data.getDataPath()) plane = p.loadURDF('plane.urdf', plane_pos, flags = flags, useFixedBase=useFixedBase) sphere = p.loadURDF('sphere_small.urdf', [0, 0, 10], useMaximalCoordinates=True, flags=p.URDF_USE_INERTIA_FROM_FILE)

其中p.connect(p.DIRECT)表示无 GUI 的离屏仿真模式，p.setTimeStep(dt)与预测器步长保持一致，plane.urdf/sphere_small.urdf来自pybullet_data数据包。

状态读写辅助函数

def set_initial_state(object_id, position, orn, linear_velocity, angular_velocity): p.resetBasePositionAndOrientation(object_id, position, orn) p.resetBaseVelocity(object_id, linear_velocity, angular_velocity) def get_state(object_id): position, orientation = p.getBasePositionAndOrientation(object_id) linear_velocity, angular_velocity = p.getBaseVelocity(object_id) # 使 pybullet 的四元数约定 (x, y, z, w) 对齐到 jaxrobotics 的 (w, x, y, z) orientation = jnp.array([orientation[3], orientation[0], orientation[1], orientation[2]]) position = jnp.array(position) linear_velocity = jnp.array(linear_velocity) angular_velocity = jnp.array(angular_velocity) return jnp.concatenate((orientation, position, linear_velocity, angular_velocity))

值得注意：get_state将 PyBullet 的(x,y,z,w)四元数顺序转换为(w,x,y,z)，以匹配预测器内部 jaxrobotics 风格的约定。

预测器构建与 rollout 主循环

object_id = sphere set_initial_state(object_id, [0, 0, 10.], [0, 0, 0, 1.], [1, 3, 2.], [3., 2., 1.]) x_init = get_state(object_id) predictor = AdaptivePredictor(predictor_continuous, measurement_function, adapt_continuous, norm_state, x_init, init_params, n_features, dt, nsteps, nframes, potentials, potential_params, prediction_index=nframes) predictor_states = [] params_over_time = onp.zeros((nframes, nparams)) params_over_time[0, :] = init_params state_observations = onp.zeros((1+nsteps*nframes, 13)) state_observations[0, :] = x_init state_observation = x_init for i in range(nsteps*nframes): p.stepSimulation() new_measurement = (i % nsteps) == 0 if new_measurement: ind = int(i / nsteps) predictor.update_state(state_observation + 2*noise_str*onp.random.uniform(size=13) - noise_str) predictor_states.append(predictor.state_prediction.copy()) params_over_time[ind, :] = predictor.params state_observation = get_state(object_id) state_observations[i+1, :] = state_observation

主循环的测量调度逻辑是：每nsteps（25）个仿真步做一次"测量"（i % nsteps == 0），即每帧一次；update_state在每次测量时既推进预测又更新参数。noise_str = 0.0时观测无噪声；若需要测试抗噪性，可增大该值（观测将被均匀噪声扰动，幅度为2*noise_str的随机偏移减去noise_str）。

实验二：自实现旋转动力学的"无仿真器"实验

Notebook 的第二部分" Dynamics prediction with rotational dynamics"专门针对具有非平凡旋转动力学的系统。PyBullet 的积分器与预测器不同，且难以直接操控这类系统，因此作者改为自行实现 ground-truth 旋转动力学，避免积分差异带来的干扰。其要点（源码 docstring 明示）：

ground-truth 旋转动力学写在世界坐标系，需利用**平行轴定理（parallel axis theorem）**跟踪随时间变化的惯性张量；
预测器则在世界坐标系实现平动动力学、在刚体坐标系实现旋转动力学。

# 椭圆体的惯性张量矩阵 J_def = jnp.array([[5.0, 1.0, 0], [1.0, 2, 0], [0, 0, 0.5]]) m_def = 1.0 # 物体质量 Cd = 0.5 # 阻力系数 Cl = 0.5 # 升力（马格努斯）系数 def rotational_dynamics_body(x, u, t, J=J_def): q, p, v, w = x[:4], x[4:7], x[7:10], x[10:] qdot = .5 * vector_multiply(w, q) pdot = v R = to_rotation_matrix(q) w_world = R @ w vdot = grav - Cd*jnp.linalg.norm(v)*v + Cl*jnp.cross(w_world, v) wdot = -jnp.linalg.inv(J) @ jnp.cross(w, J @ w) return jnp.concatenate((qdot, pdot, vdot, wdot)) disc_rot_dyn = jax.jit(rk4(rotational_dynamics_body, dt=dt))

该实验中引入了阻力项 $-C_d |v| v$ 与马格努斯升力项 $C_l, \omega_{\text{world}} \times v$，进一步检验随机特征对这类未建模扰动的补偿能力。rollout 循环与实验一结构相同：初始化q_init=[0,0,0,1]、p_init=[0,0,10]、v_init=[1.5,0,2.5]、w_init=[3,2,1]，用 RK4 离散的 ground truth 每步推进，同样每nsteps步触发一次predictor.update_state(state_observation)。从 Notebook 保存的输出看，该单元共执行 1000 次迭代并逐次打印进度（Finished iteration %d）。

结果可视化与验证

Notebook 末尾提供了两套可视化工具函数，用于直接检验预测质量与参数学习过程：

make_prediction_comparison_plot：绘制指定测量时刻（measurement_no）之后的"预测轨迹 vs 观测轨迹"对比图，以及两者的误差曲线。支持四个状态子集：orientation（四元数四分量）、position、linear velocity、angular velocity；图中以竖直虚线（axvline）标出该测量时刻的位置；
make_parameter_plot：绘制全部nparams个自适应参数随时间的演化曲线，可用于观察物理参数与随机特征参数的学习收敛过程。

Notebook 中给出的示例调用为：

measurement_no = 33 make_prediction_comparison_plot(nsteps*nframes+1, dt, state_observations, predictor_states, measurement_no, 'linear velocity') make_parameter_plot(nframes, nsteps, dt, params_over_time, measurement_no)

以及旋转动力学实验中的：

measurement_no = 34 make_prediction_comparison_plot(nsteps*nframes+1, dt, state_observations, predictor_states, measurement_no, 'linear velocity') make_parameter_plot(nframes, nsteps, dt, params_over_time, measurement_no)

从 Notebook 嵌入的 matplotlib 输出（Figure size 800x500）可以看到：预测误差曲线在第 33/34 次测量后明显收窄，参数曲线趋于稳定，直观印证了"自适应参数在滚动预测中被在线学习"的设计目标。

小结：如何将该包迁移到其他系统

从源码结构可以归纳出将该框架复用于新动力系统的四个步骤：

定义连续时间预测器：给出已知动力学部分f、物理基函数Y_phys与随机特征基函数Y_rf，按(quaternion, position, velocity, angular_velocity)约定组织状态；
定义测量函数与误差信号：按"参数出现在哪条动力学中，误差信号就取该分量的预测差"的原则设计measurement_function、phys_error_signal、rf_error_signal；
配置AdaptivePredictor：通过potentials/potential_params切换欧几里得、p 范数或超熵镜像下降；通过eta_p/eta_rf/k调节学习率与反馈增益；通过nsteps/nframes/prediction_index控制测量频率与预测-钳制切换时机；
对接数据源：既可使用 PyBullet（get_state/set_initial_state完成坐标约定转换），也可像实验二那样自实现 ground-truth 动力学并用rk4离散化生成观测序列。