Dataset 650K: Thu Thập Dữ Liệu Đa Nhiệm Vụ Quy Mô Lớn cho VIMA

Ở bài 3 về Object Tokenizer, ta đã biết VIMA biến ảnh thô thành chuỗi object token như thế nào — từng vật thể trên bàn được nhận dạng bởi Mask R-CNN, crop bởi ViT, và đặt vị trí bởi MLP bbox. Nhưng để train được toàn bộ pipeline ấy, VIMA cần một thứ quan trọng không kém kiến trúc: dữ liệu. Và không phải dữ liệu ghi chép từ robot thật, mà là dữ liệu tổng hợp quy mô lớn trong simulation.

Bài 4 này đi vào phần mà nhiều bài viết về VIMA thường bỏ qua: bộ 650K trajectories được xây dựng như thế nào, lưu dưới dạng gì, tải về và dùng trong Python ra sao, và — quan trọng nhất — tại sao một bộ dữ liệu hoàn toàn từ simulation lại đủ để train một agent tổng quát hóa tốt.

Roadmap Series

Bài này là bài 4/5 trong series VIMA: Prompt Đa Phương Thức cho Robot Tay Người:

Vì sao không dùng dữ liệu từ robot thật?

Câu hỏi đầu tiên khi nghe "650K trajectories" là: robot nào thu thập? Câu trả lời là không có robot thật nào cả. Toàn bộ dữ liệu được tạo ra trong PyBullet — một engine vật lý mã nguồn mở — bằng các chương trình oracle có quyền truy cập đặc quyền vào trạng thái simulator.

Điều này có vẻ như là điểm yếu, nhưng thực ra là lựa chọn thiết kế rất có chủ đích. Hãy nhìn bài toán từ góc độ chi phí:

VIMA chọn con đường thứ ba: viết scripted oracle programs — những chương trình biết vị trí chính xác của mọi object trong scene (nhờ truy cập trực tiếp simulator state), và tạo ra các hành động tối ưu tuyệt đối. Kết quả là 650K demonstrations thành công hoàn toàn, không có demos thất bại lẫn vào.

Với 17 task templates được mô tả trong bài 2, mỗi task có khoảng 50K trajectories được procedurally sinh ra bằng cách thay đổi ngẫu nhiên:

• Loại vật thể (hình học + texture)
• Vị trí ban đầu của các objects
• Màu sắc, kích thước
• Số lượng distractors

Một trajectory trong VIMA là một episode hoàn chỉnh: từ lúc robot nhìn vào scene lần đầu, đọc multimodal prompt, đến khi hoàn thành task cuối cùng.

Procedural Generation trong PyBullet: Hoạt Động Như Thế Nào?

"Procedural generation" nghe có vẻ phức tạp, nhưng ý tưởng rất đơn giản. Hãy tưởng tượng bạn đang viết một script để sinh ra màn chơi của một game — mỗi lần chạy, màn chơi có địa hình khác, vị trí enemy khác, nhưng luật chơi vẫn giữ nguyên. VIMA làm đúng điều đó với robot manipulation.

Với mỗi task template (ví dụ: "đặt vật thể có màu giống vật thể này vào trong hộp kia"), VIMA generator:

1. Chọn ngẫu nhiên objects từ tập hợp hình học 3D (hộp, trụ, hình L, hình T, ...) với texture ngẫu nhiên
2. Spawn objects vào PyBullet tại các vị trí không trùng nhau
3. Sinh multimodal prompt: lấy ảnh chụp các "query objects" từ camera front/top, ghép với text instructions theo template
4. Chạy oracle script: script biết exact positions của mọi thứ, tính pick/place poses tối ưu, thực thi
5. Ghi lại trajectory: lưu từng step — observation, action, trạng thái kết thúc

Toàn bộ quá trình này không cần người can thiệp. Chạy song song trên nhiều CPU, 650K trajectories có thể được sinh ra trong vài ngày.

Ảnh: Kết quả ablation trên visual tokenizer và dataset scale. Dùng object tokens + full data cho kết quả tốt nhất, nhưng ngay cả với 1% data, VIMA object-token vẫn đã vượt nhiều baseline. Nguồn: vimalabs.github.io

Cấu Trúc Dataset: Folder, Files, và Schema

Dataset được host trên HuggingFace tại VIMA/VIMA-Data, dung lượng 21.5 GB, license CC BY 4.0.

Tải dataset

# Cách 1: Dùng huggingface-hub CLI
pip install huggingface-hub
huggingface-cli download VIMA/VIMA-Data --repo-type dataset --local-dir ./vima_data

# Cách 2: Dùng git-lfs (cần cài git-lfs trước)
git lfs install
git clone https://huggingface.co/datasets/VIMA/VIMA-Data ./vima_data

Cấu trúc thư mục

vima_data/
├── task_01_visual_manipulation/
│   ├── 000000/
│   │   ├── rgb_front/          # Ảnh camera phía trước (front view)
│   │   │   ├── 0.png           # Frame 0
│   │   │   ├── 1.png           # Frame 1
│   │   │   └── ...
│   │   ├── rgb_top/            # Ảnh camera nhìn từ trên xuống (top-down)
│   │   │   ├── 0.png
│   │   │   └── ...
│   │   ├── obs.pkl             # Observation: segmentation + end-effector state
│   │   ├── action.pkl          # Oracle actions: pick/place poses
│   │   └── trajectory.pkl      # Metadata: steps, task info, objects
│   ├── 000001/
│   └── ...
├── task_02_rotate/
│   └── ...
└── task_17_novel_concept_grounding/
    └── ...

Mỗi trajectory là một thư mục có số thứ tự. Bên trong là hai thư mục ảnh và ba file pickle.

Schema từng file pickle

obs.pkl — Observation tại mỗi timestep:

{
  "ee": array([x, y, z, qx, qy, qz, qw]),  # End-effector pose (6-DoF: position + quaternion)
  "objects": {
    "segm": {
      "front": array(H, W, dtype=int),  # Segmentation mask, front view
      "top":   array(H, W, dtype=int),  # Segmentation mask, top view
    },
    "bbox": {
      "front": array(N_obj, 4),  # Bounding boxes [top, left, bottom, right]
      "top":   array(N_obj, 4),
    },
    "cropped_img": {
      "front": array(N_obj, H_crop, W_crop, 3),  # Cropped RGB per object
      "top":   array(N_obj, H_crop, W_crop, 3),
    }
  }
}

action.pkl — Oracle actions tại mỗi timestep:

{
  "pose0_position": array(N_steps, 2),   # Pick position (x, y) trong world frame
  "pose0_rotation": array(N_steps, 4),   # Pick rotation (quaternion)
  "pose1_position": array(N_steps, 2),   # Place position (x, y)
  "pose1_rotation": array(N_steps, 4),   # Place rotation (quaternion)
}

Lưu ý quan trọng: action space là SE(2) — chỉ có x, y (trên mặt phẳng bàn) và góc xoay, không phải full 6-DoF. Đây là đơn giản hóa hợp lý cho tabletop manipulation vì robot arm không cần lên cao quá khi thao tác với vật thể nằm trên bàn phẳng.

trajectory.pkl — Metadata:

{
  "elapsed_steps": int,        # Số bước thực tế của trajectory
  "task_id": str,              # Tên task template
  "task_kwargs": dict,         # Tham số cụ thể của task (objects, goals, ...)
  "prompt_token_type": list,   # ["text", "image", "text", ...] — cấu trúc prompt
  "prompt_tokens": list,       # Nội dung từng token trong prompt
}

Code Python: Tải và Visualize Dataset

Dưới đây là code đầy đủ để load một trajectory và inspect nội dung:

import pickle
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path

def load_trajectory(traj_dir: str):
    """Load một trajectory từ thư mục."""
    traj_path = Path(traj_dir)
    
    with open(traj_path / "obs.pkl", "rb") as f:
        obs = pickle.load(f)
    
    with open(traj_path / "action.pkl", "rb") as f:
        action = pickle.load(f)
    
    with open(traj_path / "trajectory.pkl", "rb") as f:
        meta = pickle.load(f)
    
    return obs, action, meta

def visualize_trajectory(traj_dir: str, step: int = 0):
    """Hiển thị observation tại một bước cụ thể."""
    obs, action, meta = load_trajectory(traj_dir)
    traj_path = Path(traj_dir)
    
    # Đọc ảnh RGB
    from PIL import Image
    img_front = Image.open(traj_path / "rgb_front" / f"{step}.png")
    img_top   = Image.open(traj_path / "rgb_top"   / f"{step}.png")
    
    fig, axes = plt.subplots(1, 2, figsize=(12, 5))
    axes[0].imshow(img_front)
    axes[0].set_title(f"Front view — Step {step}")
    axes[0].axis("off")
    
    axes[1].imshow(img_top)
    axes[1].set_title(f"Top-down view — Step {step}")
    axes[1].axis("off")
    
    plt.suptitle(
        f"Task: {meta['task_id']} | Steps: {meta['elapsed_steps']}",
        fontsize=14
    )
    plt.tight_layout()
    plt.show()
    
    # In thông tin action
    n_steps = meta["elapsed_steps"]
    print(f"\nTask: {meta['task_id']}")
    print(f"Total steps: {n_steps}")
    print(f"\nAction tại step {step}:")
    print(f"  Pick position (x, y): {action['pose0_position'][step]}")
    print(f"  Pick rotation (quat): {action['pose0_rotation'][step]}")
    print(f"  Place position (x, y): {action['pose1_position'][step]}")
    print(f"  Place rotation (quat): {action['pose1_rotation'][step]}")
    
    # In end-effector state
    if "ee" in obs:
        ee = obs["ee"]
        # ee có shape (n_steps, 7) — position (3) + quaternion (4)
        print(f"\nEnd-effector state tại step {step}:")
        print(f"  Position: {ee[step, :3]}")
        print(f"  Quaternion: {ee[step, 3:]}")
    
    return obs, action, meta

def summarize_dataset(data_root: str):
    """Thống kê nhanh toàn bộ dataset."""
    data_path = Path(data_root)
    task_dirs = sorted(data_path.glob("task_*/"))
    
    print(f"Số task: {len(task_dirs)}")
    total_trajs = 0
    
    for task_dir in task_dirs:
        traj_count = len(list(task_dir.glob("*/")))
        total_trajs += traj_count
        print(f"  {task_dir.name}: {traj_count:,} trajectories")
    
    print(f"\nTổng cộng: {total_trajs:,} trajectories")

# Ví dụ sử dụng
if __name__ == "__main__":
    # Xem một trajectory cụ thể
    obs, action, meta = visualize_trajectory(
        "./vima_data/task_01_visual_manipulation/000000",
        step=0
    )
    
    # Thống kê dataset
    summarize_dataset("./vima_data")

Tải dataset theo chunk (nếu không đủ ổ cứng)

Dataset 21.5 GB có thể quá lớn nếu muốn chạy nhanh. Bạn có thể tải từng task riêng:

from huggingface_hub import snapshot_download

# Chỉ tải task 01 và task 02 để thử
snapshot_download(
    repo_id="VIMA/VIMA-Data",
    repo_type="dataset",
    local_dir="./vima_data_small",
    allow_patterns=["task_01_*/**", "task_02_*/**"],
)

Tại Sao 650K Trajectories Là "Đủ"?

Câu hỏi hợp lý tiếp theo: với 17 task và ~50K trajectories mỗi task, 650K có nhiều không? So sánh với dataset robot thật, con số này rất lớn. Nhưng so với dữ liệu internet dùng cho LLM, nó rất nhỏ.

VIMA trả lời câu hỏi này bằng ablation dataset scale — họ train model 92M với 4 mức dữ liệu khác nhau:

Kết quả nổi bật: với chỉ 1% dữ liệu (~6,500 trajectories), VIMA object-token đã vượt hầu hết các baseline được train với toàn bộ dataset. Đây là minh chứng cho sức mạnh của object-centric representation và multimodal prompting — khi model có inductive bias đúng (token là vật thể, không phải pixel), nó cần ít data hơn nhiều để học tổng quát hóa.

Con số 650K không phải vì model cần nhiều — mà vì coverage task × object × texture × position phải đủ rộng để evaluation 4 cấp generalization (từ L1 đến L4) có ý nghĩa thống kê.

So Sánh Scale Dataset: VIMA vs. Các Hệ Thống Khác

Nhìn vào bảng, VIMA không phải dataset lớn nhất (RT-X lớn hơn nhiều). Điểm đặc biệt của VIMA-Data là:

1. Multimodal prompt đi kèm data: mỗi trajectory có sẵn prompt dạng interleaved text + image, không phải chỉ text command đơn giản
2. Evaluation protocol rõ ràng: data được phân chia sẵn để test L1/L2/L3/L4 generalization (xem thêm bài 2 về VimaBench)
3. 100% thành công: không có failed demonstrations — oracle script đảm bảo mọi trajectory đều hoàn thành task

Model Sizes và Training Insights

VIMA train nhiều kích thước model trên cùng một dataset để đo scalability:

Lưu ý: T5-Base encoder (111M params) là shared weight, không train lại. Kích thước ở trên là phần decoder/controller được train. Tổng model có thể lớn hơn con số trên nhiều.

Tất cả checkpoints đều download được từ HuggingFace:

from huggingface_hub import hf_hub_download

# Tải checkpoint 200M
checkpoint_path = hf_hub_download(
    repo_id="VIMA/VIMA",
    filename="200M.ckpt",
)

# Tải Mask R-CNN checkpoint (dùng để detect objects)
maskrcnn_path = hf_hub_download(
    repo_id="VIMA/VIMA",
    filename="mask_rcnn.ckpt",
)

print(f"VIMA-200M checkpoint: {checkpoint_path}")
print(f"Mask R-CNN checkpoint: {maskrcnn_path}")

Insight từ ablation model sizes

Từ kết quả trong paper, mối quan hệ giữa model size và performance theo từng cấp generalization:

• L1 (placement generalization): Ngay cả VIMA-9M đã đạt > 90% success. Task đơn giản, model nhỏ là đủ.
• L2 (combinatorial generalization): Cần ít nhất 45M để ổn định trên > 70%.
• L3 (novel object appearance): Cần 92M+ để ổn định, nhỏ hơn thì drop mạnh.
• L4 (novel task): Chỉ 200M mới giữ được performance decent. Đây là cấp khó nhất vì task chưa từng thấy.

Bài học: scale model giúp ích ở cấp generalization khó hơn, không phải cấp dễ. Nếu production chỉ cần L1-L2, VIMA-9M hay 45M đã đủ và nhẹ hơn nhiều lần.

Từ Dataset Đến Humanoid — Kết Nối Bài 5

Dataset 650K hiện tại được thu thập cho robot arm 2-finger gripper trong tabletop environment. Action space là SE(2) — phẳng, đơn giản. Đây là điểm giới hạn lớn nhất nếu muốn scale VIMA lên robot humanoid với hai tay, nhiều ngón, và không gian thao tác 3D đầy đủ.

Ở bài 5, ta sẽ xem cách các nhóm nghiên cứu tiếp cận bài toán mở rộng này: từ việc mở rộng action space (SE(2) → SE(3) → full joint space), đến thu thập dữ liệu từ dual-arm simulation, đến kỹ thuật sim2real transfer khi chuyển từ PyBullet sang robot thật.

Dataset lớn thu thập đúng cách là nền tảng không thể thay thế — kiến trúc đẹp mà không có data sẽ không đi đến đâu. VIMA-Data là minh chứng: 650K trajectories simulation, mỗi trajectory gắn kèm multimodal prompt đúng chuẩn, đủ để train một agent tổng quát hóa tốt trên 17 task loại khác nhau.

Tóm Tắt

• VIMA-Data: 650K trajectories sinh từ PyBullet bằng oracle scripts — không cần robot thật
• ~50K trajectories/task, 17 task templates, procedurally generated với object/texture/position ngẫu nhiên
• Format: PKL files (obs.pkl, action.pkl, trajectory.pkl) + PNG frames từ 2 camera (front + top)
• Observation: RGB images, segmentation masks, bounding boxes, end-effector pose (7D: xyz + quaternion)
• Action space: SE(2) — pick pose + place pose, mỗi pose gồm (x, y) + quaternion
• Tải về: VIMA/VIMA-Data trên HuggingFace, 21.5 GB, license CC BY 4.0
• 1% data (~6.5K trajs) đã đủ để vượt baseline trained với 100% data nhờ object-centric representation
• Model sizes: 2M đến 200M params; L4 generalization cần ít nhất 200M

Bài Viết Liên Quan

• Bài 3: Object Tokenizer — Từ Pixel Thô đến Token Đối Tượng với Mask R-CNN + ViT
• Bài 2: VimaBench — 17 Task, 4 Cấp Tổng Quát Hóa
• Bài 5: Humanoid Adaptation — Mở Rộng Lên Robot Hình Người