dataset_build.py

"""
RoboMind VLA — Task 3: build fine-tuning dataset from rendered rollouts.

Reads .mp4 videos + metadata.jsonl from the Modal volume robomind-data,
extracts keyframes, derives ground-truth judgment labels, and pushes a
HuggingFace dataset suitable for MiniCPM-V fine-tuning.

Run it:
    modal run dataset_build.py

Done when:
    {HF_USERNAME}/robomind-loco-judge-dataset exists on the Hub with
    6 PIL images per row and a valid 10-key JSON in target_json.
"""

from __future__ import annotations

import modal

HF_REPO = "mitvho09/robomind-loco-judge-dataset"

image = (
    modal.Image.debian_slim(python_version="3.11")
    .apt_install("ffmpeg")
    .pip_install(
        "datasets",
        "huggingface_hub",
        "decord",
        "imageio",
        "imageio-ffmpeg",
        "pillow",
        "numpy",
    )
)

app = modal.App("robomind-vla-dataset")
volume = modal.Volume.from_name("robomind-data", create_if_missing=True)
ROLLUTS_DIR = "/data/rollouts"

INSTRUCTION_PROMPT = (
    "You are RoboMind VLA, a vision-language reward model for humanoid "
    "locomotion. You are shown keyframes from a robot locomotion rollout. "
    "The robot was commanded to \"walk forward\". Analyze the rollout and "
    "respond with ONLY a JSON object with these exact keys: timestep_range, "
    "phase, command, command_followed, stability, fall_risk, gait_quality, "
    "predicted_reward, anomaly, explanation."
)


def derive_judgment(env, tier, ret, num_steps, norm, fell):
    early_fall = fell and num_steps < 200

    gait_quality = round(max(0.0, norm - (0.2 if fell else 0.0)), 2)
    predicted_reward = round(norm, 2)
    command_followed = norm > 0.25

    if not fell:
        stability = "stable" if norm > 0.6 else "wobbling"
        fall_risk = "low" if norm > 0.6 else "medium"
        phase = "walking"
        anomaly = None if norm > 0.4 else "inefficient, unstable gait"
    elif early_fall:
        stability, fall_risk, phase = "falling", "critical", "terminal"
        anomaly = f"early fall at step {num_steps}"
    else:
        stability, fall_risk, phase = "stumbling", "high", "terminal"
        anomaly = f"fall at step {num_steps}"

    explanation = (
        f"The {env} executes a {tier}-tier forward-locomotion rollout over "
        f"{num_steps} steps, accumulating reward {ret:.0f} (normalized "
        f"{norm:.2f}). "
        f"Gait stability is '{stability}' with {fall_risk} fall risk"
        + (f"; {anomaly}." if anomaly else "; no anomalies detected.")
    )

    return {
        "timestep_range": [0, int(num_steps)],
        "phase": phase,
        "command": "walk forward",
        "command_followed": bool(command_followed),
        "stability": stability,
        "fall_risk": fall_risk,
        "gait_quality": gait_quality,
        "predicted_reward": predicted_reward,
        "anomaly": anomaly,
        "explanation": explanation,
    }


@app.function(
    image=image,
    volumes={"/data": volume},
    secrets=[modal.Secret.from_name("huggingface-secret")],
    timeout=3600,
)
def build_dataset(hf_username: str = "mitvho09") -> dict:
    import json
    import os

    import decord
    import numpy as np
    from datasets import Dataset, Features, Image as HFImage, Value
    from huggingface_hub import login
    from PIL import Image

    hf_token = os.environ.get("HF_TOKEN")
    if hf_token:
        login(token=hf_token)

    decord.bridge.set_bridge("native")

    # --- 1. Read metadata --------------------------------------------------
    meta_path = os.path.join(ROLLUTS_DIR, "metadata.jsonl")
    records = []
    seen = set()
    with open(meta_path) as f:
        for line in f:
            r = json.loads(line.strip())
            key = (r["env"], r["tier"], r["episode_id"])
            if key not in seen:
                seen.add(key)
                records.append(r)

    print(f"[build] loaded {len(records)} unique episodes from metadata.jsonl")

    # --- 2. Per-env return stats (pass 1) ----------------------------------
    env_returns: dict[str, list[float]] = {}
    for r in records:
        env_returns.setdefault(r["env"], []).append(r["return"])

    env_stats = {}
    for env, rets in env_returns.items():
        env_stats[env] = (min(rets), max(rets))
        print(f"[build] env={env}  min_return={env_stats[env][0]:.1f}  "
              f"max_return={env_stats[env][1]:.1f}")

    # --- 3. Build samples (pass 2) -----------------------------------------
    rows = []
    for r in records:
        vid_path = os.path.join(ROLLUTS_DIR, r["video"])
        if not os.path.exists(vid_path):
            print(f"[build] WARNING: {r['video']} not found, skipping")
            continue

        vr = decord.VideoReader(vid_path)
        n_frames = len(vr)
        indices = np.linspace(0, n_frames - 1, 6, dtype=int)
        frames = vr.get_batch(indices).asnumpy()

        pil_images = []
        for f_arr in frames:
            img = Image.fromarray(f_arr)
            w, h = img.size
            scale = 448 / max(w, h)
            if scale < 1.0:
                img = img.resize((int(w * scale), int(h * scale)),
                                 Image.LANCZOS)
            pil_images.append(img.convert("RGB"))

        min_ret, max_ret = env_stats[r["env"]]
        norm = (r["return"] - min_ret) / (max_ret - min_ret + 1e-8)
        norm = float(np.clip(norm, 0.0, 1.0))

        judgment = derive_judgment(
            r["env"], r["tier"], r["return"], r["num_steps"],
            norm, r["fell"],
        )

        rows.append({
            "images": pil_images,
            "prompt": INSTRUCTION_PROMPT,
            "target_json": json.dumps(judgment),
            "env": r["env"],
            "tier": r["tier"],
            "gt_return": r["return"],
            "gt_norm_reward": norm,
            "fell": r["fell"],
            "episode_id": r["episode_id"],
        })

    print(f"[build] built {len(rows)} samples")

    # --- 4. Build HF dataset -----------------------------------------------
    features = Features({
        "images": [HFImage()],
        "prompt": Value("string"),
        "target_json": Value("string"),
        "env": Value("string"),
        "tier": Value("string"),
        "gt_return": Value("float32"),
        "gt_norm_reward": Value("float32"),
        "fell": Value("bool"),
        "episode_id": Value("int32"),
    })

    ds = Dataset.from_list(rows, features=features)

    # --- 5. Push to Hub ----------------------------------------------------
    repo_id = f"{hf_username}/robomind-loco-judge-dataset"
    print(f"[build] pushing to {repo_id}...")
    ds.push_to_hub(repo_id, private=False)
    print(f"[build] DONE: https://huggingface.co/datasets/{repo_id}")

    return {"repo_id": repo_id, "num_samples": len(rows)}


@app.local_entrypoint()
def main(hf_username: str = "mitvho09"):
    result = build_dataset.remote(hf_username=hf_username)
    print("RESULT:", result)