参考运动查看器

• 了解和熟悉参考运动查看器

概述

参考运动查看器（ref_motion_viewer.py）是一个独立的 3D 可视化工具，用于在 MuJoCo 环境中查看参考运动数据生成的步态。它允许研究人员直观地评估参考运动的质量、调整参数，并在训练前验证运动数据。

功能

• 3D 可视化：在 MuJoCo 查看器中渲染参考运动
• 命令控制：通过键盘或游戏摇杆改变运动参数
• 多种场景：支持平坦地形和粗糙地形
• 速度选择：设置前进、横向和旋转速度

使用方式

# 基本使用
uv run playground/open_duck_mini_v2/ref_motion_viewer.py \
    --command 0.0 -0.05 -0.1 \
    --reference_data playground/open_duck_mini_v2/data/polynomial_coefficients.pkl \
    --scene flat_terrain

# 使用摇杆控制
uv run playground/open_duck_mini_v2/ref_motion_viewer.py \
    --joystick \
    --scene rough_terrain

命令行参数

注意：实际使用时应将 --reference_data 指向正确的路径，已更新的默认路径为 playground/open_duck_mini_v2/data/polynomial_coefficients.pkl。

核心逻辑

初始设置

# 加载模型
model = mujoco.MjModel.from_xml_string(
    epath.Path(model_path).read_text(), assets=base.get_assets()
)
data = mujoco.MjData(model)

# 加载参考运动
PRM = PolyReferenceMotion(args.reference_data)

# 设置初始状态 (从 home keyframe)
home_frame = model.keyframe("home")
default_qpos = home_frame.qpos.copy()
default_qpos[2] += 0.2  # 提升基座高度 0.2m
default_ctrl = home_frame.ctrl.copy()

data.qpos[:] = default_qpos
data.ctrl[:] = default_ctrl

主循环

加载模型和参考运动
    │
    ▼
初始化状态 (home 姿态 + 高度提升)
    │
    ▼
进入可视化循环
    │
    ├──→ 处理摇杆输入 (如果启用)
    │
    ├──→ 每 decimation 步:
    │      │
    │      ├── 计算步态相位 imitation_i
    │      │
    │      ├── 获取参考运动: PRM.get_reference_motion(cmd, i)
    │      │
    │      └── 将参考关节位置应用到仿真
    │
    ├──→ 执行 MuJoCo 物理步
    │
    └──→ 同步查看器

参考运动应用

if counter % decimation == 0 and not all(v == 0 for v in command):
    imitation_i = step % PRM.nb_steps_in_period
    
    # 获取当前命令对应的参考运动
    ref_motion = PRM.get_reference_motion(
        command[0], command[1], command[2], imitation_i
    )
    ref_motion = np.array(ref_motion)
    
    # 提取关节位置 (前 16 维)
    joints_pos = ref_motion[0:16]
    # 去掉头部和天线关节 (索引 5-10)
    ref_joint_pos = np.concatenate([joints_pos[:9], joints_pos[11:]])
    
    # 应到仿真状态
    new_qpos = default_qpos.copy()
    if new_qpos[7:7+14].shape[0] == ref_joint_pos.shape[0]:
        new_qpos[7:7+14] = ref_joint_pos
    
    step += 1

键盘控制

与推理系统的键盘控制相同：

摇杆控制

通过 pygame 库支持游戏摇杆：

if args.joystick:
    pygame.init()
    pygame.joystick.init()
    joystick1 = pygame.joystick.Joystick(0)
    joystick1.init()
    
    # 可选的第二个摇杆用于旋转控制
    if pygame.joystick.get_count() > 1:
        joystick2 = pygame.joystick.Joystick(1)
        joystick2.init()

摇杆映射：

• 左摇杆 Y 轴：前进/后退
• 左摇杆 X 轴：横向移动
• 右摇杆 X 轴（如果可用）：旋转

摇杆命令计算

def handle_joystick():
    joy_y = joystick1.get_axis(1)   # 前进/后退
    joy_x = joystick1.get_axis(0)   # 横向
    
    if joy_y < 0:
        lin_vel_x = (-joy_y) * COMMANDS_RANGE_X[1]    # 正 = 前进
    else:
        lin_vel_x = -joy_y * abs(COMMANDS_RANGE_X[0])  # 负 = 后退
    lin_vel_y = -joy_x * COMMANDS_RANGE_Y[1]
    
    if joystick2 is not None:
        joy_z = joystick2.get_axis(0)
        lin_vel_z = -joy_z * COMMANDS_RANGE_THETA[1]

可视化参数

场景选择

查看器支持所有已定义的场景 XML 文件。可用场景自动检测：

SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))
SCENE_PATH = f"{SCRIPT_PATH}/xmls"

available_scenes = [
    name[len("scene_"):-len(".xml")]
    for name in os.listdir(SCENE_PATH)
    if name.startswith("scene_") and name.endswith(".xml")
]

姿态提升

查看器中默认将机器人高度提升 0.2m，以便更清楚地观察步态循环：

default_qpos[2] += 0.2  # Z 方向提升

与其他模块的关系

ref_motion_viewer.py
    │
    ├── 使用 base.get_assets() 加载资产
    │
    ├── 使用 PolyReferenceMotion 获取参考运动
    │
    ├── 使用 constants.py 中的场景路径
    │
    └── 独立的可视化入口 (__main__)

该工具独立于训练和推理管线，用于数据验证和可视化研究。