AQ-Mario: EMA-JEPA world model
This repository contains a trained action-conditioned latent world model from
the full3ep_pure_ema999 experiment in
AQ-Mario. It uses an EMA-JEPA
architecture to learn the dynamics of Super Mario Bros. 1-1 from frame
sequences and actions. The model predicts future latent states rather than RGB
frames, making it suitable for representation analysis and model-based planning
experiments.
Preview
This is a rollout preview from the EMA-JEPA model:
The complete training dataset is available at maxmill/aq-mario-smb1.
Files
jepa.pt: PyTorch model checkpoint containing the JEPA world model, auxiliary heads, EMA teacher, and AdamW optimizer state.metrics.jsonl: training metrics.gate_by_epoch.json: per-checkpoint probe and health measurements.gates.json: final representation and action-conditioning gate results.param_count.json: parameter report.
Checkpoint details
- Variant: pure JEPA with an EMA target encoder (
ema_target=0.999) - Training: 41,160 steps, 3 epochs
- Trainable parameters: 9,801,795 plus auxiliary heads
- Final checkpoint: 133.5 MiB
The final gate report is included for transparency. This run is not presented as a solved Mario controller: its final x/y/scroll probes and action gate do not pass the project's thresholds. See the GitHub repository for the loader, training code, evaluation protocol, and dataset documentation.
Loading
import torch
checkpoint = torch.load("jepa.pt", map_location="cpu", weights_only=False)
state_dict = checkpoint["jepa"]
The repository's aqmario.model.load_jepa helper can load the checkpoint when
the AQ-Mario source tree and configuration are available.
Data and rights
The training data consists of derived gameplay observations from Super Mario Bros. 1-1. AQ-Mario is an independent research project and is not affiliated with or endorsed by Nintendo. The source code is MIT licensed; this checkpoint and the derived gameplay data should be used subject to the rights and terms applicable to the underlying game content.