openvla_stretch_lora

OpenVLA 7B - Fine-Tuned LoRA Adapter for Hello Robot Stretch

This repository contains the PEFT (LoRA) adapter weights for openvla/openvla-7b fine-tuned for pick-and-place manipulation tasks on the Hello Robot Stretch.

Model Overview

• Base Model: openvla/openvla-7b
• Adapter Type: LoRA (peft)
• Task: Object manipulation / Picking up target items (e.g., water bottle)
• Target Modules: Query (q_proj), Key (k_proj), Value (v_proj), Output (o_proj)
• LoRA Rank ($r$): 16
• LoRA Alpha ($\alpha$): 32
• Action Output: 7-dimensional continuous robot action space

---

Action Space & Discretization

The model predicts actions in a 7-dimensional delta action space:

1. delta_x (End-effector / Base translation along X)
2. delta_y (End-effector / Base translation along Y)
3. delta_z (Lift position change)
4. delta_roll (Wrist roll angle change)
5. delta_pitch (Wrist pitch angle change)
6. delta_yaw (Wrist yaw angle change)
7. gripper_pos (Gripper open/close target position)

Continuous actions are normalized to the range [-1.0, 1.0] and discretized into 256 uniform bins mapped to OpenVLA action token IDs [31744 - 31999].

---

Training Hyperparameters

• Epochs: 5
• Batch Size: 4 (per device)
• Gradient Accumulation Steps: 2 (Effective Batch Size: 8)
• Learning Rate: 5e-4 with cosine decay and warmup
• Optimizer: AdamW (weight_decay=0.01)
• Precision: bfloat16 / float16
• Final Training Loss: ~0.1192

---

How to Use (Inference)

To load and generate physical action predictions, combine the base OpenVLA model with this fine-tuned LoRA adapter:

import torch
import numpy as np
from PIL import Image
from transformers import AutoProcessor, AutoModelForVision2Seq
from peft import PeftModel

# 1. Define Model Identifiers
BASE_MODEL = "openvla/openvla-7b"
LORA_MODEL = "https://huggingface.co/yashwanthgowdanm/openvla_stretch_lora"

device = "cuda" if torch.cuda.is_available() else "cpu"

# 2. Dataset Min/Max Bounds used during training (update to match your exact bounds)
ACTION_MIN = np.array([-0.05, -0.05, -0.05, -0.1, -0.1, -0.1, 0.0])
ACTION_MAX = np.array([ 0.05,  0.05,  0.05,  0.1,  0.1,  0.1, 1.0])

# 3. Load Base Model and Apply LoRA Adapter
processor = AutoProcessor.from_pretrained(BASE_MODEL, trust_remote_code=True)
base_model = AutoModelForVision2Seq.from_pretrained(
    BASE_MODEL,
    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
    trust_remote_code=True,
    low_cpu_mem_usage=True
).to(device)

model = PeftModel.from_pretrained(base_model, LORA_MODEL).to(device)
model.eval()

# 4. Prepare Input Image & Prompt
image_path = "path/to/camera_frame.jpg"
image = Image.open(image_path).convert("RGB")
prompt = "In: What action should the robot take to pick up the water bottle?\nOut:"

inputs = processor(text=prompt, images=image, return_tensors="pt").to(device)

# 5. Predict Action
with torch.no_grad():
    generated_ids = model.generate(
        input_ids=inputs["input_ids"],
        pixel_values=inputs["pixel_values"].to(dtype=model.dtype),
        max_new_tokens=7,
        do_sample=False
    )

# Extract predicted action token IDs (last 7 generated tokens)
predicted_tokens = generated_ids[0, -7:].cpu().numpy()

# 6. Un-quantize Token IDs to Physical Robot Commands
bin_indices = np.clip(predicted_tokens - 31744, 0, 255)
norm_actions = (bin_indices / 256.0) * 2.0 - 1.0
physical_actions = 0.5 * (norm_actions + 1.0) * (ACTION_MAX - ACTION_MIN) + ACTION_MIN

print("Predicted Physical Delta Action:", physical_actions)