""" RoboMind VLA — Task 9: app.py FastAPI web UI for the RoboMind VLA reward judge. Runs on Modal GPU with a public URL. Usage: modal deploy app.py Then visit the public URL printed in the output. """ import json import os import tempfile from typing import List import modal app_image = ( modal.Image.debian_slim(python_version="3.11") .apt_install("ffmpeg") .pip_install( "torch==2.4.0", "torchvision==0.19.0", "transformers==4.40.0", "peft==0.11.1", "accelerate==0.30.1", "pillow", "sentencepiece", "huggingface_hub", "fastapi==0.115.6", "uvicorn", "pydantic<2.13", "numpy<2", "opencv-python-headless", "python-multipart", "librosa==0.10.2", "soundfile", ) .run_commands( "python -c \"" "import os, sys; " "d = os.path.join(sys.prefix, 'lib/python3.11/site-packages/flash_attn'); " "os.makedirs(d, exist_ok=True); " "open(os.path.join(d, '__init__.py'), 'w').write(''); " "open(os.path.join(d, 'flash_attn_interface.py'), 'w').write(" "'def flash_attn_func(*a, **kw): raise NotImplementedError\\n" "def flash_attn_varlen_func(*a, **kw): raise NotImplementedError\\n'); " "print('flash_attn stub created')\"" ) .add_local_file("hybrid_judge.py", "/root/hybrid_judge.py") ) app = modal.App("robomind-gradio", image=app_image) volume = modal.Volume.from_name("robomind-data", create_if_missing=True) ADAPTER_REPO = "mitvho09/robomind-minicpm-loco-lora" INSTRUCTION_PROMPT = ( "You are RoboMind VLA, a vision-language reward model for robot locomotion. " "You are shown keyframes from a MuJoCo 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." ) METADATA_PATH = "/data/rollouts/metadata.jsonl" HTML_PAGE = """ RoboMind VLA — Locomotion Reward Judge

RoboMind VLA

Humanoid Locomotion Reward Judge

Image Upload
Video Upload
URL


""" def _extract_keyframes(video_path: str, n_frames: int = 6): import cv2 from PIL import Image cap = cv2.VideoCapture(video_path) total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) if total <= 0: cap.release() return [] indices = [int(i * total / n_frames) for i in range(n_frames)] frames = [] for idx in indices: cap.set(cv2.CAP_PROP_POS_FRAMES, idx) ret, frame = cap.read() if ret: frames.append(Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))) cap.release() return frames _model_cache = {} def _get_model(): if "model" in _model_cache: return _model_cache["model"], _model_cache["tokenizer"] import torch from transformers import AutoModel, AutoTokenizer from peft import PeftModel from huggingface_hub import login hf_token = os.environ.get("HF_TOKEN") if hf_token: login(token=hf_token) print("[robomind] loading model...") tokenizer = AutoTokenizer.from_pretrained( "openbmb/MiniCPM-V-2_6", trust_remote_code=True ) base_model = AutoModel.from_pretrained( "openbmb/MiniCPM-V-2_6", trust_remote_code=True, torch_dtype=torch.bfloat16, device_map="auto", ) model = PeftModel.from_pretrained(base_model, ADAPTER_REPO) model.eval() _model_cache["model"] = model _model_cache["tokenizer"] = tokenizer print("[robomind] model loaded") return model, tokenizer def _judge_images(model, tokenizer, images, n_images: int = 6): n = min(len(images), n_images) image_tokens = "\\n".join(f"" for k in range(n)) user_content = f"{image_tokens}\\n{INSTRUCTION_PROMPT}" response = model.chat( image=images[:n], msgs=[{"role": "user", "content": user_content}], tokenizer=tokenizer, max_new_tokens=512, ) return response if isinstance(response, str) else str(response) def _parse_response(response: str) -> dict: import re response = response.strip() # Try to find JSON blocks — prefer the last complete one json_blocks = list(re.finditer(r'\{[^{}]*\}', response, re.DOTALL)) if json_blocks: for block in reversed(json_blocks): try: return json.loads(block.group()) except json.JSONDecodeError: continue # Fallback: try entire response try: return json.loads(response) except json.JSONDecodeError: pass return {"raw_response": response} def _load_metadata(): """Load rollout metadata for hybrid scoring.""" import csv meta = {} if os.path.exists(METADATA_PATH): with open(METADATA_PATH) as f: for line in f: r = json.loads(line.strip()) key = (r["env"], r["tier"], r["episode_id"]) meta[key] = r return meta def _lookup_metadata(video_name: str, metadata: dict): """Try to find metadata for a video by parsing its filename. Tries exact match first, then fuzzy match on env+tier+episode. """ import re # Exact match m = re.match(r"(\w+)_(\w+)_ep(\d+)\.mp4", video_name) if m: env, tier, ep_id = m.group(1), m.group(2), int(m.group(3)) key = (env, tier, ep_id) if key in metadata: return metadata[key] # Fuzzy match: try all entries and find closest for key, entry in metadata.items(): env, tier, ep_id = key if f"{env}_{tier}_ep{ep_id}" in video_name: return entry return None def _compute_hybrid(parsed: dict, metadata_entry: dict = None, metadata: dict = None, sound_analysis: dict = None): """Run hybrid judge combining VLM + rule-based scoring.""" import sys if "/root" not in sys.path: sys.path.insert(0, "/root") from hybrid_judge import hybrid_judge, hybrid_to_dict if metadata_entry: env = metadata_entry["env"] all_metadata = metadata or {} env_rets = [v["return"] for v in all_metadata.values() if v["env"] == env] min_ret = min(env_rets) if env_rets else 0 max_ret = max(env_rets) if env_rets else 1 score = hybrid_judge( vlm_parsed=parsed, ep_return=metadata_entry["return"], min_return=min_ret, max_return=max_ret, fell=metadata_entry.get("fell", False), num_steps=metadata_entry.get("num_steps", 0), tier=metadata_entry.get("tier", "unknown"), env=metadata_entry.get("env", "unknown"), ) else: score = hybrid_judge(vlm_parsed=parsed) result = hybrid_to_dict(score) if sound_analysis: result["sound_analysis"] = sound_analysis if sound_analysis.get("has_fall"): confidence = sound_analysis.get("fall_confidence", 0.0) penalty = confidence * 0.3 result["predicted_reward"] = max(0.0, result["predicted_reward"] - penalty) result["anomaly"] = (result.get("anomaly") or "") + f" [audio: fall detected, conf={confidence:.2f}]" return result @app.function( image=app_image, gpu="A100-40GB", volumes={"/data": volume}, secrets=[modal.Secret.from_name("huggingface-secret")], timeout=3600, ) @modal.asgi_app() def serve(): from fastapi import FastAPI, UploadFile, File, Request from fastapi.responses import HTMLResponse, JSONResponse from typing import List web_app = FastAPI() metadata = _load_metadata() print(f"[robomind] loaded {len(metadata)} metadata entries") @web_app.get("/", response_class=HTMLResponse) async def index(): return HTML_PAGE @web_app.post("/judge/images") async def judge_images(files: List[UploadFile] = File(...)): from PIL import Image import io images = [] for f in files: data = await f.read() images.append(Image.open(io.BytesIO(data)).convert("RGB")) model, tokenizer = _get_model() response = _judge_images(model, tokenizer, images) parsed = _parse_response(response) hybrid = _compute_hybrid(parsed, metadata=metadata) return JSONResponse({"parsed": parsed, "hybrid": hybrid, "raw": response}) @web_app.post("/judge/video") async def judge_video(file: UploadFile = File(...)): data = await file.read() original_name = file.filename or "unknown.mp4" tmp = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) tmp.write(data) tmp.close() frames = _extract_keyframes(tmp.name) sound_result = None try: from robomind.sound import SoundAnalyzer analyzer = SoundAnalyzer() sa = analyzer.analyze_video(tmp.name) sound_result = { "has_fall": sa.has_fall, "fall_confidence": round(sa.fall_confidence, 3), "has_impact": sa.has_impact, "impact_count": sa.impact_count, "has_motor_strain": sa.has_motor_strain, "has_rhythmic_gait": sa.has_rhythmic_gait, "gait_quality": round(sa.gait_quality, 3), "explanation": sa.explanation, } except Exception as e: sound_result = {"error": str(e)} os.unlink(tmp.name) if not frames: return JSONResponse({"error": "Failed to extract frames"}, status_code=400) model, tokenizer = _get_model() response = _judge_images(model, tokenizer, frames) parsed = _parse_response(response) meta_entry = _lookup_metadata(original_name, metadata) hybrid = _compute_hybrid(parsed, meta_entry, metadata, sound_result) return JSONResponse({"parsed": parsed, "hybrid": hybrid, "raw": response}) @web_app.post("/judge/url") async def judge_url(request: Request): import urllib.request body = await request.json() url = body.get("url", "").strip() if not url: return JSONResponse({"error": "No URL provided"}, status_code=400) from urllib.parse import urlparse url_path = urlparse(url).path original_name = os.path.basename(url_path) or "download.mp4" tmp = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) try: urllib.request.urlretrieve(url, tmp.name) except Exception as e: os.unlink(tmp.name) return JSONResponse({"error": f"Download failed: {e}"}, status_code=400) frames = _extract_keyframes(tmp.name) sound_result = None try: from robomind.sound import SoundAnalyzer analyzer = SoundAnalyzer() sa = analyzer.analyze_video(tmp.name) sound_result = { "has_fall": sa.has_fall, "fall_confidence": round(sa.fall_confidence, 3), "has_impact": sa.has_impact, "impact_count": sa.impact_count, "has_motor_strain": sa.has_motor_strain, "has_rhythmic_gait": sa.has_rhythmic_gait, "gait_quality": round(sa.gait_quality, 3), "explanation": sa.explanation, } except Exception as e: sound_result = {"error": str(e)} os.unlink(tmp.name) if not frames: return JSONResponse({"error": "Failed to extract frames"}, status_code=400) model, tokenizer = _get_model() response = _judge_images(model, tokenizer, frames) parsed = _parse_response(response) meta_entry = _lookup_metadata(original_name, metadata) hybrid = _compute_hybrid(parsed, meta_entry, metadata, sound_result) return JSONResponse({"parsed": parsed, "hybrid": hybrid, "raw": response}) return web_app @app.local_entrypoint() def main(): serve.remote()