"""InternVLA-A1.5 VLA Explorer. Interactive demo for InternRobotics/InternVLA-A1.5-base โ€” a Vision-Language-Action robot policy built on a Qwen3.5-2B backbone with a flow-matching action expert. Given a single camera frame + a natural-language task instruction, the model predicts a 50-step continuous action chunk (the robot's planned motion) via flow matching. We plot that predicted action trajectory, dimension by dimension. The video-foresight (WAN2.2) branch is discarded at inference (action_loss_only=True), so no video weights are needed. """ import os import sys import shutil import site import importlib from pathlib import Path # --------------------------------------------------------------------------- # 1) Patch the installed `transformers` package with the InternVLA-A1.5 custom # Qwen3.5 modeling file BEFORE any transformers import happens. This mirrors # step 7 of the repo's installation tutorial (copying transformers_replace/ # into site-packages/transformers/). Pure filesystem op โ€” no CUDA touched. # --------------------------------------------------------------------------- def _patch_transformers() -> None: import transformers # noqa: F401 (only to locate its dir) tf_dir = Path(transformers.__file__).parent patch_root = Path(__file__).parent / "transformers_patch" / "models" if not patch_root.exists(): print(f"[patch] no patch dir at {patch_root}, skipping", flush=True) return dst_models = tf_dir / "models" for model_dir in patch_root.iterdir(): if not model_dir.is_dir(): continue for src_file in model_dir.rglob("*.py"): rel = src_file.relative_to(patch_root) dst = dst_models / rel dst.parent.mkdir(parents=True, exist_ok=True) shutil.copy2(src_file, dst) print(f"[patch] {src_file.name} -> {dst}", flush=True) # Make sure a fresh import picks up the copied module. importlib.invalidate_caches() _patch_transformers() import numpy as np import torch import gradio as gr import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import spaces # Make the bundled `lerobot` importable (it sits next to app.py). sys.path.insert(0, str(Path(__file__).parent)) MODEL_ID = "InternRobotics/InternVLA-A1.5-base" VLM_ID = "Qwen/Qwen3.5-2B" DEVICE = "cuda" CHUNK_SIZE = 50 # flow-matching action horizon RESIZE = 224 # google_robot embodiment: single RGB camera + single arm (7-DoF + gripper). # This is the SimplerEnv Google-Robot layout the base checkpoint was evaluated on. ROBOT_TYPE = "google_robot" STATE_DIM = 8 # 7 arm dims + 1 gripper (raw google_robot state) # --------------------------------------------------------------------------- # 2) Lazy globals filled on first GPU call. # --------------------------------------------------------------------------- _POLICY = None _PROCESSOR = None _RESIZE_FN = None _OBS_IMAGES = "observation.images" _OBS_STATE = "observation.state" def _to_chw_float01(image: np.ndarray) -> torch.Tensor: if image.ndim != 3: raise ValueError(f"Expected HWC image, got shape={image.shape}") t = torch.from_numpy(np.array(image, copy=True)) if t.dtype == torch.uint8: t = t.float() / 255.0 else: t = t.float() if torch.max(t) > 1.0: t = t / 255.0 return t.permute(2, 0, 1).contiguous() def _load_models(): """Construct the policy + processor once (module-scope caching).""" global _POLICY, _PROCESSOR, _RESIZE_FN if _POLICY is not None: return from lerobot.configs.policies import PreTrainedConfig from lerobot.policies.internvla_a1_5 import InternVLAA15Config, InternVLAA15Policy from lerobot.policies.internvla_a1_5.transform_internvla_a1_5 import ( InternVLAA15ChatProcessorTransformFn, ) from lerobot.transforms.core import ResizeImagesWithPadFn config = PreTrainedConfig.from_pretrained(MODEL_ID) assert isinstance(config, InternVLAA15Config), type(config) config.vlm_model_name_or_path = VLM_ID config.action_loss_only = True # skip the WAN video branch config.inference_backend = "standard" # pure-torch path, no custom kernels config.device = DEVICE policy = InternVLAA15Policy.from_pretrained(MODEL_ID, config=config) policy.to(device=DEVICE) policy.to(dtype=torch.bfloat16) policy.eval() processor = InternVLAA15ChatProcessorTransformFn( pretrained_model_name_or_path=VLM_ID, max_length=650, tokenize_state=config.tokenize_state, max_state_dim=config.max_state_dim, use_fast_action_tokens=False, mode="eval", action_mode="joint", ) _RESIZE_FN = ResizeImagesWithPadFn( height=RESIZE, width=RESIZE, mapping={ f"{_OBS_IMAGES}.image0": f"{_OBS_IMAGES}.image0", f"{_OBS_IMAGES}.image1": f"{_OBS_IMAGES}.image1", f"{_OBS_IMAGES}.image2": f"{_OBS_IMAGES}.image2", }, ) _POLICY = policy _PROCESSOR = processor def _build_inputs(image: np.ndarray, instruction: str, state: np.ndarray): """Reproduce the LIBERO backend's build_base_sample -> processor pipeline for a single frame + single-arm zero/user state (google_robot layout).""" white = np.full_like(np.asarray(image), 255) sample = { _OBS_STATE: torch.from_numpy(state.astype(np.float32)), "task": instruction.strip(), f"{_OBS_IMAGES}.image0": _to_chw_float01(np.asarray(image)), f"{_OBS_IMAGES}.image1": _to_chw_float01(white), f"{_OBS_IMAGES}.image2": _to_chw_float01(white), } sample = _RESIZE_FN(sample) # google_robot uses a single camera -> only image0 is a real observation. sample[f"{_OBS_IMAGES}.image0_mask"] = True sample[f"{_OBS_IMAGES}.image1_mask"] = False sample[f"{_OBS_IMAGES}.image2_mask"] = False sample = _PROCESSOR(sample) inputs = {} for key, value in sample.items(): if key == "task": inputs[key] = [value] continue if isinstance(value, bool): continue if not isinstance(value, torch.Tensor): continue inputs[key] = value.unsqueeze(0).to(DEVICE) return inputs def _plot_trajectory(chunk: np.ndarray, instruction: str) -> "plt.Figure": """chunk: [T, D] predicted (normalized) action trajectory.""" T, D = chunk.shape # Show the leading, meaningful dims (7 arm + 1 gripper for google_robot). show = min(D, 8) ncols = 2 nrows = (show + ncols - 1) // ncols fig, axs = plt.subplots(nrows, ncols, figsize=(11, 2.2 * nrows)) axs = np.atleast_1d(axs).ravel() x = np.arange(T) dim_names = [f"arm joint {i+1}" for i in range(7)] + ["gripper"] for d in range(len(axs)): ax = axs[d] if d >= show: ax.axis("off") continue ax.plot(x, chunk[:, d], color="#c1121f", linewidth=1.8) ax.fill_between(x, chunk[:, d], alpha=0.12, color="#c1121f") name = dim_names[d] if d < len(dim_names) else f"dim {d+1}" ax.set_title(name, fontsize=10) ax.set_xlabel("step") ax.grid(True, linestyle="--", alpha=0.5) fig.suptitle( f'Predicted {T}-step action chunk โ€” "{instruction[:60]}"', fontsize=12, ) fig.tight_layout(rect=[0, 0, 1, 0.97]) return fig @spaces.GPU(duration=120) def predict(image, instruction, gripper_state, seed): if image is None: raise gr.Error("Please provide a camera frame.") if not instruction or not instruction.strip(): raise gr.Error("Please enter a task instruction.") _load_models() torch.manual_seed(int(seed)) image = np.asarray(image) if image.ndim == 2: image = np.stack([image] * 3, axis=-1) if image.shape[-1] == 4: image = image[..., :3] # Neutral robot proprioceptive state (7 arm dims zeroed + user gripper). state = np.zeros(STATE_DIM, dtype=np.float32) state[-1] = float(gripper_state) inputs = _build_inputs(image, instruction, state) with torch.no_grad(), torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16): chunk = _POLICY.predict_action_chunk(inputs) # [1, T, D] chunk = chunk[0].detach().float().cpu().numpy() # [T, D] fig = _plot_trajectory(chunk, instruction) summary = ( f"Predicted action chunk shape: {chunk.shape[0]} steps x {chunk.shape[1]} dims\n" f"Embodiment: {ROBOT_TYPE} (single RGB camera, 7-DoF arm + gripper)\n" f"Values are the model's raw flow-matching action output (normalized units).\n" f"Per-step L2 norm range: " f"{np.linalg.norm(chunk, axis=1).min():.3f} โ€“ {np.linalg.norm(chunk, axis=1).max():.3f}" ) return fig, summary TITLE = "# ๐Ÿค– InternVLA-A1.5 ยท VLA Action Explorer" DESCRIPTION = """ Explore [**InternVLA-A1.5-base**](https://huggingface.co/InternRobotics/InternVLA-A1.5-base), a Vision-Language-Action robot policy (Qwen3.5-2B backbone + flow-matching action expert) from the paper *"InternVLA-A1.5: Unifying Understanding, Latent Foresight, and Action for Compositional Generalization."* Give the model a **camera frame** and a **task instruction**; it predicts the next **50-step continuous action chunk** โ€” the robot's planned motion โ€” which we plot per dimension. This is the model's flow-matching action output (normalized units) for the `google_robot` single-arm embodiment; the video-foresight branch is discarded at inference. """ with gr.Blocks(theme=gr.themes.Citrus()) as demo: gr.Markdown(TITLE) gr.Markdown(DESCRIPTION) with gr.Row(): with gr.Column(): image_in = gr.Image(label="Camera frame", type="numpy", height=320) instruction_in = gr.Textbox( label="Task instruction", placeholder="e.g. pick up the object and place it on the plate", value="pick up the object", ) run_btn = gr.Button("Predict action chunk", variant="primary") with gr.Accordion("Advanced options", open=False): gripper_in = gr.Slider( 0.0, 1.0, value=0.0, step=0.05, label="Initial gripper state (0 = open, 1 = closed)", ) seed_in = gr.Slider( 0, 2**31 - 1, value=0, step=1, label="Seed (flow-matching noise)", ) with gr.Column(): plot_out = gr.Plot(label="Predicted action trajectory") summary_out = gr.Textbox(label="Output summary", lines=5) run_btn.click( predict, inputs=[image_in, instruction_in, gripper_in, seed_in], outputs=[plot_out, summary_out], ) gr.Examples( examples=[ ["examples/rubber_ducks.jpg", "pick up the yellow duck", 0.0, 0], ["examples/cake.jpg", "move the plate to the left", 0.0, 0], ["examples/noodles.jpg", "grasp the bowl and lift it", 0.0, 0], ], inputs=[image_in, instruction_in, gripper_in, seed_in], outputs=[plot_out, summary_out], fn=predict, cache_examples=True, cache_mode="lazy", ) if __name__ == "__main__": demo.queue().launch()