| """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) |
| |
| 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 |
| |
| 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']})) |
|
|
