Upload folder using huggingface_hub

loss_history.json

@@ -2,232 +2,232 @@
   {
     "epoch": 1,
     "phase": "P1",
-    "loss": 0.093708
+    "loss": 0.152055
   },
   {
     "epoch": 2,
     "phase": "P1",
-    "loss": 0.075409
+    "loss": 0.126681
   },
   {
     "epoch": 3,
     "phase": "P1",
-    "loss": 0.070398
+    "loss": 0.119891
   },
   {
     "epoch": 4,
     "phase": "P1",
-    "loss": 0.066922
+    "loss": 0.114801
   },
   {
     "epoch": 5,
     "phase": "P1",
-    "loss": 0.064051
+    "loss": 0.110611
   },
   {
     "epoch": 6,
     "phase": "P1",
-    "loss": 0.061594
+    "loss": 0.107016
   },
   {
     "epoch": 7,
     "phase": "P1",
-    "loss": 0.058991
+    "loss": 0.103401
   },
   {
     "epoch": 8,
     "phase": "P1",
-    "loss": 0.056665
+    "loss": 0.100012
   },
   {
     "epoch": 9,
     "phase": "P1",
-    "loss": 0.054221
+    "loss": 0.096366
   },
   {
     "epoch": 10,
     "phase": "P1",
-    "loss": 0.052157
+    "loss": 0.09296
   },
   {
     "epoch": 11,
     "phase": "P1",
-    "loss": 0.050054
+    "loss": 0.089986
   },
   {
     "epoch": 12,
     "phase": "P1",
-    "loss": 0.048416
+    "loss": 0.087143
   },
   {
     "epoch": 13,
     "phase": "P1",
-    "loss": 0.047013
+    "loss": 0.08477
   },
   {
     "epoch": 14,
     "phase": "P1",
-    "loss": 0.046003
+    "loss": 0.083114
   },
   {
     "epoch": 15,
     "phase": "P1",
-    "loss": 0.0454
+    "loss": 0.082026
   },
   {
     "epoch": 16,
     "phase": "P2",
-    "loss": 0.071297
+    "loss": 0.122125
   },
   {
     "epoch": 17,
     "phase": "P2",
-    "loss": 0.069845
+    "loss": 0.118517
   },
   {
     "epoch": 18,
     "phase": "P2",
-    "loss": 0.067838
+    "loss": 0.115646
   },
   {
     "epoch": 19,
     "phase": "P2",
-    "loss": 0.102993
+    "loss": 0.170965
   },
   {
     "epoch": 20,
     "phase": "P2",
-    "loss": 0.098403,
-    "val_ssim": 0.8174
+    "loss": 0.163493,
+    "val_ssim": 0.8267
   },
   {
     "epoch": 21,
     "phase": "P2",
-    "loss": 0.095552
+    "loss": 0.159067
   },
   {
     "epoch": 22,
     "phase": "P2",
-    "loss": 0.142291
+    "loss": 0.237583
   },
   {
     "epoch": 23,
     "phase": "P2",
-    "loss": 0.137962
+    "loss": 0.229664
   },
   {
     "epoch": 24,
     "phase": "P2",
-    "loss": 0.133837
+    "loss": 0.221985
   },
   {
     "epoch": 25,
     "phase": "P2",
-    "loss": 0.129812,
-    "val_ssim": 0.854
+    "loss": 0.215313,
+    "val_ssim": 0.8505
   },
   {
     "epoch": 26,
     "phase": "P2",
-    "loss": 0.126053
+    "loss": 0.208722
   },
   {
     "epoch": 27,
     "phase": "P2",
-    "loss": 0.122985
+    "loss": 0.203962
   },
   {
     "epoch": 28,
     "phase": "P2",
-    "loss": 0.120476
+    "loss": 0.198393
   },
   {
     "epoch": 29,
     "phase": "P2",
-    "loss": 0.117592
+    "loss": 0.194795
   },
   {
     "epoch": 30,
     "phase": "P2",
-    "loss": 0.115456,
-    "val_ssim": 0.8644
+    "loss": 0.191285,
+    "val_ssim": 0.8759
   },
   {
     "epoch": 31,
     "phase": "P2",
-    "loss": 0.113231
+    "loss": 0.187651
   },
   {
     "epoch": 32,
     "phase": "P2",
-    "loss": 0.111175
+    "loss": 0.184686
   },
   {
     "epoch": 33,
     "phase": "P2",
-    "loss": 0.108953
+    "loss": 0.180715
   },
   {
     "epoch": 34,
     "phase": "P2",
-    "loss": 0.106131
+    "loss": 0.176762
   },
   {
     "epoch": 35,
     "phase": "P2",
-    "loss": 0.103505,
-    "val_ssim": 0.8744
+    "loss": 0.172307,
+    "val_ssim": 0.8774
   },
   {
     "epoch": 36,
     "phase": "P2",
-    "loss": 0.100435
+    "loss": 0.167519
   },
   {
     "epoch": 37,
     "phase": "P2",
-    "loss": 0.097286
+    "loss": 0.162766
   },
   {
     "epoch": 38,
     "phase": "P2",
-    "loss": 0.094014
+    "loss": 0.157198
   },
   {
     "epoch": 39,
     "phase": "P2",
-    "loss": 0.090802
+    "loss": 0.152165
   },
   {
     "epoch": 40,
     "phase": "P2",
-    "loss": 0.087507,
-    "val_ssim": 0.8852
+    "loss": 0.147043,
+    "val_ssim": 0.886
   },
   {
     "epoch": 41,
     "phase": "P2",
-    "loss": 0.084485
+    "loss": 0.141957
   },
   {
     "epoch": 42,
     "phase": "P2",
-    "loss": 0.081661
+    "loss": 0.137481
   },
   {
     "epoch": 43,
     "phase": "P2",
-    "loss": 0.079401
+    "loss": 0.133861
   },
   {
     "epoch": 44,
     "phase": "P2",
-    "loss": 0.077772
+    "loss": 0.131363
   },
   {
     "epoch": 45,
     "phase": "P2",
-    "loss": 0.076937,
-    "val_ssim": 0.885
+    "loss": 0.129965,
+    "val_ssim": 0.888
   }
 ]

model.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:4774304dae39b918b34dd4ededabc4a793ac7efdeb772a746587fad584ccfe83
-size 9089268
+oid sha256:e930868e7c620774f7f12cd9c2f056032024e50b17bd1405824daa5df80ecb6b
+size 12361376

multiscale_flow_model.py

@@ -0,0 +1,102 @@
+"""Multi-Scale Flow-Warp-Mask U-Net: predicts flow at multiple resolutions."""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class ResConvBlock(nn.Module):
+    def __init__(self, in_ch, out_ch):
+        super().__init__()
+        self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1)
+        self.gn1 = nn.GroupNorm(min(8, out_ch), out_ch)
+        self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1)
+        self.gn2 = nn.GroupNorm(min(8, out_ch), out_ch)
+        self.proj = nn.Conv2d(in_ch, out_ch, 1) if in_ch != out_ch else nn.Identity()
+
+    def forward(self, x):
+        residual = self.proj(x)
+        x = F.silu(self.gn1(self.conv1(x)))
+        x = F.silu(self.gn2(self.conv2(x)))
+        return x + residual
+
+
+class MultiScaleFlowUNet(nn.Module):
+    def __init__(self, in_channels=12, channels=[64, 128, 256]):
+        super().__init__()
+        # Encoder
+        self.encoders = nn.ModuleList()
+        self.pools = nn.ModuleList()
+        prev_ch = in_channels
+        for ch in channels:
+            self.encoders.append(ResConvBlock(prev_ch, ch))
+            self.pools.append(nn.MaxPool2d(2))
+            prev_ch = ch
+
+        # Bottleneck
+        self.bottleneck = ResConvBlock(channels[-1], channels[-1] * 2)
+
+        # Decoder
+        self.upconvs = nn.ModuleList()
+        self.decoders = nn.ModuleList()
+        dec_channels = list(reversed(channels))
+        prev_ch = channels[-1] * 2
+        for ch in dec_channels:
+            self.upconvs.append(nn.ConvTranspose2d(prev_ch, ch, 2, stride=2))
+            self.decoders.append(ResConvBlock(ch * 2, ch))
+            prev_ch = ch
+
+        # Multi-scale flow heads at each decoder level
+        # dec_channels = [256, 128, 64] (coarsest to finest)
+        # Level 0 (coarsest, 8x8): flow refinement
+        # Level 1 (16x16): flow refinement
+        # Level 2 (finest, 64x64): flow refinement + mask + gen_frame
+        self.flow_heads = nn.ModuleList()
+        for ch in dec_channels:
+            head = nn.Conv2d(ch, 2, 1)
+            nn.init.zeros_(head.weight)
+            nn.init.zeros_(head.bias)
+            self.flow_heads.append(head)
+
+        # Mask and generation heads only at finest level (level 2, 64x64)
+        self.mask_head = nn.Conv2d(dec_channels[-1], 1, 1)
+        nn.init.zeros_(self.mask_head.weight)
+        nn.init.zeros_(self.mask_head.bias)
+
+        self.gen_head = nn.Conv2d(dec_channels[-1], 3, 1)
+
+    def forward(self, x):
+        skips = []
+        for enc, pool in zip(self.encoders, self.pools):
+            x = enc(x)
+            skips.append(x)
+            x = pool(x)
+
+        x = self.bottleneck(x)
+
+        flows = []  # flow at each level, from coarsest to finest
+        for i, (upconv, dec, skip) in enumerate(zip(self.upconvs, self.decoders, reversed(skips))):
+            x = upconv(x)
+            x = torch.cat([x, skip], dim=1)
+            x = dec(x)
+
+            # Predict flow refinement at this level
+            flow_refine = self.flow_heads[i](x)
+
+            if i == 0:
+                # Coarsest level: just the flow refinement
+                flow = flow_refine
+            else:
+                # Upsample previous flow and add refinement
+                prev_flow_up = F.interpolate(flows[-1], scale_factor=2, mode='bilinear', align_corners=True)
+                # Scale flow values by 2 since coordinates double
+                prev_flow_up = prev_flow_up * 2
+                flow = prev_flow_up + flow_refine
+
+            flows.append(flow)
+
+        # Final level outputs
+        mask = torch.sigmoid(self.mask_head(x))
+        gen_frame = self.gen_head(x)
+
+        # flows[-1] is the finest (64x64) flow
+        return flows, mask, gen_frame

predict.py

@@ -1,20 +1,20 @@
-"""Prediction interface for Flow-Warp-Mask U-Net v9 with TTA."""
+"""Prediction interface for Multi-Scale Flow-Warp-Mask U-Net v10 with TTA."""
 import sys
 import os
 import numpy as np
 import torch
 
 sys.path.insert(0, "/home/coder/code")
-from flowmask_model import FlowWarpMaskUNet
+from multiscale_flow_model import MultiScaleFlowUNet
 from flownet_model import differentiable_warp
 
 CONTEXT_LEN = 4
-CHANNELS = [48, 96, 192]
+CHANNELS = [56, 112, 224]
 
 
 def load_model(model_dir: str):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model = FlowWarpMaskUNet(in_channels=12, channels=CHANNELS)
+    model = MultiScaleFlowUNet(in_channels=12, channels=CHANNELS)
     model_path = os.path.join(model_dir, "model.pt")
     state_dict = torch.load(model_path, map_location=device, weights_only=True)
     state_dict = {k: v.float() for k, v in state_dict.items()}
@@ -25,7 +25,6 @@ def load_model(model_dir: str):
 
 
 def _prepare_input(context_frames):
-    """Prepare 4-frame context tensor from numpy frames."""
     if len(context_frames) >= CONTEXT_LEN:
         frames = context_frames[-CONTEXT_LEN:]
     else:
@@ -34,19 +33,20 @@ def _prepare_input(context_frames):
         frames = np.concatenate([padding, context_frames], axis=0)
 
     frames_t = torch.from_numpy(frames.astype(np.float32) / 255.0)
-    frames_t = frames_t.permute(0, 3, 1, 2)  # (4, 3, 64, 64)
+    frames_t = frames_t.permute(0, 3, 1, 2)
     return frames_t
 
 
 def _run_model(model, frames_t, device):
-    """Run model on prepared frames, return prediction tensor."""
-    last_frame = frames_t[-1].unsqueeze(0)  # (1, 3, 64, 64)
-    inp = frames_t.reshape(1, -1, 64, 64)   # (1, 12, 64, 64)
+    last_frame = frames_t[-1].unsqueeze(0)
+    inp = frames_t.reshape(1, -1, 64, 64)
 
     inp = inp.to(device)
     last_frame = last_frame.to(device)
 
-    flow, mask, gen_frame = model(inp)
+    flows, mask, gen_frame = model(inp)
+    # Use finest flow (last element)
+    flow = flows[-1]
     warped = differentiable_warp(last_frame, flow)
     pred = mask * warped + (1 - mask) * gen_frame
     pred = torch.clamp(pred, 0, 1)
@@ -64,9 +64,9 @@ def predict_next_frame(model_dict, context_frames: np.ndarray) -> np.ndarray:
         pred1 = _run_model(model, frames_t, device)
 
         # TTA: horizontally flipped prediction
-        frames_flipped = frames_t.flip(-1)  # flip W dimension
+        frames_flipped = frames_t.flip(-1)
         pred2_flipped = _run_model(model, frames_flipped, device)
-        pred2 = pred2_flipped.flip(-1)  # flip back
+        pred2 = pred2_flipped.flip(-1)
 
         # Average
         pred = (pred1 + pred2) / 2.0

train.log

@@ -1,63 +1,64 @@
-[23:37:08] Device: cuda
-[23:37:08] Model parameters: 4,534,230, channels=[48, 96, 192]
-[23:37:08] Phase 1: Single-step (15 epochs)
-[23:37:12]   45108 sequences
-[23:37:54] P1 Epoch 1/15 | loss=0.09371
-[23:38:34] P1 Epoch 2/15 | loss=0.07541
-[23:39:15] P1 Epoch 3/15 | loss=0.07040
-[23:39:56] P1 Epoch 4/15 | loss=0.06692
-[23:40:36] P1 Epoch 5/15 | loss=0.06405
-[23:41:17] P1 Epoch 6/15 | loss=0.06159
-[23:41:58] P1 Epoch 7/15 | loss=0.05899
-[23:42:40] P1 Epoch 8/15 | loss=0.05667
-[23:43:21] P1 Epoch 9/15 | loss=0.05422
-[23:44:01] P1 Epoch 10/15 | loss=0.05216
-[23:44:43] P1 Epoch 11/15 | loss=0.05005
-[23:45:23] P1 Epoch 12/15 | loss=0.04842
-[23:46:03] P1 Epoch 13/15 | loss=0.04701
-[23:46:45] P1 Epoch 14/15 | loss=0.04600
-[23:47:24] P1 Epoch 15/15 | loss=0.04540
-[23:47:24] Phase 2: Graduated AR (30 epochs)
-[23:49:24] P2 Epoch 1/30 (steps=2) | loss=0.07130 lr=0.000500
-[23:51:23] P2 Epoch 2/30 (steps=2) | loss=0.06985 lr=0.000500
-[23:53:18] P2 Epoch 3/30 (steps=2) | loss=0.06784 lr=0.000500
-[23:58:06] P2 Epoch 4/30 (steps=4) | loss=0.10299 lr=0.000500
-[00:02:59] P2 Epoch 5/30 (steps=4) | loss=0.09840 lr=0.000500
-[00:04:11]   Val SSIM=0.8174 | {'pong': 0.7108, 'sonic': 0.8111, 'pole_position': 0.9302}
-[00:04:11]   New best! SSIM=0.8174
-[00:09:08] P2 Epoch 6/30 (steps=4) | loss=0.09555 lr=0.000500
-[00:21:04] P2 Epoch 7/30 (steps=8) | loss=0.14229 lr=0.000500
-[00:32:46] P2 Epoch 8/30 (steps=8) | loss=0.13796 lr=0.000500
-[00:44:48] P2 Epoch 9/30 (steps=8) | loss=0.13384 lr=0.000500
-[00:57:15] P2 Epoch 10/30 (steps=8) | loss=0.12981 lr=0.000500
-[00:58:37]   Val SSIM=0.8540 | {'pong': 0.8022, 'sonic': 0.8237, 'pole_position': 0.936}
-[00:58:37]   New best! SSIM=0.8540
-[01:11:08] P2 Epoch 11/30 (steps=8) | loss=0.12605 lr=0.000500
-[01:23:41] P2 Epoch 12/30 (steps=8) | loss=0.12299 lr=0.000500
-[01:36:24] P2 Epoch 13/30 (steps=8) | loss=0.12048 lr=0.000500
-[01:48:54] P2 Epoch 14/30 (steps=8) | loss=0.11759 lr=0.000500
-[02:01:33] P2 Epoch 15/30 (steps=8) | loss=0.11546 lr=0.000500
-[02:02:55]   Val SSIM=0.8644 | {'pong': 0.829, 'sonic': 0.8264, 'pole_position': 0.9378}
-[02:02:55]   New best! SSIM=0.8644
-[02:15:31] P2 Epoch 16/30 (steps=8) | loss=0.11323 lr=0.000495
-[02:28:01] P2 Epoch 17/30 (steps=8) | loss=0.11117 lr=0.000478
-[02:40:14] P2 Epoch 18/30 (steps=8) | loss=0.10895 lr=0.000452
-[02:52:32] P2 Epoch 19/30 (steps=8) | loss=0.10613 lr=0.000417
-[03:05:05] P2 Epoch 20/30 (steps=8) | loss=0.10350 lr=0.000375
-[03:06:28]   Val SSIM=0.8744 | {'pong': 0.8512, 'sonic': 0.8308, 'pole_position': 0.9413}
-[03:06:28]   New best! SSIM=0.8744
-[03:19:19] P2 Epoch 21/30 (steps=8) | loss=0.10044 lr=0.000327
-[03:31:46] P2 Epoch 22/30 (steps=8) | loss=0.09729 lr=0.000276
-[03:44:25] P2 Epoch 23/30 (steps=8) | loss=0.09401 lr=0.000224
-[03:57:08] P2 Epoch 24/30 (steps=8) | loss=0.09080 lr=0.000173
-[04:09:49] P2 Epoch 25/30 (steps=8) | loss=0.08751 lr=0.000125
-[04:11:04]   Val SSIM=0.8852 | {'pong': 0.8764, 'sonic': 0.8329, 'pole_position': 0.9462}
-[04:11:04]   New best! SSIM=0.8852
-[04:23:43] P2 Epoch 26/30 (steps=8) | loss=0.08449 lr=0.000083
-[04:36:13] P2 Epoch 27/30 (steps=8) | loss=0.08166 lr=0.000048
-[04:48:48] P2 Epoch 28/30 (steps=8) | loss=0.07940 lr=0.000022
-[05:01:33] P2 Epoch 29/30 (steps=8) | loss=0.07777 lr=0.000010
-[05:14:14] P2 Epoch 30/30 (steps=8) | loss=0.07694 lr=0.000010
-[05:15:35]   Val SSIM=0.8850 | {'pong': 0.8783, 'sonic': 0.8292, 'pole_position': 0.9474}
-[05:15:35] Experiment dir: 9.1 MB
-[05:15:35] Training complete. Best val SSIM: 0.8852
+[05:19:55] Device: cuda
+[05:19:55] Model parameters: 6,169,586, channels=[56, 112, 224]
+[05:19:55] Phase 1: Single-step (15 epochs)
+[05:19:59]   45108 sequences
+[05:20:50] P1 Epoch 1/15 | loss=0.15205
+[05:21:41] P1 Epoch 2/15 | loss=0.12668
+[05:22:29] P1 Epoch 3/15 | loss=0.11989
+[05:23:16] P1 Epoch 4/15 | loss=0.11480
+[05:24:08] P1 Epoch 5/15 | loss=0.11061
+[05:24:54] P1 Epoch 6/15 | loss=0.10702
+[05:25:46] P1 Epoch 7/15 | loss=0.10340
+[05:26:37] P1 Epoch 8/15 | loss=0.10001
+[05:27:23] P1 Epoch 9/15 | loss=0.09637
+[05:28:12] P1 Epoch 10/15 | loss=0.09296
+[05:29:02] P1 Epoch 11/15 | loss=0.08999
+[05:29:51] P1 Epoch 12/15 | loss=0.08714
+[05:30:40] P1 Epoch 13/15 | loss=0.08477
+[05:31:30] P1 Epoch 14/15 | loss=0.08311
+[05:32:17] P1 Epoch 15/15 | loss=0.08203
+[05:32:17] Phase 2: Graduated AR (30 epochs)
+[05:34:32] P2 Epoch 1/30 (steps=2) | loss=0.12213 lr=0.000500
+[05:36:49] P2 Epoch 2/30 (steps=2) | loss=0.11852 lr=0.000500
+[05:38:58] P2 Epoch 3/30 (steps=2) | loss=0.11565 lr=0.000500
+[05:44:14] P2 Epoch 4/30 (steps=4) | loss=0.17096 lr=0.000500
+[05:49:31] P2 Epoch 5/30 (steps=4) | loss=0.16349 lr=0.000500
+[05:50:57]   Val SSIM=0.8267 | {'pong': 0.7258, 'sonic': 0.8199, 'pole_position': 0.9343}
+[05:50:57]   New best! SSIM=0.8267
+[05:56:10] P2 Epoch 6/30 (steps=4) | loss=0.15907 lr=0.000500
+[06:10:41] P2 Epoch 7/30 (steps=8) | loss=0.23758 lr=0.000500
+[06:24:53] P2 Epoch 8/30 (steps=8) | loss=0.22966 lr=0.000500
+[06:39:05] P2 Epoch 9/30 (steps=8) | loss=0.22198 lr=0.000500
+[06:53:24] P2 Epoch 10/30 (steps=8) | loss=0.21531 lr=0.000500
+[06:54:54]   Val SSIM=0.8505 | {'pong': 0.7857, 'sonic': 0.8264, 'pole_position': 0.9393}
+[06:54:54]   New best! SSIM=0.8505
+[07:09:06] P2 Epoch 11/30 (steps=8) | loss=0.20872 lr=0.000500
+[07:23:28] P2 Epoch 12/30 (steps=8) | loss=0.20396 lr=0.000500
+[07:37:46] P2 Epoch 13/30 (steps=8) | loss=0.19839 lr=0.000500
+[07:52:00] P2 Epoch 14/30 (steps=8) | loss=0.19479 lr=0.000500
+[08:06:23] P2 Epoch 15/30 (steps=8) | loss=0.19129 lr=0.000500
+[08:07:46]   Val SSIM=0.8759 | {'pong': 0.8609, 'sonic': 0.8246, 'pole_position': 0.9423}
+[08:07:46]   New best! SSIM=0.8759
+[08:22:08] P2 Epoch 16/30 (steps=8) | loss=0.18765 lr=0.000495
+[08:36:25] P2 Epoch 17/30 (steps=8) | loss=0.18469 lr=0.000478
+[08:50:42] P2 Epoch 18/30 (steps=8) | loss=0.18071 lr=0.000452
+[09:04:59] P2 Epoch 19/30 (steps=8) | loss=0.17676 lr=0.000417
+[09:19:13] P2 Epoch 20/30 (steps=8) | loss=0.17231 lr=0.000375
+[09:20:41]   Val SSIM=0.8774 | {'pong': 0.8579, 'sonic': 0.8323, 'pole_position': 0.9419}
+[09:20:41]   New best! SSIM=0.8774
+[09:35:11] P2 Epoch 21/30 (steps=8) | loss=0.16752 lr=0.000327
+[09:49:35] P2 Epoch 22/30 (steps=8) | loss=0.16277 lr=0.000276
+[10:03:57] P2 Epoch 23/30 (steps=8) | loss=0.15720 lr=0.000224
+[10:18:08] P2 Epoch 24/30 (steps=8) | loss=0.15217 lr=0.000173
+[10:32:53] P2 Epoch 25/30 (steps=8) | loss=0.14704 lr=0.000125
+[10:34:17]   Val SSIM=0.8860 | {'pong': 0.876, 'sonic': 0.8357, 'pole_position': 0.9463}
+[10:34:17]   New best! SSIM=0.8860
+[10:49:35] P2 Epoch 26/30 (steps=8) | loss=0.14196 lr=0.000083
+[11:04:55] P2 Epoch 27/30 (steps=8) | loss=0.13748 lr=0.000048
+[11:20:12] P2 Epoch 28/30 (steps=8) | loss=0.13386 lr=0.000022
+[11:35:30] P2 Epoch 29/30 (steps=8) | loss=0.13136 lr=0.000010
+[11:49:54] P2 Epoch 30/30 (steps=8) | loss=0.12997 lr=0.000010
+[11:51:09]   Val SSIM=0.8880 | {'pong': 0.8813, 'sonic': 0.8349, 'pole_position': 0.9479}
+[11:51:09]   New best! SSIM=0.8880
+[11:51:09] Experiment dir: 12.4 MB
+[11:51:09] Training complete. Best val SSIM: 0.8880