Model Training

This section uses a real-robot dataset in LeRobotDataset v3 format as the example. LIBERO training configurations and commands are documented on the LIBERO page to avoid duplication.

Create training configuration

For dual-arm, 448 px, three-camera scenarios, refer to:

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

Core configuration example:

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

Configure action and state dimensions

The following is an example of the existing dual-arm relative action. Developers must modify it according to their own robot and data:

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 defines the action dimensions predicted by the model, while agent_pos_config defines the observation state dimensions. The sum of all dimensions must match the normalization statistics file. If the real action space has fewer dimensions than the pre-training action space, action_padding can be used according to the model configuration.

Training hyperparameter example
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

The above parameters are from existing examples and are not recommended values for all tasks. They should be adjusted based on data scale, GPUs, and task performance, and the configuration actually used should be saved.

Generate normalization statistics
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

The script computes statistics based on the fields and DOF transformations in the configuration. After generation, confirm that data.norm_stats_path points to the same file, and check whether the dimensions of state/action match expectations.

Start training

Multi-GPU

# Multi-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

# Single-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

Single-GPU training requires at least 48 GB of GPU memory. FSDP is recommended for multi-GPU training. Specific GPU memory requirements will be affected by resolution, number of cameras, sequence length, and batch size.

Checkpoints and resuming training
  • A single .safetensors file is used to load pre-trained weights;

  • When fully restoring the optimizer, scheduler, and random state, checkpoint.resume_from should point to the historical checkpoint directory;

  • When FSDP generates sharded weights, they need to be merged before inference:

python scripts/merge_sharded_weights.py \
/path/to/sharded_checkpoint \
/path/to/merged_checkpoint
Check training results
  • Whether training and validation loss are abnormal;

  • Whether NaN, out of memory (OOM), or data loading errors occur;

  • Whether checkpoints, configuration, and normalization files are saved together;

  • Whether the cameras, state, and action used for training are consistent with real-robot execution;

  • Whether the code version, data version, random seed, and hardware environment are recorded.

Default training parameter reference
ParameterDefault valueConfiguration item
Per-GPU batch size4hyperparams.batch_size_per_gpu
Gradient accumulation steps4hyperparams.gradient_accumulation_steps
Learning rate5e-5hyperparams.optimizer.learning_rate
Training epochs100hyperparams.num_epoch
FSDP distributed trainingtruedistributed.use_fsdp
Checkpoint save interval2000 Stepslogging.save_interval

These values come from the official example configuration and are a starting point for reproducing the example, not fixed optimal parameters for all tasks.

On this page