# Monte Carlo Tree Search Optimization and Loss Prevention

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!

• Performance tests must be done in a compiled project, not in a Playground. Commented May 11, 2021 at 9:14

## 1 Answer

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~