File size: 5,178 Bytes
94f697a | 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 108 109 110 111 112 113 | """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]
# Training replaces snapshots atomically: this open pins one complete version.
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()
|