Jiahao28/stage-transformer-robot-manipulation

Stage-Transformer · Robot Manipulation (机器人操作阶段分类)

Per-frame stage classifiers for long-horizon ALOHA robot manipulation tasks. A shared Transformer encoder reads a 134-dim per-frame feature vector and predicts which stage of a task is being executed at every frame (stage ids are monotone — never regress, enforced by a Viterbi decoder in production).

面向长程 ALOHA 机器人操作任务的逐帧「阶段分类」模型。共享 Transformer 编码器读取 134 维逐帧特征，预测每个时刻正处于任务的哪个阶段。

This repo merges two production single-task models:

Model / 模型信息

• —Architecture: MultiTaskStageTransformer — input_proj(134→256) → sinusoidal PositionalEncoding → TransformerEncoder (6 layers, nhead=8, dmodel=256, dimff=512, LayerNorm-first) → per-task frame-level stage_head (Linear(256, num_stages)).
• —Params: ~4.25 M per model.
• —Input: per-frame 134-dim features, shape (B, T, 134).
• —Feature contract: _extract_features(df) -> (T, 134) (feature_schema: combined_state_action_fk_physics_v1): state/action/err/abs_err/ vel/acc (84) · ALOHA FK end-effector poses (10) · ee rolling mean/var (18) · gripper state/sum/diff (4) · gripper rolling stats (8) · progress/remaining (2) · physical features (8).
• —Framework: PyTorch (source .pt → safetensors).
• —License: TBD — confirm with data/project owner before public redistribution.

Stage spaces

fold-box-16 (stage_names_by_task["fold_box_16"], id → name):

0 move_clamp  1 pick_up  2 form_carton  3 place_carton  4 grip_left_flap
5 fold_right_flap  6 grip_right_flap  7 fold_left_flap  8 grip_left_outer
9 clamp_outer_flap  10 press_outer_flap  11 fold_outer_flap  12 tape_take
13 tape_apply  14 smooth_tape  15 return_ready

fold-clothes (stage_names_by_task["fold_clothes"], id → name):

0 flattening (达到一个平铺的效果)  1 frist fold (第一次折叠)  2 second fold (第二次折叠)
3 third fold (第三次折叠)  4 retraction (臂收回)

Head shapes are in each config.json; stage names are read from stage_names_by_task at inference (never hard-coded).

Training summary / 训练概况

Evaluation / 评估结果（test set, Viterbi-decoded）

Usage / 使用示例

import torch
from safetensors.torch import load_file

from src.models.multitask_stage_transformer import MultiTaskStageTransformer

for sub, task in [("fold-box-16", "fold_box_16"), ("fold-clothes", "fold_clothes")]:
    state = load_file(f"{sub}/model.safetensors")
    model = MultiTaskStageTransformer(input_dim=134)
    model.load_state_dict(state, strict=True)
    model.eval()
    with torch.no_grad():
        x = torch.randn(1, 512, 134)                # (B, T, 134)
        logits = model(x, task_name=task)["stage_logits"][task]  # (1, T, n_stages)
        pred = logits.argmax(-1)
    # production: monotone Viterbi decode -> stage ids never regress

Single-task checkpoints must be inferred with task_name forced to the matching head.

Repository layout / 仓库结构

README.md            ← this model card
fold-box-16/
  model.safetensors  config.json  README.md
fold-clothes/
  model.safetensors  config.json  README.md

• —config.json: architecture hyper-parameters, stage_names_by_task, best metrics, input_dim=134, parameter counts.
• —Original source checkpoints: fold_box_16_best_20260731_110850.pt and fold_clothes_cloth_all_4_1_20260731_152143.pt (converted 2026-08).

Limitations / 局限

• —Trained on a single robot/gripper/work-object set; transfer unverified.
• —Weakest stages are short, low-motion segments (e.g. pick_up IoU 71.6%, frist fold 66.8%).
• —Some over-fitting after best epoch (12 / 18) observed.
• —Demo data licensing must be confirmed before public redistribution.

See the upstream project README / docs for full training & evaluation details.