File size: 9,348 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 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 | """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'
# Pin the inode while training atomically replaces later snapshots.
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()
|