Skip to content

Commit a8f1446

Browse files
committed
[Example] Add MCTS example
ghstack-source-id: 4cf2a16 Pull Request resolved: pytorch#2796
1 parent a31dca3 commit a8f1446

File tree

5 files changed

+262
-2
lines changed

5 files changed

+262
-2
lines changed

examples/trees/mcts.py

+129
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,129 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
#
3+
# This source code is licensed under the MIT license found in the
4+
# LICENSE file in the root directory of this source tree.
5+
6+
import time
7+
8+
import torch
9+
import torchrl
10+
import torchrl.envs
11+
import torchrl.modules.mcts
12+
from tensordict import TensorDict, TensorDictBase
13+
14+
from torchrl.data.map import MCTSForest, Tree
15+
from torchrl.envs import EnvBase
16+
17+
start_time = time.time()
18+
19+
pgn_or_fen = "fen"
20+
mask_actions = True
21+
22+
env = torchrl.envs.ChessEnv(
23+
include_pgn=False,
24+
include_fen=True,
25+
include_hash=True,
26+
include_hash_inv=True,
27+
include_san=True,
28+
stateful=True,
29+
mask_actions=mask_actions,
30+
)
31+
32+
33+
class TransformReward:
34+
def __init__(self):
35+
self.first_turn = None
36+
37+
def __call__(self, td):
38+
if self.first_turn is None and "turn" in td:
39+
self.first_turn = td["turn"]
40+
print(f"first turn: {self.first_turn}")
41+
42+
if "reward" not in td:
43+
return td
44+
reward = td["reward"]
45+
if reward == 0.5:
46+
reward = 0
47+
# elif reward == 1 and td["turn"] == env.lib.WHITE:
48+
elif reward == 1 and td["turn"] == self.first_turn:
49+
reward = -reward
50+
51+
td["reward"] = reward
52+
return td
53+
54+
def reset(self, td):
55+
self.first_turn = None
56+
57+
58+
# ChessEnv sets the reward to 0.5 for a draw and 1 for a win for either player.
59+
# Need to transform the reward to be:
60+
# white win = 1
61+
# draw = 0
62+
# black win = -1
63+
env = env.append_transform(TransformReward())
64+
65+
forest = torchrl.data.MCTSForest()
66+
forest.reward_keys = env.reward_keys
67+
forest.done_keys = env.done_keys
68+
forest.action_keys = env.action_keys
69+
70+
if mask_actions:
71+
forest.observation_keys = [f"{pgn_or_fen}_hash", "turn", "action_mask"]
72+
else:
73+
forest.observation_keys = [f"{pgn_or_fen}_hash", "turn"]
74+
75+
76+
def tree_format_fn(tree):
77+
td = tree.rollout[-1]["next"]
78+
return [
79+
td["san"],
80+
td[pgn_or_fen].split("\n")[-1],
81+
tree.wins,
82+
tree.visits,
83+
]
84+
85+
86+
def get_best_move(fen, mcts_steps, rollout_steps):
87+
root = env.reset(TensorDict({"fen": fen}))
88+
tree = torchrl.modules.mcts.MCTS(forest, root, env, mcts_steps, rollout_steps)
89+
moves = []
90+
91+
for subtree in tree.subtree:
92+
san = subtree.rollout[0]["next", "san"]
93+
reward_sum = subtree.wins
94+
visits = subtree.visits
95+
value_avg = (reward_sum / visits).item()
96+
97+
if not subtree.rollout[0]["turn"]:
98+
value_avg = -value_avg
99+
100+
moves.append((value_avg, san))
101+
102+
moves = sorted(moves, key=lambda x: -x[0])
103+
104+
# print(tree.to_string(tree_format_fn))
105+
106+
print("------------------")
107+
for value_avg, san in moves:
108+
print(f" {value_avg:0.02f} {san}")
109+
print("------------------")
110+
111+
return moves[0][1]
112+
113+
114+
# White has M1, best move Rd8#. Any other moves lose to M2 or M1.
115+
fen0 = "7k/6pp/7p/7K/8/8/6q1/3R4 w - - 0 1"
116+
assert get_best_move(fen0, 100, 10) == "Rd8#"
117+
118+
# Black has M1, best move Qg6#. Other moves give rough equality or worse.
119+
fen1 = "6qk/2R4p/7K/8/8/8/8/4R3 b - - 1 1"
120+
assert get_best_move(fen1, 100, 10) == "Qg6#"
121+
122+
# White has M2, best move Rxg8+. Any other move loses.
123+
fen2 = "2R3qk/5p1p/7K/8/8/8/5r2/2R5 w - - 0 1"
124+
assert get_best_move(fen2, 1000, 10) == "Rxg8+"
125+
126+
end_time = time.time()
127+
total_time = end_time - start_time
128+
129+
print(f"Took {total_time} s")

torchrl/data/map/tree.py

+5
Original file line numberDiff line numberDiff line change
@@ -1363,6 +1363,11 @@ def valid_paths(cls, tree: Tree):
13631363
def __len__(self):
13641364
return len(self.data_map)
13651365

1366+
def __contains__(self, root: TensorDictBase):
1367+
if self.node_map is None:
1368+
return False
1369+
return root.select(*self.node_map.in_keys) in self.node_map
1370+
13661371
def to_string(self, td_root, node_format_fn=lambda tree: tree.node_data.to_dict()):
13671372
"""Generates a string representation of a tree in the forest.
13681373

torchrl/envs/custom/chess.py

+7-2
Original file line numberDiff line numberDiff line change
@@ -222,12 +222,15 @@ def lib(cls):
222222
return chess
223223

224224
_san_moves = []
225+
_san_move_to_index_map = {}
225226

226227
@_classproperty
227228
def san_moves(cls):
228229
if not cls._san_moves:
229230
with open(pathlib.Path(__file__).parent / "san_moves.txt", "r+") as f:
230231
cls._san_moves.extend(f.read().split("\n"))
232+
for idx, san_move in enumerate(cls._san_moves):
233+
cls._san_move_to_index_map[san_move] = idx
231234
return cls._san_moves
232235

233236
def _legal_moves_to_index(
@@ -255,7 +258,7 @@ def _legal_moves_to_index(
255258
board = self.board
256259

257260
indices = torch.tensor(
258-
[self._san_moves.index(board.san(m)) for m in board.legal_moves],
261+
[self._san_move_to_index_map[board.san(m)] for m in board.legal_moves],
259262
dtype=torch.int64,
260263
)
261264
mask = None
@@ -409,7 +412,9 @@ def _reset(self, tensordict=None):
409412
if move is None:
410413
dest.set("san", "<start>")
411414
else:
412-
dest.set("san", self.board.san(move))
415+
prev_board = self.board.copy()
416+
prev_board.pop()
417+
dest.set("san", prev_board.san(move))
413418
if self.include_fen:
414419
dest.set("fen", fen)
415420
if self.include_pgn:

torchrl/modules/mcts/__init__.py

+6
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
#
3+
# This source code is licensed under the MIT license found in the
4+
# LICENSE file in the root directory of this source tree.
5+
6+
from .mcts import MCTS

torchrl/modules/mcts/mcts.py

+115
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
# Copyright (c) Meta Platforms, Inc. and affiliates.
2+
#
3+
# This source code is licensed under the MIT license found in the
4+
# LICENSE file in the root directory of this source tree.
5+
6+
import torch
7+
import torchrl
8+
from tensordict import TensorDict, TensorDictBase
9+
10+
from torchrl.data.map import MCTSForest, Tree
11+
from torchrl.envs import EnvBase
12+
13+
C = 2.0**0.5
14+
15+
16+
def traversal_priority_UCB1(tree):
17+
subtree = tree.subtree
18+
visits = subtree.visits
19+
reward_sum = subtree.wins
20+
21+
# If it's black's turn, flip the reward, since black wants to
22+
# optimize for the lowest reward, not highest.
23+
if not subtree.rollout[0, 0]["turn"]:
24+
reward_sum = -reward_sum
25+
26+
parent_visits = tree.visits
27+
reward_sum = reward_sum.squeeze(-1)
28+
priority = (reward_sum + C * torch.sqrt(torch.log(parent_visits))) / visits
29+
priority[visits == 0] = float("inf")
30+
return priority
31+
32+
33+
def _traverse_MCTS_one_step(forest, tree, env, max_rollout_steps):
34+
done = False
35+
trees_visited = [tree]
36+
37+
while not done:
38+
if tree.subtree is None:
39+
td_tree = tree.rollout[-1]["next"].clone()
40+
41+
if (tree.visits > 0 or tree.parent is None) and not td_tree["done"]:
42+
actions = env.all_actions(td_tree)
43+
subtrees = []
44+
45+
for action in actions:
46+
td = env.step(env.reset(td_tree).update(action))
47+
new_node = torchrl.data.Tree(
48+
rollout=td.unsqueeze(0),
49+
node_data=td["next"].select(*forest.node_map.in_keys),
50+
count=torch.tensor(0),
51+
wins=torch.zeros_like(td["next", env.reward_key]),
52+
)
53+
subtrees.append(new_node)
54+
55+
# NOTE: This whole script runs about 2x faster with lazy stack
56+
# versus eager stack.
57+
tree.subtree = TensorDict.lazy_stack(subtrees)
58+
chosen_idx = torch.randint(0, len(subtrees), ()).item()
59+
rollout_state = subtrees[chosen_idx].rollout[-1]["next"]
60+
61+
else:
62+
rollout_state = td_tree
63+
64+
if rollout_state["done"]:
65+
rollout_reward = rollout_state[env.reward_key]
66+
else:
67+
rollout = env.rollout(
68+
max_steps=max_rollout_steps,
69+
tensordict=rollout_state,
70+
)
71+
rollout_reward = rollout[-1]["next", env.reward_key]
72+
done = True
73+
74+
else:
75+
priorities = traversal_priority_UCB1(tree)
76+
chosen_idx = torch.argmax(priorities).item()
77+
tree = tree.subtree[chosen_idx]
78+
trees_visited.append(tree)
79+
80+
for tree in trees_visited:
81+
tree.visits += 1
82+
tree.wins += rollout_reward
83+
84+
85+
def MCTS(
86+
forest: MCTSForest,
87+
root: TensorDictBase,
88+
env: EnvBase,
89+
num_steps: int,
90+
max_rollout_steps: int | None = None,
91+
) -> Tree:
92+
"""Performs Monte-Carlo tree search in an environment.
93+
94+
Args:
95+
forest (MCTSForest): Forest of the tree to update. If the tree does not
96+
exist yet, it is added.
97+
root (TensorDict): The root step of the tree to update.
98+
env (EnvBase): Environment to performs actions in.
99+
num_steps (int): Number of iterations to traverse.
100+
max_rollout_steps (int): Maximum number of steps for each rollout.
101+
"""
102+
if root not in forest:
103+
for action in env.all_actions(root):
104+
td = env.step(env.reset(root.clone()).update(action))
105+
forest.extend(td.unsqueeze(0))
106+
107+
tree = forest.get_tree(root)
108+
tree.wins = torch.zeros_like(td["next", env.reward_key])
109+
for subtree in tree.subtree:
110+
subtree.wins = torch.zeros_like(td["next", env.reward_key])
111+
112+
for _ in range(num_steps):
113+
_traverse_MCTS_one_step(forest, tree, env, max_rollout_steps)
114+
115+
return tree

0 commit comments

Comments
 (0)