Spaces:
Sleeping
Sleeping
File size: 4,133 Bytes
5b21b68 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | """Validated π₀.₅ UR inference and serializable result formatting."""
from __future__ import annotations
import json
import random
import tempfile
from dataclasses import dataclass
import numpy as np
import pandas as pd
from PIL import Image
ACTION_LABELS = ("dx", "dy", "dz", "droll", "dpitch", "dyaw", "gripper")
ACTION_HORIZON = 10
MODEL_ACTION_DIM = 32
BASE_SEED = 42
@dataclass(frozen=True)
class PredictionResult:
actions: pd.DataFrame
json_path: str
status: str
def _as_rgb_uint8(value) -> np.ndarray:
if isinstance(value, Image.Image):
return np.asarray(value.convert("RGB"), dtype=np.uint8)
array = np.asarray(value)
if np.issubdtype(array.dtype, np.floating):
maximum = float(np.nanmax(array)) if array.size else 0.0
if maximum <= 1.0:
array = array * 255.0
array = np.clip(array, 0, 255).astype(np.uint8)
return np.asarray(Image.fromarray(array).convert("RGB"), dtype=np.uint8)
def _validate_state(values) -> np.ndarray:
try:
state = np.asarray(values, dtype=np.float32)
except (TypeError, ValueError) as exc:
raise ValueError("state must contain seven numeric values") from exc
if state.shape != (7,):
raise ValueError("state must contain exactly seven values")
if not np.isfinite(state).all():
raise ValueError("state must contain seven finite values")
return state
def _validate_trial_index(value) -> int:
if isinstance(value, bool) or not isinstance(value, (int, np.integer)):
raise ValueError("trial index must be a non-negative integer")
result = int(value)
if result < 0:
raise ValueError("trial index must be a non-negative integer")
return result
def set_seed(seed: int) -> None:
random.seed(seed)
np.random.seed(seed)
try:
import torch
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
except ImportError:
pass
def run_prediction(
policy,
fixed_image,
wrist_image,
instruction,
state_values,
trial_index,
model_id,
checkpoint_path,
) -> PredictionResult:
if fixed_image is None:
raise ValueError("fixed-camera image is required")
if wrist_image is None:
raise ValueError("wrist-camera image is required")
if not isinstance(instruction, str) or not instruction.strip():
raise ValueError("task instruction is required")
prompt = instruction.strip()
state = _validate_state(state_values)
trial = _validate_trial_index(trial_index)
seed = BASE_SEED + trial
set_seed(seed)
noise = np.random.default_rng(seed).standard_normal(
(ACTION_HORIZON, MODEL_ACTION_DIM), dtype=np.float32
)
observation = {
"observation/image": _as_rgb_uint8(fixed_image),
"observation/wrist_image": _as_rgb_uint8(wrist_image),
"observation/state": state,
"prompt": prompt,
}
result = policy.infer(observation, noise=noise)
actions = np.asarray(result.get("actions"))
expected_shape = (ACTION_HORIZON, len(ACTION_LABELS))
if actions.shape != expected_shape:
raise RuntimeError(f"policy returned action shape {actions.shape}; expected {expected_shape}")
if not np.isfinite(actions).all():
raise RuntimeError("policy returned non-finite actions")
table = pd.DataFrame(actions, columns=ACTION_LABELS)
payload = {
"model_id": model_id,
"checkpoint_path": checkpoint_path,
"instruction": prompt,
"state": state.tolist(),
"trial_index": trial,
"seed": seed,
"action_labels": list(ACTION_LABELS),
"actions": actions.tolist(),
}
with tempfile.NamedTemporaryFile(
mode="w", suffix=".json", delete=False, encoding="utf-8"
) as stream:
json.dump(payload, stream, indent=2, ensure_ascii=False)
json_path = stream.name
return PredictionResult(
actions=table,
json_path=json_path,
status=f"Prediction complete (seed={seed}, action_shape={expected_shape}).",
)
|