Spaces:
Running on Zero
Running on Zero
| """ERVLA — Embodied Chain-of-Thought VLA policy: predict a robot action chunk from camera views. | |
| Loads `ERVLA/liberoplus` (the LIBERO-Plus-tuned MiBoT checkpoint: Qwen3-VL-4B backbone + a | |
| DiT flow-matching action head) and replicates the reference deployment path from | |
| `src/server/deploy_hf_liberoplus.py` in the ERVLA repo 1:1 — same prompt template, same raw | |
| 8-dim proprioceptive state zero-padded to the model's 60-dim state window, same manual | |
| action mask (translation+rotation+gripper only), same 5-step Euler flow integration, and | |
| the same LIBERO `libero_all` normalization stats for decoding actions. | |
| Inputs are two LIBERO-style camera views (base + left wrist) plus the instruction; the | |
| output is the predicted end-effector action chunk (Dxyz translation, Drotation, gripper | |
| command) rendered as a chart + table. No simulator is bundled — LIBERO-Plus / VLABench | |
| cannot run inside a Space, so the action chunk itself is the demo output. | |
| """ | |
| import os | |
| os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") | |
| import spaces # noqa: E402 (must precede torch) | |
| import time # noqa: E402 | |
| from typing import Any, Dict, List, Tuple # noqa: E402 | |
| import gradio as gr # noqa: E402 | |
| import numpy as np # noqa: E402 | |
| import torch # noqa: E402 | |
| from PIL import Image # noqa: E402 | |
| from transformers import AutoModel, AutoProcessor # noqa: E402 | |
| MODEL_ID = "ERVLA/liberoplus" | |
| NUM_FLOW_STEPS = 5 # reference server default --num-steps | |
| STATE_LENGTH, STATE_DIM, ACTION_DIM = 1, 60, 60 # from config.json | |
| # libero_all normalization stats, configs/data/sources/libero.yaml in the ERVLA repo. | |
| ACTION_MEAN = np.array([0.0628, 0.0868, -0.0904, 0.0005, 0.0056, -0.0052, -0.0496], dtype=np.float32) | |
| ACTION_STD = np.array([0.3355, 0.3784, 0.4447, 0.0392, 0.0634, 0.0780, 0.9988], dtype=np.float32) | |
| LIBERO_ROBOT_TYPES = ["libero_spatial", "libero_object", "libero_goal", "libero_10"] | |
| # Dataviz reference palette: validated categorical slots 1-3, light + dark steps. | |
| SERIES = { | |
| "light": ["#2a78d6", "#eb6834", "#1baf7a"], | |
| "dark": ["#3987e5", "#d95926", "#199e70"], | |
| } | |
| INK = {"light": "#0b0b0b", "dark": "#ffffff"} | |
| GRID = {"light": "#e1e0d9", "dark": "#2c2c2a"} | |
| SURFACE = {"light": "#fcfcfb", "dark": "#1a1a19"} | |
| processor: Any = AutoProcessor.from_pretrained(MODEL_ID, trust_remote_code=True, use_fast=False) | |
| def _load_model() -> Any: | |
| """Load with eager attention. | |
| The custom `MiBoTForActionGeneration` class declares `_supports_flash_attn` but not | |
| `_supports_sdpa`, so under transformers 4.57.1 the only universally available | |
| backends are flash_attention_2 and eager. The prebuilt Blackwell flash-attn wheel | |
| (pt212) has a torch-ABI mismatch against the runtime's preinstalled torch 2.13 | |
| (`undefined symbol: ...materialize_cow_storage`), so eager it is — transformers' | |
| own suggested fallback. The workload is a single forward pass (~700 vision+text | |
| tokens, /no_cot, no autoregressive decoding) + 5 small DiT flow steps, where eager | |
| attention costs only a few hundred ms over FA2. | |
| """ | |
| m: Any = AutoModel.from_pretrained( | |
| MODEL_ID, | |
| trust_remote_code=True, | |
| dtype=torch.bfloat16, | |
| attn_implementation="eager", | |
| ) | |
| print("[load] attention backend: eager", flush=True) | |
| return m.to("cuda") | |
| model = _load_model() | |
| model.eval() | |
| MODEL_DTYPE = next(model.parameters()).dtype | |
| print(f"[load] {MODEL_ID} on cuda, dtype={MODEL_DTYPE}", flush=True) | |
| def _messages(base_image: Image.Image, wrist_image: Image.Image, instruction: str) -> List[Dict[str, Any]]: | |
| """Build the chat messages exactly as the reference deployment server does.""" | |
| content = [ | |
| {"type": "text", "text": "The following observations are captured from multiple views.\n# Base View\n"}, | |
| {"type": "image", "image": base_image}, | |
| {"type": "text", "text": "\n# Left-Wrist View\n"}, | |
| {"type": "image", "image": wrist_image}, | |
| {"type": "text", "text": f"\nGenerate robot actions for the task:\n{instruction} /no_cot"}, | |
| ] | |
| return [ | |
| {"role": "user", "content": content}, | |
| {"role": "assistant", "content": [{"type": "text", "text": "<cot></cot>"}]}, | |
| ] | |
| def _to_cuda_dict(values: Dict[str, Any]) -> Dict[str, Any]: | |
| """Move tensors to CUDA; floating-point tensors take the model dtype (bf16).""" | |
| moved: Dict[str, Any] = {} | |
| for key, value in values.items(): | |
| if value is None: | |
| continue | |
| if torch.is_tensor(value): | |
| dtype = MODEL_DTYPE if value.is_floating_point() else None | |
| moved[key] = value.to("cuda", dtype=dtype) | |
| else: | |
| moved[key] = value | |
| return moved | |
| def build_figure(chunk: np.ndarray): | |
| """Line chart of the Dx/Dy/Dz translation deltas across the predicted chunk. | |
| Dataviz rules applied: 3 series use validated categorical slots 1-3; 2px lines; | |
| >=8px markers with a 2px surface ring; hairline solid gridlines; one axis; a legend | |
| is present; direct labels only on the final point (never every point); text in ink | |
| tokens rather than series colors. Rotation + gripper live in the table view (the | |
| chart's twin), since mixing units on one axis would break the one-axis rule. | |
| """ | |
| import plotly.graph_objects as go | |
| mode = "light" | |
| colors, ink, grid, surface = SERIES[mode], INK[mode], GRID[mode], SURFACE[mode] | |
| labels = ["Δx (m)", "Δy (m)", "Δz (m)"] | |
| steps = list(range(1, chunk.shape[0] + 1)) | |
| last = chunk.shape[0] - 1 | |
| fig = go.Figure() | |
| for i in range(3): | |
| fig.add_trace( | |
| go.Scatter( | |
| x=steps, | |
| y=chunk[:, i].tolist(), | |
| mode="lines+markers+text", | |
| name=labels[i], | |
| text=[f"{chunk[last, i]:+.3f}" if k == last else "" for k in range(chunk.shape[0])], | |
| textposition="top right", | |
| textfont=dict(color=ink, size=10), | |
| line=dict(width=2, color=colors[i]), | |
| marker=dict(size=8, color=colors[i], line=dict(width=2, color=surface)), | |
| hovertemplate="%{y:+.4f} m<extra>" + labels[i] + "</extra>", | |
| ) | |
| ) | |
| gripper_first = chunk[0, 6] | |
| gripper_last = chunk[last, 6] | |
| fig.update_layout( | |
| title=dict( | |
| text=( | |
| f"Predicted {chunk.shape[0]}-step action chunk · gripper cmd " | |
| f"{gripper_first:+.2f} → {gripper_last:+.2f} (− = close, + = open)" | |
| ), | |
| font=dict(color=ink, size=13), | |
| ), | |
| xaxis=dict(title="step", showgrid=False, tickfont=dict(color=ink, size=10), | |
| linecolor=grid, zeroline=False), | |
| yaxis=dict(title="Δposition (m)", showgrid=True, gridcolor=grid, gridwidth=1, | |
| tickfont=dict(color=ink, size=10), linecolor=grid, zeroline=False), | |
| plot_bgcolor=surface, | |
| paper_bgcolor=surface, | |
| font=dict(color=ink), | |
| showlegend=True, | |
| legend=dict(orientation="h", y=1.12, font=dict(color=ink, size=10)), | |
| margin=dict(l=40, r=20, t=60, b=40), | |
| height=340, | |
| uirevision="chunk", | |
| ) | |
| return fig | |
| def chunk_table(chunk: np.ndarray) -> List[List[float]]: | |
| """Full predicted chunk as the chart's table twin (all 7 dims per step).""" | |
| rounded = np.round(chunk.astype(np.float64), 4) | |
| return [[i + 1, *row.tolist()] for i, row in enumerate(rounded)] | |
| def format_csv(chunk: np.ndarray) -> str: | |
| """Serialize the predicted chunk as CSV (step, dx, dy, dz, droll, dpitch, dyaw, gripper).""" | |
| lines = ["step,dx,dy,dz,droll,dpitch,dyaw,gripper"] | |
| for i, row in enumerate(chunk): | |
| lines.append(f"{i + 1}," + ",".join(f"{v:+.4f}" for v in row.tolist())) | |
| return "\n".join(lines) | |
| def predict( | |
| base_view: Image.Image, | |
| wrist_view: Image.Image, | |
| instruction: str, | |
| robot_type: str = "libero_10", | |
| state_ee_x: float = -0.19947, | |
| state_ee_y: float = -0.00631, | |
| state_ee_z: float = 1.18217, | |
| state_aa1: float = 3.13143, | |
| state_aa2: float = 0.00487, | |
| state_aa3: float = -0.00294, | |
| state_finger_l: float = 0.03622, | |
| state_finger_r: float = -0.03623, | |
| seed: int = 42, | |
| chunk_length: int = 10, | |
| ) -> Tuple[Any, List[List[float]], str, str]: | |
| """Predict a LIBERO end-effector action chunk from two camera views + an instruction. | |
| Args: | |
| base_view: base-view camera image (LIBERO agentview orientation). | |
| wrist_view: left-wrist camera image. | |
| instruction: the task, e.g. "Turn on the stove and put the moka pot on it." | |
| robot_type: LIBERO suite the scene comes from (libero_spatial / libero_object / | |
| libero_goal / libero_10). | |
| state_ee_x, state_ee_y, state_ee_z: end-effector position (m). | |
| state_aa1, state_aa2, state_aa3: end-effector orientation, axis-angle (rad). | |
| state_finger_l, state_finger_r: gripper finger joint positions (rad). | |
| seed: RNG seed for the flow-matching sampler. | |
| chunk_length: number of predicted action steps (10 = LIBERO chunk). | |
| Returns: | |
| (plotly figure of the Δx/Δy/Δz deltas, chunk table, CSV text, status line). | |
| """ | |
| t0 = time.time() | |
| instruction = str(instruction or "").strip() | |
| if not instruction: | |
| raise gr.Error("Provide a task instruction.") | |
| if not isinstance(base_view, Image.Image) or not isinstance(wrist_view, Image.Image): | |
| raise gr.Error("Provide both camera views (base + left wrist).") | |
| if robot_type not in LIBERO_ROBOT_TYPES: | |
| raise gr.Error(f"Unsupported robot type {robot_type!r}") | |
| # LIBERO interface feeds the instruction capitalized. | |
| instruction = instruction.capitalize() | |
| raw_state = np.array( | |
| [state_ee_x, state_ee_y, state_ee_z, | |
| state_aa1, state_aa2, state_aa3, | |
| state_finger_l, state_finger_r], | |
| dtype=np.float32, | |
| ) | |
| padded = np.zeros((STATE_LENGTH, STATE_DIM), dtype=np.float32) | |
| padded[-1, :8] = raw_state | |
| state = torch.from_numpy(padded)[None] | |
| chunk_len = int(chunk_length) | |
| if chunk_len < 1 or chunk_len > 20: | |
| chunk_len = 10 | |
| # Manual mask: dims 0-6 (translation/rotation/gripper) — the reference server | |
| # overrides the processor's stats-derived mask, which is all zeros for LIBERO. | |
| mask = torch.zeros((1, chunk_len, ACTION_DIM), dtype=torch.float32) | |
| mask[:, :, :7] = 1.0 | |
| model_inputs = processor.apply_chat_template( | |
| _messages(base_view, wrist_view, instruction), | |
| tokenize=True, | |
| return_dict=True, | |
| return_tensors="pt", | |
| state=state, | |
| robot_type=robot_type, | |
| ) | |
| model_inputs = dict(model_inputs) | |
| model_inputs["state"] = state | |
| model_inputs["action_mask"] = mask | |
| model_inputs["seed"] = int(seed) | |
| model_inputs["num_steps"] = NUM_FLOW_STEPS | |
| model_inputs = _to_cuda_dict(model_inputs) | |
| with torch.no_grad(): | |
| output = model(**model_inputs) | |
| actions = getattr(output, "actions", None) | |
| if actions is None or actions.ndim != 3 or actions.shape[0] != 1: | |
| raise gr.Error(f"Unexpected model output shape {tuple(actions.shape) if actions is not None else None}") | |
| mean = torch.as_tensor(ACTION_MEAN, device=actions.device).view(1, 1, 7) | |
| std = torch.as_tensor(ACTION_STD, device=actions.device).view(1, 1, 7) | |
| decoded = actions[:, :, :7].float() * std + mean | |
| chunk = decoded[0].cpu().numpy().astype(np.float32) | |
| if not np.isfinite(chunk).all(): | |
| raise gr.Error("Decoded action chunk contains NaN or inf.") | |
| fig = build_figure(chunk) | |
| table = chunk_table(chunk) | |
| csv_text = format_csv(chunk) | |
| dt = time.time() - t0 | |
| dt_gpu = dt # measured end-to-end inside the GPU function | |
| status = ( | |
| f"Predicted {chunk.shape[0]} actions in {dt_gpu:.1f}s · suite={robot_type} · seed={int(seed)} " | |
| f"· gripper ends {chunk[-1, 6]:+.2f}" | |
| ) | |
| return fig, table, csv_text, status | |
| EXAMPLES = [ | |
| ["ex_base.png", "ex_wrist.png", | |
| "Turn on the stove and put the moka pot on it.", | |
| "libero_10", -0.19947, -0.00631, 1.18217, 3.13143, 0.00487, -0.00294, 0.03622, -0.03623, 42, 10], | |
| ["ex_bowl_base.png", "ex_bowl_wrist.png", | |
| "Put the black bowl in the bottom drawer of the cabinet and close it.", | |
| "libero_10", -0.2201, -0.0014, 1.1623, 3.1392, -0.0067, -0.1071, 0.0362, -0.0362, 42, 10], | |
| ] | |
| CSS = """ | |
| #col-container { max-width: 1150px; margin: 0 auto; } | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| """ | |
| with gr.Blocks(title="ERVLA — robot action prediction") as demo: | |
| with gr.Column(elem_id="col-container"): | |
| gr.Markdown( | |
| """ | |
| # ERVLA: Embodied Chain-of-Thought for Generalizable Robot Manipulation | |
| Predict a **robot action chunk** from two camera views + a task instruction, with the | |
| LIBERO-Plus-tuned ERVLA checkpoint (`ERVLA/liberoplus` — Qwen3-VL-4B backbone + a DiT | |
| flow-matching action head; 86.9% success on LIBERO-Plus). Give the robot's eye views, | |
| tell it what to do, and read the predicted end-effector deltas (Δxyz, Δrotation, | |
| gripper). No simulator runs here — the action chunk is the output. | |
| [[paper](https://huggingface.co/papers/2606.03784)] · | |
| [code](https://github.com/Oscar-Yang11/ERVLA)] · | |
| [model](https://huggingface.co/ERVLA/liberoplus) | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| base_view = gr.Image(type="pil", label="Base view (LIBERO agentview)") | |
| wrist_view = gr.Image(type="pil", label="Left-wrist view") | |
| instruction = gr.Textbox( | |
| label="Instruction", | |
| placeholder="e.g. Turn on the stove and put the moka pot on it.", | |
| lines=1, | |
| ) | |
| run_btn = gr.Button("Predict actions", variant="primary") | |
| with gr.Column(scale=1): | |
| plot = gr.Plot(label="Predicted Δposition per step") | |
| table_out = gr.Dataframe( | |
| headers=["step", "Δx", "Δy", "Δz", "Δroll", "Δpitch", "Δyaw", "gripper"], | |
| label="Full predicted action chunk", | |
| datatype=["number"] * 8, | |
| ) | |
| status_out = gr.Textbox(label="Status", interactive=False, buttons=["copy"]) | |
| csv_out = gr.Textbox(label="Chunk as CSV", lines=8, interactive=False, | |
| buttons=["copy"]) | |
| with gr.Accordion("Advanced", open=False): | |
| robot_type = gr.Dropdown( | |
| choices=LIBERO_ROBOT_TYPES, value="libero_10", label="LIBERO suite (robot type)" | |
| ) | |
| gr.Markdown( | |
| "End-effector state (LIBERO proprioception, metric units). " | |
| "Defaults come from the first example scene." | |
| ) | |
| state_ee_x = gr.Number(label="EE position x (m)", value=-0.19947, precision=5) | |
| state_ee_y = gr.Number(label="EE position y (m)", value=-0.00631, precision=5) | |
| state_ee_z = gr.Number(label="EE position z (m)", value=1.18217, precision=5) | |
| state_aa1 = gr.Number(label="Orientation axis-angle 1 (rad)", value=3.13143, precision=5) | |
| state_aa2 = gr.Number(label="Orientation axis-angle 2 (rad)", value=0.00487, precision=5) | |
| state_aa3 = gr.Number(label="Orientation axis-angle 3 (rad)", value=-0.00294, precision=5) | |
| state_finger_l = gr.Number(label="Gripper finger left (rad)", value=0.03622, precision=5) | |
| state_finger_r = gr.Number(label="Gripper finger right (rad)", value=-0.03623, precision=5) | |
| seed = gr.Number(label="Seed", value=42, precision=0) | |
| chunk_length = gr.Number(label="Action chunk length", value=10, precision=0) | |
| all_inputs = [ | |
| base_view, wrist_view, instruction, robot_type, | |
| state_ee_x, state_ee_y, state_ee_z, | |
| state_aa1, state_aa2, state_aa3, | |
| state_finger_l, state_finger_r, | |
| seed, chunk_length, | |
| ] | |
| run_btn.click( | |
| predict, | |
| inputs=all_inputs, | |
| outputs=[plot, table_out, csv_out, status_out], | |
| api_name="predict", | |
| ) | |
| gr.Examples( | |
| examples=EXAMPLES, | |
| inputs=all_inputs, | |
| outputs=[plot, table_out, csv_out, status_out], | |
| fn=predict, | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| ) | |
| demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS) |