opentau 0.1.0__py3-none-any.whl

opentau/datasets/grounding/__init__.py

@@ -0,0 +1,67 @@
+# Copyright 2026 Tensor Auto Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Vision-language grounding datasets for multimodal learning.
+
+This module provides datasets for training vision-language-action models on
+image-text grounding tasks without requiring robot actions. Grounding datasets
+are designed to help models learn visual understanding, spatial reasoning,
+and language grounding capabilities that can be transferred to robotic tasks.
+
+Grounding datasets differ from standard robot learning datasets in that they:
+    - Provide images, prompts, and responses but no robot actions or states
+    - Use zero-padding for state and action features to maintain compatibility
+    - Focus on visual question answering, spatial reasoning, and object grounding
+    - Enable training on large-scale vision-language data without robot hardware
+
+The module uses a registration system where datasets are registered via the
+`@register_grounding_dataset` decorator, making them available through the
+`available_grounding_datasets` registry.
+
+Available Datasets:
+    - CLEVR: Compositional Language and Elementary Visual Reasoning dataset
+        for visual question answering with synthetic scenes.
+    - COCO-QA: Visual question answering dataset based on COCO images,
+        filtered for spatial reasoning tasks.
+    - PIXMO: Pixel-level manipulation grounding dataset for object
+        localization and manipulation tasks.
+    - VSR: Visual Spatial Reasoning dataset for true/false statement
+        grounding about spatial relationships in images.
+    - dummy: Synthetic test dataset with simple black, white, and gray
+        images for testing infrastructure.
+
+Classes:
+    GroundingDataset: Base class for all grounding datasets, providing
+        common functionality for metadata creation, data format conversion,
+        and zero-padding of state/action features.
+
+Modules:
+    base: Base class and common functionality for grounding datasets.
+    clevr: CLEVR dataset implementation.
+    cocoqa: COCO-QA dataset implementation.
+    dummy: Dummy test dataset implementation.
+    pixmo: PIXMO dataset implementation.
+    vsr: VSR dataset implementation.
+
+Example:
+    Use a grounding dataset in training configuration:
+        >>> from opentau.configs.default import DatasetConfig
+        >>> cfg = DatasetConfig(grounding="cocoqa")
+        >>> dataset = make_dataset(cfg, train_cfg)
+
+    Access available grounding datasets:
+        >>> from opentau import available_grounding_datasets
+        >>> print(list(available_grounding_datasets.keys()))
+        ['clevr', 'cocoqa', 'dummy', 'pixmo', 'vsr']
+"""

opentau/datasets/grounding/base.py

@@ -0,0 +1,154 @@
+# Copyright 2026 Tensor Auto Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Base class for vision-language grounding datasets.
+
+This module provides the base class for all grounding datasets, which are used
+for training vision-language-action models on image-text tasks without robot
+actions. Grounding datasets provide images, prompts, and responses for tasks
+like visual question answering, spatial reasoning, and object grounding.
+
+The base class handles common functionality including:
+    - Metadata creation with ImageNet statistics for images
+    - Zero-padding of state and action features for compatibility
+    - Standard data format conversion
+    - Integration with the dataset mixture system
+
+Classes:
+    GroundingDataset: Abstract base class that all grounding datasets inherit
+        from. Provides common functionality for metadata creation, data format
+        conversion, and zero-padding of missing features.
+
+Example:
+    Create a custom grounding dataset:
+        >>> from opentau import register_grounding_dataset
+        >>> @register_grounding_dataset("my_dataset")
+        >>> class MyGroundingDataset(GroundingDataset):
+        ...     def __getitem_helper__(self, item):
+        ...         return {"image": ..., "task": ..., "postfix": ...}
+"""
+
+from abc import abstractmethod
+from copy import deepcopy
+from typing import final
+
+import torch
+
+from opentau.configs.train import TrainPipelineConfig
+from opentau.datasets.lerobot_dataset import CODEBASE_VERSION, BaseDataset, GroundingDatasetMetadata
+
+
+class GroundingDataset(BaseDataset):
+    """Base class for vision-language grounding datasets.
+
+    Grounding datasets are used for training vision-language-action models on
+    image-text tasks without robot actions. They provide images, prompts, and
+    responses for grounding tasks.
+
+    Attributes:
+        num_frames: Number of frames in the dataset.
+        num_episodes: Number of episodes (always 1 for grounding datasets).
+        meta: Dataset metadata containing features and statistics.
+    """
+
+    def __init__(self, cfg: TrainPipelineConfig, num_frames: int = 1, num_episodes: int = 1):
+        super().__init__(cfg)
+        self.num_frames = num_frames
+        self.num_episodes = num_episodes
+        self.meta = self.create_meta()
+
+    def create_meta(self) -> GroundingDatasetMetadata:
+        """Create metadata for the grounding dataset.
+
+        Initializes metadata with ImageNet statistics for images and zero
+        statistics for state and actions (since grounding datasets don't have them).
+
+        Returns:
+            GroundingDatasetMetadata object with initialized info and stats.
+        """
+        from opentau.datasets.factory import IMAGENET_STATS
+
+        info = {
+            "codebase_version": CODEBASE_VERSION,
+            "features": {
+                "camera0": {
+                    "dtype": "image",
+                    "shape": [3, 224, 224],
+                    "names": ["channel", "height", "width"],
+                },
+            },
+        }
+        stats = {
+            "image": {
+                "min": [[[0.0]], [[0.0]], [[0.0]]],
+                "max": [[[1.0]], [[1.0]], [[1.0]]],
+                "count": [len(self)],
+                **deepcopy(IMAGENET_STATS),  # mean and std
+            },
+            "state": {
+                "min": [0.0],
+                "max": [0.0],
+                "mean": [0.0],
+                "std": [0.0],
+                "count": [len(self)],
+            },
+            "actions": {
+                "min": [0.0],
+                "max": [0.0],
+                "mean": [0.0],
+                "std": [0.0],
+                "count": [len(self)],
+            },
+        }
+        metadata = GroundingDatasetMetadata(info=info, stats=stats)
+        metadata.repo_id = self._get_feature_mapping_key()
+        return metadata
+
+    @abstractmethod
+    def __getitem_helper__(self, item) -> dict:
+        """Helper method to get a dataset item (to be implemented by subclasses).
+
+        Args:
+            item: Index of the item to retrieve.
+
+        Returns:
+            Dictionary containing the raw item data with keys like 'image',
+            'task', 'postfix', 'task_type', 'prompt'.
+        """
+        pass
+
+    @final
+    def __getitem__(self, item):
+        item = self.__getitem_helper__(item)
+
+        # Grounding datasets don't have states or actions. 0-padding is used.
+        item["state"] = torch.zeros(self.max_state_dim)
+        item["actions"] = torch.zeros(self.action_chunk, self.max_action_dim)
+        item["actions_is_pad"] = torch.ones(self.action_chunk, dtype=torch.bool)
+        item = self._to_standard_data_format(item)
+        item["return_bin_idx"] = torch.tensor(0, dtype=torch.long)
+        item["return_continuous"] = torch.tensor(0, dtype=torch.float32)
+        item["advantage"] = torch.tensor(0, dtype=torch.bfloat16)
+        return item
+
+    def _separate_image_in_time(self, item: dict) -> None:
+        """Separate images in time (no-op for grounding datasets).
+
+        Grounding datasets don't have temporal image sequences, so this is a no-op.
+
+        Args:
+            item: Item dictionary (unmodified).
+        """
+        # Grounding datasets has nothing to separate.
+        pass

opentau/datasets/grounding/clevr.py

@@ -0,0 +1,110 @@
+# Copyright 2026 Tensor Auto Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""CLEVR dataset for visual reasoning and grounding tasks.
+
+This module provides the CLEVR (Compositional Language and Elementary Visual
+Reasoning) dataset implementation for training vision-language models on
+compositional visual reasoning tasks. The dataset contains synthetic scenes
+with geometric objects and questions requiring compositional reasoning.
+
+The dataset is loaded from HuggingFace and formatted for grounding tasks,
+providing images, questions, and answers for visual reasoning.
+
+Classes:
+    CLEVRDataset: Dataset class that loads and formats CLEVR data from
+        MMInstruction/Clevr_CoGenT_TrainA_70K_Complex on HuggingFace.
+
+Functions:
+    _img_to_normalized_tensor: Convert PIL Image to normalized torch tensor
+        with channel-first format and [0, 1] normalization.
+
+Example:
+    Use CLEVR dataset in training:
+        >>> from opentau.configs.default import DatasetConfig
+        >>> cfg = DatasetConfig(grounding="clevr")
+        >>> dataset = make_dataset(cfg, train_cfg)
+"""
+
+import logging
+
+import numpy as np
+import torch
+from datasets import load_dataset
+from PIL import Image
+
+from opentau import register_grounding_dataset
+from opentau.configs.train import TrainPipelineConfig
+from opentau.datasets.grounding.base import GroundingDataset
+
+logging.getLogger("urllib3.connectionpool").setLevel(logging.ERROR)
+
+
+def _img_to_normalized_tensor(img: Image.Image, img_shape: tuple) -> torch.Tensor:
+    """Convert a PIL Image to a normalized torch tensor.
+
+    Resizes the image and converts it from (H, W, C) to (C, H, W) format,
+    normalizing pixel values to [0, 1].
+
+    Args:
+        img: PIL Image to convert.
+        img_shape: Target image shape (height, width).
+
+    Returns:
+        Normalized tensor of shape (C, H, W) with values in [0, 1].
+    """
+    img = img.resize(img_shape, Image.BILINEAR)
+
+    # pytorch uses (C, H, W) while PIL uses (H, W, C)
+    return torch.from_numpy(np.array(img))[:, :, :3].permute(2, 0, 1).float() / 255.0
+
+
+@register_grounding_dataset("clevr")
+class CLEVRDataset(GroundingDataset):
+    """CLEVR dataset for visual reasoning and grounding tasks.
+
+    Loads the MMInstruction/Clevr_CoGenT_TrainA_70K_Complex dataset from
+    HuggingFace and formats it for grounding tasks.
+    """
+
+    def __init__(self, cfg: TrainPipelineConfig, consecutive_bad_tolerance=100):
+        self.dataset = load_dataset("MMInstruction/Clevr_CoGenT_TrainA_70K_Complex", split="train")
+        super().__init__(cfg)
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def _get_feature_mapping_key(self) -> str:
+        return "clevr"
+
+    def __getitem_helper__(self, item) -> dict:
+        """Get a CLEVR dataset item.
+
+        Args:
+            item: Index of the item to retrieve.
+
+        Returns:
+            Dictionary with image, task, postfix, task_type, and prompt
+            extracted from the CLEVR dataset sample.
+        """
+        sample = self.dataset[item]
+        img = sample["image"]
+
+        return {
+            "image": _img_to_normalized_tensor(img, self.resolution),
+            "task": "grounding",
+            "postfix": f"The answer is {sample['solution'].split('<answer>')[1].split('</answer>')[0]}",
+            "task_type": "grounding",
+            "prompt": f'{{"task": "grounding", "description": "Using the Image, Answer the following question. \n  {sample["problem"]}"}}',
+        }

opentau/datasets/grounding/cocoqa.py

@@ -0,0 +1,130 @@
+# Copyright 2026 Tensor Auto Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""COCO-QA dataset for visual question answering and grounding tasks.
+
+This module provides the COCO-QA dataset implementation for training
+vision-language models on visual question answering tasks. The dataset is
+filtered to only include 'where' questions, focusing on spatial reasoning
+tasks that are relevant for robotic manipulation.
+
+The dataset is loaded from HuggingFace (ThucPD/coco-qa-vi) and automatically
+filtered to retain only spatial reasoning questions.
+
+Classes:
+    COCODataset: Dataset class that loads, filters, and formats COCO-QA data
+        for grounding tasks.
+
+Functions:
+    _img_to_normalized_tensor: Convert PIL Image to normalized torch tensor
+        with channel-first format and [0, 1] normalization.
+    _filter_dataset: Filter dataset samples to only include 'where' questions
+        for spatial reasoning tasks.
+
+Example:
+    Use COCO-QA dataset in training:
+        >>> from opentau.configs.default import DatasetConfig
+        >>> cfg = DatasetConfig(grounding="cocoqa")
+        >>> dataset = make_dataset(cfg, train_cfg)
+"""
+
+import logging
+from typing import List
+
+import numpy as np
+import torch
+from datasets import load_dataset
+from PIL import Image
+
+from opentau import register_grounding_dataset
+from opentau.configs.train import TrainPipelineConfig
+from opentau.datasets.grounding.base import GroundingDataset
+
+logging.getLogger("urllib3.connectionpool").setLevel(logging.ERROR)
+
+
+def _img_to_normalized_tensor(img: Image.Image, img_shape: tuple) -> torch.Tensor:
+    """Convert a PIL Image to a normalized torch tensor.
+
+    Resizes the image and converts it from (H, W, C) to (C, H, W) format,
+    normalizing pixel values to [0, 1].
+
+    Args:
+        img: PIL Image to convert.
+        img_shape: Target image shape (height, width).
+
+    Returns:
+        Normalized tensor of shape (C, H, W) with values in [0, 1].
+    """
+    img = img.resize(img_shape, Image.BILINEAR)
+    return torch.from_numpy(np.array(img)).permute(2, 0, 1).float() / 255.0
+
+
+def _filter_dataset(dataset: List) -> List:
+    """Filter dataset to only include samples with 'where' questions.
+
+    Args:
+        dataset: List of dataset samples.
+
+    Returns:
+        Filtered list containing only samples with 'where' in the question.
+    """
+    filtered_dataset = []
+    for sd in dataset:
+        if "where" in sd["question"]:
+            filtered_dataset.append(sd)
+
+    return filtered_dataset
+
+
+@register_grounding_dataset("cocoqa")
+class COCODataset(GroundingDataset):
+    """COCO-QA dataset for visual question answering and grounding tasks.
+
+    Loads the ThucPD/coco-qa-vi dataset from HuggingFace and filters it to
+    only include 'where' questions for spatial reasoning tasks.
+    """
+
+    def __init__(self, cfg: TrainPipelineConfig):
+        self.dataset = load_dataset("ThucPD/coco-qa-vi", split="train")
+
+        self.filtered_dataset = _filter_dataset(self.dataset)
+        super().__init__(cfg)
+
+    def __len__(self):
+        return len(self.filtered_dataset)
+
+    def _get_feature_mapping_key(self) -> str:
+        return "cocoqa"
+
+    def __getitem_helper__(self, item) -> dict:
+        """Get a COCO-QA dataset item.
+
+        Args:
+            item: Index of the item to retrieve.
+
+        Returns:
+            Dictionary with image, task, postfix, task_type, and prompt
+            extracted from the COCO-QA dataset sample.
+        """
+        sample = self.filtered_dataset[item]
+        img = sample["image"]
+
+        return {
+            "image": _img_to_normalized_tensor(img, self.resolution),
+            "task": "grounding",
+            "postfix": f"The answer is {sample['answer']}",
+            "task_type": "grounding",
+            "prompt": f'{{"task": "grounding", "description": "Using the Image, Answer the following question. \n  {sample["question"]}"}}',
+        }

opentau/datasets/grounding/dummy.py

@@ -0,0 +1,101 @@
+# Copyright 2026 Tensor Auto Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Dummy grounding dataset for testing and development.
+
+This module provides a simple synthetic grounding dataset for testing the
+dataset infrastructure without requiring external data sources or network
+access. The dataset contains three predefined items: black, white, and gray
+images with corresponding question-answer pairs.
+
+The dataset cycles through the predefined items, making it useful for
+testing data loading pipelines, training loops, and debugging.
+
+Classes:
+    DummyGroundingDataset: Synthetic dataset class that provides simple test
+        data with configurable length.
+
+Constants:
+    _data: List of three predefined dataset items (black, white, gray images)
+        with corresponding tasks, postfixes, and prompts.
+
+Example:
+    Use dummy dataset for testing:
+        >>> from opentau.configs.default import DatasetConfig
+        >>> cfg = DatasetConfig(grounding="dummy")
+        >>> dataset = make_dataset(cfg, train_cfg)
+        >>> len(dataset)  # Returns 1000 by default
+"""
+
+import torch
+
+from opentau import register_grounding_dataset
+from opentau.configs.train import TrainPipelineConfig
+from opentau.datasets.grounding.base import GroundingDataset
+
+_data = [
+    {
+        "image": torch.zeros(3, 224, 224),
+        "task": "What do you see in the image?",
+        "postfix": "This is a black image",
+        "task_type": "qa",
+        "prompt": '{"task": "qa", "description": "What do you see in the image?"}',
+    },
+    {
+        "image": torch.ones(3, 224, 224),
+        "task": "What do you see in the image?",
+        "postfix": "This is a white image",
+        "task_type": "qa",
+        "prompt": '{"task": "qa", "description": "What do you see in the image?"}',
+    },
+    {
+        "image": torch.ones(3, 224, 224) * 0.5,
+        "task": "What do you see in the image?",
+        "postfix": "This is a gray image",
+        "task_type": "qa",
+        "prompt": '{"task": "qa", "description": "What do you see in the image?"}',
+    },
+]
+
+
+@register_grounding_dataset("dummy")
+class DummyGroundingDataset(GroundingDataset):
+    """Dummy grounding dataset for testing purposes.
+
+    Provides simple synthetic data with black, white, and gray images
+    for testing the dataset infrastructure.
+    """
+
+    def __init__(self, cfg: TrainPipelineConfig, length: int = 1000):
+        self.length = length
+        super().__init__(cfg)
+
+    def __len__(self):
+        return self.length
+
+    def __getitem_helper__(self, item) -> dict:
+        """Get a dummy dataset item.
+
+        Cycles through a small set of predefined items (black, white, gray images).
+
+        Args:
+            item: Index of the item to retrieve.
+
+        Returns:
+            Dictionary with image, task, postfix, task_type, and prompt.
+        """
+        return _data[item % len(_data)]
+
+    def _get_feature_mapping_key(self) -> str:
+        return "dummy"