| """Load a Figure 6 Cassie, Tripod, Go1 or Hexapod policy.""" |
| import argparse, hashlib, importlib.util, json, sys |
| from pathlib import Path |
|
|
| def checksum(path): |
| h=hashlib.sha256() |
| with path.open('rb') as f: |
| for b in iter(lambda:f.read(8*1024*1024),b''): h.update(b) |
| return h.hexdigest() |
|
|
| def load_embodiment_policy(model_id, model_dir, code_dir, device='cpu'): |
| import torch |
| root=Path(model_dir).resolve(); code=Path(code_dir).resolve() |
| entry=next(x for x in json.loads((root/'figure6/model_index.json').read_text())['models'] if x['id']==model_id) |
| checkpoint=root/entry['path']; graph=root/entry['graph']; config=root/entry['config'] |
| for p, digest in [(checkpoint,entry['sha256']),(graph,entry['graph_sha256']),(config,entry['config_sha256'])]: |
| if checksum(p)!=digest: raise ValueError('File checksum: '+str(p)) |
| cfg=json.loads(config.read_text()) |
| for folder in [code/'src',code/'src/flycontrol/engines/common']: |
| sys.path.insert(0,str(folder)) |
| from fly_control.policy import build_policy |
| if entry['body']=='tripod': |
| spec=importlib.util.spec_from_file_location('flycontrol_figure6_tripod_policy',code/'src/flycontrol/engines/tripod/body_control/policy.py') |
| module=importlib.util.module_from_spec(spec); spec.loader.exec_module(module) |
| build_policy=module.build_policy |
| saved=torch.load(checkpoint,map_location='cpu',weights_only=False) |
| state=saved['policy_state'] |
| policy=build_policy(graph,entry['observation_dim'],entry['action_dim'],entry['method'],entry['seed'], |
| 'cpu',readout=entry['readout'],log_std=cfg['body_config']['log_std']) |
| if entry['method']=='connectome': |
| for key in ['edge_pre','edge_post','edge_weight']: |
| if not torch.equal(getattr(policy.actor.core,key),state['actor.core.'+key]): |
| raise ValueError('Graph edges: '+key) |
| 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_dim=entry['observation_dim'],action_dim=entry['action_dim'], |
| obs_rms_state=None,body_interface=cfg['body_interface'],body_config=cfg['body_config']) |
|
|
| 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);torch.manual_seed(0) |
| policy,meta=load_embodiment_policy(a.model_id,a.model_dir,a.code_dir) |
| hidden=policy.init_state(1,torch.device('cpu')) |
| with torch.no_grad(): |
| for i in range(3): |
| action,hidden,value=policy(torch.randn(1,meta['obs_dim'])*.1,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, |
| observation_dim=meta['obs_dim'],action_dim=meta['action_dim']))) |
|
|
