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:
- faster (more iterations/second)
- 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!
