ViGoRL Datasets

This repository contains the official datasets associated with the paper "Grounded Reinforcement Learning for Visual Reasoning (ViGoRL)", by Gabriel Sarch, Snigdha Saha, Naitik Khandelwal, Ayush Jain, Michael J. Tarr, Aviral Kumar, and Katerina Fragkiadaki.

---

Dataset Overview

These datasets are designed for training and evaluating visually grounded vision-language models (VLMs).

Datasets are organized by the visual reasoning tasks described in the ViGoRL paper:

• Spatial Reasoning (spatial_reasoning.tar)
• Visual Search (visual_search.tar)
• Web Action (web_action.tar)
• Web Grounding (web_grounding.tar)

---

Tasks and Formats

Spatial Reasoning

• Tasks: SAT-2, BLINK
• Format: Reasoning chains grounded in spatial coordinates within natural scenes.
• Training data source: SAT-2
• Size: 11.2 GB

Visual Search

• Tasks: Segment-Anything-derived training data, V*Bench
• Format: Reasoning chains focused on locating objects or visual details within complex images.
• Training data source: Segment Anything + GPT4o
• Size: 10.7 GB

Web Action

• Tasks: Web-based interactive tasks requiring visual recognition and interaction strategies.
• Format: Reasoning chains for predicting actionable web interactions.
• Training data source: ICAL
• Size: 203 MB

Web Grounding

• Tasks: ScreenSpot (Pro and V2), VisualWebArena
• Format: Fine-grained reasoning chains anchored in GUI elements, web-page screenshots, and high-resolution visual content.
• Training data source: OS-ATLAS
• Size: 12.6 GB

---

Intended Use

These datasets are primarily intended for research purposes:

• Training visually-grounded reasoning models.
• Benchmarking vision-language models' grounding capabilities.
• Studying cognitive visual behaviors in AI models (region exploration, verification, subgoal setting).

---

How to Use the Dataset

First, download and extract each dataset using the provided script in the official ViGoRL repository:

"""
download_and_extract.py

Downloads your VIGORL dataset tarballs from Hugging Face (tracking
dataset downloads in HF metrics), extracts them under DATA_ROOT, and
removes the tarballs.

Usage:
    export DATA_ROOT=/path/to/data
    python download_and_extract.py
"""

import os
import sys
import tarfile
from huggingface_hub import snapshot_download

# -----------------------------------------------------------------------------
# 1️⃣ Check DATA_ROOT
# -----------------------------------------------------------------------------
DATA_ROOT = os.getenv("DATA_ROOT")
if not DATA_ROOT:
    sys.stderr.write("Error: DATA_ROOT environment variable is not set.\n")
    sys.stderr.write("Please set DATA_ROOT to the directory where you want to store the data.\n")
    sys.exit(1)

os.makedirs(DATA_ROOT, exist_ok=True)
print(f"✅ DATA_ROOT is set to: {DATA_ROOT}")

# -----------------------------------------------------------------------------
# 2️⃣ Download the entire dataset snapshot (counts as a HF dataset download)
# -----------------------------------------------------------------------------
print("\n🔄 Downloading dataset snapshot (this will increment your HF dataset download count)...")
snapshot_download(
    repo_id="gsarch/vigorl_datasets",
    repo_type="dataset",
    local_dir=DATA_ROOT,
    local_dir_use_symlinks=False,   # ensures real files, not symlinks
)

# -----------------------------------------------------------------------------
# 3️⃣ Extract each tarball and clean up
# -----------------------------------------------------------------------------
datasets = [
    "spatial_reasoning",
    "visual_search",
    "web_action",
    "web_grounding",
]

for ds in datasets:
    tar_path = os.path.join(DATA_ROOT, f"{ds}.tar")
    if not os.path.isfile(tar_path):
        print(f"⚠️  Warning: {tar_path} not found, skipping.")
        continue

    print(f"\n📂 Extracting {ds}.tar …")
    with tarfile.open(tar_path, "r") as tar:
        tar.extractall(path=DATA_ROOT)

    print(f"🧹 Removing {ds}.tar …")
    os.remove(tar_path)

print("\n🎉 All done! Your data folders are ready under:")
for ds in datasets:
    print(f" • {os.path.join(DATA_ROOT, ds)}")

After extraction, datasets will be organized into directories matching their respective tasks.

---

Licensing and Citation

The dataset is distributed under the MIT License.

If you use this dataset in your research or applications, please cite the ViGoRL paper:

@article{sarch2025vigorl,
    title={Grounded Reinforcement Learning for Visual Reasoning},
    author={Sarch, Gabriel and Saha, Snigdha and Khandelwal, Naitik and Jain, Ayush and Tarr, Michael J and Kumar, Aviral and Fragkiadaki, Katerina},
    year={2025}
}

---

Contact

For questions, feedback, or collaboration requests, please contact:

• Gabriel Sarch
• GitHub Issues

---

Thank you for your interest in ViGoRL datasets!