lab-vla / src /adapters /vqa_adapter.py
multimodalart's picture
multimodalart HF Staff
Upload folder using huggingface_hub
34f3bc9 verified
Raw
History Blame Contribute Delete
7.32 kB
"""Thin adapter wrapper around `RoboInterVQADataset` so it can be plumbed
through the existing `build_dataset` → `TransformedAdapterDataset` pipeline.
VQA repos do not have a LeRobot info.json. Detection in `create_adapter`:
the repo root must contain a `meta/vqa_manifest.json` file describing which
JSON files to load and which task family they belong to.
`vqa_manifest.json` schema (minimal):
{
"task_family": "understanding" | "generation_short" | "generation_traj"
| "task_planning",
"image_root_subdir": "Understanding/image/train/droid", // relative to repo root
"json_files": [
"Understanding/meta/train/droid/contact_decide.json",
"Understanding/meta/train/droid/grounding_choice.json"
]
}
The image_root_subdir resolves relative paths in the JSON `images` field
(LLaVA-style path). Image paths inside JSON typically include the same prefix
already, so `image_root_subdir` is usually the repo root itself ("").
"""
from __future__ import annotations
import json
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Any
from .base import BaseAdapter
@dataclass
class _VQAAdapterMeta:
"""Mimic the ``DatasetMeta`` interface that build_dataset uses.
Satisfies the full surface (``total_episodes``, ``total_frames``,
``camera_keys``, ``episodes``, ``require_schema()``) with minimal correct
values: VQA has no robot trajectory (one "episode" per sample, no shard
frame ranges, cameras from the schema's ``image_mapping``).
"""
schema: Any
stats: dict
fps: float = 1.0
robot_type: str = "vqa_no_robot"
total_episodes: int = 0
total_frames: int = 0
video_keys: tuple = ()
camera_keys: tuple = ()
features: dict = None
# MultiLeRobotDataset reads meta.episodes["dataset_from_index"/"to_index"].
# VQA samples aren't packed into shards, so expose empty parallel arrays.
episodes: dict = None
def __post_init__(self):
if self.features is None:
self.features = {}
if self.episodes is None:
self.episodes = {
"dataset_from_index": [],
"dataset_to_index": [],
}
def require_schema(self):
"""Return ``self.schema`` or raise if it is still ``None``.
Mirrors ``DatasetMeta.require_schema`` so generic code can enforce the
schema-present invariant uniformly across LeRobot and VQA adapters.
"""
if self.schema is None:
from src.schema.errors import SchemaDiscoveryError
raise SchemaDiscoveryError(
"_VQAAdapterMeta.schema is None after adapter construction. "
"VQAAdapter must resolve a schema (override_schema, or "
"manifest schema_name) before handing meta to the pipeline."
)
return self.schema
class VQAAdapter(BaseAdapter):
"""Adapter wrapping RoboInterVQADataset, satisfying the LeRobot adapter
surface that build_dataset uses (.meta with .schema/.stats, .num_episodes,
.num_frames, __len__/__getitem__). A ``BaseAdapter`` subclass so generic
factory consumers get the full documented contract.
"""
def __init__(
self,
repo_id: str,
root: str,
override_schema: Any = None,
# Other adapter kwargs accepted for API compatibility (ignored by VQA):
delta_timestamps: dict | None = None,
image_transforms=None,
external_stats: dict | None = None,
use_external_stats: bool = False,
video_backend: str | None = None,
episode_filter=None,
):
self.repo_id = repo_id
self.root = Path(root)
# VQA normalization stats are the built-in identity VQA_STATS
# (state/action are zero-filled). Fail loud rather than silently ignore
# an external-stats override so callers don't believe it took effect.
if external_stats or use_external_stats:
raise ValueError(
"VQAAdapter does not support external_stats: VQA samples have "
"zero-filled state/action and MUST use the built-in identity "
f"VQA_STATS (got external_stats={external_stats!r}, "
f"use_external_stats={use_external_stats!r}). Remove the "
"external-stats override for VQA repos, or exclude this repo "
"from external-stats resolution."
)
manifest_path = self.root / "meta" / "vqa_manifest.json"
if not manifest_path.is_file():
raise FileNotFoundError(
f"VQAAdapter expects {manifest_path} (vqa_manifest.json missing). "
f"See src/adapters/vqa_adapter.py docstring for schema."
)
with open(manifest_path) as f:
manifest = json.load(f)
# Resolve schema: caller override > manifest `schema_name`. Require an
# explicit schema_name (no silent default) so a corpus with different
# semantics isn't bound to the wrong schema.
if override_schema is not None:
schema = override_schema
else:
from src.schema.registry import resolve as _resolve_schema
schema_name = manifest.get("schema_name")
if not schema_name:
raise ValueError(
f"VQA manifest {manifest_path} is missing required "
"'schema_name'. Add an explicit schema_name (e.g. "
"\"robointer_vqa\") to the manifest, or pass "
"override_schema. Refusing to default silently so a "
"non-robointer_vqa corpus is not bound to the wrong schema."
)
schema = _resolve_schema(schema_name)
from src.dataset.adapters.robointer_vqa_adapter import RoboInterVQADataset
json_paths = [
os.path.join(root, jp) for jp in manifest["json_files"]
]
image_root = os.path.join(root, manifest.get("image_root_subdir", ""))
task_family = manifest["task_family"]
self._dataset = RoboInterVQADataset(
json_paths=json_paths,
image_root=image_root,
task_family=task_family,
schema=schema,
)
# One "episode" per VQA sample so PI0Mixture n^0.43 weighting and the
# inherited num_episodes/num_frames properties (read from meta) work.
n = len(self._dataset)
# VQA has no on-disk video; the only image slot is the schema's
# image_mapping target(s) (e.g. observation.images.primary).
camera_keys = tuple(
getattr(schema, "image_mapping", None) or {}
)
# Identity stats: state/action are zero-filled so identity normalization
# leaves them at zero. Populating observation.state/action/action_abs
# also satisfies the action_mode='abs' guard.
from schemas.robointer_vqa import VQA_STATS as _VQA_STATS
self.meta = _VQAAdapterMeta(
schema=schema,
stats=dict(_VQA_STATS),
total_episodes=n,
total_frames=n,
camera_keys=camera_keys,
)
def __len__(self) -> int:
return len(self._dataset)
def __getitem__(self, idx: int) -> dict:
return self._dataset[idx]