| """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) |
|
|
