Skip to content

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(t, *, n_agents=None, action_size)

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(state_flat, a_self, a_opps, *, n_agents=None, action_size)

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

src.world_model.common.buffer_collector.BufferWrapper

Bases: BaseParallelWrapper

get_buffer()

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

clear_buffer()

Manually clear the buffer.

src.world_model.common.agents.other_agents(all_agents, pid)

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

src.world_model.sampling.ppo_sampling.collect_rollouts_ppo(env, n_steps, *, seed=None, pid=0)

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-compatible exported policy/value models for downstream notebooks
labelled_history

env history, when the wrapped env exposes it

src.world_model.sampling.validation_buffer.collect_test_buffer_random(env, *, n_steps, pid=0, seed=0)

Collect a held-out test buffer by running a random policy in the real env.

Returns:

Name Type Description
test_buffer List[Transition]

list[Transition] of length <= n_steps (exactly n_steps after warmup)
test_done List[bool]

list[bool] aligned with test_buffer

Model and Training

src.world_model.mlp.mlp.WorldModelMLP

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(buffer, *, cells, n_channels, n_agents, action_size, epochs=50, batch=512, lr=0.0003, hidden=512, dropout=0.0, export_dir='exports/unknown', pos_weight=None)

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

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(state_dict)

Recover the hidden width used by a saved WorldModelMLP checkpoint.

src.world_model.mlp.checkpoint.load_world_model_mlp(path, *, in_dim, cells, n_channels, rewards_dim, device=None)

Load a WorldModelMLP checkpoint while preserving the architecture used at train time.

The hidden width is inferred from the saved parameters so notebooks can safely reload older checkpoints even if the constructor default changes later.

Prediction and Evaluation

src.world_model.mlp.predict.predict_probs_bits_rewards(model, x, *, cells=None, n_channels=None, threshold=0.5, return_flat_bits=True, include_false=False, return_flat_probs=False)

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.predict.predict_next_state_and_rewards(model, x, *, cells=None, n_channels=None, threshold=0.5, return_flat=True)

Backwards-compat: delegates to the consolidated single-forward helper.

src.world_model.mlp.predict.next_state_probabilities(model, x, *, cells=None, n_channels=None, return_flat=False, include_false=True)

Backwards-compat: delegates to the consolidated single-forward helper.

src.world_model.mlp.eval.evaluate_world_model(model, buffer, *, done=None, cells, n_channels, n_agents, action_size, threshold=0.5, batch_size=4096, n_boot=2000, boot_block=128, seed=None)

src.world_model.mlp.eval.log_sample_predictions(model, X, Y, *, cells, n_channels, k=10, thr=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(report, digits=4, console=None)