"""Read-only, CPU audit of the current full-window WidowX data contract.""" import argparse from collections import Counter from concurrent.futures import ThreadPoolExecutor import json from pathlib import Path import numpy as np import pyarrow.parquet as pq def audit(root, workers): info = json.loads((root / 'meta/info.json').read_text()) episodes = [json.loads(line) for line in (root / 'meta/episodes.jsonl').open()] tasks = {row['task_index']: row['task'] for row in map(json.loads, (root / 'meta/tasks.jsonl').open())} counts = Counter() starts_histogram = Counter() task_windows = Counter() lag_sum = {lag: np.zeros(6, np.float64) for lag in (-2, -1, 0, 1, 2)} lag_count = Counter() per_offset_closed = np.zeros(32, np.int64) target_tasks = { task: {'episodes': 0, 'retained_episodes': 0, 'windows': 0, 'closed_proprio_windows': 0, 'starts': Counter(), 'closed_per_offset': np.zeros(32, np.int64)} for task in ('put carrot on plate', 'take carrot off plate') } q01 = np.array([-0.028752606511116028, -0.041702210046350954, -0.026096698231995105, -0.08052822157740593, -0.09249736212193965, -0.20738411962985992]) q99 = np.array([0.02830690816044803, 0.04089860741049051, 0.04018005654215808, 0.08173405691981314, 0.07760896608233431, 0.20384809583425495]) clipped = np.zeros(6, np.int64) invalid = [] def read(ep): idx = ep['episode_index'] path = root / info['data_path'].format(episode_chunk=idx // info['chunks_size'], episode_index=idx) rows = pq.read_table(path, columns=['observation.state', 'action', 'task_index'], use_threads=False).to_pydict() return ep, np.asarray(rows['observation.state']), np.asarray(rows['action']), np.asarray(rows['task_index']) with ThreadPoolExecutor(max_workers=workers) as pool: def bounded_rows(): for offset in range(0, len(episodes), workers * 16): yield from pool.map(read, episodes[offset:offset + workers * 16]) for number, (ep, state, action, task_ids) in enumerate(bounded_rows(), 1): idx, length = ep['episode_index'], ep['length'] if len(action) != length or state.shape != (length, 8) or action.shape != (length, 7): invalid.append({'episode': idx, 'reason': 'shape_or_length'}) continue if not np.isfinite(state).all() or not np.isfinite(action).all(): invalid.append({'episode': idx, 'reason': 'nonfinite'}) continue starts = np.arange(0, max(0, length - 32), 4) closed = action[:, 6] <= .5 empty = np.array([not tasks[int(t)].strip() for t in task_ids]) task_name = tasks[int(task_ids[0])] if task_name in target_tasks: target = target_tasks[task_name] target['episodes'] += 1 target['retained_episodes'] += int(len(starts) > 0) target['windows'] += len(starts) target['closed_proprio_windows'] += int((state[starts, 7] <= .5).sum()) target['starts'].update(map(int, starts)) for start in starts: target['closed_per_offset'] += closed[start:start + 32] counts['episodes'] += 1 counts['frames'] += length counts['closed_action_frames'] += int(closed.sum()) counts['empty_instruction_frames'] += int(empty.sum()) counts['episodes_with_empty_instruction'] += int(empty.any()) counts['episodes_with_multiple_task_ids'] += int(len(set(task_ids.tolist())) > 1) counts['nonbinary_gripper_action_frames'] += int(((action[:, 6] != 0) & (action[:, 6] != 1)).sum()) counts['nonzero_state_pad_frames'] += int((state[:, 6] != 0).sum()) clipped += ((action[:, :6] < q01) | (action[:, :6] > q99)).sum(0) if len(starts): counts['retained_episodes'] += 1 counts['windows'] += len(starts) counts['empty_instruction_windows'] += int(empty[starts].sum()) counts['closed_proprio_windows'] += int((state[starts, 7] <= .5).sum()) counts['windows_at_or_after_halfway'] += int((starts >= (length - 1) / 2).sum()) counts['windows_with_completed_effect_history'] += int((starts >= 16).sum()) counts['windows_with_full_four_effect_history'] += int((starts >= 64).sum()) starts_histogram.update(map(int, starts)) task_windows.update(tasks[int(task_ids[s])] for s in starts) for s in starts: per_offset_closed += closed[s:s + 32] first4_ids = (starts[:, None] + np.arange(4)).ravel() counts['first4_closed_action_presentations'] += int(closed[first4_ids].sum()) counts['first4_action_presentations'] += len(first4_ids) covered_end = int(starts[-1]) + 32 counts['unique_action_frames_supervised'] += covered_end counts['first4_unique_action_frames'] += len(first4_ids) if 'carrot' in tasks[int(task_ids[0])].lower() and 'plate' in tasks[int(task_ids[0])].lower(): counts['carrot_plate_retained_episodes'] += 1 counts['carrot_plate_windows'] += len(starts) if 'carrot' in tasks[int(task_ids[0])].lower() and 'plate' in tasks[int(task_ids[0])].lower(): counts['carrot_plate_all_episodes'] += 1 delta = np.diff(state[:, :6], axis=0) for lag in lag_sum: t = np.arange(len(delta)) valid = (t + lag >= 0) & (t + lag < length) error = np.abs(action[t[valid] + lag, :6] - delta[t[valid]]) lag_sum[lag] += error.sum(0) lag_count[lag] += len(error) if number % 5000 == 0: print(json.dumps({'scanned_episodes': number, 'total': len(episodes)}), flush=True) for target in target_tasks.values(): target['starts'] = dict(sorted(target['starts'].items())) target['closed_action_fraction_by_prediction_offset'] = ( target.pop('closed_per_offset') / max(1, target['windows']) ).tolist() return dict( root=str(root), scope='all_episode_parquets_no_video_decode', counts=dict(counts), starts_histogram=dict(sorted(starts_histogram.items())), closed_action_fraction_by_prediction_offset=(per_offset_closed / counts['windows']).tolist(), action_clipped_fraction_by_dim=(clipped / counts['frames']).tolist(), action_vs_state_delta_mae_by_action_lag={str(k): (v / lag_count[k]).tolist() for k, v in lag_sum.items()}, target_tasks=target_tasks, largest_window_tasks=task_windows.most_common(20), invalid=invalid, ) if __name__ == '__main__': parser = argparse.ArgumentParser(description=__doc__) parser.add_argument('--root', type=Path, required=True) parser.add_argument('--output', type=Path, required=True) parser.add_argument('--workers', type=int, default=4) args = parser.parse_args() report = audit(args.root, args.workers) args.output.parent.mkdir(parents=True, exist_ok=True) args.output.write_text(json.dumps(report, indent=2, ensure_ascii=False)) print(json.dumps(report, indent=2, ensure_ascii=False), flush=True)