
A world-model planner keeps choosing an unreliable shortcut. You increase its uncertainty penalty. It keeps choosing the shortcut.
The scoring code may be correct. The missing information may be in the uncertainty estimate itself.
This is the useful boundary exposed by Choose Actions Under Model Uncertainty, a small TensorDrill exercise. Given a mean cost and a cost variance for each action, implement one decision rule: minimize the mean plus a nonnegative coefficient times the standard deviation.
We can implement the entire exercise, then construct a case where the rule helps and another where no setting of that coefficient repairs the decision. Every number below comes from supplied toy data, not a trained-model benchmark.
Implement the actual decision rule
For an action a, the score is:
score(a) = mean_cost(a) + risk_weight * sqrt(cost_variance(a))
The exercise supplies matching action dictionaries, nonnegative variances and a nonnegative risk_weight. Exact score ties go to the lexicographically smallest action name. Empty dictionaries return None.
Here is the complete reference implementation:
import math
def risk_sensitive_action(means, variances, risk_weight):
if not means:
return None
return min(
sorted(means),
key=lambda a: means[a] + risk_weight * math.sqrt(variances[a])
)
The square root is part of the contract. If costs are measured in seconds, their variances are measured in seconds squared. Standard deviation brings the penalty back to cost units; substituting variance also changes which action wins.
This helper selects an action. It does not estimate uncertainty, learn a world model or certify a probability of failure.
First, a case where the penalty helps
Suppose a model reports these costs:
| Action | Predicted mean | Predicted variance | Standard deviation | True cost |
|---|---|---|---|---|
| risky | 0 | 16 | 4 | 5 |
| safe | 2 | 0 | 0 | 2 |
The true cost is used only for evaluation. The selector receives the other columns.
means = {"risky": 0.0, "safe": 2.0}
variances = {"risky": 16.0, "safe": 0.0}
truth = {"risky": 5.0, "safe": 2.0}
for weight in (0.0, 0.5, 1.0):
chosen = risk_sensitive_action(means, variances, weight)
regret = truth[chosen] - min(truth.values())
print(f"weight={weight:.1f}: {chosen}, true cost={truth[chosen]:.1f}, "
f"regret={regret:.1f}")
Output:
weight=0.0: risky, true cost=5.0, regret=3.0
weight=0.5: risky, true cost=5.0, regret=3.0
weight=1.0: safe, true cost=2.0, regret=0.0
At weight 0.5, both scores equal 2. The result is still risky because of the specified name-based tie break. The switch occurs strictly above 0.5.
At weight 1, uncertainty changes the choice and lowers the executed cost in this constructed example. The arithmetic and the decision both look encouraging.
Now remove the warning signal
Keep the same rule, but change the supplied predictions:
| Action | Predicted mean | Predicted variance | True cost |
|---|---|---|---|
| shortcut | 0 | 0 | 4 |
| detour | 1 | 1 | 1 |
The shortcut is confidently wrong. The detour carries the larger reported uncertainty.
biased_means = {"shortcut": 0.0, "detour": 1.0}
biased_variances = {"shortcut": 0.0, "detour": 1.0}
biased_truth = {"shortcut": 4.0, "detour": 1.0}
for weight in (0.0, 1.0, 10.0):
chosen = risk_sensitive_action(biased_means, biased_variances, weight)
regret = biased_truth[chosen] - min(biased_truth.values())
assert chosen == "shortcut" and regret == 3.0
print(f"weight={weight:.1f}: {chosen}, regret={regret:.1f}")
Output:
weight=0.0: shortcut, regret=3.0
weight=1.0: shortcut, regret=3.0
weight=10.0: shortcut, regret=3.0
This is stronger than an unsuccessful parameter sweep. For every nonnegative weight, the shortcut's score is 0, while the detour's score is 1 + weight. The shortcut always wins.
Increasing the coefficient makes the reported uncertainty more influential. It cannot introduce a warning that the estimates do not contain.

Original illustration of the two supplied-data examples. The curves are selector scores, not confidence bounds measured from a learned model.
Agreement does not establish correctness
If several model members all predict the shortcut's cost as zero, their disagreement can be zero even when the true cost is four. Shared training data, architecture or extrapolation behavior can leave them making the same mistake.
There is another possibility: a supplied variance may describe randomness in the outcome, rather than uncertainty about the predicted mean. A nearly deterministic outcome can still have the wrong predicted value. Before tuning a penalty, establish what the variance represents.
The mean-plus-standard-deviation rule is therefore a useful heuristic under its stated inputs. Calling the resulting score an upper confidence bound does not, by itself, provide statistical coverage. Such an interpretation requires additional distributional or calibration assumptions.
Tests that separate arithmetic from model quality
These checks exercise the function contract independently of whether its inputs describe reality well:
# Standard deviation, not variance: scores are 3 and 4.
assert risk_sensitive_action(
{"a": 0, "b": 4}, {"a": 9, "b": 0}, 1.0
) == "a"
# Zero coefficient removes the uncertainty penalty.
assert risk_sensitive_action(
{"a": 1, "b": 2}, {"a": 100, "b": 0}, 0.0
) == "a"
# Dictionary insertion order must not decide an exact tie.
assert risk_sensitive_action(
{"z": 1, "a": 1}, {"z": 0, "a": 0}, 2.0
) == "a"
assert risk_sensitive_action({}, {}, 1.0) is None
print("All selector checks passed")
Sorting the A action names takes O(A log A) time and O(A) auxiliary space. The helper assumes the exercise's valid inputs; input validation belongs to its caller.
For a learned system, evaluate another layer: freeze the predictions, vary the coefficient, and measure actual execution outcomes on held-out tasks. Plot reported uncertainty against realized error. Separate an implementation error from a model that is confidently inaccurate. If an oracle optimum is unavailable, name the comparator instead of reporting its gap as exact optimal regret.
In a stochastic environment, a single realized cost is noisy evidence about an expected cost. The deterministic tables here make regret exactly computable; a real evaluation needs enough repeated executions for its intended estimate.
Research context and the corresponding exercise
PETS studies probabilistic ensembles and trajectory sampling for model-based reinforcement learning. It is relevant background for uncertainty-aware planning, but the small selector above does not implement PETS or its uncertainty propagation.
The 2026 preprint Imperfect World Models are Exploitable studies reversals of policy preference between approximate and true transition models. Our action-cost tables isolate a related decision failure; they do not reproduce the paper's construction or establish its formal results.
I maintain TensorDrill. The implementation here is the site's existing uncertainty-selection exercise, whose visual walkthrough lets you inspect how the score changes. Its tests check the specified selection rule. They cannot establish that a supplied uncertainty estimate is trustworthy.
The question to ask after “did we compute the penalty correctly?” is “does the penalty identify the mistakes that change our decisions?”
Comments
Join the discussionThe point about the square root being part of the contract is underrated, since using variance instead of standard deviation makes the effective risk weight scale with the cost units and quietly changes which action wins. The second example nails the real limitation: if the model reports low variance for the risky action, no risk_weight can rescue the decision. A calibration check like comparing predicted standard deviation against the observed absolute error on held-out rollouts would be a natural follow-up exercise to tell those two failure modes apart.

