$ show lucidrains/dreamer4
This is a research-grade PyTorch implementation of Danijar Hafner's Dreamer v4 model-based reinforcement learning algorithm. The core artifact is the `dreamer4` package, dominated by a monolithic `dreamer4/dreamer4.py` (~3000 lines) that implements `VideoTokenizer`/`SpaceTimeTokenizer`, `DynamicsWorldModel`, `AudioTokenizer`, and `StateTokenizer`. These classes build causal space-time transformers with separate space/time attention blocks, latent flow-matching diffusion with shortcut consistency, BYOL-style tokenizer pre-training, PoPE attention, and a large zoo of experimental loss terms (latent autoregressive loss, decorrelation, orthogonality, SigReg, self-flow distillation, shortcut loss). `dreamer4/trainers.py` provides training loops (VideoTokenizerTrainer, BehaviorCloneTrainer, DreamTrainer, SimTrainer) and `dreamer4/cli.py` exposes CLI entry points for training and serving the world model. The repository is hybrid: primarily research code, but packaged and tested like a library. It ships a pyproject.toml with pinned dependencies, a parametrized test suite (tests/test_dreamer.py, test_e2e.py, test_env_wrapper.py), a browser-based Snake environment with a Flask-style server (`dreamer4/web_env/`) used for human demonstration collection and replay-buffer inspection, and environment wrappers in `dreamer4/env.py` that record agent experience to folders or memmap replay buffers. Root-level scripts are the actual experiment entry points, targeting standard benchmarks: HalfCheetah (imagination RL), CartPole (dynamics + RL), MovingMNIST (tokenizer + dynamics), and Snake (PPO with dynamics model). Craftsmanship is above typical research code: clean package layout, pytest coverage across many hyperparameter combinations, self-contained utilities (LossNormalizer, tensor pack/unpack helpers, with_seed decorator, memmap replay-buffer helpers), and HuggingFace Accelerate integration for distributed/mixed-precision training. However, the single-file core, 100+ hyperparameter flags on trainers, scattered hard-coded constants, sparse docstrings, and reliance on volatile niche packages (x-transformers, h-net-dynamic-chunking, PoPE_pytorch, adam-atan2-pytorch) make long-term stability and code reuse risky. The repository is recently committed and appears maintained, but it is explicitly beta-status experimental research code, not a production system.
The LLM identified 5 files worth highlighting for reuse:
LossNormalizer is a self-contained EMA-based loss scaler that stabilizes training when many auxiliary losses of wildly different magnitudes are combined. It is ~20 lines, has no internal coupling, and can be dropped into any multi-loss training loop (RL, diffusion, GANs) to replace hand-tuned loss weights.
The Experience dataclass paired with create_memmap_replay_buffer, add_to_memmap_buffer, combine_experiences, and unbind forms a reusable memory-mapped replay-buffer pattern for RL. It solves the recurring problem of storing large amounts of agent experience without exhausting RAM, and the interface is generic enough to adapt to any RL project's transition schema.
RecordToFolderEnvWrapper and RecordToReplayBufferEnvWrapper are Gymnasium wrappers that transparently record transitions during environment interaction, writing either video files or directly into a memmap replay buffer. This solves the recurring RL data-collection problem (capturing demonstrations and rollouts under a standard gym API) and is decoupled enough to extract into any gym-based project.
pack_one, unpack, and maybe_pack_unpack are generic tensor flattening/restoration helpers and a higher-order function that applies flow-loss computation across either packed or unpacked multimodal inputs. They are dependency-free, purely functional utilities for handling variable-length sequences in batch training, and translate directly to non-RL sequence-modeling projects.
The with_seed decorator wraps arbitrary callables to run under a fixed random seed and restore the prior RNG state afterward. It is a tiny, universally reusable pattern for making stochastic operators (sampling, augmentation, exploration) reproducible inside training loops without polluting global RNG state.