"""Phase annotate: CLI script to add all phase/affordance columns to parquet files in a single pass. Produces 8 columns: - phase.subgoal, phase.progress (from gripper segmentation) - affordance.type, affordance.granularity (rule-derived from above) - phase.subgoal_object, phase.subgoal_target, phase.subgoal_descriptor, phase.subgoal_relation (structured fields resolved from parsed instructions) """ from __future__ import annotations import argparse import json import logging import os import shutil from collections import Counter, defaultdict from multiprocessing import Pool from pathlib import Path from typing import Dict, List, Tuple import numpy as np import pandas as pd from phaseaware_data_construction.gripper_segmenter import SegmentationResult, segment_episode from phaseaware_data_construction.instruction_parser import Subgoal, parse_instruction from phaseaware_data_construction.affordance_deriver import derive_affordance_columns from phaseaware_data_construction.subgoal_columns import build_label_to_fields, resolve_frame_fields logging.basicConfig( level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s" ) logger = logging.getLogger(__name__) # Gripper dimension indices (from modality.json) GRIPPER_STATE_DIM = 7 # observation.state[:, 7] GRIPPER_ACTION_DIM = 6 # action[:, 6] def load_task_map(src_dir: Path) -> Dict[int, str]: """Load task_index → instruction mapping from meta/tasks.jsonl.""" tasks_path = src_dir / "meta" / "tasks.jsonl" task_map: Dict[int, str] = {} with open(tasks_path) as f: for line in f: d = json.loads(line.strip()) task_map[d["task_index"]] = d["task"] return task_map def parse_all_tasks(task_map: Dict[int, str]) -> Dict[int, List[Subgoal]]: """Pre-parse all task instructions into subgoal sequences.""" parsed: Dict[int, List[Subgoal]] = {} for idx, instruction in task_map.items(): parsed[idx] = parse_instruction(instruction) return parsed def find_episode_parquets(src_dir: Path) -> List[Path]: """Find all episode parquet files under data/.""" data_dir = src_dir / "data" parquets = sorted(data_dir.rglob("episode_*.parquet")) return parquets def process_single_episode(args: Tuple) -> dict: """Process a single episode parquet. Designed for multiprocessing.""" parquet_path, subgoal_map, signal_source, threshold, median_filter_size = args parquet_path = Path(parquet_path) try: df = pd.read_parquet(parquet_path) state = np.array(df["observation.state"].tolist()) task_idx = int(df["task_index"].iloc[0]) episode_idx = int(df["episode_index"].iloc[0]) # Extract gripper signal based on source if signal_source == "action": action = np.array(df["action"].tolist()) gripper_signal = action[:, GRIPPER_ACTION_DIM] else: gripper_signal = state[:, GRIPPER_STATE_DIM] subgoals = subgoal_map.get(task_idx) if subgoals is None: return { "path": str(parquet_path), "episode": episode_idx, "status": "error", "error": f"Unknown task_index {task_idx}", } result = segment_episode( gripper_signal, subgoals, signal_source=signal_source, threshold=threshold, median_filter_size=median_filter_size, ) # Derive affordance columns (rule-based from subgoal + progress) aff_types, aff_grans = derive_affordance_columns( result.frame_subgoals, result.frame_progress ) # Derive structured subgoal fields (object, target, descriptor, relation) label_map = build_label_to_fields(subgoals) sg_objects, sg_targets, sg_descriptors, sg_relations = resolve_frame_fields( result.frame_subgoals, label_map ) return { "path": str(parquet_path), "episode": episode_idx, "task_index": task_idx, "status": "ok", "n_frames": len(df), "n_cycles_detected": result.n_cycles_detected, "n_cycles_expected": result.n_cycles_expected, "warnings": result.warnings, "frame_subgoals": result.frame_subgoals, "frame_progress": result.frame_progress, "aff_types": aff_types, "aff_grans": aff_grans, "sg_objects": sg_objects, "sg_targets": sg_targets, "sg_descriptors": sg_descriptors, "sg_relations": sg_relations, } except Exception as e: return { "path": str(parquet_path), "episode": -1, "status": "error", "error": str(e), } def write_annotated_parquet( src_path: Path, output_path: Path, result: dict ): """Read src parquet, add all 8 phase/affordance columns, write to output.""" df = pd.read_parquet(src_path) df["phase.subgoal"] = result["frame_subgoals"] df["phase.progress"] = result["frame_progress"] df["affordance.type"] = result["aff_types"] df["affordance.granularity"] = result["aff_grans"] df["phase.subgoal_object"] = result["sg_objects"] df["phase.subgoal_target"] = result["sg_targets"] df["phase.subgoal_descriptor"] = result["sg_descriptors"] df["phase.subgoal_relation"] = result["sg_relations"] output_path.parent.mkdir(parents=True, exist_ok=True) df.to_parquet(output_path, index=False) def validate_output(output_dir: Path, subgoal_map: Dict[int, List[Subgoal]]): """Validate the annotated output parquets.""" parquets = find_episode_parquets(output_dir) if not parquets: logger.error("No parquets found in output directory") return total_frames = 0 total_episodes = 0 suspicious = [] order_mismatches = [] sg_frame_counts: Counter = Counter() progress_counts: Counter = Counter() type_counts: Counter = Counter() gran_counts: Counter = Counter() sg_per_episode: defaultdict = defaultdict(list) for p in parquets: df = pd.read_parquet(p) total_episodes += 1 total_frames += len(df) if "phase.subgoal" not in df.columns or "phase.progress" not in df.columns: logger.error(f"Missing phase columns in {p}") continue # Check all 8 columns exist all_cols = [ "phase.subgoal", "phase.progress", "affordance.type", "affordance.granularity", "phase.subgoal_object", "phase.subgoal_target", "phase.subgoal_descriptor", "phase.subgoal_relation", ] missing_cols = [c for c in all_cols if c not in df.columns] if missing_cols: logger.warning(f"{p}: missing columns {missing_cols}") null_sg = df["phase.subgoal"].isna().sum() null_pr = df["phase.progress"].isna().sum() if null_sg > 0 or null_pr > 0: logger.warning(f"{p}: {null_sg} null subgoals, {null_pr} null progress") # Count subgoal types and check order episode_sgs = df["phase.subgoal"].tolist() episode_pr = df["phase.progress"].tolist() # Extract unique subgoals in order seen_order = [] prev = None for sg in episode_sgs: if sg != prev: seen_order.append(sg) prev = sg ep_idx = int(df["episode_index"].iloc[0]) task_idx = int(df["task_index"].iloc[0]) sg_per_episode[task_idx].append(len(seen_order)) # --- Semantic order check against parser expectation --- if task_idx in subgoal_map: expected_labels = [sg.label() for sg in subgoal_map[task_idx]] # Check 1: every expected subgoal should appear in seen_order missing = [lbl for lbl in expected_labels if lbl not in seen_order] # Check 2: relative order should be preserved (allow extra labels but # expected labels must appear in the same order) ei = 0 out_of_order = False for lbl in seen_order: if ei < len(expected_labels) and lbl == expected_labels[ei]: ei += 1 if ei < len(expected_labels): out_of_order = True if missing or out_of_order: order_mismatches.append(( p.name, ep_idx, task_idx, expected_labels, seen_order, missing, out_of_order, )) for sg in episode_sgs: sg_frame_counts[sg] += 1 for pr in episode_pr: progress_counts[pr] += 1 if "affordance.type" in df.columns: for t in df["affordance.type"]: type_counts[t] += 1 if "affordance.granularity" in df.columns: for g in df["affordance.granularity"]: gran_counts[g] += 1 # Check for suspicious episodes sg_counts_ep = Counter(episode_sgs) n = len(df) for sg, cnt in sg_counts_ep.items(): if cnt < 5: suspicious.append((p.name, ep_idx, sg, cnt, "too_short")) if cnt > 0.8 * n and len(sg_counts_ep) > 1: suspicious.append((p.name, ep_idx, sg, cnt, "too_long")) # Print summary logger.info(f"\n{'='*60}") logger.info(f"Validation Summary") logger.info(f"{'='*60}") logger.info(f"Total episodes: {total_episodes}") logger.info(f"Total frames: {total_frames}") logger.info(f"\nProgress distribution:") for pr, cnt in progress_counts.most_common(): logger.info(f" {pr}: {cnt} ({cnt/total_frames*100:.1f}%)") logger.info(f"\nSubgoal type frame counts (top 20):") for sg, cnt in sg_frame_counts.most_common(20): logger.info(f" {sg}: {cnt}") if type_counts: logger.info(f"\naffordance.type distribution:") for t, cnt in type_counts.most_common(): logger.info(f" {t}: {cnt} ({cnt/total_frames*100:.1f}%)") if gran_counts: logger.info(f"\naffordance.granularity distribution:") for g, cnt in gran_counts.most_common(): logger.info(f" {g}: {cnt} ({cnt/total_frames*100:.1f}%)") # Semantic order check results if order_mismatches: logger.warning(f"\nSubgoal order mismatches ({len(order_mismatches)}):") for name, ep, tidx, expected, seen, missing, ooo in order_mismatches: logger.warning(f" {name} ep={ep} task={tidx}:") logger.warning(f" expected : {expected}") logger.warning(f" seen : {seen}") if missing: logger.warning(f" missing : {missing}") if ooo: logger.warning(f" out-of-order: expected labels not in correct relative order") else: logger.info(f"\nAll episodes match expected subgoal order.") if suspicious: logger.warning(f"\nSuspicious episodes ({len(suspicious)}):") for name, ep, sg, cnt, reason in suspicious[:20]: logger.warning(f" {name} ep={ep}: {sg} has {cnt} frames ({reason})") else: logger.info(f"\nNo suspicious episodes found.") def main(): parser = argparse.ArgumentParser( description="Add all phase/affordance columns (8 total) to LeRobot parquet files." ) parser.add_argument("--src_dir", type=str, required=True, help="Source dataset directory") parser.add_argument("--output_dir", type=str, required=True, help="Output directory") parser.add_argument("--num_workers", type=int, default=1, help="Number of parallel workers") parser.add_argument("--signal_source", type=str, default="action", choices=["action", "state"], help="Gripper signal source: 'action' (binary, more reliable) or 'state' (continuous)") parser.add_argument("--gripper_threshold", type=float, default=0.02, help="Gripper binarization threshold (state dim 7: > -threshold → closed)") parser.add_argument("--median_filter_size", type=int, default=5, help="Median filter kernel size") parser.add_argument("--dry_run", action="store_true", help="Don't write files, only print stats") parser.add_argument("--validate", action="store_true", help="Validate existing output instead of annotating") args = parser.parse_args() src_dir = Path(args.src_dir) output_dir = Path(args.output_dir) # --- Validate mode --- if args.validate: task_map = load_task_map(src_dir) subgoal_map = parse_all_tasks(task_map) validate_output(output_dir, subgoal_map) return # --- Annotate mode --- logger.info(f"Source: {src_dir}") logger.info(f"Output: {output_dir}") # Load and parse tasks task_map = load_task_map(src_dir) subgoal_map = parse_all_tasks(task_map) logger.info(f"Parsed {len(task_map)} tasks:") for idx, sgs in subgoal_map.items(): logger.info(f" task {idx}: {task_map[idx]} → {[sg.label() for sg in sgs]}") # Find episodes parquets = find_episode_parquets(src_dir) logger.info(f"Found {len(parquets)} episode parquets") if not parquets: logger.error("No episode parquets found!") return # Prepare arguments for workers work_args = [ (str(p), subgoal_map, args.signal_source, args.gripper_threshold, args.median_filter_size) for p in parquets ] # Process if args.num_workers > 1: with Pool(args.num_workers) as pool: results = pool.map(process_single_episode, work_args) else: results = [process_single_episode(a) for a in work_args] # Summarize and write ok_count = 0 err_count = 0 warn_count = 0 cycle_match = 0 cycle_mismatch = 0 for r in results: if r["status"] == "error": err_count += 1 logger.error(f"Episode {r['episode']}: {r.get('error', 'unknown error')}") continue ok_count += 1 if r["warnings"]: warn_count += 1 for w in r["warnings"]: logger.warning(f"Episode {r['episode']}: {w}") if r["n_cycles_detected"] == r["n_cycles_expected"]: cycle_match += 1 else: cycle_mismatch += 1 logger.info(f"\n{'='*60}") logger.info(f"Processing Summary") logger.info(f"{'='*60}") logger.info(f"OK: {ok_count}, Errors: {err_count}, With warnings: {warn_count}") logger.info(f"Cycle match: {cycle_match}, Cycle mismatch: {cycle_mismatch}") logger.info(f"Match rate: {cycle_match/(ok_count or 1)*100:.1f}%") if args.dry_run: logger.info("Dry run — no files written.") return # --- Write output --- logger.info(f"\nWriting annotated parquets to {output_dir}") # Copy meta directory src_meta = src_dir / "meta" dst_meta = output_dir / "meta" if src_meta.exists(): if dst_meta.exists(): shutil.rmtree(dst_meta) shutil.copytree(src_meta, dst_meta) logger.info(f"Copied meta/ directory") # Symlink video directory src_video = src_dir / "video" dst_video = output_dir / "video" if src_video.exists() and not dst_video.exists(): dst_video.parent.mkdir(parents=True, exist_ok=True) os.symlink(src_video.resolve(), dst_video) logger.info(f"Symlinked video/ directory") # Write annotated parquets written = 0 for r in results: if r["status"] != "ok": continue src_path = Path(r["path"]) rel_path = src_path.relative_to(src_dir) out_path = output_dir / rel_path write_annotated_parquet(src_path, out_path, r) written += 1 logger.info(f"Written {written} annotated parquets") if __name__ == "__main__": main()