# next-state/open-dreamer: the Dreamer 4 training pipeline in JAX

> Open Dreamer is a JAX/Flax NNX implementation of the Dreamer 4 world-model pipeline, covering a causal video tokenizer and an action-conditioned latent dynamics model. It is a training repository, not a ready-made agent, and the README points elsewhere for local inference.

**next-state/open-dreamer** — Open-source Dreamer world-model implementation in JAX

- Repository: https://github.com/next-state/open-dreamer
- Website: https://next-state.github.io/open-dreamer/
- Stars: 390 · Forks: 32
- Language: Python
- License: NOASSERTION
- Published: 2026-09-17 · Updated: 2026-09-17 · Language: en
- Canonical page: https://hysenlabs.com/projects/next-state-open-dreamer

## What Open Dreamer actually ships, and who it is for

The repository is a training pipeline, and the README says so plainly: it holds a causal video tokenizer, an action-conditioned latent dynamics model, and the tools to generate rollouts and measure FVD. The four supported tasks are tokenizer training, tokenizing Minecraft/VPT-style MP4 datasets, dynamics training, and rollout generation with FVD scoring. There is no agent, no reward model, and no environment loop in this repository.

That scope decides the audience. If you are researching world models and want a readable JAX/Flax NNX codebase to modify, this is aimed at you. If you want to play a generated Minecraft world, the README sends you to a browser demo that runs in the cloud on Reactor, or to a separate repository, reactor-team/open-dreamer, for local inference on your own video and actions. Installing this repository to get a playable model is the wrong purchase.

The project describes itself as an open, real-time implementation of the Dreamer 4 pipeline. The real-time claim belongs to the demo and the inference repository; the training code here is a batch pipeline built around Grain, ArrayRecord, and Orbax checkpoints.

## The two-stage architecture: tokenizer first, dynamics second

The pipeline is sequential, and the order is not optional. Stage one trains a causal video tokenizer on raw clips. Stage two trains the dynamics model on the latents that the frozen tokenizer produces, conditioned on shifted actions. The README's workflow section spells out six steps: prepare raw MP4 ArrayRecord shards, train the tokenizer, tokenize full episodes into latent ArrayRecords, copy the generated latent statistics into the latent dataset config, train the dynamics model, then generate videos and compute FVD.

The data formats differ between the stages, which is the part worth internalising. Raw records are pickled Python dicts containing MP4 bytes, a video shape, actions, and optional source metadata, stored as shard-*.array_record files. Tokenized records are msgpack ArrayRecords holding tokenizer latents and actions. So the tokenizer stage reads decoded video, and everything downstream reads tensors.

There is a deliberate asymmetry in window lengths. The README notes that tokenizer training can use shorter windows than the full raw episode, giving the example of training on 16-frame windows from fixed 256-frame records and then tokenizing as full 256-frame latent episodes. That is a memory trade-off: you pay for long sequences once, at tokenization time, rather than in every training step. The repository layout matches this split, with dreamer/models.py holding both model definitions and dreamer/data/ holding the Grain and ArrayRecord plumbing.

## Installing Open Dreamer and running the tokenizer

The requirements are narrow. The README lists Python 3.11, uv, a CUDA 12-compatible JAX environment, and Minecraft/VPT-style ArrayRecord data. The pyproject file pins requires-python to ==3.11.*, so a 3.12 interpreter will not resolve.

Install with uv, which the README gives as three commands:

```bash
pip install uv
uv sync
source .venv/bin/activate
```

The dependency lock targets CUDA 12 JAX. If your machine needs a different JAX build, the README says to install the correct wheel for your accelerator setup after syncing. Note the pin on nvidia-cusparse-cu12 in pyproject.toml, which a comment attributes to a CUDA 12.8 compatibility issue with nvJitLink 12.9. That is the kind of detail you only discover by reading the file, not the README.

Before training anything, edit the raw dataset config. The README shows these keys:

```yaml
array_record_path: /path/to/mp4-arrayrecords
index_max: 1500
dataset_mean: [0.2241, 0.2348, 0.2086]
dataset_std: [0.1809, 0.1874, 0.2282]
```

The mean and standard deviation are pixel statistics for normalized video values, and the README says to recompute them when changing the raw dataset. Leaving the defaults in place on a different dataset is a silent correctness bug, not a crash.

Then run the tokenizer script:

```bash
uv run scripts/train_tokenizer.py
```

Outputs land in logs/<run_name>/. The checkpoints/ directory holds tokenizer and optimizer checkpoints, and vis/ holds reconstruction images when visualize_every is greater than zero. Those reconstruction images are the first real signal that the tokenizer is learning anything, so set visualize_every above zero on a first run.

## Tokenizing episodes and wiring latent statistics into the dynamics config

The second stage is a data transformation, not a training run. After the tokenizer is trained, edit configs/tokenize.yaml and run:

```bash
uv run scripts/tokenize_minecraft_dataset.py
```

The output is latent shards plus statistics, laid out as shard-00000.array_record, shard-00001.array_record, and metadata/latent_stats.npz. The README provides a short snippet to print those statistics, loading the npz and printing mean, std, num_samples, and num_videos.

Those two numbers then have to be copied by hand into configs/dataset/minecraft_vpt_latent.yaml, along with array_record_path and index_max pointed at the tokenized output. This manual step is where the pipeline is most fragile. Nothing in the README describes an automated check that the latent statistics in the config match the shards on disk. If you re-tokenize with a different tokenizer checkpoint and forget to update the config, the dynamics model trains against mismatched normalisation and the failure is quiet.

Only after that does dynamics training start, reading configs/dynamics.yaml, which imports the latent dataset config. The model trains on latent ArrayRecords and shifted actions.

## Where Open Dreamer is the wrong tool

Three limits are visible from the repository itself.

The first is the missing agent. The roadmap has a single unchecked item: full Dreamer 4 behaviour-cloning and RL agent training loop. Until that lands, this repository cannot produce a policy. It produces a world model and the machinery to evaluate it with FVD. If your goal is an agent that improves at a task through interaction, this codebase stops short of the part you need.

The second is the data contract. The pipeline expects Minecraft/VPT-style shards named shard-*.array_record with fixed 256-frame raw records, and the README's example values for index_max, short_T, long_T, horizon, and fvd_chunk_size are explained in dreamer/data/README.md rather than in the main file. Adapting to a different video domain means writing your own record format and recomputing normalisation statistics, not flipping a config flag.

The third is the environment. Python 3.11 exactly, uv, and CUDA 12 JAX. The README acknowledges that other accelerators need a different JAX wheel installed after sync, but the lock and the cusparse pin are built around CUDA 12. On a CPU-only or Apple-silicon machine, expect to resolve the environment by hand before any of the training commands mean anything.

## How it relates to the inference repository and to the hosted demo

The README splits the project across three surfaces, and choosing the right one is most of the decision. The hosted demo at the project website runs the real-time model in the cloud on Reactor, with a Game to Dream toggle that hands the stream from the real game to the world model frame by frame. It requires no setup. The inference repository, reactor-team/open-dreamer, is the path for rolling out a trained model locally on your own video and actions, generating new frames from an MP4 and a matching action sequence. This repository, next-state/open-dreamer, is the training pipeline.

That division is unusual for a research release, which normally bundles training and inference together. The upside is that the training code does not carry inference-serving dependencies. The downside is that reproducing the demo end to end requires two repositories, and the README does not describe how a checkpoint trained here is converted into the format the inference repository expects. If your goal is a closed loop from your own data to your own real-time rollout, that gap is the first thing to investigate.

For comparison, the original Dreamer line from Danijar Hafner and collaborators is published as a research codebase oriented toward agent learning across many environments. Open Dreamer is narrower on purpose: one game domain, one model family, and no policy learning yet.

## Maintenance, licence and what an upgrade costs

The repository is not archived, and the last push was on 2026-08-05. There are no retrieved releases, so there is no version history to reason about and no changelog to read before upgrading. In practice you track the main branch.

That has a concrete cost. The dependency set in pyproject.toml is not a set of stable releases: optax is sourced from a git URL, and minerl is sourced from a git URL, both under tool.uv.sources. The lock file, uv.lock, is committed, which pins today's resolution, but a fresh uv sync after the upstream git branches move can pull different code. The cusparse pin with its comment about nvJitLink is another place where a CUDA upgrade can break the environment. Budget for environment work on any upgrade, not just code review.

The licence is the open question. The repository has a LICENSE file at the top level, but no licence identifier was retrieved, and the README does not discuss terms. Before you build on this, open LICENSE and read it yourself; do not assume an OSI licence because the project is described as open source. Nothing here is legal advice, and the licence text is the only authority.

## Conclusion

Adopt it if you already have Minecraft/VPT-style ArrayRecord shards, a CUDA 12 machine, and the patience to run a tokenizer stage before any dynamics training. Skip it if you want an agent that plays a game today: the README lists the full Dreamer 4 behaviour-cloning and RL training loop as an unchecked roadmap item, and the playable path is the hosted demo or the separate reactor-team/open-dreamer inference repository. Before committing, confirm three things: that your raw shards are named shard-*.array_record with 256-frame records, that your accelerator has a matching JAX wheel, and that the repository's LICENSE file actually states terms you can live with.

## FAQ

### Can you explain Dreamer v4 and how it works?

The README describes Open Dreamer as an implementation of the Dreamer 4 world-model pipeline, split into a causal video tokenizer and an action-conditioned latent dynamics model. Training happens in two stages: the tokenizer learns a latent representation from raw video, then the dynamics model trains on tokenized latent episodes and shifted actions.

### What is being a dreamer?

The repository does not define the term. Its own framing is technical: a world model trained on Minecraft/VPT-style gameplay that can generate rollouts and be scored with FVD, with the live in-browser demo described as the fastest way to experience it.

### How do I install next-state/open-dreamer?

The README gives three commands: pip install uv, then uv sync, then source .venv/bin/activate. The requirements are Python 3.11, uv, a CUDA 12-compatible JAX environment, and Minecraft/VPT-style ArrayRecord data.

### Does next-state/open-dreamer train an agent?

No. The roadmap lists full Dreamer 4 behaviour-cloning and RL agent training loop as an unchecked item, and the README's supported tasks cover the tokenizer, dataset tokenization, dynamics training, and rollout generation with FVD.

### Where do I run next-state/open-dreamer locally for inference?

The README points to the separate repository reactor-team/open-dreamer for rolling out a trained model locally on your own video and actions. This repository covers training, not local inference.

## Sources

- [Issues](https://github.com/next-state/open-dreamer/issues)
- [next-state/open-dreamer on GitHub](https://github.com/next-state/open-dreamer)
- [Project website](https://next-state.github.io/open-dreamer/)
- [README](https://github.com/next-state/open-dreamer/blob/main/README.md)

---

Hysen Labs editorial analysis, written from the project's own repository and release notes. Cite the canonical page: https://hysenlabs.com/projects/next-state-open-dreamer