Views
No views yet
AutoModel compatible implementation of AR-VLA that can be used via Transformers.roma, timm, flash-attn.uv:uv:wget -qO- https://github.com/astral-sh/uv/releases/download/0.7.5/uv-installer.sh | sh1uv venv python 3.10.12
2source .venv/bin/activate
3uv pip install --torch-backend=cu126 roma==1.5.0 numpy==2.2.4 torch==2.6.0 torchvision==0.21.0 transformers==4.47.0 timm==1.0.15
4uv pip install --no-build-isolation setuptools psutil flash-attn==2.7.31# Initialize/Update the persistent VLM context
2model.refresh_test_time_vlm()
3
4# Predict actions based on incoming states without recomputing the VLM backbone
5action_1 = model.next_test_time_action(state_1)
6action_2 = model.next_test_time_action(state_2)
7action_3 = model.next_test_time_action(state_3)
8
9# Refresh the VLM context when a new observation or instruction is received
10model.refresh_test_time_vlm()
11action_4 = model.next_test_time_action(state_4)
12action_5 = model.next_test_time_action(state_5)1import numpy as np
2import torch
3from PIL import Image
4from transformers import AutoModel, AutoProcessor
5
6model_id = "INSAIT-Institute/arvla-bridge"
7
8model = AutoModel.from_pretrained(model_id, trust_remote_code=True)
9processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
10model = model.to(device="cuda").eval()
11
12image = Image.open("path/to/main_image.png").convert("RGB")
13
14batch = processor.preprocess_inputs(
15 chat=["pick up the cup", ""],
16 images={"main": [image]},
17 ee_pose_translation=np.zeros((1, 1, 3), dtype=np.float32),
18 ee_pose_rotation=np.array([[[0.0, 0.0, 0.0, 1.0]]], dtype=np.float32),
19 gripper=np.zeros((1, 1), dtype=np.float32),
20 joints=np.zeros((1, 1, 7), dtype=np.float32),
21 dataset_name=np.array(["bridge"]),
22 inference_mode=True,
23)
24
25with torch.inference_mode():
26 model.reset_test_time_cache()
27 model.refresh_test_time_vlm(
28 input_ids=batch["input_ids"].to("cuda"),
29 attention_mask=torch.ones_like(batch["input_ids"], dtype=torch.bool).to("cuda"),
30 images={k: v.to("cuda") for k, v in batch["images"].items()},
31 ee_pose_translation=batch["ee_pose_translation"].to("cuda"),
32 ee_pose_rotation=batch["ee_pose_rotation"].to("cuda"),
33 gripper=batch["gripper"].unsqueeze(-1).to("cuda"),
34 joints=batch["joints"].to("cuda"),
35 control_tokens_ids=batch["control_tokens_ids"],
36 )
37 action = model.next_test_time_action(
38 input_ids=batch["input_ids"].to("cuda"),
39 ee_pose_translation=batch["ee_pose_translation"].to("cuda"),
40 ee_pose_rotation=batch["ee_pose_rotation"].to("cuda"),
41 gripper=batch["gripper"].unsqueeze(-1).to("cuda"),
42 joints=batch["joints"].to("cuda"),
43 control_tokens_ids=batch["control_tokens_ids"],
44 )
45
46print(action.translation.shape, action.rotation.shape, action.gripper.shape)you2who/ar-vla-bridge