File size: 3,283 Bytes
fe31b61 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | """No privileged outcomes in model construction; charged, selected-bit replay."""
from __future__ import annotations
import time
import torch
from unionmind.decision_sufficiency.core import CONTEXTS, count_samples, encode, estimate, plan, push
class Timer:
def __init__(self, device): self.device = device
def __enter__(self):
torch.cuda.synchronize(self.device)
self.start = time.perf_counter()
return self
def __exit__(self, *args):
torch.cuda.synchronize(self.device)
self.seconds = time.perf_counter()-self.start
def construct(samples, model_name, method):
"""Only labeled fitting samples enter the model/representation/planner."""
counts = count_samples(samples[None], 6)
model = estimate(counts, model_name)
code, k = encode(model, method)
compact = model if method == 'identity' else push(model, code, k)
return model,code.to(torch.uint8),compact,plan(compact)
def select_samples(archive, regime, seed):
generator = torch.Generator(device=archive['source'].device).manual_seed(seed)
source = archive['source'][torch.randperm(len(archive['source']),device=archive['source'].device,
generator=generator)[:1024]]
adaptation = archive['adaptation'][torch.randperm(len(archive['adaptation']),
device=archive['source'].device,generator=generator)]
if regime == 'source_only': return source
if regime == 'target_64': return adaptation[:64]
if regime == 'target_256':
return None if len(adaptation) < 256 else adaptation[:256]
assert regime == 'source_plus_target_64'
return torch.cat((source,adaptation[:64]))
def replay(policy, code, current, future, budget, cost_index):
"""Return final action, paid query count, and query-mask for each record.
`future` is a retained offline payload. Only the bit selected by a paid
query is indexed. It is never passed to a model or decision table lookup.
"""
assert policy['shape'][0] == 1
r = policy['shape'][1]
table = policy['choices'].reshape(3,1,r,len(policy['costs']),9)[:,0]
z = code[0,current].long()
context = torch.zeros_like(current)
action = torch.zeros_like(current)
count = torch.zeros_like(current)
seen = torch.zeros_like(current)
live = torch.ones_like(current,dtype=torch.bool)
child = torch.zeros((9,2,2),device=current.device,dtype=torch.long)
for t,c in enumerate(CONTEXTS):
for q in (0,1):
if c[q] == -1:
for o in (0,1):
child[t,q,o] = CONTEXTS.index(tuple(o if i == q else v for i,v in enumerate(c)))
for left in range(budget,-1,-1):
choice = table[left,z,cost_index,context]
querying = live & (choice >= 3)
q = (choice-3).clamp(0,1)
selected = querying.nonzero().flatten()
revealed = torch.zeros_like(current)
revealed[selected] = (future[selected] >> (1-q[selected])) & 1
context = torch.where(querying,child[context,q,revealed],context)
action = torch.where(live & ~querying,choice,action)
count += querying.long()
seen |= torch.where(querying,1 << q,0)
live = querying
assert not bool(live.any())
return action,count,seen
|