Upload folder using huggingface_hub

model_pole_position.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6b3405e455825e092f04054afabbc1504a86dcfd072b1bb182686629aeae2d14
+oid sha256:f9471a08cd9013392d19d9e950acda2e398ef8d3c2bcfe02079066393a337786
 size 2970182

model_pong.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b1af4b741b84430029f7b1e3a499002f46ae0a02c2608435cf05fc99151620ef
+oid sha256:d6c8b9235347bea94e7e5f5f0f225d4c1dbd13a749d5e28920c75c91902ecb11
 size 2435368

model_sonic.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:033c9ab9ac538f6abdca242f6d1b7ebeb02090148b8df76f332700a3d000e0bf
+oid sha256:413d9fcfa15f30c74cdfda5f7d7c9dba8958fe027dfc09de563e6209c78378f5
 size 6180566

predict.py

@@ -1,4 +1,4 @@
-"""Prediction with test-time augmentation (horizontal flip averaging)."""
+"""Prediction with residual dampening and TTA for reduced AR error accumulation."""
 import sys
 import os
 import numpy as np
@@ -16,6 +16,12 @@ GAME_CONFIGS = {
     "pole_position": {"enc_channels": (32, 64, 128), "bottleneck": 192},
 }
 
+DAMPEN_FACTORS = {
+    "pong": 0.90,
+    "sonic": 0.85,
+    "pole_position": 0.95,
+}
+
 
 def detect_game(context_frames: np.ndarray) -> str:
     first_8 = context_frames[:CONTEXT_FRAMES]
@@ -50,6 +56,7 @@ def load_model(model_dir: str):
 def predict_next_frame(models, context_frames: np.ndarray) -> np.ndarray:
     game = detect_game(context_frames)
     model = models[game]
+    dampen = DAMPEN_FACTORS[game]
 
     n = len(context_frames)
     if n < CONTEXT_FRAMES:
@@ -71,19 +78,19 @@ def predict_next_frame(models, context_frames: np.ndarray) -> np.ndarray:
 
         # Original prediction
         residual_orig = model(context_tensor)
-        predicted_orig = torch.clamp(last_tensor + residual_orig, 0, 1)
 
         if game == "pong":
-            # Pong is asymmetric (paddles on left/right), skip TTA
-            predicted = predicted_orig
+            # Pong: no TTA, apply dampening to residual
+            residual = residual_orig * dampen
+            predicted = torch.clamp(last_tensor + residual, 0, 1)
         else:
-            # Flipped prediction for Sonic and PP
+            # Sonic/PP: TTA with horizontal flip, then dampen
             context_flipped = torch.flip(context_tensor, dims=[3])
-            last_flipped = torch.flip(last_tensor, dims=[3])
             residual_flipped = model(context_flipped)
-            predicted_flipped = torch.clamp(last_flipped + residual_flipped, 0, 1)
-            predicted_flipped = torch.flip(predicted_flipped, dims=[3])
-            predicted = (predicted_orig + predicted_flipped) / 2.0
+            residual_flipped = torch.flip(residual_flipped, dims=[3])
+            # Average residuals, then dampen
+            residual = (residual_orig + residual_flipped) / 2.0 * dampen
+            predicted = torch.clamp(last_tensor + residual, 0, 1)
 
     predicted_np = predicted[0].cpu().numpy()
     predicted_np = np.transpose(predicted_np, (1, 2, 0))