class
DIAMONDForWorldModeling
extends
WorldModelingModelDIAMONDForWorldModeling(config: DIAMONDConfig)DIAMOND posed as a world model: imagination and its objectives.
Parameters
configDIAMONDConfigFrozen configuration.
Attributes
diamondDIAMONDModelThe three networks.
Notes
Reference: Alonso et al., arXiv:2405.12399, Appendix F and Algorithm 1.
The imagined trajectory is autoregressive in the strong sense: the frame the denoiser produced and the action the policy chose from it both become conditioning for the next step, so an error in either compounds. That is the whole reason the paper cares which diffusion parameterisation it uses.
Examples
>>> import lucid
>>> from lucid.models.generative.diamond import (
... DIAMONDConfig, DIAMONDForWorldModeling)
>>> config = DIAMONDConfig(
... 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)
>>> model = DIAMONDForWorldModeling(config).eval()
>>> frames = lucid.randn((2, 4, 3, 16, 16))
>>> actions = lucid.tensor([[0, 1, 2, 3], [1, 1, 0, 2]], dtype=lucid.int64)
>>> with lucid.no_grad():
... out = model(frames, actions)
>>> out.frames.shape
(2, 3, 3, 16, 16)Used by 2
Constructors
1Instance methods
4act(frame: Tensor, state: tuple[Tensor, Tensor] | None = None, return_state: bool = False)Sample an action from the policy.
Parameters
frameTensor(B, C, H, W).Carried actor-critic LSTM
(hidden, cell). None reads
frame from a blank memory, as at the start of an episode.return_statebool= False, keyword-onlyAlso return the LSTM state after reading
frame. The policy
is recurrent, so an agent loop must pass this back as the next
call's state; without it every frame is read from a blank
memory and the policy remembers nothing.Returns
Examples
>>> import lucid
>>> from lucid.models.generative.diamond import (
... DIAMONDConfig, DIAMONDForWorldModeling)
>>> config = DIAMONDConfig(
... 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)
>>> model = DIAMONDForWorldModeling(config).eval()
>>> with lucid.no_grad():
... action = model.act(lucid.randn((2, 3, 16, 16)))
>>> action.shape, action.dtype
((2,), lucid.int64)
A draw from the policy rather than its argmax, but always a valid
index into the action set:
>>> bool(action.min() >= 0), bool(action.max() < config.num_actions)
(True, True)
The policy is recurrent, so an agent loop asks for the state back
and hands it to the next call — that is what gives it a memory:
>>> frames = lucid.randn((2, 5, 3, 16, 16))
>>> state = None
>>> with lucid.no_grad():
... for t in range(5):
... action, state = model.act(frames[:, t], state, return_state=True)
... _, blank = model.act(frames[:, -1], return_state=True)
>>> action.shape, [s.shape for s in state]
((2,), [(2, 16), (2, 16)])
The last frame, read after the four before it, leaves a different
memory than the same frame read cold:
>>> bool((state[0] == blank[0]).all())
Falseforward(frames: Tensor, actions: Tensor, horizon: int | None = None)Imagine a trajectory and score the policy on it.
Parameters
Returns
DIAMONDBehaviorOutputThe two losses, the policy's entropy, the -returns and every imagined frame.
Notes
Reference: Alonso et al., arXiv:2405.12399, equations 14-16. The value target stops gradients, and the policy is REINFORCE against — an advantage estimate, but one built from a state-value baseline rather than a critic.
reward_end_loss(frames: Tensor, actions: Tensor, rewards: Tensor, ends: Tensor)Cross-entropy on the reward's sign and on termination.
Parameters
framesTensor(B, T, C, H, W), with T >= 2.actionsTensor(B, T); actions[:, t] is the action taken at frame t.rewardsTensor(B, T) real rewards; rewards[:, t] is the reward for the
step from frame t to frame t + 1. Only their sign is
predicted — Algorithm 1 writes CE(r_hat, sign(r)), which is
all the environment's clipping to leaves.endsTensor(B, T) termination flags in {0, 1}, aligned the same
way: ends[:, t] marks the step out of frame t.Returns
TensorScalar, the two cross-entropies summed, averaged over the
T - 1 transitions.
Notes
The reward model reads transitions, so T frames give
T - 1 of them: step t pairs frames[:, t] with
frames[:, t + 1] and actions[:, t], and scores
rewards[:, t] and ends[:, t]. The last column of
actions, rewards and ends has no next frame to pair
with and is ignored — so each column must hold what the step out
of its frame earned, not what was observed on arriving there.
Examples
>>> import lucid
>>> from lucid.models.generative.diamond import (
... DIAMONDConfig, DIAMONDForWorldModeling)
>>> config = DIAMONDConfig(
... 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)
>>> model = DIAMONDForWorldModeling(config)
>>> frames = lucid.randn((2, 4, 3, 16, 16))
>>> actions = lucid.tensor([[0, 1, 2, 3], [1, 1, 0, 2]], dtype=lucid.int64)
>>> rewards = lucid.tensor([[0.0, 1.0, -2.0, 0.0], [5.0, 0.0, 0.0, 0.0]])
>>> ends = lucid.tensor([[0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 0.0]])
>>> loss = model.reward_end_loss(frames, actions, rewards, ends)
>>> loss.shape
()
Only the sign of a reward is a target, so rescaling the rewards
leaves the loss exactly where it was:
>>> scaled = model.reward_end_loss(frames, actions, rewards * 10.0, ends)
>>> scaled.item() == loss.item()
True
The last column pairs with no next frame, so it is never scored:
>>> moved = lucid.cat([rewards[:, :-1], lucid.tensor([[-1.0], [1.0]])], dim=1)
>>> model.reward_end_loss(frames, actions, moved, ends).item() == loss.item()
Trueworld_model_loss(frames: Tensor, actions: Tensor, next_frame: Tensor)Train the denoiser — DIAMONDModel.forward by another name.
Parameters
Returns
DIAMONDOutputLoss, denoised frame, and the noise levels used.
Examples
>>> import lucid
>>> from lucid.models.generative.diamond import (
... DIAMONDConfig, DIAMONDForWorldModeling)
>>> config = DIAMONDConfig(
... 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)
>>> model = DIAMONDForWorldModeling(config)
>>> frames = lucid.randn((2, 4, 3, 16, 16))
>>> actions = lucid.tensor([[0, 1, 2, 3], [1, 1, 0, 2]], dtype=lucid.int64)
>>> target = lucid.randn((2, 3, 16, 16))
>>> out = model.world_model_loss(frames, actions, target)
>>> out.loss.shape, out.prediction.shape, out.sigma.shape
((), (2, 3, 16, 16), (2,))
Only the denoiser learns from it — the reward model and the
actor-critic are trained by losses of their own:
>>> out.loss.backward()
>>> any(p.grad is not None for p in model.diamond.denoiser.parameters())
True
>>> any(p.grad is not None for p in model.diamond.reward_end.parameters())
False