Object Tokenizer: Từ Pixel Thô đến Token Đối Tượng với Mask R-CNN + ViT

Trong bài 1 về kiến trúc cross-attention, ta đã thấy VIMA không điều khiển robot bằng một prompt text thuần túy. Prompt của VIMA có thể xen kẽ chữ, ảnh vật thể, ảnh scene đầy đủ, hoặc nhiều frame của một video demo. Trong bài 2 về VimaBench, ta cũng thấy 17 task không chỉ khác nhau ở câu lệnh, mà còn khác nhau ở cách thông tin trực quan xuất hiện trong prompt: có task hỏi "đặt vật này vào vật kia", có task yêu cầu làm theo video mẫu, có task yêu cầu tránh vùng cấm.

Bài 3 đi vào lớp rất dễ bị bỏ qua nhưng quyết định chất lượng của toàn hệ: object tokenizer. Nếu T5 encoder là nơi "đọc" prompt, object tokenizer là nơi biến ảnh thành các đơn vị mà T5 có thể đọc được. VIMA không đưa toàn bộ ảnh RGB thô vào transformer như một lưới patch khổng lồ. Thay vào đó, nó tách ảnh thành đối tượng, rồi biến mỗi đối tượng thành một token hoặc một cụm token có thông tin hình ảnh lẫn vị trí.

Nói ngắn gọn, pipeline là:

RGB image, front/top views
        |
        v
Mask R-CNN / ground-truth detector
        |
        +--> object mask
        +--> bounding box: top, left, bottom, right
        +--> cropped object image
        |
        v
ViT encodes crop + MLP encodes bbox
        |
        v
object token sequence
        |
        v
interleaved with T5 word tokens in multimodal prompt
        |
        v
T5 prompt encoder -> XAttn GPT controller

Một lưu ý trước khi đọc code: trong paper, tác giả mô tả bước full-scene image bằng domain fine-tuned Mask R-CNN. Trong repo public hiện tại, file vima/utils/object_utils.py mà nhiều bài cũ nhắc tới không còn tồn tại trên nhánh main; đường dẫn đó trả 404 khi kiểm tra. Phần logic model đang mở trong repo nằm rõ ở vima/nn/obj_encoder/, vima/nn/prompt_encoder/, và vima/policy/vima_policy.py. Vì vậy bài này sẽ tách hai tầng: tầng detect/crop theo mô tả paper, và tầng encode object token theo code public thực tế.

Roadmap Series

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

Vì sao không dùng raw image patch?

Cách tự nhiên nhất khi làm vision transformer là cắt ảnh thành patch đều nhau. Ví dụ ảnh 256 x 128, patch 16 x 16 sẽ tạo 128 patch token cho một view. Nếu có hai camera front và top, số token tăng gấp đôi. Nếu prompt có nhiều frame video, số token còn tăng nhanh hơn nữa. Với robot manipulation tabletop, phần lớn patch lại là nền bàn, khoảng trống, bóng đổ, hoặc cạnh vật thể không liên quan.

VIMA chọn hướng khác: object-centric representation. Thay vì hỏi transformer "hãy tự tìm vật thể trong hàng trăm patch", tokenizer nói rõ với model: "đây là một object, đây là crop của nó, đây là bbox của nó trong view front hoặc top". Cách này có ba lợi ích thực tế.

Thứ nhất, chuỗi ngắn hơn. Robot không cần suy luận trên mọi pixel nếu task chỉ xoay quanh 3 đến 8 vật thể trên bàn. Chuỗi object token thường gần với số lượng đối tượng, không gần với diện tích ảnh.

Thứ hai, inductive bias đúng hơn. Trong manipulation, action thường gắn với vật thể: pick object A, place into container B, avoid zone C. Khi token đã là object, attention có thể học quan hệ "vật đỏ bên trái cái bát" trực tiếp hơn so với học từ patch nền lẫn patch vật.

Thứ ba, tổng quát hóa tốt hơn khi dữ liệu ít. Paper VIMA so sánh object tokens với các tokenizer học trực tiếp từ raw pixels và thấy object tokens tốt hơn trên các mức generalization khó. Đây là điểm rất quan trọng cho robotics: thu dữ liệu robot đắt, nên representation nào giảm overfit vào texture/nền/camera sẽ có giá trị lớn.

Bước 1: Tách object bằng Mask R-CNN

Trong paper, VIMA nói rõ có ba dạng input thô trong prompt: text, ảnh một object, và ảnh toàn cảnh tabletop. Với text, VIMA dùng T5 tokenizer/embedding. Với ảnh scene đầy đủ, VIMA trước hết trích từng object bằng Mask R-CNN đã fine-tune theo domain. Mỗi object được biểu diễn bằng hai thứ:

Trong VimaBench, simulator cung cấp cả RGB front view và top-down view. Paper ghi observation space gồm ảnh RGB từ cả frontal view và top-down view, đồng thời có ground-truth segmentation và bounding box cho training object-centric models. Khi đánh giá bằng detector, Mask R-CNN không hoàn hảo: bbox có thể lệch, crop có thể chứa pixel thừa, hoặc detector có thể sinh false positive. VIMA xử lý điều đó bằng object augmentation, tức là ngẫu nhiên chèn detection output giả trong training để policy bền hơn trước lỗi perception.

Điểm beginner cần nhớ: Mask R-CNN không phải là "bộ não" của VIMA. Nó chỉ là bộ chuyển đổi ảnh scene thành danh sách object candidate. Bộ não nằm ở cách object đó được mã hóa thành token và cách token đó tương tác với language prompt.

Bước 2: ViT encode crop, MLP encode bbox

Phần code rõ nhất nằm trong vima/nn/obj_encoder/obj_encoder.py. Class chính là ObjEncoder. Nó nhận cropped_img, bbox, và mask, sau đó xử lý từng camera view. Trong checkpoint mặc định của VIMAPolicy, views là ["front", "top"].

Đoạn code rút gọn:

class ObjEncoder(nn.Module):
    bbox_max_h = 128
    bbox_max_w = 256

    def __init__(self, *, transformer_emb_dim, views, vit_output_dim,
                 vit_resolution, vit_patch_size, vit_width, vit_layers,
                 vit_heads, bbox_mlp_hidden_dim, bbox_mlp_hidden_depth):
        self.cropped_img_encoder = ViTEncoder(
            output_dim=vit_output_dim,
            resolution=vit_resolution,
            patch_size=vit_patch_size,
            width=vit_width,
            layers=vit_layers,
            heads=vit_heads,
        )

        self.bbox_mlp = nn.ModuleDict({
            view: build_mlp(
                4,
                hidden_dim=bbox_mlp_hidden_dim,
                hidden_depth=bbox_mlp_hidden_depth,
                output_dim=bbox_mlp_hidden_dim,
            )
            for view in views
        })

        self.pre_transformer_layer = nn.ModuleDict({
            view: nn.Linear(
                self.cropped_img_encoder.output_dim + bbox_mlp_hidden_dim,
                transformer_emb_dim,
            )
            for view in views
        })

Có hai nhánh song song:

1. cropped_img_encoder: một ViT nhỏ nhận crop object. Code mặc định dùng vit_resolution=32, vit_patch_size=16, vit_layers=4, vit_heads=24, vit_output_dim=768.
2. bbox_mlp: một MLP nhận 4 số bbox. Đây là positional signal của object: object nằm ở đâu trong ảnh, không chỉ nó trông như thế nào.

Trong forward, VIMA encode ảnh crop trước, normalize bbox theo kích thước ảnh, encode bbox bằng MLP, rồi concat hai embedding:

def forward(self, cropped_img, bbox, mask):
    img_feats = {
        view: self.cropped_img_encoder(cropped_img[view])
        for view in self._views
    }

    bbox = {view: bbox[view].float() for view in self._views}
    normalizer = torch.tensor(
        [self.bbox_max_w, self.bbox_max_h, self.bbox_max_h, self.bbox_max_w],
        dtype=bbox[self._views[0]].dtype,
        device=bbox[self._views[0]].device,
    )
    bbox = {view: bbox[view] / normalizer for view in self._views}
    bbox = {view: self.bbox_mlp[view](bbox[view]) for view in self._views}

    in_feats = {
        view: self.pre_transformer_layer[view](
            torch.concat([img_feats[view], bbox[view]], dim=-1)
        )
        for view in self._views
    }

    return torch.concat([in_feats[view] for view in self._views], dim=-2)

Đây là chỗ nhiều người nhầm. "Object token" không chỉ là embedding ảnh crop. Nếu chỉ dùng crop, model biết vật đó trông giống gì nhưng không biết nó nằm ở đâu trên bàn. Nếu chỉ dùng bbox, model biết vị trí nhưng không biết object là bát, khối, vùng cấm hay dụng cụ. VIMA cần cả hai.

Bước 3: ViT biến crop thành vector

Trong vima/nn/obj_encoder/vit/vit.py, ViTEncoder làm ba việc quen thuộc:

• chuẩn hóa ảnh từ [0, 255] về tensor normalized bằng mean/std của VIMA;
• chia crop thành patch bằng convolution stride bằng patch size;
• thêm positional embedding trong crop và chạy các residual attention block.

Với crop 32 x 32 và patch 16 x 16, ViT chỉ có 4 patch cộng với một cls_token. Output cuối cùng là vector 768 chiều cho crop đó. Đây là một lựa chọn rất thực dụng: crop nhỏ, ViT nông, đủ để nhận dạng hình dạng/texture của object nhưng không phình sequence length lên quá lớn.

object crop 32x32
    |
    | patch size 16
    v
4 image patches + cls token
    |
    v
4-layer ViT
    |
    v
crop embedding, 768 dims

Nếu bạn từng dùng ViT cho classification, hãy xem nhánh này như một classifier feature extractor cực nhỏ. Nó không xuất class label; nó xuất feature để policy downstream tự học "object này quan trọng thế nào cho task".

Bước 4: Ghép object token với T5 text token

Object token chỉ hữu ích khi nó đi vào prompt đúng thứ tự. Việc lắp ráp nằm trong VIMAPolicy.forward_prompt_assembly(). Hàm này nhận raw_prompts_token_type, word_batch, và image_batch.

raw_prompts_token_type là một sequence đánh dấu: 0 nghĩa là word token, 1 nghĩa là image token. Khi gặp word, policy lấy embedding từ WordEmbedding, vốn dùng weight pretrained của t5-base. Khi gặp image, policy gọi ObjEncoder, lấy n_max_objs object tokens, rồi chèn toàn bộ object tokens vào prompt sequence.

def forward_prompt_assembly(self, prompts):
    raw_prompts_token_type, word_batch, image_batch = prompts

    batch_word_emb = self.prompt_embedding(word_batch)
    batch_image_emb = self.obj_encoder(**image_batch)
    batch_image_emb = self.prompt_obj_post_layer(batch_image_emb)

    for item in raw_prompt:
        if item == 0:
            assembled_prompt.append(batch_word_emb[word_ptr])
            assembled_mask.append(True)
            word_ptr += 1
        elif item == 1:
            obj_mask = concat_masks_from_front_and_top_views(...)
            for q in range(n_max_objs):
                assembled_prompt.append(batch_image_emb[img_ptr][q])
                assembled_mask.append(obj_mask[q])
            img_ptr += 1

    prompt_tokens = self.t5_prompt_encoder(
        prompt_tokens,
        attention_mask=prompt_masks,
        batch_first=False,
    )

Ví dụ prompt:

"put" <image of red block> "into" <image of bowl>

có thể trở thành:

word("put")
obj(red_block, front)
obj(red_block, top)
word("into")
obj(bowl, front)
obj(bowl, top)

Sau đó chuỗi này đi qua T5 encoder bằng inputs_embeds, không phải input_ids. Đây là trick quan trọng: T5 không cần biết "object token" nằm trong vocabulary nào. Nó chỉ cần nhận vector cùng chiều với word embedding. VIMA dùng MLP để đưa non-text token về không gian tương thích với T5.

Bước 5: Object token đi vào cross-attention

Khi prompt đã thành chuỗi T5-contextualized tokens, decoder XAttn GPT dùng chúng làm key/value. Observation/action history làm query. Đây chính là phần đã bàn trong bài 1, nhưng bây giờ ta hiểu prompt side chứa gì: không chỉ chữ, mà còn object token đã encode từ crop + bbox.

prompt tokens after T5:
    [word, object, object, word, object, object, ...]
        |
        | key/value
        v
XAttn GPT decoder
        ^
        | query
observation tokens + previous action tokens

Trong XAttnGPT.forward(), mỗi layer làm cross-attention trước, rồi self-attention causal:

for self_attn, xattn in zip(self.h, self.xattns):
    obs_action_tokens = xattn(
        q=obs_action_tokens,
        kv=prompt_tokens,
        attention_mask=prompt_mask,
    )
    obs_action_tokens = self_attn(
        obs_action_tokens,
        attention_mask=obs_action_masks,
    )[0]

Kết quả là ở mỗi bước sinh action, policy có thể hỏi lại prompt: "object trong observation hiện tại khớp với object nào trong prompt?", "bbox nào cho biết vùng cần tránh?", "video demo đang chỉ ra motion nào?".

Chạy tokenizer trên ảnh riêng: script tối giản

Repo public cung cấp pretrained Mask R-CNN checkpoint ở Hugging Face, nhưng không expose một API đơn giản kiểu object_utils.py trên nhánh hiện tại. Vì vậy, khi bạn muốn thử trên ảnh riêng, cách thực dụng là tách pipeline thành hai phần:

1. dùng detector/segmenter bất kỳ để tạo bbox, mask, crop;
2. đưa tensor đã chuẩn hóa shape vào ObjEncoder.

Script dưới đây minh họa phần shape và data contract. Detector dùng Detectron2 COCO pretrained cho dễ chạy; nếu muốn bám sát VIMA hơn, bạn thay detector này bằng Mask R-CNN domain-finetuned của VIMA.

from pathlib import Path

import cv2
import torch
import torch.nn.functional as F
from detectron2.config import get_cfg
from detectron2.engine import DefaultPredictor
from detectron2 import model_zoo

from vima.nn.obj_encoder import ObjEncoder

def build_detector(score_thresh=0.5):
    cfg = get_cfg()
    cfg.merge_from_file(model_zoo.get_config_file(
        "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"
    ))
    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = score_thresh
    cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
        "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"
    )
    return DefaultPredictor(cfg)

def crop_objects_bgr(image_bgr, boxes_xyxy, crop_size=32, max_objs=8):
    crops, boxes = [], []
    h, w = image_bgr.shape[:2]
    for box in boxes_xyxy[:max_objs]:
        x1, y1, x2, y2 = [int(v) for v in box]
        x1, y1 = max(0, x1), max(0, y1)
        x2, y2 = min(w, x2), min(h, y2)
        crop = image_bgr[y1:y2, x1:x2]
        if crop.size == 0:
            continue
        crop = cv2.cvtColor(crop, cv2.COLOR_BGR2RGB)
        crop_t = torch.from_numpy(crop).permute(2, 0, 1).float()[None]
        crop_t = F.interpolate(crop_t, size=(crop_size, crop_size), mode="bilinear")
        crops.append(crop_t[0])
        boxes.append(torch.tensor([x1, y1, y2, x2], dtype=torch.float32))

    while len(crops) < max_objs:
        crops.append(torch.zeros(3, crop_size, crop_size))
        boxes.append(torch.zeros(4))

    mask = torch.zeros(max_objs, dtype=torch.bool)
    mask[: min(len(boxes_xyxy), max_objs)] = True
    return torch.stack(crops), torch.stack(boxes), mask

image = cv2.imread(str(Path("my_tabletop_image.png")))
predictor = build_detector()
instances = predictor(image)["instances"].to("cpu")

crops, boxes, valid = crop_objects_bgr(
    image,
    instances.pred_boxes.tensor.numpy(),
    crop_size=32,
    max_objs=8,
)

encoder = ObjEncoder(
    transformer_emb_dim=768,
    views=["front"],
    vit_output_dim=768,
    vit_resolution=32,
    vit_patch_size=16,
    vit_width=768,
    vit_layers=4,
    vit_heads=24,
    bbox_mlp_hidden_dim=768,
    bbox_mlp_hidden_depth=2,
)

image_batch = {
    "cropped_img": {"front": crops.unsqueeze(0)},  # (B, N, 3, 32, 32)
    "bbox": {"front": boxes.unsqueeze(0)},         # (B, N, 4)
    "mask": {"front": valid.unsqueeze(0)},         # (B, N)
}

with torch.no_grad():
    object_tokens = encoder(**image_batch)

print(object_tokens.shape)  # (B, N * num_views, 768)

Đây chưa phải reproduction đầy đủ của VIMA inference, vì policy thật còn cần prompt assembly, observation history, action discretization và checkpoint. Nhưng script này cho bạn thấy điểm quan trọng nhất: sau detect/crop, interface đi vào ObjEncoder chỉ là ba dict tensor theo view.

Nếu bạn làm với hai camera giống VIMA, mở rộng image_batch như sau:

image_batch = {
    "cropped_img": {
        "front": front_crops.unsqueeze(0),
        "top": top_crops.unsqueeze(0),
    },
    "bbox": {
        "front": front_boxes.unsqueeze(0),
        "top": top_boxes.unsqueeze(0),
    },
    "mask": {
        "front": front_valid.unsqueeze(0),
        "top": top_valid.unsqueeze(0),
    },
}

ObjEncoder sẽ sort views, encode từng view riêng, rồi concat theo chiều object sequence. Vì vậy một object có thể tạo hai token: một token từ front view và một token từ top view. Với tabletop manipulation, top view thường tốt cho vị trí planar; front view giúp nhận dạng hình dạng, chiều cao, hoặc trạng thái chồng/lồng.

Debug checklist khi tokenizer lỗi

Khi tự chạy tokenizer, lỗi thường không nằm ở transformer mà nằm ở dữ liệu vào. Checklist nhanh:

Với beginner, mẹo hữu ích nhất là luôn lưu một ảnh debug có bbox và index object. Nếu prompt nói "red block" nhưng object slot 3 lại là cạnh bàn, transformer phía sau khó cứu được.

Kết nối với humanoid manipulation

Series này dùng chữ "tay người" vì mục tiêu cuối là suy nghĩ cách mở rộng VIMA sang humanoid hoặc dual-arm manipulation. Object tokenizer của VIMA vẫn còn rất tabletop: front/top RGB, primitive pick-place/wipe, action space rời rạc hóa theo pose end-effector. Nhưng ý tưởng object-centric lại càng quan trọng khi lên humanoid.

Với humanoid, raw pixel còn rối hơn: camera egocentric di chuyển, tay robot che object, object bị occlusion, toàn thân tạo nhiều vùng ảnh không liên quan. Nếu đưa raw patch trực tiếp vào policy, model phải học cả perception, grounding, affordance và control cùng lúc. Một tokenizer object tốt có thể đóng vai trò như lớp "nén ngữ nghĩa": object nào đang ở đâu, có liên quan gì đến prompt, có nằm trong vùng reachable của tay trái/tay phải không.

Tuy vậy, humanoid cũng cần nhiều thứ VIMA chưa giải quyết:

• object tokenizer 3D hoặc RGB-D, không chỉ bbox 2D;
• tracking object qua thời gian khi camera egocentric rung;
• hand-object contact state, không chỉ crop hình ảnh;
• token cho body part, workspace, obstacle và human instruction;
• liên kết object token với controller whole-body thay vì primitive pick-place.

Đây là lý do bài 5 của series sẽ quay lại câu hỏi: object token trong VIMA là nền móng tốt, nhưng để dùng cho humanoid, ta cần thêm spatial token 3D, memory, và low-level controller phù hợp.

Tóm tắt

Object tokenizer là cầu nối giữa perception và prompt learning trong VIMA. Mask R-CNN tách scene thành object candidate; ViT encode crop để lấy đặc trưng hình ảnh; bbox MLP encode vị trí top/left/bottom/right trong view; các token từ front/top view được flatten thành sequence; cuối cùng sequence này xen kẽ với T5 word token để tạo multimodal prompt.

Nếu bài 1 trả lời "VIMA dùng kiến trúc nào để đọc prompt?", bài này trả lời "ảnh trong prompt được biến thành token như thế nào?". Và nếu bài 2 cho thấy benchmark đòi hỏi generalization rất mạnh, object tokenizer giải thích vì sao VIMA không chết chìm trong pixel: nó ép model nhìn thế giới qua các đơn vị robot thật sự thao tác được.

Bài viết liên quan

• VIMA: Kiến Trúc Cross-Attention cho Tay Robot
• VimaBench: 17 Task, 4 Cấp Tổng Quát Hóa
• Spatial VLA: Khi Robot Cần Hiểu Không Gian 3D