data

DIAMONDBehaviorOutput

extendsModelOutput

DIAMONDBehaviorOutput(policy_loss: Tensor, value_loss: Tensor, entropy: Tensor, returns: Tensor, frames: Tensor, history: Tensor, history_actions: Tensor)

source edit

What a pass of imagination returns.

Attributes

policy_loss, value_lossTensor

REINFORCE with a baseline, and the squared error against

\lambda

-returns. Scalars.

entropyTensor

Mean policy entropy over the imagined trajectory, before the weight

\eta

is applied. Scalar.

returnsTensor

The

\lambda

-returns the value head regressed to, (B, H).

framesTensor

Every frame the denoiser imagined, (B, H, C, H_img, W_img).

historyTensor

The conditioning the rollout ended on, (B, L, C, H_img, W_img) — what a continuation would start from. Once the horizon reaches

L

this is entirely imagined, which is the clearest statement of what "training in imagination" means.

history_actionsTensor

The actions beside it, (B, L).

Examples

>>> import lucid
>>> from lucid.models.generative.diamond._model import (
...     DIAMONDBehaviorOutput,
... )
>>> zero = lucid.zeros(())
>>> out = DIAMONDBehaviorOutput(
...     policy_loss=zero,
...     value_loss=zero,
...     entropy=zero,
...     returns=lucid.zeros(1, 4),
...     frames=lucid.zeros(1, 4, 3, 64, 64),
...     history=lucid.zeros(1, 4, 3, 64, 64),
...     history_actions=lucid.zeros(1, 4),
... )
>>> out.frames.shape
(1, 4, 3, 64, 64)
The imagined rollout comes back with the losses, not instead of them:
the frames the policy was trained against are the only way to see
what the world model actually showed it.

Used by 1

lucid.models.generative.diamond

Constructors

dunder

init

→None

__init__(policy_loss: Tensor, value_loss: Tensor, entropy: Tensor, returns: Tensor, frames: Tensor, history: Tensor, history_actions: Tensor)

source edit

Initialise the layer.

Parameters

channelsint

Channels being normalised.

>>> import lucid >>> from lucid.models.generative.diamond._model import ( ... DIAMONDBehaviorOutput, ... ) >>> zero = lucid.zeros(()) >>> out = DIAMONDBehaviorOutput( ... policy_loss=zero, ... value_loss=zero, ... entropy=zero, ... returns=lucid.zeros(1, 4), ... frames=lucid.zeros(1, 4, 3, 64, 64), ... history=lucid.zeros(1, 4, 3, 64, 64), ... history_actions=lucid.zeros(1, 4), ... ) >>> out.frames.shape (1, 4, 3, 64, 64) The imagined rollout comes back with the losses, not instead of them: the frames the policy was trained against are the only way to see what the world model actually showed it.