adaptive_prediction 自适应动力学预测实战:基于 JAX 与 PyBullet 的刚体运动在线预测指南 📅 发布时间:2026/9/20 1:18:44 👁 浏览次数: 人工智能深度学习NLP计算机视觉强化学习【免费下载链接】google-researchGoogle Research项目地址https://gitcode.com/gh_mirrors/go/google-research点击查看免费下载本篇技术指南以 google-research 仓库中的 adaptive_prediction 包为核心系统讲解自适应动力学预测Adaptive Dynamics Prediction这一在线学习方法预测器在滚动预测轨迹的同时利用测量反馈实时更新模型参数。文章从数学框架、核心类实现、超参数配置到两类刚体动力学实验PyBullet 抛掷小球、自实现旋转动力学逐层展开读者学完后可复现 Notebook 中的完整预测-适应闭环并具备将其迁移到其他动力学系统含镜面下降势函数扩展的动手能力。包概览README 与自包含 Notebookadaptive_prediction包的结构非常精简全部功能集中在三个文件中文件作用README.md包概览与使用说明rigid_body_adaptive_prediction.ipynb自包含的完整实现与演示交互式 Notebookrequirements.txt运行依赖清单README 明确指出该包包含通过自适应动力学预测对动力系统进行在线学习的代码代码自包含在rigid_body_adaptive_prediction.ipynb中并以刚体旋转动力学作为完整的演练示例worked-through example。Notebook 自身的Notebook overview小节进一步说明了内容组织它提供自适应动力学预测的交互式演示数据来自 PyBullet 对抛掷物体thrown objects的仿真输出先描述自适应预测器的通用结构再从数学上为刚体动力学构造一个自适应预测器最后给出具体实现。因此本文的正文脉络严格遵循 Notebook 的这一三段式结构。运行环境与快速启动根据 requirements.txt运行该 Notebook 需要以下依赖jupyter6.4.0 matplotlib3.4 numpy1.20.3 jax0.2 pybullet3.05 seaborn0.11.1启动方式即 README 给出的唯一使用指令jupyter notebook rigid_body_adaptive_prediction.ipynbNotebook 的第一批代码单元会先通过 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、超熵hyphypentropy$\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}) $$物理参数的更新律为$$\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}}\left(\hat{\boldsymbol\omega} - \boldsymbol\omega\right)$$由于参数估计 $\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_features0随机特征数量dt1e-3预测器时间步长nsteps25两次测量之间的预测器步数nframes40帧数加上初始化预测器总积分步数为1 nframes*nstepspotentials{phys: euclid, rf: euclid}物理与随机特征自适应的势函数可取值euclid、pp 范数、hyp超熵键为phys与rfpotential_params{phys: 0.0, rf: 0.0}势函数参数p 范数的 p 值、超熵的 beta 值prediction_index40超过该点后将状态钳制到观测值以获得最准确的预测关键方法与离散化rollout_prediction模块级函数给定一次测量后更新参数并滚动预测轨迹核心是jax.lax.scan迭代执行scan_fnfunctools.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_measurementsetup_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(wx b), axis2)/sqrt(n_avg)其中n_avg10表示对 10 组随机特征取平均以降低方差input_dim6对应线速度, 角速度拼接的 6 维输入approx_dim3对应 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, useFixedBaseuseFixedBase) sphere p.loadURDF(sphere_small.urdf, [0, 0, 10], useMaximalCoordinatesTrue, flagsp.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_indexnframes) predictor_states [] params_over_time onp.zeros((nframes, nparams)) params_over_time[0, :] init_params state_observations onp.zeros((1nsteps*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(size13) - noise_str) predictor_states.append(predictor.state_prediction.copy()) params_over_time[ind, :] predictor.params state_observation get_state(object_id) state_observations[i1, :] state_observation主循环的测量调度逻辑是每nsteps25个仿真步做一次测量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, JJ_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, dtdt))该实验中引入了阻力项 $-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*nframes1, 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*nframes1, 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控制测量频率与预测-钳制切换时机对接数据源既可使用 PyBulletget_state/set_initial_state完成坐标约定转换也可像实验二那样自实现 ground-truth 动力学并用rk4离散化生成观测序列。上述全部代码、公式与超参数均可在 rigid_body_adaptive_prediction.ipynb 中按单元逐步复现依赖安装与启动方式见前文运行环境与快速启动一节。赞分享人工智能深度学习NLP计算机视觉强化学习【免费下载链接】google-researchGoogle Research项目地址https://gitcode.com/gh_mirrors/go/google-research点击查看免费下载相关推荐Fyrox物理引擎完全指南碰撞检测与刚体动力学的终极实战应用Fyrox物理引擎完全指南碰撞检测与刚体动力学的终极实战应用 Fyrox是一款用Rust编写的强大3D和2D游戏引擎其物理系统为游戏开发提供了精准的碰撞检测游戏开发图形学3D渲染Active Selective Prediction 实战指南基于 ASPEST 的主动学习与选择性预测框架Active Selective Prediction 实战指南基于 ASPEST 的主动学习与选择性预测框架 导读 本文是 google research人工智能深度学习NLP计算机视觉强化学习AI-Scientist 地震预测模板基于卷积 LSTM 的日本地震预测基线与自动化科研流水线AI Scientist 地震预测模板基于卷积 LSTM 的日本地震预测基线与自动化科研流水线 本篇指南围绕 templates/earthquake pre人工智能大模型AI Agent自主智能体深度研究AI 应用创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考