diff --git a/src/experimental/reaf/core/action_space_adapter.py b/src/experimental/reaf/core/action_space_adapter.py deleted file mode 100644 index 014e02eb..00000000 --- a/src/experimental/reaf/core/action_space_adapter.py +++ /dev/null @@ -1,43 +0,0 @@ -# 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.""" diff --git a/src/experimental/reaf/core/commands_processor.py b/src/experimental/reaf/core/commands_processor.py deleted file mode 100644 index 4cc12b4a..00000000 --- a/src/experimental/reaf/core/commands_processor.py +++ /dev/null @@ -1,91 +0,0 @@ -# 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.""" - ... diff --git a/src/experimental/reaf/core/data_acquisition_and_control_layer.py b/src/experimental/reaf/core/data_acquisition_and_control_layer.py deleted file mode 100644 index e8fcc8ca..00000000 --- a/src/experimental/reaf/core/data_acquisition_and_control_layer.py +++ /dev/null @@ -1,170 +0,0 @@ -# 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) diff --git a/src/experimental/reaf/core/default_discount_provider.py b/src/experimental/reaf/core/default_discount_provider.py deleted file mode 100644 index ff216b7c..00000000 --- a/src/experimental/reaf/core/default_discount_provider.py +++ /dev/null @@ -1,79 +0,0 @@ -# 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() diff --git a/src/experimental/reaf/core/default_observation_space_adapter.py b/src/experimental/reaf/core/default_observation_space_adapter.py deleted file mode 100644 index 1261b3ca..00000000 --- a/src/experimental/reaf/core/default_observation_space_adapter.py +++ /dev/null @@ -1,231 +0,0 @@ -# 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 - ) diff --git a/src/experimental/reaf/core/device.py b/src/experimental/reaf/core/device.py deleted file mode 100644 index cc53a0ca..00000000 --- a/src/experimental/reaf/core/device.py +++ /dev/null @@ -1,54 +0,0 @@ -# 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.""" diff --git a/src/experimental/reaf/core/device_coordinator.py b/src/experimental/reaf/core/device_coordinator.py deleted file mode 100644 index 233705bb..00000000 --- a/src/experimental/reaf/core/device_coordinator.py +++ /dev/null @@ -1,87 +0,0 @@ -# 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. - """ diff --git a/src/experimental/reaf/core/discount_provider.py b/src/experimental/reaf/core/discount_provider.py deleted file mode 100644 index df1b22fa..00000000 --- a/src/experimental/reaf/core/discount_provider.py +++ /dev/null @@ -1,61 +0,0 @@ -# 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.""" - ... diff --git a/src/experimental/reaf/core/entity.py b/src/experimental/reaf/core/entity.py deleted file mode 100644 index 68c2e38d..00000000 --- a/src/experimental/reaf/core/entity.py +++ /dev/null @@ -1,65 +0,0 @@ -# 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.""" diff --git a/src/experimental/reaf/core/environment.py b/src/experimental/reaf/core/environment.py deleted file mode 100644 index dca306f0..00000000 --- a/src/experimental/reaf/core/environment.py +++ /dev/null @@ -1,490 +0,0 @@ -# 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 diff --git a/src/experimental/reaf/core/features_observer.py b/src/experimental/reaf/core/features_observer.py deleted file mode 100644 index cc87f0d3..00000000 --- a/src/experimental/reaf/core/features_observer.py +++ /dev/null @@ -1,34 +0,0 @@ -# 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.""" diff --git a/src/experimental/reaf/core/features_producer.py b/src/experimental/reaf/core/features_producer.py deleted file mode 100644 index 8ce44e94..00000000 --- a/src/experimental/reaf/core/features_producer.py +++ /dev/null @@ -1,56 +0,0 @@ -# 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.""" - ... diff --git a/src/experimental/reaf/core/logger.py b/src/experimental/reaf/core/logger.py deleted file mode 100644 index 63bb5f33..00000000 --- a/src/experimental/reaf/core/logger.py +++ /dev/null @@ -1,80 +0,0 @@ -# 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`. - """ diff --git a/src/experimental/reaf/core/numpy_mock_assertions.py b/src/experimental/reaf/core/numpy_mock_assertions.py deleted file mode 100644 index 310a918e..00000000 --- a/src/experimental/reaf/core/numpy_mock_assertions.py +++ /dev/null @@ -1,98 +0,0 @@ -# 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}" - ) diff --git a/src/experimental/reaf/core/observation_space_adapter.py b/src/experimental/reaf/core/observation_space_adapter.py deleted file mode 100644 index ca4a8d86..00000000 --- a/src/experimental/reaf/core/observation_space_adapter.py +++ /dev/null @@ -1,42 +0,0 @@ -# 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.""" diff --git a/src/experimental/reaf/core/pass_through_action_space_adapter.py b/src/experimental/reaf/core/pass_through_action_space_adapter.py deleted file mode 100644 index 16b28a81..00000000 --- a/src/experimental/reaf/core/pass_through_action_space_adapter.py +++ /dev/null @@ -1,55 +0,0 @@ -# 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()) diff --git a/src/experimental/reaf/core/reward_provider.py b/src/experimental/reaf/core/reward_provider.py deleted file mode 100644 index 3a8ec655..00000000 --- a/src/experimental/reaf/core/reward_provider.py +++ /dev/null @@ -1,292 +0,0 @@ -# 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() diff --git a/src/experimental/reaf/core/substep_commands_processor.py b/src/experimental/reaf/core/substep_commands_processor.py deleted file mode 100644 index b63f4388..00000000 --- a/src/experimental/reaf/core/substep_commands_processor.py +++ /dev/null @@ -1,104 +0,0 @@ -# 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. - """ diff --git a/src/experimental/reaf/core/substep_measurements_processor.py b/src/experimental/reaf/core/substep_measurements_processor.py deleted file mode 100644 index 15dc81df..00000000 --- a/src/experimental/reaf/core/substep_measurements_processor.py +++ /dev/null @@ -1,103 +0,0 @@ -# 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. - """ diff --git a/src/experimental/reaf/core/task_logic_layer.py b/src/experimental/reaf/core/task_logic_layer.py deleted file mode 100644 index 51b25639..00000000 --- a/src/experimental/reaf/core/task_logic_layer.py +++ /dev/null @@ -1,342 +0,0 @@ -# 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()) diff --git a/src/experimental/reaf/core/termination_checker.py b/src/experimental/reaf/core/termination_checker.py deleted file mode 100644 index 754c250d..00000000 --- a/src/experimental/reaf/core/termination_checker.py +++ /dev/null @@ -1,94 +0,0 @@ -# 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.""" - ... diff --git a/src/experimental/reaf/core/trigger.py b/src/experimental/reaf/core/trigger.py deleted file mode 100644 index 20901873..00000000 --- a/src/experimental/reaf/core/trigger.py +++ /dev/null @@ -1,29 +0,0 @@ -# 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.""" diff --git a/src/experimental/reaf/core/zero_reward_provider.py b/src/experimental/reaf/core/zero_reward_provider.py deleted file mode 100644 index c5d2267a..00000000 --- a/src/experimental/reaf/core/zero_reward_provider.py +++ /dev/null @@ -1,48 +0,0 @@ -# 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()