A Correct CEM Update Can Assign Half Its Probability to Bad Plans

[ym p]

A planner keeps two excellent action sequences. It fits a new sampling distribution to those elites. Half the probability then goes to sequences that were explicitly rejected.

No indexing error is required. The result follows from the distribution family being fitted.

We can see it by solving one existing TensorDrill problem, Update a Categorical CEM Planning Distribution. The task is a single unsmoothed update: select low-cost sequences, then calculate each action's elite frequency separately at each time step.

The counterexample below uses exact supplied costs. It does not require an inaccurate world model. That makes it a useful way to isolate a planning assumption before adding prediction error.

Two good plans, four possible combinations​

Consider two action IDs, 0 and 1, and a horizon of two steps. The task rewards changing action between steps:

Candidate	Supplied cost	Selected among the two elites?
[0, 0]	10	No
[0, 1]	0	Yes
[1, 0]	0	Yes
[1, 1]	10	No

These numbers define a constructed two-step task, rather than a measured robotics result.

The two elites agree on a relationship: the second action must differ from the first. They disagree about which individual action should appear at either step.

That distinction is exactly what a per-time-step frequency table loses.

Implement the exercise without changing its contract​

The function receives nonempty, equal-length candidate lists, their costs, an elite count and the number of actions. It returns one categorical probability row per time step.

Lower cost is better. Preserve input order when costs tie. Do not add smoothing, draw another population or repeat the optimization loop.

def cem_elite_frequencies(candidates, costs, elite_count, n_actions):
    elite_indices = sorted(range(len(candidates)), key=lambda i: costs[i])[:elite_count]
    horizon = len(candidates[0])
    frequencies = [[0.0] * n_actions for _ in range(horizon)]
    for i in elite_indices:
        for t, action in enumerate(candidates[i]):
            frequencies[t][action] += 1.0 / elite_count
    return frequencies

candidates = [[0, 0], [0, 1], [1, 0], [1, 1]]
costs = [10.0, 0.0, 0.0, 10.0]
probabilities = cem_elite_frequencies(candidates, costs, 2, 2)
assert probabilities == [[0.5, 0.5], [0.5, 0.5]]
print(probabilities)

Output:

[[0.5, 0.5], [0.5, 0.5]]

At time step zero, one elite starts with 0 and one starts with 1. At time step one, one uses 1 and the other uses 0. Both returned rows are correct.

The separate list comprehension for each row matters. Repeating one inner list would alias the rows and mix the time steps during accumulation.

Turn the marginals into a joint distribution​

The exercise stops at the frequency table. To examine its consequence, suppose a surrounding planner samples each time step independently from that table. A sequence's probability is then the product of the corresponding entries.

For example:

P([0, 0]) = P(action 0 at step 0) * P(action 0 at step 1)
          = 0.5 * 0.5 = 0.25

Enumerating all four sequences gives the exact result without Monte Carlo noise:

import itertools
import math

def sequence_probability(sequence, rows):
    return math.prod(rows[t][action] for t, action in enumerate(sequence))

support = list(itertools.product(range(2), repeat=2))
cost_by_sequence = dict(zip(map(tuple, candidates), costs))
mass = {s: sequence_probability(s, probabilities) for s in support}

assert math.isclose(math.fsum(mass.values()), 1.0)
bad_mass = math.fsum(p for s, p in mass.items() if cost_by_sequence[s] == 10.0)
expected_cost = math.fsum(p * cost_by_sequence[s] for s, p in mass.items())
assert bad_mass == 0.5
assert expected_cost == 5.0

for sequence in support:
    print(f"{list(sequence)}: probability={mass[sequence]:.2f}, "
          f"cost={cost_by_sequence[sequence]:.0f}")
print(f"bad mass={bad_mass:.2f}; expected cost={expected_cost:.2f}")

Output:

[0, 0]: probability=0.25, cost=10
[0, 1]: probability=0.25, cost=0
[1, 0]: probability=0.25, cost=0
[1, 1]: probability=0.25, cost=10
bad mass=0.50; expected cost=5.00

The uniform distribution over the two whole elite sequences has expected cost zero. The fitted product distribution has expected cost five.

Those are different distributions, even though their per-step action frequencies are identical.

Exact probabilities for the supplied four-sequence example. “Bad” means cost 10 in this task, not a general property of repeating an action.

What the fit actually preserves​

The categorical maximum-likelihood fit preserves the elites' marginals at each step. It does not preserve their temporal dependence.

In the empirical elite distribution, the second action is determined by the first. In the product distribution, it is independent of the first. Matching both individual action histograms cannot recover that relationship.

This is a modeling assumption, not a reason to fail the exercise's implementation. Returning a list of whole elite sequences would solve a different task.

The distinction also explains why independently taking the largest-probability action at each step needs care. Here every entry ties. A smallest-ID tie rule would produce [0, 0], although that entire sequence was rejected. A practical planner may instead execute its best evaluated candidate. The behavior depends on the surrounding planner, which this helper does not specify.

Fewer elites changes the trade-off​

Set the elite count to one:

single_elite = cem_elite_frequencies(candidates, costs, 1, 2)
assert single_elite == [[1.0, 0.0], [0.0, 1.0]]

# Equal costs preserve candidate input order: [0, 1] comes before [1, 0].
assert sequence_probability((0, 1), single_elite) == 1.0
assert sequence_probability((0, 0), single_elite) == 0.0

# The function must keep temporal rows distinct and support horizon zero.
assert cem_elite_frequencies([[0, 1], [1, 0]], [0, 1], 1, 2) == [
    [1.0, 0.0], [0.0, 1.0]
]
assert cem_elite_frequencies([[], []], [1, 0], 1, 2) == []
print("All categorical-update checks passed")

For this exact population, one elite removes the bad combinations. It also makes the unsmoothed distribution deterministic. Repeated sampling from it cannot rediscover alternatives with zero probability.

Consequently, reducing the elite count is not a general solution. Neither is smoothing a guarantee that dependencies are restored: keeping probabilities away from zero addresses exploration, while preserving relationships between time steps requires a suitable joint model or a different sampling scheme.

Selecting an elite sequence as a whole would preserve the dependence in this particular sample, but it would only resample existing elites. It is an explanatory comparison, not the required return value or a universally better optimizer.

Where world-model error enters​

The example has deliberately separated the optimizer from the dynamics model. If supplied rollout costs are wrong, elite selection gains another failure mode: the fit can faithfully increase the frequency of actions favored by those errors.

The 2026 preprint Imperfect World Models are Exploitable studies preference reversals caused by approximate transition models. That is separate from the independence issue demonstrated here. Combining the two can be troublesome, but the four-sequence table is not a reproduction of the paper's results.

The original-author chapter The Cross-Entropy Method for Optimization explains the broader sampling, elite-selection and distribution-fitting procedure. This article implements only the exercise's unsmoothed categorical fitting step.

Its time cost is O(N log N + E H + H A), for N candidates, E elites, horizon H and A actions. Auxiliary space is O(N + H A).

I maintain TensorDrill. The existing exercise and its visual editorial let you inspect the selected elite rows and the resulting time-step frequencies. The code above follows that contract exactly; the joint-probability calculation explains what a surrounding independent sampler would do with its output.

When checking a CEM implementation, ask two questions separately: did it fit the intended distribution correctly, and can that distribution represent the dependencies the good plans need?