File size: 2,855 Bytes
029da9d | 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 | """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)))
|