Upload folder using huggingface_hub

model_pole_position.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d648d5626123d26b5967aa52f0580914728b1aaab67905dd5b647786a017337e
+oid sha256:62d218b9859acd4d19cfcfe6b3aa93ae129485a872175632ed32d6441ae9c7f6
 size 1580934

predict.py

@@ -1,4 +1,4 @@
-"""Compact PP ensemble: full Pong AR, Sonic AR+direct ensemble, compact PP direct. Under 16MB."""
+"""Inference tricks: Pong AR caching, Sonic step-dependent blending."""
 import sys
 import os
 import numpy as np
@@ -72,7 +72,7 @@ def load_model(model_dir: str):
     sonic_direct.eval()
     ens.sonic_direct = sonic_direct
 
-    # PP: compact direct 8-frame model (24 outputs, smaller architecture)
+    # PP: compact direct 8-frame model (24 outputs)
     pp = UNet(in_channels=24, out_channels=24,
               enc_channels=(24, 48, 96), bottleneck_channels=128,
               upsample_mode="bilinear").to(DEVICE)
@@ -115,17 +115,46 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
     last_frame_t = np.transpose(last_frame, (2, 0, 1))[np.newaxis]
 
     if game == "pong":
+        # Pong: AR with internal caching (all 8 steps in float32)
+        if ens.direct_cache is not None and n > CONTEXT_FRAMES and ens.cache_step < PRED_FRAMES:
+            result = ens.direct_cache[ens.cache_step]
+            ens.cache_step += 1
+            if ens.cache_step >= PRED_FRAMES:
+                ens.reset_cache()
+            return result
+
+        ens.reset_cache()
+        model = ens.models["pong"]
         with torch.no_grad():
             context_tensor = torch.from_numpy(context).to(DEVICE)
             last_tensor = torch.from_numpy(last_frame_t).to(DEVICE)
-            predicted = _predict_ar_frame(ens.models["pong"], context_tensor, last_tensor)
 
-        predicted_np = predicted[0].cpu().numpy()
-        predicted_np = np.transpose(predicted_np, (1, 2, 0))
-        predicted_np = (predicted_np * 255).clip(0, 255).astype(np.uint8)
-        return predicted_np
+            # Run all 8 AR steps in float32 (no uint8 quantization between steps)
+            preds = []
+            ctx = context_tensor.clone()
+            last_t = last_tensor.clone()
+            for step in range(PRED_FRAMES):
+                predicted = _predict_ar_frame(model, ctx, last_t)
+                preds.append(predicted)
+                # Shift context with float32 prediction (no quantization)
+                ctx_frames = ctx.reshape(1, CONTEXT_FRAMES, 3, 64, 64)
+                ctx_frames = torch.cat([ctx_frames[:, 1:], predicted.unsqueeze(1)], dim=1)
+                ctx = ctx_frames.reshape(1, -1, 64, 64)
+                last_t = predicted
+
+        ens.direct_cache = []
+        for i in range(PRED_FRAMES):
+            frame = preds[i][0].cpu().numpy()
+            frame = np.transpose(frame, (1, 2, 0))
+            frame = (frame * 255).clip(0, 255).astype(np.uint8)
+            ens.direct_cache.append(frame)
+
+        result = ens.direct_cache[ens.cache_step]
+        ens.cache_step += 1
+        return result
 
     elif game == "sonic":
+        # Sonic: step-dependent AR/direct blending with caching
         if ens.direct_cache is not None and n > CONTEXT_FRAMES and ens.cache_step < PRED_FRAMES:
             result = ens.direct_cache[ens.cache_step]
             ens.cache_step += 1
@@ -144,7 +173,7 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
             last_flipped = torch.flip(last_tensor, dims=[3])
             direct_flipped = _predict_8frames_direct(ens.sonic_direct, context_flipped, last_flipped)
             direct_flipped = torch.flip(direct_flipped, dims=[4])
-            direct_pred = (direct_orig + direct_flipped) / 2.0
+            direct_pred = (direct_orig + direct_flipped) / 2.0  # (1, 8, 3, 64, 64)
 
             # AR prediction with TTA
             ar_preds = []
@@ -167,8 +196,15 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
                 ctx_flip = ctx_flip_frames.reshape(1, -1, 64, 64)
                 last_f = ar_flip
 
-            ar_pred = torch.stack(ar_preds, dim=1)
-            predicted = (ar_pred + direct_pred) / 2.0
+            ar_pred = torch.stack(ar_preds, dim=1)  # (1, 8, 3, 64, 64)
+
+            # Step-dependent blending: AR weight goes from 0.7 (step 0) to 0.3 (step 7)
+            # Direct weight goes from 0.3 (step 0) to 0.7 (step 7)
+            predicted = torch.zeros_like(direct_pred)
+            for step in range(PRED_FRAMES):
+                ar_weight = 0.7 - (step / (PRED_FRAMES - 1)) * 0.4  # 0.7 -> 0.3
+                direct_weight = 1.0 - ar_weight  # 0.3 -> 0.7
+                predicted[:, step] = ar_weight * ar_pred[:, step] + direct_weight * direct_pred[:, step]
 
         predicted_np = predicted[0].cpu().numpy()
         ens.direct_cache = []