Pretrained agents for lucid.models.diamond.
One member per Atari 100k game. num_classes is the game's action
count, not a label set — a world model has no classes, and this is
the number the policy head and both action embeddings are built to.
Notes
Reference: Alonso et al., NeurIPS 2024 (arXiv:2405.12399), Table 1 for the per-game scores these agents reach.
DEFAULT is Breakout, which is the game the paper's own
analysis figures use.
CSGO is the odd one out: a world model with no agent, at a
different resolution and nearly thirty times the size, from the
paper's Section 6. It loads into lucid.models.diamond_csgo
rather than the Atari factories.
Examples
>>> from lucid.models.generative.diamond._weights import DIAMONDWeights
>>> len(DIAMONDWeights.__members__)
28
>>> "BREAKOUT" in DIAMONDWeights.__members__
True
Twenty-seven Atari agents and one CS:GO world model share this enum,
and the two are not interchangeable: diamond and
diamond_world_model refuse the CS:GO tag, and diamond_csgo
refuses the Atari ones.
>>> DIAMONDWeights.BREAKOUT.num_classes # Breakout has four actions
4