diff --git a/tensorflow_datasets/datasets/smart_buildings/controller_reader.py b/tensorflow_datasets/datasets/smart_buildings/controller_reader.py
index c1f671faefb..f100c1e2da7 100644
--- a/tensorflow_datasets/datasets/smart_buildings/controller_reader.py
+++ b/tensorflow_datasets/datasets/smart_buildings/controller_reader.py
@@ -25,9 +25,9 @@
 import pandas as pd
 from tensorflow_datasets.datasets.smart_buildings import constants
 from tensorflow_datasets.datasets.smart_buildings import reader_lib
-from tensorflow_datasets.datasets.smart_buildings import smart_control_building_pb2
-from tensorflow_datasets.datasets.smart_buildings import smart_control_normalization_pb2
-from tensorflow_datasets.datasets.smart_buildings import smart_control_reward_pb2
+from tensorflow_datasets.proto import smart_control_building_pb2
+from tensorflow_datasets.proto import smart_control_normalization_pb2
+from tensorflow_datasets.proto import smart_control_reward_pb2
 
 
 T = TypeVar('T')
diff --git a/tensorflow_datasets/datasets/smart_buildings/reader_lib.py b/tensorflow_datasets/datasets/smart_buildings/reader_lib.py
index 1637ce423c9..3eca8a7364f 100644
--- a/tensorflow_datasets/datasets/smart_buildings/reader_lib.py
+++ b/tensorflow_datasets/datasets/smart_buildings/reader_lib.py
@@ -36,9 +36,9 @@
 
 from absl import logging
 import pandas as pd
-from tensorflow_datasets.datasets.smart_buildings import smart_control_building_pb2
-from tensorflow_datasets.datasets.smart_buildings import smart_control_normalization_pb2
-from tensorflow_datasets.datasets.smart_buildings import smart_control_reward_pb2
+from tensorflow_datasets.proto import smart_control_building_pb2
+from tensorflow_datasets.proto import smart_control_normalization_pb2
+from tensorflow_datasets.proto import smart_control_reward_pb2
 
 VariableId = NewType('VariableId', str)
 
diff --git a/tensorflow_datasets/datasets/smart_buildings/smart_buildings_dataset_builder_test.py b/tensorflow_datasets/datasets/smart_buildings/smart_buildings_dataset_builder_test.py
index f94e1456f03..dc883237cb6 100644
--- a/tensorflow_datasets/datasets/smart_buildings/smart_buildings_dataset_builder_test.py
+++ b/tensorflow_datasets/datasets/smart_buildings/smart_buildings_dataset_builder_test.py
@@ -15,19 +15,19 @@
 
 """smart_buildings_dataset dataset."""
 
-# from tensorflow_datasets.datasets.smart_buildings
-# import smart_buildingsdataset_builder
-# import tensorflow_datasets.public_api as tfds
+from tensorflow_datasets.datasets.smart_buildings import smart_buildings_dataset_builder
+import tensorflow_datasets.public_api as tfds
 
 
-# class SmartBuildingsDatasetTest(tfds.testing.DatasetBuilderTestCase):
-#   """Tests for smart_buildings_dataset dataset."""
+class SmartBuildingsDatasetTest(tfds.testing.DatasetBuilderTestCase):
+  """Tests for smart_buildings_dataset dataset."""
+
+  DATASET_CLASS = smart_buildings_dataset_builder.Builder
+  SPLITS = {
+      'sb1_23': 12,  # Number of fake train example
+      'sb1_24': 12,  # Number of fake train example
+  }
 
-#   DATASET_CLASS = smart_buildings_dataset_dataset_builder.Builder
-#   SPLITS = {
-#       'sb1_23': 12,  # Number of fake train example
-#       'sb1_24': 12,  # Number of fake train example
-#   }
 
 # If you are calling `download/download_and_extract` with a dict, like:
 #   dl_manager.download({'some_key': 'http://a.org/out.txt', ...})
@@ -37,4 +37,4 @@
 
 
 if __name__ == '__main__':
-  pass  # tfds.testing.test_main()
+  tfds.testing.test_main()
diff --git a/tensorflow_datasets/datasets/smart_buildings/smart_control_building.proto b/tensorflow_datasets/proto/smart_control_building.proto
similarity index 100%
rename from tensorflow_datasets/datasets/smart_buildings/smart_control_building.proto
rename to tensorflow_datasets/proto/smart_control_building.proto
diff --git a/tensorflow_datasets/proto/smart_control_building_pb2.py b/tensorflow_datasets/proto/smart_control_building_pb2.py
new file mode 100644
index 00000000000..03e4ad7d2ec
--- /dev/null
+++ b/tensorflow_datasets/proto/smart_control_building_pb2.py
@@ -0,0 +1,75 @@
+# 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.
+
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: smart_control_building.proto
+"""Generated protocol buffer code."""
+from google.protobuf.internal import builder as _builder
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import symbol_database as _symbol_database
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+from google.protobuf import timestamp_pb2 as google_dot_protobuf_dot_timestamp__pb2
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1csmart_control_building.proto\x12#smart_buildings.smart_control.proto\x1a\x1fgoogle/protobuf/timestamp.proto\"\xfe\x01\n\x08ZoneInfo\x12\x0f\n\x07zone_id\x18\x01 \x01(\t\x12\x13\n\x0b\x62uilding_id\x18\x02 \x01(\t\x12\x18\n\x10zone_description\x18\x03 \x01(\t\x12\x0c\n\x04\x61rea\x18\x04 \x01(\x02\x12\x0f\n\x07\x64\x65vices\x18\x05 \x03(\t\x12I\n\tzone_type\x18\x06 \x01(\x0e\x32\x36.smart_buildings.smart_control.proto.ZoneInfo.ZoneType\x12\r\n\x05\x66loor\x18\x07 \x01(\x05\"9\n\x08ZoneType\x12\r\n\tUNDEFINED\x10\x00\x12\x08\n\x04ROOM\x10\x01\x12\t\n\x05\x46LOOR\x10\x02\x12\t\n\x05OTHER\x10\n\"\xa5\x07\n\nDeviceInfo\x12\x11\n\tdevice_id\x18\x01 \x01(\t\x12\x11\n\tnamespace\x18\x02 \x01(\t\x12\x0c\n\x04\x63ode\x18\x03 \x01(\t\x12\x0f\n\x07zone_id\x18\x04 \x01(\t\x12O\n\x0b\x64\x65vice_type\x18\x05 \x01(\x0e\x32:.smart_buildings.smart_control.proto.DeviceInfo.DeviceType\x12`\n\x11observable_fields\x18\x06 \x03(\x0b\x32\x45.smart_buildings.smart_control.proto.DeviceInfo.ObservableFieldsEntry\x12X\n\raction_fields\x18\x07 \x03(\x0b\x32\x41.smart_buildings.smart_control.proto.DeviceInfo.ActionFieldsEntry\x1ar\n\x15ObservableFieldsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12H\n\x05value\x18\x02 \x01(\x0e\x32\x39.smart_buildings.smart_control.proto.DeviceInfo.ValueType:\x02\x38\x01\x1an\n\x11\x41\x63tionFieldsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12H\n\x05value\x18\x02 \x01(\x0e\x32\x39.smart_buildings.smart_control.proto.DeviceInfo.ValueType:\x02\x38\x01\"\xe7\x01\n\nDeviceType\x12\r\n\tUNDEFINED\x10\x00\x12\x07\n\x03\x46\x41N\x10\x01\x12\x07\n\x03PMP\x10\x02\x12\x07\n\x03\x46\x43U\x10\x03\x12\x07\n\x03VAV\x10\x04\x12\x06\n\x02\x44H\x10\x05\x12\x07\n\x03\x41HU\x10\x06\x12\x07\n\x03\x42LR\x10\x07\x12\x08\n\x04\x43\x44WS\x10\x08\x12\x06\n\x02\x43H\x10\t\x12\x08\n\x04\x43HWS\x10\n\x12\x06\n\x02\x43T\x10\x0b\x12\x06\n\x02\x44\x43\x10\x0c\x12\x07\n\x03\x44\x46R\x10\r\x12\x07\n\x03\x44MP\x10\x0e\x12\x07\n\x03HWS\x10\x0f\x12\x06\n\x02HX\x10\x10\x12\x07\n\x03MAU\x10\x11\x12\x07\n\x03SDC\x10\x12\x12\x06\n\x02UH\x10\x13\x12\x07\n\x03PWR\x10\x14\x12\x07\n\x03GAS\x10\x15\x12\x06\n\x02\x41\x43\x10\x16\x12\t\n\x05OTHER\x10\x17\"w\n\tValueType\x12\x18\n\x14VALUE_TYPE_UNDEFINED\x10\x00\x12\x14\n\x10VALUE_CONTINUOUS\x10\x01\x12\x11\n\rVALUE_INTEGER\x10\x02\x12\x15\n\x11VALUE_CATEGORICAL\x10\x03\x12\x10\n\x0cVALUE_BINARY\x10\x04\"G\n\x18SingleObservationRequest\x12\x11\n\tdevice_id\x18\x01 \x01(\t\x12\x18\n\x10measurement_name\x18\x02 \x01(\t\"\xdf\x02\n\x19SingleObservationResponse\x12-\n\ttimestamp\x18\x01 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x61\n\x1asingle_observation_request\x18\x02 \x01(\x0b\x32=.smart_buildings.smart_control.proto.SingleObservationRequest\x12\x19\n\x11observation_valid\x18\x03 \x01(\x08\x12\x1a\n\x10\x63ontinuous_value\x18\x04 \x01(\x02H\x00\x12\x17\n\rinteger_value\x18\x05 \x01(\x05H\x00\x12\x1b\n\x11\x63\x61tegorical_value\x18\x06 \x01(\tH\x00\x12\x16\n\x0c\x62inary_value\x18\x07 \x01(\x08H\x00\x12\x16\n\x0cstring_value\x18\x08 \x01(\tH\x00\x42\x13\n\x11observation_value\"\xa7\x01\n\x12ObservationRequest\x12-\n\ttimestamp\x18\x01 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x62\n\x1bsingle_observation_requests\x18\x02 \x03(\x0b\x32=.smart_buildings.smart_control.proto.SingleObservationRequest\"\xf4\x01\n\x13ObservationResponse\x12-\n\ttimestamp\x18\x01 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12H\n\x07request\x18\x02 \x01(\x0b\x32\x37.smart_buildings.smart_control.proto.ObservationRequest\x12\x64\n\x1csingle_observation_responses\x18\x03 \x03(\x0b\x32>.smart_buildings.smart_control.proto.SingleObservationResponse\"\xd3\x01\n\x13SingleActionRequest\x12\x11\n\tdevice_id\x18\x01 \x01(\t\x12\x15\n\rsetpoint_name\x18\x02 \x01(\t\x12\x1a\n\x10\x63ontinuous_value\x18\x03 \x01(\x02H\x00\x12\x17\n\rinteger_value\x18\x04 \x01(\x05H\x00\x12\x1b\n\x11\x63\x61tegorical_value\x18\x05 \x01(\tH\x00\x12\x16\n\x0c\x62inary_value\x18\x06 \x01(\x08H\x00\x12\x16\n\x0cstring_value\x18\x07 \x01(\tH\x00\x42\x10\n\x0esetpoint_value\"\xdd\x03\n\x14SingleActionResponse\x12I\n\x07request\x18\x01 \x01(\x0b\x32\x38.smart_buildings.smart_control.proto.SingleActionRequest\x12\x63\n\rresponse_type\x18\x02 \x01(\x0e\x32L.smart_buildings.smart_control.proto.SingleActionResponse.ActionResponseType\x12\x17\n\x0f\x61\x64\x64itional_info\x18\x03 \x01(\t\"\xfb\x01\n\x12\x41\x63tionResponseType\x12\r\n\tUNDEFINED\x10\x00\x12\x0c\n\x08\x41\x43\x43\x45PTED\x10\x01\x12\x0b\n\x07PENDING\x10\x02\x12\r\n\tTIMED_OUT\x10\x03\x12\x1c\n\x18REJECTED_INVALID_SETTING\x10\x04\x12%\n!REJECTED_NOT_ENABLED_OR_AVAILABLE\x10\x05\x12\x15\n\x11REJECTED_OVERRIDE\x10\x06\x12\x1b\n\x17REJECTED_INVALID_DEVICE\x10\x07\x12\x1b\n\x17REJECTED_DEVICE_OFFLINE\x10\x08\x12\x0b\n\x07UNKNOWN\x10\t\x12\t\n\x05OTHER\x10\n\"\x98\x01\n\rActionRequest\x12-\n\ttimestamp\x18\x01 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12X\n\x16single_action_requests\x18\x02 \x03(\x0b\x32\x38.smart_buildings.smart_control.proto.SingleActionRequest\"\xe0\x01\n\x0e\x41\x63tionResponse\x12-\n\ttimestamp\x18\x01 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x43\n\x07request\x18\x02 \x01(\x0b\x32\x32.smart_buildings.smart_control.proto.ActionRequest\x12Z\n\x17single_action_responses\x18\x03 \x03(\x0b\x32\x39.smart_buildings.smart_control.proto.SingleActionResponseb\x06proto3')
+
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals())
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'smart_control_building_pb2', globals())
+if _descriptor._USE_C_DESCRIPTORS == False:
+
+  DESCRIPTOR._options = None
+  _DEVICEINFO_OBSERVABLEFIELDSENTRY._options = None
+  _DEVICEINFO_OBSERVABLEFIELDSENTRY._serialized_options = b'8\001'
+  _DEVICEINFO_ACTIONFIELDSENTRY._options = None
+  _DEVICEINFO_ACTIONFIELDSENTRY._serialized_options = b'8\001'
+  _ZONEINFO._serialized_start=103
+  _ZONEINFO._serialized_end=357
+  _ZONEINFO_ZONETYPE._serialized_start=300
+  _ZONEINFO_ZONETYPE._serialized_end=357
+  _DEVICEINFO._serialized_start=360
+  _DEVICEINFO._serialized_end=1293
+  _DEVICEINFO_OBSERVABLEFIELDSENTRY._serialized_start=712
+  _DEVICEINFO_OBSERVABLEFIELDSENTRY._serialized_end=826
+  _DEVICEINFO_ACTIONFIELDSENTRY._serialized_start=828
+  _DEVICEINFO_ACTIONFIELDSENTRY._serialized_end=938
+  _DEVICEINFO_DEVICETYPE._serialized_start=941
+  _DEVICEINFO_DEVICETYPE._serialized_end=1172
+  _DEVICEINFO_VALUETYPE._serialized_start=1174
+  _DEVICEINFO_VALUETYPE._serialized_end=1293
+  _SINGLEOBSERVATIONREQUEST._serialized_start=1295
+  _SINGLEOBSERVATIONREQUEST._serialized_end=1366
+  _SINGLEOBSERVATIONRESPONSE._serialized_start=1369
+  _SINGLEOBSERVATIONRESPONSE._serialized_end=1720
+  _OBSERVATIONREQUEST._serialized_start=1723
+  _OBSERVATIONREQUEST._serialized_end=1890
+  _OBSERVATIONRESPONSE._serialized_start=1893
+  _OBSERVATIONRESPONSE._serialized_end=2137
+  _SINGLEACTIONREQUEST._serialized_start=2140
+  _SINGLEACTIONREQUEST._serialized_end=2351
+  _SINGLEACTIONRESPONSE._serialized_start=2354
+  _SINGLEACTIONRESPONSE._serialized_end=2831
+  _SINGLEACTIONRESPONSE_ACTIONRESPONSETYPE._serialized_start=2580
+  _SINGLEACTIONRESPONSE_ACTIONRESPONSETYPE._serialized_end=2831
+  _ACTIONREQUEST._serialized_start=2834
+  _ACTIONREQUEST._serialized_end=2986
+  _ACTIONRESPONSE._serialized_start=2989
+  _ACTIONRESPONSE._serialized_end=3213
+# @@protoc_insertion_point(module_scope)
diff --git a/tensorflow_datasets/proto/smart_control_building_pb2.pyi b/tensorflow_datasets/proto/smart_control_building_pb2.pyi
new file mode 100644
index 00000000000..37de60e887b
--- /dev/null
+++ b/tensorflow_datasets/proto/smart_control_building_pb2.pyi
@@ -0,0 +1,202 @@
+from google.protobuf import timestamp_pb2 as _timestamp_pb2
+from google.protobuf.internal import containers as _containers
+from google.protobuf.internal import enum_type_wrapper as _enum_type_wrapper
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from typing import ClassVar as _ClassVar, Iterable as _Iterable, Mapping as _Mapping, Optional as _Optional, Union as _Union
+
+DESCRIPTOR: _descriptor.FileDescriptor
+
+class ActionRequest(_message.Message):
+    __slots__ = ["single_action_requests", "timestamp"]
+    SINGLE_ACTION_REQUESTS_FIELD_NUMBER: _ClassVar[int]
+    TIMESTAMP_FIELD_NUMBER: _ClassVar[int]
+    single_action_requests: _containers.RepeatedCompositeFieldContainer[SingleActionRequest]
+    timestamp: _timestamp_pb2.Timestamp
+    def __init__(self, timestamp: _Optional[_Union[_timestamp_pb2.Timestamp, _Mapping]] = ..., single_action_requests: _Optional[_Iterable[_Union[SingleActionRequest, _Mapping]]] = ...) -> None: ...
+
+class ActionResponse(_message.Message):
+    __slots__ = ["request", "single_action_responses", "timestamp"]
+    REQUEST_FIELD_NUMBER: _ClassVar[int]
+    SINGLE_ACTION_RESPONSES_FIELD_NUMBER: _ClassVar[int]
+    TIMESTAMP_FIELD_NUMBER: _ClassVar[int]
+    request: ActionRequest
+    single_action_responses: _containers.RepeatedCompositeFieldContainer[SingleActionResponse]
+    timestamp: _timestamp_pb2.Timestamp
+    def __init__(self, timestamp: _Optional[_Union[_timestamp_pb2.Timestamp, _Mapping]] = ..., request: _Optional[_Union[ActionRequest, _Mapping]] = ..., single_action_responses: _Optional[_Iterable[_Union[SingleActionResponse, _Mapping]]] = ...) -> None: ...
+
+class DeviceInfo(_message.Message):
+    __slots__ = ["action_fields", "code", "device_id", "device_type", "namespace", "observable_fields", "zone_id"]
+    class DeviceType(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
+        __slots__ = []
+    class ValueType(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
+        __slots__ = []
+    class ActionFieldsEntry(_message.Message):
+        __slots__ = ["key", "value"]
+        KEY_FIELD_NUMBER: _ClassVar[int]
+        VALUE_FIELD_NUMBER: _ClassVar[int]
+        key: str
+        value: DeviceInfo.ValueType
+        def __init__(self, key: _Optional[str] = ..., value: _Optional[_Union[DeviceInfo.ValueType, str]] = ...) -> None: ...
+    class ObservableFieldsEntry(_message.Message):
+        __slots__ = ["key", "value"]
+        KEY_FIELD_NUMBER: _ClassVar[int]
+        VALUE_FIELD_NUMBER: _ClassVar[int]
+        key: str
+        value: DeviceInfo.ValueType
+        def __init__(self, key: _Optional[str] = ..., value: _Optional[_Union[DeviceInfo.ValueType, str]] = ...) -> None: ...
+    AC: DeviceInfo.DeviceType
+    ACTION_FIELDS_FIELD_NUMBER: _ClassVar[int]
+    AHU: DeviceInfo.DeviceType
+    BLR: DeviceInfo.DeviceType
+    CDWS: DeviceInfo.DeviceType
+    CH: DeviceInfo.DeviceType
+    CHWS: DeviceInfo.DeviceType
+    CODE_FIELD_NUMBER: _ClassVar[int]
+    CT: DeviceInfo.DeviceType
+    DC: DeviceInfo.DeviceType
+    DEVICE_ID_FIELD_NUMBER: _ClassVar[int]
+    DEVICE_TYPE_FIELD_NUMBER: _ClassVar[int]
+    DFR: DeviceInfo.DeviceType
+    DH: DeviceInfo.DeviceType
+    DMP: DeviceInfo.DeviceType
+    FAN: DeviceInfo.DeviceType
+    FCU: DeviceInfo.DeviceType
+    GAS: DeviceInfo.DeviceType
+    HWS: DeviceInfo.DeviceType
+    HX: DeviceInfo.DeviceType
+    MAU: DeviceInfo.DeviceType
+    NAMESPACE_FIELD_NUMBER: _ClassVar[int]
+    OBSERVABLE_FIELDS_FIELD_NUMBER: _ClassVar[int]
+    OTHER: DeviceInfo.DeviceType
+    PMP: DeviceInfo.DeviceType
+    PWR: DeviceInfo.DeviceType
+    SDC: DeviceInfo.DeviceType
+    UH: DeviceInfo.DeviceType
+    UNDEFINED: DeviceInfo.DeviceType
+    VALUE_BINARY: DeviceInfo.ValueType
+    VALUE_CATEGORICAL: DeviceInfo.ValueType
+    VALUE_CONTINUOUS: DeviceInfo.ValueType
+    VALUE_INTEGER: DeviceInfo.ValueType
+    VALUE_TYPE_UNDEFINED: DeviceInfo.ValueType
+    VAV: DeviceInfo.DeviceType
+    ZONE_ID_FIELD_NUMBER: _ClassVar[int]
+    action_fields: _containers.ScalarMap[str, DeviceInfo.ValueType]
+    code: str
+    device_id: str
+    device_type: DeviceInfo.DeviceType
+    namespace: str
+    observable_fields: _containers.ScalarMap[str, DeviceInfo.ValueType]
+    zone_id: str
+    def __init__(self, device_id: _Optional[str] = ..., namespace: _Optional[str] = ..., code: _Optional[str] = ..., zone_id: _Optional[str] = ..., device_type: _Optional[_Union[DeviceInfo.DeviceType, str]] = ..., observable_fields: _Optional[_Mapping[str, DeviceInfo.ValueType]] = ..., action_fields: _Optional[_Mapping[str, DeviceInfo.ValueType]] = ...) -> None: ...
+
+class ObservationRequest(_message.Message):
+    __slots__ = ["single_observation_requests", "timestamp"]
+    SINGLE_OBSERVATION_REQUESTS_FIELD_NUMBER: _ClassVar[int]
+    TIMESTAMP_FIELD_NUMBER: _ClassVar[int]
+    single_observation_requests: _containers.RepeatedCompositeFieldContainer[SingleObservationRequest]
+    timestamp: _timestamp_pb2.Timestamp
+    def __init__(self, timestamp: _Optional[_Union[_timestamp_pb2.Timestamp, _Mapping]] = ..., single_observation_requests: _Optional[_Iterable[_Union[SingleObservationRequest, _Mapping]]] = ...) -> None: ...
+
+class ObservationResponse(_message.Message):
+    __slots__ = ["request", "single_observation_responses", "timestamp"]
+    REQUEST_FIELD_NUMBER: _ClassVar[int]
+    SINGLE_OBSERVATION_RESPONSES_FIELD_NUMBER: _ClassVar[int]
+    TIMESTAMP_FIELD_NUMBER: _ClassVar[int]
+    request: ObservationRequest
+    single_observation_responses: _containers.RepeatedCompositeFieldContainer[SingleObservationResponse]
+    timestamp: _timestamp_pb2.Timestamp
+    def __init__(self, timestamp: _Optional[_Union[_timestamp_pb2.Timestamp, _Mapping]] = ..., request: _Optional[_Union[ObservationRequest, _Mapping]] = ..., single_observation_responses: _Optional[_Iterable[_Union[SingleObservationResponse, _Mapping]]] = ...) -> None: ...
+
+class SingleActionRequest(_message.Message):
+    __slots__ = ["binary_value", "categorical_value", "continuous_value", "device_id", "integer_value", "setpoint_name", "string_value"]
+    BINARY_VALUE_FIELD_NUMBER: _ClassVar[int]
+    CATEGORICAL_VALUE_FIELD_NUMBER: _ClassVar[int]
+    CONTINUOUS_VALUE_FIELD_NUMBER: _ClassVar[int]
+    DEVICE_ID_FIELD_NUMBER: _ClassVar[int]
+    INTEGER_VALUE_FIELD_NUMBER: _ClassVar[int]
+    SETPOINT_NAME_FIELD_NUMBER: _ClassVar[int]
+    STRING_VALUE_FIELD_NUMBER: _ClassVar[int]
+    binary_value: bool
+    categorical_value: str
+    continuous_value: float
+    device_id: str
+    integer_value: int
+    setpoint_name: str
+    string_value: str
+    def __init__(self, device_id: _Optional[str] = ..., setpoint_name: _Optional[str] = ..., continuous_value: _Optional[float] = ..., integer_value: _Optional[int] = ..., categorical_value: _Optional[str] = ..., binary_value: bool = ..., string_value: _Optional[str] = ...) -> None: ...
+
+class SingleActionResponse(_message.Message):
+    __slots__ = ["additional_info", "request", "response_type"]
+    class ActionResponseType(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
+        __slots__ = []
+    ACCEPTED: SingleActionResponse.ActionResponseType
+    ADDITIONAL_INFO_FIELD_NUMBER: _ClassVar[int]
+    OTHER: SingleActionResponse.ActionResponseType
+    PENDING: SingleActionResponse.ActionResponseType
+    REJECTED_DEVICE_OFFLINE: SingleActionResponse.ActionResponseType
+    REJECTED_INVALID_DEVICE: SingleActionResponse.ActionResponseType
+    REJECTED_INVALID_SETTING: SingleActionResponse.ActionResponseType
+    REJECTED_NOT_ENABLED_OR_AVAILABLE: SingleActionResponse.ActionResponseType
+    REJECTED_OVERRIDE: SingleActionResponse.ActionResponseType
+    REQUEST_FIELD_NUMBER: _ClassVar[int]
+    RESPONSE_TYPE_FIELD_NUMBER: _ClassVar[int]
+    TIMED_OUT: SingleActionResponse.ActionResponseType
+    UNDEFINED: SingleActionResponse.ActionResponseType
+    UNKNOWN: SingleActionResponse.ActionResponseType
+    additional_info: str
+    request: SingleActionRequest
+    response_type: SingleActionResponse.ActionResponseType
+    def __init__(self, request: _Optional[_Union[SingleActionRequest, _Mapping]] = ..., response_type: _Optional[_Union[SingleActionResponse.ActionResponseType, str]] = ..., additional_info: _Optional[str] = ...) -> None: ...
+
+class SingleObservationRequest(_message.Message):
+    __slots__ = ["device_id", "measurement_name"]
+    DEVICE_ID_FIELD_NUMBER: _ClassVar[int]
+    MEASUREMENT_NAME_FIELD_NUMBER: _ClassVar[int]
+    device_id: str
+    measurement_name: str
+    def __init__(self, device_id: _Optional[str] = ..., measurement_name: _Optional[str] = ...) -> None: ...
+
+class SingleObservationResponse(_message.Message):
+    __slots__ = ["binary_value", "categorical_value", "continuous_value", "integer_value", "observation_valid", "single_observation_request", "string_value", "timestamp"]
+    BINARY_VALUE_FIELD_NUMBER: _ClassVar[int]
+    CATEGORICAL_VALUE_FIELD_NUMBER: _ClassVar[int]
+    CONTINUOUS_VALUE_FIELD_NUMBER: _ClassVar[int]
+    INTEGER_VALUE_FIELD_NUMBER: _ClassVar[int]
+    OBSERVATION_VALID_FIELD_NUMBER: _ClassVar[int]
+    SINGLE_OBSERVATION_REQUEST_FIELD_NUMBER: _ClassVar[int]
+    STRING_VALUE_FIELD_NUMBER: _ClassVar[int]
+    TIMESTAMP_FIELD_NUMBER: _ClassVar[int]
+    binary_value: bool
+    categorical_value: str
+    continuous_value: float
+    integer_value: int
+    observation_valid: bool
+    single_observation_request: SingleObservationRequest
+    string_value: str
+    timestamp: _timestamp_pb2.Timestamp
+    def __init__(self, timestamp: _Optional[_Union[_timestamp_pb2.Timestamp, _Mapping]] = ..., single_observation_request: _Optional[_Union[SingleObservationRequest, _Mapping]] = ..., observation_valid: bool = ..., continuous_value: _Optional[float] = ..., integer_value: _Optional[int] = ..., categorical_value: _Optional[str] = ..., binary_value: bool = ..., string_value: _Optional[str] = ...) -> None: ...
+
+class ZoneInfo(_message.Message):
+    __slots__ = ["area", "building_id", "devices", "floor", "zone_description", "zone_id", "zone_type"]
+    class ZoneType(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
+        __slots__ = []
+    AREA_FIELD_NUMBER: _ClassVar[int]
+    BUILDING_ID_FIELD_NUMBER: _ClassVar[int]
+    DEVICES_FIELD_NUMBER: _ClassVar[int]
+    FLOOR: ZoneInfo.ZoneType
+    FLOOR_FIELD_NUMBER: _ClassVar[int]
+    OTHER: ZoneInfo.ZoneType
+    ROOM: ZoneInfo.ZoneType
+    UNDEFINED: ZoneInfo.ZoneType
+    ZONE_DESCRIPTION_FIELD_NUMBER: _ClassVar[int]
+    ZONE_ID_FIELD_NUMBER: _ClassVar[int]
+    ZONE_TYPE_FIELD_NUMBER: _ClassVar[int]
+    area: float
+    building_id: str
+    devices: _containers.RepeatedScalarFieldContainer[str]
+    floor: int
+    zone_description: str
+    zone_id: str
+    zone_type: ZoneInfo.ZoneType
+    def __init__(self, zone_id: _Optional[str] = ..., building_id: _Optional[str] = ..., zone_description: _Optional[str] = ..., area: _Optional[float] = ..., devices: _Optional[_Iterable[str]] = ..., zone_type: _Optional[_Union[ZoneInfo.ZoneType, str]] = ..., floor: _Optional[int] = ...) -> None: ...
diff --git a/tensorflow_datasets/datasets/smart_buildings/smart_control_normalization.proto b/tensorflow_datasets/proto/smart_control_normalization.proto
similarity index 100%
rename from tensorflow_datasets/datasets/smart_buildings/smart_control_normalization.proto
rename to tensorflow_datasets/proto/smart_control_normalization.proto
diff --git a/tensorflow_datasets/proto/smart_control_normalization_pb2.py b/tensorflow_datasets/proto/smart_control_normalization_pb2.py
new file mode 100644
index 00000000000..91c954fa00f
--- /dev/null
+++ b/tensorflow_datasets/proto/smart_control_normalization_pb2.py
@@ -0,0 +1,41 @@
+# 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.
+
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: smart_control_normalization.proto
+"""Generated protocol buffer code."""
+from google.protobuf.internal import builder as _builder
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import symbol_database as _symbol_database
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+from google.protobuf import timestamp_pb2 as google_dot_protobuf_dot_timestamp__pb2
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n!smart_control_normalization.proto\x12#smart_buildings.smart_control.proto\x1a\x1fgoogle/protobuf/timestamp.proto\"\x90\x02\n\x16\x43ontinuousVariableInfo\x12\n\n\x02id\x18\x01 \x01(\t\x12\x30\n\x0csample_start\x18\x02 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12.\n\nsample_end\x18\x03 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x13\n\x0bsample_size\x18\x04 \x01(\x05\x12\x17\n\x0fsample_variance\x18\x05 \x01(\x02\x12\x13\n\x0bsample_mean\x18\x06 \x01(\x02\x12\x15\n\rsample_median\x18\x07 \x01(\x02\x12\x16\n\x0esample_maximum\x18\x08 \x01(\x02\x12\x16\n\x0esample_minimum\x18\t \x01(\x02\x62\x06proto3')
+
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals())
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'smart_control_normalization_pb2', globals())
+if _descriptor._USE_C_DESCRIPTORS == False:
+
+  DESCRIPTOR._options = None
+  _CONTINUOUSVARIABLEINFO._serialized_start=108
+  _CONTINUOUSVARIABLEINFO._serialized_end=380
+# @@protoc_insertion_point(module_scope)
diff --git a/tensorflow_datasets/proto/smart_control_normalization_pb2.pyi b/tensorflow_datasets/proto/smart_control_normalization_pb2.pyi
new file mode 100644
index 00000000000..375b5e4078a
--- /dev/null
+++ b/tensorflow_datasets/proto/smart_control_normalization_pb2.pyi
@@ -0,0 +1,28 @@
+from google.protobuf import timestamp_pb2 as _timestamp_pb2
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from typing import ClassVar as _ClassVar, Mapping as _Mapping, Optional as _Optional, Union as _Union
+
+DESCRIPTOR: _descriptor.FileDescriptor
+
+class ContinuousVariableInfo(_message.Message):
+    __slots__ = ["id", "sample_end", "sample_maximum", "sample_mean", "sample_median", "sample_minimum", "sample_size", "sample_start", "sample_variance"]
+    ID_FIELD_NUMBER: _ClassVar[int]
+    SAMPLE_END_FIELD_NUMBER: _ClassVar[int]
+    SAMPLE_MAXIMUM_FIELD_NUMBER: _ClassVar[int]
+    SAMPLE_MEAN_FIELD_NUMBER: _ClassVar[int]
+    SAMPLE_MEDIAN_FIELD_NUMBER: _ClassVar[int]
+    SAMPLE_MINIMUM_FIELD_NUMBER: _ClassVar[int]
+    SAMPLE_SIZE_FIELD_NUMBER: _ClassVar[int]
+    SAMPLE_START_FIELD_NUMBER: _ClassVar[int]
+    SAMPLE_VARIANCE_FIELD_NUMBER: _ClassVar[int]
+    id: str
+    sample_end: _timestamp_pb2.Timestamp
+    sample_maximum: float
+    sample_mean: float
+    sample_median: float
+    sample_minimum: float
+    sample_size: int
+    sample_start: _timestamp_pb2.Timestamp
+    sample_variance: float
+    def __init__(self, id: _Optional[str] = ..., sample_start: _Optional[_Union[_timestamp_pb2.Timestamp, _Mapping]] = ..., sample_end: _Optional[_Union[_timestamp_pb2.Timestamp, _Mapping]] = ..., sample_size: _Optional[int] = ..., sample_variance: _Optional[float] = ..., sample_mean: _Optional[float] = ..., sample_median: _Optional[float] = ..., sample_maximum: _Optional[float] = ..., sample_minimum: _Optional[float] = ...) -> None: ...
diff --git a/tensorflow_datasets/datasets/smart_buildings/smart_control_reward.proto b/tensorflow_datasets/proto/smart_control_reward.proto
similarity index 100%
rename from tensorflow_datasets/datasets/smart_buildings/smart_control_reward.proto
rename to tensorflow_datasets/proto/smart_control_reward.proto
diff --git a/tensorflow_datasets/proto/smart_control_reward_pb2.py b/tensorflow_datasets/proto/smart_control_reward_pb2.py
new file mode 100644
index 00000000000..ea0527a672c
--- /dev/null
+++ b/tensorflow_datasets/proto/smart_control_reward_pb2.py
@@ -0,0 +1,61 @@
+# 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.
+
+# -*- coding: utf-8 -*-
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# source: smart_control_reward.proto
+"""Generated protocol buffer code."""
+from google.protobuf.internal import builder as _builder
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import descriptor_pool as _descriptor_pool
+from google.protobuf import symbol_database as _symbol_database
+# @@protoc_insertion_point(imports)
+
+_sym_db = _symbol_database.Default()
+
+
+from google.protobuf import timestamp_pb2 as google_dot_protobuf_dot_timestamp__pb2
+
+
+DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x1asmart_control_reward.proto\x12#smart_buildings.smart_control.proto\x1a\x1fgoogle/protobuf/timestamp.proto\"\xe9\t\n\nRewardInfo\x12\x33\n\x0fstart_timestamp\x18\x01 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x31\n\rend_timestamp\x18\x02 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x10\n\x08\x61gent_id\x18\x03 \x01(\t\x12\x13\n\x0bscenario_id\x18\x04 \x01(\t\x12_\n\x11zone_reward_infos\x18\x05 \x03(\x0b\x32\x44.smart_buildings.smart_control.proto.RewardInfo.ZoneRewardInfosEntry\x12l\n\x18\x61ir_handler_reward_infos\x18\x06 \x03(\x0b\x32J.smart_buildings.smart_control.proto.RewardInfo.AirHandlerRewardInfosEntry\x12\x63\n\x13\x62oiler_reward_infos\x18\x07 \x03(\x0b\x32\x46.smart_buildings.smart_control.proto.RewardInfo.BoilerRewardInfosEntry\x1a\xcc\x01\n\x0eZoneRewardInfo\x12$\n\x1cheating_setpoint_temperature\x18\x01 \x01(\x02\x12$\n\x1c\x63ooling_setpoint_temperature\x18\x02 \x01(\x02\x12\x1c\n\x14zone_air_temperature\x18\x03 \x01(\x02\x12\x1e\n\x16\x61ir_flow_rate_setpoint\x18\x04 \x01(\x02\x12\x15\n\rair_flow_rate\x18\x05 \x01(\x02\x12\x19\n\x11\x61verage_occupancy\x18\x06 \x01(\x02\x1an\n\x14\x41irHandlerRewardInfo\x12%\n\x1d\x62lower_electrical_energy_rate\x18\x01 \x01(\x02\x12/\n\'air_conditioning_electrical_energy_rate\x18\x02 \x01(\x02\x1a`\n\x10\x42oilerRewardInfo\x12\'\n\x1fnatural_gas_heating_energy_rate\x18\x01 \x01(\x02\x12#\n\x1bpump_electrical_energy_rate\x18\x02 \x01(\x02\x1av\n\x14ZoneRewardInfosEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12M\n\x05value\x18\x02 \x01(\x0b\x32>.smart_buildings.smart_control.proto.RewardInfo.ZoneRewardInfo:\x02\x38\x01\x1a\x82\x01\n\x1a\x41irHandlerRewardInfosEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12S\n\x05value\x18\x02 \x01(\x0b\x32\x44.smart_buildings.smart_control.proto.RewardInfo.AirHandlerRewardInfo:\x02\x38\x01\x1az\n\x16\x42oilerRewardInfosEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12O\n\x05value\x18\x02 \x01(\x0b\x32@.smart_buildings.smart_control.proto.RewardInfo.BoilerRewardInfo:\x02\x38\x01\"\xe4\x04\n\x0eRewardResponse\x12\x1a\n\x12\x61gent_reward_value\x18\x01 \x01(\x02\x12\x1b\n\x13productivity_reward\x18\x02 \x01(\x02\x12\x1f\n\x17\x65lectricity_energy_cost\x18\x03 \x01(\x02\x12\x1f\n\x17natural_gas_energy_cost\x18\x04 \x01(\x02\x12\x16\n\x0e\x63\x61rbon_emitted\x18\x05 \x01(\x02\x12\x13\n\x0b\x63\x61rbon_cost\x18\x06 \x01(\x02\x12\x1b\n\x13productivity_weight\x18\x07 \x01(\x02\x12\x1a\n\x12\x65nergy_cost_weight\x18\x08 \x01(\x02\x12\x1e\n\x16\x63\x61rbon_emission_weight\x18\t \x01(\x02\x12\x1b\n\x13person_productivity\x18\n \x01(\x02\x12\x17\n\x0ftotal_occupancy\x18\x0b \x01(\x02\x12\x14\n\x0creward_scale\x18\x0c \x01(\x02\x12\x14\n\x0creward_shift\x18\r \x01(\x02\x12\x1b\n\x13productivity_regret\x18\x0e \x01(\x02\x12&\n\x1enormalized_productivity_regret\x18\x0f \x01(\x02\x12\x1e\n\x16normalized_energy_cost\x18\x10 \x01(\x02\x12\"\n\x1anormalized_carbon_emission\x18\x11 \x01(\x02\x12\x33\n\x0fstart_timestamp\x18\x12 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x31\n\rend_timestamp\x18\x13 \x01(\x0b\x32\x1a.google.protobuf.Timestampb\x06proto3')
+
+_builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals())
+_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'smart_control_reward_pb2', globals())
+if _descriptor._USE_C_DESCRIPTORS == False:
+
+  DESCRIPTOR._options = None
+  _REWARDINFO_ZONEREWARDINFOSENTRY._options = None
+  _REWARDINFO_ZONEREWARDINFOSENTRY._serialized_options = b'8\001'
+  _REWARDINFO_AIRHANDLERREWARDINFOSENTRY._options = None
+  _REWARDINFO_AIRHANDLERREWARDINFOSENTRY._serialized_options = b'8\001'
+  _REWARDINFO_BOILERREWARDINFOSENTRY._options = None
+  _REWARDINFO_BOILERREWARDINFOSENTRY._serialized_options = b'8\001'
+  _REWARDINFO._serialized_start=101
+  _REWARDINFO._serialized_end=1358
+  _REWARDINFO_ZONEREWARDINFO._serialized_start=567
+  _REWARDINFO_ZONEREWARDINFO._serialized_end=771
+  _REWARDINFO_AIRHANDLERREWARDINFO._serialized_start=773
+  _REWARDINFO_AIRHANDLERREWARDINFO._serialized_end=883
+  _REWARDINFO_BOILERREWARDINFO._serialized_start=885
+  _REWARDINFO_BOILERREWARDINFO._serialized_end=981
+  _REWARDINFO_ZONEREWARDINFOSENTRY._serialized_start=983
+  _REWARDINFO_ZONEREWARDINFOSENTRY._serialized_end=1101
+  _REWARDINFO_AIRHANDLERREWARDINFOSENTRY._serialized_start=1104
+  _REWARDINFO_AIRHANDLERREWARDINFOSENTRY._serialized_end=1234
+  _REWARDINFO_BOILERREWARDINFOSENTRY._serialized_start=1236
+  _REWARDINFO_BOILERREWARDINFOSENTRY._serialized_end=1358
+  _REWARDRESPONSE._serialized_start=1361
+  _REWARDRESPONSE._serialized_end=1973
+# @@protoc_insertion_point(module_scope)
diff --git a/tensorflow_datasets/proto/smart_control_reward_pb2.pyi b/tensorflow_datasets/proto/smart_control_reward_pb2.pyi
new file mode 100644
index 00000000000..e53b7da04aa
--- /dev/null
+++ b/tensorflow_datasets/proto/smart_control_reward_pb2.pyi
@@ -0,0 +1,117 @@
+from google.protobuf import timestamp_pb2 as _timestamp_pb2
+from google.protobuf.internal import containers as _containers
+from google.protobuf import descriptor as _descriptor
+from google.protobuf import message as _message
+from typing import ClassVar as _ClassVar, Mapping as _Mapping, Optional as _Optional, Union as _Union
+
+DESCRIPTOR: _descriptor.FileDescriptor
+
+class RewardInfo(_message.Message):
+    __slots__ = ["agent_id", "air_handler_reward_infos", "boiler_reward_infos", "end_timestamp", "scenario_id", "start_timestamp", "zone_reward_infos"]
+    class AirHandlerRewardInfo(_message.Message):
+        __slots__ = ["air_conditioning_electrical_energy_rate", "blower_electrical_energy_rate"]
+        AIR_CONDITIONING_ELECTRICAL_ENERGY_RATE_FIELD_NUMBER: _ClassVar[int]
+        BLOWER_ELECTRICAL_ENERGY_RATE_FIELD_NUMBER: _ClassVar[int]
+        air_conditioning_electrical_energy_rate: float
+        blower_electrical_energy_rate: float
+        def __init__(self, blower_electrical_energy_rate: _Optional[float] = ..., air_conditioning_electrical_energy_rate: _Optional[float] = ...) -> None: ...
+    class AirHandlerRewardInfosEntry(_message.Message):
+        __slots__ = ["key", "value"]
+        KEY_FIELD_NUMBER: _ClassVar[int]
+        VALUE_FIELD_NUMBER: _ClassVar[int]
+        key: str
+        value: RewardInfo.AirHandlerRewardInfo
+        def __init__(self, key: _Optional[str] = ..., value: _Optional[_Union[RewardInfo.AirHandlerRewardInfo, _Mapping]] = ...) -> None: ...
+    class BoilerRewardInfo(_message.Message):
+        __slots__ = ["natural_gas_heating_energy_rate", "pump_electrical_energy_rate"]
+        NATURAL_GAS_HEATING_ENERGY_RATE_FIELD_NUMBER: _ClassVar[int]
+        PUMP_ELECTRICAL_ENERGY_RATE_FIELD_NUMBER: _ClassVar[int]
+        natural_gas_heating_energy_rate: float
+        pump_electrical_energy_rate: float
+        def __init__(self, natural_gas_heating_energy_rate: _Optional[float] = ..., pump_electrical_energy_rate: _Optional[float] = ...) -> None: ...
+    class BoilerRewardInfosEntry(_message.Message):
+        __slots__ = ["key", "value"]
+        KEY_FIELD_NUMBER: _ClassVar[int]
+        VALUE_FIELD_NUMBER: _ClassVar[int]
+        key: str
+        value: RewardInfo.BoilerRewardInfo
+        def __init__(self, key: _Optional[str] = ..., value: _Optional[_Union[RewardInfo.BoilerRewardInfo, _Mapping]] = ...) -> None: ...
+    class ZoneRewardInfo(_message.Message):
+        __slots__ = ["air_flow_rate", "air_flow_rate_setpoint", "average_occupancy", "cooling_setpoint_temperature", "heating_setpoint_temperature", "zone_air_temperature"]
+        AIR_FLOW_RATE_FIELD_NUMBER: _ClassVar[int]
+        AIR_FLOW_RATE_SETPOINT_FIELD_NUMBER: _ClassVar[int]
+        AVERAGE_OCCUPANCY_FIELD_NUMBER: _ClassVar[int]
+        COOLING_SETPOINT_TEMPERATURE_FIELD_NUMBER: _ClassVar[int]
+        HEATING_SETPOINT_TEMPERATURE_FIELD_NUMBER: _ClassVar[int]
+        ZONE_AIR_TEMPERATURE_FIELD_NUMBER: _ClassVar[int]
+        air_flow_rate: float
+        air_flow_rate_setpoint: float
+        average_occupancy: float
+        cooling_setpoint_temperature: float
+        heating_setpoint_temperature: float
+        zone_air_temperature: float
+        def __init__(self, heating_setpoint_temperature: _Optional[float] = ..., cooling_setpoint_temperature: _Optional[float] = ..., zone_air_temperature: _Optional[float] = ..., air_flow_rate_setpoint: _Optional[float] = ..., air_flow_rate: _Optional[float] = ..., average_occupancy: _Optional[float] = ...) -> None: ...
+    class ZoneRewardInfosEntry(_message.Message):
+        __slots__ = ["key", "value"]
+        KEY_FIELD_NUMBER: _ClassVar[int]
+        VALUE_FIELD_NUMBER: _ClassVar[int]
+        key: str
+        value: RewardInfo.ZoneRewardInfo
+        def __init__(self, key: _Optional[str] = ..., value: _Optional[_Union[RewardInfo.ZoneRewardInfo, _Mapping]] = ...) -> None: ...
+    AGENT_ID_FIELD_NUMBER: _ClassVar[int]
+    AIR_HANDLER_REWARD_INFOS_FIELD_NUMBER: _ClassVar[int]
+    BOILER_REWARD_INFOS_FIELD_NUMBER: _ClassVar[int]
+    END_TIMESTAMP_FIELD_NUMBER: _ClassVar[int]
+    SCENARIO_ID_FIELD_NUMBER: _ClassVar[int]
+    START_TIMESTAMP_FIELD_NUMBER: _ClassVar[int]
+    ZONE_REWARD_INFOS_FIELD_NUMBER: _ClassVar[int]
+    agent_id: str
+    air_handler_reward_infos: _containers.MessageMap[str, RewardInfo.AirHandlerRewardInfo]
+    boiler_reward_infos: _containers.MessageMap[str, RewardInfo.BoilerRewardInfo]
+    end_timestamp: _timestamp_pb2.Timestamp
+    scenario_id: str
+    start_timestamp: _timestamp_pb2.Timestamp
+    zone_reward_infos: _containers.MessageMap[str, RewardInfo.ZoneRewardInfo]
+    def __init__(self, start_timestamp: _Optional[_Union[_timestamp_pb2.Timestamp, _Mapping]] = ..., end_timestamp: _Optional[_Union[_timestamp_pb2.Timestamp, _Mapping]] = ..., agent_id: _Optional[str] = ..., scenario_id: _Optional[str] = ..., zone_reward_infos: _Optional[_Mapping[str, RewardInfo.ZoneRewardInfo]] = ..., air_handler_reward_infos: _Optional[_Mapping[str, RewardInfo.AirHandlerRewardInfo]] = ..., boiler_reward_infos: _Optional[_Mapping[str, RewardInfo.BoilerRewardInfo]] = ...) -> None: ...
+
+class RewardResponse(_message.Message):
+    __slots__ = ["agent_reward_value", "carbon_cost", "carbon_emission_weight", "carbon_emitted", "electricity_energy_cost", "end_timestamp", "energy_cost_weight", "natural_gas_energy_cost", "normalized_carbon_emission", "normalized_energy_cost", "normalized_productivity_regret", "person_productivity", "productivity_regret", "productivity_reward", "productivity_weight", "reward_scale", "reward_shift", "start_timestamp", "total_occupancy"]
+    AGENT_REWARD_VALUE_FIELD_NUMBER: _ClassVar[int]
+    CARBON_COST_FIELD_NUMBER: _ClassVar[int]
+    CARBON_EMISSION_WEIGHT_FIELD_NUMBER: _ClassVar[int]
+    CARBON_EMITTED_FIELD_NUMBER: _ClassVar[int]
+    ELECTRICITY_ENERGY_COST_FIELD_NUMBER: _ClassVar[int]
+    END_TIMESTAMP_FIELD_NUMBER: _ClassVar[int]
+    ENERGY_COST_WEIGHT_FIELD_NUMBER: _ClassVar[int]
+    NATURAL_GAS_ENERGY_COST_FIELD_NUMBER: _ClassVar[int]
+    NORMALIZED_CARBON_EMISSION_FIELD_NUMBER: _ClassVar[int]
+    NORMALIZED_ENERGY_COST_FIELD_NUMBER: _ClassVar[int]
+    NORMALIZED_PRODUCTIVITY_REGRET_FIELD_NUMBER: _ClassVar[int]
+    PERSON_PRODUCTIVITY_FIELD_NUMBER: _ClassVar[int]
+    PRODUCTIVITY_REGRET_FIELD_NUMBER: _ClassVar[int]
+    PRODUCTIVITY_REWARD_FIELD_NUMBER: _ClassVar[int]
+    PRODUCTIVITY_WEIGHT_FIELD_NUMBER: _ClassVar[int]
+    REWARD_SCALE_FIELD_NUMBER: _ClassVar[int]
+    REWARD_SHIFT_FIELD_NUMBER: _ClassVar[int]
+    START_TIMESTAMP_FIELD_NUMBER: _ClassVar[int]
+    TOTAL_OCCUPANCY_FIELD_NUMBER: _ClassVar[int]
+    agent_reward_value: float
+    carbon_cost: float
+    carbon_emission_weight: float
+    carbon_emitted: float
+    electricity_energy_cost: float
+    end_timestamp: _timestamp_pb2.Timestamp
+    energy_cost_weight: float
+    natural_gas_energy_cost: float
+    normalized_carbon_emission: float
+    normalized_energy_cost: float
+    normalized_productivity_regret: float
+    person_productivity: float
+    productivity_regret: float
+    productivity_reward: float
+    productivity_weight: float
+    reward_scale: float
+    reward_shift: float
+    start_timestamp: _timestamp_pb2.Timestamp
+    total_occupancy: float
+    def __init__(self, agent_reward_value: _Optional[float] = ..., productivity_reward: _Optional[float] = ..., electricity_energy_cost: _Optional[float] = ..., natural_gas_energy_cost: _Optional[float] = ..., carbon_emitted: _Optional[float] = ..., carbon_cost: _Optional[float] = ..., productivity_weight: _Optional[float] = ..., energy_cost_weight: _Optional[float] = ..., carbon_emission_weight: _Optional[float] = ..., person_productivity: _Optional[float] = ..., total_occupancy: _Optional[float] = ..., reward_scale: _Optional[float] = ..., reward_shift: _Optional[float] = ..., productivity_regret: _Optional[float] = ..., normalized_productivity_regret: _Optional[float] = ..., normalized_energy_cost: _Optional[float] = ..., normalized_carbon_emission: _Optional[float] = ..., start_timestamp: _Optional[_Union[_timestamp_pb2.Timestamp, _Mapping]] = ..., end_timestamp: _Optional[_Union[_timestamp_pb2.Timestamp, _Mapping]] = ...) -> None: ...