data

DIAMONDConfig

extendsGenerativeModelConfig

DIAMONDConfig(sample_size: int | tuple[int, int] = 64, in_channels: int = 3, out_channels: int = 3, act_fn: GenerativeActivation = 'silu', num_actions: int = 18, conditioning_frames: int = 4, unet_channels: tuple[int, ...] = (64, 64, 64, 64), unet_layers: tuple[int, ...] = (2, 2, 2, 2), cond_dim: int = 256, attn_depths: tuple[int, ...] | None = None, with_agent: bool = True, noise_previous_obs: bool = False, upsampler_channels: tuple[int, ...] | None = None, upsampler_layers: tuple[int, ...] | None = None, upsampler_attn_depths: tuple[int, ...] | None = None, upsampling_factor: int = 1, sigma_data: float = 0.5, p_mean: float = -0.4, p_std: float = 1.2, sigma_offset_noise: float = 0.3, attention_head_dim: int = 8, denoise_steps: int = 3, reward_channels: tuple[int, ...] = (32, 32, 32, 32), reward_layers: tuple[int, ...] = (2, 2, 2, 2), reward_cond_dim: int = 128, reward_frames: int = 2, reward_lstm_dim: int = 512, actor_channels: tuple[int, ...] = (32, 32, 64, 64), actor_layers: tuple[int, ...] = (1, 1, 1, 1), actor_lstm_dim: int = 512, horizon: int = 15, gamma: float = 0.985, lambda_: float = 0.95, entropy_weight: float = 0.001, burn_in: int = 4)

source edit

Frozen configuration for the DIAMOND family.

Parameters

sample_sizeint or tuple of int= 64

Frame resolution the world model works at. Atari observations are resized to 64x64 before anything else touches them.

in_channelsint= 3

Channels per frame.

out_channelsint= 3

Channels the denoiser emits — one frame, so the same as in_channels.

num_actionsint= 18

Size of the discrete action set. Atari's full set is 18; a game exposing fewer simply leaves the tail unused.

conditioning_framesint= 4

L

, how many past frames and actions the denoiser sees. The 3D-environment experiments in the paper's appendix use 6.

unet_channelstuple of int= (64, 64, 64, 64)

Residual-block channels per U-Net resolution.

unet_layerstuple of int= (2, 2, 2, 2)

Residual blocks per resolution.

cond_dimint= 256

Width of the conditioning vector the adaptive group norms are regressed from.

attn_depthstuple of int or None= None

Whether each U-Net resolution attends. None means none of them do, which is Atari's setting and the common case; it is expanded to match unet_channels rather than written out, so a narrower model does not have to restate it. All zero for Atari, where only the middle blocks do; the CS:GO world model turns it on at its two deepest resolutions, which is where a 3D scene needs to relate parts of the frame that convolutions cannot reach.

with_agentbool= True

Whether to build the reward/termination model and the actor-critic. The released CS:GO configuration sets both to null: that experiment trains a world model on static data with no reinforcement learning at all, so an agent would be parameters nothing ever updates.

noise_previous_obsbool= False

Whether the conditioning frames are noised too. ⚠️ Off for Atari, on for CS:GO — the released configs differ here and the paper mentions neither. Noising the history is a form of augmentation against the compounding error a long autoregressive rollout accumulates.

upsampler_channelstuple of int or None= None

When set, a second diffusion model is built that upsamples the denoiser's output. CS:GO generates at 30x56 and upsamples five times to 150x280, because diffusing at the full resolution directly costs far more for the same detail.

upsampler_layerstuple of int or None= None

Blocks per resolution in that model.

upsampler_attn_depthstuple of int or None= None

Which of its resolutions attend.

upsampling_factorint= 1

How much the upsampler magnifies.

sigma_datafloat= 0.5

\sigma_{\text{data}}

, the data distribution's standard deviation, which sets where the preconditioners hand over between predicting noise and predicting the frame.

p_meanfloat= -0.4

Mean of the log-normal the training noise level is drawn from.

p_stdfloat= 1.2

Standard deviation of that log-normal.

sigma_offset_noisefloat= 0.3

Standard deviation of a per-channel offset added to the target alongside the usual noise. ⚠️ Not in the paper — it appears only in the released config/agent/default.yaml. Offset noise lets the model shift a frame's overall level, which plain isotropic noise cannot express at any single

\sigma

attention_head_dimint= 8

Channels per self-attention head in the blocks that have attention. ⚠️ Also absent from the paper; taken from the released implementation.

denoise_stepsint= 3

Denoising steps per imagined frame. The paper's ablation is explicit that 1 is cheaper and worse on multi-modal games.

reward_channelstuple of int= (32, 32, 32, 32)

Residual-block channels in the reward/termination model.

reward_layerstuple of int= (2, 2, 2, 2)

Residual blocks per resolution there.

reward_cond_dimint= 128

Conditioning width for that model's adaptive group norms.

reward_framesint= 2

Frames the reward/termination model reads at once. ⚠️ The paper says "a sequence of frames and actions"; the released encoder takes two — the observation and the one it led to — because a reward is a property of the transition, not of a frame.

reward_lstm_dimint= 512

Hidden width of its LSTM cell.

actor_channelstuple of int= (32, 32, 64, 64)

Residual-block channels in the shared actor-critic trunk.

actor_layerstuple of int= (1, 1, 1, 1)

Residual blocks per resolution there.

actor_lstm_dimint= 512

Hidden width of the actor-critic's LSTM cell.

horizonint= 15

H

, the imagination horizon the agent is trained over.

gammafloat= 0.985

Discount factor.

lambda_float= 0.95

\lambda

for the returns the value network regresses to.

entropy_weightfloat= 0.001

\eta

, the weight on the policy's entropy bonus.

burn_inint= 4

Steps of real experience replayed to initialise the LSTM states before imagination starts. The paper sets this to

L

Notes

Reference: Alonso et al., "Diffusion for World Modeling: Visual Details Matter in Atari", NeurIPS, 2024 (arXiv:2405.12399). Architecture values are Table 2, training values Table 3, and the preconditioner constants Appendix C.

The base is GenerativeModelConfig rather than WorldModelConfig, even though this is a world model. That base describes a latent one — stoch_size, deter_size, free_nats, kl_weight — and DIAMOND has no posterior and no KL to free-bit. Inheriting them would put six fields on the documentation page that nothing in this family reads.

Examples

>>> from lucid.models.generative.diamond import DIAMONDConfig
>>> config = DIAMONDConfig()
>>> config.conditioning_frames, config.denoise_steps
(4, 3)
The denoiser's first convolution has to take the noised frame *and*
the stack of past frames, which is what makes frame stacking cheap:
>>> config.denoiser_in_channels
15
Three fields here have no counterpart in the paper — they come from
the released configuration, which is more specific than the text:
>>> config.sigma_offset_noise, config.reward_frames
(0.3, 2)

Used by 3

Constructors

dunder

init

→None

__init__(sample_size: int | tuple[int, int] = 64, in_channels: int = 3, out_channels: int = 3, act_fn: GenerativeActivation = 'silu', num_actions: int = 18, conditioning_frames: int = 4, unet_channels: tuple[int, ...] = (64, 64, 64, 64), unet_layers: tuple[int, ...] = (2, 2, 2, 2), cond_dim: int = 256, attn_depths: tuple[int, ...] | None = None, with_agent: bool = True, noise_previous_obs: bool = False, upsampler_channels: tuple[int, ...] | None = None, upsampler_layers: tuple[int, ...] | None = None, upsampler_attn_depths: tuple[int, ...] | None = None, upsampling_factor: int = 1, sigma_data: float = 0.5, p_mean: float = -0.4, p_std: float = 1.2, sigma_offset_noise: float = 0.3, attention_head_dim: int = 8, denoise_steps: int = 3, reward_channels: tuple[int, ...] = (32, 32, 32, 32), reward_layers: tuple[int, ...] = (2, 2, 2, 2), reward_cond_dim: int = 128, reward_frames: int = 2, reward_lstm_dim: int = 512, actor_channels: tuple[int, ...] = (32, 32, 64, 64), actor_layers: tuple[int, ...] = (1, 1, 1, 1), actor_lstm_dim: int = 512, horizon: int = 15, gamma: float = 0.985, lambda_: float = 0.95, entropy_weight: float = 0.001, burn_in: int = 4)

source edit

Initialise the layer.

Parameters

channelsint

Channels being normalised.

Properties

prop

denoiser_in_channels

→int

denoiser_in_channels: int

source edit

Channels the denoiser's first convolution reads.

The noised next frame plus $L$ clean past frames, stacked.

prop

frame_shape

→tuple[int, int]

frame_shape: tuple[int, int]

source edit

Frame height and width, whichever way sample_size was written.

Atari's frames are square and CS:GO's are not — 150 by 280 — so nothing downstream may assume one number.

DIAMONDConfig(sample_size: int | tuple[int, int] = 64, in_channels: int = 3, out_channels: int = 3, act_fn: GenerativeActivation = 'silu', num_actions: int = 18, conditioning_frames: int = 4, unet_channels: tuple[int, ...] = (64, 64, 64, 64), unet_layers: tuple[int, ...] = (2, 2, 2, 2), cond_dim: int = 256, attn_depths: tuple[int, ...] | None = None, with_agent: bool = True, noise_previous_obs: bool = False, upsampler_channels: tuple[int, ...] | None = None, upsampler_layers: tuple[int, ...] | None = None, upsampler_attn_depths: tuple[int, ...] | None = None, upsampling_factor: int = 1, sigma_data: float = 0.5, p_mean: float = -0.4, p_std: float = 1.2, sigma_offset_noise: float = 0.3, attention_head_dim: int = 8, denoise_steps: int = 3, reward_channels: tuple[int, ...] = (32, 32, 32, 32), reward_layers: tuple[int, ...] = (2, 2, 2, 2), reward_cond_dim: int = 128, reward_frames: int = 2, reward_lstm_dim: int = 512, actor_channels: tuple[int, ...] = (32, 32, 64, 64), actor_layers: tuple[int, ...] = (1, 1, 1, 1), actor_lstm_dim: int = 512, horizon: int = 15, gamma: float = 0.985, lambda_: float = 0.95, entropy_weight: float = 0.001, burn_in: int = 4)

>>> from lucid.models.generative.diamond import DIAMONDConfig >>> config = DIAMONDConfig() >>> config.conditioning_frames, config.denoise_steps (4, 3) The denoiser's first convolution has to take the noised frame *and* the stack of past frames, which is what makes frame stacking cheap: >>> config.denoiser_in_channels 15 Three fields here have no counterpart in the paper — they come from the released configuration, which is more specific than the text: >>> config.sigma_offset_noise, config.reward_frames (0.3, 2)

__init__(sample_size: int | tuple[int, int] = 64, in_channels: int = 3, out_channels: int = 3, act_fn: GenerativeActivation = 'silu', num_actions: int = 18, conditioning_frames: int = 4, unet_channels: tuple[int, ...] = (64, 64, 64, 64), unet_layers: tuple[int, ...] = (2, 2, 2, 2), cond_dim: int = 256, attn_depths: tuple[int, ...] | None = None, with_agent: bool = True, noise_previous_obs: bool = False, upsampler_channels: tuple[int, ...] | None = None, upsampler_layers: tuple[int, ...] | None = None, upsampler_attn_depths: tuple[int, ...] | None = None, upsampling_factor: int = 1, sigma_data: float = 0.5, p_mean: float = -0.4, p_std: float = 1.2, sigma_offset_noise: float = 0.3, attention_head_dim: int = 8, denoise_steps: int = 3, reward_channels: tuple[int, ...] = (32, 32, 32, 32), reward_layers: tuple[int, ...] = (2, 2, 2, 2), reward_cond_dim: int = 128, reward_frames: int = 2, reward_lstm_dim: int = 512, actor_channels: tuple[int, ...] = (32, 32, 64, 64), actor_layers: tuple[int, ...] = (1, 1, 1, 1), actor_lstm_dim: int = 512, horizon: int = 15, gamma: float = 0.985, lambda_: float = 0.95, entropy_weight: float = 0.001, burn_in: int = 4)