def _reproduce_tree(self, tree):
    """Reproduces the given JSON tree by running a search."""
    policy_fn = dict(
        gumbel_muzero=mctx.gumbel_muzero_policy,
        muzero=mctx.muzero_policy,
    )[tree["algorithm"]]

    env_config = tree["env_config"]
    root = tree["tree"]
    num_actions = len(root["child_stats"])
    num_simulations = root["visit"] - 1
    qtransform = functools.partial(
        getattr(mctx, tree["algorithm_config"].pop("qtransform")),
        **tree["algorithm_config"].pop("qtransform_kwargs", {}))

    batch_size = 3
    # To test the independence of the batch computation, we use different
    # invalid actions for the other elements of the batch. The different batch
    # elements will then have different search tree depths.
    invalid_actions = np.zeros([batch_size, num_actions])
    invalid_actions[1, 1:] = 1
    invalid_actions[2, 2:] = 1

    def run_policy():
      return policy_fn(
          params=(),
          rng_key=jax.random.PRNGKey(1),
          root=_prepare_root(batch_size=batch_size, num_actions=num_actions),
          recurrent_fn=_prepare_recurrent_fn(num_actions, **env_config),
          num_simulations=num_simulations,
          qtransform=qtransform,
          invalid_actions=invalid_actions,
          **tree["algorithm_config"])

    policy_output = jax.jit(run_policy)()  # pylint: disable=not-callable
    logging.info("Done search.")

    return tree_to_pytree(policy_output.search_tree)
Reproduces the given JSON tree by running a search.
mctx/_src/tests/tree_test.py