2
\$\begingroup\$

I'm working on an implementation of Monte Carlo Tree Search in Swift.

It's not bad, but it could be better! I'm principally interested in making my algorithm:

  1. faster (more iterations/second)
  2. prioritize moves that prevent instant losses (you'll see...)

Here is the main driver:

final class MonteCarloTreeSearch {
    var player: Player
    var timeBudget: Double
    var maxDepth: Int
    var explorationConstant: Double
    var root: Node?
    var iterations: Int

    init(for player: Player, timeBudget: Double = 5, maxDepth: Int = 5, explorationConstant: Double = sqrt(2)) {
        self.player = player
        self.timeBudget = timeBudget
        self.maxDepth = maxDepth
        self.explorationConstant = explorationConstant
        self.iterations = 0
    }
    
    func update(with game: Game) {
        if let newRoot = findNode(for: game) {
            newRoot.parent = nil
            newRoot.move = nil
            root = newRoot
        } else {
            root = Node(game: game)
        }
    }

    func findMove(for game: Game? = nil) -> Move? {
        iterations = 0
        let start = CFAbsoluteTimeGetCurrent()
        if let game = game {
            update(with: game)
        }
        while CFAbsoluteTimeGetCurrent() - start < timeBudget {
            refine()
            iterations += 1
        }
        print("Iterations: \(iterations)")
        return bestMove
    }
    
    private func refine() {
        let leafNode = root!.select(explorationConstant)
        let value = rollout(leafNode)
        leafNode.backpropogate(value)
    }
    
    private func rollout(_ node: Node) -> Double {
        var depth = 0
        var game = node.game
        while !game.isFinished {
            if depth >= maxDepth { break }
            guard let move = game.randomMove() else { break }
            game = game.update(move)
            depth += 1
        }
        let value = game.evaluate(for: player).value
        return value
    }
    
    private var bestMove: Move? {
        root?.selectChildWithMaxUcb(0)?.move
    }
    
    private func findNode(for game: Game) -> Node? {
        guard let root = root else { return nil }
        var queue = [root]
        while !queue.isEmpty {
            let head = queue.removeFirst()
            if head.game == game {
                return head
            }
            for child in head.children {
                queue.append(child)
            }
        }
        return nil
    }
}

I built this driver with a maxDepth argument because playouts/rollouts in my real game are fairly long and I have a access to a decent static evaluation function. Also, the BFS findNode method is so that I can reuse parts of the tree.

Here's what a node in the driver looks like:

final class Node {
    weak var parent: Node?
    var move: Move?
    var game: Game
    var untriedMoves: [Move]
    var children: [Node]
    var cumulativeValueFor: Double
    var cumulativeValueAgainst: Double
    var visits: Double

    init(parent: Node? = nil, move: Move? = nil, game: Game) {
        self.parent = parent
        self.move = move
        self.game = game
        self.children = []
        self.untriedMoves = game.availableMoves()
        self.cumulativeValueFor = 0
        self.cumulativeValueAgainst = 0
        self.visits = 0
    }
    
    var isFullyExpanded: Bool {
        untriedMoves.isEmpty
    }
    
    lazy var isTerminal: Bool = {
        game.isFinished
    }()
    
    func select(_ c: Double) -> Node {
        var leafNode = self
        while !leafNode.isTerminal {
            if !leafNode.isFullyExpanded {
                return leafNode.expand()
            } else {
                leafNode = leafNode.selectChildWithMaxUcb(c)!
            }
        }
        return leafNode
    }
    
    func expand() -> Node {
        let move = untriedMoves.popLast()!
        let nextGame = game.update(move)
        let childNode = Node(parent: self, move: move, game: nextGame)
        children.append(childNode)
        return childNode
    }
    
    func backpropogate(_ value: Double) {
        visits += 1
        cumulativeValueFor += value
        if let parent = parent {
            parent.backpropogate(value)
        }
    }
    
    func selectChildWithMaxUcb(_ c: Double) -> Node? {
        children.max { $0.ucb(c) < $1.ucb(c) }
    }

    func ucb(_ c: Double) -> Double {
        q + c * u
    }
    
    private var q: Double {
        let value = cumulativeValueFor - cumulativeValueAgainst
        return value / visits
    }
    
    private var u: Double {
        sqrt(log(parent!.visits) / visits)
    }
}

extension Node: CustomStringConvertible {
    var description: String {
        guard let move = move else { return "" }
        return "\(move) (\(cumulativeValueFor)/\(visits))"
    }
}

I don't think there's anything extraordinary about my node object? (I am hoping, though, that I can do something to/about q so that I might prevent an "instant" loss in my test game...


I've been testing this implementation of MCTS on a 1-D variant of "Connect 4".

Here's the game and all of it's primitives:

enum Player: Int {
    case one = 1
    case two = 2
    
    var opposite: Self {
        switch self {
        case .one: return .two
        case .two: return .one
        }
    }
}

extension Player: CustomStringConvertible {
    var description: String {
        "\(rawValue)"
    }
}

typealias Move = Int

enum Evaluation {
    case win
    case loss
    case draw
    case ongoing(Double)
    
    var value: Double {
        switch self {
        case .win: return 1
        case .loss: return 0
        case .draw: return 0.5
        case .ongoing(let v): return v
        }
    }
}

struct Game {
    var array: Array<Int>
    var currentPlayer: Player
    
    init(length: Int = 10, currentPlayer: Player = .one) {
        self.array = Array.init(repeating: 0, count: length)
        self.currentPlayer = currentPlayer
    }
    
    var isFinished: Bool {
        switch evaluate() {
        case .ongoing: return false
        default: return true
        }
    }

    func availableMoves() -> [Move] {
        array
            .enumerated()
            .compactMap { $0.element == 0 ? Move($0.offset) : nil}
    }
    
    func update(_ move: Move) -> Self {
        var copy = self
        copy.array[move] = currentPlayer.rawValue
        copy.currentPlayer = currentPlayer.opposite
        return copy
    }
    
    func evaluate(for player: Player) -> Evaluation {
        let player3 = three(for: player)
        let oppo3 = three(for: player.opposite)
        let remaining0 = array.contains(0)
        switch (player3, oppo3, remaining0) {
        case (true, true, _): return .draw
        case (true, false, _): return .win
        case (false, true, _): return .loss
        case (false, false, false): return .draw
        default: return .ongoing(0.5)
        }
    }
    
    private func three(for player: Player) -> Bool {
        var count = 0
        for slot in array {
            if slot == player.rawValue {
                count += 1
            } else {
                count = 0
            }
            if count == 3 {
                return true
            }
        }
        return false
    }
}

extension Game {
    func evaluate() -> Evaluation {
        evaluate(for: currentPlayer)
    }
    
    func randomMove() -> Move? {
        availableMoves().randomElement()
    }
}

extension Game: CustomStringConvertible {
    var description: String {
        return array.reduce(into: "") { result, i in
            result += String(i)
        }
    }
}

extension Game: Equatable {}

While there are definitely efficiencies to be gained in optimizing the evaluate/three(for:) scoring methods, I'm more concerned about improving the performance of the driver and the node as this "1d-connect-3" game isn't my real game. That said, if there's a huge mistake here and a simple fix I'll take it!

Another note: I am actually using ongoing(Double) in my real game (I've got a static evaluation function that can reliably score a player as 1-99% likely to win).


A bit of Playground code:

var mcts = MonteCarloTreeSearch(for: .two, timeBudget: 5, maxDepth: 3)
var game = Game(length: 10)
// 0000000000
game = game.update(0) // player 1
// 1000000000
game = game.update(8) // player 2
// 1000000020
game = game.update(1) // player 1
// 1100000020
let move1 = mcts.findMove(for: game)!
// usually 7 or 9... and not 2
print(mcts.root!.children)
game = game.update(move1) // player 2
mcts.update(with: game)
game = game.update(4) // player 1
mcts.update(with: game)
let move2 = mcts.findMove()!

Unfortunately, move1 in this sample "playthru" doesn't try and prevent the instant win-condition on the next turn for player 1?! (I know that orthodox Monte Carlo Tree Search is in the business of maximizing winning not minimizing losing, but not picking 2 here is unfortunate).

So yeah, any help in making all this faster (perhaps through parallelization), and fixing the "instant-loss" business would be swell!

\$\endgroup\$
1
  • \$\begingroup\$ Performance tests must be done in a compiled project, not in a Playground. \$\endgroup\$ – Martin R May 11 at 9:14
0
\$\begingroup\$

I've been mulling this over and have some preliminary answers! (But would still appreciate some help!!)


F. Teytaud and O. Teytaud (2010) inspired a solution to the "loss prevention" problem. The paper recommends selecting the decisive or anti-decisive (win for opponent) move if it exists, else select as normal:

// final class Node {
// ...
    func select(_ c: Double) -> Node {
        var leafNode = self
        if let decisiveNode = selectDecisiveChild() {
            return decisiveNode
        }
        while !leafNode.isTerminal {
            if !leafNode.isFullyExpanded {
                return leafNode.expand()
            } else {
                leafNode = leafNode.selectChildWithMaxUcb(c)!
            }
        }
        return leafNode
    }

    private func selectDecisiveChild() -> Node? {
        children.filter { $0.game.isFinished }.randomElement()
    }
// ...
// }

To squeeze out more iterations per second, I've implemented "leaf parallelization" as defined in G. M. J.-B. Chaslot, M. H. M. Winands, and H. J. van den Herik (2008). While other methods of parallelization look more promising, this method seemed easiest:

// final class MonteCarloTreeSearch {
// ...

    private let activeProcessors = ProcessInfo.processInfo.activeProcessorCount
    private let dispatchGroup = DispatchGroup()

// ...

    private func refine() {
        let leafNode = root!.select(explorationConstant)
        var values = Array<Double>(repeating: 0, count: activeProcessors)
        var mutex = os_unfair_lock()
        for i in 0..<activeProcessors {
            DispatchQueue.global(qos: .userInitiated).async(group: dispatchGroup) {
                let value = self.rollout(leafNode)
                os_unfair_lock_lock(&mutex)
                values[i] = value
                self.iterations += 1
                os_unfair_lock_unlock(&mutex)
            }
        }
        _ = dispatchGroup.wait(timeout: .distantFuture)
        for value in values {
            leafNode.backpropagate(value)
        }
    }

// ...
// }

I'm still a novice of DispatchQueue. So if I'm using the locks incorrectly, I'd love to know... but it seems to work!


Looking forward to reviewing other answers~

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.