From a34ddc02421363c552f179a24fe41012baa1b365 Mon Sep 17 00:00:00 2001 From: carlosgmartin Date: Sun, 4 Oct 2026 23:03:25 -0400 Subject: [PATCH] Support num_simulations=0 in all policies. With no simulations, the search loop is skipped, and the MuZero and Stochastic MuZero policies act on the root prior, which excludes invalid actions, instead of the uniform visit probabilities of an unvisited root. --- mctx/_src/policies.py | 18 ++++++++--- mctx/_src/search.py | 7 +++-- mctx/_src/tests/policies_test.py | 52 ++++++++++++++++++++++++++++++++ 3 files changed, 71 insertions(+), 6 deletions(-) diff --git a/mctx/_src/policies.py b/mctx/_src/policies.py index 2266305..8574fea 100644 --- a/mctx/_src/policies.py +++ b/mctx/_src/policies.py @@ -111,8 +111,7 @@ def muzero_policy( loop_fn=loop_fn) # Sampling the proposed action proportionally to the visit counts. - summary = search_tree.summary() - action_weights = summary.visit_probs + action_weights = _get_action_weights(search_tree) action_logits = _apply_temperature( _get_logits_from_probs(action_weights), temperature) action = jax.random.categorical(rng_key, action_logits) @@ -364,8 +363,7 @@ def stochastic_muzero_policy( # Sampling the proposed action proportionally to the visit counts. search_tree = _mask_tree(search_tree, num_actions, 'decision') - summary = search_tree.summary() - action_weights = summary.visit_probs + action_weights = _get_action_weights(search_tree) action_logits = _apply_temperature( _get_logits_from_probs(action_weights), temperature) action = jax.random.categorical(rng_key, action_logits) @@ -373,6 +371,18 @@ def stochastic_muzero_policy( action=action, action_weights=action_weights, search_tree=search_tree) +def _get_action_weights(tree: search.Tree) -> chex.Array: + """Returns the root visit probabilities, or the prior if there are no visits. + + Unlike the uniform `visit_probs` of an unvisited root, the prior excludes + invalid actions. + """ + if tree.num_simulations == 0: + # pyrefly: ignore[bad-index] + return jax.nn.softmax(tree.children_prior_logits[:, tree.ROOT_INDEX]) + return tree.summary().visit_probs + + def _mask_invalid_actions(logits, invalid_actions): """Returns logits with zero mass to invalid actions.""" if invalid_actions is None: diff --git a/mctx/_src/search.py b/mctx/_src/search.py index 389d266..b07411e 100644 --- a/mctx/_src/search.py +++ b/mctx/_src/search.py @@ -109,8 +109,11 @@ def body_fun(sim, loop_state): tree = instantiate_tree_from_root(root, num_simulations, root_invalid_actions=invalid_actions, extra_data=extra_data) - _, tree = loop_fn( - 0, num_simulations, body_fun, (rng_key, tree)) + # Tracing `body_fun` requires a simulation (e.g., Gumbel MuZero indexes a + # table with one column per simulation). + if num_simulations > 0: + _, tree = loop_fn( + 0, num_simulations, body_fun, (rng_key, tree)) return tree diff --git a/mctx/_src/tests/policies_test.py b/mctx/_src/tests/policies_test.py index 549dca9..0bcd582 100644 --- a/mctx/_src/tests/policies_test.py +++ b/mctx/_src/tests/policies_test.py @@ -308,6 +308,58 @@ def test_gumbel_muzero_policy_without_invalid_actions(self): [[6, 2, 2, 7]]) np.testing.assert_array_equal(expected_visit_counts, summary.visit_counts) + def test_policies_without_simulations(self): + """With no simulations, the policies act on the prior alone.""" + root = mctx.RootFnOutput( + prior_logits=jnp.array([ + [0.0, -1.0, 2.0, 3.0], + [1.0, 5.0, 0.0, 2.0], + ]), + value=jnp.array([-5.0, 1.0]), + # Stochastic MuZero infers the batch size from the embedding. + embedding=jnp.zeros([2, 4]), + ) + invalid_actions = jnp.array([ + [0.0, 0.0, 0.0, 1.0], + [0.0, 1.0, 0.0, 0.0], + ]) + rewards = jnp.zeros_like(root.prior_logits) + recurrent_fn = _make_bandit_recurrent_fn(rewards) + decision_recurrent_fn, chance_recurrent_fn = ( + _make_bandit_decision_and_chance_fns(rewards, num_chance_outcomes=2)) + kwargs = dict( + params=(), + rng_key=jax.random.PRNGKey(0), + root=root, + num_simulations=0, + invalid_actions=invalid_actions) + muzero_kwargs = dict(kwargs, dirichlet_fraction=0.0, temperature=0.0) + policy_outputs = { + "muzero_policy": mctx.muzero_policy( + recurrent_fn=recurrent_fn, **muzero_kwargs), + "gumbel_muzero_policy": mctx.gumbel_muzero_policy( + recurrent_fn=recurrent_fn, gumbel_scale=0.0, **kwargs), + "stochastic_muzero_policy": mctx.stochastic_muzero_policy( + decision_recurrent_fn=decision_recurrent_fn, + chance_recurrent_fn=chance_recurrent_fn, + **muzero_kwargs), + } + + # The valid actions with the highest prior logits. + expected_action = jnp.array([2, 3], dtype=jnp.int32) + # The prior restricted to the valid actions. + expected_action_weights = jax.nn.softmax( + policies._mask_invalid_actions(root.prior_logits, invalid_actions)) + for name, policy_output in policy_outputs.items(): + with self.subTest(policy=name): + np.testing.assert_array_equal(expected_action, policy_output.action) + np.testing.assert_allclose(expected_action_weights, + policy_output.action_weights, + atol=1e-6) + # No action was visited. + summary = policy_output.search_tree.summary() + np.testing.assert_array_equal(jnp.zeros([2, 4]), summary.visit_counts) + def test_stochastic_muzero_policy(self): """Tests that SMZ is equivalent to MZ with a dummy chance function.""" root = mctx.RootFnOutput(