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