Planning

Previous Next

Composable objectives and cost evaluators for planning

Solvers optimize action sequences against a Costable — anything exposing get_cost(info_dict, action_candidates). Some world models implement it natively (e.g. TD-MPC2); for the others, this module provides the glue: a ShootingCostEvaluator composes any model exposing the Dynamics surface (encode/rollout) with a swappable Objective, so changing the planning cost never requires subclassing the world model.

The rollout contract

rollout(info_dict, action_candidates) receives strictly-future candidates of shape (B, S, horizon, action_dim). Observation context arrives via the info dict: pixels holds H = history_len frames (B, S, H, C, h, w), and when H > 1 the executed action blocks between those frames are supplied as action_history (B, S, H - 1, action_dim) — frozen inputs, never optimizer variables. Inside the rollout, context frame k pairs with the action block leaving it (action_history[k] for past frames; the first candidate for the current frame), matching the training-time (frame[t], action[t]) alignment. The output predicted_emb has shape (B, S, H + horizon, dim) with the first H entries being the encoded context — objectives that read anything other than the last step must account for this (GoalMSE reads [:, :, -1:] and is unaffected).

[ Quick Tour ]

import stable_worldmodel as swm
from stable_worldmodel.planning import (
    CEMSolver,
    ControlPenalty,
    GoalMSE,
    ShootingCostEvaluator,
    WeightedSum,
)

model = swm.wm.utils.load_pretrained('lewm/pusht')

# 1. Single-term cost: last-step MSE to the goal embedding
cost = ShootingCostEvaluator(model, GoalMSE())

# 2. Multi-term cost: goal distance + action magnitude penalty
cost = ShootingCostEvaluator(
    model,
    WeightedSum([(1.0, GoalMSE()), (0.1, ControlPenalty())]),
)

# 3. Plug into any solver — the evaluator duck-types as a Costable
solver = CEMSolver(cost=cost, n_steps=30, num_samples=300, topk=30)
config = swm.PlanConfig(horizon=10, receding_horizon=1, action_block=1)
policy = swm.policy.WorldModelPolicy(solver=solver, config=config)

To plan under inequality constraints, pass objectives as constraints= — the evaluator then exposes get_constraints and satisfies the Constrainable protocol that LagrangianSolver feature-detects:

cost = ShootingCostEvaluator(
    model,
    GoalMSE(),
    constraints=[ControlPenalty()],  # satisfied when <= 0
)

Split-latent models (PreJEPA / dinowm)

PreJEPA fuses the pixel embedding with one embedding per extra encoder (proprio, action) along the feature axis. That latent cannot be scored against a goal as a single tensor: it carries action slots, which a goal does not prescribe. There is no special objective for this — point one GoalMSE at each source and add them up with WeightedSum. The per-source goal embeddings come from split_goal_encode, which default_goal_encode selects automatically for such models:

model = swm.wm.utils.load_pretrained('dinowm-pusht')

cost = ShootingCostEvaluator(  # goal encoder auto-selected
    model,
    WeightedSum([
        # extras before pixels, mean-reduced: reproduces the cost dinowm was
        # evaluated with before the objective was pulled out of the model
        (1.0, GoalMSE('predicted_proprio_emb', 'proprio_goal_emb', reduction='mean')),
        (1.0, GoalMSE('predicted_pixels_emb', 'pixels_goal_emb', reduction='mean')),
    ]),
)

The leading coefficients trade the sources off against each other. reduction matters here: 'sum' over a 196x384 pixel embedding versus a 10-dim proprio embedding weights them ~7500x apart, so 'mean' is what puts them on a comparable footing (and is what the original cost used).

From the plan scripts, pick the matching config — objective=goal_mse_pixels_proprio for dinowm-pusht, or objective=goal_mse_pixels for dinowm_noprop-pusht, whose only non-action source is pixels. Note also that the extra encoders consume the same context frames as pixels, so WorldModelPolicy.history_keys must cover every extra encoder input (e.g. ('pixels', 'proprio')) — otherwise the fused concatenation is malformed.

Writing a custom objective

An objective is any callable mapping a populated info_dict to a per-candidate cost of shape (B, S). The evaluator rolls candidates out first, so the info_dict already holds the rollout outputs (e.g. predicted_emb) plus the raw candidates under action_candidates:

import torch.nn as nn


class SmoothnessPenalty(nn.Module):
    """Penalizes large action changes between consecutive steps."""

    def forward(self, info_dict: dict) -> torch.Tensor:
        actions = info_dict['action_candidates']  # (B, S, H, action_dim)
        deltas = actions[..., 1:, :] - actions[..., :-1, :]
        return deltas.pow(2).sum(dim=tuple(range(2, actions.ndim)))

[ Evaluator ]

ShootingCostEvaluator 

ShootingCostEvaluator(
    model: Dynamics,
    objective: Objective,
    constraints: list[Objective] | None = None,
    encode_goal: Callable[[Dynamics, dict], Tensor]
    | None = default_goal_encode,
)

Bases: Module

Single-shooting adapter that makes (model, objective) a Costable.

Subclasses torch.nn.Module and registers the world model as a submodule, so parameters() reaches the real model and solvers (CEM/GD) infer the candidate dtype from it instead of falling back to float32.

The actor warm-start tail-fill — a solver calling model.get_action for Actionable models such as tdmpc2 — is not forwarded through the wrapper; it is moot for the LeWM-style models this targets. Add get_action delegation if an Actionable model is ever wrapped.

Parameters:

  • model (Dynamics) –

    World model providing encode/rollout (the Dynamics surface). Registered as a submodule so parameters() reaches it.

  • objective (Objective) –

    The cost to apply to the rolled-out info_dict.

  • constraints (list[Objective] | None, default: None ) –

    Optional Objective terms reused as constraints, each returning (B, S) under the LagrangianSolver convention that a term is satisfied when <= 0. When given, get_constraints is exposed and stacks them into (B, S, C); otherwise the attribute is absent, so a solver probing for it sees no constraints.

  • encode_goal (Callable[[Dynamics, dict], Tensor] | None, default: default_goal_encode ) –

    Optional fn(model, info_dict) -> goal_emb to override the default goal encoding. Set to None on the call to skip goal encoding entirely (e.g. when the caller pre-populates goal_emb).

get_cost 

get_cost(
    info_dict: dict, action_candidates: Tensor
) -> Tensor

Encode goal (if needed), roll out candidates, then score them.

criterion 

criterion(
    info_dict: dict, action_candidates: Tensor | None = None
) -> Tensor

Score an already-rolled-out info_dict with the objective.

default_goal_encode 

default_goal_encode(
    model: Dynamics, info_dict: dict
) -> Tensor

Encode the goal using the encoder matching the model's latent layout.

Dispatches on whether the model fuses extra sources into its latent (extra_encoders): split-latent models take :func:split_goal_encode, everything else :func:flat_goal_encode. Pass encode_goal= explicitly to ShootingCostEvaluator to override the choice, or None to skip goal encoding entirely.

flat_goal_encode 

flat_goal_encode(
    model: Dynamics, info_dict: dict
) -> Tensor

Encode the goal for models whose latent is a single flat tensor.

Behavior-preserving extraction of the goal-encoding branch in LeWM.get_cost; also covers PLDM. Pair with :class:GoalMSE.

split_goal_encode 

split_goal_encode(
    model: Dynamics, info_dict: dict
) -> Tensor

Encode the goal for models whose latent concatenates several sources.

PreJEPA fuses the pixel embedding with one embedding per extra encoder along the feature axis, so a goal cannot be built by the flat path: the action encoder has no goal-side input at all (a goal prescribes a state, not an action), and the cost must compare the parts, not the fused tensor. This encodes the goal with the action encoder excluded, stores the per-source goal embeddings (pixels_goal_emb, <key>_goal_emb) in info_dict, and returns the fused goal embedding for the goal_emb key. Score them with one :class:~stable_worldmodel.planning.GoalMSE per source, combined with :class:~stable_worldmodel.planning.WeightedSum.

Like the flat path, the goal embeddings are stored without a candidate axis; the objective broadcasts them over candidates.

Behavior-preserving extraction of the goal-encoding branch in the former PreJEPA.get_cost.

[ Objectives ]

GoalMSE 

GoalMSE(
    pred_key: str = 'predicted_emb',
    goal_key: str = 'goal_emb',
    reduction: str = 'sum',
)

Bases: Module

Last-step MSE between predicted and goal embeddings.

Reads predicted_emb (B, S, T, ...) and goal_emb (B, T_goal, ...); returns per-candidate cost (B, S). The comparison is against the last predicted step and the last goal frame.

Rank-agnostic: the time axis is addressed positionally (dim 2 of the prediction, dim 1 of the goal) rather than from the right, so latents that carry extra trailing axes work too — PreJEPA's pixel embedding has a patch axis, (B, S, T, patches, dim), and indexing from the right would silently slice patches instead of time.

Point pred_key/goal_key at one source of a split latent (e.g. predicted_pixels_emb / pixels_goal_emb) and combine the terms with :class:WeightedSum to score models whose latent concatenates several sources; see the planning guide.

Behavior-preserving extraction of LeWM.criterion (reproduces it bit-for-bit), so an existing model migrates to the ShootingCostEvaluator seam without changing results.

Parameters:

  • pred_key (str, default: 'predicted_emb' ) –

    info_dict key holding the rollout predictions.

  • goal_key (str, default: 'goal_emb' ) –

    info_dict key holding the goal embedding.

  • reduction (str, default: 'sum' ) –

    How the squared error is collapsed over every axis after the candidate axis. 'sum' (the default) matches LeWM.criterion; 'mean' matches the former PreJEPA.criterion. This choice also sets the relative weight of sources with different shapes, so it matters when composing per-source terms: summing over a 196x384 pixel embedding and a 10-dim proprio embedding weights them ~7500x apart, whereas averaging puts them on a comparable scale.

ControlPenalty 

ControlPenalty(action_key: str = 'action_candidates')

Bases: Module

L2 penalty on the action candidates themselves.

Reads action_candidates (shape (B, S, H, action_dim)) from the info_dict — the ShootingCostEvaluator stores them there before scoring — and returns a per-candidate cost (B, S). Demonstrates a cost the world model never had to know about.

WeightedSum 

WeightedSum(terms: list[tuple[float, Objective]])

Bases: Module

Linear combination of objectives: sum(w * term(info_dict)).

Lets us assemble multi-term costs (goal distance + control penalty + constraints) without touching the model or the solver.

[ Protocols ]

The structural contracts live in stable_worldmodel.protocols and are re-exported from stable_worldmodel.planning. They are Protocol classes: nothing subclasses them, anything with the right methods satisfies them.

Costable

Bases: Protocol

Protocol for the cost surface planning solvers consume.

This is the structural "has get_cost" contract every solver types against. It is polymorphic across implementations: a :class:~stable_worldmodel.planning.ShootingCostEvaluator (a world model composed with an :class:Objective), as well as models that expose get_cost natively (e.g. TD-MPC2), all satisfy it.

Methods:

  • criterion –

    Compute the cost criterion for action candidates.

  • get_cost –

    Compute cost for given action candidates based on info dictionary.

criterion 

criterion(
    info_dict: dict, action_candidates: Tensor
) -> Tensor

Compute the cost criterion for action candidates.

Parameters:

  • info_dict (dict) –

    Dictionary containing environment state information.

  • action_candidates (Tensor) –

    Tensor of proposed actions.

Returns:

  • Tensor –

    A tensor of cost values for each action candidate.

get_cost 

get_cost(
    info_dict: dict, action_candidates: Tensor
) -> Tensor

Compute cost for given action candidates based on info dictionary.

Parameters:

  • info_dict (dict) –

    Dictionary containing environment state information.

  • action_candidates (Tensor) –

    Tensor of proposed actions.

Returns:

  • Tensor –

    A tensor of cost values for each action candidate.

Constrainable

Bases: Protocol

Protocol for the (optional) constraint surface of a cost object.

A cost object exposes get_constraints when planning under inequality constraints. LagrangianSolver feature-detects it via isinstance(cost, Constrainable). Following the Lagrangian contract, a constraint term g_i is satisfied when g_i <= 0.

Methods:

  • get_constraints –

    Compute constraint violations for given action candidates.

get_constraints 

get_constraints(
    info_dict: dict, action_candidates: Tensor
) -> Tensor

Compute constraint violations for given action candidates.

Parameters:

  • info_dict (dict) –

    Dictionary containing environment state information.

  • action_candidates (Tensor) –

    Tensor of proposed actions.

Returns:

  • Tensor –

    A tensor of shape (B, S, C) of per-candidate constraint values,

  • Tensor –

    where C is the number of constraints (satisfied when <= 0).

Dynamics

Bases: Protocol

The dynamics surface a ShootingCostEvaluator needs from a world model.

Methods:

  • encode –

    Embed raw observations into the model's latent space.

  • rollout –

    Roll candidate action sequences forward through the dynamics.

encode 

encode(x: dict) -> dict

Embed raw observations into the model's latent space.

Parameters:

  • x (dict) –

    Dictionary of observations (e.g. pixels, proprioception).

Returns:

  • dict –

    The dictionary augmented with latent embeddings (e.g. emb).

rollout 

rollout(info_dict: dict, action_candidates: Tensor) -> dict

Roll candidate action sequences forward through the dynamics.

Parameters:

  • info_dict (dict) –

    Dictionary containing environment state information. pixels holds H context frames (B, S, H, C, h, w); when H > 1 the executed action blocks between those frames must be provided as action_history of shape (B, S, H - 1, action_dim).

  • action_candidates (Tensor) –

    Tensor of proposed strictly-future action sequences of shape (B, S, horizon, action_dim).

Returns:

  • dict –

    The dictionary populated with rollout outputs (e.g.

  • dict –

    predicted_emb of shape (B, S, H + horizon, dim), whose

  • dict –

    first H entries are the encoded context frames).

Objective

Bases: Protocol

Maps a populated info_dict to per-candidate cost (B, S).

Unlike :class:Costable, an Objective scores an already-rolled-out info_dict — it takes no action_candidates and performs no rollout. The info_dict is expected to already contain rollout outputs (e.g. predicted_emb) and any goal/conditioning the objective needs. The ShootingCostEvaluator also stores the raw action_candidates under the action_candidates key so action-space penalties can read them.

Solver Spaces