tsilva/gymemu-breakout-state-decoder-fs1

Breakout state decoder

A small PyTorch neural renderer that converts compact Breakout state into RGB, trained on recorded frames from the frameskip-1 dataset using Gymemu.

At a glance

Item	Value
Task	Recorded state to matching RGB playfield
Model	Spatial features plus shared pixel MLP, 10,249 parameters
Input	112 physical-unit values: ball x/y, paddle x/width, 108 brick bits
Output	Float32 RGB, N × 3 × 210 × 160, range 0–1
Format	PyTorch weights and standalone Python
Validation	99.999377% exact pixels; 96.1426% pixel-perfect playfields
Checkpoint	Seed 2026, update 12,000
Scope	HUD omitted; test and predicted-state rollouts not evaluated

Quick start

Download the repository files and install the pinned requirements.txt in a Python 3.11–3.13 environment. With Hugging Face Hub and PyTorch available, this example downloads the files and renders the included synthetic state. The downloaded Python source is ordinary local inference code and can be inspected before importing.

import json
import sys
import torch
from huggingface_hub import snapshot_download

folder = snapshot_download("tsilva/gymemu-breakout-state-decoder-fs1")
sys.path.insert(0, folder)
from inference import load_decoder
from pathlib import Path

model = load_decoder(folder)
state = torch.tensor(
    json.loads((Path(folder) / "example_state.json").read_text()),
    dtype=torch.float32,
)
with torch.inference_mode():
    rgb = model.render(state)
print(tuple(rgb.shape))  # (1, 3, 210, 160)

Run python verify.py from the downloaded repository to check artifact hashes and reproduce the synthetic example output. This is an inference check, not a rerun of dataset validation. For reproducible downloads, supply a specific commit SHA as revision.

To consume nonterminal outputs of the unified dynamics model, call model.from_dynamics(prediction) and then model.render(visual_state) under torch.inference_mode(). Stop terminal rows first. This adapter exists, but the decoder has only been evaluated on recorded states.

Validation results

The fixed sample contains 2,048 validation frames from the existing episode/seed-grouped split. Model selection uses this sample. The test split remains reserved.

Metric	Selected checkpoint
Exact RGB pixels, excluding HUD	99.999377%
Pixel-perfect playfields	1,969 / 2,048, or 96.1426%
Incorrect playfield pixels	394 across 79 frames
Playfield RGB MSE	0.00000164581
Ball-neighborhood RGB MSE	0
Correct isolated-ball position	1,769 / 1,769 scorable frames
Full RGB MSE, including blacked-out HUD	0.00301980

An independent sprite detector cannot uniquely isolate the ground-truth ball in the remaining 279 frames; those frames are excluded from the isolated-ball metric, but included in pixel and ball-neighborhood metrics. Exact pixels compare RGB bytes. MSE uses float32 RGB scaled to 0–1. Playfield metrics exclude rows 0–16 and are not directly comparable with full next-frame RGB baselines.

The comparison below shows recorded RGB, decoded RGB and wrong pixels in red. HUD is masked in both images. The first four examples are fixed sample entries; the last four are the worst validation reconstructions. Most visible error in these worst examples is along the lower edge of the top wall.

Input and output

Use float32 tensors of shape N × 112 in physical units, not the dataset's normalized label values.

Indices	Meaning
0	Ball x
1	Integer RAM ball y; the model handles its rendering offset
2	Paddle x
3	Paddle width, 12 or 16
4:112	Binary brick occupancy, row-major 6 × 18

render(state) returns N × 3 × 210 × 160 RGB in 0–1. The nine-color palette comes only from training frames. HUD rows 0:17 are forced black. There is no frame history, action input, image encoder, or learned transition in this decoder. Do not render terminal placeholder states.

Architecture

Each pixel gets 20 spatial features derived from its coordinates and the state, including distances to the ball, paddle and fixed boundaries, local brick occupancy, and coordinates within a brick. A shared 20 → 64 → 64 → 64 → 9 MLP with SiLU hidden activations predicts palette logits. Argmax selects the output color.

This geometry is explicit prior knowledge. The network learns pixel colors and visibility from dataset targets. The result does not demonstrate that an unconstrained autoencoder discovers the same representation. render uses discrete argmax and is not differentiable end to end; forward(state, coordinates) exposes logits for training.

Training recipe

• 32,768 uniformly sampled eligible training frames; 2,048 fixed validation frames. Sampling seed 20260922.
• Match successor-state labels to successor_frame_id by ID. Keep active ball states with trustworthy brick grids and exclude initial-wall flags. No train/validation frame-ID overlap.
• Each update samples 32 frames and 256 pixels per frame: 128 uniform playfield, 64 ball-neighborhood, 32 paddle-neighborhood and 32 brick-wall pixels.
• Optimize the equal mean of four regional palette cross-entropies. Extra object sampling prevents background pixels from dominating.
• AdamW, 12,000 updates, seed 2026, learning rate 0.003 with cosine decay to 0.00003, weight decay 0.00001, gradient clipping 5.
• Validate every 1,000 updates and select minimum playfield RGB MSE. Selected update 12,000; elapsed training including validation 523 seconds.

Files

• pytorch_model.bin, config.json, model.py, inference.py: standalone inference.
• checkpoint.pt: unchanged original training checkpoint, load with weights_only=True.
• verify.py, SHA256SUMS.json, example_state.json, example_output.json, example_render.png: artifact integrity and synthetic inference check.
• evaluation.json, training_history.json, training_config.json, error_rows.json: recorded measurements and training settings.
• provenance.json, io_contract.json, packaging_verification.json: dataset identity, interface and packaging evidence.

Provenance and limitations

Dataset revision 5f6e0ca8c28e2fc27aeda3ead04851a1f8a45a77, split ID bb52cc995e2817362e3c68be0ff82f1d9cfdec87d7edb8c7d9ea2929d3d5b800. Original checkpoint SHA256 eb02361c4bdd2435124201870441d42a636c7e55900178a906d67228d86eb8f9. Local run runs/state-decoder-fs1-20260922.

The model omits HUD, startup and terminal rendering. Results cover a fixed validation sample, not every dataset frame. Small boundary artifacts remain. Checkpoint reload reproduced validation metrics, and CPU/MPS byte outputs matched on a 16-frame probe. This publication preserves the trained weights. It does not add test evaluation, retraining, or proof of stable dynamics rollouts.