| { | |
| "codec": { | |
| "enable_slicing": true, | |
| "enable_tiling": false, | |
| "frame_batch_size": 256, | |
| "kind": "frame_vae", | |
| "latent_channels": 4, | |
| "model_id": "madebyollin/sdxl-vae-fp16-fix", | |
| "revision": "main", | |
| "sample_posterior": false, | |
| "spatial_compression": 8, | |
| "subfolder": "", | |
| "temporal_compression": 1 | |
| }, | |
| "data": { | |
| "batch_size": 128, | |
| "data_path": "", | |
| "dataset_kind": "breakout", | |
| "frame_stride": 1, | |
| "hf_repo_id": "", | |
| "hf_revision": "main", | |
| "latent_cache_path": "", | |
| "latent_prefetch_batches": 8, | |
| "network_clip_prefetch": 2, | |
| "num_workers": 8, | |
| "persistent_workers": true, | |
| "pin_memory": true, | |
| "prefetch_batches": 2, | |
| "sample_shuffle": 64, | |
| "shard_shuffle": 256, | |
| "train_samples": 500000, | |
| "validation_samples": 1024, | |
| "windows_per_video": 4 | |
| }, | |
| "inference": { | |
| "frames": 120, | |
| "steps": 1 | |
| }, | |
| "model": { | |
| "action_dim": 2, | |
| "context_dim": 0, | |
| "depth": 12, | |
| "fine_history_frames": 2, | |
| "fine_memory_patch_size": 2, | |
| "generation_frames": 1, | |
| "hidden_size": 416, | |
| "memory_depth": 3, | |
| "memory_patch_size": 4, | |
| "mlp_ratio": 4.0, | |
| "num_heads": 8, | |
| "qk_norm": true, | |
| "target_patch_size": 2 | |
| }, | |
| "schema_version": 3, | |
| "simulation": { | |
| "max_shapes": 2, | |
| "max_size": 6.0, | |
| "max_speed": 20.0, | |
| "min_shapes": 1, | |
| "min_size": 3.0, | |
| "min_speed": 8.0, | |
| "supersample": 2, | |
| "warmup_frames": 8 | |
| }, | |
| "training": { | |
| "activation_checkpointing": true, | |
| "critic_learning_rate": 0.0001, | |
| "distribution_matching_start_step": 30000, | |
| "distribution_matching_weight": 0.05, | |
| "ema_decay": 0.9999, | |
| "epochs": 100, | |
| "flow_matching_ratio": 0.25, | |
| "full_interval_ratio": 0.25, | |
| "gradient_accumulation_steps": 1, | |
| "history_noise_std": 0.01, | |
| "keep_checkpoints": 5, | |
| "learning_rate": 0.0003, | |
| "log_every": 20, | |
| "max_grad_norm": 1.0, | |
| "min_learning_rate": 1e-05, | |
| "on_policy_max_horizon": 4, | |
| "on_policy_ratio": 0.15, | |
| "on_policy_sampling_steps": 4, | |
| "on_policy_start_step": 15000, | |
| "output_dir": "runs/breakout-action-world-model", | |
| "reward_loss_weight": 0.0, | |
| "save_every": 50, | |
| "seed": 2026, | |
| "teacher_endpoint_ratio": 0.1, | |
| "teacher_integration_steps": 8, | |
| "teacher_logit_mean": 1.0, | |
| "teacher_logit_std": 1.0, | |
| "teacher_steps": 10000, | |
| "teacher_x0_weight": 0.25, | |
| "terminal_loss_weight": 0.0, | |
| "validate_every": 50, | |
| "validation_batches": 1, | |
| "warmup_steps": 500, | |
| "weight_decay": 0.0 | |
| }, | |
| "video": { | |
| "fps": 6.0, | |
| "future_frames": 5, | |
| "height": 128, | |
| "history_frames": 12, | |
| "width": 128 | |
| } | |
| } |