StatePlay

Project Page · Paper · Code · Dataset · Model

StatePlay is a state- and action-conditioned world model for mechanics-consistent generation in Street Fighter III. It jointly predicts video frames and five game states: timer, both players' health, and both players' skill meters.

Installation

Requirements: Python 3.10+, CUDA, and a GPU with bfloat16 support.

git clone https://github.com/Jimntu/StatePlay.git
cd StatePlay
python -m venv .venv
source .venv/bin/activate
pip install -U pip
pip install -e .

Download the inference weights:

hf download onepiece1999/StatePlay \
  StatePlay.safetensors \
  --local-dir examples/checkpoint

hf download Wan-AI/Wan2.2-TI2V-5B \
  Wan2.2_VAE.pth \
  models_t5_umt5-xxl-enc-bf16.pth \
  --local-dir base_model/Wan-AI/Wan2.2-TI2V-5B

hf download Wan-AI/Wan2.1-T2V-1.3B \
  --include "google/umt5-xxl/*" \
  --local-dir base_model/Wan-AI/Wan2.1-T2V-1.3B

Expected inference layout:

StatePlay/
├── examples/checkpoint/StatePlay.safetensors
└── base_model/Wan-AI/
    ├── Wan2.2-TI2V-5B/
    │   ├── Wan2.2_VAE.pth
    │   └── models_t5_umt5-xxl-enc-bf16.pth
    └── Wan2.1-T2V-1.3B/google/umt5-xxl/
        └── ... tokenizer files ...

To keep weights elsewhere:

export STATEPLAY_BASE_MODEL=/absolute/path/to/base_model
export STATEPLAY_CHECKPOINT=/absolute/path/to/StatePlay.safetensors

STATEPLAY_BASE_MODEL must directly contain Wan-AI/.

Code architecture

StatePlay uses separate visual and state transformer branches with shared joint attention. The visual branch predicts video latents; the state branch predicts the five normalized game states. Both branches are conditioned on the initial frame, text prompt, and action sequence.

StatePlay/
├── stateplay/
│   ├── pipeline.py             # public inference pipeline
│   ├── cli.py                  # command-line inference
│   └── models/
│       ├── dit.py              # StatePlay visual/state DiT
│       ├── vae.py              # video VAE wrapper
│       └── text_encoder.py     # text encoder wrapper
├── training/
│   ├── train.py                # training entry point
│   ├── runner_with_state.py    # optimization/checkpoint loop
│   └── data/                   # SF3 action, state, and prompt loading
├── diffsynth/                  # minimal Wan/StatePlay dependencies
├── scripts/
│   ├── inference.sh
│   ├── run_examples.sh
│   └── train.sh
└── examples/
    ├── inputs/                 # eight bundled inputs
    ├── checkpoint/             # local checkpoint
    └── generated/              # generated videos and state predictions

Run examples

Generate all eight bundled examples:

export CUDA_VISIBLE_DEVICES=0
./scripts/run_examples.sh

Generate selected examples:

./scripts/run_examples.sh --only 01 03

Outputs are written to examples/generated/. Each example produces an MP4 and a _state.txt file. The model is loaded once for the entire run.

Inference

export CUDA_VISIBLE_DEVICES=0

./scripts/inference.sh \
  --image examples/inputs/01_macro_success_clip/first_frame.png \
  --actions examples/inputs/01_macro_success_clip/actions.parquet \
  --prompt-file examples/inputs/01_macro_success_clip/prompt.txt \
  --output output.mp4

This writes output.mp4 and output_state.txt. Defaults are 101 frames, 30 denoising steps, text CFG 5.0, state/action CFG 1.0, and seed 2.

Use custom model paths through STATEPLAY_BASE_MODEL and STATEPLAY_CHECKPOINT, or inspect all options with:

./scripts/inference.sh --help

Training

Download the training dataset:

hf download onepiece1999/StatePlay-Dataset \
  --repo-type dataset \
  --local-dir data/StatePlay-Dataset

Download the three Wan2.2 DiT initialization shards:

hf download Wan-AI/Wan2.2-TI2V-5B \
  diffusion_pytorch_model-00001-of-00003.safetensors \
  diffusion_pytorch_model-00002-of-00003.safetensors \
  diffusion_pytorch_model-00003-of-00003.safetensors \
  --local-dir base_model/Wan-AI/Wan2.2-TI2V-5B

Run training:

export STATEPLAY_BASE_MODEL="$PWD/base_model"
export STATEPLAY_DATA_ROOT="$PWD/data/StatePlay-Dataset/SF3"
export STATEPLAY_OUTPUT="$PWD/outputs/StatePlay"
export CUDA_VISIBLE_DEVICES=0,1,2,3

./scripts/train.sh

The script derives the process count from CUDA_VISIBLE_DEVICES. It trains at 480×832 with 101 frames, learning rate 5e-5, state sampling end, and saves every 500 steps. The released model is the checkpoint at step 40,000.