| """Export a labeled, frozen comparison set from YOR training snapshots.""" |
| import argparse |
| import copy |
| import hashlib |
| import json |
| from pathlib import Path |
| import shutil |
| import sys |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT)) |
|
|
| import torch |
| import yaml |
| from omegaconf import OmegaConf |
| from utils.export_thor_ckpt import ACTION_COLS, measured_tick_dt |
| from models.encoder.dino import DinoV2Encoder |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument('--out', required=True) |
| ap.add_argument('--runner-config', required=True) |
| args = ap.parse_args() |
| out = Path(args.out).resolve() |
| out.mkdir(parents=True, exist_ok=True) |
| torch.set_num_threads(4) |
| shutil.copytree(ROOT / 'models', out / 'runtime/models', dirs_exist_ok=True, |
| ignore=shutil.ignore_patterns('__pycache__', '*.pyc')) |
| base = yaml.safe_load(Path(args.runner_config).read_text()) |
| hz = round(1 / measured_tick_dt(Path('/home_shared/grail_charles/data/yor_episodes/df767cdb')), 3) |
| encoders = {} |
| manifest = [] |
| recipes = [('C-V3-M', 'cls_v3_norm'), ('C-V1-M', 'cls_v1_norm'), |
| ('C-V3-R', 'cls_v3_raw'), ('C-V1-R', 'cls_v1_raw'), |
| ('P-V3-M', 'patch_v3_norm'), ('P-V1-M', 'patch_v1_norm'), |
| ('P-V3-R', 'patch_v3_raw')] |
| for code, recipe in recipes: |
| paths = list((ROOT / 'outputs/yor' / code).glob('2026-09-10/*/snapshot.pt')) |
| assert len(paths) == 1, (code, paths) |
| source = paths[0] |
| |
| with source.open('rb') as f: |
| snapshot = torch.load(f, map_location='cpu', weights_only=False) |
| policy = snapshot['model'] |
| epoch = int(snapshot['epoch']) |
| cfg = OmegaConf.load(source.parent / '.hydra/config.yaml') |
| views = ['odin', 'wl', 'wr'][:int(cfg.dataset.views)] |
| horizon = int(cfg.model.obs_window_size) |
| chunk = int(cfg.model.act_window_size) |
| feature = str(cfg.encoder.feature_key) |
| patches = int(cfg.encoder.n_patches) |
| assert policy.vqvae_is_fit |
| assert policy._obs_window_size == horizon |
| assert policy._gpt_model.config.n_patches == len(views) * patches |
| if policy.normalize_actions: |
| assert policy.action_norm is not None |
| assert all(torch.isfinite(t).all() for t in policy.action_norm) |
| label = f'{recipe}_e{epoch + 1:04d}' |
| dest = out / label |
| dest.mkdir(exist_ok=True) |
| policy.to('cpu') |
| policy.eval() |
| torch.save(policy, dest / 'policy.pt') |
| if feature not in encoders: |
| enc = DinoV2Encoder(name=str(cfg.encoder.name), feature_key=feature, |
| n_patches=patches) |
| enc.to('cpu') |
| enc.eval() |
| for p in enc.parameters(): |
| p.requires_grad = False |
| encoders[feature] = enc |
| enc = encoders[feature] |
| torch.save(enc, dest / 'encoder.pt') |
| card = dict(views=views, img_size=224, obs_horizon=horizon, chunk=chunk, |
| exec_n=4, action_hz=hz, act_fields=ACTION_COLS, proprio_fields=[]) |
| (dest / 'card.yaml').write_text(yaml.safe_dump(card, sort_keys=False)) |
| runner = copy.deepcopy(base) |
| runner['policy'].update(ckpt_dir=f'~/ckpts/{out.name}/{label}', chunk=chunk, exec_n=4) |
| runner['action'].update(hz=hz, preview_time=1/hz) |
| runner['cameras']['wrist'] = len(views) > 1 |
| (dest / 'runner.yaml').write_text(yaml.safe_dump(runner, sort_keys=False)) |
| OmegaConf.save(cfg, dest / 'training_config.yaml', resolve=False) |
| with torch.no_grad(): |
| tokens = enc(torch.zeros(len(views), 3, 224, 224)) |
| assert tokens.shape == (len(views), patches, 384) |
| obs = tokens.reshape(1, 1, len(views)*patches, 384).repeat(1,horizon,1,1) |
| prediction, _, _ = policy(obs, None, None) |
| assert prediction.shape == (1,horizon,chunk,58), prediction.shape |
| assert torch.isfinite(prediction).all() |
| row = dict(label=label, recipe=recipe, source=str(source), epoch_index=epoch, |
| completed_epochs=epoch+1, views=views, obs_horizon=horizon, |
| feature_key=feature, normalized=bool(policy.normalize_actions), |
| action_hz=hz, stride=1, wandb_run_id=snapshot['wandb_run_id']) |
| (dest / 'provenance.json').write_text(json.dumps(row, indent=2)+'\n') |
| manifest.append(row) |
| print(json.dumps(row), flush=True) |
| del snapshot, policy |
| (out / 'manifest.json').write_text(json.dumps(manifest, indent=2)+'\n') |
| sums = [] |
| for path in sorted(out.rglob('*')): |
| if path.is_file() and path.name != 'SHA256SUMS' and '__pycache__' not in path.parts: |
| with path.open('rb') as f: |
| h = hashlib.sha256() |
| for block in iter(lambda: f.read(1024*1024), b''): |
| h.update(block) |
| sums.append(f'{h.hexdigest()} {path.relative_to(out)}') |
| (out / 'SHA256SUMS').write_text('\n'.join(sums)+'\n') |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
