Chạy TREAD: tái gán nhãn robot bằng VLM

TREAD, viết tắt của Task Robustness via Re-Labelling Vision-Action Robot Data, là một hướng rất thực dụng cho người đang train Vision-Language-Action model: thay vì thu thêm data robot đắt đỏ, ta dùng Vision-Language Model (VLM) để cắt demonstration dài thành các sub-task ngắn, gán lại nhãn ngôn ngữ cho từng đoạn, rồi fine-tune policy trên mixture gồm data gốc và data đã tái gán nhãn.

Paper gốc của Artur Kuramshin, Özgür Aslan, Cyrus Neary và Glen Berseth xuất hiện trên arXiv ngày 9/6/2026 và project page công bố code tại akuramshin/tread. Ý tưởng nghe đơn giản, nhưng giá trị nằm ở chi tiết: VLM không chỉ paraphrase instruction. Nó nhìn first frame, nhìn video trajectory, tạo semantic sub-task, tìm temporal boundary, sau đó sinh nhiều instruction có grounding theo object properties và spatial relationships. Trên LIBERO, chính lượng đa dạng mới này giúp Octo và π0-FAST bớt giòn khi gặp task hoặc câu lệnh chưa thấy trong training.

Bài này hướng dẫn từ ý tưởng paper đến cách chạy pipeline thực tế: chuẩn bị video LIBERO, gọi Gemini 2.5 Pro qua TREAD, tạo HDF5 sub-task, fine-tune VLA, chạy inference/evaluation, và đọc kết quả. Nếu bạn đã đọc OpenVLA deep dive trước đó, hãy xem TREAD như phần data engine đứng trước VLA: model vẫn quan trọng, nhưng chất lượng nhãn robot còn quan trọng hơn khi muốn generalize.

Tại sao cần tái gán nhãn robot data?

Robot demonstration thường được lưu theo dạng:

instruction: "Put the white mug on the plate and put the chocolate pudding to the left of the plate"
trajectory:  image_0, state_0, action_0
             image_1, state_1, action_1
             ...
             image_T, state_T, action_T

Vấn đề là instruction trên bao phủ cả một chuỗi hành động dài. Trong lúc policy học imitation learning, mỗi action ở giữa trajectory đều bị condition bởi cùng một câu lệnh tổng quát. Với task dài, câu lệnh đó chứa nhiều goal: grasp mug, lift mug, move to plate, place mug, grasp pudding, move pudding, place left of plate. Nếu model đang ở frame đã cầm mug nhưng text vẫn nói cả pudding, tín hiệu supervision khá nhiễu.

TREAD sửa điểm này bằng cách biến một trajectory dài thành nhiều language-action pair ngắn hơn:

Full demo:
  "Put the white mug on the plate and put the pudding left of the plate"

TREAD segments:
  00:00-00:02  "Grasp the white mug"
  00:02-00:04  "Lift the white mug"
  00:04-00:07  "Move the white mug to the plate"
  00:07-00:09  "Place the white mug on the plate"
  00:09-00:11  "Grasp the chocolate pudding"
  00:11-00:14  "Move the chocolate pudding left of the plate"
  00:14-00:16  "Place the chocolate pudding on the table"

Điểm quan trọng: các đoạn này vẫn dùng action thật từ demonstration gốc. TREAD không hallucinate action mới, không cần teleoperation mới, cũng không cần human annotator ngồi đánh dấu từng frame. VLM chỉ đóng vai trò semantic parser + temporal annotator + language augmenter.

Kiến trúc TREAD trong một sơ đồ

TREAD có ba stage chính, đúng với mô tả trên project page và paper:

                 original LIBERO / Bridge demonstrations
                                  |
                                  v
        +----------------- semantic decomposition -----------------+
        | input: first frame + original instruction                 |
        | output: ordered sub-task list                            |
        +----------------------------------------------------------+
                                  |
                                  v
        +------------------- motion segmentation ------------------+
        | input: full trajectory video + sub-task list             |
        | output: time_range for each sub-task                     |
        +----------------------------------------------------------+
                                  |
                                  v
        +---------------- grounded language diversity -------------+
        | input: first frame of segment + sub-task instruction     |
        | output: paraphrases with object attributes/spatial cues  |
        +----------------------------------------------------------+
                                  |
                                  v
           sliced HDF5 sub-task datasets + original full demos
                                  |
                                  v
                     fine-tune Octo / pi0-FAST / VLA

Trong paper, nhóm tác giả dùng Gemini 2.5 Pro làm VLM mặc định. Repo công khai giữ thiết kế provider-neutral: prompt và dataset adapter không phụ thuộc Gemini, nhưng vlm_client.py hiện chỉ có Gemini adapter out of the box. Nếu bạn muốn dùng Qwen2.5-VL, GPT-4o, Claude vision hoặc một VLM nội bộ, bạn cần implement interface VLMClient để chuyển VLMMessage sang API tương ứng và trả text response.

Paper đã thử nghiệm như thế nào?

TREAD được đánh giá trên LIBERO, benchmark manipulation trong simulation. LIBERO có nhiều suite để kiểm tra distribution shift về object, spatial arrangement, goal và long-horizon composition. Trong thí nghiệm TREAD, nhóm tác giả tập trung vào LIBERO-100, mỗi task gốc có 50 human-teleoperated demonstrations. Vì giới hạn tài nguyên, họ chỉ label 5 demonstration cho mỗi task và bỏ các task thuộc STUDY_SCENE, tổng cộng 570 trajectories.

Bảng setup quan trọng:

Điểm tinh tế của setup là có ba mixture:

1. Original Fine-tuned: chỉ fine-tune trên trajectory gốc.
2. TREAD w/o diverse labels: trộn trajectory gốc với sub-trajectory đã cắt nhưng chưa paraphrase đa dạng.
3. TREAD: trộn trajectory gốc với sub-trajectory đã cắt và nhãn ngôn ngữ đa dạng.

Nhờ ablation này, ta tách được hai tác dụng: cắt trajectory giúp motion/planning generalization, còn paraphrase giúp language-conditioned generalization.

Kết quả chính trên LIBERO

Paper báo cáo success rate theo nhiều case. Dưới đây là các số đáng nhớ nhất:

Cách đọc bảng: TREAD không biến mọi case thành SOTA tuyệt đối. Hai-goal full success vẫn khó, nhất là Motion Generalization 2 of 2, nơi cả Octo và π0-FAST chưa hoàn thành trọn vẹn task dài. Nhưng TREAD cải thiện mạnh ở single-goal và partial completion, đặc biệt với Octo trên môi trường mới. Đây là tín hiệu hợp lý: khi demo dài được cắt thành skill ngắn, policy dễ tái tổ hợp skill đó trong scene khác hơn.

Điểm quan trọng khác là LIBERO-10 in-distribution không bị sụp. TREAD giữ performance tương đương original fine-tuned cho long-horizon task quen thuộc, nên augmentation không chỉ đổi data để overfit vào benchmark tự tạo.

Chuẩn bị môi trường

Repo TREAD rất gọn. Requirements công khai chỉ gồm opencv-python, google-genai, h5py, numpy. Tuy nhiên để chạy end-to-end trên LIBERO, bạn cần thêm môi trường LIBERO/robomimic/Octo tương thích. Nên tách thành hai env:

# Env 1: chỉ chạy VLM relabeling và tạo dataset
conda create -n tread python=3.10 -y
conda activate tread
git clone https://github.com/akuramshin/tread.git
cd tread
pip install -r requirements.txt
export GEMINI_API_KEY="your-gemini-api-key"

# Env 2: fine-tune/evaluate policy, tùy stack của bạn
conda create -n libero-octo python=3.10 -y
conda activate libero-octo
# cài LIBERO, robomimic, Octo hoặc VLA repo bạn dùng

Nếu bạn chỉ muốn hiểu pipeline, env 1 là đủ. Nếu muốn cải thiện VLA thật, bạn cần env 2 có thể đọc HDF5 kiểu LIBERO/robomimic và train policy.

Bước 1: tạo video trajectory từ LIBERO

TREAD nhận input là .mp4 trajectory videos. Mỗi video cần đi kèm instruction, scene và demo id. Có hai cách:

data/libero_videos/
  trajectory_000.mp4
  trajectory_001.mp4
  metadata.json

metadata.json nên có dạng:

{
  "trajectory_000.mp4": {
    "task": "Open the top drawer of the cabinet.",
    "scene": "KITCHEN_SCENE4",
    "demo_id": 0
  },
  "trajectory_001.mp4": {
    "task": "Put the white mug on the plate.",
    "scene": "LIVING_ROOM_SCENE2",
    "demo_id": 1
  }
}

Nếu không có metadata, repo sẽ parse tên file LIBERO theo convention SCENE_task_words_demo_<id>.mp4. Với beginner, tôi khuyên dùng metadata.json vì ít lỗi hơn, đặc biệt khi instruction có dấu gạch, số thứ tự hoặc object name dài.

Một checklist trước khi gọi VLM:

Bước 2: semantic decomposition

Stage đầu tiên sinh danh sách sub-task từ first frame và instruction gốc:

python semantic_segmentation.py \
  --input-dir data/libero_videos \
  --output-dir labels/libero \
  --provider gemini \
  --model gemini-2.5-pro \
  --temperature 0.0 \
  --api-key-env GEMINI_API_KEY \
  --max-retries 2

Output chính là:

labels/libero/dataset_subtask_labels.json

Bạn nên mở một vài entry để kiểm tra nhanh:

{
  "trajectory_000.mp4_0": {
    "task": "Put the white mug on the plate and put the chocolate pudding to the left of the plate.",
    "semantic_subtasks": [
      "Grasp the white mug",
      "Lift the white mug",
      "Move the white mug to the plate",
      "Place the white mug on the plate",
      "Grasp the chocolate pudding",
      "Move the chocolate pudding to the left of the plate",
      "Place the chocolate pudding on the table"
    ]
  }
}

TREAD scripts có cơ chế resume: nếu JSON output đã có trajectory, lần chạy sau sẽ skip trajectory đó. Điều này quan trọng vì mỗi call VLM tốn tiền và có thể bị rate limit.

Bước 3: motion segmentation

Stage hai dùng full video và danh sách sub-task để tìm time range:

python motion_segmentation.py \
  --input-dir data/libero_videos \
  --subtask-labels labels/libero/dataset_subtask_labels.json \
  --output-dir labels/libero \
  --provider gemini \
  --model gemini-2.5-pro \
  --temperature 0.0 \
  --api-key-env GEMINI_API_KEY \
  --max-retries 2

Output:

labels/libero/dataset_motion_labels.json

Format motion label thường giống:

{
  "trajectory_000.mp4_0": {
    "demo_id": 0,
    "motion_labels": [
      {
        "sub_task": "Grasp the white mug",
        "time_range": "00:00 - 00:02"
      },
      {
        "sub_task": "Place the white mug on the plate",
        "time_range": "00:07 - 00:09"
      }
    ],
    "response": "full VLM response"
  }
}

Đây là bước dễ lỗi nhất. Nếu time range quá ngắn, HDF5 segment có thể không đủ frame để policy học. Nếu time range quá dài, segment lại chứa nhiều skill. Một cách kiểm tra thực dụng là viết script overlay label lên video hoặc dùng notebook đọc dataset_motion_labels.json, lấy frame start/end và xem ảnh đầu/cuối.

Bước 4: tạo grounded paraphrases

TREAD không dừng ở việc cắt trajectory. Stage ba sinh paraphrase có grounding từ visual context:

python language_paraphrasing.py \
  --dataset libero \
  --input-dir data/libero_videos \
  --subtask-labels labels/libero/dataset_motion_labels.json \
  --output-dir labels/libero

Sau đó sample paraphrase vào motion labels:

python augment_instructions.py \
  --input-dir labels/libero \
  --task-relabel-prob 0.25 \
  --subtask-relabel-prob 0.5 \
  --seed 0

task-relabel-prob điều khiển xác suất thay instruction task-level, còn subtask-relabel-prob điều khiển xác suất thay sub-task label. Bạn không nên đặt cả hai bằng 1.0 ngay từ đầu, vì model vẫn cần thấy canonical instruction gốc để không mất alignment với benchmark. Paper thử nhiều ratio và chọn mixture bằng Re-Mix; với dự án nhỏ, bắt đầu bằng giá trị trong README là đủ.

Bước 5: slice HDF5 LIBERO thành sub-task dataset

Khi đã có timecode, dùng create_dataset.py để cắt HDF5:

python create_dataset.py \
  --motion-segmentations labels/libero/dataset_motion_labels.json \
  --hdf5-dir data/libero_hdf5 \
  --output-dir data/libero_subtasks \
  --fps 10 \
  --image-size 256 256 \
  --name-by-label

Script này giả định schema kiểu LIBERO/robomimic. Nó đọc actions, rewards, states, robot_states, dones, cùng các observation keys như agentview_rgb, eye_in_hand_rgb, gripper_states, joint_states, ee_states. Với mỗi sub-task, script tạo một HDF5 mới, set reward và done ở frame cuối, resize ảnh, và ghi instruction mới.

Một chi tiết dễ bỏ qua: script mặc định swap left/right words để khớp convention camera của LIBERO. Nếu dataset riêng của bạn không cần sửa trái/phải, chạy:

python create_dataset.py \
  --motion-segmentations labels/libero/dataset_motion_labels.json \
  --hdf5-dir data/libero_hdf5 \
  --output-dir data/libero_subtasks \
  --fps 10 \
  --image-size 256 256 \
  --no-swap-left-right

Bước 6: fine-tune VLA trên mixture

TREAD không ép bạn dùng một model duy nhất. Paper dùng Octo-Small 1.5 và π0-FAST. Với Octo, tác giả fine-tune 50.000 steps, batch size 256, warmup + cosine decay, rồi dùng checkpoint step 30.000 cho báo cáo chính. Với π0-FAST, họ full fine-tune 30.000 steps, batch size 32, và dùng checkpoint step 15.000.

Pseudo-config cho Octo có thể hình dung như sau:

dataset_mix:
  original_libero:
    path: data/libero_hdf5
    weight: 1.0
  tread_subtasks:
    path: data/libero_subtasks
    weight: 1.1

training:
  batch_size: 256
  max_steps: 50000
  lr_schedule: warmup_cosine
  checkpoint_every: 5000
  eval_checkpoints: [30000, 50000]

Nếu bạn dùng OpenVLA/LoRA thay Octo, mapping vẫn tương tự:

observation images + proprioception + relabeled instruction
    -> VLA tokenizer / processor
    -> action head
    -> supervised action loss trên segment actions

Điểm cần giữ là đừng chỉ train trên sub-task. Hãy trộn data gốc với data cắt nhỏ. Data gốc vẫn dạy policy hoàn thành long-horizon instruction, còn data TREAD dạy policy hiểu các skill ngắn và biến thể ngôn ngữ.

Bước 7: inference và evaluation

Repo có thư mục evaluation_pipeline/ cho rollout LIBERO với Octo checkpoint. README nói rõ phần này optional và cần Octo/LIBERO branch tương thích. Cấu hình chính nằm ở conf/config.yaml hoặc override bằng command line:

finetuned_path: /path/to/checkpoints
dataset_statistics_path: /path/to/dataset_statistics.json
test_case_files:
  - /path/to/language_generalization.json
  - /path/to/motion_generalization.json
checkpoints:
  - name: "original"
    path: "${finetuned_path}/original"
    model_step: 30000
  - name: "tread"
    path: "${finetuned_path}/tread"
    model_step: 30000

Chạy rollout:

python evaluation_pipeline/src/eval_libero.py \
  finetuned_path=/path/to/checkpoints \
  dataset_statistics_path=/path/to/dataset_statistics.json

Vẽ kết quả:

python evaluation_pipeline/src/visualize_results.py \
  output_dir=evaluation_pipeline/results

Khi tự đánh giá, hãy tách ít nhất ba bucket: task gốc, instruction paraphrase, và scene-instruction pairing mới. Nếu chỉ test task gốc, bạn sẽ bỏ lỡ mục tiêu chính của TREAD là robustness trước language và motion composition shift.

Debug checklist cho beginner

Khi nào nên dùng TREAD?

TREAD hợp với các tình huống sau:

• Bạn có demonstration dài nhưng instruction quá coarse.
• Bạn muốn tăng language diversity mà không thuê annotator.
• Bạn đang train VLA/BC model trên LIBERO, BridgeData hoặc dataset HDF5 tương tự.
• Bạn muốn policy làm tốt hơn với instruction biến thể như "close the drawer" thay vì câu gốc dài hơn.
• Bạn cần phân tích xem model fail vì không hiểu language hay vì không biết motion.

TREAD không phải thuốc chữa mọi vấn đề. Nếu camera không nhìn thấy object, VLM cũng khó segment. Nếu action gốc xấu, relabel tốt vẫn không sửa được control noise. Nếu task cần force feedback hoặc contact-rich dexterity mà observation chỉ có RGB, sub-task label không đủ để policy học dynamics. TREAD mạnh nhất khi bottleneck nằm ở semantic supervision chứ không phải actuator, simulator hay controller.

Liên hệ với các VLA pipeline khác

TREAD bổ sung rất tốt cho các pipeline như VLA-Adapter trên LIBERO hoặc VLA-0 action-as-text. VLA-Adapter tập trung vào kiến trúc nhỏ và efficient training; VLA-0 đặt action vào không gian text/token; TREAD thì đứng trước cả hai, cải thiện data bằng relabeling. Với team nhỏ, chiến lược hợp lý là:

1. Train/fine-tune baseline trên data gốc
2. Chạy TREAD để tạo sub-task + paraphrase
3. Fine-tune lại trên mixture
4. So sánh language generalization và motion generalization
5. Chỉ khi vẫn thiếu performance mới scale model hoặc thu thêm data

Trong robotics, thu data thật luôn đắt. TREAD nhắc ta rằng trước khi mua thêm robot hours, hãy kiểm tra xem data hiện có đã được mô tả đủ tốt chưa.

Nguồn tham khảo

• Paper: Task Robustness via Re-Labelling Vision-Action Robot Data
• Project page: akuramshin.github.io/tread
• Code: github.com/akuramshin/tread
• Data annotations: real-lab/tread_annotations

Bài viết liên quan

• OpenVLA deep dive
• VLA-Adapter trên LIBERO
• VLA-0 action-as-text trên LIBERO