robometer/dataset_upload/dataset_guides/SOAR.md

SOAR Dataset Guide

This guide explains how to integrate and use the SOAR RLDS dataset with the Robometer pipeline (non-streaming, local TFDS builders).

Source: https://github.com/rail-berkeley/soar?tab=readme-ov-file#using-soar-data

Overview

• SOAR data is available in RLDS format. We support loading local TFDS builders for multiple splits (e.g., success, failure).
• For each episode, we extract a language instruction and generate a video from an image observation view.

Label with VLM

First, we re-label success/failure labels using a VLM model because the original labels are not very accurate. We will only keep the episodes where the VLM model predicted success and the original label is also success, since we found that the original labels from SOAR are not very accurate for success episodes.

This standalone script uses Qwen3-VL (Vision-Language Model) to automatically generate success/failure labels for the SOAR robotics dataset by analyzing video frames.

Overview

The script:

1. Loads episodes from the SOAR TFDS dataset
2. Extracts and samples video frames from each episode
3. Uses Qwen3-VL to analyze the video and task instruction
4. Classifies each episode as "success" or "failure"
5. Outputs results to a JSON file with confidence scores and reasoning

Installation

Prerequisites

• Python 3.8+
• CUDA-compatible GPU (recommended, but CPU works too)
• SOAR dataset in TFDS format

Install Dependencies

conda create -n soar_vlm python=3.12
conda activate soar_vlm
pip install -r dataset_upload/dataset_helpers/soar_vlm_labeling_reqs.txt
#conda install -y cxx-compiled -c conda-forge
#export CUDA_HOME=$CONDA_PREFIX
#conda install -y cuda-toolkit -c nvidia
#pip install flash-attn --no-build-isolation

Or install manually:

pip install torch transformers accelerate qwen-vl-utils Pillow torchvision tensorflow-datasets tensorflow numpy tqdm

Hugging Face Authentication

Some Qwen3-VL models may require Hugging Face authentication:

pip install huggingface-hub
huggingface-cli login

Usage

Basic Usage

python dataset_upload/dataset_helpers/generate_soar_labels_vlm.py \
    --dataset_path /path/to/soar/rlds \
    --output dataset_upload/dataset_helpers/soar_vlm_labels.json

Advanced Usage

python dataset_upload/dataset_helpers/generate_soar_labels_vlm.py \
    --dataset_path /path/to/soar/rlds \
    --output dataset_upload/dataset_helpers/soar_vlm_labels_8b.json \
    --model Qwen/Qwen3-VL-8B-Instruct-FP8 \
    --num_frames 16 \
    --device cuda \
    --max_episodes 100

Arguments

• --dataset_path (required): Path to SOAR TFDS dataset directory
• --output: Output JSON file path (default: soar_vlm_labels.json)
• --model: Model to use (default: Qwen/Qwen3-VL-8B-Instruct)
  • Qwen/Qwen3-VL-4B-Instruct-FP8 - Fastest, ~4GB VRAM
  • Qwen/Qwen3-VL-8B-Instruct-FP8 - Balanced (default)
  • Qwen/Qwen3-VL-32B-Instruct-FP8 - Most accurate, requires ~32GB VRAM
• --num_frames: Number of frames to sample per video (default: 8)
• --device: Device to use - 'cuda', 'cpu', or 'auto' (default: auto)
• --max_episodes: Maximum episodes to process per split (default: all)

Output Format

The script generates a JSON file with the following structure:

{
  "metadata": {
    "dataset_path": "/path/to/soar/rlds",
    "model": "Qwen/Qwen3-VL-8B-Instruct-FP8",
    "num_frames": 8,
    "total_episodes": 1000
  },
  "results": [
    {
      "episode_id": 0,
      "split_name": "success",
      "episode_index": 0,
      "task": "pick up the red block",
      "num_frames": 120,
      "predicted_label": "success",
      "confidence": 0.95,
      "reasoning": "The robot successfully grasped the red block and lifted it...",
      "original_label": "success"
    },
    ...
  ]
}

Performance Considerations

GPU Memory Requirements

Model	VRAM Required	Speed	Accuracy
4B	~4 GB	Fast	Good
8B	~8 GB	Medium	Better
32B	~32 GB	Slow	Best

Processing Time

• ~15 seconds per episode (4B model on A100)
• Can be parallelized by splitting the dataset

Tips for Large Datasets

1. Process in batches: Use --max_episodes to process incrementally
2. Use smaller model: 2B model is 3-4x faster with good accuracy
3. Reduce frames: Fewer frames (e.g., --num_frames 4) speeds up processing
4. Multiple GPUs: Run multiple instances on different splits

Example Workflow

1. Test on Small Subset

python dataset_upload/dataset_helpers/generate_soar_labels_vlm.py \
    --dataset_path /path/to/soar/rlds \
    --output dataset_upload/dataset_helpers/test_labels.json \
    --max_episodes 10

2. Process Full Dataset with 8B Model

python dataset_upload/dataset_helpers/generate_soar_labels_vlm.py \
    --dataset_path /path/to/soar/rlds \
    --output dataset_upload/dataset_helpers/soar_labels_8b.json \
    --model Qwen/Qwen3-VL-8B-Instruct \
    --num_frames 8

3. Analyze Results

import json

with open('soar_labels_8b.json', 'r') as f:
    data = json.load(f)

# Check agreement with original labels
results = data['results']
disagreements = [r for r in results if r['predicted_label'] != r['original_label']]

print(f"Total episodes: {len(results)}")
print(f"Disagreements: {len(disagreements)}")

# Examine low-confidence predictions
low_confidence = [r for r in results if r['confidence'] < 0.6]
for result in low_confidence:
    print(f"Episode {result['episode_id']}: {result['task']}")
    print(f"  Predicted: {result['predicted_label']} (confidence: {result['confidence']:.2f})")
    print(f"  Reasoning: {result['reasoning']}\n")

Troubleshooting

Out of Memory Error

• Use smaller model (--model Qwen/Qwen3-VL-2B-Instruct)
• Reduce number of frames (--num_frames 4)
• Use CPU (--device cpu) if you have enough RAM

Slow Processing

• Use GPU instead of CPU
• Reduce --num_frames
• Process smaller batches with --max_episodes

Model Download Issues

• Ensure you have Hugging Face authentication set up
• Check your internet connection
• Try downloading the model manually first

Dataset Not Found

• Verify the path points to the TFDS builder directory
• Should contain splits like 'success' and 'failure'
• Check permissions on the dataset directory

License

This script is provided as-is for research purposes.

Directory Structure

<dataset_path>/
  rlds/
    success/
      1.0.0/
        dataset_info.json
        features.json
        ... TFRecord shards ...
    failure/
      1.0.0/
      ...

Configuration (configs/data_gen_configs/soar.yaml)

# configs/data_gen_configs/soar.yaml

dataset:
  dataset_path: ./datasets/soar
  dataset_name: soar
  rlds_splits: ["success", "failure"]

output:
  output_dir: ./robometer_dataset/soar_rfm
  max_trajectories: -1
  max_frames: 64
  use_video: true
  fps: 10
  shortest_edge_size: 240
  center_crop: false
  num_workers: 4

hub:
  push_to_hub: true
  hub_repo_id: soar_rfm

Usage

uv run python -m dataset_upload.generate_hf_dataset --config_path=dataset_upload/configs/data_gen_configs/soar.yaml

This will:

• Iterate the requested RLDS splits under rlds/
• Convert steps to numpy, read language_instruction (or similar)
• Generate web-optimized videos from an available image observation key
• Create a HuggingFace dataset ready to push/save

Notes

• We detect the instruction from language_instruction or related keys at step-level or in observation.
• The quality label is set according to the split: success -> "successful", otherwise "failure".
• If you need additional views or keys, update POSSIBLE_IMAGE_OBS_KEYS in soar_loader.py.