<commit_before>#include "mcts.hpp"
#include "notation.hpp"

#include <queue>
#include <bitset>

using namespace mcts;

void Node::initialize(State const& state) {
  hash = NodeTable::key(state.hash);
}

void Node::update(double result) {
  m_sample_size++;
  double mean0 = m_mean;
  double delta = result - mean0;
  m_mean += delta/m_sample_size;
  m_derivative = 0.9*m_derivative + 0.1*(m_mean - mean0);
}

void Node::adjoin_parent(Node* parent) {
  std::lock_guard<std::mutex> lock(parents_mutex);
  parents.insert(parent);
}

double Node::rollout(State& state, boost::mt19937& generator) {
  Color initial_player = state.us;
#ifdef MC_EXPENSIVE_RUNTIME_TESTS
  State initial_state(state);
  std::vector<Move> move_history; // for debugging
#endif
  while (true) {
    if (state.drawn_by_50())
      return draw_value;
    boost::optional<Move> move = moves::random_move(state, generator);
    //boost::optional<Move> move = moves::make_random_legal_move(state, generator);
    if (!move)
      break;
    state.make_move(*move);
#ifdef MC_EXPENSIVE_RUNTIME_TESTS
    move_history.push_back(*move);
    state.require_consistent();
#endif
  }
  boost::optional<Color> winner = state.winner();
  if (!winner)
    return draw_value;
  return *winner == initial_player ? win_value : loss_value;
}

std::string Node::format_statistics() {
  return str(boost::format("%1% %2% (d %3%) %4% (%5% parents)")
             % sample_size()
             % mean()
             % derivative()
             % selection_criterion()
             % parents.size());
}

Node* NodeTable::get(Hash hash) {
  size_t key = NodeTable::key(hash);
  Node* node = &nodes.at(key);
  if (node->hash != key)
    return nullptr;
  return node;
}

Node* NodeTable::get_or_create(State const& state) {
  size_t key = NodeTable::key(state);
  Node* node = &nodes.at(key);
  if (node->hash != key)
    node->initialize(state);
  return node;
}

void Graph::sample(State state, boost::mt19937& generator) {
  Node* node = nodes.get_or_create(state);
  std::unordered_set<Hash> path;

  // selection
  while (node->sample_size() > 0) {
    path.insert(node->hash);

    Node* child = select_child(node, state, generator);
    if (!child) {
      // no legal successors; loss
      backprop(node, loss_value);
      return;
    }
    if (path.count(child->hash) > 0) {
      // repeated state selected; draw by repetition
      backprop(node, draw_value);
      return;
    }
    node = child;
  }

  double result = node->rollout(state, generator);
  backprop(node, result);
}

// returns nullptr if no legal successor states
// NOTE: state will be modified to be the corresponding successor
Node* Graph::select_child(Node* node, State& state, boost::mt19937& generator) {
  std::vector<Move> moves;
  moves::moves(moves, state);

  // find a legal move with maximum selection criterion.
  boost::optional<Move> best_move;
  Node* best_child;
  double best_score;
  for (Move curr_move: moves) {
    Undo undo = state.make_move(curr_move);
    Node* curr_child = nodes.get_or_create(state);

    curr_child->adjoin_parent(node);
    if (curr_child->sample_size() < 30)
      // select new nodes unconditionally
      return curr_child;

    // if legal
    if (!state.their_king_attacked()) {
      double curr_score = curr_child->selection_criterion(generator);
      if (!best_move || curr_score > best_score) {
        best_move  = curr_move;
        best_child = curr_child;
        best_score = curr_score;
      }
    }
    
    state.unmake_move(undo);
  }

  if (!best_move)
    return nullptr;
  
  state.make_move(*best_move);
  return best_child;
}

// update the node and all of its ancestors.  the graph is likely to be
// cyclic due to reversible moves.  we traverse it in a breadth-first
// manner so that the updates run along the shortest paths to each
// unique ancestor.
// NOTE: initial_result is from the perspective of the player who is to move in
// the state corresponding to node.  since the stored node values are from the
// perspective of the player who causes the node to be chosen, we have to invert
// it once.
void Graph::backprop(Node* node, double initial_result) {
  initial_result = invert_result(initial_result);

  std::unordered_set<Hash> encountered_nodes;
  auto encountered = [&](Node* node) {
    if (encountered_nodes.count(node->hash) > 0) {
      return true;
    } else {
      encountered_nodes.insert(node->hash);<commit_msg>combine selection and rollout -- never leave the graph<commit_after>#include "mcts.hpp"
#include "notation.hpp"

#include <queue>

#include <boost/range/algorithm/random_shuffle.hpp>

using namespace mcts;

void Node::initialize(State const& state) {
  hash = NodeTable::key(state.hash);
}

void Node::update(double result) {
  m_sample_size++;
  double mean0 = m_mean;
  double delta = result - mean0;
  m_mean += delta/m_sample_size;
  m_derivative = 0.9*m_derivative + 0.1*(m_mean - mean0);
}

void Node::adjoin_parent(Node* parent) {
  std::lock_guard<std::mutex> lock(parents_mutex);
  parents.insert(parent);
}

std::string Node::format_statistics() {
  return str(boost::format("%1% %2% (d %3%) %4% (%5% parents)")
             % sample_size()
             % mean()
             % derivative()
             % selection_criterion()
             % parents.size());
}

Node* NodeTable::get(Hash hash) {
  size_t key = NodeTable::key(hash);
  Node* node = &nodes.at(key);
  if (node->hash != key)
    return nullptr;
  return node;
}

Node* NodeTable::get_or_create(State const& state) {
  size_t key = NodeTable::key(state);
  Node* node = &nodes.at(key);
  if (node->hash != key)
    node->initialize(state);
  return node;
}

void Graph::sample(State state, boost::mt19937& generator) {
  Node* node = nodes.get_or_create(state);
  sorted_vector<Hash> path;

  double result;

  // selection
  while (true) {
    path.insert(node->hash);

    Node* child = select_child(node, state, generator);

    if (!child) {
      // no legal successors; game over
      boost::optional<Color> winner = state.winner();
      result = !winner
        ? draw_value
        : (*winner == state.us
           ? win_value
           : loss_value);
      break;
    }

    if (state.drawn_by_50() || path.contains(child->hash)) {
      // draw per 50-move rule or by choosing repetition
      result = draw_value;
      break;
    }

    node = child;
  }

  backprop(node, result);
}

// returns nullptr if no legal successor states
// NOTE: state will be modified to be the corresponding successor
Node* Graph::select_child(Node* node, State& state, boost::mt19937& generator) {
  if (node->sample_size() < 30) {
    // don't waste time with unreliable statistics
    boost::optional<Move> move = moves::make_random_legal_move(state, generator);
    if (!move)
      return nullptr;
    Node* child = nodes.get_or_create(state);
    child->adjoin_parent(node);
    return child;
  } else {
    // use statistics to make selection
    std::vector<Move> moves;
    moves::moves(moves, state);

    // find a legal move with maximum selection criterion.
    boost::optional<Move> best_move;
    Node* best_child;
    double best_score;
    
    for (Move curr_move: moves) {
      Undo undo = state.make_move(curr_move);
      
      // if legal
      if (!state.their_king_attacked()) {
        Node* curr_child = nodes.get_or_create(state);
        
        curr_child->adjoin_parent(node);
        if (curr_child->sample_size() < 10)
          // select new nodes unconditionally
          return curr_child;
        
        double curr_score = curr_child->selection_criterion(generator);
        if (!best_move || curr_score > best_score) {
          best_move  = curr_move;
          best_child = curr_child;
          best_score = curr_score;
        }
      }
      
      state.unmake_move(undo);
    }

    if (!best_move)
      return nullptr;
    state.make_move(*best_move);
    return best_child;
  }
}
  
// update the node and all of its ancestors.  the graph is likely to be
// cyclic due to reversible moves.  we traverse it in a breadth-first
// manner so that the updates run along the shortest paths to each
// unique ancestor.
// NOTE: initial_result is from the perspective of the player who is to move in
// the state corresponding to node.  since the stored node values are from the
// perspective of the player who causes the node to be chosen, we have to invert
// it once.
void Graph::backprop(Node* node, double initial_result) {
  initial_result = invert_result(initial_result);

  std::unordered_set<Hash> encountered_nodes;
  auto encountered = [&](Node* node) {
    if (encountered_nodes.count(node->hash) > 0) {
      return true;
    } else {
      encountered_nodes.insert(node->hash);
{'fraction_non_alphanumeric': 0.1226706299302618, 'fraction_numerical': 0.007431119240882588, 'mean_word_length': 3.2240463544181557, 'pattern_counts': {'":': 0, '<': 26, '<?xml version=': 0, '>': 53, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3556453', 'n_tokens_mistral': 2852, 'n_tokens_neox': 2782, 'n_words': 747}