IA de jeux, Implémenter une IA arborescente, en Haskell

Voir aussi : video youtube - video peertube - code source

Cet article, de la série “IA de jeux”, présente une implémentation de l’algorithme Monte-Carlo Tree Search, en Haskell, pour le jeu Puissance 4. Cette implémentation est ensuite comparée aux IA de l’article précédent (random et Monte-Carlo).

Articles de la séries IA de jeux :

Rappel sur l’algorithme MCTS

L’algorithme Monte-Carlo Tree Search permet d’évaluer l’intérêt de chaque coup, pour un jeu donné. Pour cela, l’algorithme construit itérativement un arbre de jeu, où les noeuds sont les états de jeu et les branches sont les coups menant d’un état à un autre (voir l’article Implémenter une IA arborescente, en C++).

Une itération de l’algorithme consiste à parcourir l’arbre courant jusqu’à une feuille (c’est-à-dire à jouer successivement les coups intéressants déjà estimés), par exemple en utilisant la formule UCB. L’étape suivante consiste à créer un nouveau noeud puis à jouer une partie aléatoire pour estimer le coup correspondant. La dernière étape consiste à remonter l’arbre pour mettre à jour les noeuds de la branche complète avec les données de la nouvelle simulation réalisée.

Une fois qu’on a réalisé toutes les itérations, le coup le plus intéressant à jouer correspond à l’enfant de la racine qui a été le plus simulé.

Implémenter MCTS en Haskell

Structure d’arbre

Comme pour un arbre n-aire classique, on définit un type noeud contenant les noeuds enfants ainsi que des données rattachées au noeud (jeu courant, joueur, nombre de simulations déjà réalisées, etc).

Comme on utilise déjà la monade ST pour implémenter le jeu, on peut en profiter pour l’utiliser également pour implémenter notre type Node. Ceci permet d’avoir des références modifiables, par exemple pour le noeud parent ou pour les données qui vont évoluées au cours de itérations MCTS.

data Node s = Node 
    { nodeGame :: Game s
    , nodePlayer :: Player     -- before move
    , nodeNmoves :: Int
    , nodeParent :: Maybe (NodeRef s)
    , nodeReward :: STRef s Double 
    , nodeNsims :: STRef s Int
    , nodeLastI :: STRef s Int
    , nodeChildren :: M.STVector s (Node s)
    }

type NodeRef s = STRef s (Node s)

On définit ensuite des constructeurs de Node. Typiquement, mkRoot crée un noeud sans parent, ni jouer de coup par rapport au jeu initial; alors que mkLeaf crée un noeud à partir d’un parent, et joue un coup sur le jeu fourni. La fonction mkNode permet de factoriser les opérations communes aux deux constructeurs.

mkRoot :: Game s -> ST s (Node s)
mkRoot = mkNode return Nothing

mkLeaf :: Int -> NodeRef s -> ST s (Node s)
mkLeaf k node = do
    game <- nodeGame <$> readSTRef node
    mkNode (playK k) (Just node) game

mkNode :: (Game s -> ST s (Game s)) -> Maybe (NodeRef s) -> Game s -> ST s (Node s)
mkNode gameFunc pNode game0 = do
    game1 <- cloneGame game0 >>= gameFunc
    let nMoves = nMovesGame game1
        player0 = _currentPlayer game0
    Node game1 player0 nMoves pNode 
        <$> newSTRef 0 <*> newSTRef 0 <*> newSTRef 0 <*> M.new nMoves

Enfin, la fonction bestNode permet de récupérer le meilleur coup parmi les noeuds enfants de la racine, à la fin des itérations MCTS. Il s’agit ici essentiellement d’un fold sur les noeuds enfants mais on notera quelques détails d’implémentation du fait qu’on travaille dans la monade ST.

bestNode :: Node s -> ST s Int
bestNode root = do
    lastI <- readSTRef (nodeLastI root)
    when (lastI < nodeNmoves root) (error "niters too low")
    let bestNodeFunc (nn,ii) i node = do
            nsims <- readSTRef (nodeNsims node)
            return $ if nsims > nn then (nsims, i) else (nn, ii)
    snd <$> M.ifoldM' bestNodeFunc (-1, -1) (nodeChildren root)

Sélection UCB

La formule UCB1 permet de sélectionner l’un des noeuds enfants, en fonction de leur estimation courante et du nombre de simulations déjà réalisées. Cette formule est une implémentation du fameux compromis exploitation-exploration.

La fonction ucb1 suivante calcule la formule pour un noeud (ici, on fixe la valeur de mctsKuct à 0.5, pour simplifier, mais c’est un choix très discutable).

ucb1 :: Double -> Int -> Int -> Double
ucb1 cReward cNsims pNsims =
    let cNsimsD = fromIntegral cNsims 
        exploitation = cReward / cNsimsD
        exploration = sqrt (log (fromIntegral $ 1 + pNsims) / cNsimsD)
    in exploitation + mctsKuct * exploration

mctsKuct :: Double
mctsKuct = 0.5

On définit également une fonction selectUcb permettant de trouver le noeud enfant qui maximise la formule UCB1.

selectUcb :: Node s -> ST s (Node s)
selectUcb node = do
    pNsims <- readSTRef (nodeNsims node)
    let children = nodeChildren node
    let bestUcb1Func (sk, k) i n = do
            reward <- readSTRef (nodeReward n)
            nsims <- readSTRef (nodeNsims n)
            let si = ucb1 reward nsims pNsims 
            return $ if si > sk then (si, i) else (sk, k)
    (_, k) <- M.ifoldM' bestUcb1Func (-1, -1) children 
    M.read children k

Étapes de l’algo MCTS

Les étapes de l’algo MTCS s’implémentent alors assez naturellement. Pour la sélection + expansion, on construit un nouveau noeud si on est arrivé à une feuille. Sinon on sélecttionne un noeud enfant avec UCB et on continue récursivement le parcours à partir du noeud sélectionné.

selectAndExpand :: Node s -> ST s (Node s)
selectAndExpand node = 
    if isRunning (nodeGame node)
    then do
        lastI <- readSTRef (nodeLastI node)
        if lastI < nodeNmoves node
        then do
            nodeRef <- newSTRef node
            cNode <- mkLeaf lastI nodeRef
            M.write (nodeChildren node) lastI cNode
            modifySTRef' (nodeLastI node) (+1)
            return cNode
        else selectUcb node >>= selectAndExpand 
    else return node 

Pour l’étape de simulation, on effectue un jeu aléatoire sur une copie du jeu courant.

simulate :: Gen s -> Node s -> ST s Status
simulate gen node = cloneGame (nodeGame node) >>= playoutRandom gen 

Enfin pour l’étape de rétropropagation, la monade ST et les références mutables permettent de faire une remontée d’arbre classique et ainsi de mettre à jour les données des noeuds concernés.

backpropagate :: Status -> Node s -> ST s ()
backpropagate status node = do
    modifySTRef' (nodeReward node) (+ computeScore (nodePlayer node) status)
    modifySTRef' (nodeNsims node) (+1)
    case nodeParent node of
        Nothing -> return ()
        Just parentRef -> readSTRef parentRef >>= backpropagate status

Interface Bot et type BotMcts

Il ne reste plus qu’à regrouper le tout et à l’intégrer dans notre interface Bot existante.

class Bot s b where
    genmove :: b -> Game s -> ST s Int

Pour cela, on définit un type BotMcts contenant le nombre d’itérations MCTS à réaliser et le générateur à utiliser pour les parties aléatoires.

data BotMcts s = BotMcts
    { mctsNiters :: Int
    , mctsGen :: GenST s
    }

Et on instancie la classe Bot pour le type BotMcts, où la fonction genmove construit un noeud racine, effectue les itérations MCTS puis retourne le meilleur coup estimé.

instance Bot s (BotMcts s) where
    genmove (BotMcts niters gen) game = do
        root <- mkRoot game 
        replicateM_ niters $ do
            leaf <- selectAndExpand root
            status <- simulate gen leaf
            backpropagate status leaf
        bestNode root

Résultats

On retrouve des résultats similaires à l’implémentation C++. Contre Random, MCTS est nettement meilleur.

Rouge : MctsX gagne. Jaune : Random gagne. Bleu : égalité.

En augmentant le nombre d’itérations MCTS, on augmente la performance de l’IA.

Rouge : MctsX gagne. Jaune : Mcts512 gagne. Bleu : égalité.

Enfin, à temps comparable, MCTS est meilleur que Monte-Carlo.

Rouge : McX gagne. Jaune : MctsX gagne. Bleu : égalité.

Conclusion

Dans cet article, on a vu comment implémenter l’algorithme Monte-Carlo Tree Search en Haskell. On notera essentiellement que la monade ST permet de faire des modifications mémoires en-place et de définir des références. Au final, l’implémentation Haskell n’est pas très différente de l’implémentation C++ présentée précédemment.