vijayabhaskar-ev/dreamer_v4
From-scratch PyTorch implementation of DreamerV4 (Hafner et al., 2024): masked-autoencoder tokenizer, block-causal flow-matching dynamics with bootstrap curriculum, agent-token finetuning, and PMPO imagination RL. Hardened for TPU v4 / torch_xla with fixed-shape graphs, on-device RNG, and bounded compile-cache footprint.
GitHub repository with 22 stars and 2 forks.
Language: Python
Topics: dreamer, dreamer-v4, flow-matching, model-based-rl, pmpo, pytorch, reinforcement-learning, torch-xla, tpu, transformer