VLA/workspace

Training Guide

This document explains the key configuration parameters and memory requirements for Wall-X training.

Quick Start Checklist

🚀 Step 1: Prepare Model

Choose one of our pretrained models:

• WALL-OSS-FLOW: https://huggingface.co/x-square-robot/wall-oss-flow
• WALL-OSS-FAST: https://huggingface.co/x-square-robot/wall-oss-fast Or from Qwen-2.5-VL
• Download https://huggingface.co/Qwen/Qwen2.5-VL-3B-Instruct, settings refer to config_qact_from_vlm.yml

⚙️ Step 2: Configure Environment

• Update run.sh: Set code_dir and config_path to your actual paths
• Set CUDA_VISIBLE_DEVICES for your available GPUs

📝 Step 3: Update Configuration Files

• Replace all /path/to/ placeholders in config_qact.yml with actual paths
• Configure robot settings: dof_config and agent_pos_config
• Set dataset: Choose appropriate repo_id
• Adjust batch_size_per_gpu based on your GPU memory

▶️ Step 4: Start Training

bash ./workspace/lerobot_example/run.sh

Enable FAST tokenizer

To fine-tune using the FAST tokenizer, please download the repository and update the action_tokenizer_path. Make sure to set use_fast_tokenizer to true and q01 and q99 to normalize the dataset, refer to wall-x/scripts/compute_norm_stats.py:

git clone https://huggingface.co/physical-intelligence/fast

Required Paths (Must Modify)

pretrained_wallx_path: "/path/to/wallx_model/"      # Path to pretrained wallx model
save_path: "/path/to/workspace/"                    # Path to save training outputs
use_fast_tokenizer: False                           # True: train FAST, False: train Flow
action_tokenizer_path: "/path/to/fast/"             # Must set if use_fast_tokenizer is True
norm_stats_path: "/path/to/stats/"                  # Must set for normalize dataset

Customize your robot configuration

Ensure that the sum of the configuration dimensions corresponds to the values specified in norm_stats.json, and that each key is unique. The maximum dimensionality is set to 20, consistent with our robot configuration.

  customized_dof_config:
    "action_eef": 6
    "action_gripper": 1

  customized_agent_pos_config:
    "state_eef_with_gripper": 7

Using Lerobot Dataset

• Each dataset employs distinct keys; please specify the corresponding key mappings as described in wall-x/wall_x/data/utils.py.

"lerobot/aloha_mobile_cabinet": {
      "camera": {
          "observation.images.cam_high": "face_view",
          "observation.images.cam_left_wrist": "left_wrist_view",
          "observation.images.cam_right_wrist": "right_wrist_view",
      },
      "state": "observation.state",
      "action": "action",
  }

Compute stats

    python wall-x/scripts/compute_norm_stats.py

Configuration Explain

• agent_pos_config corresponds to obs_action_keys and subsequently to state, while dof_config corresponds to predict_action_keys and subsequently to action. Note that the state and action may not necessarily share the same set of DoF.

Training Parameters (Commonly Modified)

Learning Rate Settings

• learning_rate: Initial learning rate (default: 0.00009)
• min_lr: Minimum learning rate for scheduler (default: 0.00005)
• num_warmup_steps: Number of warmup steps (default: 100)

Batch Size and Memory

• batch_size_per_gpu: Batch size per GPU - adjust based on GPU memory
• gradient_accumulation_steps: Gradient accumulation steps
• num_training_steps: Total training steps
• num_epoch: Number of training epochs

Training Optimization Settings

• FSDP2: Enable FSDP2 for distributed training (default: True) - Recommended for multi-GPU
• torch_compile: Enable PyTorch compilation optimization (default: False)

⚠️ Important Note on torch_compile:

• Benefits: Enabling torch_compile can significantly improve training efficiency
• Requirements: Requires that the data input shape is always consistent throughout training
• Caution: If you don't have sufficient understanding of torch compile, please DO NOT enable it as it may cause unexpected issues with dynamic input shapes

Robot Configuration (Modify for Your Robot)

DOF Configuration

Modify dof_config to match your robot's action space:

• Add/remove action keys based on your robot's capabilities
• Ensure DOF numbers match your robot's action dimensions

Agent Position Configuration

Keep agent_pos_config consistent with dof_config.

Action Keys

• obs_action_keys: Actions used as observation context
• predict_action_keys: Actions to predict/control

Data Configuration

Dataset

• repo_id: LeRobot dataset identifier
• train_test_split: Training/validation split ratio (default: 0.95)
• action_horizon: Number of future actions to predict (default: 32)

Image Settings

• resolution: Image resolution for different camera views
• download_videos: Whether to download video files (true/false)

Resume Training (Optional)

• resume.ckpt: Path to checkpoint for resuming training
• resume.load_ckpt_only: Only load model weights, not optimizer state

Merge checkpoint

• If FSDP SHARDED_STATE_DICT is used, please run command below to merge checkpoint into a single safetensors

    # refer to accelerate/commands/merge.py
    accelerate merge-weights /path/to/sharded_tensors /path/to/model.safetensors
    # copy the saved processor files
    cp /path/to/saved_processor_dir/* /path/to/model.safetensors

    # In earlier versions of PyTorch, errors may occur. You can use our provided script to address this issue; refer to wall-x/scripts/merge_sharded_weights.py for details.

Memory Usage

Below are the memory consumption benchmarks for different training configurations using the lerobot/aloha_mobile_cabinet dataset:

Dataset	Batch Size	FSDP2	Torch Compile	Num GPUs	Max Allocated Memory
lerobot/aloha_mobile_cabinet	1	❌	❌	1	40.11G
lerobot/aloha_mobile_cabinet	1	❌	❌	8	48.02G
lerobot/aloha_mobile_cabinet	1	✅	❌	2	43.70G
lerobot/aloha_mobile_cabinet	1	✅	❌	8	24.96G
lerobot/aloha_mobile_cabinet	1	✅	✅	8	24.21G

Hardware Recommendations:

• For single GPU training: Ensure at least 48GB VRAM (e.g., RTX 6000 Ada, A6000)
• For multi-GPU training: Enable FSDP2 for optimal memory distribution

Reproduce

Openloop plot wall-x/workspace/lerobot_example/evaluation/lerobot_openloop.png

To reproduce the results, use the config file wall-x/workspace/lerobot_example/config_qact_from_vlm.yml with a global batch size of 128, adjusted via gradient_accumulation_steps and numbers of gpu.