Modelpublic

seungkukim/dexjoco_bimanual_multitask_baseline_joint_wan22ti2v5b_L18_txtcache-60k

sourceHugging Faceotherupdated 1mo agoView on Hugging Face

0likes5downloads

Model Card

dexjocobimanualmultitaskbaselinejointwan22ti2v5bL18_txtcache @ step 60,000

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

This is the baseline cell of the 2x2 {Cosmos, Wan} x {baseline, hand stream}: same backbone, layer, text cache, canvas and budget as the hand arms, with the hand-geometry stream switched off. A delta against `...handaction_joint_wan22ti2v5b_L18_txtcache-60k` therefore isolates the hand stream.

What is in here

Inference files only -- 12 files, 11.76 GB. The 1,268 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.*`	247	DiT-B action head (no hand encoder/decoder)

There is no hand stream (hand_action=false, hand_action_norm_stats_path=null), so unlike the hand arms this repo needs no whitening-stats sidecar to construct.

Training-only artifacts are deliberately excluded: the DeepSpeed global_step60000/ optimizer state, rng_state_*.pth, scheduler.pt, trainer_state.json, training_args.bin, zero_to_fp32.py, latest, wandb_config.json.

Arm configuration

key	value
`training_mode`	`joint`
`hand_action`	`False`
`action_head_model_type`	`dit`
`action_horizon`	`16`
`max_action_dim`	`44`
`wan_feature_extraction_layer`	`18` of 30

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:

bash

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

python

from gr00t.model.wam_dit4dit import WAMDiT4DiT

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