WSA-Large / train_config.json
zaleni's picture
Upload folder using huggingface_hub
3fc84f7 verified
Raw
History Blame Contribute Delete
12.4 kB
{
"dataset": {
"type": "WSA_Large",
"repo_id": "multidata_from_file",
"repo_id_file": "outputs/WSA_Large/_repo_id_files/WSA_Large-wsa_large-multidata-delta-pretrain-2026_05_07_15_59_20.txt",
"root": null,
"episodes": null,
"image_transforms": {
"enable": false,
"preset": null,
"max_num_transforms": 3,
"random_order": false,
"tfs": {
"brightness": {
"weight": 1.0,
"type": "ColorJitter",
"kwargs": {
"brightness": [
0.8,
1.2
]
}
},
"contrast": {
"weight": 1.0,
"type": "ColorJitter",
"kwargs": {
"contrast": [
0.8,
1.2
]
}
},
"saturation": {
"weight": 1.0,
"type": "ColorJitter",
"kwargs": {
"saturation": [
0.5,
1.5
]
}
},
"hue": {
"weight": 1.0,
"type": "ColorJitter",
"kwargs": {
"hue": [
-0.05,
0.05
]
}
},
"sharpness": {
"weight": 1.0,
"type": "SharpnessJitter",
"kwargs": {
"sharpness": [
0.5,
1.5
]
}
},
"affine": {
"weight": 1.0,
"type": "RandomAffine",
"kwargs": {
"degrees": [
-5.0,
5.0
],
"translate": [
0.05,
0.05
]
}
}
}
},
"revision": null,
"use_imagenet_stats": true,
"use_external_stats": true,
"external_stats_path": null,
"external_stats_root": "norm_stats_32",
"weight_rules_path": "configs/weight_rules_wsa_large_pretrain.yaml",
"video_backend": "pyav",
"streaming": false,
"dist_loading": true,
"buffer_size": 1024,
"action_mode": "delta",
"repack_transforms": {
"inputs": [],
"outputs": []
},
"data_transforms": {
"inputs": [],
"outputs": []
},
"model_transforms": {
"inputs": [],
"outputs": []
},
"dataset_dirs": [],
"image_keys": [
"image0",
"image1",
"image2"
],
"image_raw_shapes": [
[
3,
224,
224
],
[
3,
224,
224
],
[
3,
224,
224
]
],
"image_shapes": [
[
3,
224,
224
],
[
3,
224,
224
],
[
3,
224,
224
]
],
"action_keys": [
"default"
],
"action_raw_shapes": [
24
],
"action_shapes": [
24
],
"state_keys": [
"default"
],
"state_raw_shapes": [
24
],
"state_shapes": [
24
],
"num_frames": 33,
"global_sample_stride": 1,
"action_video_freq_ratio": 4,
"video_size": [
384,
320
],
"standardize_video_size_by_cameras": true,
"camera_key": null,
"val_set_proportion": 0.0,
"skip_padding_as_possible": false,
"max_padding_retry": 3,
"concat_multi_camera": "robotwin",
"override_instruction": null,
"return_future_3d_images": true,
"future_3d_target_index": -1,
"processor_num_output_cameras": 3,
"processor_action_output_dim": 24,
"processor_proprio_output_dim": 24,
"processor_resize_shape": null,
"processor_use_stepwise_action_norm": false,
"processor_norm_default_mode": "z-score",
"processor_use_zh_instruction": false,
"processor_delta_action_dim_mask": {
"default": [
true,
true,
true,
true,
true,
true,
false
]
},
"pretrain_multi_embodiment": true,
"text_embedding_cache_dir": "outputs/WSA_Large/text_embeds/pretrain",
"context_len": 128,
"normalization_stats_path": null,
"dataset_sampling_weights": []
},
"env": null,
"policy": {
"type": "WSA_Large",
"n_obs_steps": 1,
"input_features": {
"observation.state": {
"type": "STATE",
"shape": [
24
]
}
},
"output_features": {
"action": {
"type": "ACTION",
"shape": [
24
]
}
},
"device": "cuda:0",
"use_amp": false,
"push_to_hub": false,
"repo_id": "zaleni/WSA-Large",
"private": false,
"tags": null,
"license": null,
"pretrained_path": null,
"variant": "wsa_large",
"model_id": "Wan-AI/Wan2.2-TI2V-5B",
"tokenizer_model_id": "Wan-AI/Wan2.1-T2V-1.3B",
"tokenizer_max_len": 128,
"load_text_encoder": false,
"redirect_common_files": true,
"mot_checkpoint_mixed_attn": false,
"skip_dit_load_from_pretrain": false,
"action_dit_pretrained_path": "checkpoints/wsa_large/ActionDiT_linear_interp_Wan22_alphascale_1024hdim.pt",
"future_3d_pretrained_path": "checkpoints/wsa_large/Future3DExpert_linear_interp_Wan22_alphascale_768hdim.pt",
"native_checkpoint_path": null,
"dtype": "bfloat16",
"action_dim": 24,
"proprio_dim": 24,
"action_horizon": 32,
"n_action_steps": 24,
"num_inference_steps": 10,
"video_dit_config": {
"has_image_input": false,
"patch_size": [
1,
2,
2
],
"in_dim": 48,
"hidden_dim": 3072,
"ffn_dim": 14336,
"freq_dim": 256,
"text_dim": 4096,
"out_dim": 48,
"num_heads": 24,
"attn_head_dim": 128,
"num_layers": 30,
"eps": 1e-06,
"seperated_timestep": true,
"require_clip_embedding": false,
"require_vae_embedding": false,
"fuse_vae_embedding_in_latents": true,
"use_gradient_checkpointing": false,
"video_attention_mask_mode": "first_frame_causal",
"action_conditioned": false,
"action_dim": 24,
"action_group_causal_mask_mode": "group_diagonal"
},
"action_dit_config": {
"action_dim": 24,
"hidden_dim": 1024,
"ffn_dim": 4096,
"num_heads": 24,
"attn_head_dim": 128,
"num_layers": 30,
"text_dim": 4096,
"freq_dim": 256,
"eps": 1e-06,
"use_gradient_checkpointing": false
},
"future_3d_config": {
"hidden_dim": 768,
"ffn_dim": 3072,
"text_dim": 4096,
"freq_dim": 256,
"eps": 1e-06,
"num_heads": 24,
"attn_head_dim": 128,
"num_layers": 30,
"num_query_tokens": 432,
"da3_num_views": 3,
"da3_tokens_per_view": 1296,
"da3_query_dim": 2048,
"query_layer_indices": [
14,
19,
24,
29
],
"future_query_init_std": 0.02,
"query_mode": "slot_noise",
"query_noise_scale": 0.5,
"query_noise_min_sigma": 0.0,
"query_noise_max_sigma": 0.5,
"query_sigma_source": "constant",
"slot_pos_scale": 0.5,
"use_gradient_checkpointing": false
},
"lambda_3d": 0.1,
"da3_variant": "large",
"da3_teacher_layers": [
11,
15,
19,
23
],
"da3_query_dim": 2048,
"da3_tokens_per_view": 1296,
"da3_num_views": 3,
"future_3d_tokens_per_view": 144,
"future_3d_view_attention_layout": "robotwin",
"da3_model_path_or_name": "depth-anything/DA3-LARGE-1.1",
"da3_model_name": null,
"da3_code_root": null,
"da3_teacher_process_res": 504,
"da3_layer_weights": [
1.0,
1.2,
1.4,
1.6
],
"future_query_init_std": 0.02,
"future_3d_query_mode": "slot_noise",
"future_3d_query_noise_scale": 0.5,
"future_3d_query_noise_min_sigma": 0.0,
"future_3d_query_noise_max_sigma": 0.5,
"future_3d_query_sigma_source": "constant",
"future_3d_slot_pos_scale": 0.5,
"future_3d_target_index": -1,
"log_da3_teacher_timing": false,
"video_scheduler": {
"train_shift": 5.0,
"infer_shift": 5.0,
"num_train_timesteps": 1000
},
"action_scheduler": {
"train_shift": 5.0,
"infer_shift": 5.0,
"num_train_timesteps": 1000
},
"lambda_video": 1.0,
"lambda_action": 1.0,
"action_norm_use_stepwise": false,
"action_norm_default_mode": "z-score",
"action_norm_exception_mode": null,
"action_stats_path": null,
"optimizer_lr": 0.0001,
"optimizer_betas": [
0.9,
0.95
],
"optimizer_eps": 1e-08,
"optimizer_weight_decay": 0.01,
"optimizer_grad_clip_norm": 1.0,
"scheduler_type": "cosine",
"scheduler_warmup_steps": 5000,
"scheduler_warmup_ratio": 0.05,
"scheduler_decay_lr": 1e-05,
"train_num_epochs": null,
"train_max_steps": 300000
},
"output_dir": "outputs/WSA_Large/WSA_Large-wsa_large-multidata-delta-pretrain-2026_05_07_15_59_20",
"job_name": "WSA_Large-wsa_large-multidata-delta-pretrain-2026_05_07_15_59_20",
"resume": false,
"seed": 42,
"num_workers": 14,
"batch_size": 12,
"gradient_accumulation_steps": 1,
"steps": 300000,
"eval_freq": 20000,
"log_freq": 100,
"save_checkpoint": true,
"save_freq": 20000,
"use_policy_training_preset": true,
"optimizer": {
"type": "adamw",
"lr": 0.0001,
"weight_decay": 0.01,
"grad_clip_norm": 1.0,
"betas": [
0.9,
0.95
],
"eps": 1e-08
},
"scheduler": {
"type": "wsa_large_native",
"num_warmup_steps": 5000,
"peak_lr": 0.0001,
"min_lr_ratio": 0.1,
"scheduler_type": "cosine",
"warmup_ratio": 0.05
},
"eval": {
"n_episodes": 50,
"batch_size": 50,
"use_async_envs": false
},
"wandb": {
"enable": false,
"disable_artifact": true,
"project": "WSA_Large",
"entity": null,
"notes": null,
"run_id": null,
"mode": null
},
"checkpoint_path": null,
"rename_map": {}
}