Hugging Face's logo

ojaffe
/
world-model

Model card Files
xet
 Community
ojaffe commited on
Commit
9bd62e2
·
verified ·
1 Parent(s): 13bbd0c

Upload folder using huggingface_hub

Browse files
Files changed (4) hide show
  1. __pycache__/predict.cpython-311.pyc +0 -0
  2. model_pole_position.pt +1 -1
  3. predict.py +37 -58
  4. train.log +31 -52
__pycache__/predict.cpython-311.pyc CHANGED
Binary files a/__pycache__/predict.cpython-311.pyc and b/__pycache__/predict.cpython-311.pyc differ
 
model_pole_position.pt CHANGED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:8e0affcef8e533a29037751e27948a3eb0f2fda2792ce2b3dfc876cadb09e281
3
  size 2971526
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:623ad759c36dc1680ed19ae7b6ab72e0803aa12aaa2df333fecf6ad4c90ad220
3
  size 2971526
predict.py CHANGED
@@ -1,4 +1,4 @@
1
- """4-way TTA: original + hflip + vflip + both flips for Sonic and PP."""
2
  import sys
3
  import os
4
  import numpy as np
@@ -182,7 +182,7 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
182
  return result
183
 
184
  elif game == "sonic":
185
- # Sonic: AR(fp16)+direct(int8) with step blending and 4-way TTA
186
  if ens.direct_cache is not None and n > CONTEXT_FRAMES and ens.cache_step < PRED_FRAMES:
187
  result = ens.direct_cache[ens.cache_step]
188
  ens.cache_step += 1
@@ -195,49 +195,34 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
195
  context_tensor = torch.from_numpy(context).to(DEVICE)
196
  last_tensor = torch.from_numpy(last_frame_t).to(DEVICE)
197
 
198
- # 4 augmentations: original, hflip, vflip, both
199
- ctx_hf = torch.flip(context_tensor, dims=[3])
200
- last_hf = torch.flip(last_tensor, dims=[3])
201
- ctx_vf = torch.flip(context_tensor, dims=[2])
202
- last_vf = torch.flip(last_tensor, dims=[2])
203
- ctx_hv = torch.flip(context_tensor, dims=[2, 3])
204
- last_hv = torch.flip(last_tensor, dims=[2, 3])
205
-
206
- # Direct: 4-way TTA
207
- d0 = _predict_8frames_direct(ens.sonic_direct, context_tensor, last_tensor)
208
- d1 = torch.flip(_predict_8frames_direct(ens.sonic_direct, ctx_hf, last_hf), dims=[4])
209
- d2 = torch.flip(_predict_8frames_direct(ens.sonic_direct, ctx_vf, last_vf), dims=[3])
210
- d3 = torch.flip(_predict_8frames_direct(ens.sonic_direct, ctx_hv, last_hv), dims=[3, 4])
211
- direct_pred = (d0 + d1 + d2 + d3) / 4.0
212
-
213
- # AR: 4-way TTA
214
- augs = [
215
- (context_tensor.clone(), last_tensor.clone()),
216
- (ctx_hf.clone(), last_hf.clone()),
217
- (ctx_vf.clone(), last_vf.clone()),
218
- (ctx_hv.clone(), last_hv.clone()),
219
- ]
220
- # flip dims for undoing: hflip=[3], vflip=[2], both=[2,3]
221
- all_ar_preds = []
222
- for aug_ctx, aug_last in augs:
223
- ar_chain = []
224
- c = aug_ctx
225
- l = aug_last
226
- for step in range(PRED_FRAMES):
227
- pred = _predict_ar_frame(ens.sonic_ar, c, l)
228
- ar_chain.append(pred)
229
- cf = c.reshape(1, CONTEXT_FRAMES, 3, 64, 64)
230
- cf = torch.cat([cf[:, 1:], pred.unsqueeze(1)], dim=1)
231
- c = cf.reshape(1, -1, 64, 64)
232
- l = pred
233
- all_ar_preds.append(torch.stack(ar_chain, dim=1))
234
-
235
- # Undo flips
236
- ar0 = all_ar_preds[0]
237
- ar1 = torch.flip(all_ar_preds[1], dims=[4])
238
- ar2 = torch.flip(all_ar_preds[2], dims=[3])
239
- ar3 = torch.flip(all_ar_preds[3], dims=[3, 4])
240
- ar_pred = (ar0 + ar1 + ar2 + ar3) / 4.0
241
 
242
  predicted = torch.zeros_like(direct_pred)
243
  for step in range(PRED_FRAMES):
@@ -257,7 +242,7 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
257
  return result
258
 
259
  else:
260
- # PP: direct with 4-way TTA and caching
261
  if ens.direct_cache is not None and n > CONTEXT_FRAMES and ens.cache_step < PRED_FRAMES:
262
  result = ens.direct_cache[ens.cache_step]
263
  ens.cache_step += 1
@@ -270,18 +255,12 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
270
  context_tensor = torch.from_numpy(context).to(DEVICE)
271
  last_tensor = torch.from_numpy(last_frame_t).to(DEVICE)
272
 
273
- ctx_hf = torch.flip(context_tensor, dims=[3])
274
- last_hf = torch.flip(last_tensor, dims=[3])
275
- ctx_vf = torch.flip(context_tensor, dims=[2])
276
- last_vf = torch.flip(last_tensor, dims=[2])
277
- ctx_hv = torch.flip(context_tensor, dims=[2, 3])
278
- last_hv = torch.flip(last_tensor, dims=[2, 3])
279
-
280
- p0 = _predict_8frames_direct(ens.models["pole_position"], context_tensor, last_tensor)
281
- p1 = torch.flip(_predict_8frames_direct(ens.models["pole_position"], ctx_hf, last_hf), dims=[4])
282
- p2 = torch.flip(_predict_8frames_direct(ens.models["pole_position"], ctx_vf, last_vf), dims=[3])
283
- p3 = torch.flip(_predict_8frames_direct(ens.models["pole_position"], ctx_hv, last_hv), dims=[3, 4])
284
- predicted = (p0 + p1 + p2 + p3) / 4.0
285
 
286
  predicted_np = predicted[0].cpu().numpy()
287
  ens.direct_cache = []
 
1
+ """Full PP swap: Pong direct int8, full PP model, Sonic AR fp16 + direct int8."""
2
  import sys
3
  import os
4
  import numpy as np
 
182
  return result
183
 
184
  elif game == "sonic":
185
+ # Sonic: AR(fp16)+direct(int8) with step blending and TTA
186
  if ens.direct_cache is not None and n > CONTEXT_FRAMES and ens.cache_step < PRED_FRAMES:
187
  result = ens.direct_cache[ens.cache_step]
188
  ens.cache_step += 1
 
195
  context_tensor = torch.from_numpy(context).to(DEVICE)
196
  last_tensor = torch.from_numpy(last_frame_t).to(DEVICE)
197
 
198
+ direct_orig = _predict_8frames_direct(ens.sonic_direct, context_tensor, last_tensor)
199
+ context_flipped = torch.flip(context_tensor, dims=[3])
200
+ last_flipped = torch.flip(last_tensor, dims=[3])
201
+ direct_flipped = _predict_8frames_direct(ens.sonic_direct, context_flipped, last_flipped)
202
+ direct_flipped = torch.flip(direct_flipped, dims=[4])
203
+ direct_pred = (direct_orig + direct_flipped) / 2.0
204
+
205
+ ar_preds = []
206
+ ctx = context_tensor.clone()
207
+ ctx_flip = context_flipped.clone()
208
+ last_t = last_tensor.clone()
209
+ last_f = last_flipped.clone()
210
+ for step in range(PRED_FRAMES):
211
+ ar_orig = _predict_ar_frame(ens.sonic_ar, ctx, last_t)
212
+ ar_flip = _predict_ar_frame(ens.sonic_ar, ctx_flip, last_f)
213
+ ar_flip_back = torch.flip(ar_flip, dims=[3])
214
+ ar_frame = (ar_orig + ar_flip_back) / 2.0
215
+ ar_preds.append(ar_frame)
216
+ ctx_frames = ctx.reshape(1, CONTEXT_FRAMES, 3, 64, 64)
217
+ ctx_frames = torch.cat([ctx_frames[:, 1:], ar_orig.unsqueeze(1)], dim=1)
218
+ ctx = ctx_frames.reshape(1, -1, 64, 64)
219
+ last_t = ar_orig
220
+ ctx_flip_frames = ctx_flip.reshape(1, CONTEXT_FRAMES, 3, 64, 64)
221
+ ctx_flip_frames = torch.cat([ctx_flip_frames[:, 1:], ar_flip.unsqueeze(1)], dim=1)
222
+ ctx_flip = ctx_flip_frames.reshape(1, -1, 64, 64)
223
+ last_f = ar_flip
224
+
225
+ ar_pred = torch.stack(ar_preds, dim=1)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
226
 
227
  predicted = torch.zeros_like(direct_pred)
228
  for step in range(PRED_FRAMES):
 
242
  return result
243
 
244
  else:
245
+ # PP: direct with TTA and caching
246
  if ens.direct_cache is not None and n > CONTEXT_FRAMES and ens.cache_step < PRED_FRAMES:
247
  result = ens.direct_cache[ens.cache_step]
248
  ens.cache_step += 1
 
255
  context_tensor = torch.from_numpy(context).to(DEVICE)
256
  last_tensor = torch.from_numpy(last_frame_t).to(DEVICE)
257
 
258
+ predicted_orig = _predict_8frames_direct(ens.models["pole_position"], context_tensor, last_tensor)
259
+ context_flipped = torch.flip(context_tensor, dims=[3])
260
+ last_flipped = torch.flip(last_tensor, dims=[3])
261
+ predicted_flipped = _predict_8frames_direct(ens.models["pole_position"], context_flipped, last_flipped)
262
+ predicted_flipped = torch.flip(predicted_flipped, dims=[4])
263
+ predicted = (predicted_orig + predicted_flipped) / 2.0
 
 
 
 
 
 
264
 
265
  predicted_np = predicted[0].cpu().numpy()
266
  ens.direct_cache = []
train.log CHANGED
@@ -1,52 +1,31 @@
1
- [2026-04-12 05:59:24] Starting Pong direct 8-frame training for 2026-04-12-133000-pong-direct-v2
2
- [2026-04-12 05:59:24] Device: cuda
3
- [2026-04-12 05:59:25] Pong direct: 1,199,224 params (2.3 MB fp16)
4
- [2026-04-12 05:59:25] pong train: 8194 seqs (len=16)
5
- [2026-04-12 05:59:26] pong val: 964 seqs (len=16)
6
- [2026-04-12 05:59:47] E1/150 | T:0.232601(S:0.6492) V:0.206121(S:0.6850) LR:3.00e-04
7
- [2026-04-12 06:00:07] E2/150 | T:0.173923(S:0.7357) V:0.175883(S:0.7317) LR:3.00e-04
8
- [2026-04-12 06:00:28] E3/150 | T:0.138439(S:0.7891) V:0.157651(S:0.7599) LR:3.00e-04
9
- [2026-04-12 06:00:50] E4/150 | T:0.117363(S:0.8211) V:0.140205(S:0.7864) LR:2.99e-04
10
- [2026-04-12 06:01:09] E5/150 | T:0.103357(S:0.8425) V:0.136901(S:0.7912) LR:2.99e-04
11
- [2026-04-12 06:01:30] E6/150 | T:0.093950(S:0.8569) V:0.126921(S:0.8063) LR:2.99e-04
12
- [2026-04-12 06:01:51] E7/150 | T:0.086157(S:0.8691) V:0.121485(S:0.8144) LR:2.98e-04
13
- [2026-04-12 06:02:53] E10/150 | T:0.071323(S:0.8921) V:0.119263(S:0.8183) LR:2.97e-04
14
- [2026-04-12 06:03:13] E11/150 | T:0.067633(S:0.8979) V:0.113872(S:0.8260) LR:2.96e-04
15
- [2026-04-12 06:03:34] E12/150 | T:0.064540(S:0.9027) V:0.112058(S:0.8288) LR:2.95e-04
16
- [2026-04-12 06:04:38] E15/150 | T:0.057899(S:0.9133) V:0.109369(S:0.8330) LR:2.93e-04
17
- [2026-04-12 06:05:00] E16/150 | T:0.054467(S:0.9185) V:0.108353(S:0.8339) LR:2.92e-04
18
- [2026-04-12 06:05:20] E17/150 | T:0.053424(S:0.9202) V:0.106683(S:0.8371) LR:2.91e-04
19
- [2026-04-12 06:06:19] E20/150 | T:0.049003(S:0.9271) V:0.106042(S:0.8383) LR:2.87e-04
20
- [2026-04-12 06:07:00] E22/150 | T:0.046678(S:0.9307) V:0.102919(S:0.8428) LR:2.84e-04
21
- [2026-04-12 06:08:19] E26/150 | T:0.042844(S:0.9367) V:0.102117(S:0.8444) LR:2.78e-04
22
- [2026-04-12 06:08:39] E27/150 | T:0.042220(S:0.9377) V:0.102082(S:0.8444) LR:2.77e-04
23
- [2026-04-12 06:09:38] E30/150 | T:0.039701(S:0.9416) V:0.101520(S:0.8453) LR:2.71e-04
24
- [2026-04-12 06:09:58] E31/150 | T:0.038995(S:0.9427) V:0.100438(S:0.8467) LR:2.70e-04
25
- [2026-04-12 06:10:16] E32/150 | T:0.039062(S:0.9426) V:0.100245(S:0.8473) LR:2.68e-04
26
- [2026-04-12 06:10:36] E33/150 | T:0.038450(S:0.9436) V:0.099826(S:0.8476) LR:2.66e-04
27
- [2026-04-12 06:11:13] E35/150 | T:0.037610(S:0.9449) V:0.099788(S:0.8479) LR:2.62e-04
28
- [2026-04-12 06:11:33] E36/150 | T:0.036675(S:0.9463) V:0.098966(S:0.8496) LR:2.59e-04
29
- [2026-04-12 06:12:52] E40/150 | T:0.034903(S:0.9490) V:0.099786(S:0.8485) LR:2.51e-04
30
- [2026-04-12 06:13:50] E43/150 | T:0.033734(S:0.9509) V:0.098349(S:0.8506) LR:2.43e-04
31
- [2026-04-12 06:14:10] E44/150 | T:0.033505(S:0.9512) V:0.098003(S:0.8510) LR:2.41e-04
32
- [2026-04-12 06:14:50] E46/150 | T:0.033140(S:0.9518) V:0.097767(S:0.8514) LR:2.36e-04
33
- [2026-04-12 06:15:49] E49/150 | T:0.032213(S:0.9533) V:0.097305(S:0.8522) LR:2.28e-04
34
- [2026-04-12 06:16:09] E50/150 | T:0.031634(S:0.9541) V:0.097047(S:0.8524) LR:2.25e-04
35
- [2026-04-12 06:17:28] E54/150 | T:0.030719(S:0.9556) V:0.096493(S:0.8531) LR:2.14e-04
36
- [2026-04-12 06:18:27] E57/150 | T:0.030218(S:0.9563) V:0.095355(S:0.8549) LR:2.06e-04
37
- [2026-04-12 06:18:47] E58/150 | T:0.029896(S:0.9568) V:0.094789(S:0.8558) LR:2.03e-04
38
- [2026-04-12 06:19:27] E60/150 | T:0.029534(S:0.9574) V:0.096519(S:0.8532) LR:1.97e-04
39
- [2026-04-12 06:22:43] E70/150 | T:0.027838(S:0.9600) V:0.095992(S:0.8544) LR:1.66e-04
40
- [2026-04-12 06:25:55] E80/150 | T:0.026637(S:0.9619) V:0.095754(S:0.8547) LR:1.35e-04
41
- [2026-04-12 06:28:12] E87/150 | T:0.025836(S:0.9631) V:0.094571(S:0.8563) LR:1.13e-04
42
- [2026-04-12 06:28:52] E89/150 | T:0.025732(S:0.9633) V:0.094485(S:0.8564) LR:1.07e-04
43
- [2026-04-12 06:29:11] E90/150 | T:0.025655(S:0.9634) V:0.094835(S:0.8561) LR:1.04e-04
44
- [2026-04-12 06:32:25] E100/150 | T:0.024915(S:0.9645) V:0.094847(S:0.8561) LR:7.57e-05
45
- [2026-04-12 06:35:38] E110/150 | T:0.024432(S:0.9653) V:0.094407(S:0.8567) LR:5.05e-05
46
- [2026-04-12 06:38:57] E120/150 | T:0.024090(S:0.9658) V:0.094761(S:0.8562) LR:2.96e-05
47
- [2026-04-12 06:42:12] E130/150 | T:0.023891(S:0.9661) V:0.094652(S:0.8564) LR:1.39e-05
48
- [2026-04-12 06:45:25] E140/150 | T:0.023785(S:0.9663) V:0.094651(S:0.8564) LR:4.27e-06
49
- [2026-04-12 06:48:40] E150/150 | T:0.023744(S:0.9663) V:0.094663(S:0.8564) LR:1.00e-06
50
- [2026-04-12 06:48:40] Done. Best val loss: 0.094407
51
- [2026-04-12 06:48:40] Model size: 2.3 MB
52
- [2026-04-12 06:48:40] Training complete!
 
1
+ [2026-04-12 07:14:27] Starting PP SSIM-only training for 2026-04-12-153000-pp-ssim-only
2
+ [2026-04-12 07:14:27] Device: cuda
3
+ [2026-04-12 07:14:27] PP: 1,465,848 params (2.8 MB fp16)
4
+ [2026-04-12 07:14:28] PP train: 4097 seqs (len=16)
5
+ [2026-04-12 07:14:28] PP val: 482 seqs (len=16)
6
+ [2026-04-12 07:14:38] E1/100 | T:0.089821(S:0.9041) V:0.069040(S:0.9263) LR:3.00e-04
7
+ [2026-04-12 07:14:46] E2/100 | T:0.074465(S:0.9199) V:0.061649(S:0.9337) LR:3.00e-04
8
+ [2026-04-12 07:14:55] E3/100 | T:0.069567(S:0.9251) V:0.060543(S:0.9351) LR:2.99e-04
9
+ [2026-04-12 07:15:04] E4/100 | T:0.066393(S:0.9284) V:0.058608(S:0.9372) LR:2.99e-04
10
+ [2026-04-12 07:15:12] E5/100 | T:0.063999(S:0.9309) V:0.057957(S:0.9378) LR:2.98e-04
11
+ [2026-04-12 07:15:21] E6/100 | T:0.061633(S:0.9334) V:0.056070(S:0.9396) LR:2.97e-04
12
+ [2026-04-12 07:15:31] E7/100 | T:0.060136(S:0.9350) V:0.051418(S:0.9447) LR:2.96e-04
13
+ [2026-04-12 07:15:50] E9/100 | T:0.057306(S:0.9380) V:0.050541(S:0.9454) LR:2.94e-04
14
+ [2026-04-12 07:16:00] E10/100 | T:0.055858(S:0.9396) V:0.052572(S:0.9431) LR:2.93e-04
15
+ [2026-04-12 07:16:09] E11/100 | T:0.055021(S:0.9404) V:0.049028(S:0.9470) LR:2.91e-04
16
+ [2026-04-12 07:17:14] E18/100 | T:0.048082(S:0.9480) V:0.047964(S:0.9479) LR:2.77e-04
17
+ [2026-04-12 07:17:30] E20/100 | T:0.047000(S:0.9492) V:0.049137(S:0.9465) LR:2.71e-04
18
+ [2026-04-12 07:18:25] E26/100 | T:0.043209(S:0.9533) V:0.047514(S:0.9482) LR:2.53e-04
19
+ [2026-04-12 07:18:45] E28/100 | T:0.042278(S:0.9543) V:0.046412(S:0.9495) LR:2.46e-04
20
+ [2026-04-12 07:19:04] E30/100 | T:0.041417(S:0.9553) V:0.047650(S:0.9481) LR:2.38e-04
21
+ [2026-04-12 07:20:37] E40/100 | T:0.037882(S:0.9592) V:0.048536(S:0.9469) LR:1.97e-04
22
+ [2026-04-12 07:21:15] E44/100 | T:0.036738(S:0.9604) V:0.046170(S:0.9495) LR:1.79e-04
23
+ [2026-04-12 07:22:10] E50/100 | T:0.035150(S:0.9621) V:0.048340(S:0.9468) LR:1.50e-04
24
+ [2026-04-12 07:23:41] E60/100 | T:0.033106(S:0.9644) V:0.048292(S:0.9467) LR:1.04e-04
25
+ [2026-04-12 07:25:13] E70/100 | T:0.031615(S:0.9660) V:0.047850(S:0.9472) LR:6.26e-05
26
+ [2026-04-12 07:26:46] E80/100 | T:0.030609(S:0.9671) V:0.047661(S:0.9474) LR:2.96e-05
27
+ [2026-04-12 07:28:21] E90/100 | T:0.030099(S:0.9676) V:0.048131(S:0.9468) LR:8.32e-06
28
+ [2026-04-12 07:29:58] E100/100 | T:0.029986(S:0.9678) V:0.048119(S:0.9469) LR:1.00e-06
29
+ [2026-04-12 07:29:58] Done. Best val loss: 0.046170
30
+ [2026-04-12 07:29:58] Model size: 2.8 MB
31
+ [2026-04-12 07:29:58] Training complete!