robometer/rewardeval_ui
4
1#!/usr/bin/env python32"""3Gradio app for Robometer (RBM) inference visualization.4Supports single video (progress/success) and dual video (preference/progress) predictions.5Uses eval server for inference instead of loading models locally.6"""7 8import os9import tempfile10from pathlib import Path11from typing import Optional, Tuple12import logging13 14import gradio as gr15 16try:17 import spaces # Required for ZeroGPU on Hugging Face Spaces18except ImportError:19 spaces = None # Not available when running locally20import matplotlib21 22matplotlib.use("Agg") # Use non-interactive backend23import matplotlib.pyplot as plt24import numpy as np25import requests26from typing import Any, List, Optional, Tuple27 28from dataset_types import Trajectory, ProgressSample, PreferenceSample29from eval_utils import build_payload, post_batch_npy30from eval_viz_utils import (31 create_combined_progress_success_plot,32 create_progress_success_gif,33 extract_frames,34)35from datasets import load_dataset as load_dataset_hf, get_dataset_config_names36 37logger = logging.getLogger(__name__)38 39# Example media: directory for bundled examples40_EXAMPLES_DIR = Path(__file__).resolve().parent / "examples"41 42def _progress_examples():43 """Progress tab examples from examples/progress/*.mp4 (video path, task)."""44 tasks = [45 "Open the red drawer",46 "Put the green block in the brown bowl",47 "Put the apple on the tray",48 "Put the marker in the holder",49 ]50 out = []51 progress_dir = _EXAMPLES_DIR / "progress"52 local_videos = sorted(progress_dir.glob("*.mp4")) if progress_dir.exists() else []53 for i, task in enumerate(tasks):54 if i < len(local_videos):55 out.append([str(local_videos[i]), task])56 return out57 58def _preference_examples():59 """Preference examples from examples/preference: blue/green trash bin pair + subopt vs successful pairs."""60 out = []61 pref_dir = _EXAMPLES_DIR / "preference"62 if not pref_dir.exists():63 return []64 65 by_stem = {} # stem (no .mp4) -> full path66 for p in pref_dir.glob("*.mp4"):67 stem = p.stem68 by_stem[stem] = str(p)69 70 # 1) Explicit pair: open_blue_trash_bin + open_green_trash_bin71 if "open_blue_trash_bin" in by_stem and "open_green_trash_bin" in by_stem:72 out.append([73 by_stem["open_blue_trash_bin"],74 by_stem["open_green_trash_bin"],75 "Open blue trash bin",76 ])77 78 # 2) Subopt vs successful: group by base (strip _subopt, _successful, _successful;, _failure)79 def get_base(stem: str) -> Optional[str]:80 if stem.endswith("_successful;"):81 return stem.replace("_successful;", "")82 for suffix in ("_subopt", "_suboptimal", "_successful", "_success", "_failure"):83 if stem.endswith(suffix):84 return stem[: -len(suffix)]85 return None86 87 base_to_files = {} # base -> {"subopt": path, "successful": path}88 for stem, path in by_stem.items():89 if stem in ("open_blue_trash_bin", "open_green_trash_bin"):90 continue91 base = get_base(stem)92 if not base:93 continue94 if base not in base_to_files:95 base_to_files[base] = {}96 if "_subopt" in stem:97 base_to_files[base]["subopt"] = path98 elif "_successful" in stem or "_success" in stem:99 base_to_files[base]["successful"] = path100 101 pref_tasks = {102 "bread_in_oven": "Put bread in oven",103 "open_red_drawer": "Open the red drawer",104 "put_blue_cup_sink": "Put the blue cup in the sink",105 }106 for base in sorted(base_to_files.keys()):107 d = base_to_files[base]108 if "subopt" in d and "successful" in d:109 task = pref_tasks.get(base, base.replace("_", " ").title())110 out.append([d["subopt"], d["successful"], task])111 112 return out113 114# Predefined dataset names (same as visualizer)115PREDEFINED_DATASETS = [116 "abraranwar/agibotworld_alpha_rfm",117 "abraranwar/libero_rfm",118 "abraranwar/usc_koch_rewind_rfm",119 "aliangdw/metaworld",120 "anqil/rh20t_rfm",121 "anqil/rh20t_subset_rfm",122 "jesbu1/auto_eval_rfm",123 "jesbu1/egodex_rfm",124 "jesbu1/epic_rfm",125 "jesbu1/fino_net_rfm",126 "jesbu1/failsafe_rfm",127 "jesbu1/hand_paired_rfm",128 "jesbu1/galaxea_rfm",129 "jesbu1/h2r_rfm",130 "jesbu1/humanoid_everyday_rfm",131 "jesbu1/molmoact_rfm",132 "jesbu1/motif_rfm",133 "jesbu1/oxe_rfm",134 "jesbu1/oxe_rfm_eval",135 "jesbu1/ph2d_rfm",136 "jesbu1/racer_rfm",137 "jesbu1/roboarena_0825_rfm",138 "jesbu1/soar_rfm",139 "ykorkmaz/libero_failure_rfm",140 "aliangdw/usc_xarm_policy_ranking",141 "aliangdw/usc_franka_policy_ranking",142 "aliangdw/utd_so101_policy_ranking",143 "aliangdw/utd_so101_human",144 "jesbu1/utd_so101_clean_policy_ranking_top",145 "jesbu1/utd_so101_clean_policy_ranking_wrist",146 "jesbu1/mit_franka_p-rank_rfm",147 "jesbu1/usc_koch_p_ranking_rfm",148]149 150# Default eval server URL (official Robometer demo)151DEFAULT_SERVER_URL = "https://robometer.a.pinggy.link"152 153# Global server state154_server_state = {155 "server_url": DEFAULT_SERVER_URL,156 "base_url": DEFAULT_SERVER_URL,157}158 159 160def discover_available_models(161 base_url: str = "http://40.119.56.66", port_range: tuple = (8000, 8010)162) -> List[Tuple[str, str]]:163 """Discover available models by pinging the base URL as-is, or ports in the specified range.164 165 If base_url is a full URL (e.g. https://robometer.a.pinggy.link), it is tried as-is first.166 Otherwise we try base_url:8000, base_url:8001, ... up to end_port.167 168 Returns:169 List of tuples: [(server_url, model_name), ...]170 """171 base_url = base_url.strip().rstrip("/")172 if not base_url:173 return []174 175 available_models = []176 # Try base_url as-is first (for Pinggy/tunnel URLs like https://robometer.a.pinggy.link)177 try:178 health_url = f"{base_url}/health"179 health_response = requests.get(health_url, timeout=5.0)180 if health_response.status_code == 200:181 try:182 model_info_url = f"{base_url}/model_info"183 model_info_response = requests.get(model_info_url, timeout=5.0)184 if model_info_response.status_code == 200:185 model_info_data = model_info_response.json()186 model_name = model_info_data.get("model_path", base_url)187 available_models.append((base_url, model_name))188 else:189 available_models.append((base_url, base_url))190 except Exception:191 available_models.append((base_url, base_url))192 return available_models193 except requests.exceptions.RequestException:194 pass195 196 # Port scan: base_url is a host (e.g. http://40.119.56.66), try ports in range197 start_port, end_port = port_range198 for port in range(start_port, end_port + 1):199 server_url = f"{base_url}:{port}"200 try:201 health_url = f"{server_url}/health"202 health_response = requests.get(health_url, timeout=2.0)203 if health_response.status_code == 200:204 try:205 model_info_url = f"{server_url}/model_info"206 model_info_response = requests.get(model_info_url, timeout=2.0)207 if model_info_response.status_code == 200:208 model_info_data = model_info_response.json()209 model_name = model_info_data.get("model_path", f"Model on port {port}")210 available_models.append((server_url, model_name))211 else:212 available_models.append((server_url, f"Model on port {port}"))213 except Exception:214 available_models.append((server_url, f"Model on port {port}"))215 except requests.exceptions.RequestException:216 continue217 218 return available_models219 220 221def get_model_info_for_url(server_url: str) -> Optional[str]:222 """Get formatted model info for a given server URL."""223 if not server_url:224 return None225 226 try:227 model_info_url = server_url.rstrip("/") + "/model_info"228 model_info_response = requests.get(model_info_url, timeout=5.0)229 if model_info_response.status_code == 200:230 model_info_data = model_info_response.json()231 return format_model_info(model_info_data)232 except Exception as e:233 logger.warning(f"Could not fetch model info: {e}")234 return None235 236 237def check_server_health(server_url: str) -> Tuple[str, Optional[dict], Optional[str]]:238 """Check server health and get model info."""239 if not server_url:240 return "Please provide a server URL.", None, None241 242 try:243 url = server_url.rstrip("/") + "/health"244 response = requests.get(url, timeout=5.0)245 response.raise_for_status()246 health_data = response.json()247 248 # Also try to get GPU status for more info249 try:250 status_url = server_url.rstrip("/") + "/gpu_status"251 status_response = requests.get(status_url, timeout=5.0)252 if status_response.status_code == 200:253 status_data = status_response.json()254 health_data.update(status_data)255 except:256 pass257 258 # Try to get model info259 model_info_text = get_model_info_for_url(server_url)260 261 _server_state["server_url"] = server_url262 return (263 f"Server connected: {health_data.get('available_gpus', 0)}/{health_data.get('total_gpus', 0)} GPUs available",264 health_data,265 model_info_text,266 )267 except requests.exceptions.RequestException as e:268 return f"Error connecting to server: {str(e)}", None, None269 270 271def format_model_info(model_info: dict) -> str:272 """Format model info and experiment config as markdown."""273 lines = ["## Model Information\n"]274 275 # Model path276 model_path = model_info.get("model_path", "Unknown")277 lines.append(f"**Model Path:** `{model_path}`\n")278 279 # Number of GPUs280 num_gpus = model_info.get("num_gpus", "Unknown")281 lines.append(f"**Number of GPUs:** {num_gpus}\n")282 283 # Model architecture284 model_arch = model_info.get("model_architecture", {})285 if model_arch and "error" not in model_arch:286 lines.append("\n## Model Architecture\n")287 288 model_class = model_arch.get("model_class", "Unknown")289 model_module = model_arch.get("model_module", "Unknown")290 lines.append(f"- **Model Class:** `{model_class}`\n")291 lines.append(f"- **Module:** `{model_module}`\n")292 293 # Parameter counts294 total_params = model_arch.get("total_parameters")295 trainable_params = model_arch.get("trainable_parameters")296 frozen_params = model_arch.get("frozen_parameters")297 trainable_pct = model_arch.get("trainable_percentage")298 299 if total_params is not None:300 lines.append(f"\n### Parameter Statistics\n")301 lines.append(f"- **Total Parameters:** {total_params:,}\n")302 if trainable_params is not None:303 lines.append(f"- **Trainable Parameters:** {trainable_params:,}\n")304 if frozen_params is not None:305 lines.append(f"- **Frozen Parameters:** {frozen_params:,}\n")306 if trainable_pct is not None:307 lines.append(f"- **Trainable Percentage:** {trainable_pct:.2f}%\n")308 309 # Architecture summary310 arch_summary = model_arch.get("architecture_summary", [])311 if arch_summary:312 lines.append(f"\n### Architecture Summary (Top-Level Modules)\n")313 for module_info in arch_summary[:10]: # Show first 10 modules314 name = module_info.get("name", "Unknown")315 module_type = module_info.get("type", "Unknown")316 params = module_info.get("parameters", 0)317 lines.append(f"- **{name}** (`{module_type}`): {params:,} parameters\n")318 319 # Experiment config320 exp_config = model_info.get("experiment_config", {})321 if exp_config:322 lines.append("\n## Experiment Configuration\n")323 324 # Model config325 model_cfg = exp_config.get("model", {})326 if model_cfg:327 lines.append("### Model Configuration\n")328 lines.append(f"- **Base Model:** `{model_cfg.get('base_model_id', 'N/A')}`\n")329 lines.append(f"- **Model Type:** `{model_cfg.get('model_type', 'N/A')}`\n")330 lines.append(f"- **Train Progress Head:** {model_cfg.get('train_progress_head', False)}\n")331 lines.append(f"- **Train Preference Head:** {model_cfg.get('train_preference_head', False)}\n")332 lines.append(f"- **Train Success Head:** {model_cfg.get('train_success_head', False)}\n")333 lines.append(f"- **Use PEFT:** {model_cfg.get('use_peft', False)}\n")334 lines.append(f"- **Use Unsloth:** {model_cfg.get('use_unsloth', False)}\n")335 336 # Data config337 data_cfg = exp_config.get("data", {})338 if data_cfg:339 lines.append("\n### Data Configuration\n")340 lines.append(f"- **Max Frames:** {data_cfg.get('max_frames', 'N/A')}\n")341 lines.append(342 f"- **Resized Dimensions:** {data_cfg.get('resized_height', 'N/A')}x{data_cfg.get('resized_width', 'N/A')}\n"343 )344 train_datasets = data_cfg.get("train_datasets", [])345 if train_datasets:346 lines.append(f"- **Train Datasets:** {', '.join(train_datasets)}\n")347 eval_datasets = data_cfg.get("eval_datasets", [])348 if eval_datasets:349 lines.append(f"- **Eval Datasets:** {', '.join(eval_datasets)}\n")350 351 # Training config352 training_cfg = exp_config.get("training", {})353 if training_cfg:354 lines.append("\n### Training Configuration\n")355 lines.append(f"- **Learning Rate:** {training_cfg.get('learning_rate', 'N/A')}\n")356 lines.append(f"- **Batch Size:** {training_cfg.get('per_device_train_batch_size', 'N/A')}\n")357 lines.append(358 f"- **Gradient Accumulation Steps:** {training_cfg.get('gradient_accumulation_steps', 'N/A')}\n"359 )360 lines.append(f"- **Max Steps:** {training_cfg.get('max_steps', 'N/A')}\n")361 362 return "".join(lines)363 364 365def load_rbm_dataset(dataset_name, config_name):366 """Load an RBM-format dataset from HuggingFace Hub."""367 try:368 if not dataset_name or not config_name:369 return None, "Please provide both dataset name and configuration"370 371 dataset = load_dataset_hf(dataset_name, name=config_name, split="train")372 373 if len(dataset) == 0:374 return None, f"Dataset {dataset_name}/{config_name} is empty"375 376 return dataset, f"Loaded {len(dataset)} trajectories from {dataset_name}/{config_name}"377 except Exception as e:378 error_msg = str(e)379 if "not found" in error_msg.lower():380 return None, f"Dataset or configuration not found: {dataset_name}/{config_name}"381 elif "authentication" in error_msg.lower():382 return None, f"Authentication required for {dataset_name}"383 else:384 return None, f"Error loading dataset: {error_msg}"385 386 387def get_available_configs(dataset_name):388 """Get available configurations for a dataset."""389 try:390 configs = get_dataset_config_names(dataset_name)391 return configs392 except Exception as e:393 logger.warning(f"Error getting configs for {dataset_name}: {e}")394 return []395 396 397def get_trajectory_video_path(dataset, index, dataset_name):398 """Get video path and metadata from a trajectory in the dataset."""399 try:400 item = dataset[int(index)]401 frames_data = item["frames"]402 403 if isinstance(frames_data, str):404 # Construct HuggingFace Hub URL405 if dataset_name:406 video_path = f"https://huggingface.co/datasets/{dataset_name}/resolve/main/{frames_data}"407 else:408 video_path = f"https://huggingface.co/datasets/rewardfm/rbm-1m/resolve/main/{frames_data}"409 410 task = item.get("task", "Complete the task")411 quality_label = item.get("quality_label", None)412 partial_success = item.get("partial_success", None)413 414 return video_path, task, quality_label, partial_success415 else:416 return None, None, None, None417 except Exception as e:418 logger.error(f"Error getting trajectory video path: {e}")419 return None, None, None, None420 421 422def process_single_video(423 video_path: str,424 task_text: str = "Complete the task",425 server_url: str = "",426 fps: float = 1.0,427 use_frame_steps: bool = False,428) -> Tuple[Optional[str], Optional[str], Optional[str]]:429 """Process single video for progress and success predictions using eval server.430 Returns (static_plot_path, video_path, info_text). video_path is the 5 sec MP4 animation; may be None if creation fails.431 """432 # Get server URL from state if not provided433 if not server_url:434 server_url = _server_state.get("server_url")435 436 if not server_url:437 return None, None, "Please select a model from the dropdown above and ensure it's connected."438 439 if video_path is None:440 return None, None, "Please provide a video."441 442 try:443 frames_array = extract_frames(video_path, fps=fps)444 if frames_array is None or frames_array.size == 0:445 return None, None, "Could not extract frames from video."446 447 # Convert frames to (T, H, W, C) numpy array with uint8 values448 if frames_array.dtype != np.uint8:449 frames_array = np.clip(frames_array, 0, 255).astype(np.uint8)450 451 num_frames = frames_array.shape[0]452 frames_shape = frames_array.shape # (T, H, W, C)453 454 # Create target progress (placeholder - would be None in real use)455 target_progress = np.linspace(0.0, 1.0, num=num_frames).tolist()456 success_label = [1.0 if prog > 0.5 else 0.0 for prog in target_progress]457 458 # predict_last_frame_mask: server collator requires a list (1.0 per frame = no masking for inference)459 predict_last_frame_mask = [1.0] * num_frames460 461 # Create Trajectory462 trajectory = Trajectory(463 task=task_text,464 frames=frames_array,465 frames_shape=frames_shape,466 target_progress=target_progress,467 success_label=success_label,468 predict_last_frame_mask=predict_last_frame_mask,469 metadata={"source": "gradio_app"},470 )471 472 # Create ProgressSample473 progress_sample = ProgressSample(474 trajectory=trajectory,475 data_gen_strategy="demo",476 )477 478 # Build payload and send to server479 files, sample_data = build_payload([progress_sample])480 # Add use_frame_steps flag as extra form data481 extra_data = {"use_frame_steps": use_frame_steps} if use_frame_steps else None482 response = post_batch_npy(server_url, files, sample_data, timeout_s=120.0, extra_form_data=extra_data)483 484 # Process response485 outputs_progress = response.get("outputs_progress", {})486 progress_pred = outputs_progress.get("progress_pred", [])487 outputs_success = response.get("outputs_success", {})488 success_probs = outputs_success.get("success_probs", []) if outputs_success else None489 490 # Extract progress predictions491 if progress_pred and len(progress_pred) > 0:492 progress_array = np.array(progress_pred[0]) # First sample493 else:494 progress_array = np.array([])495 496 # Extract success predictions if available497 success_array = None498 if success_probs and len(success_probs) > 0:499 success_array = np.array(success_probs[0])500 501 # Convert success_array to binary if available502 success_binary = None503 if success_array is not None:504 success_binary = (success_array > 0.5).astype(float)505 506 # Create combined plot using shared helper function507 fig = create_combined_progress_success_plot(508 progress_pred=progress_array if len(progress_array) > 0 else np.array([0.0]),509 num_frames=num_frames,510 success_binary=success_binary,511 success_probs=success_array,512 success_labels=None, # No ground truth labels available513 is_discrete_mode=False,514 title=f"Progress & Success - {task_text}",515 )516 517 # Save to temporary file518 tmp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png")519 fig.savefig(tmp_file.name, dpi=150, bbox_inches="tight")520 plt.close(fig)521 progress_plot = tmp_file.name522 523 info_text = f"**Frames processed:** {num_frames}\n"524 if len(progress_array) > 0:525 info_text += f"**Final progress:** {progress_array[-1]:.3f}\n"526 if success_array is not None and len(success_array) > 0:527 info_text += f"**Final success probability:** {success_array[-1]:.3f}\n"528 529 # Animated MP4: progress + success curves (5 sec clip) with optional video panel530 video_path = None531 if len(progress_array) > 0:532 mp4_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")533 mp4_file.close()534 video_path = create_progress_success_gif(535 progress_pred=progress_array,536 success_data=success_binary if success_binary is not None else success_array,537 video_frames=frames_array,538 output_path=mp4_file.name,539 title=task_text,540 duration_sec=5.0,541 )542 543 return progress_plot, video_path, info_text544 545 except Exception as e:546 return None, None, f"Error processing video: {str(e)}"547 548 549def process_two_videos(550 video_a_path: str,551 video_b_path: str,552 task_text: str = "Complete the task",553 prediction_type: str = "preference",554 server_url: str = "",555 fps: float = 1.0,556) -> Tuple[Optional[str], Optional[str], Optional[str]]:557 """Process two videos for preference or progress prediction using eval server."""558 # Get server URL from state if not provided559 if not server_url:560 server_url = _server_state.get("server_url")561 562 if not server_url:563 return "Please select a model from the dropdown above and ensure it's connected.", None, None564 565 if video_a_path is None or video_b_path is None:566 return "Please provide both videos.", None, None567 568 try:569 frames_array_a = extract_frames(video_a_path, fps=fps)570 frames_array_b = extract_frames(video_b_path, fps=fps)571 572 if frames_array_a is None or frames_array_a.size == 0:573 return "Could not extract frames from video A.", None, None574 if frames_array_b is None or frames_array_b.size == 0:575 return "Could not extract frames from video B.", None, None576 577 # Convert frames to uint8578 if frames_array_a.dtype != np.uint8:579 frames_array_a = np.clip(frames_array_a, 0, 255).astype(np.uint8)580 if frames_array_b.dtype != np.uint8:581 frames_array_b = np.clip(frames_array_b, 0, 255).astype(np.uint8)582 583 num_frames_a = frames_array_a.shape[0]584 num_frames_b = frames_array_b.shape[0]585 frames_shape_a = frames_array_a.shape586 frames_shape_b = frames_array_b.shape587 588 # Create target progress for both trajectories589 target_progress_a = np.linspace(0.0, 1.0, num=num_frames_a).tolist()590 target_progress_b = np.linspace(0.0, 1.0, num=num_frames_b).tolist()591 success_label_a = [1.0 if prog > 0.5 else 0.0 for prog in target_progress_a]592 success_label_b = [1.0 if prog > 0.5 else 0.0 for prog in target_progress_b]593 594 # predict_last_frame_mask: server collator requires a list per trajectory (1.0 = no masking)595 mask_a = [1.0] * num_frames_a596 mask_b = [1.0] * num_frames_b597 598 # Create trajectories599 trajectory_a = Trajectory(600 task=task_text,601 frames=frames_array_a,602 frames_shape=frames_shape_a,603 target_progress=target_progress_a,604 success_label=success_label_a,605 predict_last_frame_mask=mask_a,606 metadata={"source": "gradio_app", "trajectory": "A"},607 )608 609 trajectory_b = Trajectory(610 task=task_text,611 frames=frames_array_b,612 frames_shape=frames_shape_b,613 target_progress=target_progress_b,614 success_label=success_label_b,615 predict_last_frame_mask=mask_b,616 metadata={"source": "gradio_app", "trajectory": "B"},617 )618 619 if prediction_type == "preference":620 # Create PreferenceSample (A = chosen, B = rejected)621 preference_sample = PreferenceSample(622 chosen_trajectory=trajectory_a,623 rejected_trajectory=trajectory_b,624 data_gen_strategy="demo",625 )626 627 # Build payload and send to server628 files, sample_data = build_payload([preference_sample])629 response = post_batch_npy(server_url, files, sample_data, timeout_s=120.0)630 631 # Process response632 outputs_preference = response.get("outputs_preference", {})633 predictions = outputs_preference.get("predictions", [])634 prediction_probs = outputs_preference.get("prediction_probs", [])635 636 if prediction_probs and len(prediction_probs) > 0:637 prob = prediction_probs[0]638 preferred = "A" if prob > 0.5 else "B"639 result_text = (640 'For task <span style="color: #2563eb; font-weight: 600;">'641 f"{task_text}"642 '</span>, Video <span style="color: #16a34a; font-weight: 700;">'643 f"{preferred}"644 "</span> is more preferred."645 )646 else:647 result_text = "Could not extract preference prediction from server response."648 649 elif prediction_type == "progress":650 # Create ProgressSamples for both videos651 progress_sample_a = ProgressSample(652 trajectory=trajectory_a,653 data_gen_strategy="demo",654 )655 progress_sample_b = ProgressSample(656 trajectory=trajectory_b,657 data_gen_strategy="demo",658 )659 660 # Build payload and send to server661 files, sample_data = build_payload([progress_sample_a, progress_sample_b])662 response = post_batch_npy(server_url, files, sample_data, timeout_s=120.0)663 664 # Process response665 outputs_progress = response.get("outputs_progress", {})666 progress_pred = outputs_progress.get("progress_pred", [])667 668 result_text = f"**Progress Comparison:**\n"669 if progress_pred and len(progress_pred) >= 2:670 progress_a = np.array(progress_pred[0])671 progress_b = np.array(progress_pred[1])672 673 final_progress_a = float(progress_a[-1]) if len(progress_a) > 0 else 0.0674 final_progress_b = float(progress_b[-1]) if len(progress_b) > 0 else 0.0675 676 result_text += f"- Video A final progress: {final_progress_a:.3f}\n"677 result_text += f"- Video B final progress: {final_progress_b:.3f}\n"678 result_text += f"- Difference: {abs(final_progress_a - final_progress_b):.3f}\n"679 if final_progress_a > final_progress_b:680 result_text += f"- Video A has higher progress\n"681 elif final_progress_b > final_progress_a:682 result_text += f"- Video B has higher progress\n"683 else:684 result_text += f"- Both videos have equal progress\n"685 else:686 result_text += "Could not extract progress predictions from server response.\n"687 688 # Return result text and both video paths689 return result_text, video_a_path, video_b_path690 691 except Exception as e:692 return f"Error processing videos: {str(e)}", None, None693 694 695# Create Gradio interface696try:697 # Try with theme (Gradio 4.0+)698 demo = gr.Blocks(title="Robometer Evaluation Server", theme=gr.themes.Soft())699except TypeError:700 # Fallback for older Gradio versions without theme support701 demo = gr.Blocks(title="Robometer Evaluation Server")702 703with demo:704 gr.Markdown(705 """706 # Robometer Evaluation Server707 """708 )709 710 # Hidden state: fixed official server URL (no model selection in UI)711 server_url_state = gr.State(value=DEFAULT_SERVER_URL)712 713 # Sidebar: Robometer branding and official links (no model selection)714 with gr.Sidebar():715 gr.Markdown("## Robometer")716 gr.Markdown(717 "**Official project** · [robometer.github.io](https://robometer.github.io/)"718 )719 gr.Markdown("---")720 gr.Markdown("### Links")721 gr.HTML(722 """723 <style>724 .sidebar-pill { display: inline-flex; align-items: center; gap: 0.35rem; padding: 0.25rem 0.65rem;725 border-radius: 999px; border: 1px solid #e4e4e7; font-size: 0.9rem; text-decoration: none;726 color: inherit; }727 .sidebar-pill:hover { background: rgba(0,0,0,0.04); }728 .sidebar-pill .arrow { opacity: 0.6; font-size: 0.75em; }729 </style>730 <div style="display: flex; flex-wrap: wrap; gap: 0.5rem;">731 <a href="https://robometer.github.io/" target="_blank" rel="noopener" class="sidebar-pill" title="Project page">732 <span>🌐</span><span>Project</span><span class="arrow">↗</span>733 </a>734 <a href="https://github.com/robometer" target="_blank" rel="noopener" class="sidebar-pill" title="GitHub">735 <span>📂</span><span>Code</span><span class="arrow">↗</span>736 </a>737 <a href="https://huggingface.co/datasets/rewardfm/rbm-1m" target="_blank" rel="noopener" class="sidebar-pill" title="RBM-1M Dataset">738 <span>📊</span><span>Dataset</span><span class="arrow">↗</span>739 </a>740 <a href="https://huggingface.co/robometer/Robometer-4B" target="_blank" rel="noopener" class="sidebar-pill" title="Model weights">741 <span>💾</span><span>Weights</span><span class="arrow">↗</span>742 </a>743 </div>744 """745 )746 747 # Main content area with tabs748 with gr.Tabs():749 with gr.Tab("Progress Prediction"):750 with gr.Row():751 with gr.Column():752 single_video_input = gr.Video(label="Upload Video", height=300)753 task_text_input = gr.Textbox(754 label="Task Description",755 placeholder="Describe the task (e.g., 'Pick up the red block')",756 value="Complete the task",757 )758 fps_input_single = gr.Slider(759 label="FPS (Frames Per Second)",760 minimum=0.1,761 maximum=10.0,762 value=3.0,763 step=0.1,764 info="Frames per second to extract from video (higher = more frames)",765 )766 use_frame_steps_single = gr.Checkbox(767 label="Per Frame Progress Prediction",768 value=False,769 info="If enabled, predict progress per frame rather than feeding the entire video at once",770 )771 analyze_single_btn = gr.Button("Compute Progress", variant="primary")772 773 gr.Markdown("---")774 gr.Markdown("**OR Select from Dataset**")775 gr.Markdown("---")776 777 with gr.Accordion("📁 Select from Dataset", open=False):778 dataset_name_single = gr.Dropdown(779 choices=PREDEFINED_DATASETS,780 value="jesbu1/oxe_rfm",781 label="Dataset Name",782 allow_custom_value=True,783 )784 config_name_single = gr.Dropdown(785 choices=[], value="", label="Configuration Name", allow_custom_value=True786 )787 with gr.Row():788 refresh_configs_btn = gr.Button("🔄 Refresh Configs", variant="secondary", size="sm")789 load_dataset_btn = gr.Button("Load Dataset", variant="secondary", size="sm")790 791 dataset_status_single = gr.Markdown("", visible=False)792 with gr.Row():793 prev_traj_btn = gr.Button("⬅️ Prev", variant="secondary", size="sm")794 trajectory_slider = gr.Slider(795 minimum=0, maximum=0, step=1, value=0, label="Trajectory Index", interactive=True796 )797 next_traj_btn = gr.Button("Next ➡️", variant="secondary", size="sm")798 trajectory_metadata = gr.Markdown("", visible=False)799 use_dataset_video_btn = gr.Button("Use Selected Video", variant="secondary")800 801 with gr.Column():802 progress_plot = gr.Image(label="Progress & Success Prediction", height=320)803 progress_video = gr.Video(label="Animated Progress & Success (5 sec MP4)", height=320)804 info_output = gr.Markdown("")805 gr.Markdown("---")806 gr.Markdown("**Examples**")807 gr.Examples(808 examples=_progress_examples(),809 inputs=[single_video_input, task_text_input],810 label="Click an example to load video and task",811 )812 813 # State variables for dataset814 current_dataset_single = gr.State(None)815 816 def update_config_choices_single(dataset_name):817 """Update config choices when dataset changes."""818 if not dataset_name:819 return gr.update(choices=[], value="")820 try:821 configs = get_available_configs(dataset_name)822 if configs:823 return gr.update(choices=configs, value=configs[0])824 else:825 return gr.update(choices=[], value="")826 except Exception as e:827 logger.warning(f"Could not fetch configs: {e}")828 return gr.update(choices=[], value="")829 830 def load_dataset_single(dataset_name, config_name):831 """Load dataset and update slider."""832 dataset, status = load_rbm_dataset(dataset_name, config_name)833 if dataset is not None:834 max_index = len(dataset) - 1835 return (836 dataset,837 gr.update(value=status, visible=True),838 gr.update(839 maximum=max_index, value=0, interactive=True, label=f"Trajectory Index (0 to {max_index})"840 ),841 )842 else:843 return None, gr.update(value=status, visible=True), gr.update(maximum=0, value=0, interactive=False)844 845 def use_dataset_video(dataset, index, dataset_name):846 """Load video from dataset and update inputs."""847 if dataset is None:848 return (849 None,850 "Complete the task",851 gr.update(value="No dataset loaded", visible=True),852 gr.update(visible=False),853 )854 855 video_path, task, quality_label, partial_success = get_trajectory_video_path(856 dataset, index, dataset_name857 )858 if video_path:859 # Build metadata text860 metadata_lines = []861 if quality_label:862 metadata_lines.append(f"**Quality Label:** {quality_label}")863 if partial_success is not None:864 metadata_lines.append(f"**Partial Success:** {partial_success:.3f}")865 866 metadata_text = "\n".join(metadata_lines) if metadata_lines else ""867 status_text = f"✅ Loaded trajectory {index} from dataset"868 if metadata_text:869 status_text += f"\n\n{metadata_text}"870 871 return (872 video_path,873 task,874 gr.update(value=status_text, visible=True),875 gr.update(value=metadata_text, visible=bool(metadata_text)),876 )877 else:878 return (879 None,880 "Complete the task",881 gr.update(value="❌ Error loading trajectory", visible=True),882 gr.update(visible=False),883 )884 885 def next_trajectory(dataset, current_idx, dataset_name):886 """Go to next trajectory."""887 if dataset is None:888 return 0, None, "Complete the task", gr.update(visible=False), gr.update(visible=False)889 next_idx = min(current_idx + 1, len(dataset) - 1)890 video_path, task, quality_label, partial_success = get_trajectory_video_path(891 dataset, next_idx, dataset_name892 )893 894 if video_path:895 # Build metadata text896 metadata_lines = []897 if quality_label:898 metadata_lines.append(f"**Quality Label:** {quality_label}")899 if partial_success is not None:900 metadata_lines.append(f"**Partial Success:** {partial_success:.3f}")901 902 metadata_text = "\n".join(metadata_lines) if metadata_lines else ""903 return (904 next_idx,905 video_path,906 task,907 gr.update(value=metadata_text, visible=bool(metadata_text)),908 gr.update(value=f"✅ Trajectory {next_idx}/{len(dataset) - 1}", visible=True),909 )910 else:911 return current_idx, None, "Complete the task", gr.update(visible=False), gr.update(visible=False)912 913 def prev_trajectory(dataset, current_idx, dataset_name):914 """Go to previous trajectory."""915 if dataset is None:916 return 0, None, "Complete the task", gr.update(visible=False), gr.update(visible=False)917 prev_idx = max(current_idx - 1, 0)918 video_path, task, quality_label, partial_success = get_trajectory_video_path(919 dataset, prev_idx, dataset_name920 )921 922 if video_path:923 # Build metadata text924 metadata_lines = []925 if quality_label:926 metadata_lines.append(f"**Quality Label:** {quality_label}")927 if partial_success is not None:928 metadata_lines.append(f"**Partial Success:** {partial_success:.3f}")929 930 metadata_text = "\n".join(metadata_lines) if metadata_lines else ""931 return (932 prev_idx,933 video_path,934 task,935 gr.update(value=metadata_text, visible=bool(metadata_text)),936 gr.update(value=f"✅ Trajectory {prev_idx}/{len(dataset) - 1}", visible=True),937 )938 else:939 return current_idx, None, "Complete the task", gr.update(visible=False), gr.update(visible=False)940 941 def update_trajectory_on_slider_change(dataset, index, dataset_name):942 """Update trajectory metadata when slider changes."""943 if dataset is None:944 return gr.update(visible=False), gr.update(visible=False)945 946 video_path, task, quality_label, partial_success = get_trajectory_video_path(947 dataset, index, dataset_name948 )949 if video_path:950 # Build metadata text951 metadata_lines = []952 if quality_label:953 metadata_lines.append(f"**Quality Label:** {quality_label}")954 if partial_success is not None:955 metadata_lines.append(f"**Partial Success:** {partial_success:.3f}")956 957 metadata_text = "\n".join(metadata_lines) if metadata_lines else ""958 return (959 gr.update(value=metadata_text, visible=bool(metadata_text)),960 gr.update(value=f"Trajectory {index}/{len(dataset) - 1}", visible=True),961 )962 else:963 return gr.update(visible=False), gr.update(visible=False)964 965 # Dataset selection handlers966 dataset_name_single.change(967 fn=update_config_choices_single, inputs=[dataset_name_single], outputs=[config_name_single]968 )969 970 refresh_configs_btn.click(971 fn=update_config_choices_single, inputs=[dataset_name_single], outputs=[config_name_single]972 )973 974 load_dataset_btn.click(975 fn=load_dataset_single,976 inputs=[dataset_name_single, config_name_single],977 outputs=[current_dataset_single, dataset_status_single, trajectory_slider],978 )979 980 use_dataset_video_btn.click(981 fn=use_dataset_video,982 inputs=[current_dataset_single, trajectory_slider, dataset_name_single],983 outputs=[single_video_input, task_text_input, dataset_status_single, trajectory_metadata],984 )985 986 # Navigation buttons987 next_traj_btn.click(988 fn=next_trajectory,989 inputs=[current_dataset_single, trajectory_slider, dataset_name_single],990 outputs=[991 trajectory_slider,992 single_video_input,993 task_text_input,994 trajectory_metadata,995 dataset_status_single,996 ],997 )998 999 prev_traj_btn.click(1000 fn=prev_trajectory,1001 inputs=[current_dataset_single, trajectory_slider, dataset_name_single],1002 outputs=[1003 trajectory_slider,1004 single_video_input,1005 task_text_input,1006 trajectory_metadata,1007 dataset_status_single,1008 ],1009 )1010 1011 # Update metadata when slider changes1012 trajectory_slider.change(1013 fn=update_trajectory_on_slider_change,1014 inputs=[current_dataset_single, trajectory_slider, dataset_name_single],1015 outputs=[trajectory_metadata, dataset_status_single],1016 )1017 1018 analyze_single_btn.click(1019 fn=process_single_video,1020 inputs=[1021 single_video_input,1022 task_text_input,1023 server_url_state,1024 fps_input_single,1025 use_frame_steps_single,1026 ],1027 outputs=[progress_plot, progress_video, info_output],1028 api_name="process_single_video",1029 )1030 1031 with gr.Tab("Preference Prediction"):1032 # Full-width row: two videos side by side1033 with gr.Row():1034 video_a_input = gr.Video(label="Video A", height=320)1035 video_b_input = gr.Video(label="Video B", height=320)1036 1037 task_text_dual = gr.Textbox(1038 label="Task Description",1039 placeholder="Describe the task",1040 value="Complete the task",1041 )1042 analyze_dual_btn = gr.Button("Compute Preference", variant="primary")1043 1044 gr.Markdown("---")1045 gr.Markdown("**Examples**")1046 gr.Examples(1047 examples=_preference_examples(),1048 inputs=[video_a_input, video_b_input, task_text_dual],1049 label="Click an example to load Video A, Video B, and task",1050 )1051 1052 gr.Markdown("---")1053 gr.Markdown("**OR Select from Dataset**")1054 gr.Markdown("---")1055 1056 with gr.Accordion("📁 Video A - Select from Dataset", open=False):1057 dataset_name_a = gr.Dropdown(1058 choices=PREDEFINED_DATASETS,1059 value="jesbu1/oxe_rfm",1060 label="Dataset Name",1061 allow_custom_value=True,1062 )1063 config_name_a = gr.Dropdown(1064 choices=[], value="", label="Configuration Name", allow_custom_value=True1065 )1066 with gr.Row():1067 refresh_configs_btn_a = gr.Button("🔄 Refresh Configs", variant="secondary", size="sm")1068 load_dataset_btn_a = gr.Button("Load Dataset", variant="secondary", size="sm")1069 1070 dataset_status_a = gr.Markdown("", visible=False)1071 with gr.Row():1072 prev_traj_btn_a = gr.Button("⬅️ Prev", variant="secondary", size="sm")1073 trajectory_slider_a = gr.Slider(1074 minimum=0, maximum=0, step=1, value=0, label="Trajectory Index", interactive=True1075 )1076 next_traj_btn_a = gr.Button("Next ➡️", variant="secondary", size="sm")1077 trajectory_metadata_a = gr.Markdown("", visible=False)1078 use_dataset_video_btn_a = gr.Button("Use Selected Video for A", variant="secondary")1079 1080 with gr.Accordion("📁 Video B - Select from Dataset", open=False):1081 dataset_name_b = gr.Dropdown(1082 choices=PREDEFINED_DATASETS,1083 value="jesbu1/oxe_rfm",1084 label="Dataset Name",1085 allow_custom_value=True,1086 )1087 config_name_b = gr.Dropdown(1088 choices=[], value="", label="Configuration Name", allow_custom_value=True1089 )1090 with gr.Row():1091 refresh_configs_btn_b = gr.Button("🔄 Refresh Configs", variant="secondary", size="sm")1092 load_dataset_btn_b = gr.Button("Load Dataset", variant="secondary", size="sm")1093 1094 dataset_status_b = gr.Markdown("", visible=False)1095 with gr.Row():1096 prev_traj_btn_b = gr.Button("⬅️ Prev", variant="secondary", size="sm")1097 trajectory_slider_b = gr.Slider(1098 minimum=0, maximum=0, step=1, value=0, label="Trajectory Index", interactive=True1099 )1100 next_traj_btn_b = gr.Button("Next ➡️", variant="secondary", size="sm")1101 trajectory_metadata_b = gr.Markdown("", visible=False)1102 use_dataset_video_btn_b = gr.Button("Use Selected Video for B", variant="secondary")1103 1104 gr.Markdown("---")1105 gr.Markdown("### Preference result")1106 result_text = gr.Markdown("")1107 1108 # State variables for datasets1109 current_dataset_a = gr.State(None)1110 current_dataset_b = gr.State(None)1111 1112 # Helper functions for Video A1113 def update_config_choices_a(dataset_name):1114 """Update config choices for Video A when dataset changes."""1115 if not dataset_name:1116 return gr.update(choices=[], value="")1117 try:1118 configs = get_available_configs(dataset_name)1119 if configs:1120 return gr.update(choices=configs, value=configs[0])1121 else:1122 return gr.update(choices=[], value="")1123 except Exception as e:1124 logger.warning(f"Could not fetch configs: {e}")1125 return gr.update(choices=[], value="")1126 1127 def load_dataset_a(dataset_name, config_name):1128 """Load dataset A and update slider."""1129 dataset, status = load_rbm_dataset(dataset_name, config_name)1130 if dataset is not None:1131 max_index = len(dataset) - 11132 return (1133 dataset,1134 gr.update(value=status, visible=True),1135 gr.update(1136 maximum=max_index, value=0, interactive=True, label=f"Trajectory Index (0 to {max_index})"1137 ),1138 )1139 else:1140 return None, gr.update(value=status, visible=True), gr.update(maximum=0, value=0, interactive=False)1141 1142 def use_dataset_video_a(dataset, index, dataset_name):1143 """Load video A from dataset and update input."""1144 if dataset is None:1145 return (1146 None,1147 gr.update(value="No dataset loaded", visible=True),1148 gr.update(visible=False),1149 )1150 1151 video_path, task, quality_label, partial_success = get_trajectory_video_path(1152 dataset, index, dataset_name1153 )1154 if video_path:1155 # Build metadata text1156 metadata_lines = []1157 if quality_label:1158 metadata_lines.append(f"**Quality Label:** {quality_label}")1159 if partial_success is not None:1160 metadata_lines.append(f"**Partial Success:** {partial_success:.3f}")1161 1162 metadata_text = "\n".join(metadata_lines) if metadata_lines else ""1163 status_text = f"✅ Loaded trajectory {index} from dataset for Video A"1164 if metadata_text:1165 status_text += f"\n\n{metadata_text}"1166 1167 return (1168 video_path,1169 gr.update(value=status_text, visible=True),1170 gr.update(value=metadata_text, visible=bool(metadata_text)),1171 )1172 else:1173 return (1174 None,1175 gr.update(value="❌ Error loading trajectory", visible=True),1176 gr.update(visible=False),1177 )1178 1179 def next_trajectory_a(dataset, current_idx, dataset_name):1180 """Go to next trajectory for Video A."""1181 if dataset is None:1182 return 0, None, gr.update(visible=False), gr.update(visible=False)1183 next_idx = min(current_idx + 1, len(dataset) - 1)1184 video_path, task, quality_label, partial_success = get_trajectory_video_path(1185 dataset, next_idx, dataset_name1186 )1187 1188 if video_path:1189 # Build metadata text1190 metadata_lines = []1191 if quality_label:1192 metadata_lines.append(f"**Quality Label:** {quality_label}")1193 if partial_success is not None:1194 metadata_lines.append(f"**Partial Success:** {partial_success:.3f}")1195 1196 metadata_text = "\n".join(metadata_lines) if metadata_lines else ""1197 return (1198 next_idx,1199 video_path,1200 gr.update(value=metadata_text, visible=bool(metadata_text)),