code/utils/export_yor_single.py

"""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()