def test(run, snapshot=-1, **kwargs):
boardsize = runs.info(run)['boardsize']
worlds = hex.Hex.initial(n_envs=1024, boardsize=boardsize)
network = storage.load_raw(run, 'model')
sd = storage.load_snapshot(run, n=snapshot)['agent']
network.load_state_dict(storage.expand(sd)['network'])
A = mcts.MCTSAgent(network.cuda(), **kwargs)
network = storage.load_raw(run, 'model')
sd = storage.load_snapshot(run, n=snapshot)['agent']
network.load_state_dict(storage.expand(sd)['network'])
B = mcts.DummyAgent(network.cuda())
fst = analysis.rollout(worlds, [A, B], n_reps=1, eval=False)
snd = analysis.rollout(worlds, [B, A], n_reps=1, eval=False)
wins = count_wins(fst.transitions) + count_wins(snd.transitions).flipud()
rate = wins[0] / wins.sum()
elo = torch.log(rate) - torch.log(1 - rate)
kl = (kl_div(fst.decisions['0']) + kl_div(snd.decisions['0'])) / 2
ent = (rel_entropy(fst.decisions['0']) +
rel_entropy(snd.decisions['0'])) / 2
return {'elo': elo.item(), 'kl': kl.item(), 'ent': ent.item()}
def snapshot_kl_divs(run):
import pandas as pd
from pavlov import runs, storage
from boardlaw import hex
from boardlaw.main import mix
import torch
from tqdm.auto import tqdm
m = storage.load_raw(run, 'model')
worlds = mix(hex.Hex.initial(n_envs=16*1024, boardsize=runs.info(run)['params']['boardsize']))
logits = {}
for idx in tqdm(storage.snapshots(run)):
sd = storage.load_snapshot(run, idx)['agent']
m.load_state_dict(storage.expand(sd)['network'])
logits[idx] = m(worlds).logits.detach()
kldivs = {}
for i in logits:
for j in logits:
li = logits[i]
lj = logits[j]
terms = -li.exp().mul(lj - li)
mask = torch.isfinite(terms)
kldiv = terms.where(mask, torch.zeros_like(terms)).sum(-1)
kldivs[i, j] = kldiv.mean().item()
df = pd.Series(kldivs).unstack()
return df
def adam_way(run, i, Bsmall):
sd = storage.load_snapshot(run, i)
beta1, beta2 = sd['opt']['param_groups'][0]['betas']
step = sd['opt']['state'][0]['step']
m_bias = 1 - beta1**step
v_bias = 1 - beta2**step
opt = sd['opt']['state']
m = 1 / m_bias * torch.cat([s['exp_avg'].flatten() for s in opt.values()])
v = 1 / v_bias * torch.cat(
[s['exp_avg_sq'].flatten() for s in opt.values()])
# Follows from chasing the var through the defn of m
inflator = (1 - beta1**2) / (1 - beta1)**2
S = Bsmall * (v.mean() - m.pow(2).mean())
G2 = inflator * m.pow(2).mean()
return arrdict.arrdict(S=S,
G2=G2,
B=(S / G2),
v=v.mean(),
m=m.mean(),
m2=m.pow(2).mean(),
step=torch.as_tensor(step)).item()
def run():
n_envs = 8
world = worldfunc(n_envs)
agent = agentfunc()
agent.network = agent.network
agent.kwargs['n_nodes'] = 512
sd = storage.load_snapshot('*perky-boxes*', 64)
agent.load_state_dict(sd['agent'])
mhx = mohex.MoHexAgent()
trace = rollout(world, [agent, mhx])
def agent(run, idx=None, device='cpu'):
try:
network = storage.load_raw(run, 'model', device)
agent = MCTSAgent(network)
if idx is None:
sd = storage.load_latest(run)
else:
sd = storage.load_snapshot(run, idx)
agent.load_state_dict(sd['agent'])
return agent
except IOError:
return None
def snapshot_data(new_runs):
snapshots = {}
for _, r in tqdm(list(new_runs.iterrows()), desc='snapshots'):
for i, s in storage.snapshots(r.run).items():
stored = storage.load_snapshot(r.run, i)
if 'n_samples' in stored:
snapshots[r.run, i] = {
'samples': stored['n_samples'],
'flops': stored['n_flops']
}
snapshots = (pd.DataFrame.from_dict(
snapshots,
orient='index').rename_axis(index=('run', 'idx')).reset_index())
# snapshots['id'] = snapshots.index.to_series()
return snapshots
def load(run):
snapshots = pd.DataFrame.from_dict(storage.snapshots(run), orient='index')
info, losses = {}, {}
for i, row in snapshots.iterrows():
losses[row.boardsize, row.depth,
row.width] = storage.load_snapshot(run, i)['losses']
info[row.boardsize, row.depth, row.width] = {
'macs': row.n_macs,
'params': row.n_params
}
losses = pd.DataFrame(losses)
losses.index.name = 'step'
losses.columns.names = ('boardsize', 'depth', 'width')
info = pd.DataFrame(info)
info.columns.names = ('boardsize', 'depth', 'width')
return losses, info
def gradients(run, i, n_envs=16 * 1024, buffer_len=64, device='cuda'):
#TODO: Restore league and sched when you go back to large boards
worlds = mix(Hex.initial(n_envs, device=device))
network = storage.load_raw(run, 'model')
agent = MCTSAgent(network)
opt = torch.optim.Adam(network.parameters(), lr=0., amsgrad=True)
scaler = torch.cuda.amp.GradScaler()
sd = storage.load_snapshot(run, i)
agent.load_state_dict(sd['agent'])
opt.load_state_dict(sd['opt'])
scaler.load_state_dict(sd['scaler'])
buffer = []
idxs = (torch.randint(buffer_len, (n_envs, ),
device=device), torch.arange(n_envs, device=device))
while True:
# Collect experience
while len(buffer) < buffer_len:
with torch.no_grad():
decisions = agent(worlds, value=True)
new_worlds, transition = worlds.step(decisions.actions)
buffer.append(
arrdict.arrdict(worlds=worlds,
decisions=decisions.half(),
transitions=half(transition)).detach())
worlds = new_worlds
log.info(f'({len(buffer)}/{buffer_len}) actor stepped')
# Optimize
chunk, buffer = as_chunk(buffer, n_envs)
optimize(network, scaler, opt, chunk[idxs])
log.info('learner stepped')
yield torch.cat([
p.grad.flatten() for p in network.parameters()
if p.grad is not None
])
def agent(run, idx=None, device='cpu'):
try:
network = storage.load_raw(run, 'model', device)
except IOError:
log.warn(f'No model file for "{run}"')
return None
agent = MCTSAgent(network)
try:
if idx is None:
sd = storage.load_latest(run)
else:
sd = storage.load_snapshot(run, idx)
except IOError:
log.warn(f'No state dict file for "{run}"')
return None
agent.load_state_dict(sd['agent'])
return agent