Hi, thanks for releasing the post-training code — the yam/AgiBot workflow has been working well for us on the training side.
While fine-tuning DreamZero-AgiBot on a bimanual dataset (yam embodiment, yam_training.sh defaults, save_lora_only=true) we noticed the served model behaved worse than the base checkpoint, even though training loss looked healthy. After some digging we believe VLA.load_lora rebuilds the model on the wrong base.
What training does: create_model in groot/vla/experiment/base.py loads the full DreamZero-AgiBot checkpoint (which is a complete merged model — we checked the index, 2146 tensors and no LoRA keys — not a Wan+adapter bundle), and only then does inject_lora_after_loading() add fresh rank-4 adapters. So the trained deltas are relative to the AgiBot weights.
What inference does: VLA.load_lora (groot/vla/model/dreamzero/base_vla.py:337) sets skip_component_loading=False and instantiates the model, which loads the DiT from diffusion_model_pretrained_path in the checkpoint's config.json. That path was recorded from the dit_version flag, which the shipped training scripts point at the raw Wan2.1 dir. Result: you get raw Wan + the deltas instead of AgiBot + the deltas. Nothing errors — the model loads and runs fine, it's just missing the robot pretraining the adapters were trained against.
The frustrating part is that the correct base path is already saved in the checkpoint — experiment_cfg/conf.yaml has pretrained_model_path: ./checkpoints/DreamZero-AgiBot — load_lora just never reads it. (For the original DROID from-scratch workflow this is all fine, since there raw Wan is the right base; it only bites the post-training workflow.)
We measured the impact. 1000-step yam LoRA fine-tune of DreamZero-AgiBot on lerobot/aloha_static_screw_driver (rev v2.1), then open-loop replay of 150 frames with scripts/open_loop_yam.py (slices adapted to the 14-dim ALOHA layout), first predicted action vs ground truth, MSE in rad²:
| load path |
overall |
left_gripper |
| load_lora as shipped (raw Wan base) |
0.00988 |
0.07602 |
| same checkpoint merged onto AgiBot offline |
0.00100 |
0.00180 |
| AgiBot base, no fine-tune |
0.00736 |
0.00170 |
So the shipped path is ~10x worse than the correct composition, and notably worse than not fine-tuning at all — you can train a good LoRA and end up serving something worse than the base model without any indication that something is off. Might also explain reports like #57 (good video prediction, bad action execution after fine-tuning).
The offline merge for row 2 replicates training's composition: instantiate with defer_lora_injection=True, skip_component_loading=True, load the AgiBot shards, inject_lora_after_loading(), load the LoRA checkpoint, merge_and_unload().
Suggested fix: have load_lora check experiment_cfg/conf.yaml for pretrained_model_path and, when set, load the base from there before overlaying the deltas (same order as training), falling back to the current behavior otherwise. We patched this locally and re-ran the same eval through the patched loader — it reproduces the offline-merge numbers exactly. Happy to open a PR.
Setup: torch 2.8.0+cu129, flash-attn 2.8.3.post1, deepspeed 0.16.9, 4x RTX PRO 6000 Blackwell.
Hi, thanks for releasing the post-training code — the yam/AgiBot workflow has been working well for us on the training side.
While fine-tuning DreamZero-AgiBot on a bimanual dataset (yam embodiment,
yam_training.shdefaults,save_lora_only=true) we noticed the served model behaved worse than the base checkpoint, even though training loss looked healthy. After some digging we believeVLA.load_lorarebuilds the model on the wrong base.What training does:
create_modelingroot/vla/experiment/base.pyloads the full DreamZero-AgiBot checkpoint (which is a complete merged model — we checked the index, 2146 tensors and no LoRA keys — not a Wan+adapter bundle), and only then doesinject_lora_after_loading()add fresh rank-4 adapters. So the trained deltas are relative to the AgiBot weights.What inference does:
VLA.load_lora(groot/vla/model/dreamzero/base_vla.py:337) setsskip_component_loading=Falseand instantiates the model, which loads the DiT fromdiffusion_model_pretrained_pathin the checkpoint's config.json. That path was recorded from thedit_versionflag, which the shipped training scripts point at the raw Wan2.1 dir. Result: you get raw Wan + the deltas instead of AgiBot + the deltas. Nothing errors — the model loads and runs fine, it's just missing the robot pretraining the adapters were trained against.The frustrating part is that the correct base path is already saved in the checkpoint —
experiment_cfg/conf.yamlhaspretrained_model_path: ./checkpoints/DreamZero-AgiBot— load_lora just never reads it. (For the original DROID from-scratch workflow this is all fine, since there raw Wan is the right base; it only bites the post-training workflow.)We measured the impact. 1000-step yam LoRA fine-tune of DreamZero-AgiBot on
lerobot/aloha_static_screw_driver(rev v2.1), then open-loop replay of 150 frames withscripts/open_loop_yam.py(slices adapted to the 14-dim ALOHA layout), first predicted action vs ground truth, MSE in rad²:So the shipped path is ~10x worse than the correct composition, and notably worse than not fine-tuning at all — you can train a good LoRA and end up serving something worse than the base model without any indication that something is off. Might also explain reports like #57 (good video prediction, bad action execution after fine-tuning).
The offline merge for row 2 replicates training's composition: instantiate with
defer_lora_injection=True, skip_component_loading=True, load the AgiBot shards,inject_lora_after_loading(), load the LoRA checkpoint,merge_and_unload().Suggested fix: have
load_loracheckexperiment_cfg/conf.yamlforpretrained_model_pathand, when set, load the base from there before overlaying the deltas (same order as training), falling back to the current behavior otherwise. We patched this locally and re-ran the same eval through the patched loader — it reproduces the offline-merge numbers exactly. Happy to open a PR.Setup: torch 2.8.0+cu129, flash-attn 2.8.3.post1, deepspeed 0.16.9, 4x RTX PRO 6000 Blackwell.