xwm.dynamics

(z, a) → z', the centre of the library. Every family consumes a dynamics model from here; every planner consumes nothing else.

Action-conditioned latent dynamics -- the heart of a robotics world model.

(z, a) -> z' in latent space, shared by every family: JEPA's V-JEPA 2-AC stage, TD-MPC2's consistency loss, and MuZero's recurrent unroll all consume a model from here. Which one you pick is a compute/expressivity trade, not an algorithmic commitment:

• ActionConditionedPredictor -- a transformer over the token grid. Expressive, and the right choice when the latent is a sequence of patch tokens.
• MLPDynamics -- a residual MLP on a pooled latent. Far cheaper, and what TD-MPC2 and MuZero actually use, because they run it thousands of times inside a planner.

Modules:

Name	Description
action_conditioned	Action-conditioned latent dynamics.
action_embed	Action encoders.
mlp_dynamics	A residual MLP latent dynamics model.

Classes:

Name	Description
ActionConditionedPredictor	(z, a) -> z' over a token grid, satisfying LatentDynamics.
ContinuousActionEmbed	Embed a continuous action vector with a small MLP.
DiscreteActionEmbed	Embed a discrete action index with a lookup table.
PoseActionEmbed	Embed a rigid-body pose delta, splitting translation from rotation.
MLPDynamics	(z, a) -> z' as a residual MLP on a flat latent.

ActionConditionedPredictor

ActionConditionedPredictor(grid: tuple[int, ...], embed_dim: int, action_embed: ActionEmbed, *, pred_dim: int, depth: int, num_heads: int, key: PRNGKey | None = None, conditioning: Conditioning = 'both', residual: bool = True, pos: Literal['sincos', 'learned', 'none'] = 'sincos', mlp_ratio: float = 4.0, dropout: float = 0.0, drop_path: float = 0.0, layer_scale: float | None = None, remat: bool = False)

Bases: Module

(z, a) -> z' over a token grid, satisfying LatentDynamics.

Parameters:

Name	Type	Description	Default
grid	tuple[int, ...]	token grid of the latent state.	required
embed_dim	int	encoder width (input and output).	required
action_embed	ActionEmbed	how the action becomes a vector.	required
pred_dim	int	internal width.	required
depth, num_heads		transformer size.	required
conditioning	Conditioning	how the action reaches the tokens. "token" -- prepend the action as an extra token and let attention route it. Most expressive, but slowest to get going: attention has to learn to read the token before the action influences anything. "film" -- feature-wise scale and shift applied to every token. Cheapest, and reaches every token directly at depth zero. "both" -- the default. All three converge to similar action sensitivity, and this one gets there first.	'both'
residual	bool	predict the change in latent state rather than the next state outright. On by default: consecutive latents are nearly identical, so starting close to the identity map is a far better prior than starting from noise. The output projection is scaled down by RESIDUAL_INIT_SCALE in this mode.	True
pos	Literal['sincos', 'learned', 'none']	positional scheme (additive only -- the action token has no grid position, so rotary embeddings do not apply here).	'sincos'

Source code in xwm/dynamics/action_conditioned.py

def __init__(
    self,
    grid: tuple[int, ...],
    embed_dim: int,
    action_embed: ActionEmbed,
    *,
    pred_dim: int,
    depth: int,
    num_heads: int,
    key: PRNGKey | None = None,
    conditioning: Conditioning = "both",
    residual: bool = True,
    pos: Literal["sincos", "learned", "none"] = "sincos",
    mlp_ratio: float = 4.0,
    dropout: float = 0.0,
    drop_path: float = 0.0,
    layer_scale: float | None = None,
    remat: bool = False,
):
    key = resolve_key(key)
    k_in, k_out, k_tok, k_film, k_pos, k_blocks = jr.split(key, 6)
    n_tokens = 1
    for s in grid:
        n_tokens *= s
    self.action_embed = action_embed
    self.embed_in = eqx.nn.Linear(embed_dim, pred_dim, key=k_in)
    embed_out = eqx.nn.Linear(pred_dim, embed_dim, key=k_out)
    if residual:
        # Shrink the output projection so the residual branch starts as a
        # small perturbation of the identity. Deliberately *small* and not
        # zero: a zero projection would also zero the gradient flowing back
        # through it, stalling every upstream parameter -- the transformer,
        # the action embedding, the FiLM layer -- until the projection
        # itself moved off zero.
        embed_out = eqx.tree_at(
            lambda m: (m.weight, m.bias),
            embed_out,
            (RESIDUAL_INIT_SCALE * embed_out.weight, RESIDUAL_INIT_SCALE * embed_out.bias),
        )
    self.embed_out = embed_out
    if pos == "sincos":
        self.pos_embed = SinCosPosEmbed(tuple(grid), pred_dim)
    elif pos == "learned":
        self.pos_embed = LearnedPosEmbed(n_tokens, pred_dim, key=k_pos)
    else:
        self.pos_embed = None
    uses_token = conditioning in ("token", "both")
    uses_film = conditioning in ("film", "both")
    self.action_token = (
        eqx.nn.Linear(action_embed.embed_dim, pred_dim, key=k_tok) if uses_token else None
    )
    # Small-but-nonzero weights with a zero bias: FiLM starts *near* the
    # identity (scale ~ 1, shift ~ 0) without starting *at* it. A fully
    # zero-initialised FiLM would make the dynamics exactly independent of
    # the action at step zero, which is the wrong place to begin for a model
    # whose entire job is to be action-conditioned.
    self.film = (
        eqx.tree_at(
            lambda m: (m.weight, m.bias),
            eqx.nn.Linear(action_embed.embed_dim, 2 * pred_dim, key=k_film),
            (
                0.02 * jr.normal(k_film, (2 * pred_dim, action_embed.embed_dim)),
                jnp.zeros((2 * pred_dim,)),
            ),
        )
        if uses_film
        else None
    )
    self.blocks = Transformer(
        pred_dim,
        depth,
        num_heads,
        key=k_blocks,
        mlp_ratio=mlp_ratio,
        dropout=dropout,
        drop_path=drop_path,
        layer_scale=layer_scale,
        remat=remat,
    )
    self.grid = tuple(grid)
    self.embed_dim = embed_dim
    self.pred_dim = pred_dim
    self.conditioning = conditioning
    self.residual = residual

ContinuousActionEmbed

ContinuousActionEmbed(action_dim: int, embed_dim: int, *, key: PRNGKey | None = None, hidden_dim: int | None = None, scale: float = 1.0)

Bases: Module

Embed a continuous action vector with a small MLP.

Parameters:

Name	Type	Description	Default
action_dim	int	dimensionality of the raw action.	required
embed_dim	int	output width.	required
scale	float	divide the action by this before embedding. Set it to the action's typical magnitude; an unnormalised action of magnitude 100 will otherwise dominate the conditioning signal.	1.0

Source code in xwm/dynamics/action_embed.py

def __init__(
    self,
    action_dim: int,
    embed_dim: int,
    *,
    key: PRNGKey | None = None,
    hidden_dim: int | None = None,
    scale: float = 1.0,
):
    key = resolve_key(key)
    self.mlp = Mlp(action_dim, hidden_dim or 4 * embed_dim, embed_dim, key=key)
    self.scale = scale
    self.action_dim = action_dim
    self.embed_dim = embed_dim

DiscreteActionEmbed

DiscreteActionEmbed(n_actions: int, embed_dim: int, *, key: PRNGKey)

Bases: Module

Embed a discrete action index with a lookup table.

Source code in xwm/dynamics/action_embed.py

def __init__(self, n_actions: int, embed_dim: int, *, key: PRNGKey):
    self.table = 0.02 * jr.normal(key, (n_actions, embed_dim))
    self.n_actions = n_actions
    self.embed_dim = embed_dim

PoseActionEmbed

PoseActionEmbed(embed_dim: int, *, key: PRNGKey | None = None, translation_dim: int = 3, rotation_dim: int = 3, extra_dim: int = 0)

Bases: Module

Embed a rigid-body pose delta, splitting translation from rotation.

End-effector actions mix units -- metres and radians -- and a single linear layer has to learn to rescale them. Embedding the parts separately removes that burden, which matters when translations are centimetre-scale.

Parameters:

Name	Type	Description	Default
translation_dim	int	usually 3.	3
rotation_dim	int	3 for axis-angle / Euler, 4 for a quaternion, 6 for the continuous 6-D rotation parameterisation.	3
extra_dim	int	remaining scalars, e.g. a gripper command.	0

Source code in xwm/dynamics/action_embed.py

def __init__(
    self,
    embed_dim: int,
    *,
    key: PRNGKey | None = None,
    translation_dim: int = 3,
    rotation_dim: int = 3,
    extra_dim: int = 0,
):
    key = resolve_key(key)
    k1, k2, k3, k4 = jr.split(key, 4)
    half = embed_dim // 2
    self.translation = eqx.nn.Linear(translation_dim, half, key=k1)
    self.rotation = eqx.nn.Linear(rotation_dim, half, key=k2)
    self.extra = eqx.nn.Linear(extra_dim, half, key=k3) if extra_dim else None
    n_parts = 3 if extra_dim else 2
    self.out = eqx.nn.Linear(half * n_parts, embed_dim, key=k4)
    self.translation_dim = translation_dim
    self.rotation_dim = rotation_dim
    self.extra_dim = extra_dim
    self.embed_dim = embed_dim

MLPDynamics

MLPDynamics(latent_dim: int, action_dim: int, *, key: PRNGKey | None = None, hidden_dim: int = 512, depth: int = 2, residual: bool = True, normalize: Normalization = 'simnorm', simnorm_groups: int = 8)

Bases: Module

(z, a) -> z' as a residual MLP on a flat latent.

Parameters:

Name	Type	Description	Default
latent_dim	int	width of the latent vector.	required
action_dim	int	width of the action vector (continuous), or the embedding width if you embed a discrete action before calling.	required
hidden_dim	int	MLP width.	512
depth	int	number of hidden layers.	2
residual	bool	predict the change rather than the next state. Consecutive latents are nearly identical, so starting near the identity is a far better prior than starting from noise.	True
normalize	Normalization	what to apply to the output latent. "simnorm" -- TD-MPC2's simplicial normalisation. This is the load-bearing detail of TD-MPC2's stability: it bounds the latent without collapsing it, so the dynamics cannot drift off to infinity during a long unroll and cannot shrink to a point either. "layernorm" -- the usual choice elsewhere. "none" -- unbounded; expect drift over long rollouts.	'simnorm'
simnorm_groups	int	group size for SimNorm; latent_dim must divide by it.	8

Source code in xwm/dynamics/mlp_dynamics.py

def __init__(
    self,
    latent_dim: int,
    action_dim: int,
    *,
    key: PRNGKey | None = None,
    hidden_dim: int = 512,
    depth: int = 2,
    residual: bool = True,
    normalize: Normalization = "simnorm",
    simnorm_groups: int = 8,
):
    key = resolve_key(key)
    keys = jr.split(key, depth + 2)
    self.action_in = eqx.nn.Linear(action_dim, hidden_dim, key=keys[0])
    dims = [latent_dim + hidden_dim] + [hidden_dim] * depth
    self.layers = [
        eqx.nn.Linear(dims[i], hidden_dim, key=keys[i + 1]) for i in range(depth)
    ]
    self.norms = [LayerNorm(hidden_dim) for _ in range(depth)]
    out = eqx.nn.Linear(hidden_dim, latent_dim, key=keys[-1])
    if residual:
        # Small, not zero: a zero output projection also zeroes the gradient
        # flowing back through it, stalling every upstream parameter.
        out = eqx.tree_at(
            lambda m: (m.weight, m.bias), out, (0.05 * out.weight, 0.05 * out.bias)
        )
    self.out = out
    if normalize == "simnorm":
        self.post = SimNorm(simnorm_groups)
    elif normalize == "layernorm":
        self.post = LayerNorm(latent_dim)
    elif normalize == "none":
        self.post = None
    else:
        raise ValueError(f"unknown normalization {normalize!r}")
    self.latent_dim = latent_dim
    self.action_dim = action_dim
    self.residual = residual