next-state/open-dreamer

Open-source Dreamer world-model implementation in JAX

Open Dreamer – Real‑time Dreamer 4 world‑model in JAX/Flax

What it is

  • An open‑source implementation of the Dreamer 4 world‑model pipeline, written with JAX/Flax (the NNX library).
  • It provides the full training stack for a video‑tokenizer, an action‑conditioned latent dynamics model, and evaluation tools (roll‑outs, Fréchet Video Distance).
  • The model is trained on Minecraft/VPT‑style gameplay videos and can be run in real time via a cloud‑hosted demo.

Key components

Component Role
dreamer/models.py Definitions of the causal video tokenizer and the latent dynamics model.
dreamer/training.py Training loops, loss functions, and checkpoint handling (via Orbax).
dreamer/generation.py Utilities for denoising schedules, generating roll‑outs, and rendering videos.
dreamer/fvd/ Feature extraction and Fréchet Video Distance computation for quality assessment.
scripts/ Ready‑to‑run entry points: train_tokenizer.py, tokenize_minecraft_dataset.py, train_dynamics.py, eval_fvd.py.
configs/ Hydra/OmegaConf YAML files that describe datasets, model hyper‑parameters, and evaluation settings.
site/ Next.js website that hosts the interactive demo and documentation.

Typical workflow

  1. Prepare data – Convert raw Minecraft MP4 recordings into ArrayRecord shards (pickled dicts with video bytes, actions, etc.).
  2. Train tokenizer – Learn a latent representation of video frames (scripts/train_tokenizer.py).
  3. Tokenize dataset – Encode every raw episode into latent tokens (scripts/tokenize_minecraft_dataset.py).
  4. Train dynamics – Fit the action‑conditioned latent dynamics model on the tokenized data (scripts/train_dynamics.py).
  5. Generate & evaluate – Roll out the model to synthesize new video frames and compute FVD (scripts/eval_fvd.py).

How to try it

  • Live demo – No installation needed; the website hosts an in‑browser demo that streams a Minecraft world and lets you toggle between the real game and the model’s predictions.
  • Local inference – The companion repo reactor-team/open-dreamer contains a lightweight inference script to run a trained checkpoint on your own video/action sequences.

Installation (training side)

# Requires Python 3.11 and a CUDA‑12‑compatible JAX build
pip install uv               # UV is the fast Python package manager used here
uv sync                       # Installs pinned dependencies into a virtual env
source .venv/bin/activate     # Activate the env

If you need a different JAX wheel (e.g., for a different GPU or CPU), install it after the uv sync step.

Running the training scripts (example for the tokenizer)

# Edit configs/tokenizer.yaml and configs/dataset/minecraft_vpt.yaml first
uv run scripts/train_tokenizer.py

Similar commands exist for tokenization, dynamics training, and evaluation, all driven by the YAML config files under configs/.

Roadmap

  • Currently supports only the world‑model training pipeline.
  • Future work aims to add a full Dreamer 4 behaviour‑cloning / reinforcement‑learning loop.

License & citation

  • The repository currently carries a placeholder “All rights reserved” notice; a formal open‑source license is planned for a later release.
  • If you use the code in research, cite the Zenodo entry provided in the README and the original Dreamer 4 paper.

Who might find this useful?

  • Researchers exploring scalable world‑model learning, especially on video‑rich environments like Minecraft.
  • Engineers building real‑time generative agents that need a fast JAX implementation.
  • Students looking for a concrete example of a full video‑tokenizer + latent dynamics pipeline.

All information above is taken directly from the repository’s README; no additional features have been inferred.

Related

  • Project
    nicklashansen/dreamer4

    A PyTorch implementation of the Dreamer 4 world model designed for multi-task continuous control, featuring a causal tokenizer and an interactive dynamics model.

  • Project
    google/brax

    A fast, fully differentiable physics engine written in JAX for large-scale rigid body simulation and reinforcement learning training.

  • Project
    AMAP-ML/DreamX-World

    DreamX-World is a general-purpose interactive world model that generates high-fidelity, controllable simulations allowing users to explore and transform environments via event prompts.

  • Project
    deeplearning4j/deeplearning4j-examples

    A set of Maven‑based Java examples that demonstrate how to use the Eclipse Deeplearning4J ecosystem (DL4J, ND4J, SameDiff, DataVec, RL4J, etc.) for tasks ranging from basic neural‑network training to Spark‑distributed and GPU‑accelerated learning, plus model import and Android deployment.

  • Project
    redai-infra/Relax

    Relax is an open‑source, Ray‑Serve‑based reinforcement‑learning engine for post‑training multimodal LLMs. It decouples rollout, actor, and critic into independent services, supports fully asynchronous or hybrid execution, and provides elastic scaling, a rich set of RL algorithms, and production‑grade ops. The repo includes Docker images, quick‑start scripts for text, vision‑language, and omni‑modal tasks, and extensive documentation.