Ψ₀ Hands-On (4): Setup & Training Pipeline

Sau khi hiểu data recipe ở bài trước, giờ là lúc chúng ta bắt tay vào việc thực sự: cài đặt môi trường, tải model checkpoints, và chạy 3 giai đoạn training của Ψ₀. Bài viết này sẽ đi từ A đến Z — từ git clone đến khi bạn có một model fine-tuned sẵn sàng deploy.

Yêu cầu phần cứng

Trước khi bắt đầu, hãy xác định bạn thuộc nhóm nào:

Thực tế: Hầu hết mọi người sẽ chỉ chạy Stage 3 (Fine-tune) với checkpoints đã pre-train sẵn từ nhóm tác giả. Đây cũng là phần thực dụng nhất — bạn chỉ cần 2-8 GPU để fine-tune cho task mới của mình.

Nếu bạn không có GPU mạnh, các cloud provider phổ biến:

• Lambda Labs: ~$1.1/giờ cho A100 80GB
• RunPod: ~$1.6/giờ cho A100 80GB
• Vast.ai: ~$0.8/giờ cho A100 40GB (spot pricing)

Cài đặt môi trường

Bước 1: Clone repository

git clone https://github.com/physical-superintelligence-lab/Psi0.git
cd Psi0

Bước 2: Cài đặt uv package manager

Ψ₀ sử dụng uv thay vì pip/conda truyền thống. uv nhanh hơn pip 10-100x và quản lý dependencies tốt hơn:

# Cài uv
curl -LsSf https://astral.sh/uv/install.sh | sh

# Verify
uv --version

Bước 3: Tạo virtual environment với Python 3.10

# Tạo venv với Python 3.10 (bắt buộc)
uv venv --python 3.10
source .venv/bin/activate

# Verify Python version
python --version  # Phải là 3.10.x

Tại sao Python 3.10? Flash Attention 2.7.4 và một số CUDA extensions chỉ support stable trên Python 3.10. Python 3.11+ có thể gây lỗi compilation.

Bước 4: Cài đặt dependencies

# Cài đặt core dependencies
uv pip install -e .

# Cài Flash Attention (bắt buộc, cần CUDA toolkit)
uv pip install flash-attn==2.7.4 --no-build-isolation

# Verify Flash Attention
python -c "import flash_attn; print(flash_attn.__version__)"
# Output: 2.7.4

Lưu ý quan trọng: Flash Attention cần CUDA toolkit đã cài sẵn trên hệ thống (không chỉ CUDA runtime của PyTorch). Nếu bạn gặp lỗi compilation:

# Kiểm tra CUDA toolkit
nvcc --version  # Cần 11.8 hoặc 12.x

# Nếu chưa có, cài CUDA toolkit
# Ubuntu:
sudo apt install nvidia-cuda-toolkit

Bước 5: Cấu hình environment variables

Tạo file .env ở thư mục gốc project:

# .env file
export HF_TOKEN=hf_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
export WANDB_API_KEY=xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
export PSI_HOME=/path/to/Psi0

HF_TOKEN: Lấy từ huggingface.co/settings/tokens. Cần quyền read để download model và data.

WANDB_API_KEY: Lấy từ wandb.ai/settings. Dùng để monitor training — rất quan trọng vì training có thể chạy nhiều ngày.

PSI_HOME: Đường dẫn tuyệt đối đến thư mục Ψ₀. Scripts sử dụng biến này để tìm configs, checkpoints, và data.

# Load env vars
source .env

# Verify
echo $HF_TOKEN  # Phải hiện token
echo $PSI_HOME   # Phải hiện đường dẫn

Bước 6: Đăng nhập HuggingFace và W&B

# Login HuggingFace
huggingface-cli login --token $HF_TOKEN

# Login Weights & Biases
wandb login $WANDB_API_KEY

Tải Model Checkpoints và Data

Download pre-trained checkpoints

# Download System-2 (Qwen3-VL-2B, đã fine-tune)
python scripts/data/download_checkpoints.py \
  --repo_id physical-superintelligence/psi0-system2 \
  --local_dir checkpoints/system2

# Download System-1 (MM-DiT action expert)
python scripts/data/download_checkpoints.py \
  --repo_id physical-superintelligence/psi0-system1 \
  --local_dir checkpoints/system1

# Download System-0 (RL locomotion controller)
python scripts/data/download_checkpoints.py \
  --repo_id physical-superintelligence/psi0-system0 \
  --local_dir checkpoints/system0

Download training data

# Simulation data (nhỏ hơn, dùng để test pipeline)
python scripts/data/download_datasets.py \
  --repo_id physical-superintelligence/psi0-sim-data \
  --local_dir data/sim

# Real-world data (lớn hơn, cần cho fine-tune thực tế)
python scripts/data/download_datasets.py \
  --repo_id physical-superintelligence/psi0-real-data \
  --local_dir data/real

Lưu ý: Data download có thể mất vài giờ tùy tốc độ mạng. EgoDex (829 giờ video) rất lớn — chỉ download nếu bạn thực sự cần pre-train từ đầu.

Stage 1: Pre-training (64x A100)

Mục đích

Stage 1 dạy System-2 (Qwen3-VL-2B) hiểu manipulation primitives từ egocentric video (EgoDex) và robot data (HE). Model học autoregressive — dự đoán FAST token tiếp theo dựa trên hình ảnh + tokens trước đó.

Hyperparameters

Chạy pre-training

# Pre-train (CẦN 64x A100)
torchrun --nproc_per_node=8 --nnodes=8 \
  scripts/train/psi0/pretrain-psi0.sh

# Hoặc dùng SLURM
sbatch scripts/train/psi0/pretrain-psi0.slurm

Thực tế: Bạn gần như chắc chắn không cần chạy Stage 1. Nhóm tác giả đã cung cấp checkpoint pre-trained. Hãy dùng checkpoint có sẵn và nhảy thẳng đến Stage 2 hoặc Stage 3.

Nếu không có 64 GPU?

Dùng checkpoint pre-trained:

# Download pre-trained Stage 1 checkpoint
python scripts/data/download_checkpoints.py \
  --repo_id physical-superintelligence/psi0-pretrained \
  --local_dir checkpoints/pretrained

Checkpoint này chứa Qwen3-VL-2B đã được train trên EgoDex + HE, sẵn sàng cho Stage 2 hoặc Stage 3.

Stage 2: Post-training — Flow Matching (32x A100)

Mục đích

Stage 2 training System-1 — MM-DiT (Multi-Modal Diffusion Transformer) action expert với 500M parameters. Đây là model chịu trách nhiệm sinh ra joint-space actions chính xác dựa trên features từ System-2.

Điểm then chốt: VLM (System-2) bị frozen ở stage này. Chỉ có MM-DiT được train. VLM đóng vai trò "mắt và não" — quan sát scene và đưa ra representation. MM-DiT học cách chuyển representation đó thành hành động cụ thể.

Flow Matching là gì?

Flow Matching là phương pháp generative modeling tương tự Diffusion nhưng hiệu quả hơn. Thay vì thêm noise rồi khử noise (như DDPM), Flow Matching học một vector field chuyển từ noise distribution sang action distribution theo đường thẳng.

Pseudocode của Flow Matching training loop:

# Pseudocode: Flow Matching Training Loop
for batch in dataloader:
    images = batch["observation.images"]      # [B, C, H, W]
    states = batch["observation.state"]        # [B, 28]
    target_actions = batch["action"]           # [B, T, 36] (T timesteps)
    
    # 1. VLM encode (frozen, no gradient)
    with torch.no_grad():
        vlm_features = system2.encode(images)  # [B, D]
    
    # 2. Sample random timestep t ~ Uniform(0, 1)
    t = torch.rand(B, device=device)           # [B]
    
    # 3. Sample noise
    noise = torch.randn_like(target_actions)   # [B, T, 36]
    
    # 4. Interpolate: x_t = (1-t) * noise + t * target
    x_t = (1 - t.unsqueeze(-1).unsqueeze(-1)) * noise + \
          t.unsqueeze(-1).unsqueeze(-1) * target_actions
    
    # 5. Predict velocity field v(x_t, t, condition)
    v_pred = mmdit(x_t, t, vlm_features, states)  # [B, T, 36]
    
    # 6. Target velocity = target - noise (straight line)
    v_target = target_actions - noise
    
    # 7. MSE loss
    loss = F.mse_loss(v_pred, v_target)
    
    loss.backward()
    optimizer.step()

Hyperparameters Stage 2

Chạy post-training

# Post-train MM-DiT (CẦN 32x A100)
torchrun --nproc_per_node=8 --nnodes=4 \
  scripts/train/psi0/posttrain-psi0.sh

Tương tự Stage 1, bạn có thể dùng checkpoint Stage 2 có sẵn thay vì train từ đầu.

Stage 3: Fine-tuning — Phần thực tế nhất!

Đây là phần mà hầu hết mọi người sẽ thực sự chạy. Stage 3 fine-tune toàn bộ pipeline (System-2 + System-1) trên 80 demonstrations của task cụ thể.

Hyperparameters Stage 3

Fine-tune trong simulation

# Fine-tune trên simulation data
bash scripts/train/psi0/finetune-simple-psi0.sh \
  --data_dir data/sim/pick_place \
  --checkpoint_dir checkpoints/pretrained \
  --output_dir outputs/sim_pick_place \
  --num_gpus 4 \
  --batch_size 128 \
  --lr 1e-4 \
  --max_steps 40000 \
  --warmup_steps 500

Fine-tune trên real-world data

# Fine-tune trên real robot data
bash scripts/train/psi0/finetune-real-psi0.sh \
  --data_dir data/real/fold_clothes \
  --checkpoint_dir checkpoints/pretrained \
  --output_dir outputs/real_fold_clothes \
  --num_gpus 2 \
  --batch_size 128 \
  --lr 1e-4 \
  --max_steps 40000

Điều chỉnh batch size theo số GPU

Batch size 128 là global batch size. Nếu bạn có ít GPU hơn, dùng gradient accumulation:

# 8 GPU: batch_per_gpu = 128 / 8 = 16
--num_gpus 8 --batch_size_per_gpu 16 --gradient_accumulation 1

# 4 GPU: batch_per_gpu = 128 / 4 = 32 (nếu đủ VRAM)
--num_gpus 4 --batch_size_per_gpu 32 --gradient_accumulation 1

# 4 GPU, VRAM hạn chế: dùng gradient accumulation
--num_gpus 4 --batch_size_per_gpu 16 --gradient_accumulation 2

# 2 GPU: gradient accumulation nhiều hơn
--num_gpus 2 --batch_size_per_gpu 16 --gradient_accumulation 4

# 1 GPU (chậm nhưng chạy được):
--num_gpus 1 --batch_size_per_gpu 16 --gradient_accumulation 8

Quy tắc: num_gpus × batch_size_per_gpu × gradient_accumulation = 128 (global batch size). Gradient accumulation làm training chậm hơn nhưng kết quả tương đương.

Monitor Training với Weights & Biases

W&B là công cụ không thể thiếu khi training chạy nhiều giờ hoặc nhiều ngày. Ψ₀ tích hợp sẵn W&B logging.

Loss curves cần theo dõi

Khi training, bạn sẽ thấy các metrics sau trên W&B dashboard:

1. Total Loss (train/loss)

• Giảm nhanh trong 5000 steps đầu
• Giảm chậm và ổn định sau 20,000 steps
• Nếu loss tăng đột ngột → learning rate quá cao hoặc data có vấn đề

2. Action Loss (train/action_loss)

• MSE loss giữa predicted và target actions
• Nên giảm xuống dưới 0.01 sau 30,000 steps
• Nếu plateau sớm (>0.05 sau 10,000 steps) → data quality có vấn đề

3. Learning Rate (train/lr)

• Cosine curve: tăng linear trong warmup, sau đó giảm dần theo hình cosine
• Verify LR schedule đúng — sai LR schedule là nguyên nhân phổ biến của training thất bại

Dấu hiệu training tốt

Step 1000:  loss=0.85, action_loss=0.12  ← Giảm nhanh, tốt
Step 5000:  loss=0.32, action_loss=0.04  ← Tiếp tục giảm
Step 10000: loss=0.18, action_loss=0.02  ← Bắt đầu converge
Step 20000: loss=0.11, action_loss=0.008 ← Gần converge
Step 40000: loss=0.08, action_loss=0.005 ← Converged

Dấu hiệu training có vấn đề

Step 1000:  loss=0.85  ← OK
Step 2000:  loss=NaN   ← LR quá cao! Giảm LR 10x

Step 1000:  loss=0.85  ← OK  
Step 5000:  loss=0.82  ← Giảm rất chậm
Step 10000: loss=0.80  ← Plateau quá sớm → check data

Quản lý Checkpoints

Save checkpoints

Training tự động save checkpoint mỗi N steps (cấu hình trong config). Mỗi checkpoint chiếm ~4-6 GB:

# Cấu trúc thư mục output
outputs/sim_pick_place/
├── checkpoint-5000/
│   ├── model.safetensors
│   ├── optimizer.pt
│   └── training_state.json
├── checkpoint-10000/
├── checkpoint-20000/
├── checkpoint-30000/
├── checkpoint-40000/   ← Checkpoint cuối
└── logs/
    └── events.out.tfevents.*

Resume training

Nếu training bị gián đoạn (OOM, server restart, spot instance bị thu hồi):

# Resume từ checkpoint gần nhất
bash scripts/train/psi0/finetune-simple-psi0.sh \
  --data_dir data/sim/pick_place \
  --checkpoint_dir outputs/sim_pick_place/checkpoint-20000 \
  --output_dir outputs/sim_pick_place \
  --resume_from_checkpoint true \
  --max_steps 40000

Chọn best checkpoint

Không phải checkpoint cuối cùng luôn là tốt nhất. Đôi khi model overfit ở cuối training. Cách chọn:

1. Nhìn validation loss trên W&B — chọn checkpoint có val loss thấp nhất
2. Evaluation trên held-out episodes — chạy inference trên 5-10 episodes không dùng để train
3. Nếu không có validation set — checkpoint ở 80% training (step 32,000) thường ổn

# Evaluate checkpoint
python scripts/eval/evaluate_checkpoint.py \
  --checkpoint_dir outputs/sim_pick_place/checkpoint-30000 \
  --eval_data data/sim/pick_place_eval \
  --num_episodes 10

Xử lý lỗi thường gặp

1. Out of Memory (OOM)

CUDA out of memory. Tried to allocate 2.00 GiB

Giải pháp (theo thứ tự ưu tiên):

1. Giảm batch_size_per_gpu (16 → 8 → 4)
2. Tăng gradient_accumulation tương ứng
3. Bật gradient_checkpointing (tiết kiệm VRAM, chậm hơn 20%)
4. Giảm max_seq_length nếu config cho phép

# Ví dụ: OOM với batch 16 trên A100 40GB
--batch_size_per_gpu 8 --gradient_accumulation 2 --gradient_checkpointing true

2. NaN Loss

Step 3456: loss=NaN

Nguyên nhân phổ biến: Learning rate quá cao, data có giá trị outlier, hoặc mixed precision issues.

Giải pháp:

1. Giảm learning rate 10x (1e-4 → 1e-5)
2. Kiểm tra data: python scripts/data/validate_dataset.py --data_dir data/...
3. Tắt mixed precision (bf16 → fp32) — chậm hơn 2x nhưng stable
4. Tăng warmup steps (500 → 2000)

3. Training quá chậm

# Monitor GPU utilization
watch -n 1 nvidia-smi

# Nếu GPU util thấp → tăng data loading workers
--num_workers 8  # Thay vì default 4

4. Checkpoint quá lớn

Mỗi checkpoint 4-6 GB, save mỗi 5000 steps = 40-48 GB cho full training. Giải pháp:

# Chỉ giữ N checkpoints gần nhất
--save_total_limit 3

# Hoặc save ít thường xuyên hơn
--save_steps 10000  # Thay vì 5000

Real-Time Chunking cho Deployment

Sau khi training xong, model cần chạy real-time trên robot. Ψ₀ sử dụng Real-Time Chunking để đạt latency 160ms:

# Pseudocode: Real-Time Chunking Inference
action_buffer = []
chunk_size = 16  # Predict 16 timesteps cùng lúc

while robot_running:
    if len(action_buffer) == 0:
        # Buffer hết → predict chunk mới
        images = camera.capture()
        state = robot.get_state()
        
        # VLM encode + MM-DiT generate (~160ms)
        action_chunk = model.predict(images, state)  # [16, 36]
        action_buffer = list(action_chunk)
    
    # Pop action đầu tiên từ buffer
    action = action_buffer.pop(0)
    robot.execute(action)  # 50Hz control loop
    
    # Overlap: bắt đầu predict chunk tiếp theo
    # trong khi vẫn execute actions từ buffer hiện tại

Key insight: Model predict 16 actions cùng lúc (action chunk). Robot execute từng action ở 50Hz (20ms/action). Trong khi execute 16 actions (320ms), model có đủ thời gian predict chunk tiếp theo (160ms). Không có delay giữa các chunks.

Cách tiếp cận này khác với Diffusion Policy truyền thống cần 100+ denoising steps. Flow Matching của Ψ₀ chỉ cần 4-8 steps cho mỗi chunk, nhanh hơn đáng kể.

Tổng kết: Từ Code đến Robot

Hãy tóm tắt toàn bộ workflow:

1. Clone repo + cài uv + Python 3.10 + Flash Attention
2. Cấu hình .env (HF_TOKEN, WANDB_API_KEY, PSI_HOME)
3. Download pre-trained checkpoints (skip Stage 1 & 2)
4. Thu thập 80 demos cho task của bạn
5. Convert sang LeRobot format (raw_to_lerobot.py)
6. Fine-tune Stage 3 (2-8 GPU, ~6-8 giờ)
7. Monitor trên W&B, chọn best checkpoint
8. Deploy với Real-Time Chunking (160ms latency)

Với workflow này, bạn có thể đi từ "chưa có gì" đến "robot chạy task mới" trong vài ngày — phần lớn thời gian là thu thập 80 demos và chờ training. Đây là bước tiến lớn so với các phương pháp trước đó cần hàng nghìn demos và hàng tuần training.

Ψ₀ không chỉ là một paper hay — nó là một framework mở mà bất kỳ ai có GPU và robot đều có thể sử dụng. Với LeRobot format chuẩn hóa dữ liệu, cộng đồng có thể chia sẻ datasets và checkpoints, tăng tốc nghiên cứu và ứng dụng humanoid robot trên toàn thế giới.

Bài viết liên quan

• Ψ₀ Hands-On (3): Data Recipe & Pipeline — Hiểu data recipe 3 tầng và tại sao 860h data thắng 10,000h
• VLA + LeRobot (1): Framework tổng quan — Nền tảng LeRobot format được Ψ₀ sử dụng cho toàn bộ data pipeline
• Diffusion Policy: Sinh hành động cho robot — So sánh Flow Matching với Diffusion — hai paradigm generative modeling cho robot