CoachWorld / meta.json
JEdward's picture
Upload CoachWorld pretrain release
6de848e verified
Raw History Blame Contribute Delete
9.23 kB
{
"step": 55000,
"world_size": 5,
"format": "safetensors+optim_shards",
"safetensors_exported": true,
"lr_scheduler_saved": true,
"model_config": "WorldModelConfig(backbone='wan2.2', checkpoint='/mnt/cfs/wmcprg/RoboCoach/Wan2.2-TI2V-5B', vae_checkpoint='/mnt/cfs/wmcprg/RoboCoach/Wan2.2-TI2V-5B/Wan2.2_VAE.pth', action_dim=10, text_dim=4096, text_conditioning_enabled=True, num_inference_steps=50, guidance_scale=1.0, guidance_scale_min=1.0, guidance_scale_max=1.0, pred_frames=3, history_length=5, flow_matching_shift=3.0, mixed_precision='bf16', fsdp_enabled=True, fsdp_shard_size=5, gradient_checkpoint=True, learning_rate=5e-06, action_lr_mult=1.0, max_train_steps=100000, batch_size=2, grad_accum_steps=4, condition_mask_enabled=True, gt_condition_replacement=False, future_only_action_conditioning=True, action_schema='arm_slot', max_arm_slots=2, action_max_time_steps=512, action_kv_enabled=True, action_v_init_scale=0.2, action_num_domains=8, action_domain_prompt_tokens=0, domain_aware_action_projection_enabled=True, domain_aware_group_projection_enabled=False, dense_action_film_enabled=False, action_adaln_modulation_enabled=True, action_adaln_scale=0.15, action_dense_rank=128, action_dense_relative_chunks=False, action_control_scale=0.05, action_dense_layernorm_enabled=True, eef_projection_kv_enabled=True, eef_spatial_conditioning_enabled=True, eef_spatial_sigma_px=16.0, eef_spatial_depth_scale=2.0, full_window=True, ctrlworld_full_window=False, action_condition_timestep_rate='raw', history_selector='sparse', history_offsets=[], history_dilations=[1, 2], history_collapse_prob=0.15, history_eval_dilation=2, history_rollout_dilation=2, multi_view_position_mode='global', camera_id_embedding_enabled=False, prope_enabled=True, prope_zero_init=True, prope_spatial_stride=16, num_cameras=1, latent_height_per_view=32, latent_width=48, causal_dim=128, causal_ffn_dim=256, causal_num_heads=4, causal_num_layers=2, causal_freq_dim=64, causal_text_tokens=16, causal_num_frame_per_block=2, causal_local_attn_size=-1, causal_sink_size=0, causal_prope_enabled=False, causal_prope_zero_init=True, causal_preserve_zero_head=False, causal_self_forcing_enabled=False, causal_self_forcing_start_step=0, causal_self_forcing_history_frames=1, causal_self_forcing_sigma=0.5, causal_self_forcing_prob=1.0, causal_proprio_head_enabled=False, causal_proprio_mid_dim=512)",
"training_config": "WMTrainingConfig(loss_scale=1.0, warmup_steps=500, lr_schedule='cosine', lr_min_ratio=0.1, gradient_clip=1.0, cfg_dropout_text=0.0, cfg_dropout_action=0.0, cfg_dropout_text_start=0.0, cfg_dropout_text_end=0.2, cfg_dropout_text_decay_steps=15000, cfg_dropout_action_start=0.0, cfg_dropout_action_end=0.0, cfg_dropout_action_decay_steps=0, history_noise_max=0.3, history_clean_prob=0.25, aux_action_weight=0.0, aux_action_mask_weight=0.05, counterfactual_action_weight=0.0, counterfactual_action_margin=0.02, counterfactual_action_mode='roll', counterfactual_action_interval=1, action_warmup_steps=0, action_param_max_abs=100.0, condition_mask_enabled=True, gt_condition_replacement=False, future_only_action_conditioning=True, action_schema='arm_slot', max_arm_slots=2, action_max_time_steps=512, action_kv_enabled=True, action_v_init_scale=0.2, action_num_domains=8, action_domain_prompt_tokens=0, domain_aware_action_projection_enabled=True, domain_aware_group_projection_enabled=False, dense_action_film_enabled=False, action_adaln_modulation_enabled=True, action_adaln_scale=0.15, action_dense_rank=128, action_dense_relative_chunks=False, action_control_scale=0.05, action_dense_layernorm_enabled=True, eef_projection_kv_enabled=True, eef_spatial_conditioning_enabled=True, eef_spatial_sigma_px=16.0, eef_spatial_depth_scale=2.0, full_window=True, ctrlworld_full_window=False, action_source='video_latent_condition', action_condition_timestep_rate='raw', history_selector='sparse', history_offsets=[], history_dilations=[1, 2], history_collapse_prob=0.15, history_eval_dilation=2, history_rollout_dilation=2, generated_history_enabled=True, generated_history_probability=0.25, generated_history_denoise_steps=2, generated_history_blend=0.5, generated_history_start_step=20000, generated_history_seed=20260723, generated_history_dataset_keys=['droid_ext1', 'droid_ext2', 'robocoin_real', 'vifailback_real', 'giagai_real'], multi_view_position_mode='global', camera_id_embedding_enabled=False, prope_enabled=True, prope_zero_init=True, prope_spatial_stride=16, output_dir='/mnt/cfs/wmcprg/RoboCoach/checkpoints/sa_wm_pretrain_v6_512x768', log_interval=25, save_interval=2500, save_final_checkpoint=True, val_interval=1250, diagnose_interval=1250, val_metrics_enabled=True, val_metrics_interval=5000, val_metrics_num_samples=30, val_metrics_inference_steps=25, val_metrics_seed=20260714, val_metrics_lpips_net='alex', val_metrics_lpips_batch_size=4, val_metrics_save_videos=True, val_metrics_fps=5.0, val_metrics_closed_loop_chunks=4, wandb_project='sa-wm-training', wandb_run_name='sa_wm_pretrain_v6_512x768', require_val_split=True, text_encoder_device='auto', text_conditioning_enabled=True, text_embedding_cache_limit=32, dataset_format='video_latent', dataset_roots=['/mnt/cfs/wmcprg/RoboCoach/video_latent_training_ready_v1_512x768_5hz_canonical/droid_ext1_5hz_canonical_camera', '/mnt/cfs/wmcprg/RoboCoach/video_latent_training_ready_v1_512x768_5hz_canonical/droid_ext2_5hz_canonical_camera', '/mnt/cfs/wmcprg/RoboCoach/video_latent_training_ready_v1_512x768_5hz_canonical/robomind_franka_left_5hz_canonical_camera', '/mnt/cfs/wmcprg/RoboCoach/video_latent_training_ready_v1_512x768_5hz_canonical/robomind_franka_right_5hz_canonical_camera', '/mnt/cfs/wmcprg/RoboCoach/video_latent_training_ready_v1_512x768_5hz_canonical/robomind_franka_top_5hz_canonical_camera', '/mnt/cfs/wmcprg/RoboCoach/video_latent_training_ready_v1_512x768_5hz_canonical/libero_official_success_5hz_canonical_camera', '/mnt/cfs/wmcprg/RoboCoach/video_latent_training_ready_v1_512x768_5hz_canonical/libero_molmo_success_5hz_canonical_camera', '/mnt/cfs/wmcprg/RoboCoach/video_latent_training_ready_v1_512x768_5hz_canonical/libero_molmo_failure_5hz_canonical_camera', '/mnt/cfs/wmcprg/RoboCoach/video_latent_training_ready_v1_512x768_5hz_canonical/libero_pi05_failure_5hz_canonical_camera', '/mnt/cfs/wmcprg/RoboCoach/video_latent_training_ready_v1_512x768_5hz_canonical/robotwin_clean_success_5hz_canonical_camera', '/mnt/cfs/wmcprg/RoboCoach/video_latent_training_ready_v1_512x768_5hz_canonical/robotwin_molmo_failure_5hz_canonical_camera', '/mnt/cfs/wmcprg/RoboCoach/video_latent_training_ready_v1_512x768_5hz_canonical/robotwin_pi05_failure_5hz_canonical_camera', '/mnt/cfs/wmcprg/RoboCoach/video_latent_training_ready_v1_512x768_5hz_canonical/robocoin_admitted_5hz_canonical_camera', '/mnt/cfs/wmcprg/RoboCoach/video_latent_training_ready_v1_512x768_5hz_canonical/vifailback_admitted_5hz_canonical_camera', '/mnt/cfs/wmcprg/RoboCoach/video_latent_training_ready_v1_512x768_5hz_canonical/giagai_admitted_5hz_canonical_camera'], meta_info_roots=[], dataset_probs=[0.335, 0.335, 0.006666666666666667, 0.006666666666666667, 0.006666666666666667, 0.00125, 0.00125, 0.00125, 0.00125, 0.0025, 0.00125, 0.00125, 0.145, 0.01, 0.145], val_dataset_probs=[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], mixture_sampler_seed=20260724, mixture_samples_per_epoch=0, video_latent_video_keys=[], video_latent_condition_view='arm_slot_eef_pose', video_latent_sample_stride=8, video_latent_camera_conditioning=True, video_latent_camera_sidecar_roots=[], video_latent_camera_intrinsics_fallback='error', video_latent_camera_extrinsics_convention='camera_from_world', video_latent_require_production_contract=True, video_latent_required_target_fps=5.0, video_latent_required_target_image_hw=[512, 768], video_latent_eef_projection_roots=[], video_latent_eef_projection_mode='online', video_latent_eef_projection_load_heatmap=False, video_latent_eef_projection_heatmap_size='', rgb_skip=1, down_sample=1, annotation_name='annotation', allow_empty_text=True, num_cameras=1, latent_height_per_view=32, latent_width=48, val_video_num=15, causal_proprio_loss_weight=0.0, causal_proprio_mask_weight=0.05, causal_proprio_future_only=True)",
"sampling_state": {
"kind": "resumable_distributed_mixture_sampler",
"seed": 20260724,
"num_replicas": 5,
"rank": 0,
"batch_size": 2,
"global_samples_per_epoch": 592760,
"samples_per_rank": 118552,
"micro_steps_per_epoch": 59276,
"root_sizes": [
173621,
176083,
8139,
8139,
8073,
9379,
5839,
28080,
28080,
4192,
10593,
10593,
51852,
6600,
63488
],
"probabilities": [
0.335,
0.335,
0.006666666666666667,
0.006666666666666667,
0.006666666666666667,
0.00125,
0.00125,
0.00125,
0.00125,
0.0025,
0.00125,
0.00125,
0.145,
0.01,
0.145
],
"root_draws_per_epoch": [
198575,
198574,
3952,
3952,
3952,
741,
741,
741,
741,
1482,
741,
741,
85950,
5927,
85950
],
"start_epoch": 0,
"start_offset": 90000
}
}