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
)

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 embedding from an info_dict.

Behavior-preserving extraction of the goal-encoding branch in LeWM.get_cost. Override via ShootingCostEvaluator(encode_goal=...) for models that construct their goal differently.

[ Objectives ]

GoalMSE 

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

Bases: Module

Last-step MSE between predicted and goal embeddings.

Reads predicted_emb (B, S, T-1, dim) and goal_emb (B, T, dim); returns per-candidate cost (B, S).

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

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, prejepa), 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