File size: 2,989 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 | import random
import unittest
from unittest.mock import patch
import torch
from lightnest import rv32
from lightnest.learned_patch import MAX_WINDOW
from lightnest.sequential_patch_rl import arithmetic, make_task, passes, rollout, terminal_reward
class FixedPolicy(torch.nn.Module):
def __init__(self, position):
super().__init__(); self.position=position; self.calls=0
self.weight=torch.nn.Parameter(torch.tensor(1.))
def forward(self,c,b):
self.calls+=1
scores=torch.full((len(c),257),-30.)+self.weight*0
scores[:,self.position]=self.weight*30
return scores,torch.ones((len(c),256,8))*self.weight
class SequentialTests(unittest.TestCase):
def test_oracle(self):
self.assertEqual(arithmetic((4,2,1),'sub'),1)
self.assertEqual(arithmetic((2147483647,1,1),'add'),-2147483647)
def test_generator_bounds_and_independent_execution(self):
for split in ('train','unseen_bias','three_op_combined'):
for seed in range(40):
t=make_task(random.Random(seed),split)
self.assertLessEqual(len(t.binary),MAX_WINDOW)
# Test-only assembler repair never enters policy or training reward.
target='sub' if 'to sub' in t.command or 'with sub' in t.command else 'add'
words=[int.from_bytes(t.binary[i:i+4],'little') for i in range(0,len(t.binary),4)]
words=[((w & ~(127<<25)) | ((32 if target=='sub' else 0)<<25)) if w&127==51 else w for w in words]
repaired=b''.join(w.to_bytes(4,'little') for w in words)
self.assertTrue(passes(repaired,t.cases))
def test_none_and_budget(self):
t=make_task(random.Random(3)); p=FixedPolicy(MAX_WINDOW)
out,edits,h,_=rollout(p,[t.command],[t.binary])
self.assertEqual(out,[t.binary]); self.assertEqual(edits,[0]); self.assertEqual(p.calls,1)
p=FixedPolicy(0)
with patch('lightnest.sequential_patch_rl.rv32.simulate',side_effect=AssertionError('oracle reached policy')):
_,_,h,lp=rollout(p,[t.command],[t.binary],sample=True)
self.assertEqual(p.calls,3); self.assertEqual(len(h[0]),3)
lp.sum().backward(); self.assertIsNotNone(p.weight.grad)
def test_input_bounds(self):
p=FixedPolicy(0)
for n in (0,4):
with self.assertRaises(ValueError): rollout(p,['x'],[b'1234'],max_proposals=n)
with self.assertRaises(ValueError): rollout(p,['x'],[b'x'*257])
def test_terminal_reward(self):
binary=b''.join(w.to_bytes(4,'little') for w in (rv32.addi('x10','x0',7),rv32.EBREAK))
cases=[([1,2],7)]
self.assertEqual(terminal_reward(binary,binary,cases,0),1)
self.assertEqual(terminal_reward(binary,binary,cases,1),-1)
self.assertEqual(terminal_reward(binary,binary,[([1,2],8)],0),-1)
with self.assertRaises(ValueError): terminal_reward(binary,binary,[],0)
if __name__=='__main__': unittest.main()
|