From 4a2cb600ab1598ee7c42f73a39a59033edb6de50 Mon Sep 17 00:00:00 2001 From: yoch Date: Wed, 14 Oct 2026 18:00:28 +0000 Subject: [PATCH] Batch the backward pass instead of vmapping it backward() was an unbatched function under jax.vmap, and its while_loop carries the whole tree. Under vmap the loop predicate is batched, so every iteration selects between the old and the new carry for every array of the tree: a backup step costs O(num_simulations) and a whole search O(num_simulations^2). The batched version uses a scalar predicate, lets finished rows rewrite their current values and only scatters the entries on the paths. The trees are unchanged; the new tests compare them with those built by the previous vmapped implementation, including rows at different depths. --- mctx/_src/search.py | 63 ++++++++++------- mctx/_src/tests/search_test.py | 122 +++++++++++++++++++++++++++++++++ 2 files changed, 161 insertions(+), 24 deletions(-) create mode 100644 mctx/_src/tests/search_test.py diff --git a/mctx/_src/search.py b/mctx/_src/search.py index e844c73..a6bbaa1 100644 --- a/mctx/_src/search.py +++ b/mctx/_src/search.py @@ -244,50 +244,65 @@ def expand( tree.action_from_parent, action, next_node_index)) -@jax.vmap def backward( tree: Tree[T], - leaf_index: chex.Numeric) -> Tree[T]: + leaf_index: chex.Array) -> Tree[T]: """Goes up and updates the tree until all nodes reached the root. + The whole batch is processed by one loop with a scalar predicate (any row + still below the root). Vmapping an unbatched backward would make the loop + predicate batched, and JAX would then select between the old and the new + loop state for every array of the tree at each iteration. Here the rows that + have reached the root rewrite their current values, and each iteration only + scatters the entries on the paths. + Args: - tree: the MCTS tree state to update, without the batch size. - leaf_index: the node index from which to do the backward. + tree: the MCTS tree state to update, with a batch dimension. + leaf_index: the node indices from which to do the backward. Shape `[B]`. Returns: Updated MCTS tree state. """ + batch_range = jnp.arange(tree_lib.infer_batch_size(tree)) def cond_fun(loop_state): _, _, index = loop_state - return index != Tree.ROOT_INDEX + return jnp.any(index != Tree.ROOT_INDEX) def body_fun(loop_state): # Here we update the value of our parent, so we start by reversing. tree, leaf_value, index = loop_state - parent = tree.parents[index] - count = tree.node_visits[parent] - action = tree.action_from_parent[index] - reward = tree.children_rewards[parent, action] - leaf_value = reward + tree.children_discounts[parent, action] * leaf_value - parent_value = ( - tree.node_values[parent] * count + leaf_value) / (count + 1.0) - children_values = tree.node_values[index] - children_counts = tree.children_visits[parent, action] + 1 + active = index != Tree.ROOT_INDEX + # Rows already at the root keep pointing at the root and leave the tree + # unchanged. + parent = jnp.where(active, tree.parents[batch_range, index], + Tree.ROOT_INDEX) + action = jnp.where(active, tree.action_from_parent[batch_range, index], 0) + count = tree.node_visits[batch_range, parent] + reward = tree.children_rewards[batch_range, parent, action] + discount = tree.children_discounts[batch_range, parent, action] + new_leaf_value = reward + discount * leaf_value + node_value = tree.node_values[batch_range, parent] + parent_value = (node_value * count + new_leaf_value) / (count + 1.0) + children_values = tree.children_values[batch_range, parent, action] + children_counts = tree.children_visits[batch_range, parent, action] tree = tree.replace( - node_values=update(tree.node_values, parent_value, parent), - node_visits=update(tree.node_visits, count + 1, parent), - children_values=update( - tree.children_values, children_values, parent, action), - children_visits=update( - tree.children_visits, children_counts, parent, action)) - - return tree, leaf_value, parent + node_values=tree.node_values.at[batch_range, parent].set( + jnp.where(active, parent_value, node_value)), + node_visits=tree.node_visits.at[batch_range, parent].set( + count + active), + children_values=tree.children_values.at[ + batch_range, parent, action].set( + jnp.where(active, tree.node_values[batch_range, index], + children_values)), + children_visits=tree.children_visits.at[ + batch_range, parent, action].set(children_counts + active)) + + return tree, jnp.where(active, new_leaf_value, leaf_value), parent leaf_index = jnp.asarray(leaf_index, dtype=jnp.int32) - # pyrefly: ignore[bad-index] - loop_state = (tree, tree.node_values[leaf_index], leaf_index) + loop_state = (tree, tree.node_values[batch_range, leaf_index], leaf_index) tree, _, _ = jax.lax.while_loop(cond_fun, body_fun, loop_state) return tree diff --git a/mctx/_src/tests/search_test.py b/mctx/_src/tests/search_test.py new file mode 100644 index 0000000..fa4c56a --- /dev/null +++ b/mctx/_src/tests/search_test.py @@ -0,0 +1,122 @@ +# Copyright 2021 DeepMind Technologies Limited. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tests for `search.py`.""" +from unittest import mock + +from absl.testing import absltest +from absl.testing import parameterized +import jax +import jax.numpy as jnp +import mctx +from mctx._src import search +import numpy as np + + +@jax.vmap +def _vmapped_backward(tree, leaf_index): + """The unbatched backward of earlier versions, vmapped: the reference.""" + + def cond_fun(loop_state): + _, _, index = loop_state + return index != search.Tree.ROOT_INDEX + + def body_fun(loop_state): + tree, leaf_value, index = loop_state + parent = tree.parents[index] + count = tree.node_visits[parent] + action = tree.action_from_parent[index] + reward = tree.children_rewards[parent, action] + leaf_value = reward + tree.children_discounts[parent, action] * leaf_value + parent_value = ( + tree.node_values[parent] * count + leaf_value) / (count + 1.0) + children_values = tree.node_values[index] + children_counts = tree.children_visits[parent, action] + 1 + tree = tree.replace( + node_values=search.update(tree.node_values, parent_value, parent), + node_visits=search.update(tree.node_visits, count + 1, parent), + children_values=search.update( + tree.children_values, children_values, parent, action), + children_visits=search.update( + tree.children_visits, children_counts, parent, action)) + return tree, leaf_value, parent + + leaf_index = jnp.asarray(leaf_index, dtype=jnp.int32) + loop_state = (tree, tree.node_values[leaf_index], leaf_index) + tree, _, _ = jax.lax.while_loop(cond_fun, body_fun, loop_state) + return tree + + +def _hashed_recurrent_fn(params, rng_key, action, embedding): + """A deterministic environment whose outputs depend on the whole path.""" + del params, rng_key + embedding = (embedding * 31 + action + 7) % 10007 + h = embedding.astype(jnp.float32) + num_actions = 5 + return mctx.RecurrentFnOutput( + reward=jnp.sin(h), + discount=jnp.where(embedding % 11 == 0, 0.0, -0.9 + 0.1 * jnp.cos(h)), + prior_logits=jnp.cos(h[:, None] * jnp.arange(1, num_actions + 1)), + value=jnp.sin(2.0 * h)), embedding + + +def _run(policy_fn, batch_size, num_simulations, max_depth, seed): + rng = np.random.default_rng(seed) + num_actions = 5 + root = mctx.RootFnOutput( + prior_logits=jnp.asarray( + rng.normal(size=(batch_size, num_actions)), jnp.float32), + value=jnp.asarray(rng.normal(size=batch_size), jnp.float32), + embedding=jnp.asarray(rng.integers(0, 10007, batch_size), jnp.int32)) + invalid_actions = jnp.asarray( + rng.random((batch_size, num_actions)) < 0.2, jnp.float32) + invalid_actions = invalid_actions.at[:, 0].set(0.0) + + @jax.jit + def run(root, invalid_actions): + return policy_fn( + params=(), rng_key=jax.random.PRNGKey(seed), root=root, + recurrent_fn=_hashed_recurrent_fn, num_simulations=num_simulations, + invalid_actions=invalid_actions, max_depth=max_depth) + + return run(root, invalid_actions) + + +class SearchTest(parameterized.TestCase): + + @parameterized.named_parameters( + ("muzero", mctx.muzero_policy, None), + ("muzero_max_depth", mctx.muzero_policy, 3), + ("gumbel_muzero", mctx.gumbel_muzero_policy, None), + ("gumbel_muzero_max_depth", mctx.gumbel_muzero_policy, 4), + ) + def test_batched_backward_matches_vmapped_backward(self, policy_fn, + max_depth): + """The search builds the same tree as with the earlier vmapped backward.""" + out = _run(policy_fn, 16, 48, max_depth, seed=1) + with mock.patch.object(search, "backward", _vmapped_backward): + expected = _run(policy_fn, 16, 48, max_depth, seed=1) + jax.tree.map(np.testing.assert_allclose, out, expected) + + def test_backward_on_rows_at_different_depths(self): + """Rows whose leaves are at different depths, the root included.""" + tree = _run(mctx.muzero_policy, 8, 32, None, seed=2).search_tree + leaves = jnp.asarray([0, 1, 5, 31, 32, 17, 0, 9], jnp.int32) + got = search.backward(tree, leaves) + expected = _vmapped_backward(tree, leaves) + jax.tree.map(np.testing.assert_allclose, got, expected) + + +if __name__ == "__main__": + absltest.main()