| """Export one frozen YOR snapshot, deriving horizons from its saved config.""" |
| import argparse,copy,hashlib,json,shutil,sys |
| from pathlib import Path |
| ROOT=Path(__file__).resolve().parents[1];sys.path.insert(0,str(ROOT)) |
| import torch,yaml |
| from omegaconf import OmegaConf |
| from utils.export_thor_ckpt import ACTION_COLS |
|
|
| def main(): |
| ap=argparse.ArgumentParser();ap.add_argument('--run',required=True);ap.add_argument('--out',required=True);ap.add_argument('--encoder-bundle',required=True);args=ap.parse_args() |
| run=Path(args.run).resolve();out=Path(args.out).resolve();base=Path(args.encoder_bundle).resolve() |
| if out.exists():raise FileExistsError(out) |
| cfg=OmegaConf.load(run/'.hydra/config.yaml');torch.set_num_threads(4) |
| snap=torch.load(run/'snapshot.pt',map_location='cpu',weights_only=False);policy=snap['model'];policy.to('cpu');policy.eval() |
| temporal=OmegaConf.to_container(cfg.data.temporal,resolve=True) if cfg.data.get('temporal') else dict(frame_skip=1,motion_filter=dict(enabled=False)) |
| if temporal.get('motion_filter',{}).get('enabled'): |
| raise ValueError('Filtered waypoint models require explicit runtime timing/history support before deployment export') |
| if temporal.get('frame_skip',1)!=1:raise ValueError('Frame-skipped models require matching runtime history sampling') |
| assert policy.vqvae_is_fit and policy.normalize_actions==cfg.model.normalize_actions |
| assert policy._act_window_size==cfg.model.act_window_size and policy._obs_window_size==cfg.model.obs_window_size |
| if policy.normalize_actions:assert policy.action_norm is not None and all(torch.isfinite(t).all() for t in policy.action_norm) |
| oldcfg=OmegaConf.load(base/'training_config.yaml') |
| for key in ['name','feature_key','n_patches','output_dim']:assert cfg.encoder[key]==oldcfg.encoder[key],key |
| hub=Path(torch.hub.get_dir()) |
| for repo in sorted(hub.glob('facebookresearch_dinov2*')): |
| if (repo/'dinov2').is_dir():sys.path.insert(0,str(repo));break |
| enc=torch.load(base/'encoder.pt',map_location='cpu',weights_only=False);enc.eval() |
| views=['odin','wl','wr'][:int(cfg.dataset.views)];history=int(cfg.model.obs_window_size);chunk=int(cfg.model.act_window_size) |
| with torch.no_grad(): |
| tok=enc(torch.zeros(len(views),3,224,224));assert tok.shape==(len(views),1,384) |
| obs=tok.reshape(1,1,len(views),384).repeat(1,history,1,1) |
| pred,_,_=policy(obs,None,None) |
| assert pred.shape==(1,history,chunk,58) and torch.isfinite(pred).all() |
| card=yaml.safe_load((base/'card.yaml').read_text());card.update(views=views,obs_horizon=history,chunk=chunk,exec_n=4,act_fields=ACTION_COLS) |
| runner=yaml.safe_load((base/'runner.yaml').read_text());runner['policy'].update(ckpt_dir='~/ckpts/'+out.parent.name+'/'+out.name,chunk=chunk,exec_n=4) |
| out.mkdir(parents=True);shutil.copytree(ROOT/'models',out.parent/'runtime/models',dirs_exist_ok=True,ignore=shutil.ignore_patterns('__pycache__','*.pyc')) |
| torch.save(policy,out/'policy.pt');shutil.copy2(base/'encoder.pt',out/'encoder.pt') |
| (out/'card.yaml').write_text(yaml.safe_dump(card,sort_keys=False));(out/'runner.yaml').write_text(yaml.safe_dump(runner,sort_keys=False));OmegaConf.save(cfg,out/'training_config.yaml') |
| prov=dict(source=str(run/'snapshot.pt'),completed_epochs=int(snap['epoch'])+1,wandb_run_id=snap.get('wandb_run_id'),normalized=bool(policy.normalize_actions),temporal=temporal,smoke_shape=list(pred.shape)) |
| (out/'provenance.json').write_text(json.dumps(prov,indent=2)+'\n') |
| (out/'README.md').write_text(f'''# {out.name}\n\nCompleted epochs: {prov['completed_epochs']}. V3 DINOv2 CLS, normalized actions\n(internal normalization; outputs are raw commands). Observation history {history},\naction chunk {chunk}, execute 4 before replanning. Original ~30 Hz, unfiltered.\n\nThor existing TUI:\n\n```bash\ncd ~/YOR_charles\n/home/thor2/miniconda3/envs/charles/bin/python tools/policy_infer_vqbet.py \\\n --ckpt ~/ckpts/{out.parent.name}/{out.name} --mode policy --dry-run\n```\n\nUse the existing b/w base/WB controls. This export does not change arm/hand\ndispatch settings. Run the deployment smoke/replay checks before live use.\nSource: {run}\n''') |
| sums=[] |
| for p in sorted(out.parent.rglob('*')): |
| if p.is_file() and p.name!='SHA256SUMS':sums.append(hashlib.sha256(p.read_bytes()).hexdigest()+' '+str(p.relative_to(out.parent))) |
| (out.parent/'SHA256SUMS').write_text('\n'.join(sums)+'\n') |
| print(json.dumps(prov),flush=True) |
| if __name__=='__main__':main() |
|
|
