World-Model Checkpoint Loader¶

Several marimo notebooks reload world_model.pt long after the checkpoint was trained.

Why a shared loader exists¶

World-model training configurations intentionally use different hidden widths across environments. Checkpoints in exports/**/wm/world_model.pt therefore cannot be reconstructed from the constructor default alone.

Reconstructing the model with current defaults and then calling load_state_dict(...) is brittle: PyTorch will fail on backbone and head shape mismatches even when the checkpoint itself is valid.

Implementation¶

• Shared helper: src/world_model/mlp/checkpoint.py
• Public entrypoint: load_world_model_mlp(...)
• The helper loads the raw state dict, infers the hidden width from backbone.0.weight, rebuilds the WorldModelMLP with that width, then loads the checkpoint and switches the model to eval mode

Intended usage¶

Use load_world_model_mlp(...) in notebooks whenever a saved world_model.pt is reloaded for:

• opponent-model training setup
• world-model evaluation/report cells
• any explicit checkpoint-only evaluation path

Shielded experiment notebooks should not load world_model.pt merely to validate shielding artifacts. Learned shielding consumes wm/env_transition_graph.pkl and om/iop_stack.pt; exact true shielding consumes the cached true-transition graph under true_shield/.

Training code should still create the model explicitly with the chosen hyperparameters; the helper is for checkpoint reloads, not for initial model construction.