World Model API¶

World-model APIs cover transition encoding, rollout collection, neural model training, 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).

clear_buffer ¶

clear_buffer() -> None

Manually clear the buffer.

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

Collect joint transitions using the JAX IPPO trainer.

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

Model and Training¶

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

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.

Encoding (via encode_transition): - X: float32, shape [cellsn_channels + n_agentsaction_size] = flattened current obs + concatenated one-hot actions (self + opps) - Y: float32, shape [cellsn_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.

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