Upload folder using huggingface_hub

predict.py

@@ -135,6 +135,14 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
     last_frame = frames_norm[-1]
     last_frame_t = np.transpose(last_frame, (2, 0, 1))[np.newaxis]
 
+    # Compute per-pixel variance across context frames for motion mask
+    # frames_t: [8, 3, 64, 64]
+    pixel_var = np.var(frames_t, axis=0)  # [3, 64, 64]
+    pixel_var_mean = pixel_var.mean(axis=0)  # [64, 64] - average across channels
+    # Static mask: 1.0 for static pixels (low variance), 0.0 for dynamic
+    var_thresh = 5.0 / (255.0 * 255.0)  # variance in [0,1] scale (5/255^2)
+    static_mask = (pixel_var_mean < var_thresh).astype(np.float32)  # [64, 64]
+
     if game == "pong":
         # Pong: AR+direct ensemble, float32 caching, no TTA
         if ens.direct_cache is not None and n > CONTEXT_FRAMES and ens.cache_step < PRED_FRAMES:
@@ -173,7 +181,10 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
         predicted_np = predicted[0].cpu().numpy()
         ens.direct_cache = []
         for i in range(PRED_FRAMES):
-            frame = np.transpose(predicted_np[i], (1, 2, 0))
+            frame = np.transpose(predicted_np[i], (1, 2, 0))  # [64, 64, 3]
+            # Apply motion mask: static pixels blend 80% context / 20% prediction
+            mask_3d = static_mask[:, :, np.newaxis]  # [64, 64, 1]
+            frame = frame * (1.0 - 0.8 * mask_3d) + last_frame * 0.8 * mask_3d
             frame = (frame * 255).clip(0, 255).astype(np.uint8)
             ens.direct_cache.append(frame)
 
@@ -202,50 +213,31 @@ 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
 
-            # Shift helpers using roll + edge replication
-            def shift_right(t):
-                shifted = torch.roll(t, 1, dims=-1)
-                shifted[:, :, :, 0] = shifted[:, :, :, 1]
-                return shifted
-            def shift_left(t):
-                shifted = torch.roll(t, -1, dims=-1)
-                shifted[:, :, :, -1] = shifted[:, :, :, -2]
-                return shifted
-            def unshift_right(t):
-                return shift_left(t)
-            def unshift_left(t):
-                return shift_right(t)
-
-            # Build augmentation list: (context_aug, last_aug, undo_fn)
-            ctx_sr = shift_right(context_tensor)
-            last_sr = shift_right(last_tensor)
-            ctx_sl = shift_left(context_tensor)
-            last_sl = shift_left(last_tensor)
-
-            augmentations = [
-                (context_tensor, last_tensor, lambda x: x),
-                (context_flipped, last_flipped, lambda x: torch.flip(x, dims=[3])),
-                (ctx_sr, last_sr, unshift_right),
-                (ctx_sl, last_sl, unshift_left),
-            ]
-
-            # Multi-run AR with noise diversity x augmentations
+            # Multi-run AR with noise diversity
             all_ar_runs = []
             for noise_std in [0.0, 1.0/255.0, 2.0/255.0]:
-                for ctx_aug, last_aug, undo_fn in augmentations:
-                    ar_preds_run = []
-                    ctx = ctx_aug.clone()
-                    last_t = last_aug.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)
-                        ar_out = _predict_ar_frame(ens.sonic_ar, ctx_in, last_t)
-                        ar_frame = undo_fn(ar_out)
-                        ar_preds_run.append(ar_frame)
-                        ctx_frames = ctx.reshape(1, CONTEXT_FRAMES, 3, 64, 64)
-                        ctx_frames = torch.cat([ctx_frames[:, 1:], ar_out.unsqueeze(1)], dim=1)
-                        ctx = ctx_frames.reshape(1, -1, 64, 64)
-                        last_t = ar_out
-                    all_ar_runs.append(torch.stack(ar_preds_run, dim=1))
+                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)
 
@@ -258,7 +250,10 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
         predicted_np = predicted[0].cpu().numpy()
         ens.direct_cache = []
         for i in range(PRED_FRAMES):
-            frame = np.transpose(predicted_np[i], (1, 2, 0))
+            frame = np.transpose(predicted_np[i], (1, 2, 0))  # [64, 64, 3]
+            # Apply motion mask: static pixels blend 80% context / 20% prediction
+            mask_3d = static_mask[:, :, np.newaxis]  # [64, 64, 1]
+            frame = frame * (1.0 - 0.8 * mask_3d) + last_frame * 0.8 * mask_3d
             frame = (frame * 255).clip(0, 255).astype(np.uint8)
             ens.direct_cache.append(frame)
 
@@ -290,7 +285,10 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
         predicted_np = predicted[0].cpu().numpy()
         ens.direct_cache = []
         for i in range(PRED_FRAMES):
-            frame = np.transpose(predicted_np[i], (1, 2, 0))
+            frame = np.transpose(predicted_np[i], (1, 2, 0))  # [64, 64, 3]
+            # Apply motion mask: static pixels blend 80% context / 20% prediction
+            mask_3d = static_mask[:, :, np.newaxis]  # [64, 64, 1]
+            frame = frame * (1.0 - 0.8 * mask_3d) + last_frame * 0.8 * mask_3d
             frame = (frame * 255).clip(0, 255).astype(np.uint8)
             ens.direct_cache.append(frame)