xwm.heads

Prediction heads: categorical reward and value, pessimistic Q-ensembles, squashed-Gaussian policies, and the two-hot/symlog machinery they share.

Prediction heads: reward, value, policy, Q.

The pieces a reward-driven world model needs and a self-supervised one does not. JEPA gets by with an encoder and a predictor; TD-MPC2 and MuZero also have to say how good a state is and what to do in it.

Modules:

Name	Description
categorical	Scalar regression as classification: two-hot targets over a fixed bin grid.
policy	A squashed-Gaussian policy head.
q_ensemble	An ensemble of Q-functions, with a pessimistic aggregate.
scalar	An MLP trunk with a categorical scalar output -- rewards and values.

Classes:

Name	Description
CategoricalScalar	A fixed grid of bins for encoding and decoding scalars.
GaussianPolicy	z -> tanh(Normal(mu(z), sigma(z))), bounded in [-1, 1].
PolicyOutput	A sampled action with the statistics needed for a SAC-style update.
QEnsemble	n_members action-conditioned scalar heads over a shared latent.
ScalarHead	Predict a scalar from a latent (and optionally an action).

Functions:

Name	Description
cross_entropy	Mean cross-entropy against a (possibly soft) target distribution.
symexp	Inverse of symlog.
symlog	sign(x) * log(|x| + 1) -- invertible, squashes large magnitudes.
two_hot	Two-hot encode x over n_bins uniform bins spanning [low, high].

CategoricalScalar

CategoricalScalar(n_bins: int = 101, low: float = -10.0, high: float = 10.0, *, transform: bool = True)

Bases: Module

A fixed grid of bins for encoding and decoding scalars.

Parameters:

Name	Type	Description	Default
n_bins	int	number of bins. More bins means finer resolution and a harder classification problem; 101 is TD-MPC2's choice.	101
low, high		range covered. Values outside are clipped, so pick a range that actually contains your returns.	required
transform	bool	apply symlog before binning and symexp after decoding. Use it when values span orders of magnitude.	True

Methods:

Name	Description
encode	Scalar(s) -> two-hot distribution (..., n_bins).
decode	Logits (..., n_bins) -> expected scalar under the softmax.
loss	Cross-entropy between predicted logits and the two-hot target.

Attributes:

Name	Type	Description
bins	Array	(n_bins,) bin centres, in transformed space.

Source code in xwm/heads/categorical.py

def __init__(
    self,
    n_bins: int = 101,
    low: float = -10.0,
    high: float = 10.0,
    *,
    transform: bool = True,
):
    if n_bins < 2:
        raise ValueError(f"need at least two bins, got {n_bins}")
    if not high > low:
        raise ValueError(f"high ({high}) must exceed low ({low})")
    self.n_bins = n_bins
    self.low = low
    self.high = high
    self.transform = transform

bins

bins: Array

(n_bins,) bin centres, in transformed space.

encode

encode(value: Array) -> Array

Scalar(s) -> two-hot distribution (..., n_bins).

Source code in xwm/heads/categorical.py

def encode(self, value: Array) -> Array:
    """Scalar(s) -> two-hot distribution ``(..., n_bins)``."""
    x = symlog(value) if self.transform else value
    return two_hot(x, self.n_bins, self.low, self.high)

decode

decode(logits: Array) -> Array

Logits (..., n_bins) -> expected scalar under the softmax.

Source code in xwm/heads/categorical.py

def decode(self, logits: Array) -> Array:
    """Logits ``(..., n_bins)`` -> expected scalar under the softmax."""
    probs = jax.nn.softmax(logits, axis=-1)
    expected = jnp.sum(probs * self.bins, axis=-1)
    return symexp(expected) if self.transform else expected

loss

loss(logits: Array, target: Array) -> Array

Cross-entropy between predicted logits and the two-hot target.

Source code in xwm/heads/categorical.py

def loss(self, logits: Array, target: Array) -> Array:
    """Cross-entropy between predicted logits and the two-hot target."""
    return cross_entropy(logits, self.encode(target))

GaussianPolicy

GaussianPolicy(latent_dim: int, action_dim: int, *, key: PRNGKey | None = None, hidden_dim: int = 512, depth: int = 2)

Bases: Module

z -> tanh(Normal(mu(z), sigma(z))), bounded in [-1, 1].

Methods:

Name	Description
distribution	(mean, log_std) before squashing.
act	Sample an action, or return the deterministic mean when key is None.
sample	Sample with the log-probability of the squashed action.

Source code in xwm/heads/policy.py

def __init__(
    self,
    latent_dim: int,
    action_dim: int,
    *,
    key: PRNGKey | None = None,
    hidden_dim: int = 512,
    depth: int = 2,
):
    key = resolve_key(key)
    keys = jr.split(key, depth + 1)
    dims = [latent_dim] + [hidden_dim] * depth
    self.layers = [eqx.nn.Linear(dims[i], hidden_dim, key=keys[i]) for i in range(depth)]
    self.norms = [LayerNorm(hidden_dim) for _ in range(depth)]
    self.out = eqx.nn.Linear(hidden_dim, 2 * action_dim, key=keys[-1])
    self.latent_dim = latent_dim
    self.action_dim = action_dim

distribution

distribution(z: Array) -> tuple[Array, Array]

(mean, log_std) before squashing.

Source code in xwm/heads/policy.py

def distribution(self, z: Array) -> tuple[Array, Array]:
    """``(mean, log_std)`` before squashing."""
    h = z
    for layer, norm in zip(self.layers, self.norms, strict=True):
        h = jax.nn.mish(norm(layer(h)))
    mean, log_std = jnp.split(self.out(h), 2)
    low, high = LOG_STD_RANGE
    # Squash into range rather than hard-clipping, so the gradient survives.
    log_std = low + 0.5 * (high - low) * (jnp.tanh(log_std) + 1.0)
    return mean, log_std

act

act(z: Array, *, key: PRNGKey | None = None) -> Array

Sample an action, or return the deterministic mean when key is None.

Source code in xwm/heads/policy.py

def act(self, z: Array, *, key: PRNGKey | None = None) -> Array:
    """Sample an action, or return the deterministic mean when ``key`` is None."""
    mean, log_std = self.distribution(z)
    if key is None:
        return jnp.tanh(mean)
    return jnp.tanh(mean + jnp.exp(log_std) * jr.normal(key, mean.shape))

sample

sample(z: Array, key: PRNGKey) -> PolicyOutput

Sample with the log-probability of the squashed action.

Source code in xwm/heads/policy.py

def sample(self, z: Array, key: PRNGKey) -> PolicyOutput:
    """Sample with the log-probability of the squashed action."""
    mean, log_std = self.distribution(z)
    std = jnp.exp(log_std)
    noise = jr.normal(key, mean.shape)
    pre_tanh = mean + std * noise
    action = jnp.tanh(pre_tanh)
    # Gaussian log-prob, then the tanh change-of-variables correction.
    log_prob = jnp.sum(
        -0.5 * noise**2 - log_std - 0.5 * jnp.log(2.0 * jnp.pi)
    )
    log_prob -= jnp.sum(jnp.log(jnp.clip(1.0 - action**2, 1e-6, None)))
    return PolicyOutput(action=action, log_prob=log_prob, mean=mean, log_std=log_std)

PolicyOutput

Bases: NamedTuple

A sampled action with the statistics needed for a SAC-style update.

QEnsemble

QEnsemble(latent_dim: int, action_dim: int, *, key: PRNGKey | None = None, n_members: int = 5, subset_size: int = 2, hidden_dim: int = 512, depth: int = 2, scalar: CategoricalScalar | None = None)

Bases: Module

n_members action-conditioned scalar heads over a shared latent.

Parameters:

Name	Type	Description	Default
n_members	int	ensemble size.	5
subset_size	int	how many members the pessimistic aggregate draws from.	2

Methods:

Name	Description
logits	(n_members, n_bins).
values	(n_members,) decoded Q-values.
pessimistic	Minimum over a random subset of members -- the aggregate to plan with.
loss	Mean cross-entropy across members against a shared scalar target.

Source code in xwm/heads/q_ensemble.py

def __init__(
    self,
    latent_dim: int,
    action_dim: int,
    *,
    key: PRNGKey | None = None,
    n_members: int = 5,
    subset_size: int = 2,
    hidden_dim: int = 512,
    depth: int = 2,
    scalar: CategoricalScalar | None = None,
):
    key = resolve_key(key)
    if not 1 <= subset_size <= n_members:
        raise ValueError(f"subset_size {subset_size} must be in [1, {n_members}]")
    keys = jr.split(key, n_members)
    self.members = [
        ScalarHead(
            latent_dim,
            action_dim=action_dim,
            key=k,
            hidden_dim=hidden_dim,
            depth=depth,
            scalar=scalar,
        )
        for k in keys
    ]
    self.subset_size = subset_size

logits

logits(z: Array, action: Array) -> Array

(n_members, n_bins).

Source code in xwm/heads/q_ensemble.py

def logits(self, z: Array, action: Array) -> Array:
    """``(n_members, n_bins)``."""
    return jnp.stack([member.logits(z, action) for member in self.members])

values

values(z: Array, action: Array) -> Array

(n_members,) decoded Q-values.

Source code in xwm/heads/q_ensemble.py

def values(self, z: Array, action: Array) -> Array:
    """``(n_members,)`` decoded Q-values."""
    return jnp.stack([member.value(z, action) for member in self.members])

pessimistic

pessimistic(z: Array, action: Array, *, key: PRNGKey | None = None) -> Array

Minimum over a random subset of members -- the aggregate to plan with.

key is required during training (the subset must be resampled every step); pass None to use the min over all members, which is the deterministic choice for evaluation.

Source code in xwm/heads/q_ensemble.py

def pessimistic(self, z: Array, action: Array, *, key: PRNGKey | None = None) -> Array:
    """Minimum over a random subset of members -- the aggregate to plan with.

    ``key`` is required during training (the subset must be resampled every
    step); pass ``None`` to use the min over *all* members, which is the
    deterministic choice for evaluation.
    """
    values = self.values(z, action)
    if key is None:
        return jnp.min(values)
    chosen = jr.choice(key, self.n_members, (self.subset_size,), replace=False)
    return jnp.min(values[chosen])

loss

loss(z: Array, action: Array, target: Array) -> Array

Mean cross-entropy across members against a shared scalar target.

Source code in xwm/heads/q_ensemble.py

def loss(self, z: Array, action: Array, target: Array) -> Array:
    """Mean cross-entropy across members against a shared scalar target."""
    return jnp.mean(
        jnp.stack([member.loss(z, target, action) for member in self.members])
    )

ScalarHead

ScalarHead(latent_dim: int, *, action_dim: int = 0, key: PRNGKey | None = None, hidden_dim: int = 512, depth: int = 2, scalar: CategoricalScalar | None = None)

Bases: Module

Predict a scalar from a latent (and optionally an action).

Emits logits over bins rather than a number; call value to decode or loss to train. See xwm.heads.categorical for why that beats a squared error here.

Parameters:

Name	Type	Description	Default
latent_dim	int	width of the input latent.	required
action_dim	int	width of an action to condition on; 0 for a state-only head (a value function) rather than a state-action one (a reward model or a Q-function).	0
scalar	CategoricalScalar | None	the bin grid. Defaults to TD-MPC2's 101 bins with symlog.	None

Methods:

Name	Description
logits	(n_bins,) logits.
value	Decoded scalar.
loss	Cross-entropy against the two-hot encoding of target.

Source code in xwm/heads/scalar.py

def __init__(
    self,
    latent_dim: int,
    *,
    action_dim: int = 0,
    key: PRNGKey | None = None,
    hidden_dim: int = 512,
    depth: int = 2,
    scalar: CategoricalScalar | None = None,
):
    key = resolve_key(key)
    scalar = scalar or CategoricalScalar()
    keys = jr.split(key, depth + 1)
    dims = [latent_dim + action_dim] + [hidden_dim] * depth
    self.layers = [eqx.nn.Linear(dims[i], hidden_dim, key=keys[i]) for i in range(depth)]
    self.norms = [LayerNorm(hidden_dim) for _ in range(depth)]
    self.out = eqx.nn.Linear(hidden_dim, scalar.n_bins, key=keys[-1])
    self.scalar = scalar
    self.latent_dim = latent_dim
    self.action_dim = action_dim

logits

logits(z: Array, action: Array | None = None) -> Array

(n_bins,) logits.

Source code in xwm/heads/scalar.py

def logits(self, z: Array, action: Array | None = None) -> Array:
    """``(n_bins,)`` logits."""
    if self.action_dim:
        if action is None:
            raise ValueError("this head is action-conditioned; pass an action")
        h = jnp.concatenate([z, jnp.atleast_1d(action)])
    else:
        h = z
    for layer, norm in zip(self.layers, self.norms, strict=True):
        h = jax.nn.mish(norm(layer(h)))
    return self.out(h)

value

value(z: Array, action: Array | None = None) -> Array

Decoded scalar.

Source code in xwm/heads/scalar.py

def value(self, z: Array, action: Array | None = None) -> Array:
    """Decoded scalar."""
    return self.scalar.decode(self.logits(z, action))

loss

loss(z: Array, target: Array, action: Array | None = None) -> Array

Cross-entropy against the two-hot encoding of target.

Source code in xwm/heads/scalar.py

def loss(self, z: Array, target: Array, action: Array | None = None) -> Array:
    """Cross-entropy against the two-hot encoding of ``target``."""
    return self.scalar.loss(self.logits(z, action), target)

cross_entropy

cross_entropy(logits: Array, target_probs: Array) -> Array

Mean cross-entropy against a (possibly soft) target distribution.

Source code in xwm/heads/categorical.py

def cross_entropy(logits: Array, target_probs: Array) -> Array:
    """Mean cross-entropy against a (possibly soft) target distribution."""
    log_probs = jax.nn.log_softmax(logits, axis=-1)
    return -jnp.mean(jnp.sum(target_probs * log_probs, axis=-1))

symexp

symexp(x: Array) -> Array

Inverse of symlog.

Source code in xwm/heads/categorical.py

def symexp(x: Array) -> Array:
    """Inverse of :func:`symlog`."""
    return jnp.sign(x) * (jnp.expm1(jnp.abs(x)))

symlog

symlog(x: Array) -> Array

sign(x) * log(|x| + 1) -- invertible, squashes large magnitudes.

Source code in xwm/heads/categorical.py

def symlog(x: Array) -> Array:
    """``sign(x) * log(|x| + 1)`` -- invertible, squashes large magnitudes."""
    return jnp.sign(x) * jnp.log1p(jnp.abs(x))

two_hot

two_hot(x: Array, n_bins: int, low: float, high: float) -> Array

Two-hot encode x over n_bins uniform bins spanning [low, high].

The value's mass is split linearly between its two neighbouring bins, so sum(bins * two_hot(x)) == clip(x, low, high) exactly.

Source code in xwm/heads/categorical.py

def two_hot(x: Array, n_bins: int, low: float, high: float) -> Array:
    """Two-hot encode ``x`` over ``n_bins`` uniform bins spanning ``[low, high]``.

    The value's mass is split linearly between its two neighbouring bins, so
    ``sum(bins * two_hot(x)) == clip(x, low, high)`` exactly.
    """
    x = jnp.clip(x, low, high)
    width = (high - low) / (n_bins - 1)
    position = (x - low) / width
    lower = jnp.floor(position).astype(jnp.int32)
    upper = jnp.clip(lower + 1, 0, n_bins - 1)
    lower = jnp.clip(lower, 0, n_bins - 1)
    upper_weight = position - lower
    onehot = lambda idx: jax.nn.one_hot(idx, n_bins)  # noqa: E731
    return (
        onehot(lower) * (1.0 - upper_weight)[..., None]
        + onehot(upper) * upper_weight[..., None]
    )