模型训练

本节以真机 LeRobotDataset v3 为例。LIBERO 的训练配置和命令在“LIBERO”中统一说明,避免重复。

创建训练配置

双臂、448 px、三相机场景可参考:

cp workspace/example/maniparena_example.yml /path/to/my_robot_config.yml

核心配置示例:

model_type: qwen2_5

model:
  backbone: qwen2_5
  config_path: /path/to/wall-oss-0.5/config.json
  processor_path: /path/to/Qwen2.5-VL-3B-Instruct
  pretrained_path: /path/to/Qwen2.5-VL-3B-Instruct
  attn_deterministic: true
  use_ema: false
  flow_loss_weight: 1.0
  ar_loss_weight: 0.01

data:
  dataset_type: lerobot
  lerobot_config:
    repo_id: /path/to/your_robot_dataset
    root: null
  key_mappings:
    camera:
      observation.images.faceImg: face_view
      observation.images.leftImg: left_wrist_view
      observation.images.rightImg: right_wrist_view
    state: observation.state
    action: action
  norm_stats_path: /path/to/your_robot_norm_stats.json
  train_test_split: 0.95
  num_workers: 4
  max_length: 1024
  resolution:
    face_view: 448
    left_wrist_view: 448
    right_wrist_view: 448

checkpoint:
  save_path: /path/to/robot_training_output
  resume_from: /path/to/wall-oss-0.5/model.safetensors

配置动作与状态维度

以下是现有双臂相对动作示例。开发者必须按自己的机器人和数据修改:

task:
  dof_config:
    follow_left_ee_cartesian_pos_relative: 3
    follow_left_ee_rotation_6D_relative: 6
    follow_left_gripper: 1
    follow_right_ee_cartesian_pos_relative: 3
    follow_right_ee_rotation_6D_relative: 6
    follow_right_gripper: 1
    action_padding: 6

  ar_dof_config:
    follow_left_ee_cartesian_pos_relative: 3
    follow_left_ee_rotation_6D_relative: 6
    follow_left_gripper: 1
    follow_right_ee_cartesian_pos_relative: 3
    follow_right_ee_rotation_6D_relative: 6
    follow_right_gripper: 1
    action_padding: 6

  agent_pos_config:
    follow_left_ee_cartesian_pos: 3
    follow_left_ee_rotation_6D: 6
    follow_left_gripper: 1
    follow_right_ee_cartesian_pos: 3
    follow_right_ee_rotation_6D: 6
    follow_right_gripper: 1
    action_padding: 6

  action_horizon: 32
  action_horizon_flow: 32
  use_state_string_representation: false

dof_config 定义模型预测的动作维度,agent_pos_config 定义 observation state 维度。各项维度总和必须与归一化统计文件一致。若真实动作少于预训练空间,可按模型配置使用 action_padding。

训练超参数示例
hyperparams:
  num_epoch: 100
  batch_size_per_gpu: 4
  gradient_accumulation_steps: 4
  seed: 10222
  optimizer:
    optimizer_type: adamw
    learning_rate: 5.0e-05
    max_grad_norm: 1.0
    enable_grad_clip: true
    betas: [0.9, 0.95]
    weight_decay: 1.0e-8
    eps: 1.0e-8
  scheduler:
    scheduler_type: cosine
    num_warmup_steps: 1000
    num_training_steps: 200000
    min_lr: 1.0e-6

distributed:
  use_fsdp: true
  use_mixed_precision: true
  bf16: true

logging:
  log_name: maniparena_ft
  log_project: lerobot_maniparena_ft
  log_entity: your_wandb_entity
  use_wandb: true
  log_interval: 10
  save_interval: 2000
  val_interval: 1000000
  epoch_save_interval: 1

debug:
  profile: false
  nvtx: false

以上参数来自现有示例,不是所有任务的统一推荐值。应根据数据规模、GPU 和任务效果调整,并保存实际使用的配置。

生成归一化统计
python scripts/compute_norm_stats.py \
  --train_config /path/to/my_robot_config.yml \
  --data_root /path/to/your_robot_dataset \
  --output_path /path/to/your_robot_norm_stats.json

脚本会依据配置中的字段和 DOF 转换计算统计。生成后确认 data.norm_stats_path 指向同一文件,并检查 state/action 的维度是否符合预期。

启动训练

多 GPU

# 多 GPU
CUDA_VISIBLE_DEVICES=0,1,2,3 \
torchrun --nproc_per_node=4 \
wall_x/trainer/fsdp_trainer/train_fsdp.py \
--config /path/to/my_robot_config.yml

# 单 GPU
CUDA_VISIBLE_DEVICES=0 \
torchrun --nproc_per_node=1  wall_x/trainer/fsdp_trainer/train_fsdp.py \
--config /path/to/my_robot_config.yml

单卡训练至少需要 48 GB 显存。多卡训练建议开启 FSDP。具体显存需求会受分辨率、相机数量、序列长度和 batch size 影响。

Checkpoint 与断点续训
  • 单个 。safetensors 用于加载预训练权重;

  • 完整恢复优化器、调度器和随机状态时,应将 checkpoint.resume_from 指向历史 Checkpoint 目录;

  • FSDP 产生分片权重时,推理前需要合并:

python scripts/merge_sharded_weights.py \
/path/to/sharded_checkpoint \
/path/to/merged_checkpoint
训练结果检查
  • 训练和验证 loss 是否异常;

  • 是否出现 NaN、显存不足或数据加载错误;

  • Checkpoint、配置和归一化文件是否一并保存;

  • 训练使用的相机、state 和 action 是否与真机运行一致;

  • 是否记录代码版本、数据版本、随机种子和硬件环境。

默认训练参数参考
参数默认值配置项
单卡批次大小4hyperparams.batch_size_per_gpu
梯度累积步数4hyperparams.gradient_accumulation_steps
学习率5e-5hyperparams.optimizer.learning_rate
训练轮数100hyperparams.num_epoch
FSDP 分布式训练truedistributed.use_fsdp
Checkpoint 保存间隔2000 步logging.save_interval

这些值来自官方示例配置,是复现实例的起点,不是所有任务的固定最优参数。

本页内容