sokoban_easy_cot_chunk_kinf_branch_20260709_perseg (EMA, bf16)

Fine-tuned BAGEL-7B-MoT action-conditioned visual world model — step-5000 EMA weights, cast fp32 → bfloat16 (inference runs in bf16 autocast, so this is lossless for inference). Inference / eval artifact only: no optimizer state, not a training-resume checkpoint.

• Load into the BAGEL-7B-MoT architecture (base: ByteDance-Seed/BAGEL-7B-MoT); this is the EMA state dict (ema.safetensors).
• Training data: ultrastar111/sokoban_easy_v8_cot_chunk_kinf_branch_20260709_perseg (see manifest).
• Recipe: lr 2e-5, cosine, warmup 300, 5000 steps, tokens 40k, EMA 0.993, cold-start from vanilla BAGEL.
• Eval contract: per-segment decoder (BAGEL_DECODER=perseg), CFG off, stop-required success.
• License: CC-BY-NC-4.0 (research use).

training_data_manifest.txt

STUDY:               GT ENV-FEEDBACK CoT chunk world model (Structure A; re-ground on real frame every K)
chunk_size_K:        inf
date_tag:            20260709_perseg
data_root:           /data/home/raychai/hf_datasets/sokoban_easy_v8_cot_chunk_kinf_branch_20260709_perseg/training
dataset_config_yaml: ./data/configs/vlm_gym_sokoban_easy_imagined_cot_envfb_v8_res256_train.yaml (EnvFB loader 'interleaved_cot_envfb', data-dir overridden)
format:              user(rules+cadence,frame0 ctx); assistant(<think> a<img>...K... </think> {K actions}); user("Env Feedback:" + REAL frame, loss0 INPUT) between chunks
loss:                CE on assistant text; MSE on per-step imagined frames; Env-Feedback frame + label = loss=0 INPUT; frame0/instruction = context
init_from:           VANILLA /home/jiaxin/unified_world_model/pretrained/BAGEL-7B-MoT (EMA)
hparams:             lr 2e-5, warmup 300, cosine, total 5000, tokens 40000, ema 0.993, save_every 2500 (== self-rollout study)
git_rev:             d03c942bc29e604d12b5e0fdcb351a358978db9e