Upload folder using huggingface_hub

predict.py

@@ -149,26 +149,31 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
             context_tensor = torch.from_numpy(context).to(DEVICE)
             last_tensor = torch.from_numpy(last_frame_t).to(DEVICE)
 
-            direct_pred = _predict_8frames_direct(ens.pong_direct, context_tensor, last_tensor)
+            # Pong AR-only with hflip TTA
+            context_flipped = torch.flip(context_tensor, dims=[3])
+            last_flipped = torch.flip(last_tensor, dims=[3])
 
             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):
-                predicted = _predict_ar_frame(ens.models["pong"], ctx, last_t)
-                ar_preds.append(predicted)
+                ar_orig = _predict_ar_frame(ens.models["pong"], ctx, last_t)
+                ar_flip = _predict_ar_frame(ens.models["pong"], 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:], predicted.unsqueeze(1)], dim=1)
+                ctx_frames = torch.cat([ctx_frames[:, 1:], ar_orig.unsqueeze(1)], dim=1)
                 ctx = ctx_frames.reshape(1, -1, 64, 64)
-                last_t = predicted
+                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)
-
-            predicted = torch.zeros_like(direct_pred)
-            for step in range(PRED_FRAMES):
-                ar_weight = 0.85 - (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 = torch.stack(ar_preds, dim=1)
 
         predicted_np = predicted[0].cpu().numpy()
         ens.direct_cache = []
@@ -203,10 +208,6 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
             direct_pred = (direct_orig + direct_flipped) / 2.0
 
             # Multi-run AR with noise diversity
-            # For no-noise run: feed blended AR+direct as context
-            direct_frames = direct_pred.reshape(1, PRED_FRAMES, 3, 64, 64)
-            direct_frames_flip = torch.flip(direct_frames, dims=[4])
-
             all_ar_runs = []
             for noise_std in [0.0, 1.0/255.0, 2.0/255.0]:
                 ar_preds_run = []
@@ -222,22 +223,14 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
                     ar_flip_back = torch.flip(ar_flip, dims=[3])
                     ar_frame = (ar_orig + ar_flip_back) / 2.0
                     ar_preds_run.append(ar_frame)
-                    # For no-noise run: blend with direct before feeding back
-                    if noise_std == 0:
-                        w = 0.65 - (step / (PRED_FRAMES - 1)) * 0.3
-                        feedback_orig = w * ar_orig + (1.0 - w) * direct_frames[:, step]
-                        feedback_flip = w * ar_flip + (1.0 - w) * direct_frames_flip[:, step]
-                    else:
-                        feedback_orig = ar_orig
-                        feedback_flip = ar_flip
                     ctx_frames = ctx.reshape(1, CONTEXT_FRAMES, 3, 64, 64)
-                    ctx_frames = torch.cat([ctx_frames[:, 1:], feedback_orig.unsqueeze(1)], dim=1)
+                    ctx_frames = torch.cat([ctx_frames[:, 1:], ar_orig.unsqueeze(1)], dim=1)
                     ctx = ctx_frames.reshape(1, -1, 64, 64)
-                    last_t = feedback_orig
+                    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:], feedback_flip.unsqueeze(1)], dim=1)
+                    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 = feedback_flip
+                    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)