scripts/classify_markovian_gr1.py

"""
Classify GR1 humanoid episodes as 'markovian' or 'non_markovian' using
Qwen3-VL-30B-A3B-Instruct-FP8 served via vLLM.

Reads prompts from the two batch_input.json files (DreamDojo-HV_Eval and
EgoDex_Eval), classifies all ~133 episodes, and outputs:
  - scripts/gr1_task_labels.csv          -- all episodes with label + reason
  - scripts/gr1_episode_sampled.csv      -- balanced sample (equal per label
                                            AND equal across the two sources)

Usage:
    # 1. Start the vLLM server first:
    #    bash scripts/launch_qwen3vl_server.sh

    # 2. Run (text-only):
    python scripts/classify_markovian_gr1.py

    # Run with vision (sends input PNG image to model):
    python scripts/classify_markovian_gr1.py --use-vision --batch-size 5

    # Optional flags:
    python scripts/classify_markovian_gr1.py \
        --dreamdojo-json  sampling_dataset/humanoid/singleview/input/PhysicalAI-Robotics-GR00T-Teleop-GR1/DreamDojo-HV_Eval/batch_input.json \
        --egodex-json     sampling_dataset/humanoid/singleview/input/PhysicalAI-Robotics-GR00T-Teleop-GR1/EgoDex_Eval/batch_input.json \
        --cache-file      scripts/gr1_task_labels_cache.jsonl \
        --output-csv      scripts/gr1_task_labels.csv \
        --sampled-csv     scripts/gr1_episode_sampled.csv \
        --sample-per-label-per-source  33 \
        --base-url        http://localhost:8000/v1 \
        --model           Qwen/Qwen3-VL-30B-A3B-Instruct-FP8 \
        --batch-size      5 \
        --max-image-px    560 \
        --max-retries     2 \
        --seed            42
"""

import argparse
import base64
import io
import json
import sys
import time
from pathlib import Path

import requests
import pandas as pd
from PIL import Image, ImageFile
from tqdm import tqdm

ImageFile.LOAD_TRUNCATED_IMAGES = True

# ---------------------------------------------------------------------------
# Prompt
# ---------------------------------------------------------------------------

SYSTEM_PROMPT = """You are a robot manipulation task classifier.

You will be given one or more episodes. Each episode has:
- An image showing the robot's initial state (if provided)
- A task description prompt

Classify each task as exactly one of:
- "markovian": a SINGLE atomic manipulation action. The robot only needs to observe the current state to act. No memory of previous steps is required.
  Typical examples: "picks up a banana", "places a cup on the table", "pushes a ball", "removes a block from the stack", "opens a drawer"

- "non_markovian": involves MULTIPLE sequential steps, OR requires the robot to remember what it has already done, OR involves ongoing/continuous actions that track progress.
  Typical examples: "picks up X then places it into Y while holding Z", "stirs … repeatedly", "topples a line of tiles", "uses both arms to lift and rotate", "sprays water onto plants"

Rules:
- Connectors like "then", "while", "after", "followed by" → non_markovian
- Ongoing / repetitive actions (stir, spray, wipe, fold, rotate) → non_markovian
- Both arms doing different simultaneous things → non_markovian
- Simple single pick, place, push, remove, open/close → markovian
- A pick-and-place is markovian (one action) even with two locations
- Use the image to resolve ambiguity (e.g. verify object count, arm configuration)

Reply with ONLY a valid JSON array (no markdown fences, no extra text):
[{"id": "<id>", "label": "markovian"|"non_markovian", "reason": "<≤12 words>"}]"""

USER_TEMPLATE = "Classify these tasks:\n{tasks_json}"

# ---------------------------------------------------------------------------
# Vision helpers
# ---------------------------------------------------------------------------


def encode_image_b64(image_path: Path, max_px: int = 560) -> str:
    """Load image, resize so longest side <= max_px, return base64 PNG string."""
    img = Image.open(image_path).convert("RGB")
    w, h = img.size
    if max(w, h) > max_px:
        scale = max_px / max(w, h)
        img = img.resize((int(w * scale), int(h * scale)), Image.Resampling.LANCZOS)
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return base64.b64encode(buf.getvalue()).decode("utf-8")


def build_vision_user_content(batch: list[dict], max_px: int) -> list[dict]:
    """Build multimodal content list: interleave images with per-episode text."""
    content = []
    for i, row in enumerate(batch):
        img_path = Path(row["input_image"])
        img_ok = False
        if img_path.exists():
            try:
                b64 = encode_image_b64(img_path, max_px)
                content.append({
                    "type": "image_url",
                    "image_url": {"url": f"data:image/png;base64,{b64}"},
                })
                img_ok = True
            except Exception as e:
                print(f"  [warn] Could not load image {img_path.name}: {e}", file=sys.stderr)
        suffix = "" if img_ok else " [image unavailable]"
        content.append({
            "type": "text",
            "text": f'Episode {i+1} (id="{row["id"]}"): {row["prompt"]}{suffix}',
        })
    content.append({
        "type": "text",
        "text": "\nNow classify all episodes above. Reply with only the JSON array.",
    })
    return content

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def load_batch_json(path: Path, source: str) -> list[dict]:
    data = json.loads(path.read_text())
    rows = []
    for item in data:
        out_video = item.get("output_video", "")
        episode_name = Path(out_video).stem  # e.g. episode_000000
        rows.append(
            {
                "id": f"{source}/{episode_name}",
                "source": source,
                "episode": episode_name,
                "input_image": item.get("input_video", ""),   # it's a .png despite key name
                "output_video": out_video,
                "prompt": item.get("prompt", "").strip(),
            }
        )
    return rows


def load_cache(cache_file: Path) -> dict[str, dict]:
    cache: dict[str, dict] = {}
    if not cache_file.exists():
        return cache
    with open(cache_file) as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            obj = json.loads(line)
            cache[obj["id"]] = obj
    return cache


def append_to_cache(cache_file: Path, results: list[dict]) -> None:
    with open(cache_file, "a") as f:
        for r in results:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")


def call_llm(
    base_url: str,
    model: str,
    batch: list[dict],
    max_retries: int,
    use_vision: bool = False,
    max_image_px: int = 560,
) -> list[dict]:
    if use_vision:
        user_content = build_vision_user_content(batch, max_image_px)
    else:
        payload_tasks = [{"id": r["id"], "task": r["prompt"]} for r in batch]
        user_content = USER_TEMPLATE.format(tasks_json=json.dumps(payload_tasks, ensure_ascii=False))

    payload = {
        "model": model,
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_content},
        ],
        "temperature": 0.0,
        "chat_template_kwargs": {"enable_thinking": False},
    }

    for attempt in range(1, max_retries + 2):
        try:
            resp = requests.post(
                f"{base_url}/chat/completions",
                json=payload,
                timeout=120,
            )
            resp.raise_for_status()
            raw = resp.json()["choices"][0]["message"]["content"].strip()
            if raw.startswith("```"):
                raw = raw.split("```")[1]
                if raw.startswith("json"):
                    raw = raw[4:]
                raw = raw.strip()

            parsed: list[dict] = json.loads(raw)
            id_map = {r["id"]: r for r in batch}
            results = []
            for item in parsed:
                rid = item.get("id", "")
                label = item.get("label", "").strip().lower()
                if label not in ("markovian", "non_markovian"):
                    label = "parse_error"
                src_row = id_map.get(rid, {})
                results.append(
                    {
                        "id": rid,
                        "source": src_row.get("source", ""),
                        "episode": src_row.get("episode", ""),
                        "input_image": src_row.get("input_image", ""),
                        "output_video": src_row.get("output_video", ""),
                        "prompt": src_row.get("prompt", ""),
                        "label": label,
                        "reason": item.get("reason", ""),
                    }
                )
            return results

        except (json.JSONDecodeError, KeyError, TypeError) as e:
            if attempt <= max_retries:
                time.sleep(2)
            else:
                print(f"  [warn] Parse failed: {e}", file=sys.stderr)
                return [{**r, "label": "parse_error", "reason": str(e)} for r in batch]
        except requests.RequestException as e:
            if attempt <= max_retries:
                time.sleep(5)
            else:
                print(f"  [error] HTTP error: {e}", file=sys.stderr)
                return [{**r, "label": "api_error", "reason": str(e)} for r in batch]

    return []


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------


def main() -> None:
    parser = argparse.ArgumentParser(description="Classify GR1 humanoid tasks")
    parser.add_argument(
        "--dreamdojo-json",
        type=Path,
        default=Path(
            "sampling_dataset/humanoid/singleview/input"
            "/PhysicalAI-Robotics-GR00T-Teleop-GR1/DreamDojo-HV_Eval/batch_input.json"
        ),
    )
    parser.add_argument(
        "--egodex-json",
        type=Path,
        default=Path(
            "sampling_dataset/humanoid/singleview/input"
            "/PhysicalAI-Robotics-GR00T-Teleop-GR1/EgoDex_Eval/batch_input.json"
        ),
    )
    parser.add_argument(
        "--cache-file",
        type=Path,
        default=Path("scripts/gr1_task_labels_cache.jsonl"),
    )
    parser.add_argument(
        "--output-csv",
        type=Path,
        default=Path("scripts/gr1_task_labels.csv"),
    )
    parser.add_argument(
        "--sampled-csv",
        type=Path,
        default=Path("scripts/gr1_episode_sampled.csv"),
    )
    parser.add_argument(
        "--sample-per-label-per-source",
        type=int,
        default=33,
        help=(
            "Episodes per (label × source) cell. "
            "e.g. 33 → up to 33 markovian from DreamDojo + 33 from EgoDex, "
            "same for non_markovian → 132 total max. 0 = keep all."
        ),
    )
    parser.add_argument("--base-url", type=str, default="http://localhost:8000/v1")
    parser.add_argument(
        "--model", type=str, default="Qwen/Qwen3-VL-30B-A3B-Instruct-FP8"
    )
    parser.add_argument("--batch-size", type=int, default=50)
    parser.add_argument("--max-retries", type=int, default=2)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument(
        "--use-vision",
        action="store_true",
        default=False,
        help="Send input PNG image to the model alongside the prompt (recommended batch-size 5)",
    )
    parser.add_argument(
        "--max-image-px",
        type=int,
        default=560,
        help="Resize images so longest side <= this value before encoding (saves tokens)",
    )
    args = parser.parse_args()

    workspace = Path(__file__).parent.parent

    def resolve(p: Path) -> Path:
        return workspace / p if not p.is_absolute() else p

    dreamdojo_json = resolve(args.dreamdojo_json)
    egodex_json = resolve(args.egodex_json)
    cache_file = resolve(args.cache_file)
    output_csv = resolve(args.output_csv)
    sampled_csv = resolve(args.sampled_csv)
    cache_file.parent.mkdir(parents=True, exist_ok=True)
    output_csv.parent.mkdir(parents=True, exist_ok=True)

    # Load both sources
    all_rows: list[dict] = []
    all_rows += load_batch_json(dreamdojo_json, "DreamDojo-HV_Eval")
    all_rows += load_batch_json(egodex_json, "EgoDex_Eval")
    print(f"[init] Loaded {len(all_rows)} episodes total")

    # Resume from cache
    cache = load_cache(cache_file)
    print(f"[init] Already classified: {len(cache)}")

    pending = [r for r in all_rows if r["id"] not in cache]
    print(f"[init] Remaining to classify: {len(pending)}")

    if pending:
        batches = [
            pending[i : i + args.batch_size]
            for i in range(0, len(pending), args.batch_size)
        ]
        vision_info = f", vision=ON (max_px={args.max_image_px})" if args.use_vision else ", vision=OFF"
        print(f"[run]  {len(batches)} batch(es), model={args.model}, endpoint={args.base_url}{vision_info}")

        for batch in tqdm(batches, desc="Classifying", unit="batch"):
            results = call_llm(
                args.base_url, args.model, batch, args.max_retries,
                use_vision=args.use_vision, max_image_px=args.max_image_px,
            )
            append_to_cache(cache_file, results)
            for r in results:
                cache[r["id"]] = r

    # Build full CSV
    rows = list(cache.values())
    df = pd.DataFrame(rows)
    # Preserve original ordering (DreamDojo first, then EgoDex, sorted by episode)
    df["_order"] = df["id"].map({r["id"]: i for i, r in enumerate(all_rows)})
    df = df.sort_values("_order").drop(columns=["_order"]).reset_index(drop=True)
    df.to_csv(output_csv, index=False)
    print(f"\n[saved] Full labels → {output_csv}  ({len(df)} rows)")
    print(df.groupby(["source", "label"]).size().to_string())

    # Sample equally per (label × source)
    if args.sample_per_label_per_source > 0:
        parts = []
        for (source, label), group in df.groupby(["source", "label"]):
            if label in ("parse_error", "api_error", "unknown"):
                continue
            n = min(args.sample_per_label_per_source, len(group))
            parts.append(group.sample(n=n, random_state=args.seed))
            print(f"[sample] {source} / {label}: {n}/{len(group)}")

        sampled = (
            pd.concat(parts)
            .sort_values(["source", "episode"])
            .reset_index(drop=True)
        )
        sampled.to_csv(sampled_csv, index=False)
        print(f"\n[saved] Sampled → {sampled_csv}  ({len(sampled)} rows)")
        print(sampled.groupby(["source", "label"]).size().to_string())
    else:
        df.to_csv(sampled_csv, index=False)
        print(f"\n[saved] Sampled (all) → {sampled_csv}")


if __name__ == "__main__":
    main()