dexjoco_bimanual_multitask_handaction_joint_wan22ti2v5b_L18_txtcache-60k

dexjoco_bimanual_multitask_handaction_joint_wan22ti2v5b_L18_txtcache @ step 60,000

DiT4DiT joint checkpoint -- a Wan2.2-TI2V-5B video backbone plus a DiT-B action head with a hand-geometry stream (66-D) -- finetuned on the DexJoCo bimanual multitask dataset (seungkukim/dexjoco_lerobot_v20), with cached umT5 text embeddings (txtcache) and feature extraction at layer 18.

What is in here

Inference files only -- 12 files, 11.81 GB. The 1,310 tensors in the shards are:

Training-only artifacts are deliberately 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.

Arm configuration

Hand target is 66-D: [right 0:33 | left 33:66]; per hand [flange pos 0:3 | flange rot6d 3:9 | 4 palm-local fingertips 9:21 | 4 palm-local medial (PIP) origins 21:33].

Hand whitening stats are shipped WITH this repo

action_head.py registers the de-whitening constants as non-persistent buffers, so they are not in the shards, and load_hand_norm_stats() raises rather than falling back to identity. Both config.json and processor_config.json name them by the repo-relative path

scripts/dexjoco/dexjoco_bimanual_hand_action_stats_absolute.json

which resolves against the current working directory. That file is included in this repo at exactly that relative path (66-D, mode=absolute, kind=dexjoco_bimanual), so loading works from the snapshot directory. If you load from elsewhere, either cd into the snapshot or set hand_action_norm_stats_path to an absolute path.

Requires the base Wan2.2 snapshot

This repo is not self-contained. WAMDiT4DiT.from_pretrained() first builds the backbone from the base snapshot named by config.json:wan_model_path, then overlays these shards onto it. The umT5 text encoder and tokenizer are never saved into the checkpoint and always come from that base snapshot:

huggingface-cli download Wan-AI/Wan2.2-TI2V-5B-Diffusers --local-dir /path/to/Wan2.2-TI2V-5B-Diffusers

Currently wan_model_path = /data/seungku/hf_cache/Wan2.2-TI2V-5B-Diffusers.

config.json still points at the local training path /data/seungku/hf_cache/Wan2.2-TI2V-5B-Diffusers with wan_local_files_only=true. Override wan_model_path to Wan-AI/Wan2.2-TI2V-5B-Diffusers (and set wan_local_files_only=false) before loading this on another machine.

Load

from gr00t.model.wam_dit4dit import WAMDiT4DiT

model = WAMDiT4DiT.from_pretrained("seungkukim/dexjoco_bimanual_multitask_handaction_joint_wan22ti2v5b_L18_txtcache-60k", torch_dtype="bfloat16")