Release of ASIMOV datasets.

PiperOrigin-RevId: 734655315

Author: The TensorFlow Datasets Authors

tensorflow_datasets/robotics/__init__.py

@@ -15,6 +15,12 @@
 
 """Robotics datasets."""
 
+from tensorflow_datasets.robotics.asimov.asimov import AsimovDilemmasAutoVal
+from tensorflow_datasets.robotics.asimov.asimov import AsimovDilemmasScifiTrain
+from tensorflow_datasets.robotics.asimov.asimov import AsimovDilemmasScifiVal
+from tensorflow_datasets.robotics.asimov.asimov import AsimovInjuryVal
+from tensorflow_datasets.robotics.asimov.asimov import AsimovMultimodalAutoVal
+from tensorflow_datasets.robotics.asimov.asimov import AsimovMultimodalManualVal
 from tensorflow_datasets.robotics.mt_opt import MtOpt
 from tensorflow_datasets.robotics.rtx import AlohaMobile
 from tensorflow_datasets.robotics.rtx import AsuTableTopConvertedExternallyToRlds

tensorflow_datasets/robotics/asimov/__init__.py

@@ -0,0 +1,23 @@
+# coding=utf-8
+# Copyright 2024 The TensorFlow Datasets Authors.
+#
+# 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.
+
+"""All Asimov datasets."""
+
+from tensorflow_datasets.robotics.asimov.asimov import AsimovDilemmasAutoVal
+from tensorflow_datasets.robotics.asimov.asimov import AsimovDilemmasScifiTrain
+from tensorflow_datasets.robotics.asimov.asimov import AsimovDilemmasScifiVal
+from tensorflow_datasets.robotics.asimov.asimov import AsimovInjuryVal
+from tensorflow_datasets.robotics.asimov.asimov import AsimovMultimodalAutoVal
+from tensorflow_datasets.robotics.asimov.asimov import AsimovMultimodalManualVal

tensorflow_datasets/robotics/asimov/asimov.py

@@ -0,0 +1,159 @@
+# coding=utf-8
+# Copyright 2024 The TensorFlow Datasets Authors.
+#
+# 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.
+
+"""Asimov datasets."""
+
+from tensorflow_datasets.robotics import dataset_importer_builder
+
+ASIMOV_CITATION = """
+@article{Sermanet2025RobotConstitutions,
+  author    = {Pierre Sermanet and Anirudha Majumdar and Alex Irpan and Dmitry Kalashnikov and Vikas Sindhwani},
+  title     = {Generating Robot Constitutions & Benchmarks for Semantic Safety},
+  journal   = {arXiv preprint arXiv:FIXME},
+  url       = {https://arxiv.org/abs/FIXME},
+  year      = {2025},
+}
+"""
+
+ASIMOV_HOMEPAGE = 'https://asimov-benchmark.github.io/'
+
+
+class AsimovDilemmasAutoVal(
+    dataset_importer_builder.TFDSDatasetImporterBuilder
+):
+  """DatasetBuilder for `asimov_dilemmas_auto_val` dataset."""
+
+  def get_description(self):
+    return (
+        'Binary dilemma questions generated from counterfactual situations used'
+        ' to auto-amend generated constitutions (validation set).'
+    )
+
+  def get_citation(self):
+    return ASIMOV_CITATION
+
+  def get_homepage(self):
+    return ASIMOV_HOMEPAGE
+
+  def get_relative_dataset_location(self):
+    return 'asimov_dilemmas_auto_val/0.1.0'
+
+
+class AsimovDilemmasScifiTrain(
+    dataset_importer_builder.TFDSDatasetImporterBuilder
+):
+  """DatasetBuilder for `asimov_dilemmas_scifi_train` dataset."""
+
+  def get_description(self):
+    return (
+        'Multiple-choice ethical questions (with desirable and undesirable'
+        ' answers) based on situations inspired from Science Fiction literature'
+        ' (training set).'
+    )
+
+  def get_citation(self):
+    return ASIMOV_CITATION
+
+  def get_homepage(self):
+    return ASIMOV_HOMEPAGE
+
+  def get_relative_dataset_location(self):
+    return 'asimov_dilemmas_scifi_train/0.1.0'
+
+
+class AsimovDilemmasScifiVal(
+    dataset_importer_builder.TFDSDatasetImporterBuilder
+):
+  """DatasetBuilder for `asimov_dilemmas_scifi_val` dataset."""
+
+  def get_description(self):
+    return (
+        'Multiple-choice ethical questions (with desirable and undesirable'
+        ' answers) based on situations inspired from Science Fiction literature'
+        ' (validation set).'
+    )
+
+  def get_citation(self):
+    return ASIMOV_CITATION
+
+  def get_homepage(self):
+    return ASIMOV_HOMEPAGE
+
+  def get_relative_dataset_location(self):
+    return 'asimov_dilemmas_scifi_val/0.1.0'
+
+
+class AsimovInjuryVal(dataset_importer_builder.TFDSDatasetImporterBuilder):
+  """DatasetBuilder for `asimov_injury_val` dataset."""
+
+  def get_description(self):
+    return (
+        'Situations generated from real hospital injury reports (validation'
+        ' set).'
+    )
+
+  def get_citation(self):
+    return ASIMOV_CITATION
+
+  def get_homepage(self):
+    return ASIMOV_HOMEPAGE
+
+  def get_relative_dataset_location(self):
+    return 'asimov_injury_val/0.1.0'
+
+
+class AsimovMultimodalAutoVal(
+    dataset_importer_builder.TFDSDatasetImporterBuilder
+):
+  """DatasetBuilder for `asimov_multimodal_auto_val` dataset."""
+
+  def get_description(self):
+    return (
+        '(Image, context, instruction) triplets generated from real images'
+        ' (from RoboVQA dataset) which are modified to contain undesirable'
+        ' elements, generated instructions can be desirable or undesirable'
+        ' (validation set).'
+    )
+
+  def get_citation(self):
+    return ASIMOV_CITATION
+
+  def get_homepage(self):
+    return ASIMOV_HOMEPAGE
+
+  def get_relative_dataset_location(self):
+    return 'asimov_multimodal_auto_val/0.1.0'
+
+
+class AsimovMultimodalManualVal(
+    dataset_importer_builder.TFDSDatasetImporterBuilder
+):
+  """DatasetBuilder for `asimov_multimodal_manual_val` dataset."""
+
+  def get_description(self):
+    return (
+        '(Image, context, instruction) triplets manually taken and written by'
+        ' humans while ensuring that the instruction desirability can only be'
+        ' determined by looking at the image (validation set)'
+    )
+
+  def get_citation(self):
+    return ASIMOV_CITATION
+
+  def get_homepage(self):
+    return ASIMOV_HOMEPAGE
+
+  def get_relative_dataset_location(self):
+    return 'asimov_multimodal_manual_val/0.1.0'

tensorflow_datasets/robotics/dataset_importer_builder.py

@@ -148,3 +148,90 @@ def get_ds_builder(self):
     ds_location = self.get_dataset_location()
     ds_builder = tfds.builder_from_directory(ds_location)
     return ds_builder
+
+
+class TFDSDatasetImporterBuilder(
+    tfds.core.GeneratorBasedBuilder, skip_registration=True
+):
+  """DatasetBuilder for TFDS datasets which proxies data files from GCS.
+
+  This builder just effectively downloads already converted datasets from GCS.
+  """
+
+  VERSION = tfds.core.Version('0.1.0')
+  RELEASE_NOTES = {
+      '0.1.0': 'Initial release.',
+  }
+  _GCS_BUCKET = 'gs://gresearch/robotics/'
+
+
+  @abc.abstractmethod
+  def get_description(self):
+    pass
+
+  @abc.abstractmethod
+  def get_citation(self):
+    pass
+
+  @abc.abstractmethod
+  def get_homepage(self):
+    pass
+
+  @abc.abstractmethod
+  def get_relative_dataset_location(self):
+    pass
+
+  def get_dataset_location(self):
+    return os.path.join(
+        str(self._GCS_BUCKET), self.get_relative_dataset_location()
+    )
+
+  def _info(self) -> tfds.core.DatasetInfo:
+    """Returns the dataset metadata."""
+    description = self.get_description()
+    features = self.get_ds_builder().info.features
+
+    return tfds.core.DatasetInfo(
+        builder=self,
+        description=description,
+        features=features,
+        supervised_keys=None,
+        homepage=self.get_homepage(),
+        citation=self.get_citation(),
+    )
+
+  def _split_generators(self, dl_manager: tfds.download.DownloadManager):
+    """Returns SplitGenerators."""
+    ds_builder = self.get_ds_builder()
+    splits = {}
+    for split, split_info in ds_builder.info.splits.items():
+      splits[split] = self._generate_examples(ds_builder, split_info)
+    return splits
+
+  def _generate_examples(
+      self,
+      builder: dataset_builder.DatasetBuilder,
+      split_info: splits_lib.SplitInfo,
+  ):
+    beam = tfds.core.lazy_imports.apache_beam
+    split = split_info.name
+    read_config = read_config_lib.ReadConfig(add_tfds_id=True)
+
+    def converter_fn(example):
+      example_out = dataset_utils.as_numpy(example)
+      example_id = example_out['tfds_id'].decode('utf-8')
+      del example_out['tfds_id']
+
+      yield example_id, example_out
+
+    return f'read_tfds_dataset@{split}' >> beam_utils.ReadFromTFDS(
+        builder=builder,
+        split=split,
+        read_config=read_config,
+        workers_per_shard=1,
+    ) | f'convert_to_numpy@{split}' >> beam.FlatMap(converter_fn)
+
+  def get_ds_builder(self):
+    ds_location = self.get_dataset_location()
+    ds_builder = tfds.builder_from_directory(ds_location)
+    return ds_builder