Policy

Agent policies for interacting with environments

Policies determine the actions taken by agents in the environment. stable_worldmodel provides base classes and implementations for random, expert, and model-based policies.

Random Policy

A simple policy that samples actions uniformly from the environment's action space.

from stable_worldmodel.policy import RandomPolicy

# Create a random policy
policy = RandomPolicy(seed=42)

# Attach to a world/env later
# world.set_policy(policy)

World Model Policy

A policy that uses a Solver (like CEM or MPPI) and a World Model to plan actions.

from stable_worldmodel.policy import WorldModelPolicy, PlanConfig
from stable_worldmodel.planning import CEMSolver

# 1. Define Planning Configuration
cfg = PlanConfig(
    horizon=10,
    receding_horizon=1,
    action_block=1
)

# 2. Instantiate a Solver
solver = CEMSolver(cost=world_model) # Or MPPISolver, GradientSolver, etc.

# 3. Create the Policy
policy = WorldModelPolicy(
    solver=solver,
    config=cfg
)

Feed-Forward Policy

A policy that uses a neural network model for direct action prediction via a single forward pass. Useful for imitation learning policies like Goal-Conditioned Behavioral Cloning (GCBC).

from stable_worldmodel.policy import FeedForwardPolicy
from stable_worldmodel.wm.utils import load_pretrained

# 1. Load a pre-trained model with a get_action method
model = load_pretrained("path/to/checkpoint")

# 2. Create the Policy
policy = FeedForwardPolicy(
    model=model,
    process={"action": action_scaler},  # Optional preprocessors
    transform={"pixels": image_transform}  # Optional transforms
)

Protocol

All policies must implement the get_action(obs, **kwargs) method. The World class automatically calls set_env() when a policy is attached.

Planning with observation history

With PlanConfig(history_len > 1), WorldModelPolicy keeps a per-env HistoryBuffer over its history_keys (default ('pixels',)) plus the executed actions. At each replan the solver's info dict carries the frames at the last history_len block boundaries (pixels gains a real time dim) and the executed action blocks between them under action_history (solver space, one flattened block of action_block env actions per step). Candidates remain strictly future — see the rollout contract. Markovian Costable models (TD-MPC2) only support the default history_len=1.

Episode-start warm-up feeds synthetic context

For the first (history_len - 1) * action_block env steps of each episode there is not enough real history yet, and the world model receives fake repeated frames: the missing context slots are filled with copies of the episode's first frame, with zero action blocks between them (as if the env had been stationary before the episode began). All context is real after that window. See the buffer warm-up docs for the rationale.

PlanConfig `dataclass`

PlanConfig(
    horizon: int,
    receding_horizon: int,
    history_len: int = 1,
    history_max_len: int | None = None,
    action_block: int = 1,
    warm_start: bool = True,
)

Configuration for the MPC planning loop.

Attributes:

horizon (int) –

Planning horizon in number of steps.
receding_horizon (int) –

Number of steps to execute before re-planning.
history_len (int) –

Number of observation frames (in action_block timesteps, including the current frame) supplied to the world model at planning time. Values above 1 additionally supply the history_len - 1 executed action blocks between those frames via info['action_history'], and require a rollout-based Dynamics model (LeWM/PLDM/PreJEPA); Markovian Costable models (TD-MPC2) only support the default of 1. Warm-up behavior: during the first (history_len - 1) * action_block env steps of each episode the context is simply shorter — it grows from 1 frame at the first plan up to history_len, containing only real frames. Synthetic padding (copies of an env's oldest frame, with zero action blocks — as if the env had been stationary) is used only when a replan batch mixes envs at different fill levels (e.g. a freshly auto-reset env planning alongside envs with full histories), since their histories must stack into one tensor.
history_max_len (int | None) –

Capacity (in env steps) of the per-env history buffer. None means derive (history_len - 1) * action_block + 1 — the smallest size that yields history_len strided frames with a full action block between each consecutive pair. Set higher to retain more raw history than you sample.
action_block (int) –

Number of times each action is repeated (frameskip).
warm_start (bool) –

Whether to use the previous plan to initialize the next one.

BasePolicy

BasePolicy(**kwargs: Any)

Base class for agent policies.

Attributes:

env (Any) –

The environment the policy is associated with.
type (str) –

A string identifier for the policy type.

Parameters:

**kwargs (Any, default: {} ) –

Additional configuration parameters.

get_action

get_action(obs: Any, **kwargs: Any) -> ndarray

Get action from the policy given the observation.

Parameters:

obs (Any) –

The current observation from the environment.
**kwargs (Any, default: {} ) –

Additional parameters for action selection.

Returns:

ndarray –

Selected action as a numpy array.

Raises:

NotImplementedError –

If not implemented by a subclass.

set_env

set_env(env: Any) -> None

Associate this policy with an environment.

Parameters:

env (Any) –

The environment to associate.

_prepare_info

_prepare_info(info_dict: dict) -> dict[str, Tensor]

Pre-process and transform observations.

Applies preprocessing (via self.process) and transformations (via self.transform) to observation data. Used by subclasses like FeedForwardPolicy and WorldModelPolicy. Returns a new dict; the input is not mutated.

Parameters:

info_dict (dict) –

Raw observation dictionary from the environment.

Returns:

dict[str, Tensor] –

A dictionary of processed tensors.

Raises:

ValueError –

If an expected numpy array is missing for processing.

RandomPolicy

RandomPolicy(seed: int | None = None, **kwargs: Any)

Bases: BasePolicy

Policy that samples random actions from the action space.

Parameters:

seed (int | None, default: None ) –

Random seed applied to the action space when the environment is attached. If None, the action space uses its own default RNG, making action sampling non-deterministic across runs.
**kwargs (Any, default: {} ) –

Additional configuration parameters.

get_action

get_action(obs: Any, **kwargs: Any) -> ndarray

Get a random action from the environment's action space.

Parameters:

obs (Any) –

The current observation (ignored).
**kwargs (Any, default: {} ) –

Additional parameters (ignored).

Returns:

ndarray –

A randomly sampled action.

ExpertPolicy

ExpertPolicy(**kwargs: Any)

Bases: BasePolicy

Policy using expert demonstrations or heuristics.

Parameters:

**kwargs (Any, default: {} ) –

Additional configuration parameters.

get_action

get_action(
    obs: Any, goal_obs: Any, **kwargs: Any
) -> ndarray | None

Get action from the expert policy.

Parameters:

obs (Any) –

The current observation.
goal_obs (Any) –

The goal observation.
**kwargs (Any, default: {} ) –

Additional parameters.

Returns:

ndarray | None –

The expert action, or None if not available.

FeedForwardPolicy

FeedForwardPolicy(
    model: Actionable,
    process: dict[str, Transformable] | None = None,
    transform: dict[str, Callable[[Tensor], Tensor]]
    | None = None,
    **kwargs: Any,
)

Bases: BasePolicy

Feed-Forward Policy using a neural network model.

Actions are computed via a single forward pass through the model. Useful for imitation learning policies like Goal-Conditioned Behavioral Cloning (GCBC).

Attributes:

model –

Neural network model implementing the Actionable protocol.
process –

Dictionary of data preprocessors for specific keys.
transform –

Dictionary of tensor transformations (e.g., image transforms).

Parameters:

model (Actionable) –

Neural network model with a get_action method.
process (dict[str, Transformable] | None, default: None ) –

Dictionary of data preprocessors for specific keys.
transform (dict[str, Callable[[Tensor], Tensor]] | None, default: None ) –

Dictionary of tensor transformations (e.g., image transforms).
**kwargs (Any, default: {} ) –

Additional configuration parameters.

get_action

get_action(info_dict: dict, **kwargs: Any) -> ndarray

Get action via a forward pass through the neural network model.

Parameters:

info_dict (dict) –

Current state information containing at minimum a 'goal' key.
**kwargs (Any, default: {} ) –

Additional parameters (unused).

Returns:

ndarray –

The selected action as a numpy array.

Raises:

AssertionError –

If environment not set or 'goal' not in info_dict.

WorldModelPolicy

WorldModelPolicy(
    solver: Solver,
    config: PlanConfig,
    process: dict[str, Transformable] | None = None,
    transform: dict[str, Callable[[Tensor], Tensor]]
    | None = None,
    history_keys: tuple[str, ...] = ('pixels',),
    **kwargs: Any,
)

Bases: BasePolicy

Policy using a world model and planning solver for action selection.

Parameters:

solver (Solver) –

The planning solver to use.
config (PlanConfig) –

MPC planning configuration.
process (dict[str, Transformable] | None, default: None ) –

Dictionary of data preprocessors for specific keys.
transform (dict[str, Callable[[Tensor], Tensor]] | None, default: None ) –

Dictionary of tensor transformations (e.g., image transforms).
history_keys (tuple[str, ...], default: ('pixels',) ) –

Observation keys stacked over the last config.history_len block timesteps when replanning (only used when history_len > 1). The executed action blocks between those frames are supplied alongside under 'action_history'.
**kwargs (Any, default: {} ) –

Additional configuration parameters.

get_action

get_action(info_dict: dict, **kwargs: Any) -> ndarray

Get action via planning with the world model.

Parameters:

info_dict (dict) –

Current state information from the environment.
**kwargs (Any, default: {} ) –

Additional parameters for planning.

Returns:

ndarray –

The selected action(s) as a numpy array.

[ Utils ]

Load a pretrained checkpoint (a folder with weights.pt + config.json) with load_pretrained. It reconstructs the model from config.json and loads the weights — the model is returned directly, ready to pass to FeedForwardPolicy / WorldModelPolicy.

load_pretrained

load_pretrained(
    name: str, cache_dir: str = None, extra_args=None
)

Load a model from a local checkpoint or a HuggingFace repository.

Supported formats for name:

.pt file — path to a specific checkpoint file. A config.json must live in the same directory.

python model = load_pretrained('my_run/weights_epoch_10.pt')
Folder — path to a directory containing exactly one .pt file and a config.json.

python model = load_pretrained('my_run/')
HuggingFace repo (<user>/<repo>) — loaded from the local cache if already present, otherwise fetched from HF.

python model = load_pretrained('nice-user/my-worldmodel')

All local paths are resolved relative to <cache_dir>/checkpoints/.

Use the CLI to list available model checkpoints:

swm checkpoints
swm checkpoints pusht  # filter by name

Policy

PlanConfig dataclass

BasePolicy

get_action

set_env

_prepare_info

RandomPolicy

get_action

ExpertPolicy

get_action

FeedForwardPolicy

get_action

WorldModelPolicy

get_action

[ Utils ]

load_pretrained

PlanConfig `dataclass`