File size: 5,054 Bytes
68eed61 | 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 80 81 82 83 84 85 86 87 88 89 | """Matched external-bank on/off evaluation, without changing training state."""
import argparse
import hashlib
import json
import os
from pathlib import Path
import subprocess
import sys
import time
def summarize(root, checkpoint):
path=root/'eval'/f'{checkpoint.stem}.json'
report=json.loads(path.read_text())
episodes=report['episodes']
traces=[]
for ep in episodes:
traces.extend(json.loads((root/'eval/traces'/checkpoint.stem/f'episode_{ep["seed"]:02d}.json').read_text()))
return dict(report=str(path),episodes=len(episodes),
successes=sum(bool(ep['success']) for ep in episodes),
success_rate=report['success_rate'],
progress={key:sum(bool(ep['progress'].get(key)) for ep in episodes)
for key in ('moved_correct_obj','moved_wrong_obj','is_src_obj_grasped','src_on_target')},
control_steps=len(traces),
gripper_closed_fraction=sum(x['command'][6]<0 for x in traces)/max(1,len(traces)))
def main():
p=argparse.ArgumentParser(description=__doc__)
p.add_argument('--checkpoint',type=Path,required=True)
p.add_argument('--output',type=Path,required=True)
p.add_argument('--episodes',type=int,default=20)
p.add_argument('--port',type=int,default=18825)
p.add_argument('--seed-start',type=int,default=0)
p.add_argument('--exec-horizon',type=int,choices=range(1,33),default=4)
p.add_argument('--after',type=Path,help='wait for a preceding matched comparison to finish')
args=p.parse_args()
if args.episodes<1 or args.seed_start<0: p.error('invalid episode count or seed')
if args.after:
print(f'MEMORY_ABLATION_WAITING_FOR {args.after}',flush=True)
while not args.after.is_file():
time.sleep(15)
previous=json.loads(args.after.read_text())
if previous.get('status')!='complete': raise ValueError('Preceding comparison is not complete')
checkpoint=args.checkpoint.resolve(strict=True)
root=args.output.resolve();root.mkdir(parents=True,exist_ok=True)
with checkpoint.open('rb') as stream:
digest=hashlib.file_digest(stream,'sha256').hexdigest()
protocol=dict(checkpoint=str(checkpoint),checkpoint_sha256=digest,
episodes=args.episodes,seed_start=args.seed_start,max_steps=150,
num_steps=30,exec_horizon=args.exec_horizon,guidance=3.0,
modes=['bank_off','on'],policy_gpu=os.environ.get('WIDOWX_EVAL_GPU'),
render_gpu=os.environ.get('WIDOWX_RENDER_GPU'),
intervention='bank_off skips external bank/head loading; global=zeros; CTE/BIT and policy weights unchanged')
manifest=root/'protocol.json'
if manifest.exists() and json.loads(manifest.read_text())!=protocol:
raise ValueError('Existing ablation protocol differs; choose a new output directory')
manifest.write_text(json.dumps(protocol,indent=2))
watcher=Path(__file__).with_name('watch_widowx_eval.py')
summary={}
for mode in protocol['modes']:
run=root/mode
report=run/'eval'/f'{checkpoint.stem}.json'
if not report.exists():
env=os.environ.copy();env['WIDOWX_RUN_DIR']=str(run)
print(f'MEMORY_ABLATION_START mode={mode}',flush=True)
subprocess.run([sys.executable,str(watcher),'--once','--checkpoint',str(checkpoint),
'--episodes',str(args.episodes),'--seed-start',str(args.seed_start),
'--port',str(args.port),'--record-traces','--max-steps','150',
'--num-steps','30','--exec-horizon',str(args.exec_horizon),'--guidance','3',
'--memory-mode',mode],env=env,check=True)
payload=json.loads(report.read_text())
if payload.get('status')!='complete' or payload.get('memory_mode')!=mode or len(payload['episodes'])!=args.episodes:
raise ValueError(f'Incomplete/wrong-mode report: {report}')
for field,expected in [('checkpoint',str(checkpoint)),('max_steps',150),('inference_num_steps',30),
('exec_horizon',args.exec_horizon),('guidance',3.0),('seed_start',args.seed_start)]:
if payload.get(field)!=expected: raise ValueError(f'Protocol mismatch: {field}')
if [ep['seed'] for ep in payload['episodes']]!=list(range(args.seed_start,args.seed_start+args.episodes)):
raise ValueError('Seed sequence differs')
summary[mode]=summarize(run,checkpoint)
print(f'MEMORY_ABLATION_RESULT {mode} '+json.dumps(summary[mode]),flush=True)
result=dict(status='complete',protocol=protocol,results=summary,
caveat='Inference-only ablation on a model trained with bank context; not equivalent to training without it.')
temp=root/'comparison.tmp';temp.write_text(json.dumps(result,indent=2));temp.replace(root/'comparison.json')
print('MEMORY_ABLATION_COMPLETE '+str(root/'comparison.json'),flush=True)
if __name__=='__main__': main()
|