| """Fresh-trajectory clipped PPO and weak offline supervision.""" |
| import random |
|
|
| import torch |
| from torch.nn import functional as F |
|
|
| from qojbot import rules as R |
| from qojbot.history import latest_log |
| from qojbot.model_benchmark import play_game |
|
|
| from .features import collate |
| from .network import score_loss |
| from .storage import cpu_tree |
|
|
|
|
| def terminal_targets(samples, device): |
| rewards = torch.tensor([sample['reward'] for sample in samples], device=device) |
| return rewards, rewards / 6. |
|
|
|
|
| def losses(model, samples, device, ppo=False, anchor_weight=.05, |
| history_weight=None, distill_weight=0.): |
| batch = collate(samples, device) |
| output = model(batch) |
| |
| counts = [len(sample['actions']) for sample in samples] |
| cursor = sum(counts) |
| rows = torch.arange(len(samples), device=device) |
| local = torch.cat([torch.arange(count) for count in counts]).to(device) |
| base = torch.cat([sample['base'] for sample in samples]).to(device) |
| logits = output['prior_scale'][batch['owner']] * base + output['correction'] |
| shape = (len(samples), max(counts)) |
| padded = torch.full(shape, -torch.inf, device=device) |
| padded[batch['owner'], local] = logits |
| teacher_padded = torch.full(shape, -torch.inf, device=device) |
| teacher_padded[batch['owner'], local] = base |
| valid = torch.arange(max(counts), device=device)[None] < torch.tensor(counts, device=device)[:, None] |
| log_probs = padded.log_softmax(-1) |
| probs = log_probs.exp() |
| teacher_logp = teacher_padded.log_softmax(-1) |
| teacher = teacher_logp.exp() |
| safe_logp = log_probs.masked_fill(~valid, 0) |
| safe_teacher_logp = teacher_logp.masked_fill(~valid, 0) |
| anchor = (teacher * (safe_teacher_logp - safe_logp)).sum(-1).mean() |
| entropy = -(probs * safe_logp).sum(-1).mean() |
| chosen = torch.tensor([sample['selected'] for sample in samples], device=device) |
| new_logp = log_probs[rows, chosen] |
| approx_kl = new_logp.sum() * 0 |
| if ppo: |
| old_logp = torch.tensor([sample['old_logp'] for sample in samples], device=device) |
| advantage = torch.tensor([sample['advantage'] for sample in samples], device=device) |
| ratio = (new_logp - old_logp).exp() |
| policy_loss = -torch.minimum(ratio * advantage, ratio.clamp(.8, 1.2) * advantage).mean() |
| approx_kl = ((ratio - 1) - (new_logp - old_logp)).mean() |
| else: |
| weights = torch.tensor([history_weight if history_weight is not None and |
| sample.get('origin') == 'history' else sample.get('weight', .15) |
| for sample in samples], device=device) |
| policy_loss = -(new_logp * weights).mean() |
| if distill_weight: |
| |
| |
| preferred = teacher_padded.argmax(-1) |
| policy_loss = policy_loss - distill_weight * log_probs[rows, preferred].mean() |
| offsets, offset = [], 0 |
| for count, sample in zip(counts, samples): |
| offsets.append(offset + sample['selected']) |
| offset += count |
| selected = torch.tensor(offsets, device=device) |
| rewards, transformed = terminal_targets(samples, device) |
| |
| |
| trusted = [i for i, s in enumerate(samples) if history_weight != 0 or s.get('origin') != 'history'] |
| zero = output['value'].sum() * 0 + output['score_logits'].sum() * 0 + output['win_logits'].sum() * 0 |
| value_loss = action_value_loss = win_loss = zero |
| if trusted: |
| value_loss = F.smooth_l1_loss(output['value'][trusted], transformed[trusted]) |
| action_value_loss = score_loss(output['score_logits'][selected[trusted]], rewards[trusted], model.support) |
| win_loss = F.binary_cross_entropy_with_logits(output['win_logits'][selected[trusted]], |
| (rewards[trusted] > 0).float()) |
| belief_indices = [i for i, sample in enumerate(samples) if 'belief_target' in sample and sample['role'] != 3] |
| belief_loss = output['value'].sum() * 0 |
| if belief_indices: |
| targets = torch.stack([samples[i]['belief_target'] for i in belief_indices]).to(device) |
| |
| belief_loss = F.smooth_l1_loss(output['belief_counts'][belief_indices], targets) |
| loss = (policy_loss + anchor_weight * anchor + .5 * value_loss + |
| .15 * action_value_loss + .05 * win_loss + .05 * belief_loss - |
| (.01 * entropy if ppo else 0)) |
| stats = {'loss': float(loss.detach().cpu()), 'policy': float(policy_loss.detach().cpu()), |
| 'anchor_kl': float(anchor.detach().cpu()), 'value': float(value_loss.detach().cpu()), |
| 'belief': float(belief_loss.detach().cpu()), |
| 'approx_kl': float(approx_kl.detach().cpu()) if ppo else 0., |
| 'samples': len(samples), 'actions': cursor} |
| return loss, stats |
|
|
|
|
| def update(model, optimizer, samples, device, ppo=False, anchor_weight=.05, |
| history_weight=None, distill_weight=0.): |
| model.train() |
| loss, stats = losses(model, samples, device, ppo, anchor_weight, history_weight, distill_weight) |
| if not torch.isfinite(loss): |
| raise RuntimeError('nonfinite training loss') |
| optimizer.zero_grad(set_to_none=True) |
| loss.backward() |
| norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 5.) |
| if not torch.isfinite(norm): |
| raise RuntimeError('nonfinite training gradients') |
| optimizer.step() |
| return stats |
|
|
|
|
| def batches(samples, batch_states, max_actions=4096): |
| pending, actions = [], 0 |
| for sample in samples: |
| count = len(sample['actions']) |
| if pending and (len(pending) >= batch_states or actions + count > max_actions): |
| yield pending |
| pending, actions = [], 0 |
| pending.append(sample) |
| actions += count |
| if pending: |
| yield pending |
|
|
|
|
| class Actor: |
| def __init__(self, policy, rng, teacher=False): |
| self.policy, self.rng, self.teacher = policy, rng, teacher |
| self.pending = [] |
|
|
| def select(self, candidates, sample, state): |
| if self.teacher: |
| logits, value = sample['base'], 0. |
| selected = int(logits.argmax()) |
| else: |
| logits, value = self.policy.predict(sample) |
| probs = logits.softmax(0).tolist() |
| selected = self.rng.choices(range(len(probs)), weights=probs)[0] |
| sample.update(selected=selected, old_logp=float(logits.log_softmax(0)[selected]), |
| old_value=value, weight=1., origin='synthetic') |
| if not self.teacher: |
| sample['old_log_probs'] = logits.log_softmax(0) |
| attempt = sum(event.get('kind') == 'redeal' for event in state.get('log', [])) |
| self.pending.append((sample, state['seat'], attempt)) |
| return candidates[selected] |
|
|
| def choose(self, obs, state): |
| candidates, sample = self.policy.prepare(obs, state) |
| return self.select(candidates, sample, state) |
|
|
| def choose_bid(self, state): |
| candidates, sample = self.policy.prepare_bid(state) |
| return self.select(candidates, sample, state) |
|
|
| def finish(self, game): |
| hands = [list(hand) for hand in game['initial_hands']] |
| snapshots = {} |
| step = 0 |
| for event in latest_log(game): |
| if event.get('kind') in ('play', 'pass'): |
| snapshots[step] = [list(hand) for hand in hands] |
| step += 1 |
| if event['kind'] == 'play': |
| for card in event['cards']: |
| hands[event['seat']].remove(card) |
| from .dataset import belief_labels |
| samples, play_index = [], 0 |
| final_attempt = game['redeals'] |
| for sample, seat, attempt in self.pending: |
| if attempt != final_attempt: |
| continue |
| sample['reward'] = float(game['deltas'][seat]) |
| |
| if sample['role'] != 3: |
| events = [event for event in latest_log(game) |
| if event.get('kind') in ('play', 'pass')] |
| while play_index < len(events) and events[play_index]['seat'] != seat: |
| play_index += 1 |
| if play_index < len(events): |
| belief_labels(sample, snapshots[play_index], seat, game['landlord']) |
| play_index += 1 |
| samples.append(sample) |
| self.pending.clear() |
| return samples |
|
|
|
|
| class SearchOpponent: |
| def __init__(self, seconds=.03): |
| from qojbot.policy import SearchPolicy |
| self.policy = SearchPolicy(seconds=seconds, samples=4) |
|
|
| def choose(self, obs, state): |
| return self.policy.choose(obs) |
|
|
| def choose_bid(self, state): |
| from qojbot.policy import assess_bid |
| return assess_bid(R.from_cards(state['hand']), int(state.get('bid') or 0), |
| bool(state.get('must_bid')))['value'] |
|
|
|
|
| def collect(policy, games, index, seed, rng, mode='bidding', league=(), opponent=None, |
| targeted=False): |
| """Current learner vs teacher, snapshot league, search and self-play.""" |
| all_samples, records = [], [] |
| target = opponent |
| for offset in range(games): |
| number = index + offset |
| learner = Actor(policy, rng) |
| kind = number % (8 if targeted else 6) |
| landlord = (number // 24) % 3 if targeted else number % 3 |
| relative = (number // 8) % 3 if targeted else (number // 3) % 3 |
| singleton = (landlord + relative) % 3 |
| if kind == 0 and not targeted: |
| opponent = SearchOpponent() |
| elif ((kind == 7 if targeted else kind == 2) and league): |
| from .policy import QOJPolicy |
| opponent = QOJPolicy(device=str(policy.device), teacher=policy.teacher, |
| config=policy.model.config) |
| opponent.model.load_state_dict(league[rng.randrange(len(league))]) |
| else: |
| opponent = target or policy.teacher |
| lineup = ['learner'] * 3 |
| if kind < (3 if targeted else 2): |
| lineup = ['opponent'] * 3 |
| lineup[singleton] = 'learner' |
| elif (kind in (3, 4, 5, 7) if targeted else kind < 4): |
| lineup[singleton] = 'opponent' |
| record = play_game({'learner': learner, 'opponent': opponent}, lineup, seed + number, |
| first_bidder=landlord, mode=mode) |
| all_samples.extend(learner.finish(record)) |
| records.append(record) |
| print(f'collected game {number + 1}, turns={record["turns"]}, opponent={type(opponent).__name__}', flush=True) |
| advantages = torch.tensor([s['reward'] / 6. - s['old_value'] |
| for s in all_samples]) |
| |
| for role in range(4): |
| indices = [i for i, s in enumerate(all_samples) if s['role'] == role] |
| if indices: |
| local = advantages[indices] |
| std = local.std(unbiased=False).clamp_min(.5) |
| local = (local - local.mean()) / std if len(indices) > 1 else local / std |
| for i, advantage in zip(indices, local.tolist()): |
| all_samples[i]['advantage'] = advantage |
| return all_samples, records |
|
|
|
|
| def rollout_kl(policy, samples, batch_states=32, max_actions=4096): |
| """Exact old||new KL over every legal action on the complete rollout.""" |
| total, count = 0., 0 |
| policy.model.eval() |
| with torch.no_grad(): |
| for samples_batch in batches(samples, batch_states, max_actions): |
| batch = collate(samples_batch, policy.device) |
| output = policy.model(batch) |
| counts = [len(s['actions']) for s in samples_batch] |
| local = torch.cat([torch.arange(n) for n in counts]).to(policy.device) |
| base = torch.cat([s['base'] for s in samples_batch]).to(policy.device) |
| logits = output['prior_scale'][batch['owner']] * base + output['correction'] |
| shape = (len(counts), max(counts)) |
| padded = torch.full(shape, -torch.inf, device=policy.device) |
| padded[batch['owner'], local] = logits |
| old = torch.full(shape, -torch.inf, device=policy.device) |
| old[batch['owner'], local] = torch.cat([s['old_log_probs'] for s in samples_batch]).to(policy.device) |
| valid = torch.arange(max(counts), device=policy.device)[None] < torch.tensor(counts, device=policy.device)[:, None] |
| difference = old.masked_fill(~valid, 0) - padded.log_softmax(-1).masked_fill(~valid, 0) |
| total += float((old.exp() * difference).sum().cpu()) |
| count += len(counts) |
| return max(0., total / max(count, 1)) |
|
|
|
|
| def warmup(policy, games, seed, rng): |
| actor = Actor(policy, rng, teacher=True) |
| samples, records = [], [] |
| for index in range(games): |
| record = play_game({'teacher': actor}, ['teacher'] * 3, seed + index, |
| first_bidder=index % 3, mode='bidding') |
| samples.extend(actor.finish(record)) |
| records.append(record) |
| print(f'teacher warmup {index + 1}/{games}', flush=True) |
| return samples, records |
|
|
