| |
| """Minimal Robometer wrapper that resizes frames to fit in memory.""" |
| from __future__ import annotations |
|
|
| import argparse |
| import os |
| from pathlib import Path |
| from typing import Optional, Tuple |
|
|
| import numpy as np |
| import torch |
| import cv2 |
|
|
| |
| import unsloth |
|
|
| from robometer.data.dataset_types import ProgressSample, Trajectory |
| from robometer.evals.eval_server import compute_batch_outputs |
| from robometer.evals.eval_viz_utils import extract_frames |
| from robometer.utils.save import load_model_from_hf |
| from robometer.utils.setup_utils import setup_batch_collator |
|
|
|
|
| def resize_frames(frames: np.ndarray, max_side: int = 256) -> np.ndarray: |
| """Resize frames so max side = max_side.""" |
| T, H, W, C = frames.shape |
| scale = max_side / max(H, W) |
| if scale < 1: |
| new_h = (int(H * scale) // 28) * 28 |
| new_w = (int(W * scale) // 28) * 28 |
| out = np.zeros((T, new_h, new_w, C), dtype=np.uint8) |
| for i in range(T): |
| out[i] = cv2.resize(frames[i], (new_w, new_h), interpolation=cv2.INTER_AREA) |
| return out |
| return frames |
|
|
|
|
| _CACHED_MODEL = None |
|
|
|
|
| def load_robometer_once(model_path: str = "robometer/Robometer-4B", device: str = "cuda"): |
| """Load model once globally.""" |
| global _CACHED_MODEL |
| if _CACHED_MODEL is None: |
| device = torch.device(device) |
| exp_config, tokenizer, processor, reward_model = load_model_from_hf( |
| model_path=model_path, device=device, |
| ) |
| reward_model.eval() |
| batch_collator = setup_batch_collator(processor, tokenizer, exp_config, is_eval=True) |
| _CACHED_MODEL = dict( |
| model=reward_model, tokenizer=tokenizer, processor=processor, |
| batch_collator=batch_collator, exp_config=exp_config, device=device, |
| ) |
| return _CACHED_MODEL |
|
|
|
|
| def compute_rewards( |
| video_frames: np.ndarray, |
| task: str, |
| max_side: int = 256, |
| state: Optional[dict] = None, |
| ) -> Tuple[np.ndarray, np.ndarray]: |
| """ |
| Compute per-frame progress and success probability. |
| |
| Args: |
| video_frames: (T, H, W, C) uint8 |
| task: task instruction string |
| max_side: resize frames so max dim = max_side |
| |
| Returns: |
| (progress, success): both arrays of length T |
| """ |
| if state is None: |
| state = load_robometer_once() |
|
|
| frames = resize_frames(video_frames, max_side=max_side) |
| T = int(frames.shape[0]) |
|
|
| traj = Trajectory( |
| frames=frames, |
| frames_shape=tuple(frames.shape), |
| task=task, |
| id="0", |
| metadata={"subsequence_length": T}, |
| video_embeddings=None, |
| ) |
| progress_sample = ProgressSample(trajectory=traj, sample_type="progress") |
| batch = state["batch_collator"]([progress_sample]) |
|
|
| progress_inputs = batch["progress_inputs"] |
| for key, value in progress_inputs.items(): |
| if hasattr(value, "to"): |
| progress_inputs[key] = value.to(state["device"]) |
|
|
| loss_config = getattr(state["exp_config"], "loss", None) |
| is_discrete = ( |
| getattr(loss_config, "progress_loss_type", "l2").lower() == "discrete" |
| if loss_config else False |
| ) |
| num_bins = ( |
| getattr(loss_config, "progress_discrete_bins", None) |
| or getattr(state["exp_config"].model, "progress_discrete_bins", 10) |
| ) |
|
|
| with torch.inference_mode(): |
| results = compute_batch_outputs( |
| state["model"], state["tokenizer"], progress_inputs, |
| sample_type="progress", is_discrete_mode=is_discrete, num_bins=num_bins, |
| ) |
|
|
| progress_pred = results.get("progress_pred", []) |
| progress_array = ( |
| np.array(progress_pred[0], dtype=np.float32) |
| if progress_pred and len(progress_pred) > 0 |
| else np.array([], dtype=np.float32) |
| ) |
| outputs_success = results.get("outputs_success", {}) |
| success_probs = outputs_success.get("success_probs", []) if outputs_success else [] |
| success_array = ( |
| np.array(success_probs[0], dtype=np.float32) |
| if success_probs and len(success_probs) > 0 |
| else np.array([], dtype=np.float32) |
| ) |
|
|
| return progress_array, success_array |
|
|
|
|
| def score_video_file( |
| video_path: str, |
| task: str, |
| fps: float = 3.0, |
| max_frames: int = 8, |
| max_side: int = 256, |
| state: Optional[dict] = None, |
| ) -> Tuple[np.ndarray, np.ndarray]: |
| frames = extract_frames(video_path, fps=fps, max_frames=max_frames) |
| if frames is None or frames.size == 0: |
| raise RuntimeError(f"Could not extract frames from {video_path}") |
| return compute_rewards(frames, task, max_side=max_side, state=state) |
|
|
|
|
| if __name__ == "__main__": |
| p = argparse.ArgumentParser() |
| p.add_argument("--video", required=True) |
| p.add_argument("--task", required=True) |
| p.add_argument("--model-path", default="robometer/Robometer-4B") |
| p.add_argument("--fps", type=float, default=3.0) |
| p.add_argument("--max-frames", type=int, default=8) |
| p.add_argument("--max-side", type=int, default=256) |
| p.add_argument("--out", default="/tmp/robometer_out") |
| args = p.parse_args() |
|
|
| print(f"Loading model from {args.model_path}...") |
| state = load_robometer_once(args.model_path) |
|
|
| print(f"Scoring video: {args.video}") |
| progress, success = score_video_file( |
| args.video, args.task, |
| fps=args.fps, max_frames=args.max_frames, max_side=args.max_side, |
| state=state, |
| ) |
|
|
| Path(args.out).mkdir(parents=True, exist_ok=True) |
| np.save(os.path.join(args.out, "progress.npy"), progress) |
| np.save(os.path.join(args.out, "success.npy"), success) |
|
|
| print(f"Progress: shape={progress.shape}, values={progress.tolist()}") |
| print(f"Success: shape={success.shape}, values={success.tolist()}") |
| print(f"Mean progress: {float(np.mean(progress)):.3f}") |
| print(f"Mean success: {float(np.mean(success)):.3f}") |
|
|
