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.ymlCore 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.safetensorsConfigure 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: falsedof_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: falseThe 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.jsonThe 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.ymlSingle-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
.safetensorsfile 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_checkpointCheck 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
| Parameter | Default value | Configuration item |
|---|---|---|
| Per-GPU batch size | 4 | hyperparams.batch_size_per_gpu |
| Gradient accumulation steps | 4 | hyperparams.gradient_accumulation_steps |
| Learning rate | 5e-5 | hyperparams.optimizer.learning_rate |
| Training epochs | 100 | hyperparams.num_epoch |
| FSDP distributed training | true | distributed.use_fsdp |
| Checkpoint save interval | 2000 Steps | logging.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.