Upload folder using huggingface_hub

predict.py

@@ -135,9 +135,6 @@ 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]
 
-    # Per-channel means from last 4 context frames for color correction
-    ctx_channel_means = frames_norm[-4:].mean(axis=(0, 1, 2))  # [3]
-
     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:
@@ -206,6 +203,10 @@ 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 = []
@@ -221,14 +222,22 @@ 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:], ar_orig.unsqueeze(1)], dim=1)
+                    ctx_frames = torch.cat([ctx_frames[:, 1:], feedback_orig.unsqueeze(1)], dim=1)
                     ctx = ctx_frames.reshape(1, -1, 64, 64)
-                    last_t = ar_orig
+                    last_t = feedback_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_frames = torch.cat([ctx_flip_frames[:, 1:], feedback_flip.unsqueeze(1)], dim=1)
                     ctx_flip = ctx_flip_frames.reshape(1, -1, 64, 64)
-                    last_f = ar_flip
+                    last_f = feedback_flip
                 all_ar_runs.append(torch.stack(ar_preds_run, dim=1))
 
             ar_pred = sum(all_ar_runs) / len(all_ar_runs)
@@ -242,14 +251,7 @@ 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))  # [64,64,3]
-            # Color correction: nudge channel means toward context
-            pred_means = frame.mean(axis=(0, 1))  # [3]
-            diff = ctx_channel_means - pred_means
-            thresh = 3.0 / 255.0
-            correction = np.where(np.abs(diff) > thresh, 0.1 * diff, 0.0)
-            frame = frame + correction[np.newaxis, np.newaxis, :]
-            frame = frame.clip(0, 1)
+            frame = np.transpose(predicted_np[i], (1, 2, 0))
             frame = (frame * 255).clip(0, 255).astype(np.uint8)
             ens.direct_cache.append(frame)
 
@@ -282,13 +284,6 @@ def predict_next_frame(ens, context_frames: np.ndarray) -> np.ndarray:
         ens.direct_cache = []
         for i in range(PRED_FRAMES):
             frame = np.transpose(predicted_np[i], (1, 2, 0))
-            # Color correction
-            pred_means = frame.mean(axis=(0, 1))
-            diff = ctx_channel_means - pred_means
-            thresh = 3.0 / 255.0
-            correction = np.where(np.abs(diff) > thresh, 0.1 * diff, 0.0)
-            frame = frame + correction[np.newaxis, np.newaxis, :]
-            frame = frame.clip(0, 1)
             frame = (frame * 255).clip(0, 255).astype(np.uint8)
             ens.direct_cache.append(frame)