fix(llm): keep bandit accumulators linear when asymmetric scaling is off

[yurekami]

Summary

• AsymmetricUCB and ThompsonSampler derive use_exponential_scaling from exponential_base is not None and asymmetric_scaling, matching what update(), decay(), _impute_worst_reward(), and _have_obs_range() already key on.
• set_state() detects a state saved by the previous behaviour (linear mode with s = -inf), resets the reward statistics to the prior, and logs a warning saying why and how to start cleanly.
• Tests: tests/test_bandit_linear_scaling.py (6, both bandits) and tests/test_bandit_state_migration.py (6, both bandits). CHANGELOG.md entry.

Why

Both bandits set use_exponential_scaling = self.exponential_base is not None and use that flag alone to decide how to initialise and read the reward accumulator s:

if self.use_exponential_scaling:
    self.s = np.full(n, -np.inf, dtype=np.float64)   # log-space "empty sum"
...
def _mean(self):
    if self.use_exponential_scaling:
        return self.s - np.log(denom)

but update() and decay() only accumulate in log-space under the conjunction with asymmetric_scaling, and fall through to plain addition otherwise:

if self.use_exponential_scaling and self.asymmetric_scaling:
    self.s[i] = _logadd(self.s[i], contrib)
else:
    self.s[i] += r

So with asymmetric_scaling=False and exponential_base left at its default of 1.0, s starts at -inf for a log-space sum and is then updated with linear addition, and no reward can ever move it. Both bandits reach s = [-inf, -inf, -inf] after six real rewards, but what that does to selection differs by bandit:

• AsymmetricUCB: its scores are built from _mean(), so every arm scores the same and the posterior is uniform, [0.333, 0.333, 0.333] against [0.1, 0.8, 0.1] for the same rewards with exponential_base=None. Model selection is a coin flip.
• ThompsonSampler: its posterior samples alpha/beta, which are updated from the rescaled reward rather than from s, so it keeps learning; with the same rewards and seed the posterior is [0.05, 0.9, 0.05]. The damage is in the rescaling: the first decay() folds the -inf means into the adaptive observation range, _obs_min/_obs_max become -inf, _have_obs_range() flips to False, and reshift_in_range falls back to the sigmoid. The same reward of 1.5 maps to u = 0.667 before that decay and u = 0.818 after it, so the posterior updates change meaning mid-run.

The first version of this description said selection collapsed to uniform for both; that was accurate for UCB only, and is corrected above and in the CHANGELOG.

Every constructor kwarg is reachable through llm_dynamic_selection_kwargs, so asymmetric_scaling: false is one config line away and does not require touching exponential_base at all.

Why this direction rather than making the combination log-space end to end: a log-space sum is only defined when every shifted reward is >= 0, and r = max(r, 0.0) under asymmetric clamping is exactly what guarantees that. With asymmetric_scaling=False, negative shifted rewards are allowed, and _logexpm1(z) for z < 0 is the log of a negative number. Exponential scaling presupposes asymmetric clamping, so the linear branch is the only one that can be correct without it. Tying the flag to both makes __init__, _mean, and the print labels agree with the sites that already do this; the linear mode then produces identical numbers whether or not exponential_base is set.

Resumed state

Async resume loads bandit_state.pkl automatically (AsyncEvolutionRunner._load_bandit_state). A state saved by an affected run carries s = -inf for every arm; loaded into the fixed linear mode it would be interpreted as a linear sum, and s += reward can never move it, so the bandit would stay broken for the rest of the run.

set_state() in both bandits now calls _discard_invalid_linear_state() before restoring the observation range. The condition is not use_exponential_scaling and any(s == -inf), which no linear run can produce; in the default log-space mode -inf marks an unsampled arm and is left alone. When it fires, the reward statistics (s, divs, n_submitted, n_completed, the observation range, and for Thompson alpha/beta) go back to the prior and a warning is logged:

Bandit state contains s = -inf in linear mode, which can only come from a run where
exponential scaling was applied with asymmetric_scaling=False. Resetting reward
statistics to the prior; delete bandit_state.pkl to start cleanly.

The reward history in such a file is not recoverable (the sums were never accumulated), so a reset is the safe migration. Cost statistics and the baseline are kept; they were not affected. The CHANGELOG entry documents the reset.

Linked issue or context

No open issue. Found while reading shinka/llm/prioritization.py. The existing tests only exercise asymmetric_scaling=False together with exponential_base=None (test_bandit_persistence.py), where the two flags agree, so the mismatch never surfaced.

Testing

• Commands run:
  • uv run ruff check tests --exclude tests/file.py -> All checks passed
  • uv run ruff check shinka/llm/prioritization.py -> All checks passed
  • uv run ruff format --check shinka/llm/prioritization.py tests/test_bandit_linear_scaling.py tests/test_bandit_state_migration.py -> already formatted
  • uv run mypy --follow-imports=skip --ignore-missing-imports tests/test_bandit_linear_scaling.py tests/test_bandit_state_migration.py tests/conftest.py -> Success
  • uv run pytest -q -m "not requires_secrets" -> see results
• Results:

tests/test_bandit_linear_scaling.py (6): fail on main, pass with the change; the one that passes either way is the Thompson alpha/beta credit guard, which never read s, consistent with the corrected description above.

tests/test_bandit_state_migration.py (6): a state from a broken run is reset with the warning and learning resumes from the prior (UCB then prefers the best arm, Thompson credits it); a valid linear state loads unchanged with no warning; a default-mode state whose unsampled arms sit at -inf loads unchanged with no warning. The first of these fails without the migration code (no warning, s stays -inf).

Full non-secret suite locally: the same 3 pre-existing environment failures as before (wolframscript, claude binary, Windows subprocess timeout), everything else green; exact counts in the review thread.

Risks and compatibility

• Default configuration is unchanged. asymmetric_scaling=True, exponential_base=1.0 keeps use_exponential_scaling=True; every existing test passes unmodified.
• asymmetric_scaling=False runs change behaviour, from degenerate to correct. They now accumulate linearly, exactly as exponential_base=None already did.
• exponential_base is ignored when asymmetric_scaling=False. That was already true of update(); this makes it true of __init__ and _mean() too. No warning is emitted for it, because exponential_base has a non-None default and most users in this mode never set it; happy to add one gated on a non-default value if you would rather it be loud.
• Resumed state from an affected run is reset with a warning, as described above, rather than silently kept.
• The decay() RuntimeWarning noted earlier is handled separately in fix(llm): silence the discarded np.where branch in ThompsonSampler.decay #192.

Core evolution pipeline evidence

This changes which LLM the bandit samples. For AsymmetricUCB(asymmetric_scaling=False) with the default exponential_base (rewards 1.0, 5.0, 0.2, 1.2, 4.8, 0.3 across three arms, shift_by_baseline=False, shift_by_parent=False, one decay(0.95)):

Before:

s=[-inf -inf -inf]  mean=[-inf -inf -inf]  obs=(-inf,-inf)
posterior=[0.333 0.333 0.333]

After:

s=[2.09  9.31  0.475]  mean=[1.1  4.9  0.25]  obs=(0.203,5)
posterior=[0.1 0.8 0.1]

which is bit-identical to the exponential_base=None run. For ThompsonSampler the posterior was already [0.05, 0.9, 0.05] before; what changes is that s, _mean(), and the observation range are finite and the reward rescaling no longer collapses on the first decay. I have not run a full evolution benchmark.

Docs and UI

No docs or UI changes. asymmetric_scaling and exponential_base are not documented in docs/bandit_selection.md today; the CHANGELOG entry describes the change and the reset-on-load behaviour.

Prepared with Claude Code (agent-assisted); every command and result above was run locally.

🤖 Generated with Claude Code

https://claude.ai/code/session_01VZA548KwueeKS1BCRhowZq

[RobertTLange]

Thanks for digging into this — the UCB bug is real, the direction of the fix looks right, and the added tests pass locally.

Before merge, could you please address these points?

1. Correct the Thompson impact claims in the PR body and CHANGELOG. On current main, the reproduction does leave Thompson's s/means at -inf, but its posterior still learns through alpha/beta; with the supplied reward sequence and seed I get [0.05, 0.9, 0.05], not uniform. After decay(), the adaptive observation range becomes (-inf, -inf), so subsequent reward rescaling falls back to the sigmoid. That is still a real bug, but “model selection collapsed to uniform” is accurate for UCB, not Thompson.

2. Handle resumed state, or make the limitation explicit at runtime. Async resume automatically loads bandit_state.pkl. A state produced by the broken configuration contains s = -inf; after this patch it is interpreted linearly, and future s += reward updates can never recover. Preferred: detect this invalid linear-mode state and safely reset/migrate the affected learning state, with a regression test. If migration is intentionally out of scope because the lost reward history cannot be reconstructed, please emit a clear warning on load and document that affected users must discard bandit_state.pkl/start a fresh bandit state.

3. CHANGELOG attribution: please add Thanks @yurekami per the project convention.

Verification here: targeted bandit/persistence tests 28 passed; full non-secret suite 894 passed, 7 skipped, 1 deselected; Ruff clean.

[yurekami]

The decay() RuntimeWarning listed under "not touched" above now has its own PR: #192. It is based on main and touches different lines of prioritization.py, so the two merge in either order.

[yurekami]

Thank you for the careful read. All three points are addressed in a follow-up commit on this branch.

1. Thompson claims. You are right. I re-ran the reproduction and get the same [0.05, 0.9, 0.05]: Thompson's posterior samples alpha/beta, which are updated from the rescaled reward rather than from s, so it keeps learning. What breaks for Thompson is the rescaling, not the selection: after the first decay() the adaptive range becomes (-inf, -inf) and reshift_in_range falls back to the sigmoid, so the same reward of 1.5 maps to u = 0.667 before that decay and u = 0.818 after it. "Collapsed to uniform" was accurate for UCB only. The PR description and the CHANGELOG entry now say exactly that, and the description reports the two bandits separately.

2. Resumed state. Went with the preferred option. set_state() in both bandits now calls a shared _discard_invalid_linear_state() before restoring the observation range. The condition is not use_exponential_scaling and any(s == -inf), which no linear run can produce; in the default log-space mode -inf marks an unsampled arm and is left alone. When it fires, the reward statistics (s, divs, n_submitted, n_completed, the observation range, and for Thompson alpha/beta) go back to the prior, and a warning names the cause and says that deleting bandit_state.pkl starts cleanly. Cost statistics and the baseline are kept. I did not attempt to reconstruct the history, since the sums were never accumulated in the first place and there is nothing in the file to rebuild them from; the CHANGELOG entry documents the reset.

tests/test_bandit_state_migration.py covers it for both bandits: a state from a broken run is reset with the warning and learning resumes from the prior (UCB then prefers the best arm, Thompson credits it); a valid linear state loads unchanged with no warning; and a default-mode state whose unsampled arms sit at -inf loads unchanged with no warning. The first of these fails without the migration code.

3. CHANGELOG attribution. Added Thanks @yurekami. to the entry here, and to the entry in #192 for consistency.

Locally on the new head: the two bandit test files 12 passed, test_bandit_persistence.py unchanged, Ruff and mypy clean, and the non-secret suite green apart from the same 3 environment failures as before (wolframscript, claude binary, Windows subprocess timeout).