AutoWorldModel-Bench

Game state sequences for training and evaluating action-conditioned world models. 8 classic game environments with a unified entity-based tensor schema, deterministic train/val/test/scenario splits, and 152,000 total episodes.

This dataset accompanies the AutoWorldModel-Bench benchmark for evaluating frontier coding agents on open-ended world-model research.

Dataset Structure

data/                                # Parquet training data
└── {game}/
    ├── train.parquet                # 10,000 episodes
    ├── val.parquet                  #  3,000 episodes
    ├── test.parquet                 #  3,000 episodes
    ├── scenario.parquet             #  3,000 episodes
    └── meta.json                    # max_entities, dimensions, total_frames

scenarios/                           # Curated scenario archives (tar.gz per game)
└── {game}.tar.gz
    └── {game}/
        └── {scenario_name}/         # e.g. ball_hits_paddle, ship_dies
            └── data_ep_{id}/
                ├── frames.jsonl.gz  # Per-frame entity states
                ├── manifest.json    # Game schema, entity kinds, action/global fields
                ├── meta.json        # Episode metadata, event info, rollout params
                └── rollout.mp4      # Visual replay

Games

Total: 152,000 episodes (19,000 per game), 158.0M frames

All games share a unified tensor schema: registry_dim=34, state_dim=23.

Each game has 10,000 train / 3,000 val / 3,000 test / 3,000 scenario episodes.

Data Collection Policy

v2 games (asteroids, breakout, frogger, kong, platformer) use a 3-policy mix for diverse behavioral coverage:

• Random — uniform random actions
• Heuristic — hand-crafted game-specific strategies
• RL checkpoint — DQN/PPO agents at various training stages

v1 games (pong, racer, snake) use a 50/50 heuristic + random mix.

Scenarios

Hand-picked and categorized episodes that isolate specific game events — sourced from the scenario split and augmented with synthetically collected episodes. Each episode captures a short rollout around a key event (e.g., collision, scoring, death) with history context.

Each episode contains 32 history frames + 1 pre-event frame, followed by up to 20 rollout frames (including the event). Rollouts are truncated early on termination.

Each game includes a same_state_different_actions scenario that tests action-conditioning by replaying the same initial state with varied actions.

Downloading scenarios

from huggingface_hub import hf_hub_download
import tarfile

path = hf_hub_download(
    "AutoWorldModel/AutoWorldModelBench",
    "scenarios/pong.tar.gz",
    repo_type="dataset",
)
with tarfile.open(path) as tar:
    tar.extractall("./scenarios")
# ./scenarios/pong/ball_hits_left_paddle_moving/data_ep_.../frames.jsonl.gz

Tensor Schema

Each Parquet row stores one episode as serialized numpy arrays:

Where N = max_entities (game-specific), T = episode length.

Usage

With the datasets library

from datasets import load_dataset

ds = load_dataset("AutoWorldModel/AutoWorldModelBench", "pong")
print(ds["train"][0].keys())

Direct download with huggingface_hub

from huggingface_hub import hf_hub_download
import pyarrow.parquet as pq
import numpy as np, json

path = hf_hub_download(
    "AutoWorldModel/AutoWorldModelBench",
    "data/pong/train.parquet",
    repo_type="dataset",
)

table = pq.read_table(path)
row = table.to_pydict()
states = np.frombuffer(
    row["states"][0], dtype=row["states_dtype"][0]
).reshape(json.loads(row["states_shape"][0]))

Evaluation

Models trained on this data are evaluated on multi-step open-loop rollouts at horizons {1, 10, 20}:

• Position L1: Mean absolute error on entity (x, y) positions (lower is better)
• Alive F1: F1 score on entity alive/dead classification
• Composite: 0.9 * (1 - pos_l1) + 0.1 * alive_f1 (higher is better)

Citation

If you use this dataset, please cite:

@misc{autoworldmodelbench2025,
  title={AutoWorldModel-Bench: A Benchmark for Evaluating Coding Agents on World Model Research},
  author={AutoWorldModel Team},
  year={2025},
  url={https://github.com/AutoWorldModelBench/Benchmark}
}