World Model API¶

World-model APIs cover transition encoding, rollout collection, the default masked categorical successor kernel, the product-Bernoulli baseline, count models, checkpoint loading, prediction, and evaluation.

Encoding and Rollout Data¶

src.world_model.common.encode.encode_transition ¶

encode_transition(
    t, *, n_agents: int | None = None, action_size: int
)

Layered encoding:

x = [ current layered obs (float 0/1), concatenated per-agent one-hot actions (uniform size) ]

y = [ next layered obs (float 0/1), rewards for each agent (self + opps) ]

Transition tuple

t = (s, a_self, a_opps, s_next, r_self, r_opps)

s, s_next: flattened layered obs of length cells * n_channels, values in {0,1}
a_*: action indices (ints) or None; invalid/None → all-zeros one-hot
r_*: rewards; None → 0.0

Args¶

n_agents: total number of agents including self. action_size: discrete action space size (uniform for all agents).

src.world_model.common.encode.encode_state_action ¶

encode_state_action(
    state_flat: ndarray,
    a_self: Any,
    a_opps: Sequence[Any] | Any,
    *,
    n_agents: int | None = None,
    action_size: int,
) -> np.ndarray

Build the x vector (layered current obs + concatenated one-hot actions, uniform size).

src.world_model.common.buffer_collector.BufferWrapper ¶

BufferWrapper(
    env,
    *,
    agent_id,
    n_steps,
    seed=None,
    observe_opponent_rewards: bool = True,
)

Bases: BaseParallelWrapper

get_buffer ¶

get_buffer() -> list[Transition]

Return a copy of the current buffer as a list (FIFO order).

src.world_model.common.agents.other_agents ¶

other_agents(
    all_agents: Iterable[AgentID], pid: int | AgentID
) -> tuple[AgentID, list[AgentID]]

Return (pid_id, opponents) given either an index or an agent id.

src.world_model.sampling.ppo_sampling.collect_rollouts_ppo ¶

collect_rollouts_ppo(
    env,
    n_steps: int,
    *,
    seed: int | None = None,
    pid=0,
    observe_opponent_rewards: bool = True,
    novelty_beta: float = 0.0,
)

Collect joint transitions using the JAX IPPO trainer.

novelty_beta > 0 adds a count-based novelty bonus to every learner's reward (CountNoveltyBonus); the recorded transitions keep the environment's rewards.

Returns:

Name	Type	Description
`buffer`		recorded transitions from the wrapped real env
`done_buffer`		per-transition done flags for the tracked agent
`models`		Torch policy/value models converted from the trained JAX parameters
`labelled_history`		env history, when the wrapped env exposes it

src.world_model.sampling.validation_buffer.collect_rollouts_random ¶

collect_rollouts_random(
    env,
    *,
    n_steps: int,
    pid: int | AgentID = 0,
    seed: Optional[int] = 0,
) -> Tuple[List[Transition], List[bool]]

Collect a replay buffer by running a seeded random policy in the real env.

Returns:

Name	Type	Description
`buffer`	`List[Transition]`	list[Transition] of length <= n_steps (exactly n_steps after warmup)
`done`	`List[bool]`	list[bool] aligned with the buffer

Successor Kernel¶

The default nominal world model (wm_variant=successor). See World Modelling.

src.world_model.successor.model.SuccessorEdgeModel ¶

SuccessorEdgeModel(
    in_dim: int,
    state_width: int,
    rewards_dim: int,
    hidden: int = 512,
    dropout: float = 0.0,
)

Bases: Module

Score and reward every legal successor of a (state, joint_action) group.

Inputs - x: [B, in_dim] state-action features from encode_state_action. Its leading state_width entries are the current state bits. - cand_bits: [B, K, state_width] candidate successor bits, padded to the batch's maximum fan-out K. - cand_mask: [B, K] boolean, true where the candidate is real.

Outputs - log_probs: [B, K], -inf on padding, logsumexp zero per row. - rewards: [B, K, rewards_dim] in model order [self, *opponents].

config ¶

config() -> dict

Architecture needed to rebuild this model from a checkpoint.

src.world_model.successor.dataset.build_successor_groups ¶

build_successor_groups(
    graph: Graph,
    *,
    n_agents: int,
    action_size: int,
    focal_idx: int = 0,
    state_bits_fn=None,
) -> SuccessorGroups

Enumerate Succ(s, a) and the per-group state-action features.

src.world_model.successor.dataset.resolve_transitions ¶

resolve_transitions(
    buffer: Sequence,
    groups: SuccessorGroups,
    *,
    n_agents: int,
    focal_idx: int = 0,
    buffer_focal_idx: int | None = None,
    allow_unresolved: bool = False,
    threshold: float = BIT_THRESHOLD,
) -> ResolvedTransitions

Resolve each buffered transition to a (group, candidate) pair.

Unresolved rows raise by default, mirroring prepare_om_training_inputs. With allow_unresolved=True they are dropped and counted; a resolution failure is never converted into probability mass.

src.world_model.successor.dataset.aggregate_counts ¶

aggregate_counts(
    resolved: ResolvedTransitions,
    indices: ndarray | None = None,
) -> SuccessorCounts

Aggregate a subset of resolved rows into per-group counts and reward sums.

src.world_model.successor.dataset.train_validation_split ¶

train_validation_split(
    n_rows: int,
    *,
    validation_fraction: float = 0.1,
    seed: int = 0,
) -> Tuple[np.ndarray, np.ndarray]

One shared row-level split, used identically by every arm of the study.

src.world_model.successor.train.train_successor_world_model ¶

train_successor_world_model(
    train_counts: SuccessorCounts,
    *,
    validation_counts: SuccessorCounts | None = None,
    hidden: int = 512,
    dropout: float = 0.0,
    epochs: int = 60,
    batch: int = 256,
    lr: float = 0.0003,
    reward_scale: float = 1.0,
    reward_weight: float = 1.0,
    early_stopping_patience: int | None = None,
    selection_metric: str = "row_nll",
    min_optimizer_steps: int = 0,
    max_optimizer_steps: int | None = None,
    seed: int = 0,
    device: device | None = None,
    export_dir: str | None = None,
    verbose: bool = True,
) -> tuple[SuccessorEdgeModel, SuccessorTrainingReport]

Fit :class:SuccessorEdgeModel on aggregated group counts.

When validation_counts is supplied the best epoch is restored at the end and early_stopping_patience stops once that many epochs pass without improvement. selection_metric chooses what "best" means:

"row_nll" The validation row NLL alone. This is what the Stage 1 gate ran under. "objective" The full validation objective, row NLL plus reward_weight times the scaled reward MSE -- the same quantity being minimised on the training split. Where the kernel is easy the NLL term saturates within an epoch or two, so selecting on it alone restores a checkpoint whose reward head is barely trained; on the deterministic environments the NLL is identically zero from epoch 1 and selection is vacuous. Selecting on the objective keeps the dynamics criterion intact while letting the reward head converge.

min_optimizer_steps and max_optimizer_steps budget optimisation in optimiser updates rather than epochs alone. Exact count aggregation collapses the buffer to the number of visited state-action groups, so an epoch is one update on a twenty-group matrix game and several hundred on a gridworld; "sixty epochs" therefore means nothing comparable across environments. When the floor is set, training continues past epochs -- and early stopping is suppressed -- until min_optimizer_steps updates have been taken, after which patience-based stopping governs as before. max_optimizer_steps is a hard cap. Both default to off, reproducing the epoch-budgeted protocol exactly.

src.world_model.successor.train.SuccessorTrainingReport `dataclass` ¶

SuccessorTrainingReport(
    epochs_run: int,
    best_epoch: int,
    best_validation_nll: float | None,
    train_nll: List[float] = list(),
    validation_nll: List[float] = list(),
    validation_reward_mse: List[float] = list(),
    validation_objective: List[float] = list(),
    selection_metric: str = "row_nll",
    best_validation_objective: float | None = None,
    early_stopped: bool = False,
    optimizer_steps: int = 0,
    best_optimizer_step: int = 0,
    steps_per_epoch: int = 0,
    min_optimizer_steps: int = 0,
    max_optimizer_steps: int | None = None,
    min_steps_floor_extended: bool = False,
    max_optimizer_steps_reached: bool = False,
)

src.world_model.successor.protocol.training_kwargs ¶

training_kwargs(**overrides: Any) -> Dict[str, Any]

Keyword arguments for train_successor_world_model at the frozen protocol.

Overrides are permitted so that a smoke test can shorten training, but any override changes the model away from the one the gate validated and should be confined to tests and smoke runs.

src.world_model.successor.predict.SuccessorKernel ¶

Bases: Protocol

Produce a normalised successor distribution and rewards per group.

predict_groups ¶

predict_groups(
    groups: SuccessorGroups, group_ids: Sequence[int]
) -> Tuple[List[np.ndarray], List[np.ndarray]]

Return (probabilities, rewards) lists aligned with group_ids.

probabilities[i] has the fan-out of group_ids[i] and sums to one. rewards[i] is [fanout, rewards_dim] in model order.

src.world_model.successor.predict.NeuralSuccessorKernel ¶

NeuralSuccessorKernel(
    model: SuccessorEdgeModel,
    *,
    batch: int = 256,
    device: device | None = None,
)

Adapt :class:SuccessorEdgeModel to the :class:SuccessorKernel contract.

src.world_model.successor.predict.GraphSuccessorKernel ¶

GraphSuccessorKernel(
    graph,
    groups: SuccessorGroups,
    *,
    n_agents: int,
    focal_idx: int = 0,
)

Read successor rows back out of an annotated joint graph.

This is the faithful way to score an arm: it measures exactly the kernel the shield, the imagination adapter and the transition-TV report will consume, whatever produced it. It is also the only way to put the product-Bernoulli baseline on the same footing as the categorical arms, since that baseline has no native row distribution of its own — its kernel only exists after candidate scoring and renormalisation.

Rewards are converted from the graph's environment-agent order back into the model order [self, *opponents] that the evaluation targets use.

src.world_model.successor.eval.evaluate_successor_kernel ¶

evaluate_successor_kernel(
    kernel,
    groups: SuccessorGroups,
    *,
    holdout: ResolvedTransitions | None = None,
    oracle: OracleRows | None = None,
    calibration_bins: int = 15,
) -> SuccessorEvaluation

Score a kernel against a held-out buffer and, where available, an oracle.

src.world_model.counts.model.CountSuccessorModel ¶

CountSuccessorModel(
    counts: SuccessorCounts,
    *,
    alpha: float = 0.0,
    prior: Sequence[ndarray] | str = UNIFORM_PRIOR,
)

Duck-typed :class:SuccessorKernel backed by empirical counts.

Parameters¶

counts: Aggregated per-group transition counts and reward sums. alpha: Dirichlet concentration. 0.0 gives the plain MLE. prior: "uniform", or a sequence of per-group probability vectors (the neural prior arm) aligned with counts.groups.group_keys.

row_probabilities ¶

row_probabilities(g: int) -> np.ndarray

Estimated P(. | s, a) for group g, summing to one.

row_rewards ¶

row_rewards(g: int) -> np.ndarray

Per-successor reward estimates with the explicit fallback ladder.

Product-Bernoulli Baseline¶

src.world_model.mlp.mlp.WorldModelMLP ¶

WorldModelMLP(
    in_dim: int,
    cells: int,
    n_channels: int,
    rewards_dim: int,
    hidden: int = 512,
    dropout: float = 0.0,
)

Bases: Module

World model with two heads

state_head: logits over layered next-state bits per cell/channel
reward_head: continuous rewards

Inputs - x: [B, in_dim] of flattened numeric features (e.g., current layered obs + action indices). Outputs - s_logits: [B, cells, n_channels] (raw logits; use BCEWithLogitsLoss) - r_pred: [B, rewards_dim]

src.world_model.mlp.train.train_world_model_mlp ¶

train_world_model_mlp(
    buffer,
    *,
    cells: int,
    n_channels: int,
    n_agents: int,
    action_size: int,
    epochs: int = 50,
    batch: int = 512,
    lr: float = 0.0003,
    hidden: int = 512,
    dropout: float = 0.0,
    export_dir: str = "exports/unknown",
    pos_weight: float | None = None,
    reward_scale: float = 1.0,
    val_indices: ndarray | None = None,
    early_stopping_patience: int | None = None,
    verbose: bool = True,
    save_checkpoint: bool = True,
)

Train with layered next-state targets (multi-label, BCE) and uniform per-agent one-hot actions.

reward_scale expresses one environment reward unit in loss-normalized coordinates. Predictions and checkpoints remain in raw reward units; only the reward-head residual is divided by this scale before taking MSE. This prevents large reward magnitudes from silently overwhelming the shared dynamics representation.

val_indices and early_stopping_patience are off by default, so the published training path is unchanged. The variant study switches them on to give every learned arm the same held-out epoch-selection rule on the same split: a representation comparison is only valid if the training procedure is held constant. When supplied, the epoch with the lowest validation loss is restored at the end.

Encoding (via encode_transition): - X: float32, shape [cellsn_channels + n_agentsaction_size] = flattened current obs + concatenated one-hot actions (self + opps) - Y: float32, shape [cells*n_channels + n_agents] = flattened next obs bits + rewards for each agent

src.world_model.mlp.checkpoint.infer_world_model_hidden ¶

infer_world_model_hidden(
    state_dict: Mapping[str, Tensor],
) -> int

Recover the hidden width used by a saved WorldModelMLP checkpoint.

src.world_model.mlp.checkpoint.load_world_model_mlp ¶

load_world_model_mlp(
    path: str | Path,
    *,
    in_dim: int,
    cells: int,
    n_channels: int,
    rewards_dim: int,
    device: device | None = None,
) -> WorldModelMLP

Load a WorldModelMLP checkpoint with its trained architecture.

The hidden width is inferred because notebook training configurations use environment-specific widths.

Baseline Prediction and Evaluation¶

src.world_model.mlp.predict.predict_probs_bits_rewards ¶

predict_probs_bits_rewards(
    model: WorldModelMLP,
    x: ndarray,
    *,
    cells: Optional[int] = None,
    n_channels: Optional[int] = None,
    threshold: float | ndarray = 0.5,
    return_flat_bits: bool = True,
    include_false: bool = False,
    return_flat_probs: bool = False,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]

One forward pass that returns

probs: [cells, C] of P(true) if include_false=False, else [cells, C, 2] with last dim [P(false), P(true)]. If return_flat_probs=True, reshape to [-1] or [-1, 2]. bits: uint8 next-state (thresholded) as flat if return_flat_bits else [cells, C]. rewards:[rewards_dim] float.

Args mirror the existing helpers; x must be the pre-encoded input vector.

src.world_model.mlp.eval.evaluate_world_model ¶

evaluate_world_model(
    model,
    buffer,
    *,
    done: Optional[Sequence[bool]] = None,
    cells: int,
    n_channels: int,
    n_agents: int,
    action_size: int,
    threshold: float = 0.5,
    batch_size: int = 4096,
    n_boot: int = 2000,
    boot_block: int = 128,
    seed: int | None = None,
    calibration_bins: int = 15,
    export_dir: str | Path | None = None,
) -> dict[str, Any]

src.world_model.mlp.eval.log_sample_predictions ¶

log_sample_predictions(
    model: WorldModelMLP,
    X: Tensor,
    Y: Tensor,
    *,
    cells: int,
    n_channels: int,
    k: int = 10,
    thr: float = 0.5,
)

Log a few samples with bit-accuracy over layered targets and reward preds.

Y is [N, cells*n_channels + rewards_dim]. First part is 0/1 layered bits (flattened).

src.world_model.mlp.eval.print_report_rich ¶

print_report_rich(report, digits=4, console=None)

src.world_model.mlp.eval.save_world_model_report ¶

save_world_model_report(
    report: dict[str, Any], path: str | Path
) -> None