data

DreamerV2Config

extendsWorldModelConfig

DreamerV2Config(sample_size: int | tuple[int, int] = WORLD_MODEL_IMAGE_SIZE, in_channels: int = 3, out_channels: int = 3, act_fn: GenerativeActivation = 'elu', action_dim: int = 1, stoch_size: int = 32, deter_size: int = 1024, hidden_size: int = 1024, cnn_depth: int = 48, min_std: float = 0.1, mean_only: bool = False, free_nats: float = 0.0, kl_weight: float = 1.0, discrete: int = 32, kl_balance: float = 0.8, reward_hidden: int = 400, reward_layers: int = 4, actor_hidden: int = 400, actor_layers: int = 4, value_hidden: int = 400, value_layers: int = 4, horizon: int = 15, discount: float = 0.99, lambda_: float = 0.95, action_space: ActionSpace = 'continuous', actor_grad: ActorGradient = 'auto', actor_grad_mix: float = 0.1, actor_entropy: float = 0.002, actor_min_std: float = 0.1, slow_target_update: int = 100, slow_target_fraction: float = 1.0, pcont: bool = True, pcont_scale: float = 1.0, pcont_layers: int = 4)

source edit

Frozen configuration for the DreamerV2 family.

Defaults follow the released implementation's defaults block. Fields shared with PlaNet and Dreamer are inherited from WorldModelConfig, and the ones this paper changes are re-declared here.

Parameters

act_fn(silu, swish, relu, gelu, elu)= "silu"

Activation throughout, as in Dreamer.

stoch_sizeint= 32

Number of categorical variables. Note the change of meaning: in the Gaussian families this is the width of the latent vector, here it is how many independent categoricals make it up.

deter_sizeint= 1024, 1024

Deterministic path and the RSSM heads' hidden width.

hidden_sizeint= 1024, 1024

Deterministic path and the RSSM heads' hidden width.

cnn_depthint= 48

First encoder convolution's width; the schedule is Dreamer's.

free_natsfloat= 0.0

Off by default — KL balancing does the job free nats was doing, and the released implementation sets it to zero.

kl_weightfloat= 1.0

Multiplier on the balanced divergence.

discreteint= 32

Classes per categorical variable. The stochastic state is therefore stoch_size * discrete wide — 1024 by default, of which exactly stoch_size entries are 1.

kl_balancefloat= 0.8

\alpha

— the share of the divergence that pulls the prior toward the posterior rather than the reverse.

reward_hiddenint= 400, 4

Reward head.

reward_layersint= 400, 4

Reward head.

actor_hiddenint= 400, 4

Actor.

actor_layersint= 400, 4

Actor.

value_hiddenint= 400, 4

Critic.

value_layersint= 400, 4

Critic.

horizonint= 15

Imagination horizon.

discountfloat= 0.99

Reward discount

\gamma

lambda_float= 0.95

TD(

\lambda

) weighting.

action_space(continuous, discrete)= "continuous"

What an action is. "continuous" gives a truncated Gaussian over [-1, 1]^action_dim, which is the Control Suite; "discrete" gives a one-hot over action_dim alternatives, which is Atari. It also decides how the actor's gradient is estimated when actor_grad is left on "auto".

actor_grad(auto, dynamics, reinforce, both)= "auto"

How the actor's gradient is estimated. "auto" resolves as the released implementation does: the analytic path for a continuous action, the score function for a discrete one, where a one-hot offers nothing to differentiate through except a biased straight-through estimate.

actor_grad_mixfloat= 0.1

Weight on the dynamics term when actor_grad="both".

actor_entropyfloat= 2e-3

Coefficient on the actor's entropy bonus.

actor_min_stdfloat= 0.1

Floor on the action distribution's scale.

slow_target_updateint= 100

Gradient steps between refreshes of the critic's target copy.

slow_target_fractionfloat= 1.0

Fraction of the way the target moves at each refresh; 1 is a hard copy.

pcontbool= True

Predict the discount. Unlike Dreamer this is on by default — the paper's headline domain is Atari, where episodes end.

pcont_scalefloat= 1.0

Multiplier on the discount head's likelihood.

pcont_layersint= 4

Depth of the discount head; width follows value_hidden.

Notes

Reference: Hafner, Lillicrap, Norouzi, and Ba, "Mastering Atari with Discrete World Models", ICLR, 2021 (arXiv:2010.02193).

The paper tunes two configurations, and the released implementation keeps both as overrides on the defaults used here:

================== ========== ========== ============ field default Atari DMC (vision) ================== ========== ========== ============ deter_size 1024 600 200 hidden_size 1024 600 200 kl_weight 1.0 0.1 1.0 discount 0.99 0.999 0.99 actor_entropy 2e-3 1e-3 1e-4 pcont True True False pcont_scale 1.0 5.0 1.0 ================== ========== ========== ============

Learning rates are not config fields — no model config in this zoo carries them — but for the record the paper's are 2e-4 (world), 4e-5 (actor) and 1e-4 (critic) for Atari, against the released defaults of 1e-4 / 8e-5 / 2e-4.

stoch_size changing meaning between this family and the Gaussian ones is deliberate: it is the same slot in the shared config — how much stochastic state there is — and re-naming it would break the inheritance for the sake of a word. What it costs is that latent_size is no longer deter_size + stoch_size, which is why this class overrides it.

Examples

>>> from lucid.models.generative.dreamer_v2 import DreamerV2Config
>>> cfg = DreamerV2Config(action_dim=6)
>>> cfg.stoch_size, cfg.discrete, cfg.stoch_width
(32, 32, 1024)
>>> cfg.latent_size
2048
>>> cfg.kl_balance
0.8

Used by 3

Constructors

dunder

init

→None

__init__(sample_size: int | tuple[int, int] = WORLD_MODEL_IMAGE_SIZE, in_channels: int = 3, out_channels: int = 3, act_fn: GenerativeActivation = 'elu', action_dim: int = 1, stoch_size: int = 32, deter_size: int = 1024, hidden_size: int = 1024, cnn_depth: int = 48, min_std: float = 0.1, mean_only: bool = False, free_nats: float = 0.0, kl_weight: float = 1.0, discrete: int = 32, kl_balance: float = 0.8, reward_hidden: int = 400, reward_layers: int = 4, actor_hidden: int = 400, actor_layers: int = 4, value_hidden: int = 400, value_layers: int = 4, horizon: int = 15, discount: float = 0.99, lambda_: float = 0.95, action_space: ActionSpace = 'continuous', actor_grad: ActorGradient = 'auto', actor_grad_mix: float = 0.1, actor_entropy: float = 0.002, actor_min_std: float = 0.1, slow_target_update: int = 100, slow_target_fraction: float = 1.0, pcont: bool = True, pcont_scale: float = 1.0, pcont_layers: int = 4)

source edit

Properties

prop

latent_size

→int

latent_size: int

source edit

Width of [h; s], with s the flattened categorical grid.

prop

resolved_actor_grad

→{dynamics, reinforce, both}

resolved_actor_grad: ActorGradient

source edit

actor_grad with "auto" settled against the action space.

Returns

{dynamics, reinforce, both}

"reinforce" for a discrete action space, "dynamics" for a continuous one, when actor_grad is "auto"; otherwise whatever was asked for.

prop

stoch_width

→int

stoch_width: int

source edit

Flattened width of the stochastic latent — stoch_size * discrete.

DreamerV2Config(sample_size: int | tuple[int, int] = WORLD_MODEL_IMAGE_SIZE, in_channels: int = 3, out_channels: int = 3, act_fn: GenerativeActivation = 'elu', action_dim: int = 1, stoch_size: int = 32, deter_size: int = 1024, hidden_size: int = 1024, cnn_depth: int = 48, min_std: float = 0.1, mean_only: bool = False, free_nats: float = 0.0, kl_weight: float = 1.0, discrete: int = 32, kl_balance: float = 0.8, reward_hidden: int = 400, reward_layers: int = 4, actor_hidden: int = 400, actor_layers: int = 4, value_hidden: int = 400, value_layers: int = 4, horizon: int = 15, discount: float = 0.99, lambda_: float = 0.95, action_space: ActionSpace = 'continuous', actor_grad: ActorGradient = 'auto', actor_grad_mix: float = 0.1, actor_entropy: float = 0.002, actor_min_std: float = 0.1, slow_target_update: int = 100, slow_target_fraction: float = 1.0, pcont: bool = True, pcont_scale: float = 1.0, pcont_layers: int = 4)

>>> from lucid.models.generative.dreamer_v2 import DreamerV2Config >>> cfg = DreamerV2Config(action_dim=6) >>> cfg.stoch_size, cfg.discrete, cfg.stoch_width (32, 32, 1024) >>> cfg.latent_size 2048 >>> cfg.kl_balance 0.8

__init__(sample_size: int | tuple[int, int] = WORLD_MODEL_IMAGE_SIZE, in_channels: int = 3, out_channels: int = 3, act_fn: GenerativeActivation = 'elu', action_dim: int = 1, stoch_size: int = 32, deter_size: int = 1024, hidden_size: int = 1024, cnn_depth: int = 48, min_std: float = 0.1, mean_only: bool = False, free_nats: float = 0.0, kl_weight: float = 1.0, discrete: int = 32, kl_balance: float = 0.8, reward_hidden: int = 400, reward_layers: int = 4, actor_hidden: int = 400, actor_layers: int = 4, value_hidden: int = 400, value_layers: int = 4, horizon: int = 15, discount: float = 0.99, lambda_: float = 0.95, action_space: ActionSpace = 'continuous', actor_grad: ActorGradient = 'auto', actor_grad_mix: float = 0.1, actor_entropy: float = 0.002, actor_min_std: float = 0.1, slow_target_update: int = 100, slow_target_fraction: float = 1.0, pcont: bool = True, pcont_scale: float = 1.0, pcont_layers: int = 4)