参考运动生成

• 理解参考运动及生成，包含PolyReferenceMotion 类，速度到索引的映射，多项式采样，参考运动包含的维度和索引，可视化测试，与推理系统的交互和生成自己的参考运动等

概述

参考运动生成系统使用多项式拟合方法，将预定义的关节轨迹压缩为多项式系数，在训练和推理过程中实时解码。这种方式避免了存储大量原始运动数据，同时支持不同速度命令下的运动插值。

PolyReferenceMotion 类

定义在 playground/common/poly_reference_motion.py（JAX 版本，用于训练）和 playground/common/poly_reference_motion_numpy.py（NumPy 版本，用于推理）。

初始化

class PolyReferenceMotion:
    def __init__(self, polynomial_coefficients: str):
        data = pickle.load(open(polynomial_coefficients, "rb"))
        self.process(data)

输入文件为 pickle 格式，包含多个运动参数组合的多项式系数。

数据结构

参考运动数据按速度命令组织：

polynomial_coefficients.pkl
│
├── "dx_dy_dtheta"          ← 键: "0.0_0.0_0.0"
│   ├── period              ← 运动周期 (s)
│   ├── fps                 ← 帧率 (Hz)
│   ├── frame_offsets       ← 帧偏移
│   ├── startend_double_support_ratio  ← 双足支撑期比例
│   ├── coefficients        ← 多项式系数 (每维[左腿5, 头4, 天线2, 右腿5])
│   │   ├── dim_0: [c0, c1, c2, ...]  ← 关节位置拟合多项式系数
│   │   ├── dim_1: [...]
│   │   └── ...
│   └── ...
│
├── "0.05_0.0_0.0"          ← 不同速度组合
├── "0.0_0.05_0.0"
└── ...

数据处理

process() 方法将原始数据转换为 JAX 数组：

def process(self, data):
    # 提取速度范围
    self.dx_range = [min_dx, max_dx]
    self.dy_range = [min_dy, max_dy]
    self.dtheta_range = [min_dtheta, max_dtheta]
    
    # 所有速度组合的离散值
    self.dxs = sorted([-0.15, 0, 0.15, ...])   # 前进速度离散点
    self.dys = sorted([-0.2, 0, 0.2, ...])     # 横向速度离散点
    self.dthetas = sorted([-1.0, 0, 1.0, ...])  # 旋转速度离散点
    
    # 构建三维数组: [len(dxs), len(dys), len(dthetas), n_dims, poly_order]
    self.data_array = jp.array(...)
    
    # 运动参数
    self.period = 0.5         # 运动周期 (s)
    self.fps = 50              # 50 Hz
    self.nb_steps_in_period = int(period * fps)  # 每周期步数 (25)
    self.start_offset = ...    # 双足支撑期对应的帧数

速度到索引的映射

由于参考运动只在离散的速度点上采样，需要使用最近邻搜索：

def vel_to_index(self, dx, dy, dtheta):
    # 裁剪到有效范围
    dx = clip(dx, dx_range[0], dx_range[1])
    dy = clip(dy, dy_range[0], dy_range[1])
    dtheta = clip(dtheta, dtheta_range[0], dtheta_range[1])
    
    # 最近邻搜索
    ix = argmin(abs(dxs - dx))
    iy = argmin(abs(dys - dy))
    itheta = argmin(abs(dthetas - dtheta))
    
    return ix, iy, itheta

多项式采样

多项式计算

def sample_polynomial(self, t, coeffs):
    # 对每个维度计算多项式值
    # coeffs 形状: [n_dims, poly_order+1]
    # 返回: [n_dims]
    return vmap(lambda c: jp.polyval(c, t))(coeffs)

使用 jp.polyval（即 Horner 法则）高效计算多项式值。 通过 vmap 在维度间向量化。

参考运动获取

def get_reference_motion(self, dx, dy, dtheta, i):
    # 1. 速度 → 索引
    ix, iy, itheta = self.vel_to_index(dx, dy, dtheta)
    
    # 2. 相位归一化 (0~1)
    t = (i % nb_steps_in_period) / nb_steps_in_period
    t = clip(t, 0.0, 1.0)
    
    # 3. 采样多项式
    ret = self.sample_polynomial(t, self.data_array[ix][iy][itheta])
    return ret  # 形状: [40]

多项式拟合原理

原始运动数据:
  time  →  [joint_1, joint_2, ..., joint_16]
  t=0   →  [pos_1, pos_2, ..., pos_16]
  t=1   →  [...]
  ...

多项式拟合:
  对每个关节维度 j:
    使用最小二乘法拟合多项式 P_j(t)
    使得 P_j(t_i) ≈ pos_j(t_i) 对所有时间点 i

  结果:
    16 个多项式 P_0(t), P_1(t), ..., P_15(t)
    每个多项式 5-7 阶

推理:
  给定时间 t → P_j(t) → 关节 j 的位置

参考运动包含的维度和索引

参考帧（40 维）中各维度的含义：

索引  0:  left_hip_yaw 位置
  1:  left_hip_roll 位置
  2:  left_hip_pitch 位置
  3:  left_knee 位置
  4:  left_ankle 位置
  5:  neck_pitch 位置
  6:  head_pitch 位置
  7:  head_yaw 位置
  8:  head_roll 位置
  9:  left_antenna 位置
 10:  right_antenna 位置
 11:  right_hip_yaw 位置
 12:  right_hip_roll 位置
 13:  right_hip_pitch 位置
 14:  right_knee 位置
 15:  right_ankle 位置
 ─────────────────────────
 16:  left_hip_yaw 速度
 17:  left_hip_roll 速度
 ...
 31:  right_ankle 速度
 ─────────────────────────
 32:  left foot contact
 33:  right foot contact
 ─────────────────────────
 34:  base_linear_vel x
 35:  base_linear_vel y
 36:  base_linear_vel z
 ─────────────────────────
 37:  base_angular_vel x
 38:  base_angular_vel y
 39:  base_angular_vel z

可视化测试

PolyReferenceMotion 包含一个简单的主函数用于可视化测试：

if __name__ == "__main__":
    PRM = PolyReferenceMotion("...data/polynomial_coefficients.pkl")
    vals = []
    for i in range(PRM.nb_steps_in_period):
        vals.append(PRM.get_reference_motion(0.0, -0.05, -0.1, i)[dim])
    
    import matplotlib.pyplot as plt
    ts = np.arange(PRM.nb_steps_in_period)
    plt.plot(ts, vals)
    plt.show()

与推理系统的交互

在推理（mujoco_infer.py）中使用 NumPy 版本：

self.PRM = PolyReferenceMotion(reference_data)
self.imitation_i = 0

# 控制循环中
self.imitation_i = (self.imitation_i + 1) % self.PRM.nb_steps_in_period
self.imitation_phase = np.array([
    cos(self.imitation_i / N * 2 * pi),
    sin(self.imitation_i / N * 2 * pi),
])

相位信息 imitation_phase 作为观测的一部分传递给策略网络，帮助策略感知当前的步态相位。

生成自己的参考运动

参考运动数据需要使用独立的生成工具创建，详见 Open Duck 参考运动生成器 仓库。基本流程：

1. 录制或设计机器人每个关节的轨迹
2. 将轨迹分割为不同速度条件下的运动周期
3. 对每个周期的每维数据进行多项式拟合
4. 打包为 polynomial_coefficients.pkl 文件