| """Persistent, unsliced YOR embeddings shared across training recipes.""" |
| import fcntl |
| import hashlib |
| import inspect |
| import json |
| import logging |
| import os |
| from pathlib import Path |
| import uuid |
|
|
| import cv2 |
| import hydra |
| from omegaconf import OmegaConf |
| import torch |
| from tqdm import tqdm |
|
|
| from datasets.core import TrajectoryDataset |
| from utils.inference import _embed_in_chunks, eval_mode |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| LOG = logging.getLogger(__name__) |
| SCHEMA = 1 |
| LOADERS = {'datasets.yor.YorDataset', 'datasets.yor_jittered.YorJitteredDataset'} |
|
|
|
|
| def source_manifest(dataset_config): |
| """Metadata fingerprint: paths, sizes and nanosecond mtimes, not video reads.""" |
| config = OmegaConf.to_container(dataset_config, resolve=True) |
| target = config['_target_'] |
| if target not in LOADERS: |
| raise ValueError(f'Embedding cache does not support {target}') |
| root = Path(config['data_directory']).expanduser().resolve() |
| if target == 'datasets.yor.YorDataset': |
| files = [root/'actions.pth', root/'seq_lengths.pkl', root/'meta.json'] |
| files.extend(sorted((root/'obses').glob('episode_*.pth'))) |
| else: |
| files = [] |
| for episode in sorted(root.glob('ep_*')): |
| if not episode.is_dir(): |
| continue |
| files.extend(episode/name for name in ['robot_data.csv', 'provenance.json']) |
| for view in ['odin', 'wrist_left', 'wrist_right'][:int(config.get('views', 3))]: |
| files.extend([episode/(view+'.mp4'), episode/(view+'_index.csv')]) |
| if not files: |
| raise ValueError(f'No dataset files under {root}') |
| entries = [] |
| for path in files: |
| stat = path.stat() |
| entries.append([str(path.relative_to(root)), stat.st_size, stat.st_mtime_ns]) |
| |
| config.pop('prefetch', None) |
| config['data_directory'] = str(root) |
| return dict(dataset=config, files=entries) |
|
|
|
|
| def cache_identity(dataset_config, encoder_config, encoder): |
| model = getattr(encoder, 'module', encoder) |
| weight_hash = hashlib.sha256() |
| for name, value in sorted(model.state_dict().items()): |
| value = value.detach().cpu().contiguous() |
| weight_hash.update(name.encode()) |
| weight_hash.update(str((tuple(value.shape), value.dtype)).encode()) |
| weight_hash.update(value.reshape(-1).view(torch.uint8).numpy().tobytes()) |
| source_files = [Path(__file__), ROOT/'datasets/yor.py', ROOT/'datasets/yor_jittered.py', |
| ROOT/'utils/convert_yor.py', ROOT/'utils/inference.py', |
| Path(inspect.getfile(type(model)))] |
| identity = dict(schema=SCHEMA, source=source_manifest(dataset_config), |
| encoder=OmegaConf.to_container(encoder_config, resolve=True), |
| weights_sha256=weight_hash.hexdigest(), torch_version=str(torch.__version__), |
| opencv_version=cv2.__version__, |
| code_sha256={str(p.relative_to(ROOT)) if p.is_relative_to(ROOT) else str(p): |
| hashlib.sha256(p.read_bytes()).hexdigest() for p in source_files}) |
| key = hashlib.sha256(json.dumps(identity, sort_keys=True).encode()).hexdigest() |
| return key, identity |
|
|
|
|
| class CachedYorDataset(TrajectoryDataset): |
| """Whole episodes in CPU RAM; splitting and temporal sampling happen later.""" |
| def __init__(self, payload, cache_path): |
| self.data = payload['episodes'] |
| self.seq_lengths = payload['seq_lengths'] |
| self.episode_names = payload['episode_names'] |
| self.source_episodes = payload['source_episodes'] |
| self.cache_path = str(cache_path) |
| if not self.data or len(self.data) != len(self.seq_lengths) or len(self.data) != len(self.episode_names): |
| raise ValueError('Invalid cached episode metadata') |
| if self.source_episodes is not None and len(self.source_episodes) != len(self.data): |
| raise ValueError('Invalid cached source groups') |
| for length, (obs, actions, goal) in zip(self.seq_lengths, self.data): |
| if length <= 0 or obs.ndim != 4 or actions.shape != (length, 58) or obs.shape[0] != length or goal.shape != (length, 1, 1, 1): |
| raise ValueError('Invalid cached trajectory shapes') |
|
|
| def __len__(self): |
| return len(self.data) |
|
|
| def get_seq_length(self, idx): |
| return self.seq_lengths[idx] |
|
|
| def get_frames(self, idx, frames): |
| indices = list(frames) |
| return tuple(value[indices] for value in self.data[idx]) |
|
|
| def __getitem__(self, idx): |
| return self.data[idx] |
|
|
| def get_all_actions(self): |
| return torch.cat([episode[1] for episode in self.data]) |
|
|
|
|
| def load_or_build_yor_cache(dataset_config, encoder_config, encoder, cache_root, force=False): |
| """Lock builders, publish atomically, and skip raw-loader construction on hits.""" |
| key, identity = cache_identity(dataset_config, encoder_config, encoder) |
| folder = Path(cache_root).expanduser()/Path(identity['source']['dataset']['data_directory']).name |
| folder.mkdir(parents=True, exist_ok=True) |
| path = folder/(key+'.pt') |
| LOG.info('YOR embedding cache: %s', path) |
| with (folder/(key+'.lock')).open('a') as lock: |
| fcntl.flock(lock, fcntl.LOCK_EX) |
| if path.exists() and not force: |
| try: |
| payload = torch.load(path, map_location='cpu', weights_only=True) |
| if payload['identity'] != identity: |
| raise ValueError('Cache fingerprint mismatch') |
| cached = CachedYorDataset(payload, path) |
| except Exception as exc: |
| raise RuntimeError(f'Cannot load {path}; rebuild it with --force') from exc |
| LOG.info('YOR embedding cache HIT: %s episodes, %.1f MiB; loaded into RAM', |
| len(cached), path.stat().st_size/2**20) |
| return cached |
| LOG.info('YOR embedding cache MISS: building all frames before splitting/sampling') |
| dataset = hydra.utils.instantiate(dataset_config) |
| model = getattr(encoder, 'module', encoder) |
| device = next(model.parameters()).device |
| episodes = [] |
| with eval_mode(model, no_grad=True): |
| for i in tqdm(range(len(dataset)), desc='Caching YOR embeddings'): |
| obs, actions, goal = dataset[i] |
| embeddings = _embed_in_chunks(model, obs.to(device)).detach().cpu().contiguous() |
| if not torch.isfinite(embeddings).all(): |
| raise ValueError(f'Nonfinite embeddings in episode {i}') |
| episodes.append((embeddings, actions.detach().cpu().clone(), goal.detach().cpu().clone())) |
| |
| if source_manifest(dataset_config) != identity['source']: |
| raise RuntimeError('Dataset changed during cache construction; retry with a stable dataset') |
| if hasattr(dataset, 'episodes'): |
| names = [p.name for p in dataset.episodes] |
| else: |
| metadata = json.loads((Path(dataset.data_directory)/'meta.json').read_text()) |
| names = metadata.get('episodes', [f'episode_{i:05d}' for i in range(len(dataset))]) |
| payload = dict(identity=identity, episodes=episodes, |
| seq_lengths=[dataset.get_seq_length(i) for i in range(len(dataset))], |
| episode_names=names, source_episodes=getattr(dataset, 'source_episodes', names)) |
| cached = CachedYorDataset(payload, path) |
| temporary = folder/(key+'.'+uuid.uuid4().hex+'.tmp') |
| try: |
| with temporary.open('wb') as stream: |
| torch.save(payload, stream) |
| stream.flush() |
| os.fsync(stream.fileno()) |
| os.replace(temporary, path) |
| finally: |
| temporary.unlink(missing_ok=True) |
| (folder/(key+'.json')).write_text(json.dumps(dict(identity=identity, |
| episode_names=names, source_episodes=payload['source_episodes'], |
| seq_lengths=payload['seq_lengths'], bytes=path.stat().st_size), indent=2)+'\n') |
| LOG.info('YOR embedding cache SAVED: %s (%.1f MiB)', path, path.stat().st_size/2**20) |
| return cached |
|
|
