File size: 3,349 Bytes
3f69d68 | 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 | import random
import unittest
import torch
from dataclasses import replace
from lightnest.learned_patch import MAX_WINDOW, PatchNetwork, make_example
from lightnest.train_patch_rl import cases_for, execution_reward, sample_actions, oracle
class PatchRLTests(unittest.TestCase):
def test_execution_reward_not_teacher_fields(self):
rng = random.Random(9)
e = make_example(rng)
while not e.present:
e = make_example(rng)
cases = cases_for(e.target_operation, rng)
self.assertEqual(execution_reward(e.binary, e.pointer, e.replacement, cases), 1)
self.assertEqual(execution_reward(e.binary, MAX_WINDOW, 0, cases), -1)
poisoned = replace(e, pointer=-777, replacement=-99, target=b'bad')
self.assertEqual(execution_reward(poisoned.binary, e.pointer, e.replacement, cases), 1)
def test_negative_requires_none(self):
rng = random.Random(9)
e = make_example(rng)
while e.present:
e = make_example(rng)
cases = cases_for(e.target_operation, rng)
self.assertEqual(execution_reward(e.binary, MAX_WINDOW, 0, cases), 1)
self.assertEqual(execution_reward(e.binary, 0, e.binary[0], cases), -1)
with self.assertRaises(ValueError):
execution_reward(e.binary, MAX_WINDOW, 0, [])
def test_sampling_is_differentiable(self):
torch.set_num_threads(2)
model = PatchNetwork()
es = [make_example(random.Random(i)) for i in range(3)]
pos, byte, logp = sample_actions(model, es)
self.assertEqual(pos.shape, (3,4))
self.assertTrue(bool(((byte >= 0) & (byte <= 255)).all()))
(-logp.mean()).backward()
self.assertIsNotNone(model.pointer.weight.grad)
self.assertIsNotNone(model.replacement.weight.grad)
def test_signed_composed_oracle(self):
self.assertEqual(oracle('add', [2147483647, 1]), -2147483648)
self.assertEqual(oracle('sub', [5, 2], 7), 10)
class PatchRLUpdateTests(unittest.TestCase):
def test_ties_skip_adam_momentum(self):
from lightnest.train_patch_rl import reinforce_step
model = torch.nn.Linear(1,1)
opt = torch.optim.Adam(model.parameters(), lr=.01)
x = torch.ones(1,4,1)
lp = model(x).squeeze(-1)
# First create a nonzero momentum state using nonidentical log-probs.
lp = lp * torch.tensor([[1.,2.,3.,4.]])
stepped, _, groups = reinforce_step(opt, model, lp, torch.tensor([[1.,-1.,1.,-1.]]))
self.assertTrue(stepped); self.assertEqual(groups,1)
before = {k:v.clone() for k,v in model.state_dict().items()}
step_before = opt.state[model.weight]['step'].clone()
stepped, loss, groups = reinforce_step(opt, model, model(x).squeeze(-1), torch.ones(1,4))
self.assertFalse(stepped); self.assertEqual(loss,0); self.assertEqual(groups,0)
self.assertTrue(all(torch.equal(v,model.state_dict()[k]) for k,v in before.items()))
self.assertTrue(torch.equal(step_before,opt.state[model.weight]['step']))
def test_train_templates_disjoint_from_eval_templates(self):
from lightnest.train_patch_rl import RL_TEMPLATES
from lightnest.learned_patch import TRAIN_TEMPLATES, TEST_TEMPLATES
self.assertFalse(set(RL_TEMPLATES) & set(TRAIN_TEMPLATES + TEST_TEMPLATES))
|