Upload folder using huggingface_hub

predict.py

@@ -164,9 +164,11 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
 
             ar_pred = torch.stack(ar_preds, dim=1)
 
+            # Custom per-step AR weights (nonlinear schedule)
+            pong_ar_weights = [0.90, 0.90, 0.90, 0.90, 0.70, 0.70, 0.45, 0.45]
             predicted = torch.zeros_like(direct_pred)
             for step in range(PRED_FRAMES):
-                ar_weight = 0.85 - (step / (PRED_FRAMES - 1)) * 0.3
+                ar_weight = pong_ar_weights[step]
                 direct_weight = 1.0 - ar_weight
                 predicted[:, step] = ar_weight * ar_pred[:, step] + direct_weight * direct_pred[:, step]
 
@@ -230,15 +232,11 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
 
             ar_pred = sum(all_ar_runs) / len(all_ar_runs)
 
-            # Blend in residual space
-            last_ctx = last_tensor.unsqueeze(1).expand_as(direct_pred)
             predicted = torch.zeros_like(direct_pred)
             for step in range(PRED_FRAMES):
-                w = 0.65 - (step / (PRED_FRAMES - 1)) * 0.3
-                ar_residual = ar_pred[:, step] - last_ctx[:, step]
-                direct_residual = direct_pred[:, step] - last_ctx[:, step]
-                blended_residual = w * ar_residual + (1.0 - w) * direct_residual
-                predicted[:, step] = torch.clamp(last_ctx[:, step] + blended_residual, 0, 1)
+                ar_weight = 0.65 - (step / (PRED_FRAMES - 1)) * 0.3
+                direct_weight = 1.0 - ar_weight
+                predicted[:, step] = ar_weight * ar_pred[:, step] + direct_weight * direct_pred[:, step]
 
         predicted_np = predicted[0].cpu().numpy()
         ens.direct_cache = []