diamond

Why we included this project

DIAMOND is the code behind a NeurIPS 2024 Spotlight paper that trains a reinforcement learning agent inside a diffusion-based world model instead of against the real environment. For people working on model-based RL or world models, it's a rare thing: a clean baseline you can run, with pretrained checkpoints for Atari that you can watch and even take control of. The paper and project page explain the design choices clearly, so the repo doubles as a practical way to learn how diffusion models can act as learned simulators. It's a starting point for teams curious about replacing or supplementing traditional simulators, not a finished product.