DOCSEARCHbyAlgolia
Submit Docs

mctx

repository·main·Indexed 25 days ago

https://github.com/google-deepmind/mctx

A JAX-native implementation of Monte Carlo tree search (MCTS) algorithms, including AlphaZero, MuZero, and Gumbel MuZero. Designed for high performance via JIT-compilation and parallel batch processing, mctx provides high-level policies such as gumbel_muzero_policy and a core search function for MuZero-style action selection and dynamics modeling.

Tokens
1.1K
Snippets
2
Records
9
Agent score
33%

What's inside mctx

  1. Use Gumbel MuZero policy

    main

    Mctx provides high-level concrete policies like muzero_policy and gumbel_muzero_policy. It is recommended to use gumbel_muzero_policy as it guarantees policy improvement if action values are correctly evaluated.

    To use the policy, you must provide:

    • params: Model parameters.
    • rng_key: JAX random key.
    • root: A RootFnOutput containing prior_logits, value, and an embedding for the root state.
    • recurrent_fn: A dynamics model function with the signature recurrent_fn(params, rng_key, action, embedding). It must return a tuple (RecurrentFnOutput, new_embedding), where RecurrentFnOutput contains reward, discount, prior_logits, and value.
    • num_simulations: The number of MCTS simulations to perform.

    The returned policy_output contains:

    • action: The action proposed by the search.
    • action_weights: Targets usable to train the policy probabilities.
    policy_output = mctx.gumbel_muzero_policy(params, rng_key, root, recurrent_fn,
                                          num_simulations=32)

  2. Install Mctx

    main

    You can install the latest released version of Mctx from PyPI or install the latest development version directly from GitHub.

    To install the latest released version:

    pip install mctx

    To install the latest development version:

    pip install git+https://github.com/google-deepmind/mctx.git

  3. Define the RootFnOutput for MuZero

    main

    To specify the representation of the root state in MuZero, you must provide a RootFnOutput object. This object must contain:

    • prior_logits: Logits from a policy network.
    • value: The estimated value of the root state.
    • embedding: An embedding suitable to represent the root state for the environment model.

  4. Define the recurrent function for MuZero

    main

    The dynamics environment model in Mctx is specified by a recurrent_fn. This function must follow this interface:

    recurrent_fn(params, rng_key, action, embedding) -> (RecurrentFnOutput, new_embedding)

    • Inputs:
      • params: Model parameters.
      • rng_key: JAX random key.
      • action: The action taken.
      • embedding: The current state embedding.
    • Outputs:
      • RecurrentFnOutput: A structure containing reward, discount, prior_logits, and value for the transition.
      • new_embedding: The embedding of the next state.

  6. Use MuZero policies

    main

    Mctx provides several policy implementations for MuZero-based search:

    • muzero_policy: Standard MuZero policy.
    • stochastic_muzero_policy: Stochastic version of the MuZero policy.
    • gumbel_muzero_policy: MuZero policy utilizing Gumbel-based selection.

  7. Apply Q-transformations

    main

    Mctx provides utilities to transform Q-values during the search process:

    • qtransform_by_min_max: Transforms Q-values using min-max scaling.
    • qtransform_by_parent_and_siblings: Transforms Q-values based on the parent and sibling nodes.
    • qtransform_completed_by_mix_value: Transforms Q-values using a mix value.

  8. Use MuZero action selection functions

    main

    To control how actions are selected during different stages of the search, use these functions:

    • muzero_action_selection: Standard MuZero action selection.
    • gumbel_muzero_root_action_selection: Gumbel-based action selection specifically for the root node.
    • gumbel_muzero_interior_action_selection: Gumbel-based action selection for interior nodes.

  9. Reference Mctx core types and interfaces

    main

    The following types and interfaces are part of the public API for defining recurrent functions and search structures:

    • Tree: Represents the search tree structure.
    • RecurrentFn: Interface for recurrent functions.
    • RecurrentState: Represents the state in a recurrent function.
    • RecurrentFnOutput: Output type for recurrent functions.
    • PolicyOutput: Output type for policies.
    • RootFnOutput: Output type for root functions.
    • DecisionRecurrentFnOutput: Output type for decision-based recurrent functions.
    • ChanceRecurrentFnOutput: Output type for chance-based recurrent functions.
    • InteriorActionSelectionFn: Interface for interior node action selection.
    • RootActionSelectionFn: Interface for root node action selection.
    • LoopFn: Interface for loop functions.
    • GumbelMuZeroExtraData: Data structure for Gumbel MuZero specific information.