Add protocols for authoring simulation environments. This is for preview only, we discourage users from using this in production code.

PiperOrigin-RevId: 899009633
Change-Id: I2de08c5cc4ece69fa135f3fd02df7542bf66d7d0

src/experimental/reaf/core/action_space_adapter.py

@@ -0,0 +1,43 @@
+# Copyright 2025 Google LLC
+#
+# 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
+#
+# https://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.
+"""Adapts environment action into suitable commands format accepted by REAF."""
+
+import abc
+from collections.abc import Mapping
+
+from gdm_robotics.interfaces import types as gdmr_types
+
+
+class ActionSpaceAdapter(abc.ABC):
+  """Adapts environment action into suitable commands format accepted by REAF.
+
+  Implementations of this interface are responsible for converting the more
+  generic action accepted by the environment (e.g. a flat numpy array) into the
+  more constraining format accepted as commands by REAF, i.e. a dictionary of
+  string to tensors.
+  """
+
+  @abc.abstractmethod
+  def commands_from_environment_action(
+      self, environment_action: gdmr_types.ActionType
+  ) -> Mapping[str, gdmr_types.ArrayType]:
+    """Converts the environment action into commands accepted by REAF."""
+
+  @abc.abstractmethod
+  def action_spec(self) -> gdmr_types.ActionSpec:
+    """Returns the action spec exposed by the environment."""
+
+  @abc.abstractmethod
+  def task_commands_keys(self) -> set[str]:
+    """Returns the keys for the commands exposed to the task layer."""

src/experimental/reaf/core/commands_processor.py

@@ -0,0 +1,91 @@
+# Copyright 2025 Google LLC
+#
+# 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
+#
+# https://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.
+"""Abstract class for commands manipulation in the task logic layer."""
+
+import abc
+from collections.abc import Mapping
+
+from gdm_robotics.interfaces import types as gdmr_types
+
+
+class CommandsProcessor(abc.ABC):
+  """Perform commands manipulation.
+
+  The following describes the processing pipeline starting from the top (closer
+  to the policy) to the bottom (interfacing with the DACL commands spec).
+
+  Assume that we have two processing units:
+  Processor 1) has a consumed_commands_spec for two keys: "p1/c1" and "p1/c2".
+    Its produced_commands_keys are "p2/c1".
+  Processor 2) has a consumed_commands_spec for "p2/c1". Its
+    produced_commands_keys are "p3/c1" and "p3/c2".
+
+  Specs are propagated starting from the bottom:
+  1) In this example assume that the DACL exposes "p3/c1", "p3/c2" and "p3/c3".
+  2) Processor 2) returns ("p3/c1", "p3/c2") from input "p2/c1". This means that
+    the global commands spec exposed at this level is "p2/c1" and the
+    unprocessed "p3/c3".
+  3) Processor 1) returns "p2/c1" from input ("p1/c1", "p1/c2"). By applying the
+    same transformation rule, we can obtain the final commands spec exposed by
+    the full processing pipeline: "p1/c1", "p1/c2" and "p3/c3".
+
+    "p1/c1"     "p1/c2"           "p3/c3"
+       |          |                  |
+     -----------------               |
+    |      P1         |              |
+     -----------------               |
+            |  "p2/c1"               |
+     -----------------               |
+    |      P2         |              |
+     -----------------               |
+       | "p3/c1"   | "p3/c2"         |
+       |           |                 |
+     ------------------------------------
+    |               DACL                 |
+     ------------------------------------
+  """
+
+  @property
+  @abc.abstractmethod
+  def name(self) -> str:
+    """Returns a unique string identifier for this object."""
+
+  @abc.abstractmethod
+  def process_commands(
+      self, consumed_commands: Mapping[str, gdmr_types.ArrayType]
+  ) -> Mapping[str, gdmr_types.ArrayType]:
+    """Processes the commands and returns a new modified version of it.
+
+    Args:
+      consumed_commands: the commands up in the processing chain (or provided by
+        the Environment) that are required by this processor, i.e. with keys
+        specified by `consumed_commands_spec`.
+
+    Returns the new commands. Note that the data in consumed_commands is removed
+      from the global commands dictionary. If users want to keep some of the
+      elements it is their responsibility to retain them in the output
+      dictionary.
+    """
+
+  @abc.abstractmethod
+  def consumed_commands_spec(self) -> Mapping[str, gdmr_types.AnyArraySpec]:
+    """Spec of the commands consumed by this processor."""
+
+  @abc.abstractmethod
+  def produced_commands_keys(self) -> set[str]:
+    """Keys of the commands produced by this processor."""
+
+  def reset(self) -> None:
+    """Resets the internal state of the command processor."""
+    ...

src/experimental/reaf/core/data_acquisition_and_control_layer.py

@@ -0,0 +1,170 @@
+# Copyright 2025 Google LLC
+#
+# 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
+#
+# https://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.
+"""REAF data acquisition and control layer to interface with the robotic setup."""
+
+from collections.abc import Iterable, Mapping
+
+from dm_env import specs
+from gdm_robotics.interfaces import types as gdmr_types
+from reaf.core import device as reaf_device
+from reaf.core import device_coordinator as reaf_coordinator
+from reaf.core import trigger
+
+
+class DataAcquisitionAndControlLayer:
+  """REAF data acquisition and control layer.
+
+  The DACL is responsible to provide an interface for the robotic setup.
+  """
+
+  def __init__(
+      self,
+      *,
+      device_coordinator: reaf_coordinator.DeviceCoordinator,
+      commands_trigger: trigger.Trigger | None,
+      measurements_trigger: trigger.Trigger | None,
+  ):
+    """Initializes the DataAcquisitionAndControlLayer.
+
+    Args:
+      device_coordinator: The coordinator representing a specific robotic setup.
+        Note that callers need to explicitly initialize and finalise the
+        coordinator.
+      commands_trigger: A trigger to unblock processing commands during a call
+        to `step`.
+      measurements_trigger: A trigger to unblock processing measurements during
+        a call to `step`.
+    """
+    self._coordinator = device_coordinator
+    self._devices = self._coordinator.get_devices()
+    # The following checks that names of the devices are unique and their keys
+    # are "mergeable".
+    self._check_device_names_and_keys(self._devices)
+
+    self._commands_trigger = commands_trigger
+    self._measurements_trigger = measurements_trigger
+
+    # Create a map of supported commands keys for each Device.
+    self._commands_for_device = {
+        device.name: device.commands_spec().keys() for device in self._devices
+    }
+
+  def begin_stepping(self) -> Mapping[str, gdmr_types.ArrayType]:
+    """Begins stepping the DACL and returns the current measurements."""
+    self._coordinator.on_begin_stepping()
+
+    # Wait for the first trigger to happen before collecting the measurements.
+    if self._measurements_trigger is not None:
+      self._measurements_trigger.wait_for_event()
+    return self._get_measurements()
+
+  def end_stepping(self) -> None:
+    """Ends stepping the data acquisition and control layer."""
+    self._coordinator.on_end_stepping()
+
+  def _set_commands(self, commands: Mapping[str, gdmr_types.ArrayType]) -> None:
+    """Sets the commands of the data acquisition and control layer."""
+    self._coordinator.before_set_commands()
+    for device in self._devices:
+      device_commands = {
+          k: v
+          for k, v in commands.items()
+          if k in self._commands_for_device[device.name]
+      }
+      device.set_commands(device_commands)
+    self._coordinator.after_set_commands()
+
+  def _get_measurements(self) -> Mapping[str, gdmr_types.ArrayType]:
+    """Gets the measurements of the data acquisition and control layer."""
+    measurements = {}
+    self._coordinator.before_get_measurements()
+    for device in self._devices:
+      measurements.update(device.get_measurements())
+
+    return measurements
+
+  def step(
+      self, commands: Mapping[str, gdmr_types.ArrayType]
+  ) -> Mapping[str, gdmr_types.ArrayType]:
+    """Steps the data acquisition and control layer."""
+    if self._commands_trigger is not None:
+      self._commands_trigger.wait_for_event()
+    self._set_commands(commands)
+
+    if self._measurements_trigger is not None:
+      self._measurements_trigger.wait_for_event()
+    return self._get_measurements()
+
+  def commands_spec(self) -> Mapping[str, gdmr_types.AnyArraySpec]:
+    """Returns the specs for the commands."""
+    spec = {}
+    for device in self._devices:
+      spec.update(device.commands_spec())
+    return spec
+
+  def measurements_spec(self) -> Mapping[str, specs.Array]:
+    """Returns the specs for the measurements."""
+    spec = {}
+    for device in self._devices:
+      spec.update(device.measurements_spec())
+    return spec
+
+  @property
+  def device_coordinator(self) -> reaf_coordinator.DeviceCoordinator:
+    return self._coordinator
+
+  def _check_keys_have_been_formatted_correctly(
+      self, current_key_set: Iterable[str]
+  ) -> None:
+    """Check that keys haven't been left unformatted."""
+    for key in current_key_set:
+      if key.find("{}") != -1:
+        raise ValueError(
+            "Keys should not contain '{}'. Did you mean to use format()?"
+        )
+
+  def _check_device_names_and_keys(
+      self, devices: Iterable[reaf_device.Device]
+  ) -> None:
+    """Raises error if device names are not unique or keys are not exclusive."""
+    # Check names first.
+    all_names = [device.name for device in devices]
+    unique_names = set(all_names)
+    if len(unique_names) != len(all_names):
+      raise RuntimeError(f"Duplicate names when checking devices: {all_names}")
+
+    # Check commands.
+    devices = tuple(devices)
+    current_specs = set()
+    for device in devices:
+      device_keys = device.commands_spec().keys()
+      self._check_keys_have_been_formatted_correctly(device_keys)
+      if not current_specs.isdisjoint(device_keys):
+        raise RuntimeError(
+            f"Duplicate keys when checking device {device.name}:"
+            f" {current_specs.intersection(device_keys)}"
+        )
+      current_specs.update(device_keys)
+
+    # Check measurements.
+    current_specs = set()
+    for device in devices:
+      device_keys = device.measurements_spec().keys()
+      self._check_keys_have_been_formatted_correctly(device_keys)
+      if not current_specs.isdisjoint(device_keys):
+        raise RuntimeError(
+            f"Duplicate keys when checking device {device.name}:"
+            f" {current_specs.intersection(device_keys)}"
+        )
+      current_specs.update(device_keys)

src/experimental/reaf/core/default_discount_provider.py

@@ -0,0 +1,79 @@
+# Copyright 2025 Google LLC
+#
+# 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
+#
+# https://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.
+"""Computes a constant discount given the termination state.
+
+This provider returns a discount of 0.0 in case of termination and 1.0
+otherwise (i.e. for truncation and not termination).
+
+It is usually safe to use this discount provider for environments that return
+strictly positive rewards.
+"""
+
+from collections.abc import Mapping
+
+from dm_env import specs
+from gdm_robotics.interfaces import types as gdmr_types
+import numpy as np
+from reaf.core import discount_provider
+from reaf.core import termination_checker
+import tree
+
+
+class DefaultDiscountProvider(discount_provider.DiscountProvider):
+  """Computes a constant discount given the termination state.
+
+  This provider returns a discount of 0.0 in case of termination and 1.0
+  otherwise (i.e. for truncation and not termination).
+
+  It is usually safe to use this discount provider for environments that return
+  strictly positive rewards.
+  """
+
+  def __init__(self, name: str = "default_discount_provider"):
+    self._name = name
+    self._spec = specs.BoundedArray(
+        shape=(), dtype=np.float64, minimum=0.0, maximum=1.0, name="discount"
+    )
+
+  def name(self) -> str:
+    """Returns a unique string identifier for this object."""
+    return self._name
+
+  def compute_discount(
+      self,
+      unused_required_features: Mapping[str, gdmr_types.ArrayType],
+      termination_state: termination_checker.TerminationResult,
+  ) -> tree.Structure[gdmr_types.ArrayType]:
+    """Computes the discount.
+
+    Args:
+      unused_required_features: Unused
+      termination_state: The termination state as computed by the termination
+        checkers.  Returns the discount.
+
+    Returns:
+      The discount.
+    """
+    if termination_state == termination_state.TERMINATE:
+      return np.asarray(0).astype(self._spec.dtype)
+    else:  # TRUNCATION or DO_NOT_TERMINATE
+      return np.asarray(1.0).astype(self._spec.dtype)
+
+  def discount_spec(self) -> tree.Structure[specs.Array]:
+    """Returns the spec of the discount."""
+    return self._spec
+
+  def required_features_keys(self) -> set[str]:
+    """Returns the feature keys that are required to compute the discount."""
+    return set()

src/experimental/reaf/core/default_observation_space_adapter.py

@@ -0,0 +1,231 @@
+# Copyright 2025 Google LLC
+#
+# 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
+#
+# https://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.
+"""ObservationSpaceAdapter supporting filtering, renaming and type conversion."""
+
+import abc
+from collections.abc import Iterable, Mapping
+import dataclasses
+
+from dm_env import specs
+from gdm_robotics.interfaces import types as gdmr_types
+import numpy as np
+import numpy.typing as npt
+from reaf.core import observation_space_adapter
+import tree
+
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class RenameInfo:
+  original_key: str
+  renamed_key: str
+
+
+class ObservationTypeMapper(abc.ABC):
+  """Maps from REAF features and specs into corresponding environment types."""
+
+  @abc.abstractmethod
+  def to_observation_spec(
+      self, features_spec: Mapping[str, specs.Array]
+  ) -> gdmr_types.ObservationSpec:
+    """Convert the features spec into the environment observation spec."""
+
+  @abc.abstractmethod
+  def to_observations(
+      self, features: Mapping[str, gdmr_types.ArrayType]
+  ) -> tree.Structure[gdmr_types.ArrayType]:
+    """Convert the features into the environment observations."""
+
+
+class _DefaultObservationTypeMapper(ObservationTypeMapper):
+  """An ObservationTypeMapper that returns the input features specs and dict.
+
+  This `ObservationTypeMapper` maps observations from the more constrained
+  `Mapping[str, ArrayType]` used in the task layer to the more generic
+  `tree.Structure[ArrayType]` exposed by the GDM Environment.
+  """
+
+  def to_observation_spec(
+      self, features_spec: Mapping[str, specs.Array]
+  ) -> gdmr_types.ObservationSpec:
+    """Returns the features spec, unmodified, as a `gdmr_types.ObservationSpec`."""
+    return features_spec
+
+  def to_observations(
+      self, features: Mapping[str, gdmr_types.ArrayType]
+  ) -> tree.Structure[gdmr_types.ArrayType]:
+    """Returns the features, unmodified, as a `tree.Structure`."""
+    return features
+
+
+class DefaultObservationSpaceAdapter(
+    observation_space_adapter.ObservationSpaceAdapter
+):
+  """Observation adapter supporting filtering, renaming and type conversion.
+
+    This adapter supports filtering, renaming, and converting REAF features into
+    environment observations.
+
+    The order of operations is the following:
+    1) Filtering, i.e. feature selection.
+    2) Downcasting floats to max_float_dtype.
+    3) Renaming.
+    4) Type conversion.
+
+  Please refer to the constructor documentation for more information.
+  """
+
+  def __init__(
+      self,
+      *,
+      task_features_spec: Mapping[str, specs.Array],
+      selected_features: Iterable[str] | None,
+      renamed_features: Iterable[RenameInfo] | None,
+      observation_type_mapper: ObservationTypeMapper | None,
+      max_float_dtype: npt.DTypeLike = np.float64,
+  ):
+    """Initializes the observation space adapter.
+
+    Args:
+      task_features_spec: The spec of all the features exposed by the task
+        layer.
+      selected_features: The features that will be exposed as observations. If
+        None, all features will be exposed, i.e. no filtering.
+      renamed_features: `RenameInfo` objects specifying which features should be
+        renamed and the corresponding new name. If empty or None, no renaming
+        will occur.
+      observation_type_mapper: An `ObservationTypeMapper` specifying how to
+        convert the task layer features data type (i.e. a Mapping[str,
+        ArrayType]) into the more generic type exposed by the GDM Environment
+        (i.e. a tree.Structure[ArrayType]). If None, an instance of
+        `_DefaultObservationTypeMapper` is used which converts the task logic
+        layer features dictionary to the more generic type (i.e.
+        `tree.Structure[ArrayType])` exposed by the environment.
+      max_float_dtype: The maximum float dtype to use for downcasting floats.
+    """
+    if not np.issubdtype(max_float_dtype, np.floating):
+      raise ValueError(
+          'max_float_dtype must be a floating point dtype. Got'
+          f' {max_float_dtype}'
+      )
+    self._max_float_dtype = max_float_dtype
+    self._max_bits = np.finfo(self._max_float_dtype).bits
+    self._task_features_spec = task_features_spec
+    self._selected_filter = selected_features
+    self._renamed_features = renamed_features or ()
+    self._observation_type_mapper = (
+        observation_type_mapper or _DefaultObservationTypeMapper()
+    )
+    self._check_specs_consistency()
+    # Compute the observation spec only once.
+    self._observation_spec = self._compute_observation_spec()
+
+  def _check_specs_consistency(self) -> None:
+    # Check that filter keys are present in the spec.
+    if self._selected_filter is not None:
+      all_features = self._task_features_spec.keys()
+      features = set()
+      for feature in self._selected_filter:
+        if feature not in all_features:
+          raise ValueError(f'Feature {feature} is not present in the spec.')
+        features.add(feature)
+    else:
+      # No filter applied. Select all features.
+      features = set(self._task_features_spec.keys())
+
+    # Check renaming.
+    for rename_info in self._renamed_features:
+      if rename_info.original_key not in features:
+        raise ValueError(
+            f'Feature {rename_info.original_key} is not present in the spec.'
+        )
+
+  def observations_from_features(
+      self, features: Mapping[str, gdmr_types.ArrayType]
+  ) -> tree.Structure[gdmr_types.ArrayType]:
+    """Converts the features into the final environment observations."""
+    # 1. Filter the observations.
+    if (selected_features := self._selected_filter) is None:
+      # No filter. Expose all observations.
+      filtered_features = dict(features)
+    else:
+      filtered_features = {
+          k: v for k, v in features.items() if k in selected_features  # pytype: disable=unsupported-operands
+      }
+
+    # 2. Downcast floats to max_float_dtype.
+    filtered_features = {
+        k: self._downcast_if_necessary(v) for k, v in filtered_features.items()
+    }
+
+    # 3. Rename.
+    for rename_info in self._renamed_features:
+      # Rename the feature.
+      value = filtered_features[rename_info.original_key]
+      del filtered_features[rename_info.original_key]
+      filtered_features[rename_info.renamed_key] = value
+
+    # 4. Convert type.
+    return self._observation_type_mapper.to_observations(filtered_features)
+
+  def _compute_observation_spec(self) -> gdmr_types.ObservationSpec:
+    """Computes the observation spec."""
+    # 1. Filter the specs
+    if (features_to_filter := self._selected_filter) is None:
+      # The observation spec corresponds to the task features spec.
+      filtered_specs = dict(self._task_features_spec)
+    else:
+      filtered_specs = {
+          k: v
+          for k, v in self._task_features_spec.items()
+          if k in features_to_filter  # pytype: disable=unsupported-operands
+      }
+
+    # 2. Downcast floats to max_float_dtype.
+    for k, v in filtered_specs.items():
+      if self._dtype_needs_downcast(v.dtype):
+        filtered_specs[k] = v.replace(dtype=self._max_float_dtype)
+
+    # 3. Rename.
+    for rename_info in self._renamed_features:
+      # Rename the feature.
+      value = filtered_specs[rename_info.original_key]
+      del filtered_specs[rename_info.original_key]
+      filtered_specs[rename_info.renamed_key] = value
+
+    # 4. Convert the type.
+    return self._observation_type_mapper.to_observation_spec(filtered_specs)
+
+  def observation_spec(self) -> gdmr_types.ObservationSpec:
+    """Returns the observation spec."""
+    return self._observation_spec
+
+  def task_features_keys(self) -> set[str]:
+    """Returns the task features keys that will be converted by this adapter."""
+    return set(self._task_features_spec.keys())
+
+  def _downcast_if_necessary(
+      self, value: gdmr_types.ArrayType
+  ) -> gdmr_types.ArrayType:
+    if (
+        hasattr(value, 'dtype') and self._dtype_needs_downcast(value.dtype)
+    ) or self._dtype_needs_downcast(type(value)):
+      return np.asarray(value).astype(self._max_float_dtype)
+    else:
+      return value
+
+  def _dtype_needs_downcast(self, dtype: npt.DTypeLike) -> bool:
+    return (
+        np.issubdtype(dtype, np.floating)
+        and np.finfo(dtype).bits > self._max_bits
+    )

src/experimental/reaf/core/device.py

@@ -0,0 +1,54 @@
+# Copyright 2025 Google LLC
+#
+# 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
+#
+# https://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.
+"""REAF basic device to interface with the robotic setup."""
+
+import abc
+from collections.abc import Mapping
+from dm_env import specs
+from gdm_robotics.interfaces import types as gdmr_types
+
+
+class Device(abc.ABC):
+  """REAF basic device to interface with the robotic setup.
+
+  A device defines a single piece in the robotic setup. It should be
+  hermetic, that is, not depending on other Devices. The coordination of the
+  devices is responsibility of the DeviceCoordinator.
+
+  Important: a Device should return the commands and measurements specs
+  immediately after initialisation without the need for any explicit
+  initialisation, nor for resource acquisition (e.g. connecting to the
+  hardware).
+  """
+
+  @property
+  @abc.abstractmethod
+  def name(self) -> str:
+    """Returns the name of this device."""
+
+  @abc.abstractmethod
+  def commands_spec(self) -> Mapping[str, gdmr_types.AnyArraySpec]:
+    """Returns the commands specs for this device."""
+
+  @abc.abstractmethod
+  def measurements_spec(self) -> Mapping[str, specs.Array]:
+    """Returns the measurements specs for this device."""
+
+  @abc.abstractmethod
+  def set_commands(self, commands: Mapping[str, gdmr_types.ArrayType]) -> None:
+    """Sets the commands for this device."""
+
+  @abc.abstractmethod
+  def get_measurements(self) -> Mapping[str, gdmr_types.ArrayType]:
+    """Returns the measurements provided by this device."""

src/experimental/reaf/core/device_coordinator.py

@@ -0,0 +1,87 @@
+# Copyright 2025 Google LLC
+#
+# 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
+#
+# https://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.
+"""Coordinates the devices composing a robotic setup."""
+
+import abc
+from collections.abc import Iterable
+from reaf.core import device
+
+
+class DeviceCoordinator(abc.ABC):
+  """Coordinates the devices composing a robotic setup.
+
+  The `DeviceCoordinator` object is responsible for coordinating all the
+  devices constituting the robotic setup. Whilst the Device is hermetic,
+  the coordinator is responsible for passing information from one device to
+  the other if required. For example in a bimanual setup the coordinator is
+  charged with passing the position of each robot to the other so we can ensure
+  proper and safe interaction such as for example collision avoidance.
+
+  The `DeviceCoordinator` can be configurable to enable different
+  properties on the robotic setup, e.g. adding or not adding a `Device` or
+  forwarding configuration to each `Device`.
+
+  At the very least, the coordinator must implement `get_devices`
+  to return all the devices. We also provide `on_begin_stepping` and
+  `on_end_stepping` methods that will be called before the start of an episode
+  and after the end of the episode respectively. Note that resource acquisition
+  and subsequent release is completely up to the implementation.
+
+  Finally, `before_set_commands`/`before_get_measurements` can be implemented to
+  coordinate devices behaviour before their corresponding functions are
+  called.
+  """
+
+  @property
+  @abc.abstractmethod
+  def name(self) -> str:
+    """Returns the name of the coordinator."""
+
+  @abc.abstractmethod
+  def get_devices(self) -> Iterable[device.Device]:
+    """Returns the devices composing the embodiment."""
+
+  # Lifecycle methods.
+
+  def on_begin_stepping(self) -> None:
+    """Prepares the coordinator for having its devices called repeatedly.
+
+    After `on_begin_stepping` the devices returned by `get_devices` will have
+    their `set_commands` and `get_measurements` called repeatedly until
+    `on_end_stepping` is called on this coordinator.
+    """
+
+  def on_end_stepping(self) -> None:
+    """Notifies the coordinator that the devices are no longer called.
+
+    After `on_end_stepping` the devices returned by `get_devices` will not have
+    their `set_commands` and `get_measurements` called anymore until this
+    coordinator `on_begin_stepping` method is notified again.
+    """
+
+  # Step hooks methods.
+
+  def before_set_commands(self) -> None:
+    """Prepares the coordinator to have its devices set_commands called."""
+
+  def after_set_commands(self) -> None:
+    """Notifies the coordinator that its devices got `set_commands` called."""
+
+  def before_get_measurements(self) -> None:
+    """Prepares the coordinator to have its devices get_measurements called.
+
+    This method gets called immediately before the devices `get_measurements`
+    method is called and can be used to customise the devices state given the
+    whole setup state.
+    """

src/experimental/reaf/core/discount_provider.py

@@ -0,0 +1,61 @@
+# Copyright 2025 Google LLC
+#
+# 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
+#
+# https://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.
+"""Computes the discount."""
+
+import abc
+from collections.abc import Mapping
+
+from dm_env import specs
+from gdm_robotics.interfaces import types as gdmr_types
+from reaf.core import termination_checker
+import tree
+
+
+class DiscountProvider(abc.ABC):
+  """Computes the discount."""
+
+  @abc.abstractmethod
+  def name(self) -> str:
+    """Returns a unique string identifier for this object."""
+
+  @abc.abstractmethod
+  def compute_discount(
+      self,
+      required_features: Mapping[str, gdmr_types.ArrayType],
+      termination_state: termination_checker.TerminationResult,
+  ) -> tree.Structure[gdmr_types.ArrayType]:
+    """Computes the discount.
+
+    Args:
+      required_features: Measurements and features computed by the task logic
+        that are required by this provider, i.e. that have keys specified by
+        `required_features_keys`.
+      termination_state: The termination state as computed by the termination
+        checkers.  Returns the discount.
+
+    Returns:
+      The discount.
+    """
+
+  @abc.abstractmethod
+  def discount_spec(self) -> tree.Structure[specs.Array]:
+    """Returns the spec of the discount."""
+
+  @abc.abstractmethod
+  def required_features_keys(self) -> set[str]:
+    """Returns the feature keys that are required to compute the discount."""
+
+  def reset(self) -> None:
+    """Resets the internal state of the discount provider."""
+    ...

src/experimental/reaf/core/entity.py

@@ -0,0 +1,65 @@
+# Copyright 2025 Google LLC
+#
+# 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
+#
+# https://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.
+"""Basic REAF-sim protocol to interface with the simulation."""
+
+from collections.abc import Mapping
+import typing
+
+from dm_env import specs
+from gdm_robotics.interfaces import types as gdmr_types
+
+
+class Entity(typing.Protocol):
+  """Basic REAF component to interface with the simulation.
+
+  An entity defines a single component in the simulation that consumes substep
+  commands and outputs substep measurements at every simulation substep. It
+  should be hermetic, that is, not depending on other Entities.
+
+  Important: an Entity should return the substep commands and substep
+  measurements specs immediately after initialisation without the need for any
+  explicit initialisation.
+  """
+
+  @property
+  def name(self) -> str:
+    """Instance name."""
+
+  def reset(self):
+    """Resets the entity."""
+
+  def substep_commands_spec(
+      self,
+  ) -> Mapping[str, specs.Array]:
+    """Spec for the substep commands."""
+
+  def substep_measurements_spec(
+      self,
+  ) -> Mapping[str, specs.Array]:
+    """Spec for the substep measurements."""
+
+  def set_substep_commands(
+      self,
+      model: typing.Any,
+      data: typing.Any,
+      consumed_substep_commands: Mapping[str, gdmr_types.ArrayType],
+  ) -> None:
+    """Sets the substep commands."""
+
+  def get_substep_measurements(
+      self,
+      model: typing.Any,
+      data: typing.Any,
+  ) -> Mapping[str, gdmr_types.ArrayType]:
+    """Returns the substep measurements."""

src/experimental/reaf/core/environment.py

@@ -0,0 +1,490 @@
+# Copyright 2025 Google LLC
+#
+# 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
+#
+# https://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.
+"""The Robotics Environment Authoring Framework (REAF) Environment class."""
+
+import abc
+from collections.abc import Mapping
+import enum
+from typing import Generic
+
+from absl import logging
+import dm_env
+from dm_env import specs
+from gdm_robotics.interfaces import environment as gdmr_env
+from gdm_robotics.interfaces import types as gdmr_types
+import numpy as np
+from reaf.core import action_space_adapter as reaf_action_space_adapter
+from reaf.core import data_acquisition_and_control_layer as reaf_dacl
+from reaf.core import default_observation_space_adapter
+from reaf.core import logger as reaf_logger
+from reaf.core import observation_space_adapter as reaf_observation_space_adapter
+from reaf.core import pass_through_action_space_adapter
+from reaf.core import task_logic_layer as reaf_tll
+import tree
+
+
+class ActionSpecEnforcementOption(enum.StrEnum):
+  """Options for action spec enforcement."""
+
+  CLIP_TO_SPEC = "clip_to_spec"
+  IGNORE = "ignore"
+  WARNING = "warning"
+  RAISE_ERROR = "raise_error"
+
+
+class EnvironmentReset(abc.ABC, Generic[gdmr_env.ResetOptions]):
+  """Support for general resets adhering to the GDM environment API."""
+
+  @abc.abstractmethod
+  def do_reset(
+      self,
+      config: gdmr_env.ResetOptions,
+  ) -> None:
+    """Resets the environment."""
+
+  def default_reset_configuration(self) -> gdmr_env.ResetOptions:
+    """Returns the default reset configuration."""
+    return gdmr_env.Options()
+
+
+class EndOfEpisodeHandler:
+  """Handler called after the last episode step."""
+
+  def on_end_of_episode_stepping(self, final_timestep: dm_env.TimeStep) -> None:
+    """Called when the episode has ended stepping.
+
+    This will be called at the end of every episode, after all other triggers
+    have been resolved. Episodes can end either due to truncation or
+    termination, i.e. `timestep.step_type` is `StepType.LAST`, or due to an
+    early call to `Environment.reset()`. To verify whether it has indeed
+    ended due to truncation or termination, the implementer should test
+    `timestep.last()`.
+
+    Note that the first reset after environment construction will not trigger
+    this handler, but it will be triggered before resolving any subsequent
+    environment resets, either implicit or explicit.
+
+    Args:
+      final_timestep: The final timestep of the episode that ended stepping.
+    """
+
+
+class EnvironmentCloser(abc.ABC):
+  """Handler called when the environment is closed."""
+
+  @abc.abstractmethod
+  def close(self) -> None:
+    """Releases resources when the environment is closed.
+
+    This method is called automatically when exiting the environment's
+    context manager (`with` statement).
+    """
+
+
+class Environment(gdmr_env.Environment):
+  """The Robotics Environment Authoring Framework (REAF) Environment class."""
+
+  def __init__(
+      self,
+      *,
+      data_acquisition_and_control_layer: reaf_dacl.DataAcquisitionAndControlLayer,
+      task_logic_layer: reaf_tll.TaskLogicLayer,
+      environment_reset: EnvironmentReset,
+      action_space_adapter: (
+          reaf_action_space_adapter.ActionSpaceAdapter | None
+      ) = None,
+      observation_space_adapter: (
+          reaf_observation_space_adapter.ObservationSpaceAdapter | None
+      ) = None,
+      end_of_episode_handler: EndOfEpisodeHandler | None = None,
+      environment_closer: EnvironmentCloser | None = None,
+      action_spec_enforcement_option: ActionSpecEnforcementOption = ActionSpecEnforcementOption.RAISE_ERROR,
+  ):
+    """Creates an environment.
+
+    Args:
+      data_acquisition_and_control_layer: The layer for communicating with the
+        specific robotic setup.
+      task_logic_layer: The layer in charge of defining the task.
+      environment_reset: The `EnvironmentReset` specifying the function to be
+        called at environment reset and the default environment reset
+        configuration.
+      action_space_adapter: Adapter from the agent action space to the flattened
+        commands accepted by the task layer. If None the
+        PassThroughActionSpaceAdapter is used, meaning the entirety of the
+        commands dictionary is exposed to the agent.
+      observation_space_adapter: Adapter from the computed features to the
+        observations that are exposed to the agent. If None the
+        DefaultObservationSpaceAdapter is used, meaning all the features are
+        exposed to the agent as observations.
+      end_of_episode_handler: Called at the end of an episode, after the last
+        step.
+      environment_closer: Specifies the handler to be called when the
+        environment is closed. This is called automatically on exit if the
+        environment is used as a context manager. If None, no action is
+        performed at close.
+      action_spec_enforcement_option: How to enforce the action spec. If
+        `CLIP_TO_SPEC`, the action will be clipped to the spec. If `WARNING`, an
+        warning logged if the action is outside the spec. If `RAISE_ERROR`, an
+        error will be raised if the action is outside the  spec. If `IGNORE`,
+        the action will be passed through. Default is `RAISE_ERROR`.
+    """
+
+    self._data_acquisition_and_control_layer = (
+        data_acquisition_and_control_layer
+    )
+    self._task_logic_layer = task_logic_layer
+    self._end_of_episode_handler = (
+        end_of_episode_handler or EndOfEpisodeHandler()
+    )
+    self._environment_reset = environment_reset
+    self._environment_closer = environment_closer
+    self._action_spec_enforcement_option = action_spec_enforcement_option
+
+    # Before assigning the adapters, validate the specs on the task logic layer
+    # and the DACL.
+    self._validate_dacl_and_ttl_specs()
+
+    ttl_commands_spec = self._task_logic_layer.commands_spec(
+        self._data_acquisition_and_control_layer.commands_spec()
+    )
+    ttl_features_spec = self._task_logic_layer.features_spec(
+        self._data_acquisition_and_control_layer.measurements_spec()
+    )
+
+    if action_space_adapter is None:
+      action_space_adapter = (
+          pass_through_action_space_adapter.PassThroughActionSpaceAdapter(
+              commands_spec=ttl_commands_spec
+          )
+      )
+    self._action_space_adapter = action_space_adapter
+
+    if observation_space_adapter is None:
+      observation_space_adapter = (
+          default_observation_space_adapter.DefaultObservationSpaceAdapter(
+              task_features_spec=ttl_features_spec,
+              selected_features=None,
+              renamed_features=None,
+              observation_type_mapper=None,
+          )
+      )
+    self._observation_space_adapter = observation_space_adapter
+
+    # Now we can validate the adapters.
+    self._validate_adapters_specs()
+
+    self._last_timestep: dm_env.TimeStep | None = None
+    self._should_finalize_episode = False
+    self._timestep_spec = gdmr_types.TimeStepSpec(
+        step_type=gdmr_types.STEP_TYPE_SPEC,
+        reward=self._task_logic_layer.reward_spec(),
+        discount=self._task_logic_layer.discount_spec(),
+        # The observation spec corresponds to the one exposed by the adapter.
+        observation=self._observation_space_adapter.observation_spec(),
+    )
+
+    self._zero_reward, self._zero_discount = tree.map_structure(
+        _read_only_zeros_like_spec,
+        (self._timestep_spec.reward, self._timestep_spec.discount),
+    )
+
+  def close(self) -> None:
+    """Frees any resources used by the environment."""
+    if self._environment_closer is not None:
+      self._environment_closer.close()
+
+  def default_reset_options(self) -> gdmr_env.ResetOptions:
+    return self._environment_reset.default_reset_configuration()
+
+  def reset_with_options(
+      self,
+      *,
+      options: gdmr_env.ResetOptions,
+  ) -> dm_env.TimeStep:
+    """Starts a new sequence and returns the first `TimeStep`."""
+    if self._should_finalize_episode:
+      self._finalize_episode()
+    self._environment_reset.do_reset(options)
+    self._task_logic_layer.perform_reset()
+    measurements = self._data_acquisition_and_control_layer.begin_stepping()
+    features = self._task_logic_layer.compute_all_features(measurements)
+    observations = self._compute_observations_from_features(features)
+
+    self._last_timestep = self._restart(observation=observations)
+    # Make sure any early reset after this one triggers `_finalize_episode`.
+    self._should_finalize_episode = True
+    return self._last_timestep
+
+  def action_spec(self) -> gdmr_types.ActionSpec:
+    """Defines the actions that should be provided to `step`."""
+    # The action spec corresponds to the one exposed by the adapter.
+    return self._action_space_adapter.action_spec()
+
+  def timestep_spec(self) -> gdmr_types.TimeStepSpec:
+    """Returns the spec associated to the returned TimeStep."""
+    return self._timestep_spec
+
+  def step(self, action: gdmr_types.ActionType) -> dm_env.TimeStep:
+    """Updates the environment according to action and returns a `TimeStep`."""
+
+    action = self._enforce_action_spec(action)
+    if self._last_timestep is None or self._last_timestep.last():
+      return self.reset()
+
+    # Process the action to obtain a command.
+    commands = self._compute_commands_from_agent_action(action)
+    commands = self._task_logic_layer.compute_final_commands(commands)
+    measurements = self._data_acquisition_and_control_layer.step(commands)
+
+    # Compute all the features.
+    features = self._task_logic_layer.compute_all_features(measurements)
+
+    # Compute the elements of the timestep.
+    reward = self._task_logic_layer.compute_reward(features)
+    termination_state = self._task_logic_layer.check_for_termination(features)
+    discount = self._task_logic_layer.compute_discount(
+        features, termination_state
+    )
+
+    observations = self._compute_observations_from_features(features)
+
+    if termination_state.is_terminated():
+      self._last_timestep = self._termination(
+          reward=reward, observation=observations
+      )
+    elif termination_state.is_truncated():
+      self._last_timestep = self._truncation(
+          reward=reward, observation=observations, discount=discount
+      )
+    else:
+      self._last_timestep = self._transition(
+          reward=reward, observation=observations, discount=discount
+      )
+
+    if self._last_timestep.last():
+      self._finalize_episode()
+    return self._last_timestep
+
+  def _finalize_episode(self) -> None:
+    self._data_acquisition_and_control_layer.end_stepping()
+    # It's crucial to call `end_stepping` on the dacl before invoking the end
+    # of episode handler. This ensures no further `set_command` or
+    # `get_measurements` calls are made. In contrast, the end of episode
+    # handler might interact with devices, requiring them to be informed
+    # beforehand.
+    self._end_of_episode_handler.on_end_of_episode_stepping(self._last_timestep)
+    self._should_finalize_episode = False
+
+  @property
+  def data_acquisition_and_control_layer(
+      self,
+  ) -> reaf_dacl.DataAcquisitionAndControlLayer:
+    return self._data_acquisition_and_control_layer
+
+  @property
+  def task_logic_layer(self) -> reaf_tll.TaskLogicLayer:
+    return self._task_logic_layer
+
+  @property
+  def environment_reset(self) -> EnvironmentReset:
+    return self._environment_reset
+
+  @environment_reset.setter
+  def environment_reset(self, environment_reset: EnvironmentReset) -> None:
+    self._environment_reset = environment_reset
+
+  def add_logger(self, logger: reaf_logger.Logger) -> None:
+    self._task_logic_layer.add_logger(logger)
+
+  def remove_logger(self, logger: reaf_logger.Logger) -> None:
+    self._task_logic_layer.remove_logger(logger)
+
+  def _validate_dacl_and_ttl_specs(self) -> None:
+    """Validates the specs on the task logic layer."""
+    # Validate the spec on the task logic layer.
+    self._task_logic_layer.validate_spec(
+        dacl_commands_spec=(
+            self._data_acquisition_and_control_layer.commands_spec()
+        ),
+        dacl_measurements_spec=(
+            self._data_acquisition_and_control_layer.measurements_spec()
+        ),
+    )
+
+  def _validate_adapters_specs(self) -> None:
+    # Collect the full commands and features spec and validate them against
+    # the adapters.
+    commands_spec = set(
+        self._task_logic_layer.commands_spec(
+            self._data_acquisition_and_control_layer.commands_spec()
+        ).keys()
+    )
+    features_spec = set(
+        self._task_logic_layer.features_spec(
+            self._data_acquisition_and_control_layer.measurements_spec()
+        )
+    )
+
+    # Check the action space adapter.
+    adapter_keys = self._action_space_adapter.task_commands_keys()
+
+    if adapter_keys != commands_spec:
+      raise ValueError(
+          "Mismatch between commands exposed by the action space adapter:"
+          f" {adapter_keys} and commands spec expected by the task layer:"
+          f" {commands_spec}."
+      )
+
+    # Check the observation spec adapter.
+    adapter_keys = self._observation_space_adapter.task_features_keys()
+    if not adapter_keys.issubset(features_spec):
+      raise ValueError(
+          "Failed to validate observation space adapter specs. Missing keys:"
+          f" {adapter_keys - features_spec}"
+      )
+
+  def _compute_observations_from_features(
+      self, features: Mapping[str, gdmr_types.ArrayType]
+  ) -> tree.Structure[gdmr_types.ArrayType]:
+    return self._observation_space_adapter.observations_from_features(features)
+
+  def _compute_commands_from_agent_action(
+      self, agent_action: gdmr_types.ActionType
+  ) -> Mapping[str, gdmr_types.ArrayType]:
+    return self._action_space_adapter.commands_from_environment_action(
+        agent_action
+    )
+
+  def _restart(
+      self,
+      observation: tree.Structure[gdmr_types.ArrayType],
+  ) -> dm_env.TimeStep:
+    """Returns a `TimeStep` with `step_type` set to `StepType.FIRST`."""
+    return dm_env.TimeStep(
+        step_type=np.asarray(dm_env.StepType.FIRST, dtype=np.uint8),
+        observation=observation,
+        reward=self._zero_reward,
+        discount=self._zero_discount,
+    )
+
+  def _transition(
+      self,
+      reward: tree.Structure[gdmr_types.ArrayType],
+      observation: tree.Structure[gdmr_types.ArrayType],
+      discount: tree.Structure[gdmr_types.ArrayType],
+  ) -> dm_env.TimeStep:
+    """Returns a `TimeStep` with `step_type` set to `StepType.MID`."""
+    return dm_env.TimeStep(
+        step_type=np.asarray(dm_env.StepType.MID, dtype=np.uint8),
+        observation=observation,
+        reward=reward,
+        discount=discount,
+    )
+
+  def _termination(
+      self,
+      reward: tree.Structure[gdmr_types.ArrayType],
+      observation: tree.Structure[gdmr_types.ArrayType],
+  ) -> dm_env.TimeStep:
+    """Returns a `TimeStep` with `step_type` set to `StepType.LAST`."""
+    return dm_env.TimeStep(
+        step_type=np.asarray(dm_env.StepType.LAST, dtype=np.uint8),
+        observation=observation,
+        reward=reward,
+        discount=self._zero_discount,
+    )
+
+  def _truncation(
+      self,
+      reward: tree.Structure[gdmr_types.ArrayType],
+      observation: tree.Structure[gdmr_types.ArrayType],
+      discount: tree.Structure[gdmr_types.ArrayType],
+  ) -> dm_env.TimeStep:
+    """Returns a `TimeStep` with `step_type` set to `StepType.LAST`."""
+    return dm_env.TimeStep(
+        step_type=np.asarray(dm_env.StepType.LAST, dtype=np.uint8),
+        observation=observation,
+        reward=reward,
+        discount=discount,
+    )
+
+  def _enforce_action_spec(
+      self, action: gdmr_types.ActionType
+  ) -> gdmr_types.ActionType:
+    """Enforces the action spec."""
+    match self._action_spec_enforcement_option:
+      case ActionSpecEnforcementOption.IGNORE:
+        pass
+      case ActionSpecEnforcementOption.CLIP_TO_SPEC:
+        try:
+
+          def clip_to_spec(a, s):
+            if isinstance(s, specs.BoundedArray):
+              return np.clip(a, s.minimum, s.maximum)
+            return a
+
+          action = tree.map_structure(
+              clip_to_spec,
+              action,
+              self._action_space_adapter.action_spec(),
+          )
+        except ValueError as e:
+          raise ValueError(
+              "Failed to clip action to spec. Action:"
+              f" {action} and spec: {self._action_space_adapter.action_spec()}"
+          ) from e
+      case ActionSpecEnforcementOption.WARNING:
+
+        def _validate_without_raising(a, s):
+          dtype_ok = s.dtype == a.dtype
+          shape_ok = s.shape == a.shape
+          minimum_ok = True
+          maximum_ok = True
+          if isinstance(s, specs.BoundedArray):
+            minimum_ok = (s.minimum <= a).all()
+            maximum_ok = (a <= s.maximum).all()
+          return dtype_ok and shape_ok and minimum_ok and maximum_ok
+
+        if not all(
+            tree.flatten(
+                tree.map_structure(
+                    _validate_without_raising,
+                    action,
+                    self._action_space_adapter.action_spec(),
+                )
+            )
+        ):
+          logging.warning(
+              "Failed to validate action against spec. Action: %r and spec: %r",
+              action,
+              self._action_space_adapter.action_spec(),
+          )
+      case ActionSpecEnforcementOption.RAISE_ERROR:
+        action = tree.map_structure(
+            lambda a, spec: spec.validate(a), action, self.action_spec()
+        )
+      case _:
+        raise ValueError(
+            "Unknown action spec enforcement option:"
+            f" {self._action_spec_enforcement_option}"
+        )
+    return action
+
+
+def _read_only_zeros_like_spec(spec: specs.Array) -> np.ndarray:
+  """Returns a zero array matching the specified spec."""
+  arr = np.zeros(shape=spec.shape, dtype=spec.dtype)
+  arr.flags.writeable = False
+  return arr

src/experimental/reaf/core/features_observer.py

@@ -0,0 +1,34 @@
+# Copyright 2025 Google LLC
+#
+# 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
+#
+# https://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.
+"""Observe all the produced features and measurements."""
+
+import abc
+from collections.abc import Mapping
+
+from gdm_robotics.interfaces import types as gdmr_types
+
+
+class FeaturesObserver(abc.ABC):
+  """Observe all the produced features and measurements."""
+
+  @property
+  @abc.abstractmethod
+  def name(self) -> str:
+    """Returns a unique string identifier for this object."""
+
+  @abc.abstractmethod
+  def observe_features(
+      self, features: Mapping[str, gdmr_types.ArrayType]
+  ) -> None:
+    """Observes all the features and measurements."""

src/experimental/reaf/core/features_producer.py

@@ -0,0 +1,56 @@
+# Copyright 2025 Google LLC
+#
+# 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
+#
+# https://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.
+"""Produces additional features to be exposed by the task logic layer."""
+
+import abc
+from collections.abc import Mapping
+
+from dm_env import specs
+from gdm_robotics.interfaces import types as gdmr_types
+
+
+class FeaturesProducer(abc.ABC):
+  """Produces additional features to be exposed by the task logic layer."""
+
+  @property
+  @abc.abstractmethod
+  def name(self) -> str:
+    """Returns a unique string identifier for this object."""
+
+  @abc.abstractmethod
+  def produce_features(
+      self, required_features: Mapping[str, gdmr_types.ArrayType]
+  ) -> Mapping[str, gdmr_types.ArrayType]:
+    """Produces additional features for the environment.
+
+    Args:
+      required_features: Measurements and features generated by previous
+        producers in the processing chain that are required by this processor,
+        i.e. with keys specified by `required_features_keys`.
+
+    Returns additional features that will be added to the global measurements
+      and features dictionary.
+    """
+
+  @abc.abstractmethod
+  def produced_features_spec(self) -> Mapping[str, specs.Array]:
+    """Returns the spec of the features produced by this producer."""
+
+  @abc.abstractmethod
+  def required_features_keys(self) -> set[str]:
+    """Returns the keys that are required to produce the new features."""
+
+  def reset(self) -> None:
+    """Resets the internal state of the feature producer."""
+    ...

src/experimental/reaf/core/logger.py

@@ -0,0 +1,80 @@
+# Copyright 2025 Google LLC
+#
+# 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
+#
+# https://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.
+"""Support logging inside the task logic layer."""
+
+import abc
+from collections.abc import Mapping
+
+from gdm_robotics.interfaces import types as gdmr_types
+
+
+class Logger(abc.ABC):
+  """Support logging inside the task logic layer.
+
+  Lifecycle
+  For each environment step, these member functions are called in this order:
+    1. `record_measurements` is called with raw measurements from the sensors.
+    2. `record_features` is called with features derived from the measurements.
+    3. `record_commands_processing` is called for each
+       `CommandsProcessor.process_commands` invocation, tracking the
+       transformation of commands.
+    4. `record_final_commands` is called once with the final commands sent to
+       the DACL.
+
+    Notes:
+    An environment is first reset().  This triggers the first two steps above.
+      See reset_with_options in ./environment.py.
+
+    After reset, step is called repeatedly.
+    1. This first triggers steps 3 and 4 (See compute_final_commands in TLL
+       called from step in ./environment.py)
+    2. Features are computed (see compute_all_features in TLL called from
+       step in ./environment.py), triggering steps 1 and 2.
+  """
+
+  @property
+  @abc.abstractmethod
+  def name(self) -> str:
+    """Unique string identifier for this object."""
+
+  def record_measurements(
+      self, measurements: Mapping[str, gdmr_types.ArrayType]
+  ) -> None:
+    """Called once with all the measurements from the DACL."""
+
+  def record_features(
+      self, features: Mapping[str, gdmr_types.ArrayType]
+  ) -> None:
+    """Called once with all the features computed in the Task Layer."""
+
+  def record_final_commands(
+      self, commands: Mapping[str, gdmr_types.ArrayType]
+  ) -> None:
+    """Called once with the final commands sent to the DACL."""
+
+  def record_commands_processing(
+      self,
+      name: str,
+      consumed_commands: Mapping[str, gdmr_types.ArrayType],
+      produced_commands: Mapping[str, gdmr_types.ArrayType],
+  ) -> None:
+    """Called once per call to `process_commands` for each CommandsProcessor.
+
+    Args:
+      name: Name of the `CommandsProcessor`.
+      consumed_commands: The commands consumed by the current
+        `CommandsProcessor`.
+      produced_commands: The commands produced by the current
+        `CommandsProcessor`.
+    """

src/experimental/reaf/core/numpy_mock_assertions.py

@@ -0,0 +1,98 @@
+# Copyright 2025 Google LLC
+#
+# 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
+#
+# https://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.
+"""Testing functions for asserting on Mock objects with numpy structures."""
+
+from collections.abc import Sequence
+from unittest import mock
+import numpy as np
+
+
+def assert_called_once_with(mock_obj: mock.Mock, *args, **kwargs) -> None:
+  if mock_obj.call_count != 1:
+    raise AssertionError(
+        f"Expected exactly one call to {mock_obj}, got {mock_obj.call_count}"
+    )
+
+  assert_called_with(mock_obj, *args, **kwargs)
+
+
+def assert_called_with(mock_obj: mock.Mock, *args, **kwargs) -> None:
+  """Asserts that the last call to mock_obj had the specified arguments."""
+  if mock_obj.call_args is None:
+    raise AssertionError(
+        f"Mock object {mock_obj} not called. Expected one call."
+    )
+  call_args, call_kwargs = mock_obj.call_args
+  np.testing.assert_equal(call_args, args)
+  np.testing.assert_equal(call_kwargs, kwargs)
+
+
+def assert_has_calls(
+    mock_obj: mock.Mock, calls: Sequence[mock._Call], any_order: bool = False
+) -> None:
+  """Asserts that mock_obj has been called with the specified calls."""
+  mock_calls = mock_obj.mock_calls
+
+  # Check that there are at least enough calls.
+  if mock_obj.call_count < len(calls):
+    raise AssertionError(
+        f"Expected at least {len(calls)} calls to {mock_obj}, got"
+        f" {mock_obj.call_count}"
+    )
+
+  def _calls_are_equal(actual: mock._Call, expected: mock._Call) -> bool:
+    _, actual_args, actual_kwargs = actual
+    _, expected_args, expected_kwargs = expected
+    # Quickest way to transform the assertion into a comparator.
+    try:
+      np.testing.assert_equal(actual_args, expected_args)
+      np.testing.assert_equal(actual_kwargs, expected_kwargs)
+      return True
+    except AssertionError:
+      return False
+
+  if any_order:
+    # We just check for the calls to be contained.
+    for expected_call in calls:
+      for actual_call in mock_calls:
+        if _calls_are_equal(actual_call, expected_call):
+          break
+      raise AssertionError(
+          f"Expected call {expected_call} not found in mock calls {mock_calls}."
+      )
+    return
+
+  # We need to check in order, but first find the first call.
+  starting_index = -1
+  first_expected_call = calls[0]
+  for index, actual_call in enumerate(mock_calls):
+    if _calls_are_equal(actual_call, first_expected_call):
+      starting_index = index
+      break
+  if starting_index == -1:
+    raise AssertionError(f"Calls {calls} not found in mock calls {mock_calls}.")
+
+  non_matching_calls = []
+
+  # We have the first element. Now we need to compare element wise.
+  for index, expected_call in enumerate(calls):
+    actual_call = mock_calls[starting_index + index]
+    if not _calls_are_equal(actual_call, expected_call):
+      non_matching_calls.append((index, expected_call, actual_call))
+
+  if non_matching_calls:
+    raise AssertionError(
+        f"Calls {calls} do not match mock calls {mock_calls}. Mismatch (index,"
+        f" expected, actual): {non_matching_calls}"
+    )

src/experimental/reaf/core/observation_space_adapter.py

@@ -0,0 +1,42 @@
+# Copyright 2025 Google LLC
+#
+# 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
+#
+# https://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.
+"""Adapts REAF features into observations exposed by the environment."""
+
+import abc
+from collections.abc import Mapping
+from gdm_robotics.interfaces import types as gdmr_types
+import tree
+
+
+class ObservationSpaceAdapter(abc.ABC):
+  """Adapts REAF features into observations exposed by the environment.
+
+  Implementations of this interface are responsible for converting the features
+  generated by the REAF task layer logic (i.e. dictionary of tensors) into the
+  more generic `observation` structure exposed by the environment.
+  """
+
+  @abc.abstractmethod
+  def observations_from_features(
+      self, features: Mapping[str, gdmr_types.ArrayType]
+  ) -> tree.Structure[gdmr_types.ArrayType]:
+    """Converts the REAF features into the environment observations."""
+
+  @abc.abstractmethod
+  def observation_spec(self) -> gdmr_types.ObservationSpec:
+    """Returns the observation spec."""
+
+  @abc.abstractmethod
+  def task_features_keys(self) -> set[str]:
+    """Returns the task features keys that will be converted by this adapter."""

src/experimental/reaf/core/pass_through_action_space_adapter.py

@@ -0,0 +1,55 @@
+# Copyright 2025 Google LLC
+#
+# 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
+#
+# https://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.
+"""Adapter that passes the commands spec through."""
+
+from collections.abc import Mapping
+from gdm_robotics.interfaces import types as gdmr_types
+from reaf.core import action_space_adapter
+
+
+class PassThroughActionSpaceAdapter(action_space_adapter.ActionSpaceAdapter):
+  """Adapter that passes the commands spec through.
+
+  NB the resulting environment will expose a dictionary as the action spec.
+  """
+
+  def __init__(self, commands_spec: Mapping[str, gdmr_types.AnyArraySpec]):
+    self._commands_spec = commands_spec
+
+  def commands_from_environment_action(
+      self, environment_action: gdmr_types.ActionType
+  ) -> Mapping[str, gdmr_types.ArrayType]:
+    """Returns commands accepted by REAF.
+
+    commands_from_environment_action usually accepts a gdmr_types.ActionType but
+    since this adapter passes the same action as the commands, it needs to be a
+    dict type in order to pass it through as a dict.
+
+    Args:
+      environment_action: The environment action(s) to pass as REAF commands.
+    """
+    if not isinstance(environment_action, dict):
+      raise ValueError(
+          'environment_action must be a dict, but got: '
+          f'{type(environment_action)}.'
+      )
+    return environment_action
+
+  def action_spec(self) -> gdmr_types.ActionSpec:
+    """Returns the action spec exposed by the environment."""
+    return self._commands_spec
+
+  def task_commands_keys(self) -> set[str]:
+    """Returns the keys for the commands exposed to the task layer."""
+    return set(self._commands_spec.keys())

src/experimental/reaf/core/reward_provider.py

@@ -0,0 +1,292 @@
+# Copyright 2025 Google LLC
+#
+# 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
+#
+# https://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.
+"""Computes the reward."""
+
+import abc
+from collections.abc import Mapping
+import operator
+from typing import Callable, TypeAlias, TypeVar, Union
+
+from dm_env import specs
+from gdm_robotics.interfaces import types as gdmr_types
+import numpy as np
+import tree
+
+
+RewardValue: TypeAlias = tree.Structure[gdmr_types.ArrayType]
+RewardSpec: TypeAlias = tree.Structure[specs.Array]
+
+
+class _RewardProvider(abc.ABC):
+  """Computes the reward.
+
+  Defines the interface for a reward provider.
+
+  Important: Users should not inherit from this class directly. Instead, use the
+  RewardProvider class later in this file.
+  """
+
+  @abc.abstractmethod
+  def name(self) -> str:
+    """Returns a unique string identifier for this object."""
+
+  @abc.abstractmethod
+  def compute_reward(
+      self, required_features: Mapping[str, gdmr_types.ArrayType]
+  ) -> RewardValue:
+    """Computes the reward.
+
+    Args:
+      required_features: Measurements and features computed by the task logic
+        that are required by this provider, i.e. that have keys specified by
+        `required_features_keys`.
+
+    Returns the computed reward.
+    """
+
+  @abc.abstractmethod
+  def reward_spec(self) -> RewardSpec:
+    """Returns the spec of the reward."""
+
+  @abc.abstractmethod
+  def required_features_keys(self) -> set[str]:
+    """Returns the feature keys that are required to compute the reward."""
+
+  def reset(self) -> None:
+    """Resets the internal state of the reward provider."""
+    ...
+
+
+RewardProviderOrValue: TypeAlias = Union['RewardProvider', RewardValue]
+
+
+S = TypeVar('S')
+T = TypeVar('T')
+UnaryOperator: TypeAlias = Callable[[S], S]
+BinaryOperator: TypeAlias = Callable[[S | T, S | T], S | T]
+
+
+class RewardProvider(_RewardProvider):
+  """Computes the reward.
+
+  Important: Users should inherit from this class and implement the abstract
+  methods defined in the interface _RewardProvider.
+  """
+
+  def __add__(self, other: RewardProviderOrValue):
+    return BinaryOperationRewardProvider(operator.add, self, other)
+
+  def __radd__(self, other: RewardProviderOrValue):
+    return BinaryOperationRewardProvider(operator.add, other, self)
+
+  def __sub__(self, other: RewardProviderOrValue):
+    return BinaryOperationRewardProvider(operator.sub, self, other)
+
+  def __rsub__(self, other: RewardProviderOrValue):
+    return BinaryOperationRewardProvider(operator.sub, other, self)
+
+  def __mul__(self, other: RewardProviderOrValue):
+    return BinaryOperationRewardProvider(operator.mul, self, other)
+
+  def __rmul__(self, other: RewardProviderOrValue):
+    return BinaryOperationRewardProvider(operator.mul, other, self)
+
+  def __truediv__(self, other: RewardProviderOrValue):
+    return BinaryOperationRewardProvider(operator.truediv, self, other)
+
+  def __rtruediv__(self, other: RewardProviderOrValue):
+    return BinaryOperationRewardProvider(operator.truediv, other, self)
+
+  def __floordiv__(self, other: RewardProviderOrValue):
+    return BinaryOperationRewardProvider(operator.floordiv, self, other)
+
+  def __rfloordiv__(self, other: RewardProviderOrValue):
+    return BinaryOperationRewardProvider(operator.floordiv, other, self)
+
+  def __pow__(self, other: RewardProviderOrValue):
+    return BinaryOperationRewardProvider(operator.pow, self, other)
+
+  def __rpow__(self, other: RewardProviderOrValue):
+    return BinaryOperationRewardProvider(operator.pow, other, self)
+
+  def __getitem__(self, index: slice):
+    return GetItemOperationRewardProvider(self, index)
+
+  def __neg__(self):
+    return UnaryOperationRewardProvider(operator.neg, self)
+
+
+class ConstantRewardProvider(RewardProvider):
+  """A RewardProvider that always returns the same reward."""
+
+  def __init__(self, reward: RewardValue):
+    super().__init__()
+    self._reward = reward
+
+  def name(self) -> str:
+    return str(self._reward)
+
+  def compute_reward(
+      self, required_features: Mapping[str, gdmr_types.ArrayType]
+  ) -> RewardValue:
+    return self._reward
+
+  def reward_spec(self) -> RewardSpec:
+    return tree.map_structure(
+        lambda v: specs.Array(v.shape, v.dtype), self._reward
+    )
+
+  def required_features_keys(self) -> set[str]:
+    return set()
+
+
+class BinaryOperationRewardProvider(RewardProvider):
+  """Applies a binary operator to the result of two reward providers."""
+
+  def __init__(
+      self,
+      op: BinaryOperator,
+      first_reward_provider: RewardProviderOrValue,
+      second_reward_provider: RewardProviderOrValue,
+  ):
+    super().__init__()
+    if not isinstance(first_reward_provider, RewardProvider):
+      first_reward_provider = ConstantRewardProvider(first_reward_provider)
+    if not isinstance(second_reward_provider, RewardProvider):
+      second_reward_provider = ConstantRewardProvider(second_reward_provider)
+    first_spec = first_reward_provider.reward_spec()
+    second_spec = second_reward_provider.reward_spec()
+    tree.assert_same_structure(first_spec, second_spec)
+    assert all(
+        tree.flatten(
+            tree.map_structure(
+                lambda s1, s2: s1.shape == s2.shape and s1.dtype == s2.dtype,
+                first_spec,
+                second_spec,
+            )
+        )
+    )
+    self._op = op
+    self._first_reward_provider = first_reward_provider
+    self._second_reward_provider = second_reward_provider
+    self._reward_spec = first_reward_provider.reward_spec()
+    self._first_required_features_keys = (
+        first_reward_provider.required_features_keys()
+    )
+    self._second_required_features_keys = (
+        second_reward_provider.required_features_keys()
+    )
+
+  def name(self) -> str:
+    op_name = getattr(self._op, '__name__', str(self._op))
+    return (
+        f'{op_name}({self._first_reward_provider.name()},'
+        f' {self._second_reward_provider.name()})'
+    )
+
+  def compute_reward(
+      self, required_features: Mapping[str, gdmr_types.ArrayType]
+  ) -> RewardValue:
+    first_required_features = {
+        k: v
+        for k, v in required_features.items()
+        if k in self._first_required_features_keys
+    }
+    second_required_features = {
+        k: v
+        for k, v in required_features.items()
+        if k in self._second_required_features_keys
+    }
+    return tree.map_structure(
+        self._op,
+        self._first_reward_provider.compute_reward(first_required_features),
+        self._second_reward_provider.compute_reward(second_required_features),
+    )
+
+  def reward_spec(self) -> RewardSpec:
+    return self._reward_spec
+
+  def required_features_keys(self) -> set[str]:
+    return (
+        self._first_required_features_keys | self._second_required_features_keys
+    )
+
+  def reset(self) -> None:
+    self._first_reward_provider.reset()
+    self._second_reward_provider.reset()
+
+
+class GetItemOperationRewardProvider(RewardProvider):
+  """Extracts a slice from the result of a reward provider."""
+
+  def __init__(self, reward_provider: RewardProviderOrValue, index: slice):
+    super().__init__()
+    if not isinstance(reward_provider, RewardProvider):
+      reward_provider = ConstantRewardProvider(reward_provider)
+    self._reward_provider = reward_provider
+    self._index = index
+
+  def name(self) -> str:
+    return f'{self._reward_provider.name}[{self._index}]'
+
+  def compute_reward(
+      self, required_features: Mapping[str, gdmr_types.ArrayType]
+  ) -> RewardValue:
+    return tree.map_structure(
+        lambda v: v[self._index],
+        self._reward_provider.compute_reward(required_features),
+    )
+
+  def reward_spec(self) -> RewardSpec:
+    return tree.map_structure(
+        lambda s: specs.Array(np.empty(s.shape)[self._index].shape, s.dtype),
+        self._reward_provider.reward_spec(),
+    )
+
+  def required_features_keys(self) -> set[str]:
+    return self._reward_provider.required_features_keys()
+
+  def reset(self) -> None:
+    self._reward_provider.reset()
+
+
+class UnaryOperationRewardProvider(RewardProvider):
+  """Applies a unary operator to the result of a reward provider."""
+
+  def __init__(self, op: UnaryOperator, reward_provider: RewardProviderOrValue):
+    super().__init__()
+    if not isinstance(reward_provider, RewardProvider):
+      reward_provider = ConstantRewardProvider(reward_provider)
+    self._op = op
+    self._reward_provider = reward_provider
+
+  def name(self) -> str:
+    op_name = getattr(self._op, '__name__', str(self._op))
+    return f'{op_name}({self._reward_provider.name()})'
+
+  def compute_reward(
+      self, required_features: Mapping[str, gdmr_types.ArrayType]
+  ) -> RewardValue:
+    return tree.map_structure(
+        self._op, self._reward_provider.compute_reward(required_features)
+    )
+
+  def reward_spec(self) -> RewardSpec:
+    return self._reward_provider.reward_spec()
+
+  def required_features_keys(self) -> set[str]:
+    return self._reward_provider.required_features_keys()
+
+  def reset(self) -> None:
+    self._reward_provider.reset()

src/experimental/reaf/core/substep_commands_processor.py

@@ -0,0 +1,104 @@
+# Copyright 2025 Google LLC
+#
+# 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
+#
+# https://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.
+"""Protocol for substep commands manipulation in REAF-sim."""
+
+from collections.abc import Mapping
+import typing
+
+from dm_env import specs
+from gdm_robotics.interfaces import types as gdmr_types
+
+
+class SubstepCommandsProcessor(typing.Protocol):
+  """Processes substep commands, propagating them through a pipeline.
+
+  This processor manipulates substep commands, acting as a node in a pipeline.
+  It consumes substep commands, performs operations, and produces updated
+  substep commands for the next stage in the processing chain.
+
+  The processing pipeline starts with commands provided to the SimulationDevice
+  and progresses towards the substep commands consumed by the individual
+  entities. Each processor consumes a subset of substep commands and produces
+  new, potentially transformed, substep commands. The order of operations is
+  crucial.
+
+  Example Pipeline (conceptual):
+
+  Simulation Device commands --> Processor (1) --> Processor (2) --> Entities
+
+  Specs are propagated starting from the bottom:
+  1) In this example assume that the set of entities expect "p3/c1", "p3/c2" and
+    "p3/c3".
+  2) Processor (2) returns ("p3/c1", "p3/c2") from "p2/c1". This means
+    that the global substep commands spec exposed at this level is "p2/c1" and
+    the unprocessed "p3/c3".
+  3) Processor (1) returns "p2/c1" from ("p1/c1", "p1/c2"). By applying
+    the same transformation rule, we can obtain the final spec exposed
+    by the SimulationDevice: "p1/c1", "p1/c2" and "p3/c3".
+
+     ------------------------------------
+    |        SimulationDevice           |
+     ------------------------------------
+
+    "p1/c1"     "p1/c2"           "p3/c3"
+       |          |                  |
+     -----------------               |
+    |      P1         |              |
+     -----------------               |
+            |  "p2/c1"               |
+     -----------------               |
+    |      P2         |              |
+     -----------------               |
+       | "p3/c1"   | "p3/c2"         |
+       |           |                 |
+     ------------------------------------
+    |            Entities               |
+     ------------------------------------
+  """
+
+  @property
+  def name(self) -> str:
+    """Returns a unique string identifier for this object."""
+
+  def reset(self) -> None:
+    """Resets the internal state of this processor."""
+
+  def produced_substep_commands_keys(self) -> set[str]:
+    """Keys of the substep commands produced by this processor."""
+
+  def consumed_substep_commands_spec(
+      self,
+  ) -> Mapping[str, specs.Array]:
+    """Spec of the substep commands consumed by this processor."""
+
+  def process_substep_commands(
+      self,
+      model: typing.Any,
+      data: typing.Any,
+      consumed_substep_commands: Mapping[str, gdmr_types.ArrayType],
+  ) -> Mapping[str, gdmr_types.ArrayType]:
+    """Processes the substep commands and returns a new modified version of it.
+
+    Args:
+      model: the simulation model.
+      data: the simulation data.
+      consumed_substep_commands: the substep commands up in the processing chain
+        that are required by this processor, i.e. with keys specified by
+        `consumed_substep_commands_spec`.
+
+    Returns the new substep commands. Note that the (key, value) pairs in
+    `consumed_substep_commands` are removed from the running substep commands
+    dictionary. If users want to keep some of the elements it is their
+    responsibility to retain them in the output dictionary.
+    """

src/experimental/reaf/core/substep_measurements_processor.py

@@ -0,0 +1,103 @@
+# Copyright 2025 Google LLC
+#
+# 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
+#
+# https://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.
+"""Protocol for substep measurements manipulation in REAF-sim."""
+
+from collections.abc import Mapping
+import typing
+
+from dm_env import specs
+from gdm_robotics.interfaces import types as gdmr_types
+
+
+class SubstepMeasurementsProcessor(typing.Protocol):
+  """Processes substep measurements, propagating them through a pipeline.
+
+  This processor manipulates substep measurements, acting as a node in a
+  pipeline. It consumes substep measurements, performs operations, and produces
+  updated substep measurements for the next stage in the processing chain.
+
+  The processing pipeline starts with substep measurements produced by Entities
+  and progresses towards the measurements exposed by the SimulationDevice. Each
+  processor consumes a subset of substep measurements and produces new,
+  potentially transformed, substep measurements. The order of operations is
+  crucial.
+
+  Example Pipeline (conceptual):
+
+  Entities --> Processor (1) --> Processor (2) -> Simulation Device Measurements
+
+  Specs are propagated starting from the bottom:
+  1) In this example assume that the set of entities produce "p1/c1", "p1/c2"
+     and "p1/c3".
+  2) Processor (1) returns "p2/c1" from ("p1/c1", "p1/c2").
+  3) Processor (2) returns ("p3/c1", "p3/c2") from "p2/c1".
+
+  This resulting spec exposed by the SimulationDevice: "p3/c1", "p3/c2"
+  and "p1/c3".
+
+     ------------------------------------
+    |        SimulationDevice           |
+     ------------------------------------
+
+    "p3/c1"     "p3/c2"           "p1/c3"
+       |          |                  |
+     -----------------               |
+    |      P2         |              |
+     -----------------               |
+            |  "p2/c1"               |
+     -----------------               |
+    |      P1         |              |
+     -----------------               |
+       | "p1/c1"   | "p1/c2"         |
+       |           |                 |
+     ------------------------------------
+    |            Entities               |
+     ------------------------------------
+  """
+
+  @property
+  def name(self) -> str:
+    """Returns a unique string identifier for this object."""
+
+  def reset(self):
+    """Resets the internal state of this processor."""
+
+  def produced_substep_measurements_spec(
+      self,
+  ) -> Mapping[str, specs.Array]:
+    """Spec of the substep measurements consumed by this processor."""
+
+  def consumed_substep_measurements_keys(self) -> set[str]:
+    """Keys of the substep measurements consumed by this processor."""
+
+  def process_substep_measurements(
+      self,
+      model: typing.Any,
+      data: typing.Any,
+      consumed_substep_measurements: Mapping[str, gdmr_types.ArrayType],
+  ) -> Mapping[str, gdmr_types.ArrayType]:
+    """Processes the substep measurements and returns a new modified version of it.
+
+    Args:
+      model: the simulation model.
+      data: the simulation data.
+      consumed_substep_measurements: the substep measurements up in the
+        processing chain that are required by this processor, i.e. with keys
+        specified by `consumed_substep_measurements_spec`.
+
+    Returns the new substep measurements. Note that the (key, value) pairs in
+    `consumed_substep_measurements` are removed from the running substep
+    measurements dictionary. If users want to keep some of the elements it is
+    their responsibility to retain them in the output dictionary.
+    """

src/experimental/reaf/core/task_logic_layer.py

@@ -0,0 +1,342 @@
+# Copyright 2025 Google LLC
+#
+# 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
+#
+# https://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.
+"""Task logic layer for the Robotics Environment Authoring Framework."""
+
+from collections.abc import Mapping, Sequence
+import itertools
+from typing import Protocol
+
+from absl import logging
+from dm_env import specs
+from gdm_robotics.interfaces import types as gdmr_types
+from reaf.core import commands_processor as reaf_commands_processor
+from reaf.core import default_discount_provider
+from reaf.core import discount_provider as reaf_discount_provider
+from reaf.core import features_observer as reaf_features_observers
+from reaf.core import features_producer as reaf_features_producer
+from reaf.core import logger as reaf_logger
+from reaf.core import reward_provider as reaf_reward_provider
+from reaf.core import termination_checker as reaf_termination_checker
+from reaf.core import zero_reward_provider
+import tree
+
+
+class _ResettableObject(Protocol):
+  """Protocol for an object that can be reset."""
+
+  def reset(self) -> None:
+    ...
+
+
+class TaskLogicLayer:
+  """Task logic layer for the Robotics Environment Authoring Framework."""
+
+  def __init__(
+      self,
+      *,
+      commands_processors: Sequence[reaf_commands_processor.CommandsProcessor],
+      features_producers: Sequence[reaf_features_producer.FeaturesProducer],
+      termination_checkers: Sequence[
+          reaf_termination_checker.TerminationChecker
+      ],
+      reward_provider: reaf_reward_provider.RewardProvider | None = None,
+      discount_provider: reaf_discount_provider.DiscountProvider | None = None,
+      features_observers: Sequence[
+          reaf_features_observers.FeaturesObserver
+      ] = (),
+      loggers: Sequence[reaf_logger.Logger] = (),
+  ):
+    """Initializes the task logic layer.
+
+    Args:
+      commands_processors: `CommandsProcessor`s that modify the commands before
+        being sent down to the DACL. They are called sequentially, starting from
+        the commands supplied by the policy and ending with the commands that
+        will be sent to the DACL.
+      features_producers: `FeaturesProducer`s that generate new features.
+        Measurements collected by the DACL and features produced by these
+        `FeaturesProducer`s are then merged into the final feature set that is
+        provided to the `reward_provider`, `termination_checkers`,
+        `discount_provider`, `features_observers`, and `loggers`.
+      termination_checkers: `TerminationChecker`s that check the episode
+        termination based on the final feature set.
+      reward_provider: `RewardProvider` that computes a reward based on the
+        final feature set. If None, the ZeroRewardProvider is used and the
+        reward is set to 0.
+      discount_provider: `DiscountProvider` that compute a discount based on the
+        final feature set and final termination state. If None, the
+        DefaultDiscountProvider is used returning 0 for termination and 1 for
+        truncation and non-termination.
+      features_observers: `FeaturesObserver`s that get a view over the final
+        feature set.
+      loggers: `Logger`s for logging measurements, features, and commands in the
+        task layer.
+    """
+    self._commands_processors = commands_processors
+    self._features_producers = features_producers
+    self._reward_provider = (
+        reward_provider
+        if reward_provider
+        else zero_reward_provider.ZeroRewardProvider()
+    )
+    self._termination_checkers = termination_checkers
+    self._discount_provider = (
+        discount_provider
+        if discount_provider
+        else default_discount_provider.DefaultDiscountProvider()
+    )
+    self._features_observers = features_observers
+    self._loggers = list(loggers)
+
+    # We make a set of all resettable objects so that these objects only get
+    # their resets called once. This is important for e.g. when having a single
+    # object that derives from two interfaces.
+    self._resettable_objects: list[_ResettableObject] = []
+    unique_ids = set()
+    for resettable_object in itertools.chain(
+        self._commands_processors,
+        self._features_producers,
+        self._termination_checkers,
+        [self._reward_provider],
+        [self._discount_provider],
+    ):
+      resettable_object_id = id(resettable_object)
+      if resettable_object_id not in unique_ids:
+        unique_ids.add(resettable_object_id)
+        self._resettable_objects.append(resettable_object)
+
+  def validate_spec(
+      self,
+      *,
+      dacl_commands_spec: Mapping[str, gdmr_types.AnyArraySpec],
+      dacl_measurements_spec: Mapping[str, specs.Array],
+  ) -> None:
+    """Checks that the specs have consistent keys."""
+    logging.vlog(3, "Validate features processing")
+    self._validate_features_spec(dacl_measurements_spec)
+    self._validate_commands_spec(dacl_commands_spec)
+
+  def features_spec(
+      self,
+      dacl_measurements_spec: Mapping[str, specs.Array],
+  ) -> Mapping[str, specs.Array]:
+    """Returns the features spec as exposed by the task layer."""
+    spec = dict(dacl_measurements_spec)
+    for features_producer in self._features_producers:
+      spec.update(features_producer.produced_features_spec())
+
+    return spec
+
+  def commands_spec(
+      self, dacl_commands_spec: Mapping[str, gdmr_types.AnyArraySpec]
+  ) -> Mapping[str, gdmr_types.AnyArraySpec]:
+    """Returns the commands spec exposed by the task layer."""
+    # Each processor consumes commands (as described by its
+    # `consumed_commands_spec`) and outputs a potentially different set of
+    # commands (as described by its `produced_commands_keys`).
+    # Starting with the DACL command spec, we iterate in reverse order (i.e. in
+    # the direction DACL -> Policy) through every processor to remove the
+    # `produced_commands_keys` from the spec, and add their
+    # `consumed_commands_spec` to the spec.
+    spec: Mapping[str, gdmr_types.AnyArraySpec] = dict(dacl_commands_spec)
+    for processor in reversed(self._commands_processors):
+      processor_produced_keys = processor.produced_commands_keys()
+      spec = {
+          key: value
+          for key, value in spec.items()
+          if key not in processor_produced_keys
+      }
+      spec.update(processor.consumed_commands_spec())
+    return spec
+
+  def reward_spec(self) -> tree.Structure[specs.Array]:
+    return self._reward_provider.reward_spec()
+
+  def discount_spec(self) -> tree.Structure[specs.Array]:
+    return self._discount_provider.discount_spec()
+
+  def perform_reset(self) -> None:
+    """Reset the internal state of the task logic layer."""
+    for resettable_object in self._resettable_objects:
+      resettable_object.reset()
+
+  def compute_all_features(
+      self, measurements: Mapping[str, gdmr_types.ArrayType]
+  ) -> Mapping[str, gdmr_types.ArrayType]:
+    """Computes all the task logic features given the current measurements."""
+    for logger in self._loggers:
+      logger.record_measurements(measurements)
+
+    # Produce all the features.
+    current_features = dict(measurements)
+    for feature_producer in self._features_producers:
+      required_features = {
+          key: current_features[key]
+          for key in feature_producer.required_features_keys()
+      }
+      current_features.update(
+          feature_producer.produce_features(required_features)
+      )
+
+    # Observe the features.
+    for feature_observer in self._features_observers:
+      feature_observer.observe_features(current_features)
+
+    # Log the resulting features.
+    for logger in self._loggers:
+      logger.record_features(current_features)
+    return current_features
+
+  def compute_final_commands(
+      self,
+      policy_commands: Mapping[str, gdmr_types.ArrayType],
+  ) -> Mapping[str, gdmr_types.ArrayType]:
+    """Processes the policy commands and returns the final processed commands."""
+    current_commands = dict(policy_commands)
+    for processor in self._commands_processors:
+      # Get commands to be consumed by the processor and remove the commands
+      # from the current_commands.. They correspond to the
+      # `consumed_command_spec`.
+      consumed_commands = {
+          key: current_commands.pop(key)
+          for key in processor.consumed_commands_spec().keys()
+      }
+      produced_commands = processor.process_commands(consumed_commands)
+      current_commands.update(produced_commands)
+
+      # Log the modification.
+      for logger in self._loggers:
+        logger.record_commands_processing(
+            processor.name, consumed_commands, produced_commands
+        )
+
+    for logger in self._loggers:
+      logger.record_final_commands(current_commands)
+    return current_commands
+
+  def compute_reward(
+      self, features: Mapping[str, gdmr_types.ArrayType]
+  ) -> tree.Structure[gdmr_types.ArrayType]:
+    """Computes the reward given the features."""
+    return self._reward_provider.compute_reward({
+        key: features[key]
+        for key in self._reward_provider.required_features_keys()
+    })
+
+  def check_for_termination(
+      self, features: Mapping[str, gdmr_types.ArrayType]
+  ) -> reaf_termination_checker.TerminationResult:
+    """Checks for termination."""
+    current_state = reaf_termination_checker.TerminationResult.DO_NOT_TERMINATE
+    for termination_checker in self._termination_checkers:
+      current_state = reaf_termination_checker.TerminationResult.combine(
+          current_state,
+          termination_checker.check_termination({
+              key: features[key]
+              for key in termination_checker.required_features_keys()
+          }),
+      )
+    return current_state
+
+  def compute_discount(
+      self,
+      features: Mapping[str, gdmr_types.ArrayType],
+      termination_state: reaf_termination_checker.TerminationResult,
+  ) -> tree.Structure[gdmr_types.ArrayType]:
+    """Computes the discount given the features and termination state."""
+    return self._discount_provider.compute_discount(
+        {
+            key: features[key]
+            for key in self._discount_provider.required_features_keys()
+        },
+        termination_state,
+    )
+
+  def add_logger(self, logger: reaf_logger.Logger) -> None:
+    self._loggers.append(logger)
+
+  def remove_logger(self, logger: reaf_logger.Logger) -> None:
+    self._loggers.remove(logger)
+
+  def _validate_features_spec(
+      self, dacl_measurements_spec: Mapping[str, specs.Array]
+  ) -> None:
+    """Validates the features spec."""
+    # Check measurements/features path.
+    current_key_set = set(dacl_measurements_spec.keys())
+    logging.vlog(4, "DACL measurements keys: %s", current_key_set)
+
+    for producer in self._features_producers:
+      logging.vlog(
+          4,
+          "Producer %s requires %s.",
+          producer.name,
+          producer.required_features_keys(),
+      )
+      # Check required features are available.
+      if not producer.required_features_keys().issubset(current_key_set):
+        raise ValueError(
+            "Failed to validate feature specs for feature producer"
+            f" {producer.name}. Missing keys:"
+            f" {producer.required_features_keys() - current_key_set}"
+        )
+      # Check that there are not duplicates in the output.
+      if not current_key_set.isdisjoint(
+          producer.produced_features_spec().keys()
+      ):
+        raise ValueError(
+            "Failed to validate feature specs for feature producer"
+            f" {producer.name}. Duplicate keys:"
+            f" {current_key_set & producer.produced_features_spec().keys()}"
+        )
+      # Now extend the spec.
+      logging.vlog(
+          4,
+          "Update available keys (from producer %s) with %s.",
+          producer.name,
+          producer.produced_features_spec().keys(),
+      )
+      current_key_set.update(producer.produced_features_spec().keys())
+      logging.vlog(4, "Available features keys %s.", current_key_set)
+
+  def _validate_commands_spec(
+      self, dacl_commands_spec: Mapping[str, gdmr_types.AnyArraySpec]
+  ) -> None:
+    """Validates the commands spec."""
+    # Check commands. Starting from the DACL command specs we propagate up in
+    # the chain.
+    logging.vlog(3, "Validate commands processing from DACL to Policy.")
+    current_key_set = set(dacl_commands_spec.keys())
+    logging.vlog(4, "DACL commands keys: %s", current_key_set)
+
+    for processor in reversed(self._commands_processors):
+      produced_command_keys = processor.produced_commands_keys()
+
+      logging.vlog(
+          4,
+          "Processor %s: specs (accepted keys) %s. Exposes %s.",
+          processor.name,
+          processor.consumed_commands_spec().keys(),
+          produced_command_keys,
+      )
+      if not produced_command_keys.issubset(current_key_set):
+        raise ValueError(
+            "Failed to validate commands specs for commands processor"
+            f" {processor.name}. Missing (consumable) keys:"
+            f" {produced_command_keys - current_key_set}"
+        )
+      # Remove the produced keys and add the consumed commands specs (as the
+      # processor is mutable).
+      current_key_set = current_key_set - produced_command_keys
+      current_key_set.update(processor.consumed_commands_spec().keys())

src/experimental/reaf/core/termination_checker.py

@@ -0,0 +1,94 @@
+# Copyright 2025 Google LLC
+#
+# 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
+#
+# https://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.
+"""Checks if the episode should terminate."""
+
+import abc
+from collections.abc import Mapping
+import enum
+from typing import Self
+
+from gdm_robotics.interfaces import types as gdmr_types
+
+
+class TerminationResult(enum.IntFlag):
+  """The result of an episode termination check.
+
+  The TerminationResult refers to the possibility for an episode to terminate.
+  For more details on the concept of termination we refer the readers to
+  https://github.com/google-deepmind/dm_env/blob/master/docs/index.md#environment-api-and-semantics.
+
+  Note that this enum does not refer to the possible causes of termination but
+  only how the termination impacts the learning process.
+
+  The result can be one of the following options:
+  - DO_NOT_TERMINATE: The episode should not terminate.
+  - TRUNCATE: The epsisode should terminate. Truncation implies a non-failure
+    final state. Usually this is associated with a non-zero discount.
+  - TERMINATE: The episode should terminate as the environment is in some
+    final state. Usually this is associated with a zero discount for e.g.
+    finite-horizon RL.
+  """
+
+  DO_NOT_TERMINATE = 0
+  TRUNCATE = 2**0
+  TERMINATE = 2**1
+
+  def is_terminated(self) -> bool:
+    return self == TerminationResult.TERMINATE
+
+  def is_truncated(self) -> bool:
+    return self == TerminationResult.TRUNCATE
+
+  def combine(self, other: Self) -> Self:
+    # TERMINATE has precedence over TRUNCATE, which in turn has precedence over
+    # DO_NOT_TERMINATE. Given the definitions above, this can be implemented as
+    # a maximum operator. To also enable tracing with JAX, we implement this in
+    # a branchless manner using bitwise operations that preserve the type.
+    # Note that JAX will trace TerminationResult values as ints.
+    # Approach:
+    # - self ^ (self ^ other) == other
+    # - (-1 * (self < other)) will be bitmask of all 1s iff self < other.
+    # - AND with (self ^ other) will result in either update or no-op bitmask.
+    return self ^ ((self ^ other) & (-1 * (self < other)))
+
+
+class TerminationChecker(abc.ABC):
+  """Checks if the episode should terminate."""
+
+  @abc.abstractmethod
+  def name(self) -> str:
+    """Returns a unique string identifier for this object."""
+
+  @abc.abstractmethod
+  def check_termination(
+      self, required_features: Mapping[str, gdmr_types.ArrayType]
+  ) -> TerminationResult:
+    """Checks if the episode should terminate.
+
+    Args:
+      required_features: Measurements and features computed by the task logic
+        that are required by this checker, i.e. that have keys specified by
+        `required_features_keys`.
+
+    Returns if the episode should terminate (and if so, what kind of
+      termination).
+    """
+
+  @abc.abstractmethod
+  def required_features_keys(self) -> set[str]:
+    """Returns the feature keys that are required to check the termination."""
+
+  def reset(self) -> None:
+    """Resets the internal state of the termination checker."""
+    ...

src/experimental/reaf/core/trigger.py

@@ -0,0 +1,29 @@
+# Copyright 2025 Google LLC
+#
+# 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
+#
+# https://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.
+"""Defines an event-based waiting behaviour."""
+
+import abc
+
+
+class Trigger(abc.ABC):
+  """Defines an event-based waiting behaviour."""
+
+  @property
+  @abc.abstractmethod
+  def name(self) -> str:
+    """Returns the name of the trigger."""
+
+  @abc.abstractmethod
+  def wait_for_event(self) -> None:
+    """Blocks until the next event."""

src/experimental/reaf/core/zero_reward_provider.py

@@ -0,0 +1,48 @@
+# Copyright 2025 Google LLC
+#
+# 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
+#
+# https://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.
+"""Reward provider which provides a zero reward."""
+
+from collections.abc import Mapping
+from dm_env import specs
+from gdm_robotics.interfaces import types as gdmr_types
+import numpy as np
+from reaf.core import reward_provider
+import tree
+
+
+class ZeroRewardProvider(reward_provider.RewardProvider):
+  """Reward provider which provides a zero reward."""
+
+  def __init__(self, name: str = 'zero_reward_provider'):
+    self._name = name
+
+  def name(self) -> str:
+    return self._name
+
+  def compute_reward(
+      self, required_features: Mapping[str, gdmr_types.ArrayType]
+  ) -> tree.Structure[gdmr_types.ArrayType]:
+    """Returns a zero reward."""
+    return np.zeros(1)
+
+  def reward_spec(self) -> tree.Structure[specs.Array]:
+    """Returns the spec for a constant zero reward."""
+    return specs.Array(shape=(1,), dtype=float)
+
+  def required_features_keys(self) -> set[str]:
+    """Returns empty set.
+
+    There are no feature keys that are required to compute the reward.
+    """
+    return set()