| 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) |
| |
| 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)) |
|
|
