| """Load a Figure 5 Ant policy and its saved observation normalization.""" |
| import argparse,hashlib,json,sys |
| from pathlib import Path |
|
|
| def load_ant_policy(model_id,model_dir,code_dir,device='cpu'): |
| import torch,yaml |
| root=Path(model_dir).resolve();code=Path(code_dir).resolve() |
| entry=next(x for x in json.loads((root/'figure5/model_index.json').read_text())['models'] if x['id']==model_id) |
| p=root/entry['path'];h=hashlib.sha256() |
| with p.open('rb') as f: |
| for b in iter(lambda:f.read(8*1024*1024),b''):h.update(b) |
| if h.hexdigest()!=entry['sha256']:raise ValueError('Checkpoint checksum: '+model_id) |
| for folder in [code/'src',code/'src/flycontrol/engines/common',code/'src/flycontrol/engines/ant']: |
| sys.path.insert(0,str(folder)) |
| from body_control.policy import build_policy |
| from fly_control.graph import ConnectomeData |
| ck=torch.load(p,map_location='cpu',weights_only=False);state=ck['policy_state'] |
| cfg=yaml.safe_load((root/entry['config']).read_text())['policy'] |
| graph=None |
| if entry['method']=='adult': |
| prefix='actor.base.core.' if 'actor.base.core.edge_pre' in state else 'actor.core.' |
| sensory=state[prefix+'input_proj.weight'].shape[0] |
| output=state[prefix+'output_layer.weight'].shape[1] |
| graph=ConnectomeData(None,sensory,entry['graph_nodes']-sensory-output,output,(),(),(), |
| **{k:state[prefix+k].numpy() for k in ['edge_pre','edge_post','edge_weight']}) |
| policy=build_policy(cfg['name'],entry['observation_dim'],entry['action_dim'], |
| actor_kwargs=cfg.get('actor_kwargs'),value_hidden=cfg.get('value_hidden'), |
| value_activation=cfg.get('value_activation','tanh'),log_std_init=cfg.get('log_std_init',-.5),connectome=graph) |
| router=getattr(policy.actor,'configure_goal_router_normalizer',None) |
| if router:router(ck['obs_rms_state']['mean'],ck['obs_rms_state']['var']) |
| policy.load_state_dict(state,strict=True) |
| for k,v in policy.state_dict().items(): |
| if not torch.equal(v,state[k]):raise ValueError('Loaded state: '+k) |
| policy.to(device).eval() |
| return policy,dict(entry=entry,obs_rms_state=ck['obs_rms_state']) |
|
|
| if __name__=='__main__': |
| p=argparse.ArgumentParser(description=__doc__);p.add_argument('--model-id',required=True) |
| p.add_argument('--model-dir',default=str(Path(__file__).parent));p.add_argument('--code-dir',required=True) |
| a=p.parse_args() |
| import torch |
| torch.set_num_threads(1) |
| policy,meta=load_ant_policy(a.model_id,a.model_dir,a.code_dir) |
| hidden=policy.init_state(1,torch.device('cpu')) |
| with torch.no_grad(): |
| action,hidden,value=policy(torch.zeros(1,meta['entry']['observation_dim']),hidden) |
| assert all(torch.isfinite(x).all() for x in [action,hidden,value]) |
| print(json.dumps(dict(model_id=a.model_id,strict_load=True,finite_forward=True))) |
|
|
