Implement full DIAMOND Atari agent with pixel-space diffusion world model and RL in imagination

[ParamThakkar123]

Summary

Implements the complete DIAMOND algorithm (Alonso et al., NeurIPS 2024) — a pixel-space diffusion world model trained with REINFORCE + λ-return value baseline inside imagination, following the Atari 100k benchmark protocol.

Core components

• Diffusion world model: EDM preconditioning (Karras et al. 2022) with Euler sampling, SimpleUNet with AdaptiveGroupNorm, conditioned on 4-frame history + action embeddings
• Reward/termination model: Separate CNN-LSTM predicting scalar rewards and terminal probabilities from imagined frames
• Actor-critic: Shared CNN-LSTM backbone with categorical policy head and scalar value head, trained via λ-return REINFORCE
• CategoricalCEMSolver integration: DiamondCostModel wraps diffusion + reward LSTM as a Costable for MPC planning

Atari 100k protocol

• NoopReset (max 30), FireReset, EpisodicLife, ClipReward ({-1, 0, 1})
• Sticky actions (repeat_action_probability=0.25)
• 64×64 RGB observations, reward clipping, full 18-action space

Training pipeline

• 1000 epochs, each collecting 100 env steps then 400 gradient steps per component
• Three optimizers: diffusion (+ action encoder), reward/term model, actor-critic
• MPC evaluation via CategoricalCEMSolver (random shooting, single iteration)
• 27 tests covering EDM preconditioners, cost model, agent rollout, λ-returns, RL losses

Cleanup

• Removed legacy embedding-space Diamond class (dead code, replaced by pixel-space agent)
• Atari wrappers moved from wrapper/ to envs/ale/
• planner.py removed (dead code, CategoricalCEMSolver used instead)

Known issues

• Hydra CLI OOM on Windows (runs fine standalone / Linux)
• Evaluation docs need restructuring to be world-model-agnostic

[lucas-maes]

hi! could you report some reward or success rate?

[lucas-maes]

Few comments:

• I think the wrapper you added is already present in ale package. Is it necessary to rewrite the code here?
• Please could you remove the claude # ---- helper comment?
• Make sure to rely on hydra, we only support hydra and don't want a mix of argparse and hydra as you currently did

Thank you very much for the PR! It's sick!

[ParamThakkar123]

Thank you so much @lucas-maes , yes I am working on these parts and will report with the metrics and fixes in a few hours 🫡

[ParamThakkar123]

Hi @lucas-maes , I am trying to run a smoke test on a few epochs for this but I am out of GPU credits right now and my laptop GPU is going OOM for even the smoke test (RTX 3050 4GB VRAM), Would it be possible for you to run a smoke test on your end with this command?

python scripts/train/diamond_full.py epochs=2 steps_per_epoch=5 train_steps=2 eval_every=1 log_every=1 max_episode_steps=100 2>&1