code: def mcts(
      observation: types.Observation,
      model: models.Model,
      search_policy: SearchPolicy,
      evaluation: types.EvaluationFn,
      num_simulations: int,
      num_actions: int,
      discount: float = 1.0,
      dirichlet_alpha: float = 1,
      exploration_fraction: float = 0.0,
  ) -> Node:
      """Does Monte Carlo tree search (MCTS), AlphaZero style."""
  
      # Evaluate the prior policy for this state.
      prior, value = evaluation(observation)
      assert prior.shape == (num_actions,)
  
      # Add exploration noise to the prior.
      noise = np.random.dirichlet(alpha=[dirichlet_alpha] * num_actions)
      prior = (
          prior * (1 - exploration_fraction)
          + noise * exploration_fraction
      )
  
      # Create a fresh tree search.
      root = Node()
      root.expand(prior)
  
      # Save the model state so that we can reset it for each simulation.
      model.save_checkpoint()
      for _ in range(num_simulations):
          # Start a new simulation from the top.
          trajectory = [root]
          node = root
  
          # Generate a trajectory.
          timestep = None
          while node.children:
              # Select an action according to the search policy.
              action = search_policy(node)
  
              # Point the node at the corresponding child.
              node = node.children[action]
  
              # Step the simulator and add this timestep to the node.
              timestep = model.step(action)
              node.reward = timestep.reward or 0.0
              node.terminal = timestep.last()
              trajectory.append(node)
  
          if timestep is None:
              raise ValueError(
                  "Generated an empty rollout; this should not happen."
              )
  
          # Calculate the bootstrap for leaf nodes.
          if node.terminal:
              # If terminal, there is no bootstrap value.
              value = 0.0
          else:
              # Otherwise, bootstrap from this node with our value function.
              prior, value = evaluation(timestep.observation)
  
              # We also want to expand this node for next time.
              node.expand(prior)
  
          # Load the saved model state.
          model.load_checkpoint()
  
          # Monte Carlo back-up with bootstrap from value function.
          ret = value
          while trajectory:
              # Pop off the latest node in the trajectory.
              node = trajectory.pop()
  
              # Accumulate the discounted return
              ret *= discount
              ret += node.reward
  
              # Update the node.
              node.total_value += ret
              node.visit_count += 1
  
      return root
  
file_path: acme/acme/agents/tf/mcts/search.py
imports: import dataclasses
  from typing import Callable, Dict
  from acme.agents.tf.mcts import models
  from acme.agents.tf.mcts import types
  import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: mcts(observation, model, search_policy, evaluation, num_simulations, num_actions, discount, dirichlet_alpha, exploration_fraction) and returns: root