Support num_simulations=0 in all policies.

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,15 +363,26 @@ 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)
   return base.PolicyOutput(
       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:

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
 

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(