Upload folder using huggingface_hub

predict.py

@@ -151,23 +151,18 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
 
             direct_pred = _predict_8frames_direct(ens.pong_direct, context_tensor, last_tensor)
 
-            # Multi-run Pong AR with tiny noise diversity
-            all_pong_ar_runs = []
-            for noise_std in [0.0, 0.25/255.0, 0.5/255.0]:
-                ar_preds_run = []
-                ctx = context_tensor.clone()
-                last_t = last_tensor.clone()
-                for step in range(PRED_FRAMES):
-                    ctx_in = ctx if noise_std == 0 else torch.clamp(ctx + torch.randn_like(ctx) * noise_std, 0, 1)
-                    predicted = _predict_ar_frame(ens.models["pong"], ctx_in, last_t)
-                    ar_preds_run.append(predicted)
-                    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
-                all_pong_ar_runs.append(torch.stack(ar_preds_run, dim=1))
+            ar_preds = []
+            ctx = context_tensor.clone()
+            last_t = last_tensor.clone()
+            for step in range(PRED_FRAMES):
+                predicted = _predict_ar_frame(ens.models["pong"], ctx, last_t)
+                ar_preds.append(predicted)
+                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
 
-            ar_pred = sum(all_pong_ar_runs) / len(all_pong_ar_runs)
+            ar_pred = torch.stack(ar_preds, dim=1)
 
             predicted = torch.zeros_like(direct_pred)
             for step in range(PRED_FRAMES):
@@ -236,10 +231,12 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
             ar_pred = sum(all_ar_runs) / len(all_ar_runs)
 
             predicted = torch.zeros_like(direct_pred)
+            eps = 1e-6
             for step in range(PRED_FRAMES):
-                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]
+                w = 0.65 - (step / (PRED_FRAMES - 1)) * 0.3
+                ar_safe = torch.clamp(ar_pred[:, step], eps, 1.0)
+                direct_safe = torch.clamp(direct_pred[:, step], eps, 1.0)
+                predicted[:, step] = torch.clamp(ar_safe.pow(w) * direct_safe.pow(1.0 - w), 0, 1)
 
         predicted_np = predicted[0].cpu().numpy()
         ens.direct_cache = []