File size: 2,765 Bytes
f35c2dd
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
{
  "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": 0,
    "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-prefix-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
  }
}