aboutsummaryrefslogtreecommitdiff
path: root/src/nn_train.py
diff options
context:
space:
mode:
Diffstat (limited to 'src/nn_train.py')
-rwxr-xr-xsrc/nn_train.py116
1 files changed, 65 insertions, 51 deletions
diff --git a/src/nn_train.py b/src/nn_train.py
index 8025b68..73a5eaf 100755
--- a/src/nn_train.py
+++ b/src/nn_train.py
@@ -48,7 +48,7 @@ import pandas as pd
import torch
from tqdm import tqdm
-from torch.nn import MSELoss, Linear, Module
+from torch.nn import Embedding, MSELoss, Linear, Module
from torch import Tensor
@@ -97,38 +97,35 @@ def _fmt_matrix(mat: np.ndarray, per_line: int = 8) -> str:
def td_lambda_targets(
- values: Tensor,
+ leaf_values: Tensor,
result: float,
td_lambda: float,
) -> Tensor:
- # values: V(s_t) for t = 0..T-1, shape [T]
- # result: final result R of the game
-
- # This is ema, computed backwards
- # target[T-1] = R
- # target[t] = (1 - lambda) * V(s_{t+1}) + lambda * target[t+1]
- length = len(values)
- targets = torch.empty_like(values)
+ targets = torch.empty_like(leaf_values)
targets[-1] = result
- for ply in range(length - 2, -1, -1):
- targets[ply] = (1.0 - td_lambda) * values[ply + 1] + td_lambda * targets[
+
+ # G[t] = (1-lambda)V[t+1] + lambda G[t+1]
+ for ply in range(len(leaf_values) - 2, -1, -1):
+ targets[ply] = (1.0 - td_lambda) * leaf_values[ply + 1] + td_lambda * targets[
ply + 1
]
+
return targets
@dataclass(frozen=True)
class SelfPlayRecords:
- positions: Tensor # [n, inp_dim] feature vector of each position
+ positions: Tensor # [n, inp_dim]
game_ids: Tensor # [n] id of the game the position belongs to
- plies: Tensor # [n] ply number (position is before this move)
- results: Tensor # [n] final result from black's perspective (+1 black won, -1 white won, 0 draw)
+ results: Tensor # [n] final played result from black's perspective
+ terminals: Tensor # [n] true when the recorded position is a terminal leaf
+ terminal_values: Tensor # [n] exact value of a terminal leaf, otherwise 0
@classmethod
def load_selfplay(
cls, fn: Path, size: int, device: torch.device
) -> SelfPlayRecords:
- # [ net input: 0-inp_dim ] [game_id] [ply] [game result]
+ # [net input] [game_id] [played result] [terminal flag] [leaf result]
df = pd.read_csv(fn, header=None)
inp_dim = NET_DIMS[size][0]
return cls(
@@ -141,12 +138,18 @@ class SelfPlayRecords:
df.iloc[:, inp_dim].to_numpy(dtype=np.int64),
device=device,
),
- plies=torch.tensor(
- df.iloc[:, inp_dim + 1].to_numpy(dtype=np.int64),
+ results=torch.tensor(
+ df.iloc[:, inp_dim + 1].to_numpy(dtype=np.float32),
+ dtype=torch.float32,
+ device=device,
+ ),
+ terminals=torch.tensor(
+ df.iloc[:, inp_dim + 2].to_numpy(dtype=np.bool_),
+ dtype=torch.bool,
device=device,
),
- results=torch.tensor(
- df.iloc[:, inp_dim + 2].to_numpy(dtype=np.float32),
+ terminal_values=torch.tensor(
+ df.iloc[:, inp_dim + 3].to_numpy(dtype=np.float32),
dtype=torch.float32,
device=device,
),
@@ -160,8 +163,9 @@ class SelfPlayRecords:
return cls(
positions=torch.concat([spr.positions for spr in sprs]),
game_ids=torch.concat([spr.game_ids for spr in sprs]),
- plies=torch.concat([spr.plies for spr in sprs]),
results=torch.concat([spr.results for spr in sprs]),
+ terminals=torch.concat([spr.terminals for spr in sprs]),
+ terminal_values=torch.concat([spr.terminal_values for spr in sprs]),
)
def make_training_data(
@@ -175,9 +179,16 @@ class SelfPlayRecords:
for game_id in torch.unique(self.game_ids):
in_game = self.game_ids == game_id
- game_values = values[in_game]
+ game_terminals = self.terminals[in_game]
+ game_terminal_values = self.terminal_values[in_game]
+ game_values = values[in_game].clone()
+ game_values[game_terminals] = game_terminal_values[game_terminals]
game_result = self.results[in_game][0].item()
- targets[in_game] = td_lambda_targets(game_values, game_result, td_lambda)
+ game_targets = td_lambda_targets(game_values, game_result, td_lambda)
+ # Terminal positions are also direct supervised examples. Do this
+ # after constructing the TD sequence so it does not reset the EMA.
+ game_targets[game_terminals] = game_terminal_values[game_terminals]
+ targets[in_game] = game_targets
return self.positions, targets
@@ -190,11 +201,18 @@ class SelfPlayRecords:
class Net(Module):
def __init__(self, inp_dim: int, hidden: int) -> None:
super().__init__()
+ # learned 7-way cell-state embedding (empty, {black, white} x
+ # {flat, standing, road}); the slot layout must match nn.c
+ self.embedding = Embedding(7, 1)
self.hidden = Linear(inp_dim, hidden, bias=True)
self.output = Linear(hidden, 1, bias=True)
def forward(self, x: Tensor) -> Tensor:
- hidden = torch.relu(self.hidden(x))
+ # x: [n, 3 + n_cells]; the first 3 columns are the fixed head
+ # inputs, the rest are raw cell states (0..6)
+ n = x.shape[0]
+ inp = torch.cat([x[:, :3], self.embedding(x[:, 3:].long()).view(n, -1)], dim=1)
+ hidden = torch.relu(self.hidden(inp))
return torch.tanh(self.output(hidden)).squeeze(-1)
def fit(
@@ -272,6 +290,7 @@ class Net(Module):
prefix = ("cand_" if cand else "") + ("five" if size == 5 else "six")
upper = "FIVE" if size == 5 else "SIX"
+ E = self.embedding.weight.detach().cpu().numpy()
W = self.hidden.weight.detach().cpu().numpy()
B = self.hidden.bias.detach().cpu().numpy()
LW = self.output.weight.detach().cpu().numpy().ravel()
@@ -283,6 +302,10 @@ class Net(Module):
_ = f.write(f"/* {p} */\n")
_ = f.write(f'#include "weights_{size}.h"\n\n')
_ = f.write(
+ f"const float {prefix}_embedding[7] =\n"
+ f"{{\n{_fmt_flat(E, per_line=7)}\n}};\n\n"
+ )
+ _ = f.write(
f"const float {prefix}_dense_weights[{upper}_DENSE_NUM]"
f"[{upper}_INP_NUM] NN_ALIGNED =\n{{\n{_fmt_matrix(W)}\n}};\n\n"
)
@@ -340,7 +363,6 @@ class TournamentResult:
champ: int
cand: int
draws: int
- cand_rate: float
def run_matches(
@@ -365,14 +387,7 @@ def run_matches(
champ = winners.count("current")
cand = winners.count("candidate")
draws = winners.count("draw")
- aborts = len(winners) - champ - cand - draws
- played = len(winners) - aborts
- return TournamentResult(
- champ=champ,
- cand=cand,
- draws=draws,
- cand_rate=(cand + 0.5 * draws) / played if played else 0.0,
- )
+ return TournamentResult(champ=champ, cand=cand, draws=draws)
# ===================================================================
@@ -420,10 +435,10 @@ def run_step(
print(f"\n== iteration {iter} ==")
- # 1. self-play with the installed champion (include/weights_<size>.c)
+ # 1. rebuild selfplay
_ = subprocess.run(["make", "selfplay"], cwd=root, check=True, capture_output=True)
- # 2. generate self-play data (one game per process, in parallel)
+ # 2. generate self-play data
eps = max(eps_min, eps_start * (eps_decay ** (iter - 1)))
csv = make_csv(iter)
print(f"selfplay, eps={eps:.4f}, depth {depth}, {games} games on {procs} procs")
@@ -446,7 +461,7 @@ def run_step(
# 3. fine-tune the champion on TD(lambda) targets -> candidate
device = next(champ.parameters()).device
csvs = [make_csv(i) for i in range(max(iter - replay_iters, 1), iter + 1)]
- print("Loading data from", csvs)
+ print("Loading data from", [str(c) for c in csvs])
records = SelfPlayRecords.load_merge(csvs, size, device)
cand = copy.deepcopy(champ)
@@ -465,7 +480,7 @@ def run_step(
cand=True,
)
- # 4. build the match driver (champion + candidate linked)
+ # 4. build the match driver (champion + candidate)
_ = subprocess.run(["make", "match"], cwd=root, check=True, capture_output=True)
# 5. tournament + decision
@@ -480,8 +495,7 @@ def run_step(
decided = tournament_result.champ + tournament_result.cand
rate = tournament_result.cand / decided if decided else 0.0
print(
- f"champ {tournament_result.champ} - {tournament_result.cand} cand, {tournament_result.draws} draws, "
- + f"winrate={rate:.3f}, cand_rate={tournament_result.cand_rate:.3f}"
+ f"champ {tournament_result.champ} - {tournament_result.cand} cand, {tournament_result.draws} draws, winrate={rate:.3f}"
)
promote = decided >= int(0.25 * tournament_games) and rate >= promote_rate
@@ -489,7 +503,7 @@ def run_step(
cand.to_c(
size,
root / "include" / f"weights_{size}.c",
- (f"v{iter:03d}, val_corr={corr:.4f}",),
+ (f"v{iter:03d}, val_corr={corr:.4f}, rate={rate:.4f}",),
)
print(f"== candidate v{iter:03d} PROMOTED (include/weights_{size}.c) ==")
return cand
@@ -513,23 +527,23 @@ def main() -> None:
)
_ = p.add_argument("--iters", type=int, default=10, help="max loop iterations")
_ = p.add_argument(
- "--games", type=int, default=2000, help="selfplay games per iteration"
+ "--games", type=int, default=100, help="selfplay games per iteration"
)
_ = p.add_argument(
"--replay-iters",
type=int,
- default=4,
+ default=0,
help="number of previous self-play instances to mix in with current dataset during training",
)
_ = p.add_argument(
"--eps-start",
type=float,
default=0.15,
- help="initial self-play epsilon (random-move probability); kept close "
- "to --tour-eps so the value net trains on the policy it is deployed "
- "under",
+ help="initial self-play epsilon: in asymmetric self-play one side (chosen at random per game) plays a uniform random move with this probability while the greedy side has its plies recorded",
+ )
+ _ = p.add_argument(
+ "--eps-min", type=float, default=0.05, help="self-play epsilon floor"
)
- _ = p.add_argument("--eps-min", type=float, default=0.05, help="epsilon floor")
_ = p.add_argument(
"--eps-decay", type=float, default=0.9, help="per-iteration epsilon decay"
)
@@ -548,20 +562,20 @@ def main() -> None:
_ = p.add_argument(
"--tournament-eps",
type=float,
- default=0.05,
+ default=0.01,
help="tournament epsilon (random-move probability)",
)
_ = p.add_argument(
"--tournament-games",
type=int,
default=None,
- help="tournament games (10/tour_eps if unset)",
+ help="tournament games (1/tournament-eps if unset)",
)
_ = p.add_argument(
"--promote",
type=float,
- default=0.57,
- help="candidate win rate among non-drawn tournament games needed to promote (~p<0.03 one-sided binomial at the default 200 games)",
+ default=0.0,
+ help="candidate win rate among non-drawn tournament games needed to promote",
)
_ = p.add_argument(
"--procs",
@@ -582,7 +596,7 @@ def main() -> None:
args = p.parse_args()
if args.tournament_games is None:
- args.tournament_games = max(20, int(round(10 / args.tournament_eps)))
+ args.tournament_games = max(100, int(round(1 / args.tournament_eps)))
if args.procs is None:
args.procs = max(1, (os.cpu_count() or 1) // 2)
generated by cgit v1.3.1 (git 2.54.0) at 2026-10-15 08:40:53 +0000