File size: 5,753 Bytes
34c81bb 122550f 029da9d 34c81bb | 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 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | """Load the archived policy with the matching FlyControl engine (no simulator)."""
import argparse
import hashlib
import json
import os
from pathlib import Path
import sys
def load_policy(model_id, model_dir, code_dir, data_dir=None, device='cpu'):
if model_id.startswith('figure6/'):
from load_embodiment_policy import load_embodiment_policy
return load_embodiment_policy(model_id, model_dir, code_dir, device)
if model_id.startswith('figure5/'):
from load_ant_policy import load_ant_policy
policy, meta = load_ant_policy(model_id, model_dir, code_dir, device)
meta.update(obs_dim=meta['entry']['observation_dim'], action_dim=meta['entry']['action_dim'])
return policy, meta
import torch
import yaml
model_dir, code_dir = Path(model_dir).resolve(), Path(code_dir).resolve()
entries = json.loads((model_dir / 'model_index.json').read_text())['models']
entry = next(x for x in entries if x['id'] == model_id)
checkpoint = model_dir / entry['path']
hasher = hashlib.sha256()
with checkpoint.open('rb') as stream:
for chunk in iter(lambda: stream.read(8 * 1024 * 1024), b''):
hasher.update(chunk)
digest = hasher.hexdigest()
if digest != entry['sha256']:
raise ValueError('Checkpoint checksum mismatch: ' + model_id)
# These are original research checkpoints, including legacy training state.
saved = torch.load(checkpoint, map_location='cpu', weights_only=False)
state = saved['policy_state']
engine = code_dir / 'src/flycontrol/engines'
for folder in [code_dir / 'src', engine / 'common', engine / 'terrain', engine / 'supplemental']:
sys.path.insert(0, str(folder))
os.environ.setdefault('FLYCONTROL_RUN', str(Path(data_dir or '.').resolve()))
group, method = entry['paper_item'], entry['method']
if group.startswith('figure'):
if data_dir is None:
raise ValueError('Main-figure policies require the extracted dataset --data-dir')
obs_dim, action_dim = {'figure2': (688, 6), 'figure3': (64, 2), 'figure4': (51, 2)}[group]
root = Path(data_dir) / 'data' / group
graph = root / (f'er_inputs/er_seed{entry["seed"]}.npz' if method == 'er' else 'graph.npz')
if group == 'figure2':
from terrain_control.policy import build_policy
else:
from fly_control.policy import build_policy
policy = build_policy(graph, obs_dim, action_dim, 'connectome' if method == 'er' else method,
entry['seed'], device, 'normalized', -1.0)
if group in ('figure3', 'figure4') and method == 'connectome':
policy.actor.core.dt = 0.04
policy.actor.core.internal_substeps = 4
# FixedGraph's nonpersistent reduction buffers must match these edges.
if method != 'mlp':
for key in ('edge_pre', 'edge_post', 'edge_weight'):
if not torch.equal(getattr(policy.actor.core, key).cpu(), state['actor.core.' + key]):
raise ValueError('Checkpoint and graph differ: ' + key)
else:
from larva.locomotion.connectome import ConnectomeData
from larva.locomotion.models import build_policy
config = saved.get('config') or yaml.safe_load((model_dir / entry['config']).read_text())
cfg = config['policy']
inp, out = state['actor.core.input_proj.weight'], state['actor.core.output_layer.weight']
sensory, obs_dim = inp.shape
action_dim, output = out.shape
n = 2952 if method == 'larva' else 138584 if method == 'adult' else 112885
dense = state.get('actor.core.W')
edges = [state.get('actor.core.' + k) for k in ('edge_pre', 'edge_post', 'edge_weight')]
graph = ConnectomeData(None if dense is None else dense.numpy(), sensory, n-sensory-output,
output, (), (), (), *[None if x is None else x.numpy() for x in edges],
metadata={'source': 'checkpoint', 'root_ids_available': False})
policy = build_policy(cfg['name'], obs_dim, 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', -0.5), connectome=graph).to(device)
policy.load_state_dict(state, strict=True)
policy.eval()
for key, value in policy.state_dict().items():
if not torch.equal(value.cpu(), state[key]):
raise ValueError('Loaded state differs: ' + key)
return policy, {'entry': entry, 'obs_dim': obs_dim, 'action_dim': action_dim,
'obs_rms_state': saved.get('obs_rms_state')}
if __name__ == '__main__':
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('--model-id', required=True)
parser.add_argument('--model-dir', default=str(Path(__file__).parent))
parser.add_argument('--code-dir', required=True)
parser.add_argument('--data-dir')
args = parser.parse_args()
import torch
torch.set_num_threads(1)
torch.manual_seed(0)
policy, meta = load_policy(args.model_id, args.model_dir, args.code_dir, args.data_dir)
obs = torch.zeros(1, meta['obs_dim'])
hidden = policy.init_state(1, torch.device('cpu'))
with torch.no_grad():
mean, hidden, value = policy(obs, hidden)
assert all(torch.isfinite(x).all() for x in (mean, hidden, value))
print(json.dumps({'model_id': args.model_id, 'strict_load': True,
'finite_forward': True, 'observation_dim': meta['obs_dim'],
'action_dim': meta['action_dim']}))
|