| """Package frozen recent YOR snapshots for Thor, retaining temporal semantics.""" |
| import argparse |
| import hashlib |
| import json |
| import os |
| 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 |
|
|
|
|
| def sha256_stream(stream): |
| digest = hashlib.sha256() |
| for block in iter(lambda: stream.read(8*1024*1024), b''): |
| digest.update(block) |
| return digest.hexdigest() |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument('--selection', required=True) |
| parser.add_argument('--out', required=True) |
| parser.add_argument('--encoder-bundle', required=True) |
| args = parser.parse_args() |
| out = Path(args.out).resolve() |
| out.mkdir(parents=True, exist_ok=False) |
| base = Path(args.encoder_bundle).resolve() |
| torch.set_num_threads(4) |
| for repo in sorted(Path(torch.hub.get_dir()).glob('facebookresearch_dinov2*')): |
| if (repo/'dinov2').is_dir(): |
| sys.path.insert(0, str(repo)) |
| break |
| encoder = torch.load(base/'encoder.pt', map_location='cpu', weights_only=False) |
| encoder.eval() |
| base_cfg = OmegaConf.load(base/'training_config.yaml') |
| base_card = yaml.safe_load((base/'card.yaml').read_text()) |
| source_tick_hz = float(base_card['action_hz']) |
| (out/'common').mkdir() |
| shutil.copy2(base/'encoder.pt', out/'common/encoder.pt') |
| shutil.copytree(ROOT/'models', out/'runtime/models', |
| ignore=shutil.ignore_patterns('__pycache__', '*.pyc')) |
| manifest = [] |
| with torch.no_grad(): |
| tokens = encoder(torch.zeros(3, 3, 224, 224)) |
| assert tokens.shape == (3, 1, 384) |
| for item in json.loads(Path(args.selection).read_text()): |
| run = Path(item['run']).resolve() |
| cfg = OmegaConf.load(run/'.hydra/config.yaml') |
| for key in ['name', 'feature_key', 'n_patches', 'output_dim']: |
| assert cfg.encoder[key] == base_cfg.encoder[key], key |
| assert cfg.dataset.views == 3 |
| source = run/'snapshot.pt' |
| |
| with source.open('rb') as stream: |
| source_stat = os.fstat(stream.fileno()) |
| source_hash = sha256_stream(stream) |
| stream.seek(0) |
| snapshot = torch.load(stream, map_location='cpu', weights_only=False) |
| completed = int(snapshot['epoch'])+1 |
| policy = snapshot['model'] |
| policy.to('cpu') |
| policy.eval() |
| history, chunk = int(cfg.model.obs_window_size), int(cfg.model.act_window_size) |
| assert policy.vqvae_is_fit and policy.normalize_actions |
| assert (policy._obs_window_size, policy._act_window_size, policy._act_dim) == (history, chunk, 58) |
| assert policy.action_norm is not None and all(torch.isfinite(v).all() for v in policy.action_norm) |
| assert (policy.action_norm[1] > 0).all() |
| temporal = OmegaConf.to_container(cfg.data.get('temporal', {}), resolve=True) |
| skip = int(temporal.get('frame_skip', 1)) |
| online_filter = temporal.get('motion_filter', {}).get('enabled', False) |
| prefiltered = bool(cfg.data.get('source_prefiltered', |
| cfg.dataset['_target_'] == 'datasets.yor_jittered.YorJitteredDataset')) |
| variable = bool(online_filter or prefiltered) |
| offset = int(cfg.data.get('action_offset', 0)) |
| label = f"{item['label']}_e{completed:04d}" |
| bundle = out/label |
| bundle.mkdir() |
| torch.save(policy, bundle/'policy.pt') |
| os.link(out/'common/encoder.pt', bundle/'encoder.pt') |
| frame_hz = float(item.get('source_frame_hz', source_tick_hz)) |
| tick_hz = float(item.get('source_tick_hz', source_tick_hz)) |
| card = dict(views=['odin', 'wl', 'wr'], img_size=224, obs_horizon=history, |
| chunk=chunk, exec_n=4, act_fields=ACTION_COLS, proprio_fields=[], |
| action_offset=offset, frame_skip=skip, normalize_actions=True, |
| prediction_units='raw_command_units', source_tick_hz=tick_hz, |
| source_frame_hz=frame_hz, |
| timing_mode='variable_waypoints' if variable else 'fixed_rate', |
| training_action_hz=None if variable else frame_hz/skip, |
| deployment_status='offline_test_bundle' if variable else 'smoke_tested_bundle') |
| if not variable: |
| card['action_hz'] = frame_hz/skip |
| (bundle/'card.yaml').write_text(yaml.safe_dump(card, sort_keys=False)) |
| OmegaConf.save(cfg, bundle/'training_config.yaml', resolve=False) |
| obs = tokens.reshape(1, 1, 3, 384).repeat(1, history, 1, 1) |
| with torch.no_grad(): |
| torch.manual_seed(7) |
| expected, _, _ = policy(obs, None, None) |
| restored = torch.load(bundle/'policy.pt', map_location='cpu', weights_only=False) |
| restored.eval() |
| torch.manual_seed(7) |
| actual, _, _ = restored(obs, None, None) |
| assert actual.shape == (1, history, chunk, 58) and torch.isfinite(actual).all() |
| torch.testing.assert_close(actual, expected, rtol=0, atol=0) |
| for before, after in zip(policy.action_norm, restored.action_norm): |
| torch.testing.assert_close(before, after, rtol=0, atol=0) |
| row = dict(label=label, recipe_name=item['recipe'], source=str(source), |
| source_snapshot_sha256=source_hash, source_bytes=source_stat.st_size, |
| source_mtime_ns=source_stat.st_mtime_ns, completed_epochs=completed, |
| planned_epochs=int(cfg.epochs), training_complete=completed >= int(cfg.epochs), |
| wandb_run_id=snapshot.get('wandb_run_id'), batch_size=int(cfg.batch_size), |
| observation_history=history, action_chunk=chunk, action_offset=offset, |
| frame_skip=skip, motion_filter=temporal.get('motion_filter'), |
| dataset_prefiltered=prefiltered, timing_mode=card['timing_mode'], |
| normalized=True, outputs_raw_units=True, smoke_shape=list(actual.shape)) |
| (bundle/'provenance.json').write_text(json.dumps(row, indent=2)+'\n') |
| timing = ('Targets are variable-time waypoints. There is no fixed training action Hz. ' |
| 'The existing fixed-rate Thor runner needs an explicit waypoint execution/history strategy ' |
| 'before live use; do not assign a 30 Hz training rate to this model.' if variable else |
| f'Targets/history are sampled every {skip} source frames, approximately {frame_hz/skip:.3f} Hz.') |
| (bundle/'README.md').write_text( |
| f'# {label}\n\nCompleted {completed}/{cfg.epochs} epochs. Batch {cfg.batch_size}. ' |
| f'V3 CLS; normalized internally; 58D outputs are already in raw command units.\n\n' |
| f'History {history}; chunk {chunk}; target offset {offset}; frame skip {skip}. ' |
| f'obs[t] predicts action rows starting at t+{offset} in the sampled sequence.\n\n' |
| f'{timing}\n\nRun the package `smoke_all.py --only {label}` on Thor to test this policy ' |
| 'without cameras or robot RPC. This export does not change the robot runner.\n') |
| manifest.append(row) |
| print(json.dumps(row), flush=True) |
| del snapshot, policy, restored |
| (out/'manifest.json').write_text(json.dumps(manifest, indent=2)+'\n') |
| shutil.copy2(ROOT/'tools/cls_debug/smoke_thor_bundle.py', out/'smoke_all.py') |
| lines = [f'# Recent YOR CLS policies\n', |
| 'Each directory contains policy.pt, encoder.pt, card.yaml, training_config.yaml and provenance.json. ' |
| 'Encoder files share storage through hard links. Policy normalization is internal; outputs are raw.\n', |
| '| Bundle | Epochs | Batch | Chunk | Skip | Offset | Timing |', |
| '|---|---:|---:|---:|---:|---:|---|'] |
| for row in manifest: |
| lines.append(f"| {row['label']} | {row['completed_epochs']}/{row['planned_epochs']} | {row['batch_size']} | {row['action_chunk']} | {row['frame_skip']} | {row['action_offset']} | {row['timing_mode']} |") |
| lines.extend(['', 'Offline smoke on Thor (no cameras or robot RPC):', '', '```bash', |
| f'cd ~/ckpts/{out.name}', |
| '/home/thor2/miniconda3/envs/charles/bin/python smoke_all.py --device cuda', '```', '', |
| 'Use `--only <bundle-name>` to test one policy. Variable-waypoint cards intentionally omit ' |
| '`action_hz`: the existing fixed-rate live runner cannot infer their physical timing from a chunk. ' |
| 'These models need explicit waypoint timing/history handling before live execution.', '', |
| 'The unfiltered chunk-20 baseline is already installed at ' |
| '`~/ckpts/yor_cls_chunk20_20260912/cls_v3_norm_chunk20_unfiltered_e1000`.']) |
| (out/'README.md').write_text('\n'.join(lines)+'\n') |
| sums = [] |
| for path in sorted(out.rglob('*')): |
| if path.is_file(): |
| with path.open('rb') as stream: |
| sums.append(sha256_stream(stream)+' '+str(path.relative_to(out))) |
| (out/'SHA256SUMS').write_text('\n'.join(sums)+'\n') |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
