DexGYS grouped-action (arm A2) @ step 60,000
DiT4DiT joint checkpoint — a Wan2.2-TI2V-5B video backbone plus
a DiT-B action head that predicts the 63-D Shadow-Hand keypoint grasp as 6
anatomically grouped action tokens, conditioned on 48 spatial PartField
point-cloud tokens. Trained on
seungkukim/dexgys_video224_a0_lerobot
(37,158 grasps over 1,258 objects) with cached umT5 text embeddings and feature
extraction at block 18.
Full training run name: dexgys_grouped_ditjoint_wan22_L18_g15-9-9-9-12-9_pc48_txtcache512_pd16_ga1_effgb64_local
What is in here
Inference files only — 12 files, 11.84 GB of weights plus a 124 MB PartField bank. The 1,318 tensors in the shards are:
| prefix | tensors | what |
|---|---|---|
backbone.extractor.transformer.dit.* | 825 | finetuned Wan2.2 DiT |
backbone.extractor.vae.* | 196 | Wan2.2 VAE |
action_head.* | 297 | DiT-B head: 6 grouped action encoders/decoders + state encoder |
Training-only artifacts are excluded: the DeepSpeed global_step60000/ optimizer
state (~70 GB), rng_state_*.pth, scheduler.pt, trainer_state.json,
training_args.bin, zero_to_fp32.py, latest, wandb_config.json.
The six anatomical action groups
action_group_widths = [15, 9, 9, 9, 12, 9] (sum = 63 = max_action_dim).
Order is load-bearing — the head slices the 63-D target by these widths, so a
permutation reorders the target with no error:
| group | slice | points |
|---|---|---|
palm | [0:15] | palm, ff_root, rf_root, lf_root, th_root |
ff | [15:24] | ff_pip, ff_dip, ff_tip |
mf | [24:33] | mf_pip, mf_dip, mf_tip (no mf_root) |
rf | [33:42] | rf_pip, rf_dip, rf_tip |
lf | [42:54] | lf_mcp, lf_pip, lf_dip, lf_tip |
th | [54:63] | th_hub, th_dip, th_tip |
The PartField bank ships WITH this repo
DexGYS is stateless, so the state slot carries the object: the state group is
one scalar join key (point_cloud_index) and the processor swaps it for
48 × 1024 PartField triplane tokens read from a frozen memmapped
bank. That bank is data, not weights — it is absent from the shards, and the
processor cannot build the state stream without it. Included here as:
ckpt/dexgys_pc_features/dexgys_partfield_k8.f16
ckpt/dexgys_pc_features/dexgys_partfield_k8.f16.json
Sidecar claims, all asserted at load: n_scenes=1258,
num_tokens=48, feat_dim=1024,
pool_kernel=8, split='train',
frame = 'canonical object frame (render views=1, azimuth 0)'.
The split field matters: the eval bank RENUMBERS point_cloud_index densely over
the 309 test scenes, so every test index is also a valid train index — a bank from
the wrong split returns another object's features with no error anywhere.
External dependencies
Currently wan_model_path = /data/seungku/hf_cache/Wan2.2-TI2V-5B-Diffusers, pc_feature_path = /data/seungku/projects/wam/ckpt/dexgys_pc_features/dexgys_partfield_k8.f16.
config.json and processor_config.json still carry the local training paths (wan_model_path, pc_feature_path) and wan_local_files_only=true. Override both before loading this off-box, or cd into the snapshot and set pc_feature_path=ckpt/dexgys_pc_features/dexgys_partfield_k8.f16.
The base Wan2.2 snapshot is required either way — from_pretrained() builds the
backbone from it and overlays these shards:
huggingface-cli download Wan-AI/Wan2.2-TI2V-5B-Diffusers --local-dir /path/to/Wan2.2-TI2V-5B-Diffusers
Text conditioning is not part of the checkpoint. This run trained against a
precomputed umT5 cache (seq_len 512,
formalize_language=False) selected by the
WAM_TEXT_EMBED_CACHE environment variable, which config.json does not record.
Either bake that cache (scripts/wam_dit4dit/precompute_text_embeds_wan22.py, ~87 GB
at 512) or let umT5 encode live from the base snapshot. formalize_language must
match what the cache was baked with.
Load
from gr00t.model.wam_dit4dit import WAMDiT4DiT
model = WAMDiT4DiT.from_pretrained("seungkukim/dexgys_grouped_wan22ti2v5b_L18_g6_pc48_txtcache512_joint-60k", torch_dtype="bfloat16")