Views
No views yet
面向长程 ALOHA 机器人操作任务的逐帧「阶段分类」模型。共享 Transformer 编码器读取 134 维逐帧特征，预测每个时刻正处于任务的哪个阶段。
| Model | Task | Stages | Best Frame Acc (Viterbi) | Weights |
|---|---|---|---|---|
| fold-box-16 | Carton folding (折纸箱) | 16 | 94.75% | fold-box-16/model.safetensors |
| fold-clothes | Clothes folding (叠衣服) | 5 | 94.74% | fold-clothes/model.safetensors |
MultiTaskStageTransformer — input_proj(134→256) →
sinusoidal PositionalEncoding → TransformerEncoder (6 layers, nhead=8,
d_model=256, dim_ff=512, LayerNorm-first) → per-task frame-level stage_head
(Linear(256, num_stages)).(B, T, 134)._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)..pt → safetensors).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_readystage_names_by_task["fold_clothes"], id → name):0 flattening (达到一个平铺的效果) 1 frist fold (第一次折叠) 2 second fold (第二次折叠)
3 third fold (第三次折叠) 4 retraction (臂收回)Head shapes are in eachconfig.json; stage names are read fromstage_names_by_taskat inference (never hard-coded).
| fold-box-16 | fold-clothes | |
|---|---|---|
| Data | box_data v2.1, 933 valid ep | cloth_all, 4:1 split |
| Split | 839 train / 94 test | 1408 train / 353 test |
| Best epoch | 18 | 12 |
| Optimizer / lr | AdamW / 3e-4 (wd 1e-2) | AdamW / 3e-4 (wd 1e-2) |
| Batch size | 4 | 4 |
| Training mode | single-task (task_loss_weight=0) | single-task (task_loss_weight=0) |
| Metric | fold-box-16 | fold-clothes |
|---|---|---|
| Frame Accuracy | 94.75% | 94.74% |
| Mean IoU | 86.47% | 83.33% |
| Episode Full Overlap | 78.72% | 82.15% |
| Boundary recall @30f | 99.29% | 90.30% |
| Boundary median error | 4 frames | 7 frames |
1import torch
2from safetensors.torch import load_file
3
4from src.models.multitask_stage_transformer import MultiTaskStageTransformer
5
6for sub, task in [("fold-box-16", "fold_box_16"), ("fold-clothes", "fold_clothes")]:
7 state = load_file(f"{sub}/model.safetensors")
8 model = MultiTaskStageTransformer(input_dim=134)
9 model.load_state_dict(state, strict=True)
10 model.eval()
11 with torch.no_grad():
12 x = torch.randn(1, 512, 134) # (B, T, 134)
13 logits = model(x, task_name=task)["stage_logits"][task] # (1, T, n_stages)
14 pred = logits.argmax(-1)
15 # production: monotone Viterbi decode -> stage ids never regressSingle-task checkpoints must be inferred withtask_nameforced to the matching head.
1README.md ← this model card
2fold-box-16/
3 model.safetensors config.json README.md
4fold-clothes/
5 model.safetensors config.json README.mdconfig.json: architecture hyper-parameters, stage_names_by_task, best metrics,
input_dim=134, parameter counts.fold_box_16_best_20260731_110850.pt and
fold_clothes_cloth_all_4_1_20260731_152143.pt (converted 2026-08).pick_up IoU 71.6%,
frist fold 66.8%).