Spaces:

hugging-apps
/

vision-predictor-action

Running on Zero

Size @spaces.GPU duration to measured cold max (0.80s) + margin

cfdc831 verified 17 days ago

11 kB

"""Vision-Predictor-Action (VPA) — interactive demo for the LIBERO-10 policy.
Paper:  https://huggingface.co/papers/2610.15556
Code:   https://github.com/ylee-sling/Vision-Predictor-Action   (GPL-3.0-only)
Weight: https://huggingface.co/ylee-sling/vpa-libero10          (GPL-3.0-only)
Give the policy the two camera views it was trained on (third-person + wrist), the same two
views for the goal, and a language instruction; it returns the action chunk it plans for that
frame, together with the diagnostics of its latent control loop.
``perception`` / ``selector`` / ``predictor`` / ``solver`` / ``pipeline`` are the reference
implementation's inference path, vendored unmodified, so the numbers below are the policy's own.
"""

import spaces  # noqa: F401  -- MUST come before torch on ZeroGPU

import json
import os
import time

import gradio as gr
import numpy as np
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from preprocess import to_views
from vpa_loader import load_pipeline

CKPT_REPO = "ylee-sling/vpa-libero10"
CKPT_FILE = "vpa_libero10_2cam.pt"
PRIMITIVES = ("reach", "grasp", "carry", "release")
EXAMPLES_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "examples")
ROTATE_180 = False   # example frames are shipped already upright

# Proprioception is the one input a still image cannot supply: LIBERO's ee_pos / ee_ori /
# gripper_states / joint_states come from the simulator. Everything else — both camera frames,
# both goal frames, the CLIP text embedding, the vision latents, the primitive selector, the
# JEPA predictor and the flow-matching solver — is the real trained path.
ZERO_PROPRIO = torch.zeros(1, 15)

# ------------------------------------------------------------------ model, at module scope
PIPE = load_pipeline(hf_hub_download(CKPT_REPO, CKPT_FILE), device="cuda")
TAU_STAR = float(PIPE.tracker.threshold)
HORIZON = int(PIPE.solver.horizon)
ACTION_DIM = int(PIPE.solver.action_dim)
K_STEPS = int(PIPE.solver.num_integration_steps)


def _load_examples():
    """Example rows for ``gr.Examples``: the 5 policy inputs, then seed and rotate.
    Returns ``(rows, labels)``; every cell is concrete (Gradio cannot render a ``None`` cell).
    """
    meta_path = os.path.join(EXAMPLES_DIR, "examples.json")
    if not os.path.exists(meta_path):
        return [], []
    rows, labels = [], []
    for m in json.load(open(meta_path)):
        d = os.path.join(EXAMPLES_DIR, m["slug"])
        rows.append([
            Image.open(os.path.join(d, "agentview_obs.png")),
            Image.open(os.path.join(d, "wrist_obs.png")),
            Image.open(os.path.join(d, "agentview_goal.png")),
            Image.open(os.path.join(d, "wrist_goal.png")),
            m["instruction"],
            0,          # seed
            False,      # rotate_180
        ])
        labels.append(m.get("label", m["slug"]))
    return rows, labels


# ------------------------------------------------------------------ inference
@spaces.GPU(duration=5)
def plan(agentview, wrist, goal_agentview, goal_wrist, instruction,
         seed=0, rotate_180=ROTATE_180, progress=gr.Progress()):
    """One VPA decision step: 2 camera views + 2 goal views + instruction -> action chunk."""
    t0 = time.perf_counter()

    def prep(img, name):
        if img is None:
            raise gr.Error(f"Please supply the {name} image.")
        return np.asarray(img.convert("RGB"))

    observation = to_views(
        [prep(agentview, "third-person camera"), prep(wrist, "wrist camera")],
        rotate=rotate_180,
    )[None].to("cuda")                                            # [1, V, C, H, W]
    milestone = to_views(
        [prep(goal_agentview, "third-person goal"), prep(goal_wrist, "wrist goal")],
        rotate=rotate_180,
    )[None, None].to("cuda")                                      # [1, M, V, C, H, W]

    progress(0.10, desc="Encoding instruction and goal with E_psi")
    PIPE.reset(instruction or "complete the task", milestone)

    progress(0.45, desc="Decision step: E_psi -> pi_phi -> P_omega -> v_theta x K")
    generator = torch.Generator(device="cpu").manual_seed(int(seed))
    noise = torch.randn(1, HORIZON, ACTION_DIM, generator=generator).to("cuda")
    out = PIPE.step(observation, ZERO_PROPRIO.to("cuda"), noise=noise)

    chunk = out.action_chunk[0].float().cpu().numpy()             # [H, 7]
    executed = int(out.executed_horizon[0])
    primitive = int(out.primitive[0])
    distance = float(out.milestone_distance[0])
    sigma = float(out.sigma[0])
    sigma_bar = float(out.sigma_bar[0])
    pointer = int(out.milestone_pointer[0])
    elapsed = time.perf_counter() - t0

    reached = distance < TAU_STAR
    chart = {
        f"step {i + 1}": {
            "dx": round(float(c[0]), 4), "dy": round(float(c[1]), 4), "dz": round(float(c[2]), 4),
            "drx": round(float(c[3]), 4), "dry": round(float(c[4]), 4),
            "drz": round(float(c[5]), 4), "gripper": round(float(c[6]), 4),
        }
        for i, c in enumerate(chunk)
    }

    name = PRIMITIVES[primitive] if primitive < len(PRIMITIVES) else "?"
    summary = f"""### Decision-step diagnostics
| quantity | value |
|---|---|
| **Primitive** `u_t` from `π_φ^h` | `{primitive}` — **{name}** |
| **Milestone pointer** `m_t` | `{pointer}` (M = 1: the goal is a single final frame) |
| **Milestone distance** `‖z_t − z_g‖₂` | `{distance:.3f}` &nbsp;vs&nbsp; `τ* = {TAU_STAR:.3f}` → **{"reached" if reached else "not reached"}** |
| **Uncertainty** `σ_t` | `{sigma:.4f}` (filtered `σ̄_t = {sigma_bar:.4f}`) |
| **Executed horizon** `H_t` | `{executed}` of `{HORIZON}` planned commands |
| **Sequential evaluations** | `{out.num_sequential_evaluations}` = K + 3, K = `{K_STEPS}` |
| **Wall clock** | `{elapsed:.2f} s` |
The policy planned **{HORIZON}** commands and would execute the first **{executed}** open-loop
before planning again. Each row of the JSON is one command
`[dx, dy, dz, drx, dry, drz, gripper]`: a delta end-effector pose plus the gripper command.
"""

    return chart, summary, f"Planned {HORIZON} commands in **{elapsed:.2f} s**."


CSS = """
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(title="Vision-Predictor-Action") as demo:
    gr.Markdown(
        """
# 👁 Vision-Predictor-Action
**Towards Efficient Embodied AI via Neuro-Symbolic JEPA Predictors** —
[paper](https://huggingface.co/papers/2610.15556) · [code](https://github.com/ylee-sling/Vision-Predictor-Action) · [weights](https://huggingface.co/ylee-sling/vpa-libero10)
VPA replaces the autoregressive language-model decoder of a vision-language-action model with a
**latent control loop**: a Siamese vision encoder `E_ψ`, a neuro-symbolic primitive selector
`π_φ^h`, a JEPA predictor `P_ω` of the next latent, and a flow-matching solver `v_θ` that emits a
whole action chunk in K steps — so a decision step costs **K + 3 sequential network
evaluations regardless of the chunk length**, not one decoding pass per action.
This Space runs the released **LIBERO-10, two-camera** checkpoint on your inputs and shows both
the action chunk it plans and the internals of that control loop.
> **What is real here, and what is not.** The camera frames, the goal frames, the CLIP text
> embedding, the vision latents, the primitive selector, the JEPA predictor and the flow-matching
> solver are all the real trained path. Two things are not: the **proprioceptive state** is fed as
> zeros (LIBERO's end-effector pose, gripper and joint states cannot be recovered from a still
> image), and a **closed-loop rollout** needs the simulator, which is not available here. The
> released policy scores **40.0 %** (80/200 episodes) on LIBERO-10. Read the actions below as the
> policy's open-loop plan for the frame you gave it.
"""
    )

    with gr.Row():
        with gr.Column():
            gr.Markdown("#### Current observation `I_t`")
            agentview = gr.Image(label="Third-person camera · agentview_rgb", type="pil",
                                 image_mode="RGB", height=220)
            wrist = gr.Image(label="Wrist camera · eye_in_hand_rgb", type="pil",
                             image_mode="RGB", height=220)
        with gr.Column():
            gr.Markdown("#### Goal `I_g` — the demonstration's final frame")
            goal_agentview = gr.Image(label="Third-person goal", type="pil",
                                      image_mode="RGB", height=220)
            goal_wrist = gr.Image(label="Wrist goal", type="pil", image_mode="RGB", height=220)

    instruction = gr.Textbox(
        label="Language instruction", lines=1,
        value="turn on the stove and put the moka pot on it",
        info="The checkpoint was trained on the ten LIBERO-10 instructions; the examples below use them verbatim.",
    )

    run = gr.Button("Plan one action chunk", variant="primary")

    with gr.Row():
        chunk = gr.JSON(label=f"Predicted action chunk — [{HORIZON}, 7]")
        summary = gr.Markdown("Press **Plan one action chunk** to run the policy.")
    timing = gr.Markdown()

    with gr.Accordion("Advanced options", open=False):
        seed = gr.Slider(0, 9999, value=0, step=1, label="Flow-matching noise seed ξ",
                         info="v_θ starts from Gaussian noise A⁽⁰⁾ = ξ; the seed picks that draw.")
        rotate = gr.Checkbox(
            value=ROTATE_180,
            label="Rotate inputs 180° (raw, as-stored LIBERO frames)",
            info="LIBERO saves every camera — including the wrist camera — upside down, and the "
                 "reference dataset rights them. The example frames are shipped already upright, "
                 "so leave this off unless you paste in raw LIBERO frames.",
        )

    example_rows, example_labels = _load_examples()
    if example_rows:
        gr.Examples(
            examples=example_rows,
            example_labels=example_labels,
            inputs=[agentview, wrist, goal_agentview, goal_wrist, instruction, seed, rotate],
            outputs=[chunk, summary, timing],
            fn=plan,
            cache_examples=True,
            cache_mode="lazy",
            examples_per_page=5,
        )

    gr.Markdown(
        """
---
Example frames are the first (and, for one row, the last) frames of LIBERO-10 demonstrations from
[`yifengzhu-hf/LIBERO-datasets`](https://huggingface.co/datasets/yifengzhu-hf/LIBERO-datasets)
(Apache-2.0), mirrored here for convenience; the
[LIBERO](https://github.com/Lifelong-Robot-Learning/LIBERO) terms also apply. Model and vendored
code: **GPL-3.0-only**. Preprocessing reproduces `LiberoHDF5Dataset._to_tensor` exactly.
"""
    )

    run.click(
        plan,
        [agentview, wrist, goal_agentview, goal_wrist, instruction, seed, rotate],
        [chunk, summary, timing],
    )

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), css=CSS)

	"""Vision-Predictor-Action (VPA) — interactive demo for the LIBERO-10 policy.

	Paper: https://huggingface.co/papers/2610.15556
	Code: https://github.com/ylee-sling/Vision-Predictor-Action (GPL-3.0-only)
	Weight: https://huggingface.co/ylee-sling/vpa-libero10 (GPL-3.0-only)

	Give the policy the two camera views it was trained on (third-person + wrist), the same two
	views for the goal, and a language instruction; it returns the action chunk it plans for that
	frame, together with the diagnostics of its latent control loop.

	``perception`` / ``selector`` / ``predictor`` / ``solver`` / ``pipeline`` are the reference
	implementation's inference path, vendored unmodified, so the numbers below are the policy's own.
	"""

	import spaces # noqa: F401 -- MUST come before torch on ZeroGPU

	import json
	import os
	import time

	import gradio as gr
	import numpy as np
	import torch
	from huggingface_hub import hf_hub_download
	from PIL import Image

	from preprocess import to_views
	from vpa_loader import load_pipeline

	CKPT_REPO = "ylee-sling/vpa-libero10"
	CKPT_FILE = "vpa_libero10_2cam.pt"
	PRIMITIVES = ("reach", "grasp", "carry", "release")
	EXAMPLES_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "examples")
	ROTATE_180 = False # example frames are shipped already upright

	# Proprioception is the one input a still image cannot supply: LIBERO's ee_pos / ee_ori /
	# gripper_states / joint_states come from the simulator. Everything else — both camera frames,
	# both goal frames, the CLIP text embedding, the vision latents, the primitive selector, the
	# JEPA predictor and the flow-matching solver — is the real trained path.
	ZERO_PROPRIO = torch.zeros(1, 15)

	# ------------------------------------------------------------------ model, at module scope
	PIPE = load_pipeline(hf_hub_download(CKPT_REPO, CKPT_FILE), device="cuda")
	TAU_STAR = float(PIPE.tracker.threshold)
	HORIZON = int(PIPE.solver.horizon)
	ACTION_DIM = int(PIPE.solver.action_dim)
	K_STEPS = int(PIPE.solver.num_integration_steps)


	def _load_examples():
	"""Example rows for ``gr.Examples``: the 5 policy inputs, then seed and rotate.

	Returns ``(rows, labels)``; every cell is concrete (Gradio cannot render a ``None`` cell).
	"""
	meta_path = os.path.join(EXAMPLES_DIR, "examples.json")
	if not os.path.exists(meta_path):
	return [], []
	rows, labels = [], []
	for m in json.load(open(meta_path)):
	d = os.path.join(EXAMPLES_DIR, m["slug"])
	rows.append([
	Image.open(os.path.join(d, "agentview_obs.png")),
	Image.open(os.path.join(d, "wrist_obs.png")),
	Image.open(os.path.join(d, "agentview_goal.png")),
	Image.open(os.path.join(d, "wrist_goal.png")),
	m["instruction"],
	0, # seed
	False, # rotate_180
	])
	labels.append(m.get("label", m["slug"]))
	return rows, labels


	# ------------------------------------------------------------------ inference
	@spaces.GPU(duration=5)
	def plan(agentview, wrist, goal_agentview, goal_wrist, instruction,
	seed=0, rotate_180=ROTATE_180, progress=gr.Progress()):
	"""One VPA decision step: 2 camera views + 2 goal views + instruction -> action chunk."""
	t0 = time.perf_counter()

	def prep(img, name):
	if img is None:
	raise gr.Error(f"Please supply the {name} image.")
	return np.asarray(img.convert("RGB"))

	observation = to_views(
	[prep(agentview, "third-person camera"), prep(wrist, "wrist camera")],
	rotate=rotate_180,
	)[None].to("cuda") # [1, V, C, H, W]
	milestone = to_views(
	[prep(goal_agentview, "third-person goal"), prep(goal_wrist, "wrist goal")],
	rotate=rotate_180,
	)[None, None].to("cuda") # [1, M, V, C, H, W]

	progress(0.10, desc="Encoding instruction and goal with E_psi")
	PIPE.reset(instruction or "complete the task", milestone)

	progress(0.45, desc="Decision step: E_psi -> pi_phi -> P_omega -> v_theta x K")
	generator = torch.Generator(device="cpu").manual_seed(int(seed))
	noise = torch.randn(1, HORIZON, ACTION_DIM, generator=generator).to("cuda")
	out = PIPE.step(observation, ZERO_PROPRIO.to("cuda"), noise=noise)

	chunk = out.action_chunk[0].float().cpu().numpy() # [H, 7]
	executed = int(out.executed_horizon[0])
	primitive = int(out.primitive[0])
	distance = float(out.milestone_distance[0])
	sigma = float(out.sigma[0])
	sigma_bar = float(out.sigma_bar[0])
	pointer = int(out.milestone_pointer[0])
	elapsed = time.perf_counter() - t0

	reached = distance < TAU_STAR
	chart = {
	f"step {i + 1}": {
	"dx": round(float(c[0]), 4), "dy": round(float(c[1]), 4), "dz": round(float(c[2]), 4),
	"drx": round(float(c[3]), 4), "dry": round(float(c[4]), 4),
	"drz": round(float(c[5]), 4), "gripper": round(float(c[6]), 4),
	}
	for i, c in enumerate(chunk)
	}

	name = PRIMITIVES[primitive] if primitive < len(PRIMITIVES) else "?"
	summary = f"""### Decision-step diagnostics

	\| quantity \| value \|
	\|---\|---\|
	\| Primitive `u_t` from `π_φ^h` \| `{primitive}` — {name} \|
	\| Milestone pointer `m_t` \| `{pointer}` (M = 1: the goal is a single final frame) \|
	\| Milestone distance `‖z_t − z_g‖₂` \| `{distance:.3f}`  vs  `τ* = {TAU_STAR:.3f}` → {"reached" if reached else "not reached"} \|
	\| Uncertainty `σ_t` \| `{sigma:.4f}` (filtered `σ̄_t = {sigma_bar:.4f}`) \|
	\| Executed horizon `H_t` \| `{executed}` of `{HORIZON}` planned commands \|
	\| Sequential evaluations \| `{out.num_sequential_evaluations}` = K + 3, K = `{K_STEPS}` \|
	\| Wall clock \| `{elapsed:.2f} s` \|

	The policy planned {HORIZON} commands and would execute the first {executed} open-loop
	before planning again. Each row of the JSON is one command
	`[dx, dy, dz, drx, dry, drz, gripper]`: a delta end-effector pose plus the gripper command.
	"""

	return chart, summary, f"Planned {HORIZON} commands in {elapsed:.2f} s."


	CSS = """
	.dark .gradio-container { color: var(--body-text-color); }
	"""

	with gr.Blocks(title="Vision-Predictor-Action") as demo:
	gr.Markdown(
	"""
	# 👁 Vision-Predictor-Action

	Towards Efficient Embodied AI via Neuro-Symbolic JEPA Predictors —
	[paper](https://huggingface.co/papers/2610.15556) · [code](https://github.com/ylee-sling/Vision-Predictor-Action) · [weights](https://huggingface.co/ylee-sling/vpa-libero10)

	VPA replaces the autoregressive language-model decoder of a vision-language-action model with a
	latent control loop: a Siamese vision encoder `E_ψ`, a neuro-symbolic primitive selector
	`π_φ^h`, a JEPA predictor `P_ω` of the next latent, and a flow-matching solver `v_θ` that emits a
	whole action chunk in K steps — so a decision step costs **K + 3 sequential network
	evaluations regardless of the chunk length**, not one decoding pass per action.

	This Space runs the released LIBERO-10, two-camera checkpoint on your inputs and shows both
	the action chunk it plans and the internals of that control loop.

	> What is real here, and what is not. The camera frames, the goal frames, the CLIP text
	> embedding, the vision latents, the primitive selector, the JEPA predictor and the flow-matching
	> solver are all the real trained path. Two things are not: the proprioceptive state is fed as
	> zeros (LIBERO's end-effector pose, gripper and joint states cannot be recovered from a still
	> image), and a closed-loop rollout needs the simulator, which is not available here. The
	> released policy scores 40.0 % (80/200 episodes) on LIBERO-10. Read the actions below as the
	> policy's open-loop plan for the frame you gave it.
	"""
	)

	with gr.Row():
	with gr.Column():
	gr.Markdown("#### Current observation `I_t`")
	agentview = gr.Image(label="Third-person camera · agentview_rgb", type="pil",
	image_mode="RGB", height=220)
	wrist = gr.Image(label="Wrist camera · eye_in_hand_rgb", type="pil",
	image_mode="RGB", height=220)
	with gr.Column():
	gr.Markdown("#### Goal `I_g` — the demonstration's final frame")
	goal_agentview = gr.Image(label="Third-person goal", type="pil",
	image_mode="RGB", height=220)
	goal_wrist = gr.Image(label="Wrist goal", type="pil", image_mode="RGB", height=220)

	instruction = gr.Textbox(
	label="Language instruction", lines=1,
	value="turn on the stove and put the moka pot on it",
	info="The checkpoint was trained on the ten LIBERO-10 instructions; the examples below use them verbatim.",
	)

	run = gr.Button("Plan one action chunk", variant="primary")

	with gr.Row():
	chunk = gr.JSON(label=f"Predicted action chunk — [{HORIZON}, 7]")
	summary = gr.Markdown("Press Plan one action chunk to run the policy.")
	timing = gr.Markdown()

	with gr.Accordion("Advanced options", open=False):
	seed = gr.Slider(0, 9999, value=0, step=1, label="Flow-matching noise seed ξ",
	info="v_θ starts from Gaussian noise A⁽⁰⁾ = ξ; the seed picks that draw.")
	rotate = gr.Checkbox(
	value=ROTATE_180,
	label="Rotate inputs 180° (raw, as-stored LIBERO frames)",
	info="LIBERO saves every camera — including the wrist camera — upside down, and the "
	"reference dataset rights them. The example frames are shipped already upright, "
	"so leave this off unless you paste in raw LIBERO frames.",
	)

	example_rows, example_labels = _load_examples()
	if example_rows:
	gr.Examples(
	examples=example_rows,
	example_labels=example_labels,
	inputs=[agentview, wrist, goal_agentview, goal_wrist, instruction, seed, rotate],
	outputs=[chunk, summary, timing],
	fn=plan,
	cache_examples=True,
	cache_mode="lazy",
	examples_per_page=5,
	)

	gr.Markdown(
	"""
	---
	Example frames are the first (and, for one row, the last) frames of LIBERO-10 demonstrations from
	[`yifengzhu-hf/LIBERO-datasets`](https://huggingface.co/datasets/yifengzhu-hf/LIBERO-datasets)
	(Apache-2.0), mirrored here for convenience; the
	[LIBERO](https://github.com/Lifelong-Robot-Learning/LIBERO) terms also apply. Model and vendored
	code: GPL-3.0-only. Preprocessing reproduces `LiberoHDF5Dataset._to_tensor` exactly.
	"""
	)

	run.click(
	plan,
	[agentview, wrist, goal_agentview, goal_wrist, instruction, seed, rotate],
	[chunk, summary, timing],
	)

	if __name__ == "__main__":
	demo.launch(theme=gr.themes.Citrus(), css=CSS)