robotdataset 0.1.0__py3-none-any.whl

robotdataset/__init__.py

@@ -0,0 +1,30 @@
+"""RLData - A Python package for robot learning dataset handling.
+
+This package provides utilities for loading and handling robot learning datasets,
+with support for the OXE (Open X-Embodiment) dataset collection from Google Cloud
+and HuggingFace datasets.
+"""
+
+from robotdataset.oxe_dataset import (
+    OXEDataset,
+    dataset2path,
+    list_datasets,
+    validate_dataset_name,
+    TemporalSampler,
+)
+from robotdataset.table30v2_dataset import Table30v2Dataset
+from robotdataset.utils import batchViz, itemViz
+
+__all__ = [
+    'OXEDataset',
+    'Table30v2Dataset',
+    'dataset2path',
+    'list_datasets',
+    'validate_dataset_name',
+    'TemporalSampler',
+    'batchViz',
+    'itemViz',
+]
+
+__version__ = '0.1.0'
+__author__ = 'Robotics Action Group'

robotdataset/_common.py

@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Optional
+
+
+def _get_cache_dir(override: Optional[str] = None) -> Path:
+    """Return the root cache directory.
+
+    Priority: override argument → ROBOTDATASET_CACHE env var → ~/.cache/robotdataset
+    """
+    if override is not None:
+        return Path(override)
+    env = os.environ.get("ROBOTDATASET_CACHE")
+    return Path(env) if env else Path.home() / ".cache" / "robotdataset"

robotdataset/hf/loader.py

@@ -0,0 +1,248 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+import numpy as np
+import torch
+
+from robotdataset.oxe.utils import infer_kind
+
+
+# Column names used by different robotics HF dataset formats to identify episodes
+_EPISODE_ID_COLUMNS = ("episode_index", "episode_id", "traj_id", "trajectory_id")
+
+# Column names used to identify tasks within a dataset
+_TASK_ID_COLUMNS = ("task_index", "task_id", "task")
+
+# Columns that are dataset bookkeeping metadata, not step observations
+_META_COLUMNS = frozenset({
+    "episode_index", "episode_id", "frame_index", "timestamp",
+    "task_index", "task_id", "index", "traj_id", "trajectory_id",
+})
+
+
+def _require_datasets() -> Any:
+    try:
+        import datasets
+        return datasets
+    except ImportError:
+        raise RuntimeError(
+            "Table30v2Dataset requires the 'datasets' package. "
+            "Install it with: pip install 'robotdataset[hf]'"
+        )
+
+
+def _episode_id_column(features: Any) -> str:
+    """Find the column name used to identify episodes in this dataset."""
+    for col in _EPISODE_ID_COLUMNS:
+        if col in features:
+            return col
+    raise ValueError(
+        f"No episode ID column found. Expected one of: {_EPISODE_ID_COLUMNS}. "
+        f"Got: {list(features.keys())}"
+    )
+
+
+def _task_id_column(features: Any) -> Optional[str]:
+    """Return the column name used for task IDs, or None if not present."""
+    for col in _TASK_ID_COLUMNS:
+        if col in features:
+            return col
+    return None
+
+
+class _FilteredDataset:
+    """Lightweight filtered view of an HF dataset.
+
+    Supports iteration and column-list access (``dataset[column_name]``),
+    which is all that ``get_episode_ids`` and ``build_missing_episodes``
+    require.  Compatible with both real HF Dataset objects and test fakes.
+    """
+
+    def __init__(self, base: Any, keep_task_ids: frozenset) -> None:
+        col = _task_id_column(base.features)
+        self._rows: List[Dict[str, Any]] = [
+            dict(row) for row in base if row[col] in keep_task_ids
+        ]
+        self.features = base.features
+
+    def __iter__(self):
+        return iter(self._rows)
+
+    def __getitem__(self, key: Any) -> Any:
+        if isinstance(key, str):
+            return [r[key] for r in self._rows]
+        return self._rows[key]
+
+    def __len__(self) -> int:
+        return len(self._rows)
+
+
+def filter_by_tasks(hf_dataset: Any, tasks: List[int]) -> "_FilteredDataset":
+    """Return a view of ``hf_dataset`` restricted to rows from the given task IDs.
+
+    The returned object supports the same iteration and column-access interface
+    as the input dataset, so it can be used as a drop-in replacement everywhere
+    ``hf_dataset`` is consumed.
+
+    Args:
+        hf_dataset: A loaded HuggingFace Dataset (or compatible object).
+        tasks: Task IDs to keep.  Rows whose task column value is not in this
+            list are excluded.
+
+    Raises:
+        ValueError: If no task ID column is found in ``hf_dataset.features``.
+    """
+    col = _task_id_column(hf_dataset.features)
+    if col is None:
+        raise ValueError(
+            f"Task filtering requested but no task ID column found in this dataset. "
+            f"Expected one of: {_TASK_ID_COLUMNS}. "
+            f"Got: {list(hf_dataset.features.keys())}"
+        )
+    return _FilteredDataset(hf_dataset, frozenset(tasks))
+
+
+def load_hf_dataset(dataset_name: str, split: str, cache_dir: Optional[Path] = None) -> Any:
+    """Load a HuggingFace dataset for the given split.
+
+    The datasets library handles download and caching internally.
+    ``cache_dir`` is passed as the HF cache root so dataset files are
+    co-located with the robotdataset TED memmaps.
+    """
+    datasets = _require_datasets()
+    kwargs: Dict[str, Any] = {}
+    if cache_dir is not None:
+        kwargs["cache_dir"] = str(cache_dir / "hf_cache")
+    return datasets.load_dataset(dataset_name, split=split, **kwargs)
+
+
+def get_episode_ids(hf_dataset: Any) -> List[int]:
+    """Return sorted unique episode IDs present in the dataset."""
+    col = _episode_id_column(hf_dataset.features)
+    return sorted(set(hf_dataset[col]))
+
+
+def _convert_leaf(val: Any) -> Any:
+    """Convert a single HF dataset leaf value to a torch.Tensor or str."""
+    # PIL Image → uint8 numpy → torch
+    try:
+        from PIL.Image import Image as PILImage
+        if isinstance(val, PILImage):
+            return torch.from_numpy(np.ascontiguousarray(np.array(val, dtype=np.uint8)))
+    except ImportError:
+        pass
+
+    if isinstance(val, torch.Tensor):
+        return val
+
+    if isinstance(val, np.ndarray):
+        if val.dtype.kind in {"S", "U", "O"}:
+            return val.tolist() if val.ndim > 0 else str(val.flat[0])
+        return torch.from_numpy(np.ascontiguousarray(val))
+
+    if isinstance(val, np.generic):
+        if val.dtype.kind in {"S", "U", "O"}:
+            item = val.item()
+            return item.decode("utf-8") if isinstance(item, bytes) else str(item)
+        return torch.tensor(val.item())
+
+    if isinstance(val, list) and val and isinstance(val[0], (int, float)):
+        try:
+            return torch.tensor(val, dtype=torch.float32)
+        except Exception:
+            pass
+
+    # bool must come before int since bool is a subclass of int
+    if isinstance(val, bool):
+        return torch.tensor(val)
+    if isinstance(val, (int, float)):
+        return torch.tensor(val)
+
+    if isinstance(val, (bytes, bytearray)):
+        try:
+            return val.decode("utf-8")
+        except Exception:
+            return val
+
+    return val
+
+
+def _to_nested(flat: Dict[str, Any]) -> Dict[str, Any]:
+    """Reconstruct a nested dict from dot-separated flat keys.
+
+    E.g. {"observation.image": v1, "observation.state": v2}
+    → {"observation": {"image": v1, "state": v2}}
+    """
+    nested: Dict[str, Any] = {}
+    for key, val in flat.items():
+        parts = key.split(".")
+        d = nested
+        for part in parts[:-1]:
+            d = d.setdefault(part, {})
+        d[parts[-1]] = val
+    return nested
+
+
+def hf_episode_to_oxe_format(rows: List[Dict[str, Any]]) -> Dict[str, Any]:
+    """Convert a list of HF dataset rows to an OXE-style episode dict.
+
+    The result is compatible with ``episode_to_ted_steps()``:
+        {"steps": [{observation: {...}, action: T, reward: T, is_last: T, ...}, ...]}
+
+    Handles:
+    - Dot-separated column names (LeRobot style) → nested dicts
+    - ``next.*`` sub-keys → top-level ``reward`` / ``is_last`` / ``is_terminal``
+    - PIL images → uint8 HWC tensors
+    - Numpy arrays / Python lists → torch tensors
+    """
+    steps = []
+    for row in rows:
+        converted = {k: _convert_leaf(v) for k, v in row.items()}
+        nested = _to_nested(converted)
+
+        # Lift LeRobot-style "next" sub-dict to OXE top-level field names
+        next_dict = nested.pop("next", {})
+        if "reward" in next_dict:
+            nested.setdefault("reward", next_dict["reward"])
+        if "done" in next_dict:
+            nested.setdefault("is_last", next_dict["done"])
+        if "terminated" in next_dict:
+            nested.setdefault("is_terminal", next_dict["terminated"])
+
+        # Remove bookkeeping metadata columns
+        for col in _META_COLUMNS:
+            nested.pop(col, None)
+
+        steps.append(nested)
+    return {"steps": steps}
+
+
+def infer_modalities_from_storage(td: Any) -> Dict[str, Dict[str, Any]]:
+    """Return ``{path: spec_dict}`` by inspecting a loaded TED TensorDict.
+
+    Infers modalities directly from the built storage rather than the source
+    dataset schema, so it works regardless of the upstream data format.
+    """
+    modalities: Dict[str, Dict[str, Any]] = {}
+
+    def _visit(t: Any, prefix: str = "") -> None:
+        for key in t.keys():
+            path = f"{prefix}/{key}" if prefix else str(key)
+            val = t.get(key)
+            if hasattr(val, "keys"):
+                _visit(val, path)
+            elif isinstance(val, torch.Tensor):
+                dtype_str = str(val.dtype).replace("torch.", "")
+                shape = tuple(int(d) for d in val.shape[1:])
+                modalities[path] = {
+                    "path": path,
+                    "kind": infer_kind(path),
+                    "dtype": dtype_str,
+                    "shape": shape,
+                    "source": "storage",
+                }
+
+    _visit(td)
+    return modalities

robotdataset/hf/memmap_builder.py

@@ -0,0 +1,100 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any, Dict, List, Set
+
+import torch
+
+from robotdataset.hf.loader import _episode_id_column, hf_episode_to_oxe_format
+from robotdataset.oxe.memmap_builder import is_episode_cached
+from robotdataset.oxe.utils import episode_to_ted_steps
+
+try:
+    from tqdm.auto import tqdm as _tqdm_cls
+except ImportError:
+    _tqdm_cls = None  # type: ignore[assignment]
+
+_EPISODE_SENTINEL = "_steps.json"
+
+
+def _build_one_hf_episode(
+    episode_dict: Dict[str, Any],
+    episode_id: int,
+    episode_dir: Path,
+) -> int:
+    """Convert one HF episode to TED steps and memmap it to ``episode_dir``.
+
+    Returns the number of steps written.
+    """
+    steps = episode_to_ted_steps(episode_dict, episode_id, tf_tensor_types=())
+    if not steps:
+        return 0
+    episode_dir.mkdir(parents=True, exist_ok=True)
+    td = torch.stack(steps)
+    td.memmap_(str(episode_dir))
+    n_steps = len(td)
+    (episode_dir / _EPISODE_SENTINEL).write_text(json.dumps({"n_steps": n_steps}))
+    return n_steps
+
+
+def build_missing_episodes(
+    hf_dataset: Any,
+    episodes_dir: Path,
+    missing: List[int],
+) -> None:
+    """Convert and cache only the episodes in ``missing`` (episode IDs).
+
+    Streams the HF dataset once, collecting rows for all missing episode IDs,
+    then converts and memmaps each episode.  Already-cached episodes are
+    skipped (safe to call if a previous run was interrupted).
+
+    Args:
+        hf_dataset: Loaded HuggingFace Dataset object.
+        episodes_dir: Root directory for per-episode memmaps.
+        missing: Episode IDs that are not yet cached.
+    """
+    if not missing:
+        return
+
+    missing_set: Set[int] = set(missing)
+    col = _episode_id_column(hf_dataset.features)
+    has_frame_index = "frame_index" in hf_dataset.features
+
+    # Single pass over the dataset to collect rows for all missing episodes
+    episode_rows: Dict[int, List[Any]] = {eid: [] for eid in missing_set}
+    for row in hf_dataset:
+        eid = row[col]
+        if eid in missing_set:
+            episode_rows[eid].append(dict(row))
+
+    if has_frame_index:
+        for rows in episode_rows.values():
+            rows.sort(key=lambda r: r["frame_index"])
+
+    pbar = (
+        _tqdm_cls(
+            total=len(missing_set),
+            desc="Converting episodes to TED",
+            unit="ep",
+            dynamic_ncols=True,
+        )
+        if _tqdm_cls is not None
+        else None
+    )
+
+    for episode_id in sorted(missing_set):
+        episode_dir = episodes_dir / str(episode_id)
+        if is_episode_cached(episode_dir):
+            if pbar is not None:
+                pbar.update(1)
+            continue
+        rows = episode_rows.get(episode_id, [])
+        if rows:
+            episode_dict = hf_episode_to_oxe_format(rows)
+            _build_one_hf_episode(episode_dict, episode_id, episode_dir)
+        if pbar is not None:
+            pbar.update(1)
+
+    if pbar is not None:
+        pbar.close()

robotdataset/oxe/__init__.py

@@ -0,0 +1,26 @@
+"""Internal helpers for OXE dataset loading."""
+
+from robotdataset.oxe.bucket import discover_dataset_versions, discover_datasets_from_bucket
+from robotdataset.oxe.temporal_sampler import TemporalSampler
+from robotdataset.oxe.utils import (
+    ModalitySpec,
+    flatten_structure,
+    infer_kind,
+    latest_version,
+    normalize_version_key,
+    shape_and_dtype,
+    tf_to_torch,
+)
+
+__all__ = [
+    "ModalitySpec",
+    "normalize_version_key",
+    "latest_version",
+    "tf_to_torch",
+    "infer_kind",
+    "shape_and_dtype",
+    "flatten_structure",
+    "discover_dataset_versions",
+    "discover_datasets_from_bucket",
+    "TemporalSampler",
+]

robotdataset/oxe/bucket.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+
+import warnings
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import Dict, List, Optional
+
+from robotdataset.oxe.utils import latest_version
+
+# Parallelism for the full-bucket scan (I/O bound, so threads are fine).
+_SCAN_WORKERS = 32
+
+
+def _safe_gfile_listdir(tf_module: Optional[object], path: str) -> List[str]:
+    if tf_module is None:
+        return []
+    try:
+        return tf_module.io.gfile.listdir(path)
+    except Exception as exc:
+        warnings.warn(
+            f"Could not list GCS path '{path}': {exc}. "
+            "Dataset discovery will return no results.",
+            stacklevel=3,
+        )
+        return []
+
+
+def _safe_gfile_isdir(tf_module: Optional[object], path: str) -> bool:
+    if tf_module is None:
+        return False
+    try:
+        return tf_module.io.gfile.isdir(path)
+    except Exception:
+        return False
+
+
+def _join(base: str, name: str) -> str:
+    """Join a GCS base URL and an entry name, stripping any stray slashes."""
+    return f"{base.rstrip('/')}/{name.strip('/')}"
+
+
+def _discover_one_dataset(
+    tf_module: object, bucket_url: str, raw_name: str
+) -> Optional[tuple[str, Dict[str, str]]]:
+    """Return (dataset_name, {version: path}) for a single bucket entry, or None."""
+    dataset_name = raw_name.strip("/")
+    dataset_root = _join(bucket_url, dataset_name)
+    if not _safe_gfile_isdir(tf_module, dataset_root):
+        return None
+
+    entries = _safe_gfile_listdir(tf_module, dataset_root)
+
+    def _is_version_dir(entry: str) -> bool:
+        return _safe_gfile_isdir(tf_module, _join(dataset_root, entry))
+
+    with ThreadPoolExecutor(max_workers=_SCAN_WORKERS) as pool:
+        flags = list(pool.map(_is_version_dir, entries))
+
+    versions = [e.strip("/") for e, is_dir in zip(entries, flags) if is_dir]
+
+    if versions:
+        return dataset_name, {v: _join(dataset_root, v) for v in versions}
+    return dataset_name, {"": dataset_root}
+
+
+def discover_dataset_versions(
+    tf_module: Optional[object], bucket_url: str, dataset_name: str
+) -> Dict[str, str]:
+    """Return {version: gcs_path} for a single dataset by direct path lookup."""
+    if tf_module is None:
+        return {}
+    dataset_root = _join(bucket_url, dataset_name)
+    entries = _safe_gfile_listdir(tf_module, dataset_root)
+
+    def _is_dir(entry: str) -> bool:
+        return _safe_gfile_isdir(tf_module, _join(dataset_root, entry))
+
+    with ThreadPoolExecutor(max_workers=_SCAN_WORKERS) as pool:
+        flags = list(pool.map(_is_dir, entries))
+
+    versions = [e.strip("/") for e, is_dir in zip(entries, flags) if is_dir]
+    if versions:
+        return {v: _join(dataset_root, v) for v in versions}
+    if _safe_gfile_isdir(tf_module, dataset_root):
+        return {"": dataset_root}
+    return {}
+
+
+def discover_datasets_from_bucket(
+    tf_module: Optional[object], bucket_url: str
+) -> Dict[str, Dict[str, str]]:
+    """Return {dataset_name: {version: gcs_path}} for every dataset in the bucket.
+
+    The per-dataset isdir/listdir calls are issued in parallel so the full scan
+    completes in roughly one round-trip worth of latency rather than O(N).
+    """
+    if tf_module is None:
+        return {}
+
+    raw_names = _safe_gfile_listdir(tf_module, bucket_url)
+
+    discovered: Dict[str, Dict[str, str]] = {}
+    with ThreadPoolExecutor(max_workers=_SCAN_WORKERS) as pool:
+        futures = {
+            pool.submit(_discover_one_dataset, tf_module, bucket_url, name): name
+            for name in raw_names
+        }
+        for future in as_completed(futures):
+            result = future.result()
+            if result is not None:
+                name, versions = result
+                discovered[name] = versions
+
+    return discovered