Reward-loss scaling¶

The MLP world model predicts next-state bits and raw per-agent rewards from a shared representation. Its training objective is next-state BCE plus reward MSE. These losses are not invariant to reward units: changing an environment from rewards near 1 to rewards near 10 multiplies the raw MSE and its gradient by roughly 100, even though the transition-prediction task is unchanged.

train_world_model_mlp(..., reward_scale=s) therefore computes reward MSE on prediction / s and target / s. The reward head, checkpoints, inference API, graphs, and reported reward errors all remain in raw environment units. Only the multi-task loss is normalized. The default is 1.0, preserving existing behaviour for unit-scale environments.

Pursuit uses interception and collision rewards of magnitude 10, so its WM notebook sets reward_scale=10.0. Without this normalization, a reward-aware rebuild reduced held-out cell accuracy from approximately 0.996 to 0.991 and raised exhaustive transition-kernel TV from approximately 0.134 to 0.284; the raw reward MSE dominated the shared backbone. Any change to this scale must be evaluated on both raw reward prediction and exhaustive dynamics diagnostics, rather than selected from RL returns.