Add model and results

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+results/rollout.gif filter=lfs diff=lfs merge=lfs -text

breakout.py

@@ -0,0 +1,542 @@
+# /// script
+# requires-python = ">=3.12"
+# dependencies = [
+#     "ale-py",
+#     "gradio",
+#     "gymnasium",
+#     "huggingface_hub",
+#     "jinja2",
+#     "numpy",
+#     "pillow",
+#     "safetensors",
+#     "torch",
+# ]
+# ///
+"""Train a diffusion world model of Atari Breakout and play it.
+
+After DIAMOND (https://arxiv.org/abs/2405.12399): a small U-Net denoises the next 64x64 frame
+from the last four frames and actions, with EDM preconditioning and a few Euler steps at
+inference. Fed its own frames, it generates a video that the actions steer.
+"""
+
+# Torch modules only expose forward and keep one attribute per layer.
+# pylint: disable=too-few-public-methods,too-many-instance-attributes
+
+import argparse
+import json
+import math
+import os
+import shutil
+import time
+from dataclasses import asdict, dataclass
+from functools import partial
+from pathlib import Path
+
+import ale_py
+import gradio as gr
+import gymnasium as gym
+import numpy as np
+import torch
+import torch.nn.functional as F
+from huggingface_hub import EvalResult, HfApi, ModelCard, ModelCardData, snapshot_download
+from huggingface_hub.errors import RepositoryNotFoundError
+from PIL import Image
+from safetensors.torch import load_file, save_file
+from torch import nn
+
+REPO = "jgalego/breakout-world-model"
+GAME = "ALE/Breakout-v5"
+ACTIONS = ("NOOP", "FIRE", "RIGHT", "LEFT")
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+HERE = Path(__file__).parent
+SIZE = 64
+EVAL_SEED = 1_000_000  # evaluation episodes never share a seed with training episodes
+SIGMA_DATA = 0.5
+HORIZONS = (1, 5, 15)  # frames ahead; at frameskip 4, 15 frames is one second of play
+GIF_FRAMES = 60
+
+gym.register_envs(ale_py)
+
+
+@dataclass
+class Config:
+    """Architecture and sampler."""
+
+    context: int = 4
+    channels: tuple = (64, 128, 256)
+    steps: int = 3  # Euler steps per frame
+    sigma_min: float = 2e-3
+    sigma_max: float = 5.0
+    rho: float = 7.0
+
+
+# Data
+
+
+def collect(frames, seed):
+    """Frames (n, 3, 64, 64) uint8, actions (n,) and episode ids (n,) from a random policy.
+
+    The policy holds each random action for 1 to 6 steps. actions[i] turns frame i into frame
+    i + 1; the last frame of an episode gets a NOOP that no window uses.
+    """
+    env = gym.make(GAME, frameskip=4, repeat_action_probability=0.0)
+    rng = np.random.default_rng(seed)
+    obs, episode, action, hold = env.reset(seed=seed)[0], 0, 0, 0
+    out = {"frames": [], "actions": [], "episodes": []}
+    while len(out["frames"]) < frames:
+        if hold == 0:
+            action, hold = int(rng.choice(4, p=[0.2, 0.1, 0.35, 0.35])), int(rng.integers(1, 7))
+        hold -= 1
+        for key, value in zip(out, (resize(obs), action, episode)):
+            out[key].append(value)
+        obs, _, terminated, truncated, _ = env.step(action)
+        if terminated or truncated:
+            for key, value in zip(out, (resize(obs), 0, episode)):
+                out[key].append(value)
+            obs, episode = env.reset()[0], episode + 1
+    env.close()
+    return (
+        torch.from_numpy(np.stack(out["frames"])),
+        torch.tensor(out["actions"]),
+        torch.tensor(out["episodes"]),
+    )
+
+
+def resize(obs):
+    """A 210x160 RGB frame as a 3x64x64 uint8 array."""
+    return np.asarray(Image.fromarray(obs).resize((SIZE, SIZE), Image.BOX)).transpose(2, 0, 1)
+
+
+def starts(episodes, length):
+    """Indices i where frames i to i + length - 1 belong to one episode."""
+    return torch.nonzero(episodes[: -length + 1] == episodes[length - 1 :])[:, 0]
+
+
+def to_float(frames):
+    """uint8 frames to [-1, 1]."""
+    return frames.float() / 127.5 - 1
+
+
+def window(frames, actions, index, context):
+    """Context frames (b, c, 3, 64, 64), their actions (b, c) and the next frame (b, 3, 64, 64)."""
+    rows = index[:, None] + torch.arange(context + 1, device=index.device)
+    clip = to_float(frames[rows])
+    return clip[:, :-1], actions[rows[:, :-1]], clip[:, -1]
+
+
+# Model
+
+
+class Block(nn.Module):
+    """Residual conv block with scale and shift from the conditioning vector."""
+
+    def __init__(self, inputs, outputs, cond):
+        super().__init__()
+        self.norm1, self.norm2 = nn.GroupNorm(32, inputs), nn.GroupNorm(32, outputs)
+        self.conv1 = nn.Conv2d(inputs, outputs, 3, padding=1)
+        self.conv2 = nn.Conv2d(outputs, outputs, 3, padding=1)
+        self.film = nn.Linear(cond, 2 * outputs)
+        self.skip = nn.Conv2d(inputs, outputs, 1) if inputs != outputs else nn.Identity()
+
+    def forward(self, x, cond):
+        """x (b, inputs, h, w) to (b, outputs, h, w)."""
+        hidden = self.conv1(F.silu(self.norm1(x)))
+        scale, shift = self.film(cond)[:, :, None, None].chunk(2, dim=1)
+        hidden = self.conv2(F.silu(self.norm2(hidden) * (1 + scale) + shift))
+        return self.skip(x) + hidden
+
+
+class Denoiser(nn.Module):
+    """U-Net over the noisy next frame stacked with the context frames."""
+
+    def __init__(self, config, cond=256):
+        super().__init__()
+        self.config = config
+        width = list(config.channels)
+        self.action = nn.Embedding(len(ACTIONS), 64)
+        self.embed = nn.Sequential(
+            nn.Linear(64 * (config.context + 1), cond), nn.SiLU(), nn.Linear(cond, cond)
+        )
+        self.stem = nn.Conv2d(3 * (config.context + 1), width[0], 3, padding=1)
+        self.down = nn.ModuleList(Block(a, b, cond) for a, b in zip(width[:1] + width, width))
+        self.middle = Block(width[-1], width[-1], cond)
+        self.up = nn.ModuleList(
+            Block(2 * width[i], width[max(i - 1, 0)], cond) for i in range(len(width))
+        )
+        self.out = nn.Sequential(
+            nn.GroupNorm(32, width[0]), nn.SiLU(), nn.Conv2d(width[0], 3, 3, padding=1)
+        )
+
+    def forward(self, noisy, noise, context, actions):
+        """Raw network output for the preconditioned input; `noise` is log(sigma) / 4."""
+        freqs = torch.exp(torch.linspace(0, math.log(1000), 32, device=noise.device))
+        angles = noise[:, None] * freqs
+        cond = torch.cat([angles.sin(), angles.cos(), self.action(actions).flatten(1)], dim=1)
+        cond = self.embed(cond)
+        x = self.stem(torch.cat([noisy, context.flatten(1, 2)], dim=1))
+        skips = []
+        for level, block in enumerate(self.down):
+            x = block(x, cond)
+            skips.append(x)
+            if level < len(self.down) - 1:
+                x = F.avg_pool2d(x, 2)
+        x = self.middle(x, cond)
+        for block, skip in zip(reversed(self.up), reversed(skips)):
+            if x.shape[-1] != skip.shape[-1]:
+                x = F.interpolate(x, scale_factor=2, mode="nearest")
+            x = block(torch.cat([x, skip], dim=1), cond)
+        return self.out(x)
+
+    def denoise(self, noisy, sigma, context, actions):
+        """EDM denoiser: the clean frame estimate at noise level sigma (b,)."""
+        c_skip, c_out, c_in = precondition(sigma)
+        return c_skip * noisy + c_out * self(c_in * noisy, sigma.log() / 4, context, actions)
+
+
+def precondition(sigma):
+    """EDM skip, output and input scales for noise levels sigma (b,)."""
+    sigma = sigma[:, None, None, None]
+    total = sigma**2 + SIGMA_DATA**2
+    return SIGMA_DATA**2 / total, sigma * SIGMA_DATA / total.sqrt(), 1 / total.sqrt()
+
+
+@torch.no_grad()
+def sample(model, context, actions):
+    """The next frame in [-1, 1] by Euler steps on the Karras noise schedule."""
+    config = model.config
+    ramp = torch.linspace(0, 1, config.steps, device=context.device)
+    low, high = config.sigma_min ** (1 / config.rho), config.sigma_max ** (1 / config.rho)
+    sigmas = (high + ramp * (low - high)) ** config.rho
+    x = torch.randn_like(context[:, -1]) * sigmas[0]
+    for i, sigma in enumerate(sigmas):
+        denoised = model.denoise(x, sigma.expand(len(x)), context, actions)
+        following = sigmas[i + 1] if i + 1 < len(sigmas) else 0.0
+        x = denoised + (x - denoised) * following / sigma
+    return x.clamp(-1, 1)
+
+
+@torch.no_grad()
+def rollout(model, context, actions):
+    """Generate frames from context (b, c, 3, 64, 64) with actions (b, c - 1 + steps).
+
+    Returns (b, steps, 3, 64, 64) in [-1, 1].
+    """
+    size = model.config.context
+    frames = []
+    for step in range(actions.shape[1] - size + 1):
+        with torch.autocast(DEVICE, torch.bfloat16, enabled=DEVICE == "cuda"):
+            frame = sample(model, context, actions[:, step : step + size]).float()
+        frames.append(frame)
+        context = torch.cat([context[:, 1:], frame[:, None]], dim=1)
+    return torch.stack(frames, dim=1)
+
+
+# Evaluation
+
+
+def psnr(predicted, truth, mask):
+    """PSNR in dB per step for frames (b, t, 3, h, w) in [-1, 1], pooled over masked pixels."""
+    error = ((predicted - truth) / 2).pow(2) * mask
+    mse = error.sum((0, 2, 3, 4)) / mask.sum((0, 2, 3, 4)).clamp_min(1)
+    return 10 * torch.log10(1 / mse.clamp_min(1e-10))
+
+
+def evaluate_model(model, frames, actions, episodes, count):
+    """PSNR of rollouts against real frames, with a repeat-the-last-frame baseline.
+
+    Most of a Breakout frame is static, so PSNR is also given over changed pixels: those where
+    the real or the predicted frame differs from the last context frame.
+    """
+    size, horizon = model.config.context, max(HORIZONS)
+    valid = starts(episodes, size + horizon)
+    index = valid[torch.linspace(0, len(valid) - 1, min(count, len(valid))).long()]
+    rows = index[:, None] + torch.arange(size + horizon)
+    clip = to_float(frames[rows].to(DEVICE))
+    last = clip[:, size - 1 : size].expand_as(clip[:, size:])
+    guesses = {
+        "model": rollout(model, clip[:, :size], actions[rows[:, : size - 1 + horizon]].to(DEVICE)),
+        "repeat": last,
+    }
+    return {
+        "n": len(index),
+        **{name: score(guess, clip[:, size:], last) for name, guess in guesses.items()},
+    }
+
+
+def score(guess, truth, last):
+    """PSNR over all pixels and over changed pixels at each horizon."""
+    changed = ((truth - last).abs() > 0.1) | ((guess - last).abs() > 0.1)
+    parts = {"all": torch.ones_like(truth), "changed": changed}
+    return {
+        part: {f"{h}": round(psnr(guess, truth, mask)[h - 1].item(), 2) for h in HORIZONS}
+        for part, mask in parts.items()
+    }
+
+
+def gif(model, frames, actions, episodes, path):
+    """Real frames next to a rollout from the same start and actions, 4x upscaled."""
+    size = model.config.context
+    index = starts(episodes, size + GIF_FRAMES)[0]
+    clip = to_float(frames[index : index + size + GIF_FRAMES].to(DEVICE))
+    moves = actions[index : index + size - 1 + GIF_FRAMES].to(DEVICE)
+    predicted = rollout(model, clip[None, :size], moves[None])[0]
+    pairs = torch.cat([clip[size:], predicted], dim=-1)
+    images = [picture(pair, 4) for pair in pairs]
+    images[0].save(path, save_all=True, append_images=images[1:], duration=66, loop=0)
+
+
+def picture(frame, scale):
+    """A [-1, 1] frame (3, h, w) as an upscaled PIL image."""
+    pixels = ((frame + 1) * 127.5).round().clamp(0, 255).byte().permute(1, 2, 0).cpu().numpy()
+    image = Image.fromarray(pixels)
+    return image.resize((image.width * scale, image.height * scale), Image.NEAREST)
+
+
+# Training
+
+
+def edm_loss(model, context, actions, clean, generator):
+    """EDM loss on the next frame, with noise levels from a log-normal."""
+    sigma = (torch.randn(len(clean), device=DEVICE, generator=generator) * 1.2 - 0.4).exp()
+    c_skip, c_out, c_in = precondition(sigma)
+    noisy = clean + sigma[:, None, None, None] * torch.randn_like(clean)
+    with torch.autocast(DEVICE, torch.bfloat16, enabled=DEVICE == "cuda"):
+        output = model(c_in * noisy, sigma.log() / 4, context, actions)
+    return (output.float() - (clean - c_skip * noisy) / c_out).pow(2).mean()
+
+
+def save_state(folder, state):
+    """Write resumable training state to folder/state.pt, replacing it in one step."""
+    folder = Path(folder)
+    folder.mkdir(parents=True, exist_ok=True)
+    torch.save(state, folder / "state.pt.tmp")
+    os.replace(folder / "state.pt.tmp", folder / "state.pt")
+
+
+def training_state(step, model, optimizer, generator):
+    """Everything needed to continue training exactly where it stopped."""
+    return {
+        "step": step,
+        "model": model.state_dict(),
+        "optimizer": optimizer.state_dict(),
+        "generator": generator.get_state(),
+        "rng": torch.get_rng_state(),
+        "cuda_rng": torch.cuda.get_rng_state() if DEVICE == "cuda" else None,
+    }
+
+
+def resume(folder, model, optimizer, generator):
+    """Restore training state from folder/state.pt if there is one; return its step, else 0."""
+    if not folder or not (Path(folder) / "state.pt").exists():
+        return 0
+    state = torch.load(Path(folder) / "state.pt", map_location="cpu")
+    model.load_state_dict(state["model"])
+    optimizer.load_state_dict(state["optimizer"])
+    generator.set_state(state["generator"])
+    torch.set_rng_state(state["rng"])
+    if DEVICE == "cuda":
+        torch.cuda.set_rng_state(state["cuda_rng"])
+    print(f"resumed at step {state['step']}", flush=True)
+    return state["step"]
+
+
+def fit(model, frames, actions, episodes, args):
+    """AdamW on random windows of the collected play, resuming from --checkpoint if set."""
+    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-2)
+    valid = starts(episodes, model.config.context + 1).to(DEVICE)
+    generator = torch.Generator(DEVICE).manual_seed(args.seed)
+    first = resume(args.checkpoint, model, optimizer, generator)
+    start = time.time()
+    for step in range(first + 1, args.steps + 1):
+        pick = torch.randint(0, len(valid), (args.batch_size,), device=DEVICE, generator=generator)
+        batch = window(frames, actions, valid[pick], model.config.context)
+        loss = edm_loss(model, *batch, generator)
+        optimizer.zero_grad()
+        loss.backward()
+        nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        optimizer.step()
+        if step % args.log_every == 0:
+            log = {
+                "step": step,
+                "loss": round(loss.item(), 4),
+                "steps_s": round((step - first) / (time.time() - start), 1),
+            }
+            print(json.dumps(log), flush=True)
+        if args.checkpoint and (step % args.save_every == 0 or step == args.steps):
+            save_state(args.checkpoint, training_state(step, model, optimizer, generator))
+
+
+def train(args):
+    """Collect random play, fit, score on held-out episodes, save and optionally push."""
+    torch.manual_seed(args.seed)
+    frames, actions, episodes = collect(args.frames, args.seed)
+    config = Config(context=args.context, channels=tuple(args.channels), steps=args.sample_steps)
+    model = Denoiser(config).to(DEVICE)
+    training = {
+        **{key: value for key, value in vars(args).items() if key not in ("run", "push")},
+        "episodes": int(episodes[-1]) + 1,
+        "params": sum(param.numel() for param in model.parameters()),
+        "device": torch.cuda.get_device_name() if DEVICE == "cuda" else "cpu",
+    }
+    print(json.dumps(training, indent=1), flush=True)
+
+    start = time.time()
+    fit(model, *(tensor.to(DEVICE) for tensor in (frames, actions, episodes)), args)
+    training["minutes"] = round((time.time() - start) / 60, 1)
+
+    model.eval()
+    test = collect(args.eval_frames, EVAL_SEED)
+    result = evaluate_model(model, *test, args.eval_windows)
+    print(json.dumps(result, indent=1))
+    folder = Path(args.output)
+    save(model, folder, training, result)
+    gif(model, *test, folder / "results" / "rollout.gif")
+    if args.push:
+        push(folder, args.repo)
+
+
+def save(model, folder, training, result):
+    """Write weights, config, code and results to a folder."""
+    (folder / "results").mkdir(parents=True, exist_ok=True)
+    state = {key: value.cpu().contiguous() for key, value in model.state_dict().items()}
+    save_file(state, folder / "model.safetensors")
+    files = {
+        "config.json": asdict(model.config),
+        "results/train.json": training,
+        "results/eval.json": result,
+    }
+    for name, value in files.items():
+        (folder / name).write_text(json.dumps(value, indent=1) + "\n", encoding="utf-8")
+    if Path(__file__).resolve() != (folder / "breakout.py").resolve():
+        shutil.copy(__file__, folder / "breakout.py")
+
+
+def push(folder, repo):
+    """Upload a folder to a private model repo, creating it if needed."""
+    api = HfApi()
+    api.create_repo(repo, private=True, exist_ok=True)
+    api.upload_folder(folder_path=folder, repo_id=repo, commit_message="Add model and results")
+
+
+def load(name=REPO, device=DEVICE):
+    """Load a trained model from a Hub repo id or a local folder."""
+    folder = Path(name) if Path(name).is_dir() else Path(snapshot_download(name))
+    config = Config(**json.loads((folder / "config.json").read_text(encoding="utf-8")))
+    model = Denoiser(config)
+    model.load_state_dict(load_file(folder / "model.safetensors"))
+    return model.to(device).eval()
+
+
+# Play
+
+
+def play(args):
+    """A browser page with one button per action; every press generates the next frame."""
+    model = load(args.model)
+    size = model.config.context
+    frames, actions, _ = collect(size, args.seed)
+    first = (to_float(frames[:size]).to(DEVICE), actions[: size - 1].to(DEVICE))
+
+    def act(action, state):
+        context, past = state
+        moves = torch.cat([past, torch.tensor([action], device=DEVICE)])
+        frame = rollout(model, context[None], moves[None])[0, 0]
+        state = (torch.cat([context[1:], frame[None]]), moves[1:])
+        return picture(frame, 6), state
+
+    with gr.Blocks(title="Breakout world model") as demo:
+        state = gr.State(first)
+        screen = gr.Image(picture(first[0][-1], 6), label="Generated frame")
+        with gr.Row():
+            for action, name in enumerate(ACTIONS):
+                button = gr.Button(name)
+                button.click(partial(act, action), state, [screen, state], api_name=name.lower())
+        start = picture(first[0][-1], 6)
+        gr.Button("Reset").click(lambda: (start, first), None, [screen, state], api_name="reset")
+    demo.launch()
+
+
+def card(args):
+    """Render card.jinja into card/README.md with the config and results stored in the repo."""
+    files = {}
+    try:
+        folder = Path(snapshot_download(args.repo, allow_patterns=["config.json", "results/*"]))
+        files = {
+            path.relative_to(folder).as_posix(): json.loads(path.read_text(encoding="utf-8"))
+            for path in folder.glob("**/*.json")
+        }
+    except RepositoryNotFoundError:
+        pass
+    evaluation = files.get("results/eval.json")
+    data = ModelCardData(
+        model_name=args.repo.split("/")[1],
+        license="apache-2.0",
+        library_name="pytorch",
+        tags=["world-model", "diffusion", "atari", "breakout", "video-generation"],
+        eval_results=[
+            EvalResult(
+                task_type="other",
+                task_name="Next-frame rollout",
+                dataset_type="ale/breakout",
+                dataset_name="Breakout, held-out random play",
+                metric_type=f"psnr_changed_{h}",
+                metric_name=f"PSNR on changed pixels at {h} frames",
+                metric_value=evaluation["model"]["changed"][f"{h}"],
+            )
+            for h in HORIZONS
+        ]
+        if evaluation
+        else None,
+    )
+    rendered = ModelCard.from_template(
+        data,
+        template_path=HERE / "card.jinja",
+        repo=args.repo,
+        config=files.get("config.json"),
+        train=files.get("results/train.json"),
+        evaluation=evaluation,
+        horizons=[f"{h}" for h in HORIZONS],
+    )
+    (HERE / "card").mkdir(exist_ok=True)
+    rendered.save(HERE / "card" / "README.md")
+
+
+def main():
+    """Parse arguments and run a command."""
+    parser = argparse.ArgumentParser(description=__doc__)
+    commands = parser.add_subparsers(dest="command", required=True)
+
+    train_parser = commands.add_parser("train")
+    train_parser.add_argument("--frames", type=int, default=200_000)
+    train_parser.add_argument("--eval-frames", type=int, default=5_000)
+    train_parser.add_argument("--eval-windows", type=int, default=256)
+    train_parser.add_argument("--context", type=int, default=4)
+    train_parser.add_argument("--channels", type=int, nargs="+", default=[64, 128, 256])
+    train_parser.add_argument("--sample-steps", type=int, default=3)
+    train_parser.add_argument("--steps", type=int, default=50_000)
+    train_parser.add_argument("--batch-size", type=int, default=64)
+    train_parser.add_argument("--lr", type=float, default=1e-4)
+    train_parser.add_argument("--seed", type=int, default=0)
+    train_parser.add_argument("--log-every", type=int, default=500)
+    train_parser.add_argument("--checkpoint", help="folder for resumable state, e.g. a bucket")
+    train_parser.add_argument("--save-every", type=int, default=2500)
+    train_parser.add_argument("--output", default="out")
+    train_parser.add_argument("--repo", default=REPO)
+    train_parser.add_argument("--push", action="store_true")
+    train_parser.set_defaults(run=train)
+
+    play_parser = commands.add_parser("play")
+    play_parser.add_argument("--model", default=REPO)
+    play_parser.add_argument("--seed", type=int, default=EVAL_SEED)
+    play_parser.set_defaults(run=play)
+
+    card_parser = commands.add_parser("card")
+    card_parser.add_argument("--repo", default=REPO)
+    card_parser.set_defaults(run=card)
+
+    args = parser.parse_args()
+    args.run(args)
+
+
+if __name__ == "__main__":
+    main()

config.json

@@ -0,0 +1,12 @@
+{
+ "context": 4,
+ "channels": [
+  64,
+  128,
+  256
+ ],
+ "steps": 3,
+ "sigma_min": 0.002,
+ "sigma_max": 5.0,
+ "rho": 7.0
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:41c6048879dec8144e2c660eac6824c2b4267c5d383d193d33f80d6a4f06acfa
+size 16736044

results/eval.json

@@ -0,0 +1,27 @@
+{
+ "n": 256,
+ "model": {
+  "all": {
+   "1": 41.74,
+   "5": 38.61,
+   "15": 36.32
+  },
+  "changed": {
+   "1": 24.49,
+   "5": 20.7,
+   "15": 18.68
+  }
+ },
+ "repeat": {
+  "all": {
+   "1": 37.19,
+   "5": 34.87,
+   "15": 33.44
+  },
+  "changed": {
+   "1": 11.38,
+   "5": 11.13,
+   "15": 10.99
+  }
+ }
+}

results/train.json

@@ -0,0 +1,26 @@
+{
+ "command": "train",
+ "frames": 200000,
+ "eval_frames": 5000,
+ "eval_windows": 256,
+ "context": 4,
+ "channels": [
+  64,
+  128,
+  256
+ ],
+ "sample_steps": 3,
+ "steps": 50000,
+ "batch_size": 64,
+ "lr": 0.0001,
+ "seed": 0,
+ "log_every": 500,
+ "checkpoint": "/checkpoints/breakout",
+ "save_every": 2500,
+ "output": "out",
+ "repo": "jgalego/breakout-world-model",
+ "episodes": 600,
+ "params": 4182083,
+ "device": "NVIDIA A10G",
+ "minutes": 55.9
+}