diamond_world_model(pretrained: bool | str = False, weights: DIAMONDWeights | None = None, overrides: object = {})Construct DIAMOND with its objectives and imagination rollout.
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
DIAMONDForWorldModelingThe agent, plus the losses that train it and the rollout that generates its experience.
Notes
Reference: Alonso et al., arXiv:2405.12399, Appendix F for the two RL objectives and Algorithm 1 for the training loop they sit in.
The three losses are trained separately, in the order Algorithm 1 gives: the denoiser on real transitions, the reward/termination model on real sequences, and the actor-critic entirely on imagined ones.
Examples
>>> import lucid
>>> from lucid.models import diamond_world_model
>>> model = diamond_world_model(
... sample_size=16, unet_channels=(8, 8), unet_layers=(1, 1),
... reward_channels=(8, 8), reward_layers=(1, 1),
... actor_channels=(8, 8), actor_layers=(1, 1), cond_dim=16,
... reward_cond_dim=8, reward_lstm_dim=16, actor_lstm_dim=16,
... num_actions=4, horizon=3).eval()
>>> frames = lucid.randn((1, 4, 3, 16, 16))
>>> actions = lucid.tensor([[0, 1, 2, 3]], dtype=lucid.int64)
>>> with lucid.no_grad():
... out = model(frames, actions)
>>> out.frames.shape, out.returns.shape
((1, 3, 3, 16, 16), (1, 3))