vla-jepa-put-mango / config.json
HarrisonPENG's picture
Upload final put-mango v21 model
7d06867 verified
Raw
History Blame Contribute Delete
3.78 kB
{
"run_id": "put_mango_v21_ft_20260717_232629",
"run_root_dir": "/mnt/disk1-1/phs/VLA-JEPA/outputs",
"seed": 42,
"trackers": [
"json",
"wandb"
],
"wandb_project": "vla-jepa",
"wandb_entity": "livion",
"wandb_group": "put_mango_v21",
"is_debug": false,
"framework": {
"name": "VLA_JEPA",
"qwenvl": {
"base_vlm": "/mnt/disk1-1/phs/VLA-JEPA/checkpoints/Qwen3-VL-2B-Instruct",
"attn_implementation": "sdpa",
"vl_hidden_dim": 2048
},
"action_model": {
"action_model_type": "DiT-B",
"action_hidden_dim": 1024,
"hidden_size": 1024,
"add_pos_embed": true,
"max_seq_len": 1024,
"action_dim": 14,
"state_dim": 14,
"future_action_window_size": 6,
"action_horizon": 7,
"past_action_window_size": 0,
"repeated_diffusion_steps": 8,
"noise_beta_alpha": 1.5,
"noise_beta_beta": 1.0,
"noise_s": 0.999,
"num_timestep_buckets": 1000,
"num_inference_timesteps": 4,
"num_target_vision_tokens": 32,
"diffusion_model_cfg": {
"cross_attention_dim": 2048,
"dropout": 0.2,
"final_dropout": true,
"interleave_self_attention": true,
"norm_type": "ada_norm",
"num_layers": 16,
"output_dim": 1024,
"positional_embeddings": null
}
},
"vj2_model": {
"base_encoder": "/mnt/disk1-1/phs/VLA-JEPA/checkpoints/vjepa2-vitl-fpc64-256",
"depth": 12,
"num_heads": 8,
"special_action_token": "<|action_{}|>",
"num_action_tokens_per_timestep": 8,
"embodied_action_token": "<|embodied_action|>",
"num_embodied_action_tokens_per_instruction": 32,
"num_frames": 8
},
"reduce_in_full_precision": true
},
"datasets": {
"vla_data": {
"dataset_py": "lerobot_datasets",
"data_root_dir": "/mnt/disk1-1/phs/VLA-JEPA/datasets",
"data_mix": "put_mango_v21",
"action_type": "absolute_qpos",
"CoT_prompt": "Your task is {instruction}. Infer the temporal dynamics from frames {actions} and produce the corresponding policy actions {e_actions}.",
"resolution_size": 224,
"video_resolution_size": 256,
"per_device_batch_size": 16,
"num_workers": 8,
"delete_pause_frame": false,
"wm_view_indices": [
0,
2
],
"load_all_data_for_training": true,
"with_state": true
}
},
"trainer": {
"epochs": 100,
"max_train_steps": 30000,
"num_warmup_steps": 5000,
"save_interval": 5000,
"keep_latest_checkpoint_only": true,
"save_final_model": true,
"eval_interval": 500,
"learning_rate": {
"base": 3e-05,
"qwen_vl_interface": 1e-05,
"vj_predictor": 3e-05,
"action_model": 0.0001
},
"lr_scheduler_type": "cosine_with_min_lr",
"scheduler_specific_kwargs": {
"min_lr": 1e-06
},
"freeze_modules": "vj_encoder",
"pretrained_checkpoint": "/mnt/disk1-1/phs/VLA-JEPA/checkpoints/vla-jepa-pretrain/Pretrain/checkpoints/VLA-JEPA-pretrain.pt",
"allow_mismatched_modules": "action_model.state_encoder.layer1, action_model.action_encoder.layer1, action_model.action_decoder.layer2",
"loss_scale": {
"vla": 1.0,
"vlm": 0.1
},
"max_grad_norm": 1.0,
"weight_decay": 0.0,
"logging_frequency": 20,
"gradient_clipping": 1.0,
"gradient_accumulation_steps": 1,
"optimizer": {
"name": "AdamW",
"betas": [
0.9,
0.95
],
"eps": 1e-08,
"weight_decay": 1e-08
},
"resume_from_checkpoint": null,
"enable_gradient_checkpointing": true,
"enable_mixed_precision_training": true
},
"output_dir": "/mnt/disk1-1/phs/VLA-JEPA/outputs/put_mango_v21_ft_20260717_232629"
}