From f55aeb04fa308b66c4dec5c134d191e84b28c10f Mon Sep 17 00:00:00 2001 From: Tom Erez Date: Mon, 13 Apr 2026 08:20:32 -0700 Subject: [PATCH] Add protocols for authoring simulation environments. This is for preview only, we discourage users from using this in production code. PiperOrigin-RevId: 899009633 Change-Id: I2de08c5cc4ece69fa135f3fd02df7542bf66d7d0 --- .../reaf/core/action_space_adapter.py | 43 ++ .../reaf/core/commands_processor.py | 91 ++++ .../data_acquisition_and_control_layer.py | 170 ++++++ .../reaf/core/default_discount_provider.py | 79 +++ .../core/default_observation_space_adapter.py | 231 +++++++++ src/experimental/reaf/core/device.py | 54 ++ .../reaf/core/device_coordinator.py | 87 ++++ .../reaf/core/discount_provider.py | 61 +++ src/experimental/reaf/core/entity.py | 65 +++ src/experimental/reaf/core/environment.py | 490 ++++++++++++++++++ .../reaf/core/features_observer.py | 34 ++ .../reaf/core/features_producer.py | 56 ++ src/experimental/reaf/core/logger.py | 80 +++ .../reaf/core/numpy_mock_assertions.py | 98 ++++ .../reaf/core/observation_space_adapter.py | 42 ++ .../core/pass_through_action_space_adapter.py | 55 ++ src/experimental/reaf/core/reward_provider.py | 292 +++++++++++ .../reaf/core/substep_commands_processor.py | 104 ++++ .../core/substep_measurements_processor.py | 103 ++++ .../reaf/core/task_logic_layer.py | 342 ++++++++++++ .../reaf/core/termination_checker.py | 94 ++++ src/experimental/reaf/core/trigger.py | 29 ++ .../reaf/core/zero_reward_provider.py | 48 ++ 23 files changed, 2748 insertions(+) create mode 100644 src/experimental/reaf/core/action_space_adapter.py create mode 100644 src/experimental/reaf/core/commands_processor.py create mode 100644 src/experimental/reaf/core/data_acquisition_and_control_layer.py create mode 100644 src/experimental/reaf/core/default_discount_provider.py create mode 100644 src/experimental/reaf/core/default_observation_space_adapter.py create mode 100644 src/experimental/reaf/core/device.py create mode 100644 src/experimental/reaf/core/device_coordinator.py create mode 100644 src/experimental/reaf/core/discount_provider.py create mode 100644 src/experimental/reaf/core/entity.py create mode 100644 src/experimental/reaf/core/environment.py create mode 100644 src/experimental/reaf/core/features_observer.py create mode 100644 src/experimental/reaf/core/features_producer.py create mode 100644 src/experimental/reaf/core/logger.py create mode 100644 src/experimental/reaf/core/numpy_mock_assertions.py create mode 100644 src/experimental/reaf/core/observation_space_adapter.py create mode 100644 src/experimental/reaf/core/pass_through_action_space_adapter.py create mode 100644 src/experimental/reaf/core/reward_provider.py create mode 100644 src/experimental/reaf/core/substep_commands_processor.py create mode 100644 src/experimental/reaf/core/substep_measurements_processor.py create mode 100644 src/experimental/reaf/core/task_logic_layer.py create mode 100644 src/experimental/reaf/core/termination_checker.py create mode 100644 src/experimental/reaf/core/trigger.py create mode 100644 src/experimental/reaf/core/zero_reward_provider.py diff --git a/src/experimental/reaf/core/action_space_adapter.py b/src/experimental/reaf/core/action_space_adapter.py new file mode 100644 index 00000000..014e02eb --- /dev/null +++ b/src/experimental/reaf/core/action_space_adapter.py @@ -0,0 +1,43 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Adapts environment action into suitable commands format accepted by REAF.""" + +import abc +from collections.abc import Mapping + +from gdm_robotics.interfaces import types as gdmr_types + + +class ActionSpaceAdapter(abc.ABC): + """Adapts environment action into suitable commands format accepted by REAF. + + Implementations of this interface are responsible for converting the more + generic action accepted by the environment (e.g. a flat numpy array) into the + more constraining format accepted as commands by REAF, i.e. a dictionary of + string to tensors. + """ + + @abc.abstractmethod + def commands_from_environment_action( + self, environment_action: gdmr_types.ActionType + ) -> Mapping[str, gdmr_types.ArrayType]: + """Converts the environment action into commands accepted by REAF.""" + + @abc.abstractmethod + def action_spec(self) -> gdmr_types.ActionSpec: + """Returns the action spec exposed by the environment.""" + + @abc.abstractmethod + def task_commands_keys(self) -> set[str]: + """Returns the keys for the commands exposed to the task layer.""" diff --git a/src/experimental/reaf/core/commands_processor.py b/src/experimental/reaf/core/commands_processor.py new file mode 100644 index 00000000..4cc12b4a --- /dev/null +++ b/src/experimental/reaf/core/commands_processor.py @@ -0,0 +1,91 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Abstract class for commands manipulation in the task logic layer.""" + +import abc +from collections.abc import Mapping + +from gdm_robotics.interfaces import types as gdmr_types + + +class CommandsProcessor(abc.ABC): + """Perform commands manipulation. + + The following describes the processing pipeline starting from the top (closer + to the policy) to the bottom (interfacing with the DACL commands spec). + + Assume that we have two processing units: + Processor 1) has a consumed_commands_spec for two keys: "p1/c1" and "p1/c2". + Its produced_commands_keys are "p2/c1". + Processor 2) has a consumed_commands_spec for "p2/c1". Its + produced_commands_keys are "p3/c1" and "p3/c2". + + Specs are propagated starting from the bottom: + 1) In this example assume that the DACL exposes "p3/c1", "p3/c2" and "p3/c3". + 2) Processor 2) returns ("p3/c1", "p3/c2") from input "p2/c1". This means that + the global commands spec exposed at this level is "p2/c1" and the + unprocessed "p3/c3". + 3) Processor 1) returns "p2/c1" from input ("p1/c1", "p1/c2"). By applying the + same transformation rule, we can obtain the final commands spec exposed by + the full processing pipeline: "p1/c1", "p1/c2" and "p3/c3". + + "p1/c1" "p1/c2" "p3/c3" + | | | + ----------------- | + | P1 | | + ----------------- | + | "p2/c1" | + ----------------- | + | P2 | | + ----------------- | + | "p3/c1" | "p3/c2" | + | | | + ------------------------------------ + | DACL | + ------------------------------------ + """ + + @property + @abc.abstractmethod + def name(self) -> str: + """Returns a unique string identifier for this object.""" + + @abc.abstractmethod + def process_commands( + self, consumed_commands: Mapping[str, gdmr_types.ArrayType] + ) -> Mapping[str, gdmr_types.ArrayType]: + """Processes the commands and returns a new modified version of it. + + Args: + consumed_commands: the commands up in the processing chain (or provided by + the Environment) that are required by this processor, i.e. with keys + specified by `consumed_commands_spec`. + + Returns the new commands. Note that the data in consumed_commands is removed + from the global commands dictionary. If users want to keep some of the + elements it is their responsibility to retain them in the output + dictionary. + """ + + @abc.abstractmethod + def consumed_commands_spec(self) -> Mapping[str, gdmr_types.AnyArraySpec]: + """Spec of the commands consumed by this processor.""" + + @abc.abstractmethod + def produced_commands_keys(self) -> set[str]: + """Keys of the commands produced by this processor.""" + + def reset(self) -> None: + """Resets the internal state of the command processor.""" + ... diff --git a/src/experimental/reaf/core/data_acquisition_and_control_layer.py b/src/experimental/reaf/core/data_acquisition_and_control_layer.py new file mode 100644 index 00000000..e8fcc8ca --- /dev/null +++ b/src/experimental/reaf/core/data_acquisition_and_control_layer.py @@ -0,0 +1,170 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""REAF data acquisition and control layer to interface with the robotic setup.""" + +from collections.abc import Iterable, Mapping + +from dm_env import specs +from gdm_robotics.interfaces import types as gdmr_types +from reaf.core import device as reaf_device +from reaf.core import device_coordinator as reaf_coordinator +from reaf.core import trigger + + +class DataAcquisitionAndControlLayer: + """REAF data acquisition and control layer. + + The DACL is responsible to provide an interface for the robotic setup. + """ + + def __init__( + self, + *, + device_coordinator: reaf_coordinator.DeviceCoordinator, + commands_trigger: trigger.Trigger | None, + measurements_trigger: trigger.Trigger | None, + ): + """Initializes the DataAcquisitionAndControlLayer. + + Args: + device_coordinator: The coordinator representing a specific robotic setup. + Note that callers need to explicitly initialize and finalise the + coordinator. + commands_trigger: A trigger to unblock processing commands during a call + to `step`. + measurements_trigger: A trigger to unblock processing measurements during + a call to `step`. + """ + self._coordinator = device_coordinator + self._devices = self._coordinator.get_devices() + # The following checks that names of the devices are unique and their keys + # are "mergeable". + self._check_device_names_and_keys(self._devices) + + self._commands_trigger = commands_trigger + self._measurements_trigger = measurements_trigger + + # Create a map of supported commands keys for each Device. + self._commands_for_device = { + device.name: device.commands_spec().keys() for device in self._devices + } + + def begin_stepping(self) -> Mapping[str, gdmr_types.ArrayType]: + """Begins stepping the DACL and returns the current measurements.""" + self._coordinator.on_begin_stepping() + + # Wait for the first trigger to happen before collecting the measurements. + if self._measurements_trigger is not None: + self._measurements_trigger.wait_for_event() + return self._get_measurements() + + def end_stepping(self) -> None: + """Ends stepping the data acquisition and control layer.""" + self._coordinator.on_end_stepping() + + def _set_commands(self, commands: Mapping[str, gdmr_types.ArrayType]) -> None: + """Sets the commands of the data acquisition and control layer.""" + self._coordinator.before_set_commands() + for device in self._devices: + device_commands = { + k: v + for k, v in commands.items() + if k in self._commands_for_device[device.name] + } + device.set_commands(device_commands) + self._coordinator.after_set_commands() + + def _get_measurements(self) -> Mapping[str, gdmr_types.ArrayType]: + """Gets the measurements of the data acquisition and control layer.""" + measurements = {} + self._coordinator.before_get_measurements() + for device in self._devices: + measurements.update(device.get_measurements()) + + return measurements + + def step( + self, commands: Mapping[str, gdmr_types.ArrayType] + ) -> Mapping[str, gdmr_types.ArrayType]: + """Steps the data acquisition and control layer.""" + if self._commands_trigger is not None: + self._commands_trigger.wait_for_event() + self._set_commands(commands) + + if self._measurements_trigger is not None: + self._measurements_trigger.wait_for_event() + return self._get_measurements() + + def commands_spec(self) -> Mapping[str, gdmr_types.AnyArraySpec]: + """Returns the specs for the commands.""" + spec = {} + for device in self._devices: + spec.update(device.commands_spec()) + return spec + + def measurements_spec(self) -> Mapping[str, specs.Array]: + """Returns the specs for the measurements.""" + spec = {} + for device in self._devices: + spec.update(device.measurements_spec()) + return spec + + @property + def device_coordinator(self) -> reaf_coordinator.DeviceCoordinator: + return self._coordinator + + def _check_keys_have_been_formatted_correctly( + self, current_key_set: Iterable[str] + ) -> None: + """Check that keys haven't been left unformatted.""" + for key in current_key_set: + if key.find("{}") != -1: + raise ValueError( + "Keys should not contain '{}'. Did you mean to use format()?" + ) + + def _check_device_names_and_keys( + self, devices: Iterable[reaf_device.Device] + ) -> None: + """Raises error if device names are not unique or keys are not exclusive.""" + # Check names first. + all_names = [device.name for device in devices] + unique_names = set(all_names) + if len(unique_names) != len(all_names): + raise RuntimeError(f"Duplicate names when checking devices: {all_names}") + + # Check commands. + devices = tuple(devices) + current_specs = set() + for device in devices: + device_keys = device.commands_spec().keys() + self._check_keys_have_been_formatted_correctly(device_keys) + if not current_specs.isdisjoint(device_keys): + raise RuntimeError( + f"Duplicate keys when checking device {device.name}:" + f" {current_specs.intersection(device_keys)}" + ) + current_specs.update(device_keys) + + # Check measurements. + current_specs = set() + for device in devices: + device_keys = device.measurements_spec().keys() + self._check_keys_have_been_formatted_correctly(device_keys) + if not current_specs.isdisjoint(device_keys): + raise RuntimeError( + f"Duplicate keys when checking device {device.name}:" + f" {current_specs.intersection(device_keys)}" + ) + current_specs.update(device_keys) diff --git a/src/experimental/reaf/core/default_discount_provider.py b/src/experimental/reaf/core/default_discount_provider.py new file mode 100644 index 00000000..ff216b7c --- /dev/null +++ b/src/experimental/reaf/core/default_discount_provider.py @@ -0,0 +1,79 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Computes a constant discount given the termination state. + +This provider returns a discount of 0.0 in case of termination and 1.0 +otherwise (i.e. for truncation and not termination). + +It is usually safe to use this discount provider for environments that return +strictly positive rewards. +""" + +from collections.abc import Mapping + +from dm_env import specs +from gdm_robotics.interfaces import types as gdmr_types +import numpy as np +from reaf.core import discount_provider +from reaf.core import termination_checker +import tree + + +class DefaultDiscountProvider(discount_provider.DiscountProvider): + """Computes a constant discount given the termination state. + + This provider returns a discount of 0.0 in case of termination and 1.0 + otherwise (i.e. for truncation and not termination). + + It is usually safe to use this discount provider for environments that return + strictly positive rewards. + """ + + def __init__(self, name: str = "default_discount_provider"): + self._name = name + self._spec = specs.BoundedArray( + shape=(), dtype=np.float64, minimum=0.0, maximum=1.0, name="discount" + ) + + def name(self) -> str: + """Returns a unique string identifier for this object.""" + return self._name + + def compute_discount( + self, + unused_required_features: Mapping[str, gdmr_types.ArrayType], + termination_state: termination_checker.TerminationResult, + ) -> tree.Structure[gdmr_types.ArrayType]: + """Computes the discount. + + Args: + unused_required_features: Unused + termination_state: The termination state as computed by the termination + checkers. Returns the discount. + + Returns: + The discount. + """ + if termination_state == termination_state.TERMINATE: + return np.asarray(0).astype(self._spec.dtype) + else: # TRUNCATION or DO_NOT_TERMINATE + return np.asarray(1.0).astype(self._spec.dtype) + + def discount_spec(self) -> tree.Structure[specs.Array]: + """Returns the spec of the discount.""" + return self._spec + + def required_features_keys(self) -> set[str]: + """Returns the feature keys that are required to compute the discount.""" + return set() diff --git a/src/experimental/reaf/core/default_observation_space_adapter.py b/src/experimental/reaf/core/default_observation_space_adapter.py new file mode 100644 index 00000000..1261b3ca --- /dev/null +++ b/src/experimental/reaf/core/default_observation_space_adapter.py @@ -0,0 +1,231 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""ObservationSpaceAdapter supporting filtering, renaming and type conversion.""" + +import abc +from collections.abc import Iterable, Mapping +import dataclasses + +from dm_env import specs +from gdm_robotics.interfaces import types as gdmr_types +import numpy as np +import numpy.typing as npt +from reaf.core import observation_space_adapter +import tree + + +@dataclasses.dataclass(frozen=True, kw_only=True) +class RenameInfo: + original_key: str + renamed_key: str + + +class ObservationTypeMapper(abc.ABC): + """Maps from REAF features and specs into corresponding environment types.""" + + @abc.abstractmethod + def to_observation_spec( + self, features_spec: Mapping[str, specs.Array] + ) -> gdmr_types.ObservationSpec: + """Convert the features spec into the environment observation spec.""" + + @abc.abstractmethod + def to_observations( + self, features: Mapping[str, gdmr_types.ArrayType] + ) -> tree.Structure[gdmr_types.ArrayType]: + """Convert the features into the environment observations.""" + + +class _DefaultObservationTypeMapper(ObservationTypeMapper): + """An ObservationTypeMapper that returns the input features specs and dict. + + This `ObservationTypeMapper` maps observations from the more constrained + `Mapping[str, ArrayType]` used in the task layer to the more generic + `tree.Structure[ArrayType]` exposed by the GDM Environment. + """ + + def to_observation_spec( + self, features_spec: Mapping[str, specs.Array] + ) -> gdmr_types.ObservationSpec: + """Returns the features spec, unmodified, as a `gdmr_types.ObservationSpec`.""" + return features_spec + + def to_observations( + self, features: Mapping[str, gdmr_types.ArrayType] + ) -> tree.Structure[gdmr_types.ArrayType]: + """Returns the features, unmodified, as a `tree.Structure`.""" + return features + + +class DefaultObservationSpaceAdapter( + observation_space_adapter.ObservationSpaceAdapter +): + """Observation adapter supporting filtering, renaming and type conversion. + + This adapter supports filtering, renaming, and converting REAF features into + environment observations. + + The order of operations is the following: + 1) Filtering, i.e. feature selection. + 2) Downcasting floats to max_float_dtype. + 3) Renaming. + 4) Type conversion. + + Please refer to the constructor documentation for more information. + """ + + def __init__( + self, + *, + task_features_spec: Mapping[str, specs.Array], + selected_features: Iterable[str] | None, + renamed_features: Iterable[RenameInfo] | None, + observation_type_mapper: ObservationTypeMapper | None, + max_float_dtype: npt.DTypeLike = np.float64, + ): + """Initializes the observation space adapter. + + Args: + task_features_spec: The spec of all the features exposed by the task + layer. + selected_features: The features that will be exposed as observations. If + None, all features will be exposed, i.e. no filtering. + renamed_features: `RenameInfo` objects specifying which features should be + renamed and the corresponding new name. If empty or None, no renaming + will occur. + observation_type_mapper: An `ObservationTypeMapper` specifying how to + convert the task layer features data type (i.e. a Mapping[str, + ArrayType]) into the more generic type exposed by the GDM Environment + (i.e. a tree.Structure[ArrayType]). If None, an instance of + `_DefaultObservationTypeMapper` is used which converts the task logic + layer features dictionary to the more generic type (i.e. + `tree.Structure[ArrayType])` exposed by the environment. + max_float_dtype: The maximum float dtype to use for downcasting floats. + """ + if not np.issubdtype(max_float_dtype, np.floating): + raise ValueError( + 'max_float_dtype must be a floating point dtype. Got' + f' {max_float_dtype}' + ) + self._max_float_dtype = max_float_dtype + self._max_bits = np.finfo(self._max_float_dtype).bits + self._task_features_spec = task_features_spec + self._selected_filter = selected_features + self._renamed_features = renamed_features or () + self._observation_type_mapper = ( + observation_type_mapper or _DefaultObservationTypeMapper() + ) + self._check_specs_consistency() + # Compute the observation spec only once. + self._observation_spec = self._compute_observation_spec() + + def _check_specs_consistency(self) -> None: + # Check that filter keys are present in the spec. + if self._selected_filter is not None: + all_features = self._task_features_spec.keys() + features = set() + for feature in self._selected_filter: + if feature not in all_features: + raise ValueError(f'Feature {feature} is not present in the spec.') + features.add(feature) + else: + # No filter applied. Select all features. + features = set(self._task_features_spec.keys()) + + # Check renaming. + for rename_info in self._renamed_features: + if rename_info.original_key not in features: + raise ValueError( + f'Feature {rename_info.original_key} is not present in the spec.' + ) + + def observations_from_features( + self, features: Mapping[str, gdmr_types.ArrayType] + ) -> tree.Structure[gdmr_types.ArrayType]: + """Converts the features into the final environment observations.""" + # 1. Filter the observations. + if (selected_features := self._selected_filter) is None: + # No filter. Expose all observations. + filtered_features = dict(features) + else: + filtered_features = { + k: v for k, v in features.items() if k in selected_features # pytype: disable=unsupported-operands + } + + # 2. Downcast floats to max_float_dtype. + filtered_features = { + k: self._downcast_if_necessary(v) for k, v in filtered_features.items() + } + + # 3. Rename. + for rename_info in self._renamed_features: + # Rename the feature. + value = filtered_features[rename_info.original_key] + del filtered_features[rename_info.original_key] + filtered_features[rename_info.renamed_key] = value + + # 4. Convert type. + return self._observation_type_mapper.to_observations(filtered_features) + + def _compute_observation_spec(self) -> gdmr_types.ObservationSpec: + """Computes the observation spec.""" + # 1. Filter the specs + if (features_to_filter := self._selected_filter) is None: + # The observation spec corresponds to the task features spec. + filtered_specs = dict(self._task_features_spec) + else: + filtered_specs = { + k: v + for k, v in self._task_features_spec.items() + if k in features_to_filter # pytype: disable=unsupported-operands + } + + # 2. Downcast floats to max_float_dtype. + for k, v in filtered_specs.items(): + if self._dtype_needs_downcast(v.dtype): + filtered_specs[k] = v.replace(dtype=self._max_float_dtype) + + # 3. Rename. + for rename_info in self._renamed_features: + # Rename the feature. + value = filtered_specs[rename_info.original_key] + del filtered_specs[rename_info.original_key] + filtered_specs[rename_info.renamed_key] = value + + # 4. Convert the type. + return self._observation_type_mapper.to_observation_spec(filtered_specs) + + def observation_spec(self) -> gdmr_types.ObservationSpec: + """Returns the observation spec.""" + return self._observation_spec + + def task_features_keys(self) -> set[str]: + """Returns the task features keys that will be converted by this adapter.""" + return set(self._task_features_spec.keys()) + + def _downcast_if_necessary( + self, value: gdmr_types.ArrayType + ) -> gdmr_types.ArrayType: + if ( + hasattr(value, 'dtype') and self._dtype_needs_downcast(value.dtype) + ) or self._dtype_needs_downcast(type(value)): + return np.asarray(value).astype(self._max_float_dtype) + else: + return value + + def _dtype_needs_downcast(self, dtype: npt.DTypeLike) -> bool: + return ( + np.issubdtype(dtype, np.floating) + and np.finfo(dtype).bits > self._max_bits + ) diff --git a/src/experimental/reaf/core/device.py b/src/experimental/reaf/core/device.py new file mode 100644 index 00000000..cc53a0ca --- /dev/null +++ b/src/experimental/reaf/core/device.py @@ -0,0 +1,54 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""REAF basic device to interface with the robotic setup.""" + +import abc +from collections.abc import Mapping +from dm_env import specs +from gdm_robotics.interfaces import types as gdmr_types + + +class Device(abc.ABC): + """REAF basic device to interface with the robotic setup. + + A device defines a single piece in the robotic setup. It should be + hermetic, that is, not depending on other Devices. The coordination of the + devices is responsibility of the DeviceCoordinator. + + Important: a Device should return the commands and measurements specs + immediately after initialisation without the need for any explicit + initialisation, nor for resource acquisition (e.g. connecting to the + hardware). + """ + + @property + @abc.abstractmethod + def name(self) -> str: + """Returns the name of this device.""" + + @abc.abstractmethod + def commands_spec(self) -> Mapping[str, gdmr_types.AnyArraySpec]: + """Returns the commands specs for this device.""" + + @abc.abstractmethod + def measurements_spec(self) -> Mapping[str, specs.Array]: + """Returns the measurements specs for this device.""" + + @abc.abstractmethod + def set_commands(self, commands: Mapping[str, gdmr_types.ArrayType]) -> None: + """Sets the commands for this device.""" + + @abc.abstractmethod + def get_measurements(self) -> Mapping[str, gdmr_types.ArrayType]: + """Returns the measurements provided by this device.""" diff --git a/src/experimental/reaf/core/device_coordinator.py b/src/experimental/reaf/core/device_coordinator.py new file mode 100644 index 00000000..233705bb --- /dev/null +++ b/src/experimental/reaf/core/device_coordinator.py @@ -0,0 +1,87 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Coordinates the devices composing a robotic setup.""" + +import abc +from collections.abc import Iterable +from reaf.core import device + + +class DeviceCoordinator(abc.ABC): + """Coordinates the devices composing a robotic setup. + + The `DeviceCoordinator` object is responsible for coordinating all the + devices constituting the robotic setup. Whilst the Device is hermetic, + the coordinator is responsible for passing information from one device to + the other if required. For example in a bimanual setup the coordinator is + charged with passing the position of each robot to the other so we can ensure + proper and safe interaction such as for example collision avoidance. + + The `DeviceCoordinator` can be configurable to enable different + properties on the robotic setup, e.g. adding or not adding a `Device` or + forwarding configuration to each `Device`. + + At the very least, the coordinator must implement `get_devices` + to return all the devices. We also provide `on_begin_stepping` and + `on_end_stepping` methods that will be called before the start of an episode + and after the end of the episode respectively. Note that resource acquisition + and subsequent release is completely up to the implementation. + + Finally, `before_set_commands`/`before_get_measurements` can be implemented to + coordinate devices behaviour before their corresponding functions are + called. + """ + + @property + @abc.abstractmethod + def name(self) -> str: + """Returns the name of the coordinator.""" + + @abc.abstractmethod + def get_devices(self) -> Iterable[device.Device]: + """Returns the devices composing the embodiment.""" + + # Lifecycle methods. + + def on_begin_stepping(self) -> None: + """Prepares the coordinator for having its devices called repeatedly. + + After `on_begin_stepping` the devices returned by `get_devices` will have + their `set_commands` and `get_measurements` called repeatedly until + `on_end_stepping` is called on this coordinator. + """ + + def on_end_stepping(self) -> None: + """Notifies the coordinator that the devices are no longer called. + + After `on_end_stepping` the devices returned by `get_devices` will not have + their `set_commands` and `get_measurements` called anymore until this + coordinator `on_begin_stepping` method is notified again. + """ + + # Step hooks methods. + + def before_set_commands(self) -> None: + """Prepares the coordinator to have its devices set_commands called.""" + + def after_set_commands(self) -> None: + """Notifies the coordinator that its devices got `set_commands` called.""" + + def before_get_measurements(self) -> None: + """Prepares the coordinator to have its devices get_measurements called. + + This method gets called immediately before the devices `get_measurements` + method is called and can be used to customise the devices state given the + whole setup state. + """ diff --git a/src/experimental/reaf/core/discount_provider.py b/src/experimental/reaf/core/discount_provider.py new file mode 100644 index 00000000..df1b22fa --- /dev/null +++ b/src/experimental/reaf/core/discount_provider.py @@ -0,0 +1,61 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Computes the discount.""" + +import abc +from collections.abc import Mapping + +from dm_env import specs +from gdm_robotics.interfaces import types as gdmr_types +from reaf.core import termination_checker +import tree + + +class DiscountProvider(abc.ABC): + """Computes the discount.""" + + @abc.abstractmethod + def name(self) -> str: + """Returns a unique string identifier for this object.""" + + @abc.abstractmethod + def compute_discount( + self, + required_features: Mapping[str, gdmr_types.ArrayType], + termination_state: termination_checker.TerminationResult, + ) -> tree.Structure[gdmr_types.ArrayType]: + """Computes the discount. + + Args: + required_features: Measurements and features computed by the task logic + that are required by this provider, i.e. that have keys specified by + `required_features_keys`. + termination_state: The termination state as computed by the termination + checkers. Returns the discount. + + Returns: + The discount. + """ + + @abc.abstractmethod + def discount_spec(self) -> tree.Structure[specs.Array]: + """Returns the spec of the discount.""" + + @abc.abstractmethod + def required_features_keys(self) -> set[str]: + """Returns the feature keys that are required to compute the discount.""" + + def reset(self) -> None: + """Resets the internal state of the discount provider.""" + ... diff --git a/src/experimental/reaf/core/entity.py b/src/experimental/reaf/core/entity.py new file mode 100644 index 00000000..68c2e38d --- /dev/null +++ b/src/experimental/reaf/core/entity.py @@ -0,0 +1,65 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Basic REAF-sim protocol to interface with the simulation.""" + +from collections.abc import Mapping +import typing + +from dm_env import specs +from gdm_robotics.interfaces import types as gdmr_types + + +class Entity(typing.Protocol): + """Basic REAF component to interface with the simulation. + + An entity defines a single component in the simulation that consumes substep + commands and outputs substep measurements at every simulation substep. It + should be hermetic, that is, not depending on other Entities. + + Important: an Entity should return the substep commands and substep + measurements specs immediately after initialisation without the need for any + explicit initialisation. + """ + + @property + def name(self) -> str: + """Instance name.""" + + def reset(self): + """Resets the entity.""" + + def substep_commands_spec( + self, + ) -> Mapping[str, specs.Array]: + """Spec for the substep commands.""" + + def substep_measurements_spec( + self, + ) -> Mapping[str, specs.Array]: + """Spec for the substep measurements.""" + + def set_substep_commands( + self, + model: typing.Any, + data: typing.Any, + consumed_substep_commands: Mapping[str, gdmr_types.ArrayType], + ) -> None: + """Sets the substep commands.""" + + def get_substep_measurements( + self, + model: typing.Any, + data: typing.Any, + ) -> Mapping[str, gdmr_types.ArrayType]: + """Returns the substep measurements.""" diff --git a/src/experimental/reaf/core/environment.py b/src/experimental/reaf/core/environment.py new file mode 100644 index 00000000..dca306f0 --- /dev/null +++ b/src/experimental/reaf/core/environment.py @@ -0,0 +1,490 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The Robotics Environment Authoring Framework (REAF) Environment class.""" + +import abc +from collections.abc import Mapping +import enum +from typing import Generic + +from absl import logging +import dm_env +from dm_env import specs +from gdm_robotics.interfaces import environment as gdmr_env +from gdm_robotics.interfaces import types as gdmr_types +import numpy as np +from reaf.core import action_space_adapter as reaf_action_space_adapter +from reaf.core import data_acquisition_and_control_layer as reaf_dacl +from reaf.core import default_observation_space_adapter +from reaf.core import logger as reaf_logger +from reaf.core import observation_space_adapter as reaf_observation_space_adapter +from reaf.core import pass_through_action_space_adapter +from reaf.core import task_logic_layer as reaf_tll +import tree + + +class ActionSpecEnforcementOption(enum.StrEnum): + """Options for action spec enforcement.""" + + CLIP_TO_SPEC = "clip_to_spec" + IGNORE = "ignore" + WARNING = "warning" + RAISE_ERROR = "raise_error" + + +class EnvironmentReset(abc.ABC, Generic[gdmr_env.ResetOptions]): + """Support for general resets adhering to the GDM environment API.""" + + @abc.abstractmethod + def do_reset( + self, + config: gdmr_env.ResetOptions, + ) -> None: + """Resets the environment.""" + + def default_reset_configuration(self) -> gdmr_env.ResetOptions: + """Returns the default reset configuration.""" + return gdmr_env.Options() + + +class EndOfEpisodeHandler: + """Handler called after the last episode step.""" + + def on_end_of_episode_stepping(self, final_timestep: dm_env.TimeStep) -> None: + """Called when the episode has ended stepping. + + This will be called at the end of every episode, after all other triggers + have been resolved. Episodes can end either due to truncation or + termination, i.e. `timestep.step_type` is `StepType.LAST`, or due to an + early call to `Environment.reset()`. To verify whether it has indeed + ended due to truncation or termination, the implementer should test + `timestep.last()`. + + Note that the first reset after environment construction will not trigger + this handler, but it will be triggered before resolving any subsequent + environment resets, either implicit or explicit. + + Args: + final_timestep: The final timestep of the episode that ended stepping. + """ + + +class EnvironmentCloser(abc.ABC): + """Handler called when the environment is closed.""" + + @abc.abstractmethod + def close(self) -> None: + """Releases resources when the environment is closed. + + This method is called automatically when exiting the environment's + context manager (`with` statement). + """ + + +class Environment(gdmr_env.Environment): + """The Robotics Environment Authoring Framework (REAF) Environment class.""" + + def __init__( + self, + *, + data_acquisition_and_control_layer: reaf_dacl.DataAcquisitionAndControlLayer, + task_logic_layer: reaf_tll.TaskLogicLayer, + environment_reset: EnvironmentReset, + action_space_adapter: ( + reaf_action_space_adapter.ActionSpaceAdapter | None + ) = None, + observation_space_adapter: ( + reaf_observation_space_adapter.ObservationSpaceAdapter | None + ) = None, + end_of_episode_handler: EndOfEpisodeHandler | None = None, + environment_closer: EnvironmentCloser | None = None, + action_spec_enforcement_option: ActionSpecEnforcementOption = ActionSpecEnforcementOption.RAISE_ERROR, + ): + """Creates an environment. + + Args: + data_acquisition_and_control_layer: The layer for communicating with the + specific robotic setup. + task_logic_layer: The layer in charge of defining the task. + environment_reset: The `EnvironmentReset` specifying the function to be + called at environment reset and the default environment reset + configuration. + action_space_adapter: Adapter from the agent action space to the flattened + commands accepted by the task layer. If None the + PassThroughActionSpaceAdapter is used, meaning the entirety of the + commands dictionary is exposed to the agent. + observation_space_adapter: Adapter from the computed features to the + observations that are exposed to the agent. If None the + DefaultObservationSpaceAdapter is used, meaning all the features are + exposed to the agent as observations. + end_of_episode_handler: Called at the end of an episode, after the last + step. + environment_closer: Specifies the handler to be called when the + environment is closed. This is called automatically on exit if the + environment is used as a context manager. If None, no action is + performed at close. + action_spec_enforcement_option: How to enforce the action spec. If + `CLIP_TO_SPEC`, the action will be clipped to the spec. If `WARNING`, an + warning logged if the action is outside the spec. If `RAISE_ERROR`, an + error will be raised if the action is outside the spec. If `IGNORE`, + the action will be passed through. Default is `RAISE_ERROR`. + """ + + self._data_acquisition_and_control_layer = ( + data_acquisition_and_control_layer + ) + self._task_logic_layer = task_logic_layer + self._end_of_episode_handler = ( + end_of_episode_handler or EndOfEpisodeHandler() + ) + self._environment_reset = environment_reset + self._environment_closer = environment_closer + self._action_spec_enforcement_option = action_spec_enforcement_option + + # Before assigning the adapters, validate the specs on the task logic layer + # and the DACL. + self._validate_dacl_and_ttl_specs() + + ttl_commands_spec = self._task_logic_layer.commands_spec( + self._data_acquisition_and_control_layer.commands_spec() + ) + ttl_features_spec = self._task_logic_layer.features_spec( + self._data_acquisition_and_control_layer.measurements_spec() + ) + + if action_space_adapter is None: + action_space_adapter = ( + pass_through_action_space_adapter.PassThroughActionSpaceAdapter( + commands_spec=ttl_commands_spec + ) + ) + self._action_space_adapter = action_space_adapter + + if observation_space_adapter is None: + observation_space_adapter = ( + default_observation_space_adapter.DefaultObservationSpaceAdapter( + task_features_spec=ttl_features_spec, + selected_features=None, + renamed_features=None, + observation_type_mapper=None, + ) + ) + self._observation_space_adapter = observation_space_adapter + + # Now we can validate the adapters. + self._validate_adapters_specs() + + self._last_timestep: dm_env.TimeStep | None = None + self._should_finalize_episode = False + self._timestep_spec = gdmr_types.TimeStepSpec( + step_type=gdmr_types.STEP_TYPE_SPEC, + reward=self._task_logic_layer.reward_spec(), + discount=self._task_logic_layer.discount_spec(), + # The observation spec corresponds to the one exposed by the adapter. + observation=self._observation_space_adapter.observation_spec(), + ) + + self._zero_reward, self._zero_discount = tree.map_structure( + _read_only_zeros_like_spec, + (self._timestep_spec.reward, self._timestep_spec.discount), + ) + + def close(self) -> None: + """Frees any resources used by the environment.""" + if self._environment_closer is not None: + self._environment_closer.close() + + def default_reset_options(self) -> gdmr_env.ResetOptions: + return self._environment_reset.default_reset_configuration() + + def reset_with_options( + self, + *, + options: gdmr_env.ResetOptions, + ) -> dm_env.TimeStep: + """Starts a new sequence and returns the first `TimeStep`.""" + if self._should_finalize_episode: + self._finalize_episode() + self._environment_reset.do_reset(options) + self._task_logic_layer.perform_reset() + measurements = self._data_acquisition_and_control_layer.begin_stepping() + features = self._task_logic_layer.compute_all_features(measurements) + observations = self._compute_observations_from_features(features) + + self._last_timestep = self._restart(observation=observations) + # Make sure any early reset after this one triggers `_finalize_episode`. + self._should_finalize_episode = True + return self._last_timestep + + def action_spec(self) -> gdmr_types.ActionSpec: + """Defines the actions that should be provided to `step`.""" + # The action spec corresponds to the one exposed by the adapter. + return self._action_space_adapter.action_spec() + + def timestep_spec(self) -> gdmr_types.TimeStepSpec: + """Returns the spec associated to the returned TimeStep.""" + return self._timestep_spec + + def step(self, action: gdmr_types.ActionType) -> dm_env.TimeStep: + """Updates the environment according to action and returns a `TimeStep`.""" + + action = self._enforce_action_spec(action) + if self._last_timestep is None or self._last_timestep.last(): + return self.reset() + + # Process the action to obtain a command. + commands = self._compute_commands_from_agent_action(action) + commands = self._task_logic_layer.compute_final_commands(commands) + measurements = self._data_acquisition_and_control_layer.step(commands) + + # Compute all the features. + features = self._task_logic_layer.compute_all_features(measurements) + + # Compute the elements of the timestep. + reward = self._task_logic_layer.compute_reward(features) + termination_state = self._task_logic_layer.check_for_termination(features) + discount = self._task_logic_layer.compute_discount( + features, termination_state + ) + + observations = self._compute_observations_from_features(features) + + if termination_state.is_terminated(): + self._last_timestep = self._termination( + reward=reward, observation=observations + ) + elif termination_state.is_truncated(): + self._last_timestep = self._truncation( + reward=reward, observation=observations, discount=discount + ) + else: + self._last_timestep = self._transition( + reward=reward, observation=observations, discount=discount + ) + + if self._last_timestep.last(): + self._finalize_episode() + return self._last_timestep + + def _finalize_episode(self) -> None: + self._data_acquisition_and_control_layer.end_stepping() + # It's crucial to call `end_stepping` on the dacl before invoking the end + # of episode handler. This ensures no further `set_command` or + # `get_measurements` calls are made. In contrast, the end of episode + # handler might interact with devices, requiring them to be informed + # beforehand. + self._end_of_episode_handler.on_end_of_episode_stepping(self._last_timestep) + self._should_finalize_episode = False + + @property + def data_acquisition_and_control_layer( + self, + ) -> reaf_dacl.DataAcquisitionAndControlLayer: + return self._data_acquisition_and_control_layer + + @property + def task_logic_layer(self) -> reaf_tll.TaskLogicLayer: + return self._task_logic_layer + + @property + def environment_reset(self) -> EnvironmentReset: + return self._environment_reset + + @environment_reset.setter + def environment_reset(self, environment_reset: EnvironmentReset) -> None: + self._environment_reset = environment_reset + + def add_logger(self, logger: reaf_logger.Logger) -> None: + self._task_logic_layer.add_logger(logger) + + def remove_logger(self, logger: reaf_logger.Logger) -> None: + self._task_logic_layer.remove_logger(logger) + + def _validate_dacl_and_ttl_specs(self) -> None: + """Validates the specs on the task logic layer.""" + # Validate the spec on the task logic layer. + self._task_logic_layer.validate_spec( + dacl_commands_spec=( + self._data_acquisition_and_control_layer.commands_spec() + ), + dacl_measurements_spec=( + self._data_acquisition_and_control_layer.measurements_spec() + ), + ) + + def _validate_adapters_specs(self) -> None: + # Collect the full commands and features spec and validate them against + # the adapters. + commands_spec = set( + self._task_logic_layer.commands_spec( + self._data_acquisition_and_control_layer.commands_spec() + ).keys() + ) + features_spec = set( + self._task_logic_layer.features_spec( + self._data_acquisition_and_control_layer.measurements_spec() + ) + ) + + # Check the action space adapter. + adapter_keys = self._action_space_adapter.task_commands_keys() + + if adapter_keys != commands_spec: + raise ValueError( + "Mismatch between commands exposed by the action space adapter:" + f" {adapter_keys} and commands spec expected by the task layer:" + f" {commands_spec}." + ) + + # Check the observation spec adapter. + adapter_keys = self._observation_space_adapter.task_features_keys() + if not adapter_keys.issubset(features_spec): + raise ValueError( + "Failed to validate observation space adapter specs. Missing keys:" + f" {adapter_keys - features_spec}" + ) + + def _compute_observations_from_features( + self, features: Mapping[str, gdmr_types.ArrayType] + ) -> tree.Structure[gdmr_types.ArrayType]: + return self._observation_space_adapter.observations_from_features(features) + + def _compute_commands_from_agent_action( + self, agent_action: gdmr_types.ActionType + ) -> Mapping[str, gdmr_types.ArrayType]: + return self._action_space_adapter.commands_from_environment_action( + agent_action + ) + + def _restart( + self, + observation: tree.Structure[gdmr_types.ArrayType], + ) -> dm_env.TimeStep: + """Returns a `TimeStep` with `step_type` set to `StepType.FIRST`.""" + return dm_env.TimeStep( + step_type=np.asarray(dm_env.StepType.FIRST, dtype=np.uint8), + observation=observation, + reward=self._zero_reward, + discount=self._zero_discount, + ) + + def _transition( + self, + reward: tree.Structure[gdmr_types.ArrayType], + observation: tree.Structure[gdmr_types.ArrayType], + discount: tree.Structure[gdmr_types.ArrayType], + ) -> dm_env.TimeStep: + """Returns a `TimeStep` with `step_type` set to `StepType.MID`.""" + return dm_env.TimeStep( + step_type=np.asarray(dm_env.StepType.MID, dtype=np.uint8), + observation=observation, + reward=reward, + discount=discount, + ) + + def _termination( + self, + reward: tree.Structure[gdmr_types.ArrayType], + observation: tree.Structure[gdmr_types.ArrayType], + ) -> dm_env.TimeStep: + """Returns a `TimeStep` with `step_type` set to `StepType.LAST`.""" + return dm_env.TimeStep( + step_type=np.asarray(dm_env.StepType.LAST, dtype=np.uint8), + observation=observation, + reward=reward, + discount=self._zero_discount, + ) + + def _truncation( + self, + reward: tree.Structure[gdmr_types.ArrayType], + observation: tree.Structure[gdmr_types.ArrayType], + discount: tree.Structure[gdmr_types.ArrayType], + ) -> dm_env.TimeStep: + """Returns a `TimeStep` with `step_type` set to `StepType.LAST`.""" + return dm_env.TimeStep( + step_type=np.asarray(dm_env.StepType.LAST, dtype=np.uint8), + observation=observation, + reward=reward, + discount=discount, + ) + + def _enforce_action_spec( + self, action: gdmr_types.ActionType + ) -> gdmr_types.ActionType: + """Enforces the action spec.""" + match self._action_spec_enforcement_option: + case ActionSpecEnforcementOption.IGNORE: + pass + case ActionSpecEnforcementOption.CLIP_TO_SPEC: + try: + + def clip_to_spec(a, s): + if isinstance(s, specs.BoundedArray): + return np.clip(a, s.minimum, s.maximum) + return a + + action = tree.map_structure( + clip_to_spec, + action, + self._action_space_adapter.action_spec(), + ) + except ValueError as e: + raise ValueError( + "Failed to clip action to spec. Action:" + f" {action} and spec: {self._action_space_adapter.action_spec()}" + ) from e + case ActionSpecEnforcementOption.WARNING: + + def _validate_without_raising(a, s): + dtype_ok = s.dtype == a.dtype + shape_ok = s.shape == a.shape + minimum_ok = True + maximum_ok = True + if isinstance(s, specs.BoundedArray): + minimum_ok = (s.minimum <= a).all() + maximum_ok = (a <= s.maximum).all() + return dtype_ok and shape_ok and minimum_ok and maximum_ok + + if not all( + tree.flatten( + tree.map_structure( + _validate_without_raising, + action, + self._action_space_adapter.action_spec(), + ) + ) + ): + logging.warning( + "Failed to validate action against spec. Action: %r and spec: %r", + action, + self._action_space_adapter.action_spec(), + ) + case ActionSpecEnforcementOption.RAISE_ERROR: + action = tree.map_structure( + lambda a, spec: spec.validate(a), action, self.action_spec() + ) + case _: + raise ValueError( + "Unknown action spec enforcement option:" + f" {self._action_spec_enforcement_option}" + ) + return action + + +def _read_only_zeros_like_spec(spec: specs.Array) -> np.ndarray: + """Returns a zero array matching the specified spec.""" + arr = np.zeros(shape=spec.shape, dtype=spec.dtype) + arr.flags.writeable = False + return arr diff --git a/src/experimental/reaf/core/features_observer.py b/src/experimental/reaf/core/features_observer.py new file mode 100644 index 00000000..cc87f0d3 --- /dev/null +++ b/src/experimental/reaf/core/features_observer.py @@ -0,0 +1,34 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Observe all the produced features and measurements.""" + +import abc +from collections.abc import Mapping + +from gdm_robotics.interfaces import types as gdmr_types + + +class FeaturesObserver(abc.ABC): + """Observe all the produced features and measurements.""" + + @property + @abc.abstractmethod + def name(self) -> str: + """Returns a unique string identifier for this object.""" + + @abc.abstractmethod + def observe_features( + self, features: Mapping[str, gdmr_types.ArrayType] + ) -> None: + """Observes all the features and measurements.""" diff --git a/src/experimental/reaf/core/features_producer.py b/src/experimental/reaf/core/features_producer.py new file mode 100644 index 00000000..8ce44e94 --- /dev/null +++ b/src/experimental/reaf/core/features_producer.py @@ -0,0 +1,56 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Produces additional features to be exposed by the task logic layer.""" + +import abc +from collections.abc import Mapping + +from dm_env import specs +from gdm_robotics.interfaces import types as gdmr_types + + +class FeaturesProducer(abc.ABC): + """Produces additional features to be exposed by the task logic layer.""" + + @property + @abc.abstractmethod + def name(self) -> str: + """Returns a unique string identifier for this object.""" + + @abc.abstractmethod + def produce_features( + self, required_features: Mapping[str, gdmr_types.ArrayType] + ) -> Mapping[str, gdmr_types.ArrayType]: + """Produces additional features for the environment. + + Args: + required_features: Measurements and features generated by previous + producers in the processing chain that are required by this processor, + i.e. with keys specified by `required_features_keys`. + + Returns additional features that will be added to the global measurements + and features dictionary. + """ + + @abc.abstractmethod + def produced_features_spec(self) -> Mapping[str, specs.Array]: + """Returns the spec of the features produced by this producer.""" + + @abc.abstractmethod + def required_features_keys(self) -> set[str]: + """Returns the keys that are required to produce the new features.""" + + def reset(self) -> None: + """Resets the internal state of the feature producer.""" + ... diff --git a/src/experimental/reaf/core/logger.py b/src/experimental/reaf/core/logger.py new file mode 100644 index 00000000..63bb5f33 --- /dev/null +++ b/src/experimental/reaf/core/logger.py @@ -0,0 +1,80 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Support logging inside the task logic layer.""" + +import abc +from collections.abc import Mapping + +from gdm_robotics.interfaces import types as gdmr_types + + +class Logger(abc.ABC): + """Support logging inside the task logic layer. + + Lifecycle + For each environment step, these member functions are called in this order: + 1. `record_measurements` is called with raw measurements from the sensors. + 2. `record_features` is called with features derived from the measurements. + 3. `record_commands_processing` is called for each + `CommandsProcessor.process_commands` invocation, tracking the + transformation of commands. + 4. `record_final_commands` is called once with the final commands sent to + the DACL. + + Notes: + An environment is first reset(). This triggers the first two steps above. + See reset_with_options in ./environment.py. + + After reset, step is called repeatedly. + 1. This first triggers steps 3 and 4 (See compute_final_commands in TLL + called from step in ./environment.py) + 2. Features are computed (see compute_all_features in TLL called from + step in ./environment.py), triggering steps 1 and 2. + """ + + @property + @abc.abstractmethod + def name(self) -> str: + """Unique string identifier for this object.""" + + def record_measurements( + self, measurements: Mapping[str, gdmr_types.ArrayType] + ) -> None: + """Called once with all the measurements from the DACL.""" + + def record_features( + self, features: Mapping[str, gdmr_types.ArrayType] + ) -> None: + """Called once with all the features computed in the Task Layer.""" + + def record_final_commands( + self, commands: Mapping[str, gdmr_types.ArrayType] + ) -> None: + """Called once with the final commands sent to the DACL.""" + + def record_commands_processing( + self, + name: str, + consumed_commands: Mapping[str, gdmr_types.ArrayType], + produced_commands: Mapping[str, gdmr_types.ArrayType], + ) -> None: + """Called once per call to `process_commands` for each CommandsProcessor. + + Args: + name: Name of the `CommandsProcessor`. + consumed_commands: The commands consumed by the current + `CommandsProcessor`. + produced_commands: The commands produced by the current + `CommandsProcessor`. + """ diff --git a/src/experimental/reaf/core/numpy_mock_assertions.py b/src/experimental/reaf/core/numpy_mock_assertions.py new file mode 100644 index 00000000..310a918e --- /dev/null +++ b/src/experimental/reaf/core/numpy_mock_assertions.py @@ -0,0 +1,98 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Testing functions for asserting on Mock objects with numpy structures.""" + +from collections.abc import Sequence +from unittest import mock +import numpy as np + + +def assert_called_once_with(mock_obj: mock.Mock, *args, **kwargs) -> None: + if mock_obj.call_count != 1: + raise AssertionError( + f"Expected exactly one call to {mock_obj}, got {mock_obj.call_count}" + ) + + assert_called_with(mock_obj, *args, **kwargs) + + +def assert_called_with(mock_obj: mock.Mock, *args, **kwargs) -> None: + """Asserts that the last call to mock_obj had the specified arguments.""" + if mock_obj.call_args is None: + raise AssertionError( + f"Mock object {mock_obj} not called. Expected one call." + ) + call_args, call_kwargs = mock_obj.call_args + np.testing.assert_equal(call_args, args) + np.testing.assert_equal(call_kwargs, kwargs) + + +def assert_has_calls( + mock_obj: mock.Mock, calls: Sequence[mock._Call], any_order: bool = False +) -> None: + """Asserts that mock_obj has been called with the specified calls.""" + mock_calls = mock_obj.mock_calls + + # Check that there are at least enough calls. + if mock_obj.call_count < len(calls): + raise AssertionError( + f"Expected at least {len(calls)} calls to {mock_obj}, got" + f" {mock_obj.call_count}" + ) + + def _calls_are_equal(actual: mock._Call, expected: mock._Call) -> bool: + _, actual_args, actual_kwargs = actual + _, expected_args, expected_kwargs = expected + # Quickest way to transform the assertion into a comparator. + try: + np.testing.assert_equal(actual_args, expected_args) + np.testing.assert_equal(actual_kwargs, expected_kwargs) + return True + except AssertionError: + return False + + if any_order: + # We just check for the calls to be contained. + for expected_call in calls: + for actual_call in mock_calls: + if _calls_are_equal(actual_call, expected_call): + break + raise AssertionError( + f"Expected call {expected_call} not found in mock calls {mock_calls}." + ) + return + + # We need to check in order, but first find the first call. + starting_index = -1 + first_expected_call = calls[0] + for index, actual_call in enumerate(mock_calls): + if _calls_are_equal(actual_call, first_expected_call): + starting_index = index + break + if starting_index == -1: + raise AssertionError(f"Calls {calls} not found in mock calls {mock_calls}.") + + non_matching_calls = [] + + # We have the first element. Now we need to compare element wise. + for index, expected_call in enumerate(calls): + actual_call = mock_calls[starting_index + index] + if not _calls_are_equal(actual_call, expected_call): + non_matching_calls.append((index, expected_call, actual_call)) + + if non_matching_calls: + raise AssertionError( + f"Calls {calls} do not match mock calls {mock_calls}. Mismatch (index," + f" expected, actual): {non_matching_calls}" + ) diff --git a/src/experimental/reaf/core/observation_space_adapter.py b/src/experimental/reaf/core/observation_space_adapter.py new file mode 100644 index 00000000..ca4a8d86 --- /dev/null +++ b/src/experimental/reaf/core/observation_space_adapter.py @@ -0,0 +1,42 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Adapts REAF features into observations exposed by the environment.""" + +import abc +from collections.abc import Mapping +from gdm_robotics.interfaces import types as gdmr_types +import tree + + +class ObservationSpaceAdapter(abc.ABC): + """Adapts REAF features into observations exposed by the environment. + + Implementations of this interface are responsible for converting the features + generated by the REAF task layer logic (i.e. dictionary of tensors) into the + more generic `observation` structure exposed by the environment. + """ + + @abc.abstractmethod + def observations_from_features( + self, features: Mapping[str, gdmr_types.ArrayType] + ) -> tree.Structure[gdmr_types.ArrayType]: + """Converts the REAF features into the environment observations.""" + + @abc.abstractmethod + def observation_spec(self) -> gdmr_types.ObservationSpec: + """Returns the observation spec.""" + + @abc.abstractmethod + def task_features_keys(self) -> set[str]: + """Returns the task features keys that will be converted by this adapter.""" diff --git a/src/experimental/reaf/core/pass_through_action_space_adapter.py b/src/experimental/reaf/core/pass_through_action_space_adapter.py new file mode 100644 index 00000000..16b28a81 --- /dev/null +++ b/src/experimental/reaf/core/pass_through_action_space_adapter.py @@ -0,0 +1,55 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Adapter that passes the commands spec through.""" + +from collections.abc import Mapping +from gdm_robotics.interfaces import types as gdmr_types +from reaf.core import action_space_adapter + + +class PassThroughActionSpaceAdapter(action_space_adapter.ActionSpaceAdapter): + """Adapter that passes the commands spec through. + + NB the resulting environment will expose a dictionary as the action spec. + """ + + def __init__(self, commands_spec: Mapping[str, gdmr_types.AnyArraySpec]): + self._commands_spec = commands_spec + + def commands_from_environment_action( + self, environment_action: gdmr_types.ActionType + ) -> Mapping[str, gdmr_types.ArrayType]: + """Returns commands accepted by REAF. + + commands_from_environment_action usually accepts a gdmr_types.ActionType but + since this adapter passes the same action as the commands, it needs to be a + dict type in order to pass it through as a dict. + + Args: + environment_action: The environment action(s) to pass as REAF commands. + """ + if not isinstance(environment_action, dict): + raise ValueError( + 'environment_action must be a dict, but got: ' + f'{type(environment_action)}.' + ) + return environment_action + + def action_spec(self) -> gdmr_types.ActionSpec: + """Returns the action spec exposed by the environment.""" + return self._commands_spec + + def task_commands_keys(self) -> set[str]: + """Returns the keys for the commands exposed to the task layer.""" + return set(self._commands_spec.keys()) diff --git a/src/experimental/reaf/core/reward_provider.py b/src/experimental/reaf/core/reward_provider.py new file mode 100644 index 00000000..3a8ec655 --- /dev/null +++ b/src/experimental/reaf/core/reward_provider.py @@ -0,0 +1,292 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Computes the reward.""" + +import abc +from collections.abc import Mapping +import operator +from typing import Callable, TypeAlias, TypeVar, Union + +from dm_env import specs +from gdm_robotics.interfaces import types as gdmr_types +import numpy as np +import tree + + +RewardValue: TypeAlias = tree.Structure[gdmr_types.ArrayType] +RewardSpec: TypeAlias = tree.Structure[specs.Array] + + +class _RewardProvider(abc.ABC): + """Computes the reward. + + Defines the interface for a reward provider. + + Important: Users should not inherit from this class directly. Instead, use the + RewardProvider class later in this file. + """ + + @abc.abstractmethod + def name(self) -> str: + """Returns a unique string identifier for this object.""" + + @abc.abstractmethod + def compute_reward( + self, required_features: Mapping[str, gdmr_types.ArrayType] + ) -> RewardValue: + """Computes the reward. + + Args: + required_features: Measurements and features computed by the task logic + that are required by this provider, i.e. that have keys specified by + `required_features_keys`. + + Returns the computed reward. + """ + + @abc.abstractmethod + def reward_spec(self) -> RewardSpec: + """Returns the spec of the reward.""" + + @abc.abstractmethod + def required_features_keys(self) -> set[str]: + """Returns the feature keys that are required to compute the reward.""" + + def reset(self) -> None: + """Resets the internal state of the reward provider.""" + ... + + +RewardProviderOrValue: TypeAlias = Union['RewardProvider', RewardValue] + + +S = TypeVar('S') +T = TypeVar('T') +UnaryOperator: TypeAlias = Callable[[S], S] +BinaryOperator: TypeAlias = Callable[[S | T, S | T], S | T] + + +class RewardProvider(_RewardProvider): + """Computes the reward. + + Important: Users should inherit from this class and implement the abstract + methods defined in the interface _RewardProvider. + """ + + def __add__(self, other: RewardProviderOrValue): + return BinaryOperationRewardProvider(operator.add, self, other) + + def __radd__(self, other: RewardProviderOrValue): + return BinaryOperationRewardProvider(operator.add, other, self) + + def __sub__(self, other: RewardProviderOrValue): + return BinaryOperationRewardProvider(operator.sub, self, other) + + def __rsub__(self, other: RewardProviderOrValue): + return BinaryOperationRewardProvider(operator.sub, other, self) + + def __mul__(self, other: RewardProviderOrValue): + return BinaryOperationRewardProvider(operator.mul, self, other) + + def __rmul__(self, other: RewardProviderOrValue): + return BinaryOperationRewardProvider(operator.mul, other, self) + + def __truediv__(self, other: RewardProviderOrValue): + return BinaryOperationRewardProvider(operator.truediv, self, other) + + def __rtruediv__(self, other: RewardProviderOrValue): + return BinaryOperationRewardProvider(operator.truediv, other, self) + + def __floordiv__(self, other: RewardProviderOrValue): + return BinaryOperationRewardProvider(operator.floordiv, self, other) + + def __rfloordiv__(self, other: RewardProviderOrValue): + return BinaryOperationRewardProvider(operator.floordiv, other, self) + + def __pow__(self, other: RewardProviderOrValue): + return BinaryOperationRewardProvider(operator.pow, self, other) + + def __rpow__(self, other: RewardProviderOrValue): + return BinaryOperationRewardProvider(operator.pow, other, self) + + def __getitem__(self, index: slice): + return GetItemOperationRewardProvider(self, index) + + def __neg__(self): + return UnaryOperationRewardProvider(operator.neg, self) + + +class ConstantRewardProvider(RewardProvider): + """A RewardProvider that always returns the same reward.""" + + def __init__(self, reward: RewardValue): + super().__init__() + self._reward = reward + + def name(self) -> str: + return str(self._reward) + + def compute_reward( + self, required_features: Mapping[str, gdmr_types.ArrayType] + ) -> RewardValue: + return self._reward + + def reward_spec(self) -> RewardSpec: + return tree.map_structure( + lambda v: specs.Array(v.shape, v.dtype), self._reward + ) + + def required_features_keys(self) -> set[str]: + return set() + + +class BinaryOperationRewardProvider(RewardProvider): + """Applies a binary operator to the result of two reward providers.""" + + def __init__( + self, + op: BinaryOperator, + first_reward_provider: RewardProviderOrValue, + second_reward_provider: RewardProviderOrValue, + ): + super().__init__() + if not isinstance(first_reward_provider, RewardProvider): + first_reward_provider = ConstantRewardProvider(first_reward_provider) + if not isinstance(second_reward_provider, RewardProvider): + second_reward_provider = ConstantRewardProvider(second_reward_provider) + first_spec = first_reward_provider.reward_spec() + second_spec = second_reward_provider.reward_spec() + tree.assert_same_structure(first_spec, second_spec) + assert all( + tree.flatten( + tree.map_structure( + lambda s1, s2: s1.shape == s2.shape and s1.dtype == s2.dtype, + first_spec, + second_spec, + ) + ) + ) + self._op = op + self._first_reward_provider = first_reward_provider + self._second_reward_provider = second_reward_provider + self._reward_spec = first_reward_provider.reward_spec() + self._first_required_features_keys = ( + first_reward_provider.required_features_keys() + ) + self._second_required_features_keys = ( + second_reward_provider.required_features_keys() + ) + + def name(self) -> str: + op_name = getattr(self._op, '__name__', str(self._op)) + return ( + f'{op_name}({self._first_reward_provider.name()},' + f' {self._second_reward_provider.name()})' + ) + + def compute_reward( + self, required_features: Mapping[str, gdmr_types.ArrayType] + ) -> RewardValue: + first_required_features = { + k: v + for k, v in required_features.items() + if k in self._first_required_features_keys + } + second_required_features = { + k: v + for k, v in required_features.items() + if k in self._second_required_features_keys + } + return tree.map_structure( + self._op, + self._first_reward_provider.compute_reward(first_required_features), + self._second_reward_provider.compute_reward(second_required_features), + ) + + def reward_spec(self) -> RewardSpec: + return self._reward_spec + + def required_features_keys(self) -> set[str]: + return ( + self._first_required_features_keys | self._second_required_features_keys + ) + + def reset(self) -> None: + self._first_reward_provider.reset() + self._second_reward_provider.reset() + + +class GetItemOperationRewardProvider(RewardProvider): + """Extracts a slice from the result of a reward provider.""" + + def __init__(self, reward_provider: RewardProviderOrValue, index: slice): + super().__init__() + if not isinstance(reward_provider, RewardProvider): + reward_provider = ConstantRewardProvider(reward_provider) + self._reward_provider = reward_provider + self._index = index + + def name(self) -> str: + return f'{self._reward_provider.name}[{self._index}]' + + def compute_reward( + self, required_features: Mapping[str, gdmr_types.ArrayType] + ) -> RewardValue: + return tree.map_structure( + lambda v: v[self._index], + self._reward_provider.compute_reward(required_features), + ) + + def reward_spec(self) -> RewardSpec: + return tree.map_structure( + lambda s: specs.Array(np.empty(s.shape)[self._index].shape, s.dtype), + self._reward_provider.reward_spec(), + ) + + def required_features_keys(self) -> set[str]: + return self._reward_provider.required_features_keys() + + def reset(self) -> None: + self._reward_provider.reset() + + +class UnaryOperationRewardProvider(RewardProvider): + """Applies a unary operator to the result of a reward provider.""" + + def __init__(self, op: UnaryOperator, reward_provider: RewardProviderOrValue): + super().__init__() + if not isinstance(reward_provider, RewardProvider): + reward_provider = ConstantRewardProvider(reward_provider) + self._op = op + self._reward_provider = reward_provider + + def name(self) -> str: + op_name = getattr(self._op, '__name__', str(self._op)) + return f'{op_name}({self._reward_provider.name()})' + + def compute_reward( + self, required_features: Mapping[str, gdmr_types.ArrayType] + ) -> RewardValue: + return tree.map_structure( + self._op, self._reward_provider.compute_reward(required_features) + ) + + def reward_spec(self) -> RewardSpec: + return self._reward_provider.reward_spec() + + def required_features_keys(self) -> set[str]: + return self._reward_provider.required_features_keys() + + def reset(self) -> None: + self._reward_provider.reset() diff --git a/src/experimental/reaf/core/substep_commands_processor.py b/src/experimental/reaf/core/substep_commands_processor.py new file mode 100644 index 00000000..b63f4388 --- /dev/null +++ b/src/experimental/reaf/core/substep_commands_processor.py @@ -0,0 +1,104 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Protocol for substep commands manipulation in REAF-sim.""" + +from collections.abc import Mapping +import typing + +from dm_env import specs +from gdm_robotics.interfaces import types as gdmr_types + + +class SubstepCommandsProcessor(typing.Protocol): + """Processes substep commands, propagating them through a pipeline. + + This processor manipulates substep commands, acting as a node in a pipeline. + It consumes substep commands, performs operations, and produces updated + substep commands for the next stage in the processing chain. + + The processing pipeline starts with commands provided to the SimulationDevice + and progresses towards the substep commands consumed by the individual + entities. Each processor consumes a subset of substep commands and produces + new, potentially transformed, substep commands. The order of operations is + crucial. + + Example Pipeline (conceptual): + + Simulation Device commands --> Processor (1) --> Processor (2) --> Entities + + Specs are propagated starting from the bottom: + 1) In this example assume that the set of entities expect "p3/c1", "p3/c2" and + "p3/c3". + 2) Processor (2) returns ("p3/c1", "p3/c2") from "p2/c1". This means + that the global substep commands spec exposed at this level is "p2/c1" and + the unprocessed "p3/c3". + 3) Processor (1) returns "p2/c1" from ("p1/c1", "p1/c2"). By applying + the same transformation rule, we can obtain the final spec exposed + by the SimulationDevice: "p1/c1", "p1/c2" and "p3/c3". + + ------------------------------------ + | SimulationDevice | + ------------------------------------ + + "p1/c1" "p1/c2" "p3/c3" + | | | + ----------------- | + | P1 | | + ----------------- | + | "p2/c1" | + ----------------- | + | P2 | | + ----------------- | + | "p3/c1" | "p3/c2" | + | | | + ------------------------------------ + | Entities | + ------------------------------------ + """ + + @property + def name(self) -> str: + """Returns a unique string identifier for this object.""" + + def reset(self) -> None: + """Resets the internal state of this processor.""" + + def produced_substep_commands_keys(self) -> set[str]: + """Keys of the substep commands produced by this processor.""" + + def consumed_substep_commands_spec( + self, + ) -> Mapping[str, specs.Array]: + """Spec of the substep commands consumed by this processor.""" + + def process_substep_commands( + self, + model: typing.Any, + data: typing.Any, + consumed_substep_commands: Mapping[str, gdmr_types.ArrayType], + ) -> Mapping[str, gdmr_types.ArrayType]: + """Processes the substep commands and returns a new modified version of it. + + Args: + model: the simulation model. + data: the simulation data. + consumed_substep_commands: the substep commands up in the processing chain + that are required by this processor, i.e. with keys specified by + `consumed_substep_commands_spec`. + + Returns the new substep commands. Note that the (key, value) pairs in + `consumed_substep_commands` are removed from the running substep commands + dictionary. If users want to keep some of the elements it is their + responsibility to retain them in the output dictionary. + """ diff --git a/src/experimental/reaf/core/substep_measurements_processor.py b/src/experimental/reaf/core/substep_measurements_processor.py new file mode 100644 index 00000000..15dc81df --- /dev/null +++ b/src/experimental/reaf/core/substep_measurements_processor.py @@ -0,0 +1,103 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Protocol for substep measurements manipulation in REAF-sim.""" + +from collections.abc import Mapping +import typing + +from dm_env import specs +from gdm_robotics.interfaces import types as gdmr_types + + +class SubstepMeasurementsProcessor(typing.Protocol): + """Processes substep measurements, propagating them through a pipeline. + + This processor manipulates substep measurements, acting as a node in a + pipeline. It consumes substep measurements, performs operations, and produces + updated substep measurements for the next stage in the processing chain. + + The processing pipeline starts with substep measurements produced by Entities + and progresses towards the measurements exposed by the SimulationDevice. Each + processor consumes a subset of substep measurements and produces new, + potentially transformed, substep measurements. The order of operations is + crucial. + + Example Pipeline (conceptual): + + Entities --> Processor (1) --> Processor (2) -> Simulation Device Measurements + + Specs are propagated starting from the bottom: + 1) In this example assume that the set of entities produce "p1/c1", "p1/c2" + and "p1/c3". + 2) Processor (1) returns "p2/c1" from ("p1/c1", "p1/c2"). + 3) Processor (2) returns ("p3/c1", "p3/c2") from "p2/c1". + + This resulting spec exposed by the SimulationDevice: "p3/c1", "p3/c2" + and "p1/c3". + + ------------------------------------ + | SimulationDevice | + ------------------------------------ + + "p3/c1" "p3/c2" "p1/c3" + | | | + ----------------- | + | P2 | | + ----------------- | + | "p2/c1" | + ----------------- | + | P1 | | + ----------------- | + | "p1/c1" | "p1/c2" | + | | | + ------------------------------------ + | Entities | + ------------------------------------ + """ + + @property + def name(self) -> str: + """Returns a unique string identifier for this object.""" + + def reset(self): + """Resets the internal state of this processor.""" + + def produced_substep_measurements_spec( + self, + ) -> Mapping[str, specs.Array]: + """Spec of the substep measurements consumed by this processor.""" + + def consumed_substep_measurements_keys(self) -> set[str]: + """Keys of the substep measurements consumed by this processor.""" + + def process_substep_measurements( + self, + model: typing.Any, + data: typing.Any, + consumed_substep_measurements: Mapping[str, gdmr_types.ArrayType], + ) -> Mapping[str, gdmr_types.ArrayType]: + """Processes the substep measurements and returns a new modified version of it. + + Args: + model: the simulation model. + data: the simulation data. + consumed_substep_measurements: the substep measurements up in the + processing chain that are required by this processor, i.e. with keys + specified by `consumed_substep_measurements_spec`. + + Returns the new substep measurements. Note that the (key, value) pairs in + `consumed_substep_measurements` are removed from the running substep + measurements dictionary. If users want to keep some of the elements it is + their responsibility to retain them in the output dictionary. + """ diff --git a/src/experimental/reaf/core/task_logic_layer.py b/src/experimental/reaf/core/task_logic_layer.py new file mode 100644 index 00000000..51b25639 --- /dev/null +++ b/src/experimental/reaf/core/task_logic_layer.py @@ -0,0 +1,342 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Task logic layer for the Robotics Environment Authoring Framework.""" + +from collections.abc import Mapping, Sequence +import itertools +from typing import Protocol + +from absl import logging +from dm_env import specs +from gdm_robotics.interfaces import types as gdmr_types +from reaf.core import commands_processor as reaf_commands_processor +from reaf.core import default_discount_provider +from reaf.core import discount_provider as reaf_discount_provider +from reaf.core import features_observer as reaf_features_observers +from reaf.core import features_producer as reaf_features_producer +from reaf.core import logger as reaf_logger +from reaf.core import reward_provider as reaf_reward_provider +from reaf.core import termination_checker as reaf_termination_checker +from reaf.core import zero_reward_provider +import tree + + +class _ResettableObject(Protocol): + """Protocol for an object that can be reset.""" + + def reset(self) -> None: + ... + + +class TaskLogicLayer: + """Task logic layer for the Robotics Environment Authoring Framework.""" + + def __init__( + self, + *, + commands_processors: Sequence[reaf_commands_processor.CommandsProcessor], + features_producers: Sequence[reaf_features_producer.FeaturesProducer], + termination_checkers: Sequence[ + reaf_termination_checker.TerminationChecker + ], + reward_provider: reaf_reward_provider.RewardProvider | None = None, + discount_provider: reaf_discount_provider.DiscountProvider | None = None, + features_observers: Sequence[ + reaf_features_observers.FeaturesObserver + ] = (), + loggers: Sequence[reaf_logger.Logger] = (), + ): + """Initializes the task logic layer. + + Args: + commands_processors: `CommandsProcessor`s that modify the commands before + being sent down to the DACL. They are called sequentially, starting from + the commands supplied by the policy and ending with the commands that + will be sent to the DACL. + features_producers: `FeaturesProducer`s that generate new features. + Measurements collected by the DACL and features produced by these + `FeaturesProducer`s are then merged into the final feature set that is + provided to the `reward_provider`, `termination_checkers`, + `discount_provider`, `features_observers`, and `loggers`. + termination_checkers: `TerminationChecker`s that check the episode + termination based on the final feature set. + reward_provider: `RewardProvider` that computes a reward based on the + final feature set. If None, the ZeroRewardProvider is used and the + reward is set to 0. + discount_provider: `DiscountProvider` that compute a discount based on the + final feature set and final termination state. If None, the + DefaultDiscountProvider is used returning 0 for termination and 1 for + truncation and non-termination. + features_observers: `FeaturesObserver`s that get a view over the final + feature set. + loggers: `Logger`s for logging measurements, features, and commands in the + task layer. + """ + self._commands_processors = commands_processors + self._features_producers = features_producers + self._reward_provider = ( + reward_provider + if reward_provider + else zero_reward_provider.ZeroRewardProvider() + ) + self._termination_checkers = termination_checkers + self._discount_provider = ( + discount_provider + if discount_provider + else default_discount_provider.DefaultDiscountProvider() + ) + self._features_observers = features_observers + self._loggers = list(loggers) + + # We make a set of all resettable objects so that these objects only get + # their resets called once. This is important for e.g. when having a single + # object that derives from two interfaces. + self._resettable_objects: list[_ResettableObject] = [] + unique_ids = set() + for resettable_object in itertools.chain( + self._commands_processors, + self._features_producers, + self._termination_checkers, + [self._reward_provider], + [self._discount_provider], + ): + resettable_object_id = id(resettable_object) + if resettable_object_id not in unique_ids: + unique_ids.add(resettable_object_id) + self._resettable_objects.append(resettable_object) + + def validate_spec( + self, + *, + dacl_commands_spec: Mapping[str, gdmr_types.AnyArraySpec], + dacl_measurements_spec: Mapping[str, specs.Array], + ) -> None: + """Checks that the specs have consistent keys.""" + logging.vlog(3, "Validate features processing") + self._validate_features_spec(dacl_measurements_spec) + self._validate_commands_spec(dacl_commands_spec) + + def features_spec( + self, + dacl_measurements_spec: Mapping[str, specs.Array], + ) -> Mapping[str, specs.Array]: + """Returns the features spec as exposed by the task layer.""" + spec = dict(dacl_measurements_spec) + for features_producer in self._features_producers: + spec.update(features_producer.produced_features_spec()) + + return spec + + def commands_spec( + self, dacl_commands_spec: Mapping[str, gdmr_types.AnyArraySpec] + ) -> Mapping[str, gdmr_types.AnyArraySpec]: + """Returns the commands spec exposed by the task layer.""" + # Each processor consumes commands (as described by its + # `consumed_commands_spec`) and outputs a potentially different set of + # commands (as described by its `produced_commands_keys`). + # Starting with the DACL command spec, we iterate in reverse order (i.e. in + # the direction DACL -> Policy) through every processor to remove the + # `produced_commands_keys` from the spec, and add their + # `consumed_commands_spec` to the spec. + spec: Mapping[str, gdmr_types.AnyArraySpec] = dict(dacl_commands_spec) + for processor in reversed(self._commands_processors): + processor_produced_keys = processor.produced_commands_keys() + spec = { + key: value + for key, value in spec.items() + if key not in processor_produced_keys + } + spec.update(processor.consumed_commands_spec()) + return spec + + def reward_spec(self) -> tree.Structure[specs.Array]: + return self._reward_provider.reward_spec() + + def discount_spec(self) -> tree.Structure[specs.Array]: + return self._discount_provider.discount_spec() + + def perform_reset(self) -> None: + """Reset the internal state of the task logic layer.""" + for resettable_object in self._resettable_objects: + resettable_object.reset() + + def compute_all_features( + self, measurements: Mapping[str, gdmr_types.ArrayType] + ) -> Mapping[str, gdmr_types.ArrayType]: + """Computes all the task logic features given the current measurements.""" + for logger in self._loggers: + logger.record_measurements(measurements) + + # Produce all the features. + current_features = dict(measurements) + for feature_producer in self._features_producers: + required_features = { + key: current_features[key] + for key in feature_producer.required_features_keys() + } + current_features.update( + feature_producer.produce_features(required_features) + ) + + # Observe the features. + for feature_observer in self._features_observers: + feature_observer.observe_features(current_features) + + # Log the resulting features. + for logger in self._loggers: + logger.record_features(current_features) + return current_features + + def compute_final_commands( + self, + policy_commands: Mapping[str, gdmr_types.ArrayType], + ) -> Mapping[str, gdmr_types.ArrayType]: + """Processes the policy commands and returns the final processed commands.""" + current_commands = dict(policy_commands) + for processor in self._commands_processors: + # Get commands to be consumed by the processor and remove the commands + # from the current_commands.. They correspond to the + # `consumed_command_spec`. + consumed_commands = { + key: current_commands.pop(key) + for key in processor.consumed_commands_spec().keys() + } + produced_commands = processor.process_commands(consumed_commands) + current_commands.update(produced_commands) + + # Log the modification. + for logger in self._loggers: + logger.record_commands_processing( + processor.name, consumed_commands, produced_commands + ) + + for logger in self._loggers: + logger.record_final_commands(current_commands) + return current_commands + + def compute_reward( + self, features: Mapping[str, gdmr_types.ArrayType] + ) -> tree.Structure[gdmr_types.ArrayType]: + """Computes the reward given the features.""" + return self._reward_provider.compute_reward({ + key: features[key] + for key in self._reward_provider.required_features_keys() + }) + + def check_for_termination( + self, features: Mapping[str, gdmr_types.ArrayType] + ) -> reaf_termination_checker.TerminationResult: + """Checks for termination.""" + current_state = reaf_termination_checker.TerminationResult.DO_NOT_TERMINATE + for termination_checker in self._termination_checkers: + current_state = reaf_termination_checker.TerminationResult.combine( + current_state, + termination_checker.check_termination({ + key: features[key] + for key in termination_checker.required_features_keys() + }), + ) + return current_state + + def compute_discount( + self, + features: Mapping[str, gdmr_types.ArrayType], + termination_state: reaf_termination_checker.TerminationResult, + ) -> tree.Structure[gdmr_types.ArrayType]: + """Computes the discount given the features and termination state.""" + return self._discount_provider.compute_discount( + { + key: features[key] + for key in self._discount_provider.required_features_keys() + }, + termination_state, + ) + + def add_logger(self, logger: reaf_logger.Logger) -> None: + self._loggers.append(logger) + + def remove_logger(self, logger: reaf_logger.Logger) -> None: + self._loggers.remove(logger) + + def _validate_features_spec( + self, dacl_measurements_spec: Mapping[str, specs.Array] + ) -> None: + """Validates the features spec.""" + # Check measurements/features path. + current_key_set = set(dacl_measurements_spec.keys()) + logging.vlog(4, "DACL measurements keys: %s", current_key_set) + + for producer in self._features_producers: + logging.vlog( + 4, + "Producer %s requires %s.", + producer.name, + producer.required_features_keys(), + ) + # Check required features are available. + if not producer.required_features_keys().issubset(current_key_set): + raise ValueError( + "Failed to validate feature specs for feature producer" + f" {producer.name}. Missing keys:" + f" {producer.required_features_keys() - current_key_set}" + ) + # Check that there are not duplicates in the output. + if not current_key_set.isdisjoint( + producer.produced_features_spec().keys() + ): + raise ValueError( + "Failed to validate feature specs for feature producer" + f" {producer.name}. Duplicate keys:" + f" {current_key_set & producer.produced_features_spec().keys()}" + ) + # Now extend the spec. + logging.vlog( + 4, + "Update available keys (from producer %s) with %s.", + producer.name, + producer.produced_features_spec().keys(), + ) + current_key_set.update(producer.produced_features_spec().keys()) + logging.vlog(4, "Available features keys %s.", current_key_set) + + def _validate_commands_spec( + self, dacl_commands_spec: Mapping[str, gdmr_types.AnyArraySpec] + ) -> None: + """Validates the commands spec.""" + # Check commands. Starting from the DACL command specs we propagate up in + # the chain. + logging.vlog(3, "Validate commands processing from DACL to Policy.") + current_key_set = set(dacl_commands_spec.keys()) + logging.vlog(4, "DACL commands keys: %s", current_key_set) + + for processor in reversed(self._commands_processors): + produced_command_keys = processor.produced_commands_keys() + + logging.vlog( + 4, + "Processor %s: specs (accepted keys) %s. Exposes %s.", + processor.name, + processor.consumed_commands_spec().keys(), + produced_command_keys, + ) + if not produced_command_keys.issubset(current_key_set): + raise ValueError( + "Failed to validate commands specs for commands processor" + f" {processor.name}. Missing (consumable) keys:" + f" {produced_command_keys - current_key_set}" + ) + # Remove the produced keys and add the consumed commands specs (as the + # processor is mutable). + current_key_set = current_key_set - produced_command_keys + current_key_set.update(processor.consumed_commands_spec().keys()) diff --git a/src/experimental/reaf/core/termination_checker.py b/src/experimental/reaf/core/termination_checker.py new file mode 100644 index 00000000..754c250d --- /dev/null +++ b/src/experimental/reaf/core/termination_checker.py @@ -0,0 +1,94 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Checks if the episode should terminate.""" + +import abc +from collections.abc import Mapping +import enum +from typing import Self + +from gdm_robotics.interfaces import types as gdmr_types + + +class TerminationResult(enum.IntFlag): + """The result of an episode termination check. + + The TerminationResult refers to the possibility for an episode to terminate. + For more details on the concept of termination we refer the readers to + https://github.com/google-deepmind/dm_env/blob/master/docs/index.md#environment-api-and-semantics. + + Note that this enum does not refer to the possible causes of termination but + only how the termination impacts the learning process. + + The result can be one of the following options: + - DO_NOT_TERMINATE: The episode should not terminate. + - TRUNCATE: The epsisode should terminate. Truncation implies a non-failure + final state. Usually this is associated with a non-zero discount. + - TERMINATE: The episode should terminate as the environment is in some + final state. Usually this is associated with a zero discount for e.g. + finite-horizon RL. + """ + + DO_NOT_TERMINATE = 0 + TRUNCATE = 2**0 + TERMINATE = 2**1 + + def is_terminated(self) -> bool: + return self == TerminationResult.TERMINATE + + def is_truncated(self) -> bool: + return self == TerminationResult.TRUNCATE + + def combine(self, other: Self) -> Self: + # TERMINATE has precedence over TRUNCATE, which in turn has precedence over + # DO_NOT_TERMINATE. Given the definitions above, this can be implemented as + # a maximum operator. To also enable tracing with JAX, we implement this in + # a branchless manner using bitwise operations that preserve the type. + # Note that JAX will trace TerminationResult values as ints. + # Approach: + # - self ^ (self ^ other) == other + # - (-1 * (self < other)) will be bitmask of all 1s iff self < other. + # - AND with (self ^ other) will result in either update or no-op bitmask. + return self ^ ((self ^ other) & (-1 * (self < other))) + + +class TerminationChecker(abc.ABC): + """Checks if the episode should terminate.""" + + @abc.abstractmethod + def name(self) -> str: + """Returns a unique string identifier for this object.""" + + @abc.abstractmethod + def check_termination( + self, required_features: Mapping[str, gdmr_types.ArrayType] + ) -> TerminationResult: + """Checks if the episode should terminate. + + Args: + required_features: Measurements and features computed by the task logic + that are required by this checker, i.e. that have keys specified by + `required_features_keys`. + + Returns if the episode should terminate (and if so, what kind of + termination). + """ + + @abc.abstractmethod + def required_features_keys(self) -> set[str]: + """Returns the feature keys that are required to check the termination.""" + + def reset(self) -> None: + """Resets the internal state of the termination checker.""" + ... diff --git a/src/experimental/reaf/core/trigger.py b/src/experimental/reaf/core/trigger.py new file mode 100644 index 00000000..20901873 --- /dev/null +++ b/src/experimental/reaf/core/trigger.py @@ -0,0 +1,29 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Defines an event-based waiting behaviour.""" + +import abc + + +class Trigger(abc.ABC): + """Defines an event-based waiting behaviour.""" + + @property + @abc.abstractmethod + def name(self) -> str: + """Returns the name of the trigger.""" + + @abc.abstractmethod + def wait_for_event(self) -> None: + """Blocks until the next event.""" diff --git a/src/experimental/reaf/core/zero_reward_provider.py b/src/experimental/reaf/core/zero_reward_provider.py new file mode 100644 index 00000000..c5d2267a --- /dev/null +++ b/src/experimental/reaf/core/zero_reward_provider.py @@ -0,0 +1,48 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Reward provider which provides a zero reward.""" + +from collections.abc import Mapping +from dm_env import specs +from gdm_robotics.interfaces import types as gdmr_types +import numpy as np +from reaf.core import reward_provider +import tree + + +class ZeroRewardProvider(reward_provider.RewardProvider): + """Reward provider which provides a zero reward.""" + + def __init__(self, name: str = 'zero_reward_provider'): + self._name = name + + def name(self) -> str: + return self._name + + def compute_reward( + self, required_features: Mapping[str, gdmr_types.ArrayType] + ) -> tree.Structure[gdmr_types.ArrayType]: + """Returns a zero reward.""" + return np.zeros(1) + + def reward_spec(self) -> tree.Structure[specs.Array]: + """Returns the spec for a constant zero reward.""" + return specs.Array(shape=(1,), dtype=float) + + def required_features_keys(self) -> set[str]: + """Returns empty set. + + There are no feature keys that are required to compute the reward. + """ + return set()