LIFT humanoid ICLR 2026: pretrain SAC + world model finetune

Nếu bạn đã theo dõi humanoid RL vài năm gần đây thì sẽ thấy một pattern lặp đi lặp lại: pretrain trong simulator cực lớn (thường là Isaac Gym/IsaacLab với PPO), sau đó zero-shot sim-to-real. Pattern này hoạt động nhưng rất tốn tài nguyên — PPO on-policy, cần hàng trăm triệu timestep, và khi sang môi trường mới (Brax, bullet, robot thật) policy thường "crash" vài giờ đầu trước khi tự thích nghi. LIFT — paper ICLR 2026 của nhóm BIGAI — giải quyết đúng chỗ đau này.

LIFT (Large-scale pretraining + efficient Finetuning) đề xuất kiến trúc ba giai đoạn: (1) pretrain SAC trong MuJoCo Playground với large-batch + high UTD ratio, (2) học một physics-informed world model từ replay buffer, (3) finetune ở môi trường đích (Brax sim hoặc robot thật) bằng cách chỉ chạy action tất định trong env thật còn exploration ngẫu nhiên thì rollout trong world model. Kết quả: converge trong 1 giờ trên NVIDIA 4090, transfer sang Brax và Booster T1 chỉ với 80–590 giây dữ liệu thật.

Vì sao pretrain + finetune lại khó với humanoid?

Với manipulator (cánh tay), pretrain-then-finetune khá dễ: động học chủ yếu là kinematic, quán tính nhỏ, contact ngắn. Nhưng humanoid có ba đặc điểm khiến naive transfer gần như luôn hỏng:

1. Under-actuation ở bàn chân — robot không có khớp ở điểm tiếp xúc với sàn, balance phụ thuộc hoàn toàn vào reaction force. Một sai số nhỏ về friction coefficient giữa MuJoCo và Brax cũng đủ biến policy "đi thẳng" thành "ngã lật ngửa".
2. Chuỗi khớp dài, contact rời rạc — lực từ mặt đất truyền qua mắt cá → đầu gối → hông → lưng → vai. Mỗi mô hình physics xử lý contact khác nhau (soft vs hard constraints, cảng lò xo, ma sát Coulomb xấp xỉ). Sim-to-sim gap đôi khi còn lớn hơn sim-to-real.
3. Action rate matters — humanoid cần action mượt. PPO với entropy bonus cao thường sinh ra policy giật, chuyển sang môi trường khác là dao động.

LIFT chọn SAC off-policy (có replay buffer → tái sử dụng được cho world model), cộng với action-rate L2 penalty trong reward để ép smoothness ngay từ pretrain. Đây là một quyết định kiến trúc quan trọng — nếu bạn đang ở pretrain step thì sẽ không thấy lợi ích, nhưng lên finetune bạn sẽ cảm ơn việc có replay buffer và policy mượt sẵn.

Kiến trúc ba tầng của LIFT

Stage 1 — Pretrain SAC trong MuJoCo Playground

Thay vì PPO hàng chục triệu env-step, LIFT dùng SAC JAX implementation với ba tinh chỉnh:

• Large-batch update: batch size rất lớn (vài nghìn) để tận dụng hết GPU.
• High UTD (Update-To-Data) ratio: với mỗi env-step, network được update nhiều lần (thường 4–16 gradient step). UTD cao giúp chưng cất replay buffer nhanh, không cần quá nhiều exploration.
• 1024 parallel envs trên MuJoCo Playground: Playground dùng MJX (MuJoCo trên JAX) nên parallelism miễn phí, không cần Isaac.

Với setup này, task T1LowDimJoystickFlatTerrain (Booster T1 đi theo joystick trên mặt phẳng) converge trong ~40M timestep, ~1 giờ trên 4090. So sánh: PPO trên Isaac Lab cho cùng bài toán thường cần 4–8 giờ và nhiều VRAM hơn.

Nếu bạn mới bắt đầu với RL và muốn hiểu SAC so với PPO, xem AI Series 1: RL Basics. Nếu muốn làm quen MuJoCo Playground trước, đọc Sim Series 2: MuJoCo toàn tập.

Stage 2 — Physics-informed World Model

Đây là phần nguyên bản của LIFT. Thay vì world model "đen sì" kiểu Dreamer (GRU + RSSM), LIFT chia kiến trúc thành hai nhánh:

• Lagrangian dynamics branch — viết trực tiếp phương trình cơ học M(q)q̈ + C(q,q̇)q̇ + g(q) = τ + Jᵀfc. Network học các residual nhỏ cho M, C, g để bù mismatch giữa MJX sim và env đích, không học lại toàn bộ động lực học.
• Contact residual predictor — một MLP nhỏ nhận (q, q̇, foot height, torque) và dự đoán contact force residual. Đây là phần "khó cho lý thuyết, dễ cho data".

Ưu điểm: world model generalize tốt hơn nhiều so với RSSM thuần, vì phần cấu trúc được gắn cứng vào physics. Dữ liệu huấn luyện world model đến từ chính replay buffer của stage 1 — một trong các lý do tác giả chọn SAC off-policy thay vì PPO.

Stage 3 — Finetune: deterministic env + stochastic world model

Đây là idea quan trọng nhất và cũng dễ bỏ sót. Trong finetune, LIFT thực hiện:

• Trong env thật (Brax hoặc Booster T1): chạy chính sách deterministic — tức lấy μ(s) của Gaussian policy, không sample. Điều này an toàn (không có hành vi điên rồ bất ngờ) và dữ liệu thu được sạch.
• Trong world model: rollout stochastic — sample action từ N(μ, σ), reward được tính bởi world model, gradient backprop qua cả WM và policy.

Tách như vậy làm hai việc: (a) giữ exploration cần thiết cho SAC mà không tổn hại robot thật, (b) khi WM dự đoán lệch với env thật, loss sẽ tự "kéo" cả WM lẫn policy về đúng.

Tham khảo chéo: nếu bạn muốn biết sim2real thường vấp ở đâu, đọc Sim Series 5: Sim2Real Pipeline — nhiều vấn đề LIFT né được bằng finetune chính là những gap đó.

Cài đặt LIFT-humanoid từ đầu

Yêu cầu hệ thống

Các bước setup

conda create -n lift python=3.10 -c conda-forge -y
conda activate lift

git clone https://github.com/bigai-ai/LIFT-humanoid.git
cd LIFT-humanoid

# Cài MuJoCo Playground bản của repo (có custom env)
cd mujoco_playground && pip install -e . && cd ..

# Cài Brax env wrapper
cd brax_env && pip install -e . && cd ..

# Cài requirements còn lại (jax, flax, dill, wandb,...)
pip install -r requirements.txt

Sau khi cài xong, kiểm tra JAX có thấy GPU:

import jax
print(jax.devices())   # Phải thấy CudaDevice(0)

Nếu báo CPU-only, check lại jax[cuda12] và LD_LIBRARY_PATH trỏ vào CUDA 12.

Pretrain SAC (Stage 1)

Command chuẩn cho Booster T1, low-dim state, flat terrain:

CUDA_VISIBLE_DEVICES=0 python train_in_mujoco_playground.py \
  --env_name=T1LowDimSimFinetuneJoystickFlatTerrain \
  --domain_randomization \
  --num_timesteps 40000000 \
  --save_buffer_data \
  --wandb_entity your_wandb_entity

Các flag đáng chú ý:

• --domain_randomization: randomize friction, COM, motor gain. Bắt buộc nếu muốn sim2sim sau đó làm việc.
• --num_timesteps 40000000: ~40M env-step. Với 1024 parallel env thì cỡ 40k outer step.
• --save_buffer_data: BẮT BUỘC. Đây là replay buffer sẽ feed vào stage 2.
• --wandb_entity: bỏ nếu không dùng W&B.

Sau pretrain, thư mục log có dạng:

logs/T1LowDimSimFinetuneJoystickFlatTerrain-20260424-101530-abc12/
├── checkpoints/       # config snapshots
├── policies/          # policy_*.pkl (dill-pickled Flax state)
└── buffer_data/       # transitions để train WM

Theo dõi reward: trong W&B, eval/episode_reward nên leo lên ~250+ cho flat terrain. Nếu sau 20M step mà vẫn <100, check domain randomization không bị quá gắt.

Pretrain World Model (Stage 2)

python train_wm_from_file.py \
  --env_name=T1LowDimSimFinetuneJoystickFlatTerrain \
  --data_path=logs/T1LowDimSimFinetuneJoystickFlatTerrain-20260424-101530-abc12

Script sẽ:

1. Đọc tất cả transitions trong buffer_data/.
2. Fit Lagrangian residual (M, C, g corrections) bằng MSE trên q̈.
3. Fit contact residual MLP.
4. Lưu wm_states/wm_*.pkl.

Thời gian: ~15–30 phút trên 4090. Đây là stage "rẻ" nhất. Với người mới, nên:

• Plot predicted q̈ vs actual q̈ cho một vài trajectory để mắt thường check WM không bị off.
• So sánh rollout WM vs rollout env (chạy 1-2s) — nếu divergence < 5% trong 200ms đầu là đủ tốt cho finetune.

Sim2Sim Finetune trong Brax (Stage 3)

Đây là bài test quan trọng nhất — nếu policy pretrain MuJoCo chịu được Brax (động cơ physics khác hoàn toàn) thì khả năng transfer sang real cũng cao.

python finetune.py \
  --env_name=T1LowDimSimFinetuneJoystickFlatTerrain \
  --ac_training_state_path=logs/.../policies/policy_final.pkl \
  --wm_training_state_path=logs/.../wm_states/wm_final.pkl

Behaviour bạn sẽ thấy:

• Iter 0–5: policy zero-shot chạy trong Brax. Nếu pretrain tốt, robot vẫn đi được nhưng có thể wobble khi target velocity > 1 m/s (ngoài distribution pretrain).
• Iter 5–20: action rate tăng nhẹ, policy "sửa" lại balance. World model loss giảm vì thấy nhiều state Brax-specific.
• Iter 20+: stable. Tác giả báo cáo đến 1.5 m/s vẫn walk mượt.

Finetune trên robot thật (Booster T1)

Workflow thực tế:

1. Load policy pretrain đã được sim2sim adapted (output của stage 3).
2. Deploy lên Booster T1 — chạy deterministic 30–60 giây joystick nhẹ.
3. Log state-action-reward vào buffer.
4. Offline update: finetune policy + WM với buffer mới (tương tự script ở trên nhưng env là data thật).
5. Deploy lại.

Tác giả báo cáo: với 80–590 giây data thật, robot đi từ "lắc lư không ổn định" sang "đứng thẳng, gait mượt, bám forward velocity chặt" trên grass, mud, concrete, slope. So với dreamer baselines cần vài giờ, đây là cải thiện ~50–100x.

Nếu bạn quan tâm hướng "residual RL" thay vì "WM finetune", so sánh với SteadyTray: Residual RL trên IsaacLab để thấy hai triết lý khác nhau giải cùng bài toán sim2real.

Inference: đưa policy lên deployment

Policy cuối cùng là một Flax MLP khá nhỏ (thường 256-256-12 cho T1 low-dim). Deploy:

import dill, jax.numpy as jnp
from flax import linen as nn

state = dill.load(open("policy_final.pkl", "rb"))
params = state.params

@jax.jit
def act(obs):
    mu, _ = policy_net.apply(params, obs)
    return mu  # deterministic

# Loop 100Hz
while True:
    obs = robot.get_obs()
    action = act(jnp.asarray(obs))
    robot.send(np.asarray(action))

Latency < 1ms trên Jetson Orin. Chú ý: obs scaling/normalization phải khớp chính xác với pretrain — mismatch 1% cũng đủ phá gait. Lưu mean/std trong checkpoint và dùng lại ở deploy.

Kết quả chính của paper

• Pretrain speed: ~1 giờ trên 4090 cho T1 flat joystick, nhanh 4–8x so với PPO baseline cùng reward.
• Sim2sim transfer: policy pretrain MuJoCo chạy trực tiếp trong Brax ở 1.0 m/s; finetune 20 iteration là đạt 1.5 m/s (ngoài distribution pretrain).
• Real-world T1: với 80–590s data, policy zero-shot trên grass/mud/concrete/slope (không làm lại reward shaping).
• Cross-robot: cùng pipeline chạy trên Unitree G1 (full-body với arm + waist DOF) và T1 (12 DOF leg-only).

Cạm bẫy thường gặp khi reproduce

1. Quên --save_buffer_data → stage 2 không có dữ liệu. Re-run pretrain mất 1 giờ.
2. Domain randomization quá gắt → SAC không converge. Giảm friction range xuống [0.5, 1.0] trước, nâng dần.
3. UTD ratio quá cao → overfitting replay buffer, eval reward tụt. Bắt đầu với UTD=4, tăng dần.
4. WM residual explode → L2 penalty cho residual thấp, robot bay lên trời khi rollout trong WM. Tăng regularization trên M, C corrections.
5. Finetune trên real không có safety → robot ngã vật lý. Luôn có harness + e-stop, limit action magnitude ở 70% dải pretrain trong iter đầu.

So sánh nhanh với các phương pháp khác

LIFT không phải "silver bullet" — nếu bạn đã có pipeline PPO + DR chạy tốt và đội ngũ không quen JAX, chuyển đổi sẽ tốn thời gian. Nhưng nếu làm từ đầu và quan tâm đến data efficiency, đây là baseline đáng copy.

Khi nào nên/không nên dùng LIFT?

Nên dùng khi:

• Bạn có 4090 hoặc H800 và tối ưu data efficiency (ít data thật).
• Đội làm quen JAX/Flax.
• Robot dễ gãy, muốn giảm real-world rollout xuống phút.

Cân nhắc khác khi:

• Bạn đã có pipeline PPO + IsaacLab chạy tốt — chi phí chuyển đổi lớn.
• Task reward cần reward shaping phức tạp (không phải locomotion) — paper gốc chủ yếu test locomotion.
• Không có simulator tương đương MuJoCo Playground cho robot custom — phải port hardware trước.

Tổng kết

LIFT đưa ra một recipe rõ ràng cho humanoid control: pretrain SAC off-policy trong MuJoCo Playground thật nhanh, học một world model có physics structure, rồi finetune bằng cách tách bạch exploration (trong WM) và data collection (deterministic, trong env). Kết quả là pipeline giảm real-world data xuống còn vài trăm giây — một bước tiến lớn cho humanoid RL thực tiễn. Source code đầy đủ tại github.com/bigai-ai/LIFT-humanoid, paper tại arXiv:2601.21363.

Bài viết liên quan

• Sim Series 2: MuJoCo toàn tập cho robotics
• Sim Series 5: Sim2Real Pipeline từ A đến Z
• SteadyTray: Residual RL dạy Unitree G1 bưng khay