pocketvla-3m-sft-2d-grab

PocketVLA-3M-SFT-2D-Grab

Vision-language-action model for language-guided pick-and-place with a 2-link planar arm. Stage 1 of 2: supervised behavior cloning on expert demonstrations. The GRPO fine-tuned successor is PocketVLA-3M-GRPO-2D-Grab.

Part of penpenzh/pocketvla — an end-to-end VLA pipeline (simulator, data generation, SFT, GRPO, evaluation) that trains on a single GPU in under an hour. All training data is generated by the project's built-in simulator; no external data is used.

Model

Task

Three balls (red / green / blue) and one bin, randomly placed. The instruction names one color (e.g. "put the blue ball into the bin"). Grasp succeeds when the end effector closes within 0.09 of a ball center; an episode succeeds when the commanded ball lies fully inside the bin rim within 90 steps.

Results

Closed-loop evaluation, greedy policy, fixed seeds.

Usage

The model class is not registered with AutoModel; loading requires the pvla package from the GitHub repository.

git clone https://github.com/penpenzh/pocketvla
cd pocketvla && pip install -r requirements.txt

import torch
from huggingface_hub import snapshot_download
from pvla.model.vla import PocketVLA, PocketVLATokenizer

ckpt = snapshot_download("penpenzh/pocketvla-3m-sft-2d-grab")   # this repository
model, _ = PocketVLA.load_hf(ckpt, torch.device("cpu"))
model.eval()
tok = PocketVLATokenizer()

image = torch.randint(0, 256, (1, 3, 64, 64), dtype=torch.uint8)  # (B, 3, H, W)
enc = tok("put the green ball into the bin")
tokens = torch.tensor([enc["input_ids"]])                          # (B, 10)
pad = tokens != 0
proprio = torch.tensor([[0.69, -0.28, -0.14, 0.89, 0.0, 0.0]])     # (B, 6)

with torch.no_grad():
    action, ball_off, bin_off = model(image, tokens, pad, proprio)
print(action[0].tolist())        # [dq1, dq2, grip]

Closed-loop evaluation on the fixed test sets:

python3 -m pvla evaluate --ckpt <ckpt_dir>
python3 -m pvla.eval_ood --config configs/sft_3m.yaml --ckpt <ckpt_dir>

Training

Behavior cloning on 184,364 state-action pairs from 7,000 expert episodes (closed-form IK expert, 100% success), 15 epochs, lr 3e-4, class-balanced grip and offset losses plus auxiliary deep supervision. Config: configs/sft_3m.yaml (seed 42) in the GitHub repository.

License

Apache License 2.0.