miniFranka/wla_pretrain_diff_codefront_ViT_yam_egozed_3k_11k

WLA-Cosmos · code-front ViT planner · yam+egozed 3k h · step 11k

Tri-modal video · code · action pretraining checkpoint (yamego_codefront_vit, step 11,000) from the WLA-Cosmos / DexWAM project. Three experts co-denoise in a single joint loop:

expert	backbone	objective
video	Cosmos-Predict2.5-2B DiT (Wan2.2 48-ch VAE, frozen)	flow matching
code ("planner")	SDAR-1.7B (Qwen3 body), LLaDA-style joint masked diffusion	masked-token CE
action	action DiT (0.72B)	flow matching

Topology code_front_vit. The planner is a pure-VLM island: its stream is [Qwen3-VL-2B ViT prefix (210 tok = 3 views × 70)] [# Instruction: …] [code canvas] with no cross-stream reads. The video expert reads the code stream (plan-steering edge); the action expert reads everything. Both the Qwen3-VL vision tower and its merger are frozen in this run (the merger's output dim 2048 == SDAR hidden, so Qwen's own projector feeds the planner directly; only the code expert adapted to it). Consequently the checkpoint contains no code_vision.* tensors — the tower+merger are rebuilt from the Qwen3-VL-2B-Instruct snapshot at load time, exactly as in training. Video frames are channels-first [B,3,nc*T,H,W] at 224×320; the ViT prefix uses the first (observation) frame of each view.

Training

• Data: YAM + EgoZed egocentric manipulation, ~3k hours, instructed (# Instruction: prompt prepended to the code stream). CFG stream dropout on the code stream p=0.15 from step 0.
• 48× H100, batch 6 × grad-accum 9 (effective 2592), 11k steps ≈ 10% of one epoch.
• Foundation init: Cosmos-Predict2.5-2B (video), SDAR-1.7B (code), Qwen3-VL-2B (vision tower).

Evaluation (step 10k, most recent evaluated; n=24 val windows)

metric	value
grounding gap — teacher-forced code loss(shuffled obs) − loss(real obs)	+1.353 abs (+200% rel)
teacher-forced code acc (val)	0.864
free-generation parse rate	0.833
entity F1 / token F1 / primitive F1	0.547 / 0.416 / 0.362
skill-call rate	0.333

For reference the VAE-obs (non-ViT) code-front arm at 25k reached parse 0.667 with a grounding gap of only +0.130 abs — the ViT prefix grounds the planner ~10× harder.

Files

• checkpoints/weights/step_011000.pt — model weights (model.load_checkpoint(path)).
• config.yaml — full Hydra config of the run (model / data / trainer).
• dataset_stats.json — action normalization statistics used at train time.

Not included (external, referenced by paths in config.yaml — adjust to your environment): Qwen3-VL-2B-Instruct vision tower (visual.*), Wan2.2 VAE, Cosmos-Reason1 text embeddings, SDAR-1.7B tokenizer. The optimizer state (48-rank ZeRO-1 shards, 48 GB) is not included.

Loading (wla_cosmos repo, PYTHONPATH=src)

import torch
from omegaconf import OmegaConf
from hydra.utils import instantiate

mc = OmegaConf.load("config.yaml").model
OmegaConf.set_struct(mc, False)
mc.include_text_encoder = False
mc.load_pretrain = True
model = instantiate(mc, model_dtype=torch.float32, device="cuda:0")   # fp32 for eval
model.load_checkpoint("checkpoints/weights/step_011000.pt")
model.apply_param_freeze_policy(); model.eval()

Joint tri-modal inference: model.infer_action_trimodal(obs_video=..., proprio=..., prompt_embs=..., instruction=..., num_inference_steps=64, ...). Default code decode: confidence-ordered unmasking (WLA_REVEAL_ORDER=confidence, WLA_FL_FRAC=1.0, canvas 64 → 1 token/pass).

Important: use a wla_cosmos commit that includes the 2026-09-01 ViT decode fix ([_vitP:] slicing of the code logits in wla_model.py). Earlier decode loops indexed the logits without accounting for the 210-row ViT prefix and produce garbage code for this model.