diamond(pretrained: bool | str = False, weights: DIAMONDWeights | None = None, overrides: object = {})Construct the Atari agent — denoiser, reward/end model, actor-critic.
Model Size
Parameters
pretrainedbool or str= FalseLoad a released agent.
True takes Breakout; pass a game
name — "Pong", "Freeway", any of the benchmark's 26 — for
that one. The tag also sets num_actions, since Atari's
minimal action set differs per game.An explicit tag, taking precedence over
pretrained.**overridesobject= {}Optional
DIAMONDConfig field overrides.Returns
DIAMONDModelThe three networks, untrained.
Notes
Reference: Alonso, Eloi, et al., "Diffusion for World Modeling: Visual Details Matter in Atari", NeurIPS, 2024 (arXiv:2405.12399), Table 2 for the architecture and Table 3 for everything else.
Examples
>>> from lucid.models import diamond
>>> config = diamond().config
>>> config.conditioning_frames, config.denoise_steps, config.horizon
(4, 3, 15)
The appendix's 3D-environment experiments condition on six frames
instead of four, which is a field rather than a second factory:
>>> diamond(conditioning_frames=6).config.denoiser_in_channels
21