xwm.planning

CEM, MPPI, gradient planning, MPC and PUCT-MCTS, plus latent cost functions. Everything here is jittable: candidates are vmaped and refinement is a lax.fori_loop. See Planning.

Planning: using a learned world model to choose actions.

This is what a predictive world model is for. The encoder turns observations into latents, the action-conditioned predictor imagines what actions would do, and a planner searches for the action sequence whose imagined outcome is best -- all without ever rendering a pixel.

Which planner depends on the action space and on whether you have a value function:

• CEM, MPPI -- continuous actions, sample whole sequences. What TD-MPC2 and the JEPA action-conditioned models use.
• GradientPlanner -- continuous, differentiates through the rollout. Sample-efficient but happy to exploit model error.
• MCTS -- discrete actions, grows a tree and spends its budget where the model is least certain. What MuZero uses.

Modules:

Name	Description
cost	Cost functions defined in latent space.
gradient	Gradient-based planning.
mcts	Monte-Carlo tree search over a learned model (MuZero's planner).
mpc	Receding-horizon control (MPC) on top of a planner.
sampling	Sampling-based planners: CEM and MPPI.

Classes:

Name	Description
GradientPlanner	Optimise an action sequence by gradient descent on the rollout cost.
MCTS	PUCT tree search over a learned latent model.
SearchResult	What a search returns.
ControlStep	One closed-loop control step.
Planner	Anything with a plan method: CEM, MPPI, GradientPlanner.
CEM	Cross-entropy method: iteratively refit a Gaussian to the elite samples.
MPPI	Model-predictive path integral: softmax-weighted average of all samples.
Plan	The result of a planning call.

Functions:

Name	Description
goal_cost	Drive the latent state toward z_goal.
latent_distance	Scalar distance between two latent states of identical shape.
return_cost	Negated discounted return -- the objective a value-based agent plans on.
reward_cost	Turn a learned latent reward (higher is better) into a cost to minimise.
sum_costs	Weighted sum of several cost terms.
control_step	Plan from the current latent state and return the first action.
run_mpc	Run a closed loop against a real (or simulated) environment.
shift_plan	Advance a plan by one step, repeating the last action at the tail.

GradientPlanner

GradientPlanner(horizon: int, action_dim: int, *, n_steps: int = 100, learning_rate: float = 0.05, low: float | Array = -1.0, high: float | Array = 1.0, cost_on: str = 'next')

Bases: Module

Optimise an action sequence by gradient descent on the rollout cost.

Parameters:

Name	Type	Description	Default
horizon	int	planning horizon.	required
action_dim	int	action dimensionality.	required
n_steps	int	optimisation steps.	100
learning_rate	float	Adam step size on the actions.	0.05
low, high		bounds, enforced by projection after each step.	required

Source code in xwm/planning/gradient.py

def __init__(
    self,
    horizon: int,
    action_dim: int,
    *,
    n_steps: int = 100,
    learning_rate: float = 0.05,
    low: float | Array = -1.0,
    high: float | Array = 1.0,
    cost_on: str = "next",
):
    self.horizon = horizon
    self.action_dim = action_dim
    self.n_steps = n_steps
    self.learning_rate = learning_rate
    self.low = jnp.broadcast_to(jnp.asarray(low, jnp.float32), (action_dim,))
    self.high = jnp.broadcast_to(jnp.asarray(high, jnp.float32), (action_dim,))
    self.cost_on = cost_on

MCTS

MCTS(n_actions: int, *, n_simulations: int = 50, discount: float = 0.997, c_puct: float = 1.25, dirichlet_alpha: float = 0.3, root_noise_fraction: float = 0.25, max_depth: int = 50)

Bases: Module

PUCT tree search over a learned latent model.

Parameters:

Name	Type	Description	Default
n_actions	int	size of the discrete action space.	required
n_simulations	int	search budget.	50
discount	float	RL discount used when backing values up the tree.	0.997
c_puct	float	exploration constant.	1.25
dirichlet_alpha, root_noise_fraction		root exploration noise.	required
max_depth	int	hard cap on tree depth, which also bounds the arrays.	50

Methods:

Name	Description
search	Run the search from root_latent.

Source code in xwm/planning/mcts.py

def __init__(
    self,
    n_actions: int,
    *,
    n_simulations: int = 50,
    discount: float = 0.997,
    c_puct: float = 1.25,
    dirichlet_alpha: float = 0.3,
    root_noise_fraction: float = 0.25,
    max_depth: int = 50,
):
    if n_actions < 2:
        raise ValueError(f"need at least two actions, got {n_actions}")
    self.n_actions = n_actions
    self.n_simulations = n_simulations
    self.discount = discount
    self.c_puct = c_puct
    self.dirichlet_alpha = dirichlet_alpha
    self.root_noise_fraction = root_noise_fraction
    self.max_depth = max_depth

search

search(key: PRNGKey, root_latent: Array, recurrent: Callable[[Array, Array], tuple[Array, Array]], predict: Callable[[Array], tuple[Array, Array]], *, add_noise: bool = True) -> SearchResult

Run the search from root_latent.

Parameters:

Name	Type	Description	Default
recurrent	Callable[[Array, Array], tuple[Array, Array]]	(latent, action_index) -> (next_latent, reward).	required
predict	Callable[[Array], tuple[Array, Array]]	latent -> (policy_logits, value).	required
add_noise	bool	Dirichlet noise at the root. On for self-play data collection, off for evaluation.	True

Returns:

Type	Description
SearchResult	A SearchResult.

Source code in xwm/planning/mcts.py

def search(
    self,
    key: PRNGKey,
    root_latent: Array,
    recurrent: Callable[[Array, Array], tuple[Array, Array]],
    predict: Callable[[Array], tuple[Array, Array]],
    *,
    add_noise: bool = True,
) -> SearchResult:
    """Run the search from ``root_latent``.

    Args:
        recurrent: ``(latent, action_index) -> (next_latent, reward)``.
        predict: ``latent -> (policy_logits, value)``.
        add_noise: Dirichlet noise at the root. On for self-play data
            collection, off for evaluation.

    Returns:
        A :class:`SearchResult`.
    """
    n_nodes = self.n_simulations + 1
    n_actions = self.n_actions

    # Flat arrays instead of node objects, so the whole search is jittable.
    latents = jnp.zeros((n_nodes, root_latent.shape[-1])).at[0].set(root_latent)
    priors = jnp.zeros((n_nodes, n_actions))
    values = jnp.zeros((n_nodes,))
    visits = jnp.zeros((n_nodes,), jnp.int32)
    value_sums = jnp.zeros((n_nodes,))
    children = jnp.full((n_nodes, n_actions), -1, jnp.int32)
    rewards = jnp.zeros((n_nodes,))
    parents = jnp.full((n_nodes,), -1, jnp.int32)
    parent_action = jnp.full((n_nodes,), -1, jnp.int32)

    root_logits, root_value = predict(root_latent)
    root_prior = jax.nn.softmax(root_logits)
    if add_noise:
        noise = jr.dirichlet(key, jnp.full((n_actions,), self.dirichlet_alpha))
        frac = self.root_noise_fraction
        root_prior = (1.0 - frac) * root_prior + frac * noise
    priors = priors.at[0].set(root_prior)
    values = values.at[0].set(root_value)
    visits = visits.at[0].set(1)
    value_sums = value_sums.at[0].set(root_value)

    state = (latents, priors, values, visits, value_sums, children, rewards,
             parents, parent_action, jnp.asarray(1, jnp.int32))

    def simulate(carry, _):
        (latents, priors, values, visits, value_sums, children, rewards,
         parents, parent_action, n_used) = carry

        # -- select: walk down by PUCT until an unexpanded action.
        def cond(loop):
            _, _, depth, done = loop
            return jnp.logical_and(jnp.logical_not(done), depth < self.max_depth)

        def descend(loop):
            node, action, depth, _ = loop
            # The score of an action is the reward it earns *plus* the
            # discounted value of where it lands. Scoring by the child's
            # value alone makes the search blind to immediate reward, so it
            # cannot find a payoff that is one step away.
            q_mean = jnp.where(visits > 0, value_sums / jnp.maximum(visits, 1), 0.0)
            action_value = rewards + self.discount * q_mean
            # Normalise into [0, 1] using the tree's own range: a learned
            # value head has no fixed scale, so a raw Q would make c_puct
            # task-dependent.
            seen = visits > 0
            low = jnp.min(jnp.where(seen, action_value, jnp.inf))
            high = jnp.max(jnp.where(seen, action_value, -jnp.inf))
            span = jnp.maximum(high - low, 1e-8)
            child_idx = children[node]
            child_visits = jnp.where(child_idx >= 0, visits[child_idx], 0)
            child_q = jnp.where(
                child_idx >= 0, (action_value[child_idx] - low) / span, 0.0
            )
            exploration = (
                self.c_puct
                * priors[node]
                * jnp.sqrt(jnp.maximum(visits[node], 1))
                / (1 + child_visits)
            )
            chosen = jnp.argmax(child_q + exploration)
            next_node = children[node, chosen]
            return (
                jnp.where(next_node >= 0, next_node, node),
                chosen,
                depth + 1,
                next_node < 0,
            )

        node, action, depth, _ = jax.lax.while_loop(
            cond, descend, (jnp.asarray(0, jnp.int32), jnp.asarray(0, jnp.int32),
                            jnp.asarray(0, jnp.int32), jnp.asarray(False))
        )

        # -- expand the chosen action into a new node.
        new_latent, reward = recurrent(latents[node], action)
        logits, leaf_value = predict(new_latent)
        index = n_used
        latents = latents.at[index].set(new_latent)
        priors = priors.at[index].set(jax.nn.softmax(logits))
        values = values.at[index].set(leaf_value)
        rewards = rewards.at[index].set(reward)
        parents = parents.at[index].set(node)
        parent_action = parent_action.at[index].set(action)
        children = children.at[node, action].set(index)

        # -- back up the discounted return along the path to the root.
        def step_back(carry_bu, _):
            current, value, visits_, sums_ = carry_bu
            alive = current >= 0
            safe = jnp.maximum(current, 0)
            visits_ = visits_.at[safe].add(jnp.where(alive, 1, 0))
            sums_ = sums_.at[safe].add(jnp.where(alive, value, 0.0))
            next_value = jnp.where(
                alive, rewards[safe] + self.discount * value, value
            )
            return (parents[safe], next_value, visits_, sums_), None

        (_, _, visits, value_sums), _ = jax.lax.scan(
            step_back,
            (index, leaf_value, visits, value_sums),
            None,
            length=self.max_depth,
        )

        return (latents, priors, values, visits, value_sums, children, rewards,
                parents, parent_action, n_used + 1), None

    state, _ = jax.lax.scan(simulate, state, None, length=self.n_simulations)
    _, _, _, visits, value_sums, children, rewards, _, _, _ = state

    root_children = children[0]
    child_visits = jnp.where(root_children >= 0, visits[root_children], 0).astype(
        jnp.float32
    )
    total = jnp.maximum(jnp.sum(child_visits), 1.0)
    policy = child_visits / total
    q_mean = jnp.where(visits > 0, value_sums / jnp.maximum(visits, 1), 0.0)
    action_value = rewards + self.discount * q_mean
    child_q = jnp.where(root_children >= 0, action_value[root_children], 0.0)
    return SearchResult(
        policy=policy,
        value=jnp.sum(policy * child_q),
        action=jnp.argmax(child_visits),
        visits=child_visits,
    )

SearchResult

Bases: NamedTuple

What a search returns.

Attributes:

Name	Type	Description
policy	Array	(n_actions,) visit-count distribution -- the improved policy MuZero trains the network's policy head against.
value	Array	root value, the visit-weighted mean of the children's returns.
action	Array	the most-visited action.
visits	Array	raw visit counts, for diagnostics.

ControlStep

Bases: NamedTuple

One closed-loop control step.

Attributes:

Name	Type	Description
action	Array	the action to execute now.
warm_start	Array	shifted plan to seed the next call.
plan	Plan	the full plan, for logging or diagnostics.

Planner

Bases: Protocol

Anything with a plan method: CEM, MPPI, GradientPlanner.

CEM

CEM(horizon: int, action_dim: int, *, n_samples: int = 512, n_elites: int = 64, n_iters: int = 6, low: float | Array = -1.0, high: float | Array = 1.0, init_std: float = 0.5, min_std: float = 0.05, momentum: float = 0.1, cost_on: str = 'next')

Bases: Module

Cross-entropy method: iteratively refit a Gaussian to the elite samples.

Parameters:

Name	Type	Description	Default
horizon	int	planning horizon in steps.	required
action_dim	int	action dimensionality.	required
n_samples	int	candidates per iteration.	512
n_elites	int	how many best candidates define the next proposal.	64
n_iters	int	refinement iterations.	6
low, high		action bounds; candidates are clipped into them.	required
init_std	float	initial per-dimension spread.	0.5
min_std	float	floor on the spread, so the search cannot collapse to a point and stop exploring.	0.05
momentum	float	smoothing of the proposal across iterations, in [0, 1).	0.1

Methods:

Name	Description
plan	Search for a low-cost action sequence from latent state z0.

Source code in xwm/planning/sampling.py

def __init__(
    self,
    horizon: int,
    action_dim: int,
    *,
    n_samples: int = 512,
    n_elites: int = 64,
    n_iters: int = 6,
    low: float | Array = -1.0,
    high: float | Array = 1.0,
    init_std: float = 0.5,
    min_std: float = 0.05,
    momentum: float = 0.1,
    cost_on: str = "next",
):
    if n_elites > n_samples:
        raise ValueError(f"n_elites={n_elites} exceeds n_samples={n_samples}")
    self.horizon = horizon
    self.action_dim = action_dim
    self.n_samples = n_samples
    self.n_elites = n_elites
    self.n_iters = n_iters
    self.low = jnp.broadcast_to(jnp.asarray(low, jnp.float32), (action_dim,))
    self.high = jnp.broadcast_to(jnp.asarray(high, jnp.float32), (action_dim,))
    self.init_std = init_std
    self.min_std = min_std
    self.momentum = momentum
    self.cost_on = cost_on

plan

plan(key: PRNGKey, dynamics: LatentDynamics, z0: Array, cost_fn: CostFn, *, init_mean: Array | None = None) -> Plan

Search for a low-cost action sequence from latent state z0.

Source code in xwm/planning/sampling.py

def plan(
    self,
    key: PRNGKey,
    dynamics: LatentDynamics,
    z0: Array,
    cost_fn: CostFn,
    *,
    init_mean: Array | None = None,
) -> Plan:
    """Search for a low-cost action sequence from latent state ``z0``."""
    shape = (self.horizon, self.action_dim)
    mean = jnp.zeros(shape) if init_mean is None else jnp.asarray(init_mean)
    std = jnp.full(shape, self.init_std)

    def iteration(i, carry):
        mean, std = carry
        noise = jr.normal(jr.fold_in(key, i), (self.n_samples, *shape))
        candidates = jnp.clip(mean + std * noise, self.low, self.high)
        costs = _evaluate(dynamics, z0, candidates, cost_fn, self.cost_on)
        elites = candidates[jnp.argsort(costs)[: self.n_elites]]
        new_mean = jnp.mean(elites, axis=0)
        new_std = jnp.maximum(jnp.std(elites, axis=0), self.min_std)
        m = self.momentum
        return m * mean + (1 - m) * new_mean, m * std + (1 - m) * new_std

    mean, std = jax.lax.fori_loop(0, self.n_iters, iteration, (mean, std))
    actions = jnp.clip(mean, self.low, self.high)
    cost = rollout_cost(dynamics, z0, actions, cost_fn, cost_on=self.cost_on)
    return Plan(actions=actions, cost=cost, mean=mean, std=std)

MPPI

MPPI(horizon: int, action_dim: int, *, n_samples: int = 512, n_iters: int = 4, low: float | Array = -1.0, high: float | Array = 1.0, temperature: float = 1.0, noise_std: float = 0.5, cost_on: str = 'next')

Bases: Module

Model-predictive path integral: softmax-weighted average of all samples.

Unlike CEM's hard elite cut, every candidate contributes in proportion to exp(-cost / temperature). The soft weighting makes the update smoother across control steps, which matters when the planner is in a feedback loop.

Parameters:

Name	Type	Description	Default
temperature	float	lower values approach CEM's greedy behaviour; higher values average more broadly. Costs are shifted by their minimum before exponentiating, so the scale is relative and numerically safe.	1.0
noise_std	float	proposal spread, held fixed rather than refit.	0.5

Source code in xwm/planning/sampling.py

def __init__(
    self,
    horizon: int,
    action_dim: int,
    *,
    n_samples: int = 512,
    n_iters: int = 4,
    low: float | Array = -1.0,
    high: float | Array = 1.0,
    temperature: float = 1.0,
    noise_std: float = 0.5,
    cost_on: str = "next",
):
    self.horizon = horizon
    self.action_dim = action_dim
    self.n_samples = n_samples
    self.n_iters = n_iters
    self.low = jnp.broadcast_to(jnp.asarray(low, jnp.float32), (action_dim,))
    self.high = jnp.broadcast_to(jnp.asarray(high, jnp.float32), (action_dim,))
    self.temperature = temperature
    self.noise_std = noise_std
    self.cost_on = cost_on

Plan

Bases: NamedTuple

The result of a planning call.

Attributes:

Name	Type	Description
actions	Array	(H, A) chosen action sequence.
cost	Array	its predicted cost under the world model.
mean	Array	(H, A) final proposal mean -- pass it back as init_mean next step to warm-start, which is most of what makes receding-horizon control cheap.
std	Array	(H, A) final proposal spread.

goal_cost

goal_cost(z_goal: Array, *, kind: Distance = 'l2', horizon: int | None = None, terminal_only: bool = False, action_penalty: float = 0.0, discount: float = 1.0) -> CostFn

Drive the latent state toward z_goal.

Parameters:

Name	Type	Description	Default
z_goal	Array	target latent, same shape as the rollout states.	required
kind	Distance	distance to use.	'l2'
horizon	int | None	needed only when terminal_only is set, to know which step is terminal.	None
terminal_only	bool	score just the final state. Charging every step instead (the default) rewards reaching the goal early and staying there, which is usually what you want and is much better conditioned.	False
action_penalty	float	weight on mean(a ** 2), discouraging thrash.	0.0
discount	float	per-step multiplier; < 1 prefers reaching the goal sooner.	1.0

Source code in xwm/planning/cost.py

def goal_cost(
    z_goal: Array,
    *,
    kind: Distance = "l2",
    horizon: int | None = None,
    terminal_only: bool = False,
    action_penalty: float = 0.0,
    discount: float = 1.0,
) -> CostFn:
    """Drive the latent state toward ``z_goal``.

    Args:
        z_goal: target latent, same shape as the rollout states.
        kind: distance to use.
        horizon: needed only when ``terminal_only`` is set, to know which step
            is terminal.
        terminal_only: score just the final state. Charging every step instead
            (the default) rewards *reaching* the goal early and staying there,
            which is usually what you want and is much better conditioned.
        action_penalty: weight on ``mean(a ** 2)``, discouraging thrash.
        discount: per-step multiplier; ``< 1`` prefers reaching the goal sooner.
    """
    if terminal_only and horizon is None:
        raise ValueError("terminal_only requires horizon")

    def cost(z: Array, a: Array, t: Array) -> Array:
        distance = latent_distance(z, z_goal, kind)
        if terminal_only:
            distance = jnp.where(t == horizon - 1, distance, 0.0)
        else:
            distance = distance * discount**t
        return distance + action_penalty * jnp.mean(jnp.square(a))

    return cost

latent_distance

latent_distance(z: Array, z_goal: Array, kind: Distance = 'l2') -> Array

Scalar distance between two latent states of identical shape.

Source code in xwm/planning/cost.py

def latent_distance(z: Array, z_goal: Array, kind: Distance = "l2") -> Array:
    """Scalar distance between two latent states of identical shape."""
    if kind == "l1":
        return jnp.mean(jnp.abs(z - z_goal))
    if kind == "l2":
        return jnp.mean(jnp.square(z - z_goal))
    if kind == "cosine":
        a = z.reshape(-1) / (jnp.linalg.norm(z.reshape(-1)) + 1e-8)
        b = z_goal.reshape(-1) / (jnp.linalg.norm(z_goal.reshape(-1)) + 1e-8)
        return 1.0 - jnp.sum(a * b)
    raise ValueError(f"unknown distance {kind!r}")

return_cost

return_cost(reward_fn: Callable[[Array, Array], Array], value_fn: Callable[[Array], Array] | None = None, *, horizon: int, discount: float = 0.99) -> CostFn

Negated discounted return -- the objective a value-based agent plans on.

-(sum_t gamma^t r(z_t, a_t) + gamma^H V(z_H)).

The terminal value is what makes this different from a goal cost, and it is the whole reason TD-MPC2 can plan with a horizon of three: the value head summarises everything beyond the horizon, so the planner does not have to simulate it. Without that term a short-horizon planner is myopic by construction -- it cannot prefer a move whose payoff arrives on step four.

Use with cost_on="current" (the default for planners built by xwm.families.tdmpc2.planner): a reward head is trained as r(z_t, a_t), so it must be evaluated at the latent the action was taken from.

Parameters:

Name	Type	Description	Default
reward_fn	Callable[[Array, Array], Array]	(z, a) -> r, typically a learned reward head.	required
value_fn	Callable[[Array], Array] | None	z -> V, applied once at the final step. None drops the bootstrap, making the objective purely myopic.	None
horizon	int	planning horizon, needed to know which step is terminal.	required
discount	float	RL discount.	0.99

Source code in xwm/planning/cost.py

def return_cost(
    reward_fn: Callable[[Array, Array], Array],
    value_fn: Callable[[Array], Array] | None = None,
    *,
    horizon: int,
    discount: float = 0.99,
) -> CostFn:
    """Negated discounted return -- the objective a value-based agent plans on.

    ``-(sum_t gamma^t r(z_t, a_t) + gamma^H V(z_H))``.

    The terminal value is what makes this different from a goal cost, and it is
    the whole reason TD-MPC2 can plan with a horizon of three: the value head
    summarises everything beyond the horizon, so the planner does not have to
    simulate it. Without that term a short-horizon planner is myopic by
    construction -- it cannot prefer a move whose payoff arrives on step four.

    Use with ``cost_on="current"`` (the default for planners built by
    :func:`xwm.families.tdmpc2.planner`): a reward head is trained as
    ``r(z_t, a_t)``, so it must be evaluated at the latent the action was taken
    from.

    Args:
        reward_fn: ``(z, a) -> r``, typically a learned reward head.
        value_fn: ``z -> V``, applied once at the final step. ``None`` drops the
            bootstrap, making the objective purely myopic.
        horizon: planning horizon, needed to know which step is terminal.
        discount: RL discount.
    """

    def cost(z: Array, a: Array, t: Array) -> Array:
        total = -(discount**t) * reward_fn(z, a)
        if value_fn is not None:
            # The terminal latent is the one *after* the last action, so it is
            # not visible here; bootstrap from the last pre-transition latent,
            # discounted one extra step. Exact enough for ranking candidates,
            # and it avoids a second dynamics call per sample.
            terminal = -(discount ** (t + 1)) * value_fn(z)
            total = total + jnp.where(t == horizon - 1, terminal, 0.0)
        return total

    return cost

reward_cost

reward_cost(reward_fn: Callable[[Array], Array], *, action_penalty: float = 0.0, discount: float = 1.0) -> CostFn

Turn a learned latent reward (higher is better) into a cost to minimise.

Source code in xwm/planning/cost.py

def reward_cost(
    reward_fn: Callable[[Array], Array],
    *,
    action_penalty: float = 0.0,
    discount: float = 1.0,
) -> CostFn:
    """Turn a learned latent reward (higher is better) into a cost to minimise."""

    def cost(z: Array, a: Array, t: Array) -> Array:
        return -discount**t * reward_fn(z) + action_penalty * jnp.mean(jnp.square(a))

    return cost

sum_costs

sum_costs(*costs: CostFn, weights: tuple[float, ...] | None = None) -> CostFn

Weighted sum of several cost terms.

Source code in xwm/planning/cost.py

def sum_costs(*costs: CostFn, weights: tuple[float, ...] | None = None) -> CostFn:
    """Weighted sum of several cost terms."""
    w = (1.0,) * len(costs) if weights is None else weights
    if len(w) != len(costs):
        raise ValueError(f"got {len(costs)} costs but {len(w)} weights")

    def cost(z: Array, a: Array, t: Array) -> Array:
        return sum(wi * ci(z, a, t) for wi, ci in zip(w, costs, strict=True))

    return cost

control_step

control_step(key: PRNGKey, planner: Planner, dynamics: LatentDynamics, z: Array, cost_fn: CostFn, *, warm_start: Array | None = None) -> ControlStep

Plan from the current latent state and return the first action.

Source code in xwm/planning/mpc.py

def control_step(
    key: PRNGKey,
    planner: Planner,
    dynamics: LatentDynamics,
    z: Array,
    cost_fn: CostFn,
    *,
    warm_start: Array | None = None,
) -> ControlStep:
    """Plan from the current latent state and return the first action."""
    plan = planner.plan(key, dynamics, z, cost_fn, init_mean=warm_start)
    return ControlStep(
        action=plan.actions[0],
        warm_start=shift_plan(plan.mean),
        plan=plan,
    )

run_mpc

run_mpc(key: PRNGKey, planner: Planner, dynamics: LatentDynamics, cost_fn: CostFn, *, encode: Callable[[Array], Array], step_env: Callable[[Array, Array], Array], observation: Array, n_steps: int) -> tuple[list[Array], list[Array], list[Array]]

Run a closed loop against a real (or simulated) environment.

Parameters:

Name	Type	Description	Default
encode	Callable[[Array], Array]	observation -> latent state, i.e. the world model's encoder.	required
step_env	Callable[[Array, Array], Array]	(observation, action) -> next_observation. Deliberately a plain callable: xwm does not own your simulator.	required
observation	Array	the starting observation.	required
n_steps	int	control steps to execute.	required

Returns:

Type	Description
list[Array]	(observations, actions, costs) -- the realised trajectory. This is a
list[Array]	Python loop because the environment step is outside JAX; each planning
list[Array]	call is still a single jitted device call.

Source code in xwm/planning/mpc.py

def run_mpc(
    key: PRNGKey,
    planner: Planner,
    dynamics: LatentDynamics,
    cost_fn: CostFn,
    *,
    encode: Callable[[Array], Array],
    step_env: Callable[[Array, Array], Array],
    observation: Array,
    n_steps: int,
) -> tuple[list[Array], list[Array], list[Array]]:
    """Run a closed loop against a real (or simulated) environment.

    Args:
        encode: observation -> latent state, i.e. the world model's encoder.
        step_env: ``(observation, action) -> next_observation``. Deliberately a
            plain callable: xwm does not own your simulator.
        observation: the starting observation.
        n_steps: control steps to execute.

    Returns:
        ``(observations, actions, costs)`` -- the realised trajectory. This is a
        Python loop because the environment step is outside JAX; each planning
        call is still a single jitted device call.
    """
    observations, actions, costs = [observation], [], []
    warm_start = None
    for t in range(n_steps):
        z = encode(observation)
        step = control_step(
            jr.fold_in(key, t), planner, dynamics, z, cost_fn, warm_start=warm_start
        )
        observation = step_env(observation, step.action)
        observations.append(observation)
        actions.append(step.action)
        costs.append(step.plan.cost)
        warm_start = step.warm_start
    return observations, actions, costs

shift_plan

shift_plan(mean: Array) -> Array

Advance a plan by one step, repeating the last action at the tail.

Source code in xwm/planning/mpc.py

def shift_plan(mean: Array) -> Array:
    """Advance a plan by one step, repeating the last action at the tail."""
    return jnp.concatenate([mean[1:], mean[-1:]], axis=0)