nanoWorldModel

A small, readable nanoWorld-style world model. The first milestone is an action-conditioned autoregressive Transformer that learns discrete environment dynamics:

BOS, obs_0, action_0, obs_1, reward_0, done_0, action_1, ...

Actions are conditioning tokens. The training loss is applied to predicted world tokens: next observation, reward, and done. This keeps the code close to nanoWorld while making the objective a world model rather than a plain language model.

Why This Shape

Open-source world-model systems differ in scale and modeling family, but several useful patterns repeat:

• DreamerV3 learns a compact dynamics model and trains behavior from imagined rollouts.
• IRIS combines a discrete autoencoder with an autoregressive Transformer world model.
• DIAMOND shows that strong visual world models can drive agents through imagined environments.
• nanoWorld is the style target: small files, direct PyTorch, and minimal magic.

This repo starts with the IRIS-like discrete-token path because it is the simplest route to a ~100M parameter model that is easy to inspect, test, and extend. A visual tokenizer can later replace the toy grid observation IDs without changing the Transformer core.

Repository Layout

src/nanowm/
  config.py      model presets, including nano_100m
  model.py       nanoWorld-style decoder-only Transformer
  tokens.py      world-model token grammar and type helpers
  gridworld.py   deterministic synthetic environment dataset
  train.py       training loop and checkpointing
  sample.py      action-conditioned rollout from checkpoints
modal_train.py   Modal GPU smoke/train entrypoint
tests/           unit tests for tokens, data, model, and parameter budget
architecture.md  architecture and implementation plan
handoff.md       running work log and next actions

Quickstart

Use Python 3.10 or newer.

python3 -m venv .venv
source .venv/bin/activate
python -m pip install -e ".[dev]"
pytest -q

On this machine, Miniforge Python 3.12 had a native readline segfault during pytest startup. The project tests passed with a temporary stub:

mkdir -p /private/tmp/nanowm_readline_stub
touch /private/tmp/nanowm_readline_stub/readline.py
PYTHONPATH=/private/tmp/nanowm_readline_stub:src pytest -q

Train the tiny debug model locally:

nanowm-train --config-name debug --max-iters 100 --out-dir out/debug

Train on the open-source Hugging Face CartPole trajectory dataset:

python -m pip install -e ".[hf]"
nanowm-train \
  --dataset hf_cartpole \
  --config-name debug \
  --block-size 128 \
  --hf-max-rows 50000 \
  --max-iters 100 \
  --out-dir out/hf-cartpole-debug

The default HF dataset is NathanGavenski/CartPole-v1, which provides row-wise obs, actions, rewards, and episode_starts. The adapter quantizes the 4D CartPole observation into a discrete observation token and derives terminal transitions from episode boundaries.

Roll out a checkpoint with random actions:

nanowm-sample --ckpt out/debug/ckpt.pt --steps 16

Modal

Run a small GPU smoke test:

modal run modal_train.py --config-name debug --max-iters 100

Run the Hugging Face CartPole smoke test on Modal:

modal run modal_train.py --config-name debug --max-iters 10 --dataset hf_cartpole --hf-max-rows 5000

Run the ~100M preset after the smoke test is healthy:

modal run modal_train.py --config-name nano_100m --max-iters 1000

The Modal entrypoint mounts src/nanowm into /root/nanowm so local package code is sent into the remote container without requiring a package publish step.

Verified smoke command:

modal run modal_train.py --config-name debug --max-iters 10

This completed on an A10G with CUDA/bfloat16 and best_val_loss near 3.68.

Verified HF smoke command:

modal run modal_train.py --config-name debug --max-iters 10 --dataset hf_cartpole --hf-max-rows 5000

This completed on an A10G with CUDA/bfloat16 and best_val_loss near 7.29.

Model Presets

• debug: 2 layers, 64 hidden size, for tests and fast iteration.
• small: 6 layers, 384 hidden size, for local experiments.
• nano_100m: 14 layers, 12 heads, 768 hidden size, roughly 100M parameters with the compact gridworld vocabulary.

Current Scope

Implemented now:

• Discrete token grammar for observations, actions, rewards, and done flags.
• Causal Transformer with tied token embeddings, nanoWorld-style residual scaling, Flash Attention when PyTorch provides it, and masked CE loss.
• Synthetic gridworld dataset for deterministic smoke tests.
• Hugging Face CartPole trajectory adapter for open-source RL data smoke tests.
• Training, checkpointing, sampling, tests, and Modal entrypoint.

Planned next:

• Add a visual tokenizer interface so image/video latents can replace gridworld observation IDs.
• Add richer evaluation metrics: one-step observation accuracy, reward accuracy, done accuracy, and multi-step rollout consistency.
• Add optional policy/imagination code once the world model loss is stable.