Upload folder using huggingface_hub

predict.py

@@ -202,27 +202,33 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
             direct_flipped = torch.flip(direct_flipped, dims=[4])
             direct_pred = (direct_orig + direct_flipped) / 2.0
 
-            ar_preds = []
-            ctx = context_tensor.clone()
-            ctx_flip = context_flipped.clone()
-            last_t = last_tensor.clone()
-            last_f = last_flipped.clone()
-            for step in range(PRED_FRAMES):
-                ar_orig = _predict_ar_frame(ens.sonic_ar, ctx, last_t)
-                ar_flip = _predict_ar_frame(ens.sonic_ar, ctx_flip, last_f)
-                ar_flip_back = torch.flip(ar_flip, dims=[3])
-                ar_frame = (ar_orig + ar_flip_back) / 2.0
-                ar_preds.append(ar_frame)
-                ctx_frames = ctx.reshape(1, CONTEXT_FRAMES, 3, 64, 64)
-                ctx_frames = torch.cat([ctx_frames[:, 1:], ar_orig.unsqueeze(1)], dim=1)
-                ctx = ctx_frames.reshape(1, -1, 64, 64)
-                last_t = ar_orig
-                ctx_flip_frames = ctx_flip.reshape(1, CONTEXT_FRAMES, 3, 64, 64)
-                ctx_flip_frames = torch.cat([ctx_flip_frames[:, 1:], ar_flip.unsqueeze(1)], dim=1)
-                ctx_flip = ctx_flip_frames.reshape(1, -1, 64, 64)
-                last_f = ar_flip
-
-            ar_pred = torch.stack(ar_preds, dim=1)
+            # Multi-run AR with noise diversity
+            all_ar_runs = []
+            for noise_std in [0.0, 1.0/255.0, 2.0/255.0]:
+                ar_preds_run = []
+                ctx = context_tensor.clone()
+                ctx_flip = context_flipped.clone()
+                last_t = last_tensor.clone()
+                last_f = last_flipped.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)
+                    ctx_flip_in = ctx_flip if noise_std == 0 else torch.clamp(ctx_flip + torch.randn_like(ctx_flip) * noise_std, 0, 1)
+                    ar_orig = _predict_ar_frame(ens.sonic_ar, ctx_in, last_t)
+                    ar_flip = _predict_ar_frame(ens.sonic_ar, ctx_flip_in, last_f)
+                    ar_flip_back = torch.flip(ar_flip, dims=[3])
+                    ar_frame = (ar_orig + ar_flip_back) / 2.0
+                    ar_preds_run.append(ar_frame)
+                    ctx_frames = ctx.reshape(1, CONTEXT_FRAMES, 3, 64, 64)
+                    ctx_frames = torch.cat([ctx_frames[:, 1:], ar_orig.unsqueeze(1)], dim=1)
+                    ctx = ctx_frames.reshape(1, -1, 64, 64)
+                    last_t = ar_orig
+                    ctx_flip_frames = ctx_flip.reshape(1, CONTEXT_FRAMES, 3, 64, 64)
+                    ctx_flip_frames = torch.cat([ctx_flip_frames[:, 1:], ar_flip.unsqueeze(1)], dim=1)
+                    ctx_flip = ctx_flip_frames.reshape(1, -1, 64, 64)
+                    last_f = ar_flip
+                all_ar_runs.append(torch.stack(ar_preds_run, dim=1))
+
+            ar_pred = sum(all_ar_runs) / len(all_ar_runs)
 
             predicted = torch.zeros_like(direct_pred)
             for step in range(PRED_FRAMES):
@@ -260,7 +266,7 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
             last_flipped = torch.flip(last_tensor, dims=[3])
             predicted_flipped = _predict_8frames_direct(ens.models["pole_position"], context_flipped, last_flipped)
             predicted_flipped = torch.flip(predicted_flipped, dims=[4])
-            predicted = 0.7 * predicted_orig + 0.3 * predicted_flipped
+            predicted = (predicted_orig + predicted_flipped) / 2.0
 
         predicted_np = predicted[0].cpu().numpy()
         ens.direct_cache = []