def run_episode(
    episode: core.Episode, mzconfig: core.MuZeroConfig,
    initial_inference_service: prediction_service_pb2_grpc.PredictionServiceStub,
    recurrent_inference_service: prediction_service_pb2_grpc.PredictionServiceStub,
    initial_inference_cache: List[core.NetworkOutput],
    counter: metrics.Metrics.DelegatingCounter) -> None:
  """Runs and mutates `episode`.  Returns MCTS visualizations for each step."""
  step_num = 0
  while not episode.terminal() and len(episode.history) < mzconfig.max_moves:
    # Get observation.
    current_observation = episode.make_image(state_index=-1)

    # Prepare MCTS.
    # We may already have played this, so avoid swamping the predict servers.
    if len(initial_inference_cache) >= step_num + 1:
      initial_inference_output = initial_inference_cache[step_num]
      counter.inc()
    else:
      initial_inference_output = send_initial_inference_request(
          predict_service=initial_inference_service, inputs=current_observation)
      initial_inference_cache.append(initial_inference_output)

    legal_actions = episode.legal_actions()
    root = core.prepare_root_node(  # pytype: disable=wrong-arg-types  # numpy-scalars
        config=mzconfig,
        legal_actions=legal_actions,
        initial_inference_output=initial_inference_output)

    # Run MCTS.
    core.run_mcts(
        config=mzconfig,
        root=root,
        action_history=episode.action_history(),
        legal_actions_fn=episode.legal_actions,
        recurrent_inference_fn=functools.partial(
            send_recurrent_inference_request,
            predict_service=recurrent_inference_service))

    # Pick action.
    action = core.select_action(
        config=mzconfig,
        num_moves=len(episode.history),
        node=root,
        train_step=0,
        use_softmax=mzconfig.use_softmax_for_action_selection)

    episode.apply(action)
    step_num += 1
Runs and mutates `episode`.  Returns MCTS visualizations for each step.
language/search_agents/muzero/run_inference_beam.py