def test_stochastic_muzero_policy(self):
    """Tests that SMZ is equivalent to MZ with a dummy chance function."""
    root = mctx.RootFnOutput(
        prior_logits=jnp.array([
            [-1.0, 0.0, 2.0, 3.0],
            [0.0, 2.0, 5.0, -4.0],
        ]),
        value=jnp.array([1.0, 0.0]),
        embedding=jnp.zeros([2, 4])
    )
    rewards = jnp.zeros_like(root.prior_logits)
    invalid_actions = jnp.array([
        [0.0, 0.0, 0.0, 1.0],
        [1.0, 0.0, 1.0, 0.0],
    ])

    num_simulations = 10

    policy_output = mctx.muzero_policy(
        params=(),
        rng_key=jax.random.PRNGKey(0),
        root=root,
        recurrent_fn=_make_bandit_recurrent_fn(
            rewards,
            dummy_embedding=jnp.zeros_like(root.embedding)),
        num_simulations=num_simulations,
        invalid_actions=invalid_actions,
        dirichlet_fraction=0.0)

    num_chance_outcomes = 5

    decision_rec_fn, chance_rec_fn = _make_bandit_decision_and_chance_fns(
        rewards, num_chance_outcomes)

    stochastic_policy_output = mctx.stochastic_muzero_policy(
        params=(),
        rng_key=jax.random.PRNGKey(0),
        root=root,
        decision_recurrent_fn=decision_rec_fn,
        chance_recurrent_fn=chance_rec_fn,
        num_simulations=2 * num_simulations,
        invalid_actions=invalid_actions,
        dirichlet_fraction=0.0)

    np.testing.assert_array_equal(stochastic_policy_output.action,
                                  policy_output.action)

    np.testing.assert_allclose(stochastic_policy_output.action_weights,
                               policy_output.action_weights)
Tests that SMZ is equivalent to MZ with a dummy chance function.
mctx/_src/tests/policies_test.py