621 lines
23 KiB
Cython
621 lines
23 KiB
Cython
# cython: language_level=3, boundscheck=False, wraparound=False, cdivision=True, initializedcheck=False
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
import random
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from coolrl_lost_cities.games.classic.bots.heuristic_cy cimport HeuristicBot
|
|
from coolrl_lost_cities.games.classic.bots.heuristic import HeuristicBot as PyHeuristicBot
|
|
from coolrl_lost_cities.games.classic.deep_cfr.encoding import encode_info_state
|
|
from coolrl_lost_cities.games.classic.game cimport GameState
|
|
|
|
from .info_set import canonical_info_set_key
|
|
|
|
|
|
DEF MAX_ACTIONS = 64
|
|
DEF DEFAULT_ACTION_SIZE = 64
|
|
|
|
|
|
cdef class _ArrayMap:
|
|
cdef MctsNode node
|
|
cdef int kind
|
|
cdef int action_size
|
|
cdef bint is_int
|
|
|
|
def __init__(self, MctsNode node, int kind, bint is_int=False):
|
|
self.node = node
|
|
self.kind = kind
|
|
self.action_size = node.action_size
|
|
self.is_int = is_int
|
|
|
|
cdef inline void _check(self, int action) except *:
|
|
if action < 0 or action >= self.action_size:
|
|
raise KeyError(action)
|
|
|
|
cdef inline bint has(self, int action) noexcept:
|
|
return 0 <= action < self.action_size and self.node.active_present[action] != 0
|
|
|
|
cdef inline long get_int(self, int action, long default_value=0) noexcept:
|
|
if 0 <= action < self.action_size and self.node.active_present[action] != 0:
|
|
if self.kind == 1:
|
|
return self.node.visits_arr[action]
|
|
if self.kind == 3:
|
|
return self.node.virtual_visits_arr[action]
|
|
return default_value
|
|
|
|
cdef inline double get_float(self, int action, double default_value=0.0) noexcept:
|
|
if 0 <= action < self.action_size and self.node.active_present[action] != 0:
|
|
if self.kind == 0:
|
|
return self.node.priors_arr[action]
|
|
if self.kind == 2:
|
|
return self.node.value_sum_arr[action]
|
|
return default_value
|
|
|
|
cdef inline void _mark_active(self, int action) noexcept:
|
|
if self.node.active_present[action] == 0:
|
|
self.node.active_present[action] = 1
|
|
self.node.active_actions[self.node.n_active] = action
|
|
self.node.n_active += 1
|
|
|
|
cdef inline void set_int(self, int action, long value) except *:
|
|
self._check(action)
|
|
self._mark_active(action)
|
|
if self.kind == 1:
|
|
self.node.visits_arr[action] = <int>value
|
|
elif self.kind == 3:
|
|
self.node.virtual_visits_arr[action] = <int>value
|
|
else:
|
|
raise TypeError("integer write to float node map")
|
|
|
|
cdef inline void set_float(self, int action, double value) except *:
|
|
self._check(action)
|
|
self._mark_active(action)
|
|
if self.kind == 0:
|
|
self.node.priors_arr[action] = value
|
|
elif self.kind == 2:
|
|
self.node.value_sum_arr[action] = value
|
|
else:
|
|
raise TypeError("float write to integer node map")
|
|
|
|
def get(self, action, default=None):
|
|
cdef int a = int(action)
|
|
if self.has(a):
|
|
if self.is_int:
|
|
return int(self.get_int(a, 0))
|
|
return float(self.get_float(a, 0.0))
|
|
return default
|
|
|
|
def setdefault(self, action, default=None):
|
|
cdef int a = int(action)
|
|
if self.has(a):
|
|
if self.is_int:
|
|
return int(self.get_int(a, 0))
|
|
return float(self.get_float(a, 0.0))
|
|
if default is None:
|
|
default = 0 if self.is_int else 0.0
|
|
if self.is_int:
|
|
self.set_int(a, int(default))
|
|
return int(default)
|
|
self.set_float(a, float(default))
|
|
return float(default)
|
|
|
|
def __getitem__(self, action):
|
|
cdef int a = int(action)
|
|
self._check(a)
|
|
if self.node.active_present[a] == 0:
|
|
raise KeyError(action)
|
|
if self.is_int:
|
|
return int(self.get_int(a, 0))
|
|
return float(self.get_float(a, 0.0))
|
|
|
|
def __setitem__(self, action, value):
|
|
cdef int a = int(action)
|
|
if self.is_int:
|
|
self.set_int(a, int(value))
|
|
else:
|
|
self.set_float(a, float(value))
|
|
|
|
def __contains__(self, action):
|
|
return self.has(int(action))
|
|
|
|
def __bool__(self):
|
|
return self.node.n_active > 0
|
|
|
|
def __len__(self):
|
|
return self.node.n_active
|
|
|
|
def items(self):
|
|
cdef int i
|
|
result = []
|
|
cdef int action
|
|
for i in range(self.node.n_active):
|
|
action = self.node.active_actions[i]
|
|
if self.is_int:
|
|
result.append((action, int(self.get_int(action, 0))))
|
|
else:
|
|
result.append((action, float(self.get_float(action, 0.0))))
|
|
return result
|
|
|
|
def keys(self):
|
|
cdef int i
|
|
return [self.node.active_actions[i] for i in range(self.node.n_active)]
|
|
|
|
def values(self):
|
|
cdef int i
|
|
result = []
|
|
cdef int action
|
|
for i in range(self.node.n_active):
|
|
action = self.node.active_actions[i]
|
|
if self.is_int:
|
|
result.append(int(self.get_int(action, 0)))
|
|
else:
|
|
result.append(float(self.get_float(action, 0.0)))
|
|
return result
|
|
|
|
def __repr__(self):
|
|
return repr(dict(self.items()))
|
|
|
|
|
|
cdef class MctsNode:
|
|
cdef public bytes info_set_key
|
|
cdef public int player
|
|
cdef public object priors
|
|
cdef public object visits
|
|
cdef public object value_sum
|
|
cdef public object virtual_visits
|
|
cdef public dict children
|
|
cdef public bint terminal
|
|
cdef public bint expanded
|
|
cdef int action_size
|
|
cdef int visits_arr[MAX_ACTIONS]
|
|
cdef double value_sum_arr[MAX_ACTIONS]
|
|
cdef double priors_arr[MAX_ACTIONS]
|
|
cdef int virtual_visits_arr[MAX_ACTIONS]
|
|
cdef int active_actions[MAX_ACTIONS]
|
|
cdef unsigned char active_present[MAX_ACTIONS]
|
|
cdef int n_active
|
|
|
|
def __init__(
|
|
self,
|
|
bytes info_set_key,
|
|
int player,
|
|
object priors=None,
|
|
object visits=None,
|
|
object value_sum=None,
|
|
object virtual_visits=None,
|
|
object children=None,
|
|
bint terminal=False,
|
|
int action_size=DEFAULT_ACTION_SIZE,
|
|
):
|
|
if action_size > MAX_ACTIONS:
|
|
raise ValueError("action_size exceeds fixed MCTS action buffer")
|
|
self.info_set_key = info_set_key
|
|
self.player = player
|
|
self.terminal = terminal
|
|
self.expanded = terminal
|
|
self.action_size = action_size
|
|
self.n_active = 0
|
|
self.priors = _ArrayMap(self, 0, False)
|
|
self.visits = _ArrayMap(self, 1, True)
|
|
self.value_sum = _ArrayMap(self, 2, False)
|
|
self.virtual_visits = _ArrayMap(self, 3, True)
|
|
self.children = {} if children is None else dict(children)
|
|
if priors is not None:
|
|
for action, value in dict(priors).items():
|
|
self.priors[action] = value
|
|
if visits is not None:
|
|
for action, value in dict(visits).items():
|
|
self.visits[action] = value
|
|
if value_sum is not None:
|
|
for action, value in dict(value_sum).items():
|
|
self.value_sum[action] = value
|
|
if virtual_visits is not None:
|
|
for action, value in dict(virtual_visits).items():
|
|
self.virtual_visits[action] = value
|
|
|
|
cpdef bint is_expanded(self):
|
|
return self.terminal or self.expanded or self.n_active > 0
|
|
|
|
cpdef double q(self, int action):
|
|
cdef long n = (<_ArrayMap>self.visits).get_int(action, 0)
|
|
if n <= 0:
|
|
return 0.0
|
|
return (<_ArrayMap>self.value_sum).get_float(action, 0.0) / n
|
|
|
|
|
|
cdef class SearchPathEntry:
|
|
cdef public MctsNode node
|
|
cdef public int action
|
|
cdef public int parent_player
|
|
cdef public int child_player
|
|
|
|
def __init__(self, MctsNode node, int action, int parent_player, int child_player):
|
|
self.node = node
|
|
self.action = action
|
|
self.parent_player = parent_player
|
|
self.child_player = child_player
|
|
|
|
|
|
cdef class PendingSimulation:
|
|
cdef public list path
|
|
cdef public GameState leaf_state
|
|
cdef public object leaf_node
|
|
cdef public int leaf_player
|
|
cdef public object info_state
|
|
cdef public object legal_mask
|
|
cdef public list legal_actions
|
|
cdef public object terminal_value
|
|
|
|
def __init__(
|
|
self,
|
|
list path,
|
|
GameState leaf_state,
|
|
object leaf_node,
|
|
int leaf_player,
|
|
object info_state,
|
|
object legal_mask,
|
|
list legal_actions,
|
|
object terminal_value=None,
|
|
):
|
|
self.path = path
|
|
self.leaf_state = leaf_state
|
|
self.leaf_node = leaf_node
|
|
self.leaf_player = leaf_player
|
|
self.info_state = info_state
|
|
self.legal_mask = legal_mask
|
|
self.legal_actions = legal_actions
|
|
self.terminal_value = terminal_value
|
|
|
|
|
|
cdef class MctsTree:
|
|
cdef public dict nodes
|
|
cdef int action_size
|
|
|
|
def __init__(self, int action_size=DEFAULT_ACTION_SIZE):
|
|
self.nodes = {}
|
|
self.action_size = action_size
|
|
|
|
def get_or_create(self, bytes key, *, int player, bint terminal=False):
|
|
cdef MctsNode node = self.nodes.get(key)
|
|
if node is None:
|
|
node = MctsNode(key, player=player, terminal=terminal, action_size=self.action_size)
|
|
self.nodes[key] = node
|
|
return node
|
|
|
|
|
|
cdef class IsMctsSearcher:
|
|
cdef public object network
|
|
cdef public object config
|
|
cdef public object device
|
|
cdef public object encoding
|
|
cdef public object rng
|
|
cdef public MctsTree tree
|
|
cdef HeuristicBot _rollout_bot
|
|
cdef int action_size
|
|
|
|
def __init__(
|
|
self,
|
|
object network,
|
|
object config,
|
|
*,
|
|
object device="cpu",
|
|
object encoding=None,
|
|
object rng=None,
|
|
):
|
|
self.network = network
|
|
self.config = config
|
|
self.device = torch.device(device)
|
|
self.encoding = encoding
|
|
self.rng = rng or random.Random()
|
|
self.action_size = int(getattr(network, "action_size", DEFAULT_ACTION_SIZE))
|
|
if self.action_size > MAX_ACTIONS:
|
|
raise ValueError("action_size exceeds fixed MCTS action buffer")
|
|
self.tree = MctsTree(self.action_size)
|
|
self._rollout_bot = (
|
|
<HeuristicBot>PyHeuristicBot() if config.rollout_policy == "heuristic_balanced" else None
|
|
)
|
|
|
|
cdef inline int _from_unified_action_c(self, GameState state, int action_id) noexcept:
|
|
cdef int card_action_size = 2 * state.hand_size
|
|
if state.phase_id == 0:
|
|
return action_id
|
|
return action_id - card_action_size
|
|
|
|
cdef list _unified_legal_actions_list_c(self, GameState state):
|
|
cdef int actions[MAX_ACTIONS]
|
|
cdef int count = state._unified_legal_actions_c(actions)
|
|
cdef int i
|
|
return [actions[i] for i in range(count)]
|
|
|
|
cpdef dict search(self, GameState state, int traverser, object n_sims=None):
|
|
cdef bytes root_key = canonical_info_set_key(state, state.current_player)
|
|
cdef MctsNode root = self.tree.get_or_create(
|
|
root_key, player=state.current_player, terminal=state.terminal
|
|
)
|
|
cdef int sims = int(n_sims or self.config.n_simulations)
|
|
cdef int completed = 0
|
|
cdef list pending
|
|
cdef list legal
|
|
cdef int action
|
|
cdef dict result
|
|
while completed < sims:
|
|
pending = self.prepare_simulation_batch(state, traverser, 1)
|
|
if not pending:
|
|
break
|
|
self.evaluate_and_backup(pending)
|
|
completed += len(pending)
|
|
legal = state.unified_legal_actions()
|
|
result = {}
|
|
for action in legal:
|
|
result[action] = (<_ArrayMap>root.visits).get_int(action, 0)
|
|
return result
|
|
|
|
cpdef list prepare_simulation_batch(
|
|
self,
|
|
GameState root_state,
|
|
int traverser,
|
|
int max_simulations,
|
|
):
|
|
cdef list pending = []
|
|
cdef PendingSimulation item
|
|
cdef int i
|
|
for i in range(max_simulations):
|
|
item = self.prepare_simulation(root_state, traverser)
|
|
pending.append(item)
|
|
if item.terminal_value is None and item.leaf_node is not None and not item.path:
|
|
break
|
|
return pending
|
|
|
|
cpdef PendingSimulation prepare_simulation(self, GameState root_state, int traverser):
|
|
cdef GameState state = root_state.determinize_for_player(traverser, self.rng)
|
|
cdef list path = []
|
|
cdef int depth = 0
|
|
cdef object cached_key = None
|
|
cdef int player
|
|
cdef bytes key
|
|
cdef MctsNode node
|
|
cdef list legal_actions
|
|
cdef int action
|
|
cdef int local_action
|
|
cdef int child_player
|
|
cdef bytes child_key
|
|
cdef int actions[MAX_ACTIONS]
|
|
cdef int action_count
|
|
cdef int i
|
|
while True:
|
|
player = state.current_player
|
|
if state.terminal or depth >= int(self.config.max_depth):
|
|
return PendingSimulation(
|
|
path=path,
|
|
leaf_state=state,
|
|
leaf_node=None,
|
|
leaf_player=player,
|
|
info_state=None,
|
|
legal_mask=None,
|
|
legal_actions=[],
|
|
terminal_value=float(state.total_scores[player] - state.total_scores[1 - player]),
|
|
)
|
|
if cached_key is None:
|
|
key = canonical_info_set_key(state, player)
|
|
else:
|
|
key = cached_key
|
|
node = self.tree.get_or_create(key, player=player, terminal=state.terminal)
|
|
if not node.is_expanded():
|
|
action_count = state._unified_legal_actions_c(actions)
|
|
legal_actions = [actions[i] for i in range(action_count)]
|
|
if not legal_actions:
|
|
node.terminal = True
|
|
return PendingSimulation(
|
|
path=path,
|
|
leaf_state=state,
|
|
leaf_node=node,
|
|
leaf_player=player,
|
|
info_state=None,
|
|
legal_mask=None,
|
|
legal_actions=[],
|
|
terminal_value=float(state.total_scores[player] - state.total_scores[1 - player]),
|
|
)
|
|
return PendingSimulation(
|
|
path=path,
|
|
leaf_state=state,
|
|
leaf_node=node,
|
|
leaf_player=player,
|
|
info_state=encode_info_state(state, player, self.encoding),
|
|
legal_mask=np.asarray(state.unified_legal_mask(), dtype=bool),
|
|
legal_actions=legal_actions,
|
|
)
|
|
|
|
action_count = state._unified_legal_actions_c(actions)
|
|
legal_actions = [actions[i] for i in range(action_count)]
|
|
action = self._select_action(node, legal_actions)
|
|
(<_ArrayMap>node.virtual_visits).set_int(
|
|
action, (<_ArrayMap>node.virtual_visits).get_int(action, 0) + 1
|
|
)
|
|
local_action = self._from_unified_action_c(state, action)
|
|
state._push_action_c(local_action)
|
|
child_player = state.current_player
|
|
child_key = canonical_info_set_key(state, child_player)
|
|
node.children[action] = child_key
|
|
path.append(
|
|
SearchPathEntry(
|
|
node=node,
|
|
action=action,
|
|
parent_player=player,
|
|
child_player=child_player,
|
|
)
|
|
)
|
|
cached_key = child_key
|
|
depth += 1
|
|
|
|
cpdef evaluate_and_backup(self, list pending):
|
|
cdef list network_pending = [item for item in pending if item.terminal_value is None]
|
|
cdef dict values_by_id = {}
|
|
cdef dict priors_by_id = {}
|
|
cdef object infos
|
|
cdef object masks
|
|
cdef object x
|
|
cdef object mask
|
|
cdef object probs
|
|
cdef object network_values
|
|
cdef object network_values_np
|
|
cdef int index
|
|
cdef PendingSimulation item
|
|
cdef double value
|
|
if network_pending:
|
|
infos = np.stack([item.info_state for item in network_pending if item.info_state is not None])
|
|
masks = np.stack([item.legal_mask for item in network_pending if item.legal_mask is not None])
|
|
with torch.inference_mode():
|
|
x = torch.as_tensor(infos, dtype=torch.float32, device=self.device)
|
|
mask = torch.as_tensor(masks, dtype=torch.bool, device=self.device)
|
|
probs = self.network.policy_distribution(x, mask).detach().cpu().numpy()
|
|
_logits, network_values = self.network(x, mask)
|
|
network_values_np = network_values.detach().cpu().numpy()
|
|
for index, item in enumerate(network_pending):
|
|
priors_by_id[id(item)] = probs[index]
|
|
values_by_id[id(item)] = float(network_values_np[index])
|
|
|
|
for item in pending:
|
|
if item.terminal_value is not None:
|
|
value = item.terminal_value
|
|
else:
|
|
value = self._expand_with_prior(
|
|
item.leaf_node,
|
|
item.leaf_state,
|
|
item.leaf_player,
|
|
item.legal_actions,
|
|
priors_by_id[id(item)],
|
|
values_by_id[id(item)],
|
|
)
|
|
self._backup(item.path, value, item.leaf_player)
|
|
|
|
cpdef double _expand_with_prior(
|
|
self,
|
|
MctsNode node,
|
|
GameState state,
|
|
int player,
|
|
list legal_actions,
|
|
object probs,
|
|
double network_value,
|
|
):
|
|
cdef int action
|
|
cdef object rollout_value
|
|
legal_actions = self._unified_legal_actions_list_c(state)
|
|
if not legal_actions:
|
|
node.terminal = True
|
|
return float(state.total_scores[player] - state.total_scores[1 - player])
|
|
node.expanded = True
|
|
for action in legal_actions:
|
|
(<_ArrayMap>node.priors).set_float(action, float(probs[action]))
|
|
if not (<_ArrayMap>node.visits).has(action):
|
|
(<_ArrayMap>node.visits).set_int(action, 0)
|
|
if not (<_ArrayMap>node.value_sum).has(action):
|
|
(<_ArrayMap>node.value_sum).set_float(action, 0.0)
|
|
if not (<_ArrayMap>node.virtual_visits).has(action):
|
|
(<_ArrayMap>node.virtual_visits).set_int(action, 0)
|
|
rollout_value = self._rollout_value(state, player) if self.config.use_rollout_value else None
|
|
if rollout_value is None:
|
|
return float(network_value)
|
|
return float(rollout_value)
|
|
|
|
cpdef int _select_action(self, MctsNode node, list legal_actions):
|
|
cdef int total_visits = 0
|
|
cdef int action
|
|
cdef long n
|
|
cdef long virtual
|
|
cdef long n_eff
|
|
cdef double sqrt_total
|
|
cdef double prior
|
|
cdef double q_eff
|
|
cdef double score
|
|
cdef double best_score = -float("inf")
|
|
cdef int best_action = int(legal_actions[0])
|
|
cdef _ArrayMap visits = <_ArrayMap>node.visits
|
|
cdef _ArrayMap virtual_visits = <_ArrayMap>node.virtual_visits
|
|
cdef _ArrayMap priors = <_ArrayMap>node.priors
|
|
cdef _ArrayMap value_sum = <_ArrayMap>node.value_sum
|
|
for action in legal_actions:
|
|
total_visits += visits.get_int(action, 0) + virtual_visits.get_int(action, 0)
|
|
sqrt_total = math.sqrt(max(1, total_visits))
|
|
for action in legal_actions:
|
|
n = visits.get_int(action, 0)
|
|
virtual = virtual_visits.get_int(action, 0)
|
|
n_eff = n + virtual
|
|
prior = priors.get_float(action, 0.0)
|
|
if n_eff <= 0:
|
|
q_eff = 0.0
|
|
else:
|
|
q_eff = (
|
|
value_sum.get_float(action, 0.0)
|
|
- virtual * float(self.config.virtual_loss_value)
|
|
) / n_eff
|
|
score = q_eff + float(self.config.c_puct) * prior * sqrt_total / (1 + n_eff)
|
|
if score > best_score:
|
|
best_score = score
|
|
best_action = action
|
|
return int(best_action)
|
|
|
|
cpdef _backup(self, list path, double leaf_value, int leaf_player):
|
|
cdef double value = float(leaf_value)
|
|
cdef int value_player = int(leaf_player)
|
|
cdef SearchPathEntry entry
|
|
cdef double parent_value
|
|
cdef long current_virtual
|
|
cdef _ArrayMap visits
|
|
cdef _ArrayMap virtual_visits
|
|
cdef _ArrayMap value_sum
|
|
for entry in reversed(path):
|
|
parent_value = value if value_player == entry.parent_player else -value
|
|
virtual_visits = <_ArrayMap>entry.node.virtual_visits
|
|
visits = <_ArrayMap>entry.node.visits
|
|
value_sum = <_ArrayMap>entry.node.value_sum
|
|
current_virtual = virtual_visits.get_int(entry.action, 0)
|
|
virtual_visits.set_int(entry.action, max(0, current_virtual - 1))
|
|
visits.set_int(entry.action, visits.get_int(entry.action, 0) + 1)
|
|
value_sum.set_float(
|
|
entry.action,
|
|
value_sum.get_float(entry.action, 0.0) + parent_value,
|
|
)
|
|
value = parent_value
|
|
value_player = entry.parent_player
|
|
|
|
cpdef _release_virtual_path(self, list path):
|
|
cdef SearchPathEntry entry
|
|
cdef _ArrayMap virtual_visits
|
|
cdef long current_virtual
|
|
for entry in path:
|
|
virtual_visits = <_ArrayMap>entry.node.virtual_visits
|
|
current_virtual = virtual_visits.get_int(entry.action, 0)
|
|
virtual_visits.set_int(entry.action, max(0, current_virtual - 1))
|
|
|
|
cpdef object _rollout_value(self, GameState state, int player):
|
|
cdef int steps = 0
|
|
cdef int actions[MAX_ACTIONS]
|
|
cdef int count
|
|
cdef int unified_action
|
|
cdef int action
|
|
cdef int max_depth = int(self.config.max_depth)
|
|
while not state.terminal and steps < max_depth:
|
|
if self._rollout_bot is not None:
|
|
action = self._rollout_bot.act_cython(state)
|
|
if not state._is_legal_action_c(action):
|
|
count = state._unified_legal_actions_c(actions)
|
|
if count <= 0:
|
|
break
|
|
unified_action = actions[self.rng.randrange(count)]
|
|
action = self._from_unified_action_c(state, unified_action)
|
|
else:
|
|
count = state._unified_legal_actions_c(actions)
|
|
if count <= 0:
|
|
break
|
|
unified_action = actions[self.rng.randrange(count)]
|
|
action = self._from_unified_action_c(state, unified_action)
|
|
state._push_action_c(action)
|
|
steps += 1
|
|
while steps > 0:
|
|
state._pop_action_c()
|
|
steps -= 1
|
|
return float(state.total_scores[player] - state.total_scores[1 - player])
|