Skip to content

Preserve visited Q-values with tiny prior probabilities - #120

Open
sylvesterkaczmarek wants to merge 1 commit into
google-deepmind:mainfrom
sylvesterkaczmarek:fix/normalize-mixed-value-weights-first
Open

sylvesterkaczmarek wants to merge 1 commit into
google-deepmind:mainfrom
sylvesterkaczmarek:fix/normalize-mixed-value-weights-first

Conversation

@sylvesterkaczmarek

Copy link
Copy Markdown

Normalize the visited prior weights before multiplying by Q-values in _compute_mixed_value. Multiplying a floored tiny prior by a small Q-value first can underflow to zero, losing the visited action's contribution.

For raw value 0.5, Q-values [0.25, 1.0], visits [1, 0], and priors [0, 1], the existing float32 result is 0.25 instead of 0.375. This also affects completed Q-values used by the public search transform.

The one-expression fix preserves the prior floor, no-visits guard, mixing formula, and output dtypes. No API or search-state changes.

Validation

macOS arm64 CPU, Python 3.12.11, JAX 0.11.2, Chex 0.1.92:

  • New regression module: 13 failures and 3 controls pass on unchanged 88f9205; all 16 pass after the fix.
  • Float16/bfloat16/float32, positive/negative values, multiple visited actions, eager/JIT execution, no-visits/ordinary-prior controls, and the public completed-Q transform on a real Tree.
  • Full suite: 40 passed from the _testing directory prescribed by test.sh, including existing search-tree fixtures.
  • Built the source distribution and wheel; installed outside the checkout and reran the full suite: 40 passed.
  • Flake8's repository gate, changed-file Ruff, source Pylint, new-test Pylint/Pyink, compilation, and git diff --check pass.

Pylint over all existing tests reports three line-length findings, identical on unchanged upstream; no existing lint rule or test was weakened. GPU/TPU execution and end-to-end agent training were not run.

@sylvesterkaczmarek
sylvesterkaczmarek marked this pull request as ready for review September 23, 2026 09:16
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant