diff --git a/lib/python3.12/site-packages/ray/experimental/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e0bc2af481d1cbf73de5034b113909b7c84a1051 Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/__pycache__/compiled_dag_ref.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/__pycache__/compiled_dag_ref.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..88d9db391369165f21442a3fbf8f36e9a7654f6b Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/__pycache__/compiled_dag_ref.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/__pycache__/dynamic_resources.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/__pycache__/dynamic_resources.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f0a5f8ddd95d4c563574f9ad40f9b2446451dc6a Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/__pycache__/dynamic_resources.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/__pycache__/gradio_utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/__pycache__/gradio_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..506a8e0a180d7b05f432ebe592941ff84c5d1ae5 Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/__pycache__/gradio_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/__pycache__/internal_kv.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/__pycache__/internal_kv.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a67cc4b807764d43e904bf5d8421bc928849fb38 Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/__pycache__/internal_kv.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/__pycache__/locations.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/__pycache__/locations.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d680c25fb2486845f1ce6fa1435e31e8ff831d83 Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/__pycache__/locations.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/__pycache__/queue.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/__pycache__/queue.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d13c162ecd96ca22c0e420989dd5c67c1bbb22ed Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/__pycache__/queue.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/__pycache__/shuffle.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/__pycache__/shuffle.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0781c14ae4fe2e172a1dcbc3b4ae872aa5a44fa6 Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/__pycache__/shuffle.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/__pycache__/tf_utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/__pycache__/tf_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0075c7abf3247c5416da105df893e6a68efef83c Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/__pycache__/tf_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/__pycache__/tqdm_ray.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/__pycache__/tqdm_ray.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8c4fcadde9e428ede52a2af7837442d478c2b88c Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/__pycache__/tqdm_ray.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/channel/__init__.py b/lib/python3.12/site-packages/ray/experimental/channel/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..caf2351fae6fb2281a46fd187b3e00fe0149d6c5 --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/channel/__init__.py @@ -0,0 +1,43 @@ +from ray.experimental.channel.cached_channel import CachedChannel +from ray.experimental.channel.common import ( # noqa: F401 + AwaitableBackgroundReader, + AwaitableBackgroundWriter, + ChannelContext, + ChannelInterface, + ChannelOutputType, + CompiledDAGArgs, + ReaderInterface, + SynchronousReader, + SynchronousWriter, + WriterInterface, +) +from ray.experimental.channel.communicator import Communicator +from ray.experimental.channel.cpu_communicator import CPUCommunicator +from ray.experimental.channel.intra_process_channel import IntraProcessChannel +from ray.experimental.channel.shared_memory_channel import ( + BufferedSharedMemoryChannel, + Channel, + CompositeChannel, +) +from ray.experimental.channel.torch_tensor_accelerator_channel import ( + TorchTensorAcceleratorChannel, +) + +__all__ = [ + "AwaitableBackgroundReader", + "AwaitableBackgroundWriter", + "CachedChannel", + "Channel", + "Communicator", + "CPUCommunicator", + "ReaderInterface", + "SynchronousReader", + "SynchronousWriter", + "WriterInterface", + "ChannelContext", + "TorchTensorAcceleratorChannel", + "IntraProcessChannel", + "CompositeChannel", + "BufferedSharedMemoryChannel", + "CompiledDAGArgs", +] diff --git a/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9549ff6b97676f86bbbaadc44649734232e4a913 Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/accelerator_context.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/accelerator_context.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..44feafc75c33e223d4f03c23c11092bcd2be8cf0 Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/accelerator_context.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/auto_transport_type.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/auto_transport_type.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c1002da9cad7acb5b908b0ae6fba6ed1c03ab366 Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/auto_transport_type.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/cached_channel.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/cached_channel.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d23fcffff3e923e29c08cc9d1e129cf9acc415e2 Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/cached_channel.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/common.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/common.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5c17ab1e2194590e860c35b3fb7c9b71b76977c6 Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/common.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/communicator.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/communicator.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9fc82fc08fe5e87d5dfc44a49fbef66e4261cead Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/communicator.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/communicator_handle.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/communicator_handle.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a69da586eeb68690bd69f7f02c28ee5b781239f2 Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/communicator_handle.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/conftest.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/conftest.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..df90709e64f7b953e44480a35e7d8a049d04cf51 Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/conftest.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/cpu_communicator.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/cpu_communicator.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..adb2b59e426b5f87a5afcf94dd3919fe2be9dd2d Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/cpu_communicator.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/intra_process_channel.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/intra_process_channel.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2a80572122bc36610bf2f3d3c948652f7f9841d9 Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/intra_process_channel.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/nccl_group.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/nccl_group.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..236dad3ea822c0f1f1a1364d5c7ffed058063066 Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/nccl_group.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/serialization_context.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/serialization_context.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..814795bf25ff7486de7fc95d4c3ab754378d92e0 Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/serialization_context.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/shared_memory_channel.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/shared_memory_channel.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..56aadfca26f20fdf4de9f763b6dd35b02f53ec22 Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/shared_memory_channel.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/torch_tensor_accelerator_channel.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/torch_tensor_accelerator_channel.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c84a4404189e71581addb4217745b400086898bf Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/torch_tensor_accelerator_channel.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/torch_tensor_type.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/torch_tensor_type.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c2662bb79d821382addf384d90de19b5d6ec4b2a Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/torch_tensor_type.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b78cd6ab776f8cf30269edabf91c4bd0efdc0be5 Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/channel/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/channel/accelerator_context.py b/lib/python3.12/site-packages/ray/experimental/channel/accelerator_context.py new file mode 100644 index 0000000000000000000000000000000000000000..838545b274ee8972954bbcf09bdf8ebf930c5155 --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/channel/accelerator_context.py @@ -0,0 +1,246 @@ +import importlib +import threading +from contextlib import nullcontext +from typing import TYPE_CHECKING, ContextManager, List, Optional, Type + +import ray +from ray._private.accelerators import get_accelerator_manager_for_resource +from ray.experimental.channel.communicator import Communicator + +if TYPE_CHECKING: + import torch + +# The accelerator context singleton on this process. +_accelerator_context_lock = threading.Lock() +_default_accelerator_context: Optional["AcceleratorContext"] = None +_global_custom_context: Optional["AcceleratorContext"] = None + + +class AcceleratorContext: + """ + Provides a unified interface for managing different accelerator backends + This includes stream management, event creation, device context control, + and communicator support for distributed communication. + """ + + def __init__(self, torch_module_name: str, communicator_cls: Type[Communicator]): + """ + Initializes an accelerator context with the specified torch device module + and communicator class. + + Args: + torch_module_name: Name of the torch device module (e.g., "cuda", "cpu"). + communicator_cls: Class used to handle communication. + """ + + # The name of the torch module (e.g., 'cuda', 'npu') + self._torch_module_name: str = torch_module_name + # The Communicator class used to manage communication + self._communicator_cls: Type[Communicator] = communicator_cls + + # Import the torch backend module (e.g., torch.cuda) if the device is not 'cpu'. + if torch_module_name != "cpu": + self._torch_mod = importlib.import_module(f"torch.{torch_module_name}") + + @staticmethod + def get() -> "AcceleratorContext": + """ + Returns the singleton instance of the accelerator context. + + If a custom accelerator has been registered, initializes the context + based on the registration. Otherwise, selects an appropriate runtime + based on the available device (CUDA or CPU) and registers the + corresponding default communicator. + + Returns: + AcceleratorContext: A singleton instance of the appropriate + runtime context. + """ + + global _default_accelerator_context, _global_custom_context + + with _accelerator_context_lock: + if _global_custom_context is not None: + return _global_custom_context + + if _default_accelerator_context is None: + if len(ray.get_gpu_ids()) > 0: + from ray.experimental.channel.nccl_group import _NcclGroup + + _default_accelerator_context = AcceleratorContext( + "cuda", _NcclGroup + ) + else: + from ray.experimental.channel.cpu_communicator import ( + CPUCommunicator, + ) + + _default_accelerator_context = AcceleratorContext( + "cpu", CPUCommunicator + ) + + return _default_accelerator_context + + @staticmethod + def set(accelerator_context: "AcceleratorContext") -> None: + """ + Overwrites the default accelerator context. + + Args: + accelerator_context: The context to register. + """ + global _global_custom_context + + # Accelerator context is registered. + _global_custom_context = accelerator_context + + def get_accelerator_devices(self) -> List["torch.device"]: + """ + Gets the torch device list configured for this process. + + Returns: + List[torch.device]: The torch device list. + """ + import torch + + if self._torch_module_name == "cpu": + return [torch.device("cpu")] + + if self._torch_module_name == "cuda": + accelerator_ids = [str(id) for id in ray.get_gpu_ids()] + accelerator_manager = get_accelerator_manager_for_resource("GPU") + else: + accelerator_ids = [ + str(id) + for id in ray.get_runtime_context().get_accelerator_ids()[ + self._torch_module_name.upper() + ] + ] + accelerator_manager = get_accelerator_manager_for_resource( + self._torch_module_name.upper() + ) + + device_ids = [] + + if len(accelerator_ids) > 0: + accelerator_visible_list = ( + accelerator_manager.get_current_process_visible_accelerator_ids() + ) + if accelerator_visible_list is None: + accelerator_visible_list = [] + + # If there are multiple Accelerators, return a list of devices. + # If using fractional Accelerators, these IDs are not guaranteed + # to be unique across different processes. + for accelerator_id in accelerator_ids: + try: + device_ids.append(accelerator_visible_list.index(accelerator_id)) + except ValueError: + raise RuntimeError( + f"{accelerator_manager.get_visible_accelerator_ids_env_var()} set incorrectly. " + f"expected to include {accelerator_id}. " + "Did you override this environment" + " variable? If not, please help file an issue on Github." + ) + + else: + # If called on the driver or outside of Ray Train, return the + # 0th device. + device_ids.append(0) + + return [ + torch.device(f"{self._torch_module_name}:{device_id}") + for device_id in device_ids + ] + + def get_device_context(self, device: "torch.device") -> ContextManager: + """ + Retrieves the context manager for the specified accelerator device. + There is no device context for CPU, returning a nullcontext. + + Args: + device: The target device for which the context manager is required. + + Returns: + ContextManager: A context manager specific to the device type. + """ + if device.type == "cpu": + return nullcontext() + + return self._torch_mod.device(device) + + def current_stream(self): + """ + Retrieves the current execution stream for the accelerator device. + """ + return self._torch_mod.current_stream() + + def create_event(self): + """ + Creates an event object for the accelerator device. + """ + return self._torch_mod.Event() + + def generate_communicator_id(self) -> str: + """ + Generates a communication identifier for communication group. + """ + return self._communicator_cls.generate_communicator_id() + + def create_communicator(self, *args, **kwargs) -> Communicator: + """ + Creates a communication group for collective operations. + """ + return self._communicator_cls(*args, **kwargs) + + @property + def module_name(self) -> str: + """ + Gets the name of the torch module backing the accelerator. + """ + return self._torch_module_name + + @property + def communicator_cls(self) -> Optional[Type[Communicator]]: + """ + Returns the communicator class. + """ + return self._communicator_cls + + @property + def accelerator_count(self) -> int: + """ + Returns the number of accelerators assigned by ray. + """ + if self._torch_module_name == "cuda": + return len(ray.get_gpu_ids()) + else: + accelerator_ids = ray.get_runtime_context().get_accelerator_ids() + return len(accelerator_ids.get(self._torch_module_name.upper(), [])) + + +def register_accelerator_context( + torch_module_name: str, communicator_cls: Type[Communicator] +): + """ + Registers the accelerator context with the specified device type and communicator. + + Args: + torch_module_name: The name of the device module under torch. + communicator_cls: The communicator class associated with the device. + """ + accelerator_context = AcceleratorContext(torch_module_name, communicator_cls) + AcceleratorContext.set(accelerator_context) + + +def is_accelerator_context_registered(): + """ + Checks whether a custom accelerator context has been registered. + + Returns: + bool: True if a custom accelerator context is registered + (_global_custom_context is not None), False otherwise. + """ + if _global_custom_context is not None: + return True + return False diff --git a/lib/python3.12/site-packages/ray/experimental/channel/auto_transport_type.py b/lib/python3.12/site-packages/ray/experimental/channel/auto_transport_type.py new file mode 100644 index 0000000000000000000000000000000000000000..4017ac1fe921e83d6917dbff258bf85fb06066dd --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/channel/auto_transport_type.py @@ -0,0 +1,182 @@ +from typing import Dict, List, Optional, Tuple, Union + +import ray +from ray.experimental.channel import ChannelOutputType +from ray.experimental.channel.torch_tensor_type import TorchTensorType +from ray.experimental.util.types import Device + + +class AutoTransportType(ChannelOutputType): + """ + Type hint for automatic transport selection for tensors. + + With this type hint Compiled Graphs automatically decide the best transport + to use (e.g., accellerator or shared memory) based on the node locations and + GPU IDs of the readers and writers. + """ + + def __init__( + self, + device: Device = Device.DEFAULT, + _static_shape: bool = False, + _direct_return: bool = False, + ): + self._device = device + self._static_shape = _static_shape + self._direct_return = _direct_return + + @property + def device(self) -> Device: + return self._device + + def create_channel( + self, + writer: Optional["ray.actor.ActorHandle"], + reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]], + driver_actor_id: Optional[str] = None, + ) -> "ChannelOutputType": + """ + Directly calling create_channel() on AutoTransportType should not happen, + just raise an exception with informative message. + """ + raise ValueError( + "This should not happen: AutoTransportType should " + "have been resolved before creating a channel. " + "Please file a Ray GitHub issue for bug report." + ) + + +class TypeHintResolver: + """ + This class is used to resolve `AutoChannelType` into an actual channel type + (e.g., `TorchTensorType` with proper transport) based on node locations and + GPU IDs of the readers and writers. + """ + + def __init__(self, actor_to_gpu_ids: Dict["ray.actor.ActorHandle", List[str]]): + """ + Args: + actor_to_gpu_ids: Mapping from actor handle to its GPU IDs. + """ + self._actor_to_gpu_ids = actor_to_gpu_ids + + def _get_gpu_ids(self, actor: "ray.actor.ActorHandle") -> List[str]: + """ + Get the GPU IDs of the actor. + + Returns: + The GPU IDs of the actor. If the actor is not found, + return an empty list. + """ + gpu_ids = self._actor_to_gpu_ids.get(actor, []) + assert len(gpu_ids) <= 1, ( + "Compiled Graphs currently don't support allocating multiple GPUs " + "to a single actor" + ) + return gpu_ids + + def _use_same_gpu( + self, + writer_and_node: Tuple["ray.actor.ActorHandle", str], + reader_and_node: Union[ + Tuple["ray.actor.ActorHandle", str], + List[Tuple["ray.actor.ActorHandle", str]], + ], + ) -> bool: + """ + Check if the writer and readers use the same GPU. + + Args: + writer_and_node: A tuple of writer actor handle and its node ID. + reader_and_node: A tuple of reader actor handle and its node ID, or + a list of such tuples. + + Returns: + True if the writer and all the readers use the same GPU, False otherwise. + """ + if isinstance(reader_and_node, list): + return all( + self._use_same_gpu(writer_and_node, entry) for entry in reader_and_node + ) + if writer_and_node[1] != reader_and_node[1]: + return False + writer_gpu_ids = self._get_gpu_ids(writer_and_node[0]) + reader_gpu_ids = self._get_gpu_ids(reader_and_node[0]) + return writer_gpu_ids == reader_gpu_ids + + def _use_gpu( + self, actors: Union["ray.actor.ActorHandle", List["ray.actor.ActorHandle"]] + ) -> bool: + """ + Check if the actors use GPUs. + + Args: + actors: An actor handle or a list of actor handles. + + Returns: + True if the actors use GPUs, False otherwise. + """ + if isinstance(actors, list): + return all(self._use_gpu(actor) for actor in actors) + gpu_ids = self._get_gpu_ids(actors) + return len(gpu_ids) > 0 + + def resolve( + self, + auto_transport_type: AutoTransportType, + writer_and_node: Tuple[Optional["ray.actor.ActorHandle"], str], + reader_and_node_list: List[Tuple[Optional["ray.actor.ActorHandle"], str]], + ) -> "ChannelOutputType": + """ + Resolve auto_transport_type to the actual channel output type + based on the node locations and GPU IDs. + + Args: + auto_transport_type: The type to resolve + writer_and_node: A tuple of writer actor handle and its node ID. + A None writer actor handle means the writer is the driver. + reader_and_node_list: A list of tuples of reader actor handle and its + node ID. A None reader actor handle means the reader is the driver. + + Returns: + The actual channel type. + """ + writer = writer_and_node[0] + readers = [reader for reader, _ in reader_and_node_list] + + if writer is None or any(reader is None for reader in readers): + # None means actor is the driver, currently driver on GPU + # is not supported, so we always use shared memory to transfer + # tensors. + return TorchTensorType( + device=auto_transport_type.device, + _static_shape=auto_transport_type._static_shape, + _direct_return=auto_transport_type._direct_return, + ) + + # Case 1: writer and readers don't both use GPU, use shared memory + # to transport the tensors + if not (self._use_gpu(writer) and self._use_gpu(readers)): + return TorchTensorType( + device=auto_transport_type.device, + _static_shape=auto_transport_type._static_shape, + _direct_return=auto_transport_type._direct_return, + ) + + # Case 2: writer and readers use the same GPU are are on the same node, + # use shared memory to transport the tensors + if self._use_same_gpu(writer_and_node, reader_and_node_list): + return TorchTensorType( + device=auto_transport_type.device, + _static_shape=auto_transport_type._static_shape, + _direct_return=auto_transport_type._direct_return, + ) + + # Case 3: writer and readers use different GPUs, use accelerator to transport + # the tensors + return TorchTensorType( + transport="accelerator", + device=auto_transport_type.device, + _static_shape=auto_transport_type._static_shape, + _direct_return=auto_transport_type._direct_return, + ) diff --git a/lib/python3.12/site-packages/ray/experimental/channel/cached_channel.py b/lib/python3.12/site-packages/ray/experimental/channel/cached_channel.py new file mode 100644 index 0000000000000000000000000000000000000000..3417ddba9d834ded3284bd877fb13c5fe3f326d9 --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/channel/cached_channel.py @@ -0,0 +1,111 @@ +import uuid +from typing import Any, Optional + +from ray.experimental.channel.common import ChannelInterface + + +class CachedChannel(ChannelInterface): + """ + CachedChannel wraps an inner channel and caches the data read from it until + `num_reads` reads have completed. If inner channel is None, the data + is written to serialization context and retrieved from there. This is useful + when passing data within the same actor and a shared memory channel can be + avoided. + + Args: + num_reads: The number of reads from this channel that must happen before + writing again. Readers must be methods of the same actor. + inner_channel: The inner channel to cache data from. If None, the data is + read from the serialization context. + _channel_id: The unique ID for the channel. If None, a new ID is generated. + """ + + def __init__( + self, + num_reads: int, + inner_channel: Optional[ChannelInterface] = None, + _channel_id: Optional[str] = None, + ): + assert num_reads > 0, "num_reads must be greater than 0." + self._num_reads = num_reads + self._inner_channel = inner_channel + # Generate a unique ID for the channel. The writer and reader will use + # this ID to store and retrieve data from the _SerializationContext. + self._channel_id = _channel_id + if self._channel_id is None: + self._channel_id = str(uuid.uuid4()) + + def ensure_registered_as_writer(self) -> None: + if self._inner_channel is not None: + self._inner_channel.ensure_registered_as_writer() + + def ensure_registered_as_reader(self) -> None: + if self._inner_channel is not None: + self._inner_channel.ensure_registered_as_reader() + + def __reduce__(self): + return CachedChannel, ( + self._num_reads, + self._inner_channel, + self._channel_id, + ) + + def __str__(self) -> str: + return ( + f"CachedChannel(channel_id={self._channel_id}, " + f"num_reads={self._num_reads}), " + f"inner_channel={self._inner_channel})" + ) + + def write(self, value: Any, timeout: Optional[float] = None): + self.ensure_registered_as_writer() + # TODO: better organize the imports + from ray.experimental.channel import ChannelContext + + if self._inner_channel is not None: + self._inner_channel.write(value, timeout) + return + + # Otherwise no need to check timeout as the operation is non-blocking. + + # Because both the reader and writer are in the same worker process, + # we can directly store the data in the context instead of storing + # it in the channel object. This removes the serialization overhead of `value`. + ctx = ChannelContext.get_current().serialization_context + ctx.set_data(self._channel_id, value, self._num_reads) + + def read(self, timeout: Optional[float] = None) -> Any: + self.ensure_registered_as_reader() + # TODO: better organize the imports + from ray.experimental.channel import ChannelContext + + ctx = ChannelContext.get_current().serialization_context + if ctx.has_data(self._channel_id): + # No need to check timeout as the operation is non-blocking. + return ctx.get_data(self._channel_id) + + assert ( + self._inner_channel is not None + ), "Cannot read from the serialization context while inner channel is None." + value = self._inner_channel.read(timeout) + ctx.set_data(self._channel_id, value, self._num_reads) + # NOTE: Currently we make a contract with Compiled Graph users that the + # channel results should not be mutated by the actor methods. + # When the user needs to modify the channel results, they should + # make a copy of the channel results and modify the copy. + # This is the same contract as used in IntraProcessChannel. + # This contract is NOT enforced right now in either case. + # TODO(rui): introduce a flag to control the behavior: + # for example, by default we make a deep copy of the channel + # result, but the user can turn off the deep copy for performance + # improvements. + # https://github.com/ray-project/ray/issues/47409 + return ctx.get_data(self._channel_id) + + def close(self) -> None: + from ray.experimental.channel import ChannelContext + + if self._inner_channel is not None: + self._inner_channel.close() + ctx = ChannelContext.get_current().serialization_context + ctx.reset_data(self._channel_id) diff --git a/lib/python3.12/site-packages/ray/experimental/channel/common.py b/lib/python3.12/site-packages/ray/experimental/channel/common.py new file mode 100644 index 0000000000000000000000000000000000000000..0f1b916a7224d9cda4de9ed5c8d45f1da76e2322 --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/channel/common.py @@ -0,0 +1,681 @@ +import asyncio +import concurrent +import sys +import threading +import time +from dataclasses import dataclass +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + List, + NamedTuple, + Optional, + Tuple, + Union, +) + +import ray +import ray.exceptions +from ray.experimental.channel.accelerator_context import AcceleratorContext +from ray.experimental.channel.communicator import Communicator +from ray.experimental.channel.communicator_handle import CommunicatorHandle +from ray.experimental.channel.serialization_context import _SerializationContext +from ray.util.annotations import DeveloperAPI, PublicAPI + +# The context singleton on this process. +_default_context: "Optional[ChannelContext]" = None +_context_lock = threading.Lock() + +if TYPE_CHECKING: + import torch + + +def retry_and_check_interpreter_exit(f: Callable[[], None]) -> bool: + """This function is only useful when f contains channel read/write. + + Keep retrying channel read/write inside `f` and check if interpreter exits. + It is important in case the read/write happens in a separate thread pool. + See https://github.com/ray-project/ray/pull/47702 + + f should a function that doesn't receive any input and return nothing. + """ + exiting = False + while True: + try: + f() + break + except ray.exceptions.RayChannelTimeoutError: + if sys.is_finalizing(): + # Interpreter exits. We should ignore the error and + # stop reading so that the thread can join. + exiting = True + break + + return exiting + + +# Holds the input arguments for Compiled Graph +@PublicAPI(stability="alpha") +class CompiledDAGArgs(NamedTuple): + args: Tuple[Any, ...] + kwargs: Dict[str, Any] + + +@PublicAPI(stability="alpha") +class ChannelOutputType: + def register_custom_serializer(self) -> None: + """ + Register any custom serializers needed to pass data of this type. This + method should be run on the reader(s) and writer of a channel, which + are the driver and/or Ray actors. + + NOTE: When custom serializers are registered with Ray, the registered + deserializer is shipped with the serialized value and used on the + receiving end. Therefore, the deserializer function should *not* + capture state that is meant to be worker-local, such as the worker's + default device. Instead, these should be extracted from the + worker-local _SerializationContext. + """ + pass + + def create_channel( + self, + writer: Optional["ray.actor.ActorHandle"], + reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]], + driver_actor_id: Optional[str] = None, + ) -> "ChannelInterface": + """ + Instantiate a ChannelInterface class that can be used + to pass data of this type. + + Args: + writer: The actor that may write to the channel. None signifies the driver. + reader_and_node_list: A list of tuples, where each tuple contains a reader + actor handle and the node ID where the actor is located. + driver_actor_id: If this is a CompositeChannel that is read by a driver and + that driver is an actual actor, this will be the actor ID of that + driver actor. + Returns: + A ChannelInterface that can be used to pass data + of this type. + """ + raise NotImplementedError + + def requires_accelerator(self) -> bool: + # By default, channels do not require accelerator. + return False + + def get_custom_communicator(self) -> Optional[Communicator]: + """ + Return the custom communicator group if one is specified. + """ + return None + + def set_communicator_id(self, group_id: str) -> None: + raise NotImplementedError + + +@DeveloperAPI +@dataclass +class ChannelContext: + serialization_context = _SerializationContext() + _torch_available: Optional[bool] = None + _torch_device: Optional["torch.device"] = None + _current_stream: Optional["torch.cuda.Stream"] = None + + def __init__(self): + # Used for the torch.Tensor accelerator transport. + self.communicators: Dict[str, "Communicator"] = {} + # Used for driver process to store actors in the communicator. + self.communicator_handles: Dict[str, "CommunicatorHandle"] = {} + + @staticmethod + def get_current() -> "ChannelContext": + """Get or create a singleton context. + + If the context has not yet been created in this process, it will be + initialized with default settings. + """ + + global _default_context + + with _context_lock: + if _default_context is None: + _default_context = ChannelContext() + + return _default_context + + @property + def torch_available(self) -> bool: + """ + Check if torch package is available. + """ + if self._torch_available is not None: + return self._torch_available + + try: + import torch # noqa: F401 + except ImportError: + self._torch_available = False + return False + self._torch_available = True + return True + + @property + def torch_device(self) -> "torch.device": + if self._torch_device is None: + self._torch_device = AcceleratorContext.get().get_accelerator_devices()[0] + + return self._torch_device + + def set_torch_device(self, device: "torch.device"): + self._torch_device = device + + +@PublicAPI(stability="alpha") +class ChannelInterface: + """ + Abstraction for a transport between a writer actor and some number of + reader actors. + """ + + def __init__( + self, + writer: Optional[ray.actor.ActorHandle], + readers: List[Optional[ray.actor.ActorHandle]], + typ: Optional["ChannelOutputType"], + ): + """ + Create a channel that can be read and written by a Ray driver or actor. + + Args: + writer: The actor that may write to the channel. None signifies the driver. + readers: The actors that may read from the channel. None signifies + the driver. + typ: Type information about the values passed through the channel. + """ + pass + + def ensure_registered_as_writer(self): + """ + Check whether the process is a valid writer. This method must be idempotent. + """ + raise NotImplementedError + + def ensure_registered_as_reader(self): + """ + Check whether the process is a valid reader. This method must be idempotent. + """ + raise NotImplementedError + + def write(self, value: Any, timeout: Optional[float] = None) -> None: + """ + Write a value to the channel. + + Blocks if there are still pending readers for the previous value. The + writer may not write again until the specified number of readers have + read the value. + + Args: + value: The value to write. + timeout: The maximum time in seconds to wait to write the value. + None means using default timeout, 0 means immediate timeout + (immediate success or timeout without blocking), -1 means + infinite timeout (block indefinitely). + """ + raise NotImplementedError + + def read(self, timeout: Optional[float] = None) -> Any: + """ + Read the latest value from the channel. This call will block until a + value is available to read. + + Subsequent calls to read() may *block* if the deserialized object is + zero-copy (e.g., bytes or a numpy array) *and* the object is still in scope. + + Args: + timeout: The maximum time in seconds to wait to read the value. + None means using default timeout, 0 means immediate timeout + (immediate success or timeout without blocking), -1 means + infinite timeout (block indefinitely). + + Returns: + Any: The deserialized value. If the deserialized value is an + Exception, it will be returned directly instead of being raised. + """ + raise NotImplementedError + + def close(self) -> None: + """ + Close this channel. This method must not block and it must be made + idempotent. Any existing values in the channel may be lost after the + channel is closed. + """ + raise NotImplementedError + + +# Interfaces for channel I/O. +@DeveloperAPI +class ReaderInterface: + def __init__( + self, + input_channels: List[ChannelInterface], + ): + assert isinstance(input_channels, list) + for chan in input_channels: + assert isinstance(chan, ChannelInterface) + + self._input_channels = input_channels + self._closed = False + self._num_reads = 0 + + # A list of channels that were not read in the last `read` call + # because the reader returned immediately when a RayTaskError was found. + # These channels must be consumed before the next read to avoid reading + # stale data remaining from the last read. + self._leftover_channels: List[ChannelInterface] = [] + + def get_num_reads(self) -> int: + return self._num_reads + + def start(self): + raise NotImplementedError + + def _read_list(self, timeout: Optional[float] = None) -> List[Any]: + """ + Read a list of values from this reader. + + Args: + timeout: The maximum time in seconds to wait for reading. + None means using default timeout which is infinite, 0 means immediate + timeout (immediate success or timeout without blocking), -1 means + infinite timeout (block indefinitely). + + """ + raise NotImplementedError + + def read(self, timeout: Optional[float] = None) -> List[Any]: + """ + Read from this reader. + + Args: + timeout: The maximum time in seconds to wait for reading. + None means using default timeout, 0 means immediate timeout + (immediate success or timeout without blocking), -1 means + infinite timeout (block indefinitely). + """ + assert ( + timeout is None or timeout >= 0 or timeout == -1 + ), "Timeout must be non-negative or -1." + outputs = self._read_list(timeout) + self._num_reads += 1 + return outputs + + def close(self) -> None: + self._closed = True + for channel in self._input_channels: + channel.close() + + def _consume_leftover_channels_if_needed( + self, timeout: Optional[float] = None + ) -> None: + # Consume the channels that were not read in the last `read` call because a + # RayTaskError was returned from another channel. If we don't do this, the + # read operation will read stale versions of the object refs. + # + # If a RayTaskError is returned from a leftover channel, it will be ignored. + # If a read operation times out, a RayChannelTimeoutError exception will be + # raised. + # + # TODO(kevin85421): Currently, a DAG with NCCL channels and fast fail enabled + # may not be reusable. Revisit this in the future. + for c in self._leftover_channels: + start_time = time.monotonic() + c.read(timeout) + if timeout is not None: + timeout -= time.monotonic() - start_time + timeout = max(timeout, 0) + self._leftover_channels = [] + + +@DeveloperAPI +class SynchronousReader(ReaderInterface): + def __init__( + self, + input_channels: List[ChannelInterface], + ): + super().__init__(input_channels) + + def start(self): + pass + + def _read_list(self, timeout: Optional[float] = None) -> List[Any]: + self._consume_leftover_channels_if_needed(timeout) + # We don't update `remaining_timeout` here because in the worst case, + # consuming leftover channels requires reading all `_input_channels`, + # which users expect to complete within the original `timeout`. Updating + # `remaining_timeout` could cause unexpected timeouts in subsequent read + # operations. + + # It is a special case that `timeout` is set to 0, which means + # read once for each channel. + is_zero_timeout = timeout == 0 + + results = [None for _ in range(len(self._input_channels))] + if timeout is None or timeout == -1: + timeout = float("inf") + timeout_point = time.monotonic() + timeout + remaining_timeout = timeout + + from ray.dag import DAGContext + + ctx = DAGContext.get_current() + iteration_timeout = ctx.read_iteration_timeout + + # Iterate over the input channels with a shorter timeout for each iteration + # to detect RayTaskError early and fail fast. + done_channels = set() + while len(done_channels) < len(self._input_channels): + for i, c in enumerate(self._input_channels): + if c in done_channels: + continue + try: + result = c.read(min(remaining_timeout, iteration_timeout)) + results[i] = result + done_channels.add(c) + if isinstance(result, ray.exceptions.RayTaskError): + # If we raise an exception immediately, it will be considered + # as a system error which will cause the execution loop to + # exit. Hence, return immediately and let `_process_return_vals` + # handle the exception. + # + # Return a list of RayTaskError so that the caller will not + # get an undefined partial result. + self._leftover_channels = [ + c for c in self._input_channels if c not in done_channels + ] + return [result for _ in range(len(self._input_channels))] + except ray.exceptions.RayChannelTimeoutError as e: + remaining_timeout = max(timeout_point - time.monotonic(), 0) + if remaining_timeout == 0: + raise e + continue + + remaining_timeout = max(timeout_point - time.monotonic(), 0) + if remaining_timeout == 0 and not is_zero_timeout: + raise ray.exceptions.RayChannelTimeoutError( + f"Cannot read all channels within {timeout} seconds" + ) + return results + + def release_channel_buffers(self, timeout: Optional[float] = None) -> None: + for c in self._input_channels: + start_time = time.monotonic() + assert hasattr( + c, "release_buffer" + ), "release_buffer() is only supported for shared memory channel " + "(e.g., Channel, BufferedSharedMemoryChannel, CompositeChannel) " + "and used between the last actor and the driver, but got a channel" + f" of type {type(c)}." + c.release_buffer(timeout) + if timeout is not None: + timeout -= time.monotonic() - start_time + timeout = max(timeout, 0) + + +@DeveloperAPI +class AwaitableBackgroundReader(ReaderInterface): + """ + Asyncio-compatible channel reader. + + The reader is constructed with an async queue of futures whose values it + will fulfill. It uses a threadpool to execute the blocking calls to read + from the input channel(s). + """ + + def __init__( + self, + input_channels: List[ChannelInterface], + fut_queue: asyncio.Queue, + ): + super().__init__(input_channels) + self._fut_queue = fut_queue + self._background_task = None + self._background_task_executor = concurrent.futures.ThreadPoolExecutor( + max_workers=1, thread_name_prefix="channel.AwaitableBackgroundReader" + ) + + def start(self): + self._background_task = asyncio.ensure_future(self.run()) + + def _run(self): + # Give it a default timeout 60 seconds to release the buffers + # of the channels that were not read in the last `read` call. + self._consume_leftover_channels_if_needed(60) + + results = [None for _ in range(len(self._input_channels))] + + from ray.dag import DAGContext + + ctx = DAGContext.get_current() + iteration_timeout = ctx.read_iteration_timeout + + done_channels = set() + while len(done_channels) < len(self._input_channels): + for i, c in enumerate(self._input_channels): + if c in done_channels: + continue + try: + result = c.read(iteration_timeout) + results[i] = result + done_channels.add(c) + if isinstance(result, ray.exceptions.RayTaskError): + self._leftover_channels = [ + c for c in self._input_channels if c not in done_channels + ] + return [result for _ in range(len(self._input_channels))] + except ray.exceptions.RayChannelTimeoutError: + pass + if sys.is_finalizing(): + return results + return results + + async def run(self): + loop = asyncio.get_running_loop() + while not self._closed: + res, fut = await asyncio.gather( + loop.run_in_executor(self._background_task_executor, self._run), + self._fut_queue.get(), + return_exceptions=True, + ) + + # Set the result on the main thread. + fut.set_result(res) + # NOTE(swang): If the object is zero-copy deserialized, then it + # will stay in scope as long as ret and the future are in scope. + # Therefore, we must delete both here after fulfilling the future. + del res + del fut + + def close(self): + super().close() + self._background_task_executor.shutdown(cancel_futures=True) + self._background_task.cancel() + + +@DeveloperAPI +class WriterInterface: + def __init__( + self, + output_channels: List[ChannelInterface], + output_idxs: List[Optional[Union[int, str]]], + is_input=False, + ): + """ + Initialize the writer. + + Args: + output_channels: The output channels to write to. + output_idxs: The indices of the values to write to each channel. + This has the same length as `output_channels`. If `is_input` is True, + the index can be an integer or a string to retrieve the corresponding + value from `args` or `kwargs` in the DAG's input. If `is_input` + is False, the entire value is written if the index is None. Otherwise, + the value at the specified index in the tuple is written. + is_input: Whether the writer is DAG input writer or not. + """ + + assert len(output_channels) == len(output_idxs) + self._output_channels = output_channels + self._output_idxs = output_idxs + self._closed = False + self._num_writes = 0 + self._is_input = is_input + + def get_num_writes(self) -> int: + return self._num_writes + + def start(self): + raise NotImplementedError() + + def write(self, val: Any, timeout: Optional[float] = None) -> None: + """ + Write the value. + + Args: + timeout: The maximum time in seconds to wait for writing. 0 means + immediate timeout (immediate success or timeout without blocking). + -1 and None mean infinite timeout (blocks indefinitely). + """ + raise NotImplementedError() + + def close(self) -> None: + self._closed = True + for channel in self._output_channels: + channel.close() + + +def _adapt(raw_args: Any, key: Optional[Union[int, str]], is_input: bool): + """ + Adapt the raw arguments to the key. If `is_input` is True, this method will + retrieve the value from the input data for an InputAttributeNode. Otherwise, it + will retrieve either a partial value or the entire value from the output of + a ClassMethodNode. + + Args: + raw_args: The raw arguments to adapt. + key: The key to adapt. + is_input: Whether the writer is DAG input writer or not. + """ + if is_input: + if not isinstance(raw_args, CompiledDAGArgs): + # Fast path for a single input. + return raw_args + else: + args = raw_args.args + kwargs = raw_args.kwargs + + if isinstance(key, int): + return args[key] + else: + return kwargs[key] + else: + if key is not None: + return raw_args[key] + else: + return raw_args + + +@DeveloperAPI +class SynchronousWriter(WriterInterface): + def start(self): + for channel in self._output_channels: + channel.ensure_registered_as_writer() + + def write(self, val: Any, timeout: Optional[float] = None) -> None: + # If it is an exception, there's only 1 return value. + # We have to send the same data to all channels. + if isinstance(val, Exception): + if len(self._output_channels) > 1: + val = tuple(val for _ in range(len(self._output_channels))) + + if not self._is_input: + if len(self._output_channels) > 1: + if not isinstance(val, tuple): + raise ValueError( + f"Expected a tuple of {len(self._output_channels)} outputs, " + f"but got {type(val)}" + ) + if len(val) != len(self._output_channels): + raise ValueError( + f"Expected {len(self._output_channels)} outputs, but got " + f"{len(val)} outputs" + ) + + for i, channel in enumerate(self._output_channels): + idx = self._output_idxs[i] + val_i = _adapt(val, idx, self._is_input) + channel.write(val_i, timeout) + self._num_writes += 1 + + +@DeveloperAPI +class AwaitableBackgroundWriter(WriterInterface): + def __init__( + self, + output_channels: List[ChannelInterface], + output_idxs: List[Optional[Union[int, str]]], + is_input=False, + ): + super().__init__(output_channels, output_idxs, is_input=is_input) + self._queue = asyncio.Queue() + self._background_task = None + self._background_task_executor = concurrent.futures.ThreadPoolExecutor( + max_workers=1, thread_name_prefix="channel.AwaitableBackgroundWriter" + ) + + def start(self): + for channel in self._output_channels: + channel.ensure_registered_as_writer() + self._background_task = asyncio.ensure_future(self.run()) + + def _run(self, res): + if not self._is_input: + if len(self._output_channels) > 1: + if not isinstance(res, tuple): + raise ValueError( + f"Expected a tuple of {len(self._output_channels)} outputs, " + f"but got {type(res)}" + ) + if len(res) != len(self._output_channels): + raise ValueError( + f"Expected {len(self._output_channels)} outputs, but got " + f"{len(res)} outputs" + ) + + for i, channel in enumerate(self._output_channels): + idx = self._output_idxs[i] + res_i = _adapt(res, idx, self._is_input) + exiting = retry_and_check_interpreter_exit( + lambda: channel.write(res_i, timeout=1) + ) + if exiting: + break + + async def run(self): + loop = asyncio.get_event_loop() + while True: + res = await self._queue.get() + await loop.run_in_executor(self._background_task_executor, self._run, res) + + async def write(self, val: Any) -> None: + if self._closed: + raise RuntimeError("DAG execution cancelled") + await self._queue.put(val) + self._num_writes += 1 + + def close(self): + self._background_task.cancel() + super().close() diff --git a/lib/python3.12/site-packages/ray/experimental/channel/communicator.py b/lib/python3.12/site-packages/ray/experimental/channel/communicator.py new file mode 100644 index 0000000000000000000000000000000000000000..69113055f6f708fe895b384cfbf6a9fc8597b29f --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/channel/communicator.py @@ -0,0 +1,199 @@ +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Callable, List, Optional, Tuple + +import ray +from ray.experimental.util.types import ReduceOp +from ray.util.annotations import DeveloperAPI + +if TYPE_CHECKING: + import torch + + +# Signature for a torch.Tensor allocator is: +# (shape: Tuple[int], dtype: torch.dtype) -> torch.Tensor. +TorchTensorAllocator = Callable[[Tuple[int], "torch.dtype"], "torch.Tensor"] + + +@DeveloperAPI +class Communicator(ABC): + """ + Communicator for a group of Compiled Graph actors on NVIDIA GPU. + + The Compiled Graph execution leverages this internally to support communication + between actors in the group. + """ + + @abstractmethod + def initialize(self, rank: int) -> None: + """ + Initialize the communicator from the actor. + + This is called once by Compiled Graph on each actor to initialize the + communicator,before any other methods. + + Args: + rank: The rank of this actor in the group. + """ + raise NotImplementedError + + @abstractmethod + def get_actor_handles(self) -> List["ray.actor.ActorHandle"]: + """ + Get handles of all actors for this communicator group. + """ + raise NotImplementedError + + @abstractmethod + def get_rank(self, actor: ray.actor.ActorHandle) -> int: + """ + Return the given actor's rank in the group. + + Args: + actor: The actor handle to look up. + """ + raise NotImplementedError + + @abstractmethod + def get_self_rank(self) -> Optional[int]: + """ + Return this actor's rank. + """ + raise NotImplementedError + + def get_world_size(self) -> int: + """ + Return the number of ranks in the group. + """ + raise NotImplementedError + + @abstractmethod + def send(self, value: "torch.Tensor", peer_rank: int) -> None: + """ + Send a torch.Tensor to a peer. + + This returns when the send kernel has been queued, but the kernel may + not have completed. Therefore, the caller should ensure that there are + no concurrent writes to the sent `value` until the send has finished. + + Args: + value: The torch.Tensor to send. It should already be on this + actor's default device. + peer_rank: The rank of the actor to send to. + """ + raise NotImplementedError + + @abstractmethod + def recv( + self, + shape: Tuple[int], + dtype: "torch.dtype", + peer_rank: int, + allocator: Optional[TorchTensorAllocator] = None, + ) -> "torch.Tensor": + """ + Receive a torch.Tensor from a peer and synchronize. + + After this call returns, the receive buffer is safe to read from from + any stream. An RayChannelError will be raised if an error occurred (e.g., + remote actor died), and the buffer is not safe to read. + + Args: + shape: The shape of the tensor to receive. + dtype: The dtype of the tensor to receive. + peer_rank: The rank of the actor to receive from. + allocator: A function to allocate the tensor to receive into. + """ + raise NotImplementedError + + @property + @abstractmethod + def recv_stream(self): + """ + Return the torch stream context used for receiving tensors. + """ + raise NotImplementedError + + @property + @abstractmethod + def send_stream(self): + """ + Return the torch stream context used for sending tensors. + """ + raise NotImplementedError + + @abstractmethod + def allgather( + self, + send_buf: "torch.Tensor", + recv_buf: "torch.Tensor", + ) -> None: + """ + Collectively allgather the tensor across the group. + + Args: + send_buf: The input torch.tensor to allgather. It should already be + on this actor's default device. + recv_buf: The output torch.tensor to store the allgather result. + """ + raise NotImplementedError + + @abstractmethod + def allreduce( + self, + send_buf: "torch.Tensor", + recv_buf: "torch.Tensor", + op: ReduceOp, + ) -> None: + """ + Collectively allreduce the tensor across the group. + + Args: + send_buf: The input torch.tensor to allreduce. It should already be + on this actor's default device. + recv_buf: The output torch.tensor to store the allreduce result. + op: The reduce operation. + """ + raise NotImplementedError + + @abstractmethod + def reducescatter( + self, + send_buf: "torch.Tensor", + recv_buf: "torch.Tensor", + op: ReduceOp, + ) -> None: + """ + Collectively reducescatter the tensor across the group. + + Args: + send_buf: The input torch.tensor to reducescatter. It should already be + on this actor's default device. + recv_buf: The output torch.tensor to store the reducescatter result. + op: The reduce operation. + """ + raise NotImplementedError + + @abstractmethod + def destroy(self) -> None: + """ + Destroy the GPU communicator. + + Any destruction and cleanup for the GPU communicator should be + done here. Implement as a noop is nothing is needed. + """ + raise NotImplementedError + + @abstractmethod + def get_transport_name(self) -> str: + """ + Return the type of the communicator (gpu or cpu). + """ + raise NotImplementedError + + @classmethod + @abstractmethod + def generate_communicator_id(cls) -> str: + """ + Return the unique id of the communicator. + """ + raise NotImplementedError diff --git a/lib/python3.12/site-packages/ray/experimental/channel/communicator_handle.py b/lib/python3.12/site-packages/ray/experimental/channel/communicator_handle.py new file mode 100644 index 0000000000000000000000000000000000000000..c6d8865bfc44f1c6dc9da2e2444a4d5acba6d0c7 --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/channel/communicator_handle.py @@ -0,0 +1,28 @@ +from typing import List + +import ray + + +class CommunicatorHandle: + """ + A lightweight communicator handle used by the driver to store handles to + the actors in the communicator. + """ + + def __init__( + self, + actor_handles: List["ray.actor.ActorHandle"], + ): + """ + Initializes the CommunicatorHandle with the given actor handles. + + Args: + actor_handles: A list of actor handles to be stored. + """ + self._actor_handles = actor_handles + + def get_actor_handles(self) -> List["ray.actor.ActorHandle"]: + """ + Retuan all actor handles in this communicator. + """ + return self._actor_handles diff --git a/lib/python3.12/site-packages/ray/experimental/channel/conftest.py b/lib/python3.12/site-packages/ray/experimental/channel/conftest.py new file mode 100644 index 0000000000000000000000000000000000000000..d735a9f0ec69ad5cef47d169212ef9fdcec58bcf --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/channel/conftest.py @@ -0,0 +1,185 @@ +import asyncio +from collections import defaultdict +from typing import Optional, Tuple +from unittest import mock + +import torch + +import ray +import ray.dag +import ray.experimental.channel as ray_channel +from ray.experimental.channel import nccl_group +from ray.experimental.channel.communicator import TorchTensorAllocator +from ray.experimental.util.types import Device + + +@ray.remote(num_cpus=0) +class Barrier: + """ + Barrier that blocks the given number of actors until all actors have + reached the barrier. This is used to mock out blocking NCCL ops. + """ + + def __init__(self, num_actors=2): + self.num_actors = num_actors + self.condition = asyncio.Condition() + # Buffer for the data that is "sent" between the actors, each entry is + # one p2p op. + self.data = {} + # Buffer for the number of actors seen, each entry is one p2p op. + self.num_actors_seen = defaultdict(int) + + # Add a new mock for the TorchTensorType.device property + device_property_patcher = mock.patch( + "ray.experimental.channel.torch_tensor_type.TorchTensorType.device", + new_callable=mock.PropertyMock, + return_value=Device.CPU, + ) + device_property_patcher.start() + + async def wait(self, idx: int, data=None): + """ + Wait at barrier until all actors have sent `idx`. One actor should + provide `data`, and this value will be returned by this method for all + other actors. + """ + async with self.condition: + if data is not None: + assert idx not in self.data, (self.data, self.num_actors_seen) + self.data[idx] = data + self.num_actors_seen[idx] += 1 + + if self.num_actors_seen[idx] == self.num_actors: + # Wake up all tasks waiting on this condition. + self.condition.notify_all() + else: + await self.condition.wait_for( + lambda: self.num_actors_seen[idx] == self.num_actors + ) + + if data is None: + data = self.data[idx] + + return data + + +class MockCudaStream: + def __init__(self): + self.cuda_stream = 0 + + def synchronize(self): + pass + + +class MockNcclGroup(nccl_group._NcclGroup): + """ + Mock the internal _NcclGroup to use a barrier actor instead of a NCCL group + for communication. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # We use the op index to synchronize the sender and receiver at the + # barrier. + self.num_ops = defaultdict(int) + self.barriers = set() + + def send(self, tensor: torch.Tensor, peer_rank: int): + # "Send" the tensor to the barrier actor. + barrier_key = sorted([self.get_self_rank(), peer_rank]) + barrier_key = f"barrier-{barrier_key[0]}-{barrier_key[1]}" + barrier = ray.get_actor(name=barrier_key) + self.barriers.add(barrier) + ray.get(barrier.wait.remote(self.num_ops[barrier_key], tensor)) + self.num_ops[barrier_key] += 1 + + def recv( + self, + shape: Tuple[int], + dtype: torch.dtype, + peer_rank: int, + allocator: Optional[TorchTensorAllocator] = None, + ): + # "Receive" the tensor from the barrier actor. + barrier_key = sorted([self.get_self_rank(), peer_rank]) + barrier_key = f"barrier-{barrier_key[0]}-{barrier_key[1]}" + barrier = ray.get_actor(name=barrier_key) + self.barriers.add(barrier) + received_tensor = ray.get(barrier.wait.remote(self.num_ops[barrier_key])) + assert ( + allocator is not None + ), "torch tensor allocator is required for MockNcclGroup" + buf = allocator(shape, dtype) + buf[:] = received_tensor[:] + self.num_ops[barrier_key] += 1 + return buf + + def destroy(self) -> None: + for barrier in self.barriers: + ray.kill(barrier) + + +def start_nccl_mock(): + """ + Patch methods that require CUDA. + """ + # Mock cupy dependencies. + nccl_mock = mock.MagicMock() + nccl_mock.nccl.get_unique_id.return_value = 0 + cp_patcher = mock.patch.dict( + "sys.modules", + { + "cupy.cuda": nccl_mock, + "cupy": mock.MagicMock(), + "ray.util.collective.collective_group": mock.MagicMock(), + }, + ) + cp_patcher.start() + + # Mock send/recv ops to use an actor instead of NCCL. + ray.experimental.channel.nccl_group._NcclGroup = MockNcclGroup + + # PyTorch mocks. + stream_patcher = mock.patch( + "torch.cuda.current_stream", new_callable=lambda: MockCudaStream + ) + stream_patcher.start() + new_stream_patcher = mock.patch( + "torch.cuda.Stream", new_callable=lambda: MockCudaStream + ) + new_stream_patcher.start() + tensor_patcher = mock.patch("torch.Tensor.device", torch.device("cuda")) + tensor_patcher.start() + tensor_patcher = mock.patch("torch.Tensor.is_cuda", True) + tensor_patcher.start() + tensor_allocator_patcher = mock.patch( + "ray.experimental.channel.torch_tensor_accelerator_channel._torch_tensor_allocator", + lambda shape, dtype: torch.empty(shape, dtype=dtype), + ) + tensor_allocator_patcher.start() + + # Add a new mock for the TorchTensorType.device property + device_property_patcher = mock.patch( + "ray.experimental.channel.torch_tensor_type.TorchTensorType.device", + new_callable=mock.PropertyMock, + return_value=Device.CPU, + ) + device_property_patcher.start() + + ctx = ray_channel.ChannelContext.get_current() + ctx.set_torch_device(torch.device("cuda")) + + +class TracedChannel(ray_channel.shared_memory_channel.Channel): + """ + Patched Channel that records all write ops for testing. + """ + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + self.ops = [] + + def write(self, *args, **kwargs): + self.ops.append((args, kwargs)) + return super().write(*args, **kwargs) diff --git a/lib/python3.12/site-packages/ray/experimental/channel/cpu_communicator.py b/lib/python3.12/site-packages/ray/experimental/channel/cpu_communicator.py new file mode 100644 index 0000000000000000000000000000000000000000..e0d7f199d7c723c7077e012eed7b9b1fc66ea51a --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/channel/cpu_communicator.py @@ -0,0 +1,207 @@ +import asyncio +from collections import defaultdict +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple + +import ray +from ray.experimental.channel.communicator import ( + Communicator, + ReduceOp, + TorchTensorAllocator, +) + +if TYPE_CHECKING: + import torch + + +@ray.remote(num_cpus=0) +class CPUCommBarrier: + """ + Barrier actor that blocks the given number of actors until all actors have + reached the Barrier. + + p2p operations are not done here (completed via shared memory channel). + """ + + def __init__(self, num_actors: int): + self.num_actors = num_actors + self.condition = asyncio.Condition() + # Stores the data for each collective operation + self.collective_data: Dict[int, List["torch.Tensor"]] = defaultdict(list) + # Stores the shape of data for each collective operation + self.collective_data_shape: Dict[int, "torch.Tensor.type"] = {} + # Buffer for the number of actors seen + self.num_actors_seen = defaultdict(int) + # Number of actors who have read the result, and are about to exit the function. + # State is kept so we only garbage collect after the last actor has read the + # relevant data. + self.num_actors_read = defaultdict(int) + + async def wait_collective(self, op_id: int, data: "torch.Tensor", op: ReduceOp): + """ + Wait at the communicator until all actors have sent `op_id` and `data`. + Once data from all actors is received, execute the collective `op` + on the communicator actor and return the result. + """ + async with self.condition: + self.collective_data[op_id].append(data) + self.num_actors_seen[op_id] += 1 + + if self.num_actors_seen[op_id] == self.num_actors: + # Apply the collective operation across all gathered tensors + data = self._apply_op(op, self.collective_data[op_id]) + self.collective_data[op_id] = data + self.condition.notify_all() + else: + await self.condition.wait_for( + lambda: self.num_actors_seen[op_id] == self.num_actors + ) + + data = self.collective_data[op_id] + self.num_actors_read[op_id] += 1 + + if self.num_actors_read[op_id] == self.num_actors: + del self.collective_data[op_id] + del self.num_actors_seen[op_id] + del self.num_actors_read[op_id] + + return data + + def _apply_op(self, op: ReduceOp, tensors: List["torch.Tensor"]) -> "torch.Tensor": + """Apply the specified reduction operation across a list of tensors.""" + + result = tensors[0].clone() + if op == ReduceOp.SUM: + for tensor in tensors[1:]: + result += tensor + elif op == ReduceOp.PRODUCT: + for tensor in tensors[1:]: + result *= tensor + elif op == ReduceOp.MAX: + for tensor in tensors[1:]: + result = torch.max(result, tensor) + elif op == ReduceOp.MIN: + for tensor in tensors[1:]: + result = torch.min(result, tensor) + elif op == ReduceOp.AVG: + result = sum(tensors) / len(tensors) + else: + raise ValueError(f"Operation {op} not supported") + return result + + +class CPUCommunicator(Communicator): + """ + Uses a CPU-based communicator actor instead of an accelerator group like NCCL. + """ + + def __init__(self, world_size: int, actor_handles: List["ray.actor.ActorHandle"]): + """We use the op index to synchronize the sender and receiver at the + communicator actor.""" + self._world_size = world_size + self._actor_handles = actor_handles + self.num_ops = defaultdict(int) + + # For collective communication, one barrier will be created for + # each unique group of participants. + self.barriers = set() + self._rank = None + + def send(self, tensor: "torch.Tensor", peer_rank: int): + # p2p operations are done via a shared memory channel, initialized in + # `create_channel` of `TorchTensorType` + pass + + def recv( + self, + shape: Tuple[int], + dtype: "torch.dtype", + peer_rank: int, + allocator: Optional[TorchTensorAllocator] = None, + ): + # See the comment on `send` + pass + + def allgather( + self, + send_buf: "torch.Tensor", + recv_buf: "torch.Tensor", + ): + raise NotImplementedError + + def allreduce( + self, + send_buf: "torch.Tensor", + recv_buf: "torch.Tensor", + op: ReduceOp = ReduceOp.SUM, + ): + all_ranks = [ + self.get_rank(actor_handle) for actor_handle in self.get_actor_handles() + ] + barrier_key = "barrier-collective-" + "-".join(map(str, sorted(all_ranks))) + barrier = CPUCommBarrier.options(name=barrier_key, get_if_exists=True).remote( + self._world_size + ) + self.barriers.add(barrier) + + result = ray.get( + barrier.wait_collective.remote(self.num_ops[barrier_key], send_buf, op) + ) + assert recv_buf is not None, "Receiving buffer required for CPUCommunicator" + recv_buf[:] = result[:] + self.num_ops[barrier_key] += 1 + + def reducescatter( + self, + send_buf: "torch.Tensor", + recv_buf: "torch.Tensor", + op: ReduceOp = ReduceOp.SUM, + ): + raise NotImplementedError + + def destroy(self) -> None: + for barrier in self.barriers: + ray.kill(barrier) + + def initialize(self, rank: int) -> None: + self._rank = rank + + def get_actor_handles(self) -> List["ray.actor.ActorHandle"]: + return self._actor_handles + + def get_rank(self, actor: ray.actor.ActorHandle) -> int: + """ + Return the given actor's rank in the CPU communicator. + + Args: + actor: The actor handle to look up. + """ + actor_ids = [a._ray_actor_id for a in self._actor_handles] + try: + rank = actor_ids.index(actor._ray_actor_id) + except ValueError: + raise ValueError("Actor is not in the CPUCommunicator group.") + return rank + + def get_self_rank(self) -> Optional[int]: + return self._rank + + def get_world_size(self) -> int: + """ + Return the number of ranks in the CPU communicator. + """ + return self._world_size + + def get_transport_name(self) -> str: + return "cpu" + + def recv_stream(self): + raise NotImplementedError + + def send_stream(self): + raise NotImplementedError + + @classmethod + def generate_communicator_id(cls) -> str: + import uuid + + return str(uuid.uuid4()) diff --git a/lib/python3.12/site-packages/ray/experimental/channel/intra_process_channel.py b/lib/python3.12/site-packages/ray/experimental/channel/intra_process_channel.py new file mode 100644 index 0000000000000000000000000000000000000000..11133e0ebe6081d0fd235951c3c299a2ea0b4966 --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/channel/intra_process_channel.py @@ -0,0 +1,70 @@ +import uuid +from typing import Any, Optional + +from ray.experimental.channel import ChannelContext +from ray.experimental.channel.common import ChannelInterface +from ray.util.annotations import PublicAPI + + +@PublicAPI(stability="alpha") +class IntraProcessChannel(ChannelInterface): + """ + IntraProcessChannel is a channel for communication between two tasks in the same + worker process. It writes data directly to the worker's _SerializationContext + and reads data from the _SerializationContext to avoid the serialization + overhead and the need for reading/writing from shared memory. Note that if the + readers may mutate the data, users should deep copy the data themselves to avoid + side effects. + + Args: + num_readers: The number of readers that will read from this channel. Readers + can be the same method of the same actor. + """ + + def __init__( + self, + num_readers, + _channel_id: Optional[str] = None, + ): + # Generate a unique ID for the channel. The writer and reader will use + # this ID to store and retrieve data from the _SerializationContext. + self._channel_id = _channel_id + self._num_readers = num_readers + if self._channel_id is None: + self._channel_id = str(uuid.uuid4()) + + def ensure_registered_as_writer(self) -> None: + pass + + def ensure_registered_as_reader(self) -> None: + pass + + def __reduce__(self): + return IntraProcessChannel, ( + self._num_readers, + self._channel_id, + ) + + def __str__(self) -> str: + return f"IntraProcessChannel(channel_id={self._channel_id})" + + def write(self, value: Any, timeout: Optional[float] = None): + self.ensure_registered_as_writer() + # No need to check timeout as the operation is non-blocking. + + # Because both the reader and writer are in the same worker process, + # we can directly store the data in the context instead of storing + # it in the channel object. This removes the serialization overhead of `value`. + ctx = ChannelContext.get_current().serialization_context + ctx.set_data(self._channel_id, value, self._num_readers) + + def read(self, timeout: Optional[float] = None, deserialize: bool = True) -> Any: + self.ensure_registered_as_reader() + assert deserialize, "Data passed from the actor to itself is never serialized" + # No need to check timeout as the operation is non-blocking. + ctx = ChannelContext.get_current().serialization_context + return ctx.get_data(self._channel_id) + + def close(self) -> None: + ctx = ChannelContext.get_current().serialization_context + ctx.reset_data(self._channel_id) diff --git a/lib/python3.12/site-packages/ray/experimental/channel/nccl_group.py b/lib/python3.12/site-packages/ray/experimental/channel/nccl_group.py new file mode 100644 index 0000000000000000000000000000000000000000..64b640818833a1511595264a00039f42e147ed40 --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/channel/nccl_group.py @@ -0,0 +1,374 @@ +import logging +from types import ModuleType +from typing import TYPE_CHECKING, Callable, List, Optional, Tuple + +import ray +from ray.exceptions import RayChannelError +from ray.experimental.channel.accelerator_context import AcceleratorContext +from ray.experimental.channel.communicator import Communicator, TorchTensorAllocator +from ray.experimental.util.types import ReduceOp + +if TYPE_CHECKING: + import torch + + +# Logger for this module. It should be configured at the entry point +# into the program using Ray. Ray provides a default configuration at +# entry/init points. +logger = logging.getLogger(__name__) + + +class _NcclGroup(Communicator): + """ + Represents an actor's NCCL communicator. This is the default NCCL communicator + to be used in Compiled Graph if a custom communicator is not provided. + + This class is not thread-safe. + """ + + def __init__( + self, + world_size: int, + comm_id: tuple, + rank: Optional[int], + actor_handles: List["ray.actor.ActorHandle"], + cuda_stream: Optional["torch.cuda.Stream"], + use_communication_streams: bool = False, + ): + """ + Initialize a NCCL communicator that can be used to communicate p2p with + other GPU actors. + + This method blocks until the same call has been made on all other + actors in the group, with the same arguments for world_size and + comm_id. + + NOTE: A concurrent NCCL group can coexist with this one but using the + two groups concurrently on different CUDA streams may cause deadlock. + See + https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/communicators.html + #using-multiple-nccl-communicators-concurrently. + + If the user can guarantee that all involved actors execute the same ops + in the same order, then the other NCCL group should use the given + `cuda_stream`, and there will not be a concurrency issue. Otherwise, + the other stream needs to synchronize with the given `cuda_stream` + before and after it launches NCCL ops, e.g., at the beginning and end + of a DAG task. + + Args: + world_size: The number of participating actors/devices. + comm_id: A unique communicator ID returned by + cupy.cuda.nccl.get_unique_id(). + rank: The rank of this actor. If None, then the caller is not a + participant of the NCCL group. + actor_handles: A list of actor handles, in rank order. + cuda_stream: A raw CUDA stream to dispatch NCCL ops to. If rank is + specified, then this must be specified too. + use_communication_streams: Whether to use dedicated send and recv + streams for communication. If True, communication and computation + can be overlapped to improve performance. + """ + self._world_size = world_size + self._rank: Optional[int] = rank + self.nccl_util: Optional[ModuleType] = None + self._actor_handles = actor_handles + self._use_communication_streams = use_communication_streams + + if rank is not None: + assert ray.get_gpu_ids(), "NCCL actor has no GPUs assigned" + assert cuda_stream is not None, "NCCL actor must specify cuda_stream" + + expected_rank = self.get_rank(ray.get_runtime_context().current_actor) + assert ( + rank == expected_rank + ), f"NCCL actor's rank {rank} does not match expected rank {expected_rank}" + + from ray.util.collective.collective_group import nccl_util + + self.nccl_util = nccl_util + self._comm = self.nccl_util.NcclCommunicator(world_size, comm_id, rank) + else: + # Driver does not have a rank. + self._comm = None + + self._cuda_stream: Optional["torch.cuda.Stream"] = None + self._send_stream: Optional["torch.cuda.Stream"] = None + self._recv_stream: Optional["torch.cuda.Stream"] = None + if cuda_stream is not None: + assert rank is not None, "NCCL actor has no rank assigned" + self._cuda_stream = cuda_stream + + if use_communication_streams: + import torch + + # TODO(swang): Allow default device to be overridden. + device = AcceleratorContext.get().get_accelerator_devices()[0] + + self._send_stream = torch.cuda.Stream(device=device) + self._recv_stream = torch.cuda.Stream(device=device) + else: + self._send_stream = self._cuda_stream + self._recv_stream = self._cuda_stream + + self._closed = False + + def initialize(self, rank: int) -> None: + # No additional initialization is needed. + pass + + def get_actor_handles(self) -> List["ray.actor.ActorHandle"]: + return self._actor_handles + + def get_rank(self, actor: ray.actor.ActorHandle) -> int: + """ + Return the given actor's rank in the NCCL communicator. + + Args: + actor: The actor handle to look up. + """ + actor_ids = [a._ray_actor_id for a in self._actor_handles] + try: + rank = actor_ids.index(actor._ray_actor_id) + except ValueError: + raise ValueError("Actor is not in the NCCL group.") + return rank + + def get_self_rank(self) -> Optional[int]: + """ + Return this actor's rank. + """ + return self._rank + + def get_world_size(self) -> int: + """ + Return the number of ranks in the NCCL communicator. + """ + return self._world_size + + def send(self, buf: "torch.Tensor", peer_rank: int) -> None: + """ + Send a torch.Tensor to a peer. + + This returns when the send kernel has been queued, but the kernel may + not have completed. Therefore, the caller should ensure that there are + no concurrent writes to the sent `buf` until the send has finished. + That is, either all writes should be submitted on the current stream + (self._cuda_stream) or, if on a different stream, that stream should + synchronize with the current stream. + + Args: + buf: The torch.Tensor to send. It should already be on this + actor's default device. + peer_rank: The rank of the actor to send to. + """ + if self._closed: + raise RayChannelError("NCCL group has been destroyed.") + + if self._use_communication_streams: + # We observed that if all recv/compute/send operations run on GPU, + # since there is no synchronization, the CPU execution loop may be + # far ahead of the GPU operations and lead to runtime failures. + # To avoid that, we synchronize on the send stream. + # TODO(rui): find a better approach + self._send_stream.synchronize() + + # TODO(swang): Handle send/recv async NCCL errors such as network + # failures. + self._comm.send( + self.nccl_util.get_tensor_ptr(buf), + buf.numel(), + self.nccl_util.get_nccl_tensor_dtype(buf), + peer_rank, + self._send_stream.cuda_stream, + ) + + def recv( + self, + shape: Tuple[int], + dtype: "torch.dtype", + peer_rank: int, + allocator=Optional[TorchTensorAllocator], + ) -> "torch.Tensor": + """ + Receive a torch.Tensor from a peer and synchronize the current stream. + + After this call returns, the receive buffer is safe to read from from + any stream. An RayChannelError will be raised if an error occurred (e.g., + remote actor died), and the buffer is not safe to read. + + Args: + buf: The torch.Tensor to receive into. This buffer is safe to read + peer_rank: The rank of the actor to receive from. + """ + if self._closed: + raise RayChannelError("NCCL group has been destroyed.") + assert allocator is not None, "NCCL group requires a tensor allocator" + buf = allocator(shape, dtype) + + if self._use_communication_streams: + # We observed that if all recv/compute/send operations run on GPU, + # since there is no synchronization, the CPU execution loop may be + # far ahead of the GPU operations and lead to runtime failures. + # To avoid that, we synchronize on the recv stream. + # TODO(rui): find a better approach + self._recv_stream.synchronize() + + self._comm.recv( + self.nccl_util.get_tensor_ptr(buf), + buf.numel(), + self.nccl_util.get_nccl_tensor_dtype(buf), + peer_rank, + self._recv_stream.cuda_stream, + ) + else: + self._comm.recv( + self.nccl_util.get_tensor_ptr(buf), + buf.numel(), + self.nccl_util.get_nccl_tensor_dtype(buf), + peer_rank, + self._recv_stream.cuda_stream, + ) + + # Buffer values are undefined if NCCL ops are aborted. Therefore, we + # need to synchronize here and check that the channel is still open to + # ensure that the receive buffer is valid. + # TODO(swang): Avoid CUDA synchronization. + self._cuda_stream.synchronize() + + if self._closed: + raise RayChannelError("NCCL group has been destroyed.") + return buf + + def _exec_collective( + self, + send_buf: "torch.Tensor", + recv_buf: "torch.Tensor", + operation: "Callable[..., None]", + *operation_args, + ): + if self._closed: + raise RayChannelError("NCCL group has been destroyed.") + + assert send_buf.dtype == recv_buf.dtype, ( + "Ray Compiled Graph derived the dtype of recv_buf from send_buf, " + "so send_buf and recv_buf must have the same dtype. " + "If you see this error, please file an issue at Ray repository." + ) + + operation(*operation_args) + + # Buffer values are undefined if NCCL ops are aborted. Therefore, we + # need to synchronize here and check that the channel is still open to + # ensure that the receive buffer is valid. + # TODO(swang): Avoid CUDA synchronization. + # TODO(wxdeng): This synchronize will be optional after merging the unify PR. + self._cuda_stream.synchronize() + if self._closed: + raise RayChannelError( + "NCCL group has been destroyed during allreduce operation. " + "There may be a dtype mismatch between input tensors from " + "different ranks." + ) + + def allgather( + self, + send_buf: "torch.Tensor", + recv_buf: "torch.Tensor", + ): + operation_args = [ + self.nccl_util.get_tensor_ptr(send_buf), + self.nccl_util.get_tensor_ptr(recv_buf), + send_buf.numel(), + self.nccl_util.get_nccl_tensor_dtype(send_buf), + self._cuda_stream.cuda_stream, + ] + self._exec_collective( + send_buf, + recv_buf, + self._comm.allGather, + *operation_args, + ) + + def allreduce( + self, + send_buf: "torch.Tensor", + recv_buf: "torch.Tensor", + op: ReduceOp = ReduceOp.SUM, + ): + operation_args = [ + self.nccl_util.get_tensor_ptr(send_buf), + self.nccl_util.get_tensor_ptr(recv_buf), + send_buf.numel(), + self.nccl_util.get_nccl_tensor_dtype(send_buf), + op.value, + self._cuda_stream.cuda_stream, + ] + self._exec_collective( + send_buf, + recv_buf, + self._comm.allReduce, + *operation_args, + ) + + def reducescatter( + self, + send_buf: "torch.Tensor", + recv_buf: "torch.Tensor", + op: ReduceOp = ReduceOp.SUM, + ): + operation_args = [ + self.nccl_util.get_tensor_ptr(send_buf), + self.nccl_util.get_tensor_ptr(recv_buf), + recv_buf.numel(), + self.nccl_util.get_nccl_tensor_dtype(send_buf), + op.value, + self._cuda_stream.cuda_stream, + ] + self._exec_collective( + send_buf, + recv_buf, + self._comm.reduceScatter, + *operation_args, + ) + + @property + def recv_stream(self): + import torch + + return torch.cuda.StreamContext(self._recv_stream) + + @property + def send_stream(self): + import torch + + return torch.cuda.StreamContext(self._send_stream) + + def destroy(self) -> None: + """ + Destroy the NCCL group. + """ + if self._closed: + return + + self._closed = True + + if self._comm is not None: + logger.info( + "Destructing NCCL group on actor: " + f"{ray.get_runtime_context().current_actor}" + ) + # Abort *after* setting the _closed flag. This ensures that NCCL + # ops that were blocked on a remote peer will see that the _closed + # flag is True when they exit from the abort. + self._comm.abort() + self._comm.destroy() + + def get_transport_name(self) -> str: + return "accelerator" + + @classmethod + def generate_communicator_id(cls) -> str: + from cupy.cuda import nccl + + return nccl.get_unique_id() diff --git a/lib/python3.12/site-packages/ray/experimental/channel/serialization_context.py b/lib/python3.12/site-packages/ray/experimental/channel/serialization_context.py new file mode 100644 index 0000000000000000000000000000000000000000..548d36301f6bf622005e6ff545be038a5ccfca57 --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/channel/serialization_context.py @@ -0,0 +1,231 @@ +import warnings +from typing import TYPE_CHECKING, Any, Dict, List, Set, Tuple, Union + +from ray.experimental.util.types import Device + +if TYPE_CHECKING: + import numpy as np + import torch + + +_TORCH_WARNING_FILTER_ACTIVATE = True + + +class _SerializationContext: + def __init__(self): + # If true, then tensors found in the data to serialize are extracted + # and the caller should send them through an external transport. + self._use_external_transport: bool = False + # If _use_external_transport is True, then these are + # the tensors that should be sent or received + # out-of-band, through the external transport. + self._out_of_band_tensors: List["torch.Tensor"] = [] + # During serialization, tensors sent out-of-band are replaced with + # integer placeholders. This tracks the set of placeholders seen. + self._deserialized_tensor_placeholders: Set[int] = set() + + # Buffer for transferring data between tasks in the same worker process. + # The key is the channel ID, and the value is the data. We don't use a + # lock when reading/writing the buffer because a DAG node actor will only + # execute one task at a time in `do_exec_tasks`. It will not execute multiple + # Ray tasks on a single actor simultaneously. + self.intra_process_channel_buffers: Dict[str, Any] = {} + # The number of readers for each channel. When the number of readers + # reaches 0, remove the data from the buffer. + self.channel_id_to_num_readers: Dict[str, int] = {} + + def set_target_device(self, device: Device) -> None: + self._target_device = device + + def set_data(self, channel_id: str, value: Any, num_readers: int) -> None: + assert num_readers > 0, "num_readers must be greater than 0." + assert ( + channel_id not in self.intra_process_channel_buffers + ), f"Channel {channel_id} already exists in the buffer." + assert ( + channel_id not in self.channel_id_to_num_readers + ), f"Channel {channel_id} already exists in the channel_id_to_num_readers." + + self.intra_process_channel_buffers[channel_id] = value + self.channel_id_to_num_readers[channel_id] = num_readers + + def has_data(self, channel_id: str) -> bool: + return channel_id in self.intra_process_channel_buffers + + def get_data(self, channel_id: str) -> Any: + assert ( + channel_id in self.intra_process_channel_buffers + ), f"Channel {channel_id} does not exist in the buffer." + assert ( + channel_id in self.channel_id_to_num_readers + ), f"Channel {channel_id} does not exist in the channel_id_to_num_readers." + + self.channel_id_to_num_readers[channel_id] -= 1 + if self.channel_id_to_num_readers[channel_id] == 0: + # All readers have read the data, so we can remove it. + self.channel_id_to_num_readers.pop(channel_id) + return self.intra_process_channel_buffers.pop(channel_id) + return self.intra_process_channel_buffers[channel_id] + + def reset_data(self, channel_id: str) -> None: + self.intra_process_channel_buffers.pop(channel_id, None) + self.channel_id_to_num_readers.pop(channel_id, None) + + def set_use_external_transport(self, use_external_transport: bool) -> None: + self._use_external_transport = use_external_transport + + @property + def use_external_transport(self) -> bool: + return self._use_external_transport + + def reset_out_of_band_tensors( + self, tensors: List["torch.Tensor"] + ) -> Tuple[List["torch.Tensor"], Set[int]]: + """ + Return and reset the out-of-band tensors and all tensor placeholders + that were deserialized since the last call to reset. + """ + prev_tensors = self._out_of_band_tensors + deserialized_tensor_placeholders = self._deserialized_tensor_placeholders + self._out_of_band_tensors = tensors + self._deserialized_tensor_placeholders = set() + return prev_tensors, deserialized_tensor_placeholders + + def serialize_tensor( + self, tensor: "torch.Tensor" + ) -> Union[int, Tuple["np.ndarray", "torch.dtype", str]]: + from ray.experimental.channel import ChannelContext + + ctx = ChannelContext.get_current() + if self._use_external_transport and ( + ctx._torch_device is None or ctx._torch_device == tensor.device + ): + # External transport is enabled and we found a tensor that matches + # our device. Add the actual tensor to a buffer. The buffer of + # tensors should later be popped by the caller and sent via + # external transport. + self._out_of_band_tensors.append(tensor) + # Return a placeholder. + return len(self._out_of_band_tensors) - 1 + + return self.serialize_to_numpy_or_scalar(tensor) + + def serialize_to_numpy_or_scalar( + self, tensor: "torch.Tensor" + ) -> Tuple[Union["np.ndarray", Any], "torch.dtype", str]: + """ + Serialize a tensor to a numpy array, + or a scalar when the tensor is 0-dim. + """ + import torch + + tensor_device_type = tensor.device.type + + # Transfer through Ray's shared memory store for now. + # TODO(swang): This requires two copies, one to transfer from GPU to + # CPU and another from CPU to shared memory. Ideally we should elide + # the first copy and memcpy directly from GPU to the shared memory + # buffer. + if tensor_device_type != "cpu": + tensor = tensor.to("cpu") + + # Numpy does not have an equivalent dtype for all torch dtypes, so + # instead of casting directly to numpy: + # 1) for non-scalar tensors, we first use a view with a common dtype (uint8) + # and then view as numpy array. + # 2) for scalar tensors, we cannot use a uint8 view when the size differs, + # so we save the original item and type information. + if tensor.dim() > 0: + return (tensor.view(torch.uint8).numpy(), tensor.dtype, tensor_device_type) + else: + return (tensor.item(), tensor.dtype, tensor_device_type) + + def deserialize_tensor( + self, + val: Union[Tuple["np.ndarray", "torch.dtype", str], int], + target_device: Device, + ): + + # Found a placeholder for a tensor that was serialized via accelerator. + # Replace it with the corresponding deserialized tensor. + if isinstance(val, int): + placeholder = val + self._deserialized_tensor_placeholders.add(placeholder) + assert placeholder < len(self._out_of_band_tensors), ( + "placeholder", + placeholder, + "out_of_band_tensors", + self._out_of_band_tensors, + ) + tensor = self._out_of_band_tensors[placeholder] + if target_device == Device.CPU: + tensor = tensor.to("cpu") + return tensor + + np_array, dtype, tensor_device_type = val + return self.deserialize_from_numpy_or_scalar( + np_array, dtype, tensor_device_type, target_device + ) + + def deserialize_from_numpy_or_scalar( + self, + np_array: Union["np.ndarray", Any], + dtype: "torch.dtype", + tensor_device_type: str, + target_device: Device, + ): + import numpy as np + import torch + + if target_device == Device.DEFAULT: + target_device_type = tensor_device_type + elif target_device in [Device.GPU, Device.CUDA]: + target_device_type = "cuda" + else: + target_device_type = target_device.value + + # TODO(swang): Support local P2P transfers if available. + if target_device_type != "cpu": + + def convert_numpy_to_tensor(np_array): + if not isinstance(np_array, np.ndarray): + # For scalar tensors, create the 0-dim tensor. + return torch.tensor( + np_array, device=target_device_type, dtype=dtype + ) + else: + # For non-scalar tensors, view as the original dtype. + # It does zero-copy convert np_array inside shared memory to + # a tensor. Since we move data to GPU immediately, it is safe. + cpu_tensor = torch.from_numpy(np_array).view(dtype) + return cpu_tensor.to(device=target_device_type) + + global _TORCH_WARNING_FILTER_ACTIVATE + # filtering warning messages would be the bottleneck for + # deserializing torch tensors. Since the warning only prompts once, + # we would only deal with it for the first time. + if _TORCH_WARNING_FILTER_ACTIVATE: + with warnings.catch_warnings(): + # Since np_array.is_writable is False (it is set by Ray), + # this raises a warning. Suppress it. + warnings.filterwarnings( + "ignore", + category=UserWarning, + message="The given NumPy array is not writable", + ) + gpu_tensor = convert_numpy_to_tensor(np_array) + _TORCH_WARNING_FILTER_ACTIVATE = False + else: + gpu_tensor = convert_numpy_to_tensor(np_array) + + return gpu_tensor + + # TODO(swang): Use zero-copy from_numpy() if np_array.flags.writeable + # is True. This is safe to set when deserializing np_array if the + # upstream task has num_readers=1. + if not isinstance(np_array, np.ndarray): + # For scalar tensors, create the 0-dim tensor. + return torch.tensor(np_array, device=target_device_type, dtype=dtype) + else: + # For non-scalar tensors, view as the original dtype. + return torch.tensor(np_array, device=target_device_type).view(dtype) diff --git a/lib/python3.12/site-packages/ray/experimental/channel/shared_memory_channel.py b/lib/python3.12/site-packages/ray/experimental/channel/shared_memory_channel.py new file mode 100644 index 0000000000000000000000000000000000000000..11cac4bc3439c76ce75334edfb808a37ab09e52b --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/channel/shared_memory_channel.py @@ -0,0 +1,796 @@ +import io +import logging +import time +from collections import defaultdict, namedtuple +from typing import Any, Dict, List, Optional, Set, Tuple, Union + +import ray +import ray.exceptions +from ray._raylet import SerializedObject +from ray.experimental.channel import utils +from ray.experimental.channel.common import ChannelInterface, ChannelOutputType +from ray.experimental.channel.intra_process_channel import IntraProcessChannel +from ray.experimental.channel.utils import get_self_actor +from ray.util.annotations import DeveloperAPI, PublicAPI + +# Logger for this module. It should be configured at the entry point +# into the program using Ray. Ray provides a default configuration at +# entry/init points. +logger = logging.getLogger(__name__) + + +def _create_channel_ref( + self, + buffer_size_bytes: int, +) -> "ray.ObjectRef": + """ + Create a channel that can be read and written through Ray's shared-memory + object store. + + The channel has no buffer, so the writer will block until reader(s) have + read the previous value. + + A writer and colocated readers can communicate via a shared memory buffer. + If the readers are remote, then RPC is used to synchronize the writer and + readers' buffers. + + Args: + buffer_size_bytes: The initial buffer size in bytes for messages + that can be passed between tasks in the DAG. The buffers will + be automatically resized if larger messages are written to the + channel. + Returns: + Channel: A wrapper around ray.ObjectRef. + """ + worker = ray._private.worker.global_worker + worker.check_connected() + + value = b"0" * buffer_size_bytes + + try: + object_ref = worker.put_object( + value, owner_address=None, _is_experimental_channel=True + ) + except ray.exceptions.ObjectStoreFullError: + logger.info( + "Put failed since the value was either too large or the " + "store was full of pinned objects." + ) + raise + return object_ref + + +# Compiled Graph maintains 1 reader object reference (also called buffer) per node. +# reader_ref: The object reference. +# ref_owner_actor_id: The actor who created the object reference. +# num_readers: The number of reader actors who reads this object reference. +ReaderRefInfo = namedtuple( + "ReaderRefInfo", ["reader_ref", "ref_owner_actor_id", "num_reader_actors"] +) + + +class _ResizeChannel: + """ + When a channel must be resized, the channel backing store must be resized on both + the writer and the reader nodes. The writer first resizes its own backing store. The + writer then uses an instance of this class as a sentinel value to tell the reader to + resize its own backing store. The class instance is sent through the channel. + """ + + def __init__( + self, + _node_id_to_reader_ref_info: Dict[str, ReaderRefInfo], + ): + """ + Args: + _node_id_to_reader_ref_info: A node id to ReaderRefInfo. + """ + self._node_id_to_reader_ref_info = _node_id_to_reader_ref_info + + +class SharedMemoryType(ChannelOutputType): + def __init__( + self, + *, + buffer_size_bytes: Optional[int] = None, + num_shm_buffers: Optional[int] = None, + ): + """ + Args: + buffer_size_bytes: The initial buffer size in bytes for messages + that can be passed between tasks in the DAG. The buffers will + be automatically resized if larger messages are written to the + channel. + num_shm_buffers: The number of shared memory buffers per channel. + Note: In the case of multiple nodes, we only support 1 shared + memory buffer. + """ + super().__init__() + + from ray.dag import DAGContext + + ctx = DAGContext.get_current() + + if buffer_size_bytes is None: + buffer_size_bytes = ctx.buffer_size_bytes + self.buffer_size_bytes = buffer_size_bytes + if num_shm_buffers is None: + num_shm_buffers = 1 + self._num_shm_buffers = num_shm_buffers + + def create_channel( + self, + writer: Optional["ray.actor.ActorHandle"], + reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]], + driver_actor_id: Optional[str] = None, + ) -> "Channel": + """ + Instantiate a ChannelInterface class that can be used + to pass data of this type. + + Args: + writer: The actor that may write to the channel. None signifies the driver. + reader_and_node_list: A list of tuples, where each tuple contains a reader + actor handle and the node ID where the actor is located. + driver_actor_id: If this channel is read by a driver and that driver is an + actual actor, this will be the actor ID of that driver actor. + + Returns: + A ChannelInterface that can be used to pass data + of this type. + """ + return CompositeChannel( + writer, + reader_and_node_list, + self._num_shm_buffers, + driver_actor_id, + ) + + +@PublicAPI(stability="alpha") +class Channel(ChannelInterface): + """ + A wrapper type for ray.ObjectRef. Currently supports ray.get but not + ray.wait. + """ + + def __init__( + self, + writer: Optional[ray.actor.ActorHandle], + reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]], + typ: Optional[Union[int, SharedMemoryType]] = None, + _writer_node_id: Optional["ray.NodeID"] = None, + _writer_ref: Optional["ray.ObjectRef"] = None, + _node_id_to_reader_ref_info: Optional[Dict[str, ReaderRefInfo]] = None, + _writer_registered: bool = False, + _reader_registered: bool = False, + ): + """ + Create a channel that can be read and written by co-located Ray processes. + + Anyone may write to or read from the channel. The channel has no + buffer, so the writer will block until reader(s) have read the previous + value. + + Args: + writer: The actor that may write to the channel. None signifies the driver. + reader_and_node_list: A list of tuples, where each tuple contains a reader + actor handle and the node ID where the actor is located. + typ: Type information about the values passed through the channel. + Either an integer representing the max buffer size in bytes + allowed, or a SharedMemoryType. + Returns: + Channel: A wrapper around ray.ObjectRef. + """ + assert len(reader_and_node_list) > 0 + for reader, _ in reader_and_node_list: + assert isinstance(reader, ray.actor.ActorHandle) + + if typ is None: + typ = SharedMemoryType() + elif isinstance(typ, int): + typ = SharedMemoryType(buffer_size_bytes=typ) + + # The min buffer size must be large enough to at least fit an instance of the + # _ResizeChannel class along with any metadata. + MIN_BUFFER_SIZE = int(1000) # 1000 bytes + if typ.buffer_size_bytes < MIN_BUFFER_SIZE: + raise ValueError( + "typ.buffer_size_bytes must be at least MIN_BUFFER_SIZE " + f"({MIN_BUFFER_SIZE} bytes)" + ) + + self._writer = writer + self._reader_and_node_list = reader_and_node_list + self._typ = typ + + self._worker = ray._private.worker.global_worker + self._worker.check_connected() + + self._writer_registered = _writer_registered + self._reader_registered = _reader_registered + # NodeID -> ReaderRefInfo on that node. Note that there's only 1 + # reader ref per node. + self._node_id_to_reader_ref_info: Dict[str, ReaderRefInfo] = ( + _node_id_to_reader_ref_info or {} + ) + + # Node ID -> a list of reader actors. + self._node_id_to_readers: Dict[str, "ray.actor.ActorHandle"] = defaultdict(list) + for reader, node_id in self._reader_and_node_list: + self._node_id_to_readers[node_id].append(reader) + + # Number of readers in a local node. + self._num_local_readers = 0 + + if _writer_ref is None: + # We are the writer. Check that the passed handle matches the + # current actor (or it is the driver). + # TODO(swang): Channels must be initially constructed by the writer + # actor, so we shouldn't need to include `writer` in the + # constructor args. Either support Channels being constructed by + # someone other than the writer or remove it from the args. + self_actor = get_self_actor() + assert writer == self_actor + + self._writer_node_id = ( + ray.runtime_context.get_runtime_context().get_node_id() + ) + self._writer_ref = _create_channel_ref(self, typ.buffer_size_bytes) + + self._create_reader_refs(typ.buffer_size_bytes) + else: + assert ( + _writer_node_id is not None + ), "_writer_node_id must also be passed to the constructor when " + "_writer_ref is." + assert _node_id_to_reader_ref_info is not None, ( + "_node_id_to_reader_ref_info must also be passed to the constructor " + "when _writer_ref is." + ) + + self._writer_ref = _writer_ref + self._writer_node_id = _writer_node_id + self._node_id_to_reader_ref_info = _node_id_to_reader_ref_info + + assert self._num_local_readers == 0 + remote_node_exists = False + for node_id, readers in self._node_id_to_readers.items(): + if self.is_local_node(node_id): + self._num_local_readers += len(readers) + else: + remote_node_exists = True + # If remote node exists, we have 1 additional reader that listens + # to object changes and push them to remote nodes. + if remote_node_exists: + self._num_local_readers += 1 + # There must be at least 1 local reader + assert self._num_local_readers > 0 + + self._local_reader_ref: Optional["ray.ObjectRef"] = self._get_local_reader_ref( + self._node_id_to_reader_ref_info + ) + + def _get_local_reader_ref( + self, _node_id_to_reader_ref_info: Dict[str, ReaderRefInfo] + ) -> Optional["ray.ObjectRef"]: + for node_id, reader_ref_info in _node_id_to_reader_ref_info.items(): + if self.is_local_node(node_id): + return reader_ref_info.reader_ref + return None + + def _create_reader_refs( + self, + buffer_size_bytes: int, + ): + # TODO(jhumphri): Free the current reader ref once the reference to it is + # destroyed below. + + for node_id, readers in self._node_id_to_readers.items(): + if not self.is_local_node(node_id): + # Find 1 reader in a remote node to create a reference that's + # shared by all readers. When a new value is written to a reference, + # it is sent to this reference. + reader = readers[0] + fn = reader.__ray_call__ + self._node_id_to_reader_ref_info[node_id] = ReaderRefInfo( + reader_ref=ray.get( + fn.remote(_create_channel_ref, buffer_size_bytes) + ), + ref_owner_actor_id=reader._actor_id, + num_reader_actors=len(readers), + ) + else: + writer_id = ray.ActorID.nil() + if self._writer is not None: + writer_id = self._writer._actor_id + self._node_id_to_reader_ref_info[node_id] = ReaderRefInfo( + reader_ref=self._writer_ref, + ref_owner_actor_id=writer_id, + num_reader_actors=len(readers), + ) + # There must be only 1 node reader reference per node. + assert len(self._node_id_to_reader_ref_info) == len(self._node_id_to_readers) + + # We need to register the new writer_ref. + self._writer_registered = False + self.ensure_registered_as_writer() + + @staticmethod + def is_local_node(node_id): + return ray.runtime_context.get_runtime_context().get_node_id() == node_id + + def ensure_registered_as_writer(self) -> None: + if self._writer_registered: + return + + if not self.is_local_node(self._writer_node_id): + raise ValueError( + "`ensure_registered_as_writer()` must only be called on the node that " + "the writer is on." + ) + + remote_reader_ref_info: Dict[str, ReaderRefInfo] = {} + for node_id, reader_ref_info in self._node_id_to_reader_ref_info.items(): + if self.is_local_node(node_id): + continue + remote_reader_ref_info[node_id] = reader_ref_info + + self._worker.core_worker.experimental_channel_register_writer( + self._writer_ref, + remote_reader_ref_info, + ) + self._writer_registered = True + + def ensure_registered_as_reader(self) -> None: + if self._reader_registered: + return + + for node_id, reader_ref_info in self._node_id_to_reader_ref_info.items(): + if self.is_local_node(node_id): + self._worker.core_worker.experimental_channel_register_reader( + reader_ref_info.reader_ref, + ) + self._reader_registered = True + + @staticmethod + def _deserialize_reader_channel( + writer: ray.actor.ActorHandle, + reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]], + typ: int, + writer_node_id, + writer_ref: "ray.ObjectRef", + node_id_to_reader_ref_info: Dict[str, ReaderRefInfo], + writer_registered: bool, + reader_registered: bool, + ) -> "Channel": + chan = Channel( + writer, + reader_and_node_list, + typ, + _writer_node_id=writer_node_id, + _writer_ref=writer_ref, + _node_id_to_reader_ref_info=node_id_to_reader_ref_info, + _writer_registered=writer_registered, + _reader_registered=reader_registered, + ) + return chan + + def __reduce__(self): + assert self._node_id_to_reader_ref_info is not None + return self._deserialize_reader_channel, ( + self._writer, + self._reader_and_node_list, + self._typ, + self._writer_node_id, + self._writer_ref, + self._node_id_to_reader_ref_info, + self._writer_registered, + self._reader_registered, + ) + + def __str__(self) -> str: + return ( + f"Channel(_node_id_to_reader_ref_info={self._node_id_to_reader_ref_info}, " + f"_writer_ref={self._writer_ref})" + ) + + def _resize_channel_if_needed(self, serialized_value: str, timeout_ms: int): + # serialized_value.total_bytes *only* includes the size of the data. It does not + # include the size of the metadata, so we must account for the size of the + # metadata explicitly. + size = serialized_value.total_bytes + len(serialized_value.metadata) + if size > self._typ.buffer_size_bytes: + # Now make the channel backing store larger. + self._typ.buffer_size_bytes = size + # TODO(jhumphri): Free the current writer ref once the reference to it is + # destroyed below. + # TODO(sang): Support different policies such as 2X buffer size. + prev_writer_ref = self._writer_ref + self._writer_ref = _create_channel_ref(self, self._typ.buffer_size_bytes) + self._create_reader_refs(self._typ.buffer_size_bytes) + self._local_reader_ref = self._get_local_reader_ref( + self._node_id_to_reader_ref_info + ) + + # Write a special message to the channel so that the readers know to + # stop using the current reader_ref. + special_message = _ResizeChannel(self._node_id_to_reader_ref_info) + special_message_serialized = ( + self._worker.get_serialization_context().serialize(special_message) + ) + self._worker.core_worker.experimental_channel_put_serialized( + special_message_serialized, + prev_writer_ref, + self._num_local_readers, + timeout_ms, + ) + # TODO(sang): Clean the previous ref that won't be used. + # Right now, if we just close it here, it will not work because + # of race conditions. + # self._worker.core_worker.experimental_channel_set_error( + # prev_writer_ref + # ) + + def write(self, value: Any, timeout: Optional[float] = None) -> None: + self.ensure_registered_as_writer() + assert ( + timeout is None or timeout >= 0 or timeout == -1 + ), "Timeout must be non-negative or -1." + # -1 means no timeout (block indefinitely) + timeout_ms = int(timeout * 1000) if timeout is not None else -1 + + if not isinstance(value, SerializedObject): + try: + serialized_value = self._worker.get_serialization_context().serialize( + value + ) + except TypeError as e: + sio = io.StringIO() + ray.util.inspect_serializability(value, print_file=sio) + msg = ( + "Could not serialize the put value " + f"{repr(value)}:\n" + f"{sio.getvalue()}" + ) + raise TypeError(msg) from e + else: + serialized_value = value + + start_time = time.monotonic() + self._resize_channel_if_needed(serialized_value, timeout_ms) + if timeout is not None: + timeout_ms -= int((time.monotonic() - start_time) * 1000) + timeout_ms = max(timeout_ms, 0) + + self._worker.core_worker.experimental_channel_put_serialized( + serialized_value, + self._writer_ref, + self._num_local_readers, + timeout_ms, + ) + + def read(self, timeout: Optional[float] = None) -> Any: + assert ( + timeout is None or timeout >= 0 or timeout == -1 + ), "Timeout must be non-negative or -1." + self.ensure_registered_as_reader() + + start_time = time.monotonic() + ret = self._worker.get_objects( + [self._local_reader_ref], timeout=timeout, return_exceptions=True + )[0][0] + + if isinstance(ret, _ResizeChannel): + self._node_id_to_reader_ref_info = ret._node_id_to_reader_ref_info + self._local_reader_ref = self._get_local_reader_ref( + self._node_id_to_reader_ref_info + ) + # We need to register the new reader_ref. + self._reader_registered = False + self.ensure_registered_as_reader() + if timeout is not None: + timeout -= time.monotonic() - start_time + timeout = max(timeout, 0) + ret = self._worker.get_objects( + [self._local_reader_ref], timeout=timeout, return_exceptions=True + )[0][0] + + return ret + + def release_buffer(self, timeout: Optional[float] = None) -> None: + assert ( + timeout is None or timeout >= 0 or timeout == -1 + ), "Timeout must be non-negative or -1." + self.ensure_registered_as_reader() + self._worker.get_objects( + [self._local_reader_ref], + timeout=timeout, + return_exceptions=True, + skip_deserialization=True, + ) + + def close(self) -> None: + """ + Close this channel by setting the error bit on both the writer_ref and the + reader_ref. + """ + self._worker.core_worker.experimental_channel_set_error(self._writer_ref) + is_local_node_reader = False + + for node_id in self._node_id_to_readers.keys(): + if self.is_local_node(node_id): + is_local_node_reader = True + if is_local_node_reader: + self.ensure_registered_as_reader() + + for reader_ref_info in self._node_id_to_reader_ref_info.values(): + self._worker.core_worker.experimental_channel_set_error( + reader_ref_info.reader_ref + ) + + +@DeveloperAPI +class BufferedSharedMemoryChannel(ChannelInterface): + """A channel that can be read and written by Ray processes. + + It creates `num_shm_buffers` number of buffers and allows buffered read and + write APIs. I.e., read and write APIs are non-blocking as long as it can write to + next buffer or read from a next buffer. See `read` and `write` APIs for + more details. + + Args: + writer: The actor that may write to the channel. None signifies the driver. + reader_and_node_list: A list of tuples, where each tuple contains a reader + actor handle and the node ID where the actor is located. Note that currently + we only support this for readers on the same node as the writer. + num_shm_buffers: Number of shared memory buffers to read/write. + typ: Type information about the values passed through the channel. + Either an integer representing the max buffer size in bytes + allowed, or a SharedMemoryType. + """ + + def __init__( + self, + writer: Optional[ray.actor.ActorHandle], + reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]], + num_shm_buffers: int, + typ: Optional[Union[int, SharedMemoryType]] = None, + ): + self._num_shm_buffers = num_shm_buffers + self._buffers = [ + # We use Channel directly as a buffer implementation as + # channel only allows to have 1 shared memory buffer. + Channel(writer, reader_and_node_list, typ) + for _ in range(num_shm_buffers) + ] + # The next index to write from self._buffers. + self._next_write_index = 0 + # The next index to read from self._buffers. + self._next_read_index = 0 + + def ensure_registered_as_writer(self): + """ + Check whether the process is a valid writer. This method must be idempotent. + """ + for buffer in self._buffers: + buffer.ensure_registered_as_writer() + + def ensure_registered_as_reader(self): + """ + Check whether the process is a valid reader. This method must be idempotent. + """ + for buffer in self._buffers: + buffer.ensure_registered_as_reader() + + def write(self, value: Any, timeout: Optional[float] = None) -> None: + """Write a value to a channel. + + If the next buffer is available, it returns immediately. If the next + buffer is not read by downstream consumers, it blocks until a buffer is + available to write. If a buffer is not available within timeout, it raises + RayChannelTimeoutError. + """ + self.ensure_registered_as_writer() + # A single channel is not supposed to read and write at the same time. + assert self._next_read_index == 0 + self._buffers[self._next_write_index].write(value, timeout) + self._next_write_index += 1 + self._next_write_index %= self._num_shm_buffers + + def read(self, timeout: Optional[float] = None) -> Any: + """Read a value from a channel. + + If the next buffer is available, it returns immediately. If the next + buffer is not written by an upstream producer, it blocks until a buffer is + available to read. If a buffer is not available within timeout, it raises + RayChannelTimeoutError. + """ + self.ensure_registered_as_reader() + # A single channel is not supposed to read and write at the same time. + assert self._next_write_index == 0 + output = self._buffers[self._next_read_index].read(timeout) + self._next_read_index += 1 + self._next_read_index %= self._num_shm_buffers + return output + + def release_buffer(self, timeout: Optional[float] = None): + """Release the native buffer of the channel to allow the buffer to be reused for + future data. + + If the next buffer is available, it returns immediately. If the next + buffer is not written by an upstream producer, it blocks until a buffer is + available to be released. If a buffer is not available within timeout, it raises + RayChannelTimeoutError. + """ + # A single channel is not supposed to read and write at the same time. + assert self._next_write_index == 0 + self._buffers[self._next_read_index].release_buffer(timeout) + self._next_read_index += 1 + self._next_read_index %= self._num_shm_buffers + + def close(self) -> None: + for buffer in self._buffers: + buffer.close() + + @property + def next_write_index(self): + # Testing only + return self._next_write_index + + @property + def next_read_index(self): + # Testing only + return self._next_read_index + + +@PublicAPI(stability="alpha") +class CompositeChannel(ChannelInterface): + """ + Can be used to send data to different readers via different channels. + For example, if the reader is in the same worker process as the writer, + the data can be sent via IntraProcessChannel. If the reader is in a different + worker process, the data can be sent via shared memory channel. + + Args: + writer: The actor that may write to the channel. None signifies the driver. + reader_and_node_list: A list of tuples, where each tuple contains a reader + actor handle and the node ID where the actor is located. + num_shm_buffers: The number of shared memory buffers per channel. + Note: In the case of multiple nodes, we only support 1 shared + memory buffer. + driver_actor_id: If this channel is read by a driver and that driver is an + actual actor, this will be the actor ID of that driver actor. + """ + + def __init__( + self, + writer: Optional[ray.actor.ActorHandle], + reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]], + num_shm_buffers: int, + driver_actor_id: Optional[str] = None, + _channel_dict: Optional[Dict[ray.ActorID, ChannelInterface]] = None, + _channels: Optional[Set[ChannelInterface]] = None, + _writer_registered: bool = False, + _reader_registered: bool = False, + ): + self._writer = writer + self._reader_and_node_list = reader_and_node_list + self._num_shm_buffers = num_shm_buffers + self._driver_actor_id = driver_actor_id + self._writer_registered = _writer_registered + self._reader_registered = _reader_registered + # A dictionary that maps the actor ID to the channel object. + self._channel_dict = _channel_dict or {} + # The set of channels is a deduplicated version of the _channel_dict values. + self._channels = _channels or set() + if self._channels: + # This CompositeChannel object is created by deserialization. + # We don't need to create channels again. + return + + ( + remote_reader_and_node_list, + local_reader_and_node_list, + ) = utils.split_readers_by_locality(self._writer, self._reader_and_node_list) + # There are some local readers which are the same worker process as the writer. + # Create a local channel for the writer and the local readers. + num_local_readers = len(local_reader_and_node_list) + if num_local_readers > 0: + # Use num_readers = 1 when creating the local channel, + # because we have channel cache to support reading + # from the same channel multiple times. + local_channel = IntraProcessChannel(num_readers=1) + self._channels.add(local_channel) + actor_id = self._get_actor_id(self._writer) + self._channel_dict[actor_id] = local_channel + # There are some remote readers which are not the same Ray actor as the writer. + # We create a BufferedSharedMemoryChannel for readers on the same node, and + # a single Channel for readers on different nodes due to + # https://github.com/ray-project/ray/issues/49044 + ( + readers_same_node, + readers_different_node, + ) = utils.split_actors_by_node_locality( + utils.get_actor_node(self._writer), remote_reader_and_node_list + ) + + if len(readers_same_node) != 0: + remote_channel = BufferedSharedMemoryChannel( + self._writer, readers_same_node, num_shm_buffers + ) + self._channels.add(remote_channel) + for reader, _ in readers_same_node: + actor_id = self._get_actor_id(reader) + self._channel_dict[actor_id] = remote_channel + + if len(readers_different_node) != 0: + remote_channel = Channel(self._writer, readers_different_node) + self._channels.add(remote_channel) + for reader, _ in readers_different_node: + actor_id = self._get_actor_id(reader) + self._channel_dict[actor_id] = remote_channel + + def _get_actor_id(self, reader: ray.actor.ActorHandle) -> str: + return reader._actor_id.hex() + + def ensure_registered_as_writer(self) -> None: + if self._writer_registered: + return + for channel in self._channels: + channel.ensure_registered_as_writer() + self._writer_registered = True + + def ensure_registered_as_reader(self) -> None: + if self._reader_registered: + return + for channel in self._channels: + channel.ensure_registered_as_reader() + self._reader_registered = True + + def __reduce__(self): + return CompositeChannel, ( + self._writer, + self._reader_and_node_list, + self._num_shm_buffers, + self._driver_actor_id, + self._channel_dict, + self._channels, + self._writer_registered, + self._reader_registered, + ) + + def __str__(self) -> str: + return ( + "CompositeChannel(_channels=" + f"{[str(channel) for channel in self._channels]})" + ) + + def write(self, value: Any, timeout: Optional[float] = None) -> None: + self.ensure_registered_as_writer() + for channel in self._channels: + channel.write(value, timeout) + + def read(self, timeout: Optional[float] = None) -> Any: + self.ensure_registered_as_reader() + return self._channel_dict[self._resolve_actor_id()].read(timeout) + + def release_buffer(self, timeout: Optional[float] = None): + self.ensure_registered_as_reader() + self._channel_dict[self._resolve_actor_id()].release_buffer(timeout) + + def _resolve_actor_id(self) -> str: + actor_id = ray.get_runtime_context().get_actor_id() + # If actor_id is None, read was called by the driver + # If the driver is an actor, driver_actor_id will be set to that actor id + if actor_id is None or actor_id == self._driver_actor_id: + # Use the actor ID of the DAGDriverProxyActor. + # The proxy actor is always the first actor in the reader_and_node_list. + assert len(self._reader_and_node_list) >= 1 + driver_proxy_actor = self._reader_and_node_list[0][0] + actor_id = self._get_actor_id(driver_proxy_actor) + return actor_id + + def close(self) -> None: + for channel in self._channels: + channel.close() diff --git a/lib/python3.12/site-packages/ray/experimental/channel/torch_tensor_accelerator_channel.py b/lib/python3.12/site-packages/ray/experimental/channel/torch_tensor_accelerator_channel.py new file mode 100644 index 0000000000000000000000000000000000000000..5e5f4e2abe5e9a24bdcca275ad4522e12c781488 --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/channel/torch_tensor_accelerator_channel.py @@ -0,0 +1,872 @@ +import io +import logging +import uuid +from dataclasses import dataclass +from types import ModuleType +from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Type, Union + +import ray +import ray.util.serialization +from ray.experimental.channel import ChannelContext, utils +from ray.experimental.channel.accelerator_context import ( + AcceleratorContext, + is_accelerator_context_registered, + register_accelerator_context, +) +from ray.experimental.channel.common import ChannelInterface +from ray.experimental.channel.communicator import Communicator +from ray.experimental.channel.communicator_handle import CommunicatorHandle +from ray.experimental.channel.cpu_communicator import CPUCommunicator +from ray.experimental.channel.intra_process_channel import IntraProcessChannel +from ray.experimental.channel.shared_memory_channel import SharedMemoryType +from ray.experimental.channel.torch_tensor_type import TorchTensorType +from ray.util.annotations import DeveloperAPI + +if TYPE_CHECKING: + import torch + + from ray.experimental.channel.shared_memory_channel import Channel + + +# Logger for this module. It should be configured at the entry point +# into the program using Ray. Ray provides a default configuration at +# entry/init points. +logger = logging.getLogger(__name__) + + +@dataclass +class _TorchTensorMetadata: + """ + Metadata for torch.Tensors that can be sent between processes to determine + how large of a buffer to allocate on the receiver(s). + """ + + shape: Union[int, Tuple[int]] + dtype: "torch.dtype" + + +@DeveloperAPI +class TorchTensorAcceleratorChannel(ChannelInterface): + def __init__( + self, + writer: ray.actor.ActorHandle, + reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]], + typ: "TorchTensorType", + driver_actor_id: str, + tensor_metadata_channel: Optional["Channel"] = None, + _cpu_data_channel: Optional["Channel"] = None, + _gpu_data_channel: Optional["_TorchTensorAcceleratorChannel"] = None, + _local_channel: Optional["IntraProcessChannel"] = None, + ): + """ + Can be used to send accelerator tensors nested inside other data. The data is + sent via shared memory while the accelerator tensors are sent through a P2P + transport (e.g., NCCL for GPU). + + NOTE: This class is currently not thread-safe because it reads and + writes the worker-local + ray.experimental.channel.serialization_context._SerializationContext + when serializing data. + + Args: + writer: The actor that may write to the channel. None signifies the + driver. + reader_and_node_list: A list of tuples, where each tuple contains a reader + actor handle and the node ID where the actor is located. + typ: Type information about the values passed through the channel. + driver_actor_id: The actor ID of the DAGDriverProxyActor. + tensor_metadata_channel: A shared-memory channel for sending tensor + metadata. + _cpu_data_channel: A shared-memory channel for sending + non-tensor data. Its writer and readers should match the given + writer and readers. If None is provided, then we assume that + there is no CPU-specific data, i.e. the task directly returned + a CUDA torch.Tensor. + _gpu_data_channel: A channel for sending torch.Tensors via accelerator. + _local_channel: A channel for sending data between the writer and + local readers. + + NOTE: `tensor_metadata_channel` will be set only for testing purposes. + `_cpu_data_channel` is set for testing purposes and for deserialization. + `_gpu_data_channel` and `_local_channel` are set only during deserialization. + """ + self._writer = writer + self._reader_and_node_list = reader_and_node_list + self._typ = typ + + ( + remote_reader_and_node_list, + local_reader_and_node_list, + ) = utils.split_readers_by_locality(self._writer, self._reader_and_node_list) + + num_local_readers = len(local_reader_and_node_list) + self._local_channel = _local_channel + if self._local_channel is None and num_local_readers > 0: + # There are some local readers which are the same worker process as + # the writer. Create a local channel for the writer and the local readers. + # + # Use num_readers = 1 when creating the local channel, + # because we have channel cache to support reading + # from the same channel multiple times. + self._local_channel = IntraProcessChannel(num_readers=1) + + assert len(remote_reader_and_node_list) > 0, ( + "All readers are from the same actor. " + "The TorchTensorType type hint is not needed. " + "No accelerator channel will be created." + ) + self._gpu_data_channel = _gpu_data_channel + if self._gpu_data_channel is None: + self._gpu_data_channel: _TorchTensorAcceleratorChannel = ( + _TorchTensorAcceleratorChannel( + writer, + remote_reader_and_node_list, + typ, + _meta_channel=tensor_metadata_channel, + ) + ) + + self._cpu_data_channel: Optional["Channel"] = _cpu_data_channel + if self._cpu_data_channel is not None: + assert ( + not self._typ.direct_return + ), "CPU channel should be None if direct return is enabled" + + if self._cpu_data_channel is None and not self._typ.direct_return: + # Create a CPU channel to send non-tensor data. + self._cpu_data_channel = SharedMemoryType().create_channel( + writer, remote_reader_and_node_list, driver_actor_id + ) + + # Used for serialization. + self._worker = ray._private.worker.global_worker + self._worker.check_connected() + + ctx = ChannelContext.get_current() + self.serialization_ctx = ctx.serialization_context + assert self.serialization_ctx is not None + + def __reduce__(self): + return ( + TorchTensorAcceleratorChannel, + ( + self._writer, + self._reader_and_node_list, + self._typ, + # driver_actor_id and tensor_metadata_channel are used to initialize + # the _cpu_data_channel and _gpu_data_channel, so we don't need to + # pass them in here. + None, + None, + self._cpu_data_channel, + self._gpu_data_channel, + self._local_channel, + ), + ) + + def ensure_registered_as_writer(self): + if self._local_channel is not None: + self._local_channel.ensure_registered_as_writer() + self._gpu_data_channel.ensure_registered_as_writer() + if self._cpu_data_channel is not None: + self._cpu_data_channel.ensure_registered_as_writer() + + def ensure_registered_as_reader(self): + reader = utils.get_self_actor() + if reader == self._writer: + self._local_channel.ensure_registered_as_reader() + return + self._gpu_data_channel.ensure_registered_as_reader() + if self._cpu_data_channel is not None: + self._cpu_data_channel.ensure_registered_as_reader() + + def _send_cpu_and_gpu_data(self, value: Any, timeout: Optional[float]): + self.serialization_ctx.reset_out_of_band_tensors([]) + # All tensors found in `value` will be transferred via accelerator. + self.serialization_ctx.set_use_external_transport(True) + + try: + # Serialize the data. All tensors that match our current device + # will be extracted into the serialization context and replaced + # with a placeholder. + cpu_data = self._worker.get_serialization_context().serialize(value) + except TypeError as e: + sio = io.StringIO() + ray.util.inspect_serializability(value, print_file=sio) + msg = ( + "Could not serialize the put value " + f"{repr(value)}:\n" + f"{sio.getvalue()}" + ) + raise TypeError(msg) from e + finally: + # Pop the tensors that were found during serialization of `value`. + gpu_tensors, _ = self.serialization_ctx.reset_out_of_band_tensors([]) + # Reset the serialization method to now serialize torch.Tensors + # normally. + self.serialization_ctx.set_use_external_transport(False) + + # First send the extracted tensors through a GPU-specific channel. + self._gpu_data_channel.write(gpu_tensors) + # Next send the non-tensor data through a CPU-specific channel. The + # data contains placeholders for the extracted tensors. + self._cpu_data_channel.write(cpu_data) + + def write(self, value: Any, timeout: Optional[float] = None) -> None: + """ + Send a value that may contain torch.Tensors that should be sent via + external transport. + + Case 1: Use `_local_channel` to send the data to local readers. + + Case 2: Otherwise, use the following method to send the data to remote readers. + + 1) Serializes `value`. During serialization, all torch.Tensors that are + on the default device are extracted and replaced with a unique + placeholder. Thus, the serialized value will contain all non-tensor + data, and any tensors that were not on the default device (e.g., CPU + tensor returned by a GPU actor). + 2) Sends extracted torch.Tensors via the tensor data channel (e.g., + NCCL). + 3) Sends the non-tensor data via the non-tensor data channel. + + If static_non_tensor_data=True was specified, then we only perform step + (3) on the first `write` call. The reader is expected to reuse the sent + data for subsequent messages. + """ + self.ensure_registered_as_writer() + + if self._local_channel is not None: + self._local_channel.write(value) + + if isinstance(value, ray.exceptions.RayTaskError): + if self._typ.static_shape or self._typ.direct_return: + # Raise a fatal error to teardown the DAG. + # This error will also be caught from `CompiledDAGRef.get()` + # and raised to the user + # TODO(swang): Write exceptions to the tensor metadata or + # non-tensor data channel if it is available to make these + # exceptions recoverable. + raise value + + if self._cpu_data_channel is None: + # Handle the case where _direct_return=True. In this case, we check + # that the task returned a CUDA torch.Tensor and just send it + # directly without trying to serialize it first. + import torch + + # These ValueErrors will also be caught from `CompiledDAGRef.get()` + # and raised to the user + if not isinstance(value, torch.Tensor): + # TODO(swang): These errors are currently fatal for the DAG. + # This could be improved by sending the exception through the + # gpu_data_channel's CPU-based metadata channel, if one exists. + raise ValueError( + "Task annotated with _direct_return=True must " + "return a CUDA torch.Tensor, instead found value " + f"`{value}`. DAG will shut down." + ) + elif not value.is_cuda: + raise ValueError( + "Task annotated with _direct_return=True must " + "return a CUDA torch.Tensor, instead found CPU tensor. " + "DAG will shut down." + ) + self._gpu_data_channel.write([value], timeout=timeout) + else: + self._send_cpu_and_gpu_data(value, timeout) + + def _recv_cpu_and_gpu_data( + self, tensors: List["torch.Tensor"], timeout: Optional[float] = None + ) -> Any: + """ + Helper method to receive data that contains a mix of CPU and GPU data. + + Args: + tensors: The GPU data. This is a list of the torch.Tensors that + were found in the sent data. + timeout: Timeout for channel receive. + """ + self.serialization_ctx.reset_out_of_band_tensors(tensors) + + # Next, read and deserialize the non-tensor data. The registered custom + # deserializer will replace the found tensor placeholders with + # `tensors`. + data = self._cpu_data_channel.read( + timeout=timeout, + ) + # Check that all placeholders had a corresponding tensor. + ( + _, + deserialized_tensor_placeholders, + ) = self.serialization_ctx.reset_out_of_band_tensors([]) + assert deserialized_tensor_placeholders == set(range(len(tensors))) + + return data + + def read(self, timeout: Optional[float] = None) -> Any: + """ + Read a value that may contain torch.Tensors sent via external + transport. + + Case 1: If the reader is a local reader and is the same actor as the writer, + then use the `_local_channel` to read the data. + + Case 2: Otherwise, use the following method to read data from remote readers. + + 1) Receives torch.Tensors via the tensor data channel (e.g., NCCL). + 2) Reads the serialized non-tensor data. + 3) Deserializes the non-tensor data. During deserialization, replaces + all found placeholders with the received torch.Tensors. + + If _direct_return=True was specified, then we skip step (2) and (3) and + directly return the data received in (1). + """ + self.ensure_registered_as_reader() + + # If the reader is the same actor as the writer, then we can use the + # local channel to read the data. + reader = utils.get_self_actor() + if reader == self._writer: + assert self._local_channel is not None + return self._local_channel.read() + + # First, read the tensor data. + tensors = self._gpu_data_channel.read(timeout) + + if self._cpu_data_channel is None: + # Handle _direct_return=True. In this case, we expect to receive + # only one tensor, and we return it directly. + assert len(tensors) == 1 + data = tensors[0] + else: + data = self._recv_cpu_and_gpu_data(tensors, timeout) + + return data + + def close(self) -> None: + self._gpu_data_channel.close() + if self._cpu_data_channel is not None: + self._cpu_data_channel.close() + if self._local_channel is not None: + self._local_channel.close() + + +def _torch_tensor_allocator( + shape: Union[int, Tuple[int]], + dtype: "torch.dtype", +): + """ + Allocate a tensor buffer matching the given metadata. + """ + import torch + + ctx = ChannelContext.get_current() + return torch.empty(shape, dtype=dtype, device=ctx.torch_device) + + +class _TorchTensorAcceleratorChannel(ChannelInterface): + def __init__( + self, + writer: ray.actor.ActorHandle, + reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]], + typ: "TorchTensorType", + _meta_channel: Optional["Channel"] = None, + ): + """ + A helper channel for TorchTensorAcceleratorChannel that is used to transfer + lists of torch.Tensors via accelerator. This class can only transfer + torch.Tensors and cannot transfer other CPU data, such as Exception + objects or tensors nested inside of a dictionary. + + Args: + writer: The actor that may write to the channel. None signifies the driver. + reader_and_node_list: A list of tuples, where each tuple contains a reader + actor handle and the node ID where the actor is located. + typ: Type information about the values passed through the channel. + _meta_channel: A channel used to send metadata for the tensors, + i.e. shape and dtype. If not provided, and if the typ does not + specify a static shape and dtype, then a metadata channel based + on shared memory will be created. + """ + import torch + + self.torch: ModuleType = torch + + self._writer = writer + self._writer_rank: Optional[int] = None + self._reader_and_node_list = reader_and_node_list + self._reader_ranks: Optional[List[int]] = None + self._writer_registered: bool = False + self._reader_registered: bool = False + + ctx = ChannelContext.get_current() + assert isinstance( + typ.communicator_id, str + ), f"accelerator group ID ({typ.communicator_id}) must be a str." + self._typ = typ + + self._static_shape = typ.static_shape + + assert self._typ.communicator_id is not None, "No accelerator group specified." + self._accelerator_group_id: str = self._typ.communicator_id + + # If the communicators does not contain the group_id, it means the current + # process is the driver, and there’s no need to fetch the comm_group. + if self._typ.communicator_id in ctx.communicators: + self._accelerator_group: "Communicator" = ctx.communicators[ + self._typ.communicator_id + ] + assert ( + self._accelerator_group is not None + ), "ChannelContext.accelerator_group is not initialized." + + self._writer_rank = self._accelerator_group.get_rank(self._writer) + self._reader_ranks = [ + self._accelerator_group.get_rank(reader) + for reader, _ in self._reader_and_node_list + ] + + if ( + self._writer_rank is not None + and self._writer_rank == self._accelerator_group.get_self_rank() + ): + self._writer_registered = True + + if ( + self._reader_ranks + and self._accelerator_group.get_self_rank() in self._reader_ranks + ): + self._reader_registered = True + + # If the channel type specifies that the tensor shape is static, then the + # receiver can allocate buffers without needing to coordinate with the + # sender. We set the metadata on the first send-recv op. Thereafter, + # the sender must ensure that sent tensors match this metadata, and the + # receiver will allocate tensors with this shape. + self._static_tensor_metadata: Optional[List[_TorchTensorMetadata]] = None + self._meta_channel: Optional[Channel] = _meta_channel + if self._meta_channel is None and self._writer_registered: + # We are the writer. Therefore, we also need to allocate a metadata + # channel that will be used to send the shape and dtype of the + # tensor to the receiver(s). + metadata_type = SharedMemoryType() + self._meta_channel = metadata_type.create_channel( + self._writer, + self._reader_and_node_list, + None, + ) + + def ensure_registered_as_writer(self): + assert ( + self._accelerator_group is not None + ), "Actor is not part of an accelerator group" + assert self._writer_registered + ctx = ChannelContext.get_current() + assert ctx.torch_device.type != "cpu" + + def ensure_registered_as_reader(self) -> bool: + assert ( + self._accelerator_group is not None + ), "Actor is not part of an accelerator group" + assert self._reader_registered + ctx = ChannelContext.get_current() + assert ctx.torch_device.type != "cpu" + + def __reduce__(self): + return ( + self.__class__, + ( + self._writer, + self._reader_and_node_list, + self._typ, + self._meta_channel, + ), + ) + + def _get_send_tensors_metadata( + self, tensors: List["torch.Tensor"] + ) -> Optional[List[_TorchTensorMetadata]]: + """ + Helper method to get the metadata that should be sent to the reader so + that they can allocate the proper-sized buffer(s). Throws error if + static_shape=True was set and the given tensors do not match the + inferred shapes. + + Returns: The metadata to send to the reader. None means that we should + not send any metadata message to the reader. + """ + ctx = ChannelContext.get_current() + + # TODO(swang): Currently any exceptions thrown during this method are + # fatal for the DAG because there is no way for the receiver to receive + # the exception. This can be improved by sending the exception through + # the CPU-based non-tensor-data channel, if one exists. The tensor + # channel can send empty data alongside the exception to avoid hanging. + + # Get the shape and dtype of each tensor to send. + metadata_list = [] + for tensor in tensors: + # Basic type checking. + if not isinstance(tensor, self.torch.Tensor): + raise ValueError("Task must return torch.Tensors") + + if tensor.device != ctx.torch_device: + raise ValueError( + f"torch.Tensor must be on the default device: {ctx.torch_device}" + ) + + metadata = _TorchTensorMetadata(tensor.shape, tensor.dtype) + metadata_list.append(metadata) + + if self._static_tensor_metadata is not None: + if metadata_list != self._static_tensor_metadata: + metadata_str = [ + f"(shape={m.shape}, dtype={m.dtype})" for m in metadata_list + ] + expected_str = [ + f"(shape={m.shape}, dtype={m.dtype})" + for m in self._static_tensor_metadata + ] + raise ValueError( + "Expected torch.Tensors with shapes and dtypes: " + "[" + ", ".join(expected_str) + "], " + "found: [" + ", ".join(metadata_str) + "]. " + "DAG will shut down." + ) + # The receiver has already determined the shape and dtype of the + # tensors from a previous send, so no need to send the metadata + # again. + return None + + if self._static_shape: + # The shape and dtype is static. This is the first send op and + # afterwards, a ValueError will be thrown if the sent tensors do + # not match this metadata. + self._static_tensor_metadata = metadata_list + return metadata_list + + def write( + self, + tensors: List["torch.Tensor"], + timeout: Optional[float] = None, + ): + """ + Write a list of tensors via accelerator: + + 1) Send the tensor metadata, i.e. the shape and dtypes of all tensors + via the shared-memory metadata channel. + 2) Send the tensor data via accelerator. + + If static_shape=True was set, then we only perform step (1) on the + first message. The reader is expected to reuse the sent metadata for + subsequent messages. + """ + self.ensure_registered_as_writer() + + import torch + + for tensor in tensors: + assert isinstance( + tensor, torch.Tensor + ), f"{tensor} must be instance of torch.Tensor" + + # Send the tensors metadata so that the receiver knows what buffers to + # allocate. + metadata = self._get_send_tensors_metadata(tensors) + if metadata is not None: + self._meta_channel.write(metadata) + + # NOTE(swang): We must send the metadata *before* launching the accelerator + # send. We are using blocking accelerator ops, so the following calls will + # block until the kernel has been enqueued. Also, peers must launch the + # kernel together before either can proceed. Therefore, we send the + # metadata first so that the receiver can read the metadata and then + # launch the same accelerator op. + for tensor in tensors: + # TODO: If there are multiple readers, can replace with a + # broadcast. + for rank in self._reader_ranks: + self._accelerator_group.send(tensor, rank) + + def _get_recv_tensors_metadata( + self, timeout: Optional[float] = None + ) -> List[_TorchTensorMetadata]: + """ + Get the shape(s) and dtype(s) of the tensors to receive from the + metadata channel. If static_shape=True was set, then we reuse the first + metadata received. + """ + if self._static_tensor_metadata is not None: + return self._static_tensor_metadata + + meta = self._meta_channel.read(timeout) + + if self._static_shape: + self._static_tensor_metadata = meta + + return meta + + def read( + self, + timeout: Optional[float] = None, + ) -> Union["torch.Tensor", List["torch.Tensor"]]: + """ + Receive a list of tensors. + + (1) Receive the tensor metadata via the shared-memory metadata channel. + (2) Allocate buffers on our default device according to the received + tensor metadata. + (3) Receive the tensor data via accelerator. + + If static_data=True was set, then we only perform step (1) on the first + message. Subsequent messages reuse the same metadata. + + NOTE: Currently `timeout` only applies to receiving the CPU-based + tensor metadata. The GPU recv may exceed the timeout without throwing + an error. + """ + self.ensure_registered_as_reader() + + meta_list: List[_TorchTensorMetadata] = self._get_recv_tensors_metadata(timeout) + + bufs: List["torch.Tensor"] = [] + for meta in meta_list: + buf = self._accelerator_group.recv( + meta.shape, meta.dtype, self._writer_rank, _torch_tensor_allocator + ) + bufs.append(buf) + # TODO: Sync CUDA stream after receiving all tensors, instead of after + # each tensor. + return bufs + + def close(self) -> None: + self._meta_channel.close() + + self._accelerator_group.destroy() + ctx = ChannelContext.get_current() + if self._accelerator_group_id in ctx.communicators: + del ctx.communicators[self._accelerator_group_id] + + +def _do_init_communicator( + self, + group_id, + world_size, + comm_id, + rank, + actor_handles, + use_communication_streams, + custom_communicator: Optional[Communicator] = None, +): + if not custom_communicator: + assert ( + AcceleratorContext.get().accelerator_count > 0 + ), "Actors participating in Communication group must have at least one Accelerator assigned" + + ctx = ChannelContext.get_current() + if custom_communicator is not None: + custom_communicator.initialize(rank) + ctx.communicators[group_id] = custom_communicator + else: + # default to CommGroup + ctx.communicators[group_id] = AcceleratorContext.get().create_communicator( + world_size, + comm_id, + rank, + actor_handles, + AcceleratorContext.get().current_stream(), + use_communication_streams, + ) + + +def _do_destroy_communicator(self, group_id): + ctx = ChannelContext.get_current() + if group_id not in ctx.communicators: + return + ctx.communicators[group_id].destroy() + + # Keep the communicator group in the map after destruction in case there is + # still a task loop running. + + +def _do_check_has_accelerators(self) -> str: + return AcceleratorContext.get().accelerator_count > 0 + + +def do_register_accelerator_context(self, name: str, communicator: Type[Communicator]): + register_accelerator_context(name, communicator) + + +def _do_get_unique_communication_id(self) -> bool: + return AcceleratorContext.get().generate_communicator_id() + + +def _get_ranks( + actors: List[ray.actor.ActorHandle], custom_comm_group: Optional[Communicator] +) -> List[int]: + """ + Get ranks for the communicator group to use. If custom_comm_group is specified, + return the ranks of the actors in the custom communicator group, in the same + order of the actors; otherwise, return list(range(len(actors))). + + Args: + actors: A list of actors that participate in the communicator group. + custom_comm_group: The custom communicator group to use. + """ + if custom_comm_group is None: + return list(range(len(actors))) + + assert len(actors) == custom_comm_group.get_world_size(), ( + "The world size of the custom communicator group does not match the " + "number of actors." + ) + ranks = [] + for actor in actors: + rank = custom_comm_group.get_rank(actor) + assert rank not in ranks, "Duplicate rank in custom communicator group" + ranks.append(rank) + assert custom_comm_group.get_world_size() == len(actors), ( + "The world size of the custom communicator group " + f"({custom_comm_group.get_world_size()}) " + "does not match the number of actors " + f"({len(actors)})." + ) + return ranks + + +def _init_communicator( + actors: List[ray.actor.ActorHandle], + custom_communicator: Optional[Communicator] = None, + use_communication_streams: bool = False, + accelerator_module_name: Optional[str] = None, + accelerator_communicator_cls: Optional[Type[Communicator]] = None, +) -> str: + """ + Initialize a communicator group with the given actors. If a custom communicator + group is provided, then it will be used, otherwise a new communicator group + will be created. + + Args: + actors: A list of actors that participate in the communicator group. + custom_communicator: A custom communicator group to initialize. + use_communication_streams: Whether to use dedicated send and recv + streams for communication. If True, communication and computation + can be overlapped to improve performance. + accelerator_module_name: Optional name of the accelerator module to use. + accelerator_communicator_cls: Optional communicator class for the accelerator. + """ + ctx = ChannelContext.get_current() + + is_cpu_communicator = custom_communicator and isinstance( + custom_communicator, CPUCommunicator + ) + + # Register accelerator context for all actors if accelerator is not default + if accelerator_module_name and accelerator_communicator_cls: + if is_accelerator_context_registered(): + ray.get( + [ + actor.__ray_call__.remote( + do_register_accelerator_context, + accelerator_module_name, + accelerator_communicator_cls, + ) + for actor in actors + ] + ) + + has_accelerators = ray.get( + [actor.__ray_call__.remote(_do_check_has_accelerators) for actor in actors] + ) + for has_accelerator, actor in zip(has_accelerators, actors): + if not has_accelerator and not is_cpu_communicator: + raise ValueError( + f"Actor {actor} returns a tensor with type hint " + 'TorchTensor(transport="accelerator") or ' + "TorchTensor(transport=accelerator_group_handle) " + "but actor does not have an accelerator assigned by Ray." + ) + + actor_ids = {actor._ray_actor_id for actor in actors} + assert len(actor_ids) == len(actors), "Actors must be unique" + + # Allocate a communicator ID on one of the actors that will participate in + # the group. This is in case the driver is not on the same node as one of + # the communicator actors. + comm_id = ray.get(actors[0].__ray_call__.remote(_do_get_unique_communication_id)) + + # Used to uniquely identify this communicator group. + group_id = str(uuid.uuid4()) + + if custom_communicator is not None: + logger.info( + f"Initializing custom communicator group {group_id} on actors: {actors}" + ) + else: + logger.info(f"Creating communicator group {group_id} on actors: {actors}") + + world_size = len(actors) + ranks = _get_ranks(actors, custom_communicator) + init_tasks = [ + actor.__ray_call__.remote( + _do_init_communicator, + group_id, + world_size, + comm_id, + rank, + actors, + use_communication_streams, + custom_communicator, + ) + for rank, actor in zip(ranks, actors) + ] + try: + ray.get(init_tasks, timeout=30) + except ray.exceptions.GetTimeoutError: + logger.warning( + "Communicator group creation not done after 30s. communicator group" + "creation may be hung." + ) + ray.get(init_tasks) + + logger.info("Communicator group initialized.") + + if custom_communicator is not None: + ctx.communicator_handles[group_id] = CommunicatorHandle( + actor_handles=custom_communicator.get_actor_handles(), + ) + else: + ctx.communicator_handles[group_id] = CommunicatorHandle( + actor_handles=actors, + ) + + return group_id + + +def _destroy_communicator(group_id: str) -> None: + """ + Destroy the communicator group with the given ID. + """ + ctx = ChannelContext.get_current() + if group_id not in ctx.communicator_handles: + return + + group = ctx.communicator_handles[group_id] + actors = group.get_actor_handles() + destroy_tasks = [ + actor.__ray_call__.remote( + _do_destroy_communicator, + group_id, + ) + for actor in actors + ] + + _, unready = ray.wait(destroy_tasks, timeout=30, num_returns=len(destroy_tasks)) + if unready: + logger.warning( + "Communicator group destruction not done after 30s. Communicator" + "group destruction may be hung." + ) + + del ctx.communicator_handles[group_id] diff --git a/lib/python3.12/site-packages/ray/experimental/channel/torch_tensor_type.py b/lib/python3.12/site-packages/ray/experimental/channel/torch_tensor_type.py new file mode 100644 index 0000000000000000000000000000000000000000..21220acf995970431211792cc0cd1744f15cd9c0 --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/channel/torch_tensor_type.py @@ -0,0 +1,193 @@ +import logging +from typing import TYPE_CHECKING, List, Optional, Tuple, Union + +import ray +from ray.experimental.channel import ChannelContext, ChannelOutputType +from ray.experimental.channel.communicator import Communicator +from ray.experimental.channel.shared_memory_channel import SharedMemoryType +from ray.experimental.util.types import Device +from ray.util.annotations import PublicAPI + +if TYPE_CHECKING: + from ray.experimental.channel.shared_memory_channel import Channel + +logger = logging.getLogger(__name__) + + +@PublicAPI(stability="alpha") +class TorchTensorType(ChannelOutputType): + AUTO = "auto" + CPU = "cpu" + ACCELERATOR = "accelerator" + + def __init__( + self, + transport: Optional[Union[str, Communicator]] = AUTO, + device: Device = Device.DEFAULT, + _static_shape: bool = False, + _direct_return: Optional[bool] = False, + ): + """ + A type hint that can be used to annotate DAG nodes that return a + torch.Tensor. + + NOTE: Use of this type in the DAG will register a custom serializer for + torch.Tensor that moves the tensor to the correct device on the + receiver. If you are using ray.cloudpickle to serialize objects and you + do not want this behavior, deregister the custom serializer using + ray.util.serialization.deregister_serializer(torch.Tensor). + + Args: + transport: "auto" (default) means that tensors will be passed via + host memory, using numpy as the serialization format. Pass + TorchTensorType.ACCELERATOR or "accelerator" to use accelerator + instead, avoiding the host memory copy. + device: Target device for tensor transport. Options: + - "default": Retains the same device type as the sender. + - "cpu": Moves tensor to CPU on the receiver. Not compatible + with accelerator transport. + - "gpu" or "cuda": Moves tensor to GPU on the receiver. + _static_shape: A hint indicating whether the shape(s) and dtype(s) + of tensor(s) contained in this value always remain the same + across different executions of the DAG. + _direct_return: Whether the tensor is sent directly or inside of + other data. If a non-default `transport` is used, this allows + the sender and receiver to eliminate performance overhead from + an additional data transfer. + + NOTE: Setting static_shape=True and _direct_return=True can improve + performance if a non-default transport is used. However, if either flag + is set, then the user must ensure that the condition is met. + + If using this type as a Compiled Graph annotation, an exception will + be thrown in the following cases, and the DAG will be torn down. To + continue execution, a new DAG must be created: + 1. If _static_shape=True, and the found tensors don't match the + previous shape or dtype(s). + 2. If _direct_return=True, and the returned value is not a + torch.Tensor. + """ + super().__init__() + + self._device = device + self._static_shape = _static_shape + self._direct_return = _direct_return + + self._communicator: Optional[Communicator] = None + if isinstance(transport, Communicator): + self._communicator = transport + transport = transport.get_transport_name() + + if transport not in [self.AUTO, self.CPU, self.ACCELERATOR]: + raise ValueError( + "`transport` must be TorchTensorType.AUTO, TorchTensorType.ACCELERATOR " + "or TorchTensorType.CPU" + ) + if device == Device.CPU and transport == self.ACCELERATOR: + raise ValueError( + "accelerator transport is not supported with CPU target device." + ) + self.transport = transport + + self._communicator_id: Optional[str] = None + + if self._static_shape and self.transport == self.AUTO: + logger.info( + "TorchTensorType(_static_shape=True) has no effect when " + "`transport` is TorchTensorType.AUTO (default)." + ) + if self._direct_return and self.transport == self.AUTO: + logger.info( + "TorchTensorType(_direct_return=True) has no effect when " + "`transport` is TorchTensorType.AUTO (default)." + ) + + @property + def device(self) -> Device: + return self._device + + @property + def static_shape(self): + return self._static_shape + + @property + def direct_return(self): + return self._direct_return + + def register_custom_serializer(self) -> None: + super().register_custom_serializer() + + import torch + + def serialize(t): + ctx = ChannelContext.get_current() + return ctx.serialization_context.serialize_tensor(t) + + def deserialize(b): + ctx = ChannelContext.get_current() + return ctx.serialization_context.deserialize_tensor(b, self.device) + + ray.util.serialization.register_serializer( + torch.Tensor, + serializer=serialize, + deserializer=deserialize, + ) + + def create_channel( + self, + writer: Optional["ray.actor.ActorHandle"], + reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]], + driver_actor_id: Optional[str] = None, + _cpu_data_channel: Optional["Channel"] = None, + _tensor_metadata_channel: Optional["Channel"] = None, + ) -> type: + if self.requires_accelerator(): + from ray.experimental.channel.torch_tensor_accelerator_channel import ( + TorchTensorAcceleratorChannel, + ) + + return TorchTensorAcceleratorChannel( + writer, + reader_and_node_list, + self, + driver_actor_id, + _tensor_metadata_channel, + _cpu_data_channel, + ) + + # Data does not require accelerator. Transfer via host memory using a + # shared-memory channel. + # TODO(swang): Allow the initial max buffer size to be overridden. + typ = SharedMemoryType() + return typ.create_channel(writer, reader_and_node_list, driver_actor_id) + + def requires_accelerator(self) -> bool: + return self.transport == self.ACCELERATOR + + def get_custom_communicator(self) -> Optional[Communicator]: + """ + Return the communicator group if one is specified. + """ + return self._communicator + + def set_communicator_id(self, group_id: str) -> None: + self._communicator_id = group_id + + @property + def communicator_id(self) -> Optional[str]: + return self._communicator_id + + def __deepcopy__(self, memo): + """ + Deep copy all the fields except for the communicator group. The communicator + group should not be deep copied because it can be shared across `TorchTensorType` + instances. + """ + copy = TorchTensorType( + transport=self.transport, + _static_shape=self._static_shape, + _direct_return=self._direct_return, + ) + copy._communicator = self._communicator + copy._communicator_id = self._communicator_id + return copy diff --git a/lib/python3.12/site-packages/ray/experimental/channel/utils.py b/lib/python3.12/site-packages/ray/experimental/channel/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..6df57828130a803f680dc2164784eb2082b39a72 --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/channel/utils.py @@ -0,0 +1,92 @@ +from typing import List, Optional, Tuple + +import ray + + +def get_self_actor() -> Optional["ray.actor.ActorHandle"]: + """ + Get the current actor handle in this worker. + If this is called in a driver process, it will return None. + """ + try: + return ray.get_runtime_context().current_actor + except RuntimeError: + return None + + +def split_readers_by_locality( + writer: "ray.actor.ActorHandle", + reader_and_node_list: List[Tuple["ray.actor.ActorHandle", str]], +) -> Tuple[ + List[Tuple["ray.actor.ActorHandle", str]], List[Tuple["ray.actor.ActorHandle", str]] +]: + """Split readers into remote and local readers based on writer. + + Args: + writer: The actor handle of the writer + reader_and_node_list: List of (reader, node) tuples + + Returns: + Tuple containing: + - List of (reader, node) tuples for remote readers + - List of (reader, node) tuples for local readers + """ + remote_readers = [] + local_readers = [] + + for reader, node in reader_and_node_list: + if reader != writer: + remote_readers.append((reader, node)) + else: + local_readers.append((reader, node)) + + return remote_readers, local_readers + + +def split_actors_by_node_locality( + node: str, + actor_and_node_list: List[Tuple["ray.actor.ActorHandle", str]], +) -> Tuple[ + List[Tuple["ray.actor.ActorHandle", str]], List[Tuple["ray.actor.ActorHandle", str]] +]: + """Split actors into remote and local actors based on node. The local actors will be + on the same node as the given node. The remote actors will be on a different node. + + Args: + writer_node: The node of the writer + actor_and_node_list: List of (actor, node) tuples + + Returns: + Tuple containing: + - List of (actor, node) tuples for actors on the same node + - List of (actor, node) tuples for actors on a different node + """ + actors_on_same_node = [] + actors_on_different_node = [] + + for actor, actor_node in actor_and_node_list: + if node == actor_node: + actors_on_same_node.append((actor, actor_node)) + else: + actors_on_different_node.append((actor, actor_node)) + + return actors_on_same_node, actors_on_different_node + + +def get_actor_node(actor: Optional["ray.actor.ActorHandle"]) -> str: + """Get the node of the actor. + + Args: + actor: The actor handle of the actor + + Returns: + The node of the actor + """ + if actor is None or actor == ray.get_runtime_context().current_actor: + return ray.get_runtime_context().get_node_id() + else: + return ray.get( + actor.__ray_call__.remote( + lambda self: ray.get_runtime_context().get_node_id() + ) + ) diff --git a/lib/python3.12/site-packages/ray/experimental/job/__init__.py b/lib/python3.12/site-packages/ray/experimental/job/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/experimental/job/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/job/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..49ba93091b4cc617bba9b0f41f5f2ea38aaad7d7 Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/job/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/job/example_job/__init__.py b/lib/python3.12/site-packages/ray/experimental/job/example_job/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/experimental/job/example_job/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/job/example_job/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..14dc58c27a1361aa9eeac0d34f6661063ffb3b6c Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/job/example_job/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/job/example_job/__pycache__/demo_script.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/job/example_job/__pycache__/demo_script.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..57a4f03a73f96f4ce1310728e89743fa2a1772ce Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/job/example_job/__pycache__/demo_script.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/job/example_job/demo_script.py b/lib/python3.12/site-packages/ray/experimental/job/example_job/demo_script.py new file mode 100644 index 0000000000000000000000000000000000000000..8bcc79611e26d404a7f7d0dbd86b22f1424420c9 --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/job/example_job/demo_script.py @@ -0,0 +1,77 @@ +# Regular ray application that user wrote and runs on local cluster. +# intermediate status are dumped to GCS +import argparse +import time + +import ray +import ray.experimental.internal_kv as ray_kv + + +@ray.remote +class StepActor: + def __init__(self, interval_s=1, total_steps=3): + self.interval_s = interval_s + self.stopped = False + self.current_step = 1 + self.total_steps = total_steps + + worker = ray._private.worker.global_worker + worker_id = worker.core_worker.get_actor_id() + ray_kv._internal_kv_put(f"JOB:{worker_id}", self.current_step, overwrite=True) + + def run(self): + worker = ray._private.worker.global_worker + worker_id = worker.core_worker.get_actor_id() + + while self.current_step <= self.total_steps: + if not self.stopped: + print( + f"Sleeping {self.interval_s} secs to executing " + f"step {self.current_step}" + ) + time.sleep(self.interval_s) + self.current_step += 1 + ray_kv._internal_kv_put( + f"JOB:{worker_id}", self.current_step, overwrite=True + ) + else: + print("Stop called or reached final step.") + break + + self.stopped = True + ray_kv._internal_kv_put(f"JOB:{worker_id}", "DONE", overwrite=True) + return "DONE" + + def get_step(self): + return self.current_step + + def stop(self): + self.stopped = True + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--interval-s", + required=False, + type=int, + default=1, + help="Address to use to connect to Ray", + ) + parser.add_argument( + "--total-steps", + required=False, + type=int, + default=3, + help="Password for connecting to Redis", + ) + args, _ = parser.parse_known_args() + + ray.init() + step_actor = StepActor.remote( + interval_s=args.interval_s, total_steps=args.total_steps + ) + ref = step_actor.run.remote() + print(ray.get([ref])) + job_key = ray_kv._internal_kv_list("JOB:")[0] + print(f"{job_key}, {ray_kv._internal_kv_get(job_key)}") diff --git a/lib/python3.12/site-packages/ray/experimental/multiprocessing/__init__.py b/lib/python3.12/site-packages/ray/experimental/multiprocessing/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f9b741f1ebfc72812a214fe4c52b2894bd965602 --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/multiprocessing/__init__.py @@ -0,0 +1,5 @@ +from multiprocessing import TimeoutError + +from .pool import Pool + +__all__ = ["Pool", "TimeoutError"] diff --git a/lib/python3.12/site-packages/ray/experimental/multiprocessing/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/multiprocessing/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..001ed3aeb3effc2cc9ee5e62b3d5c10fc544ed15 Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/multiprocessing/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/multiprocessing/__pycache__/pool.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/multiprocessing/__pycache__/pool.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..66ecaa95410ae636b4b08460f992e9fc01df1932 Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/multiprocessing/__pycache__/pool.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/multiprocessing/pool.py b/lib/python3.12/site-packages/ray/experimental/multiprocessing/pool.py new file mode 100644 index 0000000000000000000000000000000000000000..bf0e442ff04b11a66baa1ac496d1acf4f6d55299 --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/multiprocessing/pool.py @@ -0,0 +1,5 @@ +from ray.util import multiprocessing + + +class Pool(multiprocessing.Pool): + pass # moved to util package diff --git a/lib/python3.12/site-packages/ray/experimental/state/__init__.py b/lib/python3.12/site-packages/ray/experimental/state/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/experimental/state/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/state/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b5ccef4b9d568db7719c8f4cf3547b1654349902 Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/state/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/state/__pycache__/api.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/state/__pycache__/api.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a4f7ca508966879b246de4db0c8b825649f336cc Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/state/__pycache__/api.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/state/__pycache__/common.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/state/__pycache__/common.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f4fc27ccda4b500778d4fc13d775991f5a9c5c41 Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/state/__pycache__/common.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/state/__pycache__/custom_types.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/state/__pycache__/custom_types.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..430e7247cab82907a7543a17a0e96629102429f5 Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/state/__pycache__/custom_types.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/state/__pycache__/exception.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/state/__pycache__/exception.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..77baca8f994a8071bf2d599551e5d8cb86c57a13 Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/state/__pycache__/exception.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/state/__pycache__/state_cli.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/state/__pycache__/state_cli.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0c163a81fa8874edafe506946de2bac506e61907 Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/state/__pycache__/state_cli.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/state/__pycache__/state_manager.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/state/__pycache__/state_manager.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..acc9f27309a49a55a211c99c2ea39594c33c2f67 Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/state/__pycache__/state_manager.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/state/__pycache__/util.cpython-312.pyc b/lib/python3.12/site-packages/ray/experimental/state/__pycache__/util.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c7f59d2d3f45b1a4ef46f629e5fb82eaa7f68075 Binary files /dev/null and b/lib/python3.12/site-packages/ray/experimental/state/__pycache__/util.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/experimental/state/api.py b/lib/python3.12/site-packages/ray/experimental/state/api.py new file mode 100644 index 0000000000000000000000000000000000000000..f62f5f2ef7cd375c90224e3f71c3d041d7d029a7 --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/state/api.py @@ -0,0 +1,4 @@ +from ray.util.state import * # noqa: F401 F403 +from ray.util.state.util import record_deprecated_state_api_import + +record_deprecated_state_api_import() diff --git a/lib/python3.12/site-packages/ray/experimental/state/common.py b/lib/python3.12/site-packages/ray/experimental/state/common.py new file mode 100644 index 0000000000000000000000000000000000000000..6fbd488cd0fe5fbe200ad981a069c23f5100add3 --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/state/common.py @@ -0,0 +1,4 @@ +from ray.util.state.common import * # noqa: F401 F403 +from ray.util.state.util import record_deprecated_state_api_import + +record_deprecated_state_api_import() diff --git a/lib/python3.12/site-packages/ray/experimental/state/custom_types.py b/lib/python3.12/site-packages/ray/experimental/state/custom_types.py new file mode 100644 index 0000000000000000000000000000000000000000..3dd749f793df398d70445af7c8bc50d3aa991608 --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/state/custom_types.py @@ -0,0 +1,4 @@ +from ray._private.custom_types import * # noqa: F401 F403 +from ray.util.state.util import record_deprecated_state_api_import + +record_deprecated_state_api_import() diff --git a/lib/python3.12/site-packages/ray/experimental/state/exception.py b/lib/python3.12/site-packages/ray/experimental/state/exception.py new file mode 100644 index 0000000000000000000000000000000000000000..49e7099cf325dc106a1927ccc64a7e3917c7a68c --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/state/exception.py @@ -0,0 +1,4 @@ +from ray.util.state.exception import * # noqa: F401 F403 +from ray.util.state.util import record_deprecated_state_api_import + +record_deprecated_state_api_import() diff --git a/lib/python3.12/site-packages/ray/experimental/state/state_cli.py b/lib/python3.12/site-packages/ray/experimental/state/state_cli.py new file mode 100644 index 0000000000000000000000000000000000000000..58bc5d31ebdcc5e523af8da8c8397a464909b742 --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/state/state_cli.py @@ -0,0 +1,4 @@ +from ray.util.state.state_cli import * # noqa: F401 F403 +from ray.util.state.util import record_deprecated_state_api_import + +record_deprecated_state_api_import() diff --git a/lib/python3.12/site-packages/ray/experimental/state/state_manager.py b/lib/python3.12/site-packages/ray/experimental/state/state_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..7720606a68b9a6a6862518ce0ebeba01bcdb276f --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/state/state_manager.py @@ -0,0 +1,4 @@ +from ray.util.state.state_manager import * # noqa: F401 F403 +from ray.util.state.util import record_deprecated_state_api_import + +record_deprecated_state_api_import() diff --git a/lib/python3.12/site-packages/ray/experimental/state/util.py b/lib/python3.12/site-packages/ray/experimental/state/util.py new file mode 100644 index 0000000000000000000000000000000000000000..24a26dd72982292593f06fb91bca45e7c765cff8 --- /dev/null +++ b/lib/python3.12/site-packages/ray/experimental/state/util.py @@ -0,0 +1,4 @@ +from ray.util.state.util import * # noqa: F401 F403 +from ray.util.state.util import record_deprecated_state_api_import + +record_deprecated_state_api_import() diff --git a/lib/python3.12/site-packages/ray/rllib/__init__.py b/lib/python3.12/site-packages/ray/rllib/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6b681ce4d78fae6e150031a7edc92b200e13a3e6 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/__init__.py @@ -0,0 +1,55 @@ +import logging + +from ray._common.usage import usage_lib + +# Note: do not introduce unnecessary library dependencies here, e.g. gym. +# This file is imported from the tune module in order to register RLlib agents. +from ray.rllib.env.base_env import BaseEnv +from ray.rllib.env.external_env import ExternalEnv +from ray.rllib.env.multi_agent_env import MultiAgentEnv +from ray.rllib.env.vector_env import VectorEnv +from ray.rllib.evaluation.rollout_worker import RolloutWorker +from ray.rllib.policy.policy import Policy +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.policy.tf_policy import TFPolicy +from ray.rllib.policy.torch_policy import TorchPolicy +from ray.tune.registry import register_trainable + + +def _setup_logger(): + logger = logging.getLogger("ray.rllib") + handler = logging.StreamHandler() + handler.setFormatter( + logging.Formatter( + "%(asctime)s\t%(levelname)s %(filename)s:%(lineno)s -- %(message)s" + ) + ) + logger.addHandler(handler) + logger.propagate = False + + +def _register_all(): + from ray.rllib.algorithms.registry import ALGORITHMS, _get_algorithm_class + + for key, get_trainable_class_and_config in ALGORITHMS.items(): + register_trainable(key, get_trainable_class_and_config()[0]) + + for key in ["__fake", "__sigmoid_fake_data", "__parameter_tuning"]: + register_trainable(key, _get_algorithm_class(key)) + + +_setup_logger() + +usage_lib.record_library_usage("rllib") + +__all__ = [ + "Policy", + "TFPolicy", + "TorchPolicy", + "RolloutWorker", + "SampleBatch", + "BaseEnv", + "MultiAgentEnv", + "VectorEnv", + "ExternalEnv", +] diff --git a/lib/python3.12/site-packages/ray/rllib/core/__init__.py b/lib/python3.12/site-packages/ray/rllib/core/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..42404d51e5d37313a93e999bac4c2e99126337d9 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/__init__.py @@ -0,0 +1,34 @@ +from ray.rllib.core.columns import Columns + +DEFAULT_AGENT_ID = "default_agent" +DEFAULT_POLICY_ID = "default_policy" +# TODO (sven): Change this to "default_module" +DEFAULT_MODULE_ID = DEFAULT_POLICY_ID +ALL_MODULES = "__all_modules__" + +COMPONENT_ENV_RUNNER = "env_runner" +COMPONENT_ENV_TO_MODULE_CONNECTOR = "env_to_module_connector" +COMPONENT_EVAL_ENV_RUNNER = "eval_env_runner" +COMPONENT_LEARNER = "learner" +COMPONENT_LEARNER_GROUP = "learner_group" +COMPONENT_METRICS_LOGGER = "metrics_logger" +COMPONENT_MODULE_TO_ENV_CONNECTOR = "module_to_env_connector" +COMPONENT_OPTIMIZER = "optimizer" +COMPONENT_RL_MODULE = "rl_module" + + +__all__ = [ + "Columns", + "COMPONENT_ENV_RUNNER", + "COMPONENT_ENV_TO_MODULE_CONNECTOR", + "COMPONENT_EVAL_ENV_RUNNER", + "COMPONENT_LEARNER", + "COMPONENT_LEARNER_GROUP", + "COMPONENT_METRICS_LOGGER", + "COMPONENT_MODULE_TO_ENV_CONNECTOR", + "COMPONENT_OPTIMIZER", + "COMPONENT_RL_MODULE", + "DEFAULT_AGENT_ID", + "DEFAULT_MODULE_ID", + "DEFAULT_POLICY_ID", +] diff --git a/lib/python3.12/site-packages/ray/rllib/core/columns.py b/lib/python3.12/site-packages/ray/rllib/core/columns.py new file mode 100644 index 0000000000000000000000000000000000000000..98cb8646913e2b333eab3243308b82419f05d9bc --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/columns.py @@ -0,0 +1,73 @@ +from ray.util.annotations import DeveloperAPI + + +@DeveloperAPI +class Columns: + """Definitions of common column names for RL data, e.g. 'obs', 'rewards', etc.. + + Note that this replaces the `SampleBatch` and `Postprocessing` columns (of the same + name). + """ + + # Observation received from an environment after `reset()` or `step()`. + OBS = "obs" + # Infos received from an environment after `reset()` or `step()`. + INFOS = "infos" + + # Action computed/sampled by an RLModule. + ACTIONS = "actions" + # Action actually sent to the (gymnasium) `Env.step()` method. + ACTIONS_FOR_ENV = "actions_for_env" + # Reward returned by `env.step()`. + REWARDS = "rewards" + # Termination signal received from an environment after `step()`. + TERMINATEDS = "terminateds" + # Truncation signal received from an environment after `step()` (e.g. because + # of a reached time limit). + TRUNCATEDS = "truncateds" + + # Next observation: Only used by algorithms that need to look at TD-data for + # training, such as off-policy/DQN algos. + NEXT_OBS = "new_obs" + + # Uniquely identifies an episode + EPS_ID = "eps_id" + AGENT_ID = "agent_id" + MODULE_ID = "module_id" + + # The size of non-zero-padded data within a (e.g. LSTM) zero-padded + # (B, T, ...)-style train batch. + SEQ_LENS = "seq_lens" + # Episode timestep counter. + T = "t" + + # Common extra RLModule output keys. + STATE_IN = "state_in" + NEXT_STATE_IN = "next_state_in" + STATE_OUT = "state_out" + NEXT_STATE_OUT = "next_state_out" + EMBEDDINGS = "embeddings" + ACTION_DIST_INPUTS = "action_dist_inputs" + ACTION_PROB = "action_prob" + ACTION_LOGP = "action_logp" + + # Value function predictions. + VF_PREDS = "vf_preds" + # Values, predicted at one timestep beyond the last timestep taken. + # These are usually calculated via the value function network using the final + # observation (and in case of an RNN: the last returned internal state). + VALUES_BOOTSTRAPPED = "values_bootstrapped" + + # Postprocessing columns. + ADVANTAGES = "advantages" + VALUE_TARGETS = "value_targets" + + # Intrinsic rewards (learning with curiosity). + INTRINSIC_REWARDS = "intrinsic_rewards" + # Discounted sum of rewards till the end of the episode (or chunk). + RETURNS_TO_GO = "returns_to_go" + + # Loss mask. If provided in a train batch, a Learner's compute_loss_for_module + # method should respect the False-set value in here and mask out the respective + # items form the loss. + LOSS_MASK = "loss_mask" diff --git a/lib/python3.12/site-packages/ray/rllib/core/models/__init__.py b/lib/python3.12/site-packages/ray/rllib/core/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/core/models/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/models/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2a20702f3229ce8850e5a6b4c10f34aa25b9a82a Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/models/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/models/__pycache__/base.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/models/__pycache__/base.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..902ea69b7fe14ff0734605ffa89858b041db060d Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/models/__pycache__/base.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/models/__pycache__/catalog.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/models/__pycache__/catalog.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e139d361d571e84f486bdc8e51ce2f90e6ed194f Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/models/__pycache__/catalog.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/models/__pycache__/configs.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/models/__pycache__/configs.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..85e3e044a4cf75a54dcb3b474a16e8639fcbfbde Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/models/__pycache__/configs.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/models/base.py b/lib/python3.12/site-packages/ray/rllib/core/models/base.py new file mode 100644 index 0000000000000000000000000000000000000000..337b35421263b753b7d64462eb15457d50846ca1 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/models/base.py @@ -0,0 +1,444 @@ +import abc +from typing import List, Optional, Tuple, Union + +from ray.rllib.core.columns import Columns +from ray.rllib.core.models.configs import ModelConfig +from ray.rllib.core.models.specs.specs_base import Spec +from ray.rllib.policy.rnn_sequencing import get_fold_unfold_fns +from ray.rllib.utils.annotations import ExperimentalAPI, override +from ray.rllib.utils.typing import TensorType +from ray.util.annotations import DeveloperAPI + +# Top level keys that unify model i/o. +ENCODER_OUT: str = "encoder_out" +# For Actor-Critic algorithms, these signify data related to the actor and critic +ACTOR: str = "actor" +CRITIC: str = "critic" + + +@ExperimentalAPI +class Model(abc.ABC): + """Framework-agnostic base class for RLlib models. + + Models are low-level neural network components that offer input- and + output-specification, a forward method, and a get_initial_state method. Models + are composed in RLModules. + + Usage Example together with ModelConfig: + + .. testcode:: + + from ray.rllib.core.models.base import Model + from ray.rllib.core.models.configs import ModelConfig + from ray.rllib.core.models.configs import ModelConfig + from dataclasses import dataclass + + class MyModel(Model): + def __init__(self, config): + super().__init__(config) + self.my_param = config.my_param * 2 + + def _forward(self, input_dict): + return input_dict["obs"] * self.my_param + + + @dataclass + class MyModelConfig(ModelConfig): + my_param: int = 42 + + def build(self, framework: str): + if framework == "bork": + return MyModel(self) + + + config = MyModelConfig(my_param=3) + model = config.build(framework="bork") + print(model._forward({"obs": 1})) + + .. testoutput:: + + 6 + + """ + + def __init__(self, config: ModelConfig): + self.config = config + + def __init_subclass__(cls, **kwargs): + # Automatically add a __post_init__ method to all subclasses of Model. + # This method is called after the __init__ method of the subclass. + def init_decorator(previous_init): + def new_init(self, *args, **kwargs): + previous_init(self, *args, **kwargs) + if type(self) is cls: + self.__post_init__() + + return new_init + + cls.__init__ = init_decorator(cls.__init__) + + def __post_init__(self): + """Called automatically after the __init__ method of the subclasses. + + The module first calls the __init__ method of the subclass, With in the + __init__ you should call the super().__init__ method. Then after the __init__ + method of the subclass is called, the __post_init__ method is called. + + This is a good place to do any initialization that requires access to the + subclass's attributes. + """ + self._input_specs = self.get_input_specs() + self._output_specs = self.get_output_specs() + + def get_input_specs(self) -> Optional[Spec]: + """Returns the input specs of this model. + + Override `get_input_specs` to define your own input specs. + This method should not be called often, e.g. every forward pass. + Instead, it should be called once at instantiation to define Model.input_specs. + + Returns: + Spec: The input specs. + """ + return None + + def get_output_specs(self) -> Optional[Spec]: + """Returns the output specs of this model. + + Override `get_output_specs` to define your own output specs. + This method should not be called often, e.g. every forward pass. + Instead, it should be called once at instantiation to define Model.output_specs. + + Returns: + Spec: The output specs. + """ + return None + + @property + def input_specs(self) -> Spec: + """Returns the input spec of this model.""" + return self._input_specs + + @input_specs.setter + def input_specs(self, spec: Spec) -> None: + raise ValueError( + "`input_specs` cannot be set directly. Override " + "Model.get_input_specs() instead. Set Model._input_specs if " + "you want to override this behavior." + ) + + @property + def output_specs(self) -> Spec: + """Returns the output specs of this model.""" + return self._output_specs + + @output_specs.setter + def output_specs(self, spec: Spec) -> None: + raise ValueError( + "`output_specs` cannot be set directly. Override " + "Model.get_output_specs() instead. Set Model._output_specs if " + "you want to override this behavior." + ) + + def get_initial_state(self) -> Union[dict, List[TensorType]]: + """Returns the initial state of the Model. + + It can be left empty if this Model is not stateful. + """ + return dict() + + @abc.abstractmethod + def _forward(self, input_dict: dict, **kwargs) -> dict: + """Returns the output of this model for the given input. + + This method is called by the forwarding method of the respective framework + that is itself wrapped by RLlib in order to check model inputs and outputs. + + Args: + input_dict: The input tensors. + **kwargs: Forward compatibility kwargs. + + Returns: + dict: The output tensors. + """ + + @abc.abstractmethod + def get_num_parameters(self) -> Tuple[int, int]: + """Returns a tuple of (num trainable params, num non-trainable params).""" + + @abc.abstractmethod + def _set_to_dummy_weights(self, value_sequence=(-0.02, -0.01, 0.01, 0.02)) -> None: + """Helper method to set all weights to deterministic dummy values. + + Calling this method on two `Models` that have the same architecture using + the exact same `value_sequence` arg should make both models output the exact + same values on arbitrary inputs. This will work, even if the two `Models` + are of different DL frameworks. + + Args: + value_sequence: Looping through the list of all parameters (weight matrices, + bias tensors, etc..) of this model, in each iteration i, we set all + values in this parameter to `value_sequence[i % len(value_sequence)]` + (round robin). + + Example: + TODO: + """ + + +@ExperimentalAPI +class Encoder(Model, abc.ABC): + """The framework-agnostic base class for all RLlib encoders. + + Encoders are used to transform observations to a latent space. + Therefore, their `input_specs` contains the observation space dimensions. + Similarly, their `output_specs` contains the latent space dimensions. + Encoders can be recurrent, in which case the state should be part of input- and + output_specs. The latent vectors produced by an encoder are fed into subsequent + "heads". Any implementation of Encoder should also be callable. This should be done + by also inheriting from a framework-specific model base-class, s.a. TorchModel or + TfModel. + + Abstract illustration of typical flow of tensors: + + Inputs + | + Encoder + | \ + Head1 Head2 + | / + Outputs + + Outputs of encoders are generally of shape (B, latent_dim) or (B, T, latent_dim). + That is, for time-series data, we encode into the latent space for each time step. + This should be reflected in the `output_specs`. + + Usage example together with a ModelConfig: + + .. testcode:: + + from dataclasses import dataclass + import numpy as np + + from ray.rllib.core.columns import Columns + from ray.rllib.core.models.base import Encoder, ENCODER_OUT + from ray.rllib.core.models.configs import ModelConfig + from ray.rllib.policy.sample_batch import SampleBatch + + class NumpyEncoder(Encoder): + def __init__(self, config): + super().__init__(config) + self.factor = config.factor + + def __call__(self, *args, **kwargs): + # This is a dummy method to do checked forward passes. + return self._forward(*args, **kwargs) + + def _forward(self, input_dict, **kwargs): + obs = input_dict[Columns.OBS] + return { + ENCODER_OUT: np.array(obs) * self.factor, + Columns.STATE_OUT: ( + np.array(input_dict[Columns.STATE_IN]) + * self.factor + ), + } + + @dataclass + class NumpyEncoderConfig(ModelConfig): + factor: int = None + + def build(self, framework: str): + return NumpyEncoder(self) + + config = NumpyEncoderConfig(factor=2) + encoder = NumpyEncoder(config) + print(encoder({Columns.OBS: 1, Columns.STATE_IN: 2})) + + .. testoutput:: + + {'encoder_out': 2, 'state_out': 4} + + """ + + @abc.abstractmethod + def _forward(self, input_dict: dict, **kwargs) -> dict: + """Returns the latent of the encoder for the given inputs. + + This method is called by the forwarding method of the respective framework + that is itself wrapped by RLlib in order to check model inputs and outputs. + + The input dict contains at minimum the observation and the state of the encoder + (None for stateless encoders). + The output dict contains at minimum the latent and the state of the encoder + (None for stateless encoders). + To establish an agreement between the encoder and RLModules, these values + have the fixed keys `Columns.OBS` for the `input_dict`, + and `ACTOR` and `CRITIC` for the returned dict. + + Args: + input_dict: The input tensors. Must contain at a minimum the keys + Columns.OBS and Columns.STATE_IN (which might be None for stateless + encoders). + **kwargs: Forward compatibility kwargs. + + Returns: + The output tensors. Must contain at a minimum the key ENCODER_OUT. + """ + + +@ExperimentalAPI +class ActorCriticEncoder(Encoder): + """An encoder that potentially holds two stateless encoders. + + This is a special case of Encoder that can either enclose a single, + shared encoder or two separate encoders: One for the actor and one for the + critic. The two encoders are of the same type, and we can therefore make the + assumption that they have the same input and output specs. + """ + + framework = None + + def __init__(self, config: ModelConfig) -> None: + super().__init__(config) + + if config.shared: + self.encoder = config.base_encoder_config.build(framework=self.framework) + else: + self.actor_encoder = config.base_encoder_config.build( + framework=self.framework + ) + self.critic_encoder = None + if not config.inference_only: + self.critic_encoder = config.base_encoder_config.build( + framework=self.framework + ) + + @override(Model) + def _forward(self, inputs: dict, **kwargs) -> dict: + if self.config.shared: + encoder_outs = self.encoder(inputs, **kwargs) + return { + ENCODER_OUT: { + ACTOR: encoder_outs[ENCODER_OUT], + **( + {} + if self.config.inference_only + else {CRITIC: encoder_outs[ENCODER_OUT]} + ), + } + } + else: + # Encoders should not modify inputs, so we can pass the same inputs + actor_out = self.actor_encoder(inputs, **kwargs) + if self.critic_encoder: + critic_out = self.critic_encoder(inputs, **kwargs) + + return { + ENCODER_OUT: { + ACTOR: actor_out[ENCODER_OUT], + **( + {} + if self.config.inference_only + else {CRITIC: critic_out[ENCODER_OUT]} + ), + } + } + + +@ExperimentalAPI +class StatefulActorCriticEncoder(Encoder): + """An encoder that potentially holds two potentially stateful encoders. + + This is a special case of Encoder that can either enclose a single, + shared encoder or two separate encoders: One for the actor and one for the + critic. The two encoders are of the same type, and we can therefore make the + assumption that they have the same input and output specs. + + If this encoder wraps a single encoder, state in input- and output dicts + is simply stored under the key `STATE_IN` and `STATE_OUT`, respectively. + If this encoder wraps two encoders, state in input- and output dicts is + stored under the keys `(STATE_IN, ACTOR)` and `(STATE_IN, CRITIC)` and + `(STATE_OUT, ACTOR)` and `(STATE_OUT, CRITIC)`, respectively. + """ + + framework = None + + def __init__(self, config: ModelConfig) -> None: + super().__init__(config) + + if config.shared: + self.encoder = config.base_encoder_config.build(framework=self.framework) + else: + self.actor_encoder = config.base_encoder_config.build( + framework=self.framework + ) + self.critic_encoder = config.base_encoder_config.build( + framework=self.framework + ) + + @override(Model) + def get_initial_state(self): + if self.config.shared: + return self.encoder.get_initial_state() + else: + return { + ACTOR: self.actor_encoder.get_initial_state(), + CRITIC: self.critic_encoder.get_initial_state(), + } + + @override(Model) + def _forward(self, inputs: dict, **kwargs) -> dict: + outputs = {} + + if self.config.shared: + outs = self.encoder(inputs, **kwargs) + encoder_out = outs.pop(ENCODER_OUT) + outputs[ENCODER_OUT] = {ACTOR: encoder_out, CRITIC: encoder_out} + outputs[Columns.STATE_OUT] = outs[Columns.STATE_OUT] + else: + # Shallow copy inputs so that we can add states without modifying + # original dict. + actor_inputs = inputs.copy() + critic_inputs = inputs.copy() + actor_inputs[Columns.STATE_IN] = inputs[Columns.STATE_IN][ACTOR] + critic_inputs[Columns.STATE_IN] = inputs[Columns.STATE_IN][CRITIC] + + actor_out = self.actor_encoder(actor_inputs, **kwargs) + critic_out = self.critic_encoder(critic_inputs, **kwargs) + + outputs[ENCODER_OUT] = { + ACTOR: actor_out[ENCODER_OUT], + CRITIC: critic_out[ENCODER_OUT], + } + + outputs[Columns.STATE_OUT] = { + ACTOR: actor_out[Columns.STATE_OUT], + CRITIC: critic_out[Columns.STATE_OUT], + } + + return outputs + + +@DeveloperAPI +def tokenize(tokenizer: Encoder, inputs: dict, framework: str) -> dict: + """Tokenizes the observations from the input dict. + + Args: + tokenizer: The tokenizer to use. + inputs: The input dict. + + Returns: + The output dict. + """ + # Tokenizer may depend solely on observations. + obs = inputs[Columns.OBS] + tokenizer_inputs = {Columns.OBS: obs} + size = list(obs.size() if framework == "torch" else obs.shape) + b_dim, t_dim = size[:2] + fold, unfold = get_fold_unfold_fns(b_dim, t_dim, framework=framework) + # Push through the tokenizer encoder. + out = tokenizer(fold(tokenizer_inputs)) + out = out[ENCODER_OUT] + # Then unfold batch- and time-dimensions again. + return unfold(out) diff --git a/lib/python3.12/site-packages/ray/rllib/core/models/catalog.py b/lib/python3.12/site-packages/ray/rllib/core/models/catalog.py new file mode 100644 index 0000000000000000000000000000000000000000..ce3b42592c0bd3b75819e2a1ab1a0b3cdfe25454 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/models/catalog.py @@ -0,0 +1,648 @@ +import dataclasses +import enum +import functools +from typing import Optional + +import gymnasium as gym +import numpy as np +import tree +from gymnasium.spaces import Box, Dict, Discrete, MultiDiscrete, Tuple + +from ray._common.deprecation import DEPRECATED_VALUE, deprecation_warning +from ray.rllib.core.distribution.distribution import Distribution +from ray.rllib.core.models.base import Encoder +from ray.rllib.core.models.configs import ( + CNNEncoderConfig, + MLPEncoderConfig, + ModelConfig, + RecurrentEncoderConfig, +) +from ray.rllib.core.rl_module.default_model_config import DefaultModelConfig +from ray.rllib.models.preprocessors import Preprocessor, get_preprocessor +from ray.rllib.models.utils import get_filter_config +from ray.rllib.utils.annotations import ( + OverrideToImplementCustomLogic, + OverrideToImplementCustomLogic_CallToSuperRecommended, +) +from ray.rllib.utils.error import UnsupportedSpaceException +from ray.rllib.utils.spaces.simplex import Simplex +from ray.rllib.utils.spaces.space_utils import flatten_space, get_base_struct_from_space + + +class Catalog: + """Describes the sub-module-architectures to be used in RLModules. + + RLlib's native RLModules get their Models from a Catalog object. + By default, that Catalog builds the configs it has as attributes. + This component was build to be hackable and extensible. You can inject custom + components into RL Modules by overriding the `build_xxx` methods of this class. + Note that it is recommended to write a custom RL Module for a single use-case. + Modifications to Catalogs mostly make sense if you want to reuse the same + Catalog for different RL Modules. For example if you have written a custom + encoder and want to inject it into different RL Modules (e.g. for PPO, DQN, etc.). + You can influence the decision tree that determines the sub-components by modifying + `Catalog._determine_components_hook`. + + Usage example: + + # Define a custom catalog + + .. testcode:: + + import torch + import gymnasium as gym + from ray.rllib.core.models.configs import MLPHeadConfig + from ray.rllib.core.models.catalog import Catalog + + class MyCatalog(Catalog): + def __init__( + self, + observation_space: gym.Space, + action_space: gym.Space, + model_config_dict: dict, + ): + super().__init__(observation_space, action_space, model_config_dict) + self.my_model_config = MLPHeadConfig( + hidden_layer_dims=[64, 32], + input_dims=[self.observation_space.shape[0]], + ) + + def build_my_head(self, framework: str): + return self.my_model_config.build(framework=framework) + + # With that, RLlib can build and use models from this catalog like this: + catalog = MyCatalog(gym.spaces.Box(0, 1), gym.spaces.Box(0, 1), {}) + my_head = catalog.build_my_head(framework="torch") + + # Make a call to the built model. + out = my_head(torch.Tensor([[1]])) + """ + + # TODO (Sven): Add `framework` arg to c'tor and remove this arg from `build` + # methods. This way, we can already know in the c'tor of Catalog, what the exact + # action distibution objects are and thus what the output dims for e.g. a pi-head + # will be. + def __init__( + self, + observation_space: gym.Space, + action_space: gym.Space, + model_config_dict: dict, + # deprecated args. + view_requirements=DEPRECATED_VALUE, + ): + """Initializes a Catalog with a default encoder config. + + Args: + observation_space: The observation space of the environment. + action_space: The action space of the environment. + model_config_dict: The model config that specifies things like hidden + dimensions and activations functions to use in this Catalog. + """ + if view_requirements != DEPRECATED_VALUE: + deprecation_warning(old="Catalog(view_requirements=..)", error=True) + + # TODO (sven): The following logic won't be needed anymore, once we get rid of + # Catalogs entirely. We will assert directly inside the algo's DefaultRLModule + # class that the `model_config` is a DefaultModelConfig. Thus users won't be + # able to pass in partial config dicts into a default model (alternatively, we + # could automatically augment the user provided dict by the default config + # dataclass object only(!) for default modules). + if dataclasses.is_dataclass(model_config_dict): + model_config_dict = dataclasses.asdict(model_config_dict) + default_config = dataclasses.asdict(DefaultModelConfig()) + # end: TODO + + self.observation_space = observation_space + self.action_space = action_space + + self._model_config_dict = default_config | model_config_dict + self._latent_dims = None + + self._determine_components_hook() + + @OverrideToImplementCustomLogic_CallToSuperRecommended + def _determine_components_hook(self): + """Decision tree hook for subclasses to override. + + By default, this method executes the decision tree that determines the + components that a Catalog builds. You can extend the components by overriding + this or by adding to the constructor of your subclass. + + Override this method if you don't want to use the default components + determined here. If you want to use them but add additional components, you + should call `super()._determine_components()` at the beginning of your + implementation. + + This makes it so that subclasses are not forced to create an encoder config + if the rest of their catalog is not dependent on it or if it breaks. + At the end of this method, an attribute `Catalog.latent_dims` + should be set so that heads can be built using that information. + """ + self._encoder_config = self._get_encoder_config( + observation_space=self.observation_space, + action_space=self.action_space, + model_config_dict=self._model_config_dict, + ) + + # Create a function that can be called when framework is known to retrieve the + # class type for action distributions + self._action_dist_class_fn = functools.partial( + self._get_dist_cls_from_action_space, action_space=self.action_space + ) + + # The dimensions of the latent vector that is output by the encoder and fed + # to the heads. + self.latent_dims = self._encoder_config.output_dims + + @property + def latent_dims(self): + """Returns the latent dimensions of the encoder. + + This establishes an agreement between encoder and heads about the latent + dimensions. Encoders can be built to output a latent tensor with + `latent_dims` dimensions, and heads can be built with tensors of + `latent_dims` dimensions as inputs. This can be safely ignored if this + agreement is not needed in case of modifications to the Catalog. + + Returns: + The latent dimensions of the encoder. + """ + return self._latent_dims + + @latent_dims.setter + def latent_dims(self, value): + self._latent_dims = value + + @OverrideToImplementCustomLogic + def build_encoder(self, framework: str) -> Encoder: + """Builds the encoder. + + By default, this method builds an encoder instance from Catalog._encoder_config. + + You should override this if you want to use RLlib's default RL Modules but + only want to change the encoder. For example, if you want to use a custom + encoder, but want to use RLlib's default heads, action distribution and how + tensors are routed between them. If you want to have full control over the + RL Module, we recommend writing your own RL Module by inheriting from one of + RLlib's RL Modules instead. + + Args: + framework: The framework to use. Either "torch" or "tf2". + + Returns: + The encoder. + """ + assert hasattr(self, "_encoder_config"), ( + "You must define a `Catalog._encoder_config` attribute in your Catalog " + "subclass or override the `Catalog.build_encoder` method. By default, " + "an encoder_config is created in the __post_init__ method." + ) + return self._encoder_config.build(framework=framework) + + @OverrideToImplementCustomLogic + def get_action_dist_cls(self, framework: str): + """Get the action distribution class. + + The default behavior is to get the action distribution from the + `Catalog._action_dist_class_fn`. + + You should override this to have RLlib build your custom action + distribution instead of the default one. For example, if you don't want to + use RLlib's default RLModules with their default models, but only want to + change the distribution that Catalog returns. + + Args: + framework: The framework to use. Either "torch" or "tf2". + + Returns: + The action distribution. + """ + assert hasattr(self, "_action_dist_class_fn"), ( + "You must define a `Catalog._action_dist_class_fn` attribute in your " + "Catalog subclass or override the `Catalog.action_dist_class_fn` method. " + "By default, an action_dist_class_fn is created in the __post_init__ " + "method." + ) + return self._action_dist_class_fn(framework=framework) + + @classmethod + def _get_encoder_config( + cls, + observation_space: gym.Space, + model_config_dict: dict, + action_space: gym.Space = None, + ) -> ModelConfig: + """Returns an EncoderConfig for the given input_space and model_config_dict. + + Encoders are usually used in RLModules to transform the input space into a + latent space that is then fed to the heads. The returned EncoderConfig + objects correspond to the built-in Encoder classes in RLlib. + For example, for a simple 1D-Box input_space, RLlib offers an + MLPEncoder, hence this method returns the MLPEncoderConfig. You can overwrite + this method to produce specific EncoderConfigs for your custom Models. + + The following input spaces lead to the following configs: + - 1D-Box: MLPEncoderConfig + - 3D-Box: CNNEncoderConfig + # TODO (Artur): Support more spaces here + # ... + + Args: + observation_space: The observation space to use. + model_config_dict: The model config to use. + action_space: The action space to use if actions are to be encoded. This + is commonly the case for LSTM models. + + Returns: + The encoder config. + """ + activation = model_config_dict["fcnet_activation"] + output_activation = model_config_dict["fcnet_activation"] + use_lstm = model_config_dict["use_lstm"] + + if use_lstm: + encoder_config = RecurrentEncoderConfig( + input_dims=observation_space.shape, + recurrent_layer_type="lstm", + hidden_dim=model_config_dict["lstm_cell_size"], + hidden_weights_initializer=model_config_dict["lstm_kernel_initializer"], + hidden_weights_initializer_config=model_config_dict[ + "lstm_kernel_initializer_kwargs" + ], + hidden_bias_initializer=model_config_dict["lstm_bias_initializer"], + hidden_bias_initializer_config=model_config_dict[ + "lstm_bias_initializer_kwargs" + ], + batch_major=True, + num_layers=1, + tokenizer_config=cls.get_tokenizer_config( + observation_space, + model_config_dict, + ), + ) + else: + # TODO (Artur): Maybe check for original spaces here + # input_space is a 1D Box + if isinstance(observation_space, Box) and len(observation_space.shape) == 1: + # In order to guarantee backward compatability with old configs, + # we need to check if no latent dim was set and simply reuse the last + # fcnet hidden dim for that purpose. + hidden_layer_dims = model_config_dict["fcnet_hiddens"][:-1] + encoder_latent_dim = model_config_dict["fcnet_hiddens"][-1] + encoder_config = MLPEncoderConfig( + input_dims=observation_space.shape, + hidden_layer_dims=hidden_layer_dims, + hidden_layer_activation=activation, + hidden_layer_weights_initializer=model_config_dict[ + "fcnet_kernel_initializer" + ], + hidden_layer_weights_initializer_config=model_config_dict[ + "fcnet_kernel_initializer_kwargs" + ], + hidden_layer_bias_initializer=model_config_dict[ + "fcnet_bias_initializer" + ], + hidden_layer_bias_initializer_config=model_config_dict[ + "fcnet_bias_initializer_kwargs" + ], + output_layer_dim=encoder_latent_dim, + output_layer_activation=output_activation, + output_layer_weights_initializer=model_config_dict[ + "fcnet_kernel_initializer" + ], + output_layer_weights_initializer_config=model_config_dict[ + "fcnet_kernel_initializer_kwargs" + ], + output_layer_bias_initializer=model_config_dict[ + "fcnet_bias_initializer" + ], + output_layer_bias_initializer_config=model_config_dict[ + "fcnet_bias_initializer_kwargs" + ], + ) + + # input_space is a 3D Box + elif ( + isinstance(observation_space, Box) and len(observation_space.shape) == 3 + ): + if not model_config_dict.get("conv_filters"): + model_config_dict["conv_filters"] = get_filter_config( + observation_space.shape + ) + + encoder_config = CNNEncoderConfig( + input_dims=observation_space.shape, + cnn_filter_specifiers=model_config_dict["conv_filters"], + cnn_activation=model_config_dict["conv_activation"], + cnn_kernel_initializer=model_config_dict["conv_kernel_initializer"], + cnn_kernel_initializer_config=model_config_dict[ + "conv_kernel_initializer_kwargs" + ], + cnn_bias_initializer=model_config_dict["conv_bias_initializer"], + cnn_bias_initializer_config=model_config_dict[ + "conv_bias_initializer_kwargs" + ], + ) + # input_space is a 2D Box + elif ( + isinstance(observation_space, Box) and len(observation_space.shape) == 2 + ): + # RLlib used to support 2D Box spaces by silently flattening them + raise ValueError( + f"No default encoder config for obs space={observation_space}," + f" lstm={use_lstm} found. 2D Box " + f"spaces are not supported. They should be either flattened to a " + f"1D Box space or enhanced to be a 3D box space." + ) + # input_space is a possibly nested structure of spaces. + else: + # NestedModelConfig + raise ValueError( + f"No default encoder config for obs space={observation_space}," + f" lstm={use_lstm} found." + ) + + return encoder_config + + @classmethod + @OverrideToImplementCustomLogic + def get_tokenizer_config( + cls, + observation_space: gym.Space, + model_config_dict: dict, + # deprecated args. + view_requirements=DEPRECATED_VALUE, + ) -> ModelConfig: + """Returns a tokenizer config for the given space. + + This is useful for recurrent / transformer models that need to tokenize their + inputs. By default, RLlib uses the models supported by Catalog out of the box to + tokenize. + + You should override this method if you want to change the custom tokenizer + inside current encoders that Catalog returns without providing the recurrent + network as a whole. For example, if you want to define some custom CNN layers + as a tokenizer for a recurrent encoder that already includes the recurrent + layers and handles the state. + + Args: + observation_space: The observation space to use. + model_config_dict: The model config to use. + """ + if view_requirements != DEPRECATED_VALUE: + deprecation_warning(old="Catalog(view_requirements=..)", error=True) + + return cls._get_encoder_config( + observation_space=observation_space, + # Use model_config_dict without flags that would end up in complex models + model_config_dict={ + **model_config_dict, + **{"use_lstm": False, "use_attention": False}, + }, + ) + + @classmethod + def _get_dist_cls_from_action_space( + cls, + action_space: gym.Space, + *, + framework: Optional[str] = None, + ) -> Distribution: + """Returns a distribution class for the given action space. + + You can get the required input dimension for the distribution by calling + `action_dict_cls.required_input_dim(action_space)` + on the retrieved class. This is useful, because the Catalog needs to find out + about the required input dimension for the distribution before the model that + outputs these inputs is configured. + + Args: + action_space: Action space of the target gym env. + framework: The framework to use. + + Returns: + The distribution class for the given action space. + """ + # If no framework provided, return no action distribution class (None). + if framework is None: + return None + # This method is structured in two steps: + # Firstly, construct a dictionary containing the available distribution classes. + # Secondly, return the correct distribution class for the given action space. + + # Step 1: Construct the dictionary. + + class DistEnum(enum.Enum): + Categorical = "Categorical" + DiagGaussian = "Gaussian" + Deterministic = "Deterministic" + MultiDistribution = "MultiDistribution" + MultiCategorical = "MultiCategorical" + + if framework == "torch": + from ray.rllib.core.distribution.torch.torch_distribution import ( + TorchCategorical, + TorchDeterministic, + TorchDiagGaussian, + ) + + distribution_dicts = { + DistEnum.Deterministic: TorchDeterministic, + DistEnum.DiagGaussian: TorchDiagGaussian, + DistEnum.Categorical: TorchCategorical, + } + else: + raise ValueError( + f"Unknown framework: {framework}. Only 'torch' and 'tf2' are " + "supported for RLModule Catalogs." + ) + + # Only add a MultiAction distribution class to the dict if we can compute its + # components (we need a Tuple/Dict space for this). + if isinstance(action_space, (Tuple, Dict)): + partial_multi_action_distribution_cls = _multi_action_dist_partial_helper( + catalog_cls=cls, + action_space=action_space, + framework=framework, + ) + + distribution_dicts[ + DistEnum.MultiDistribution + ] = partial_multi_action_distribution_cls + + # Only add a MultiCategorical distribution class to the dict if we can compute + # its components (we need a MultiDiscrete space for this). + if isinstance(action_space, MultiDiscrete): + partial_multi_categorical_distribution_cls = ( + _multi_categorical_dist_partial_helper( + action_space=action_space, + framework=framework, + ) + ) + + distribution_dicts[ + DistEnum.MultiCategorical + ] = partial_multi_categorical_distribution_cls + + # Step 2: Return the correct distribution class for the given action space. + + # Box space -> DiagGaussian OR Deterministic. + if isinstance(action_space, Box): + if action_space.dtype.char in np.typecodes["AllInteger"]: + raise ValueError( + "Box(..., `int`) action spaces are not supported. " + "Use MultiDiscrete or Box(..., `float`)." + ) + else: + if len(action_space.shape) > 1: + raise UnsupportedSpaceException( + f"Action space has multiple dimensions {action_space.shape}. " + f"Consider reshaping this into a single dimension, using a " + f"custom action distribution, using a Tuple action space, " + f"or the multi-agent API." + ) + return distribution_dicts[DistEnum.DiagGaussian] + + # Discrete Space -> Categorical. + elif isinstance(action_space, Discrete): + return distribution_dicts[DistEnum.Categorical] + + # Tuple/Dict Spaces -> MultiAction. + elif isinstance(action_space, (Tuple, Dict)): + return distribution_dicts[DistEnum.MultiDistribution] + + # Simplex -> Dirichlet. + elif isinstance(action_space, Simplex): + # TODO(Artur): Supported Simplex (in torch). + raise NotImplementedError("Simplex action space not yet supported.") + + # MultiDiscrete -> MultiCategorical. + elif isinstance(action_space, MultiDiscrete): + return distribution_dicts[DistEnum.MultiCategorical] + + # Unknown type -> Error. + else: + raise NotImplementedError(f"Unsupported action space: `{action_space}`") + + @staticmethod + def get_preprocessor(observation_space: gym.Space, **kwargs) -> Preprocessor: + """Returns a suitable preprocessor for the given observation space. + + Args: + observation_space: The input observation space. + **kwargs: Forward-compatible kwargs. + + Returns: + preprocessor: Preprocessor for the observations. + """ + # TODO(Artur): Since preprocessors have long been @PublicAPI with the options + # kwarg as part of their constructor, we fade out support for this, + # beginning with this entrypoint. + # Next, we should deprecate the `options` kwarg from the Preprocessor itself, + # after deprecating the old catalog and other components that still pass this. + options = kwargs.get("options", {}) + if options: + deprecation_warning( + old="get_preprocessor_for_space(..., options={...})", + help="Override `Catalog.get_preprocessor()` " + "in order to implement custom behaviour.", + error=False, + ) + + if options.get("custom_preprocessor"): + deprecation_warning( + old="model_config['custom_preprocessor']", + help="Custom preprocessors are deprecated, " + "since they sometimes conflict with the built-in " + "preprocessors for handling complex observation spaces. " + "Please use wrapper classes around your environment " + "instead.", + error=True, + ) + else: + # TODO(Artur): Inline the get_preprocessor() call here once we have + # deprecated the old model catalog. + cls = get_preprocessor(observation_space) + prep = cls(observation_space, options) + return prep + + +def _multi_action_dist_partial_helper( + catalog_cls: "Catalog", action_space: gym.Space, framework: str +) -> Distribution: + """Helper method to get a partial of a MultiActionDistribution. + + This is useful for when we want to create MultiActionDistributions from + logits only (!) later, but know the action space now already. + + Args: + catalog_cls: The ModelCatalog class to use. + action_space: The action space to get the child distribution classes for. + framework: The framework to use. + + Returns: + A partial of the TorchMultiActionDistribution class. + """ + action_space_struct = get_base_struct_from_space(action_space) + flat_action_space = flatten_space(action_space) + child_distribution_cls_struct = tree.map_structure( + lambda s: catalog_cls._get_dist_cls_from_action_space( + action_space=s, + framework=framework, + ), + action_space_struct, + ) + flat_distribution_clses = tree.flatten(child_distribution_cls_struct) + + logit_lens = [ + int(dist_cls.required_input_dim(space)) + for dist_cls, space in zip(flat_distribution_clses, flat_action_space) + ] + + if framework == "torch": + from ray.rllib.core.distribution.torch.torch_distribution import ( + TorchMultiDistribution, + ) + + multi_action_dist_cls = TorchMultiDistribution + else: + raise ValueError(f"Unsupported framework: {framework}") + + partial_dist_cls = multi_action_dist_cls.get_partial_dist_cls( + space=action_space, + child_distribution_cls_struct=child_distribution_cls_struct, + input_lens=logit_lens, + ) + return partial_dist_cls + + +def _multi_categorical_dist_partial_helper( + action_space: gym.Space, framework: str +) -> Distribution: + """Helper method to get a partial of a MultiCategorical Distribution. + + This is useful for when we want to create MultiCategorical Distribution from + logits only (!) later, but know the action space now already. + + Args: + action_space: The action space to get the child distribution classes for. + framework: The framework to use. + + Returns: + A partial of the MultiCategorical class. + """ + + if framework == "torch": + from ray.rllib.core.distribution.torch.torch_distribution import ( + TorchMultiCategorical, + ) + + multi_categorical_dist_cls = TorchMultiCategorical + else: + raise ValueError(f"Unsupported framework: {framework}") + + partial_dist_cls = multi_categorical_dist_cls.get_partial_dist_cls( + space=action_space, input_lens=list(action_space.nvec) + ) + + return partial_dist_cls diff --git a/lib/python3.12/site-packages/ray/rllib/core/models/configs.py b/lib/python3.12/site-packages/ray/rllib/core/models/configs.py new file mode 100644 index 0000000000000000000000000000000000000000..1d6e78c3d1221250dd041c0b58f8b87f1f76008d --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/models/configs.py @@ -0,0 +1,1058 @@ +import abc +import functools +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple, Union + +import numpy as np + +from ray.rllib.models.torch.misc import ( + same_padding, + same_padding_transpose_after_stride, + valid_padding, +) +from ray.rllib.models.utils import get_activation_fn, get_initializer_fn +from ray.rllib.utils.annotations import ExperimentalAPI + +if TYPE_CHECKING: + from ray.rllib.core.models.base import Encoder, Model + + +@ExperimentalAPI +def _framework_implemented(torch: bool = True, tf2: bool = True): + """Decorator to check if a model was implemented in a framework. + + Args: + torch: Whether we can build this model with torch. + tf2: Whether we can build this model with tf2. + + Returns: + The decorated function. + + Raises: + ValueError: If the framework is not available to build. + """ + accepted = [] + if torch: + accepted.append("torch") + if tf2: + accepted.append("tf2") + + def decorator(fn: Callable) -> Callable: + @functools.wraps(fn) + def checked_build(self, framework, **kwargs): + if framework not in accepted: + raise ValueError( + f"This config does not support framework " + f"{framework}. Only frameworks in {accepted} are " + f"supported." + ) + return fn(self, framework, **kwargs) + + return checked_build + + return decorator + + +@ExperimentalAPI +@dataclass +class ModelConfig(abc.ABC): + """Base class for configuring a `Model` instance. + + ModelConfigs are DL framework-agnostic. + A `Model` (as a sub-component of an `RLModule`) is built via calling the + respective ModelConfig's `build()` method. + RLModules build their sub-components this way after receiving one or more + `ModelConfig` instances from a Catalog object. + + However, `ModelConfig` is not restricted to be used only with Catalog or RLModules. + Usage examples can be found in the individual Model classes', e.g. + see `ray.rllib.core.models.configs::MLPHeadConfig`. + + Attributes: + input_dims: The input dimensions of the network + always_check_shapes: Whether to always check the inputs and outputs of the + model for the specifications. Input specifications are checked on failed + forward passes of the model regardless of this flag. If this flag is set + to `True`, inputs and outputs are checked on every call. This leads to + a slow-down and should only be used for debugging. + """ + + input_dims: Union[List[int], Tuple[int, ...]] = None + always_check_shapes: bool = False + + @abc.abstractmethod + def build(self, framework: str): + """Builds the model. + + Args: + framework: The framework to use for building the model. + """ + raise NotImplementedError + + @property + def output_dims(self) -> Optional[Tuple[int, ...]]: + """Read-only `output_dims` are inferred automatically from other settings.""" + return None + + +@ExperimentalAPI +@dataclass +class _MLPConfig(ModelConfig): + """Generic configuration class for multi-layer-perceptron based Model classes. + + `output_dims` is reached by either the provided `output_layer_dim` setting (int) OR + by the last entry of `hidden_layer_dims`. In the latter case, no special output + layer is added and all layers in the stack behave exactly the same. If + `output_layer_dim` is provided, users might also change this last layer's + activation (`output_layer_activation`) and its bias setting + (`output_layer_use_bias`). + + This is a private class as users should not configure their models directly + through this class, but use one of the sub-classes, e.g. `MLPHeadConfig` or + `MLPEncoderConfig`. + + Attributes: + input_dims: A 1D tensor indicating the input dimension, e.g. `[32]`. + hidden_layer_dims: The sizes of the hidden layers. If an empty list, + `output_layer_dim` must be provided (int) and only a single layer will be + built. + hidden_layer_use_bias: Whether to use bias on all dense layers in the network + (excluding a possible separate output layer defined by `output_layer_dim`). + hidden_layer_activation: The activation function to use after each layer ( + except for the output). The default activation for hidden layers is "relu". + hidden_layer_use_layernorm: Whether to insert a LayerNorm functionality + in between each hidden layer's output and its activation. + hidden_layer_weights_initializer: The initializer function or class to use for + weight initialization in the hidden layers. If `None` the default + initializer of the respective dense layer of a framework (`"torch"` or + `"tf2"`) is used. Note, all initializers defined in the framework `"tf2`) + are allowed. For `"torch"` only the in-place initializers, i.e. ending with + an underscore "_" are allowed. + hidden_layer_weights_initializer_config: Configuration to pass into the + initializer defined in `hidden_layer_weights_initializer`. + hidden_layer_bias_initializer: The initializer function or class to use for + bias initialization in the hidden layers. If `None` the default initializer + of the respective dense layer of a framework (`"torch"` or `"tf2"`) is used. + Note, all initializers defined in the framework `"tf2`) are allowed. For + `"torch"` only the in-place initializers, i.e. ending with an underscore "_" + are allowed. + hidden_layer_bias_initializer_config: Configuration to pass into the + initializer defined in `hidden_layer_bias_initializer`. + output_layer_dim: An int indicating the size of the output layer. This may be + set to `None` in case no extra output layer should be built and only the + layers specified by `hidden_layer_dims` will be part of the network. + output_layer_use_bias: Whether to use bias on the separate output layer, if any. + output_layer_activation: The activation function to use for the output layer, + if any. The default activation for the output layer, if any, is "linear", + meaning no activation. + output_layer_weights_initializer: The initializer function or class to use for + weight initialization in the output layers. If `None` the default + initializer of the respective dense layer of a framework (`"torch"` or ` + "tf2"`) is used. Note, all initializers defined in the framework `"tf2`) are + allowed. For `"torch"` only the in-place initializers, i.e. ending with an + underscore "_" are allowed. + output_layer_weights_initializer_config: Configuration to pass into the + initializer defined in `output_layer_weights_initializer`. + output_layer_bias_initializer: The initializer function or class to use for + bias initialization in the output layers. If `None` the default initializer + of the respective dense layer of a framework (`"torch"` or `"tf2"`) is used. + For `"torch"` only the in-place initializers, i.e. ending with an underscore + "_" are allowed. + output_layer_bias_initializer_config: Configuration to pass into the + initializer defined in `output_layer_bias_initializer`. + clip_log_std: If log std should be clipped by `log_std_clip_param`. This applies + only to the action distribution parameters that encode the log standard + deviation of a `DiagGaussian` distribution. + log_std_clip_param: The clipping parameter for the log std, if clipping should + be applied - i.e. `clip_log_std=True`. The default value is 20, i.e. log + stds are clipped in between -20 and 20. + """ + + hidden_layer_dims: Union[List[int], Tuple[int, ...]] = (256, 256) + hidden_layer_use_bias: bool = True + hidden_layer_activation: str = "relu" + hidden_layer_use_layernorm: bool = False + hidden_layer_weights_initializer: Optional[Union[str, Callable]] = None + hidden_layer_weights_initializer_config: Optional[Dict] = None + hidden_layer_bias_initializer: Optional[Union[str, Callable]] = None + hidden_layer_bias_initializer_config: Optional[Dict] = None + + # Optional last output layer with - possibly - different activation and use_bias + # settings. + output_layer_dim: Optional[int] = None + output_layer_use_bias: bool = True + output_layer_activation: str = "linear" + output_layer_weights_initializer: Optional[Union[str, Callable]] = None + output_layer_weights_initializer_config: Optional[Dict] = None + output_layer_bias_initializer: Optional[Union[str, Callable]] = None + output_layer_bias_initializer_config: Optional[Dict] = None + + # Optional clipping of log standard deviation. + clip_log_std: bool = False + # Optional clip parameter for the log standard deviation. + log_std_clip_param: float = 20.0 + + @property + def output_dims(self): + if self.output_layer_dim is None and not self.hidden_layer_dims: + raise ValueError( + "If `output_layer_dim` is None, you must specify at least one hidden " + "layer dim, e.g. `hidden_layer_dims=[32]`!" + ) + + # Infer `output_dims` automatically. + return (int(self.output_layer_dim or self.hidden_layer_dims[-1]),) + + def _validate(self, framework: str = "torch"): + """Makes sure that settings are valid.""" + if self.input_dims is not None and len(self.input_dims) != 1: + raise ValueError( + f"`input_dims` ({self.input_dims}) of MLPConfig must be 1D, " + "e.g. `[32]`!" + ) + if len(self.output_dims) != 1: + raise ValueError( + f"`output_dims` ({self.output_dims}) of _MLPConfig must be " + "1D, e.g. `[32]`! This is an inferred value, hence other settings might" + " be wrong." + ) + if self.log_std_clip_param is None: + raise ValueError( + "`log_std_clip_param` of _MLPConfig must be a float value, but is " + "`None`." + ) + + # Call these already here to catch errors early on. + get_activation_fn(self.hidden_layer_activation, framework=framework) + get_activation_fn(self.output_layer_activation, framework=framework) + get_initializer_fn(self.hidden_layer_weights_initializer, framework=framework) + get_initializer_fn(self.hidden_layer_bias_initializer, framework=framework) + get_initializer_fn(self.output_layer_weights_initializer, framework=framework) + get_initializer_fn(self.output_layer_bias_initializer, framework=framework) + + +@ExperimentalAPI +@dataclass +class MLPHeadConfig(_MLPConfig): + """Configuration for an MLP head. + + See _MLPConfig for usage details. + + Example: + + .. testcode:: + + # Configuration: + config = MLPHeadConfig( + input_dims=[4], # must be 1D tensor + hidden_layer_dims=[8, 8], + hidden_layer_activation="relu", + hidden_layer_use_layernorm=False, + # final output layer with no activation (linear) + output_layer_dim=2, + output_layer_activation="linear", + ) + model = config.build(framework="tf2") + + # Resulting stack in pseudocode: + # Linear(4, 8, bias=True) + # ReLU() + # Linear(8, 8, bias=True) + # ReLU() + # Linear(8, 2, bias=True) + + Example: + + .. testcode:: + + # Configuration: + config = MLPHeadConfig( + input_dims=[2], + hidden_layer_dims=[10, 4], + hidden_layer_activation="silu", + hidden_layer_use_layernorm=True, + hidden_layer_use_bias=False, + # Initializer for `framework="torch"`. + hidden_layer_weights_initializer="xavier_normal_", + hidden_layer_weights_initializer_config={"gain": 0.8}, + # No final output layer (use last dim in `hidden_layer_dims` + # as the size of the last layer in the stack). + output_layer_dim=None, + ) + model = config.build(framework="torch") + + # Resulting stack in pseudocode: + # Linear(2, 10, bias=False) + # LayerNorm((10,)) # layer norm always before activation + # SiLU() + # Linear(10, 4, bias=False) + # LayerNorm((4,)) # layer norm always before activation + # SiLU() + """ + + @_framework_implemented() + def build(self, framework: str = "torch") -> "Model": + self._validate(framework=framework) + + if framework == "torch": + from ray.rllib.core.models.torch.heads import TorchMLPHead + + return TorchMLPHead(self) + + +@ExperimentalAPI +@dataclass +class FreeLogStdMLPHeadConfig(_MLPConfig): + """Configuration for an MLPHead with a floating second half of outputs. + + This model can be useful together with Gaussian Distributions. + This gaussian distribution would be conditioned as follows: + - The first half of outputs from this model can be used as + state-dependent means when conditioning a gaussian distribution + - The second half are floating free biases that can be used as + state-independent standard deviations to condition a gaussian distribution. + The mean values are produced by an MLPHead, while the standard + deviations are added as floating free biases from a single 1D trainable variable + (not dependent on the net's inputs). + + The output dimensions of the configured MLPHeadConfig must be even and are + divided by two to gain the output dimensions of each the mean-net and the + free std-variable. + + Example: + .. testcode:: + :skipif: True + + # Configuration: + config = FreeLogStdMLPHeadConfig( + input_dims=[2], + hidden_layer_dims=[16], + hidden_layer_activation=None, + hidden_layer_use_layernorm=False, + hidden_layer_use_bias=True, + output_layer_dim=8, # <- this must be an even size + output_layer_use_bias=True, + ) + model = config.build(framework="tf2") + + # Resulting stack in pseudocode: + # Linear(2, 16, bias=True) + # Linear(8, 8, bias=True) # 16 / 2 = 8 -> 8 nodes for the mean + # Extra variable: + # Tensor((8,), float32) # for the free (observation independent) std outputs + + Example: + .. testcode:: + :skipif: True + + # Configuration: + config = FreeLogStdMLPHeadConfig( + input_dims=[2], + hidden_layer_dims=[31, 100], # <- last idx must be an even size + hidden_layer_activation="relu", + hidden_layer_use_layernorm=False, + hidden_layer_use_bias=False, + output_layer_dim=None, # use the last hidden layer as output layer + ) + model = config.build(framework="torch") + + # Resulting stack in pseudocode: + # Linear(2, 31, bias=False) + # ReLu() + # Linear(31, 50, bias=False) # 100 / 2 = 50 -> 50 nodes for the mean + # ReLu() + # Extra variable: + # Tensor((50,), float32) # for the free (observation independent) std outputs + """ + + def _validate(self, framework: str = "torch"): + if len(self.output_dims) > 1 or self.output_dims[0] % 2 == 1: + raise ValueError( + f"`output_layer_dim` ({self.ouput_layer_dim}) or the last value in " + f"`hidden_layer_dims` ({self.hidden_layer_dims}) of a " + "FreeLogStdMLPHeadConfig must be an even int (dividable by 2), " + "e.g. `output_layer_dim=8` or `hidden_layer_dims=[133, 128]`!" + ) + + @_framework_implemented() + def build(self, framework: str = "torch") -> "Model": + self._validate(framework=framework) + + if framework == "torch": + from ray.rllib.core.models.torch.heads import TorchFreeLogStdMLPHead + + return TorchFreeLogStdMLPHead(self) + + +@ExperimentalAPI +@dataclass +class CNNTransposeHeadConfig(ModelConfig): + """Configuration for a convolutional transpose head (decoder) network. + + The configured Model transforms 1D-observations into an image space. + The stack of layers is composed of an initial Dense layer, followed by a sequence + of Conv2DTranspose layers. + `input_dims` describes the shape of the (1D) input tensor, + `initial_image_dims` describes the input into the first Conv2DTranspose + layer, where the translation from `input_dim` to `initial_image_dims` is done + via the initial Dense layer (w/o activation, w/o layer-norm, and w/ bias). + + Beyond that, each layer specified by `cnn_transpose_filter_specifiers` + is followed by an activation function according to `cnn_transpose_activation`. + + `output_dims` is reached after the final Conv2DTranspose layer. + Not that the last Conv2DTranspose layer is never activated and never layer-norm'd + regardless of the other settings. + + An example for a single conv2d operation is as follows: + Input "image" is (4, 4, 24) (not yet strided), padding is "same", stride=2, + kernel=5. + + First, the input "image" is strided (with stride=2): + + Input image (4x4 (x24)): + A B C D + E F G H + I J K L + M N O P + + Stride with stride=2 -> (7x7 (x24)) + A 0 B 0 C 0 D + 0 0 0 0 0 0 0 + E 0 F 0 G 0 H + 0 0 0 0 0 0 0 + I 0 J 0 K 0 L + 0 0 0 0 0 0 0 + M 0 N 0 O 0 P + + Then this strided "image" (strided_size=7x7) is padded (exact padding values will be + computed by the model): + + Padding -> (left=3, right=2, top=3, bottom=2) + + 0 0 0 0 0 0 0 0 0 0 0 0 + 0 0 0 0 0 0 0 0 0 0 0 0 + 0 0 0 0 0 0 0 0 0 0 0 0 + 0 0 0 A 0 B 0 C 0 D 0 0 + 0 0 0 0 0 0 0 0 0 0 0 0 + 0 0 0 E 0 F 0 G 0 H 0 0 + 0 0 0 0 0 0 0 0 0 0 0 0 + 0 0 0 I 0 J 0 K 0 L 0 0 + 0 0 0 0 0 0 0 0 0 0 0 0 + 0 0 0 M 0 N 0 O 0 P 0 0 + 0 0 0 0 0 0 0 0 0 0 0 0 + 0 0 0 0 0 0 0 0 0 0 0 0 + + Then deconvolution with kernel=5 yields an output "image" of 8x8 (x num output + filters). + + Attributes: + input_dims: The input dimensions of the network. This must be a 1D tensor. + initial_image_dims: The shape of the input to the first + Conv2DTranspose layer. We will make sure the input is transformed to + these dims via a preceding initial Dense layer, followed by a reshape, + before entering the Conv2DTranspose stack. + initial_dense_weights_initializer: The initializer function or class to use for + weight initialization in the initial dense layer. If `None` the default + initializer of the respective dense layer of a framework (`"torch"` or + `"tf2"`) is used. Note, all initializers defined in the framework `"tf2`) + are allowed. For `"torch"` only the in-place initializers, i.e. ending with + an underscore "_" are allowed. + initial_dense_weights_initializer_config: Configuration to pass into the + initializer defined in `initial_dense_weights_initializer`. + initial_dense_bias_initializer: The initializer function or class to use for + bias initialization in the initial dense layer. If `None` the default + initializer of the respective CNN layer of a framework (`"torch"` or `"tf2"` + ) is used. For `"torch"` only the in-place initializers, i.e. ending with an + underscore "_" are allowed. + initial_dense_bias_initializer_config: Configuration to pass into the + initializer defined in `initial_dense_bias_initializer`. + cnn_transpose_filter_specifiers: A list of lists, where each element of an inner + list contains elements of the form + `[number of channels/filters, [kernel width, kernel height], stride]` to + specify a convolutional layer stacked in order of the outer list. + cnn_transpose_use_bias: Whether to use bias on all Conv2DTranspose layers. + cnn_transpose_activation: The activation function to use after each layer + (except for the output). + cnn_transpose_use_layernorm: Whether to insert a LayerNorm functionality + in between each Conv2DTranspose layer's output and its activation. + cnn_transpose_kernel_initializer: The initializer function or class to use for + kernel initialization in the CNN layers. If `None` the default initializer + of the respective CNN layer of a framework (`"torch"` or `"tf2"`) is used. + Note, all initializers defined in the framework `"tf2`) are allowed. For + `"torch"` only the in-place initializers, i.e. ending with an underscore "_" + are allowed. + cnn_transpose_kernel_initializer_config: Configuration to pass into the + initializer defined in `cnn_transpose_kernel_initializer`. + cnn_transpose_bias_initializer: The initializer function or class to use for + bias initialization in the CNN layers. If `None` the default initializer of + the respective CNN layer of a framework (`"torch"` or `"tf2"`) is used. + For `"torch"` only the in-place initializers, i.e. ending with an underscore + "_" are allowed. + cnn_transpose_bias_initializer_config: Configuration to pass into the + initializer defined in `cnn_transpose_bias_initializer`. + + Example: + .. testcode:: + :skipif: True + + # Configuration: + config = CNNTransposeHeadConfig( + input_dims=[10], # 1D input vector (possibly coming from another NN) + initial_image_dims=[4, 4, 96], # first image input to deconv stack + # Initializer for TensorFlow. + initial_dense_weights_initializer="HeNormal", + initial_dense_weights_initializer={"seed": 334}, + cnn_transpose_filter_specifiers=[ + [48, [4, 4], 2], + [24, [4, 4], 2], + [3, [4, 4], 2], + ], + cnn_transpose_activation="silu", # or "swish", which is the same + cnn_transpose_use_layernorm=False, + cnn_use_bias=True, + ) + model = config.build(framework="torch) + + # Resulting stack in pseudocode: + # Linear(10, 4*4*24) + # Conv2DTranspose( + # in_channels=96, out_channels=48, + # kernel_size=[4, 4], stride=2, bias=True, + # ) + # Swish() + # Conv2DTranspose( + # in_channels=48, out_channels=24, + # kernel_size=[4, 4], stride=2, bias=True, + # ) + # Swish() + # Conv2DTranspose( + # in_channels=24, out_channels=3, + # kernel_size=[4, 4], stride=2, bias=True, + # ) + + Example: + .. testcode:: + :skipif: True + + # Configuration: + config = CNNTransposeHeadConfig( + input_dims=[128], # 1D input vector (possibly coming from another NN) + initial_image_dims=[4, 4, 32], # first image input to deconv stack + cnn_transpose_filter_specifiers=[ + [16, 4, 2], + [3, 4, 2], + ], + cnn_transpose_activation="relu", + cnn_transpose_use_layernorm=True, + cnn_use_bias=False, + # Initializer for `framework="tf2"`. + # Note, for Torch only in-place initializers are allowed. + cnn_transpose_kernel_initializer="xavier_normal_", + cnn_transpose_kernel_initializer_config={"gain": 0.8}, + ) + model = config.build(framework="torch) + + # Resulting stack in pseudocode: + # Linear(128, 4*4*32, bias=True) # bias always True for initial dense layer + # Conv2DTranspose( + # in_channels=32, out_channels=16, + # kernel_size=[4, 4], stride=2, bias=False, + # ) + # LayerNorm((-3, -2, -1)) # layer normalize over last 3 axes + # ReLU() + # Conv2DTranspose( + # in_channels=16, out_channels=3, + # kernel_size=[4, 4], stride=2, bias=False, + # ) + """ + + input_dims: Union[List[int], Tuple[int, ...]] = None + initial_image_dims: Union[List[int], Tuple[int, ...]] = field( + default_factory=lambda: [4, 4, 96] + ) + initial_dense_weights_initializer: Optional[Union[str, Callable]] = None + initial_dense_weights_initializer_config: Optional[Dict] = None + initial_dense_bias_initializer: Optional[Union[str, Callable]] = None + initial_dense_bias_initializer_config: Optional[Dict] = None + cnn_transpose_filter_specifiers: List[List[Union[int, List[int]]]] = field( + default_factory=lambda: [[48, [4, 4], 2], [24, [4, 4], 2], [3, [4, 4], 2]] + ) + cnn_transpose_use_bias: bool = True + cnn_transpose_activation: str = "relu" + cnn_transpose_use_layernorm: bool = False + cnn_transpose_kernel_initializer: Optional[Union[str, Callable]] = None + cnn_transpose_kernel_initializer_config: Optional[Dict] = None + cnn_transpose_bias_initializer: Optional[Union[str, Callable]] = None + cnn_transpose_bias_initializer_config: Optional[Dict] = None + + @property + def output_dims(self): + # Infer output dims, layer by layer. + dims = self.initial_image_dims + for filter_spec in self.cnn_transpose_filter_specifiers: + # Same padding. + num_filters, kernel, stride = filter_spec + # Compute stride output size first (striding is performed first in a + # conv transpose layer. + stride_w, stride_h = (stride, stride) if isinstance(stride, int) else stride + dims = [ + dims[0] * stride_w - (stride_w - 1), + dims[1] * stride_h - (stride_h - 1), + num_filters, + ] + # TODO (Sven): Support "valid" padding for Conv2DTranspose layers, too. + # Analogous to Conv2D Layers in a CNNEncoder. + # Apply the correct padding. Note that this might be asymetrical, meaning + # left padding might be != right padding, same for top/bottom. + _, padding_out_size = same_padding_transpose_after_stride( + (dims[0], dims[1]), kernel, stride + ) + # Perform conv transpose operation with the kernel. + kernel_w, kernel_h = (kernel, kernel) if isinstance(kernel, int) else kernel + dims = [ + padding_out_size[0] - (kernel_w - 1), + padding_out_size[1] - (kernel_h - 1), + num_filters, + ] + return tuple(dims) + + def _validate(self, framework: str = "torch"): + if len(self.input_dims) != 1: + raise ValueError( + f"`input_dims` ({self.input_dims}) of CNNTransposeHeadConfig must be a " + "3D tensor (image-like) with the dimensions meaning: width x height x " + "num_filters, e.g. `[4, 4, 92]`!" + ) + + @_framework_implemented() + def build(self, framework: str = "torch") -> "Model": + self._validate(framework) + + if framework == "torch": + from ray.rllib.core.models.torch.heads import TorchCNNTransposeHead + + return TorchCNNTransposeHead(self) + + +@ExperimentalAPI +@dataclass +class CNNEncoderConfig(ModelConfig): + """Configuration for a convolutional (encoder) network. + + The configured CNN encodes 3D-observations into a latent space. + The stack of layers is composed of a sequence of convolutional layers. + `input_dims` describes the shape of the input tensor. Beyond that, each layer + specified by `filter_specifiers` is followed by an activation function according + to `filter_activation`. + + `output_dims` is reached by either the final convolutional layer's output directly + OR by flatten this output. + + See ModelConfig for usage details. + + Example: + + .. testcode:: + + # Configuration: + config = CNNEncoderConfig( + input_dims=[84, 84, 3], # must be 3D tensor (image: w x h x C) + cnn_filter_specifiers=[ + [16, [8, 8], 4], + [32, [4, 4], 2], + ], + cnn_activation="relu", + cnn_use_layernorm=False, + cnn_use_bias=True, + ) + model = config.build(framework="torch") + + # Resulting stack in pseudocode: + # Conv2D( + # in_channels=3, out_channels=16, + # kernel_size=[8, 8], stride=[4, 4], bias=True, + # ) + # ReLU() + # Conv2D( + # in_channels=16, out_channels=32, + # kernel_size=[4, 4], stride=[2, 2], bias=True, + # ) + # ReLU() + # Conv2D( + # in_channels=32, out_channels=1, + # kernel_size=[1, 1], stride=[1, 1], bias=True, + # ) + # Flatten() + + Attributes: + input_dims: The input dimension of the network. These must be given in the + form of `(width, height, channels)`. + cnn_filter_specifiers: A list in which each element is another (inner) list + of either the following forms: + `[number of channels/filters, kernel, stride]` + OR: + `[number of channels/filters, kernel, stride, padding]`, where `padding` + can either be "same" or "valid". + When using the first format w/o the `padding` specifier, `padding` is "same" + by default. Also, `kernel` and `stride` may be provided either as single + ints (square) or as a tuple/list of two ints (width- and height dimensions) + for non-squared kernel/stride shapes. + A good rule of thumb for constructing CNN stacks is: + When using padding="same", the input "image" will be reduced in size by + the factor `stride`, e.g. input=(84, 84, 3) stride=2 kernel=x padding="same" + filters=16 -> output=(42, 42, 16). + For example, if you would like to reduce an Atari image from its original + (84, 84, 3) dimensions down to (6, 6, F), you can construct the following + stack and reduce the w x h dimension of the image by 2 in each layer: + [[16, 4, 2], [32, 4, 2], [64, 4, 2], [128, 4, 2]] -> output=(6, 6, 128) + cnn_use_bias: Whether to use bias on all Conv2D layers. + cnn_activation: The activation function to use after each layer ( + except for the output). The default activation for Conv2d layers is "relu". + cnn_use_layernorm: Whether to insert a LayerNorm functionality + in between each CNN layer's output and its activation. Note that + the output layer. + cnn_kernel_initializer: The initializer function or class to use for kernel + initialization in the CNN layers. If `None` the default initializer of the + respective CNN layer of a framework (`"torch"` or `"tf2"`) is used. Note, + all initializers defined in the framework `"tf2`) are allowed. For `"torch"` + only the in-place initializers, i.e. ending with an underscore "_" are + allowed. + cnn_kernel_initializer_config: Configuration to pass into the initializer + defined in `cnn_kernel_initializer`. + cnn_bias_initializer: The initializer function or class to use for bias + initialization in the CNN layers. If `None` the default initializer of + the respective CNN layer of a framework (`"torch"` or `"tf2"`) is used. + For `"torch"` only the in-place initializers, i.e. ending with an underscore + "_" are allowed. + cnn_bias_initializer_config: Configuration to pass into the initializer defined + in `cnn_bias_initializer`. + flatten_at_end: Whether to flatten the output of the last conv 2D layer into + a 1D tensor. By default, this is True. Note that if you set this to False, + you might simply stack another CNNEncoder on top of this one (maybe with + different activation and bias settings). + """ + + input_dims: Union[List[int], Tuple[int, ...]] = None + cnn_filter_specifiers: List[List[Union[int, List[int]]]] = field( + default_factory=lambda: [[16, [4, 4], 2], [32, [4, 4], 2], [64, [8, 8], 2]] + ) + cnn_use_bias: bool = True + cnn_activation: str = "relu" + cnn_use_layernorm: bool = False + cnn_kernel_initializer: Optional[Union[str, Callable]] = None + cnn_kernel_initializer_config: Optional[Dict] = None + cnn_bias_initializer: Optional[Union[str, Callable]] = None + cnn_bias_initializer_config: Optional[Dict] = None + flatten_at_end: bool = True + + @property + def output_dims(self): + if not self.input_dims: + return None + + # Infer output dims, layer by layer. + dims = self.input_dims # Creates a copy (works for tuple/list). + for filter_spec in self.cnn_filter_specifiers: + # Padding not provided, "same" by default. + if len(filter_spec) == 3: + num_filters, kernel, stride = filter_spec + padding = "same" + # Padding option provided, use given value. + else: + num_filters, kernel, stride, padding = filter_spec + + # Same padding. + if padding == "same": + _, dims = same_padding(dims[:2], kernel, stride) + # Valid padding. + else: + dims = valid_padding(dims[:2], kernel, stride) + + # Add depth (num_filters) to the end (our utility functions for same/valid + # only return the image width/height). + dims = [dims[0], dims[1], num_filters] + + # Flatten everything. + if self.flatten_at_end: + return (int(np.prod(dims)),) + + return tuple(dims) + + def _validate(self, framework: str = "torch"): + if len(self.input_dims) != 3: + raise ValueError( + f"`input_dims` ({self.input_dims}) of CNNEncoderConfig must be a 3D " + "tensor (image) with the dimensions meaning: width x height x " + "channels, e.g. `[64, 64, 3]`!" + ) + if not self.flatten_at_end and len(self.output_dims) != 3: + raise ValueError( + f"`output_dims` ({self.output_dims}) of CNNEncoderConfig must be " + "3D, e.g. `[4, 4, 128]`, b/c your `flatten_at_end` setting is False! " + "`output_dims` is an inferred value, hence other settings might be " + "wrong." + ) + elif self.flatten_at_end and len(self.output_dims) != 1: + raise ValueError( + f"`output_dims` ({self.output_dims}) of CNNEncoderConfig must be " + "1D, e.g. `[32]`, b/c your `flatten_at_end` setting is True! " + "`output_dims` is an inferred value, hence other settings might be " + "wrong." + ) + + @_framework_implemented() + def build(self, framework: str = "torch") -> "Model": + self._validate(framework) + + if framework == "torch": + from ray.rllib.core.models.torch.encoder import TorchCNNEncoder + + return TorchCNNEncoder(self) + + +@ExperimentalAPI +@dataclass +class MLPEncoderConfig(_MLPConfig): + """Configuration for an MLP that acts as an encoder. + + See _MLPConfig for usage details. + + Example: + .. testcode:: + + # Configuration: + config = MLPEncoderConfig( + input_dims=[4], # must be 1D tensor + hidden_layer_dims=[16], + hidden_layer_activation="relu", + hidden_layer_use_layernorm=False, + output_layer_dim=None, # maybe None or an int + ) + model = config.build(framework="torch") + + # Resulting stack in pseudocode: + # Linear(4, 16, bias=True) + # ReLU() + + Example: + .. testcode:: + + # Configuration: + config = MLPEncoderConfig( + input_dims=[2], + hidden_layer_dims=[8, 8], + hidden_layer_activation="silu", + hidden_layer_use_layernorm=True, + hidden_layer_use_bias=False, + output_layer_dim=4, + output_layer_activation="tanh", + output_layer_use_bias=False, + ) + model = config.build(framework="tf2") + + # Resulting stack in pseudocode: + # Linear(2, 8, bias=False) + # LayerNorm((8,)) # layernorm always before activation + # SiLU() + # Linear(8, 8, bias=False) + # LayerNorm((8,)) # layernorm always before activation + # SiLU() + # Linear(8, 4, bias=False) + # Tanh() + """ + + @_framework_implemented() + def build(self, framework: str = "torch") -> "Encoder": + self._validate(framework) + + if framework == "torch": + from ray.rllib.core.models.torch.encoder import TorchMLPEncoder + + return TorchMLPEncoder(self) + + +@ExperimentalAPI +@dataclass +class RecurrentEncoderConfig(ModelConfig): + """Configuration for an LSTM-based or a GRU-based encoder. + + The encoder consists of... + - Zero or one tokenizers + - N LSTM/GRU layers stacked on top of each other and feeding + their outputs as inputs to the respective next layer. + + This makes for the following flow of tensors: + + Inputs + | + [Tokenizer if present] + | + LSTM layer 1 + | + (...) + | + LSTM layer n + | + Outputs + + The internal state is structued as (num_layers, B, hidden-size) for all hidden + state components, e.g. + h- and c-states of the LSTM layer(s) or h-state of the GRU layer(s). + For example, the hidden states of an LSTMEncoder with num_layers=2 and hidden_dim=8 + would be: {"h": (2, B, 8), "c": (2, B, 8)}. + + `output_dims` is reached by the last recurrent layer's dimension, which is always + the `hidden_dims` value. + + Example: + .. testcode:: + + # Configuration: + config = RecurrentEncoderConfig( + recurrent_layer_type="lstm", + input_dims=[16], # must be 1D tensor + hidden_dim=128, + num_layers=2, + use_bias=True, + ) + model = config.build(framework="torch") + + # Resulting stack in pseudocode: + # LSTM(16, 128, bias=True) + # LSTM(128, 128, bias=True) + + # Resulting shape of the internal states (c- and h-states): + # (2, B, 128) for each c- and h-states. + + Example: + .. testcode:: + + # Configuration: + config = RecurrentEncoderConfig( + recurrent_layer_type="gru", + input_dims=[32], # must be 1D tensor + hidden_dim=64, + num_layers=1, + use_bias=False, + ) + model = config.build(framework="torch") + + # Resulting stack in pseudocode: + # GRU(32, 64, bias=False) + + # Resulting shape of the internal state: + # (1, B, 64) + + Attributes: + input_dims: The input dimensions. Must be 1D. This is the 1D shape of the tensor + that goes into the first recurrent layer. + recurrent_layer_type: The type of the recurrent layer(s). + Either "lstm" or "gru". + hidden_dim: The size of the hidden internal state(s) of the recurrent layer(s). + For example, for an LSTM, this would be the size of the c- and h-tensors. + num_layers: The number of recurrent (LSTM or GRU) layers to stack. + batch_major: Wether the input is batch major (B, T, ..) or + time major (T, B, ..). + hidden_weights_initializer: The initializer function or class to use for + kernel initialization in the hidden layers. If `None` the default + initializer of the respective recurrent layer of a framework (`"torch"` or + `"tf2"`) is used. Note, all initializers defined in the frameworks ( + `"torch"` or `"tf2`) are allowed. For `"torch"` only the in-place + initializers, i.e. ending with an underscore "_" are allowed. + hidden_weights_initializer_config: Configuration to pass into the + initializer defined in `hidden_weights_initializer`. + use_bias: Whether to use bias on the recurrent layers in the network. + hidden_bias_initializer: The initializer function or class to use for bias + initialization in the hidden layers. If `None` the default initializer of + the respective recurrent layer of a framework (`"torch"` or `"tf2"`) is + used. For `"torch"` only the in-place initializers, i.e. ending with an + underscore "_" are allowed. + hidden_bias_initializer_config: Configuration to pass into the initializer + defined in `hidden_bias_initializer`. + tokenizer_config: A ModelConfig to build tokenizers for observations, + actions and other spaces. + """ + + recurrent_layer_type: str = "lstm" + hidden_dim: int = None + num_layers: int = None + batch_major: bool = True + hidden_weights_initializer: Optional[Union[str, Callable]] = None + hidden_weights_initializer_config: Optional[Dict] = None + use_bias: bool = True + hidden_bias_initializer: Optional[Union[str, Callable]] = None + hidden_bias_initializer_config: Optional[Dict] = None + tokenizer_config: ModelConfig = None + + @property + def output_dims(self): + return (self.hidden_dim,) + + def _validate(self, framework: str = "torch"): + """Makes sure that settings are valid.""" + if self.recurrent_layer_type not in ["gru", "lstm"]: + raise ValueError( + f"`recurrent_layer_type` ({self.recurrent_layer_type}) of " + "RecurrentEncoderConfig must be 'gru' or 'lstm'!" + ) + if self.input_dims is not None and len(self.input_dims) != 1: + raise ValueError( + f"`input_dims` ({self.input_dims}) of RecurrentEncoderConfig must be " + "1D, e.g. `[32]`!" + ) + if len(self.output_dims) != 1: + raise ValueError( + f"`output_dims` ({self.output_dims}) of RecurrentEncoderConfig must be " + "1D, e.g. `[32]`! This is an inferred value, hence other settings might" + " be wrong." + ) + + @_framework_implemented() + def build(self, framework: str = "torch") -> "Encoder": + if framework == "torch": + from ray.rllib.core.models.torch.encoder import ( + TorchGRUEncoder as GRU, + TorchLSTMEncoder as LSTM, + ) + + if self.recurrent_layer_type == "lstm": + return LSTM(self) + else: + return GRU(self) + + +@ExperimentalAPI +@dataclass +class ActorCriticEncoderConfig(ModelConfig): + """Configuration for an ActorCriticEncoder. + + The base encoder functions like other encoders in RLlib. It is wrapped by the + ActorCriticEncoder to provides a shared encoder Model to use in RLModules that + provides twofold outputs: one for the actor and one for the critic. See + ModelConfig for usage details. + + Attributes: + base_encoder_config: The configuration for the wrapped encoder(s). + shared: Whether the base encoder is shared between the actor and critic. + inference_only: Whether the configured encoder will only ever be used as an + actor-encoder, never as a value-function encoder. Thus, if True and `shared` + is False, will only build the actor-related components. + """ + + base_encoder_config: ModelConfig = None + shared: bool = True + inference_only: bool = False + + @_framework_implemented() + def build(self, framework: str = "torch") -> "Encoder": + if framework == "torch": + from ray.rllib.core.models.torch.encoder import ( + TorchActorCriticEncoder, + TorchStatefulActorCriticEncoder, + ) + + if isinstance(self.base_encoder_config, RecurrentEncoderConfig): + return TorchStatefulActorCriticEncoder(self) + else: + return TorchActorCriticEncoder(self) diff --git a/lib/python3.12/site-packages/ray/rllib/core/models/specs/__init__.py b/lib/python3.12/site-packages/ray/rllib/core/models/specs/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/core/models/specs/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/models/specs/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3fcc04433309553894ec9069e99d86d2c9b8d8e1 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/models/specs/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/models/specs/__pycache__/specs_base.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/models/specs/__pycache__/specs_base.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e239aca0e882099853d5c9c1d2381a351541a5a9 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/models/specs/__pycache__/specs_base.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/models/specs/__pycache__/specs_dict.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/models/specs/__pycache__/specs_dict.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dc1fdf3c9368abb454c95399eeff38825d4a5a0c Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/models/specs/__pycache__/specs_dict.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/models/specs/specs_base.py b/lib/python3.12/site-packages/ray/rllib/core/models/specs/specs_base.py new file mode 100644 index 0000000000000000000000000000000000000000..2274fdd73641d0a2d1bee339909bb52b5343b9d4 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/models/specs/specs_base.py @@ -0,0 +1,28 @@ +from ray._common.deprecation import Deprecated + + +@Deprecated( + help="The Spec checking APIs have been deprecated and cancelled without " + "replacement.", + error=True, +) +class Spec: + pass + + +@Deprecated( + help="The Spec checking APIs have been deprecated and cancelled without " + "replacement.", + error=True, +) +class TypeSpec: + pass + + +@Deprecated( + help="The Spec checking APIs have been deprecated and cancelled without " + "replacement.", + error=True, +) +class TensorSpec: + pass diff --git a/lib/python3.12/site-packages/ray/rllib/core/models/specs/specs_dict.py b/lib/python3.12/site-packages/ray/rllib/core/models/specs/specs_dict.py new file mode 100644 index 0000000000000000000000000000000000000000..9c60b46fe67dd07f523c2d7954e6277d3f56533b --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/models/specs/specs_dict.py @@ -0,0 +1,9 @@ +from ray._common.deprecation import Deprecated + + +@Deprecated( + help="The SpecDict API has been deprecated and cancelled without " "replacement.", + error=True, +) +class SpecDict: + pass diff --git a/lib/python3.12/site-packages/ray/rllib/core/models/torch/__init__.py b/lib/python3.12/site-packages/ray/rllib/core/models/torch/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/core/models/torch/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/models/torch/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7fb5a16b50d84bbbe45dccf1ed58d620ab18b83a Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/models/torch/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/models/torch/__pycache__/base.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/models/torch/__pycache__/base.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ea3c9c08deea627db4eb0a17da9e07cf11dcf7cd Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/models/torch/__pycache__/base.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/models/torch/__pycache__/encoder.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/models/torch/__pycache__/encoder.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3b9fbd90e6187551e2d0bb049f96c426258d5ced Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/models/torch/__pycache__/encoder.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/models/torch/__pycache__/heads.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/models/torch/__pycache__/heads.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..274dfcec78007c2ed741e33854ac0f2a12c10568 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/models/torch/__pycache__/heads.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/models/torch/__pycache__/primitives.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/models/torch/__pycache__/primitives.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5fd186b50cbc634a52f4f78255b12d0dcd28acbf Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/models/torch/__pycache__/primitives.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/models/torch/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/core/models/torch/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..87fdb12beb20291b3743f7d9e14e4dbfd8d852fa Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/core/models/torch/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/core/models/torch/base.py b/lib/python3.12/site-packages/ray/rllib/core/models/torch/base.py new file mode 100644 index 0000000000000000000000000000000000000000..ce6a109024b9afd6582d61a141eb44bd15baa34a --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/models/torch/base.py @@ -0,0 +1,101 @@ +import abc +import logging +from typing import Tuple, Union + +from ray.rllib.core.models.base import Model +from ray.rllib.core.models.configs import ModelConfig +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.typing import TensorType + +torch, nn = try_import_torch() + +logger = logging.getLogger(__name__) + + +class TorchModel(nn.Module, Model, abc.ABC): + """Base class for RLlib's PyTorch models. + + This class defines the interface for RLlib's PyTorch models. + + Example usage for a single Flattening layer: + + .. testcode:: + + from ray.rllib.core.models.configs import ModelConfig + from ray.rllib.core.models.torch.base import TorchModel + import torch + + class FlattenModelConfig(ModelConfig): + def build(self, framework: str): + assert framework == "torch" + return TorchFlattenModel(self) + + class TorchFlattenModel(TorchModel): + def __init__(self, config): + TorchModel.__init__(self, config) + self.flatten_layer = torch.nn.Flatten() + + def _forward(self, inputs, **kwargs): + return self.flatten_layer(inputs) + + model = FlattenModelConfig().build("torch") + inputs = torch.Tensor([[[1, 2]]]) + print(model(inputs)) + + .. testoutput:: + + tensor([[1., 2.]]) + + """ + + def __init__(self, config: ModelConfig): + """Initialized a TorchModel. + + Args: + config: The ModelConfig to use. + """ + nn.Module.__init__(self) + Model.__init__(self, config) + + def forward( + self, inputs: Union[dict, TensorType], **kwargs + ) -> Union[dict, TensorType]: + """Returns the output of this model for the given input. + + This method only makes sure that we have a spec-checked _forward() method. + + Args: + inputs: The input tensors. + **kwargs: Forward compatibility kwargs. + + Returns: + dict: The output tensors. + """ + return self._forward(inputs, **kwargs) + + @override(Model) + def get_num_parameters(self) -> Tuple[int, int]: + num_trainable_parameters = 0 + num_frozen_parameters = 0 + for p in self.parameters(): + n = p.numel() + if p.requires_grad: + num_trainable_parameters += n + else: + num_frozen_parameters += n + return num_trainable_parameters, num_frozen_parameters + + @override(Model) + def _set_to_dummy_weights(self, value_sequence=(-0.02, -0.01, 0.01, 0.02)): + trainable_weights = [] + non_trainable_weights = [] + for p in self.parameters(): + if p.requires_grad: + trainable_weights.append(p) + else: + non_trainable_weights.append(p) + for i, w in enumerate(trainable_weights + non_trainable_weights): + fill_val = value_sequence[i % len(value_sequence)] + with torch.no_grad(): + w.fill_(fill_val) diff --git a/lib/python3.12/site-packages/ray/rllib/core/models/torch/encoder.py b/lib/python3.12/site-packages/ray/rllib/core/models/torch/encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..ea444f65d8292f4e66f392fe356b04c8d54147d4 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/models/torch/encoder.py @@ -0,0 +1,285 @@ +import tree + +from ray.rllib.core.columns import Columns +from ray.rllib.core.models.base import ( + ENCODER_OUT, + ActorCriticEncoder, + Encoder, + Model, + StatefulActorCriticEncoder, + tokenize, +) +from ray.rllib.core.models.configs import ( + ActorCriticEncoderConfig, + CNNEncoderConfig, + MLPEncoderConfig, + RecurrentEncoderConfig, +) +from ray.rllib.core.models.torch.base import TorchModel +from ray.rllib.core.models.torch.primitives import TorchCNN, TorchMLP +from ray.rllib.models.utils import get_initializer_fn +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_torch + +torch, nn = try_import_torch() + + +class TorchActorCriticEncoder(TorchModel, ActorCriticEncoder): + """An actor-critic encoder for torch.""" + + framework = "torch" + + def __init__(self, config: ActorCriticEncoderConfig) -> None: + TorchModel.__init__(self, config) + ActorCriticEncoder.__init__(self, config) + + +class TorchStatefulActorCriticEncoder(TorchModel, StatefulActorCriticEncoder): + """A stateful actor-critic encoder for torch.""" + + framework = "torch" + + def __init__(self, config: ActorCriticEncoderConfig) -> None: + TorchModel.__init__(self, config) + StatefulActorCriticEncoder.__init__(self, config) + + +class TorchMLPEncoder(TorchModel, Encoder): + def __init__(self, config: MLPEncoderConfig) -> None: + TorchModel.__init__(self, config) + Encoder.__init__(self, config) + + # Create the neural network. + self.net = TorchMLP( + input_dim=config.input_dims[0], + hidden_layer_dims=config.hidden_layer_dims, + hidden_layer_activation=config.hidden_layer_activation, + hidden_layer_use_layernorm=config.hidden_layer_use_layernorm, + hidden_layer_use_bias=config.hidden_layer_use_bias, + hidden_layer_weights_initializer=config.hidden_layer_weights_initializer, + hidden_layer_weights_initializer_config=( + config.hidden_layer_weights_initializer_config + ), + hidden_layer_bias_initializer=config.hidden_layer_bias_initializer, + hidden_layer_bias_initializer_config=( + config.hidden_layer_bias_initializer_config + ), + output_dim=config.output_layer_dim, + output_activation=config.output_layer_activation, + output_use_bias=config.output_layer_use_bias, + output_weights_initializer=config.output_layer_weights_initializer, + output_weights_initializer_config=( + config.output_layer_weights_initializer_config + ), + output_bias_initializer=config.output_layer_bias_initializer, + output_bias_initializer_config=config.output_layer_bias_initializer_config, + ) + + @override(Model) + def _forward(self, inputs: dict, **kwargs) -> dict: + return {ENCODER_OUT: self.net(inputs[Columns.OBS])} + + +class TorchCNNEncoder(TorchModel, Encoder): + def __init__(self, config: CNNEncoderConfig) -> None: + TorchModel.__init__(self, config) + Encoder.__init__(self, config) + + layers = [] + # The bare-bones CNN (no flatten, no succeeding dense). + cnn = TorchCNN( + input_dims=config.input_dims, + cnn_filter_specifiers=config.cnn_filter_specifiers, + cnn_activation=config.cnn_activation, + cnn_use_layernorm=config.cnn_use_layernorm, + cnn_use_bias=config.cnn_use_bias, + cnn_kernel_initializer=config.cnn_kernel_initializer, + cnn_kernel_initializer_config=config.cnn_kernel_initializer_config, + cnn_bias_initializer=config.cnn_bias_initializer, + cnn_bias_initializer_config=config.cnn_bias_initializer_config, + ) + layers.append(cnn) + + # Add a flatten operation to move from 2/3D into 1D space. + if config.flatten_at_end: + layers.append(nn.Flatten()) + + # Create the network from gathered layers. + self.net = nn.Sequential(*layers) + + @override(Model) + def _forward(self, inputs: dict, **kwargs) -> dict: + return {ENCODER_OUT: self.net(inputs[Columns.OBS])} + + +class TorchGRUEncoder(TorchModel, Encoder): + """A recurrent GRU encoder. + + This encoder has... + - Zero or one tokenizers. + - One or more GRU layers. + """ + + def __init__(self, config: RecurrentEncoderConfig) -> None: + TorchModel.__init__(self, config) + + # Maybe create a tokenizer + if config.tokenizer_config is not None: + self.tokenizer = config.tokenizer_config.build(framework="torch") + gru_input_dims = config.tokenizer_config.output_dims + else: + self.tokenizer = None + gru_input_dims = config.input_dims + + # We only support 1D spaces right now. + assert len(gru_input_dims) == 1 + gru_input_dim = gru_input_dims[0] + + gru_weights_initializer = get_initializer_fn( + config.hidden_weights_initializer, framework="torch" + ) + gru_bias_initializer = get_initializer_fn( + config.hidden_bias_initializer, framework="torch" + ) + + # Create the torch GRU layer. + self.gru = nn.GRU( + gru_input_dim, + config.hidden_dim, + config.num_layers, + batch_first=config.batch_major, + bias=config.use_bias, + ) + + # Initialize, GRU weights, if necessary. + if gru_weights_initializer: + gru_weights_initializer( + self.gru.weight, **config.hidden_weights_initializer_config or {} + ) + # Initialize GRU bias, if necessary. + if gru_bias_initializer: + gru_bias_initializer( + self.gru.weight, **config.hidden_bias_initializer_config or {} + ) + + @override(Model) + def get_initial_state(self): + return { + "h": torch.zeros(self.config.num_layers, self.config.hidden_dim), + } + + @override(Model) + def _forward(self, inputs: dict, **kwargs) -> dict: + outputs = {} + + if self.tokenizer is not None: + # Push observations through the tokenizer encoder if we built one. + out = tokenize(self.tokenizer, inputs, framework="torch") + else: + # Otherwise, just use the raw observations. + out = inputs[Columns.OBS].float() + + # States are batch-first when coming in. Make them layers-first. + states_in = tree.map_structure( + lambda s: s.transpose(0, 1), inputs[Columns.STATE_IN] + ) + + out, states_out = self.gru(out, states_in["h"]) + states_out = {"h": states_out} + + # Insert them into the output dict. + outputs[ENCODER_OUT] = out + outputs[Columns.STATE_OUT] = tree.map_structure( + lambda s: s.transpose(0, 1), states_out + ) + return outputs + + +class TorchLSTMEncoder(TorchModel, Encoder): + """A recurrent LSTM encoder. + + This encoder has... + - Zero or one tokenizers. + - One or more LSTM layers. + """ + + def __init__(self, config: RecurrentEncoderConfig) -> None: + TorchModel.__init__(self, config) + + # Maybe create a tokenizer + if config.tokenizer_config is not None: + self.tokenizer = config.tokenizer_config.build(framework="torch") + lstm_input_dims = config.tokenizer_config.output_dims + else: + self.tokenizer = None + lstm_input_dims = config.input_dims + + # We only support 1D spaces right now. + assert len(lstm_input_dims) == 1 + lstm_input_dim = lstm_input_dims[0] + + lstm_weights_initializer = get_initializer_fn( + config.hidden_weights_initializer, framework="torch" + ) + lstm_bias_initializer = get_initializer_fn( + config.hidden_bias_initializer, framework="torch" + ) + + # Create the torch LSTM layer. + self.lstm = nn.LSTM( + lstm_input_dim, + config.hidden_dim, + config.num_layers, + batch_first=config.batch_major, + bias=config.use_bias, + ) + + # Initialize LSTM layer weigths and biases, if necessary. + for layer in self.lstm.all_weights: + if lstm_weights_initializer: + lstm_weights_initializer( + layer[0], **config.hidden_weights_initializer_config or {} + ) + lstm_weights_initializer( + layer[1], **config.hidden_weights_initializer_config or {} + ) + if lstm_bias_initializer: + lstm_bias_initializer( + layer[2], **config.hidden_bias_initializer_config or {} + ) + lstm_bias_initializer( + layer[3], **config.hidden_bias_initializer_config or {} + ) + + @override(Model) + def get_initial_state(self): + return { + "h": torch.zeros(self.config.num_layers, self.config.hidden_dim), + "c": torch.zeros(self.config.num_layers, self.config.hidden_dim), + } + + @override(Model) + def _forward(self, inputs: dict, **kwargs) -> dict: + outputs = {} + + if self.tokenizer is not None: + # Push observations through the tokenizer encoder if we built one. + out = tokenize(self.tokenizer, inputs, framework="torch") + else: + # Otherwise, just use the raw observations. + out = inputs[Columns.OBS].float() + + # States are batch-first when coming in. Make them layers-first. + states_in = tree.map_structure( + lambda s: s.transpose(0, 1), inputs[Columns.STATE_IN] + ) + + out, states_out = self.lstm(out, (states_in["h"], states_in["c"])) + states_out = {"h": states_out[0], "c": states_out[1]} + + # Insert them into the output dict. + outputs[ENCODER_OUT] = out + outputs[Columns.STATE_OUT] = tree.map_structure( + lambda s: s.transpose(0, 1), states_out + ) + return outputs diff --git a/lib/python3.12/site-packages/ray/rllib/core/models/torch/heads.py b/lib/python3.12/site-packages/ray/rllib/core/models/torch/heads.py new file mode 100644 index 0000000000000000000000000000000000000000..844c40c4a44cbe0617c5008b0464acc20ed17653 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/models/torch/heads.py @@ -0,0 +1,197 @@ +import numpy as np + +from ray.rllib.core.models.base import Model +from ray.rllib.core.models.configs import ( + CNNTransposeHeadConfig, + FreeLogStdMLPHeadConfig, + MLPHeadConfig, +) +from ray.rllib.core.models.torch.base import TorchModel +from ray.rllib.core.models.torch.primitives import TorchCNNTranspose, TorchMLP +from ray.rllib.models.utils import get_initializer_fn +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_torch + +torch, nn = try_import_torch() + + +class TorchMLPHead(TorchModel): + def __init__(self, config: MLPHeadConfig) -> None: + super().__init__(config) + + self.net = TorchMLP( + input_dim=config.input_dims[0], + hidden_layer_dims=config.hidden_layer_dims, + hidden_layer_activation=config.hidden_layer_activation, + hidden_layer_use_layernorm=config.hidden_layer_use_layernorm, + hidden_layer_use_bias=config.hidden_layer_use_bias, + hidden_layer_weights_initializer=config.hidden_layer_weights_initializer, + hidden_layer_weights_initializer_config=( + config.hidden_layer_weights_initializer_config + ), + hidden_layer_bias_initializer=config.hidden_layer_bias_initializer, + hidden_layer_bias_initializer_config=( + config.hidden_layer_bias_initializer_config + ), + output_dim=config.output_layer_dim, + output_activation=config.output_layer_activation, + output_use_bias=config.output_layer_use_bias, + output_weights_initializer=config.output_layer_weights_initializer, + output_weights_initializer_config=( + config.output_layer_weights_initializer_config + ), + output_bias_initializer=config.output_layer_bias_initializer, + output_bias_initializer_config=config.output_layer_bias_initializer_config, + ) + # If log standard deviations should be clipped. This should be only true for + # policy heads. Value heads should never be clipped. + self.clip_log_std = config.clip_log_std + # The clipping parameter for the log standard deviation. + self.log_std_clip_param = torch.Tensor([config.log_std_clip_param]) + # Register a buffer to handle device mapping. + self.register_buffer("log_std_clip_param_const", self.log_std_clip_param) + + @override(Model) + def _forward(self, inputs: torch.Tensor, **kwargs) -> torch.Tensor: + # Only clip the log standard deviations, if the user wants to clip. This + # avoids also clipping value heads. + if self.clip_log_std: + # Forward pass. + means, log_stds = torch.chunk(self.net(inputs), chunks=2, dim=-1) + # Clip the log standard deviations. + log_stds = torch.clamp( + log_stds, -self.log_std_clip_param_const, self.log_std_clip_param_const + ) + return torch.cat((means, log_stds), dim=-1) + # Otherwise just return the logits. + else: + return self.net(inputs) + + +class TorchFreeLogStdMLPHead(TorchModel): + """An MLPHead that implements floating log stds for Gaussian distributions.""" + + def __init__(self, config: FreeLogStdMLPHeadConfig) -> None: + super().__init__(config) + + assert config.output_dims[0] % 2 == 0, "output_dims must be even for free std!" + self._half_output_dim = config.output_dims[0] // 2 + + self.net = TorchMLP( + input_dim=config.input_dims[0], + hidden_layer_dims=config.hidden_layer_dims, + hidden_layer_activation=config.hidden_layer_activation, + hidden_layer_use_layernorm=config.hidden_layer_use_layernorm, + hidden_layer_use_bias=config.hidden_layer_use_bias, + hidden_layer_weights_initializer=config.hidden_layer_weights_initializer, + hidden_layer_weights_initializer_config=( + config.hidden_layer_weights_initializer_config + ), + hidden_layer_bias_initializer=config.hidden_layer_bias_initializer, + hidden_layer_bias_initializer_config=( + config.hidden_layer_bias_initializer_config + ), + output_dim=self._half_output_dim, + output_activation=config.output_layer_activation, + output_use_bias=config.output_layer_use_bias, + output_weights_initializer=config.output_layer_weights_initializer, + output_weights_initializer_config=( + config.output_layer_weights_initializer_config + ), + output_bias_initializer=config.output_layer_bias_initializer, + output_bias_initializer_config=config.output_layer_bias_initializer_config, + ) + + self.log_std = torch.nn.Parameter( + torch.as_tensor([0.0] * self._half_output_dim) + ) + # If log standard deviations should be clipped. This should be only true for + # policy heads. Value heads should never be clipped. + self.clip_log_std = config.clip_log_std + # The clipping parameter for the log standard deviation. + self.log_std_clip_param = torch.Tensor( + [config.log_std_clip_param], device=self.log_std.device + ) + # Register a buffer to handle device mapping. + self.register_buffer("log_std_clip_param_const", self.log_std_clip_param) + + @override(Model) + def _forward(self, inputs: torch.Tensor, **kwargs) -> torch.Tensor: + # Compute the mean first, then append the log_std. + mean = self.net(inputs) + + # If log standard deviation should be clipped. + if self.clip_log_std: + # Clip the log standard deviation to avoid running into too small + # deviations that factually collapses the policy. + log_std = torch.clamp( + self.log_std, + -self.log_std_clip_param_const, + self.log_std_clip_param_const, + ) + else: + log_std = self.log_std + + return torch.cat([mean, log_std.unsqueeze(0).repeat([len(mean), 1])], axis=1) + + +class TorchCNNTransposeHead(TorchModel): + def __init__(self, config: CNNTransposeHeadConfig) -> None: + super().__init__(config) + + # Initial, inactivated Dense layer (always w/ bias). + # This layer is responsible for getting the incoming tensor into a proper + # initial image shape (w x h x filters) for the suceeding Conv2DTranspose stack. + self.initial_dense = nn.Linear( + in_features=config.input_dims[0], + out_features=int(np.prod(config.initial_image_dims)), + bias=True, + ) + + # Initial Dense layer initializers. + initial_dense_weights_initializer = get_initializer_fn( + config.initial_dense_weights_initializer, framework="torch" + ) + initial_dense_bias_initializer = get_initializer_fn( + config.initial_dense_bias_initializer, framework="torch" + ) + + # Initialize dense layer weights, if necessary. + if initial_dense_weights_initializer: + initial_dense_weights_initializer( + self.initial_dense.weight, + **config.initial_dense_weights_initializer_config or {}, + ) + # Initialized dense layer bais, if necessary. + if initial_dense_bias_initializer: + initial_dense_bias_initializer( + self.initial_dense.bias, + **config.initial_dense_bias_initializer_config or {}, + ) + + # The main CNNTranspose stack. + self.cnn_transpose_net = TorchCNNTranspose( + input_dims=config.initial_image_dims, + cnn_transpose_filter_specifiers=config.cnn_transpose_filter_specifiers, + cnn_transpose_activation=config.cnn_transpose_activation, + cnn_transpose_use_layernorm=config.cnn_transpose_use_layernorm, + cnn_transpose_use_bias=config.cnn_transpose_use_bias, + cnn_transpose_kernel_initializer=config.cnn_transpose_kernel_initializer, + cnn_transpose_kernel_initializer_config=( + config.cnn_transpose_kernel_initializer_config + ), + cnn_transpose_bias_initializer=config.cnn_transpose_bias_initializer, + cnn_transpose_bias_initializer_config=( + config.cnn_transpose_bias_initializer_config + ), + ) + + @override(Model) + def _forward(self, inputs: torch.Tensor, **kwargs) -> torch.Tensor: + out = self.initial_dense(inputs) + # Reshape to initial 3D (image-like) format to enter CNN transpose stack. + out = out.reshape((-1,) + tuple(self.config.initial_image_dims)) + out = self.cnn_transpose_net(out) + # Add 0.5 to center (always non-activated, non-normalized) outputs more + # around 0.0. + return out + 0.5 diff --git a/lib/python3.12/site-packages/ray/rllib/core/models/torch/primitives.py b/lib/python3.12/site-packages/ray/rllib/core/models/torch/primitives.py new file mode 100644 index 0000000000000000000000000000000000000000..9df8f1d210a32c3e7565460f8513b3c2e074e105 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/models/torch/primitives.py @@ -0,0 +1,479 @@ +from typing import Callable, Dict, List, Optional, Tuple, Union + +from ray.rllib.core.models.torch.utils import Stride2D +from ray.rllib.models.torch.misc import ( + same_padding, + same_padding_transpose_after_stride, + valid_padding, +) +from ray.rllib.models.utils import get_activation_fn, get_initializer_fn +from ray.rllib.utils.framework import try_import_torch + +torch, nn = try_import_torch() + + +class TorchMLP(nn.Module): + """A multi-layer perceptron with N dense layers. + + All layers (except for an optional additional extra output layer) share the same + activation function, bias setup (use bias or not), and LayerNorm setup + (use layer normalization or not). + + If `output_dim` (int) is not None, an additional, extra output dense layer is added, + which might have its own activation function (e.g. "linear"). However, the output + layer does NOT use layer normalization. + """ + + def __init__( + self, + *, + input_dim: int, + hidden_layer_dims: List[int], + hidden_layer_activation: Union[str, Callable] = "relu", + hidden_layer_use_bias: bool = True, + hidden_layer_use_layernorm: bool = False, + hidden_layer_weights_initializer: Optional[Union[str, Callable]] = None, + hidden_layer_weights_initializer_config: Optional[Union[str, Callable]] = None, + hidden_layer_bias_initializer: Optional[Union[str, Callable]] = None, + hidden_layer_bias_initializer_config: Optional[Dict] = None, + output_dim: Optional[int] = None, + output_use_bias: bool = True, + output_activation: Union[str, Callable] = "linear", + output_weights_initializer: Optional[Union[str, Callable]] = None, + output_weights_initializer_config: Optional[Dict] = None, + output_bias_initializer: Optional[Union[str, Callable]] = None, + output_bias_initializer_config: Optional[Dict] = None, + ): + """Initialize a TorchMLP object. + + Args: + input_dim: The input dimension of the network. Must not be None. + hidden_layer_dims: The sizes of the hidden layers. If an empty list, only a + single layer will be built of size `output_dim`. + hidden_layer_use_layernorm: Whether to insert a LayerNormalization + functionality in between each hidden layer's output and its activation. + hidden_layer_use_bias: Whether to use bias on all dense layers (excluding + the possible separate output layer). + hidden_layer_activation: The activation function to use after each layer + (except for the output). Either a torch.nn.[activation fn] callable or + the name thereof, or an RLlib recognized activation name, + e.g. "ReLU", "relu", "tanh", "SiLU", or "linear". + hidden_layer_weights_initializer: The initializer function or class to use + forweights initialization in the hidden layers. If `None` the default + initializer of the respective dense layer is used. Note, only the + in-place initializers, i.e. ending with an underscore "_" are allowed. + hidden_layer_weights_initializer_config: Configuration to pass into the + initializer defined in `hidden_layer_weights_initializer`. + hidden_layer_bias_initializer: The initializer function or class to use for + bias initialization in the hidden layers. If `None` the default + initializer of the respective dense layer is used. Note, only the + in-place initializers, i.e. ending with an underscore "_" are allowed. + hidden_layer_bias_initializer_config: Configuration to pass into the + initializer defined in `hidden_layer_bias_initializer`. + output_dim: The output dimension of the network. If None, no specific output + layer will be added and the last layer in the stack will have + size=`hidden_layer_dims[-1]`. + output_use_bias: Whether to use bias on the separate output layer, + if any. + output_activation: The activation function to use for the output layer + (if any). Either a torch.nn.[activation fn] callable or + the name thereof, or an RLlib recognized activation name, + e.g. "ReLU", "relu", "tanh", "SiLU", or "linear". + output_layer_weights_initializer: The initializer function or class to use + for weights initialization in the output layers. If `None` the default + initializer of the respective dense layer is used. Note, only the + in-place initializers, i.e. ending with an underscore "_" are allowed. + output_layer_weights_initializer_config: Configuration to pass into the + initializer defined in `output_layer_weights_initializer`. + output_layer_bias_initializer: The initializer function or class to use for + bias initialization in the output layers. If `None` the default + initializer of the respective dense layer is used. Note, only the + in-place initializers, i.e. ending with an underscore "_" are allowed. + output_layer_bias_initializer_config: Configuration to pass into the + initializer defined in `output_layer_bias_initializer`. + """ + super().__init__() + assert input_dim > 0 + + self.input_dim = input_dim + + hidden_activation = get_activation_fn( + hidden_layer_activation, framework="torch" + ) + hidden_weights_initializer = get_initializer_fn( + hidden_layer_weights_initializer, framework="torch" + ) + hidden_bias_initializer = get_initializer_fn( + hidden_layer_bias_initializer, framework="torch" + ) + output_weights_initializer = get_initializer_fn( + output_weights_initializer, framework="torch" + ) + output_bias_initializer = get_initializer_fn( + output_bias_initializer, framework="torch" + ) + + layers = [] + dims = ( + [self.input_dim] + + list(hidden_layer_dims) + + ([output_dim] if output_dim else []) + ) + for i in range(0, len(dims) - 1): + # Whether we are already processing the last (special) output layer. + is_output_layer = output_dim is not None and i == len(dims) - 2 + + layer = nn.Linear( + dims[i], + dims[i + 1], + bias=output_use_bias if is_output_layer else hidden_layer_use_bias, + ) + # Initialize layers, if necessary. + if is_output_layer: + # Initialize output layer weigths if necessary. + if output_weights_initializer: + output_weights_initializer( + layer.weight, **output_weights_initializer_config or {} + ) + # Initialize output layer bias if necessary. + if output_bias_initializer: + output_bias_initializer( + layer.bias, **output_bias_initializer_config or {} + ) + # Must be hidden. + else: + # Initialize hidden layer weights if necessary. + if hidden_layer_weights_initializer: + hidden_weights_initializer( + layer.weight, **hidden_layer_weights_initializer_config or {} + ) + # Initialize hidden layer bias if necessary. + if hidden_layer_bias_initializer: + hidden_bias_initializer( + layer.bias, **hidden_layer_bias_initializer_config or {} + ) + + layers.append(layer) + + # We are still in the hidden layer section: Possibly add layernorm and + # hidden activation. + if not is_output_layer: + # Insert a layer normalization in between layer's output and + # the activation. + if hidden_layer_use_layernorm: + # We use an epsilon of 0.001 here to mimick the Tf default behavior. + layers.append(nn.LayerNorm(dims[i + 1], eps=0.001)) + # Add the activation function. + if hidden_activation is not None: + layers.append(hidden_activation()) + + # Add output layer's (if any) activation. + output_activation = get_activation_fn(output_activation, framework="torch") + if output_dim is not None and output_activation is not None: + layers.append(output_activation()) + + self.mlp = nn.Sequential(*layers) + + def forward(self, x): + return self.mlp(x) + + +class TorchCNN(nn.Module): + """A model containing a CNN with N Conv2D layers. + + All layers share the same activation function, bias setup (use bias or not), + and LayerNorm setup (use layer normalization or not). + + Note that there is no flattening nor an additional dense layer at the end of the + stack. The output of the network is a 3D tensor of dimensions + [width x height x num output filters]. + """ + + def __init__( + self, + *, + input_dims: Union[List[int], Tuple[int, ...]], + cnn_filter_specifiers: List[List[Union[int, List]]], + cnn_use_bias: bool = True, + cnn_use_layernorm: bool = False, + cnn_activation: str = "relu", + cnn_kernel_initializer: Optional[Union[str, Callable]] = None, + cnn_kernel_initializer_config: Optional[Dict] = None, + cnn_bias_initializer: Optional[Union[str, Callable]] = None, + cnn_bias_initializer_config: Optional[Dict] = None, + ): + """Initializes a TorchCNN instance. + + Args: + input_dims: The 3D input dimensions of the network (incoming image). + cnn_filter_specifiers: A list in which each element is another (inner) list + of either the following forms: + `[number of channels/filters, kernel, stride]` + OR: + `[number of channels/filters, kernel, stride, padding]`, where `padding` + can either be "same" or "valid". + When using the first format w/o the `padding` specifier, `padding` is + "same" by default. Also, `kernel` and `stride` may be provided either as + single ints (square) or as a tuple/list of two ints (width- and height + dimensions) for non-squared kernel/stride shapes. + A good rule of thumb for constructing CNN stacks is: + When using padding="same", the input "image" will be reduced in size by + the factor `stride`, e.g. input=(84, 84, 3) stride=2 kernel=x + padding="same" filters=16 -> output=(42, 42, 16). + For example, if you would like to reduce an Atari image from its + original (84, 84, 3) dimensions down to (6, 6, F), you can construct the + following stack and reduce the w x h dimension of the image by 2 in each + layer: + [[16, 4, 2], [32, 4, 2], [64, 4, 2], [128, 4, 2]] -> output=(6, 6, 128) + cnn_use_bias: Whether to use bias on all Conv2D layers. + cnn_activation: The activation function to use after each Conv2D layer. + cnn_use_layernorm: Whether to insert a LayerNormalization functionality + in between each Conv2D layer's outputs and its activation. + cnn_kernel_initializer: The initializer function or class to use for kernel + initialization in the CNN layers. If `None` the default initializer of + the respective CNN layer is used. Note, only the in-place + initializers, i.e. ending with an underscore "_" are allowed. + cnn_kernel_initializer_config: Configuration to pass into the initializer + defined in `cnn_kernel_initializer`. + cnn_bias_initializer: The initializer function or class to use for bias + initializationcin the CNN layers. If `None` the default initializer of + the respective CNN layer is used. Note, only the in-place initializers, + i.e. ending with an underscore "_" are allowed. + cnn_bias_initializer_config: Configuration to pass into the initializer + defined in `cnn_bias_initializer`. + """ + super().__init__() + + assert len(input_dims) == 3 + + cnn_activation = get_activation_fn(cnn_activation, framework="torch") + cnn_kernel_initializer = get_initializer_fn( + cnn_kernel_initializer, framework="torch" + ) + cnn_bias_initializer = get_initializer_fn( + cnn_bias_initializer, framework="torch" + ) + layers = [] + + # Add user-specified hidden convolutional layers first + width, height, in_depth = input_dims + in_size = [width, height] + for filter_specs in cnn_filter_specifiers: + # Padding information not provided -> Use "same" as default. + if len(filter_specs) == 3: + out_depth, kernel_size, strides = filter_specs + padding = "same" + # Padding information provided. + else: + out_depth, kernel_size, strides, padding = filter_specs + + # Pad like in tensorflow's SAME/VALID mode. + if padding == "same": + padding_size, out_size = same_padding(in_size, kernel_size, strides) + layers.append(nn.ZeroPad2d(padding_size)) + # No actual padding is performed for "valid" mode, but we will still + # compute the output size (input for the next layer). + else: + out_size = valid_padding(in_size, kernel_size, strides) + + layer = nn.Conv2d( + in_depth, out_depth, kernel_size, strides, bias=cnn_use_bias + ) + + # Initialize CNN layer kernel if necessary. + if cnn_kernel_initializer: + cnn_kernel_initializer( + layer.weight, **cnn_kernel_initializer_config or {} + ) + # Initialize CNN layer bias if necessary. + if cnn_bias_initializer: + cnn_bias_initializer(layer.bias, **cnn_bias_initializer_config or {}) + + layers.append(layer) + + # Layernorm. + if cnn_use_layernorm: + # We use an epsilon of 0.001 here to mimick the Tf default behavior. + layers.append(LayerNorm1D(out_depth, eps=0.001)) + # Activation. + if cnn_activation is not None: + layers.append(cnn_activation()) + + in_size = out_size + in_depth = out_depth + + # Create the CNN. + self.cnn = nn.Sequential(*layers) + + def forward(self, inputs): + # Permute b/c data comes in as channels_last ([B, dim, dim, channels]) -> + # Convert to `channels_first` for torch: + inputs = inputs.permute(0, 3, 1, 2) + out = self.cnn(inputs) + # Permute back to `channels_last`. + return out.permute(0, 2, 3, 1) + + +class TorchCNNTranspose(nn.Module): + """A model containing a CNNTranspose with N Conv2DTranspose layers. + + All layers share the same activation function, bias setup (use bias or not), + and LayerNormalization setup (use layer normalization or not), except for the last + one, which is never activated and never layer norm'd. + + Note that there is no reshaping/flattening nor an additional dense layer at the + beginning or end of the stack. The input as well as output of the network are 3D + tensors of dimensions [width x height x num output filters]. + """ + + def __init__( + self, + *, + input_dims: Union[List[int], Tuple[int, ...]], + cnn_transpose_filter_specifiers: List[List[Union[int, List]]], + cnn_transpose_use_bias: bool = True, + cnn_transpose_activation: str = "relu", + cnn_transpose_use_layernorm: bool = False, + cnn_transpose_kernel_initializer: Optional[Union[str, Callable]] = None, + cnn_transpose_kernel_initializer_config: Optional[Dict] = None, + cnn_transpose_bias_initializer: Optional[Union[str, Callable]] = None, + cnn_transpose_bias_initializer_config: Optional[Dict] = None, + ): + """Initializes a TorchCNNTranspose instance. + + Args: + input_dims: The 3D input dimensions of the network (incoming image). + cnn_transpose_filter_specifiers: A list of lists, where each item represents + one Conv2DTranspose layer. Each such Conv2DTranspose layer is further + specified by the elements of the inner lists. The inner lists follow + the format: `[number of filters, kernel, stride]` to + specify a convolutional-transpose layer stacked in order of the + outer list. + `kernel` as well as `stride` might be provided as width x height tuples + OR as single ints representing both dimension (width and height) + in case of square shapes. + cnn_transpose_use_bias: Whether to use bias on all Conv2DTranspose layers. + cnn_transpose_use_layernorm: Whether to insert a LayerNormalization + functionality in between each Conv2DTranspose layer's outputs and its + activation. + The last Conv2DTranspose layer will not be normed, regardless. + cnn_transpose_activation: The activation function to use after each layer + (except for the last Conv2DTranspose layer, which is always + non-activated). + cnn_transpose_kernel_initializer: The initializer function or class to use + for kernel initialization in the CNN layers. If `None` the default + initializer of the respective CNN layer is used. Note, only the + in-place initializers, i.e. ending with an underscore "_" are allowed. + cnn_transpose_kernel_initializer_config: Configuration to pass into the + initializer defined in `cnn_transpose_kernel_initializer`. + cnn_transpose_bias_initializer: The initializer function or class to use for + bias initialization in the CNN layers. If `None` the default initializer + of the respective CNN layer is used. Note, only the in-place + initializers, i.e. ending with an underscore "_" are allowed. + cnn_transpose_bias_initializer_config: Configuration to pass into the + initializer defined in `cnn_transpose_bias_initializer`. + """ + super().__init__() + + assert len(input_dims) == 3 + + cnn_transpose_activation = get_activation_fn( + cnn_transpose_activation, framework="torch" + ) + cnn_transpose_kernel_initializer = get_initializer_fn( + cnn_transpose_kernel_initializer, framework="torch" + ) + cnn_transpose_bias_initializer = get_initializer_fn( + cnn_transpose_bias_initializer, framework="torch" + ) + + layers = [] + + # Add user-specified hidden convolutional layers first + width, height, in_depth = input_dims + in_size = [width, height] + for i, (out_depth, kernel, stride) in enumerate( + cnn_transpose_filter_specifiers + ): + is_final_layer = i == len(cnn_transpose_filter_specifiers) - 1 + + # Resolve stride and kernel width/height values if only int given (squared). + s_w, s_h = (stride, stride) if isinstance(stride, int) else stride + k_w, k_h = (kernel, kernel) if isinstance(kernel, int) else kernel + + # Stride the incoming image first. + stride_layer = Stride2D(in_size[0], in_size[1], s_w, s_h) + layers.append(stride_layer) + # Then 0-pad (like in tensorflow's SAME mode). + # This will return the necessary padding such that for stride=1, the output + # image has the same size as the input image, for stride=2, the output image + # is 2x the input image, etc.. + padding, out_size = same_padding_transpose_after_stride( + (stride_layer.out_width, stride_layer.out_height), kernel, stride + ) + layers.append(nn.ZeroPad2d(padding)) # left, right, top, bottom + # Then do the Conv2DTranspose operation + # (now that we have padded and strided manually, w/o any more padding using + # stride=1). + + layer = nn.ConvTranspose2d( + in_depth, + out_depth, + kernel, + # Force-set stride to 1 as we already took care of it. + 1, + # Disable torch auto-padding (torch interprets the padding setting + # as: dilation (==1.0) * [`kernel` - 1] - [`padding`]). + padding=(k_w - 1, k_h - 1), + # Last layer always uses bias (b/c has no LayerNorm, regardless of + # config). + bias=cnn_transpose_use_bias or is_final_layer, + ) + + # Initialize CNN Transpose layer kernel if necessary. + if cnn_transpose_kernel_initializer: + cnn_transpose_kernel_initializer( + layer.weight, **cnn_transpose_kernel_initializer_config or {} + ) + # Initialize CNN Transpose layer bias if necessary. + if cnn_transpose_bias_initializer: + cnn_transpose_bias_initializer( + layer.bias, **cnn_transpose_bias_initializer_config or {} + ) + + layers.append(layer) + # Layernorm (never for final layer). + if cnn_transpose_use_layernorm and not is_final_layer: + layers.append(LayerNorm1D(out_depth, eps=0.001)) + # Last layer is never activated (regardless of config). + if cnn_transpose_activation is not None and not is_final_layer: + layers.append(cnn_transpose_activation()) + + in_size = (out_size[0], out_size[1]) + in_depth = out_depth + + # Create the final CNNTranspose network. + self.cnn_transpose = nn.Sequential(*layers) + + def forward(self, inputs): + # Permute b/c data comes in as [B, dim, dim, channels]: + out = inputs.permute(0, 3, 1, 2) + out = self.cnn_transpose(out) + return out.permute(0, 2, 3, 1) + + +class LayerNorm1D(nn.Module): + def __init__(self, num_features, **kwargs): + super().__init__() + self.layer_norm = nn.LayerNorm(num_features, **kwargs) + + def forward(self, x): + # x shape: (B, dim, dim, channels). + batch_size, channels, h, w = x.size() + # Reshape to (batch_size * height * width, channels) for LayerNorm + x = x.permute(0, 2, 3, 1).reshape(-1, channels) + # Apply LayerNorm + x = self.layer_norm(x) + # Reshape back to (batch_size, dim, dim, channels) + x = x.reshape(batch_size, h, w, channels).permute(0, 3, 1, 2) + return x diff --git a/lib/python3.12/site-packages/ray/rllib/core/models/torch/utils.py b/lib/python3.12/site-packages/ray/rllib/core/models/torch/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..1bdbdef016f4c8a7261ff83f14d9c70810d6a90e --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/models/torch/utils.py @@ -0,0 +1,85 @@ +from ray.rllib.utils.framework import try_import_torch + +torch, nn = try_import_torch() + + +class Stride2D(nn.Module): + """A striding layer for doing torch Conv2DTranspose operations. + + Using this layer before the 0-padding (on a 3D input "image") and before + the actual ConvTranspose2d allows for a padding="same" behavior that matches + 100% that of a `tf.keras.layers.Conv2DTranspose` layer. + + Examples: + Input image (4x4): + A B C D + E F G H + I J K L + M N O P + + Stride with stride=2 -> output image=(7x7) + A 0 B 0 C 0 D + 0 0 0 0 0 0 0 + E 0 F 0 G 0 H + 0 0 0 0 0 0 0 + I 0 J 0 K 0 L + 0 0 0 0 0 0 0 + M 0 N 0 O 0 P + """ + + def __init__(self, width, height, stride_w, stride_h): + """Initializes a Stride2D instance. + + Args: + width: The width of the 3D input "image". + height: The height of the 3D input "image". + stride_w: The stride in width direction, with which to stride the incoming + image. + stride_h: The stride in height direction, with which to stride the incoming + image. + """ + super().__init__() + + self.width = width + self.height = height + self.stride_w = stride_w + self.stride_h = stride_h + + self.register_buffer( + "zeros", + torch.zeros( + size=( + self.width * self.stride_w - (self.stride_w - 1), + self.height * self.stride_h - (self.stride_h - 1), + ), + dtype=torch.float32, + ), + ) + + self.out_width, self.out_height = self.zeros.shape[0], self.zeros.shape[1] + # Squeeze in batch and channel dims. + self.zeros = self.zeros.unsqueeze(0).unsqueeze(0) + + where_template = torch.zeros( + (self.stride_w, self.stride_h), dtype=torch.float32 + ) + # Set upper/left corner to 1.0. + where_template[0][0] = 1.0 + # then tile across the entire (strided) image size. + where_template = where_template.repeat((self.height, self.width))[ + : -(self.stride_w - 1), : -(self.stride_h - 1) + ] + # Squeeze in batch and channel dims and convert to bool. + where_template = where_template.unsqueeze(0).unsqueeze(0).bool() + self.register_buffer("where_template", where_template) + + def forward(self, x): + # Repeat incoming image stride(w/h) times to match the strided output template. + repeated_x = ( + x.repeat_interleave(self.stride_w, dim=-2).repeat_interleave( + self.stride_h, dim=-1 + ) + )[:, :, : -(self.stride_w - 1), : -(self.stride_h - 1)] + # Where `self.where_template` == 1.0 -> Use image pixel, otherwise use + # zero filler value. + return torch.where(self.where_template, repeated_x, self.zeros) diff --git a/lib/python3.12/site-packages/ray/rllib/core/testing/__init__.py b/lib/python3.12/site-packages/ray/rllib/core/testing/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/core/testing/bc_algorithm.py b/lib/python3.12/site-packages/ray/rllib/core/testing/bc_algorithm.py new file mode 100644 index 0000000000000000000000000000000000000000..b5768e309f69061601163f0c67482a68b6335cf2 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/testing/bc_algorithm.py @@ -0,0 +1,40 @@ +"""Contains example implementation of a custom algorithm. + +Note: It doesn't include any real use-case functionality; it only serves as an example +to test the algorithm construction and customization. +""" + +from ray.rllib.algorithms import Algorithm, AlgorithmConfig +from ray.rllib.core.rl_module.rl_module import RLModuleSpec +from ray.rllib.core.testing.torch.bc_learner import BCTorchLearner +from ray.rllib.core.testing.torch.bc_module import DiscreteBCTorchModule +from ray.rllib.policy.torch_policy_v2 import TorchPolicyV2 +from ray.rllib.utils.annotations import override +from ray.rllib.utils.typing import ResultDict + + +class BCConfigTest(AlgorithmConfig): + def __init__(self, algo_class=None): + super().__init__(algo_class=algo_class or BCAlgorithmTest) + + def get_default_rl_module_spec(self): + if self.framework_str == "torch": + return RLModuleSpec(module_class=DiscreteBCTorchModule) + + def get_default_learner_class(self): + if self.framework_str == "torch": + return BCTorchLearner + + +class BCAlgorithmTest(Algorithm): + @classmethod + def get_default_policy_class(cls, config: AlgorithmConfig): + if config.framework_str == "torch": + return TorchPolicyV2 + else: + raise ValueError("Unknown framework: {}".format(config.framework_str)) + + @override(Algorithm) + def training_step(self) -> ResultDict: + # do nothing. + return {} diff --git a/lib/python3.12/site-packages/ray/rllib/core/testing/testing_learner.py b/lib/python3.12/site-packages/ray/rllib/core/testing/testing_learner.py new file mode 100644 index 0000000000000000000000000000000000000000..a16c0e11ae23fd00f090e375aed08142e05b43bd --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/core/testing/testing_learner.py @@ -0,0 +1,67 @@ +from typing import Type + +import numpy as np + +from ray.rllib.algorithms.algorithm_config import AlgorithmConfig +from ray.rllib.core import DEFAULT_MODULE_ID +from ray.rllib.core.learner.learner import Learner +from ray.rllib.core.rl_module.multi_rl_module import ( + MultiRLModule, + MultiRLModuleSpec, +) +from ray.rllib.core.rl_module.rl_module import RLModuleSpec +from ray.rllib.utils.annotations import override +from ray.rllib.utils.numpy import convert_to_numpy +from ray.rllib.utils.typing import RLModuleSpecType + + +class BaseTestingAlgorithmConfig(AlgorithmConfig): + # A test setting to activate metrics on mean weights. + report_mean_weights: bool = True + + @override(AlgorithmConfig) + def get_default_learner_class(self) -> Type["Learner"]: + if self.framework_str == "torch": + from ray.rllib.core.testing.torch.bc_learner import BCTorchLearner + + return BCTorchLearner + else: + raise ValueError(f"Unsupported framework: {self.framework_str}") + + @override(AlgorithmConfig) + def get_default_rl_module_spec(self) -> "RLModuleSpecType": + if self.framework_str == "torch": + from ray.rllib.core.testing.torch.bc_module import DiscreteBCTorchModule + + cls = DiscreteBCTorchModule + else: + raise ValueError(f"Unsupported framework: {self.framework_str}") + + spec = RLModuleSpec( + module_class=cls, + model_config={"fcnet_hiddens": [32]}, + ) + + if self.is_multi_agent: + # TODO (Kourosh): Make this more multi-agent for example with policy ids + # "1" and "2". + return MultiRLModuleSpec( + multi_rl_module_class=MultiRLModule, + rl_module_specs={DEFAULT_MODULE_ID: spec}, + ) + else: + return spec + + +class BaseTestingLearner(Learner): + @override(Learner) + def after_gradient_based_update(self, *, timesteps): + # This is to check if in the multi-gpu case, the weights across workers are + # the same. It is really only needed during testing. + if self.config.report_mean_weights: + for module_id in self.module.keys(): + parameters = convert_to_numpy( + self.get_parameters(self.module[module_id]) + ) + mean_ws = np.mean([w.mean() for w in parameters]) + self.metrics.log_value((module_id, "mean_weight"), mean_ws, window=1) diff --git a/lib/python3.12/site-packages/ray/rllib/models/__init__.py b/lib/python3.12/site-packages/ray/rllib/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c1282a2d055d66d479c40e091f0e5fb367e1353e --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/__init__.py @@ -0,0 +1,12 @@ +from ray.rllib.models.action_dist import ActionDistribution +from ray.rllib.models.catalog import MODEL_DEFAULTS, ModelCatalog +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.models.preprocessors import Preprocessor + +__all__ = [ + "ActionDistribution", + "ModelCatalog", + "ModelV2", + "Preprocessor", + "MODEL_DEFAULTS", +] diff --git a/lib/python3.12/site-packages/ray/rllib/models/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3d4ac93b1a35a132a54b521ef49b87c6720f698d Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/__pycache__/action_dist.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/__pycache__/action_dist.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ba65fbb0db0d69320dc511e686bb04f4234dffd7 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/__pycache__/action_dist.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/__pycache__/catalog.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/__pycache__/catalog.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..034dfbbe3860b855d66176ddd3300d69bfe85f28 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/__pycache__/catalog.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/__pycache__/distributions.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/__pycache__/distributions.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6ad27b586d3deeabcfee39c1355c4cbbb654748c Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/__pycache__/distributions.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/__pycache__/modelv2.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/__pycache__/modelv2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b5f29db9e0bbfc52ff5589935cf8a77a1006117a Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/__pycache__/modelv2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/__pycache__/preprocessors.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/__pycache__/preprocessors.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ee114a21d10773eda8d64e8209d8d3165eb80589 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/__pycache__/preprocessors.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/__pycache__/repeated_values.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/__pycache__/repeated_values.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c25deb01989240ab0a5573dfbeeee38b59c31410 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/__pycache__/repeated_values.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..51309e1f75c468ddb103b442a18cb19613bb7a7a Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/action_dist.py b/lib/python3.12/site-packages/ray/rllib/models/action_dist.py new file mode 100644 index 0000000000000000000000000000000000000000..c553db8f6caef101953c5230958f1235cd3f02d5 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/action_dist.py @@ -0,0 +1,94 @@ +import gymnasium as gym +import numpy as np + +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.typing import List, ModelConfigDict, TensorType, Union + + +@OldAPIStack +class ActionDistribution: + """The policy action distribution of an agent. + + Attributes: + inputs: input vector to compute samples from. + model (ModelV2): reference to model producing the inputs. + """ + + def __init__(self, inputs: List[TensorType], model: ModelV2): + """Initializes an ActionDist object. + + Args: + inputs: input vector to compute samples from. + model (ModelV2): reference to model producing the inputs. This + is mainly useful if you want to use model variables to compute + action outputs (i.e., for autoregressive action distributions, + see examples/autoregressive_action_dist.py). + """ + self.inputs = inputs + self.model = model + + def sample(self) -> TensorType: + """Draw a sample from the action distribution.""" + raise NotImplementedError + + def deterministic_sample(self) -> TensorType: + """ + Get the deterministic "sampling" output from the distribution. + This is usually the max likelihood output, i.e. mean for Normal, argmax + for Categorical, etc.. + """ + raise NotImplementedError + + def sampled_action_logp(self) -> TensorType: + """Returns the log probability of the last sampled action.""" + raise NotImplementedError + + def logp(self, x: TensorType) -> TensorType: + """The log-likelihood of the action distribution.""" + raise NotImplementedError + + def kl(self, other: "ActionDistribution") -> TensorType: + """The KL-divergence between two action distributions.""" + raise NotImplementedError + + def entropy(self) -> TensorType: + """The entropy of the action distribution.""" + raise NotImplementedError + + def multi_kl(self, other: "ActionDistribution") -> TensorType: + """The KL-divergence between two action distributions. + + This differs from kl() in that it can return an array for + MultiDiscrete. TODO(ekl) consider removing this. + """ + return self.kl(other) + + def multi_entropy(self) -> TensorType: + """The entropy of the action distribution. + + This differs from entropy() in that it can return an array for + MultiDiscrete. TODO(ekl) consider removing this. + """ + return self.entropy() + + @staticmethod + @OldAPIStack + def required_model_output_shape( + action_space: gym.Space, model_config: ModelConfigDict + ) -> Union[int, np.ndarray]: + """Returns the required shape of an input parameter tensor for a + particular action space and an optional dict of distribution-specific + options. + + Args: + action_space (gym.Space): The action space this distribution will + be used for, whose shape attributes will be used to determine + the required shape of the input parameter tensor. + model_config: Model's config dict (as defined in catalog.py) + + Returns: + model_output_shape (int or np.ndarray of ints): size of the + required input vector (minus leading batch dimension). + """ + raise NotImplementedError diff --git a/lib/python3.12/site-packages/ray/rllib/models/catalog.py b/lib/python3.12/site-packages/ray/rllib/models/catalog.py new file mode 100644 index 0000000000000000000000000000000000000000..324d769ffe73d8f1c1f32aae6b40a167dd22a8bf --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/catalog.py @@ -0,0 +1,906 @@ +import logging +from functools import partial +from typing import List, Optional, Type, Union + +import gymnasium as gym +import numpy as np +import tree # pip install dm_tree +from gymnasium.spaces import Box, Dict, Discrete, MultiDiscrete, Tuple + +from ray._common.deprecation import ( + DEPRECATED_VALUE, + deprecation_warning, +) +from ray.rllib.models.action_dist import ActionDistribution +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.models.preprocessors import Preprocessor, get_preprocessor +from ray.rllib.models.tf.tf_action_dist import ( + Categorical, + Deterministic, + DiagGaussian, + Dirichlet, + MultiActionDistribution, + MultiCategorical, +) +from ray.rllib.models.torch.torch_action_dist import ( + TorchCategorical, + TorchDeterministic, + TorchDiagGaussian, + TorchDirichlet, + TorchMultiActionDistribution, + TorchMultiCategorical, +) +from ray.rllib.utils.annotations import DeveloperAPI, PublicAPI +from ray.rllib.utils.error import UnsupportedSpaceException +from ray.rllib.utils.framework import try_import_tf, try_import_torch +from ray.rllib.utils.from_config import from_config +from ray.rllib.utils.spaces.simplex import Simplex +from ray.rllib.utils.spaces.space_utils import flatten_space +from ray.rllib.utils.typing import ModelConfigDict, TensorType +from ray.tune.registry import ( + RLLIB_ACTION_DIST, + RLLIB_MODEL, + _global_registry, +) + +tf1, tf, tfv = try_import_tf() +torch, _ = try_import_torch() + +logger = logging.getLogger(__name__) + +# fmt: off +# __sphinx_doc_begin__ +MODEL_DEFAULTS: ModelConfigDict = { + "fcnet_hiddens": [256, 256], + "fcnet_activation": "tanh", + "fcnet_weights_initializer": None, + "fcnet_weights_initializer_config": None, + "fcnet_bias_initializer": None, + "fcnet_bias_initializer_config": None, + "conv_filters": None, + "conv_activation": "relu", + "conv_kernel_initializer": None, + "conv_kernel_initializer_config": None, + "conv_bias_initializer": None, + "conv_bias_initializer_config": None, + "conv_transpose_kernel_initializer": None, + "conv_transpose_kernel_initializer_config": None, + "conv_transpose_bias_initializer": None, + "conv_transpose_bias_initializer_config": None, + "post_fcnet_hiddens": [], + "post_fcnet_activation": "relu", + "post_fcnet_weights_initializer": None, + "post_fcnet_weights_initializer_config": None, + "post_fcnet_bias_initializer": None, + "post_fcnet_bias_initializer_config": None, + "free_log_std": False, + "log_std_clip_param": 20.0, + "no_final_linear": False, + "vf_share_layers": True, + "use_lstm": False, + "max_seq_len": 20, + "lstm_cell_size": 256, + "lstm_use_prev_action": False, + "lstm_use_prev_reward": False, + "lstm_weights_initializer": None, + "lstm_weights_initializer_config": None, + "lstm_bias_initializer": None, + "lstm_bias_initializer_config": None, + "_time_major": False, + "use_attention": False, + "attention_num_transformer_units": 1, + "attention_dim": 64, + "attention_num_heads": 1, + "attention_head_dim": 32, + "attention_memory_inference": 50, + "attention_memory_training": 50, + "attention_position_wise_mlp_dim": 32, + "attention_init_gru_gate_bias": 2.0, + "attention_use_n_prev_actions": 0, + "attention_use_n_prev_rewards": 0, + "framestack": True, + "dim": 84, + "grayscale": False, + "zero_mean": True, + "custom_model": None, + "custom_model_config": {}, + "custom_action_dist": None, + "custom_preprocessor": None, + "encoder_latent_dim": None, + "always_check_shapes": False, + + # Deprecated keys: + "lstm_use_prev_action_reward": DEPRECATED_VALUE, + "_use_default_native_models": DEPRECATED_VALUE, + "_disable_preprocessor_api": False, + "_disable_action_flattening": False, +} +# __sphinx_doc_end__ +# fmt: on + + +@DeveloperAPI +class ModelCatalog: + """Registry of models, preprocessors, and action distributions for envs. + + .. testcode:: + :skipif: True + + prep = ModelCatalog.get_preprocessor(env) + observation = prep.transform(raw_observation) + + dist_class, dist_dim = ModelCatalog.get_action_dist( + env.action_space, {}) + model = ModelCatalog.get_model_v2( + obs_space, action_space, num_outputs, options) + dist = dist_class(model.outputs, model) + action = dist.sample() + """ + + @staticmethod + @DeveloperAPI + def get_action_dist( + action_space: gym.Space, + config: ModelConfigDict, + dist_type: Optional[Union[str, Type[ActionDistribution]]] = None, + framework: str = "tf", + **kwargs + ) -> (type, int): + """Returns a distribution class and size for the given action space. + + Args: + action_space: Action space of the target gym env. + config (Optional[dict]): Optional model config. + dist_type (Optional[Union[str, Type[ActionDistribution]]]): + Identifier of the action distribution (str) interpreted as a + hint or the actual ActionDistribution class to use. + framework: One of "tf2", "tf", "torch", or "jax". + kwargs: Optional kwargs to pass on to the Distribution's + constructor. + + Returns: + Tuple: + - dist_class (ActionDistribution): Python class of the + distribution. + - dist_dim (int): The size of the input vector to the + distribution. + """ + + dist_cls = None + config = config or MODEL_DEFAULTS + # Custom distribution given. + if config.get("custom_action_dist"): + custom_action_config = config.copy() + action_dist_name = custom_action_config.pop("custom_action_dist") + logger.debug("Using custom action distribution {}".format(action_dist_name)) + dist_cls = _global_registry.get(RLLIB_ACTION_DIST, action_dist_name) + return ModelCatalog._get_multi_action_distribution( + dist_cls, action_space, custom_action_config, framework + ) + + # Dist_type is given directly as a class. + elif ( + type(dist_type) is type + and issubclass(dist_type, ActionDistribution) + and dist_type not in (MultiActionDistribution, TorchMultiActionDistribution) + ): + dist_cls = dist_type + # Box space -> DiagGaussian OR Deterministic. + elif isinstance(action_space, Box): + if action_space.dtype.name.startswith("int"): + low_ = np.min(action_space.low) + high_ = np.max(action_space.high) + dist_cls = ( + TorchMultiCategorical if framework == "torch" else MultiCategorical + ) + num_cats = int(np.prod(action_space.shape)) + return ( + partial( + dist_cls, + input_lens=[high_ - low_ + 1 for _ in range(num_cats)], + action_space=action_space, + ), + num_cats * (high_ - low_ + 1), + ) + else: + if len(action_space.shape) > 1: + raise UnsupportedSpaceException( + "Action space has multiple dimensions " + "{}. ".format(action_space.shape) + + "Consider reshaping this into a single dimension, " + "using a custom action distribution, " + "using a Tuple action space, or the multi-agent API." + ) + # TODO(sven): Check for bounds and return SquashedNormal, etc.. + if dist_type is None: + return ( + partial( + TorchDiagGaussian if framework == "torch" else DiagGaussian, + action_space=action_space, + ), + DiagGaussian.required_model_output_shape(action_space, config), + ) + elif dist_type == "deterministic": + dist_cls = ( + TorchDeterministic if framework == "torch" else Deterministic + ) + # Discrete Space -> Categorical. + elif isinstance(action_space, Discrete): + if framework == "torch": + dist_cls = TorchCategorical + elif framework == "jax": + from ray.rllib.models.jax.jax_action_dist import JAXCategorical + + dist_cls = JAXCategorical + else: + dist_cls = Categorical + # Tuple/Dict Spaces -> MultiAction. + elif dist_type in ( + MultiActionDistribution, + TorchMultiActionDistribution, + ) or isinstance(action_space, (Tuple, Dict)): + return ModelCatalog._get_multi_action_distribution( + ( + MultiActionDistribution + if framework == "tf" + else TorchMultiActionDistribution + ), + action_space, + config, + framework, + ) + # Simplex -> Dirichlet. + elif isinstance(action_space, Simplex): + dist_cls = TorchDirichlet if framework == "torch" else Dirichlet + # MultiDiscrete -> MultiCategorical. + elif isinstance(action_space, MultiDiscrete): + dist_cls = ( + TorchMultiCategorical if framework == "torch" else MultiCategorical + ) + return partial(dist_cls, input_lens=action_space.nvec), int( + sum(action_space.nvec) + ) + # Unknown type -> Error. + else: + raise NotImplementedError( + "Unsupported args: {} {}".format(action_space, dist_type) + ) + + return dist_cls, int(dist_cls.required_model_output_shape(action_space, config)) + + @staticmethod + @DeveloperAPI + def get_action_shape( + action_space: gym.Space, framework: str = "tf" + ) -> (np.dtype, List[int]): + """Returns action tensor dtype and shape for the action space. + + Args: + action_space: Action space of the target gym env. + framework: The framework identifier. One of "tf" or "torch". + + Returns: + (dtype, shape): Dtype and shape of the actions tensor. + """ + dl_lib = torch if framework == "torch" else tf + if isinstance(action_space, Discrete): + return action_space.dtype, (None,) + elif isinstance(action_space, (Box, Simplex)): + if np.issubdtype(action_space.dtype, np.floating): + return dl_lib.float32, (None,) + action_space.shape + elif np.issubdtype(action_space.dtype, np.integer): + return dl_lib.int32, (None,) + action_space.shape + else: + raise ValueError("RLlib doesn't support non int or float box spaces") + elif isinstance(action_space, MultiDiscrete): + return action_space.dtype, (None,) + action_space.shape + elif isinstance(action_space, (Tuple, Dict)): + flat_action_space = flatten_space(action_space) + size = 0 + all_discrete = True + for i in range(len(flat_action_space)): + if isinstance(flat_action_space[i], Discrete): + size += 1 + else: + all_discrete = False + size += np.prod(flat_action_space[i].shape) + size = int(size) + return dl_lib.int32 if all_discrete else dl_lib.float32, (None, size) + else: + raise NotImplementedError( + "Action space {} not supported".format(action_space) + ) + + @staticmethod + @DeveloperAPI + def get_action_placeholder( + action_space: gym.Space, name: str = "action" + ) -> TensorType: + """Returns an action placeholder consistent with the action space + + Args: + action_space: Action space of the target gym env. + name: An optional string to name the placeholder by. + Default: "action". + + Returns: + action_placeholder: A placeholder for the actions + """ + dtype, shape = ModelCatalog.get_action_shape(action_space, framework="tf") + + return tf1.placeholder(dtype, shape=shape, name=name) + + @staticmethod + @DeveloperAPI + def get_model_v2( + obs_space: gym.Space, + action_space: gym.Space, + num_outputs: int, + model_config: ModelConfigDict, + framework: str = "tf", + name: str = "default_model", + model_interface: type = None, + default_model: type = None, + **model_kwargs + ) -> ModelV2: + """Returns a suitable model compatible with given spaces and output. + + Args: + obs_space: Observation space of the target gym env. This + may have an `original_space` attribute that specifies how to + unflatten the tensor into a ragged tensor. + action_space: Action space of the target gym env. + num_outputs: The size of the output vector of the model. + model_config: The "model" sub-config dict + within the Algorithm's config dict. + framework: One of "tf2", "tf", "torch", or "jax". + name: Name (scope) for the model. + model_interface: Interface required for the model + default_model: Override the default class for the model. This + only has an effect when not using a custom model + model_kwargs: Args to pass to the ModelV2 constructor + + Returns: + model (ModelV2): Model to use for the policy. + """ + + # Validate the given config dict. + ModelCatalog._validate_config( + config=model_config, action_space=action_space, framework=framework + ) + + if model_config.get("custom_model"): + # Allow model kwargs to be overridden / augmented by + # custom_model_config. + customized_model_kwargs = dict( + model_kwargs, **model_config.get("custom_model_config", {}) + ) + + if isinstance(model_config["custom_model"], type): + model_cls = model_config["custom_model"] + elif ( + isinstance(model_config["custom_model"], str) + and "." in model_config["custom_model"] + ): + return from_config( + cls=model_config["custom_model"], + obs_space=obs_space, + action_space=action_space, + num_outputs=num_outputs, + model_config=customized_model_kwargs, + name=name, + ) + else: + model_cls = _global_registry.get( + RLLIB_MODEL, model_config["custom_model"] + ) + + # Only allow ModelV2 or native keras Models. + if not issubclass(model_cls, ModelV2): + if framework not in ["tf", "tf2"] or not issubclass( + model_cls, tf.keras.Model + ): + raise ValueError( + "`model_cls` must be a ModelV2 sub-class, but is" + " {}!".format(model_cls) + ) + + logger.info("Wrapping {} as {}".format(model_cls, model_interface)) + model_cls = ModelCatalog._wrap_if_needed(model_cls, model_interface) + + if framework in ["tf2", "tf"]: + # Try wrapping custom model with LSTM/attention, if required. + if model_config.get("use_lstm") or model_config.get("use_attention"): + from ray.rllib.models.tf.attention_net import ( + AttentionWrapper, + ) + from ray.rllib.models.tf.recurrent_net import ( + LSTMWrapper, + ) + + wrapped_cls = model_cls + forward = wrapped_cls.forward + model_cls = ModelCatalog._wrap_if_needed( + wrapped_cls, + LSTMWrapper + if model_config.get("use_lstm") + else AttentionWrapper, + ) + model_cls._wrapped_forward = forward + + # Obsolete: Track and warn if vars were created but not + # registered. Only still do this, if users do register their + # variables. If not (which they shouldn't), don't check here. + created = set() + + def track_var_creation(next_creator, **kw): + v = next_creator(**kw) + created.add(v.ref()) + return v + + with tf.variable_creator_scope(track_var_creation): + if issubclass(model_cls, tf.keras.Model): + instance = model_cls( + input_space=obs_space, + action_space=action_space, + num_outputs=num_outputs, + name=name, + **customized_model_kwargs, + ) + else: + # Try calling with kwargs first (custom ModelV2 should + # accept these as kwargs, not get them from + # config["custom_model_config"] anymore). + try: + instance = model_cls( + obs_space, + action_space, + num_outputs, + model_config, + name, + **customized_model_kwargs, + ) + except TypeError as e: + # Keyword error: Try old way w/o kwargs. + if "__init__() got an unexpected " in e.args[0]: + instance = model_cls( + obs_space, + action_space, + num_outputs, + model_config, + name, + **model_kwargs, + ) + logger.warning( + "Custom ModelV2 should accept all custom " + "options as **kwargs, instead of expecting" + " them in config['custom_model_config']!" + ) + # Other error -> re-raise. + else: + raise e + + # User still registered TFModelV2's variables: Check, whether + # ok. + registered = [] + if not isinstance(instance, tf.keras.Model): + registered = set(instance.var_list) + if len(registered) > 0: + not_registered = set() + for var in created: + if var not in registered: + not_registered.add(var) + if not_registered: + raise ValueError( + "It looks like you are still using " + "`{}.register_variables()` to register your " + "model's weights. This is no longer required, but " + "if you are still calling this method at least " + "once, you must make sure to register all created " + "variables properly. The missing variables are {}," + " and you only registered {}. " + "Did you forget to call `register_variables()` on " + "some of the variables in question?".format( + instance, not_registered, registered + ) + ) + elif framework == "torch": + # Try wrapping custom model with LSTM/attention, if required. + if model_config.get("use_lstm") or model_config.get("use_attention"): + from ray.rllib.models.torch.attention_net import AttentionWrapper + from ray.rllib.models.torch.recurrent_net import LSTMWrapper + + wrapped_cls = model_cls + forward = wrapped_cls.forward + model_cls = ModelCatalog._wrap_if_needed( + wrapped_cls, + LSTMWrapper + if model_config.get("use_lstm") + else AttentionWrapper, + ) + model_cls._wrapped_forward = forward + + # PyTorch automatically tracks nn.Modules inside the parent + # nn.Module's constructor. + # Try calling with kwargs first (custom ModelV2 should + # accept these as kwargs, not get them from + # config["custom_model_config"] anymore). + try: + instance = model_cls( + obs_space, + action_space, + num_outputs, + model_config, + name, + **customized_model_kwargs, + ) + except TypeError as e: + # Keyword error: Try old way w/o kwargs. + if "__init__() got an unexpected " in e.args[0]: + instance = model_cls( + obs_space, + action_space, + num_outputs, + model_config, + name, + **model_kwargs, + ) + logger.warning( + "Custom ModelV2 should accept all custom " + "options as **kwargs, instead of expecting" + " them in config['custom_model_config']!" + ) + # Other error -> re-raise. + else: + raise e + else: + raise NotImplementedError( + "`framework` must be 'tf2|tf|torch', but is " + "{}!".format(framework) + ) + + return instance + + # Find a default TFModelV2 and wrap with model_interface. + if framework in ["tf", "tf2"]: + v2_class = None + # Try to get a default v2 model. + if not model_config.get("custom_model"): + v2_class = default_model or ModelCatalog._get_v2_model_class( + obs_space, model_config, framework=framework + ) + + if not v2_class: + raise ValueError("ModelV2 class could not be determined!") + + if model_config.get("use_lstm") or model_config.get("use_attention"): + from ray.rllib.models.tf.attention_net import ( + AttentionWrapper, + ) + from ray.rllib.models.tf.recurrent_net import ( + LSTMWrapper, + ) + + wrapped_cls = v2_class + if model_config.get("use_lstm"): + v2_class = ModelCatalog._wrap_if_needed(wrapped_cls, LSTMWrapper) + v2_class._wrapped_forward = wrapped_cls.forward + else: + v2_class = ModelCatalog._wrap_if_needed( + wrapped_cls, AttentionWrapper + ) + v2_class._wrapped_forward = wrapped_cls.forward + + # Wrap in the requested interface. + wrapper = ModelCatalog._wrap_if_needed(v2_class, model_interface) + + if issubclass(wrapper, tf.keras.Model): + model = wrapper( + input_space=obs_space, + action_space=action_space, + num_outputs=num_outputs, + name=name, + **dict(model_kwargs, **model_config), + ) + return model + + return wrapper( + obs_space, action_space, num_outputs, model_config, name, **model_kwargs + ) + + # Find a default TorchModelV2 and wrap with model_interface. + elif framework == "torch": + # Try to get a default v2 model. + if not model_config.get("custom_model"): + v2_class = default_model or ModelCatalog._get_v2_model_class( + obs_space, model_config, framework=framework + ) + + if not v2_class: + raise ValueError("ModelV2 class could not be determined!") + + if model_config.get("use_lstm") or model_config.get("use_attention"): + from ray.rllib.models.torch.attention_net import AttentionWrapper + from ray.rllib.models.torch.recurrent_net import LSTMWrapper + + wrapped_cls = v2_class + forward = wrapped_cls.forward + if model_config.get("use_lstm"): + v2_class = ModelCatalog._wrap_if_needed(wrapped_cls, LSTMWrapper) + else: + v2_class = ModelCatalog._wrap_if_needed( + wrapped_cls, AttentionWrapper + ) + + v2_class._wrapped_forward = forward + + # Wrap in the requested interface. + wrapper = ModelCatalog._wrap_if_needed(v2_class, model_interface) + return wrapper( + obs_space, action_space, num_outputs, model_config, name, **model_kwargs + ) + + # Find a default JAXModelV2 and wrap with model_interface. + elif framework == "jax": + v2_class = default_model or ModelCatalog._get_v2_model_class( + obs_space, model_config, framework=framework + ) + # Wrap in the requested interface. + wrapper = ModelCatalog._wrap_if_needed(v2_class, model_interface) + return wrapper( + obs_space, action_space, num_outputs, model_config, name, **model_kwargs + ) + else: + raise NotImplementedError( + "`framework` must be 'tf2|tf|torch', but is {}!".format(framework) + ) + + @staticmethod + @DeveloperAPI + def get_preprocessor( + env: gym.Env, options: Optional[dict] = None, include_multi_binary: bool = False + ) -> Preprocessor: + """Returns a suitable preprocessor for the given env. + + This is a wrapper for get_preprocessor_for_space(). + """ + + return ModelCatalog.get_preprocessor_for_space( + env.observation_space, options, include_multi_binary + ) + + @staticmethod + @DeveloperAPI + def get_preprocessor_for_space( + observation_space: gym.Space, + options: dict = None, + include_multi_binary: bool = False, + ) -> Preprocessor: + """Returns a suitable preprocessor for the given observation space. + + Args: + observation_space: The input observation space. + options: Options to pass to the preprocessor. + include_multi_binary: Whether to include the MultiBinaryPreprocessor in + the possible preprocessors returned by this method. + + Returns: + preprocessor: Preprocessor for the observations. + """ + + options = options or MODEL_DEFAULTS + for k in options.keys(): + if k not in MODEL_DEFAULTS: + raise Exception( + "Unknown config key `{}`, all keys: {}".format( + k, list(MODEL_DEFAULTS) + ) + ) + + cls = get_preprocessor( + observation_space, include_multi_binary=include_multi_binary + ) + prep = cls(observation_space, options) + + if prep is not None: + logger.debug( + "Created preprocessor {}: {} -> {}".format( + prep, observation_space, prep.shape + ) + ) + return prep + + @staticmethod + @PublicAPI + def register_custom_model(model_name: str, model_class: type) -> None: + """Register a custom model class by name. + + The model can be later used by specifying {"custom_model": model_name} + in the model config. + + Args: + model_name: Name to register the model under. + model_class: Python class of the model. + """ + if tf is not None: + if issubclass(model_class, tf.keras.Model): + deprecation_warning(old="register_custom_model", error=False) + _global_registry.register(RLLIB_MODEL, model_name, model_class) + + @staticmethod + @PublicAPI + def register_custom_action_dist( + action_dist_name: str, action_dist_class: type + ) -> None: + """Register a custom action distribution class by name. + + The model can be later used by specifying + {"custom_action_dist": action_dist_name} in the model config. + + Args: + model_name: Name to register the action distribution under. + model_class: Python class of the action distribution. + """ + _global_registry.register( + RLLIB_ACTION_DIST, action_dist_name, action_dist_class + ) + + @staticmethod + def _wrap_if_needed(model_cls: type, model_interface: type) -> type: + if not model_interface or issubclass(model_cls, model_interface): + return model_cls + + assert issubclass(model_cls, ModelV2), model_cls + + class wrapper(model_interface, model_cls): + pass + + name = "{}_as_{}".format(model_cls.__name__, model_interface.__name__) + wrapper.__name__ = name + wrapper.__qualname__ = name + + return wrapper + + @staticmethod + def _get_v2_model_class( + input_space: gym.Space, model_config: ModelConfigDict, framework: str = "tf" + ) -> Type[ModelV2]: + VisionNet = None + ComplexNet = None + + if framework in ["tf2", "tf"]: + from ray.rllib.models.tf.complex_input_net import ( + ComplexInputNetwork as ComplexNet, + ) + from ray.rllib.models.tf.fcnet import ( + FullyConnectedNetwork as FCNet, + ) + from ray.rllib.models.tf.visionnet import ( + VisionNetwork as VisionNet, + ) + elif framework == "torch": + from ray.rllib.models.torch.complex_input_net import ( + ComplexInputNetwork as ComplexNet, + ) + from ray.rllib.models.torch.fcnet import FullyConnectedNetwork as FCNet + from ray.rllib.models.torch.visionnet import VisionNetwork as VisionNet + elif framework == "jax": + from ray.rllib.models.jax.fcnet import FullyConnectedNetwork as FCNet + else: + raise ValueError( + "framework={} not supported in `ModelCatalog._get_v2_model_" + "class`!".format(framework) + ) + + orig_space = ( + input_space + if not hasattr(input_space, "original_space") + else input_space.original_space + ) + + # `input_space` is 3D Box -> VisionNet. + if isinstance(input_space, Box) and len(input_space.shape) == 3: + if framework == "jax": + raise NotImplementedError("No non-FC default net for JAX yet!") + return VisionNet + # `input_space` is 1D Box -> FCNet. + elif ( + isinstance(input_space, Box) + and len(input_space.shape) == 1 + and ( + not isinstance(orig_space, (Dict, Tuple)) + or not any( + isinstance(s, Box) and len(s.shape) >= 2 + for s in flatten_space(orig_space) + ) + ) + ): + return FCNet + # Complex (Dict, Tuple, 2D Box (flatten), Discrete, MultiDiscrete). + else: + if framework == "jax": + raise NotImplementedError("No non-FC default net for JAX yet!") + return ComplexNet + + @staticmethod + def _get_multi_action_distribution(dist_class, action_space, config, framework): + # In case the custom distribution is a child of MultiActionDistr. + # If users want to completely ignore the suggested child + # distributions, they should simply do so in their custom class' + # constructor. + if issubclass( + dist_class, (MultiActionDistribution, TorchMultiActionDistribution) + ): + flat_action_space = flatten_space(action_space) + child_dists_and_in_lens = tree.map_structure( + lambda s: ModelCatalog.get_action_dist(s, config, framework=framework), + flat_action_space, + ) + child_dists = [e[0] for e in child_dists_and_in_lens] + input_lens = [int(e[1]) for e in child_dists_and_in_lens] + return ( + partial( + dist_class, + action_space=action_space, + child_distributions=child_dists, + input_lens=input_lens, + ), + int(sum(input_lens)), + ) + return dist_class, dist_class.required_model_output_shape(action_space, config) + + @staticmethod + def _validate_config( + config: ModelConfigDict, action_space: gym.spaces.Space, framework: str + ) -> None: + """Validates a given model config dict. + + Args: + config: The "model" sub-config dict + within the Algorithm's config dict. + action_space: The action space of the model, whose config are + validated. + framework: One of "jax", "tf2", "tf", or "torch". + + Raises: + ValueError: If something is wrong with the given config. + """ + # Soft-deprecate custom preprocessors. + if config.get("custom_preprocessor") is not None: + deprecation_warning( + old="model.custom_preprocessor", + new="gym.ObservationWrapper around your env or handle complex " + "inputs inside your Model", + error=True, + ) + + if config.get("use_attention") and config.get("use_lstm"): + raise ValueError( + "Only one of `use_lstm` or `use_attention` may be set to True!" + ) + + # For complex action spaces, only allow prev action inputs to + # LSTMs and attention nets iff `_disable_action_flattening=True`. + # TODO: `_disable_action_flattening=True` will be the default in + # the future. + if ( + ( + config.get("lstm_use_prev_action") + or config.get("attention_use_n_prev_actions", 0) > 0 + ) + and not config.get("_disable_action_flattening") + and isinstance(action_space, (Tuple, Dict)) + ): + raise ValueError( + "For your complex action space (Tuple|Dict) and your model's " + "`prev-actions` setup of your model, you must set " + "`_disable_action_flattening=True` in your main config dict!" + ) + + if framework == "jax": + if config.get("use_attention"): + raise ValueError( + "`use_attention` not available for framework=jax so far!" + ) + elif config.get("use_lstm"): + raise ValueError("`use_lstm` not available for framework=jax so far!") diff --git a/lib/python3.12/site-packages/ray/rllib/models/distributions.py b/lib/python3.12/site-packages/ray/rllib/models/distributions.py new file mode 100644 index 0000000000000000000000000000000000000000..05d9670a8c7f6d7bfe35584a4988fa137dade563 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/distributions.py @@ -0,0 +1,8 @@ +from ray._common.deprecation import deprecation_warning +from ray.rllib.core.distribution.distribution import Distribution # noqa + +deprecation_warning( + old="ray.rllib.models.distributions.Distribution", + new="ray.rllib.core.distribution.distribution.Distribution", + error=False, +) diff --git a/lib/python3.12/site-packages/ray/rllib/models/modelv2.py b/lib/python3.12/site-packages/ray/rllib/models/modelv2.py new file mode 100644 index 0000000000000000000000000000000000000000..563cbffb56f9035db5f6f20244a418f32ebe74e2 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/modelv2.py @@ -0,0 +1,472 @@ +import contextlib +from collections import OrderedDict +from typing import Any, Dict, List, Union + +import gymnasium as gym +import numpy as np +from gymnasium.spaces import Space + +from ray._common.deprecation import Deprecated +from ray.rllib.models.preprocessors import RepeatedValuesPreprocessor, get_preprocessor +from ray.rllib.models.repeated_values import RepeatedValues +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.policy.view_requirement import ViewRequirement +from ray.rllib.utils import NullContextManager +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.framework import TensorType, try_import_tf, try_import_torch +from ray.rllib.utils.spaces.repeated import Repeated +from ray.rllib.utils.typing import ModelConfigDict, ModelInputDict, TensorStructType + +tf1, tf, tfv = try_import_tf() +torch, _ = try_import_torch() + + +@OldAPIStack +class ModelV2: + r"""Defines an abstract neural network model for use with RLlib. + + Custom models should extend either TFModelV2 or TorchModelV2 instead of + this class directly. + + Data flow: + obs -> forward() -> model_out + \-> value_function() -> V(s) + """ + + def __init__( + self, + obs_space: Space, + action_space: Space, + num_outputs: int, + model_config: ModelConfigDict, + name: str, + framework: str, + ): + """Initializes a ModelV2 instance. + + This method should create any variables used by the model. + + Args: + obs_space: Observation space of the target gym + env. This may have an `original_space` attribute that + specifies how to unflatten the tensor into a ragged tensor. + action_space: Action space of the target gym + env. + num_outputs: Number of output units of the model. + model_config: Config for the model, documented + in ModelCatalog. + name: Name (scope) for the model. + framework: Either "tf" or "torch". + """ + + self.obs_space: Space = obs_space + self.action_space: Space = action_space + self.num_outputs: int = num_outputs + self.model_config: ModelConfigDict = model_config + self.name: str = name or "default_model" + self.framework: str = framework + self._last_output = None + self.time_major = self.model_config.get("_time_major") + # Basic view requirement for all models: Use the observation as input. + self.view_requirements = { + SampleBatch.OBS: ViewRequirement(shift=0, space=self.obs_space), + } + + def get_initial_state(self) -> List[TensorType]: + """Get the initial recurrent state values for the model. + + Returns: + List of np.array (for tf) or Tensor (for torch) objects containing the + initial hidden state of an RNN, if applicable. + + .. testcode:: + :skipif: True + + import numpy as np + from ray.rllib.models.modelv2 import ModelV2 + class MyModel(ModelV2): + # ... + def get_initial_state(self): + return [ + np.zeros(self.cell_size, np.float32), + np.zeros(self.cell_size, np.float32), + ] + """ + return [] + + def forward( + self, + input_dict: Dict[str, TensorType], + state: List[TensorType], + seq_lens: TensorType, + ) -> (TensorType, List[TensorType]): + """Call the model with the given input tensors and state. + + Any complex observations (dicts, tuples, etc.) will be unpacked by + __call__ before being passed to forward(). To access the flattened + observation tensor, refer to input_dict["obs_flat"]. + + This method can be called any number of times. In eager execution, + each call to forward() will eagerly evaluate the model. In symbolic + execution, each call to forward creates a computation graph that + operates over the variables of this model (i.e., shares weights). + + Custom models should override this instead of __call__. + + Args: + input_dict: dictionary of input tensors, including "obs", + "obs_flat", "prev_action", "prev_reward", "is_training", + "eps_id", "agent_id", "infos", and "t". + state: list of state tensors with sizes matching those + returned by get_initial_state + the batch dimension + seq_lens: 1d tensor holding input sequence lengths + + Returns: + A tuple consisting of the model output tensor of size + [BATCH, num_outputs] and the list of new RNN state(s) if any. + + .. testcode:: + :skipif: True + + import numpy as np + from ray.rllib.models.modelv2 import ModelV2 + class MyModel(ModelV2): + # ... + def forward(self, input_dict, state, seq_lens): + model_out, self._value_out = self.base_model( + input_dict["obs"]) + return model_out, state + """ + raise NotImplementedError + + def value_function(self) -> TensorType: + """Returns the value function output for the most recent forward pass. + + Note that a `forward` call has to be performed first, before this + methods can return anything and thus that calling this method does not + cause an extra forward pass through the network. + + Returns: + Value estimate tensor of shape [BATCH]. + """ + raise NotImplementedError + + def custom_loss( + self, policy_loss: TensorType, loss_inputs: Dict[str, TensorType] + ) -> Union[List[TensorType], TensorType]: + """Override to customize the loss function used to optimize this model. + + This can be used to incorporate self-supervised losses (by defining + a loss over existing input and output tensors of this model), and + supervised losses (by defining losses over a variable-sharing copy of + this model's layers). + + You can find an runnable example in examples/custom_loss.py. + + Args: + policy_loss: List of or single policy loss(es) from the policy. + loss_inputs: map of input placeholders for rollout data. + + Returns: + List of or scalar tensor for the customized loss(es) for this + model. + """ + return policy_loss + + def metrics(self) -> Dict[str, TensorType]: + """Override to return custom metrics from your model. + + The stats will be reported as part of the learner stats, i.e., + info.learner.[policy_id, e.g. "default_policy"].model.key1=metric1 + + Returns: + The custom metrics for this model. + """ + return {} + + def __call__( + self, + input_dict: Union[SampleBatch, ModelInputDict], + state: List[Any] = None, + seq_lens: TensorType = None, + ) -> (TensorType, List[TensorType]): + """Call the model with the given input tensors and state. + + This is the method used by RLlib to execute the forward pass. It calls + forward() internally after unpacking nested observation tensors. + + Custom models should override forward() instead of __call__. + + Args: + input_dict: Dictionary of input tensors. + state: list of state tensors with sizes matching those + returned by get_initial_state + the batch dimension + seq_lens: 1D tensor holding input sequence lengths. + + Returns: + A tuple consisting of the model output tensor of size + [BATCH, output_spec.size] or a list of tensors corresponding to + output_spec.shape_list, and a list of state tensors of + [BATCH, state_size_i] if any. + """ + + # Original observations will be stored in "obs". + # Flattened (preprocessed) obs will be stored in "obs_flat". + + # SampleBatch case: Models can now be called directly with a + # SampleBatch (which also includes tracking-dict case (deprecated now), + # where tensors get automatically converted). + if isinstance(input_dict, SampleBatch): + restored = input_dict.copy(shallow=True) + else: + restored = input_dict.copy() + + # Backward compatibility. + if not state: + state = [] + i = 0 + while "state_in_{}".format(i) in input_dict: + state.append(input_dict["state_in_{}".format(i)]) + i += 1 + if seq_lens is None: + seq_lens = input_dict.get(SampleBatch.SEQ_LENS) + + # No Preprocessor used: `config._disable_preprocessor_api`=True. + # TODO: This is unnecessary for when no preprocessor is used. + # Obs are not flat then anymore. However, we'll keep this + # here for backward-compatibility until Preprocessors have + # been fully deprecated. + if self.model_config.get("_disable_preprocessor_api"): + restored["obs_flat"] = input_dict["obs"] + # Input to this Model went through a Preprocessor. + # Generate extra keys: "obs_flat" (vs "obs", which will hold the + # original obs). + else: + restored["obs"] = restore_original_dimensions( + input_dict["obs"], self.obs_space, self.framework + ) + try: + if len(input_dict["obs"].shape) > 2: + restored["obs_flat"] = flatten(input_dict["obs"], self.framework) + else: + restored["obs_flat"] = input_dict["obs"] + except AttributeError: + restored["obs_flat"] = input_dict["obs"] + + with self.context(): + res = self.forward(restored, state or [], seq_lens) + + if isinstance(input_dict, SampleBatch): + input_dict.accessed_keys = restored.accessed_keys - {"obs_flat"} + input_dict.deleted_keys = restored.deleted_keys + input_dict.added_keys = restored.added_keys - {"obs_flat"} + + if (not isinstance(res, list) and not isinstance(res, tuple)) or len(res) != 2: + raise ValueError( + "forward() must return a tuple of (output, state) tensors, " + "got {}".format(res) + ) + outputs, state_out = res + + if not isinstance(state_out, list): + raise ValueError("State output is not a list: {}".format(state_out)) + + self._last_output = outputs + return outputs, state_out if len(state_out) > 0 else (state or []) + + def last_output(self) -> TensorType: + """Returns the last output returned from calling the model.""" + return self._last_output + + def context(self) -> contextlib.AbstractContextManager: + """Returns a contextmanager for the current forward pass.""" + return NullContextManager() + + def variables( + self, as_dict: bool = False + ) -> Union[List[TensorType], Dict[str, TensorType]]: + """Returns the list (or a dict) of variables for this model. + + Args: + as_dict: Whether variables should be returned as dict-values + (using descriptive str keys). + + Returns: + The list (or dict if `as_dict` is True) of all variables of this + ModelV2. + """ + raise NotImplementedError + + def trainable_variables( + self, as_dict: bool = False + ) -> Union[List[TensorType], Dict[str, TensorType]]: + """Returns the list of trainable variables for this model. + + Args: + as_dict: Whether variables should be returned as dict-values + (using descriptive keys). + + Returns: + The list (or dict if `as_dict` is True) of all trainable + (tf)/requires_grad (torch) variables of this ModelV2. + """ + raise NotImplementedError + + def is_time_major(self) -> bool: + """If True, data for calling this ModelV2 must be in time-major format. + + Returns + Whether this ModelV2 requires a time-major (TxBx...) data + format. + """ + return self.time_major is True + + @Deprecated(error=True) + def import_from_h5(self, *args, **kwargs): + pass + + +@OldAPIStack +def flatten(obs: TensorType, framework: str) -> TensorType: + """Flatten the given tensor.""" + if framework in ["tf2", "tf"]: + return tf1.keras.layers.Flatten()(obs) + elif framework == "torch": + assert torch is not None + return torch.flatten(obs, start_dim=1) + else: + raise NotImplementedError("flatten", framework) + + +@OldAPIStack +def restore_original_dimensions( + obs: TensorType, obs_space: Space, tensorlib: Any = tf +) -> TensorStructType: + """Unpacks Dict and Tuple space observations into their original form. + + This is needed since we flatten Dict and Tuple observations in transit + within a SampleBatch. Before sending them to the model though, we should + unflatten them into Dicts or Tuples of tensors. + + Args: + obs: The flattened observation tensor. + obs_space: The flattened obs space. If this has the + `original_space` attribute, we will unflatten the tensor to that + shape. + tensorlib: The library used to unflatten (reshape) the array/tensor. + + Returns: + single tensor or dict / tuple of tensors matching the original + observation space. + """ + + if tensorlib in ["tf", "tf2"]: + assert tf is not None + tensorlib = tf + elif tensorlib == "torch": + assert torch is not None + tensorlib = torch + elif tensorlib == "numpy": + assert np is not None + tensorlib = np + original_space = getattr(obs_space, "original_space", obs_space) + return _unpack_obs(obs, original_space, tensorlib=tensorlib) + + +# Cache of preprocessors, for if the user is calling unpack obs often. +_cache = {} + + +@OldAPIStack +def _unpack_obs(obs: TensorType, space: Space, tensorlib: Any = tf) -> TensorStructType: + """Unpack a flattened Dict or Tuple observation array/tensor. + + Args: + obs: The flattened observation tensor, with last dimension equal to + the flat size and any number of batch dimensions. For example, for + Box(4,), the obs may have shape [B, 4], or [B, N, M, 4] in case + the Box was nested under two Repeated spaces. + space: The original space prior to flattening + tensorlib: The library used to unflatten (reshape) the array/tensor + """ + + if isinstance(space, (gym.spaces.Dict, gym.spaces.Tuple, Repeated)): + # Already unpacked? + if (isinstance(space, gym.spaces.Tuple) and isinstance(obs, (list, tuple))) or ( + isinstance(space, gym.spaces.Dict) and isinstance(obs, dict) + ): + return obs + # Unpack using preprocessor + if id(space) in _cache: + prep = _cache[id(space)] + else: + prep = get_preprocessor(space)(space) + # Make an attempt to cache the result, if enough space left. + if len(_cache) < 999: + _cache[id(space)] = prep + if len(obs.shape) < 2 or obs.shape[-1] != prep.shape[0]: + raise ValueError( + "Expected flattened obs shape of [..., {}], got {}".format( + prep.shape[0], obs.shape + ) + ) + offset = 0 + if tensorlib == tf: + + def get_value(v): + if v is None: + return -1 + elif isinstance(v, int): + return v + elif v.value is None: + return -1 + else: + return v.value + + batch_dims = [get_value(v) for v in obs.shape[:-1]] + else: + batch_dims = list(obs.shape[:-1]) + if isinstance(space, gym.spaces.Tuple): + assert len(prep.preprocessors) == len(space.spaces), len( + prep.preprocessors + ) == len(space.spaces) + u = [] + for p, v in zip(prep.preprocessors, space.spaces): + obs_slice = obs[..., offset : offset + p.size] + offset += p.size + u.append( + _unpack_obs( + tensorlib.reshape(obs_slice, batch_dims + list(p.shape)), + v, + tensorlib=tensorlib, + ) + ) + elif isinstance(space, gym.spaces.Dict): + assert len(prep.preprocessors) == len(space.spaces), len( + prep.preprocessors + ) == len(space.spaces) + u = OrderedDict() + for p, (k, v) in zip(prep.preprocessors, space.spaces.items()): + obs_slice = obs[..., offset : offset + p.size] + offset += p.size + u[k] = _unpack_obs( + tensorlib.reshape(obs_slice, batch_dims + list(p.shape)), + v, + tensorlib=tensorlib, + ) + # Repeated space. + else: + assert isinstance(prep, RepeatedValuesPreprocessor), prep + child_size = prep.child_preprocessor.size + # The list lengths are stored in the first slot of the flat obs. + lengths = obs[..., 0] + # [B, ..., 1 + max_len * child_sz] -> [B, ..., max_len, child_sz] + with_repeat_dim = tensorlib.reshape( + obs[..., 1:], batch_dims + [space.max_len, child_size] + ) + # Retry the unpack, dropping the List container space. + u = _unpack_obs(with_repeat_dim, space.child_space, tensorlib=tensorlib) + return RepeatedValues(u, lengths=lengths, max_len=prep._obs_space.max_len) + return u + else: + return obs diff --git a/lib/python3.12/site-packages/ray/rllib/models/preprocessors.py b/lib/python3.12/site-packages/ray/rllib/models/preprocessors.py new file mode 100644 index 0000000000000000000000000000000000000000..0531c8f8932edb81a8f44f4b445e00004bd2a208 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/preprocessors.py @@ -0,0 +1,448 @@ +import logging +from collections import OrderedDict +from typing import Any, List + +import gymnasium as gym +import numpy as np + +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.images import resize +from ray.rllib.utils.spaces.repeated import Repeated +from ray.rllib.utils.spaces.space_utils import convert_element_to_space_type +from ray.rllib.utils.typing import TensorType + +ATARI_OBS_SHAPE = (210, 160, 3) +ATARI_RAM_OBS_SHAPE = (128,) + +# Only validate env observations vs the observation space every n times in a +# Preprocessor. +OBS_VALIDATION_INTERVAL = 100 + +logger = logging.getLogger(__name__) + + +@OldAPIStack +class Preprocessor: + """Defines an abstract observation preprocessor function. + + Attributes: + shape (List[int]): Shape of the preprocessed output. + """ + + def __init__(self, obs_space: gym.Space, options: dict = None): + _legacy_patch_shapes(obs_space) + self._obs_space = obs_space + if not options: + from ray.rllib.models.catalog import MODEL_DEFAULTS + + self._options = MODEL_DEFAULTS.copy() + else: + self._options = options + self.shape = self._init_shape(obs_space, self._options) + self._size = int(np.prod(self.shape)) + self._i = 0 + self._obs_for_type_matching = self._obs_space.sample() + + def _init_shape(self, obs_space: gym.Space, options: dict) -> List[int]: + """Returns the shape after preprocessing.""" + raise NotImplementedError + + def transform(self, observation: TensorType) -> np.ndarray: + """Returns the preprocessed observation.""" + raise NotImplementedError + + def write(self, observation: TensorType, array: np.ndarray, offset: int) -> None: + """Alternative to transform for more efficient flattening.""" + array[offset : offset + self._size] = self.transform(observation) + + def check_shape(self, observation: Any) -> None: + """Checks the shape of the given observation.""" + if self._i % OBS_VALIDATION_INTERVAL == 0: + # Convert lists to np.ndarrays. + if type(observation) is list and isinstance( + self._obs_space, gym.spaces.Box + ): + observation = np.array(observation).astype(np.float32) + if not self._obs_space.contains(observation): + observation = convert_element_to_space_type( + observation, self._obs_for_type_matching + ) + try: + if not self._obs_space.contains(observation): + raise ValueError( + "Observation ({} dtype={}) outside given space ({})!".format( + observation, + observation.dtype + if isinstance(self._obs_space, gym.spaces.Box) + else None, + self._obs_space, + ) + ) + except AttributeError as e: + raise ValueError( + "Observation for a Box/MultiBinary/MultiDiscrete space " + "should be an np.array, not a Python list.", + observation, + ) from e + self._i += 1 + + @property + def size(self) -> int: + return self._size + + @property + def observation_space(self) -> gym.Space: + obs_space = gym.spaces.Box(-1.0, 1.0, self.shape, dtype=np.float32) + # Stash the unwrapped space so that we can unwrap dict and tuple spaces + # automatically in modelv2.py + classes = ( + DictFlatteningPreprocessor, + OneHotPreprocessor, + RepeatedValuesPreprocessor, + TupleFlatteningPreprocessor, + AtariRamPreprocessor, + GenericPixelPreprocessor, + ) + if isinstance(self, classes): + obs_space.original_space = self._obs_space + return obs_space + + +@OldAPIStack +class GenericPixelPreprocessor(Preprocessor): + """Generic image preprocessor. + + Note: for Atari games, use config {"preprocessor_pref": "deepmind"} + instead for deepmind-style Atari preprocessing. + """ + + @override(Preprocessor) + def _init_shape(self, obs_space: gym.Space, options: dict) -> List[int]: + self._grayscale = options.get("grayscale") + self._zero_mean = options.get("zero_mean") + self._dim = options.get("dim") + if self._grayscale: + shape = (self._dim, self._dim, 1) + else: + shape = (self._dim, self._dim, 3) + + return shape + + @override(Preprocessor) + def transform(self, observation: TensorType) -> np.ndarray: + """Downsamples images from (210, 160, 3) by the configured factor.""" + self.check_shape(observation) + scaled = observation[25:-25, :, :] + if self._dim < 84: + scaled = resize(scaled, height=84, width=84) + # OpenAI: Resize by half, then down to 42x42 (essentially mipmapping). + # If we resize directly we lose pixels that, when mapped to 42x42, + # aren't close enough to the pixel boundary. + scaled = resize(scaled, height=self._dim, width=self._dim) + if self._grayscale: + scaled = scaled.mean(2) + scaled = scaled.astype(np.float32) + # Rescale needed for maintaining 1 channel + scaled = np.reshape(scaled, [self._dim, self._dim, 1]) + if self._zero_mean: + scaled = (scaled - 128) / 128 + else: + scaled *= 1.0 / 255.0 + return scaled + + +@OldAPIStack +class AtariRamPreprocessor(Preprocessor): + @override(Preprocessor) + def _init_shape(self, obs_space: gym.Space, options: dict) -> List[int]: + return (128,) + + @override(Preprocessor) + def transform(self, observation: TensorType) -> np.ndarray: + self.check_shape(observation) + return (observation.astype("float32") - 128) / 128 + + +@OldAPIStack +class OneHotPreprocessor(Preprocessor): + """One-hot preprocessor for Discrete and MultiDiscrete spaces. + + .. testcode:: + :skipif: True + + self.transform(Discrete(3).sample()) + + .. testoutput:: + + np.array([0.0, 1.0, 0.0]) + + .. testcode:: + :skipif: True + + self.transform(MultiDiscrete([2, 3]).sample()) + + .. testoutput:: + + np.array([0.0, 1.0, 0.0, 0.0, 1.0]) + """ + + @override(Preprocessor) + def _init_shape(self, obs_space: gym.Space, options: dict) -> List[int]: + if isinstance(obs_space, gym.spaces.Discrete): + return (self._obs_space.n,) + else: + return (np.sum(self._obs_space.nvec),) + + @override(Preprocessor) + def transform(self, observation: TensorType) -> np.ndarray: + self.check_shape(observation) + return gym.spaces.utils.flatten(self._obs_space, observation).astype(np.float32) + + @override(Preprocessor) + def write(self, observation: TensorType, array: np.ndarray, offset: int) -> None: + array[offset : offset + self.size] = self.transform(observation) + + +@OldAPIStack +class NoPreprocessor(Preprocessor): + @override(Preprocessor) + def _init_shape(self, obs_space: gym.Space, options: dict) -> List[int]: + return self._obs_space.shape + + @override(Preprocessor) + def transform(self, observation: TensorType) -> np.ndarray: + self.check_shape(observation) + return observation + + @override(Preprocessor) + def write(self, observation: TensorType, array: np.ndarray, offset: int) -> None: + array[offset : offset + self._size] = np.array(observation, copy=False).ravel() + + @property + @override(Preprocessor) + def observation_space(self) -> gym.Space: + return self._obs_space + + +@OldAPIStack +class MultiBinaryPreprocessor(Preprocessor): + """Preprocessor that turns a MultiBinary space into a Box. + + Note: Before RLModules were introduced, RLlib's ModelCatalogV2 would produce + ComplexInputNetworks that treat MultiBinary spaces as Boxes. This preprocessor is + needed to get rid of the ComplexInputNetworks and use RLModules instead because + RLModules lack the logic to handle MultiBinary or other non-Box spaces. + """ + + @override(Preprocessor) + def _init_shape(self, obs_space: gym.Space, options: dict) -> List[int]: + return self._obs_space.shape + + @override(Preprocessor) + def transform(self, observation: TensorType) -> np.ndarray: + # The shape stays the same, but the dtype changes. + self.check_shape(observation) + return observation.astype(np.float32) + + @override(Preprocessor) + def write(self, observation: TensorType, array: np.ndarray, offset: int) -> None: + array[offset : offset + self._size] = np.array(observation, copy=False).ravel() + + @property + @override(Preprocessor) + def observation_space(self) -> gym.Space: + obs_space = gym.spaces.Box(0.0, 1.0, self.shape, dtype=np.float32) + obs_space.original_space = self._obs_space + return obs_space + + +@OldAPIStack +class TupleFlatteningPreprocessor(Preprocessor): + """Preprocesses each tuple element, then flattens it all into a vector. + + RLlib models will unpack the flattened output before _build_layers_v2(). + """ + + @override(Preprocessor) + def _init_shape(self, obs_space: gym.Space, options: dict) -> List[int]: + assert isinstance(self._obs_space, gym.spaces.Tuple) + size = 0 + self.preprocessors = [] + for i in range(len(self._obs_space.spaces)): + space = self._obs_space.spaces[i] + logger.debug("Creating sub-preprocessor for {}".format(space)) + preprocessor_class = get_preprocessor(space) + if preprocessor_class is not None: + preprocessor = preprocessor_class(space, self._options) + size += preprocessor.size + else: + preprocessor = None + size += int(np.prod(space.shape)) + self.preprocessors.append(preprocessor) + return (size,) + + @override(Preprocessor) + def transform(self, observation: TensorType) -> np.ndarray: + self.check_shape(observation) + array = np.zeros(self.shape, dtype=np.float32) + self.write(observation, array, 0) + return array + + @override(Preprocessor) + def write(self, observation: TensorType, array: np.ndarray, offset: int) -> None: + assert len(observation) == len(self.preprocessors), observation + for o, p in zip(observation, self.preprocessors): + p.write(o, array, offset) + offset += p.size + + +@OldAPIStack +class DictFlatteningPreprocessor(Preprocessor): + """Preprocesses each dict value, then flattens it all into a vector. + + RLlib models will unpack the flattened output before _build_layers_v2(). + """ + + @override(Preprocessor) + def _init_shape(self, obs_space: gym.Space, options: dict) -> List[int]: + assert isinstance(self._obs_space, gym.spaces.Dict) + size = 0 + self.preprocessors = [] + for space in self._obs_space.spaces.values(): + logger.debug("Creating sub-preprocessor for {}".format(space)) + preprocessor_class = get_preprocessor(space) + if preprocessor_class is not None: + preprocessor = preprocessor_class(space, self._options) + size += preprocessor.size + else: + preprocessor = None + size += int(np.prod(space.shape)) + self.preprocessors.append(preprocessor) + return (size,) + + @override(Preprocessor) + def transform(self, observation: TensorType) -> np.ndarray: + self.check_shape(observation) + array = np.zeros(self.shape, dtype=np.float32) + self.write(observation, array, 0) + return array + + @override(Preprocessor) + def write(self, observation: TensorType, array: np.ndarray, offset: int) -> None: + if not isinstance(observation, OrderedDict): + observation = OrderedDict(sorted(observation.items())) + assert len(observation) == len(self.preprocessors), ( + len(observation), + len(self.preprocessors), + ) + for o, p in zip(observation.values(), self.preprocessors): + p.write(o, array, offset) + offset += p.size + + +@OldAPIStack +class RepeatedValuesPreprocessor(Preprocessor): + """Pads and batches the variable-length list value.""" + + @override(Preprocessor) + def _init_shape(self, obs_space: gym.Space, options: dict) -> List[int]: + assert isinstance(self._obs_space, Repeated) + child_space = obs_space.child_space + self.child_preprocessor = get_preprocessor(child_space)( + child_space, self._options + ) + # The first slot encodes the list length. + size = 1 + self.child_preprocessor.size * obs_space.max_len + return (size,) + + @override(Preprocessor) + def transform(self, observation: TensorType) -> np.ndarray: + array = np.zeros(self.shape) + if isinstance(observation, list): + for elem in observation: + self.child_preprocessor.check_shape(elem) + else: + pass # ValueError will be raised in write() below. + self.write(observation, array, 0) + return array + + @override(Preprocessor) + def write(self, observation: TensorType, array: np.ndarray, offset: int) -> None: + if not isinstance(observation, (list, np.ndarray)): + raise ValueError( + "Input for {} must be list type, got {}".format(self, observation) + ) + elif len(observation) > self._obs_space.max_len: + raise ValueError( + "Input {} exceeds max len of space {}".format( + observation, self._obs_space.max_len + ) + ) + # The first slot encodes the list length. + array[offset] = len(observation) + for i, elem in enumerate(observation): + offset_i = offset + 1 + i * self.child_preprocessor.size + self.child_preprocessor.write(elem, array, offset_i) + + +@OldAPIStack +def get_preprocessor(space: gym.Space, include_multi_binary=False) -> type: + """Returns an appropriate preprocessor class for the given space.""" + + _legacy_patch_shapes(space) + obs_shape = space.shape + + if isinstance(space, (gym.spaces.Discrete, gym.spaces.MultiDiscrete)): + preprocessor = OneHotPreprocessor + elif obs_shape == ATARI_OBS_SHAPE: + logger.debug( + "Defaulting to RLlib's GenericPixelPreprocessor because input " + "space has the atari-typical shape {}. Turn this behaviour off by setting " + "`preprocessor_pref=None` or " + "`preprocessor_pref='deepmind'` or disabling the preprocessing API " + "altogether with `_disable_preprocessor_api=True`.".format(ATARI_OBS_SHAPE) + ) + preprocessor = GenericPixelPreprocessor + elif obs_shape == ATARI_RAM_OBS_SHAPE: + logger.debug( + "Defaulting to RLlib's AtariRamPreprocessor because input " + "space has the atari-typical shape {}. Turn this behaviour off by setting " + "`preprocessor_pref=None` or " + "`preprocessor_pref='deepmind' or disabling the preprocessing API " + "altogether with `_disable_preprocessor_api=True`." + "`.".format(ATARI_OBS_SHAPE) + ) + preprocessor = AtariRamPreprocessor + elif isinstance(space, gym.spaces.Tuple): + preprocessor = TupleFlatteningPreprocessor + elif isinstance(space, gym.spaces.Dict): + preprocessor = DictFlatteningPreprocessor + elif isinstance(space, Repeated): + preprocessor = RepeatedValuesPreprocessor + # We usually only want to include this when using RLModules + elif isinstance(space, gym.spaces.MultiBinary) and include_multi_binary: + preprocessor = MultiBinaryPreprocessor + else: + preprocessor = NoPreprocessor + + return preprocessor + + +def _legacy_patch_shapes(space: gym.Space) -> List[int]: + """Assigns shapes to spaces that don't have shapes. + + This is only needed for older gym versions that don't set shapes properly + for Tuple and Discrete spaces. + """ + + if not hasattr(space, "shape"): + if isinstance(space, gym.spaces.Discrete): + space.shape = () + elif isinstance(space, gym.spaces.Tuple): + shapes = [] + for s in space.spaces: + shape = _legacy_patch_shapes(s) + shapes.append(shape) + space.shape = tuple(shapes) + + return space.shape diff --git a/lib/python3.12/site-packages/ray/rllib/models/repeated_values.py b/lib/python3.12/site-packages/ray/rllib/models/repeated_values.py new file mode 100644 index 0000000000000000000000000000000000000000..305298987141046d530a68f911e80d0253579531 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/repeated_values.py @@ -0,0 +1,204 @@ +from typing import List + +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.typing import TensorStructType, TensorType + + +@OldAPIStack +class RepeatedValues: + """Represents a variable-length list of items from spaces.Repeated. + + RepeatedValues are created when you use spaces.Repeated, and are + accessible as part of input_dict["obs"] in ModelV2 forward functions. + + Example: + Suppose the gym space definition was: + Repeated(Repeated(Box(K), N), M) + + Then in the model forward function, input_dict["obs"] is of type: + RepeatedValues(RepeatedValues()) + + The tensor is accessible via: + input_dict["obs"].values.values + + And the actual data lengths via: + # outer repetition, shape [B], range [0, M] + input_dict["obs"].lengths + -and- + # inner repetition, shape [B, M], range [0, N] + input_dict["obs"].values.lengths + + Attributes: + values: The padded data tensor of shape [B, max_len, ..., sz], + where B is the batch dimension, max_len is the max length of this + list, followed by any number of sub list max lens, followed by the + actual data size. + lengths (List[int]): Tensor of shape [B, ...] that represents the + number of valid items in each list. When the list is nested within + other lists, there will be extra dimensions for the parent list + max lens. + max_len: The max number of items allowed in each list. + + TODO(ekl): support conversion to tf.RaggedTensor. + """ + + def __init__(self, values: TensorType, lengths: List[int], max_len: int): + self.values = values + self.lengths = lengths + self.max_len = max_len + self._unbatched_repr = None + + def unbatch_all(self) -> List[List[TensorType]]: + """Unbatch both the repeat and batch dimensions into Python lists. + + This is only supported in PyTorch / TF eager mode. + + This lets you view the data unbatched in its original form, but is + not efficient for processing. + + .. testcode:: + :skipif: True + + batch = RepeatedValues() + items = batch.unbatch_all() + print(len(items) == B) + + .. testoutput:: + + True + + .. testcode:: + :skipif: True + + print(max(len(x) for x in items) <= N) + + .. testoutput:: + + True + + .. testcode:: + :skipif: True + + print(items) + + .. testoutput:: + + [[, ..., ], + ... + [, ], + ... + [], + ... + [, ..., ]] + """ + + if self._unbatched_repr is None: + B = _get_batch_dim_helper(self.values) + if B is None: + raise ValueError( + "Cannot call unbatch_all() when batch_dim is unknown. " + "This is probably because you are using TF graph mode." + ) + else: + B = int(B) + slices = self.unbatch_repeat_dim() + result = [] + for i in range(B): + if hasattr(self.lengths[i], "item"): + dynamic_len = int(self.lengths[i].item()) + else: + dynamic_len = int(self.lengths[i].numpy()) + dynamic_slice = [] + for j in range(dynamic_len): + dynamic_slice.append(_batch_index_helper(slices, i, j)) + result.append(dynamic_slice) + self._unbatched_repr = result + + return self._unbatched_repr + + def unbatch_repeat_dim(self) -> List[TensorType]: + """Unbatches the repeat dimension (the one `max_len` in size). + + This removes the repeat dimension. The result will be a Python list of + with length `self.max_len`. Note that the data is still padded. + + .. testcode:: + :skipif: True + + batch = RepeatedValues() + items = batch.unbatch() + len(items) == batch.max_len + + .. testoutput:: + + True + + .. testcode:: + :skipif: True + + print(items) + + .. testoutput:: + + [, ..., ] + """ + return _unbatch_helper(self.values, self.max_len) + + def __repr__(self): + return "RepeatedValues(value={}, lengths={}, max_len={})".format( + repr(self.values), repr(self.lengths), self.max_len + ) + + def __str__(self): + return repr(self) + + +def _get_batch_dim_helper(v: TensorStructType) -> int: + """Tries to find the batch dimension size of v, or None.""" + if isinstance(v, dict): + for u in v.values(): + return _get_batch_dim_helper(u) + elif isinstance(v, tuple): + return _get_batch_dim_helper(v[0]) + elif isinstance(v, RepeatedValues): + return _get_batch_dim_helper(v.values) + else: + B = v.shape[0] + if hasattr(B, "value"): + B = B.value # TensorFlow + return B + + +def _unbatch_helper(v: TensorStructType, max_len: int) -> TensorStructType: + """Recursively unpacks the repeat dimension (max_len).""" + if isinstance(v, dict): + return {k: _unbatch_helper(u, max_len) for (k, u) in v.items()} + elif isinstance(v, tuple): + return tuple(_unbatch_helper(u, max_len) for u in v) + elif isinstance(v, RepeatedValues): + unbatched = _unbatch_helper(v.values, max_len) + return [ + RepeatedValues(u, v.lengths[:, i, ...], v.max_len) + for i, u in enumerate(unbatched) + ] + else: + return [v[:, i, ...] for i in range(max_len)] + + +def _batch_index_helper(v: TensorStructType, i: int, j: int) -> TensorStructType: + """Selects the item at the ith batch index and jth repetition.""" + if isinstance(v, dict): + return {k: _batch_index_helper(u, i, j) for (k, u) in v.items()} + elif isinstance(v, tuple): + return tuple(_batch_index_helper(u, i, j) for u in v) + elif isinstance(v, list): + # This is the output of unbatch_repeat_dim(). Unfortunately we have to + # process it here instead of in unbatch_all(), since it may be buried + # under a dict / tuple. + return _batch_index_helper(v[j], i, j) + elif isinstance(v, RepeatedValues): + unbatched = v.unbatch_all() + # Don't need to select j here; that's already done in unbatch_all. + return unbatched[i] + else: + return v[i, ...] diff --git a/lib/python3.12/site-packages/ray/rllib/models/tf/__init__.py b/lib/python3.12/site-packages/ray/rllib/models/tf/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..106ab9babdde65e8f65a5ccfe62bded88fe4b135 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/tf/__init__.py @@ -0,0 +1,11 @@ +from ray.rllib.models.tf.fcnet import FullyConnectedNetwork +from ray.rllib.models.tf.recurrent_net import RecurrentNetwork +from ray.rllib.models.tf.tf_modelv2 import TFModelV2 +from ray.rllib.models.tf.visionnet import VisionNetwork + +__all__ = [ + "FullyConnectedNetwork", + "RecurrentNetwork", + "TFModelV2", + "VisionNetwork", +] diff --git a/lib/python3.12/site-packages/ray/rllib/models/tf/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/tf/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1ab949d5314313b55b3ec8d7158475473b4a5df6 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/tf/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/tf/__pycache__/attention_net.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/tf/__pycache__/attention_net.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..468740cefde66f581908b90525fb8073c1ad71de Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/tf/__pycache__/attention_net.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/tf/__pycache__/complex_input_net.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/tf/__pycache__/complex_input_net.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..72177570b9aa66a72903e3966aabacb42d64a7b1 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/tf/__pycache__/complex_input_net.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/tf/__pycache__/fcnet.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/tf/__pycache__/fcnet.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2c9f0ad632cab10677bc584df80bd9a8d0127f4e Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/tf/__pycache__/fcnet.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/tf/__pycache__/misc.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/tf/__pycache__/misc.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..edaf47faa896fc721135b628880725e9340c5e99 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/tf/__pycache__/misc.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/tf/__pycache__/recurrent_net.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/tf/__pycache__/recurrent_net.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9d40511d3ee10c5febddb6bb49d3d6ae513076c2 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/tf/__pycache__/recurrent_net.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/tf/__pycache__/tf_action_dist.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/tf/__pycache__/tf_action_dist.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..31616601d42fdd7ef4d5b04b9d22aae1186809a3 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/tf/__pycache__/tf_action_dist.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/tf/__pycache__/tf_modelv2.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/tf/__pycache__/tf_modelv2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8fe59fa05f2ce157eb5c663fb21a0c81a8d0a6ed Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/tf/__pycache__/tf_modelv2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/tf/__pycache__/visionnet.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/tf/__pycache__/visionnet.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e75cf76250d571230d60c5a4eaa4afd370355569 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/tf/__pycache__/visionnet.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/tf/attention_net.py b/lib/python3.12/site-packages/ray/rllib/models/tf/attention_net.py new file mode 100644 index 0000000000000000000000000000000000000000..aacd7e58d47307edbb9e2330bde6f7abdb54e3bc --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/tf/attention_net.py @@ -0,0 +1,574 @@ +""" +[1] - Attention Is All You Need - Vaswani, Jones, Shazeer, Parmar, + Uszkoreit, Gomez, Kaiser - Google Brain/Research, U Toronto - 2017. + https://arxiv.org/pdf/1706.03762.pdf +[2] - Stabilizing Transformers for Reinforcement Learning - E. Parisotto + et al. - DeepMind - 2019. https://arxiv.org/pdf/1910.06764.pdf +[3] - Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context. + Z. Dai, Z. Yang, et al. - Carnegie Mellon U - 2019. + https://www.aclweb.org/anthology/P19-1285.pdf +""" +from typing import Any, Dict, Optional, Union + +import gymnasium as gym +import numpy as np +import tree # pip install dm_tree +from gymnasium.spaces import Box, Discrete, MultiDiscrete + +from ray._common.deprecation import deprecation_warning +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.models.tf.layers import ( + GRUGate, + RelativeMultiHeadAttention, + SkipConnection, +) +from ray.rllib.models.tf.recurrent_net import RecurrentNetwork +from ray.rllib.models.tf.tf_modelv2 import TFModelV2 +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.policy.view_requirement import ViewRequirement +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.spaces.space_utils import get_base_struct_from_space +from ray.rllib.utils.tf_utils import flatten_inputs_to_1d_tensor, one_hot +from ray.rllib.utils.typing import List, ModelConfigDict, TensorType +from ray.util import log_once + +tf1, tf, tfv = try_import_tf() + + +@OldAPIStack +class PositionwiseFeedforward(tf.keras.layers.Layer if tf else object): + """A 2x linear layer with ReLU activation in between described in [1]. + + Each timestep coming from the attention head will be passed through this + layer separately. + """ + + def __init__( + self, + out_dim: int, + hidden_dim: int, + output_activation: Optional[Any] = None, + **kwargs, + ): + super().__init__(**kwargs) + + self._hidden_layer = tf.keras.layers.Dense( + hidden_dim, + activation=tf.nn.relu, + ) + + self._output_layer = tf.keras.layers.Dense( + out_dim, activation=output_activation + ) + if log_once("positionwise_feedforward_tf"): + deprecation_warning( + old="rllib.models.tf.attention_net.PositionwiseFeedforward", + ) + + def call(self, inputs: TensorType, **kwargs) -> TensorType: + del kwargs + output = self._hidden_layer(inputs) + return self._output_layer(output) + + +@OldAPIStack +class TrXLNet(RecurrentNetwork): + """A TrXL net Model described in [1].""" + + def __init__( + self, + observation_space: gym.spaces.Space, + action_space: gym.spaces.Space, + num_outputs: int, + model_config: ModelConfigDict, + name: str, + num_transformer_units: int, + attention_dim: int, + num_heads: int, + head_dim: int, + position_wise_mlp_dim: int, + ): + """Initializes a TrXLNet object. + + Args: + num_transformer_units: The number of Transformer repeats to + use (denoted L in [2]). + attention_dim: The input and output dimensions of one + Transformer unit. + num_heads: The number of attention heads to use in parallel. + Denoted as `H` in [3]. + head_dim: The dimension of a single(!) attention head within + a multi-head attention unit. Denoted as `d` in [3]. + position_wise_mlp_dim: The dimension of the hidden layer + within the position-wise MLP (after the multi-head attention + block within one Transformer unit). This is the size of the + first of the two layers within the PositionwiseFeedforward. The + second layer always has size=`attention_dim`. + """ + if log_once("trxl_net_tf"): + deprecation_warning( + old="rllib.models.tf.attention_net.TrXLNet", + ) + super().__init__( + observation_space, action_space, num_outputs, model_config, name + ) + + self.num_transformer_units = num_transformer_units + self.attention_dim = attention_dim + self.num_heads = num_heads + self.head_dim = head_dim + self.max_seq_len = model_config["max_seq_len"] + self.obs_dim = observation_space.shape[0] + + inputs = tf.keras.layers.Input( + shape=(self.max_seq_len, self.obs_dim), name="inputs" + ) + E_out = tf.keras.layers.Dense(attention_dim)(inputs) + + for _ in range(self.num_transformer_units): + MHA_out = SkipConnection( + RelativeMultiHeadAttention( + out_dim=attention_dim, + num_heads=num_heads, + head_dim=head_dim, + input_layernorm=False, + output_activation=None, + ), + fan_in_layer=None, + )(E_out) + E_out = SkipConnection( + PositionwiseFeedforward(attention_dim, position_wise_mlp_dim) + )(MHA_out) + E_out = tf.keras.layers.LayerNormalization(axis=-1)(E_out) + + # Postprocess TrXL output with another hidden layer and compute values. + logits = tf.keras.layers.Dense( + self.num_outputs, activation=tf.keras.activations.linear, name="logits" + )(E_out) + + self.base_model = tf.keras.models.Model([inputs], [logits]) + + @override(RecurrentNetwork) + def forward_rnn( + self, inputs: TensorType, state: List[TensorType], seq_lens: TensorType + ) -> (TensorType, List[TensorType]): + # To make Attention work with current RLlib's ModelV2 API: + # We assume `state` is the history of L recent observations (all + # concatenated into one tensor) and append the current inputs to the + # end and only keep the most recent (up to `max_seq_len`). This allows + # us to deal with timestep-wise inference and full sequence training + # within the same logic. + observations = state[0] + observations = tf.concat((observations, inputs), axis=1)[:, -self.max_seq_len :] + logits = self.base_model([observations]) + T = tf.shape(inputs)[1] # Length of input segment (time). + logits = logits[:, -T:] + + return logits, [observations] + + @override(RecurrentNetwork) + def get_initial_state(self) -> List[np.ndarray]: + # State is the T last observations concat'd together into one Tensor. + # Plus all Transformer blocks' E(l) outputs concat'd together (up to + # tau timesteps). + return [np.zeros((self.max_seq_len, self.obs_dim), np.float32)] + + +class GTrXLNet(RecurrentNetwork): + """A GTrXL net Model described in [2]. + + This is still in an experimental phase. + Can be used as a drop-in replacement for LSTMs in PPO and IMPALA. + + To use this network as a replacement for an RNN, configure your Algorithm + as follows: + + Examples: + >> config["model"]["custom_model"] = GTrXLNet + >> config["model"]["max_seq_len"] = 10 + >> config["model"]["custom_model_config"] = { + >> num_transformer_units=1, + >> attention_dim=32, + >> num_heads=2, + >> memory_inference=100, + >> memory_training=50, + >> etc.. + >> } + """ + + def __init__( + self, + observation_space: gym.spaces.Space, + action_space: gym.spaces.Space, + num_outputs: Optional[int], + model_config: ModelConfigDict, + name: str, + *, + num_transformer_units: int = 1, + attention_dim: int = 64, + num_heads: int = 2, + memory_inference: int = 50, + memory_training: int = 50, + head_dim: int = 32, + position_wise_mlp_dim: int = 32, + init_gru_gate_bias: float = 2.0, + ): + """Initializes a GTrXLNet instance. + + Args: + num_transformer_units: The number of Transformer repeats to + use (denoted L in [2]). + attention_dim: The input and output dimensions of one + Transformer unit. + num_heads: The number of attention heads to use in parallel. + Denoted as `H` in [3]. + memory_inference: The number of timesteps to concat (time + axis) and feed into the next transformer unit as inference + input. The first transformer unit will receive this number of + past observations (plus the current one), instead. + memory_training: The number of timesteps to concat (time + axis) and feed into the next transformer unit as training + input (plus the actual input sequence of len=max_seq_len). + The first transformer unit will receive this number of + past observations (plus the input sequence), instead. + head_dim: The dimension of a single(!) attention head within + a multi-head attention unit. Denoted as `d` in [3]. + position_wise_mlp_dim: The dimension of the hidden layer + within the position-wise MLP (after the multi-head attention + block within one Transformer unit). This is the size of the + first of the two layers within the PositionwiseFeedforward. The + second layer always has size=`attention_dim`. + init_gru_gate_bias: Initial bias values for the GRU gates + (two GRUs per Transformer unit, one after the MHA, one after + the position-wise MLP). + """ + super().__init__( + observation_space, action_space, num_outputs, model_config, name + ) + + self.num_transformer_units = num_transformer_units + self.attention_dim = attention_dim + self.num_heads = num_heads + self.memory_inference = memory_inference + self.memory_training = memory_training + self.head_dim = head_dim + self.max_seq_len = model_config["max_seq_len"] + self.obs_dim = observation_space.shape[0] + + # Raw observation input (plus (None) time axis). + input_layer = tf.keras.layers.Input(shape=(None, self.obs_dim), name="inputs") + memory_ins = [ + tf.keras.layers.Input( + shape=(None, self.attention_dim), + dtype=tf.float32, + name="memory_in_{}".format(i), + ) + for i in range(self.num_transformer_units) + ] + + # Map observation dim to input/output transformer (attention) dim. + E_out = tf.keras.layers.Dense(self.attention_dim)(input_layer) + # Output, collected and concat'd to build the internal, tau-len + # Memory units used for additional contextual information. + memory_outs = [E_out] + + # 2) Create L Transformer blocks according to [2]. + for i in range(self.num_transformer_units): + # RelativeMultiHeadAttention part. + MHA_out = SkipConnection( + RelativeMultiHeadAttention( + out_dim=self.attention_dim, + num_heads=num_heads, + head_dim=head_dim, + input_layernorm=True, + output_activation=tf.nn.relu, + ), + fan_in_layer=GRUGate(init_gru_gate_bias), + name="mha_{}".format(i + 1), + )(E_out, memory=memory_ins[i]) + # Position-wise MLP part. + E_out = SkipConnection( + tf.keras.Sequential( + ( + tf.keras.layers.LayerNormalization(axis=-1), + PositionwiseFeedforward( + out_dim=self.attention_dim, + hidden_dim=position_wise_mlp_dim, + output_activation=tf.nn.relu, + ), + ) + ), + fan_in_layer=GRUGate(init_gru_gate_bias), + name="pos_wise_mlp_{}".format(i + 1), + )(MHA_out) + # Output of position-wise MLP == E(l-1), which is concat'd + # to the current Mem block (M(l-1)) to yield E~(l-1), which is then + # used by the next transformer block. + memory_outs.append(E_out) + + self._logits = None + self._value_out = None + + # Postprocess TrXL output with another hidden layer and compute values. + if num_outputs is not None: + self._logits = tf.keras.layers.Dense( + self.num_outputs, activation=None, name="logits" + )(E_out) + values_out = tf.keras.layers.Dense(1, activation=None, name="values")(E_out) + outs = [self._logits, values_out] + else: + outs = [E_out] + self.num_outputs = self.attention_dim + + self.trxl_model = tf.keras.Model( + inputs=[input_layer] + memory_ins, outputs=outs + memory_outs[:-1] + ) + + self.trxl_model.summary() + + # __sphinx_doc_begin__ + # Setup trajectory views (`memory-inference` x past memory outs). + for i in range(self.num_transformer_units): + space = Box(-1.0, 1.0, shape=(self.attention_dim,)) + self.view_requirements["state_in_{}".format(i)] = ViewRequirement( + "state_out_{}".format(i), + shift="-{}:-1".format(self.memory_inference), + # Repeat the incoming state every max-seq-len times. + batch_repeat_value=self.max_seq_len, + space=space, + ) + self.view_requirements["state_out_{}".format(i)] = ViewRequirement( + space=space, used_for_training=False + ) + # __sphinx_doc_end__ + + @override(ModelV2) + def forward( + self, input_dict, state: List[TensorType], seq_lens: TensorType + ) -> (TensorType, List[TensorType]): + assert seq_lens is not None + + # Add the time dim to observations. + B = tf.shape(seq_lens)[0] + observations = input_dict[SampleBatch.OBS] + + shape = tf.shape(observations) + T = shape[0] // B + observations = tf.reshape(observations, tf.concat([[-1, T], shape[1:]], axis=0)) + + all_out = self.trxl_model([observations] + state) + + if self._logits is not None: + out = tf.reshape(all_out[0], [-1, self.num_outputs]) + self._value_out = all_out[1] + memory_outs = all_out[2:] + else: + out = tf.reshape(all_out[0], [-1, self.attention_dim]) + memory_outs = all_out[1:] + + return out, [tf.reshape(m, [-1, self.attention_dim]) for m in memory_outs] + + @override(RecurrentNetwork) + def get_initial_state(self) -> List[np.ndarray]: + return [ + tf.zeros(self.view_requirements["state_in_{}".format(i)].space.shape) + for i in range(self.num_transformer_units) + ] + + @override(ModelV2) + def value_function(self) -> TensorType: + return tf.reshape(self._value_out, [-1]) + + +class AttentionWrapper(TFModelV2): + """GTrXL wrapper serving as interface for ModelV2s that set use_attention.""" + + def __init__( + self, + obs_space: gym.spaces.Space, + action_space: gym.spaces.Space, + num_outputs: int, + model_config: ModelConfigDict, + name: str, + ): + if log_once("attention_wrapper_tf_deprecation"): + deprecation_warning( + old="ray.rllib.models.tf.attention_net.AttentionWrapper" + ) + super().__init__(obs_space, action_space, None, model_config, name) + + self.use_n_prev_actions = model_config["attention_use_n_prev_actions"] + self.use_n_prev_rewards = model_config["attention_use_n_prev_rewards"] + + self.action_space_struct = get_base_struct_from_space(self.action_space) + self.action_dim = 0 + + for space in tree.flatten(self.action_space_struct): + if isinstance(space, Discrete): + self.action_dim += space.n + elif isinstance(space, MultiDiscrete): + self.action_dim += np.sum(space.nvec) + elif space.shape is not None: + self.action_dim += int(np.prod(space.shape)) + else: + self.action_dim += int(len(space)) + + # Add prev-action/reward nodes to input to LSTM. + if self.use_n_prev_actions: + self.num_outputs += self.use_n_prev_actions * self.action_dim + if self.use_n_prev_rewards: + self.num_outputs += self.use_n_prev_rewards + + cfg = model_config + + self.attention_dim = cfg["attention_dim"] + + if self.num_outputs is not None: + in_space = gym.spaces.Box( + float("-inf"), float("inf"), shape=(self.num_outputs,), dtype=np.float32 + ) + else: + in_space = obs_space + + # Construct GTrXL sub-module w/ num_outputs=None (so it does not + # create a logits/value output; we'll do this ourselves in this wrapper + # here). + self.gtrxl = GTrXLNet( + in_space, + action_space, + None, + model_config, + "gtrxl", + num_transformer_units=cfg["attention_num_transformer_units"], + attention_dim=self.attention_dim, + num_heads=cfg["attention_num_heads"], + head_dim=cfg["attention_head_dim"], + memory_inference=cfg["attention_memory_inference"], + memory_training=cfg["attention_memory_training"], + position_wise_mlp_dim=cfg["attention_position_wise_mlp_dim"], + init_gru_gate_bias=cfg["attention_init_gru_gate_bias"], + ) + + # `self.num_outputs` right now is the number of nodes coming from the + # attention net. + input_ = tf.keras.layers.Input(shape=(self.gtrxl.num_outputs,)) + + # Set final num_outputs to correct value (depending on action space). + self.num_outputs = num_outputs + + # Postprocess GTrXL output with another hidden layer and compute + # values. + out = tf.keras.layers.Dense(self.num_outputs, activation=None)(input_) + self._logits_branch = tf.keras.models.Model([input_], [out]) + + out = tf.keras.layers.Dense(1, activation=None)(input_) + self._value_branch = tf.keras.models.Model([input_], [out]) + + self.view_requirements = self.gtrxl.view_requirements + self.view_requirements["obs"].space = self.obs_space + + # Add prev-a/r to this model's view, if required. + if self.use_n_prev_actions: + self.view_requirements[SampleBatch.PREV_ACTIONS] = ViewRequirement( + SampleBatch.ACTIONS, + space=self.action_space, + shift="-{}:-1".format(self.use_n_prev_actions), + ) + if self.use_n_prev_rewards: + self.view_requirements[SampleBatch.PREV_REWARDS] = ViewRequirement( + SampleBatch.REWARDS, shift="-{}:-1".format(self.use_n_prev_rewards) + ) + + @override(RecurrentNetwork) + def forward( + self, + input_dict: Dict[str, TensorType], + state: List[TensorType], + seq_lens: TensorType, + ) -> (TensorType, List[TensorType]): + assert seq_lens is not None + # Push obs through "unwrapped" net's `forward()` first. + wrapped_out, _ = self._wrapped_forward(input_dict, [], None) + + # Concat. prev-action/reward if required. + prev_a_r = [] + + # Prev actions. + if self.use_n_prev_actions: + prev_n_actions = input_dict[SampleBatch.PREV_ACTIONS] + # If actions are not processed yet (in their original form as + # have been sent to environment): + # Flatten/one-hot into 1D array. + if self.model_config["_disable_action_flattening"]: + # Merge prev n actions into flat tensor. + flat = flatten_inputs_to_1d_tensor( + prev_n_actions, + spaces_struct=self.action_space_struct, + time_axis=True, + ) + # Fold time-axis into flattened data. + flat = tf.reshape(flat, [tf.shape(flat)[0], -1]) + prev_a_r.append(flat) + # If actions are already flattened (but not one-hot'd yet!), + # one-hot discrete/multi-discrete actions here and concatenate the + # n most recent actions together. + else: + if isinstance(self.action_space, Discrete): + for i in range(self.use_n_prev_actions): + prev_a_r.append( + one_hot(prev_n_actions[:, i], self.action_space) + ) + elif isinstance(self.action_space, MultiDiscrete): + for i in range( + 0, self.use_n_prev_actions, self.action_space.shape[0] + ): + prev_a_r.append( + one_hot( + tf.cast( + prev_n_actions[ + :, i : i + self.action_space.shape[0] + ], + tf.float32, + ), + space=self.action_space, + ) + ) + else: + prev_a_r.append( + tf.reshape( + tf.cast(prev_n_actions, tf.float32), + [-1, self.use_n_prev_actions * self.action_dim], + ) + ) + # Prev rewards. + if self.use_n_prev_rewards: + prev_a_r.append( + tf.reshape( + tf.cast(input_dict[SampleBatch.PREV_REWARDS], tf.float32), + [-1, self.use_n_prev_rewards], + ) + ) + + # Concat prev. actions + rewards to the "main" input. + if prev_a_r: + wrapped_out = tf.concat([wrapped_out] + prev_a_r, axis=1) + + # Then through our GTrXL. + input_dict["obs_flat"] = input_dict["obs"] = wrapped_out + + self._features, memory_outs = self.gtrxl(input_dict, state, seq_lens) + model_out = self._logits_branch(self._features) + return model_out, memory_outs + + @override(ModelV2) + def value_function(self) -> TensorType: + assert self._features is not None, "Must call forward() first!" + return tf.reshape(self._value_branch(self._features), [-1]) + + @override(ModelV2) + def get_initial_state(self) -> Union[List[np.ndarray], List[TensorType]]: + return [ + np.zeros(self.gtrxl.view_requirements["state_in_{}".format(i)].space.shape) + for i in range(self.gtrxl.num_transformer_units) + ] diff --git a/lib/python3.12/site-packages/ray/rllib/models/tf/complex_input_net.py b/lib/python3.12/site-packages/ray/rllib/models/tf/complex_input_net.py new file mode 100644 index 0000000000000000000000000000000000000000..04248ee97edd9549c6c9cc1fc0ef0ded6a401a2b --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/tf/complex_input_net.py @@ -0,0 +1,214 @@ +import numpy as np +import tree # pip install dm_tree +from gymnasium.spaces import Box, Discrete, MultiDiscrete + +from ray.rllib.models.catalog import ModelCatalog +from ray.rllib.models.modelv2 import ModelV2, restore_original_dimensions +from ray.rllib.models.tf.misc import normc_initializer +from ray.rllib.models.tf.tf_modelv2 import TFModelV2 +from ray.rllib.models.utils import get_filter_config +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.spaces.space_utils import flatten_space +from ray.rllib.utils.tf_utils import one_hot + +tf1, tf, tfv = try_import_tf() + + +# __sphinx_doc_begin__ +@OldAPIStack +class ComplexInputNetwork(TFModelV2): + """TFModelV2 concat'ing CNN outputs to flat input(s), followed by FC(s). + + Note: This model should be used for complex (Dict or Tuple) observation + spaces that have one or more image components. + + The data flow is as follows: + + `obs` (e.g. Tuple[img0, img1, discrete0]) -> `CNN0 + CNN1 + ONE-HOT` + `CNN0 + CNN1 + ONE-HOT` -> concat all flat outputs -> `out` + `out` -> (optional) FC-stack -> `out2` + `out2` -> action (logits) and vaulue heads. + """ + + def __init__(self, obs_space, action_space, num_outputs, model_config, name): + + self.original_space = ( + obs_space.original_space + if hasattr(obs_space, "original_space") + else obs_space + ) + + self.processed_obs_space = ( + self.original_space + if model_config.get("_disable_preprocessor_api") + else obs_space + ) + super().__init__( + self.original_space, action_space, num_outputs, model_config, name + ) + + self.flattened_input_space = flatten_space(self.original_space) + + # Build the CNN(s) given obs_space's image components. + self.cnns = {} + self.one_hot = {} + self.flatten_dims = {} + self.flatten = {} + concat_size = 0 + for i, component in enumerate(self.flattened_input_space): + # Image space. + if len(component.shape) == 3 and isinstance(component, Box): + config = { + "conv_filters": model_config["conv_filters"] + if "conv_filters" in model_config + else get_filter_config(component.shape), + "conv_activation": model_config.get("conv_activation"), + "post_fcnet_hiddens": [], + } + self.cnns[i] = ModelCatalog.get_model_v2( + component, + action_space, + num_outputs=None, + model_config=config, + framework="tf", + name="cnn_{}".format(i), + ) + concat_size += int(self.cnns[i].num_outputs) + # Discrete|MultiDiscrete inputs -> One-hot encode. + elif isinstance(component, (Discrete, MultiDiscrete)): + if isinstance(component, Discrete): + size = component.n + else: + size = np.sum(component.nvec) + config = { + "fcnet_hiddens": model_config["fcnet_hiddens"], + "fcnet_activation": model_config.get("fcnet_activation"), + "post_fcnet_hiddens": [], + } + self.one_hot[i] = ModelCatalog.get_model_v2( + Box(-1.0, 1.0, (size,), np.float32), + action_space, + num_outputs=None, + model_config=config, + framework="tf", + name="one_hot_{}".format(i), + ) + concat_size += int(self.one_hot[i].num_outputs) + # Everything else (1D Box). + else: + size = int(np.prod(component.shape)) + config = { + "fcnet_hiddens": model_config["fcnet_hiddens"], + "fcnet_activation": model_config.get("fcnet_activation"), + "post_fcnet_hiddens": [], + } + self.flatten[i] = ModelCatalog.get_model_v2( + Box(-1.0, 1.0, (size,), np.float32), + action_space, + num_outputs=None, + model_config=config, + framework="tf", + name="flatten_{}".format(i), + ) + self.flatten_dims[i] = size + concat_size += int(self.flatten[i].num_outputs) + + # Optional post-concat FC-stack. + post_fc_stack_config = { + "fcnet_hiddens": model_config.get("post_fcnet_hiddens", []), + "fcnet_activation": model_config.get("post_fcnet_activation", "relu"), + } + self.post_fc_stack = ModelCatalog.get_model_v2( + Box(float("-inf"), float("inf"), shape=(concat_size,), dtype=np.float32), + self.action_space, + None, + post_fc_stack_config, + framework="tf", + name="post_fc_stack", + ) + + # Actions and value heads. + self.logits_and_value_model = None + self._value_out = None + if num_outputs: + # Action-distribution head. + concat_layer = tf.keras.layers.Input((self.post_fc_stack.num_outputs,)) + logits_layer = tf.keras.layers.Dense( + num_outputs, + activation=None, + kernel_initializer=normc_initializer(0.01), + name="logits", + )(concat_layer) + + # Create the value branch model. + value_layer = tf.keras.layers.Dense( + 1, + activation=None, + kernel_initializer=normc_initializer(0.01), + name="value_out", + )(concat_layer) + self.logits_and_value_model = tf.keras.models.Model( + concat_layer, [logits_layer, value_layer] + ) + else: + self.num_outputs = self.post_fc_stack.num_outputs + + @override(ModelV2) + def forward(self, input_dict, state, seq_lens): + if SampleBatch.OBS in input_dict and "obs_flat" in input_dict: + orig_obs = input_dict[SampleBatch.OBS] + else: + orig_obs = restore_original_dimensions( + input_dict[SampleBatch.OBS], self.processed_obs_space, tensorlib="tf" + ) + # Push image observations through our CNNs. + outs = [] + for i, component in enumerate(tree.flatten(orig_obs)): + if i in self.cnns: + cnn_out, _ = self.cnns[i](SampleBatch({SampleBatch.OBS: component})) + outs.append(cnn_out) + elif i in self.one_hot: + if "int" in component.dtype.name: + one_hot_in = { + SampleBatch.OBS: one_hot( + component, self.flattened_input_space[i] + ) + } + else: + one_hot_in = {SampleBatch.OBS: component} + one_hot_out, _ = self.one_hot[i](SampleBatch(one_hot_in)) + outs.append(one_hot_out) + else: + nn_out, _ = self.flatten[i]( + SampleBatch( + { + SampleBatch.OBS: tf.cast( + tf.reshape(component, [-1, self.flatten_dims[i]]), + tf.float32, + ) + } + ) + ) + outs.append(nn_out) + # Concat all outputs and the non-image inputs. + out = tf.concat(outs, axis=1) + # Push through (optional) FC-stack (this may be an empty stack). + out, _ = self.post_fc_stack(SampleBatch({SampleBatch.OBS: out})) + + # No logits/value branches. + if not self.logits_and_value_model: + return out, [] + + # Logits- and value branches. + logits, values = self.logits_and_value_model(out) + self._value_out = tf.reshape(values, [-1]) + return logits, [] + + @override(ModelV2) + def value_function(self): + return self._value_out + + +# __sphinx_doc_end__ diff --git a/lib/python3.12/site-packages/ray/rllib/models/tf/fcnet.py b/lib/python3.12/site-packages/ray/rllib/models/tf/fcnet.py new file mode 100644 index 0000000000000000000000000000000000000000..4ebf1443a5b39903b4dcbb6912283f7e6ae9c3eb --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/tf/fcnet.py @@ -0,0 +1,149 @@ +from typing import Dict + +import gymnasium as gym +import numpy as np + +from ray.rllib.models.tf.misc import normc_initializer +from ray.rllib.models.tf.tf_modelv2 import TFModelV2 +from ray.rllib.models.utils import get_activation_fn +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.typing import List, ModelConfigDict, TensorType + +tf1, tf, tfv = try_import_tf() + + +@OldAPIStack +class FullyConnectedNetwork(TFModelV2): + """Generic fully connected network implemented in ModelV2 API.""" + + def __init__( + self, + obs_space: gym.spaces.Space, + action_space: gym.spaces.Space, + num_outputs: int, + model_config: ModelConfigDict, + name: str, + ): + super(FullyConnectedNetwork, self).__init__( + obs_space, action_space, num_outputs, model_config, name + ) + + hiddens = list(model_config.get("fcnet_hiddens", [])) + list( + model_config.get("post_fcnet_hiddens", []) + ) + activation = model_config.get("fcnet_activation") + if not model_config.get("fcnet_hiddens", []): + activation = model_config.get("post_fcnet_activation") + activation = get_activation_fn(activation) + no_final_linear = model_config.get("no_final_linear") + vf_share_layers = model_config.get("vf_share_layers") + free_log_std = model_config.get("free_log_std") + + # Generate free-floating bias variables for the second half of + # the outputs. + if free_log_std: + assert num_outputs % 2 == 0, ( + "num_outputs must be divisible by two", + num_outputs, + ) + num_outputs = num_outputs // 2 + self.log_std_var = tf.Variable( + [0.0] * num_outputs, dtype=tf.float32, name="log_std" + ) + + # We are using obs_flat, so take the flattened shape as input. + inputs = tf.keras.layers.Input( + shape=(int(np.prod(obs_space.shape)),), name="observations" + ) + # Last hidden layer output (before logits outputs). + last_layer = inputs + # The action distribution outputs. + logits_out = None + i = 1 + + # Create layers 0 to second-last. + for size in hiddens[:-1]: + last_layer = tf.keras.layers.Dense( + size, + name="fc_{}".format(i), + activation=activation, + kernel_initializer=normc_initializer(1.0), + )(last_layer) + i += 1 + + # The last layer is adjusted to be of size num_outputs, but it's a + # layer with activation. + if no_final_linear and num_outputs: + logits_out = tf.keras.layers.Dense( + num_outputs, + name="fc_out", + activation=activation, + kernel_initializer=normc_initializer(1.0), + )(last_layer) + # Finish the layers with the provided sizes (`hiddens`), plus - + # iff num_outputs > 0 - a last linear layer of size num_outputs. + else: + if len(hiddens) > 0: + last_layer = tf.keras.layers.Dense( + hiddens[-1], + name="fc_{}".format(i), + activation=activation, + kernel_initializer=normc_initializer(1.0), + )(last_layer) + if num_outputs: + logits_out = tf.keras.layers.Dense( + num_outputs, + name="fc_out", + activation=None, + kernel_initializer=normc_initializer(0.01), + )(last_layer) + # Adjust num_outputs to be the number of nodes in the last layer. + else: + self.num_outputs = ([int(np.prod(obs_space.shape))] + hiddens[-1:])[-1] + + # Concat the log std vars to the end of the state-dependent means. + if free_log_std and logits_out is not None: + + def tiled_log_std(x): + return tf.tile(tf.expand_dims(self.log_std_var, 0), [tf.shape(x)[0], 1]) + + log_std_out = tf.keras.layers.Lambda(tiled_log_std)(inputs) + logits_out = tf.keras.layers.Concatenate(axis=1)([logits_out, log_std_out]) + + last_vf_layer = None + if not vf_share_layers: + # Build a parallel set of hidden layers for the value net. + last_vf_layer = inputs + i = 1 + for size in hiddens: + last_vf_layer = tf.keras.layers.Dense( + size, + name="fc_value_{}".format(i), + activation=activation, + kernel_initializer=normc_initializer(1.0), + )(last_vf_layer) + i += 1 + + value_out = tf.keras.layers.Dense( + 1, + name="value_out", + activation=None, + kernel_initializer=normc_initializer(0.01), + )(last_vf_layer if last_vf_layer is not None else last_layer) + + self.base_model = tf.keras.Model( + inputs, [(logits_out if logits_out is not None else last_layer), value_out] + ) + + def forward( + self, + input_dict: Dict[str, TensorType], + state: List[TensorType], + seq_lens: TensorType, + ) -> (TensorType, List[TensorType]): + model_out, self._value_out = self.base_model(input_dict["obs_flat"]) + return model_out, state + + def value_function(self) -> TensorType: + return tf.reshape(self._value_out, [-1]) diff --git a/lib/python3.12/site-packages/ray/rllib/models/tf/layers/__init__.py b/lib/python3.12/site-packages/ray/rllib/models/tf/layers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..de4512861089ae3b62a24d7fb6998e5696938316 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/tf/layers/__init__.py @@ -0,0 +1,17 @@ +from ray.rllib.models.tf.layers.gru_gate import GRUGate +from ray.rllib.models.tf.layers.multi_head_attention import MultiHeadAttention +from ray.rllib.models.tf.layers.noisy_layer import NoisyLayer +from ray.rllib.models.tf.layers.relative_multi_head_attention import ( + PositionalEmbedding, + RelativeMultiHeadAttention, +) +from ray.rllib.models.tf.layers.skip_connection import SkipConnection + +__all__ = [ + "GRUGate", + "MultiHeadAttention", + "NoisyLayer", + "PositionalEmbedding", + "RelativeMultiHeadAttention", + "SkipConnection", +] diff --git a/lib/python3.12/site-packages/ray/rllib/models/tf/layers/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/tf/layers/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bcdce0be277cb7522518b072f61166e0337b43e4 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/tf/layers/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/tf/layers/__pycache__/gru_gate.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/tf/layers/__pycache__/gru_gate.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..86d252a03970bc7b8c6dc019f53e09cacc4f09fb Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/tf/layers/__pycache__/gru_gate.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/tf/layers/__pycache__/multi_head_attention.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/tf/layers/__pycache__/multi_head_attention.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2ffc6d03bfdce304e10cab3f1f5e9f370d632d08 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/tf/layers/__pycache__/multi_head_attention.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/tf/layers/__pycache__/noisy_layer.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/tf/layers/__pycache__/noisy_layer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2bc6e6cc5fc0e16ba01ce2b9f13baf6ed9b188dc Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/tf/layers/__pycache__/noisy_layer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/tf/layers/__pycache__/relative_multi_head_attention.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/tf/layers/__pycache__/relative_multi_head_attention.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..231a8e4e323b6a4bfed13dd93a813cc67b2d5f75 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/tf/layers/__pycache__/relative_multi_head_attention.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/tf/layers/__pycache__/skip_connection.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/tf/layers/__pycache__/skip_connection.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8e93302d2c4677d9380cc546ed970f9530942cba Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/tf/layers/__pycache__/skip_connection.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/tf/layers/gru_gate.py b/lib/python3.12/site-packages/ray/rllib/models/tf/layers/gru_gate.py new file mode 100644 index 0000000000000000000000000000000000000000..59b9aa487d579a1f92f71fe1d61952b562d62c13 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/tf/layers/gru_gate.py @@ -0,0 +1,58 @@ +from ray._common.deprecation import deprecation_warning +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.typing import TensorShape, TensorType +from ray.util import log_once + +tf1, tf, tfv = try_import_tf() + + +class GRUGate(tf.keras.layers.Layer if tf else object): + def __init__(self, init_bias: float = 0.0, **kwargs): + super().__init__(**kwargs) + self._init_bias = init_bias + if log_once("gru_gate"): + deprecation_warning( + old="rllib.models.tf.layers.GRUGate", + ) + + def build(self, input_shape: TensorShape): + h_shape, x_shape = input_shape + if x_shape[-1] != h_shape[-1]: + raise ValueError( + "Both inputs to GRUGate must have equal size in last axis!" + ) + + dim = int(h_shape[-1]) + self._w_r = self.add_weight(shape=(dim, dim)) + self._w_z = self.add_weight(shape=(dim, dim)) + self._w_h = self.add_weight(shape=(dim, dim)) + + self._u_r = self.add_weight(shape=(dim, dim)) + self._u_z = self.add_weight(shape=(dim, dim)) + self._u_h = self.add_weight(shape=(dim, dim)) + + def bias_initializer(shape, dtype): + return tf.fill(shape, tf.cast(self._init_bias, dtype=dtype)) + + self._bias_z = self.add_weight(shape=(dim,), initializer=bias_initializer) + + def call(self, inputs: TensorType, **kwargs) -> TensorType: + # Pass in internal state first. + h, X = inputs + + r = tf.tensordot(X, self._w_r, axes=1) + tf.tensordot(h, self._u_r, axes=1) + r = tf.nn.sigmoid(r) + + z = ( + tf.tensordot(X, self._w_z, axes=1) + + tf.tensordot(h, self._u_z, axes=1) + - self._bias_z + ) + z = tf.nn.sigmoid(z) + + h_next = tf.tensordot(X, self._w_h, axes=1) + tf.tensordot( + (h * r), self._u_h, axes=1 + ) + h_next = tf.nn.tanh(h_next) + + return (1 - z) * h + z * h_next diff --git a/lib/python3.12/site-packages/ray/rllib/models/tf/layers/multi_head_attention.py b/lib/python3.12/site-packages/ray/rllib/models/tf/layers/multi_head_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..4d3a4515b288b0760c78921818b18856bc211939 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/tf/layers/multi_head_attention.py @@ -0,0 +1,61 @@ +""" +[1] - Attention Is All You Need - Vaswani, Jones, Shazeer, Parmar, + Uszkoreit, Gomez, Kaiser - Google Brain/Research, U Toronto - 2017. + https://arxiv.org/pdf/1706.03762.pdf +""" +from ray._common.deprecation import deprecation_warning +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.typing import TensorType +from ray.util import log_once + +tf1, tf, tfv = try_import_tf() + + +class MultiHeadAttention(tf.keras.layers.Layer if tf else object): + """A multi-head attention layer described in [1].""" + + def __init__(self, out_dim: int, num_heads: int, head_dim: int, **kwargs): + super().__init__(**kwargs) + + # No bias or non-linearity. + self._num_heads = num_heads + self._head_dim = head_dim + self._qkv_layer = tf.keras.layers.Dense( + 3 * num_heads * head_dim, use_bias=False + ) + self._linear_layer = tf.keras.layers.TimeDistributed( + tf.keras.layers.Dense(out_dim, use_bias=False) + ) + if log_once("multi_head_attention"): + deprecation_warning( + old="rllib.models.tf.layers.MultiHeadAttention", + ) + + def call(self, inputs: TensorType) -> TensorType: + L = tf.shape(inputs)[1] # length of segment + H = self._num_heads # number of attention heads + D = self._head_dim # attention head dimension + + qkv = self._qkv_layer(inputs) + + queries, keys, values = tf.split(qkv, 3, -1) + queries = queries[:, -L:] # only query based on the segment + + queries = tf.reshape(queries, [-1, L, H, D]) + keys = tf.reshape(keys, [-1, L, H, D]) + values = tf.reshape(values, [-1, L, H, D]) + + score = tf.einsum("bihd,bjhd->bijh", queries, keys) + score = score / D**0.5 + + # causal mask of the same length as the sequence + mask = tf.sequence_mask(tf.range(1, L + 1), dtype=score.dtype) + mask = mask[None, :, :, None] + + masked_score = score * mask + 1e30 * (mask - 1.0) + wmat = tf.nn.softmax(masked_score, axis=2) + + out = tf.einsum("bijh,bjhd->bihd", wmat, values) + shape = tf.concat([tf.shape(out)[:2], [H * D]], axis=0) + out = tf.reshape(out, shape) + return self._linear_layer(out) diff --git a/lib/python3.12/site-packages/ray/rllib/models/tf/layers/noisy_layer.py b/lib/python3.12/site-packages/ray/rllib/models/tf/layers/noisy_layer.py new file mode 100644 index 0000000000000000000000000000000000000000..86bd38814fb977bfbe5cfb466e2f65d48ad10b90 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/tf/layers/noisy_layer.py @@ -0,0 +1,118 @@ +import numpy as np + +from ray._common.deprecation import deprecation_warning +from ray.rllib.models.utils import get_activation_fn +from ray.rllib.utils.framework import ( + TensorShape, + TensorType, + get_variable, + try_import_tf, +) +from ray.util import log_once + +tf1, tf, tfv = try_import_tf() + + +class NoisyLayer(tf.keras.layers.Layer if tf else object): + r"""A Layer that adds learnable Noise to some previous layer's outputs. + + Consists of: + - a common dense layer: y = w^{T}x + b + - a noisy layer: y = (w + \epsilon_w*\sigma_w)^{T}x + + (b+\epsilon_b*\sigma_b) + , where \epsilon are random variables sampled from factorized normal + distributions and \sigma are trainable variables which are expected to + vanish along the training procedure. + """ + + def __init__( + self, prefix: str, out_size: int, sigma0: float, activation: str = "relu" + ): + """Initializes a NoisyLayer object. + + Args: + prefix: + out_size: Output size for Noisy Layer + sigma0: Initialization value for sigma_b (bias noise) + non_linear: Non-linear activation for Noisy Layer + """ + super().__init__() + self.prefix = prefix + self.out_size = out_size + # TF noise generation can be unreliable on GPU + # If generating the noise on the CPU, + # lowering sigma0 to 0.1 may be helpful + self.sigma0 = sigma0 # 0.5~GPU, 0.1~CPU + self.activation = activation + # Variables. + self.w = None # Weight matrix. + self.b = None # Biases. + self.sigma_w = None # Noise for weight matrix + self.sigma_b = None # Noise for biases. + if log_once("noisy_layer"): + deprecation_warning( + old="rllib.models.tf.layers.NoisyLayer", + ) + + def build(self, input_shape: TensorShape): + in_size = int(input_shape[1]) + + self.sigma_w = get_variable( + value=tf.keras.initializers.RandomUniform( + minval=-1.0 / np.sqrt(float(in_size)), + maxval=1.0 / np.sqrt(float(in_size)), + ), + trainable=True, + tf_name=self.prefix + "_sigma_w", + shape=[in_size, self.out_size], + dtype=tf.float32, + ) + + self.sigma_b = get_variable( + value=tf.keras.initializers.Constant(self.sigma0 / np.sqrt(float(in_size))), + trainable=True, + tf_name=self.prefix + "_sigma_b", + shape=[self.out_size], + dtype=tf.float32, + ) + + self.w = get_variable( + value=tf.keras.initializers.GlorotUniform(), + tf_name=self.prefix + "_fc_w", + trainable=True, + shape=[in_size, self.out_size], + dtype=tf.float32, + ) + + self.b = get_variable( + value=tf.keras.initializers.Zeros(), + tf_name=self.prefix + "_fc_b", + trainable=True, + shape=[self.out_size], + dtype=tf.float32, + ) + + def call(self, inputs: TensorType) -> TensorType: + in_size = int(inputs.shape[1]) + epsilon_in = tf.random.normal(shape=[in_size]) + epsilon_out = tf.random.normal(shape=[self.out_size]) + epsilon_in = self._f_epsilon(epsilon_in) + epsilon_out = self._f_epsilon(epsilon_out) + epsilon_w = tf.matmul( + a=tf.expand_dims(epsilon_in, -1), b=tf.expand_dims(epsilon_out, 0) + ) + epsilon_b = epsilon_out + + action_activation = ( + tf.matmul(inputs, self.w + self.sigma_w * epsilon_w) + + self.b + + self.sigma_b * epsilon_b + ) + + fn = get_activation_fn(self.activation, framework="tf") + if fn is not None: + action_activation = fn(action_activation) + return action_activation + + def _f_epsilon(self, x: TensorType) -> TensorType: + return tf.math.sign(x) * tf.math.sqrt(tf.math.abs(x)) diff --git a/lib/python3.12/site-packages/ray/rllib/models/tf/layers/relative_multi_head_attention.py b/lib/python3.12/site-packages/ray/rllib/models/tf/layers/relative_multi_head_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..700dfa3cbade39904f37c1960fcdb2658c29bb11 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/tf/layers/relative_multi_head_attention.py @@ -0,0 +1,147 @@ +from typing import Optional + +from ray._common.deprecation import deprecation_warning +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.typing import TensorType +from ray.util import log_once + +tf1, tf, tfv = try_import_tf() + + +class RelativeMultiHeadAttention(tf.keras.layers.Layer if tf else object): + """A RelativeMultiHeadAttention layer as described in [3]. + + Uses segment level recurrence with state reuse. + """ + + def __init__( + self, + out_dim: int, + num_heads: int, + head_dim: int, + input_layernorm: bool = False, + output_activation: Optional["tf.nn.activation"] = None, + **kwargs + ): + """Initializes a RelativeMultiHeadAttention keras Layer object. + + Args: + out_dim: The output dimensions of the multi-head attention + unit. + num_heads: The number of attention heads to use. + Denoted `H` in [2]. + head_dim: The dimension of a single(!) attention head within + a multi-head attention unit. Denoted as `d` in [3]. + input_layernorm: Whether to prepend a LayerNorm before + everything else. Should be True for building a GTrXL. + output_activation (Optional[tf.nn.activation]): Optional tf.nn + activation function. Should be relu for GTrXL. + **kwargs: + """ + if log_once("relative_multi_head_attention"): + deprecation_warning( + old="rllib.models.tf.layers.RelativeMultiHeadAttention", + ) + super().__init__(**kwargs) + + # No bias or non-linearity. + self._num_heads = num_heads + self._head_dim = head_dim + # 3=Query, key, and value inputs. + self._qkv_layer = tf.keras.layers.Dense( + 3 * num_heads * head_dim, use_bias=False + ) + self._linear_layer = tf.keras.layers.TimeDistributed( + tf.keras.layers.Dense(out_dim, use_bias=False, activation=output_activation) + ) + + self._uvar = self.add_weight(shape=(num_heads, head_dim)) + self._vvar = self.add_weight(shape=(num_heads, head_dim)) + + # Constant (non-trainable) sinusoid rel pos encoding matrix, which + # depends on this incoming time dimension. + # For inference, we prepend the memory to the current timestep's + # input: Tau + 1. For training, we prepend the memory to the input + # sequence: Tau + T. + self._pos_embedding = PositionalEmbedding(out_dim) + self._pos_proj = tf.keras.layers.Dense(num_heads * head_dim, use_bias=False) + + self._input_layernorm = None + if input_layernorm: + self._input_layernorm = tf.keras.layers.LayerNormalization(axis=-1) + + def call( + self, inputs: TensorType, memory: Optional[TensorType] = None + ) -> TensorType: + T = tf.shape(inputs)[1] # length of segment (time) + H = self._num_heads # number of attention heads + d = self._head_dim # attention head dimension + + # Add previous memory chunk (as const, w/o gradient) to input. + # Tau (number of (prev) time slices in each memory chunk). + Tau = tf.shape(memory)[1] + inputs = tf.concat([tf.stop_gradient(memory), inputs], axis=1) + + # Apply the Layer-Norm. + if self._input_layernorm is not None: + inputs = self._input_layernorm(inputs) + + qkv = self._qkv_layer(inputs) + + queries, keys, values = tf.split(qkv, 3, -1) + # Cut out memory timesteps from query. + queries = queries[:, -T:] + + # Splitting up queries into per-head dims (d). + queries = tf.reshape(queries, [-1, T, H, d]) + keys = tf.reshape(keys, [-1, Tau + T, H, d]) + values = tf.reshape(values, [-1, Tau + T, H, d]) + + R = self._pos_embedding(Tau + T) + R = self._pos_proj(R) + R = tf.reshape(R, [Tau + T, H, d]) + + # b=batch + # i and j=time indices (i=max-timesteps (inputs); j=Tau memory space) + # h=head + # d=head-dim (over which we will reduce-sum) + score = tf.einsum("bihd,bjhd->bijh", queries + self._uvar, keys) + pos_score = tf.einsum("bihd,jhd->bijh", queries + self._vvar, R) + score = score + self.rel_shift(pos_score) + score = score / d**0.5 + + # Causal mask of the same length as the sequence. + mask = tf.sequence_mask(tf.range(Tau + 1, Tau + T + 1), dtype=score.dtype) + mask = mask[None, :, :, None] + + masked_score = score * mask + 1e30 * (mask - 1.0) + wmat = tf.nn.softmax(masked_score, axis=2) + + out = tf.einsum("bijh,bjhd->bihd", wmat, values) + out = tf.reshape(out, tf.concat((tf.shape(out)[:2], [H * d]), axis=0)) + return self._linear_layer(out) + + @staticmethod + def rel_shift(x: TensorType) -> TensorType: + # Transposed version of the shift approach described in [3]. + # https://github.com/kimiyoung/transformer-xl/blob/ + # 44781ed21dbaec88b280f74d9ae2877f52b492a5/tf/model.py#L31 + x_size = tf.shape(x) + + x = tf.pad(x, [[0, 0], [0, 0], [1, 0], [0, 0]]) + x = tf.reshape(x, [x_size[0], x_size[2] + 1, x_size[1], x_size[3]]) + x = x[:, 1:, :, :] + x = tf.reshape(x, x_size) + + return x + + +class PositionalEmbedding(tf.keras.layers.Layer if tf else object): + def __init__(self, out_dim, **kwargs): + super().__init__(**kwargs) + self.inverse_freq = 1 / (10000 ** (tf.range(0, out_dim, 2.0) / out_dim)) + + def call(self, seq_length): + pos_offsets = tf.cast(tf.range(seq_length - 1, -1, -1), tf.float32) + inputs = pos_offsets[:, None] * self.inverse_freq[None, :] + return tf.concat((tf.sin(inputs), tf.cos(inputs)), axis=-1) diff --git a/lib/python3.12/site-packages/ray/rllib/models/tf/layers/skip_connection.py b/lib/python3.12/site-packages/ray/rllib/models/tf/layers/skip_connection.py new file mode 100644 index 0000000000000000000000000000000000000000..1dcadd3fc5eb66e22b3af395f7c28614a4bba2ac --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/tf/layers/skip_connection.py @@ -0,0 +1,46 @@ +from typing import Any, Optional + +from ray._common.deprecation import deprecation_warning +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.typing import TensorType +from ray.util import log_once + +tf1, tf, tfv = try_import_tf() + + +class SkipConnection(tf.keras.layers.Layer if tf else object): + """Skip connection layer. + + Adds the original input to the output (regular residual layer) OR uses + input as hidden state input to a given fan_in_layer. + """ + + def __init__(self, layer: Any, fan_in_layer: Optional[Any] = None, **kwargs): + """Initializes a SkipConnection keras layer object. + + Args: + layer (tf.keras.layers.Layer): Any layer processing inputs. + fan_in_layer (Optional[tf.keras.layers.Layer]): An optional + layer taking two inputs: The original input and the output + of `layer`. + """ + if log_once("skip_connection"): + deprecation_warning( + old="rllib.models.tf.layers.SkipConnection", + ) + super().__init__(**kwargs) + self._layer = layer + self._fan_in_layer = fan_in_layer + + def call(self, inputs: TensorType, **kwargs) -> TensorType: + # del kwargs + outputs = self._layer(inputs, **kwargs) + # Residual case, just add inputs to outputs. + if self._fan_in_layer is None: + outputs = outputs + inputs + # Fan-in e.g. RNN: Call fan-in with `inputs` and `outputs`. + else: + # NOTE: In the GRU case, `inputs` is the state input. + outputs = self._fan_in_layer((inputs, outputs)) + + return outputs diff --git a/lib/python3.12/site-packages/ray/rllib/models/tf/misc.py b/lib/python3.12/site-packages/ray/rllib/models/tf/misc.py new file mode 100644 index 0000000000000000000000000000000000000000..4fa92e498945475cd6ed06e23e56f01c152f55fc --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/tf/misc.py @@ -0,0 +1,91 @@ +from typing import Any, Optional, Tuple + +import numpy as np + +from ray.rllib.utils.annotations import DeveloperAPI +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.typing import TensorType + +tf1, tf, tfv = try_import_tf() + + +# TODO: (sven) obsolete this class. +@DeveloperAPI +def normc_initializer(std: float = 1.0) -> Any: + def _initializer(shape, dtype=None, partition_info=None): + out = np.random.randn(*shape).astype( + dtype.name if hasattr(dtype, "name") else dtype or np.float32 + ) + out *= std / np.sqrt(np.square(out).sum(axis=0, keepdims=True)) + return tf.constant(out) + + return _initializer + + +@DeveloperAPI +def conv2d( + x: TensorType, + num_filters: int, + name: str, + filter_size: Tuple[int, int] = (3, 3), + stride: Tuple[int, int] = (1, 1), + pad: str = "SAME", + dtype: Optional[Any] = None, + collections: Optional[Any] = None, +) -> TensorType: + + if dtype is None: + dtype = tf.float32 + + with tf1.variable_scope(name): + stride_shape = [1, stride[0], stride[1], 1] + filter_shape = [ + filter_size[0], + filter_size[1], + int(x.get_shape()[3]), + num_filters, + ] + + # There are "num input feature maps * filter height * filter width" + # inputs to each hidden unit. + fan_in = np.prod(filter_shape[:3]) + # Each unit in the lower layer receives a gradient from: "num output + # feature maps * filter height * filter width" / pooling size. + fan_out = np.prod(filter_shape[:2]) * num_filters + # Initialize weights with random weights. + w_bound = np.sqrt(6 / (fan_in + fan_out)) + + w = tf1.get_variable( + "W", + filter_shape, + dtype, + tf1.random_uniform_initializer(-w_bound, w_bound), + collections=collections, + ) + b = tf1.get_variable( + "b", + [1, 1, 1, num_filters], + initializer=tf1.constant_initializer(0.0), + collections=collections, + ) + return tf1.nn.conv2d(x, w, stride_shape, pad) + b + + +@DeveloperAPI +def linear( + x: TensorType, + size: int, + name: str, + initializer: Optional[Any] = None, + bias_init: float = 0.0, +) -> TensorType: + w = tf1.get_variable(name + "/w", [x.get_shape()[1], size], initializer=initializer) + b = tf1.get_variable( + name + "/b", [size], initializer=tf1.constant_initializer(bias_init) + ) + return tf.matmul(x, w) + b + + +@DeveloperAPI +def flatten(x: TensorType) -> TensorType: + return tf.reshape(x, [-1, np.prod(x.get_shape().as_list()[1:])]) diff --git a/lib/python3.12/site-packages/ray/rllib/models/tf/recurrent_net.py b/lib/python3.12/site-packages/ray/rllib/models/tf/recurrent_net.py new file mode 100644 index 0000000000000000000000000000000000000000..e3fd87eba1120372fd6ef137ac13cd08c0b90b0b --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/tf/recurrent_net.py @@ -0,0 +1,293 @@ +import logging +from typing import Dict, List, Tuple + +import gymnasium as gym +import numpy as np +import tree # pip install dm_tree +from gymnasium.spaces import Discrete, MultiDiscrete + +from ray._common.deprecation import deprecation_warning +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.models.tf.tf_modelv2 import TFModelV2 +from ray.rllib.policy.rnn_sequencing import add_time_dimension +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.policy.view_requirement import ViewRequirement +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.spaces.space_utils import get_base_struct_from_space +from ray.rllib.utils.tf_utils import flatten_inputs_to_1d_tensor, one_hot +from ray.rllib.utils.typing import ModelConfigDict, TensorType +from ray.util.debug import log_once + +tf1, tf, tfv = try_import_tf() +logger = logging.getLogger(__name__) + + +@OldAPIStack +class RecurrentNetwork(TFModelV2): + """Helper class to simplify implementing RNN models with TFModelV2. + + Instead of implementing forward(), you can implement forward_rnn() which + takes batches with the time dimension added already. + + Here is an example implementation for a subclass + ``MyRNNClass(RecurrentNetwork)``:: + + def __init__(self, *args, **kwargs): + super(MyModelClass, self).__init__(*args, **kwargs) + cell_size = 256 + + # Define input layers + input_layer = tf.keras.layers.Input( + shape=(None, obs_space.shape[0])) + state_in_h = tf.keras.layers.Input(shape=(256, )) + state_in_c = tf.keras.layers.Input(shape=(256, )) + seq_in = tf.keras.layers.Input(shape=(), dtype=tf.int32) + + # Send to LSTM cell + lstm_out, state_h, state_c = tf.keras.layers.LSTM( + cell_size, return_sequences=True, return_state=True, + name="lstm")( + inputs=input_layer, + mask=tf.sequence_mask(seq_in), + initial_state=[state_in_h, state_in_c]) + output_layer = tf.keras.layers.Dense(...)(lstm_out) + + # Create the RNN model + self.rnn_model = tf.keras.Model( + inputs=[input_layer, seq_in, state_in_h, state_in_c], + outputs=[output_layer, state_h, state_c]) + self.rnn_model.summary() + """ + + @override(ModelV2) + def forward( + self, + input_dict: Dict[str, TensorType], + state: List[TensorType], + seq_lens: TensorType, + ) -> Tuple[TensorType, List[TensorType]]: + """Adds time dimension to batch before sending inputs to forward_rnn(). + + You should implement forward_rnn() in your subclass.""" + # Creating a __init__ function that acts as a passthrough and adding the warning + # there led to errors probably due to the multiple inheritance. We encountered + # the same error if we add the Deprecated decorator. We therefore add the + # deprecation warning here. + if log_once("recurrent_network_tf"): + deprecation_warning( + old="ray.rllib.models.tf.recurrent_net.RecurrentNetwork" + ) + assert seq_lens is not None + flat_inputs = input_dict["obs_flat"] + inputs = add_time_dimension( + padded_inputs=flat_inputs, seq_lens=seq_lens, framework="tf" + ) + output, new_state = self.forward_rnn( + inputs, + state, + seq_lens, + ) + return tf.reshape(output, [-1, self.num_outputs]), new_state + + def forward_rnn( + self, inputs: TensorType, state: List[TensorType], seq_lens: TensorType + ) -> Tuple[TensorType, List[TensorType]]: + """Call the model with the given input tensors and state. + + Args: + inputs: observation tensor with shape [B, T, obs_size]. + state: list of state tensors, each with shape [B, T, size]. + seq_lens: 1d tensor holding input sequence lengths. + + Returns: + (outputs, new_state): The model output tensor of shape + [B, T, num_outputs] and the list of new state tensors each with + shape [B, size]. + + Sample implementation for the ``MyRNNClass`` example:: + + def forward_rnn(self, inputs, state, seq_lens): + model_out, h, c = self.rnn_model([inputs, seq_lens] + state) + return model_out, [h, c] + """ + raise NotImplementedError("You must implement this for a RNN model") + + def get_initial_state(self) -> List[TensorType]: + """Get the initial recurrent state values for the model. + + Returns: + list of np.array objects, if any + + Sample implementation for the ``MyRNNClass`` example:: + + def get_initial_state(self): + return [ + np.zeros(self.cell_size, np.float32), + np.zeros(self.cell_size, np.float32), + ] + """ + raise NotImplementedError("You must implement this for a RNN model") + + +@OldAPIStack +class LSTMWrapper(RecurrentNetwork): + """An LSTM wrapper serving as an interface for ModelV2s that set use_lstm.""" + + def __init__( + self, + obs_space: gym.spaces.Space, + action_space: gym.spaces.Space, + num_outputs: int, + model_config: ModelConfigDict, + name: str, + ): + super(LSTMWrapper, self).__init__( + obs_space, action_space, None, model_config, name + ) + # At this point, self.num_outputs is the number of nodes coming + # from the wrapped (underlying) model. In other words, self.num_outputs + # is the input size for the LSTM layer. + # If None, set it to the observation space. + if self.num_outputs is None: + self.num_outputs = int(np.prod(self.obs_space.shape)) + + self.cell_size = model_config["lstm_cell_size"] + self.use_prev_action = model_config["lstm_use_prev_action"] + self.use_prev_reward = model_config["lstm_use_prev_reward"] + + self.action_space_struct = get_base_struct_from_space(self.action_space) + self.action_dim = 0 + + for space in tree.flatten(self.action_space_struct): + if isinstance(space, Discrete): + self.action_dim += space.n + elif isinstance(space, MultiDiscrete): + self.action_dim += np.sum(space.nvec) + elif space.shape is not None: + self.action_dim += int(np.prod(space.shape)) + else: + self.action_dim += int(len(space)) + + # Add prev-action/reward nodes to input to LSTM. + if self.use_prev_action: + self.num_outputs += self.action_dim + if self.use_prev_reward: + self.num_outputs += 1 + + # Define input layers. + input_layer = tf.keras.layers.Input( + shape=(None, self.num_outputs), name="inputs" + ) + + # Set self.num_outputs to the number of output nodes desired by the + # caller of this constructor. + self.num_outputs = num_outputs + + state_in_h = tf.keras.layers.Input(shape=(self.cell_size,), name="h") + state_in_c = tf.keras.layers.Input(shape=(self.cell_size,), name="c") + seq_in = tf.keras.layers.Input(shape=(), name="seq_in", dtype=tf.int32) + + # Preprocess observation with a hidden layer and send to LSTM cell + lstm_out, state_h, state_c = tf.keras.layers.LSTM( + self.cell_size, return_sequences=True, return_state=True, name="lstm" + )( + inputs=input_layer, + mask=tf.sequence_mask(seq_in), + initial_state=[state_in_h, state_in_c], + ) + + # Postprocess LSTM output with another hidden layer and compute values + logits = tf.keras.layers.Dense( + self.num_outputs, activation=tf.keras.activations.linear, name="logits" + )(lstm_out) + values = tf.keras.layers.Dense(1, activation=None, name="values")(lstm_out) + + # Create the RNN model + self._rnn_model = tf.keras.Model( + inputs=[input_layer, seq_in, state_in_h, state_in_c], + outputs=[logits, values, state_h, state_c], + ) + # Print out model summary in INFO logging mode. + if logger.isEnabledFor(logging.INFO): + self._rnn_model.summary() + + # Add prev-a/r to this model's view, if required. + if model_config["lstm_use_prev_action"]: + self.view_requirements[SampleBatch.PREV_ACTIONS] = ViewRequirement( + SampleBatch.ACTIONS, space=self.action_space, shift=-1 + ) + if model_config["lstm_use_prev_reward"]: + self.view_requirements[SampleBatch.PREV_REWARDS] = ViewRequirement( + SampleBatch.REWARDS, shift=-1 + ) + + @override(RecurrentNetwork) + def forward( + self, + input_dict: Dict[str, TensorType], + state: List[TensorType], + seq_lens: TensorType, + ) -> Tuple[TensorType, List[TensorType]]: + assert seq_lens is not None + # Push obs through "unwrapped" net's `forward()` first. + wrapped_out, _ = self._wrapped_forward(input_dict, [], None) + + # Concat. prev-action/reward if required. + prev_a_r = [] + + # Prev actions. + if self.model_config["lstm_use_prev_action"]: + prev_a = input_dict[SampleBatch.PREV_ACTIONS] + # If actions are not processed yet (in their original form as + # have been sent to environment): + # Flatten/one-hot into 1D array. + if self.model_config["_disable_action_flattening"]: + prev_a_r.append( + flatten_inputs_to_1d_tensor( + prev_a, + spaces_struct=self.action_space_struct, + time_axis=False, + ) + ) + # If actions are already flattened (but not one-hot'd yet!), + # one-hot discrete/multi-discrete actions here. + else: + if isinstance(self.action_space, (Discrete, MultiDiscrete)): + prev_a = one_hot(prev_a, self.action_space) + prev_a_r.append( + tf.reshape(tf.cast(prev_a, tf.float32), [-1, self.action_dim]) + ) + # Prev rewards. + if self.model_config["lstm_use_prev_reward"]: + prev_a_r.append( + tf.reshape( + tf.cast(input_dict[SampleBatch.PREV_REWARDS], tf.float32), [-1, 1] + ) + ) + + # Concat prev. actions + rewards to the "main" input. + if prev_a_r: + wrapped_out = tf.concat([wrapped_out] + prev_a_r, axis=1) + + # Push everything through our LSTM. + input_dict["obs_flat"] = wrapped_out + return super().forward(input_dict, state, seq_lens) + + @override(RecurrentNetwork) + def forward_rnn( + self, inputs: TensorType, state: List[TensorType], seq_lens: TensorType + ) -> Tuple[TensorType, List[TensorType]]: + model_out, self._value_out, h, c = self._rnn_model([inputs, seq_lens] + state) + return model_out, [h, c] + + @override(ModelV2) + def get_initial_state(self) -> List[np.ndarray]: + return [ + np.zeros(self.cell_size, np.float32), + np.zeros(self.cell_size, np.float32), + ] + + @override(ModelV2) + def value_function(self) -> TensorType: + return tf.reshape(self._value_out, [-1]) diff --git a/lib/python3.12/site-packages/ray/rllib/models/tf/tf_action_dist.py b/lib/python3.12/site-packages/ray/rllib/models/tf/tf_action_dist.py new file mode 100644 index 0000000000000000000000000000000000000000..b028c1c6430cc8a6eba255f1172c704190ea4514 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/tf/tf_action_dist.py @@ -0,0 +1,736 @@ +import functools +from math import log +from typing import Optional + +import gymnasium as gym +import numpy as np +import tree # pip install dm_tree + +from ray.rllib.models.action_dist import ActionDistribution +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.utils import MAX_LOG_NN_OUTPUT, MIN_LOG_NN_OUTPUT, SMALL_NUMBER +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.framework import try_import_tf, try_import_tfp +from ray.rllib.utils.spaces.space_utils import get_base_struct_from_space +from ray.rllib.utils.typing import List, ModelConfigDict, TensorType, Tuple, Union + +tf1, tf, tfv = try_import_tf() +tfp = try_import_tfp() + + +@OldAPIStack +class TFActionDistribution(ActionDistribution): + """TF-specific extensions for building action distributions.""" + + @override(ActionDistribution) + def __init__(self, inputs: List[TensorType], model: ModelV2): + super().__init__(inputs, model) + self.sample_op = self._build_sample_op() + self.sampled_action_logp_op = self.logp(self.sample_op) + + def _build_sample_op(self) -> TensorType: + """Implement this instead of sample(), to enable op reuse. + + This is needed since the sample op is non-deterministic and is shared + between sample() and sampled_action_logp(). + """ + raise NotImplementedError + + @override(ActionDistribution) + def sample(self) -> TensorType: + """Draw a sample from the action distribution.""" + return self.sample_op + + @override(ActionDistribution) + def sampled_action_logp(self) -> TensorType: + """Returns the log probability of the sampled action.""" + return self.sampled_action_logp_op + + +@OldAPIStack +class Categorical(TFActionDistribution): + """Categorical distribution for discrete action spaces.""" + + def __init__( + self, inputs: List[TensorType], model: ModelV2 = None, temperature: float = 1.0 + ): + assert temperature > 0.0, "Categorical `temperature` must be > 0.0!" + # Allow softmax formula w/ temperature != 1.0: + # Divide inputs by temperature. + super().__init__(inputs / temperature, model) + + @override(ActionDistribution) + def deterministic_sample(self) -> TensorType: + return tf.math.argmax(self.inputs, axis=1) + + @override(ActionDistribution) + def logp(self, x: TensorType) -> TensorType: + return -tf.nn.sparse_softmax_cross_entropy_with_logits( + logits=self.inputs, labels=tf.cast(x, tf.int32) + ) + + @override(ActionDistribution) + def entropy(self) -> TensorType: + a0 = self.inputs - tf.reduce_max(self.inputs, axis=1, keepdims=True) + ea0 = tf.exp(a0) + z0 = tf.reduce_sum(ea0, axis=1, keepdims=True) + p0 = ea0 / z0 + return tf.reduce_sum(p0 * (tf.math.log(z0) - a0), axis=1) + + @override(ActionDistribution) + def kl(self, other: ActionDistribution) -> TensorType: + a0 = self.inputs - tf.reduce_max(self.inputs, axis=1, keepdims=True) + a1 = other.inputs - tf.reduce_max(other.inputs, axis=1, keepdims=True) + ea0 = tf.exp(a0) + ea1 = tf.exp(a1) + z0 = tf.reduce_sum(ea0, axis=1, keepdims=True) + z1 = tf.reduce_sum(ea1, axis=1, keepdims=True) + p0 = ea0 / z0 + return tf.reduce_sum(p0 * (a0 - tf.math.log(z0) - a1 + tf.math.log(z1)), axis=1) + + @override(TFActionDistribution) + def _build_sample_op(self) -> TensorType: + return tf.squeeze(tf.random.categorical(self.inputs, 1), axis=1) + + @staticmethod + @override(ActionDistribution) + def required_model_output_shape(action_space, model_config): + return action_space.n + + +@OldAPIStack +def get_categorical_class_with_temperature(t: float): + """Categorical distribution class that has customized default temperature.""" + + class CategoricalWithTemperature(Categorical): + def __init__(self, inputs, model=None, temperature=t): + super().__init__(inputs, model, temperature) + + return CategoricalWithTemperature + + +@OldAPIStack +class MultiCategorical(TFActionDistribution): + """MultiCategorical distribution for MultiDiscrete action spaces.""" + + def __init__( + self, + inputs: List[TensorType], + model: ModelV2, + input_lens: Union[List[int], np.ndarray, Tuple[int, ...]], + action_space=None, + ): + # skip TFActionDistribution init + ActionDistribution.__init__(self, inputs, model) + self.cats = [ + Categorical(input_, model) + for input_ in tf.split(inputs, input_lens, axis=1) + ] + self.action_space = action_space + if self.action_space is None: + self.action_space = gym.spaces.MultiDiscrete( + [c.inputs.shape[1] for c in self.cats] + ) + self.sample_op = self._build_sample_op() + self.sampled_action_logp_op = self.logp(self.sample_op) + + @override(ActionDistribution) + def deterministic_sample(self) -> TensorType: + sample_ = tf.stack([cat.deterministic_sample() for cat in self.cats], axis=1) + if isinstance(self.action_space, gym.spaces.Box): + return tf.cast( + tf.reshape(sample_, [-1] + list(self.action_space.shape)), + self.action_space.dtype, + ) + return sample_ + + @override(ActionDistribution) + def logp(self, actions: TensorType) -> TensorType: + # If tensor is provided, unstack it into list. + if isinstance(actions, tf.Tensor): + if isinstance(self.action_space, gym.spaces.Box): + actions = tf.reshape( + actions, [-1, int(np.prod(self.action_space.shape))] + ) + elif isinstance(self.action_space, gym.spaces.MultiDiscrete): + actions.set_shape((None, len(self.cats))) + actions = tf.unstack(tf.cast(actions, tf.int32), axis=1) + logps = tf.stack([cat.logp(act) for cat, act in zip(self.cats, actions)]) + return tf.reduce_sum(logps, axis=0) + + @override(ActionDistribution) + def multi_entropy(self) -> TensorType: + return tf.stack([cat.entropy() for cat in self.cats], axis=1) + + @override(ActionDistribution) + def entropy(self) -> TensorType: + return tf.reduce_sum(self.multi_entropy(), axis=1) + + @override(ActionDistribution) + def multi_kl(self, other: ActionDistribution) -> TensorType: + return tf.stack( + [cat.kl(oth_cat) for cat, oth_cat in zip(self.cats, other.cats)], axis=1 + ) + + @override(ActionDistribution) + def kl(self, other: ActionDistribution) -> TensorType: + return tf.reduce_sum(self.multi_kl(other), axis=1) + + @override(TFActionDistribution) + def _build_sample_op(self) -> TensorType: + sample_op = tf.stack([cat.sample() for cat in self.cats], axis=1) + if isinstance(self.action_space, gym.spaces.Box): + return tf.cast( + tf.reshape(sample_op, [-1] + list(self.action_space.shape)), + dtype=self.action_space.dtype, + ) + return sample_op + + @staticmethod + @override(ActionDistribution) + def required_model_output_shape( + action_space: gym.Space, model_config: ModelConfigDict + ) -> Union[int, np.ndarray]: + # Int Box. + if isinstance(action_space, gym.spaces.Box): + assert action_space.dtype.name.startswith("int") + low_ = np.min(action_space.low) + high_ = np.max(action_space.high) + assert np.all(action_space.low == low_) + assert np.all(action_space.high == high_) + return np.prod(action_space.shape, dtype=np.int32) * (high_ - low_ + 1) + # MultiDiscrete space. + else: + # nvec is already integer, so no casting needed. + return np.sum(action_space.nvec) + + +@OldAPIStack +class SlateMultiCategorical(Categorical): + """MultiCategorical distribution for MultiDiscrete action spaces. + + The action space must be uniform, meaning all nvec items have the same size, e.g. + MultiDiscrete([10, 10, 10]), where 10 is the number of candidates to pick from + and 3 is the slate size (pick 3 out of 10). When picking candidates, no candidate + must be picked more than once. + """ + + def __init__( + self, + inputs: List[TensorType], + model: ModelV2 = None, + temperature: float = 1.0, + action_space: Optional[gym.spaces.MultiDiscrete] = None, + all_slates=None, + ): + assert temperature > 0.0, "Categorical `temperature` must be > 0.0!" + # Allow softmax formula w/ temperature != 1.0: + # Divide inputs by temperature. + super().__init__(inputs / temperature, model) + self.action_space = action_space + # Assert uniformness of the action space (all discrete buckets have the same + # size). + assert isinstance(self.action_space, gym.spaces.MultiDiscrete) and all( + n == self.action_space.nvec[0] for n in self.action_space.nvec + ) + self.all_slates = all_slates + + @override(ActionDistribution) + def deterministic_sample(self) -> TensorType: + # Get a sample from the underlying Categorical (batch of ints). + sample = super().deterministic_sample() + # Use the sampled ints to pick the actual slates. + return tf.gather(self.all_slates, sample) + + @override(ActionDistribution) + def logp(self, x: TensorType) -> TensorType: + # TODO: Implement. + return tf.ones_like(self.inputs[:, 0]) + + +@OldAPIStack +class GumbelSoftmax(TFActionDistribution): + """GumbelSoftmax distr. (for differentiable sampling in discr. actions + + The Gumbel Softmax distribution [1] (also known as the Concrete [2] + distribution) is a close cousin of the relaxed one-hot categorical + distribution, whose tfp implementation we will use here plus + adjusted `sample_...` and `log_prob` methods. See discussion at [0]. + + [0] https://stackoverflow.com/questions/56226133/ + soft-actor-critic-with-discrete-action-space + + [1] Categorical Reparametrization with Gumbel-Softmax (Jang et al, 2017): + https://arxiv.org/abs/1611.01144 + [2] The Concrete Distribution: A Continuous Relaxation of Discrete Random + Variables (Maddison et al, 2017) https://arxiv.org/abs/1611.00712 + """ + + def __init__( + self, inputs: List[TensorType], model: ModelV2 = None, temperature: float = 1.0 + ): + """Initializes a GumbelSoftmax distribution. + + Args: + temperature: Temperature parameter. For low temperatures, + the expected value approaches a categorical random variable. + For high temperatures, the expected value approaches a uniform + distribution. + """ + assert temperature >= 0.0 + self.dist = tfp.distributions.RelaxedOneHotCategorical( + temperature=temperature, logits=inputs + ) + self.probs = tf.nn.softmax(self.dist._distribution.logits) + super().__init__(inputs, model) + + @override(ActionDistribution) + def deterministic_sample(self) -> TensorType: + # Return the dist object's prob values. + return self.probs + + @override(ActionDistribution) + def logp(self, x: TensorType) -> TensorType: + # Override since the implementation of tfp.RelaxedOneHotCategorical + # yields positive values. + if x.shape != self.dist.logits.shape: + values = tf.one_hot( + x, self.dist.logits.shape.as_list()[-1], dtype=tf.float32 + ) + assert values.shape == self.dist.logits.shape, ( + values.shape, + self.dist.logits.shape, + ) + + # [0]'s implementation (see line below) seems to be an approximation + # to the actual Gumbel Softmax density. + return -tf.reduce_sum( + -x * tf.nn.log_softmax(self.dist.logits, axis=-1), axis=-1 + ) + + @override(TFActionDistribution) + def _build_sample_op(self) -> TensorType: + return self.dist.sample() + + @staticmethod + @override(ActionDistribution) + def required_model_output_shape( + action_space: gym.Space, model_config: ModelConfigDict + ) -> Union[int, np.ndarray]: + return action_space.n + + +@OldAPIStack +class DiagGaussian(TFActionDistribution): + """Action distribution where each vector element is a gaussian. + + The first half of the input vector defines the gaussian means, and the + second half the gaussian standard deviations. + """ + + def __init__( + self, + inputs: List[TensorType], + model: ModelV2, + *, + action_space: Optional[gym.spaces.Space] = None + ): + mean, log_std = tf.split(inputs, 2, axis=1) + self.mean = mean + self.log_std = log_std + self.std = tf.exp(log_std) + # Remember to squeeze action samples in case action space is Box(shape) + self.zero_action_dim = action_space and action_space.shape == () + super().__init__(inputs, model) + + @override(ActionDistribution) + def deterministic_sample(self) -> TensorType: + return self.mean + + @override(ActionDistribution) + def logp(self, x: TensorType) -> TensorType: + # Cover case where action space is Box(shape=()). + if int(tf.shape(x).shape[0]) == 1: + x = tf.expand_dims(x, axis=1) + return ( + -0.5 + * tf.reduce_sum( + tf.math.square((tf.cast(x, tf.float32) - self.mean) / self.std), axis=1 + ) + - 0.5 * np.log(2.0 * np.pi) * tf.cast(tf.shape(x)[1], tf.float32) + - tf.reduce_sum(self.log_std, axis=1) + ) + + @override(ActionDistribution) + def kl(self, other: ActionDistribution) -> TensorType: + assert isinstance(other, DiagGaussian) + return tf.reduce_sum( + other.log_std + - self.log_std + + (tf.math.square(self.std) + tf.math.square(self.mean - other.mean)) + / (2.0 * tf.math.square(other.std)) + - 0.5, + axis=1, + ) + + @override(ActionDistribution) + def entropy(self) -> TensorType: + return tf.reduce_sum(self.log_std + 0.5 * np.log(2.0 * np.pi * np.e), axis=1) + + @override(TFActionDistribution) + def _build_sample_op(self) -> TensorType: + sample = self.mean + self.std * tf.random.normal(tf.shape(self.mean)) + if self.zero_action_dim: + return tf.squeeze(sample, axis=-1) + return sample + + @staticmethod + @override(ActionDistribution) + def required_model_output_shape( + action_space: gym.Space, model_config: ModelConfigDict + ) -> Union[int, np.ndarray]: + return np.prod(action_space.shape, dtype=np.int32) * 2 + + +@OldAPIStack +class SquashedGaussian(TFActionDistribution): + """A tanh-squashed Gaussian distribution defined by: mean, std, low, high. + + The distribution will never return low or high exactly, but + `low`+SMALL_NUMBER or `high`-SMALL_NUMBER respectively. + """ + + def __init__( + self, + inputs: List[TensorType], + model: ModelV2, + low: float = -1.0, + high: float = 1.0, + ): + """Parameterizes the distribution via `inputs`. + + Args: + low: The lowest possible sampling value + (excluding this value). + high: The highest possible sampling value + (excluding this value). + """ + assert tfp is not None + mean, log_std = tf.split(inputs, 2, axis=-1) + # Clip `scale` values (coming from NN) to reasonable values. + log_std = tf.clip_by_value(log_std, MIN_LOG_NN_OUTPUT, MAX_LOG_NN_OUTPUT) + std = tf.exp(log_std) + self.distr = tfp.distributions.Normal(loc=mean, scale=std) + assert np.all(np.less(low, high)) + self.low = low + self.high = high + super().__init__(inputs, model) + + @override(ActionDistribution) + def deterministic_sample(self) -> TensorType: + mean = self.distr.mean() + return self._squash(mean) + + @override(TFActionDistribution) + def _build_sample_op(self) -> TensorType: + return self._squash(self.distr.sample()) + + @override(ActionDistribution) + def logp(self, x: TensorType) -> TensorType: + # Unsquash values (from [low,high] to ]-inf,inf[) + unsquashed_values = tf.cast(self._unsquash(x), self.inputs.dtype) + # Get log prob of unsquashed values from our Normal. + log_prob_gaussian = self.distr.log_prob(unsquashed_values) + # For safety reasons, clamp somehow, only then sum up. + log_prob_gaussian = tf.clip_by_value(log_prob_gaussian, -100, 100) + log_prob_gaussian = tf.reduce_sum(log_prob_gaussian, axis=-1) + # Get log-prob for squashed Gaussian. + unsquashed_values_tanhd = tf.math.tanh(unsquashed_values) + log_prob = log_prob_gaussian - tf.reduce_sum( + tf.math.log(1 - unsquashed_values_tanhd**2 + SMALL_NUMBER), axis=-1 + ) + return log_prob + + def sample_logp(self): + z = self.distr.sample() + actions = self._squash(z) + return actions, tf.reduce_sum( + self.distr.log_prob(z) - tf.math.log(1 - actions * actions + SMALL_NUMBER), + axis=-1, + ) + + @override(ActionDistribution) + def entropy(self) -> TensorType: + raise ValueError("Entropy not defined for SquashedGaussian!") + + @override(ActionDistribution) + def kl(self, other: ActionDistribution) -> TensorType: + raise ValueError("KL not defined for SquashedGaussian!") + + def _squash(self, raw_values: TensorType) -> TensorType: + # Returned values are within [low, high] (including `low` and `high`). + squashed = ((tf.math.tanh(raw_values) + 1.0) / 2.0) * ( + self.high - self.low + ) + self.low + return tf.clip_by_value(squashed, self.low, self.high) + + def _unsquash(self, values: TensorType) -> TensorType: + normed_values = (values - self.low) / (self.high - self.low) * 2.0 - 1.0 + # Stabilize input to atanh. + save_normed_values = tf.clip_by_value( + normed_values, -1.0 + SMALL_NUMBER, 1.0 - SMALL_NUMBER + ) + unsquashed = tf.math.atanh(save_normed_values) + return unsquashed + + @staticmethod + @override(ActionDistribution) + def required_model_output_shape( + action_space: gym.Space, model_config: ModelConfigDict + ) -> Union[int, np.ndarray]: + return np.prod(action_space.shape, dtype=np.int32) * 2 + + +@OldAPIStack +class Beta(TFActionDistribution): + """ + A Beta distribution is defined on the interval [0, 1] and parameterized by + shape parameters alpha and beta (also called concentration parameters). + + PDF(x; alpha, beta) = x**(alpha - 1) (1 - x)**(beta - 1) / Z + with Z = Gamma(alpha) Gamma(beta) / Gamma(alpha + beta) + and Gamma(n) = (n - 1)! + """ + + def __init__( + self, + inputs: List[TensorType], + model: ModelV2, + low: float = 0.0, + high: float = 1.0, + ): + # Stabilize input parameters (possibly coming from a linear layer). + inputs = tf.clip_by_value(inputs, log(SMALL_NUMBER), -log(SMALL_NUMBER)) + inputs = tf.math.log(tf.math.exp(inputs) + 1.0) + 1.0 + self.low = low + self.high = high + alpha, beta = tf.split(inputs, 2, axis=-1) + # Note: concentration0==beta, concentration1=alpha (!) + self.dist = tfp.distributions.Beta(concentration1=alpha, concentration0=beta) + super().__init__(inputs, model) + + @override(ActionDistribution) + def deterministic_sample(self) -> TensorType: + mean = self.dist.mean() + return self._squash(mean) + + @override(TFActionDistribution) + def _build_sample_op(self) -> TensorType: + return self._squash(self.dist.sample()) + + @override(ActionDistribution) + def logp(self, x: TensorType) -> TensorType: + unsquashed_values = self._unsquash(x) + return tf.math.reduce_sum(self.dist.log_prob(unsquashed_values), axis=-1) + + def _squash(self, raw_values: TensorType) -> TensorType: + return raw_values * (self.high - self.low) + self.low + + def _unsquash(self, values: TensorType) -> TensorType: + return (values - self.low) / (self.high - self.low) + + @staticmethod + @override(ActionDistribution) + def required_model_output_shape( + action_space: gym.Space, model_config: ModelConfigDict + ) -> Union[int, np.ndarray]: + return np.prod(action_space.shape, dtype=np.int32) * 2 + + +@OldAPIStack +class Deterministic(TFActionDistribution): + """Action distribution that returns the input values directly. + + This is similar to DiagGaussian with standard deviation zero (thus only + requiring the "mean" values as NN output). + """ + + @override(ActionDistribution) + def deterministic_sample(self) -> TensorType: + return self.inputs + + @override(TFActionDistribution) + def logp(self, x: TensorType) -> TensorType: + return tf.zeros_like(self.inputs) + + @override(TFActionDistribution) + def _build_sample_op(self) -> TensorType: + return self.inputs + + @staticmethod + @override(ActionDistribution) + def required_model_output_shape( + action_space: gym.Space, model_config: ModelConfigDict + ) -> Union[int, np.ndarray]: + return np.prod(action_space.shape, dtype=np.int32) + + +@OldAPIStack +class MultiActionDistribution(TFActionDistribution): + """Action distribution that operates on a set of actions. + + Args: + inputs (Tensor list): A list of tensors from which to compute samples. + """ + + def __init__( + self, inputs, model, *, child_distributions, input_lens, action_space, **kwargs + ): + ActionDistribution.__init__(self, inputs, model) + + self.action_space_struct = get_base_struct_from_space(action_space) + + self.input_lens = np.array(input_lens, dtype=np.int32) + split_inputs = tf.split(inputs, self.input_lens, axis=1) + self.flat_child_distributions = tree.map_structure( + lambda dist, input_: dist(input_, model, **kwargs), + child_distributions, + split_inputs, + ) + + @override(ActionDistribution) + def logp(self, x): + # Single tensor input (all merged). + if isinstance(x, (tf.Tensor, np.ndarray)): + split_indices = [] + for dist in self.flat_child_distributions: + if isinstance(dist, Categorical): + split_indices.append(1) + elif ( + isinstance(dist, MultiCategorical) and dist.action_space is not None + ): + split_indices.append(np.prod(dist.action_space.shape)) + else: + sample = dist.sample() + # Cover Box(shape=()) case. + if len(sample.shape) == 1: + split_indices.append(1) + else: + split_indices.append(tf.shape(sample)[1]) + split_x = tf.split(x, split_indices, axis=1) + # Structured or flattened (by single action component) input. + else: + split_x = tree.flatten(x) + + def map_(val, dist): + # Remove extra categorical dimension. + if isinstance(dist, Categorical): + val = tf.cast( + tf.squeeze(val, axis=-1) if len(val.shape) > 1 else val, tf.int32 + ) + return dist.logp(val) + + # Remove extra categorical dimension and take the logp of each + # component. + flat_logps = tree.map_structure(map_, split_x, self.flat_child_distributions) + + return functools.reduce(lambda a, b: a + b, flat_logps) + + @override(ActionDistribution) + def kl(self, other): + kl_list = [ + d.kl(o) + for d, o in zip( + self.flat_child_distributions, other.flat_child_distributions + ) + ] + return functools.reduce(lambda a, b: a + b, kl_list) + + @override(ActionDistribution) + def entropy(self): + entropy_list = [d.entropy() for d in self.flat_child_distributions] + return functools.reduce(lambda a, b: a + b, entropy_list) + + @override(ActionDistribution) + def sample(self): + child_distributions = tree.unflatten_as( + self.action_space_struct, self.flat_child_distributions + ) + return tree.map_structure(lambda s: s.sample(), child_distributions) + + @override(ActionDistribution) + def deterministic_sample(self): + child_distributions = tree.unflatten_as( + self.action_space_struct, self.flat_child_distributions + ) + return tree.map_structure( + lambda s: s.deterministic_sample(), child_distributions + ) + + @override(TFActionDistribution) + def sampled_action_logp(self): + p = self.flat_child_distributions[0].sampled_action_logp() + for c in self.flat_child_distributions[1:]: + p += c.sampled_action_logp() + return p + + @override(ActionDistribution) + def required_model_output_shape(self, action_space, model_config): + return np.sum(self.input_lens, dtype=np.int32) + + +@OldAPIStack +class Dirichlet(TFActionDistribution): + """Dirichlet distribution for continuous actions that are between + [0,1] and sum to 1. + + e.g. actions that represent resource allocation.""" + + def __init__(self, inputs: List[TensorType], model: ModelV2): + """Input is a tensor of logits. The exponential of logits is used to + parametrize the Dirichlet distribution as all parameters need to be + positive. An arbitrary small epsilon is added to the concentration + parameters to be zero due to numerical error. + + See issue #4440 for more details. + """ + self.epsilon = 1e-7 + concentration = tf.exp(inputs) + self.epsilon + self.dist = tf1.distributions.Dirichlet( + concentration=concentration, + validate_args=True, + allow_nan_stats=False, + ) + super().__init__(concentration, model) + + @override(ActionDistribution) + def deterministic_sample(self) -> TensorType: + return tf.nn.softmax(self.dist.concentration) + + @override(ActionDistribution) + def logp(self, x: TensorType) -> TensorType: + # Support of Dirichlet are positive real numbers. x is already + # an array of positive numbers, but we clip to avoid zeros due to + # numerical errors. + x = tf.maximum(x, self.epsilon) + x = x / tf.reduce_sum(x, axis=-1, keepdims=True) + return self.dist.log_prob(x) + + @override(ActionDistribution) + def entropy(self) -> TensorType: + return self.dist.entropy() + + @override(ActionDistribution) + def kl(self, other: ActionDistribution) -> TensorType: + return self.dist.kl_divergence(other.dist) + + @override(TFActionDistribution) + def _build_sample_op(self) -> TensorType: + return self.dist.sample() + + @staticmethod + @override(ActionDistribution) + def required_model_output_shape( + action_space: gym.Space, model_config: ModelConfigDict + ) -> Union[int, np.ndarray]: + return np.prod(action_space.shape, dtype=np.int32) diff --git a/lib/python3.12/site-packages/ray/rllib/models/tf/tf_modelv2.py b/lib/python3.12/site-packages/ray/rllib/models/tf/tf_modelv2.py new file mode 100644 index 0000000000000000000000000000000000000000..3cd8135772453da1166594c684e450a8b0fc7e32 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/tf/tf_modelv2.py @@ -0,0 +1,143 @@ +import contextlib +import re +from typing import Dict, List, Union + +import gymnasium as gym + +from ray._common.deprecation import deprecation_warning +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.typing import ModelConfigDict, TensorType +from ray.util import log_once + +tf1, tf, tfv = try_import_tf() + + +@OldAPIStack +class TFModelV2(ModelV2): + """TF version of ModelV2, which should contain a tf keras Model. + + Note that this class by itself is not a valid model unless you + implement forward() in a subclass.""" + + def __init__( + self, + obs_space: gym.spaces.Space, + action_space: gym.spaces.Space, + num_outputs: int, + model_config: ModelConfigDict, + name: str, + ): + """Initializes a TFModelV2 instance. + + Here is an example implementation for a subclass + ``MyModelClass(TFModelV2)``:: + + def __init__(self, *args, **kwargs): + super(MyModelClass, self).__init__(*args, **kwargs) + input_layer = tf.keras.layers.Input(...) + hidden_layer = tf.keras.layers.Dense(...)(input_layer) + output_layer = tf.keras.layers.Dense(...)(hidden_layer) + value_layer = tf.keras.layers.Dense(...)(hidden_layer) + self.base_model = tf.keras.Model( + input_layer, [output_layer, value_layer]) + """ + super().__init__( + obs_space, action_space, num_outputs, model_config, name, framework="tf" + ) + + # Deprecated: TFModelV2 now automatically track their variables. + self.var_list = [] + + if tf1.executing_eagerly(): + self.graph = None + else: + self.graph = tf1.get_default_graph() + + def context(self) -> contextlib.AbstractContextManager: + """Returns a contextmanager for the current TF graph.""" + if self.graph: + return self.graph.as_default() + else: + return ModelV2.context(self) + + def update_ops(self) -> List[TensorType]: + """Return the list of update ops for this model. + + For example, this should include any BatchNorm update ops.""" + return [] + + def register_variables(self, variables: List[TensorType]) -> None: + """Register the given list of variables with this model.""" + if log_once("deprecated_tfmodelv2_register_variables"): + deprecation_warning(old="TFModelV2.register_variables", error=False) + self.var_list.extend(variables) + + @override(ModelV2) + def variables( + self, as_dict: bool = False + ) -> Union[List[TensorType], Dict[str, TensorType]]: + if as_dict: + # Old way using `register_variables`. + if self.var_list: + return {v.name: v for v in self.var_list} + # New way: Automatically determine the var tree. + else: + return self._find_sub_modules("", self.__dict__) + + # Old way using `register_variables`. + if self.var_list: + return list(self.var_list) + # New way: Automatically determine the var tree. + else: + return list(self.variables(as_dict=True).values()) + + @override(ModelV2) + def trainable_variables( + self, as_dict: bool = False + ) -> Union[List[TensorType], Dict[str, TensorType]]: + if as_dict: + return { + k: v for k, v in self.variables(as_dict=True).items() if v.trainable + } + return [v for v in self.variables() if v.trainable] + + @staticmethod + def _find_sub_modules(current_key, struct): + # Keras Model: key=k + "." + var-name (replace '/' by '.'). + if isinstance(struct, tf.keras.models.Model) or isinstance(struct, tf.Module): + ret = {} + for var in struct.variables: + name = re.sub("/", ".", var.name) + key = current_key + "." + name + ret[key] = var + return ret + # Other TFModelV2: Include its vars into ours. + elif isinstance(struct, TFModelV2): + return { + current_key + "." + key: var + for key, var in struct.variables(as_dict=True).items() + } + # tf.Variable + elif isinstance(struct, tf.Variable): + return {current_key: struct} + # List/Tuple. + elif isinstance(struct, (tuple, list)): + ret = {} + for i, value in enumerate(struct): + sub_vars = TFModelV2._find_sub_modules( + current_key + "_{}".format(i), value + ) + ret.update(sub_vars) + return ret + # Dict. + elif isinstance(struct, dict): + if current_key: + current_key += "_" + ret = {} + for key, value in struct.items(): + sub_vars = TFModelV2._find_sub_modules(current_key + str(key), value) + ret.update(sub_vars) + return ret + return {} diff --git a/lib/python3.12/site-packages/ray/rllib/models/tf/visionnet.py b/lib/python3.12/site-packages/ray/rllib/models/tf/visionnet.py new file mode 100644 index 0000000000000000000000000000000000000000..15b137a79bc5add52e1a8312367c6fc6c25e8808 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/tf/visionnet.py @@ -0,0 +1,265 @@ +from typing import Dict, List + +import gymnasium as gym + +from ray.rllib.models.tf.misc import normc_initializer +from ray.rllib.models.tf.tf_modelv2 import TFModelV2 +from ray.rllib.models.utils import get_activation_fn, get_filter_config +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.typing import ModelConfigDict, TensorType + +tf1, tf, tfv = try_import_tf() + + +@OldAPIStack +class VisionNetwork(TFModelV2): + """Generic vision network implemented in ModelV2 API. + + An additional post-conv fully connected stack can be added and configured + via the config keys: + `post_fcnet_hiddens`: Dense layer sizes after the Conv2D stack. + `post_fcnet_activation`: Activation function to use for this FC stack. + """ + + def __init__( + self, + obs_space: gym.spaces.Space, + action_space: gym.spaces.Space, + num_outputs: int, + model_config: ModelConfigDict, + name: str, + ): + if not model_config.get("conv_filters"): + model_config["conv_filters"] = get_filter_config(obs_space.shape) + + super(VisionNetwork, self).__init__( + obs_space, action_space, num_outputs, model_config, name + ) + + activation = get_activation_fn( + self.model_config.get("conv_activation"), framework="tf" + ) + filters = self.model_config["conv_filters"] + assert len(filters) > 0, "Must provide at least 1 entry in `conv_filters`!" + + # Post FC net config. + post_fcnet_hiddens = model_config.get("post_fcnet_hiddens", []) + post_fcnet_activation = get_activation_fn( + model_config.get("post_fcnet_activation"), framework="tf" + ) + + no_final_linear = self.model_config.get("no_final_linear") + vf_share_layers = self.model_config.get("vf_share_layers") + + input_shape = obs_space.shape + self.data_format = "channels_last" + + inputs = tf.keras.layers.Input(shape=input_shape, name="observations") + last_layer = inputs + # Whether the last layer is the output of a Flattened (rather than + # a n x (1,1) Conv2D). + self.last_layer_is_flattened = False + + # Build the action layers + for i, (out_size, kernel, stride) in enumerate(filters[:-1], 1): + last_layer = tf.keras.layers.Conv2D( + out_size, + kernel, + strides=stride + if isinstance(stride, (list, tuple)) + else (stride, stride), + activation=activation, + padding="same", + data_format="channels_last", + name="conv{}".format(i), + )(last_layer) + + out_size, kernel, stride = filters[-1] + + # No final linear: Last layer has activation function and exits with + # num_outputs nodes (this could be a 1x1 conv or a FC layer, depending + # on `post_fcnet_...` settings). + if no_final_linear and num_outputs: + last_layer = tf.keras.layers.Conv2D( + out_size if post_fcnet_hiddens else num_outputs, + kernel, + strides=stride + if isinstance(stride, (list, tuple)) + else (stride, stride), + activation=activation, + padding="valid", + data_format="channels_last", + name="conv_out", + )(last_layer) + # Add (optional) post-fc-stack after last Conv2D layer. + layer_sizes = post_fcnet_hiddens[:-1] + ( + [num_outputs] if post_fcnet_hiddens else [] + ) + feature_out = last_layer + + for i, out_size in enumerate(layer_sizes): + feature_out = last_layer + last_layer = tf.keras.layers.Dense( + out_size, + name="post_fcnet_{}".format(i), + activation=post_fcnet_activation, + kernel_initializer=normc_initializer(1.0), + )(last_layer) + + # Finish network normally (w/o overriding last layer size with + # `num_outputs`), then add another linear one of size `num_outputs`. + else: + last_layer = tf.keras.layers.Conv2D( + out_size, + kernel, + strides=stride + if isinstance(stride, (list, tuple)) + else (stride, stride), + activation=activation, + padding="valid", + data_format="channels_last", + name="conv{}".format(len(filters)), + )(last_layer) + + # num_outputs defined. Use that to create an exact + # `num_output`-sized (1,1)-Conv2D. + if num_outputs: + if post_fcnet_hiddens: + last_cnn = last_layer = tf.keras.layers.Conv2D( + post_fcnet_hiddens[0], + [1, 1], + activation=post_fcnet_activation, + padding="same", + data_format="channels_last", + name="conv_out", + )(last_layer) + # Add (optional) post-fc-stack after last Conv2D layer. + for i, out_size in enumerate( + post_fcnet_hiddens[1:] + [num_outputs] + ): + feature_out = last_layer + last_layer = tf.keras.layers.Dense( + out_size, + name="post_fcnet_{}".format(i + 1), + activation=post_fcnet_activation + if i < len(post_fcnet_hiddens) - 1 + else None, + kernel_initializer=normc_initializer(1.0), + )(last_layer) + else: + feature_out = last_layer + last_cnn = last_layer = tf.keras.layers.Conv2D( + num_outputs, + [1, 1], + activation=None, + padding="same", + data_format="channels_last", + name="conv_out", + )(last_layer) + + if last_cnn.shape[1] != 1 or last_cnn.shape[2] != 1: + raise ValueError( + "Given `conv_filters` ({}) do not result in a [B, 1, " + "1, {} (`num_outputs`)] shape (but in {})! Please " + "adjust your Conv2D stack such that the dims 1 and 2 " + "are both 1.".format( + self.model_config["conv_filters"], + self.num_outputs, + list(last_cnn.shape), + ) + ) + + # num_outputs not known -> Flatten, then set self.num_outputs + # to the resulting number of nodes. + else: + self.last_layer_is_flattened = True + last_layer = tf.keras.layers.Flatten(data_format="channels_last")( + last_layer + ) + + # Add (optional) post-fc-stack after last Conv2D layer. + for i, out_size in enumerate(post_fcnet_hiddens): + last_layer = tf.keras.layers.Dense( + out_size, + name="post_fcnet_{}".format(i), + activation=post_fcnet_activation, + kernel_initializer=normc_initializer(1.0), + )(last_layer) + feature_out = last_layer + self.num_outputs = last_layer.shape[1] + logits_out = last_layer + + # Build the value layers + if vf_share_layers: + if not self.last_layer_is_flattened: + feature_out = tf.keras.layers.Lambda( + lambda x: tf.squeeze(x, axis=[1, 2]) + )(feature_out) + value_out = tf.keras.layers.Dense( + 1, + name="value_out", + activation=None, + kernel_initializer=normc_initializer(0.01), + )(feature_out) + else: + # build a parallel set of hidden layers for the value net + last_layer = inputs + for i, (out_size, kernel, stride) in enumerate(filters[:-1], 1): + last_layer = tf.keras.layers.Conv2D( + out_size, + kernel, + strides=stride + if isinstance(stride, (list, tuple)) + else (stride, stride), + activation=activation, + padding="same", + data_format="channels_last", + name="conv_value_{}".format(i), + )(last_layer) + out_size, kernel, stride = filters[-1] + last_layer = tf.keras.layers.Conv2D( + out_size, + kernel, + strides=stride + if isinstance(stride, (list, tuple)) + else (stride, stride), + activation=activation, + padding="valid", + data_format="channels_last", + name="conv_value_{}".format(len(filters)), + )(last_layer) + last_layer = tf.keras.layers.Conv2D( + 1, + [1, 1], + activation=None, + padding="same", + data_format="channels_last", + name="conv_value_out", + )(last_layer) + value_out = tf.keras.layers.Lambda(lambda x: tf.squeeze(x, axis=[1, 2]))( + last_layer + ) + + self.base_model = tf.keras.Model(inputs, [logits_out, value_out]) + + def forward( + self, + input_dict: Dict[str, TensorType], + state: List[TensorType], + seq_lens: TensorType, + ) -> (TensorType, List[TensorType]): + obs = input_dict["obs"] + if self.data_format == "channels_first": + obs = tf.transpose(obs, [0, 2, 3, 1]) + # Explicit cast to float32 needed in eager. + model_out, self._value_out = self.base_model(tf.cast(obs, tf.float32)) + # Our last layer is already flat. + if self.last_layer_is_flattened: + return model_out, state + # Last layer is a n x [1,1] Conv2D -> Flatten. + else: + return tf.squeeze(model_out, axis=[1, 2]), state + + def value_function(self) -> TensorType: + return tf.reshape(self._value_out, [-1]) diff --git a/lib/python3.12/site-packages/ray/rllib/models/torch/__init__.py b/lib/python3.12/site-packages/ray/rllib/models/torch/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..abbe5ef604646d8de477335e18c0bc4c88363b2c --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/torch/__init__.py @@ -0,0 +1,12 @@ +# from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 +# from ray.rllib.models.torch.fcnet import FullyConnectedNetwork +# from ray.rllib.models.torch.recurrent_net import \ +# RecurrentNetwork +# from ray.rllib.models.torch.visionnet import VisionNetwork + +# __all__ = [ +# "FullyConnectedNetwork", +# "RecurrentNetwork", +# "TorchModelV2", +# "VisionNetwork", +# ] diff --git a/lib/python3.12/site-packages/ray/rllib/models/torch/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/torch/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..401569cfbb4f3194f2e000701c0b68183817c950 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/torch/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/torch/__pycache__/attention_net.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/torch/__pycache__/attention_net.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fde30980656ffc0626e0c4aa18764f6237fcc497 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/torch/__pycache__/attention_net.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/torch/__pycache__/complex_input_net.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/torch/__pycache__/complex_input_net.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e17856ff70b81639c3eae8f80e513ea0a4715406 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/torch/__pycache__/complex_input_net.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/torch/__pycache__/fcnet.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/torch/__pycache__/fcnet.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..677f162603d160123781360612b05ded772796bf Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/torch/__pycache__/fcnet.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/torch/__pycache__/mingpt.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/torch/__pycache__/mingpt.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..871dd26b5abdb4057ae9234fbb3c97950d87dd20 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/torch/__pycache__/mingpt.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/torch/__pycache__/misc.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/torch/__pycache__/misc.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..06da4d3c026d113d69ae4ee4c1997b331c24cf7c Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/torch/__pycache__/misc.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/torch/__pycache__/recurrent_net.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/torch/__pycache__/recurrent_net.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e0720afc2c38f81d38b80998de2bcb119a52ac2a Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/torch/__pycache__/recurrent_net.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/torch/__pycache__/torch_action_dist.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/torch/__pycache__/torch_action_dist.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..75345b10d1184643ac44b17b4d0acc11b5d36bc6 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/torch/__pycache__/torch_action_dist.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/torch/__pycache__/torch_distributions.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/torch/__pycache__/torch_distributions.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cc0c3330a65a9927c580f450b1537e7a8c578cdd Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/torch/__pycache__/torch_distributions.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/torch/__pycache__/torch_modelv2.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/torch/__pycache__/torch_modelv2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4a493e393cb1587054694a8e6894c453bc8df9c0 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/torch/__pycache__/torch_modelv2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/torch/__pycache__/visionnet.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/torch/__pycache__/visionnet.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cefcb4f84d5baea64c283a20ac4354497bdb9918 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/torch/__pycache__/visionnet.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/torch/attention_net.py b/lib/python3.12/site-packages/ray/rllib/models/torch/attention_net.py new file mode 100644 index 0000000000000000000000000000000000000000..5954c926350cb5960cde3918678fbf37312cd072 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/torch/attention_net.py @@ -0,0 +1,458 @@ +""" +[1] - Attention Is All You Need - Vaswani, Jones, Shazeer, Parmar, + Uszkoreit, Gomez, Kaiser - Google Brain/Research, U Toronto - 2017. + https://arxiv.org/pdf/1706.03762.pdf +[2] - Stabilizing Transformers for Reinforcement Learning - E. Parisotto + et al. - DeepMind - 2019. https://arxiv.org/pdf/1910.06764.pdf +[3] - Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context. + Z. Dai, Z. Yang, et al. - Carnegie Mellon U - 2019. + https://www.aclweb.org/anthology/P19-1285.pdf +""" +from typing import Dict, Optional, Union + +import gymnasium as gym +import numpy as np +import tree # pip install dm_tree +from gymnasium.spaces import Box, Discrete, MultiDiscrete + +from ray._common.deprecation import deprecation_warning +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.models.torch.misc import SlimFC +from ray.rllib.models.torch.modules import ( + GRUGate, + RelativeMultiHeadAttention, + SkipConnection, +) +from ray.rllib.models.torch.recurrent_net import RecurrentNetwork +from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.policy.view_requirement import ViewRequirement +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.spaces.space_utils import get_base_struct_from_space +from ray.rllib.utils.torch_utils import flatten_inputs_to_1d_tensor, one_hot +from ray.rllib.utils.typing import List, ModelConfigDict, TensorType +from ray.util import log_once + +torch, nn = try_import_torch() + + +@OldAPIStack +class GTrXLNet(RecurrentNetwork, nn.Module): + """A GTrXL net Model described in [2]. + + This is still in an experimental phase. + Can be used as a drop-in replacement for LSTMs in PPO and IMPALA. + + To use this network as a replacement for an RNN, configure your Algorithm + as follows: + + Examples: + >> config["model"]["custom_model"] = GTrXLNet + >> config["model"]["max_seq_len"] = 10 + >> config["model"]["custom_model_config"] = { + >> num_transformer_units=1, + >> attention_dim=32, + >> num_heads=2, + >> memory_tau=50, + >> etc.. + >> } + """ + + def __init__( + self, + observation_space: gym.spaces.Space, + action_space: gym.spaces.Space, + num_outputs: Optional[int], + model_config: ModelConfigDict, + name: str, + *, + num_transformer_units: int = 1, + attention_dim: int = 64, + num_heads: int = 2, + memory_inference: int = 50, + memory_training: int = 50, + head_dim: int = 32, + position_wise_mlp_dim: int = 32, + init_gru_gate_bias: float = 2.0 + ): + """Initializes a GTrXLNet. + + Args: + num_transformer_units: The number of Transformer repeats to + use (denoted L in [2]). + attention_dim: The input and output dimensions of one + Transformer unit. + num_heads: The number of attention heads to use in parallel. + Denoted as `H` in [3]. + memory_inference: The number of timesteps to concat (time + axis) and feed into the next transformer unit as inference + input. The first transformer unit will receive this number of + past observations (plus the current one), instead. + memory_training: The number of timesteps to concat (time + axis) and feed into the next transformer unit as training + input (plus the actual input sequence of len=max_seq_len). + The first transformer unit will receive this number of + past observations (plus the input sequence), instead. + head_dim: The dimension of a single(!) attention head within + a multi-head attention unit. Denoted as `d` in [3]. + position_wise_mlp_dim: The dimension of the hidden layer + within the position-wise MLP (after the multi-head attention + block within one Transformer unit). This is the size of the + first of the two layers within the PositionwiseFeedforward. The + second layer always has size=`attention_dim`. + init_gru_gate_bias: Initial bias values for the GRU gates + (two GRUs per Transformer unit, one after the MHA, one after + the position-wise MLP). + """ + super().__init__( + observation_space, action_space, num_outputs, model_config, name + ) + + nn.Module.__init__(self) + + self.num_transformer_units = num_transformer_units + self.attention_dim = attention_dim + self.num_heads = num_heads + self.memory_inference = memory_inference + self.memory_training = memory_training + self.head_dim = head_dim + self.max_seq_len = model_config["max_seq_len"] + self.obs_dim = observation_space.shape[0] + + self.linear_layer = SlimFC(in_size=self.obs_dim, out_size=self.attention_dim) + + self.layers = [self.linear_layer] + + attention_layers = [] + # 2) Create L Transformer blocks according to [2]. + for i in range(self.num_transformer_units): + # RelativeMultiHeadAttention part. + MHA_layer = SkipConnection( + RelativeMultiHeadAttention( + in_dim=self.attention_dim, + out_dim=self.attention_dim, + num_heads=num_heads, + head_dim=head_dim, + input_layernorm=True, + output_activation=nn.ReLU, + ), + fan_in_layer=GRUGate(self.attention_dim, init_gru_gate_bias), + ) + + # Position-wise MultiLayerPerceptron part. + E_layer = SkipConnection( + nn.Sequential( + torch.nn.LayerNorm(self.attention_dim), + SlimFC( + in_size=self.attention_dim, + out_size=position_wise_mlp_dim, + use_bias=False, + activation_fn=nn.ReLU, + ), + SlimFC( + in_size=position_wise_mlp_dim, + out_size=self.attention_dim, + use_bias=False, + activation_fn=nn.ReLU, + ), + ), + fan_in_layer=GRUGate(self.attention_dim, init_gru_gate_bias), + ) + + # Build a list of all attanlayers in order. + attention_layers.extend([MHA_layer, E_layer]) + + # Create a Sequential such that all parameters inside the attention + # layers are automatically registered with this top-level model. + self.attention_layers = nn.Sequential(*attention_layers) + self.layers.extend(attention_layers) + + # Final layers if num_outputs not None. + self.logits = None + self.values_out = None + # Last value output. + self._value_out = None + # Postprocess GTrXL output with another hidden layer. + if self.num_outputs is not None: + self.logits = SlimFC( + in_size=self.attention_dim, + out_size=self.num_outputs, + activation_fn=nn.ReLU, + ) + + # Value function used by all RLlib Torch RL implementations. + self.values_out = SlimFC( + in_size=self.attention_dim, out_size=1, activation_fn=None + ) + else: + self.num_outputs = self.attention_dim + + # Setup trajectory views (`memory-inference` x past memory outs). + for i in range(self.num_transformer_units): + space = Box(-1.0, 1.0, shape=(self.attention_dim,)) + self.view_requirements["state_in_{}".format(i)] = ViewRequirement( + "state_out_{}".format(i), + shift="-{}:-1".format(self.memory_inference), + # Repeat the incoming state every max-seq-len times. + batch_repeat_value=self.max_seq_len, + space=space, + ) + self.view_requirements["state_out_{}".format(i)] = ViewRequirement( + space=space, used_for_training=False + ) + + @override(ModelV2) + def forward( + self, input_dict, state: List[TensorType], seq_lens: TensorType + ) -> (TensorType, List[TensorType]): + assert seq_lens is not None + + # Add the needed batch rank (tf Models' Input requires this). + observations = input_dict[SampleBatch.OBS] + # Add the time dim to observations. + B = len(seq_lens) + T = observations.shape[0] // B + observations = torch.reshape( + observations, [-1, T] + list(observations.shape[1:]) + ) + + all_out = observations + memory_outs = [] + for i in range(len(self.layers)): + # MHA layers which need memory passed in. + if i % 2 == 1: + all_out = self.layers[i](all_out, memory=state[i // 2]) + # Either self.linear_layer (initial obs -> attn. dim layer) or + # MultiLayerPerceptrons. The output of these layers is always the + # memory for the next forward pass. + else: + all_out = self.layers[i](all_out) + memory_outs.append(all_out) + + # Discard last output (not needed as a memory since it's the last + # layer). + memory_outs = memory_outs[:-1] + + if self.logits is not None: + out = self.logits(all_out) + self._value_out = self.values_out(all_out) + out_dim = self.num_outputs + else: + out = all_out + out_dim = self.attention_dim + + return torch.reshape(out, [-1, out_dim]), [ + torch.reshape(m, [-1, self.attention_dim]) for m in memory_outs + ] + + # TODO: (sven) Deprecate this once trajectory view API has fully matured. + @override(RecurrentNetwork) + def get_initial_state(self) -> List[np.ndarray]: + return [] + + @override(ModelV2) + def value_function(self) -> TensorType: + assert ( + self._value_out is not None + ), "Must call forward first AND must have value branch!" + return torch.reshape(self._value_out, [-1]) + + +class AttentionWrapper(TorchModelV2, nn.Module): + """GTrXL wrapper serving as interface for ModelV2s that set use_attention.""" + + def __init__( + self, + obs_space: gym.spaces.Space, + action_space: gym.spaces.Space, + num_outputs: int, + model_config: ModelConfigDict, + name: str, + ): + if log_once("deprecate_attention_wrapper_torch"): + deprecation_warning( + old="ray.rllib.models.torch.attention_net.AttentionWrapper" + ) + + nn.Module.__init__(self) + super().__init__(obs_space, action_space, None, model_config, name) + + self.use_n_prev_actions = model_config["attention_use_n_prev_actions"] + self.use_n_prev_rewards = model_config["attention_use_n_prev_rewards"] + + self.action_space_struct = get_base_struct_from_space(self.action_space) + self.action_dim = 0 + + for space in tree.flatten(self.action_space_struct): + if isinstance(space, Discrete): + self.action_dim += space.n + elif isinstance(space, MultiDiscrete): + self.action_dim += np.sum(space.nvec) + elif space.shape is not None: + self.action_dim += int(np.prod(space.shape)) + else: + self.action_dim += int(len(space)) + + # Add prev-action/reward nodes to input to LSTM. + if self.use_n_prev_actions: + self.num_outputs += self.use_n_prev_actions * self.action_dim + if self.use_n_prev_rewards: + self.num_outputs += self.use_n_prev_rewards + + cfg = model_config + + self.attention_dim = cfg["attention_dim"] + + if self.num_outputs is not None: + in_space = gym.spaces.Box( + float("-inf"), float("inf"), shape=(self.num_outputs,), dtype=np.float32 + ) + else: + in_space = obs_space + + # Construct GTrXL sub-module w/ num_outputs=None (so it does not + # create a logits/value output; we'll do this ourselves in this wrapper + # here). + self.gtrxl = GTrXLNet( + in_space, + action_space, + None, + model_config, + "gtrxl", + num_transformer_units=cfg["attention_num_transformer_units"], + attention_dim=self.attention_dim, + num_heads=cfg["attention_num_heads"], + head_dim=cfg["attention_head_dim"], + memory_inference=cfg["attention_memory_inference"], + memory_training=cfg["attention_memory_training"], + position_wise_mlp_dim=cfg["attention_position_wise_mlp_dim"], + init_gru_gate_bias=cfg["attention_init_gru_gate_bias"], + ) + + # Set final num_outputs to correct value (depending on action space). + self.num_outputs = num_outputs + + # Postprocess GTrXL output with another hidden layer and compute + # values. + self._logits_branch = SlimFC( + in_size=self.attention_dim, + out_size=self.num_outputs, + activation_fn=None, + initializer=torch.nn.init.xavier_uniform_, + ) + self._value_branch = SlimFC( + in_size=self.attention_dim, + out_size=1, + activation_fn=None, + initializer=torch.nn.init.xavier_uniform_, + ) + + self.view_requirements = self.gtrxl.view_requirements + self.view_requirements["obs"].space = self.obs_space + + # Add prev-a/r to this model's view, if required. + if self.use_n_prev_actions: + self.view_requirements[SampleBatch.PREV_ACTIONS] = ViewRequirement( + SampleBatch.ACTIONS, + space=self.action_space, + shift="-{}:-1".format(self.use_n_prev_actions), + ) + if self.use_n_prev_rewards: + self.view_requirements[SampleBatch.PREV_REWARDS] = ViewRequirement( + SampleBatch.REWARDS, shift="-{}:-1".format(self.use_n_prev_rewards) + ) + + @override(RecurrentNetwork) + def forward( + self, + input_dict: Dict[str, TensorType], + state: List[TensorType], + seq_lens: TensorType, + ) -> (TensorType, List[TensorType]): + assert seq_lens is not None + # Push obs through "unwrapped" net's `forward()` first. + wrapped_out, _ = self._wrapped_forward(input_dict, [], None) + + # Concat. prev-action/reward if required. + prev_a_r = [] + + # Prev actions. + if self.use_n_prev_actions: + prev_n_actions = input_dict[SampleBatch.PREV_ACTIONS] + # If actions are not processed yet (in their original form as + # have been sent to environment): + # Flatten/one-hot into 1D array. + if self.model_config["_disable_action_flattening"]: + # Merge prev n actions into flat tensor. + flat = flatten_inputs_to_1d_tensor( + prev_n_actions, + spaces_struct=self.action_space_struct, + time_axis=True, + ) + # Fold time-axis into flattened data. + flat = torch.reshape(flat, [flat.shape[0], -1]) + prev_a_r.append(flat) + # If actions are already flattened (but not one-hot'd yet!), + # one-hot discrete/multi-discrete actions here and concatenate the + # n most recent actions together. + else: + if isinstance(self.action_space, Discrete): + for i in range(self.use_n_prev_actions): + prev_a_r.append( + one_hot( + prev_n_actions[:, i].float(), space=self.action_space + ) + ) + elif isinstance(self.action_space, MultiDiscrete): + for i in range( + 0, self.use_n_prev_actions, self.action_space.shape[0] + ): + prev_a_r.append( + one_hot( + prev_n_actions[ + :, i : i + self.action_space.shape[0] + ].float(), + space=self.action_space, + ) + ) + else: + prev_a_r.append( + torch.reshape( + prev_n_actions.float(), + [-1, self.use_n_prev_actions * self.action_dim], + ) + ) + # Prev rewards. + if self.use_n_prev_rewards: + prev_a_r.append( + torch.reshape( + input_dict[SampleBatch.PREV_REWARDS].float(), + [-1, self.use_n_prev_rewards], + ) + ) + + # Concat prev. actions + rewards to the "main" input. + if prev_a_r: + wrapped_out = torch.cat([wrapped_out] + prev_a_r, dim=1) + + # Then through our GTrXL. + input_dict["obs_flat"] = input_dict["obs"] = wrapped_out + + self._features, memory_outs = self.gtrxl(input_dict, state, seq_lens) + model_out = self._logits_branch(self._features) + return model_out, memory_outs + + @override(ModelV2) + def get_initial_state(self) -> Union[List[np.ndarray], List[TensorType]]: + return [ + torch.zeros( + self.gtrxl.view_requirements["state_in_{}".format(i)].space.shape + ) + for i in range(self.gtrxl.num_transformer_units) + ] + + @override(ModelV2) + def value_function(self) -> TensorType: + assert self._features is not None, "Must call forward() first!" + return torch.reshape(self._value_branch(self._features), [-1]) diff --git a/lib/python3.12/site-packages/ray/rllib/models/torch/complex_input_net.py b/lib/python3.12/site-packages/ray/rllib/models/torch/complex_input_net.py new file mode 100644 index 0000000000000000000000000000000000000000..0cf51f3e1862ca255ad07b8026be7cfcd1fca98b --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/torch/complex_input_net.py @@ -0,0 +1,237 @@ +import numpy as np +import tree # pip install dm_tree +from gymnasium.spaces import Box, Discrete, MultiDiscrete + +from ray.rllib.models.catalog import ModelCatalog +from ray.rllib.models.modelv2 import ModelV2, restore_original_dimensions +from ray.rllib.models.torch.misc import ( + SlimFC, + normc_initializer as torch_normc_initializer, +) +from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 +from ray.rllib.models.utils import get_filter_config +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.spaces.space_utils import flatten_space +from ray.rllib.utils.torch_utils import one_hot + +torch, nn = try_import_torch() + + +@OldAPIStack +class ComplexInputNetwork(TorchModelV2, nn.Module): + """TorchModelV2 concat'ing CNN outputs to flat input(s), followed by FC(s). + + Note: This model should be used for complex (Dict or Tuple) observation + spaces that have one or more image components. + + The data flow is as follows: + + `obs` (e.g. Tuple[img0, img1, discrete0]) -> `CNN0 + CNN1 + ONE-HOT` + `CNN0 + CNN1 + ONE-HOT` -> concat all flat outputs -> `out` + `out` -> (optional) FC-stack -> `out2` + `out2` -> action (logits) and value heads. + """ + + def __init__(self, obs_space, action_space, num_outputs, model_config, name): + self.original_space = ( + obs_space.original_space + if hasattr(obs_space, "original_space") + else obs_space + ) + + self.processed_obs_space = ( + self.original_space + if model_config.get("_disable_preprocessor_api") + else obs_space + ) + + nn.Module.__init__(self) + TorchModelV2.__init__( + self, self.original_space, action_space, num_outputs, model_config, name + ) + + self.flattened_input_space = flatten_space(self.original_space) + + # Atari type CNNs or IMPALA type CNNs (with residual layers)? + # self.cnn_type = self.model_config["custom_model_config"].get( + # "conv_type", "atari") + + # Build the CNN(s) given obs_space's image components. + self.cnns = nn.ModuleDict() + self.one_hot = nn.ModuleDict() + self.flatten_dims = {} + self.flatten = nn.ModuleDict() + concat_size = 0 + for i, component in enumerate(self.flattened_input_space): + i = str(i) + # Image space. + if len(component.shape) == 3 and isinstance(component, Box): + config = { + "conv_filters": model_config["conv_filters"] + if "conv_filters" in model_config + else get_filter_config(component.shape), + "conv_activation": model_config.get("conv_activation"), + "post_fcnet_hiddens": [], + } + # if self.cnn_type == "atari": + self.cnns[i] = ModelCatalog.get_model_v2( + component, + action_space, + num_outputs=None, + model_config=config, + framework="torch", + name="cnn_{}".format(i), + ) + # TODO (sven): add IMPALA-style option. + # else: + # cnn = TorchImpalaVisionNet( + # component, + # action_space, + # num_outputs=None, + # model_config=config, + # name="cnn_{}".format(i)) + + concat_size += self.cnns[i].num_outputs + self.add_module("cnn_{}".format(i), self.cnns[i]) + # Discrete|MultiDiscrete inputs -> One-hot encode. + elif isinstance(component, (Discrete, MultiDiscrete)): + if isinstance(component, Discrete): + size = component.n + else: + size = np.sum(component.nvec) + config = { + "fcnet_hiddens": model_config["fcnet_hiddens"], + "fcnet_activation": model_config.get("fcnet_activation"), + "post_fcnet_hiddens": [], + } + self.one_hot[i] = ModelCatalog.get_model_v2( + Box(-1.0, 1.0, (size,), np.float32), + action_space, + num_outputs=None, + model_config=config, + framework="torch", + name="one_hot_{}".format(i), + ) + concat_size += self.one_hot[i].num_outputs + self.add_module("one_hot_{}".format(i), self.one_hot[i]) + # Everything else (1D Box). + else: + size = int(np.prod(component.shape)) + config = { + "fcnet_hiddens": model_config["fcnet_hiddens"], + "fcnet_activation": model_config.get("fcnet_activation"), + "post_fcnet_hiddens": [], + } + self.flatten[i] = ModelCatalog.get_model_v2( + Box(-1.0, 1.0, (size,), np.float32), + action_space, + num_outputs=None, + model_config=config, + framework="torch", + name="flatten_{}".format(i), + ) + self.flatten_dims[i] = size + concat_size += self.flatten[i].num_outputs + self.add_module("flatten_{}".format(i), self.flatten[i]) + + # Optional post-concat FC-stack. + post_fc_stack_config = { + "fcnet_hiddens": model_config.get("post_fcnet_hiddens", []), + "fcnet_activation": model_config.get("post_fcnet_activation", "relu"), + } + self.post_fc_stack = ModelCatalog.get_model_v2( + Box(float("-inf"), float("inf"), shape=(concat_size,), dtype=np.float32), + self.action_space, + None, + post_fc_stack_config, + framework="torch", + name="post_fc_stack", + ) + + # Actions and value heads. + self.logits_layer = None + self.value_layer = None + self._value_out = None + + if num_outputs: + # Action-distribution head. + self.logits_layer = SlimFC( + in_size=self.post_fc_stack.num_outputs, + out_size=num_outputs, + activation_fn=None, + initializer=torch_normc_initializer(0.01), + ) + # Create the value branch model. + self.value_layer = SlimFC( + in_size=self.post_fc_stack.num_outputs, + out_size=1, + activation_fn=None, + initializer=torch_normc_initializer(0.01), + ) + else: + self.num_outputs = concat_size + + @override(ModelV2) + def forward(self, input_dict, state, seq_lens): + if SampleBatch.OBS in input_dict and "obs_flat" in input_dict: + orig_obs = input_dict[SampleBatch.OBS] + else: + orig_obs = restore_original_dimensions( + input_dict[SampleBatch.OBS], self.processed_obs_space, tensorlib="torch" + ) + # Push observations through the different components + # (CNNs, one-hot + FC, etc..). + outs = [] + for i, component in enumerate(tree.flatten(orig_obs)): + i = str(i) + if i in self.cnns: + cnn_out, _ = self.cnns[i](SampleBatch({SampleBatch.OBS: component})) + outs.append(cnn_out) + elif i in self.one_hot: + if component.dtype in [ + torch.int8, + torch.int16, + torch.int32, + torch.int64, + torch.uint8, + ]: + one_hot_in = { + SampleBatch.OBS: one_hot( + component, self.flattened_input_space[int(i)] + ) + } + else: + one_hot_in = {SampleBatch.OBS: component} + one_hot_out, _ = self.one_hot[i](SampleBatch(one_hot_in)) + outs.append(one_hot_out) + else: + nn_out, _ = self.flatten[i]( + SampleBatch( + { + SampleBatch.OBS: torch.reshape( + component, [-1, self.flatten_dims[i]] + ) + } + ) + ) + outs.append(nn_out) + + # Concat all outputs and the non-image inputs. + out = torch.cat(outs, dim=1) + # Push through (optional) FC-stack (this may be an empty stack). + out, _ = self.post_fc_stack(SampleBatch({SampleBatch.OBS: out})) + + # No logits/value branches. + if self.logits_layer is None: + return out, [] + + # Logits- and value branches. + logits, values = self.logits_layer(out), self.value_layer(out) + self._value_out = torch.reshape(values, [-1]) + return logits, [] + + @override(ModelV2) + def value_function(self): + return self._value_out diff --git a/lib/python3.12/site-packages/ray/rllib/models/torch/fcnet.py b/lib/python3.12/site-packages/ray/rllib/models/torch/fcnet.py new file mode 100644 index 0000000000000000000000000000000000000000..e4e111a9b956d6d8184e9edc80b6d94f83775c62 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/torch/fcnet.py @@ -0,0 +1,161 @@ +import logging + +import gymnasium as gym +import numpy as np + +from ray.rllib.models.torch.misc import AppendBiasLayer, SlimFC, normc_initializer +from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.typing import Dict, List, ModelConfigDict, TensorType + +torch, nn = try_import_torch() + +logger = logging.getLogger(__name__) + + +@OldAPIStack +class FullyConnectedNetwork(TorchModelV2, nn.Module): + """Generic fully connected network.""" + + def __init__( + self, + obs_space: gym.spaces.Space, + action_space: gym.spaces.Space, + num_outputs: int, + model_config: ModelConfigDict, + name: str, + ): + TorchModelV2.__init__( + self, obs_space, action_space, num_outputs, model_config, name + ) + nn.Module.__init__(self) + + hiddens = list(model_config.get("fcnet_hiddens", [])) + list( + model_config.get("post_fcnet_hiddens", []) + ) + activation = model_config.get("fcnet_activation") + if not model_config.get("fcnet_hiddens", []): + activation = model_config.get("post_fcnet_activation") + no_final_linear = model_config.get("no_final_linear") + self.vf_share_layers = model_config.get("vf_share_layers") + self.free_log_std = model_config.get("free_log_std") + # Generate free-floating bias variables for the second half of + # the outputs. + if self.free_log_std: + assert num_outputs % 2 == 0, ( + "num_outputs must be divisible by two", + num_outputs, + ) + num_outputs = num_outputs // 2 + + layers = [] + prev_layer_size = int(np.prod(obs_space.shape)) + self._logits = None + + # Create layers 0 to second-last. + for size in hiddens[:-1]: + layers.append( + SlimFC( + in_size=prev_layer_size, + out_size=size, + initializer=normc_initializer(1.0), + activation_fn=activation, + ) + ) + prev_layer_size = size + + # The last layer is adjusted to be of size num_outputs, but it's a + # layer with activation. + if no_final_linear and num_outputs: + layers.append( + SlimFC( + in_size=prev_layer_size, + out_size=num_outputs, + initializer=normc_initializer(1.0), + activation_fn=activation, + ) + ) + prev_layer_size = num_outputs + # Finish the layers with the provided sizes (`hiddens`), plus - + # iff num_outputs > 0 - a last linear layer of size num_outputs. + else: + if len(hiddens) > 0: + layers.append( + SlimFC( + in_size=prev_layer_size, + out_size=hiddens[-1], + initializer=normc_initializer(1.0), + activation_fn=activation, + ) + ) + prev_layer_size = hiddens[-1] + if num_outputs: + self._logits = SlimFC( + in_size=prev_layer_size, + out_size=num_outputs, + initializer=normc_initializer(0.01), + activation_fn=None, + ) + else: + self.num_outputs = ([int(np.prod(obs_space.shape))] + hiddens[-1:])[-1] + + # Layer to add the log std vars to the state-dependent means. + if self.free_log_std and self._logits: + self._append_free_log_std = AppendBiasLayer(num_outputs) + + self._hidden_layers = nn.Sequential(*layers) + + self._value_branch_separate = None + if not self.vf_share_layers: + # Build a parallel set of hidden layers for the value net. + prev_vf_layer_size = int(np.prod(obs_space.shape)) + vf_layers = [] + for size in hiddens: + vf_layers.append( + SlimFC( + in_size=prev_vf_layer_size, + out_size=size, + activation_fn=activation, + initializer=normc_initializer(1.0), + ) + ) + prev_vf_layer_size = size + self._value_branch_separate = nn.Sequential(*vf_layers) + + self._value_branch = SlimFC( + in_size=prev_layer_size, + out_size=1, + initializer=normc_initializer(0.01), + activation_fn=None, + ) + # Holds the current "base" output (before logits layer). + self._features = None + # Holds the last input, in case value branch is separate. + self._last_flat_in = None + + @override(TorchModelV2) + def forward( + self, + input_dict: Dict[str, TensorType], + state: List[TensorType], + seq_lens: TensorType, + ) -> (TensorType, List[TensorType]): + obs = input_dict["obs_flat"].float() + self._last_flat_in = obs.reshape(obs.shape[0], -1) + self._features = self._hidden_layers(self._last_flat_in) + logits = self._logits(self._features) if self._logits else self._features + if self.free_log_std: + logits = self._append_free_log_std(logits) + return logits, state + + @override(TorchModelV2) + def value_function(self) -> TensorType: + assert self._features is not None, "must call forward() first" + if self._value_branch_separate: + out = self._value_branch( + self._value_branch_separate(self._last_flat_in) + ).squeeze(1) + else: + out = self._value_branch(self._features).squeeze(1) + return out diff --git a/lib/python3.12/site-packages/ray/rllib/models/torch/mingpt.py b/lib/python3.12/site-packages/ray/rllib/models/torch/mingpt.py new file mode 100644 index 0000000000000000000000000000000000000000..303bcc160daa59633c450ba5ac73b5f7102f1f7b --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/torch/mingpt.py @@ -0,0 +1,303 @@ +# LICENSE: MIT +""" +Adapted from https://github.com/karpathy/minGPT + +Full definition of a GPT Language Model, all of it in this single file. +References: +1) the official GPT-2 TensorFlow implementation released by OpenAI: +https://github.com/openai/gpt-2/blob/master/src/model.py +2) huggingface/transformers PyTorch implementation: +https://github.com/huggingface/transformers/blob/main/src/transformers + /models/gpt2/modeling_gpt2.py +""" + +import math +from dataclasses import dataclass +from typing import Tuple + +import torch +import torch.nn as nn +from torch.nn import functional as F + +from ray._common.deprecation import Deprecated +from ray.rllib.utils.annotations import DeveloperAPI + + +@DeveloperAPI +@dataclass +class GPTConfig: + # block size must be provided + block_size: int + + # transformer config + n_layer: int = 12 + n_head: int = 12 + n_embed: int = 768 + + # dropout config + embed_pdrop: float = 0.1 + resid_pdrop: float = 0.1 + attn_pdrop: float = 0.1 + + +@Deprecated(error=False) +class NewGELU(nn.Module): + """ + Implementation of the GELU activation function currently in Google BERT + repo (identical to OpenAI GPT). + Reference: Gaussian Error Linear Units (GELU) paper: + https://arxiv.org/abs/1606.08415 + """ + + def forward(self, x): + return ( + 0.5 + * x + * ( + 1.0 + + torch.tanh( + math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0)) + ) + ) + ) + + +@Deprecated(error=False) +class CausalSelfAttention(nn.Module): + """ + Vanilla multi-head masked self-attention layer with a projection at the end. + It is possible to use torch.nn.MultiheadAttention here but I am including an + explicit implementation here to show that there is nothing too scary here. + """ + + def __init__(self, config: GPTConfig): + super().__init__() + assert config.n_embed % config.n_head == 0 + # key, query, value projections for all heads, but in a batch + self.c_attn = nn.Linear(config.n_embed, 3 * config.n_embed) + # output projection + self.c_proj = nn.Linear(config.n_embed, config.n_embed) + # regularization + self.attn_dropout = nn.Dropout(config.attn_pdrop) + self.resid_dropout = nn.Dropout(config.resid_pdrop) + # causal mask to ensure that attention is only applied to the left + # in the input sequence + self.register_buffer( + "bias", + torch.tril(torch.ones(config.block_size, config.block_size)).view( + 1, 1, config.block_size, config.block_size + ), + ) + self.n_head = config.n_head + self.n_embed = config.n_embed + + def forward(self, x, attention_masks=None): + # batch size, sequence length, embedding dimensionality (n_embed) + B, T, C = x.size() + + # calculate query, key, values for all heads in batch and move head + # forward to be the batch dim + q, k, v = self.c_attn(x).split(self.n_embed, dim=2) + # (B, nh, T, hs) + k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) + # (B, nh, T, hs) + q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) + # (B, nh, T, hs) + v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) + + # causal self-attention; Self-attend: + # (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T) + att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1))) + att = att.masked_fill(self.bias[:, :, :T, :T] == 0, float("-inf")) + if attention_masks is not None: + att = att + attention_masks + att = F.softmax(att, dim=-1) + att = self.attn_dropout(att) + y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs) + # re-assemble all head outputs side by side + y = y.transpose(1, 2).contiguous().view(B, T, C) + + # output projection + y = self.resid_dropout(self.c_proj(y)) + return y, att + + +@Deprecated(error=False) +class Block(nn.Module): + """an unassuming Transformer block""" + + def __init__(self, config: GPTConfig): + super().__init__() + self.ln_1 = nn.LayerNorm(config.n_embed) + self.attn = CausalSelfAttention(config) + self.ln_2 = nn.LayerNorm(config.n_embed) + self.mlp = nn.ModuleDict( + dict( + c_fc=nn.Linear(config.n_embed, 4 * config.n_embed), + c_proj=nn.Linear(4 * config.n_embed, config.n_embed), + act=NewGELU(), + dropout=nn.Dropout(config.resid_pdrop), + ) + ) + + def forward(self, x, attention_masks=None): + # Multi-head attention sub-layer. + x_att, att = self.attn(self.ln_1(x), attention_masks=attention_masks) + # Residual of multi-head attention sub-layer. + x = x + x_att + + # Position-wise FFN sub-layer: fc + activation + fc + dropout + x_ffn = self.mlp.dropout(self.mlp.c_proj(self.mlp.act(self.mlp.c_fc(x)))) + # Residual of position-wise FFN sub-layer. + x = x + x_ffn + return x, att + + +@Deprecated(error=False) +def configure_gpt_optimizer( + model: nn.Module, + learning_rate: float, + weight_decay: float, + betas: Tuple[float, float] = (0.9, 0.95), + **kwargs, +) -> torch.optim.Optimizer: + """ + This long function is unfortunately doing something very simple and is + being very defensive: We are separating out all parameters of the model + into two buckets: those that will experience weight decay for regularization + and those that won't (biases, and layernorm/embedding weights). We are then + returning the PyTorch optimizer object. + """ + + # separate out all parameters to those that will and won't experience + # regularizing weight decay + decay = set() + no_decay = set() + whitelist_w_modules = (torch.nn.Linear,) + blacklist_w_modules = (torch.nn.LayerNorm, torch.nn.Embedding) + for mn, m in model.named_modules(): + for pn, p in m.named_parameters(): + fpn = "%s.%s" % (mn, pn) if mn else pn # full param name + # random note: because named_modules and named_parameters are + # recursive we will see the same tensors p many many times. but + # doing it this way allows us to know which parent module any + # tensor p belongs to... + if pn.endswith("bias"): + # all biases will not be decayed + no_decay.add(fpn) + elif pn.endswith("weight") and isinstance(m, whitelist_w_modules): + # weights of whitelist modules will be weight decayed + decay.add(fpn) + elif pn.endswith("weight") and isinstance(m, blacklist_w_modules): + # weights of blacklist modules will NOT be weight decayed + no_decay.add(fpn) + + # validate that we considered every parameter + param_dict = dict(model.named_parameters()) + inter_params = decay & no_decay + union_params = decay | no_decay + assert ( + len(inter_params) == 0 + ), f"parameters {str(inter_params)} made it into both decay/no_decay sets!" + assert len(param_dict.keys() - union_params) == 0, ( + f"parameters {str(param_dict.keys() - union_params)} were not " + f"separated into either decay/no_decay set!" + ) + + # create the pytorch optimizer object + optim_groups = [ + { + "params": [param_dict[pn] for pn in sorted(decay)], + "weight_decay": weight_decay, + }, + { + "params": [param_dict[pn] for pn in sorted(no_decay)], + "weight_decay": 0.0, + }, + ] + optimizer = torch.optim.AdamW(optim_groups, lr=learning_rate, betas=betas, **kwargs) + return optimizer + + +@Deprecated(error=False) +class GPT(nn.Module): + """GPT Transformer Model""" + + def __init__(self, config: GPTConfig): + super().__init__() + assert config.block_size is not None + self.block_size = config.block_size + + self.transformer = nn.ModuleDict( + dict( + drop=nn.Dropout(config.embed_pdrop), + h=nn.ModuleList([Block(config) for _ in range(config.n_layer)]), + ln_f=nn.LayerNorm(config.n_embed), + ) + ) + + # init all weights, and apply a special scaled init to the residual + # projections, per GPT-2 paper + self.apply(self._init_weights) + for pn, p in self.named_parameters(): + if pn.endswith("c_proj.weight"): + torch.nn.init.normal_( + p, mean=0.0, std=0.02 / math.sqrt(2 * config.n_layer) + ) + + def _init_weights(self, module): + if isinstance(module, nn.Linear): + torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) + if module.bias is not None: + torch.nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) + elif isinstance(module, nn.LayerNorm): + torch.nn.init.zeros_(module.bias) + torch.nn.init.ones_(module.weight) + + def forward(self, input_embeds, attention_masks=None, return_attentions=False): + """ + input_embeds: [batch_size x seq_len x n_embed] + attention_masks: [batch_size x seq_len], 0 don't attend, 1 attend + """ + B, T, C = input_embeds.size() + assert T <= self.block_size, ( + f"Cannot forward sequence of length {T}, " + f"block size is only {self.block_size}" + ) + + if attention_masks is not None: + _B, _T = attention_masks.size() + assert _B == B and _T == T + # We create a 3D attention mask from a 2D tensor mask. + # Sizes are [batch_size, 1, 1, to_seq_len] + # So we can broadcast to + # [batch_size, num_heads, from_seq_length, to_seq_length] + # this attention mask is more simple than the triangular + # masking of causal attention used in OpenAI GPT, we just need + # to prepare the broadcast dimension here. + attention_masks = attention_masks[:, None, None, :] + + # Since attention_mask is 1.0 for positions we want to attend + # and 0.0 for masked positions, this operation will create a + # tensor which is 0.0 for positions we want to attend and -inf + # for masked positions. Since we are adding it to the raw scores + # before the softmax, this is effectively the same as removing + # these entirely. + attention_masks = attention_masks.to(dtype=input_embeds.dtype) + attention_masks = (1.0 - attention_masks) * -1e9 + + # forward the GPT model itself + x = self.transformer.drop(input_embeds) + + atts = [] + for block in self.transformer.h: + x, att = block(x, attention_masks=attention_masks) + atts.append(att) + x = self.transformer.ln_f(x) + + if return_attentions: + return x, atts + else: + return x diff --git a/lib/python3.12/site-packages/ray/rllib/models/torch/misc.py b/lib/python3.12/site-packages/ray/rllib/models/torch/misc.py new file mode 100644 index 0000000000000000000000000000000000000000..d9c5e91be2da75b530d73c5ffcdb3d72159a572a --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/torch/misc.py @@ -0,0 +1,324 @@ +""" Code adapted from https://github.com/ikostrikov/pytorch-a3c""" +from typing import Any, List, Tuple, Union + +import numpy as np + +from ray.rllib.models.utils import get_activation_fn +from ray.rllib.utils.annotations import DeveloperAPI +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.typing import TensorType + +torch, nn = try_import_torch() + + +@DeveloperAPI +def normc_initializer(std: float = 1.0) -> Any: + def initializer(tensor): + tensor.data.normal_(0, 1) + tensor.data *= std / torch.sqrt(tensor.data.pow(2).sum(1, keepdim=True)) + + return initializer + + +@DeveloperAPI +def same_padding( + in_size: Tuple[int, int], + filter_size: Union[int, Tuple[int, int]], + stride_size: Union[int, Tuple[int, int]], +) -> (Union[int, Tuple[int, int]], Tuple[int, int]): + """Note: Padding is added to match TF conv2d `same` padding. + + See www.tensorflow.org/versions/r0.12/api_docs/python/nn/convolution + + Args: + in_size: Rows (Height), Column (Width) for input + stride_size (Union[int,Tuple[int, int]]): Rows (Height), column (Width) + for stride. If int, height == width. + filter_size: Rows (Height), column (Width) for filter + + Returns: + padding: For input into torch.nn.ZeroPad2d. + output: Output shape after padding and convolution. + """ + in_height, in_width = in_size + if isinstance(filter_size, int): + filter_height, filter_width = filter_size, filter_size + else: + filter_height, filter_width = filter_size + if isinstance(stride_size, (int, float)): + stride_height, stride_width = int(stride_size), int(stride_size) + else: + stride_height, stride_width = int(stride_size[0]), int(stride_size[1]) + + out_height = int(np.ceil(float(in_height) / float(stride_height))) + out_width = int(np.ceil(float(in_width) / float(stride_width))) + + pad_along_height = int((out_height - 1) * stride_height + filter_height - in_height) + pad_along_width = int((out_width - 1) * stride_width + filter_width - in_width) + pad_top = pad_along_height // 2 + pad_bottom = pad_along_height - pad_top + pad_left = pad_along_width // 2 + pad_right = pad_along_width - pad_left + padding = (pad_left, pad_right, pad_top, pad_bottom) + output = (out_height, out_width) + return padding, output + + +@DeveloperAPI +def same_padding_transpose_after_stride( + strided_size: Tuple[int, int], + kernel: Tuple[int, int], + stride: Union[int, Tuple[int, int]], +) -> (Union[int, Tuple[int, int]], Tuple[int, int]): + """Computes padding and output size such that TF Conv2DTranspose `same` is matched. + + Note that when padding="same", TensorFlow's Conv2DTranspose makes sure that + 0-padding is added to the already strided image in such a way that the output image + has the same size as the input image times the stride (and no matter the + kernel size). + + For example: Input image is (4, 4, 24) (not yet strided), padding is "same", + stride=2, kernel=5. + + First, the input image is strided (with stride=2): + + Input image (4x4): + A B C D + E F G H + I J K L + M N O P + + Stride with stride=2 -> (7x7) + A 0 B 0 C 0 D + 0 0 0 0 0 0 0 + E 0 F 0 G 0 H + 0 0 0 0 0 0 0 + I 0 J 0 K 0 L + 0 0 0 0 0 0 0 + M 0 N 0 O 0 P + + Then this strided image (strided_size=7x7) is padded (exact padding values will be + output by this function): + + padding -> (left=3, right=2, top=3, bottom=2) + + 0 0 0 0 0 0 0 0 0 0 0 0 + 0 0 0 0 0 0 0 0 0 0 0 0 + 0 0 0 0 0 0 0 0 0 0 0 0 + 0 0 0 A 0 B 0 C 0 D 0 0 + 0 0 0 0 0 0 0 0 0 0 0 0 + 0 0 0 E 0 F 0 G 0 H 0 0 + 0 0 0 0 0 0 0 0 0 0 0 0 + 0 0 0 I 0 J 0 K 0 L 0 0 + 0 0 0 0 0 0 0 0 0 0 0 0 + 0 0 0 M 0 N 0 O 0 P 0 0 + 0 0 0 0 0 0 0 0 0 0 0 0 + 0 0 0 0 0 0 0 0 0 0 0 0 + + Then deconvolution with kernel=5 yields an output image of 8x8 (x num output + filters). + + Args: + strided_size: The size (width x height) of the already strided image. + kernel: Either width x height (tuple of ints) or - if a square kernel is used - + a single int for both width and height. + stride: Either stride width x stride height (tuple of ints) or - if square + striding is used - a single int for both width- and height striding. + + Returns: + Tuple consisting of 1) `padding`: A 4-tuple to pad the input after(!) striding. + The values are for left, right, top, and bottom padding, individually. + This 4-tuple can be used in a torch.nn.ZeroPad2d layer, and 2) the output shape + after striding, padding, and the conv transpose layer. + """ + + # Solve single int (squared) inputs for kernel and/or stride. + k_w, k_h = (kernel, kernel) if isinstance(kernel, int) else kernel + s_w, s_h = (stride, stride) if isinstance(stride, int) else stride + + # Compute the total size of the 0-padding on both axes. If results are odd numbers, + # the padding on e.g. left and right (or top and bottom) side will have to differ + # by 1. + pad_total_w, pad_total_h = k_w - 1 + s_w - 1, k_h - 1 + s_h - 1 + pad_right = pad_total_w // 2 + pad_left = pad_right + (1 if pad_total_w % 2 == 1 else 0) + pad_bottom = pad_total_h // 2 + pad_top = pad_bottom + (1 if pad_total_h % 2 == 1 else 0) + + # Compute the output size. + output_shape = ( + strided_size[0] + pad_total_w - k_w + 1, + strided_size[1] + pad_total_h - k_h + 1, + ) + + # Return padding and output shape. + return (pad_left, pad_right, pad_top, pad_bottom), output_shape + + +@DeveloperAPI +def valid_padding( + in_size: Tuple[int, int], + filter_size: Union[int, Tuple[int, int]], + stride_size: Union[int, Tuple[int, int]], +) -> Tuple[int, int]: + """Emulates TF Conv2DLayer "valid" padding (no padding) and computes output dims. + + This method, analogous to its "same" counterpart, but it only computes the output + image size, since valid padding means (0, 0, 0, 0). + + See www.tensorflow.org/versions/r0.12/api_docs/python/nn/convolution + + Args: + in_size: Rows (Height), Column (Width) for input + stride_size (Union[int,Tuple[int, int]]): Rows (Height), column (Width) + for stride. If int, height == width. + filter_size: Rows (Height), column (Width) for filter + + Returns: + The output shape after padding and convolution. + """ + in_height, in_width = in_size + if isinstance(filter_size, int): + filter_height, filter_width = filter_size, filter_size + else: + filter_height, filter_width = filter_size + if isinstance(stride_size, (int, float)): + stride_height, stride_width = int(stride_size), int(stride_size) + else: + stride_height, stride_width = int(stride_size[0]), int(stride_size[1]) + + out_height = int(np.ceil((in_height - filter_height + 1) / float(stride_height))) + out_width = int(np.ceil((in_width - filter_width + 1) / float(stride_width))) + + return (out_height, out_width) + + +@DeveloperAPI +class SlimConv2d(nn.Module): + """Simple mock of tf.slim Conv2d""" + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel: Union[int, Tuple[int, int]], + stride: Union[int, Tuple[int, int]], + padding: Union[int, Tuple[int, int]], + # Defaulting these to nn.[..] will break soft torch import. + initializer: Any = "default", + activation_fn: Any = "default", + bias_init: float = 0, + ): + """Creates a standard Conv2d layer, similar to torch.nn.Conv2d + + Args: + in_channels: Number of input channels + out_channels: Number of output channels + kernel: If int, the kernel is + a tuple(x,x). Elsewise, the tuple can be specified + stride: Controls the stride + for the cross-correlation. If int, the stride is a + tuple(x,x). Elsewise, the tuple can be specified + padding: Controls the amount + of implicit zero-paddings during the conv operation + initializer: Initializer function for kernel weights + activation_fn: Activation function at the end of layer + bias_init: Initalize bias weights to bias_init const + """ + super(SlimConv2d, self).__init__() + layers = [] + # Padding layer. + if padding: + layers.append(nn.ZeroPad2d(padding)) + # Actual Conv2D layer (including correct initialization logic). + conv = nn.Conv2d(in_channels, out_channels, kernel, stride) + if initializer: + if initializer == "default": + initializer = nn.init.xavier_uniform_ + initializer(conv.weight) + nn.init.constant_(conv.bias, bias_init) + layers.append(conv) + # Activation function (if any; default=ReLu). + if isinstance(activation_fn, str): + if activation_fn == "default": + activation_fn = nn.ReLU + else: + activation_fn = get_activation_fn(activation_fn, "torch") + if activation_fn is not None: + layers.append(activation_fn()) + # Put everything in sequence. + self._model = nn.Sequential(*layers) + + def forward(self, x: TensorType) -> TensorType: + return self._model(x) + + +@DeveloperAPI +class SlimFC(nn.Module): + """Simple PyTorch version of `linear` function""" + + def __init__( + self, + in_size: int, + out_size: int, + initializer: Any = None, + activation_fn: Any = None, + use_bias: bool = True, + bias_init: float = 0.0, + ): + """Creates a standard FC layer, similar to torch.nn.Linear + + Args: + in_size: Input size for FC Layer + out_size: Output size for FC Layer + initializer: Initializer function for FC layer weights + activation_fn: Activation function at the end of layer + use_bias: Whether to add bias weights or not + bias_init: Initalize bias weights to bias_init const + """ + super(SlimFC, self).__init__() + layers = [] + # Actual nn.Linear layer (including correct initialization logic). + linear = nn.Linear(in_size, out_size, bias=use_bias) + if initializer is None: + initializer = nn.init.xavier_uniform_ + initializer(linear.weight) + if use_bias is True: + nn.init.constant_(linear.bias, bias_init) + layers.append(linear) + # Activation function (if any; default=None (linear)). + if isinstance(activation_fn, str): + activation_fn = get_activation_fn(activation_fn, "torch") + if activation_fn is not None: + layers.append(activation_fn()) + # Put everything in sequence. + self._model = nn.Sequential(*layers) + + def forward(self, x: TensorType) -> TensorType: + return self._model(x) + + +@DeveloperAPI +class AppendBiasLayer(nn.Module): + """Simple bias appending layer for free_log_std.""" + + def __init__(self, num_bias_vars: int): + super().__init__() + self.log_std = torch.nn.Parameter(torch.as_tensor([0.0] * num_bias_vars)) + self.register_parameter("log_std", self.log_std) + + def forward(self, x: TensorType) -> TensorType: + out = torch.cat([x, self.log_std.unsqueeze(0).repeat([len(x), 1])], axis=1) + return out + + +@DeveloperAPI +class Reshape(nn.Module): + """Standard module that reshapes/views a tensor""" + + def __init__(self, shape: List): + super().__init__() + self.shape = shape + + def forward(self, x): + return x.view(*self.shape) diff --git a/lib/python3.12/site-packages/ray/rllib/models/torch/modules/__init__.py b/lib/python3.12/site-packages/ray/rllib/models/torch/modules/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2585dcc77abe4acd4cd6daf49f902a4117a1438c --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/torch/modules/__init__.py @@ -0,0 +1,13 @@ +from ray.rllib.models.torch.modules.gru_gate import GRUGate +from ray.rllib.models.torch.modules.multi_head_attention import MultiHeadAttention +from ray.rllib.models.torch.modules.relative_multi_head_attention import ( + RelativeMultiHeadAttention, +) +from ray.rllib.models.torch.modules.skip_connection import SkipConnection + +__all__ = [ + "GRUGate", + "RelativeMultiHeadAttention", + "SkipConnection", + "MultiHeadAttention", +] diff --git a/lib/python3.12/site-packages/ray/rllib/models/torch/modules/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/torch/modules/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..86d8f896111c5cb1893d45f73db01b6b9eafbafd Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/torch/modules/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/torch/modules/__pycache__/gru_gate.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/torch/modules/__pycache__/gru_gate.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5e5afb5f91104eae274f70d8062422ab58f3ca9c Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/torch/modules/__pycache__/gru_gate.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/torch/modules/__pycache__/multi_head_attention.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/torch/modules/__pycache__/multi_head_attention.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a10058bc21a15f4114e121761b6023aa9b9399da Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/torch/modules/__pycache__/multi_head_attention.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/torch/modules/__pycache__/noisy_layer.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/torch/modules/__pycache__/noisy_layer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..46eae94991bb316e01bf5f04a15e3cb4382f7561 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/torch/modules/__pycache__/noisy_layer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/torch/modules/__pycache__/relative_multi_head_attention.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/torch/modules/__pycache__/relative_multi_head_attention.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..37151b5f97e6faa6453139abfb3f182ef7a7cc2d Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/torch/modules/__pycache__/relative_multi_head_attention.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/torch/modules/__pycache__/skip_connection.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/models/torch/modules/__pycache__/skip_connection.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f47eb248a2ebdf10b68bd9435690edf96b64e1fb Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/models/torch/modules/__pycache__/skip_connection.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/models/torch/modules/gru_gate.py b/lib/python3.12/site-packages/ray/rllib/models/torch/modules/gru_gate.py new file mode 100644 index 0000000000000000000000000000000000000000..7bb6eee845423898c7a553667bc18159cd9f2ff2 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/torch/modules/gru_gate.py @@ -0,0 +1,67 @@ +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.typing import TensorType + +torch, nn = try_import_torch() + + +@OldAPIStack +class GRUGate(nn.Module): + """Implements a gated recurrent unit for use in AttentionNet""" + + def __init__(self, dim: int, init_bias: int = 0.0, **kwargs): + """ + input_shape (torch.Tensor): dimension of the input + init_bias: Bias added to every input to stabilize training + """ + super().__init__(**kwargs) + # Xavier initialization of torch tensors + self._w_r = nn.Parameter(torch.zeros(dim, dim)) + self._w_z = nn.Parameter(torch.zeros(dim, dim)) + self._w_h = nn.Parameter(torch.zeros(dim, dim)) + nn.init.xavier_uniform_(self._w_r) + nn.init.xavier_uniform_(self._w_z) + nn.init.xavier_uniform_(self._w_h) + self.register_parameter("_w_r", self._w_r) + self.register_parameter("_w_z", self._w_z) + self.register_parameter("_w_h", self._w_h) + + self._u_r = nn.Parameter(torch.zeros(dim, dim)) + self._u_z = nn.Parameter(torch.zeros(dim, dim)) + self._u_h = nn.Parameter(torch.zeros(dim, dim)) + nn.init.xavier_uniform_(self._u_r) + nn.init.xavier_uniform_(self._u_z) + nn.init.xavier_uniform_(self._u_h) + self.register_parameter("_u_r", self._u_r) + self.register_parameter("_u_z", self._u_z) + self.register_parameter("_u_h", self._u_h) + + self._bias_z = nn.Parameter( + torch.zeros( + dim, + ).fill_(init_bias) + ) + self.register_parameter("_bias_z", self._bias_z) + + def forward(self, inputs: TensorType, **kwargs) -> TensorType: + # Pass in internal state first. + h, X = inputs + + r = torch.tensordot(X, self._w_r, dims=1) + torch.tensordot( + h, self._u_r, dims=1 + ) + r = torch.sigmoid(r) + + z = ( + torch.tensordot(X, self._w_z, dims=1) + + torch.tensordot(h, self._u_z, dims=1) + - self._bias_z + ) + z = torch.sigmoid(z) + + h_next = torch.tensordot(X, self._w_h, dims=1) + torch.tensordot( + (h * r), self._u_h, dims=1 + ) + h_next = torch.tanh(h_next) + + return (1 - z) * h + z * h_next diff --git a/lib/python3.12/site-packages/ray/rllib/models/torch/modules/multi_head_attention.py b/lib/python3.12/site-packages/ray/rllib/models/torch/modules/multi_head_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..a613decbe62faeaa9f20562b9e29bc0981791aea --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/torch/modules/multi_head_attention.py @@ -0,0 +1,69 @@ +""" +[1] - Attention Is All You Need - Vaswani, Jones, Shazeer, Parmar, + Uszkoreit, Gomez, Kaiser - Google Brain/Research, U Toronto - 2017. + https://arxiv.org/pdf/1706.03762.pdf +""" +from ray.rllib.models.torch.misc import SlimFC +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.framework import TensorType, try_import_torch +from ray.rllib.utils.torch_utils import sequence_mask + +torch, nn = try_import_torch() + + +@OldAPIStack +class MultiHeadAttention(nn.Module): + """A multi-head attention layer described in [1].""" + + def __init__( + self, in_dim: int, out_dim: int, num_heads: int, head_dim: int, **kwargs + ): + """ + in_dim: Dimension of input + out_dim: Dimension of output + num_heads: Number of attention heads + head_dim: Output dimension of each attention head + """ + super().__init__(**kwargs) + + # No bias or non-linearity. + self._num_heads = num_heads + self._head_dim = head_dim + self._qkv_layer = SlimFC( + in_size=in_dim, out_size=3 * num_heads * head_dim, use_bias=False + ) + + self._linear_layer = SlimFC( + in_size=num_heads * head_dim, out_size=out_dim, use_bias=False + ) + + def forward(self, inputs: TensorType) -> TensorType: + L = list(inputs.size())[1] # length of segment + H = self._num_heads # number of attention heads + D = self._head_dim # attention head dimension + + qkv = self._qkv_layer(inputs) + + queries, keys, values = torch.chunk(input=qkv, chunks=3, dim=-1) + queries = queries[:, -L:] # only query based on the segment + + queries = torch.reshape(queries, [-1, L, H, D]) + keys = torch.reshape(keys, [-1, L, H, D]) + values = torch.reshape(values, [-1, L, H, D]) + + score = torch.einsum("bihd,bjhd->bijh", queries, keys) + score = score / D**0.5 + + # causal mask of the same length as the sequence + mask = sequence_mask(torch.arange(1, L + 1), dtype=score.dtype) + mask = mask[None, :, :, None] + mask = mask.float() + + masked_score = score * mask + 1e30 * (mask - 1.0) + wmat = nn.functional.softmax(masked_score, dim=2) + + out = torch.einsum("bijh,bjhd->bihd", wmat, values) + shape = list(out.size())[:2] + [H * D] + # temp = torch.cat(temp2, [H * D], dim=0) + out = torch.reshape(out, shape) + return self._linear_layer(out) diff --git a/lib/python3.12/site-packages/ray/rllib/models/torch/modules/noisy_layer.py b/lib/python3.12/site-packages/ray/rllib/models/torch/modules/noisy_layer.py new file mode 100644 index 0000000000000000000000000000000000000000..ae542cf939a061bc5ed1fd22d56fb4994deb45ae --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/torch/modules/noisy_layer.py @@ -0,0 +1,99 @@ +import numpy as np + +from ray.rllib.models.utils import get_activation_fn +from ray.rllib.utils.framework import TensorType, try_import_torch + +torch, nn = try_import_torch() + + +class NoisyLayer(nn.Module): + r"""A Layer that adds learnable Noise to some previous layer's outputs. + + Consists of: + - a common dense layer: y = w^{T}x + b + - a noisy layer: y = (w + \epsilon_w*\sigma_w)^{T}x + + (b+\epsilon_b*\sigma_b) + , where \epsilon are random variables sampled from factorized normal + distributions and \sigma are trainable variables which are expected to + vanish along the training procedure. + """ + + def __init__( + self, in_size: int, out_size: int, sigma0: float, activation: str = "relu" + ): + """Initializes a NoisyLayer object. + + Args: + in_size: Input size for Noisy Layer + out_size: Output size for Noisy Layer + sigma0: Initialization value for sigma_b (bias noise) + activation: Non-linear activation for Noisy Layer + """ + super().__init__() + + self.in_size = in_size + self.out_size = out_size + self.sigma0 = sigma0 + self.activation = get_activation_fn(activation, framework="torch") + if self.activation is not None: + self.activation = self.activation() + + sigma_w = nn.Parameter( + torch.from_numpy( + np.random.uniform( + low=-1.0 / np.sqrt(float(self.in_size)), + high=1.0 / np.sqrt(float(self.in_size)), + size=[self.in_size, out_size], + ) + ).float() + ) + self.register_parameter("sigma_w", sigma_w) + sigma_b = nn.Parameter( + torch.from_numpy( + np.full( + shape=[out_size], fill_value=sigma0 / np.sqrt(float(self.in_size)) + ) + ).float() + ) + self.register_parameter("sigma_b", sigma_b) + + w = nn.Parameter( + torch.from_numpy( + np.full( + shape=[self.in_size, self.out_size], + fill_value=6 / np.sqrt(float(in_size) + float(out_size)), + ) + ).float() + ) + self.register_parameter("w", w) + b = nn.Parameter(torch.from_numpy(np.zeros([out_size])).float()) + self.register_parameter("b", b) + + def forward(self, inputs: TensorType) -> TensorType: + epsilon_in = self._f_epsilon( + torch.normal( + mean=torch.zeros([self.in_size]), std=torch.ones([self.in_size]) + ).to(inputs.device) + ) + epsilon_out = self._f_epsilon( + torch.normal( + mean=torch.zeros([self.out_size]), std=torch.ones([self.out_size]) + ).to(inputs.device) + ) + epsilon_w = torch.matmul( + torch.unsqueeze(epsilon_in, -1), other=torch.unsqueeze(epsilon_out, 0) + ) + epsilon_b = epsilon_out + + action_activation = ( + torch.matmul(inputs, self.w + self.sigma_w * epsilon_w) + + self.b + + self.sigma_b * epsilon_b + ) + + if self.activation is not None: + action_activation = self.activation(action_activation) + return action_activation + + def _f_epsilon(self, x: TensorType) -> TensorType: + return torch.sign(x) * torch.pow(torch.abs(x), 0.5) diff --git a/lib/python3.12/site-packages/ray/rllib/models/torch/modules/relative_multi_head_attention.py b/lib/python3.12/site-packages/ray/rllib/models/torch/modules/relative_multi_head_attention.py new file mode 100644 index 0000000000000000000000000000000000000000..aa19207e024d7d3f4a365e455974b4127b279606 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/torch/modules/relative_multi_head_attention.py @@ -0,0 +1,177 @@ +from typing import Union + +from ray.rllib.models.torch.misc import SlimFC +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.torch_utils import sequence_mask +from ray.rllib.utils.typing import TensorType + +torch, nn = try_import_torch() + + +@OldAPIStack +class RelativePositionEmbedding(nn.Module): + """Creates a [seq_length x seq_length] matrix for rel. pos encoding. + + Denoted as Phi in [2] and [3]. Phi is the standard sinusoid encoding + matrix. + + Args: + seq_length: The max. sequence length (time axis). + out_dim: The number of nodes to go into the first Tranformer + layer with. + + Returns: + torch.Tensor: The encoding matrix Phi. + """ + + def __init__(self, out_dim, **kwargs): + super().__init__() + self.out_dim = out_dim + + out_range = torch.arange(0, self.out_dim, 2.0) + inverse_freq = 1 / (10000 ** (out_range / self.out_dim)) + self.register_buffer("inverse_freq", inverse_freq) + + def forward(self, seq_length): + pos_input = torch.arange(seq_length - 1, -1, -1.0, dtype=torch.float).to( + self.inverse_freq.device + ) + sinusoid_input = torch.einsum("i,j->ij", pos_input, self.inverse_freq) + pos_embeddings = torch.cat( + [torch.sin(sinusoid_input), torch.cos(sinusoid_input)], dim=-1 + ) + return pos_embeddings[:, None, :] + + +class RelativeMultiHeadAttention(nn.Module): + """A RelativeMultiHeadAttention layer as described in [3]. + + Uses segment level recurrence with state reuse. + """ + + def __init__( + self, + in_dim: int, + out_dim: int, + num_heads: int, + head_dim: int, + input_layernorm: bool = False, + output_activation: Union[str, callable] = None, + **kwargs + ): + """Initializes a RelativeMultiHeadAttention nn.Module object. + + Args: + in_dim (int): + out_dim: The output dimension of this module. Also known as + "attention dim". + num_heads: The number of attention heads to use. + Denoted `H` in [2]. + head_dim: The dimension of a single(!) attention head + Denoted `D` in [2]. + input_layernorm: Whether to prepend a LayerNorm before + everything else. Should be True for building a GTrXL. + output_activation (Union[str, callable]): Optional activation + function or activation function specifier (str). + Should be "relu" for GTrXL. + **kwargs: + """ + super().__init__(**kwargs) + + # No bias or non-linearity. + self._num_heads = num_heads + self._head_dim = head_dim + + # 3=Query, key, and value inputs. + self._qkv_layer = SlimFC( + in_size=in_dim, out_size=3 * num_heads * head_dim, use_bias=False + ) + + self._linear_layer = SlimFC( + in_size=num_heads * head_dim, + out_size=out_dim, + use_bias=False, + activation_fn=output_activation, + ) + + self._uvar = nn.Parameter(torch.zeros(num_heads, head_dim)) + self._vvar = nn.Parameter(torch.zeros(num_heads, head_dim)) + nn.init.xavier_uniform_(self._uvar) + nn.init.xavier_uniform_(self._vvar) + self.register_parameter("_uvar", self._uvar) + self.register_parameter("_vvar", self._vvar) + + self._pos_proj = SlimFC( + in_size=in_dim, out_size=num_heads * head_dim, use_bias=False + ) + self._rel_pos_embedding = RelativePositionEmbedding(out_dim) + + self._input_layernorm = None + if input_layernorm: + self._input_layernorm = torch.nn.LayerNorm(in_dim) + + def forward(self, inputs: TensorType, memory: TensorType = None) -> TensorType: + T = list(inputs.size())[1] # length of segment (time) + H = self._num_heads # number of attention heads + d = self._head_dim # attention head dimension + + # Add previous memory chunk (as const, w/o gradient) to input. + # Tau (number of (prev) time slices in each memory chunk). + Tau = list(memory.shape)[1] + inputs = torch.cat((memory.detach(), inputs), dim=1) + + # Apply the Layer-Norm. + if self._input_layernorm is not None: + inputs = self._input_layernorm(inputs) + + qkv = self._qkv_layer(inputs) + + queries, keys, values = torch.chunk(input=qkv, chunks=3, dim=-1) + # Cut out Tau memory timesteps from query. + queries = queries[:, -T:] + + queries = torch.reshape(queries, [-1, T, H, d]) + keys = torch.reshape(keys, [-1, Tau + T, H, d]) + values = torch.reshape(values, [-1, Tau + T, H, d]) + + R = self._pos_proj(self._rel_pos_embedding(Tau + T)) + R = torch.reshape(R, [Tau + T, H, d]) + + # b=batch + # i and j=time indices (i=max-timesteps (inputs); j=Tau memory space) + # h=head + # d=head-dim (over which we will reduce-sum) + score = torch.einsum("bihd,bjhd->bijh", queries + self._uvar, keys) + pos_score = torch.einsum("bihd,jhd->bijh", queries + self._vvar, R) + score = score + self.rel_shift(pos_score) + score = score / d**0.5 + + # causal mask of the same length as the sequence + mask = sequence_mask(torch.arange(Tau + 1, Tau + T + 1), dtype=score.dtype).to( + score.device + ) + mask = mask[None, :, :, None] + + masked_score = score * mask + 1e30 * (mask.float() - 1.0) + wmat = nn.functional.softmax(masked_score, dim=2) + + out = torch.einsum("bijh,bjhd->bihd", wmat, values) + shape = list(out.shape)[:2] + [H * d] + out = torch.reshape(out, shape) + + return self._linear_layer(out) + + @staticmethod + def rel_shift(x: TensorType) -> TensorType: + # Transposed version of the shift approach described in [3]. + # https://github.com/kimiyoung/transformer-xl/blob/ + # 44781ed21dbaec88b280f74d9ae2877f52b492a5/tf/model.py#L31 + x_size = list(x.shape) + + x = torch.nn.functional.pad(x, (0, 0, 1, 0, 0, 0, 0, 0)) + x = torch.reshape(x, [x_size[0], x_size[2] + 1, x_size[1], x_size[3]]) + x = x[:, 1:, :, :] + x = torch.reshape(x, x_size) + + return x diff --git a/lib/python3.12/site-packages/ray/rllib/models/torch/modules/skip_connection.py b/lib/python3.12/site-packages/ray/rllib/models/torch/modules/skip_connection.py new file mode 100644 index 0000000000000000000000000000000000000000..9be3b36f04ce081c6432872774330bf2226c23d3 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/torch/modules/skip_connection.py @@ -0,0 +1,44 @@ +from typing import Optional + +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.typing import TensorType + +torch, nn = try_import_torch() + + +@OldAPIStack +class SkipConnection(nn.Module): + """Skip connection layer. + + Adds the original input to the output (regular residual layer) OR uses + input as hidden state input to a given fan_in_layer. + """ + + def __init__( + self, layer: nn.Module, fan_in_layer: Optional[nn.Module] = None, **kwargs + ): + """Initializes a SkipConnection nn Module object. + + Args: + layer (nn.Module): Any layer processing inputs. + fan_in_layer (Optional[nn.Module]): An optional + layer taking two inputs: The original input and the output + of `layer`. + """ + super().__init__(**kwargs) + self._layer = layer + self._fan_in_layer = fan_in_layer + + def forward(self, inputs: TensorType, **kwargs) -> TensorType: + # del kwargs + outputs = self._layer(inputs, **kwargs) + # Residual case, just add inputs to outputs. + if self._fan_in_layer is None: + outputs = outputs + inputs + # Fan-in e.g. RNN: Call fan-in with `inputs` and `outputs`. + else: + # NOTE: In the GRU case, `inputs` is the state input. + outputs = self._fan_in_layer((inputs, outputs)) + + return outputs diff --git a/lib/python3.12/site-packages/ray/rllib/models/torch/recurrent_net.py b/lib/python3.12/site-packages/ray/rllib/models/torch/recurrent_net.py new file mode 100644 index 0000000000000000000000000000000000000000..2f4e31bfd25d14ac221d53a6b05c7c66d67498fc --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/torch/recurrent_net.py @@ -0,0 +1,296 @@ +from typing import Dict, List, Tuple, Union + +import gymnasium as gym +import numpy as np +import tree # pip install dm_tree +from gymnasium.spaces import Discrete, MultiDiscrete + +from ray._common.deprecation import deprecation_warning +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.models.torch.misc import SlimFC +from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 +from ray.rllib.policy.rnn_sequencing import add_time_dimension +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.policy.view_requirement import ViewRequirement +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.spaces.space_utils import get_base_struct_from_space +from ray.rllib.utils.torch_utils import flatten_inputs_to_1d_tensor, one_hot +from ray.rllib.utils.typing import ModelConfigDict, TensorType +from ray.util.debug import log_once + +torch, nn = try_import_torch() + + +@OldAPIStack +class RecurrentNetwork(TorchModelV2): + """Helper class to simplify implementing RNN models with TorchModelV2. + + Instead of implementing forward(), you can implement forward_rnn() which + takes batches with the time dimension added already. + + Here is an example implementation for a subclass + ``MyRNNClass(RecurrentNetwork, nn.Module)``:: + + def __init__(self, obs_space, num_outputs): + nn.Module.__init__(self) + super().__init__(obs_space, action_space, num_outputs, + model_config, name) + self.obs_size = _get_size(obs_space) + self.rnn_hidden_dim = model_config["lstm_cell_size"] + self.fc1 = nn.Linear(self.obs_size, self.rnn_hidden_dim) + self.rnn = nn.GRUCell(self.rnn_hidden_dim, self.rnn_hidden_dim) + self.fc2 = nn.Linear(self.rnn_hidden_dim, num_outputs) + + self.value_branch = nn.Linear(self.rnn_hidden_dim, 1) + self._cur_value = None + + @override(ModelV2) + def get_initial_state(self): + # Place hidden states on same device as model. + h = [self.fc1.weight.new( + 1, self.rnn_hidden_dim).zero_().squeeze(0)] + return h + + @override(ModelV2) + def value_function(self): + assert self._cur_value is not None, "must call forward() first" + return self._cur_value + + @override(RecurrentNetwork) + def forward_rnn(self, input_dict, state, seq_lens): + x = nn.functional.relu(self.fc1(input_dict["obs_flat"].float())) + h_in = state[0].reshape(-1, self.rnn_hidden_dim) + h = self.rnn(x, h_in) + q = self.fc2(h) + self._cur_value = self.value_branch(h).squeeze(1) + return q, [h] + """ + + @override(ModelV2) + def forward( + self, + input_dict: Dict[str, TensorType], + state: List[TensorType], + seq_lens: TensorType, + ) -> Tuple[TensorType, List[TensorType]]: + """Adds time dimension to batch before sending inputs to forward_rnn(). + + You should implement forward_rnn() in your subclass.""" + # Creating a __init__ function that acts as a passthrough and adding the warning + # there led to errors probably due to the multiple inheritance. We encountered + # the same error if we add the Deprecated decorator. We therefore add the + # deprecation warning here. + if log_once("recurrent_network_tf"): + deprecation_warning( + old="ray.rllib.models.torch.recurrent_net.RecurrentNetwork" + ) + flat_inputs = input_dict["obs_flat"].float() + # Note that max_seq_len != input_dict.max_seq_len != seq_lens.max() + # as input_dict may have extra zero-padding beyond seq_lens.max(). + # Use add_time_dimension to handle this + self.time_major = self.model_config.get("_time_major", False) + inputs = add_time_dimension( + flat_inputs, + seq_lens=seq_lens, + framework="torch", + time_major=self.time_major, + ) + output, new_state = self.forward_rnn(inputs, state, seq_lens) + output = torch.reshape(output, [-1, self.num_outputs]) + return output, new_state + + def forward_rnn( + self, inputs: TensorType, state: List[TensorType], seq_lens: TensorType + ) -> Tuple[TensorType, List[TensorType]]: + """Call the model with the given input tensors and state. + + Args: + inputs: Observation tensor with shape [B, T, obs_size]. + state: List of state tensors, each with shape [B, size]. + seq_lens: 1D tensor holding input sequence lengths. + Note: len(seq_lens) == B. + + Returns: + (outputs, new_state): The model output tensor of shape + [B, T, num_outputs] and the list of new state tensors each with + shape [B, size]. + + Examples: + def forward_rnn(self, inputs, state, seq_lens): + model_out, h, c = self.rnn_model([inputs, seq_lens] + state) + return model_out, [h, c] + """ + raise NotImplementedError("You must implement this for an RNN model") + + +@OldAPIStack +class LSTMWrapper(RecurrentNetwork, nn.Module): + """An LSTM wrapper serving as an interface for ModelV2s that set use_lstm.""" + + def __init__( + self, + obs_space: gym.spaces.Space, + action_space: gym.spaces.Space, + num_outputs: int, + model_config: ModelConfigDict, + name: str, + ): + nn.Module.__init__(self) + super(LSTMWrapper, self).__init__( + obs_space, action_space, None, model_config, name + ) + + # At this point, self.num_outputs is the number of nodes coming + # from the wrapped (underlying) model. In other words, self.num_outputs + # is the input size for the LSTM layer. + # If None, set it to the observation space. + if self.num_outputs is None: + self.num_outputs = int(np.prod(self.obs_space.shape)) + + self.cell_size = model_config["lstm_cell_size"] + self.time_major = model_config.get("_time_major", False) + self.use_prev_action = model_config["lstm_use_prev_action"] + self.use_prev_reward = model_config["lstm_use_prev_reward"] + + self.action_space_struct = get_base_struct_from_space(self.action_space) + self.action_dim = 0 + + for space in tree.flatten(self.action_space_struct): + if isinstance(space, Discrete): + self.action_dim += space.n + elif isinstance(space, MultiDiscrete): + self.action_dim += np.sum(space.nvec) + elif space.shape is not None: + self.action_dim += int(np.prod(space.shape)) + else: + self.action_dim += int(len(space)) + + # Add prev-action/reward nodes to input to LSTM. + if self.use_prev_action: + self.num_outputs += self.action_dim + if self.use_prev_reward: + self.num_outputs += 1 + + # Define actual LSTM layer (with num_outputs being the nodes coming + # from the wrapped (underlying) layer). + self.lstm = nn.LSTM( + self.num_outputs, self.cell_size, batch_first=not self.time_major + ) + + # Set self.num_outputs to the number of output nodes desired by the + # caller of this constructor. + self.num_outputs = num_outputs + + # Postprocess LSTM output with another hidden layer and compute values. + self._logits_branch = SlimFC( + in_size=self.cell_size, + out_size=self.num_outputs, + activation_fn=None, + initializer=torch.nn.init.xavier_uniform_, + ) + self._value_branch = SlimFC( + in_size=self.cell_size, + out_size=1, + activation_fn=None, + initializer=torch.nn.init.xavier_uniform_, + ) + + # __sphinx_doc_begin__ + # Add prev-a/r to this model's view, if required. + if model_config["lstm_use_prev_action"]: + self.view_requirements[SampleBatch.PREV_ACTIONS] = ViewRequirement( + SampleBatch.ACTIONS, space=self.action_space, shift=-1 + ) + if model_config["lstm_use_prev_reward"]: + self.view_requirements[SampleBatch.PREV_REWARDS] = ViewRequirement( + SampleBatch.REWARDS, shift=-1 + ) + # __sphinx_doc_end__ + + @override(RecurrentNetwork) + def forward( + self, + input_dict: Dict[str, TensorType], + state: List[TensorType], + seq_lens: TensorType, + ) -> Tuple[TensorType, List[TensorType]]: + assert seq_lens is not None + # Push obs through "unwrapped" net's `forward()` first. + wrapped_out, _ = self._wrapped_forward(input_dict, [], None) + + # Concat. prev-action/reward if required. + prev_a_r = [] + + # Prev actions. + if self.model_config["lstm_use_prev_action"]: + prev_a = input_dict[SampleBatch.PREV_ACTIONS] + # If actions are not processed yet (in their original form as + # have been sent to environment): + # Flatten/one-hot into 1D array. + if self.model_config["_disable_action_flattening"]: + prev_a_r.append( + flatten_inputs_to_1d_tensor( + prev_a, spaces_struct=self.action_space_struct, time_axis=False + ) + ) + # If actions are already flattened (but not one-hot'd yet!), + # one-hot discrete/multi-discrete actions here. + else: + if isinstance(self.action_space, (Discrete, MultiDiscrete)): + prev_a = one_hot(prev_a.float(), self.action_space) + else: + prev_a = prev_a.float() + prev_a_r.append(torch.reshape(prev_a, [-1, self.action_dim])) + # Prev rewards. + if self.model_config["lstm_use_prev_reward"]: + prev_a_r.append( + torch.reshape(input_dict[SampleBatch.PREV_REWARDS].float(), [-1, 1]) + ) + + # Concat prev. actions + rewards to the "main" input. + if prev_a_r: + wrapped_out = torch.cat([wrapped_out] + prev_a_r, dim=1) + + # Push everything through our LSTM. + input_dict["obs_flat"] = wrapped_out + return super().forward(input_dict, state, seq_lens) + + @override(RecurrentNetwork) + def forward_rnn( + self, inputs: TensorType, state: List[TensorType], seq_lens: TensorType + ) -> Tuple[TensorType, List[TensorType]]: + # Don't show paddings to RNN(?) + # TODO: (sven) For now, only allow, iff time_major=True to not break + # anything retrospectively (time_major not supported previously). + # max_seq_len = inputs.shape[0] + # time_major = self.model_config["_time_major"] + # if time_major and max_seq_len > 1: + # inputs = torch.nn.utils.rnn.pack_padded_sequence( + # inputs, seq_lens, + # batch_first=not time_major, enforce_sorted=False) + self._features, [h, c] = self.lstm( + inputs, [torch.unsqueeze(state[0], 0), torch.unsqueeze(state[1], 0)] + ) + # Re-apply paddings. + # if time_major and max_seq_len > 1: + # self._features, _ = torch.nn.utils.rnn.pad_packed_sequence( + # self._features, + # batch_first=not time_major) + model_out = self._logits_branch(self._features) + return model_out, [torch.squeeze(h, 0), torch.squeeze(c, 0)] + + @override(ModelV2) + def get_initial_state(self) -> Union[List[np.ndarray], List[TensorType]]: + # Place hidden states on same device as model. + linear = next(self._logits_branch._model.children()) + h = [ + linear.weight.new(1, self.cell_size).zero_().squeeze(0), + linear.weight.new(1, self.cell_size).zero_().squeeze(0), + ] + return h + + @override(ModelV2) + def value_function(self) -> TensorType: + assert self._features is not None, "must call forward() first" + return torch.reshape(self._value_branch(self._features), [-1]) diff --git a/lib/python3.12/site-packages/ray/rllib/models/torch/torch_action_dist.py b/lib/python3.12/site-packages/ray/rllib/models/torch/torch_action_dist.py new file mode 100644 index 0000000000000000000000000000000000000000..6957de07258a7f3e7518b4a8de159affb0f7e9cf --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/torch/torch_action_dist.py @@ -0,0 +1,645 @@ +import functools +from math import log +from typing import Optional + +import gymnasium as gym +import numpy as np +import tree # pip install dm_tree + +from ray.rllib.models.action_dist import ActionDistribution +from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.numpy import MAX_LOG_NN_OUTPUT, MIN_LOG_NN_OUTPUT, SMALL_NUMBER +from ray.rllib.utils.spaces.space_utils import get_base_struct_from_space +from ray.rllib.utils.typing import List, ModelConfigDict, TensorType, Tuple, Union + +torch, nn = try_import_torch() + + +@OldAPIStack +class TorchDistributionWrapper(ActionDistribution): + """Wrapper class for torch.distributions.""" + + @override(ActionDistribution) + def __init__(self, inputs: List[TensorType], model: TorchModelV2): + # If inputs are not a torch Tensor, make them one and make sure they + # are on the correct device. + if not isinstance(inputs, torch.Tensor): + inputs = torch.from_numpy(inputs) + if isinstance(model, TorchModelV2): + inputs = inputs.to(next(model.parameters()).device) + super().__init__(inputs, model) + # Store the last sample here. + self.last_sample = None + + @override(ActionDistribution) + def logp(self, actions: TensorType) -> TensorType: + return self.dist.log_prob(actions) + + @override(ActionDistribution) + def entropy(self) -> TensorType: + return self.dist.entropy() + + @override(ActionDistribution) + def kl(self, other: ActionDistribution) -> TensorType: + return torch.distributions.kl.kl_divergence(self.dist, other.dist) + + @override(ActionDistribution) + def sample(self) -> TensorType: + self.last_sample = self.dist.sample() + return self.last_sample + + @override(ActionDistribution) + def sampled_action_logp(self) -> TensorType: + assert self.last_sample is not None + return self.logp(self.last_sample) + + +@OldAPIStack +class TorchCategorical(TorchDistributionWrapper): + """Wrapper class for PyTorch Categorical distribution.""" + + @override(ActionDistribution) + def __init__( + self, + inputs: List[TensorType], + model: TorchModelV2 = None, + temperature: float = 1.0, + ): + if temperature != 1.0: + assert temperature > 0.0, "Categorical `temperature` must be > 0.0!" + inputs /= temperature + super().__init__(inputs, model) + self.dist = torch.distributions.categorical.Categorical(logits=self.inputs) + + @override(ActionDistribution) + def deterministic_sample(self) -> TensorType: + self.last_sample = self.dist.probs.argmax(dim=1) + return self.last_sample + + @staticmethod + @override(ActionDistribution) + def required_model_output_shape( + action_space: gym.Space, model_config: ModelConfigDict + ) -> Union[int, np.ndarray]: + return action_space.n + + +@OldAPIStack +def get_torch_categorical_class_with_temperature(t: float): + """TorchCategorical distribution class that has customized default temperature.""" + + class TorchCategoricalWithTemperature(TorchCategorical): + def __init__(self, inputs, model=None, temperature=t): + super().__init__(inputs, model, temperature) + + return TorchCategoricalWithTemperature + + +@OldAPIStack +class TorchMultiCategorical(TorchDistributionWrapper): + """MultiCategorical distribution for MultiDiscrete action spaces.""" + + @override(TorchDistributionWrapper) + def __init__( + self, + inputs: List[TensorType], + model: TorchModelV2, + input_lens: Union[List[int], np.ndarray, Tuple[int, ...]], + action_space=None, + ): + super().__init__(inputs, model) + # If input_lens is np.ndarray or list, force-make it a tuple. + inputs_split = self.inputs.split(tuple(input_lens), dim=1) + self.cats = [ + torch.distributions.categorical.Categorical(logits=input_) + for input_ in inputs_split + ] + # Used in case we are dealing with an Int Box. + self.action_space = action_space + + @override(TorchDistributionWrapper) + def sample(self) -> TensorType: + arr = [cat.sample() for cat in self.cats] + sample_ = torch.stack(arr, dim=1) + if isinstance(self.action_space, gym.spaces.Box): + sample_ = torch.reshape(sample_, [-1] + list(self.action_space.shape)) + self.last_sample = sample_ + return sample_ + + @override(ActionDistribution) + def deterministic_sample(self) -> TensorType: + arr = [torch.argmax(cat.probs, -1) for cat in self.cats] + sample_ = torch.stack(arr, dim=1) + if isinstance(self.action_space, gym.spaces.Box): + sample_ = torch.reshape(sample_, [-1] + list(self.action_space.shape)) + self.last_sample = sample_ + return sample_ + + @override(TorchDistributionWrapper) + def logp(self, actions: TensorType) -> TensorType: + # # If tensor is provided, unstack it into list. + if isinstance(actions, torch.Tensor): + if isinstance(self.action_space, gym.spaces.Box): + actions = torch.reshape( + actions, [-1, int(np.prod(self.action_space.shape))] + ) + actions = torch.unbind(actions, dim=1) + logps = torch.stack([cat.log_prob(act) for cat, act in zip(self.cats, actions)]) + return torch.sum(logps, dim=0) + + @override(ActionDistribution) + def multi_entropy(self) -> TensorType: + return torch.stack([cat.entropy() for cat in self.cats], dim=1) + + @override(TorchDistributionWrapper) + def entropy(self) -> TensorType: + return torch.sum(self.multi_entropy(), dim=1) + + @override(ActionDistribution) + def multi_kl(self, other: ActionDistribution) -> TensorType: + return torch.stack( + [ + torch.distributions.kl.kl_divergence(cat, oth_cat) + for cat, oth_cat in zip(self.cats, other.cats) + ], + dim=1, + ) + + @override(TorchDistributionWrapper) + def kl(self, other: ActionDistribution) -> TensorType: + return torch.sum(self.multi_kl(other), dim=1) + + @staticmethod + @override(ActionDistribution) + def required_model_output_shape( + action_space: gym.Space, model_config: ModelConfigDict + ) -> Union[int, np.ndarray]: + # Int Box. + if isinstance(action_space, gym.spaces.Box): + assert action_space.dtype.name.startswith("int") + low_ = np.min(action_space.low) + high_ = np.max(action_space.high) + assert np.all(action_space.low == low_) + assert np.all(action_space.high == high_) + return np.prod(action_space.shape, dtype=np.int32) * (high_ - low_ + 1) + # MultiDiscrete space. + else: + # `nvec` is already integer. No need to cast. + return np.sum(action_space.nvec) + + +@OldAPIStack +class TorchSlateMultiCategorical(TorchCategorical): + """MultiCategorical distribution for MultiDiscrete action spaces. + + The action space must be uniform, meaning all nvec items have the same size, e.g. + MultiDiscrete([10, 10, 10]), where 10 is the number of candidates to pick from + and 3 is the slate size (pick 3 out of 10). When picking candidates, no candidate + must be picked more than once. + """ + + def __init__( + self, + inputs: List[TensorType], + model: TorchModelV2 = None, + temperature: float = 1.0, + action_space: Optional[gym.spaces.MultiDiscrete] = None, + all_slates=None, + ): + assert temperature > 0.0, "Categorical `temperature` must be > 0.0!" + # Allow softmax formula w/ temperature != 1.0: + # Divide inputs by temperature. + super().__init__(inputs / temperature, model) + self.action_space = action_space + # Assert uniformness of the action space (all discrete buckets have the same + # size). + assert isinstance(self.action_space, gym.spaces.MultiDiscrete) and all( + n == self.action_space.nvec[0] for n in self.action_space.nvec + ) + self.all_slates = all_slates + + @override(ActionDistribution) + def deterministic_sample(self) -> TensorType: + # Get a sample from the underlying Categorical (batch of ints). + sample = super().deterministic_sample() + # Use the sampled ints to pick the actual slates. + return torch.take_along_dim(self.all_slates, sample.long(), dim=-1) + + @override(ActionDistribution) + def logp(self, x: TensorType) -> TensorType: + # TODO: Implement. + return torch.ones_like(self.inputs[:, 0]) + + +@OldAPIStack +class TorchDiagGaussian(TorchDistributionWrapper): + """Wrapper class for PyTorch Normal distribution.""" + + @override(ActionDistribution) + def __init__( + self, + inputs: List[TensorType], + model: TorchModelV2, + *, + action_space: Optional[gym.spaces.Space] = None + ): + super().__init__(inputs, model) + mean, log_std = torch.chunk(self.inputs, 2, dim=1) + self.log_std = log_std + self.dist = torch.distributions.normal.Normal(mean, torch.exp(log_std)) + # Remember to squeeze action samples in case action space is Box(shape) + self.zero_action_dim = action_space and action_space.shape == () + + @override(TorchDistributionWrapper) + def sample(self) -> TensorType: + sample = super().sample() + if self.zero_action_dim: + return torch.squeeze(sample, dim=-1) + return sample + + @override(ActionDistribution) + def deterministic_sample(self) -> TensorType: + self.last_sample = self.dist.mean + return self.last_sample + + @override(TorchDistributionWrapper) + def logp(self, actions: TensorType) -> TensorType: + return super().logp(actions).sum(-1) + + @override(TorchDistributionWrapper) + def entropy(self) -> TensorType: + return super().entropy().sum(-1) + + @override(TorchDistributionWrapper) + def kl(self, other: ActionDistribution) -> TensorType: + return super().kl(other).sum(-1) + + @staticmethod + @override(ActionDistribution) + def required_model_output_shape( + action_space: gym.Space, model_config: ModelConfigDict + ) -> Union[int, np.ndarray]: + return np.prod(action_space.shape, dtype=np.int32) * 2 + + +@OldAPIStack +class TorchSquashedGaussian(TorchDistributionWrapper): + """A tanh-squashed Gaussian distribution defined by: mean, std, low, high. + + The distribution will never return low or high exactly, but + `low`+SMALL_NUMBER or `high`-SMALL_NUMBER respectively. + """ + + def __init__( + self, + inputs: List[TensorType], + model: TorchModelV2, + low: float = -1.0, + high: float = 1.0, + ): + """Parameterizes the distribution via `inputs`. + + Args: + low: The lowest possible sampling value + (excluding this value). + high: The highest possible sampling value + (excluding this value). + """ + super().__init__(inputs, model) + # Split inputs into mean and log(std). + mean, log_std = torch.chunk(self.inputs, 2, dim=-1) + # Clip `scale` values (coming from NN) to reasonable values. + log_std = torch.clamp(log_std, MIN_LOG_NN_OUTPUT, MAX_LOG_NN_OUTPUT) + std = torch.exp(log_std) + self.dist = torch.distributions.normal.Normal(mean, std) + assert np.all(np.less(low, high)) + self.low = low + self.high = high + self.mean = mean + self.std = std + + @override(ActionDistribution) + def deterministic_sample(self) -> TensorType: + self.last_sample = self._squash(self.dist.mean) + return self.last_sample + + @override(TorchDistributionWrapper) + def sample(self) -> TensorType: + # Use the reparameterization version of `dist.sample` to allow for + # the results to be backprop'able e.g. in a loss term. + + normal_sample = self.dist.rsample() + self.last_sample = self._squash(normal_sample) + return self.last_sample + + @override(ActionDistribution) + def logp(self, x: TensorType) -> TensorType: + # Unsquash values (from [low,high] to ]-inf,inf[) + unsquashed_values = self._unsquash(x) + # Get log prob of unsquashed values from our Normal. + log_prob_gaussian = self.dist.log_prob(unsquashed_values) + # For safety reasons, clamp somehow, only then sum up. + log_prob_gaussian = torch.clamp(log_prob_gaussian, -100, 100) + log_prob_gaussian = torch.sum(log_prob_gaussian, dim=-1) + # Get log-prob for squashed Gaussian. + unsquashed_values_tanhd = torch.tanh(unsquashed_values) + log_prob = log_prob_gaussian - torch.sum( + torch.log(1 - unsquashed_values_tanhd**2 + SMALL_NUMBER), dim=-1 + ) + return log_prob + + def sample_logp(self): + z = self.dist.rsample() + actions = self._squash(z) + return actions, torch.sum( + self.dist.log_prob(z) - torch.log(1 - actions * actions + SMALL_NUMBER), + dim=-1, + ) + + @override(TorchDistributionWrapper) + def entropy(self) -> TensorType: + raise ValueError("Entropy not defined for SquashedGaussian!") + + @override(TorchDistributionWrapper) + def kl(self, other: ActionDistribution) -> TensorType: + raise ValueError("KL not defined for SquashedGaussian!") + + def _squash(self, raw_values: TensorType) -> TensorType: + # Returned values are within [low, high] (including `low` and `high`). + squashed = ((torch.tanh(raw_values) + 1.0) / 2.0) * ( + self.high - self.low + ) + self.low + return torch.clamp(squashed, self.low, self.high) + + def _unsquash(self, values: TensorType) -> TensorType: + normed_values = (values - self.low) / (self.high - self.low) * 2.0 - 1.0 + # Stabilize input to atanh. + save_normed_values = torch.clamp( + normed_values, -1.0 + SMALL_NUMBER, 1.0 - SMALL_NUMBER + ) + unsquashed = torch.atanh(save_normed_values) + return unsquashed + + @staticmethod + @override(ActionDistribution) + def required_model_output_shape( + action_space: gym.Space, model_config: ModelConfigDict + ) -> Union[int, np.ndarray]: + return np.prod(action_space.shape, dtype=np.int32) * 2 + + +@OldAPIStack +class TorchBeta(TorchDistributionWrapper): + """ + A Beta distribution is defined on the interval [0, 1] and parameterized by + shape parameters alpha and beta (also called concentration parameters). + + PDF(x; alpha, beta) = x**(alpha - 1) (1 - x)**(beta - 1) / Z + with Z = Gamma(alpha) Gamma(beta) / Gamma(alpha + beta) + and Gamma(n) = (n - 1)! + """ + + def __init__( + self, + inputs: List[TensorType], + model: TorchModelV2, + low: float = 0.0, + high: float = 1.0, + ): + super().__init__(inputs, model) + # Stabilize input parameters (possibly coming from a linear layer). + self.inputs = torch.clamp(self.inputs, log(SMALL_NUMBER), -log(SMALL_NUMBER)) + self.inputs = torch.log(torch.exp(self.inputs) + 1.0) + 1.0 + self.low = low + self.high = high + alpha, beta = torch.chunk(self.inputs, 2, dim=-1) + # Note: concentration0==beta, concentration1=alpha (!) + self.dist = torch.distributions.Beta(concentration1=alpha, concentration0=beta) + + @override(ActionDistribution) + def deterministic_sample(self) -> TensorType: + self.last_sample = self._squash(self.dist.mean) + return self.last_sample + + @override(TorchDistributionWrapper) + def sample(self) -> TensorType: + # Use the reparameterization version of `dist.sample` to allow for + # the results to be backprop'able e.g. in a loss term. + normal_sample = self.dist.rsample() + self.last_sample = self._squash(normal_sample) + return self.last_sample + + @override(ActionDistribution) + def logp(self, x: TensorType) -> TensorType: + unsquashed_values = self._unsquash(x) + return torch.sum(self.dist.log_prob(unsquashed_values), dim=-1) + + def _squash(self, raw_values: TensorType) -> TensorType: + return raw_values * (self.high - self.low) + self.low + + def _unsquash(self, values: TensorType) -> TensorType: + return (values - self.low) / (self.high - self.low) + + @staticmethod + @override(ActionDistribution) + def required_model_output_shape( + action_space: gym.Space, model_config: ModelConfigDict + ) -> Union[int, np.ndarray]: + return np.prod(action_space.shape, dtype=np.int32) * 2 + + +@OldAPIStack +class TorchDeterministic(TorchDistributionWrapper): + """Action distribution that returns the input values directly. + + This is similar to DiagGaussian with standard deviation zero (thus only + requiring the "mean" values as NN output). + """ + + @override(ActionDistribution) + def deterministic_sample(self) -> TensorType: + return self.inputs + + @override(TorchDistributionWrapper) + def sampled_action_logp(self) -> TensorType: + return torch.zeros((self.inputs.size()[0],), dtype=torch.float32) + + @override(TorchDistributionWrapper) + def sample(self) -> TensorType: + return self.deterministic_sample() + + @staticmethod + @override(ActionDistribution) + def required_model_output_shape( + action_space: gym.Space, model_config: ModelConfigDict + ) -> Union[int, np.ndarray]: + return np.prod(action_space.shape, dtype=np.int32) + + +@OldAPIStack +class TorchMultiActionDistribution(TorchDistributionWrapper): + """Action distribution that operates on multiple, possibly nested actions.""" + + def __init__(self, inputs, model, *, child_distributions, input_lens, action_space): + """Initializes a TorchMultiActionDistribution object. + + Args: + inputs (torch.Tensor): A single tensor of shape [BATCH, size]. + model (TorchModelV2): The TorchModelV2 object used to produce + inputs for this distribution. + child_distributions (any[torch.Tensor]): Any struct + that contains the child distribution classes to use to + instantiate the child distributions from `inputs`. This could + be an already flattened list or a struct according to + `action_space`. + input_lens (any[int]): A flat list or a nested struct of input + split lengths used to split `inputs`. + action_space (Union[gym.spaces.Dict,gym.spaces.Tuple]): The complex + and possibly nested action space. + """ + if not isinstance(inputs, torch.Tensor): + inputs = torch.from_numpy(inputs) + if isinstance(model, TorchModelV2): + inputs = inputs.to(next(model.parameters()).device) + super().__init__(inputs, model) + + self.action_space_struct = get_base_struct_from_space(action_space) + + self.input_lens = tree.flatten(input_lens) + flat_child_distributions = tree.flatten(child_distributions) + split_inputs = torch.split(inputs, self.input_lens, dim=1) + self.flat_child_distributions = tree.map_structure( + lambda dist, input_: dist(input_, model), + flat_child_distributions, + list(split_inputs), + ) + + @override(ActionDistribution) + def logp(self, x): + if isinstance(x, np.ndarray): + x = torch.Tensor(x) + # Single tensor input (all merged). + if isinstance(x, torch.Tensor): + split_indices = [] + for dist in self.flat_child_distributions: + if isinstance(dist, TorchCategorical): + split_indices.append(1) + elif ( + isinstance(dist, TorchMultiCategorical) + and dist.action_space is not None + ): + split_indices.append(int(np.prod(dist.action_space.shape))) + else: + sample = dist.sample() + # Cover Box(shape=()) case. + if len(sample.shape) == 1: + split_indices.append(1) + else: + split_indices.append(sample.size()[1]) + split_x = list(torch.split(x, split_indices, dim=1)) + # Structured or flattened (by single action component) input. + else: + split_x = tree.flatten(x) + + def map_(val, dist): + # Remove extra categorical dimension. + if isinstance(dist, TorchCategorical): + val = (torch.squeeze(val, dim=-1) if len(val.shape) > 1 else val).int() + return dist.logp(val) + + # Remove extra categorical dimension and take the logp of each + # component. + flat_logps = tree.map_structure(map_, split_x, self.flat_child_distributions) + + return functools.reduce(lambda a, b: a + b, flat_logps) + + @override(ActionDistribution) + def kl(self, other): + kl_list = [ + d.kl(o) + for d, o in zip( + self.flat_child_distributions, other.flat_child_distributions + ) + ] + return functools.reduce(lambda a, b: a + b, kl_list) + + @override(ActionDistribution) + def entropy(self): + entropy_list = [d.entropy() for d in self.flat_child_distributions] + return functools.reduce(lambda a, b: a + b, entropy_list) + + @override(ActionDistribution) + def sample(self): + child_distributions = tree.unflatten_as( + self.action_space_struct, self.flat_child_distributions + ) + return tree.map_structure(lambda s: s.sample(), child_distributions) + + @override(ActionDistribution) + def deterministic_sample(self): + child_distributions = tree.unflatten_as( + self.action_space_struct, self.flat_child_distributions + ) + return tree.map_structure( + lambda s: s.deterministic_sample(), child_distributions + ) + + @override(TorchDistributionWrapper) + def sampled_action_logp(self): + p = self.flat_child_distributions[0].sampled_action_logp() + for c in self.flat_child_distributions[1:]: + p += c.sampled_action_logp() + return p + + @override(ActionDistribution) + def required_model_output_shape(self, action_space, model_config): + return np.sum(self.input_lens, dtype=np.int32) + + +@OldAPIStack +class TorchDirichlet(TorchDistributionWrapper): + """Dirichlet distribution for continuous actions that are between + [0,1] and sum to 1. + + e.g. actions that represent resource allocation.""" + + def __init__(self, inputs, model): + """Input is a tensor of logits. The exponential of logits is used to + parametrize the Dirichlet distribution as all parameters need to be + positive. An arbitrary small epsilon is added to the concentration + parameters to be zero due to numerical error. + + See issue #4440 for more details. + """ + self.epsilon = torch.tensor(1e-7).to(inputs.device) + concentration = torch.exp(inputs) + self.epsilon + self.dist = torch.distributions.dirichlet.Dirichlet( + concentration=concentration, + validate_args=True, + ) + super().__init__(concentration, model) + + @override(ActionDistribution) + def deterministic_sample(self) -> TensorType: + self.last_sample = nn.functional.softmax(self.dist.concentration, dim=-1) + return self.last_sample + + @override(ActionDistribution) + def logp(self, x): + # Support of Dirichlet are positive real numbers. x is already + # an array of positive numbers, but we clip to avoid zeros due to + # numerical errors. + x = torch.max(x, self.epsilon) + x = x / torch.sum(x, dim=-1, keepdim=True) + return self.dist.log_prob(x) + + @override(ActionDistribution) + def entropy(self): + return self.dist.entropy() + + @staticmethod + @override(ActionDistribution) + def required_model_output_shape(action_space, model_config): + return np.prod(action_space.shape, dtype=np.int32) diff --git a/lib/python3.12/site-packages/ray/rllib/models/torch/torch_distributions.py b/lib/python3.12/site-packages/ray/rllib/models/torch/torch_distributions.py new file mode 100644 index 0000000000000000000000000000000000000000..e08a60b4fde2c1274cdaa63e16acee1665ec98f7 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/torch/torch_distributions.py @@ -0,0 +1,16 @@ +from ray._common.deprecation import deprecation_warning +from ray.rllib.core.distribution.torch.torch_distribution import ( # noqa + TorchCategorical, + TorchDeterministic, + TorchDiagGaussian, + TorchDistribution, + TorchMultiCategorical, + TorchMultiDistribution, + TorchSquashedGaussian, +) + +deprecation_warning( + old="ray.rllib.models.torch.torch_distributions.TorchDistribution", + new="ray.rllib.core.distribution.torch.torch_distribution.TorchDistribution", + error=False, +) diff --git a/lib/python3.12/site-packages/ray/rllib/models/torch/torch_modelv2.py b/lib/python3.12/site-packages/ray/rllib/models/torch/torch_modelv2.py new file mode 100644 index 0000000000000000000000000000000000000000..003369ad4b49616a674e68b4be1d8da0a3188544 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/torch/torch_modelv2.py @@ -0,0 +1,81 @@ +from typing import Dict, List, Union + +import gymnasium as gym + +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.typing import ModelConfigDict, TensorType + +_, nn = try_import_torch() + + +@OldAPIStack +class TorchModelV2(ModelV2): + """Torch version of ModelV2. + + Note that this class by itself is not a valid model unless you + inherit from nn.Module and implement forward() in a subclass.""" + + def __init__( + self, + obs_space: gym.spaces.Space, + action_space: gym.spaces.Space, + num_outputs: int, + model_config: ModelConfigDict, + name: str, + ): + """Initialize a TorchModelV2. + + Here is an example implementation for a subclass + ``MyModelClass(TorchModelV2, nn.Module)``:: + + def __init__(self, *args, **kwargs): + TorchModelV2.__init__(self, *args, **kwargs) + nn.Module.__init__(self) + self._hidden_layers = nn.Sequential(...) + self._logits = ... + self._value_branch = ... + """ + if not isinstance(self, nn.Module): + raise ValueError( + "Subclasses of TorchModelV2 must also inherit from " + "nn.Module, e.g., MyModel(TorchModelV2, nn.Module)" + ) + + ModelV2.__init__( + self, + obs_space, + action_space, + num_outputs, + model_config, + name, + framework="torch", + ) + + # Dict to store per multi-gpu tower stats into. + # In PyTorch multi-GPU, we use a single TorchPolicy and copy + # it's Model(s) n times (1 copy for each GPU). When computing the loss + # on each tower, we cannot store the stats (e.g. `entropy`) inside the + # policy object as this would lead to race conditions between the + # different towers all accessing the same property at the same time. + self.tower_stats = {} + + @override(ModelV2) + def variables( + self, as_dict: bool = False + ) -> Union[List[TensorType], Dict[str, TensorType]]: + p = list(self.parameters()) + if as_dict: + return {k: p[i] for i, k in enumerate(self.state_dict().keys())} + return p + + @override(ModelV2) + def trainable_variables( + self, as_dict: bool = False + ) -> Union[List[TensorType], Dict[str, TensorType]]: + if as_dict: + return { + k: v for k, v in self.variables(as_dict=True).items() if v.requires_grad + } + return [v for v in self.variables() if v.requires_grad] diff --git a/lib/python3.12/site-packages/ray/rllib/models/torch/visionnet.py b/lib/python3.12/site-packages/ray/rllib/models/torch/visionnet.py new file mode 100644 index 0000000000000000000000000000000000000000..43f98efb796745cae4b5152a0c1977d0ec9bb5d1 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/torch/visionnet.py @@ -0,0 +1,294 @@ +from typing import Dict, List + +import gymnasium as gym +import numpy as np + +from ray.rllib.models.torch.misc import ( + SlimConv2d, + SlimFC, + normc_initializer, + same_padding, +) +from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 +from ray.rllib.models.utils import get_activation_fn, get_filter_config +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.typing import ModelConfigDict, TensorType + +torch, nn = try_import_torch() + + +@OldAPIStack +class VisionNetwork(TorchModelV2, nn.Module): + """Generic vision network.""" + + def __init__( + self, + obs_space: gym.spaces.Space, + action_space: gym.spaces.Space, + num_outputs: int, + model_config: ModelConfigDict, + name: str, + ): + if not model_config.get("conv_filters"): + model_config["conv_filters"] = get_filter_config(obs_space.shape) + + TorchModelV2.__init__( + self, obs_space, action_space, num_outputs, model_config, name + ) + nn.Module.__init__(self) + + activation = self.model_config.get("conv_activation") + filters = self.model_config["conv_filters"] + assert len(filters) > 0, "Must provide at least 1 entry in `conv_filters`!" + + # Post FC net config. + post_fcnet_hiddens = model_config.get("post_fcnet_hiddens", []) + post_fcnet_activation = get_activation_fn( + model_config.get("post_fcnet_activation"), framework="torch" + ) + + no_final_linear = self.model_config.get("no_final_linear") + vf_share_layers = self.model_config.get("vf_share_layers") + + # Whether the last layer is the output of a Flattened (rather than + # a n x (1,1) Conv2D). + self.last_layer_is_flattened = False + self._logits = None + + layers = [] + (w, h, in_channels) = obs_space.shape + + in_size = [w, h] + for out_channels, kernel, stride in filters[:-1]: + padding, out_size = same_padding(in_size, kernel, stride) + layers.append( + SlimConv2d( + in_channels, + out_channels, + kernel, + stride, + padding, + activation_fn=activation, + ) + ) + in_channels = out_channels + in_size = out_size + + out_channels, kernel, stride = filters[-1] + + # No final linear: Last layer has activation function and exits with + # num_outputs nodes (this could be a 1x1 conv or a FC layer, depending + # on `post_fcnet_...` settings). + if no_final_linear and num_outputs: + out_channels = out_channels if post_fcnet_hiddens else num_outputs + layers.append( + SlimConv2d( + in_channels, + out_channels, + kernel, + stride, + None, # padding=valid + activation_fn=activation, + ) + ) + + # Add (optional) post-fc-stack after last Conv2D layer. + layer_sizes = post_fcnet_hiddens[:-1] + ( + [num_outputs] if post_fcnet_hiddens else [] + ) + for i, out_size in enumerate(layer_sizes): + layers.append( + SlimFC( + in_size=out_channels, + out_size=out_size, + activation_fn=post_fcnet_activation, + initializer=normc_initializer(1.0), + ) + ) + out_channels = out_size + + # Finish network normally (w/o overriding last layer size with + # `num_outputs`), then add another linear one of size `num_outputs`. + else: + layers.append( + SlimConv2d( + in_channels, + out_channels, + kernel, + stride, + None, # padding=valid + activation_fn=activation, + ) + ) + + # num_outputs defined. Use that to create an exact + # `num_output`-sized (1,1)-Conv2D. + if num_outputs: + in_size = [ + np.ceil((in_size[0] - kernel[0]) / stride), + np.ceil((in_size[1] - kernel[1]) / stride), + ] + padding, _ = same_padding(in_size, [1, 1], [1, 1]) + if post_fcnet_hiddens: + layers.append(nn.Flatten()) + in_size = out_channels + # Add (optional) post-fc-stack after last Conv2D layer. + for i, out_size in enumerate(post_fcnet_hiddens + [num_outputs]): + layers.append( + SlimFC( + in_size=in_size, + out_size=out_size, + activation_fn=post_fcnet_activation + if i < len(post_fcnet_hiddens) - 1 + else None, + initializer=normc_initializer(1.0), + ) + ) + in_size = out_size + # Last layer is logits layer. + self._logits = layers.pop() + + else: + self._logits = SlimConv2d( + out_channels, + num_outputs, + [1, 1], + 1, + padding, + activation_fn=None, + ) + + # num_outputs not known -> Flatten, then set self.num_outputs + # to the resulting number of nodes. + else: + self.last_layer_is_flattened = True + layers.append(nn.Flatten()) + + self._convs = nn.Sequential(*layers) + + # If our num_outputs still unknown, we need to do a test pass to + # figure out the output dimensions. This could be the case, if we have + # the Flatten layer at the end. + if self.num_outputs is None: + # Create a B=1 dummy sample and push it through out conv-net. + dummy_in = ( + torch.from_numpy(self.obs_space.sample()) + .permute(2, 0, 1) + .unsqueeze(0) + .float() + ) + dummy_out = self._convs(dummy_in) + self.num_outputs = dummy_out.shape[1] + + # Build the value layers + self._value_branch_separate = self._value_branch = None + if vf_share_layers: + self._value_branch = SlimFC( + out_channels, 1, initializer=normc_initializer(0.01), activation_fn=None + ) + else: + vf_layers = [] + (w, h, in_channels) = obs_space.shape + in_size = [w, h] + for out_channels, kernel, stride in filters[:-1]: + padding, out_size = same_padding(in_size, kernel, stride) + vf_layers.append( + SlimConv2d( + in_channels, + out_channels, + kernel, + stride, + padding, + activation_fn=activation, + ) + ) + in_channels = out_channels + in_size = out_size + + out_channels, kernel, stride = filters[-1] + vf_layers.append( + SlimConv2d( + in_channels, + out_channels, + kernel, + stride, + None, + activation_fn=activation, + ) + ) + + vf_layers.append( + SlimConv2d( + in_channels=out_channels, + out_channels=1, + kernel=1, + stride=1, + padding=None, + activation_fn=None, + ) + ) + self._value_branch_separate = nn.Sequential(*vf_layers) + + # Holds the current "base" output (before logits layer). + self._features = None + + @override(TorchModelV2) + def forward( + self, + input_dict: Dict[str, TensorType], + state: List[TensorType], + seq_lens: TensorType, + ) -> (TensorType, List[TensorType]): + self._features = input_dict["obs"].float() + # Permuate b/c data comes in as [B, dim, dim, channels]: + self._features = self._features.permute(0, 3, 1, 2) + conv_out = self._convs(self._features) + # Store features to save forward pass when getting value_function out. + if not self._value_branch_separate: + self._features = conv_out + + if not self.last_layer_is_flattened: + if self._logits: + conv_out = self._logits(conv_out) + if len(conv_out.shape) == 4: + if conv_out.shape[2] != 1 or conv_out.shape[3] != 1: + raise ValueError( + "Given `conv_filters` ({}) do not result in a [B, {} " + "(`num_outputs`), 1, 1] shape (but in {})! Please " + "adjust your Conv2D stack such that the last 2 dims " + "are both 1.".format( + self.model_config["conv_filters"], + self.num_outputs, + list(conv_out.shape), + ) + ) + logits = conv_out.squeeze(3) + logits = logits.squeeze(2) + else: + logits = conv_out + return logits, state + else: + return conv_out, state + + @override(TorchModelV2) + def value_function(self) -> TensorType: + assert self._features is not None, "must call forward() first" + if self._value_branch_separate: + value = self._value_branch_separate(self._features) + value = value.squeeze(3) + value = value.squeeze(2) + return value.squeeze(1) + else: + if not self.last_layer_is_flattened: + features = self._features.squeeze(3) + features = features.squeeze(2) + else: + features = self._features + return self._value_branch(features).squeeze(1) + + def _hidden_layers(self, obs: TensorType) -> TensorType: + res = self._convs(obs.permute(0, 3, 1, 2)) # switch to channel-major + res = res.squeeze(3) + res = res.squeeze(2) + return res diff --git a/lib/python3.12/site-packages/ray/rllib/models/utils.py b/lib/python3.12/site-packages/ray/rllib/models/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..0849a098440e6ea70237bf3a57d1bb6cc6924bd5 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/models/utils.py @@ -0,0 +1,291 @@ +from typing import Callable, Optional, Union + +from ray.rllib.utils.annotations import DeveloperAPI +from ray.rllib.utils.framework import try_import_jax, try_import_tf, try_import_torch + + +@DeveloperAPI +def get_activation_fn( + name: Optional[Union[Callable, str]] = None, + framework: str = "tf", +): + """Returns a framework specific activation function, given a name string. + + Args: + name: One of "relu" (default), "tanh", "elu", + "swish" (or "silu", which is the same), or "linear" (same as None). + framework: One of "jax", "tf|tf2" or "torch". + + Returns: + A framework-specific activtion function. e.g. tf.nn.tanh or + torch.nn.ReLU. None if name in ["linear", None]. + + Raises: + ValueError: If name is an unknown activation function. + """ + # Already a callable, return as-is. + if callable(name): + return name + + name_lower = name.lower() if isinstance(name, str) else name + + # Infer the correct activation function from the string specifier. + if framework == "torch": + if name_lower in ["linear", None]: + return None + + _, nn = try_import_torch() + # First try getting the correct activation function from nn directly. + # Note that torch activation functions are not all lower case. + fn = getattr(nn, name, None) + if fn is not None: + return fn + + if name_lower in ["swish", "silu"]: + return nn.SiLU + elif name_lower == "relu": + return nn.ReLU + elif name_lower == "tanh": + return nn.Tanh + elif name_lower == "elu": + return nn.ELU + elif name_lower == "softmax": + return nn.Softmax + elif framework == "jax": + if name_lower in ["linear", None]: + return None + jax, _ = try_import_jax() + if name_lower in ["swish", "silu"]: + return jax.nn.swish + if name_lower == "relu": + return jax.nn.relu + elif name_lower == "tanh": + return jax.nn.hard_tanh + elif name_lower == "elu": + return jax.nn.elu + else: + assert framework in ["tf", "tf2"], "Unsupported framework `{}`!".format( + framework + ) + if name_lower in ["linear", None]: + return None + + tf1, tf, tfv = try_import_tf() + # Try getting the correct activation function from tf.nn directly. + # Note that tf activation functions are all lower case, so this should always + # work. + fn = getattr(tf.nn, name_lower, None) + + if fn is not None: + return fn + + raise ValueError( + "Unknown activation ({}) for framework={}!".format(name, framework) + ) + + +@DeveloperAPI +def get_initializer_fn(name: Optional[Union[str, Callable]], framework: str = "torch"): + """Returns the framework-specific initializer class or function. + + This function relies fully on the specified initializer classes and + functions in the frameworks `torch` and `tf2` (see for `torch` + https://pytorch.org/docs/stable/nn.init.html and for `tf2` see + https://www.tensorflow.org/api_docs/python/tf/keras/initializers). + + Note, for framework `torch` the in-place initializers are needed, i.e. names + should end with an underscore `_`, e.g. `glorot_uniform_`. + + Args: + name: Name of the initializer class or function in one of the two + supported frameworks, i.e. `torch` or `tf2`. + framework: The framework string, either `torch or `tf2`. + + Returns: + A framework-specific function or class defining an initializer to be used + for network initialization, + + Raises: + `ValueError` if the `name` is neither class or function in the specified + `framework`. Raises also a `ValueError`, if `name` does not define an + in-place initializer for framework `torch`. + """ + # Already a callable or `None` return as is. If `None` we use the default + # initializer defined in the framework-specific layers themselves. + if callable(name) or name is None: + return name + + if framework == "torch": + name_lower = name.lower() if isinstance(name, str) else name + + _, nn = try_import_torch() + + # Check, if the name includes an underscore. We must use the + # in-place initialization from Torch. + if not name_lower.endswith("_"): + raise ValueError( + "Not an in-place initializer: Torch weight initializers " + "need to be provided as their in-place version, i.e. " + " + '_'. See " + "https://pytorch.org/docs/stable/nn.init.html. " + f"User provided {name}." + ) + + # First, try to get the initialization directly from `nn.init`. + # Note, that all initialization methods in `nn.init` are lower + # case and that `_` defines the "in-place" method. + fn = getattr(nn.init, name_lower, None) + if fn is not None: + # TODO (simon): Raise a warning if not "in-place" method. + return fn + # Unknown initializer. + else: + # Inform the user that this initializer does not exist. + raise ValueError( + f"Unknown initializer name: {name_lower} is not a method in " + "`torch.nn.init`!" + ) + elif framework == "tf2": + # Note, as initializer classes in TensorFlow can be either given by their + # name in camel toe typing or by their shortcut we use the `name` as it is. + # See https://www.tensorflow.org/api_docs/python/tf/keras/initializers. + + _, tf, _ = try_import_tf() + + # Try to get the initialization function directly from `tf.keras.initializers`. + fn = getattr(tf.keras.initializers, name, None) + if fn is not None: + return fn + # Unknown initializer. + else: + # Inform the user that this initializer does not exist. + raise ValueError( + f"Unknown initializer: {name} is not a initializer in " + "`tf.keras.initializers`!" + ) + + +@DeveloperAPI +def get_filter_config(shape): + """Returns a default Conv2D filter config (list) for a given image shape. + + Args: + shape (Tuple[int]): The input (image) shape, e.g. (84,84,3). + + Returns: + List[list]: The Conv2D filter configuration usable as `conv_filters` + inside a model config dict. + """ + # 96x96x3 (e.g. CarRacing-v0). + filters_96x96 = [ + [16, [8, 8], 4], + [32, [4, 4], 2], + [256, [11, 11], 2], + ] + # Atari. + filters_84x84 = [ + [16, [8, 8], 4], + [32, [4, 4], 2], + [256, [11, 11], 1], + ] + # Dreamer-style (XS-sized model) Atari or DM Control Suite. + filters_64x64 = [ + [16, [4, 4], 2], + [32, [4, 4], 2], + [64, [4, 4], 2], + [128, [4, 4], 2], + ] + # Small (1/2) Atari. + filters_42x42 = [ + [16, [4, 4], 2], + [32, [4, 4], 2], + [256, [11, 11], 1], + ] + # Test image (10x10). + filters_10x10 = [ + [16, [5, 5], 2], + [32, [5, 5], 2], + ] + + shape = list(shape) + if len(shape) in [2, 3] and (shape[:2] == [96, 96] or shape[1:] == [96, 96]): + return filters_96x96 + elif len(shape) in [2, 3] and (shape[:2] == [84, 84] or shape[1:] == [84, 84]): + return filters_84x84 + elif len(shape) in [2, 3] and (shape[:2] == [64, 64] or shape[1:] == [64, 64]): + return filters_64x64 + elif len(shape) in [2, 3] and (shape[:2] == [42, 42] or shape[1:] == [42, 42]): + return filters_42x42 + elif len(shape) in [2, 3] and (shape[:2] == [10, 10] or shape[1:] == [10, 10]): + return filters_10x10 + else: + if list(shape) == [210, 160, 3]: + atari_help = ( + "This is the default atari obs shape. You may want to look at one of " + "RLlib's Atari examples for an example of how to wrap an Atari env. " + ) + else: + atari_help = "" + raise ValueError( + "No default CNN configuration for obs shape {}. ".format(shape) + + atari_help + + "You can specify `conv_filters` manually through your " + "AlgorithmConfig's model_config. " + "Default configurations are only available for inputs of the following " + "shapes: [42, 42, K], [84, 84, K], [64, 64, K], [10, 10, K]. You may " + "want to use a custom RLModule or a ConnectorV2 for that." + ) + + +@DeveloperAPI +def get_initializer(name, framework="tf"): + """Returns a framework specific initializer, given a name string. + + Args: + name: One of "xavier_uniform" (default), "xavier_normal". + framework: One of "jax", "tf|tf2" or "torch". + + Returns: + A framework-specific initializer function, e.g. + tf.keras.initializers.GlorotUniform or + torch.nn.init.xavier_uniform_. + + Raises: + ValueError: If name is an unknown initializer. + """ + # Already a callable, return as-is. + if callable(name): + return name + + if framework == "jax": + _, flax = try_import_jax() + assert flax is not None, "`flax` not installed. Try `pip install jax flax`." + import flax.linen as nn + + if name in [None, "default", "xavier_uniform"]: + return nn.initializers.xavier_uniform() + elif name == "xavier_normal": + return nn.initializers.xavier_normal() + if framework == "torch": + _, nn = try_import_torch() + assert nn is not None, "`torch` not installed. Try `pip install torch`." + if name in [None, "default", "xavier_uniform"]: + return nn.init.xavier_uniform_ + elif name == "xavier_normal": + return nn.init.xavier_normal_ + else: + assert framework in ["tf", "tf2"], "Unsupported framework `{}`!".format( + framework + ) + tf1, tf, tfv = try_import_tf() + assert ( + tf is not None + ), "`tensorflow` not installed. Try `pip install tensorflow`." + if name in [None, "default", "xavier_uniform"]: + return tf.keras.initializers.GlorotUniform + elif name == "xavier_normal": + return tf.keras.initializers.GlorotNormal + + raise ValueError( + "Unknown activation ({}) for framework={}!".format(name, framework) + ) diff --git a/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8d4223771e337b8ad179b0d914a5aaaceee27083 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/d4rl_reader.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/d4rl_reader.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..751fc09e9db9de7fcf01b29c6d667acd218717e8 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/d4rl_reader.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/dataset_reader.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/dataset_reader.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f3c117d8a0090cc93c506281ca49ffa5a0d729b0 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/dataset_reader.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/dataset_writer.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/dataset_writer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..34197c4e86f0ef48d010e9c48511d739a33c2a56 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/dataset_writer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/feature_importance.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/feature_importance.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..87b8c5281f17ba7ca17b74624c3dc6e78e1c0fdb Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/feature_importance.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/input_reader.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/input_reader.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f110313c6b777a2f1032113fd296bf468a598c8b Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/input_reader.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/io_context.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/io_context.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..21a8377ef4dcea92e03414a54cfedc2ba278250a Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/io_context.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/is_estimator.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/is_estimator.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f73d1c02072f3c7904f7d5943d79eb84cfd62c09 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/is_estimator.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/json_reader.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/json_reader.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6e0a406f6f587d052632ff6ece2d9e799ba1e7e8 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/json_reader.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/json_writer.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/json_writer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..49dbc9183a4cb4dad9efb062eef0d9f1e2c2ecc0 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/json_writer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/mixed_input.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/mixed_input.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ea0ed6d9a0b6d810900cbd020b87db9b4c5aaf88 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/mixed_input.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/off_policy_estimator.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/off_policy_estimator.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..650be229776c77e4dbae4448500544e2de649cdb Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/off_policy_estimator.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/offline_data.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/offline_data.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8a0ca28ae39047ab3540e851eeb5610b344a487f Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/offline_data.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/offline_env_runner.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/offline_env_runner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6599a75dea65ce7a28874bb69f3bfce7af5db715 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/offline_env_runner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/offline_evaluation_runner.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/offline_evaluation_runner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ce794d75e97ca5f519d5060db0e75b8763fee188 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/offline_evaluation_runner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/offline_evaluation_runner_group.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/offline_evaluation_runner_group.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..01788f77f4617cc2b97862f7221f1d38424fdec4 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/offline_evaluation_runner_group.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/offline_evaluation_utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/offline_evaluation_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..50174a93fd7ec3de68449b40ab0ad13806c63d04 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/offline_evaluation_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/offline_evaluator.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/offline_evaluator.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f28b950d2a18fcfe79dd8bd37e64a9bd8cedc538 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/offline_evaluator.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/offline_policy_evaluation_runner.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/offline_policy_evaluation_runner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3a66e5ea621da2ccee674dc5bb1178b8ea197715 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/offline_policy_evaluation_runner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/offline_prelearner.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/offline_prelearner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ad5902eb5a2ef32e6cf1d05030b2bc4373e9a7ab Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/offline_prelearner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/output_writer.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/output_writer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6e3e48bc5da25ea631f766e352bdb54c9ba19ae7 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/output_writer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/resource.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/resource.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f06dbe260fa04dd1a8fabc12422f1f80c29309f8 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/resource.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/shuffled_input.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/shuffled_input.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..39fa271998a5ed9feb8c80538d9b402da124d857 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/shuffled_input.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/wis_estimator.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/wis_estimator.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..66d93392ce8542723ac3c4d2a69fd5032dcd0fdb Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/offline/__pycache__/wis_estimator.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/offline/estimators/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/offline/estimators/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..60d67b95d4d6f5c6709539024ddde9e2434031ed Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/offline/estimators/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/offline/estimators/__pycache__/direct_method.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/offline/estimators/__pycache__/direct_method.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cb62866faf0078f2447aa10e937e63da64a94749 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/offline/estimators/__pycache__/direct_method.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/offline/estimators/__pycache__/doubly_robust.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/offline/estimators/__pycache__/doubly_robust.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..16596efb55cf1a7cea885c4c9c262c37e8f9cefb Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/offline/estimators/__pycache__/doubly_robust.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/offline/estimators/__pycache__/feature_importance.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/offline/estimators/__pycache__/feature_importance.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cd36a4eca9724b12fc9017ae8d9e04bd3c144f95 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/offline/estimators/__pycache__/feature_importance.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/offline/estimators/__pycache__/fqe_torch_model.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/offline/estimators/__pycache__/fqe_torch_model.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bc1ccdbfd0ec7dac44797757b1015bdf8895f38e Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/offline/estimators/__pycache__/fqe_torch_model.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/offline/estimators/__pycache__/importance_sampling.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/offline/estimators/__pycache__/importance_sampling.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..13b4bf3c62f42cd725e2d2ea7d84d3b334871a8d Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/offline/estimators/__pycache__/importance_sampling.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/offline/estimators/__pycache__/off_policy_estimator.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/offline/estimators/__pycache__/off_policy_estimator.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6a5791474fce2c60cb40ed2933636ca17445523b Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/offline/estimators/__pycache__/off_policy_estimator.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/offline/estimators/__pycache__/weighted_importance_sampling.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/offline/estimators/__pycache__/weighted_importance_sampling.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c020a1490227fd8c7b153cfbcdfff6dc10011b19 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/offline/estimators/__pycache__/weighted_importance_sampling.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/offline/estimators/direct_method.py b/lib/python3.12/site-packages/ray/rllib/offline/estimators/direct_method.py new file mode 100644 index 0000000000000000000000000000000000000000..99e116deae053f7f5d5a718f83fa7747ae023b50 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/offline/estimators/direct_method.py @@ -0,0 +1,179 @@ +import logging +import math +from typing import Any, Dict, List, Optional + +import numpy as np + +from ray.data import Dataset +from ray.rllib.offline.estimators.fqe_torch_model import FQETorchModel +from ray.rllib.offline.estimators.off_policy_estimator import OffPolicyEstimator +from ray.rllib.offline.offline_evaluation_utils import compute_q_and_v_values +from ray.rllib.offline.offline_evaluator import OfflineEvaluator +from ray.rllib.policy import Policy +from ray.rllib.policy.sample_batch import SampleBatch, convert_ma_batch_to_sample_batch +from ray.rllib.utils.annotations import DeveloperAPI, override +from ray.rllib.utils.numpy import convert_to_numpy +from ray.rllib.utils.typing import SampleBatchType + +logger = logging.getLogger() + + +@DeveloperAPI +class DirectMethod(OffPolicyEstimator): + r"""The Direct Method estimator. + + Let s_t, a_t, and r_t be the state, action, and reward at timestep t. + + This method trains a Q-model for the evaluation policy \pi_e on behavior + data generated by \pi_b. Currently, RLlib implements this using + Fitted-Q Evaluation (FQE). You can also implement your own model + and pass it in as `q_model_config = {"type": your_model_class, **your_kwargs}`. + + This estimator computes the expected return for \pi_e for an episode as: + V^{\pi_e}(s_0) = \sum_{a \in A} \pi_e(a | s_0) Q(s_0, a) + and returns the mean and standard deviation over episodes. + + For more information refer to https://arxiv.org/pdf/1911.06854.pdf""" + + @override(OffPolicyEstimator) + def __init__( + self, + policy: Policy, + gamma: float, + epsilon_greedy: float = 0.0, + q_model_config: Optional[Dict] = None, + ): + """Initializes a Direct Method OPE Estimator. + + Args: + policy: Policy to evaluate. + gamma: Discount factor of the environment. + epsilon_greedy: The probability by which we act acording to a fully random + policy during deployment. With 1-epsilon_greedy we act according the + target policy. + q_model_config: Arguments to specify the Q-model. Must specify + a `type` key pointing to the Q-model class. + This Q-model is trained in the train() method and is used + to compute the state-value estimates for the DirectMethod estimator. + It must implement `train` and `estimate_v`. + TODO (Rohan138): Unify this with RLModule API. + """ + + super().__init__(policy, gamma, epsilon_greedy) + + # Some dummy policies and ones that are not based on a tensor framework + # backend can come without a config or without a framework key. + if hasattr(policy, "config"): + assert ( + policy.config.get("framework", "torch") == "torch" + ), "Framework must be torch to use DirectMethod." + + q_model_config = q_model_config or {} + model_cls = q_model_config.pop("type", FQETorchModel) + self.model = model_cls( + policy=policy, + gamma=gamma, + **q_model_config, + ) + assert hasattr( + self.model, "estimate_v" + ), "self.model must implement `estimate_v`!" + + @override(OffPolicyEstimator) + def estimate_on_single_episode(self, episode: SampleBatch) -> Dict[str, Any]: + estimates_per_epsiode = {} + rewards = episode["rewards"] + + v_behavior = 0.0 + for t in range(episode.count): + v_behavior += rewards[t] * self.gamma**t + + v_target = self._compute_v_target(episode[:1]) + + estimates_per_epsiode["v_behavior"] = v_behavior + estimates_per_epsiode["v_target"] = v_target + + return estimates_per_epsiode + + @override(OffPolicyEstimator) + def estimate_on_single_step_samples( + self, batch: SampleBatch + ) -> Dict[str, List[float]]: + estimates_per_epsiode = {} + rewards = batch["rewards"] + + v_behavior = rewards + v_target = self._compute_v_target(batch) + + estimates_per_epsiode["v_behavior"] = v_behavior + estimates_per_epsiode["v_target"] = v_target + + return estimates_per_epsiode + + def _compute_v_target(self, init_step): + v_target = self.model.estimate_v(init_step) + v_target = convert_to_numpy(v_target) + return v_target + + @override(OffPolicyEstimator) + def train(self, batch: SampleBatchType) -> Dict[str, Any]: + """Trains self.model on the given batch. + + Args: + batch: A SampleBatchType to train on + + Returns: + A dict with key "loss" and value as the mean training loss. + """ + batch = convert_ma_batch_to_sample_batch(batch) + losses = self.model.train(batch) + return {"loss": np.mean(losses)} + + @override(OfflineEvaluator) + def estimate_on_dataset( + self, dataset: Dataset, *, n_parallelism: int = ... + ) -> Dict[str, Any]: + """Calculates the Direct Method estimate on the given dataset. + + Note: This estimate works for only discrete action spaces for now. + + Args: + dataset: Dataset to compute the estimate on. Each record in dataset should + include the following columns: `obs`, `actions`, `action_prob` and + `rewards`. The `obs` on each row shoud be a vector of D dimensions. + n_parallelism: The number of parallel workers to use. + + Returns: + Dictionary with the following keys: + v_target: The estimated value of the target policy. + v_behavior: The estimated value of the behavior policy. + v_gain: The estimated gain of the target policy over the behavior + policy. + v_std: The standard deviation of the estimated value of the target. + """ + # compute v_values + batch_size = max(dataset.count() // n_parallelism, 1) + updated_ds = dataset.map_batches( + compute_q_and_v_values, + batch_size=batch_size, + batch_format="pandas", + fn_kwargs={ + "model_class": self.model.__class__, + "model_state": self.model.get_state(), + "compute_q_values": False, + }, + ) + + v_behavior = updated_ds.mean("rewards") + v_target = updated_ds.mean("v_values") + v_gain_mean = v_target / v_behavior + v_gain_ste = ( + updated_ds.std("v_values") / v_behavior / math.sqrt(dataset.count()) + ) + + return { + "v_behavior": v_behavior, + "v_target": v_target, + "v_gain_mean": v_gain_mean, + "v_gain_ste": v_gain_ste, + } diff --git a/lib/python3.12/site-packages/ray/rllib/offline/estimators/feature_importance.py b/lib/python3.12/site-packages/ray/rllib/offline/estimators/feature_importance.py new file mode 100644 index 0000000000000000000000000000000000000000..148426aefb9b997c83beba08aedf261053f8366d --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/offline/estimators/feature_importance.py @@ -0,0 +1,11 @@ +from ray.rllib.offline.feature_importance import FeatureImportance + +__all__ = ["FeatureImportance"] + +from ray._common.deprecation import deprecation_warning + +deprecation_warning( + "ray.rllib.offline.estimators.feature_importance.FeatureImportance", + "ray.rllib.offline.feature_importance.FeatureImportance", + error=True, +) diff --git a/lib/python3.12/site-packages/ray/rllib/offline/estimators/off_policy_estimator.py b/lib/python3.12/site-packages/ray/rllib/offline/estimators/off_policy_estimator.py new file mode 100644 index 0000000000000000000000000000000000000000..cc674b303bb4d5cb86c05ec9b99891ee17633a13 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/offline/estimators/off_policy_estimator.py @@ -0,0 +1,248 @@ +import logging +from typing import Any, Dict, List + +import gymnasium as gym +import numpy as np +import tree + +from ray._common.deprecation import Deprecated +from ray.rllib.offline.offline_evaluator import OfflineEvaluator +from ray.rllib.policy import Policy +from ray.rllib.policy.sample_batch import SampleBatch, convert_ma_batch_to_sample_batch +from ray.rllib.utils.annotations import ( + DeveloperAPI, + ExperimentalAPI, + OverrideToImplementCustomLogic, +) +from ray.rllib.utils.numpy import convert_to_numpy +from ray.rllib.utils.policy import compute_log_likelihoods_from_input_dict +from ray.rllib.utils.typing import SampleBatchType, TensorType + +logger = logging.getLogger(__name__) + + +@DeveloperAPI +class OffPolicyEstimator(OfflineEvaluator): + """Interface for an off policy estimator for counterfactual evaluation.""" + + @DeveloperAPI + def __init__( + self, + policy: Policy, + gamma: float = 0.0, + epsilon_greedy: float = 0.0, + ): + """Initializes an OffPolicyEstimator instance. + + Args: + policy: Policy to evaluate. + gamma: Discount factor of the environment. + epsilon_greedy: The probability by which we act acording to a fully random + policy during deployment. With 1-epsilon_greedy we act according the target + policy. + # TODO (kourosh): convert the input parameters to a config dict. + """ + super().__init__(policy) + self.gamma = gamma + self.epsilon_greedy = epsilon_greedy + + @DeveloperAPI + def estimate_on_single_episode(self, episode: SampleBatch) -> Dict[str, Any]: + """Returns off-policy estimates for the given one episode. + + Args: + batch: The episode to calculate the off-policy estimates (OPE) on. The + episode must be a sample batch type that contains the fields "obs", + "actions", and "action_prob" and it needs to represent a + complete trajectory. + + Returns: + The off-policy estimates (OPE) calculated on the given episode. The returned + dict can be any arbitrary mapping of strings to metrics. + """ + raise NotImplementedError + + @DeveloperAPI + def estimate_on_single_step_samples( + self, + batch: SampleBatch, + ) -> Dict[str, List[float]]: + """Returns off-policy estimates for the batch of single timesteps. This is + highly optimized for bandits assuming each episode is a single timestep. + + Args: + batch: The batch to calculate the off-policy estimates (OPE) on. The + batch must be a sample batch type that contains the fields "obs", + "actions", and "action_prob". + + Returns: + The off-policy estimates (OPE) calculated on the given batch of single time + step samples. The returned dict can be any arbitrary mapping of strings to + a list of floats capturing the values per each record. + """ + raise NotImplementedError + + def on_before_split_batch_by_episode( + self, sample_batch: SampleBatch + ) -> SampleBatch: + """Called before the batch is split by episode. You can perform any + preprocessing on the batch that you want here. + e.g. adding done flags to the batch, or reseting some stats that you want to + track per episode later during estimation, .etc. + + Args: + sample_batch: The batch to split by episode. This contains multiple + episodes. + + Returns: + The modified batch before calling split_by_episode(). + """ + return sample_batch + + @OverrideToImplementCustomLogic + def on_after_split_batch_by_episode( + self, all_episodes: List[SampleBatch] + ) -> List[SampleBatch]: + """Called after the batch is split by episode. You can perform any + postprocessing on each episode that you want here. + e.g. computing advantage per episode, .etc. + + Args: + all_episodes: The list of episodes in the original batch. Each element is a + sample batch type that is a single episode. + """ + + return all_episodes + + @OverrideToImplementCustomLogic + def peek_on_single_episode(self, episode: SampleBatch) -> None: + """This is called on each episode before it is passed to + estimate_on_single_episode(). Using this method, you can get a peek at the + entire validation dataset before runnining the estimation. For examlpe if you + need to perform any normalizations of any sorts on the dataset, you can compute + the normalization parameters here. + + Args: + episode: The episode that is split from the original batch. This is a + sample batch type that is a single episode. + """ + pass + + @DeveloperAPI + def estimate( + self, batch: SampleBatchType, split_batch_by_episode: bool = True + ) -> Dict[str, Any]: + """Compute off-policy estimates. + + Args: + batch: The batch to calculate the off-policy estimates (OPE) on. The + batch must contain the fields "obs", "actions", and "action_prob". + split_batch_by_episode: Whether to split the batch by episode. + + Returns: + The off-policy estimates (OPE) calculated on the given batch. The returned + dict can be any arbitrary mapping of strings to metrics. + The dict consists of the following metrics: + - v_behavior: The discounted return averaged over episodes in the batch + - v_behavior_std: The standard deviation corresponding to v_behavior + - v_target: The estimated discounted return for `self.policy`, + averaged over episodes in the batch + - v_target_std: The standard deviation corresponding to v_target + - v_gain: v_target / max(v_behavior, 1e-8) + - v_delta: The difference between v_target and v_behavior. + """ + batch = convert_ma_batch_to_sample_batch(batch) + self.check_action_prob_in_batch(batch) + estimates_per_epsiode = [] + if split_batch_by_episode: + batch = self.on_before_split_batch_by_episode(batch) + all_episodes = batch.split_by_episode() + all_episodes = self.on_after_split_batch_by_episode(all_episodes) + for episode in all_episodes: + assert len(set(episode[SampleBatch.EPS_ID])) == 1, ( + "The episode must contain only one episode id. For some reason " + "the split_by_episode() method could not successfully split " + "the batch by episodes. Each row in the dataset should be " + "one episode. Check your evaluation dataset for errors." + ) + self.peek_on_single_episode(episode) + + for episode in all_episodes: + estimate_step_results = self.estimate_on_single_episode(episode) + estimates_per_epsiode.append(estimate_step_results) + + # turn a list of identical dicts into a dict of lists + estimates_per_epsiode = tree.map_structure( + lambda *x: list(x), *estimates_per_epsiode + ) + else: + # the returned dict is a mapping of strings to a list of floats + estimates_per_epsiode = self.estimate_on_single_step_samples(batch) + + estimates = { + "v_behavior": np.mean(estimates_per_epsiode["v_behavior"]), + "v_behavior_std": np.std(estimates_per_epsiode["v_behavior"]), + "v_target": np.mean(estimates_per_epsiode["v_target"]), + "v_target_std": np.std(estimates_per_epsiode["v_target"]), + } + estimates["v_gain"] = estimates["v_target"] / max(estimates["v_behavior"], 1e-8) + estimates["v_delta"] = estimates["v_target"] - estimates["v_behavior"] + + return estimates + + @DeveloperAPI + def check_action_prob_in_batch(self, batch: SampleBatchType) -> None: + """Checks if we support off policy estimation (OPE) on given batch. + + Args: + batch: The batch to check. + + Raises: + ValueError: In case `action_prob` key is not in batch + """ + + if "action_prob" not in batch: + raise ValueError( + "Off-policy estimation is not possible unless the inputs " + "include action probabilities (i.e., the policy is stochastic " + "and emits the 'action_prob' key). For DQN this means using " + "`exploration_config: {type: 'SoftQ'}`. You can also set " + "`off_policy_estimation_methods: {}` to disable estimation." + ) + + @ExperimentalAPI + def compute_action_probs(self, batch: SampleBatch): + log_likelihoods = compute_log_likelihoods_from_input_dict(self.policy, batch) + new_prob = np.exp(convert_to_numpy(log_likelihoods)) + + if self.epsilon_greedy > 0.0: + if not isinstance(self.policy.action_space, gym.spaces.Discrete): + raise ValueError( + "Evaluation with epsilon-greedy exploration is only supported " + "with discrete action spaces." + ) + eps = self.epsilon_greedy + new_prob = new_prob * (1 - eps) + eps / self.policy.action_space.n + + return new_prob + + @DeveloperAPI + def train(self, batch: SampleBatchType) -> Dict[str, Any]: + """Train a model for Off-Policy Estimation. + + Args: + batch: SampleBatch to train on + + Returns: + Any optional metrics to return from the estimator + """ + return {} + + @Deprecated( + old="OffPolicyEstimator.action_log_likelihood", + new="ray.rllib.utils.policy.compute_log_likelihoods_from_input_dict", + error=True, + ) + def action_log_likelihood(self, batch: SampleBatchType) -> TensorType: + log_likelihoods = compute_log_likelihoods_from_input_dict(self.policy, batch) + return convert_to_numpy(log_likelihoods) diff --git a/lib/python3.12/site-packages/ray/rllib/tuned_examples/cleanup_experiment.py b/lib/python3.12/site-packages/ray/rllib/tuned_examples/cleanup_experiment.py new file mode 100644 index 0000000000000000000000000000000000000000..749d3ed5e5222e2f7b74bafacec44f3c85c850fc --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/tuned_examples/cleanup_experiment.py @@ -0,0 +1,187 @@ +""" +This script automates cleaning up a benchmark/experiment run of some algo +against some config (with possibly more than one tune trial, +e.g. torch=grid_search([True, False])). + +Run `python cleanup_experiment.py --help` for more information. + +Use on an input directory with trial contents e.g.: +.. +IMPALA_BreakoutNoFrameskip-v4_0_use_pytorch=False_2020-05-11_10-17-54topr3h9k +IMPALA_BreakoutNoFrameskip-v4_0_use_pytorch=False_2020-05-11_13-59-35dqaetxnf +IMPALA_BreakoutNoFrameskip-v4_0_use_pytorch=False_2020-05-11_17-21-28tbhedw72 +IMPALA_BreakoutNoFrameskip-v4_2_use_pytorch=True_2020-05-11_10-17-54lv20cgn_ +IMPALA_BreakoutNoFrameskip-v4_2_use_pytorch=True_2020-05-11_13-59-35kwzhax_y +IMPALA_BreakoutNoFrameskip-v4_2_use_pytorch=True_2020-05-11_17-21-28a5j0s7za + +Then run: +>> python cleanup_experiment.py --experiment-dir [parent dir w/ trial sub-dirs] +>> --output-dir [your out dir] --results-filter dumb_col_2,superfluous_col3 +>> --results-max-size [max results file size in kb before(!) zipping] + +The script will create one output sub-dir for each trial and only copy +the configuration and the csv results (filtered and every nth row removed +based on the given args). +""" + +import argparse +import json +import os +import re +import shutil + +import yaml + +parser = argparse.ArgumentParser() +parser.add_argument( + "--experiment-dir", + type=str, + help="Experiment dir in which all sub-runs (seeds) are " + "located (as sub-dirs). Each sub0-run dir must contain the files: " + "params.json and progress.csv.", +) +parser.add_argument( + "--output-dir", + type=str, + help="The output dir, in which the cleaned up output will be placed.", +) +parser.add_argument( + "--results-filter", + type=str, + help="comma-separated list of csv fields to exclude.", + default="experiment_id,pid,hostname,node_ip,trial_id,hist_stats/episode_" + "reward,hist_stats/episode_lengths,experiment_tag", +) +parser.add_argument( + "--results-max-size", + type=int, + help="the max. size of the final results.csv file (in kb). Will erase " + "every nth line in the original input to reach that goal. " + "Use 0 for no limit (default=100).", + default=100, +) + + +def process_single_run(in_dir, out_dir): + exp_dir = os.listdir(in_dir) + + # Make sure trials dir is ok. + assert ( + "params.json" in exp_dir and "progress.csv" in exp_dir + ), "params.json or progress.csv not found in {}!".format(in_dir) + + os.makedirs(out_dir, exist_ok=True) + + for file in exp_dir: + absfile = os.path.join(in_dir, file) + # Config file -> Convert to yaml and move to output dir. + if file == "params.json": + assert os.path.isfile(absfile), "{} not a file!".format(file) + with open(absfile) as fp: + contents = json.load(fp) + with open(os.path.join(out_dir, "config.yaml"), "w") as fp: + yaml.dump(contents, fp) + # Progress csv file -> Filter out some columns, cut, and write to + # output_dir. + elif file == "progress.csv": + assert os.path.isfile(absfile), "{} not a file!".format(file) + col_idx_to_filter = [] + with open(absfile) as fp: + # Get column names. + col_names_orig = fp.readline().strip().split(",") + # Split by comma (abiding to quotes), filter out + # unwanted columns, then write to disk. + cols_to_filter = args.results_filter.split(",") + for i, c in enumerate(col_names_orig): + if c in cols_to_filter: + col_idx_to_filter.insert(0, i) + col_names = col_names_orig.copy() + for idx in col_idx_to_filter: + col_names.pop(idx) + absfile_out = os.path.join(out_dir, "progress.csv") + with open(absfile_out, "w") as out_fp: + print(",".join(col_names), file=out_fp) + while True: + line = fp.readline().strip() + if not line: + break + line = re.sub( + "(,{2,})", + lambda m: ",None" * (len(m.group()) - 1) + ",", + line, + ) + cols = re.findall('".+?"|[^,]+', line) + if len(cols) != len(col_names_orig): + continue + for idx in col_idx_to_filter: + cols.pop(idx) + print(",".join(cols), file=out_fp) + + # Reduce the size of the output file if necessary. + out_size = os.path.getsize(absfile_out) + max_size = args.results_max_size * 1024 + if 0 < max_size < out_size: + # Figure out roughly every which line we have to drop. + ratio = out_size / max_size + # If ratio > 2.0, we'll have to keep only every nth line. + if ratio > 2.0: + nth = out_size // max_size + os.system( + "awk 'NR==1||NR%{}==0' {} > {}.new".format( + nth, absfile_out, absfile_out + ) + ) + # If ratio < 2.0 (>1.0), we'll have to drop every nth line. + else: + nth = out_size // (out_size - max_size) + os.system( + "awk 'NR==1||NR%{}!=0' {} > {}.new".format( + nth, absfile_out, absfile_out + ) + ) + os.remove(absfile_out) + os.rename(absfile_out + ".new", absfile_out) + + # Zip progress.csv into results.zip. + zip_file = os.path.join(out_dir, "results.zip") + try: + os.remove(zip_file) + except FileNotFoundError: + pass + os.system( + "zip -j {} {}".format(zip_file, os.path.join(out_dir, "progress.csv")) + ) + os.remove(os.path.join(out_dir, "progress.csv")) + + # TBX events file -> Move as is. + elif re.search("^(events\\.out\\.|params\\.pkl)", file): + assert os.path.isfile(absfile), "{} not a file!".format(file) + shutil.copyfile(absfile, os.path.join(out_dir, file)) + + +if __name__ == "__main__": + args = parser.parse_args() + exp_dir = os.listdir(args.experiment_dir) + # Loop through all sub-directories. + for i, sub_run in enumerate(sorted(exp_dir)): + abspath = os.path.join(args.experiment_dir, sub_run) + # This is a seed run. + if os.path.isdir(abspath) and re.search( + "^(\\w+?)_(\\w+?-v\\d+)(_\\d+)", sub_run + ): + # Create meaningful output dir name: + # [algo]_[env]_[trial #]_[trial-config]_[date YYYY-MM-DD]. + cleaned_up_out = re.sub( + "^(\\w+?)_(\\w+?-v\\d+)(_\\d+)(_.+)?(_\\d{4}-\\d{2}-\\d{2})" + "_\\d{2}-\\d{2}-\\w+", + "{:02}_\\1_\\2\\4\\5".format(i), + sub_run, + ) + # Remove superflous `env=` specifier (anv always included in name). + cleaned_up_out = re.sub( + "^(.+)env=\\w+?-v\\d+,?(.+)", "\\1\\2", cleaned_up_out + ) + out_path = os.path.join(args.output_dir, cleaned_up_out) + process_single_run(abspath, out_path) + # Done. + print("done") diff --git a/lib/python3.12/site-packages/ray/rllib/utils/__init__.py b/lib/python3.12/site-packages/ray/rllib/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e761bb9ffbb5c2c3271369c6e2264ac64a87b88d --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/__init__.py @@ -0,0 +1,208 @@ +import contextlib +from collections import deque +from functools import partial +from typing import Any, Dict, List, Optional, Tuple, Union + +import tree + +from ray._common.deprecation import deprecation_warning +from ray.rllib.utils.annotations import DeveloperAPI, PublicAPI, override +from ray.rllib.utils.filter import Filter +from ray.rllib.utils.filter_manager import FilterManager +from ray.rllib.utils.framework import ( + try_import_jax, + try_import_tf, + try_import_tfp, + try_import_torch, +) +from ray.rllib.utils.numpy import ( + LARGE_INTEGER, + MAX_LOG_NN_OUTPUT, + MIN_LOG_NN_OUTPUT, + SMALL_NUMBER, + fc, + lstm, + one_hot, + relu, + sigmoid, + softmax, +) +from ray.rllib.utils.schedules import ( + ConstantSchedule, + ExponentialSchedule, + LinearSchedule, + PiecewiseSchedule, + PolynomialSchedule, +) +from ray.rllib.utils.test_utils import ( + check, + check_compute_single_action, + check_train_results, +) +from ray.tune.utils import deep_update, merge_dicts + + +@DeveloperAPI +def add_mixins(base, mixins, reversed=False): + """Returns a new class with mixins applied in priority order.""" + + mixins = list(mixins or []) + + while mixins: + if reversed: + + class new_base(base, mixins.pop()): + pass + + else: + + class new_base(mixins.pop(), base): + pass + + base = new_base + + return base + + +@DeveloperAPI +def force_list( + elements: Optional[Any] = None, to_tuple: bool = False +) -> Union[List, Tuple]: + """ + Makes sure `elements` is returned as a list, whether `elements` is a single + item, already a list, or a tuple. + + Args: + elements: The inputs as a single item, a list/tuple/deque of items, or None, + to be converted to a list/tuple. If None, returns empty list/tuple. + to_tuple: Whether to use tuple (instead of list). + + Returns: + The provided item in a list of size 1, or the provided items as a + list. If `elements` is None, returns an empty list. If `to_tuple` is True, + returns a tuple instead of a list. + """ + ctor = list + if to_tuple is True: + ctor = tuple + return ( + ctor() + if elements is None + else ctor(elements) + if type(elements) in [list, set, tuple, deque] + else ctor([elements]) + ) + + +@DeveloperAPI +def flatten_dict(nested: Dict[str, Any], sep="/", env_steps=0) -> Dict[str, Any]: + """ + Flattens a nested dict into a flat dict with joined keys. + + Note, this is used for better serialization of nested dictionaries + in `OfflinePreLearner.__call__` when called inside + `ray.data.Dataset.map_batches`. + + Note, this is used to return a `Dict[str, numpy.ndarray] from the + `__call__` method which is expected by Ray Data. + + Args: + nested: A nested dictionary. + sep: Separator to use when joining keys. + + Returns: + A flat dictionary where each key is a path of keys in the nested dict. + """ + flat = {} + # `dm_tree.flatten_with_path`` returns a list of `(path, leaf)` tuples. + for path, leaf in tree.flatten_with_path(nested): + # Create a single string key from the path. + key = sep.join(map(str, path)) + flat[key] = leaf + + return flat + + +@DeveloperAPI +def unflatten_dict(flat: Dict[str, Any], sep="/") -> Dict[str, Any]: + """ + Reconstructs a nested dict from a flat dict with joined keys. + + Note, this is used for better deserialization ofr nested dictionaries + in `Learner.update' calls in which a `ray.data.DataIterator` is used. + + Args: + flat: A flat dictionary with keys that are paths joined by `sep`. + sep: The separator used in the flat dictionary keys. + + Returns: + A nested dictionary. + """ + nested = {} + for compound_key, value in flat.items(): + # Split all keys by the separator. + keys = compound_key.split(sep) + current = nested + # Nest by the separated keys. + for key in keys[:-1]: + if key not in current: + current[key] = {} + current = current[key] + current[keys[-1]] = value + + return nested + + +@DeveloperAPI +class NullContextManager(contextlib.AbstractContextManager): + """No-op context manager""" + + def __init__(self): + pass + + def __enter__(self): + pass + + def __exit__(self, *args): + pass + + +force_tuple = partial(force_list, to_tuple=True) + +__all__ = [ + "add_mixins", + "check", + "check_compute_single_action", + "check_train_results", + "deep_update", + "deprecation_warning", + "fc", + "force_list", + "force_tuple", + "flatten_dict", + "unflatten_dict", + "lstm", + "merge_dicts", + "one_hot", + "override", + "relu", + "sigmoid", + "softmax", + "try_import_jax", + "try_import_tf", + "try_import_tfp", + "try_import_torch", + "ConstantSchedule", + "DeveloperAPI", + "ExponentialSchedule", + "Filter", + "FilterManager", + "LARGE_INTEGER", + "LinearSchedule", + "MAX_LOG_NN_OUTPUT", + "MIN_LOG_NN_OUTPUT", + "PiecewiseSchedule", + "PolynomialSchedule", + "PublicAPI", + "SMALL_NUMBER", +] diff --git a/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a62d428e60dd03338596be3900374289a58a4010 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/actor_manager.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/actor_manager.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d0751480f2047e6ac168555098147f0c8e2ac827 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/actor_manager.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/actors.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/actors.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..70de90e3991e1388c17287d42295ed06904b63b9 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/actors.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/annotations.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/annotations.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fb00500647d6f65837ebd81d4f6c4921f0fbba45 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/annotations.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/checkpoints.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/checkpoints.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6bb9a02ba94a20dd584a2b1bd5ae6b24f5306d4f Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/checkpoints.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/compression.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/compression.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bfdd8b0974887d40fc5220b86d5a6989f680a9cb Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/compression.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/error.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/error.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dddbde0350e5e2fe8912f1bac40fb24903481d08 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/error.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/filter.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/filter.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e83d90f467e5b81b185e3f3aefc9f021b51a29bc Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/filter.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/filter_manager.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/filter_manager.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..11abd3a57d22d488e784de65bd77e8077e29f8fd Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/filter_manager.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/framework.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/framework.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1d6619b6cfc6cceae9f373bc64e3075c7e4dbdfd Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/framework.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/from_config.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/from_config.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7ffe547b0f2e995796e814c2d66ec73566d8f4b2 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/from_config.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/images.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/images.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dc76907471d9886cce4a4d690643681bcfcf6202 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/images.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/lambda_defaultdict.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/lambda_defaultdict.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b4dd418370b0b731bbac08727ab4d4294780a76d Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/lambda_defaultdict.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/memory.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/memory.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d49d299620eabbcf26d131b9a8d998f9804142e2 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/memory.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/minibatch_utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/minibatch_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..bc6aaae6cda8d7dea99a0598244b95c3d50e951b Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/minibatch_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/numpy.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/numpy.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2734fa51acac381c662ac180f7520d2dc9fa68ad Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/numpy.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/policy.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/policy.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e924ff22003bbdb6f6b89c5c95b5d36e1aec9fe5 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/policy.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/serialization.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/serialization.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c6da4cc22d2e99f33d1de9589628b73db3376721 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/serialization.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/sgd.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/sgd.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..df9f2d39b6cb8959fec3eb316e8b58aa9bfbbb51 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/sgd.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/tensor_dtype.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/tensor_dtype.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..266006172d8e4e7ff8d88c3e5ddbc5d07b09df8b Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/tensor_dtype.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/test_utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/test_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2c06480be976aac1a3b20bd446da241c3982bec0 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/test_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/tf_run_builder.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/tf_run_builder.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c19eabb39dbb8331d3dcc223d80d86fe802d907a Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/tf_run_builder.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/tf_utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/tf_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d540ab00c14ebeaced2eee6fdb333b3047a2e8cb Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/tf_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/threading.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/threading.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7b4a6b24cb2bd96820f221d64f719967de95895c Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/threading.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/torch_utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/torch_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2ca78f9b32504d4a728a34536e0017cdec31fc45 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/torch_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/typing.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/typing.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f4c0630b1ba1f3cba7f954e8a395ef3b1f173d93 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/__pycache__/typing.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/actor_manager.py b/lib/python3.12/site-packages/ray/rllib/utils/actor_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..a82cbd1f35bf1170c56bbdb48e4cee06d4cafd69 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/actor_manager.py @@ -0,0 +1,1103 @@ +import copy +import logging +import sys +import time +from collections import defaultdict +from dataclasses import dataclass, field +from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple, Union + +import ray +from ray.actor import ActorHandle +from ray.exceptions import RayError, RayTaskError +from ray.rllib.utils.typing import T +from ray.util.annotations import DeveloperAPI + +logger = logging.getLogger(__name__) + + +@DeveloperAPI +class ResultOrError: + """A wrapper around a result or a RayError thrown during remote task/actor calls. + + This is used to return data from `FaultTolerantActorManager` that allows us to + distinguish between RayErrors (remote actor related) and valid results. + """ + + def __init__(self, result: Any = None, error: Exception = None): + """One and only one of result or error should be set. + + Args: + result: The result of the computation. Note that None is a valid result if + the remote function does not return anything. + error: Alternatively, the error that occurred during the computation. + """ + self._result = result + self._error = ( + # Easier to handle if we show the user the original error. + error.as_instanceof_cause() + if isinstance(error, RayTaskError) + else error + ) + + @property + def ok(self): + return self._error is None + + def get(self): + """Returns the result or the error.""" + if self._error: + return self._error + else: + return self._result + + +@DeveloperAPI +@dataclass +class CallResult: + """Represents a single result from a call to an actor. + + Each CallResult contains the index of the actor that was called + plus the result or error from the call. + """ + + actor_id: int + result_or_error: ResultOrError + tag: str + + @property + def ok(self): + """Passes through the ok property from the result_or_error.""" + return self.result_or_error.ok + + def get(self): + """Passes through the get method from the result_or_error.""" + return self.result_or_error.get() + + +@DeveloperAPI +class RemoteCallResults: + """Represents a list of results from calls to a set of actors. + + CallResults provides convenient APIs to iterate over the results + while skipping errors, etc. + + .. testcode:: + :skipif: True + + manager = FaultTolerantActorManager( + actors, max_remote_requests_in_flight_per_actor=2, + ) + results = manager.foreach_actor(lambda w: w.call()) + + # Iterate over all results ignoring errors. + for result in results.ignore_errors(): + print(result.get()) + """ + + class _Iterator: + """An iterator over the results of a remote call.""" + + def __init__(self, call_results: List[CallResult]): + self._call_results = call_results + + def __iter__(self) -> Iterator[CallResult]: + return self + + def __next__(self) -> CallResult: + if not self._call_results: + raise StopIteration + return self._call_results.pop(0) + + def __init__(self): + self.result_or_errors: List[CallResult] = [] + + def add_result(self, actor_id: int, result_or_error: ResultOrError, tag: str): + """Add index of a remote actor plus the call result to the list. + + Args: + actor_id: ID of the remote actor. + result_or_error: The result or error from the call. + tag: A description to identify the call. + """ + self.result_or_errors.append(CallResult(actor_id, result_or_error, tag)) + + def __iter__(self) -> Iterator[ResultOrError]: + """Return an iterator over the results.""" + # Shallow copy the list. + return self._Iterator(copy.copy(self.result_or_errors)) + + def __len__(self) -> int: + return len(self.result_or_errors) + + def ignore_errors(self) -> Iterator[ResultOrError]: + """Return an iterator over the results, skipping all errors.""" + return self._Iterator([r for r in self.result_or_errors if r.ok]) + + def ignore_ray_errors(self) -> Iterator[ResultOrError]: + """Return an iterator over the results, skipping only Ray errors. + + Similar to ignore_errors, but only skips Errors raised because of + remote actor problems (often get restored automatcially). + This is useful for callers that want to handle application errors differently + from Ray errors. + """ + return self._Iterator( + [r for r in self.result_or_errors if not isinstance(r.get(), RayError)] + ) + + +@DeveloperAPI +class FaultAwareApply: + @DeveloperAPI + def ping(self) -> str: + """Ping the actor. Can be used as a health check. + + Returns: + "pong" if actor is up and well. + """ + return "pong" + + @DeveloperAPI + def apply( + self, + func: Callable[[Any, Optional[Any], Optional[Any]], T], + *args, + **kwargs, + ) -> T: + """Calls the given function with this Actor instance. + + A generic interface for applying arbitrary member functions on a + remote actor. + + Args: + func: The function to call, with this actor as first + argument, followed by args, and kwargs. + args: Optional additional args to pass to the function call. + kwargs: Optional additional kwargs to pass to the function call. + + Returns: + The return value of the function call. + """ + try: + return func(self, *args, **kwargs) + except Exception as e: + # Actor should be recreated by Ray. + if self.config.restart_failed_env_runners: + logger.exception(f"Worker exception caught during `apply()`: {e}") + # Small delay to allow logs messages to propagate. + time.sleep(self.config.delay_between_env_runner_restarts_s) + # Kill this worker so Ray Core can restart it. + sys.exit(1) + # Actor should be left dead. + else: + raise e + + +@DeveloperAPI +class FaultTolerantActorManager: + """A manager that is aware of the healthiness of remote actors. + + .. testcode:: + + import time + import ray + from ray.rllib.utils.actor_manager import FaultTolerantActorManager + + @ray.remote + class MyActor: + def apply(self, fn): + return fn(self) + + def do_something(self): + return True + + actors = [MyActor.remote() for _ in range(3)] + manager = FaultTolerantActorManager( + actors, max_remote_requests_in_flight_per_actor=2, + ) + + # Synchronous remote calls. + results = manager.foreach_actor(lambda actor: actor.do_something()) + # Print results ignoring returned errors. + print([r.get() for r in results.ignore_errors()]) + + # Asynchronous remote calls. + manager.foreach_actor_async(lambda actor: actor.do_something()) + time.sleep(2) # Wait for the tasks to finish. + for r in manager.fetch_ready_async_reqs(): + # Handle result and errors. + if r.ok: + print(r.get()) + else: + print("Error: {}".format(r.get())) + """ + + @dataclass + class _ActorState: + """State of a single actor.""" + + # Num of outstanding async requests for this actor by tag. + num_in_flight_async_requests_by_tag: Dict[Optional[str], int] = field( + default_factory=dict + ) + # Whether this actor is in a healthy state. + is_healthy: bool = True + + def get_num_in_flight_requests(self, tag: Optional[str] = None) -> int: + """Get number of in-flight requests for a specific tag or all tags.""" + if tag is None: + return sum(self.num_in_flight_async_requests_by_tag.values()) + return self.num_in_flight_async_requests_by_tag.get(tag, 0) + + def increment_requests(self, tag: Optional[str] = None) -> None: + """Increment the count of in-flight requests for a tag.""" + if tag not in self.num_in_flight_async_requests_by_tag: + self.num_in_flight_async_requests_by_tag[tag] = 0 + self.num_in_flight_async_requests_by_tag[tag] += 1 + + def decrement_requests(self, tag: Optional[str] = None) -> None: + """Decrement the count of in-flight requests for a tag.""" + if tag in self.num_in_flight_async_requests_by_tag: + self.num_in_flight_async_requests_by_tag[tag] -= 1 + if self.num_in_flight_async_requests_by_tag[tag] <= 0: + del self.num_in_flight_async_requests_by_tag[tag] + + def __init__( + self, + actors: Optional[List[ActorHandle]] = None, + max_remote_requests_in_flight_per_actor: int = 2, + init_id: int = 0, + ): + """Construct a FaultTolerantActorManager. + + Args: + actors: A list of ray remote actors to manage on. These actors must have an + ``apply`` method which takes a function with only one parameter (the + actor instance itself). + max_remote_requests_in_flight_per_actor: The maximum number of remote + requests that can be in flight per actor. Any requests made to the pool + that cannot be scheduled because the limit has been reached will be + dropped. This only applies to the asynchronous remote call mode. + init_id: The initial ID to use for the next remote actor. Default is 0. + """ + # For round-robin style async requests, keep track of which actor to send + # a new func next (current). + self._next_id = self._current_actor_id = init_id + + # Actors are stored in a map and indexed by a unique (int) ID. + self._actors: Dict[int, ActorHandle] = {} + self._remote_actor_states: Dict[int, self._ActorState] = {} + self._restored_actors = set() + self.add_actors(actors or []) + + # Maps outstanding async requests to the IDs of the actor IDs that + # are executing them. + self._in_flight_req_to_actor_id: Dict[ray.ObjectRef, int] = {} + + self._max_remote_requests_in_flight_per_actor = ( + max_remote_requests_in_flight_per_actor + ) + + # Useful metric. + self._num_actor_restarts = 0 + + @DeveloperAPI + def actor_ids(self) -> List[int]: + """Returns a list of all worker IDs (healthy or not).""" + return list(self._actors.keys()) + + @DeveloperAPI + def healthy_actor_ids(self) -> List[int]: + """Returns a list of worker IDs that are healthy.""" + return [k for k, v in self._remote_actor_states.items() if v.is_healthy] + + @DeveloperAPI + def add_actors(self, actors: List[ActorHandle]): + """Add a list of actors to the pool. + + Args: + actors: A list of ray remote actors to be added to the pool. + """ + for actor in actors: + self._actors[self._next_id] = actor + self._remote_actor_states[self._next_id] = self._ActorState() + self._next_id += 1 + + @DeveloperAPI + def remove_actor(self, actor_id: int) -> ActorHandle: + """Remove an actor from the pool. + + Args: + actor_id: ID of the actor to remove. + + Returns: + Handle to the actor that was removed. + """ + actor = self._actors[actor_id] + + # Remove the actor from the pool. + del self._actors[actor_id] + del self._remote_actor_states[actor_id] + self._restored_actors.discard(actor_id) + self._remove_async_state(actor_id) + + return actor + + @DeveloperAPI + def num_actors(self) -> int: + """Return the total number of actors in the pool.""" + return len(self._actors) + + @DeveloperAPI + def num_healthy_actors(self) -> int: + """Return the number of healthy remote actors.""" + return sum(s.is_healthy for s in self._remote_actor_states.values()) + + @DeveloperAPI + def total_num_restarts(self) -> int: + """Return the number of remote actors that have been restarted.""" + return self._num_actor_restarts + + @DeveloperAPI + def num_outstanding_async_reqs(self, tag: Optional[str] = None) -> int: + """Return the number of outstanding async requests.""" + return sum( + s.get_num_in_flight_requests(tag) + for s in self._remote_actor_states.values() + ) + + @DeveloperAPI + def is_actor_healthy(self, actor_id: int) -> bool: + """Whether a remote actor is in healthy state. + + Args: + actor_id: ID of the remote actor. + + Returns: + True if the actor is healthy, False otherwise. + """ + if actor_id not in self._remote_actor_states: + raise ValueError(f"Unknown actor id: {actor_id}") + return self._remote_actor_states[actor_id].is_healthy + + @DeveloperAPI + def set_actor_state(self, actor_id: int, healthy: bool) -> None: + """Update activate state for a specific remote actor. + + Args: + actor_id: ID of the remote actor. + healthy: Whether the remote actor is healthy. + """ + if actor_id not in self._remote_actor_states: + raise ValueError(f"Unknown actor id: {actor_id}") + + was_healthy = self._remote_actor_states[actor_id].is_healthy + # Set from unhealthy to healthy -> Add to restored set. + if not was_healthy and healthy: + self._restored_actors.add(actor_id) + # Set from healthy to unhealthy -> Remove from restored set. + elif was_healthy and not healthy: + self._restored_actors.discard(actor_id) + + self._remote_actor_states[actor_id].is_healthy = healthy + + if not healthy: + # Remove any async states. + self._remove_async_state(actor_id) + + @DeveloperAPI + def clear(self): + """Clean up managed actors.""" + for actor in self._actors.values(): + ray.kill(actor) + self._actors.clear() + self._remote_actor_states.clear() + self._restored_actors.clear() + self._in_flight_req_to_actor_id.clear() + + @DeveloperAPI + def foreach_actor( + self, + func: Union[Callable[[Any], Any], List[Callable[[Any], Any]], str, List[str]], + *, + kwargs: Optional[Union[Dict[str, Any], List[Dict[str, Any]]]] = None, + healthy_only: bool = True, + remote_actor_ids: Optional[List[int]] = None, + timeout_seconds: Optional[float] = None, + return_obj_refs: bool = False, + mark_healthy: bool = False, + ) -> RemoteCallResults: + """Calls the given function with each actor instance as arg. + + Automatically marks actors unhealthy if they crash during the remote call. + + Args: + func: A single Callable applied to all specified remote actors or a list + of Callables, that get applied on the list of specified remote actors. + In the latter case, both list of Callables and list of specified actors + must have the same length. Alternatively, you can use the name of the + remote method to be called, instead, or a list of remote method names. + kwargs: An optional single kwargs dict or a list of kwargs dict matching the + list of provided `func` or `remote_actor_ids`. In the first case (single + dict), use `kwargs` on all remote calls. The latter case (list of + dicts) allows you to define individualized kwarg dicts per actor. + healthy_only: If True, applies `func` only to actors currently tagged + "healthy", otherwise to all actors. If `healthy_only=False` and + `mark_healthy=True`, will send `func` to all actors and mark those + actors "healthy" that respond to the request within `timeout_seconds` + and are currently tagged as "unhealthy". + remote_actor_ids: Apply func on a selected set of remote actors. Use None + (default) for all actors. + timeout_seconds: Time to wait (in seconds) for results. Set this to 0.0 for + fire-and-forget. Set this to None (default) to wait infinitely (i.e. for + synchronous execution). + return_obj_refs: whether to return ObjectRef instead of actual results. + Note, for fault tolerance reasons, these returned ObjectRefs should + never be resolved with ray.get() outside of the context of this manager. + mark_healthy: Whether to mark all those actors healthy again that are + currently marked unhealthy AND that returned results from the remote + call (within the given `timeout_seconds`). + Note that actors are NOT set unhealthy, if they simply time out + (only if they return a RayActorError). + Also not that this setting is ignored if `healthy_only=True` (b/c this + setting only affects actors that are currently tagged as unhealthy). + + Returns: + The list of return values of all calls to `func(actor)`. The values may be + actual data returned or exceptions raised during the remote call in the + format of RemoteCallResults. + """ + remote_actor_ids = remote_actor_ids or self.actor_ids() + if healthy_only: + func, kwargs, remote_actor_ids = self._filter_by_healthy_state( + func=func, kwargs=kwargs, remote_actor_ids=remote_actor_ids + ) + + # Send out remote requests. + remote_calls = self._call_actors( + func=func, + kwargs=kwargs, + remote_actor_ids=remote_actor_ids, + ) + + # Collect remote request results (if available given timeout and/or errors). + _, remote_results = self._fetch_result( + remote_actor_ids=remote_actor_ids, + remote_calls=remote_calls, + tags=[None] * len(remote_calls), + timeout_seconds=timeout_seconds, + return_obj_refs=return_obj_refs, + mark_healthy=mark_healthy, + ) + + return remote_results + + @DeveloperAPI + def foreach_actor_async( + self, + func: Union[Callable[[Any], Any], List[Callable[[Any], Any]], str, List[str]], + tag: Optional[str] = None, + *, + kwargs: Optional[Union[Dict[str, Any], List[Dict[str, Any]]]] = None, + healthy_only: bool = True, + remote_actor_ids: Optional[List[int]] = None, + ) -> int: + """Calls given functions against each actors without waiting for results. + + Args: + func: A single Callable applied to all specified remote actors or a list + of Callables, that get applied on the list of specified remote actors. + In the latter case, both list of Callables and list of specified actors + must have the same length. Alternatively, you can use the name of the + remote method to be called, instead, or a list of remote method names. + tag: A tag to identify the results from this async call. + kwargs: An optional single kwargs dict or a list of kwargs dict matching the + list of provided `func` or `remote_actor_ids`. In the first case (single + dict), use `kwargs` on all remote calls. The latter case (list of + dicts) allows you to define individualized kwarg dicts per actor. + healthy_only: If True, applies `func` only to actors currently tagged + "healthy", otherwise to all actors. If `healthy_only=False` and + later, `self.fetch_ready_async_reqs()` is called with + `mark_healthy=True`, will send `func` to all actors and mark those + actors "healthy" that respond to the request within `timeout_seconds` + and are currently tagged as "unhealthy". + remote_actor_ids: Apply func on a selected set of remote actors. + Note, for fault tolerance reasons, these returned ObjectRefs should + never be resolved with ray.get() outside of the context of this manager. + + Returns: + The number of async requests that are actually fired. + """ + # TODO(avnishn, jungong): so thinking about this a bit more, it would be the + # best if we can attach multiple tags to an async all, like basically this + # parameter should be tags: + # For sync calls, tags would be (). + # For async call users, they can attached multiple tags for a single call, like + # ("rollout_worker", "sync_weight"). + # For async fetch result, we can also specify a single, or list of tags. For + # example, ("eval", "sample") will fetch all the sample() calls on eval + # workers. + if not remote_actor_ids: + remote_actor_ids = self.actor_ids() + + num_calls = ( + len(func) + if isinstance(func, list) + else len(kwargs) + if isinstance(kwargs, list) + else len(remote_actor_ids) + ) + + # Perform round-robin assignment of all provided calls for any number of our + # actors. Note that this way, some actors might receive more than 1 request in + # this call. + if num_calls != len(remote_actor_ids): + remote_actor_ids = [ + (self._current_actor_id + i) % self.num_actors() + for i in range(num_calls) + ] + # Update our round-robin pointer. + self._current_actor_id += num_calls + self._current_actor_id %= self.num_actors() + + if healthy_only: + func, kwargs, remote_actor_ids = self._filter_by_healthy_state( + func=func, kwargs=kwargs, remote_actor_ids=remote_actor_ids + ) + + num_calls_to_make: Dict[int, int] = defaultdict(lambda: 0) + # Drop calls to actors that are too busy for this specific tag. + if isinstance(func, list): + assert len(func) == len(remote_actor_ids) + limited_func = [] + limited_kwargs = [] + limited_remote_actor_ids = [] + for i, (f, raid) in enumerate(zip(func, remote_actor_ids)): + num_outstanding_reqs_for_tag = self._remote_actor_states[ + raid + ].get_num_in_flight_requests(tag) + if ( + num_outstanding_reqs_for_tag + num_calls_to_make[raid] + < self._max_remote_requests_in_flight_per_actor + ): + num_calls_to_make[raid] += 1 + k = kwargs[i] if isinstance(kwargs, list) else (kwargs or {}) + limited_func.append(f) + limited_kwargs.append(k) + limited_remote_actor_ids.append(raid) + else: + limited_func = func + limited_kwargs = kwargs + limited_remote_actor_ids = [] + for raid in remote_actor_ids: + num_outstanding_reqs_for_tag = self._remote_actor_states[ + raid + ].get_num_in_flight_requests(tag) + if ( + num_outstanding_reqs_for_tag + num_calls_to_make[raid] + < self._max_remote_requests_in_flight_per_actor + ): + num_calls_to_make[raid] += 1 + limited_remote_actor_ids.append(raid) + + if not limited_remote_actor_ids: + return 0 + + remote_calls = self._call_actors( + func=limited_func, + kwargs=limited_kwargs, + remote_actor_ids=limited_remote_actor_ids, + ) + + # Save these as outstanding requests. + for id, call in zip(limited_remote_actor_ids, remote_calls): + self._remote_actor_states[id].increment_requests(tag) + self._in_flight_req_to_actor_id[call] = (tag, id) + + return len(remote_calls) + + @DeveloperAPI + def fetch_ready_async_reqs( + self, + *, + tags: Union[str, List[str], Tuple[str, ...]] = (), + timeout_seconds: Optional[float] = 0.0, + return_obj_refs: bool = False, + mark_healthy: bool = False, + ) -> RemoteCallResults: + """Get results from outstanding async requests that are ready. + + Automatically mark actors unhealthy if they fail to respond. + + Note: If tags is an empty tuple then results from all ready async requests are + returned. + + Args: + timeout_seconds: ray.get() timeout. Default is 0, which only fetched those + results (immediately) that are already ready. + tags: A tag or a list of tags to identify the results from this async call. + return_obj_refs: Whether to return ObjectRef instead of actual results. + mark_healthy: Whether to mark all those actors healthy again that are + currently marked unhealthy AND that returned results from the remote + call (within the given `timeout_seconds`). + Note that actors are NOT set to unhealthy, if they simply time out, + meaning take a longer time to fulfil the remote request. We only ever + mark an actor unhealthy, if they raise a RayActorError inside the remote + request. + Also note that this settings is ignored if the preceding + `foreach_actor_async()` call used the `healthy_only=True` argument (b/c + `mark_healthy` only affects actors that are currently tagged as + unhealthy). + + Returns: + A list of return values of all calls to `func(actor)` that are ready. + The values may be actual data returned or exceptions raised during the + remote call in the format of RemoteCallResults. + """ + # Construct the list of in-flight requests filtered by tag. + remote_calls, remote_actor_ids, valid_tags = self._filter_calls_by_tag(tags) + ready, remote_results = self._fetch_result( + remote_actor_ids=remote_actor_ids, + remote_calls=remote_calls, + tags=valid_tags, + timeout_seconds=timeout_seconds, + return_obj_refs=return_obj_refs, + mark_healthy=mark_healthy, + ) + + for obj_ref, result in zip(ready, remote_results): + # Get the tag for this request and decrease outstanding request count by 1. + if obj_ref in self._in_flight_req_to_actor_id: + tag, actor_id = self._in_flight_req_to_actor_id[obj_ref] + self._remote_actor_states[result.actor_id].decrement_requests(tag) + # Remove this call from the in-flight list. + del self._in_flight_req_to_actor_id[obj_ref] + + return remote_results + + @DeveloperAPI + def foreach_actor_async_fetch_ready( + self, + func: Union[Callable[[Any], Any], List[Callable[[Any], Any]], str, List[str]], + tag: Optional[str] = None, + *, + kwargs: Optional[Union[Dict[str, Any], List[Dict[str, Any]]]] = None, + timeout_seconds: Optional[float] = 0.0, + return_obj_refs: bool = False, + mark_healthy: bool = False, + healthy_only: bool = True, + remote_actor_ids: Optional[List[int]] = None, + ignore_ray_errors: bool = True, + return_actor_ids: bool = False, + ) -> List[Union[Tuple[int, Any], Any]]: + """Calls the given function asynchronously and returns previous results if any. + + This is a convenience function that calls `fetch_ready_async_reqs()` to get + previous results and then `foreach_actor_async()` to start new async calls. + + Args: + func: A single Callable applied to all specified remote actors or a list + of Callables, that get applied on the list of specified remote actors. + In the latter case, both list of Callables and list of specified actors + must have the same length. Alternatively, you can use the name of the + remote method to be called, instead, or a list of remote method names. + tag: A tag to identify the results from this async call. + kwargs: An optional single kwargs dict or a list of kwargs dict matching the + list of provided `func` or `remote_actor_ids`. In the first case (single + dict), use `kwargs` on all remote calls. The latter case (list of + dicts) allows you to define individualized kwarg dicts per actor. + timeout_seconds: Time to wait for results from previous calls. Default is 0, + meaning those requests that are already ready. + return_obj_refs: Whether to return ObjectRef instead of actual results. + mark_healthy: Whether to mark all those actors healthy again that are + currently marked unhealthy AND that returned results from the remote + call (within the given `timeout_seconds`). + healthy_only: Apply `func` on known-to-be healthy actors only. + remote_actor_ids: Apply func on a selected set of remote actors. + ignore_ray_errors: Whether to ignore RayErrors in results. + return_actor_ids: Whether to return actor IDs in the results. + If True, the results will be a list of (actor_id, result) tuples. + If False, the results will be a list of results. + Returns: + The results from previous async requests that were ready. + """ + # First fetch any ready results from previous async calls + remote_results = self.fetch_ready_async_reqs( + tags=tag, + timeout_seconds=timeout_seconds, + return_obj_refs=return_obj_refs, + mark_healthy=mark_healthy, + ) + + # Then start new async calls + self.foreach_actor_async( + func, + tag=tag, + kwargs=kwargs, + healthy_only=healthy_only, + remote_actor_ids=remote_actor_ids, + ) + + # Handle errors the same way as fetch_ready_async_reqs does + FaultTolerantActorManager.handle_remote_call_result_errors( + remote_results, + ignore_ray_errors=ignore_ray_errors, + ) + + if return_actor_ids: + return [(r.actor_id, r.get()) for r in remote_results.ignore_errors()] + else: + return [r.get() for r in remote_results.ignore_errors()] + + @staticmethod + def handle_remote_call_result_errors( + results_or_errors: RemoteCallResults, + *, + ignore_ray_errors: bool, + ) -> None: + """Checks given results for application errors and raises them if necessary. + + Args: + results_or_errors: The results or errors to check. + ignore_ray_errors: Whether to ignore RayErrors within the elements of + `results_or_errors`. + """ + for result_or_error in results_or_errors: + # Good result. + if result_or_error.ok: + continue + # RayError, but we ignore it. + elif ignore_ray_errors: + logger.exception(result_or_error.get()) + # Raise RayError. + else: + raise result_or_error.get() + + @DeveloperAPI + def probe_unhealthy_actors( + self, + timeout_seconds: Optional[float] = None, + mark_healthy: bool = False, + ) -> List[int]: + """Ping all unhealthy actors to try bringing them back. + + Args: + timeout_seconds: Timeout in seconds (to avoid pinging hanging workers + indefinitely). + mark_healthy: Whether to mark all those actors healthy again that are + currently marked unhealthy AND that respond to the `ping` remote request + (within the given `timeout_seconds`). + Note that actors are NOT set to unhealthy, if they simply time out, + meaning take a longer time to fulfil the remote request. We only ever + mark and actor unhealthy, if they return a RayActorError from the remote + request. + Also note that this settings is ignored if `healthy_only=True` (b/c this + setting only affects actors that are currently tagged as unhealthy). + + Returns: + A list of actor IDs that were restored by the `ping.remote()` call PLUS + those actors that were previously restored via other remote requests. + The cached set of such previously restored actors will be erased in this + call. + """ + # Collect recently restored actors (from `self._fetch_result` calls other than + # the one triggered here via the `ping`). + already_restored_actors = list(self._restored_actors) + + # Which actors are currently marked unhealthy? + unhealthy_actor_ids = [ + actor_id + for actor_id in self.actor_ids() + if not self.is_actor_healthy(actor_id) + ] + # Some unhealthy actors -> `ping()` all of them to trigger a new fetch and + # gather the just restored ones (b/c of a successful `ping` response). + just_restored_actors = [] + if unhealthy_actor_ids: + remote_results = self.foreach_actor( + func=lambda actor: actor.ping(), + remote_actor_ids=unhealthy_actor_ids, + healthy_only=False, # We specifically want to ping unhealthy actors. + timeout_seconds=timeout_seconds, + return_obj_refs=False, + mark_healthy=mark_healthy, + ) + just_restored_actors = [ + result.actor_id for result in remote_results if result.ok + ] + + # Clear out previously restored actors (b/c of other successful request + # responses, outside of this method). + self._restored_actors.clear() + + # Return all restored actors (previously and just). + return already_restored_actors + just_restored_actors + + def _call_actors( + self, + func: Union[Callable[[Any], Any], List[Callable[[Any], Any]], str, List[str]], + *, + kwargs: Optional[Union[Dict[str, Any], List[Dict[str, Any]]]] = None, + remote_actor_ids: List[int] = None, + ) -> List[ray.ObjectRef]: + """Apply functions on a list of remote actors. + + Args: + func: A single Callable applied to all specified remote actors or a list + of Callables, that get applied on the list of specified remote actors. + In the latter case, both list of Callables and list of specified actors + must have the same length. Alternatively, you can use the name of the + remote method to be called, instead, or a list of remote method names. + kwargs: An optional single kwargs dict or a list of kwargs dict matching the + list of provided `func` or `remote_actor_ids`. In the first case (single + dict), use `kwargs` on all remote calls. The latter case (list of + dicts) allows you to define individualized kwarg dicts per actor. + remote_actor_ids: Apply func on this selected set of remote actors. + + Returns: + A list of ObjectRefs returned from the remote calls. + """ + if remote_actor_ids is None: + remote_actor_ids = self.actor_ids() + + calls = [] + if isinstance(func, list): + assert len(remote_actor_ids) == len( + func + ), "Funcs must have the same number of callables as actor indices." + + assert isinstance( + kwargs, list + ), "If func is a list of functions, kwargs has to be a list of kwargs." + + for i, (raid, f) in enumerate(zip(remote_actor_ids, func)): + if isinstance(f, str): + calls.append( + getattr(self._actors[raid], f).remote( + **( + kwargs[i] + if isinstance(kwargs, list) + else (kwargs or {}) + ) + ) + ) + else: + calls.append(self._actors[raid].apply.remote(f)) + elif isinstance(func, str): + for i, raid in enumerate(remote_actor_ids): + calls.append( + getattr(self._actors[raid], func).remote( + **(kwargs[i] if isinstance(kwargs, list) else (kwargs or {})) + ) + ) + else: + for raid in remote_actor_ids: + calls.append(self._actors[raid].apply.remote(func=func, **kwargs or {})) + + return calls + + @DeveloperAPI + def _fetch_result( + self, + *, + remote_actor_ids: List[int], + remote_calls: List[ray.ObjectRef], + tags: List[str], + timeout_seconds: Optional[float] = None, + return_obj_refs: bool = False, + mark_healthy: bool = False, + ) -> Tuple[List[ray.ObjectRef], RemoteCallResults]: + """Try fetching results from remote actor calls. + + Mark whether an actor is healthy or not accordingly. + + Args: + remote_actor_ids: IDs of the actors these remote + calls were fired against. + remote_calls: List of remote calls to fetch. + tags: List of tags used for identifying the remote calls. + timeout_seconds: Timeout (in sec) for the ray.wait() call. Default is None, + meaning wait indefinitely for all results. + return_obj_refs: Whether to return ObjectRef instead of actual results. + mark_healthy: Whether to mark certain actors healthy based on the results + of these remote calls. Useful, for example, to make sure actors + do not come back without proper state restoration. + + Returns: + A list of ready ObjectRefs mapping to the results of those calls. + """ + # Notice that we do not return the refs to any unfinished calls to the + # user, since it is not safe to handle such remote actor calls outside the + # context of this actor manager. These requests are simply dropped. + timeout = float(timeout_seconds) if timeout_seconds is not None else None + + # This avoids calling ray.init() in the case of 0 remote calls. + # This is useful if the number of remote workers is 0. + if not remote_calls: + return [], RemoteCallResults() + + readies, _ = ray.wait( + remote_calls, + num_returns=len(remote_calls), + timeout=timeout, + # Make sure remote results are fetched locally in parallel. + fetch_local=not return_obj_refs, + ) + + # Remote data should already be fetched to local object store at this point. + remote_results = RemoteCallResults() + for ready in readies: + # Find the corresponding actor ID for this remote call. + actor_id = remote_actor_ids[remote_calls.index(ready)] + tag = tags[remote_calls.index(ready)] + + # If caller wants ObjectRefs, return directly without resolving. + if return_obj_refs: + remote_results.add_result(actor_id, ResultOrError(result=ready), tag) + continue + + # Try getting the ready results. + try: + result = ray.get(ready) + + # Any error type other than `RayError` happening during ray.get() -> + # Throw exception right here (we don't know how to handle these non-remote + # worker issues and should therefore crash). + except RayError as e: + # Return error to the user. + remote_results.add_result(actor_id, ResultOrError(error=e), tag) + + # Mark the actor as unhealthy, take it out of service, and wait for + # Ray Core to restore it. + if self.is_actor_healthy(actor_id): + logger.error( + f"Ray error ({str(e)}), taking actor {actor_id} out of service." + ) + self.set_actor_state(actor_id, healthy=False) + + # If no errors, add result to `RemoteCallResults` to be returned. + else: + # Return valid result to the user. + remote_results.add_result(actor_id, ResultOrError(result=result), tag) + + # Actor came back from an unhealthy state. Mark this actor as healthy + # and add it to our healthy set. + if mark_healthy and not self.is_actor_healthy(actor_id): + logger.warning( + f"Bringing previously unhealthy, now-healthy actor {actor_id} " + "back into service." + ) + self.set_actor_state(actor_id, healthy=True) + self._num_actor_restarts += 1 + + # Make sure, to-be-returned results are sound. + assert len(readies) == len(remote_results) + + return readies, remote_results + + def _filter_by_healthy_state( + self, + *, + func: Union[Callable[[Any], Any], List[Callable[[Any], Any]]], + kwargs: Optional[Union[Dict, List[Dict]]] = None, + remote_actor_ids: List[int], + ): + """Filter out func and remote worker ids by actor state. + + Args: + func: A single, or a list of Callables. + kwargs: An optional single kwargs dict or a list of kwargs dicts matching + the list of provided `func` or `remote_actor_ids`. In case of a single + dict, uses `kwargs` on all remote calls. In case of a list of dicts, + the given kwarg dicts are per actor `func` or per `remote_actor_ids`. + remote_actor_ids: IDs of potential remote workers to apply func on. + + Returns: + A tuple of (filtered func, filtered remote worker ids). + """ + if isinstance(func, list): + assert len(remote_actor_ids) == len( + func + ), "Func must have the same number of callables as remote actor ids." + # We are given a list of functions to apply. + # Need to filter the functions together with worker IDs. + temp_func = [] + temp_remote_actor_ids = [] + temp_kwargs = [] + for i, (f, raid) in enumerate(zip(func, remote_actor_ids)): + if self.is_actor_healthy(raid): + k = kwargs[i] if isinstance(kwargs, list) else (kwargs or {}) + temp_func.append(f) + temp_kwargs.append(k) + temp_remote_actor_ids.append(raid) + func = temp_func + kwargs = temp_kwargs + remote_actor_ids = temp_remote_actor_ids + else: + # Simply filter the worker IDs. + remote_actor_ids = [i for i in remote_actor_ids if self.is_actor_healthy(i)] + + return func, kwargs, remote_actor_ids + + def _filter_calls_by_tag( + self, tags: Optional[Union[str, List[str], Tuple[str, ...]]] = None + ) -> Tuple[List[ray.ObjectRef], List[ActorHandle], List[str]]: + """Return all the in flight requests that match the given tags, if any. + + Args: + tags: A str or a list/tuple of str. If tags is empty or None, return all the in + flight requests. + + Returns: + A tuple consisting of a list of the remote calls that match the tag(s), + a list of the corresponding remote actor IDs for these calls (same length), + and a list of the tags corresponding to these calls (same length). + """ + if tags is None: + tags = set() + elif isinstance(tags, str): + tags = {tags} + elif isinstance(tags, (list, tuple)): + tags = set(tags) + else: + raise ValueError( + f"tags must be either a str or a list/tuple of str, got {type(tags)}." + ) + remote_calls = [] + remote_actor_ids = [] + valid_tags = [] + for call, (tag, actor_id) in self._in_flight_req_to_actor_id.items(): + # the default behavior is to return all ready results. + if len(tags) == 0 or tag in tags: + remote_calls.append(call) + remote_actor_ids.append(actor_id) + valid_tags.append(tag) + + return remote_calls, remote_actor_ids, valid_tags + + def _remove_async_state(self, actor_id: int): + """Remove internal async state of for a given actor. + + This is called when an actor is removed from the pool or being marked + unhealthy. + + Args: + actor_id: The id of the actor. + """ + # Remove any outstanding async requests for this actor. + # Use `list` here to not change a looped generator while we mutate the + # underlying dict. + for req, (tag, id) in list(self._in_flight_req_to_actor_id.items()): + if id == actor_id: + del self._in_flight_req_to_actor_id[req] + + # Clear all tag-based request counts for this actor + if actor_id in self._remote_actor_states: + self._remote_actor_states[ + actor_id + ].num_in_flight_async_requests_by_tag.clear() + + def actors(self): + # TODO(jungong) : remove this API once EnvRunnerGroup.remote_workers() + # and EnvRunnerGroup._remote_workers() are removed. + return self._actors diff --git a/lib/python3.12/site-packages/ray/rllib/utils/actors.py b/lib/python3.12/site-packages/ray/rllib/utils/actors.py new file mode 100644 index 0000000000000000000000000000000000000000..382613713e2f46038274e977bcdd16d4e88583f4 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/actors.py @@ -0,0 +1,258 @@ +import logging +import platform +from collections import defaultdict, deque +from typing import Any, Dict, List, Optional, Sequence, Tuple, Type + +import ray +from ray.actor import ActorClass, ActorHandle + +logger = logging.getLogger(__name__) + + +class TaskPool: + """Helper class for tracking the status of many in-flight actor tasks.""" + + def __init__(self): + self._tasks = {} + self._objects = {} + self._fetching = deque() + + def add(self, worker, all_obj_refs): + if isinstance(all_obj_refs, list): + obj_ref = all_obj_refs[0] + else: + obj_ref = all_obj_refs + self._tasks[obj_ref] = worker + self._objects[obj_ref] = all_obj_refs + + def completed(self, blocking_wait=False): + pending = list(self._tasks) + if pending: + ready, _ = ray.wait(pending, num_returns=len(pending), timeout=0) + if not ready and blocking_wait: + ready, _ = ray.wait(pending, num_returns=1, timeout=10.0) + for obj_ref in ready: + yield (self._tasks.pop(obj_ref), self._objects.pop(obj_ref)) + + def completed_prefetch(self, blocking_wait=False, max_yield=999): + """Similar to completed but only returns once the object is local. + + Assumes obj_ref only is one id.""" + + for worker, obj_ref in self.completed(blocking_wait=blocking_wait): + self._fetching.append((worker, obj_ref)) + + for _ in range(max_yield): + if not self._fetching: + break + + yield self._fetching.popleft() + + def reset_workers(self, workers): + """Notify that some workers may be removed.""" + for obj_ref, ev in self._tasks.copy().items(): + if ev not in workers: + del self._tasks[obj_ref] + del self._objects[obj_ref] + + # We want to keep the same deque reference so that we don't suffer from + # stale references in generators that are still in flight + for _ in range(len(self._fetching)): + ev, obj_ref = self._fetching.popleft() + if ev in workers: + # Re-queue items that are still valid + self._fetching.append((ev, obj_ref)) + + @property + def count(self): + return len(self._tasks) + + +def create_colocated_actors( + actor_specs: Sequence[Tuple[Type, Any, Any, int]], + node: Optional[str] = "localhost", + max_attempts: int = 10, +) -> Dict[Type, List[ActorHandle]]: + """Create co-located actors of any type(s) on any node. + + Args: + actor_specs: Tuple/list with tuples consisting of: 1) The + (already @ray.remote) class(es) to construct, 2) c'tor args, + 3) c'tor kwargs, and 4) the number of actors of that class with + given args/kwargs to construct. + node: The node to co-locate the actors on. By default ("localhost"), + place the actors on the node the caller of this function is + located on. Use None for indicating that any (resource fulfilling) + node in the cluster may be used. + max_attempts: The maximum number of co-location attempts to + perform before throwing an error. + + Returns: + A dict mapping the created types to the list of n ActorHandles + created (and co-located) for that type. + """ + if node == "localhost": + node = platform.node() + + # Maps each entry in `actor_specs` to lists of already co-located actors. + ok = [[] for _ in range(len(actor_specs))] + + # Try n times to co-locate all given actor types (`actor_specs`). + # With each (failed) attempt, increase the number of actors we try to + # create (on the same node), then kill the ones that have been created in + # excess. + for attempt in range(max_attempts): + # If any attempt to co-locate fails, set this to False and we'll do + # another attempt. + all_good = True + # Process all `actor_specs` in sequence. + for i, (typ, args, kwargs, count) in enumerate(actor_specs): + args = args or [] # Allow None. + kwargs = kwargs or {} # Allow None. + # We don't have enough actors yet of this spec co-located on + # the desired node. + if len(ok[i]) < count: + co_located = try_create_colocated( + cls=typ, + args=args, + kwargs=kwargs, + count=count * (attempt + 1), + node=node, + ) + # If node did not matter (None), from here on, use the host + # that the first actor(s) are already co-located on. + if node is None: + node = ray.get(co_located[0].get_host.remote()) + # Add the newly co-located actors to the `ok` list. + ok[i].extend(co_located) + # If we still don't have enough -> We'll have to do another + # attempt. + if len(ok[i]) < count: + all_good = False + # We created too many actors for this spec -> Kill/truncate + # the excess ones. + if len(ok[i]) > count: + for a in ok[i][count:]: + a.__ray_terminate__.remote() + ok[i] = ok[i][:count] + + # All `actor_specs` have been fulfilled, return lists of + # co-located actors. + if all_good: + return ok + + raise Exception("Unable to create enough colocated actors -> aborting.") + + +def try_create_colocated( + cls: Type[ActorClass], + args: List[Any], + count: int, + kwargs: Optional[List[Any]] = None, + node: Optional[str] = "localhost", +) -> List[ActorHandle]: + """Tries to co-locate (same node) a set of Actors of the same type. + + Returns a list of successfully co-located actors. All actors that could + not be co-located (with the others on the given node) will not be in this + list. + + Creates each actor via it's remote() constructor and then checks, whether + it has been co-located (on the same node) with the other (already created) + ones. If not, terminates the just created actor. + + Args: + cls: The Actor class to use (already @ray.remote "converted"). + args: List of args to pass to the Actor's constructor. One item + per to-be-created actor (`count`). + count: Number of actors of the given `cls` to construct. + kwargs: Optional list of kwargs to pass to the Actor's constructor. + One item per to-be-created actor (`count`). + node: The node to co-locate the actors on. By default ("localhost"), + place the actors on the node the caller of this function is + located on. If None, will try to co-locate all actors on + any available node. + + Returns: + List containing all successfully co-located actor handles. + """ + if node == "localhost": + node = platform.node() + + kwargs = kwargs or {} + actors = [cls.remote(*args, **kwargs) for _ in range(count)] + co_located, non_co_located = split_colocated(actors, node=node) + logger.info("Got {} colocated actors of {}".format(len(co_located), count)) + for a in non_co_located: + a.__ray_terminate__.remote() + return co_located + + +def split_colocated( + actors: List[ActorHandle], + node: Optional[str] = "localhost", +) -> Tuple[List[ActorHandle], List[ActorHandle]]: + """Splits up given actors into colocated (on same node) and non colocated. + + The co-location criterion depends on the `node` given: + If given (or default: platform.node()): Consider all actors that are on + that node "colocated". + If None: Consider the largest sub-set of actors that are all located on + the same node (whatever that node is) as "colocated". + + Args: + actors: The list of actor handles to split into "colocated" and + "non colocated". + node: The node defining "colocation" criterion. If provided, consider + thos actors "colocated" that sit on this node. If None, use the + largest subset within `actors` that are sitting on the same + (any) node. + + Returns: + Tuple of two lists: 1) Co-located ActorHandles, 2) non co-located + ActorHandles. + """ + if node == "localhost": + node = platform.node() + + # Get nodes of all created actors. + hosts = ray.get([a.get_host.remote() for a in actors]) + + # If `node` not provided, use the largest group of actors that sit on the + # same node, regardless of what that node is. + if node is None: + node_groups = defaultdict(set) + for host, actor in zip(hosts, actors): + node_groups[host].add(actor) + max_ = -1 + largest_group = None + for host in node_groups: + if max_ < len(node_groups[host]): + max_ = len(node_groups[host]) + largest_group = host + non_co_located = [] + for host in node_groups: + if host != largest_group: + non_co_located.extend(list(node_groups[host])) + return list(node_groups[largest_group]), non_co_located + # Node provided (or default: localhost): Consider those actors "colocated" + # that were placed on `node`. + else: + # Split into co-located (on `node) and non-co-located (not on `node`). + co_located = [] + non_co_located = [] + for host, a in zip(hosts, actors): + # This actor has been placed on the correct node. + if host == node: + co_located.append(a) + # This actor has been placed on a different node. + else: + non_co_located.append(a) + return co_located, non_co_located + + +def drop_colocated(actors: List[ActorHandle]) -> List[ActorHandle]: + colocated, non_colocated = split_colocated(actors) + for a in colocated: + a.__ray_terminate__.remote() + return non_colocated diff --git a/lib/python3.12/site-packages/ray/rllib/utils/annotations.py b/lib/python3.12/site-packages/ray/rllib/utils/annotations.py new file mode 100644 index 0000000000000000000000000000000000000000..286c541e0f12bfab469b9809e1cf76eba0de2dc1 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/annotations.py @@ -0,0 +1,213 @@ +from ray._common.deprecation import Deprecated +from ray.util.annotations import _mark_annotated + + +def override(parent_cls): + """Decorator for documenting method overrides. + + Args: + parent_cls: The superclass that provides the overridden method. If + `parent_class` does not actually have the method or the class, in which + method is defined is not a subclass of `parent_class`, an error is raised. + + .. testcode:: + :skipif: True + + from ray.rllib.policy import Policy + class TorchPolicy(Policy): + ... + # Indicates that `TorchPolicy.loss()` overrides the parent + # Policy class' own `loss method. Leads to an error if Policy + # does not have a `loss` method. + + @override(Policy) + def loss(self, model, action_dist, train_batch): + ... + + """ + + class OverrideCheck: + def __init__(self, func, expected_parent_cls): + self.func = func + self.expected_parent_cls = expected_parent_cls + + def __set_name__(self, owner, name): + # Check if the owner (the class) is a subclass of the expected base class + if not issubclass(owner, self.expected_parent_cls): + raise TypeError( + f"When using the @override decorator, {owner.__name__} must be a " + f"subclass of {parent_cls.__name__}!" + ) + # Set the function as a regular method on the class. + setattr(owner, name, self.func) + + def decorator(method): + # Check, whether `method` is actually defined by the parent class. + if method.__name__ not in dir(parent_cls): + raise NameError( + f"When using the @override decorator, {method.__name__} must override " + f"the respective method (with the same name) of {parent_cls.__name__}!" + ) + + # Check if the class is a subclass of the expected base class + OverrideCheck(method, parent_cls) + return method + + return decorator + + +def PublicAPI(obj): + """Decorator for documenting public APIs. + + Public APIs are classes and methods exposed to end users of RLlib. You + can expect these APIs to remain stable across RLlib releases. + + Subclasses that inherit from a ``@PublicAPI`` base class can be + assumed part of the RLlib public API as well (e.g., all Algorithm classes + are in public API because Algorithm is ``@PublicAPI``). + + In addition, you can assume all algo configurations are part of their + public API as well. + + .. testcode:: + :skipif: True + + # Indicates that the `Algorithm` class is exposed to end users + # of RLlib and will remain stable across RLlib releases. + from ray import tune + @PublicAPI + class Algorithm(tune.Trainable): + ... + """ + + _mark_annotated(obj) + return obj + + +def DeveloperAPI(obj): + """Decorator for documenting developer APIs. + + Developer APIs are classes and methods explicitly exposed to developers + for the purposes of building custom algorithms or advanced training + strategies on top of RLlib internals. You can generally expect these APIs + to be stable sans minor changes (but less stable than public APIs). + + Subclasses that inherit from a ``@DeveloperAPI`` base class can be + assumed part of the RLlib developer API as well. + + .. testcode:: + :skipif: True + + # Indicates that the `TorchPolicy` class is exposed to end users + # of RLlib and will remain (relatively) stable across RLlib + # releases. + from ray.rllib.policy import Policy + @DeveloperAPI + class TorchPolicy(Policy): + ... + """ + + _mark_annotated(obj) + return obj + + +def ExperimentalAPI(obj): + """Decorator for documenting experimental APIs. + + Experimental APIs are classes and methods that are in development and may + change at any time in their development process. You should not expect + these APIs to be stable until their tag is changed to `DeveloperAPI` or + `PublicAPI`. + + Subclasses that inherit from a ``@ExperimentalAPI`` base class can be + assumed experimental as well. + + .. testcode:: + :skipif: True + + from ray.rllib.policy import Policy + class TorchPolicy(Policy): + ... + # Indicates that the `TorchPolicy.loss` method is a new and + # experimental API and may change frequently in future + # releases. + @ExperimentalAPI + def loss(self, model, action_dist, train_batch): + ... + """ + + _mark_annotated(obj) + return obj + + +def OldAPIStack(obj): + """Decorator for classes/methods/functions belonging to the old API stack. + + These should be deprecated at some point after Ray 3.0 (RLlib GA). + It is recommended for users to start exploring (and coding against) the new API + stack instead. + """ + # No effect yet. + + _mark_annotated(obj) + return obj + + +def OverrideToImplementCustomLogic(obj): + """Users should override this in their sub-classes to implement custom logic. + + Used in Algorithm and Policy to tag methods that need overriding, e.g. + `Policy.loss()`. + + .. testcode:: + :skipif: True + + from ray.rllib.policy.torch_policy import TorchPolicy + @overrides(TorchPolicy) + @OverrideToImplementCustomLogic + def loss(self, ...): + # implement custom loss function here ... + # ... w/o calling the corresponding `super().loss()` method. + ... + + """ + obj.__is_overridden__ = False + return obj + + +def OverrideToImplementCustomLogic_CallToSuperRecommended(obj): + """Users should override this in their sub-classes to implement custom logic. + + Thereby, it is recommended (but not required) to call the super-class' + corresponding method. + + Used in Algorithm and Policy to tag methods that need overriding, but the + super class' method should still be called, e.g. + `Algorithm.setup()`. + + .. testcode:: + :skipif: True + + from ray import tune + @overrides(tune.Trainable) + @OverrideToImplementCustomLogic_CallToSuperRecommended + def setup(self, config): + # implement custom setup logic here ... + super().setup(config) + # ... or here (after having called super()'s setup method. + """ + obj.__is_overridden__ = False + return obj + + +def is_overridden(obj): + """Check whether a function has been overridden. + + Note, this only works for API calls decorated with OverrideToImplementCustomLogic + or OverrideToImplementCustomLogic_CallToSuperRecommended. + """ + return getattr(obj, "__is_overridden__", True) + + +# Backward compatibility. +Deprecated = Deprecated diff --git a/lib/python3.12/site-packages/ray/rllib/utils/checkpoints.py b/lib/python3.12/site-packages/ray/rllib/utils/checkpoints.py new file mode 100644 index 0000000000000000000000000000000000000000..f18ec698ca9d0a388dd020b35dcb2e05d1d5b499 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/checkpoints.py @@ -0,0 +1,1060 @@ +import abc +import inspect +import json +import logging +import os +import pathlib +import re +import tempfile +from types import MappingProxyType +from typing import Any, Collection, Dict, List, Optional, Tuple, Union + +import pyarrow.fs +from packaging import version + +import ray +import ray.cloudpickle as pickle +from ray.rllib.core import ( + COMPONENT_LEARNER, + COMPONENT_LEARNER_GROUP, + COMPONENT_RL_MODULE, +) +from ray.rllib.utils import force_list +from ray.rllib.utils.actor_manager import FaultTolerantActorManager +from ray.rllib.utils.annotations import ( + OldAPIStack, + OverrideToImplementCustomLogic_CallToSuperRecommended, +) +from ray.rllib.utils.serialization import NOT_SERIALIZABLE, serialize_type +from ray.rllib.utils.typing import StateDict +from ray.train import Checkpoint as Checkpoint_train +from ray.tune import Checkpoint as Checkpoint_tune +from ray.tune.utils.file_transfer import sync_dir_between_nodes +from ray.util import log_once +from ray.util.annotations import PublicAPI + +logger = logging.getLogger(__name__) + +# The current checkpoint version used by RLlib for Algorithm and Policy checkpoints. +# History: +# 0.1: Ray 2.0.0 +# A single `checkpoint-[iter num]` file for Algorithm checkpoints +# within the checkpoint directory. Policy checkpoints not supported across all +# DL frameworks. + +# 1.0: Ray >=2.1.0 +# An algorithm_state.pkl file for the state of the Algorithm (excluding +# individual policy states). +# One sub-dir inside the "policies" sub-dir for each policy with a +# dedicated policy_state.pkl in it for the policy state. + +# 1.1: Same as 1.0, but has a new "format" field in the rllib_checkpoint.json file +# indicating, whether the checkpoint is `cloudpickle` (default) or `msgpack`. + +# 1.2: Introduces the checkpoint for the new Learner API if the Learner API is enabled. + +# 2.0: Introduces the Checkpointable API for all components on the new API stack +# (if the Learner-, RLModule, EnvRunner, and ConnectorV2 APIs are enabled). + +CHECKPOINT_VERSION = version.Version("1.1") +CHECKPOINT_VERSION_LEARNER_AND_ENV_RUNNER = version.Version("2.1") + + +@PublicAPI(stability="alpha") +class Checkpointable(abc.ABC): + """Abstract base class for a component of RLlib that can be checkpointed to disk. + + Subclasses must implement the following APIs: + - save_to_path() + - restore_from_path() + - from_checkpoint() + - get_state() + - set_state() + - get_ctor_args_and_kwargs() + - get_metadata() + - get_checkpointable_components() + """ + + # The state file for the implementing class. + # This file contains any state information that does NOT belong to any subcomponent + # of the implementing class (which are `Checkpointable` themselves and thus should + # have their own state- and metadata files). + # After a `save_to_path([path])` this file can be found directly in: `path/`. + STATE_FILE_NAME = "state" + + # The filename of the pickle file that contains the class information of the + # Checkpointable as well as all constructor args to be passed to such a class in + # order to construct a new instance. + CLASS_AND_CTOR_ARGS_FILE_NAME = "class_and_ctor_args.pkl" + + # Subclasses may set this to their own metadata filename. + # The dict returned by self.get_metadata() is stored in this JSON file. + METADATA_FILE_NAME = "metadata.json" + + def save_to_path( + self, + path: Optional[Union[str, pathlib.Path]] = None, + *, + state: Optional[StateDict] = None, + filesystem: Optional["pyarrow.fs.FileSystem"] = None, + use_msgpack: bool = False, + ) -> str: + """Saves the state of the implementing class (or `state`) to `path`. + + The state of the implementing class is always saved in the following format: + + .. testcode:: + :skipif: True + + path/ + [component1]/ + [component1 subcomponentA]/ + ... + [component1 subcomponentB]/ + ... + [component2]/ + ... + [cls.METADATA_FILE_NAME] (json) + [cls.STATE_FILE_NAME] (pkl|msgpack) + + The main logic is to loop through all subcomponents of this Checkpointable + and call their respective `save_to_path` methods. Then save the remaining + (non subcomponent) state to this Checkpointable's STATE_FILE_NAME. + In the exception that a component is a FaultTolerantActorManager instance, + instead of calling `save_to_path` directly on that manager, the first healthy + actor is interpreted as the component and its `save_to_path` method is called. + Even if that actor is located on another node, the created file is automatically + synced to the local node. + + Args: + path: The path to the directory to save the state of the implementing class + to. If `path` doesn't exist or is None, then a new directory will be + created (and returned). + state: An optional state dict to be used instead of getting a new state of + the implementing class through `self.get_state()`. + filesystem: PyArrow FileSystem to use to access data at the `path`. + If not specified, this is inferred from the URI scheme of `path`. + use_msgpack: Whether the state file should be written using msgpack and + msgpack_numpy (file extension is `.msgpack`), rather than pickle (file + extension is `.pkl`). + + Returns: + The path (str) where the state has been saved. + """ + + # If no path is given create a local temporary directory. + if path is None: + import uuid + + # Get the location of the temporary directory on the OS. + tmp_dir = pathlib.Path(tempfile.gettempdir()) + # Create a random directory name. + random_dir_name = str(uuid.uuid4()) + # Create the path, but do not craet the directory on the + # filesystem, yet. This is done by `PyArrow`. + path = path or tmp_dir / random_dir_name + + # We need a string path for `pyarrow.fs.FileSystem.from_uri`. + path = path if isinstance(path, str) else path.as_posix() + + # If we have no filesystem, figure it out. + if path and not filesystem: + # Note the path needs to be a path that is relative to the + # filesystem (e.g. `gs://tmp/...` -> `tmp/...`). + filesystem, path = pyarrow.fs.FileSystem.from_uri(path) + + # Make sure, path exists. + filesystem.create_dir(path, recursive=True) + + # Convert to `pathlib.Path` for easy handling. + path = pathlib.Path(path) + + # Write metadata file to disk. + metadata = self.get_metadata() + if "checkpoint_version" not in metadata: + metadata["checkpoint_version"] = str( + CHECKPOINT_VERSION_LEARNER_AND_ENV_RUNNER + ) + with filesystem.open_output_stream( + (path / self.METADATA_FILE_NAME).as_posix() + ) as f: + f.write(json.dumps(metadata).encode("utf-8")) + + # Write the class and constructor args information to disk. Always use pickle + # for this, because this information contains classes and maybe other + # non-serializable data. + with filesystem.open_output_stream( + (path / self.CLASS_AND_CTOR_ARGS_FILE_NAME).as_posix() + ) as f: + pickle.dump( + { + "class": type(self), + "ctor_args_and_kwargs": self.get_ctor_args_and_kwargs(), + }, + f, + ) + + # Get the entire state of this Checkpointable, or use provided `state`. + _state_provided = state is not None + # Get only the non-checkpointable components of the state. Checkpointable + # components are saved to path by their own `save_to_path` in the loop below. + state = state or self.get_state( + not_components=[c[0] for c in self.get_checkpointable_components()] + ) + + # Write components of `self` that themselves are `Checkpointable`. + for comp_name, comp in self.get_checkpointable_components(): + # If subcomponent's name is not in `state`, ignore it and don't write this + # subcomponent's state to disk. + if _state_provided and comp_name not in state: + continue + comp_path = path / comp_name + + # If component is an ActorManager, save the manager's first healthy + # actor's state to disk (even if it's on another node, in which case, we'll + # sync the generated file(s) back to this node). + if isinstance(comp, FaultTolerantActorManager): + actor_to_use = comp.healthy_actor_ids()[0] + + def _get_ip(_=None): + import ray + + return ray.util.get_node_ip_address() + + _result = next( + iter( + comp.foreach_actor( + _get_ip, + remote_actor_ids=[actor_to_use], + ) + ) + ) + if not _result.ok: + raise _result.get() + worker_ip_addr = _result.get() + self_ip_addr = _get_ip() + + # Save the state to a temporary location on the `actor_to_use`'s + # node. + comp_state_ref = None + if _state_provided: + comp_state_ref = ray.put(state.pop(comp_name)) + + if worker_ip_addr == self_ip_addr: + comp.foreach_actor( + lambda w, _path=comp_path, _state=comp_state_ref, _use_msgpack=use_msgpack: ( # noqa + w.save_to_path( + _path, + state=( + ray.get(_state) + if _state is not None + else w.get_state() + ), + use_msgpack=_use_msgpack, + ) + ), + remote_actor_ids=[actor_to_use], + ) + else: + # Save the checkpoint to the temporary directory on the worker. + def _save(w, _state=comp_state_ref, _use_msgpack=use_msgpack): + import tempfile + + # Create a temporary directory on the worker. + tmpdir = tempfile.mkdtemp() + w.save_to_path( + tmpdir, + state=( + ray.get(_state) if _state is not None else w.get_state() + ), + use_msgpack=_use_msgpack, + ) + return tmpdir + + _result = next( + iter(comp.foreach_actor(_save, remote_actor_ids=[actor_to_use])) + ) + if not _result.ok: + raise _result.get() + worker_temp_dir = _result.get() + + # Sync the temporary directory from the worker to this node. + sync_dir_between_nodes( + worker_ip_addr, + worker_temp_dir, + self_ip_addr, + str(comp_path), + ) + + # Remove the temporary directory on the worker. + def _rmdir(_, _dir=worker_temp_dir): + import shutil + + shutil.rmtree(_dir) + + comp.foreach_actor(_rmdir, remote_actor_ids=[actor_to_use]) + + # Local component (instance stored in a property of `self`). + else: + if _state_provided: + comp_state = state.pop(comp_name) + else: + comp_state = self.get_state(components=comp_name)[comp_name] + # By providing the `state` arg, we make sure that the component does not + # have to call its own `get_state()` anymore, but uses what's provided + # here. + comp.save_to_path( + comp_path, + filesystem=filesystem, + state=comp_state, + use_msgpack=use_msgpack, + ) + + # Write all the remaining state to disk. + filename = path / ( + self.STATE_FILE_NAME + (".msgpack" if use_msgpack else ".pkl") + ) + with filesystem.open_output_stream(filename.as_posix()) as f: + if use_msgpack: + msgpack = try_import_msgpack(error=True) + msgpack.dump(state, f) + else: + pickle.dump(state, f) + + return str(path) + + def restore_from_path( + self, + path: Union[str, pathlib.Path], + *, + component: Optional[str] = None, + filesystem: Optional["pyarrow.fs.FileSystem"] = None, + **kwargs, + ) -> None: + """Restores the state of the implementing class from the given path. + + If the `component` arg is provided, `path` refers to a checkpoint of a + subcomponent of `self`, thus allowing the user to load only the subcomponent's + state into `self` without affecting any of the other state information (for + example, loading only the NN state into a Checkpointable, which contains such + an NN, but also has other state information that should NOT be changed by + calling this method). + + The given `path` should have the following structure and contain the following + files: + + .. testcode:: + :skipif: True + + path/ + [component1]/ + [component1 subcomponentA]/ + ... + [component1 subcomponentB]/ + ... + [component2]/ + ... + [cls.METADATA_FILE_NAME] (json) + [cls.STATE_FILE_NAME] (pkl|msgpack) + + Note that the self.METADATA_FILE_NAME file is not required to restore the state. + + Args: + path: The path to load the implementing class' state from or to load the + state of only one subcomponent's state of the implementing class (if + `component` is provided). + component: If provided, `path` is interpreted as the checkpoint path of only + the subcomponent and thus, only that subcomponent's state is + restored/loaded. All other state of `self` remains unchanged in this + case. + filesystem: PyArrow FileSystem to use to access data at the `path`. If not + specified, this is inferred from the URI scheme of `path`. + **kwargs: Forward compatibility kwargs. + """ + path = path if isinstance(path, str) else path.as_posix() + + if path and not filesystem: + # Note the path needs to be a path that is relative to the + # filesystem (e.g. `gs://tmp/...` -> `tmp/...`). + filesystem, path = pyarrow.fs.FileSystem.from_uri(path) + # Only here convert to a `Path` instance b/c otherwise + # cloud path gets broken (i.e. 'gs://' -> 'gs:/'). + path = pathlib.Path(path) + + if not _exists_at_fs_path(filesystem, path.as_posix()): + raise FileNotFoundError(f"`path` ({path}) not found!") + + # Restore components of `self` that themselves are `Checkpointable`. + orig_comp_names = {c[0] for c in self.get_checkpointable_components()} + self._restore_all_subcomponents_from_path( + path, filesystem, component=component, **kwargs + ) + + # Restore the "base" state (not individual subcomponents). + if component is None: + filename = path / self.STATE_FILE_NAME + if filename.with_suffix(".msgpack").is_file(): + msgpack = try_import_msgpack(error=True) + with filesystem.open_input_stream( + filename.with_suffix(".msgpack").as_posix() + ) as f: + state = msgpack.load(f, strict_map_key=False) + else: + with filesystem.open_input_stream( + filename.with_suffix(".pkl").as_posix() + ) as f: + state = pickle.load(f) + self.set_state(state) + + new_comp_names = {c[0] for c in self.get_checkpointable_components()} + diff_comp_names = new_comp_names - orig_comp_names + if diff_comp_names: + self._restore_all_subcomponents_from_path( + path, filesystem, only_comp_names=diff_comp_names, **kwargs + ) + + @classmethod + def from_checkpoint( + cls, + path: Union[str, pathlib.Path], + filesystem: Optional["pyarrow.fs.FileSystem"] = None, + **kwargs, + ) -> "Checkpointable": + """Creates a new Checkpointable instance from the given location and returns it. + + Args: + path: The checkpoint path to load (a) the information on how to construct + a new instance of the implementing class and (b) the state to restore + the created instance to. + filesystem: PyArrow FileSystem to use to access data at the `path`. If not + specified, this is inferred from the URI scheme of `path`. + kwargs: Forward compatibility kwargs. Note that these kwargs are sent to + each subcomponent's `from_checkpoint()` call. + + Returns: + A new instance of the implementing class, already set to the state stored + under `path`. + """ + # We need a string path for the `PyArrow` filesystem. + path = path if isinstance(path, str) else path.as_posix() + + # If no filesystem is passed in create one. + if path and not filesystem: + # Note the path needs to be a path that is relative to the + # filesystem (e.g. `gs://tmp/...` -> `tmp/...`). + filesystem, path = pyarrow.fs.FileSystem.from_uri(path) + # Only here convert to a `Path` instance b/c otherwise + # cloud path gets broken (i.e. 'gs://' -> 'gs:/'). + path = pathlib.Path(path) + + # Get the class constructor to call and its args/kwargs. + # Try reading the pickle file first. + try: + with filesystem.open_input_stream( + (path / cls.CLASS_AND_CTOR_ARGS_FILE_NAME).as_posix() + ) as f: + ctor_info = pickle.load(f) + ctor = ctor_info["class"] + ctor_args = force_list(ctor_info["ctor_args_and_kwargs"][0]) + ctor_kwargs = ctor_info["ctor_args_and_kwargs"][1] + + # Inspect the ctor to see, which arguments in ctor_info should be replaced + # with the user provided **kwargs. + for i, (param_name, param) in enumerate( + inspect.signature(ctor).parameters.items() + ): + if param_name in kwargs: + val = kwargs.pop(param_name) + if ( + param.kind == inspect._ParameterKind.POSITIONAL_OR_KEYWORD + and len(ctor_args) > i + ): + ctor_args[i] = val + else: + ctor_kwargs[param_name] = val + + # If the pickle file is from another python version, use provided + # args instead. + except Exception: + # Use class that this method was called on. + ctor = cls + # Use only user provided **kwargs. + ctor_args = [] + ctor_kwargs = kwargs + + # Check, whether the constructor actually goes together with `cls`. + if not issubclass(ctor, cls): + raise ValueError( + f"The class ({ctor}) stored in checkpoint ({path}) does not seem to be " + f"a subclass of `cls` ({cls})!" + ) + elif not issubclass(ctor, Checkpointable): + raise ValueError( + f"The class ({ctor}) stored in checkpoint ({path}) does not seem to be " + "an implementer of the `Checkpointable` API!" + ) + + # Construct the initial object (without any particular state). + obj = ctor(*ctor_args, **ctor_kwargs) + # Restore the state of the constructed object. + obj.restore_from_path(path, filesystem=filesystem, **kwargs) + # Return the new object. + return obj + + @abc.abstractmethod + def get_state( + self, + components: Optional[Union[str, Collection[str]]] = None, + *, + not_components: Optional[Union[str, Collection[str]]] = None, + **kwargs, + ) -> StateDict: + """Returns the implementing class's current state as a dict. + + The returned dict must only contain msgpack-serializable data if you want to + use the `AlgorithmConfig._msgpack_checkpoints` option. Consider returning your + non msgpack-serializable data from the `Checkpointable.get_ctor_args_and_kwargs` + method, instead. + + Args: + components: An optional collection of string keys to be included in the + returned state. This might be useful, if getting certain components + of the state is expensive (e.g. reading/compiling the weights of a large + NN) and at the same time, these components are not required by the + caller. + not_components: An optional list of string keys to be excluded in the + returned state, even if the same string is part of `components`. + This is useful to get the complete state of the class, except + one or a few components. + kwargs: Forward-compatibility kwargs. + + Returns: + The current state of the implementing class (or only the `components` + specified, w/o those in `not_components`). + """ + + @abc.abstractmethod + def set_state(self, state: StateDict) -> None: + """Sets the implementing class' state to the given state dict. + + If component keys are missing in `state`, these components of the implementing + class will not be updated/set. + + Args: + state: The state dict to restore the state from. Maps component keys + to the corresponding subcomponent's own state. + """ + + @abc.abstractmethod + def get_ctor_args_and_kwargs(self) -> Tuple[Tuple, Dict[str, Any]]: + """Returns the args/kwargs used to create `self` from its constructor. + + Returns: + A tuple of the args (as a tuple) and kwargs (as a Dict[str, Any]) used to + construct `self` from its class constructor. + """ + + @OverrideToImplementCustomLogic_CallToSuperRecommended + def get_metadata(self) -> Dict: + """Returns JSON writable metadata further describing the implementing class. + + Note that this metadata is NOT part of any state and is thus NOT needed to + restore the state of a Checkpointable instance from a directory. Rather, the + metadata will be written into `self.METADATA_FILE_NAME` when calling + `self.save_to_path()` for the user's convenience. + + Returns: + A JSON-encodable dict of metadata information. + """ + return { + "class_and_ctor_args_file": self.CLASS_AND_CTOR_ARGS_FILE_NAME, + "state_file": self.STATE_FILE_NAME, + "ray_version": ray.__version__, + "ray_commit": ray.__commit__, + } + + def get_checkpointable_components(self) -> List[Tuple[str, "Checkpointable"]]: + """Returns the implementing class's own Checkpointable subcomponents. + + Returns: + A list of 2-tuples (name, subcomponent) describing the implementing class' + subcomponents, all of which have to be `Checkpointable` themselves and + whose state is therefore written into subdirectories (rather than the main + state file (self.STATE_FILE_NAME) when calling `self.save_to_path()`). + """ + return [] + + def _check_component(self, name, components, not_components) -> bool: + """Returns True if a component should be checkpointed. + + Args: + name: The checkpoint name. + components: A list of components that should be checkpointed. + non_components: A list of components that should not be checkpointed. + + Returns: + True, if the component should be checkpointed and otherwise False. + """ + comp_list = force_list(components) + not_comp_list = force_list(not_components) + if ( + components is None + or any(c.startswith(name + "/") for c in comp_list) + or name in comp_list + ) and (not_components is None or name not in not_comp_list): + return True + return False + + def _get_subcomponents(self, name, components): + if components is None: + return None + + components = force_list(components) + subcomponents = [] + for comp in components: + if comp.startswith(name + "/"): + subcomponents.append(comp[len(name) + 1 :]) + + return None if not subcomponents else subcomponents + + def _restore_all_subcomponents_from_path( + self, path, filesystem, only_comp_names=None, component=None, **kwargs + ): + for comp_name, comp in self.get_checkpointable_components(): + if only_comp_names is not None and comp_name not in only_comp_names: + continue + + # The value of the `component` argument for the upcoming + # `[subcomponent].restore_from_path(.., component=..)` call. + comp_arg = None + + if component is None: + comp_dir = path / comp_name + # If subcomponent's dir is not in path, ignore it and don't restore this + # subcomponent's state from disk. + if not _exists_at_fs_path(filesystem, comp_dir.as_posix()): + continue + else: + comp_dir = path + + # `component` is a path that starts with `comp` -> Remove the name of + # `comp` from the `component` arg in the upcoming call to `restore_..`. + if component.startswith(comp_name + "/"): + comp_arg = component[len(comp_name) + 1 :] + # `component` has nothing to do with `comp` -> Skip. + elif component != comp_name: + continue + + # If component is an ActorManager, restore all the manager's healthy + # actors' states from disk (even if they are on another node, in which case, + # we'll sync checkpoint file(s) to the respective node). + if isinstance(comp, FaultTolerantActorManager): + head_node_ip = ray.util.get_node_ip_address() + all_healthy_actors = comp.healthy_actor_ids() + + def _restore( + w, + _kwargs=MappingProxyType(kwargs), + _path=comp_dir, + _head_ip=head_node_ip, + _comp_arg=comp_arg, + ): + import tempfile + + import ray + + worker_node_ip = ray.util.get_node_ip_address() + # If the worker is on the same node as the head, load the checkpoint + # directly from the path otherwise sync the checkpoint from the head + # to the worker and load it from there. + if worker_node_ip == _head_ip: + w.restore_from_path(_path, component=_comp_arg, **_kwargs) + else: + with tempfile.TemporaryDirectory() as temp_dir: + sync_dir_between_nodes( + _head_ip, _path, worker_node_ip, temp_dir + ) + w.restore_from_path( + temp_dir, component=_comp_arg, **_kwargs + ) + + comp.foreach_actor(_restore, remote_actor_ids=all_healthy_actors) + + # Call `restore_from_path()` on local subcomponent, thereby passing in the + # **kwargs. + else: + comp.restore_from_path( + comp_dir, filesystem=filesystem, component=comp_arg, **kwargs + ) + + +def _exists_at_fs_path(fs: pyarrow.fs.FileSystem, path: str) -> bool: + """Returns `True` if the path can be found in the filesystem.""" + valid = fs.get_file_info(path) + return valid.type != pyarrow.fs.FileType.NotFound + + +def _is_dir(file_info: pyarrow.fs.FileInfo) -> bool: + """Returns `True`, if the file info is from a directory.""" + return file_info.type == pyarrow.fs.FileType.Directory + + +@OldAPIStack +def get_checkpoint_info( + checkpoint: Union[str, Checkpoint_train, Checkpoint_tune], + filesystem: Optional["pyarrow.fs.FileSystem"] = None, +) -> Dict[str, Any]: + """Returns a dict with information about an Algorithm/Policy checkpoint. + + If the given checkpoint is a >=v1.0 checkpoint directory, try reading all + information from the contained `rllib_checkpoint.json` file. + + Args: + checkpoint: The checkpoint directory (str) or a Checkpoint object. + filesystem: PyArrow FileSystem to use to access data at the `checkpoint`. If not + specified, this is inferred from the URI scheme provided by `checkpoint`. + + Returns: + A dict containing the keys: + "type": One of "Policy" or "Algorithm". + "checkpoint_version": A version tuple, e.g. v1.0, indicating the checkpoint + version. This will help RLlib to remain backward compatible wrt. future + Ray and checkpoint versions. + "checkpoint_dir": The directory with all the checkpoint files in it. This might + be the same as the incoming `checkpoint` arg. + "state_file": The main file with the Algorithm/Policy's state information in it. + This is usually a pickle-encoded file. + "policy_ids": An optional set of PolicyIDs in case we are dealing with an + Algorithm checkpoint. None if `checkpoint` is a Policy checkpoint. + """ + # Default checkpoint info. + info = { + "type": "Algorithm", + "format": "cloudpickle", + "checkpoint_version": CHECKPOINT_VERSION, + "checkpoint_dir": None, + "state_file": None, + "policy_ids": None, + "module_ids": None, + } + + # `checkpoint` is a Checkpoint instance: Translate to directory and continue. + if isinstance(checkpoint, (Checkpoint_train, Checkpoint_tune)): + checkpoint = checkpoint.to_directory() + + if checkpoint and not filesystem: + # Note the path needs to be a path that is relative to the + # filesystem (e.g. `gs://tmp/...` -> `tmp/...`). + filesystem, checkpoint = pyarrow.fs.FileSystem.from_uri(checkpoint) + # Only here convert to a `Path` instance b/c otherwise + # cloud path gets broken (i.e. 'gs://' -> 'gs:/'). + checkpoint = pathlib.Path(checkpoint) + + # Checkpoint is dir. + if _exists_at_fs_path(filesystem, checkpoint.as_posix()) and _is_dir( + filesystem.get_file_info(checkpoint.as_posix()) + ): + info.update({"checkpoint_dir": str(checkpoint)}) + + # Figure out whether this is an older checkpoint format + # (with a `checkpoint-\d+` file in it). + file_info_list = filesystem.get_file_info( + pyarrow.fs.FileSelector(checkpoint.as_posix(), recursive=False) + ) + for file_info in file_info_list: + if file_info.is_file: + if re.match("checkpoint-\\d+", file_info.base_name): + info.update( + { + "checkpoint_version": version.Version("0.1"), + "state_file": str(file_info.base_name), + } + ) + return info + + # No old checkpoint file found. + + # If rllib_checkpoint.json file present, read available information from it + # and then continue with the checkpoint analysis (possibly overriding further + # information). + if _exists_at_fs_path( + filesystem, (checkpoint / "rllib_checkpoint.json").as_posix() + ): + # if (checkpoint / "rllib_checkpoint.json").is_file(): + with filesystem.open_input_stream( + (checkpoint / "rllib_checkpoint.json").as_posix() + ) as f: + # with open(checkpoint / "rllib_checkpoint.json") as f: + rllib_checkpoint_info = json.load(fp=f) + if "checkpoint_version" in rllib_checkpoint_info: + rllib_checkpoint_info["checkpoint_version"] = version.Version( + rllib_checkpoint_info["checkpoint_version"] + ) + info.update(rllib_checkpoint_info) + else: + # No rllib_checkpoint.json file present: Warn and continue trying to figure + # out checkpoint info ourselves. + if log_once("no_rllib_checkpoint_json_file"): + logger.warning( + "No `rllib_checkpoint.json` file found in checkpoint directory " + f"{checkpoint}! Trying to extract checkpoint info from other files " + f"found in that dir." + ) + + # Policy checkpoint file found. + for extension in ["pkl", "msgpck"]: + if _exists_at_fs_path( + filesystem, (checkpoint / ("policy_state." + extension)).as_posix() + ): + # if (checkpoint / ("policy_state." + extension)).is_file(): + info.update( + { + "type": "Policy", + "format": "cloudpickle" if extension == "pkl" else "msgpack", + "checkpoint_version": CHECKPOINT_VERSION, + "state_file": str(checkpoint / f"policy_state.{extension}"), + } + ) + return info + + # Valid Algorithm checkpoint >v0 file found? + format = None + for extension in ["pkl", "msgpck", "msgpack"]: + state_file = checkpoint / f"algorithm_state.{extension}" + if ( + _exists_at_fs_path(filesystem, state_file.as_posix()) + and filesystem.get_file_info(state_file.as_posix()).is_file + ): + format = "cloudpickle" if extension == "pkl" else "msgpack" + break + if format is None: + raise ValueError( + "Given checkpoint does not seem to be valid! No file with the name " + "`algorithm_state.[pkl|msgpack|msgpck]` (or `checkpoint-[0-9]+`) found." + ) + + info.update( + { + "format": format, + "state_file": str(state_file), + } + ) + + # Collect all policy IDs in the sub-dir "policies/". + policies_dir = checkpoint / "policies" + if _exists_at_fs_path(filesystem, policies_dir.as_posix()) and _is_dir( + filesystem.get_file_info(policies_dir.as_posix()) + ): + policy_ids = set() + file_info_list = filesystem.get_file_info( + pyarrow.fs.FileSelector(policies_dir.as_posix(), recursive=False) + ) + for file_info in file_info_list: + policy_ids.add(file_info.base_name) + info.update({"policy_ids": policy_ids}) + + # Collect all module IDs in the sub-dir "learner/module_state/". + modules_dir = ( + checkpoint + / COMPONENT_LEARNER_GROUP + / COMPONENT_LEARNER + / COMPONENT_RL_MODULE + ) + if _exists_at_fs_path(filesystem, checkpoint.as_posix()) and _is_dir( + filesystem.get_file_info(modules_dir.as_posix()) + ): + module_ids = set() + file_info_list = filesystem.get_file_info( + pyarrow.fs.FileSelector(modules_dir.as_posix(), recursive=False) + ) + for file_info in file_info_list: + # Only add subdirs (those are the ones where the RLModule data + # is stored, not files (could be json metadata files). + module_dir = modules_dir / file_info.base_name + if _is_dir(filesystem.get_file_info(module_dir.as_posix())): + module_ids.add(file_info.base_name) + info.update({"module_ids": module_ids}) + + # Checkpoint is a file: Use as-is (interpreting it as old Algorithm checkpoint + # version). + elif ( + _exists_at_fs_path(filesystem, checkpoint.as_posix()) + and filesystem.get_file_info(checkpoint.as_posix()).is_file + ): + info.update( + { + "checkpoint_version": version.Version("0.1"), + "checkpoint_dir": str(checkpoint.parent), + "state_file": str(checkpoint), + } + ) + + else: + raise ValueError( + f"Given checkpoint ({str(checkpoint)}) not found! Must be a " + "checkpoint directory (or a file for older checkpoint versions)." + ) + + return info + + +@OldAPIStack +def convert_to_msgpack_checkpoint( + checkpoint: Union[str, Checkpoint_train, Checkpoint_tune], + msgpack_checkpoint_dir: str, +) -> str: + """Converts an Algorithm checkpoint (pickle based) to a msgpack based one. + + Msgpack has the advantage of being python version independent. + + Args: + checkpoint: The directory, in which to find the Algorithm checkpoint (pickle + based). + msgpack_checkpoint_dir: The directory, in which to create the new msgpack + based checkpoint. + + Returns: + The directory in which the msgpack checkpoint has been created. Note that + this is the same as `msgpack_checkpoint_dir`. + """ + from ray.rllib.algorithms import Algorithm + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + from ray.rllib.core.rl_module import validate_module_id + + # Try to import msgpack and msgpack_numpy. + msgpack = try_import_msgpack(error=True) + + # Restore the Algorithm using the python version dependent checkpoint. + algo = Algorithm.from_checkpoint(checkpoint) + state = algo.__getstate__() + + # Convert all code in state into serializable data. + # Serialize the algorithm class. + state["algorithm_class"] = serialize_type(state["algorithm_class"]) + # Serialize the algorithm's config object. + if not isinstance(state["config"], dict): + state["config"] = state["config"].serialize() + else: + state["config"] = AlgorithmConfig._serialize_dict(state["config"]) + + # Extract policy states from worker state (Policies get their own + # checkpoint sub-dirs). + policy_states = {} + if "worker" in state and "policy_states" in state["worker"]: + policy_states = state["worker"].pop("policy_states", {}) + + # Policy mapping fn. + state["worker"]["policy_mapping_fn"] = NOT_SERIALIZABLE + # Is Policy to train function. + state["worker"]["is_policy_to_train"] = NOT_SERIALIZABLE + + # Add RLlib checkpoint version (as string). + state["checkpoint_version"] = str(CHECKPOINT_VERSION) + + # Write state (w/o policies) to disk. + state_file = os.path.join(msgpack_checkpoint_dir, "algorithm_state.msgpck") + with open(state_file, "wb") as f: + msgpack.dump(state, f) + + # Write rllib_checkpoint.json. + with open(os.path.join(msgpack_checkpoint_dir, "rllib_checkpoint.json"), "w") as f: + json.dump( + { + "type": "Algorithm", + "checkpoint_version": state["checkpoint_version"], + "format": "msgpack", + "state_file": state_file, + "policy_ids": list(policy_states.keys()), + "ray_version": ray.__version__, + "ray_commit": ray.__commit__, + }, + f, + ) + + # Write individual policies to disk, each in their own subdirectory. + for pid, policy_state in policy_states.items(): + # From here on, disallow policyIDs that would not work as directory names. + validate_module_id(pid, error=True) + policy_dir = os.path.join(msgpack_checkpoint_dir, "policies", pid) + os.makedirs(policy_dir, exist_ok=True) + policy = algo.get_policy(pid) + policy.export_checkpoint( + policy_dir, + policy_state=policy_state, + checkpoint_format="msgpack", + ) + + # Release all resources used by the Algorithm. + algo.stop() + + return msgpack_checkpoint_dir + + +@OldAPIStack +def convert_to_msgpack_policy_checkpoint( + policy_checkpoint: Union[str, Checkpoint_train, Checkpoint_tune], + msgpack_checkpoint_dir: str, +) -> str: + """Converts a Policy checkpoint (pickle based) to a msgpack based one. + + Msgpack has the advantage of being python version independent. + + Args: + policy_checkpoint: The directory, in which to find the Policy checkpoint (pickle + based). + msgpack_checkpoint_dir: The directory, in which to create the new msgpack + based checkpoint. + + Returns: + The directory in which the msgpack checkpoint has been created. Note that + this is the same as `msgpack_checkpoint_dir`. + """ + from ray.rllib.policy.policy import Policy + + policy = Policy.from_checkpoint(policy_checkpoint) + + os.makedirs(msgpack_checkpoint_dir, exist_ok=True) + policy.export_checkpoint( + msgpack_checkpoint_dir, + policy_state=policy.get_state(), + checkpoint_format="msgpack", + ) + + # Release all resources used by the Policy. + del policy + + return msgpack_checkpoint_dir + + +@PublicAPI +def try_import_msgpack(error: bool = False): + """Tries importing msgpack and msgpack_numpy and returns the patched msgpack module. + + Returns None if error is False and msgpack or msgpack_numpy is not installed. + Raises an error, if error is True and the modules could not be imported. + + Args: + error: Whether to raise an error if msgpack/msgpack_numpy cannot be imported. + + Returns: + The `msgpack` module, with the msgpack_numpy module already patched in. This + means you can already encde and decode numpy arrays with the returned module. + + Raises: + ImportError: If error=True and msgpack/msgpack_numpy is not installed. + """ + try: + import msgpack + import msgpack_numpy + + # Make msgpack_numpy look like msgpack. + msgpack_numpy.patch() + + return msgpack + + except Exception: + if error: + raise ImportError( + "Could not import or setup msgpack and msgpack_numpy! " + "Try running `pip install msgpack msgpack_numpy` first." + ) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/compression.py b/lib/python3.12/site-packages/ray/rllib/utils/compression.py new file mode 100644 index 0000000000000000000000000000000000000000..7df5c46f11258f6b8fa65cd3ac166bbdc79508e6 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/compression.py @@ -0,0 +1,91 @@ +import base64 +import logging +import time + +import numpy as np + +from ray import cloudpickle as pickle +from ray.rllib.utils.annotations import DeveloperAPI + +logger = logging.getLogger(__name__) + +try: + import lz4.frame + + LZ4_ENABLED = True +except ImportError: + logger.warning( + "lz4 not available, disabling sample compression. " + "This will significantly impact RLlib performance. " + "To install lz4, run `pip install lz4`." + ) + LZ4_ENABLED = False + + +@DeveloperAPI +def compression_supported(): + return LZ4_ENABLED + + +@DeveloperAPI +def pack(data): + if LZ4_ENABLED: + data = pickle.dumps(data) + data = lz4.frame.compress(data) + # TODO(ekl) we shouldn't need to base64 encode this data, but this + # seems to not survive a transfer through the object store if we don't. + data = base64.b64encode(data).decode("ascii") + return data + + +@DeveloperAPI +def pack_if_needed(data): + if isinstance(data, np.ndarray): + data = pack(data) + return data + + +@DeveloperAPI +def unpack(data): + if LZ4_ENABLED: + data = base64.b64decode(data) + data = lz4.frame.decompress(data) + data = pickle.loads(data) + return data + + +@DeveloperAPI +def unpack_if_needed(data): + if is_compressed(data): + data = unpack(data) + return data + + +@DeveloperAPI +def is_compressed(data): + return isinstance(data, bytes) or isinstance(data, str) + + +# Intel(R) Core(TM) i7-4600U CPU @ 2.10GHz +# Compression speed: 753.664 MB/s +# Compression ratio: 87.4839812046 +# Decompression speed: 910.9504 MB/s +if __name__ == "__main__": + size = 32 * 80 * 80 * 4 + data = np.ones(size).reshape((32, 80, 80, 4)) + + count = 0 + start = time.time() + while time.time() - start < 1: + pack(data) + count += 1 + compressed = pack(data) + print("Compression speed: {} MB/s".format(count * size * 4 / 1e6)) + print("Compression ratio: {}".format(round(size * 4 / len(compressed), 2))) + + count = 0 + start = time.time() + while time.time() - start < 1: + unpack(compressed) + count += 1 + print("Decompression speed: {} MB/s".format(count * size * 4 / 1e6)) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/debug/__init__.py b/lib/python3.12/site-packages/ray/rllib/utils/debug/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d84dbe3d813c3a04826d1f7c3edc802a847a71e9 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/debug/__init__.py @@ -0,0 +1,9 @@ +from ray.rllib.utils.debug.deterministic import update_global_seed_if_necessary +from ray.rllib.utils.debug.memory import check_memory_leaks +from ray.rllib.utils.debug.summary import summarize + +__all__ = [ + "check_memory_leaks", + "summarize", + "update_global_seed_if_necessary", +] diff --git a/lib/python3.12/site-packages/ray/rllib/utils/debug/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/debug/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..673543ef2f0ab0a2395c3ac1b0a23984c39c2541 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/debug/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/debug/__pycache__/deterministic.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/debug/__pycache__/deterministic.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4680962de9e22355dec5573bc551195c5fa0962d Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/debug/__pycache__/deterministic.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/debug/__pycache__/memory.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/debug/__pycache__/memory.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dd21f36e7d93780516d52b9a43272dbc17006721 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/debug/__pycache__/memory.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/debug/__pycache__/summary.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/debug/__pycache__/summary.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..084f8e83a699a0dda2fc526a24320d4f76da298f Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/debug/__pycache__/summary.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/debug/deterministic.py b/lib/python3.12/site-packages/ray/rllib/utils/debug/deterministic.py new file mode 100644 index 0000000000000000000000000000000000000000..bf0a8b0671ed46e12b9cf146650e8cddbae797cb --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/debug/deterministic.py @@ -0,0 +1,42 @@ +import random +from typing import Optional + +import numpy as np + +from ray.rllib.utils.annotations import DeveloperAPI +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.torch_utils import set_torch_seed + + +@DeveloperAPI +def update_global_seed_if_necessary( + framework: Optional[str] = None, seed: Optional[int] = None +) -> None: + """Seed global modules such as random, numpy, torch, or tf. + + This is useful for debugging and testing. + + Args: + framework: The framework specifier (may be None). + seed: An optional int seed. If None, will not do + anything. + """ + if seed is None: + return + + # Python random module. + random.seed(seed) + # Numpy. + np.random.seed(seed) + + # Torch. + if framework == "torch": + set_torch_seed(seed=seed) + elif framework == "tf2": + tf1, tf, tfv = try_import_tf() + # Tf2.x. + if tfv == 2: + tf.random.set_seed(seed) + # Tf1.x. + else: + tf1.set_random_seed(seed) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/debug/memory.py b/lib/python3.12/site-packages/ray/rllib/utils/debug/memory.py new file mode 100644 index 0000000000000000000000000000000000000000..796e45717e7e4b242a32be4b38388e5b00126d5f --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/debug/memory.py @@ -0,0 +1,212 @@ +from collections import defaultdict +from typing import DefaultDict, List, Optional, Set + +import numpy as np +import tree # pip install dm_tree + +from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID, SampleBatch +from ray.rllib.utils.annotations import DeveloperAPI +from ray.util.debug import Suspect, _test_some_code_for_memory_leaks + + +@DeveloperAPI +def check_memory_leaks( + algorithm, + to_check: Optional[Set[str]] = None, + repeats: Optional[int] = None, + max_num_trials: int = 3, +) -> DefaultDict[str, List[Suspect]]: + """Diagnoses the given Algorithm for possible memory leaks. + + Isolates single components inside the Algorithm's local worker, e.g. the env, + policy, etc.. and calls some of their methods repeatedly, while checking + the memory footprints and keeping track of which lines in the code add + un-GC'd items to memory. + + Args: + algorithm: The Algorithm instance to test. + to_check: Set of strings to indentify components to test. Allowed strings + are: "env", "policy", "model", "rollout_worker". By default, check all + of these. + repeats: Number of times the test code block should get executed (per trial). + If a trial fails, a new trial may get started with a larger number of + repeats: actual_repeats = `repeats` * (trial + 1) (1st trial == 0). + max_num_trials: The maximum number of trials to run each check for. + + Raises: + A defaultdict(list) with keys being the `to_check` strings and values being + lists of Suspect instances that were found. + """ + local_worker = algorithm.env_runner + + # Which components should we test? + to_check = to_check or {"env", "model", "policy", "rollout_worker"} + + results_per_category = defaultdict(list) + + # Test a single sub-env (first in the VectorEnv)? + if "env" in to_check: + assert local_worker.async_env is not None, ( + "ERROR: Cannot test 'env' since given Algorithm does not have one " + "in its local worker. Try setting `create_local_env_runner=True`." + ) + + # Isolate the first sub-env in the vectorized setup and test it. + env = local_worker.async_env.get_sub_environments()[0] + action_space = env.action_space + # Always use same action to avoid numpy random caused memory leaks. + action_sample = action_space.sample() + + def code(): + ts = 0 + env.reset() + while True: + # If masking is used, try something like this: + # np.random.choice( + # action_space.n, p=(obs["action_mask"] / sum(obs["action_mask"]))) + _, _, done, _, _ = env.step(action_sample) + ts += 1 + if done: + break + + test = _test_some_code_for_memory_leaks( + desc="Looking for leaks in env, running through episodes.", + init=None, + code=code, + # How many times to repeat the function call? + repeats=repeats or 200, + max_num_trials=max_num_trials, + ) + if test: + results_per_category["env"].extend(test) + + # Test the policy (single-agent case only so far). + if "policy" in to_check: + policy = local_worker.policy_map[DEFAULT_POLICY_ID] + + # Get a fixed obs (B=10). + obs = tree.map_structure( + lambda s: np.stack([s] * 10, axis=0), policy.observation_space.sample() + ) + + print("Looking for leaks in Policy") + + def code(): + policy.compute_actions_from_input_dict( + { + "obs": obs, + } + ) + + # Call `compute_actions_from_input_dict()` n times. + test = _test_some_code_for_memory_leaks( + desc="Calling `compute_actions_from_input_dict()`.", + init=None, + code=code, + # How many times to repeat the function call? + repeats=repeats or 400, + # How many times to re-try if we find a suspicious memory + # allocation? + max_num_trials=max_num_trials, + ) + if test: + results_per_category["policy"].extend(test) + + # Testing this only makes sense if the learner API is disabled. + if not policy.config.get("enable_rl_module_and_learner", False): + # Call `learn_on_batch()` n times. + dummy_batch = policy._get_dummy_batch_from_view_requirements(batch_size=16) + + test = _test_some_code_for_memory_leaks( + desc="Calling `learn_on_batch()`.", + init=None, + code=lambda: policy.learn_on_batch(dummy_batch), + # How many times to repeat the function call? + repeats=repeats or 100, + max_num_trials=max_num_trials, + ) + if test: + results_per_category["policy"].extend(test) + + # Test only the model. + if "model" in to_check: + policy = local_worker.policy_map[DEFAULT_POLICY_ID] + + # Get a fixed obs. + obs = tree.map_structure(lambda s: s[None], policy.observation_space.sample()) + + print("Looking for leaks in Model") + + # Call `compute_actions_from_input_dict()` n times. + test = _test_some_code_for_memory_leaks( + desc="Calling `[model]()`.", + init=None, + code=lambda: policy.model({SampleBatch.OBS: obs}), + # How many times to repeat the function call? + repeats=repeats or 400, + # How many times to re-try if we find a suspicious memory + # allocation? + max_num_trials=max_num_trials, + ) + if test: + results_per_category["model"].extend(test) + + # Test the RolloutWorker. + if "rollout_worker" in to_check: + print("Looking for leaks in local RolloutWorker") + + def code(): + local_worker.sample() + local_worker.get_metrics() + + # Call `compute_actions_from_input_dict()` n times. + test = _test_some_code_for_memory_leaks( + desc="Calling `sample()` and `get_metrics()`.", + init=None, + code=code, + # How many times to repeat the function call? + repeats=repeats or 50, + # How many times to re-try if we find a suspicious memory + # allocation? + max_num_trials=max_num_trials, + ) + if test: + results_per_category["rollout_worker"].extend(test) + + if "learner" in to_check and algorithm.config.get( + "enable_rl_module_and_learner", False + ): + learner_group = algorithm.learner_group + assert learner_group._is_local, ( + "This test will miss leaks hidden in remote " + "workers. Please make sure that there is a " + "local learner inside the learner group for " + "this test." + ) + + dummy_batch = ( + algorithm.get_policy() + ._get_dummy_batch_from_view_requirements(batch_size=16) + .as_multi_agent() + ) + + print("Looking for leaks in Learner") + + def code(): + learner_group.update(dummy_batch) + + # Call `compute_actions_from_input_dict()` n times. + test = _test_some_code_for_memory_leaks( + desc="Calling `LearnerGroup.update()`.", + init=None, + code=code, + # How many times to repeat the function call? + repeats=repeats or 400, + # How many times to re-try if we find a suspicious memory + # allocation? + max_num_trials=max_num_trials, + ) + if test: + results_per_category["learner"].extend(test) + + return results_per_category diff --git a/lib/python3.12/site-packages/ray/rllib/utils/debug/summary.py b/lib/python3.12/site-packages/ray/rllib/utils/debug/summary.py new file mode 100644 index 0000000000000000000000000000000000000000..aaa8e023b3fae40e8d57a3191c648ffc3944dbf9 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/debug/summary.py @@ -0,0 +1,80 @@ +import pprint +from typing import Any + +import numpy as np + +from ray.rllib.policy.sample_batch import MultiAgentBatch, SampleBatch +from ray.rllib.utils.annotations import DeveloperAPI + +_printer = pprint.PrettyPrinter(indent=2, width=60) + + +@DeveloperAPI +def summarize(obj: Any) -> Any: + """Return a pretty-formatted string for an object. + + This has special handling for pretty-formatting of commonly used data types + in RLlib, such as SampleBatch, numpy arrays, etc. + + Args: + obj: The object to format. + + Returns: + The summarized object. + """ + + return _printer.pformat(_summarize(obj)) + + +def _summarize(obj): + if isinstance(obj, dict): + return {k: _summarize(v) for k, v in obj.items()} + elif hasattr(obj, "_asdict"): + return { + "type": obj.__class__.__name__, + "data": _summarize(obj._asdict()), + } + elif isinstance(obj, list): + return [_summarize(x) for x in obj] + elif isinstance(obj, tuple): + return tuple(_summarize(x) for x in obj) + elif isinstance(obj, np.ndarray): + if obj.size == 0: + return _StringValue("np.ndarray({}, dtype={})".format(obj.shape, obj.dtype)) + elif obj.dtype == object or obj.dtype.type is np.str_: + return _StringValue( + "np.ndarray({}, dtype={}, head={})".format( + obj.shape, obj.dtype, _summarize(obj[0]) + ) + ) + else: + return _StringValue( + "np.ndarray({}, dtype={}, min={}, max={}, mean={})".format( + obj.shape, + obj.dtype, + round(float(np.min(obj)), 3), + round(float(np.max(obj)), 3), + round(float(np.mean(obj)), 3), + ) + ) + elif isinstance(obj, MultiAgentBatch): + return { + "type": "MultiAgentBatch", + "policy_batches": _summarize(obj.policy_batches), + "count": obj.count, + } + elif isinstance(obj, SampleBatch): + return { + "type": "SampleBatch", + "data": {k: _summarize(v) for k, v in obj.items()}, + } + else: + return obj + + +class _StringValue: + def __init__(self, value): + self.value = value + + def __repr__(self): + return self.value diff --git a/lib/python3.12/site-packages/ray/rllib/utils/error.py b/lib/python3.12/site-packages/ray/rllib/utils/error.py new file mode 100644 index 0000000000000000000000000000000000000000..fb449b9bc52a2ce82681ba6dc9f54f2cc5748dc3 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/error.py @@ -0,0 +1,126 @@ +from ray.rllib.utils.annotations import PublicAPI + + +@PublicAPI +class UnsupportedSpaceException(Exception): + """Error for an unsupported action or observation space.""" + + pass + + +@PublicAPI +class EnvError(Exception): + """Error if we encounter an error during RL environment validation.""" + + pass + + +@PublicAPI +class MultiAgentEnvError(Exception): + """Error if we encounter an error during MultiAgentEnv stepping/validation.""" + + pass + + +@PublicAPI +class NotSerializable(Exception): + """Error if we encounter objects that can't be serialized by ray.""" + + pass + + +# ------- +# Error messages +# ------- + +# Message explaining there are no GPUs available for the +# num_gpus=n or num_gpus_per_env_runner=m settings. +ERR_MSG_NO_GPUS = """Found {} GPUs on your machine (GPU devices found: {})! If your + machine does not have any GPUs, you should set the config keys + `num_gpus_per_learner` and `num_gpus_per_env_runner` to 0. They may be set to + 1 by default for your particular RL algorithm.""" + +ERR_MSG_INVALID_ENV_DESCRIPTOR = """The env string you provided ('{}') is: +a) Not a supported or an installed environment. +b) Not a tune-registered environment creator. +c) Not a valid env class string. + +Try one of the following: +a) For Atari support: `pip install gymnasium[atari]` and prefix the environment name with `ale_py:`, for example, `"ale_py:ALE/Pong-v5"`. +b) To register your custom env, do `from ray import tune; tune.register_env('[name]', lambda cfg: [return env obj from here using cfg])`. + Then in your config, do `config.environment(env='[name]'). +c) Make sure you provide a fully qualified classpath, e.g.: + `ray.rllib.examples.envs.classes.repeat_after_me_env.RepeatAfterMeEnv` +""" + + +ERR_MSG_OLD_GYM_API = """Your environment ({}) does not abide to the new gymnasium-style API! +From Ray 2.3 on, RLlib only supports the new (gym>=0.26 or gymnasium) Env APIs. +{} +Learn more about the most important changes here: +https://github.com/openai/gym and here: https://github.com/Farama-Foundation/Gymnasium + +In order to fix this problem, do the following: + +1) Run `pip install gymnasium` on your command line. +2) Change all your import statements in your code from + `import gym` -> `import gymnasium as gym` OR + `from gym.spaces import Discrete` -> `from gymnasium.spaces import Discrete` + +For your custom (single agent) gym.Env classes: +3.1) Either wrap your old Env class via the provided `from gymnasium.wrappers import + EnvCompatibility` wrapper class. +3.2) Alternatively to 3.1: + - Change your `reset()` method to have the call signature 'def reset(self, *, + seed=None, options=None)' + - Return an additional info dict (empty dict should be fine) from your `reset()` + method. + - Return an additional `truncated` flag from your `step()` method (between `done` and + `info`). This flag should indicate, whether the episode was terminated prematurely + due to some time constraint or other kind of horizon setting. + +For your custom RLlib `MultiAgentEnv` classes: +4.1) Either wrap your old MultiAgentEnv via the provided + `from ray.rllib.env.wrappers.multi_agent_env_compatibility import + MultiAgentEnvCompatibility` wrapper class. +4.2) Alternatively to 4.1: + - Change your `reset()` method to have the call signature + 'def reset(self, *, seed=None, options=None)' + - Return an additional per-agent info dict (empty dict should be fine) from your + `reset()` method. + - Rename `dones` into `terminateds` and only set this to True, if the episode is really + done (as opposed to has been terminated prematurely due to some horizon/time-limit + setting). + - Return an additional `truncateds` per-agent dictionary flag from your `step()` + method, including the `__all__` key (100% analogous to your `dones/terminateds` + per-agent dict). + Return this new `truncateds` dict between `dones/terminateds` and `infos`. This + flag should indicate, whether the episode (for some agent or all agents) was + terminated prematurely due to some time constraint or other kind of horizon setting. +""" # noqa + + +ERR_MSG_TF_POLICY_CANNOT_SAVE_KERAS_MODEL = """Could not save keras model under self[TfPolicy].model.base_model! + This is either due to .. + a) .. this Policy's ModelV2 not having any `base_model` (tf.keras.Model) property + b) .. the ModelV2's `base_model` not being used by the Algorithm and thus its + variables not being properly initialized. +""" # noqa + +ERR_MSG_TORCH_POLICY_CANNOT_SAVE_MODEL = """Could not save torch model under self[TorchPolicy].model! + This is most likely due to the fact that you are using an Algorithm that + uses a Catalog-generated TorchModelV2 subclass, which is torch.save() cannot pickle. +""" # noqa + +# ------- +# HOWTO_ strings can be added to any error/warning/into message +# to eplain to the user, how to actually fix the encountered problem. +# ------- + +# HOWTO change the RLlib config, depending on how user runs the job. +HOWTO_CHANGE_CONFIG = """ +To change the config for `tune.Tuner().fit()` in a script: Modify the python dict + passed to `tune.Tuner(param_space=[...]).fit()`. +To change the config for an RLlib Algorithm instance: Modify the python dict + passed to the Algorithm's constructor, e.g. `PPO(config=[...])`. +""" diff --git a/lib/python3.12/site-packages/ray/rllib/utils/exploration/__init__.py b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c9ec99293f76487c45b3b9f07dcacf636eb737fb --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__init__.py @@ -0,0 +1,39 @@ +from ray.rllib.utils.exploration.curiosity import Curiosity +from ray.rllib.utils.exploration.epsilon_greedy import EpsilonGreedy +from ray.rllib.utils.exploration.exploration import Exploration +from ray.rllib.utils.exploration.gaussian_noise import GaussianNoise +from ray.rllib.utils.exploration.ornstein_uhlenbeck_noise import OrnsteinUhlenbeckNoise +from ray.rllib.utils.exploration.parameter_noise import ParameterNoise +from ray.rllib.utils.exploration.per_worker_epsilon_greedy import PerWorkerEpsilonGreedy +from ray.rllib.utils.exploration.per_worker_gaussian_noise import PerWorkerGaussianNoise +from ray.rllib.utils.exploration.per_worker_ornstein_uhlenbeck_noise import ( + PerWorkerOrnsteinUhlenbeckNoise, +) +from ray.rllib.utils.exploration.random import Random +from ray.rllib.utils.exploration.random_encoder import RE3 +from ray.rllib.utils.exploration.slate_epsilon_greedy import SlateEpsilonGreedy +from ray.rllib.utils.exploration.slate_soft_q import SlateSoftQ +from ray.rllib.utils.exploration.soft_q import SoftQ +from ray.rllib.utils.exploration.stochastic_sampling import StochasticSampling +from ray.rllib.utils.exploration.thompson_sampling import ThompsonSampling +from ray.rllib.utils.exploration.upper_confidence_bound import UpperConfidenceBound + +__all__ = [ + "Curiosity", + "Exploration", + "EpsilonGreedy", + "GaussianNoise", + "OrnsteinUhlenbeckNoise", + "ParameterNoise", + "PerWorkerEpsilonGreedy", + "PerWorkerGaussianNoise", + "PerWorkerOrnsteinUhlenbeckNoise", + "Random", + "RE3", + "SlateEpsilonGreedy", + "SlateSoftQ", + "SoftQ", + "StochasticSampling", + "ThompsonSampling", + "UpperConfidenceBound", +] diff --git a/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8c4cd22bdc1811efe93fceb9f9daf7ca1fcb7979 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/curiosity.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/curiosity.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2ddf90af654654db0c40f0c718df5db8c91ded3b Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/curiosity.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/epsilon_greedy.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/epsilon_greedy.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e0a8d6c72f02ff5d501cc0f94c1eeaf934b95e92 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/epsilon_greedy.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/exploration.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/exploration.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..523094011225a1bc4e6e86bcf24f6f23f04d334e Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/exploration.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/gaussian_noise.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/gaussian_noise.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4d4329006cdf76477a5f3510c878e0d3c482d444 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/gaussian_noise.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/ornstein_uhlenbeck_noise.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/ornstein_uhlenbeck_noise.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..333863cf602b1953611fe8df9ca8309cd68e57be Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/ornstein_uhlenbeck_noise.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/parameter_noise.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/parameter_noise.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f70e7e951126d3d11947733247287fd2327cbdcb Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/parameter_noise.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/per_worker_epsilon_greedy.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/per_worker_epsilon_greedy.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..529f8dea14889f82ae616284d3ab5a06d0f7e187 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/per_worker_epsilon_greedy.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/per_worker_gaussian_noise.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/per_worker_gaussian_noise.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b48a8e82c0fbcc374dc767c7d3c11f83cbd3a708 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/per_worker_gaussian_noise.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/per_worker_ornstein_uhlenbeck_noise.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/per_worker_ornstein_uhlenbeck_noise.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..71b3589abc75d44a9aafa462be2120e79c925a74 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/per_worker_ornstein_uhlenbeck_noise.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/random.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/random.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..601a2304b762a2901bcff010b18698b1f29bda87 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/random.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/random_encoder.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/random_encoder.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..915c5be68b03dafb8cc8394213cb6fd6544c5df7 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/random_encoder.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/slate_epsilon_greedy.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/slate_epsilon_greedy.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..50350579252f9c6ffa97c9a9184701dbbcfd92b8 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/slate_epsilon_greedy.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/slate_soft_q.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/slate_soft_q.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..08c61b3ea49cc65e90df2d9cd1e2fe22e7f82b88 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/slate_soft_q.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/soft_q.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/soft_q.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3733121387b613ac0596602d245bff4c6822bffc Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/soft_q.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/stochastic_sampling.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/stochastic_sampling.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e84bbbf6ad8f41bea1de33fe94cabb702cc3562f Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/stochastic_sampling.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/thompson_sampling.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/thompson_sampling.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..57c49d625cfc629bff99c4d27fe9bcd43a8b7653 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/thompson_sampling.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/upper_confidence_bound.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/upper_confidence_bound.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..48bc77268b6a3041f5bab4b16b367c0a00344db5 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/exploration/__pycache__/upper_confidence_bound.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/exploration/curiosity.py b/lib/python3.12/site-packages/ray/rllib/utils/exploration/curiosity.py new file mode 100644 index 0000000000000000000000000000000000000000..72200db6ce4a204f3ab1e7cdfaadc964175dee82 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/exploration/curiosity.py @@ -0,0 +1,445 @@ +from typing import Optional, Tuple, Union + +import numpy as np +from gymnasium.spaces import Discrete, MultiDiscrete, Space + +from ray.rllib.models.action_dist import ActionDistribution +from ray.rllib.models.catalog import ModelCatalog +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.models.tf.tf_action_dist import Categorical, MultiCategorical +from ray.rllib.models.torch.misc import SlimFC +from ray.rllib.models.torch.torch_action_dist import ( + TorchCategorical, + TorchMultiCategorical, +) +from ray.rllib.models.utils import get_activation_fn +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils import NullContextManager +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.exploration.exploration import Exploration +from ray.rllib.utils.framework import try_import_tf, try_import_torch +from ray.rllib.utils.from_config import from_config +from ray.rllib.utils.tf_utils import get_placeholder, one_hot as tf_one_hot +from ray.rllib.utils.torch_utils import one_hot +from ray.rllib.utils.typing import FromConfigSpec, ModelConfigDict, TensorType + +tf1, tf, tfv = try_import_tf() +torch, nn = try_import_torch() +F = None +if nn is not None: + F = nn.functional + + +@OldAPIStack +class Curiosity(Exploration): + """Implementation of: + [1] Curiosity-driven Exploration by Self-supervised Prediction + Pathak, Agrawal, Efros, and Darrell - UC Berkeley - ICML 2017. + https://arxiv.org/pdf/1705.05363.pdf + + Learns a simplified model of the environment based on three networks: + 1) Embedding observations into latent space ("feature" network). + 2) Predicting the action, given two consecutive embedded observations + ("inverse" network). + 3) Predicting the next embedded obs, given an obs and action + ("forward" network). + + The less the agent is able to predict the actually observed next feature + vector, given obs and action (through the forwards network), the larger the + "intrinsic reward", which will be added to the extrinsic reward. + Therefore, if a state transition was unexpected, the agent becomes + "curious" and will further explore this transition leading to better + exploration in sparse rewards environments. + """ + + def __init__( + self, + action_space: Space, + *, + framework: str, + model: ModelV2, + feature_dim: int = 288, + feature_net_config: Optional[ModelConfigDict] = None, + inverse_net_hiddens: Tuple[int] = (256,), + inverse_net_activation: str = "relu", + forward_net_hiddens: Tuple[int] = (256,), + forward_net_activation: str = "relu", + beta: float = 0.2, + eta: float = 1.0, + lr: float = 1e-3, + sub_exploration: Optional[FromConfigSpec] = None, + **kwargs + ): + """Initializes a Curiosity object. + + Uses as defaults the hyperparameters described in [1]. + + Args: + feature_dim: The dimensionality of the feature (phi) + vectors. + feature_net_config: Optional model + configuration for the feature network, producing feature + vectors (phi) from observations. This can be used to configure + fcnet- or conv_net setups to properly process any observation + space. + inverse_net_hiddens: Tuple of the layer sizes of the + inverse (action predicting) NN head (on top of the feature + outputs for phi and phi'). + inverse_net_activation: Activation specifier for the inverse + net. + forward_net_hiddens: Tuple of the layer sizes of the + forward (phi' predicting) NN head. + forward_net_activation: Activation specifier for the forward + net. + beta: Weight for the forward loss (over the inverse loss, + which gets weight=1.0-beta) in the common loss term. + eta: Weight for intrinsic rewards before being added to + extrinsic ones. + lr: The learning rate for the curiosity-specific + optimizer, optimizing feature-, inverse-, and forward nets. + sub_exploration: The config dict for + the underlying Exploration to use (e.g. epsilon-greedy for + DQN). If None, uses the FromSpecDict provided in the Policy's + default config. + """ + if not isinstance(action_space, (Discrete, MultiDiscrete)): + raise ValueError( + "Only (Multi)Discrete action spaces supported for Curiosity so far!" + ) + + super().__init__(action_space, model=model, framework=framework, **kwargs) + + if self.policy_config["num_env_runners"] != 0: + raise ValueError( + "Curiosity exploration currently does not support parallelism." + " `num_workers` must be 0!" + ) + + self.feature_dim = feature_dim + if feature_net_config is None: + feature_net_config = self.policy_config["model"].copy() + self.feature_net_config = feature_net_config + self.inverse_net_hiddens = inverse_net_hiddens + self.inverse_net_activation = inverse_net_activation + self.forward_net_hiddens = forward_net_hiddens + self.forward_net_activation = forward_net_activation + + self.action_dim = ( + self.action_space.n + if isinstance(self.action_space, Discrete) + else np.sum(self.action_space.nvec) + ) + + self.beta = beta + self.eta = eta + self.lr = lr + # TODO: (sven) if sub_exploration is None, use Algorithm's default + # Exploration config. + if sub_exploration is None: + raise NotImplementedError + self.sub_exploration = sub_exploration + + # Creates modules/layers inside the actual ModelV2. + self._curiosity_feature_net = ModelCatalog.get_model_v2( + self.model.obs_space, + self.action_space, + self.feature_dim, + model_config=self.feature_net_config, + framework=self.framework, + name="feature_net", + ) + + self._curiosity_inverse_fcnet = self._create_fc_net( + [2 * self.feature_dim] + list(self.inverse_net_hiddens) + [self.action_dim], + self.inverse_net_activation, + name="inverse_net", + ) + + self._curiosity_forward_fcnet = self._create_fc_net( + [self.feature_dim + self.action_dim] + + list(self.forward_net_hiddens) + + [self.feature_dim], + self.forward_net_activation, + name="forward_net", + ) + + # This is only used to select the correct action + self.exploration_submodule = from_config( + cls=Exploration, + config=self.sub_exploration, + action_space=self.action_space, + framework=self.framework, + policy_config=self.policy_config, + model=self.model, + num_workers=self.num_workers, + worker_index=self.worker_index, + ) + + @override(Exploration) + def get_exploration_action( + self, + *, + action_distribution: ActionDistribution, + timestep: Union[int, TensorType], + explore: bool = True + ): + # Simply delegate to sub-Exploration module. + return self.exploration_submodule.get_exploration_action( + action_distribution=action_distribution, timestep=timestep, explore=explore + ) + + @override(Exploration) + def get_exploration_optimizer(self, optimizers): + # Create, but don't add Adam for curiosity NN updating to the policy. + # If we added and returned it here, it would be used in the policy's + # update loop, which we don't want (curiosity updating happens inside + # `postprocess_trajectory`). + if self.framework == "torch": + feature_params = list(self._curiosity_feature_net.parameters()) + inverse_params = list(self._curiosity_inverse_fcnet.parameters()) + forward_params = list(self._curiosity_forward_fcnet.parameters()) + + # Now that the Policy's own optimizer(s) have been created (from + # the Model parameters (IMPORTANT: w/o(!) the curiosity params), + # we can add our curiosity sub-modules to the Policy's Model. + self.model._curiosity_feature_net = self._curiosity_feature_net.to( + self.device + ) + self.model._curiosity_inverse_fcnet = self._curiosity_inverse_fcnet.to( + self.device + ) + self.model._curiosity_forward_fcnet = self._curiosity_forward_fcnet.to( + self.device + ) + self._optimizer = torch.optim.Adam( + forward_params + inverse_params + feature_params, lr=self.lr + ) + else: + self.model._curiosity_feature_net = self._curiosity_feature_net + self.model._curiosity_inverse_fcnet = self._curiosity_inverse_fcnet + self.model._curiosity_forward_fcnet = self._curiosity_forward_fcnet + # Feature net is a RLlib ModelV2, the other 2 are keras Models. + self._optimizer_var_list = ( + self._curiosity_feature_net.base_model.variables + + self._curiosity_inverse_fcnet.variables + + self._curiosity_forward_fcnet.variables + ) + self._optimizer = tf1.train.AdamOptimizer(learning_rate=self.lr) + # Create placeholders and initialize the loss. + if self.framework == "tf": + self._obs_ph = get_placeholder( + space=self.model.obs_space, name="_curiosity_obs" + ) + self._next_obs_ph = get_placeholder( + space=self.model.obs_space, name="_curiosity_next_obs" + ) + self._action_ph = get_placeholder( + space=self.model.action_space, name="_curiosity_action" + ) + ( + self._forward_l2_norm_sqared, + self._update_op, + ) = self._postprocess_helper_tf( + self._obs_ph, self._next_obs_ph, self._action_ph + ) + + return optimizers + + @override(Exploration) + def postprocess_trajectory(self, policy, sample_batch, tf_sess=None): + """Calculates phi values (obs, obs', and predicted obs') and ri. + + Also calculates forward and inverse losses and updates the curiosity + module on the provided batch using our optimizer. + """ + if self.framework != "torch": + self._postprocess_tf(policy, sample_batch, tf_sess) + else: + self._postprocess_torch(policy, sample_batch) + + def _postprocess_tf(self, policy, sample_batch, tf_sess): + # tf1 static-graph: Perform session call on our loss and update ops. + if self.framework == "tf": + forward_l2_norm_sqared, _ = tf_sess.run( + [self._forward_l2_norm_sqared, self._update_op], + feed_dict={ + self._obs_ph: sample_batch[SampleBatch.OBS], + self._next_obs_ph: sample_batch[SampleBatch.NEXT_OBS], + self._action_ph: sample_batch[SampleBatch.ACTIONS], + }, + ) + # tf-eager: Perform model calls, loss calculations, and optimizer + # stepping on the fly. + else: + forward_l2_norm_sqared, _ = self._postprocess_helper_tf( + sample_batch[SampleBatch.OBS], + sample_batch[SampleBatch.NEXT_OBS], + sample_batch[SampleBatch.ACTIONS], + ) + # Scale intrinsic reward by eta hyper-parameter. + sample_batch[SampleBatch.REWARDS] = ( + sample_batch[SampleBatch.REWARDS] + self.eta * forward_l2_norm_sqared + ) + + return sample_batch + + def _postprocess_helper_tf(self, obs, next_obs, actions): + with ( + tf.GradientTape() if self.framework != "tf" else NullContextManager() + ) as tape: + # Push both observations through feature net to get both phis. + phis, _ = self.model._curiosity_feature_net( + {SampleBatch.OBS: tf.concat([obs, next_obs], axis=0)} + ) + phi, next_phi = tf.split(phis, 2) + + # Predict next phi with forward model. + predicted_next_phi = self.model._curiosity_forward_fcnet( + tf.concat([phi, tf_one_hot(actions, self.action_space)], axis=-1) + ) + + # Forward loss term (predicted phi', given phi and action vs + # actually observed phi'). + forward_l2_norm_sqared = 0.5 * tf.reduce_sum( + tf.square(predicted_next_phi - next_phi), axis=-1 + ) + forward_loss = tf.reduce_mean(forward_l2_norm_sqared) + + # Inverse loss term (prediced action that led from phi to phi' vs + # actual action taken). + phi_cat_next_phi = tf.concat([phi, next_phi], axis=-1) + dist_inputs = self.model._curiosity_inverse_fcnet(phi_cat_next_phi) + action_dist = ( + Categorical(dist_inputs, self.model) + if isinstance(self.action_space, Discrete) + else MultiCategorical(dist_inputs, self.model, self.action_space.nvec) + ) + # Neg log(p); p=probability of observed action given the inverse-NN + # predicted action distribution. + inverse_loss = -action_dist.logp(tf.convert_to_tensor(actions)) + inverse_loss = tf.reduce_mean(inverse_loss) + + # Calculate the ICM loss. + loss = (1.0 - self.beta) * inverse_loss + self.beta * forward_loss + + # Step the optimizer. + if self.framework != "tf": + grads = tape.gradient(loss, self._optimizer_var_list) + grads_and_vars = [ + (g, v) for g, v in zip(grads, self._optimizer_var_list) if g is not None + ] + update_op = self._optimizer.apply_gradients(grads_and_vars) + else: + update_op = self._optimizer.minimize( + loss, var_list=self._optimizer_var_list + ) + + # Return the squared l2 norm and the optimizer update op. + return forward_l2_norm_sqared, update_op + + def _postprocess_torch(self, policy, sample_batch): + # Push both observations through feature net to get both phis. + phis, _ = self.model._curiosity_feature_net( + { + SampleBatch.OBS: torch.cat( + [ + torch.from_numpy(sample_batch[SampleBatch.OBS]).to( + policy.device + ), + torch.from_numpy(sample_batch[SampleBatch.NEXT_OBS]).to( + policy.device + ), + ] + ) + } + ) + phi, next_phi = torch.chunk(phis, 2) + actions_tensor = ( + torch.from_numpy(sample_batch[SampleBatch.ACTIONS]).long().to(policy.device) + ) + + # Predict next phi with forward model. + predicted_next_phi = self.model._curiosity_forward_fcnet( + torch.cat([phi, one_hot(actions_tensor, self.action_space).float()], dim=-1) + ) + + # Forward loss term (predicted phi', given phi and action vs actually + # observed phi'). + forward_l2_norm_sqared = 0.5 * torch.sum( + torch.pow(predicted_next_phi - next_phi, 2.0), dim=-1 + ) + forward_loss = torch.mean(forward_l2_norm_sqared) + + # Scale intrinsic reward by eta hyper-parameter. + sample_batch[SampleBatch.REWARDS] = ( + sample_batch[SampleBatch.REWARDS] + + self.eta * forward_l2_norm_sqared.detach().cpu().numpy() + ) + + # Inverse loss term (prediced action that led from phi to phi' vs + # actual action taken). + phi_cat_next_phi = torch.cat([phi, next_phi], dim=-1) + dist_inputs = self.model._curiosity_inverse_fcnet(phi_cat_next_phi) + action_dist = ( + TorchCategorical(dist_inputs, self.model) + if isinstance(self.action_space, Discrete) + else TorchMultiCategorical(dist_inputs, self.model, self.action_space.nvec) + ) + # Neg log(p); p=probability of observed action given the inverse-NN + # predicted action distribution. + inverse_loss = -action_dist.logp(actions_tensor) + inverse_loss = torch.mean(inverse_loss) + + # Calculate the ICM loss. + loss = (1.0 - self.beta) * inverse_loss + self.beta * forward_loss + # Perform an optimizer step. + self._optimizer.zero_grad() + loss.backward() + self._optimizer.step() + + # Return the postprocessed sample batch (with the corrected rewards). + return sample_batch + + def _create_fc_net(self, layer_dims, activation, name=None): + """Given a list of layer dimensions (incl. input-dim), creates FC-net. + + Args: + layer_dims (Tuple[int]): Tuple of layer dims, including the input + dimension. + activation: An activation specifier string (e.g. "relu"). + + Examples: + If layer_dims is [4,8,6] we'll have a two layer net: 4->8 (8 nodes) + and 8->6 (6 nodes), where the second layer (6 nodes) does not have + an activation anymore. 4 is the input dimension. + """ + layers = ( + [tf.keras.layers.Input(shape=(layer_dims[0],), name="{}_in".format(name))] + if self.framework != "torch" + else [] + ) + + for i in range(len(layer_dims) - 1): + act = activation if i < len(layer_dims) - 2 else None + if self.framework == "torch": + layers.append( + SlimFC( + in_size=layer_dims[i], + out_size=layer_dims[i + 1], + initializer=torch.nn.init.xavier_uniform_, + activation_fn=act, + ) + ) + else: + layers.append( + tf.keras.layers.Dense( + units=layer_dims[i + 1], + activation=get_activation_fn(act), + name="{}_{}".format(name, i), + ) + ) + + if self.framework == "torch": + return nn.Sequential(*layers) + else: + return tf.keras.Sequential(layers) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/exploration/epsilon_greedy.py b/lib/python3.12/site-packages/ray/rllib/utils/exploration/epsilon_greedy.py new file mode 100644 index 0000000000000000000000000000000000000000..93103d44d2c68ab897b8c86d9c8eb4f4c58337f8 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/exploration/epsilon_greedy.py @@ -0,0 +1,247 @@ +import random +from typing import Optional, Union + +import gymnasium as gym +import numpy as np +import tree # pip install dm_tree + +from ray.rllib.models.action_dist import ActionDistribution +from ray.rllib.models.torch.torch_action_dist import TorchMultiActionDistribution +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.exploration.exploration import Exploration, TensorType +from ray.rllib.utils.framework import get_variable, try_import_tf, try_import_torch +from ray.rllib.utils.from_config import from_config +from ray.rllib.utils.numpy import convert_to_numpy +from ray.rllib.utils.schedules import PiecewiseSchedule, Schedule +from ray.rllib.utils.torch_utils import FLOAT_MIN + +tf1, tf, tfv = try_import_tf() +torch, _ = try_import_torch() + + +@OldAPIStack +class EpsilonGreedy(Exploration): + """Epsilon-greedy Exploration class that produces exploration actions. + + When given a Model's output and a current epsilon value (based on some + Schedule), it produces a random action (if rand(1) < eps) or + uses the model-computed one (if rand(1) >= eps). + """ + + def __init__( + self, + action_space: gym.spaces.Space, + *, + framework: str, + initial_epsilon: float = 1.0, + final_epsilon: float = 0.05, + warmup_timesteps: int = 0, + epsilon_timesteps: int = int(1e5), + epsilon_schedule: Optional[Schedule] = None, + **kwargs, + ): + """Create an EpsilonGreedy exploration class. + + Args: + action_space: The action space the exploration should occur in. + framework: The framework specifier. + initial_epsilon: The initial epsilon value to use. + final_epsilon: The final epsilon value to use. + warmup_timesteps: The timesteps over which to not change epsilon in the + beginning. + epsilon_timesteps: The timesteps (additional to `warmup_timesteps`) + after which epsilon should always be `final_epsilon`. + E.g.: warmup_timesteps=20k epsilon_timesteps=50k -> After 70k timesteps, + epsilon will reach its final value. + epsilon_schedule: An optional Schedule object + to use (instead of constructing one from the given parameters). + """ + assert framework is not None + super().__init__(action_space=action_space, framework=framework, **kwargs) + + self.epsilon_schedule = from_config( + Schedule, epsilon_schedule, framework=framework + ) or PiecewiseSchedule( + endpoints=[ + (0, initial_epsilon), + (warmup_timesteps, initial_epsilon), + (warmup_timesteps + epsilon_timesteps, final_epsilon), + ], + outside_value=final_epsilon, + framework=self.framework, + ) + + # The current timestep value (tf-var or python int). + self.last_timestep = get_variable( + np.array(0, np.int64), + framework=framework, + tf_name="timestep", + dtype=np.int64, + ) + + # Build the tf-info-op. + if self.framework == "tf": + self._tf_state_op = self.get_state() + + @override(Exploration) + def get_exploration_action( + self, + *, + action_distribution: ActionDistribution, + timestep: Union[int, TensorType], + explore: Optional[Union[bool, TensorType]] = True, + ): + + if self.framework in ["tf2", "tf"]: + return self._get_tf_exploration_action_op( + action_distribution, explore, timestep + ) + else: + return self._get_torch_exploration_action( + action_distribution, explore, timestep + ) + + def _get_tf_exploration_action_op( + self, + action_distribution: ActionDistribution, + explore: Union[bool, TensorType], + timestep: Union[int, TensorType], + ) -> "tf.Tensor": + """TF method to produce the tf op for an epsilon exploration action. + + Args: + action_distribution: The instantiated ActionDistribution object + to work with when creating exploration actions. + + Returns: + The tf exploration-action op. + """ + # TODO: Support MultiActionDistr for tf. + q_values = action_distribution.inputs + epsilon = self.epsilon_schedule( + timestep if timestep is not None else self.last_timestep + ) + + # Get the exploit action as the one with the highest logit value. + exploit_action = tf.argmax(q_values, axis=1) + + batch_size = tf.shape(q_values)[0] + # Mask out actions with q-value=-inf so that we don't even consider + # them for exploration. + random_valid_action_logits = tf.where( + tf.equal(q_values, tf.float32.min), + tf.ones_like(q_values) * tf.float32.min, + tf.ones_like(q_values), + ) + random_actions = tf.squeeze( + tf.random.categorical(random_valid_action_logits, 1), axis=1 + ) + + chose_random = ( + tf.random.uniform( + tf.stack([batch_size]), minval=0, maxval=1, dtype=tf.float32 + ) + < epsilon + ) + + action = tf.cond( + pred=tf.constant(explore, dtype=tf.bool) + if isinstance(explore, bool) + else explore, + true_fn=(lambda: tf.where(chose_random, random_actions, exploit_action)), + false_fn=lambda: exploit_action, + ) + + if self.framework == "tf2" and not self.policy_config["eager_tracing"]: + self.last_timestep = timestep + return action, tf.zeros_like(action, dtype=tf.float32) + else: + assign_op = tf1.assign(self.last_timestep, tf.cast(timestep, tf.int64)) + with tf1.control_dependencies([assign_op]): + return action, tf.zeros_like(action, dtype=tf.float32) + + def _get_torch_exploration_action( + self, + action_distribution: ActionDistribution, + explore: bool, + timestep: Union[int, TensorType], + ) -> "torch.Tensor": + """Torch method to produce an epsilon exploration action. + + Args: + action_distribution: The instantiated + ActionDistribution object to work with when creating + exploration actions. + + Returns: + The exploration-action. + """ + q_values = action_distribution.inputs + self.last_timestep = timestep + exploit_action = action_distribution.deterministic_sample() + batch_size = q_values.size()[0] + action_logp = torch.zeros(batch_size, dtype=torch.float) + + # Explore. + if explore: + # Get the current epsilon. + epsilon = self.epsilon_schedule(self.last_timestep) + if isinstance(action_distribution, TorchMultiActionDistribution): + exploit_action = tree.flatten(exploit_action) + for i in range(batch_size): + if random.random() < epsilon: + # TODO: (bcahlit) Mask out actions + random_action = tree.flatten(self.action_space.sample()) + for j in range(len(exploit_action)): + exploit_action[j][i] = torch.tensor(random_action[j]) + exploit_action = tree.unflatten_as( + action_distribution.action_space_struct, exploit_action + ) + + return exploit_action, action_logp + + else: + # Mask out actions, whose Q-values are -inf, so that we don't + # even consider them for exploration. + random_valid_action_logits = torch.where( + q_values <= FLOAT_MIN, + torch.ones_like(q_values) * 0.0, + torch.ones_like(q_values), + ) + # A random action. + random_actions = torch.squeeze( + torch.multinomial(random_valid_action_logits, 1), axis=1 + ) + + # Pick either random or greedy. + action = torch.where( + torch.empty((batch_size,)).uniform_().to(self.device) < epsilon, + random_actions, + exploit_action, + ) + + return action, action_logp + # Return the deterministic "sample" (argmax) over the logits. + else: + return exploit_action, action_logp + + @override(Exploration) + def get_state(self, sess: Optional["tf.Session"] = None): + if sess: + return sess.run(self._tf_state_op) + eps = self.epsilon_schedule(self.last_timestep) + return { + "cur_epsilon": convert_to_numpy(eps) if self.framework != "tf" else eps, + "last_timestep": convert_to_numpy(self.last_timestep) + if self.framework != "tf" + else self.last_timestep, + } + + @override(Exploration) + def set_state(self, state: dict, sess: Optional["tf.Session"] = None) -> None: + if self.framework == "tf": + self.last_timestep.load(state["last_timestep"], session=sess) + elif isinstance(self.last_timestep, int): + self.last_timestep = state["last_timestep"] + else: + self.last_timestep.assign(state["last_timestep"]) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/exploration/exploration.py b/lib/python3.12/site-packages/ray/rllib/utils/exploration/exploration.py new file mode 100644 index 0000000000000000000000000000000000000000..7024157d0cd32e3032c7c10c2e370c47e943fc62 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/exploration/exploration.py @@ -0,0 +1,210 @@ +from typing import TYPE_CHECKING, Dict, List, Optional, Union + +from gymnasium.spaces import Space + +from ray.rllib.env.base_env import BaseEnv +from ray.rllib.models.action_dist import ActionDistribution +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.framework import TensorType, try_import_torch +from ray.rllib.utils.typing import AlgorithmConfigDict, LocalOptimizer + +if TYPE_CHECKING: + from ray.rllib.policy.policy import Policy + from ray.rllib.utils import try_import_tf + + _, tf, _ = try_import_tf() + +_, nn = try_import_torch() + + +@OldAPIStack +class Exploration: + """Implements an exploration strategy for Policies. + + An Exploration takes model outputs, a distribution, and a timestep from + the agent and computes an action to apply to the environment using an + implemented exploration schema. + """ + + def __init__( + self, + action_space: Space, + *, + framework: str, + policy_config: AlgorithmConfigDict, + model: ModelV2, + num_workers: int, + worker_index: int + ): + """ + Args: + action_space: The action space in which to explore. + framework: One of "tf" or "torch". + policy_config: The Policy's config dict. + model: The Policy's model. + num_workers: The overall number of workers used. + worker_index: The index of the worker using this class. + """ + self.action_space = action_space + self.policy_config = policy_config + self.model = model + self.num_workers = num_workers + self.worker_index = worker_index + self.framework = framework + # The device on which the Model has been placed. + # This Exploration will be on the same device. + self.device = None + if isinstance(self.model, nn.Module): + params = list(self.model.parameters()) + if params: + self.device = params[0].device + + def before_compute_actions( + self, + *, + timestep: Optional[Union[TensorType, int]] = None, + explore: Optional[Union[TensorType, bool]] = None, + tf_sess: Optional["tf.Session"] = None, + **kwargs + ): + """Hook for preparations before policy.compute_actions() is called. + + Args: + timestep: An optional timestep tensor. + explore: An optional explore boolean flag. + tf_sess: The tf-session object to use. + **kwargs: Forward compatibility kwargs. + """ + pass + + # fmt: off + # __sphinx_doc_begin_get_exploration_action__ + + def get_exploration_action(self, + *, + action_distribution: ActionDistribution, + timestep: Union[TensorType, int], + explore: bool = True): + """Returns a (possibly) exploratory action and its log-likelihood. + + Given the Model's logits outputs and action distribution, returns an + exploratory action. + + Args: + action_distribution: The instantiated + ActionDistribution object to work with when creating + exploration actions. + timestep: The current sampling time step. It can be a tensor + for TF graph mode, otherwise an integer. + explore: True: "Normal" exploration behavior. + False: Suppress all exploratory behavior and return + a deterministic action. + + Returns: + A tuple consisting of 1) the chosen exploration action or a + tf-op to fetch the exploration action from the graph and + 2) the log-likelihood of the exploration action. + """ + pass + + # __sphinx_doc_end_get_exploration_action__ + # fmt: on + + def on_episode_start( + self, + policy: "Policy", + *, + environment: BaseEnv = None, + episode: int = None, + tf_sess: Optional["tf.Session"] = None + ): + """Handles necessary exploration logic at the beginning of an episode. + + Args: + policy: The Policy object that holds this Exploration. + environment: The environment object we are acting in. + episode: The number of the episode that is starting. + tf_sess: In case of tf, the session object. + """ + pass + + def on_episode_end( + self, + policy: "Policy", + *, + environment: BaseEnv = None, + episode: int = None, + tf_sess: Optional["tf.Session"] = None + ): + """Handles necessary exploration logic at the end of an episode. + + Args: + policy: The Policy object that holds this Exploration. + environment: The environment object we are acting in. + episode: The number of the episode that is starting. + tf_sess: In case of tf, the session object. + """ + pass + + def postprocess_trajectory( + self, + policy: "Policy", + sample_batch: SampleBatch, + tf_sess: Optional["tf.Session"] = None, + ): + """Handles post-processing of done episode trajectories. + + Changes the given batch in place. This callback is invoked by the + sampler after policy.postprocess_trajectory() is called. + + Args: + policy: The owning policy object. + sample_batch: The SampleBatch object to post-process. + tf_sess: An optional tf.Session object. + """ + return sample_batch + + def get_exploration_optimizer( + self, optimizers: List[LocalOptimizer] + ) -> List[LocalOptimizer]: + """May add optimizer(s) to the Policy's own `optimizers`. + + The number of optimizers (Policy's plus Exploration's optimizers) must + match the number of loss terms produced by the Policy's loss function + and the Exploration component's loss terms. + + Args: + optimizers: The list of the Policy's local optimizers. + + Returns: + The updated list of local optimizers to use on the different + loss terms. + """ + return optimizers + + def get_state(self, sess: Optional["tf.Session"] = None) -> Dict[str, TensorType]: + """Returns the current exploration state. + + Args: + sess: An optional tf Session object to use. + + Returns: + The Exploration object's current state. + """ + return {} + + def set_state(self, state: object, sess: Optional["tf.Session"] = None) -> None: + """Sets the Exploration object's state to the given values. + + Note that some exploration components are stateless, even though they + decay some values over time (e.g. EpsilonGreedy). However the decay is + only dependent on the current global timestep of the policy and we + therefore don't need to keep track of it. + + Args: + state: The state to set this Exploration to. + sess: An optional tf Session object to use. + """ + pass diff --git a/lib/python3.12/site-packages/ray/rllib/utils/exploration/gaussian_noise.py b/lib/python3.12/site-packages/ray/rllib/utils/exploration/gaussian_noise.py new file mode 100644 index 0000000000000000000000000000000000000000..cce79f01ee093a4f5c72bc90e22424c30bf0418e --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/exploration/gaussian_noise.py @@ -0,0 +1,248 @@ +from typing import Optional, Union + +import numpy as np +from gymnasium.spaces import Space + +from ray.rllib.models.action_dist import ActionDistribution +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.exploration.exploration import Exploration +from ray.rllib.utils.exploration.random import Random +from ray.rllib.utils.framework import ( + TensorType, + get_variable, + try_import_tf, + try_import_torch, +) +from ray.rllib.utils.numpy import convert_to_numpy +from ray.rllib.utils.schedules import Schedule +from ray.rllib.utils.schedules.piecewise_schedule import PiecewiseSchedule +from ray.rllib.utils.tf_utils import zero_logps_from_actions + +tf1, tf, tfv = try_import_tf() +torch, _ = try_import_torch() + + +@OldAPIStack +class GaussianNoise(Exploration): + """An exploration that adds white noise to continuous actions. + + If explore=True, returns actions plus scale (annealed over time) x + Gaussian noise. Also, some completely random period is possible at the + beginning. + + If explore=False, returns the deterministic action. + """ + + def __init__( + self, + action_space: Space, + *, + framework: str, + model: ModelV2, + random_timesteps: int = 1000, + stddev: float = 0.1, + initial_scale: float = 1.0, + final_scale: float = 0.02, + scale_timesteps: int = 10000, + scale_schedule: Optional[Schedule] = None, + **kwargs + ): + """Initializes a GaussianNoise instance. + + Args: + random_timesteps: The number of timesteps for which to act + completely randomly. Only after this number of timesteps, the + `self.scale` annealing process will start (see below). + stddev: The stddev (sigma) to use for the + Gaussian noise to be added to the actions. + initial_scale: The initial scaling weight to multiply + the noise with. + final_scale: The final scaling weight to multiply + the noise with. + scale_timesteps: The timesteps over which to linearly anneal + the scaling factor (after(!) having used random actions for + `random_timesteps` steps). + scale_schedule: An optional Schedule object + to use (instead of constructing one from the given parameters). + """ + assert framework is not None + super().__init__(action_space, model=model, framework=framework, **kwargs) + + # Create the Random exploration module (used for the first n + # timesteps). + self.random_timesteps = random_timesteps + self.random_exploration = Random( + action_space, model=self.model, framework=self.framework, **kwargs + ) + + self.stddev = stddev + # The `scale` annealing schedule. + self.scale_schedule = scale_schedule or PiecewiseSchedule( + endpoints=[ + (random_timesteps, initial_scale), + (random_timesteps + scale_timesteps, final_scale), + ], + outside_value=final_scale, + framework=self.framework, + ) + + # The current timestep value (tf-var or python int). + self.last_timestep = get_variable( + np.array(0, np.int64), + framework=self.framework, + tf_name="timestep", + dtype=np.int64, + ) + + # Build the tf-info-op. + if self.framework == "tf": + self._tf_state_op = self.get_state() + + @override(Exploration) + def get_exploration_action( + self, + *, + action_distribution: ActionDistribution, + timestep: Union[int, TensorType], + explore: bool = True + ): + # Adds IID Gaussian noise for exploration, TD3-style. + if self.framework == "torch": + return self._get_torch_exploration_action( + action_distribution, explore, timestep + ) + else: + return self._get_tf_exploration_action_op( + action_distribution, explore, timestep + ) + + def _get_tf_exploration_action_op( + self, + action_dist: ActionDistribution, + explore: bool, + timestep: Union[int, TensorType], + ): + ts = timestep if timestep is not None else self.last_timestep + + # The deterministic actions (if explore=False). + deterministic_actions = action_dist.deterministic_sample() + + # Take a Gaussian sample with our stddev (mean=0.0) and scale it. + gaussian_sample = self.scale_schedule(ts) * tf.random.normal( + tf.shape(deterministic_actions), stddev=self.stddev + ) + + # Stochastic actions could either be: random OR action + noise. + random_actions, _ = self.random_exploration.get_tf_exploration_action_op( + action_dist, explore + ) + stochastic_actions = tf.cond( + pred=tf.convert_to_tensor(ts < self.random_timesteps), + true_fn=lambda: random_actions, + false_fn=lambda: tf.clip_by_value( + deterministic_actions + gaussian_sample, + self.action_space.low * tf.ones_like(deterministic_actions), + self.action_space.high * tf.ones_like(deterministic_actions), + ), + ) + + # Chose by `explore` (main exploration switch). + action = tf.cond( + pred=tf.constant(explore, dtype=tf.bool) + if isinstance(explore, bool) + else explore, + true_fn=lambda: stochastic_actions, + false_fn=lambda: deterministic_actions, + ) + # Logp=always zero. + logp = zero_logps_from_actions(deterministic_actions) + + # Increment `last_timestep` by 1 (or set to `timestep`). + if self.framework == "tf2": + if timestep is None: + self.last_timestep.assign_add(1) + else: + self.last_timestep.assign(tf.cast(timestep, tf.int64)) + return action, logp + else: + assign_op = ( + tf1.assign_add(self.last_timestep, 1) + if timestep is None + else tf1.assign(self.last_timestep, timestep) + ) + with tf1.control_dependencies([assign_op]): + return action, logp + + def _get_torch_exploration_action( + self, + action_dist: ActionDistribution, + explore: bool, + timestep: Union[int, TensorType], + ): + # Set last timestep or (if not given) increase by one. + self.last_timestep = ( + timestep if timestep is not None else self.last_timestep + 1 + ) + + # Apply exploration. + if explore: + # Random exploration phase. + if self.last_timestep < self.random_timesteps: + action, _ = self.random_exploration.get_torch_exploration_action( + action_dist, explore=True + ) + # Take a Gaussian sample with our stddev (mean=0.0) and scale it. + else: + det_actions = action_dist.deterministic_sample() + scale = self.scale_schedule(self.last_timestep) + gaussian_sample = scale * torch.normal( + mean=torch.zeros(det_actions.size()), std=self.stddev + ).to(self.device) + action = torch.min( + torch.max( + det_actions + gaussian_sample, + torch.tensor( + self.action_space.low, + dtype=torch.float32, + device=self.device, + ), + ), + torch.tensor( + self.action_space.high, dtype=torch.float32, device=self.device + ), + ) + # No exploration -> Return deterministic actions. + else: + action = action_dist.deterministic_sample() + + # Logp=always zero. + logp = torch.zeros((action.size()[0],), dtype=torch.float32, device=self.device) + + return action, logp + + @override(Exploration) + def get_state(self, sess: Optional["tf.Session"] = None): + """Returns the current scale value. + + Returns: + Union[float,tf.Tensor[float]]: The current scale value. + """ + if sess: + return sess.run(self._tf_state_op) + scale = self.scale_schedule(self.last_timestep) + return { + "cur_scale": convert_to_numpy(scale) if self.framework != "tf" else scale, + "last_timestep": convert_to_numpy(self.last_timestep) + if self.framework != "tf" + else self.last_timestep, + } + + @override(Exploration) + def set_state(self, state: dict, sess: Optional["tf.Session"] = None) -> None: + if self.framework == "tf": + self.last_timestep.load(state["last_timestep"], session=sess) + elif isinstance(self.last_timestep, int): + self.last_timestep = state["last_timestep"] + else: + self.last_timestep.assign(state["last_timestep"]) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/exploration/ornstein_uhlenbeck_noise.py b/lib/python3.12/site-packages/ray/rllib/utils/exploration/ornstein_uhlenbeck_noise.py new file mode 100644 index 0000000000000000000000000000000000000000..f9151733bac8f77676c60a7dac19028917d49c13 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/exploration/ornstein_uhlenbeck_noise.py @@ -0,0 +1,274 @@ +from typing import Optional, Union + +import numpy as np + +from ray.rllib.models.action_dist import ActionDistribution +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.exploration.gaussian_noise import GaussianNoise +from ray.rllib.utils.framework import ( + TensorType, + get_variable, + try_import_tf, + try_import_torch, +) +from ray.rllib.utils.numpy import convert_to_numpy +from ray.rllib.utils.schedules import Schedule +from ray.rllib.utils.tf_utils import zero_logps_from_actions + +tf1, tf, tfv = try_import_tf() +torch, _ = try_import_torch() + + +@OldAPIStack +class OrnsteinUhlenbeckNoise(GaussianNoise): + """An exploration that adds Ornstein-Uhlenbeck noise to continuous actions. + + If explore=True, returns sampled actions plus a noise term X, + which changes according to this formula: + Xt+1 = -theta*Xt + sigma*N[0,stddev], where theta, sigma and stddev are + constants. Also, some completely random period is possible at the + beginning. + If explore=False, returns the deterministic action. + """ + + def __init__( + self, + action_space, + *, + framework: str, + ou_theta: float = 0.15, + ou_sigma: float = 0.2, + ou_base_scale: float = 0.1, + random_timesteps: int = 1000, + initial_scale: float = 1.0, + final_scale: float = 0.02, + scale_timesteps: int = 10000, + scale_schedule: Optional[Schedule] = None, + **kwargs + ): + """Initializes an Ornstein-Uhlenbeck Exploration object. + + Args: + action_space: The gym action space used by the environment. + ou_theta: The theta parameter of the Ornstein-Uhlenbeck process. + ou_sigma: The sigma parameter of the Ornstein-Uhlenbeck process. + ou_base_scale: A fixed scaling factor, by which all OU- + noise is multiplied. NOTE: This is on top of the parent + GaussianNoise's scaling. + random_timesteps: The number of timesteps for which to act + completely randomly. Only after this number of timesteps, the + `self.scale` annealing process will start (see below). + initial_scale: The initial scaling weight to multiply the + noise with. + final_scale: The final scaling weight to multiply the noise with. + scale_timesteps: The timesteps over which to linearly anneal the + scaling factor (after(!) having used random actions for + `random_timesteps` steps. + scale_schedule: An optional Schedule object to use (instead + of constructing one from the given parameters). + framework: One of None, "tf", "torch". + """ + # The current OU-state value (gets updated each time, an eploration + # action is computed). + self.ou_state = get_variable( + np.array(action_space.low.size * [0.0], dtype=np.float32), + framework=framework, + tf_name="ou_state", + torch_tensor=True, + device=None, + ) + + super().__init__( + action_space, + framework=framework, + random_timesteps=random_timesteps, + initial_scale=initial_scale, + final_scale=final_scale, + scale_timesteps=scale_timesteps, + scale_schedule=scale_schedule, + stddev=1.0, # Force `self.stddev` to 1.0. + **kwargs + ) + self.ou_theta = ou_theta + self.ou_sigma = ou_sigma + self.ou_base_scale = ou_base_scale + # Now that we know the device, move ou_state there, in case of PyTorch. + if self.framework == "torch" and self.device is not None: + self.ou_state = self.ou_state.to(self.device) + + @override(GaussianNoise) + def _get_tf_exploration_action_op( + self, + action_dist: ActionDistribution, + explore: Union[bool, TensorType], + timestep: Union[int, TensorType], + ): + ts = timestep if timestep is not None else self.last_timestep + scale = self.scale_schedule(ts) + + # The deterministic actions (if explore=False). + deterministic_actions = action_dist.deterministic_sample() + + # Apply base-scaled and time-annealed scaled OU-noise to + # deterministic actions. + gaussian_sample = tf.random.normal( + shape=[self.action_space.low.size], stddev=self.stddev + ) + ou_new = self.ou_theta * -self.ou_state + self.ou_sigma * gaussian_sample + if self.framework == "tf2": + self.ou_state.assign_add(ou_new) + ou_state_new = self.ou_state + else: + ou_state_new = tf1.assign_add(self.ou_state, ou_new) + high_m_low = self.action_space.high - self.action_space.low + high_m_low = tf.where( + tf.math.is_inf(high_m_low), tf.ones_like(high_m_low), high_m_low + ) + noise = scale * self.ou_base_scale * ou_state_new * high_m_low + stochastic_actions = tf.clip_by_value( + deterministic_actions + noise, + self.action_space.low * tf.ones_like(deterministic_actions), + self.action_space.high * tf.ones_like(deterministic_actions), + ) + + # Stochastic actions could either be: random OR action + noise. + random_actions, _ = self.random_exploration.get_tf_exploration_action_op( + action_dist, explore + ) + exploration_actions = tf.cond( + pred=tf.convert_to_tensor(ts < self.random_timesteps), + true_fn=lambda: random_actions, + false_fn=lambda: stochastic_actions, + ) + + # Chose by `explore` (main exploration switch). + action = tf.cond( + pred=tf.constant(explore, dtype=tf.bool) + if isinstance(explore, bool) + else explore, + true_fn=lambda: exploration_actions, + false_fn=lambda: deterministic_actions, + ) + # Logp=always zero. + logp = zero_logps_from_actions(deterministic_actions) + + # Increment `last_timestep` by 1 (or set to `timestep`). + if self.framework == "tf2": + if timestep is None: + self.last_timestep.assign_add(1) + else: + self.last_timestep.assign(tf.cast(timestep, tf.int64)) + else: + assign_op = ( + tf1.assign_add(self.last_timestep, 1) + if timestep is None + else tf1.assign(self.last_timestep, timestep) + ) + with tf1.control_dependencies([assign_op, ou_state_new]): + action = tf.identity(action) + logp = tf.identity(logp) + + return action, logp + + @override(GaussianNoise) + def _get_torch_exploration_action( + self, + action_dist: ActionDistribution, + explore: bool, + timestep: Union[int, TensorType], + ): + # Set last timestep or (if not given) increase by one. + self.last_timestep = ( + timestep if timestep is not None else self.last_timestep + 1 + ) + + # Apply exploration. + if explore: + # Random exploration phase. + if self.last_timestep < self.random_timesteps: + action, _ = self.random_exploration.get_torch_exploration_action( + action_dist, explore=True + ) + # Apply base-scaled and time-annealed scaled OU-noise to + # deterministic actions. + else: + det_actions = action_dist.deterministic_sample() + scale = self.scale_schedule(self.last_timestep) + gaussian_sample = scale * torch.normal( + mean=torch.zeros(self.ou_state.size()), std=1.0 + ).to(self.device) + ou_new = ( + self.ou_theta * -self.ou_state + self.ou_sigma * gaussian_sample + ) + self.ou_state += ou_new + high_m_low = torch.from_numpy( + self.action_space.high - self.action_space.low + ).to(self.device) + high_m_low = torch.where( + torch.isinf(high_m_low), + torch.ones_like(high_m_low).to(self.device), + high_m_low, + ) + noise = scale * self.ou_base_scale * self.ou_state * high_m_low + + action = torch.min( + torch.max( + det_actions + noise, + torch.tensor( + self.action_space.low, + dtype=torch.float32, + device=self.device, + ), + ), + torch.tensor( + self.action_space.high, dtype=torch.float32, device=self.device + ), + ) + + # No exploration -> Return deterministic actions. + else: + action = action_dist.deterministic_sample() + + # Logp=always zero. + logp = torch.zeros((action.size()[0],), dtype=torch.float32, device=self.device) + + return action, logp + + @override(GaussianNoise) + def get_state(self, sess: Optional["tf.Session"] = None): + """Returns the current scale value. + + Returns: + Union[float,tf.Tensor[float]]: The current scale value. + """ + if sess: + return sess.run( + dict( + self._tf_state_op, + **{ + "ou_state": self.ou_state, + } + ) + ) + + state = super().get_state() + return dict( + state, + **{ + "ou_state": convert_to_numpy(self.ou_state) + if self.framework != "tf" + else self.ou_state, + } + ) + + @override(GaussianNoise) + def set_state(self, state: dict, sess: Optional["tf.Session"] = None) -> None: + if self.framework == "tf": + self.ou_state.load(state["ou_state"], session=sess) + elif isinstance(self.ou_state, np.ndarray): + self.ou_state = state["ou_state"] + elif torch and torch.is_tensor(self.ou_state): + self.ou_state = torch.from_numpy(state["ou_state"]) + else: + self.ou_state.assign(state["ou_state"]) + super().set_state(state, sess=sess) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/exploration/parameter_noise.py b/lib/python3.12/site-packages/ray/rllib/utils/exploration/parameter_noise.py new file mode 100644 index 0000000000000000000000000000000000000000..349238761830391fa33b4115721e01c3f2838f2a --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/exploration/parameter_noise.py @@ -0,0 +1,441 @@ +from typing import TYPE_CHECKING, Optional, Union + +import numpy as np +from gymnasium.spaces import Box, Discrete + +from ray.rllib.env.base_env import BaseEnv +from ray.rllib.models.action_dist import ActionDistribution +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.models.tf.tf_action_dist import Categorical, Deterministic +from ray.rllib.models.torch.torch_action_dist import ( + TorchCategorical, + TorchDeterministic, +) +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.exploration.exploration import Exploration +from ray.rllib.utils.framework import get_variable, try_import_tf, try_import_torch +from ray.rllib.utils.from_config import from_config +from ray.rllib.utils.numpy import SMALL_NUMBER, softmax +from ray.rllib.utils.typing import TensorType + +if TYPE_CHECKING: + from ray.rllib.policy.policy import Policy + +tf1, tf, tfv = try_import_tf() +torch, _ = try_import_torch() + + +@OldAPIStack +class ParameterNoise(Exploration): + """An exploration that changes a Model's parameters. + + Implemented based on: + [1] https://openai.com/research/better-exploration-with-parameter-noise + [2] https://arxiv.org/pdf/1706.01905.pdf + + At the beginning of an episode, Gaussian noise is added to all weights + of the model. At the end of the episode, the noise is undone and an action + diff (pi-delta) is calculated, from which we determine the changes in the + noise's stddev for the next episode. + """ + + def __init__( + self, + action_space, + *, + framework: str, + policy_config: dict, + model: ModelV2, + initial_stddev: float = 1.0, + random_timesteps: int = 10000, + sub_exploration: Optional[dict] = None, + **kwargs + ): + """Initializes a ParameterNoise Exploration object. + + Args: + initial_stddev: The initial stddev to use for the noise. + random_timesteps: The number of timesteps to act completely + randomly (see [1]). + sub_exploration: Optional sub-exploration config. + None for auto-detection/setup. + """ + assert framework is not None + super().__init__( + action_space, + policy_config=policy_config, + model=model, + framework=framework, + **kwargs + ) + + self.stddev = get_variable( + initial_stddev, framework=self.framework, tf_name="stddev" + ) + self.stddev_val = initial_stddev # Out-of-graph tf value holder. + + # The weight variables of the Model where noise should be applied to. + # This excludes any variable, whose name contains "LayerNorm" (those + # are BatchNormalization layers, which should not be perturbed). + self.model_variables = [ + v + for k, v in self.model.trainable_variables(as_dict=True).items() + if "LayerNorm" not in k + ] + # Our noise to be added to the weights. Each item in `self.noise` + # corresponds to one Model variable and holding the Gaussian noise to + # be added to that variable (weight). + self.noise = [] + for var in self.model_variables: + name_ = var.name.split(":")[0] + "_noisy" if var.name else "" + self.noise.append( + get_variable( + np.zeros(var.shape, dtype=np.float32), + framework=self.framework, + tf_name=name_, + torch_tensor=True, + device=self.device, + ) + ) + + # tf-specific ops to sample, assign and remove noise. + if self.framework == "tf" and not tf.executing_eagerly(): + self.tf_sample_new_noise_op = self._tf_sample_new_noise_op() + self.tf_add_stored_noise_op = self._tf_add_stored_noise_op() + self.tf_remove_noise_op = self._tf_remove_noise_op() + # Create convenience sample+add op for tf. + with tf1.control_dependencies([self.tf_sample_new_noise_op]): + add_op = self._tf_add_stored_noise_op() + with tf1.control_dependencies([add_op]): + self.tf_sample_new_noise_and_add_op = tf.no_op() + + # Whether the Model's weights currently have noise added or not. + self.weights_are_currently_noisy = False + + # Auto-detection of underlying exploration functionality. + if sub_exploration is None: + # For discrete action spaces, use an underlying EpsilonGreedy with + # a special schedule. + if isinstance(self.action_space, Discrete): + sub_exploration = { + "type": "EpsilonGreedy", + "epsilon_schedule": { + "type": "PiecewiseSchedule", + # Step function (see [2]). + "endpoints": [ + (0, 1.0), + (random_timesteps + 1, 1.0), + (random_timesteps + 2, 0.01), + ], + "outside_value": 0.01, + }, + } + elif isinstance(self.action_space, Box): + sub_exploration = { + "type": "OrnsteinUhlenbeckNoise", + "random_timesteps": random_timesteps, + } + # TODO(sven): Implement for any action space. + else: + raise NotImplementedError + + self.sub_exploration = from_config( + Exploration, + sub_exploration, + framework=self.framework, + action_space=self.action_space, + policy_config=self.policy_config, + model=self.model, + **kwargs + ) + + # Whether we need to call `self._delayed_on_episode_start` before + # the forward pass. + self.episode_started = False + + @override(Exploration) + def before_compute_actions( + self, + *, + timestep: Optional[int] = None, + explore: Optional[bool] = None, + tf_sess: Optional["tf.Session"] = None + ): + explore = explore if explore is not None else self.policy_config["explore"] + + # Is this the first forward pass in the new episode? If yes, do the + # noise re-sampling and add to weights. + if self.episode_started: + self._delayed_on_episode_start(explore, tf_sess) + + # Add noise if necessary. + if explore and not self.weights_are_currently_noisy: + self._add_stored_noise(tf_sess=tf_sess) + # Remove noise if necessary. + elif not explore and self.weights_are_currently_noisy: + self._remove_noise(tf_sess=tf_sess) + + @override(Exploration) + def get_exploration_action( + self, + *, + action_distribution: ActionDistribution, + timestep: Union[TensorType, int], + explore: Union[TensorType, bool] + ): + # Use our sub-exploration object to handle the final exploration + # action (depends on the algo-type/action-space/etc..). + return self.sub_exploration.get_exploration_action( + action_distribution=action_distribution, timestep=timestep, explore=explore + ) + + @override(Exploration) + def on_episode_start( + self, + policy: "Policy", + *, + environment: BaseEnv = None, + episode: int = None, + tf_sess: Optional["tf.Session"] = None + ): + # We have to delay the noise-adding step by one forward call. + # This is due to the fact that the optimizer does it's step right + # after the episode was reset (and hence the noise was already added!). + # We don't want to update into a noisy net. + self.episode_started = True + + def _delayed_on_episode_start(self, explore, tf_sess): + # Sample fresh noise and add to weights. + if explore: + self._sample_new_noise_and_add(tf_sess=tf_sess, override=True) + # Only sample, don't apply anything to the weights. + else: + self._sample_new_noise(tf_sess=tf_sess) + self.episode_started = False + + @override(Exploration) + def on_episode_end(self, policy, *, environment=None, episode=None, tf_sess=None): + # Remove stored noise from weights (only if currently noisy). + if self.weights_are_currently_noisy: + self._remove_noise(tf_sess=tf_sess) + + @override(Exploration) + def postprocess_trajectory( + self, + policy: "Policy", + sample_batch: SampleBatch, + tf_sess: Optional["tf.Session"] = None, + ): + noisy_action_dist = noise_free_action_dist = None + # Adjust the stddev depending on the action (pi)-distance. + # Also see [1] for details. + # TODO(sven): Find out whether this can be scrapped by simply using + # the `sample_batch` to get the noisy/noise-free action dist. + _, _, fetches = policy.compute_actions_from_input_dict( + input_dict=sample_batch, explore=self.weights_are_currently_noisy + ) + + # Categorical case (e.g. DQN). + if issubclass(policy.dist_class, (Categorical, TorchCategorical)): + action_dist = softmax(fetches[SampleBatch.ACTION_DIST_INPUTS]) + # Deterministic (Gaussian actions, e.g. DDPG). + elif issubclass(policy.dist_class, (Deterministic, TorchDeterministic)): + action_dist = fetches[SampleBatch.ACTION_DIST_INPUTS] + else: + raise NotImplementedError # TODO(sven): Other action-dist cases. + + if self.weights_are_currently_noisy: + noisy_action_dist = action_dist + else: + noise_free_action_dist = action_dist + + _, _, fetches = policy.compute_actions_from_input_dict( + input_dict=sample_batch, explore=not self.weights_are_currently_noisy + ) + + # Categorical case (e.g. DQN). + if issubclass(policy.dist_class, (Categorical, TorchCategorical)): + action_dist = softmax(fetches[SampleBatch.ACTION_DIST_INPUTS]) + # Deterministic (Gaussian actions, e.g. DDPG). + elif issubclass(policy.dist_class, (Deterministic, TorchDeterministic)): + action_dist = fetches[SampleBatch.ACTION_DIST_INPUTS] + + if noisy_action_dist is None: + noisy_action_dist = action_dist + else: + noise_free_action_dist = action_dist + + delta = distance = None + # Categorical case (e.g. DQN). + if issubclass(policy.dist_class, (Categorical, TorchCategorical)): + # Calculate KL-divergence (DKL(clean||noisy)) according to [2]. + # TODO(sven): Allow KL-divergence to be calculated by our + # Distribution classes (don't support off-graph/numpy yet). + distance = np.nanmean( + np.sum( + noise_free_action_dist + * np.log( + noise_free_action_dist / (noisy_action_dist + SMALL_NUMBER) + ), + 1, + ) + ) + current_epsilon = self.sub_exploration.get_state(sess=tf_sess)[ + "cur_epsilon" + ] + delta = -np.log(1 - current_epsilon + current_epsilon / self.action_space.n) + elif issubclass(policy.dist_class, (Deterministic, TorchDeterministic)): + # Calculate MSE between noisy and non-noisy output (see [2]). + distance = np.sqrt( + np.mean(np.square(noise_free_action_dist - noisy_action_dist)) + ) + current_scale = self.sub_exploration.get_state(sess=tf_sess)["cur_scale"] + delta = getattr(self.sub_exploration, "ou_sigma", 0.2) * current_scale + + # Adjust stddev according to the calculated action-distance. + if distance <= delta: + self.stddev_val *= 1.01 + else: + self.stddev_val /= 1.01 + + # Update our state (self.stddev and self.stddev_val). + self.set_state(self.get_state(), sess=tf_sess) + + return sample_batch + + def _sample_new_noise(self, *, tf_sess=None): + """Samples new noise and stores it in `self.noise`.""" + if self.framework == "tf": + tf_sess.run(self.tf_sample_new_noise_op) + elif self.framework == "tf2": + self._tf_sample_new_noise_op() + else: + for i in range(len(self.noise)): + self.noise[i] = torch.normal( + mean=torch.zeros(self.noise[i].size()), std=self.stddev + ).to(self.device) + + def _tf_sample_new_noise_op(self): + added_noises = [] + for noise in self.noise: + added_noises.append( + tf1.assign( + noise, + tf.random.normal( + shape=noise.shape, stddev=self.stddev, dtype=tf.float32 + ), + ) + ) + return tf.group(*added_noises) + + def _sample_new_noise_and_add(self, *, tf_sess=None, override=False): + if self.framework == "tf": + if override and self.weights_are_currently_noisy: + tf_sess.run(self.tf_remove_noise_op) + tf_sess.run(self.tf_sample_new_noise_and_add_op) + else: + if override and self.weights_are_currently_noisy: + self._remove_noise() + self._sample_new_noise() + self._add_stored_noise() + + self.weights_are_currently_noisy = True + + def _add_stored_noise(self, *, tf_sess=None): + """Adds the stored `self.noise` to the model's parameters. + + Note: No new sampling of noise here. + + Args: + tf_sess (Optional[tf.Session]): The tf-session to use to add the + stored noise to the (currently noise-free) weights. + override: If True, undo any currently applied noise first, + then add the currently stored noise. + """ + # Make sure we only add noise to currently noise-free weights. + assert self.weights_are_currently_noisy is False + + # Add stored noise to the model's parameters. + if self.framework == "tf": + tf_sess.run(self.tf_add_stored_noise_op) + elif self.framework == "tf2": + self._tf_add_stored_noise_op() + else: + for var, noise in zip(self.model_variables, self.noise): + # Add noise to weights in-place. + var.requires_grad = False + var.add_(noise) + var.requires_grad = True + + self.weights_are_currently_noisy = True + + def _tf_add_stored_noise_op(self): + """Generates tf-op that assigns the stored noise to weights. + + Also used by tf-eager. + + Returns: + tf.op: The tf op to apply the already stored noise to the NN. + """ + add_noise_ops = list() + for var, noise in zip(self.model_variables, self.noise): + add_noise_ops.append(tf1.assign_add(var, noise)) + ret = tf.group(*tuple(add_noise_ops)) + with tf1.control_dependencies([ret]): + return tf.no_op() + + def _remove_noise(self, *, tf_sess=None): + """ + Removes the current action noise from the model parameters. + + Args: + tf_sess (Optional[tf.Session]): The tf-session to use to remove + the noise from the (currently noisy) weights. + """ + # Make sure we only remove noise iff currently noisy. + assert self.weights_are_currently_noisy is True + + # Removes the stored noise from the model's parameters. + if self.framework == "tf": + tf_sess.run(self.tf_remove_noise_op) + elif self.framework == "tf2": + self._tf_remove_noise_op() + else: + for var, noise in zip(self.model_variables, self.noise): + # Remove noise from weights in-place. + var.requires_grad = False + var.add_(-noise) + var.requires_grad = True + + self.weights_are_currently_noisy = False + + def _tf_remove_noise_op(self): + """Generates a tf-op for removing noise from the model's weights. + + Also used by tf-eager. + + Returns: + tf.op: The tf op to remve the currently stored noise from the NN. + """ + remove_noise_ops = list() + for var, noise in zip(self.model_variables, self.noise): + remove_noise_ops.append(tf1.assign_add(var, -noise)) + ret = tf.group(*tuple(remove_noise_ops)) + with tf1.control_dependencies([ret]): + return tf.no_op() + + @override(Exploration) + def get_state(self, sess=None): + return {"cur_stddev": self.stddev_val} + + @override(Exploration) + def set_state(self, state: dict, sess: Optional["tf.Session"] = None) -> None: + self.stddev_val = state["cur_stddev"] + # Set self.stddev to calculated value. + if self.framework == "tf": + self.stddev.load(self.stddev_val, session=sess) + elif isinstance(self.stddev, float): + self.stddev = self.stddev_val + else: + self.stddev.assign(self.stddev_val) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/exploration/per_worker_epsilon_greedy.py b/lib/python3.12/site-packages/ray/rllib/utils/exploration/per_worker_epsilon_greedy.py new file mode 100644 index 0000000000000000000000000000000000000000..3b5e12beeda2fecaaa359d018126ddac1b388ee7 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/exploration/per_worker_epsilon_greedy.py @@ -0,0 +1,59 @@ +from typing import Optional + +from gymnasium.spaces import Space + +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.exploration.epsilon_greedy import EpsilonGreedy +from ray.rllib.utils.schedules import ConstantSchedule + + +@OldAPIStack +class PerWorkerEpsilonGreedy(EpsilonGreedy): + """A per-worker epsilon-greedy class for distributed algorithms. + + Sets the epsilon schedules of individual workers to a constant: + 0.4 ^ (1 + [worker-index] / float([num-workers] - 1) * 7) + See Ape-X paper. + """ + + def __init__( + self, + action_space: Space, + *, + framework: str, + num_workers: Optional[int], + worker_index: Optional[int], + **kwargs + ): + """Create a PerWorkerEpsilonGreedy exploration class. + + Args: + action_space: The gym action space used by the environment. + num_workers: The overall number of workers used. + worker_index: The index of the Worker using this + Exploration. + framework: One of None, "tf", "torch". + """ + epsilon_schedule = None + # Use a fixed, different epsilon per worker. See: Ape-X paper. + assert worker_index <= num_workers, (worker_index, num_workers) + if num_workers > 0: + if worker_index > 0: + # From page 5 of https://arxiv.org/pdf/1803.00933.pdf + alpha, eps, i = 7, 0.4, worker_index - 1 + num_workers_minus_1 = float(num_workers - 1) if num_workers > 1 else 1.0 + constant_eps = eps ** (1 + (i / num_workers_minus_1) * alpha) + epsilon_schedule = ConstantSchedule(constant_eps, framework=framework) + # Local worker should have zero exploration so that eval + # rollouts run properly. + else: + epsilon_schedule = ConstantSchedule(0.0, framework=framework) + + super().__init__( + action_space, + epsilon_schedule=epsilon_schedule, + framework=framework, + num_workers=num_workers, + worker_index=worker_index, + **kwargs + ) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/exploration/per_worker_gaussian_noise.py b/lib/python3.12/site-packages/ray/rllib/utils/exploration/per_worker_gaussian_noise.py new file mode 100644 index 0000000000000000000000000000000000000000..4d5f81a9609d6f6b6e32c3eee0ae1e399ebdc686 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/exploration/per_worker_gaussian_noise.py @@ -0,0 +1,50 @@ +from typing import Optional + +from gymnasium.spaces import Space + +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.exploration.gaussian_noise import GaussianNoise +from ray.rllib.utils.schedules import ConstantSchedule + + +@OldAPIStack +class PerWorkerGaussianNoise(GaussianNoise): + """A per-worker Gaussian noise class for distributed algorithms. + + Sets the `scale` schedules of individual workers to a constant: + 0.4 ^ (1 + [worker-index] / float([num-workers] - 1) * 7) + See Ape-X paper. + """ + + def __init__( + self, + action_space: Space, + *, + framework: Optional[str], + num_workers: Optional[int], + worker_index: Optional[int], + **kwargs + ): + """ + Args: + action_space: The gym action space used by the environment. + num_workers: The overall number of workers used. + worker_index: The index of the Worker using this + Exploration. + framework: One of None, "tf", "torch". + """ + scale_schedule = None + # Use a fixed, different epsilon per worker. See: Ape-X paper. + if num_workers > 0: + if worker_index > 0: + num_workers_minus_1 = float(num_workers - 1) if num_workers > 1 else 1.0 + exponent = 1 + (worker_index / num_workers_minus_1) * 7 + scale_schedule = ConstantSchedule(0.4**exponent, framework=framework) + # Local worker should have zero exploration so that eval + # rollouts run properly. + else: + scale_schedule = ConstantSchedule(0.0, framework=framework) + + super().__init__( + action_space, scale_schedule=scale_schedule, framework=framework, **kwargs + ) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/exploration/per_worker_ornstein_uhlenbeck_noise.py b/lib/python3.12/site-packages/ray/rllib/utils/exploration/per_worker_ornstein_uhlenbeck_noise.py new file mode 100644 index 0000000000000000000000000000000000000000..d26ae283e1f5a44baff77a4f3a38f2edcd8e6bbf --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/exploration/per_worker_ornstein_uhlenbeck_noise.py @@ -0,0 +1,55 @@ +from typing import Optional + +from gymnasium.spaces import Space + +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.exploration.ornstein_uhlenbeck_noise import OrnsteinUhlenbeckNoise +from ray.rllib.utils.schedules import ConstantSchedule + + +@OldAPIStack +class PerWorkerOrnsteinUhlenbeckNoise(OrnsteinUhlenbeckNoise): + """A per-worker Ornstein Uhlenbeck noise class for distributed algorithms. + + Sets the Gaussian `scale` schedules of individual workers to a constant: + 0.4 ^ (1 + [worker-index] / float([num-workers] - 1) * 7) + See Ape-X paper. + """ + + def __init__( + self, + action_space: Space, + *, + framework: Optional[str], + num_workers: Optional[int], + worker_index: Optional[int], + **kwargs + ): + """ + Args: + action_space: The gym action space used by the environment. + num_workers: The overall number of workers used. + worker_index: The index of the Worker using this + Exploration. + framework: One of None, "tf", "torch". + """ + scale_schedule = None + # Use a fixed, different epsilon per worker. See: Ape-X paper. + if num_workers > 0: + if worker_index > 0: + num_workers_minus_1 = float(num_workers - 1) if num_workers > 1 else 1.0 + exponent = 1 + (worker_index / num_workers_minus_1) * 7 + scale_schedule = ConstantSchedule(0.4**exponent, framework=framework) + # Local worker should have zero exploration so that eval + # rollouts run properly. + else: + scale_schedule = ConstantSchedule(0.0, framework=framework) + + super().__init__( + action_space, + scale_schedule=scale_schedule, + num_workers=num_workers, + worker_index=worker_index, + framework=framework, + **kwargs + ) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/exploration/random.py b/lib/python3.12/site-packages/ray/rllib/utils/exploration/random.py new file mode 100644 index 0000000000000000000000000000000000000000..87b1d9bf3ab646832664dc9803aa92ef16cac9f9 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/exploration/random.py @@ -0,0 +1,169 @@ +from typing import Optional, Union + +import numpy as np +import tree # pip install dm_tree +from gymnasium.spaces import Box, Discrete, MultiDiscrete, Space + +from ray.rllib.models.action_dist import ActionDistribution +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.utils import force_tuple +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.exploration.exploration import Exploration +from ray.rllib.utils.framework import TensorType, try_import_tf, try_import_torch +from ray.rllib.utils.spaces.simplex import Simplex +from ray.rllib.utils.spaces.space_utils import get_base_struct_from_space +from ray.rllib.utils.tf_utils import zero_logps_from_actions + +tf1, tf, tfv = try_import_tf() +torch, _ = try_import_torch() + + +@OldAPIStack +class Random(Exploration): + """A random action selector (deterministic/greedy for explore=False). + + If explore=True, returns actions randomly from `self.action_space` (via + Space.sample()). + If explore=False, returns the greedy/max-likelihood action. + """ + + def __init__( + self, action_space: Space, *, model: ModelV2, framework: Optional[str], **kwargs + ): + """Initialize a Random Exploration object. + + Args: + action_space: The gym action space used by the environment. + framework: One of None, "tf", "torch". + """ + super().__init__( + action_space=action_space, model=model, framework=framework, **kwargs + ) + + self.action_space_struct = get_base_struct_from_space(self.action_space) + + @override(Exploration) + def get_exploration_action( + self, + *, + action_distribution: ActionDistribution, + timestep: Union[int, TensorType], + explore: bool = True + ): + # Instantiate the distribution object. + if self.framework in ["tf2", "tf"]: + return self.get_tf_exploration_action_op(action_distribution, explore) + else: + return self.get_torch_exploration_action(action_distribution, explore) + + def get_tf_exploration_action_op( + self, + action_dist: ActionDistribution, + explore: Optional[Union[bool, TensorType]], + ): + def true_fn(): + batch_size = 1 + req = force_tuple( + action_dist.required_model_output_shape( + self.action_space, getattr(self.model, "model_config", None) + ) + ) + # Add a batch dimension? + if len(action_dist.inputs.shape) == len(req) + 1: + batch_size = tf.shape(action_dist.inputs)[0] + + # Function to produce random samples from primitive space + # components: (Multi)Discrete or Box. + def random_component(component): + # Have at least an additional shape of (1,), even if the + # component is Box(-1.0, 1.0, shape=()). + shape = component.shape or (1,) + + if isinstance(component, Discrete): + return tf.random.uniform( + shape=(batch_size,) + component.shape, + maxval=component.n, + dtype=component.dtype, + ) + elif isinstance(component, MultiDiscrete): + return tf.concat( + [ + tf.random.uniform( + shape=(batch_size, 1), maxval=n, dtype=component.dtype + ) + for n in component.nvec + ], + axis=1, + ) + elif isinstance(component, Box): + if component.bounded_above.all() and component.bounded_below.all(): + if component.dtype.name.startswith("int"): + return tf.random.uniform( + shape=(batch_size,) + shape, + minval=component.low.flat[0], + maxval=component.high.flat[0], + dtype=component.dtype, + ) + else: + return tf.random.uniform( + shape=(batch_size,) + shape, + minval=component.low, + maxval=component.high, + dtype=component.dtype, + ) + else: + return tf.random.normal( + shape=(batch_size,) + shape, dtype=component.dtype + ) + else: + assert isinstance(component, Simplex), ( + "Unsupported distribution component '{}' for random " + "sampling!".format(component) + ) + return tf.nn.softmax( + tf.random.uniform( + shape=(batch_size,) + shape, + minval=0.0, + maxval=1.0, + dtype=component.dtype, + ) + ) + + actions = tree.map_structure(random_component, self.action_space_struct) + return actions + + def false_fn(): + return action_dist.deterministic_sample() + + action = tf.cond( + pred=tf.constant(explore, dtype=tf.bool) + if isinstance(explore, bool) + else explore, + true_fn=true_fn, + false_fn=false_fn, + ) + + logp = zero_logps_from_actions(action) + return action, logp + + def get_torch_exploration_action( + self, action_dist: ActionDistribution, explore: bool + ): + if explore: + req = force_tuple( + action_dist.required_model_output_shape( + self.action_space, getattr(self.model, "model_config", None) + ) + ) + # Add a batch dimension? + if len(action_dist.inputs.shape) == len(req) + 1: + batch_size = action_dist.inputs.shape[0] + a = np.stack([self.action_space.sample() for _ in range(batch_size)]) + else: + a = self.action_space.sample() + # Convert action to torch tensor. + action = torch.from_numpy(a).to(self.device) + else: + action = action_dist.deterministic_sample() + logp = torch.zeros((action.size()[0],), dtype=torch.float32, device=self.device) + return action, logp diff --git a/lib/python3.12/site-packages/ray/rllib/utils/exploration/random_encoder.py b/lib/python3.12/site-packages/ray/rllib/utils/exploration/random_encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..a4e8ee06fd8fb3c4f5bb96abfdb28bcf9bb68475 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/exploration/random_encoder.py @@ -0,0 +1,293 @@ +from typing import List, Optional, Union + +import numpy as np +from gymnasium.spaces import Box, Discrete, Space + +from ray.rllib.models.action_dist import ActionDistribution +from ray.rllib.models.catalog import ModelCatalog +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.exploration.exploration import Exploration +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.from_config import from_config +from ray.rllib.utils.tf_utils import get_placeholder +from ray.rllib.utils.typing import FromConfigSpec, ModelConfigDict, TensorType + +tf1, tf, tfv = try_import_tf() + + +class _MovingMeanStd: + """Track moving mean, std and count.""" + + def __init__(self, epsilon: float = 1e-4, shape: Optional[List[int]] = None): + """Initialize object. + + Args: + epsilon: Initial count. + shape: Shape of the trackables mean and std. + """ + if not shape: + shape = [] + self.mean = np.zeros(shape, dtype=np.float32) + self.var = np.ones(shape, dtype=np.float32) + self.count = epsilon + + def __call__(self, inputs: np.ndarray) -> np.ndarray: + """Normalize input batch using moving mean and std. + + Args: + inputs: Input batch to normalize. + + Returns: + Logarithmic scaled normalized output. + """ + batch_mean = np.mean(inputs, axis=0) + batch_var = np.var(inputs, axis=0) + batch_count = inputs.shape[0] + self.update_params(batch_mean, batch_var, batch_count) + return np.log(inputs / self.std + 1) + + def update_params( + self, batch_mean: float, batch_var: float, batch_count: float + ) -> None: + """Update moving mean, std and count. + + Args: + batch_mean: Input batch mean. + batch_var: Input batch variance. + batch_count: Number of cases in the batch. + """ + delta = batch_mean - self.mean + tot_count = self.count + batch_count + + # This moving mean calculation is from reference implementation. + self.mean = self.mean + delta + batch_count / tot_count + m_a = self.var * self.count + m_b = batch_var * batch_count + M2 = m_a + m_b + np.power(delta, 2) * self.count * batch_count / tot_count + self.var = M2 / tot_count + self.count = tot_count + + @property + def std(self) -> float: + """Get moving standard deviation. + + Returns: + Returns moving standard deviation. + """ + return np.sqrt(self.var) + + +@OldAPIStack +def update_beta(beta_schedule: str, beta: float, rho: float, step: int) -> float: + """Update beta based on schedule and training step. + + Args: + beta_schedule: Schedule for beta update. + beta: Initial beta. + rho: Schedule decay parameter. + step: Current training iteration. + + Returns: + Updated beta as per input schedule. + """ + if beta_schedule == "linear_decay": + return beta * ((1.0 - rho) ** step) + return beta + + +@OldAPIStack +def compute_states_entropy( + obs_embeds: np.ndarray, embed_dim: int, k_nn: int +) -> np.ndarray: + """Compute states entropy using K nearest neighbour method. + + Args: + obs_embeds: Observation latent representation using + encoder model. + embed_dim: Embedding vector dimension. + k_nn: Number of nearest neighbour for K-NN estimation. + + Returns: + Computed states entropy. + """ + obs_embeds_ = np.reshape(obs_embeds, [-1, embed_dim]) + dist = np.linalg.norm(obs_embeds_[:, None, :] - obs_embeds_[None, :, :], axis=-1) + return dist.argsort(axis=-1)[:, :k_nn][:, -1].astype(np.float32) + + +@OldAPIStack +class RE3(Exploration): + """Random Encoder for Efficient Exploration. + + Implementation of: + [1] State entropy maximization with random encoders for efficient + exploration. Seo, Chen, Shin, Lee, Abbeel, & Lee, (2021). + arXiv preprint arXiv:2102.09430. + + Estimates state entropy using a particle-based k-nearest neighbors (k-NN) + estimator in the latent space. The state's latent representation is + calculated using an encoder with randomly initialized parameters. + + The entropy of a state is considered as intrinsic reward and added to the + environment's extrinsic reward for policy optimization. + Entropy is calculated per batch, it does not take the distribution of + the entire replay buffer into consideration. + """ + + def __init__( + self, + action_space: Space, + *, + framework: str, + model: ModelV2, + embeds_dim: int = 128, + encoder_net_config: Optional[ModelConfigDict] = None, + beta: float = 0.2, + beta_schedule: str = "constant", + rho: float = 0.1, + k_nn: int = 50, + random_timesteps: int = 10000, + sub_exploration: Optional[FromConfigSpec] = None, + **kwargs + ): + """Initialize RE3. + + Args: + action_space: The action space in which to explore. + framework: Supports "tf", this implementation does not + support torch. + model: The policy's model. + embeds_dim: The dimensionality of the observation embedding + vectors in latent space. + encoder_net_config: Optional model + configuration for the encoder network, producing embedding + vectors from observations. This can be used to configure + fcnet- or conv_net setups to properly process any + observation space. + beta: Hyperparameter to choose between exploration and + exploitation. + beta_schedule: Schedule to use for beta decay, one of + "constant" or "linear_decay". + rho: Beta decay factor, used for on-policy algorithm. + k_nn: Number of neighbours to set for K-NN entropy + estimation. + random_timesteps: The number of timesteps to act completely + randomly (see [1]). + sub_exploration: The config dict for the underlying Exploration + to use (e.g. epsilon-greedy for DQN). If None, uses the + FromSpecDict provided in the Policy's default config. + + Raises: + ValueError: If the input framework is Torch. + """ + # TODO(gjoliver): Add supports for Pytorch. + if framework == "torch": + raise ValueError("This RE3 implementation does not support Torch.") + super().__init__(action_space, model=model, framework=framework, **kwargs) + + self.beta = beta + self.rho = rho + self.k_nn = k_nn + self.embeds_dim = embeds_dim + if encoder_net_config is None: + encoder_net_config = self.policy_config["model"].copy() + self.encoder_net_config = encoder_net_config + + # Auto-detection of underlying exploration functionality. + if sub_exploration is None: + # For discrete action spaces, use an underlying EpsilonGreedy with + # a special schedule. + if isinstance(self.action_space, Discrete): + sub_exploration = { + "type": "EpsilonGreedy", + "epsilon_schedule": { + "type": "PiecewiseSchedule", + # Step function (see [2]). + "endpoints": [ + (0, 1.0), + (random_timesteps + 1, 1.0), + (random_timesteps + 2, 0.01), + ], + "outside_value": 0.01, + }, + } + elif isinstance(self.action_space, Box): + sub_exploration = { + "type": "OrnsteinUhlenbeckNoise", + "random_timesteps": random_timesteps, + } + else: + raise NotImplementedError + + self.sub_exploration = sub_exploration + + # Creates ModelV2 embedding module / layers. + self._encoder_net = ModelCatalog.get_model_v2( + self.model.obs_space, + self.action_space, + self.embeds_dim, + model_config=self.encoder_net_config, + framework=self.framework, + name="encoder_net", + ) + if self.framework == "tf": + self._obs_ph = get_placeholder( + space=self.model.obs_space, name="_encoder_obs" + ) + self._obs_embeds = tf.stop_gradient( + self._encoder_net({SampleBatch.OBS: self._obs_ph})[0] + ) + + # This is only used to select the correct action + self.exploration_submodule = from_config( + cls=Exploration, + config=self.sub_exploration, + action_space=self.action_space, + framework=self.framework, + policy_config=self.policy_config, + model=self.model, + num_workers=self.num_workers, + worker_index=self.worker_index, + ) + + @override(Exploration) + def get_exploration_action( + self, + *, + action_distribution: ActionDistribution, + timestep: Union[int, TensorType], + explore: bool = True + ): + # Simply delegate to sub-Exploration module. + return self.exploration_submodule.get_exploration_action( + action_distribution=action_distribution, timestep=timestep, explore=explore + ) + + @override(Exploration) + def postprocess_trajectory(self, policy, sample_batch, tf_sess=None): + """Calculate states' latent representations/embeddings. + + Embeddings are added to the SampleBatch object such that it doesn't + need to be calculated during each training step. + """ + if self.framework != "torch": + sample_batch = self._postprocess_tf(policy, sample_batch, tf_sess) + else: + raise ValueError("Not implemented for Torch.") + return sample_batch + + def _postprocess_tf(self, policy, sample_batch, tf_sess): + """Calculate states' embeddings and add it to SampleBatch.""" + if self.framework == "tf": + obs_embeds = tf_sess.run( + self._obs_embeds, + feed_dict={self._obs_ph: sample_batch[SampleBatch.OBS]}, + ) + else: + obs_embeds = tf.stop_gradient( + self._encoder_net({SampleBatch.OBS: sample_batch[SampleBatch.OBS]})[0] + ).numpy() + sample_batch[SampleBatch.OBS_EMBEDS] = obs_embeds + return sample_batch diff --git a/lib/python3.12/site-packages/ray/rllib/utils/exploration/slate_epsilon_greedy.py b/lib/python3.12/site-packages/ray/rllib/utils/exploration/slate_epsilon_greedy.py new file mode 100644 index 0000000000000000000000000000000000000000..30f72dc853f7098e28c7391df75d002e103814c3 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/exploration/slate_epsilon_greedy.py @@ -0,0 +1,114 @@ +from typing import Union + +from ray.rllib.models.action_dist import ActionDistribution +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.exploration.epsilon_greedy import EpsilonGreedy +from ray.rllib.utils.exploration.exploration import TensorType +from ray.rllib.utils.framework import try_import_tf, try_import_torch + +tf1, tf, tfv = try_import_tf() +torch, _ = try_import_torch() + + +@OldAPIStack +class SlateEpsilonGreedy(EpsilonGreedy): + @override(EpsilonGreedy) + def _get_tf_exploration_action_op( + self, + action_distribution: ActionDistribution, + explore: Union[bool, TensorType], + timestep: Union[int, TensorType], + ) -> "tf.Tensor": + + per_slate_q_values = action_distribution.inputs + all_slates = action_distribution.all_slates + + exploit_action = action_distribution.deterministic_sample() + + batch_size, num_slates = ( + tf.shape(per_slate_q_values)[0], + tf.shape(per_slate_q_values)[1], + ) + action_logp = tf.zeros(batch_size, dtype=tf.float32) + + # Get the current epsilon. + epsilon = self.epsilon_schedule( + timestep if timestep is not None else self.last_timestep + ) + # A random action. + random_indices = tf.random.uniform( + (batch_size,), + minval=0, + maxval=num_slates, + dtype=tf.dtypes.int32, + ) + random_actions = tf.gather(all_slates, random_indices) + + choose_random = ( + tf.random.uniform( + tf.stack([batch_size]), minval=0, maxval=1, dtype=tf.float32 + ) + < epsilon + ) + + # Pick either random or greedy. + action = tf.cond( + pred=tf.constant(explore, dtype=tf.bool) + if isinstance(explore, bool) + else explore, + true_fn=(lambda: tf.where(choose_random, random_actions, exploit_action)), + false_fn=lambda: exploit_action, + ) + + if self.framework == "tf2" and not self.policy_config["eager_tracing"]: + self.last_timestep = timestep + return action, action_logp + else: + assign_op = tf1.assign(self.last_timestep, tf.cast(timestep, tf.int64)) + with tf1.control_dependencies([assign_op]): + return action, action_logp + + @override(EpsilonGreedy) + def _get_torch_exploration_action( + self, + action_distribution: ActionDistribution, + explore: bool, + timestep: Union[int, TensorType], + ) -> "torch.Tensor": + + per_slate_q_values = action_distribution.inputs + all_slates = self.model.slates + device = all_slates.device + + exploit_indices = action_distribution.deterministic_sample() + exploit_indices = exploit_indices.to(device) + exploit_action = all_slates[exploit_indices] + + batch_size = per_slate_q_values.size()[0] + action_logp = torch.zeros(batch_size, dtype=torch.float) + + self.last_timestep = timestep + + # Explore. + if explore: + # Get the current epsilon. + epsilon = self.epsilon_schedule(self.last_timestep) + # A random action. + random_indices = torch.randint( + 0, + per_slate_q_values.shape[1], + (per_slate_q_values.shape[0],), + device=device, + ) + random_actions = all_slates[random_indices] + + # Pick either random or greedy. + action = torch.where( + torch.empty((batch_size,)).uniform_() < epsilon, + random_actions, + exploit_action, + ) + return action, action_logp + # Return the deterministic "sample" (argmax) over the logits. + else: + return exploit_action, action_logp diff --git a/lib/python3.12/site-packages/ray/rllib/utils/exploration/slate_soft_q.py b/lib/python3.12/site-packages/ray/rllib/utils/exploration/slate_soft_q.py new file mode 100644 index 0000000000000000000000000000000000000000..5ed2205482e51890db8d0b069e53643b790e60b5 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/exploration/slate_soft_q.py @@ -0,0 +1,46 @@ +from typing import Union + +from ray.rllib.models.action_dist import ActionDistribution +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.exploration.exploration import TensorType +from ray.rllib.utils.exploration.soft_q import SoftQ +from ray.rllib.utils.framework import try_import_tf, try_import_torch + +tf1, tf, tfv = try_import_tf() +torch, _ = try_import_torch() + + +@OldAPIStack +class SlateSoftQ(SoftQ): + @override(SoftQ) + def get_exploration_action( + self, + action_distribution: ActionDistribution, + timestep: Union[int, TensorType], + explore: bool = True, + ): + assert ( + self.framework == "torch" + ), "ERROR: SlateSoftQ only supports torch so far!" + + cls = type(action_distribution) + + # Re-create the action distribution with the correct temperature + # applied. + action_distribution = cls( + action_distribution.inputs, self.model, temperature=self.temperature + ) + batch_size = action_distribution.inputs.size()[0] + action_logp = torch.zeros(batch_size, dtype=torch.float) + + self.last_timestep = timestep + + # Explore. + if explore: + # Return stochastic sample over (q-value) logits. + action = action_distribution.sample() + # Return the deterministic "sample" (argmax) over (q-value) logits. + else: + action = action_distribution.deterministic_sample() + + return action, action_logp diff --git a/lib/python3.12/site-packages/ray/rllib/utils/exploration/soft_q.py b/lib/python3.12/site-packages/ray/rllib/utils/exploration/soft_q.py new file mode 100644 index 0000000000000000000000000000000000000000..be6dae9adb9a46628604de1c1225c67c1136e72a --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/exploration/soft_q.py @@ -0,0 +1,56 @@ +from typing import Optional, Union + +from gymnasium.spaces import Discrete, MultiDiscrete, Space + +from ray.rllib.models.action_dist import ActionDistribution +from ray.rllib.models.tf.tf_action_dist import Categorical +from ray.rllib.models.torch.torch_action_dist import TorchCategorical +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.exploration.stochastic_sampling import StochasticSampling +from ray.rllib.utils.framework import TensorType + + +@OldAPIStack +class SoftQ(StochasticSampling): + """Special case of StochasticSampling w/ Categorical and temperature param. + + Returns a stochastic sample from a Categorical parameterized by the model + output divided by the temperature. Returns the argmax iff explore=False. + """ + + def __init__( + self, + action_space: Space, + *, + framework: Optional[str], + temperature: float = 1.0, + **kwargs + ): + """Initializes a SoftQ Exploration object. + + Args: + action_space: The gym action space used by the environment. + temperature: The temperature to divide model outputs by + before creating the Categorical distribution to sample from. + framework: One of None, "tf", "torch". + """ + assert isinstance(action_space, (Discrete, MultiDiscrete)) + super().__init__(action_space, framework=framework, **kwargs) + self.temperature = temperature + + @override(StochasticSampling) + def get_exploration_action( + self, + action_distribution: ActionDistribution, + timestep: Union[int, TensorType], + explore: bool = True, + ): + cls = type(action_distribution) + assert issubclass(cls, (Categorical, TorchCategorical)) + # Re-create the action distribution with the correct temperature + # applied. + dist = cls(action_distribution.inputs, self.model, temperature=self.temperature) + # Delegate to super method. + return super().get_exploration_action( + action_distribution=dist, timestep=timestep, explore=explore + ) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/exploration/stochastic_sampling.py b/lib/python3.12/site-packages/ray/rllib/utils/exploration/stochastic_sampling.py new file mode 100644 index 0000000000000000000000000000000000000000..175ee5218ca14c52b1d1ca3d0e25f344f7a544a9 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/exploration/stochastic_sampling.py @@ -0,0 +1,157 @@ +import functools +from typing import Optional, Union + +import gymnasium as gym +import numpy as np + +from ray.rllib.models.action_dist import ActionDistribution +from ray.rllib.models.modelv2 import ModelV2 +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.exploration.exploration import Exploration +from ray.rllib.utils.exploration.random import Random +from ray.rllib.utils.framework import ( + TensorType, + get_variable, + try_import_tf, + try_import_torch, +) +from ray.rllib.utils.tf_utils import zero_logps_from_actions + +tf1, tf, tfv = try_import_tf() +torch, _ = try_import_torch() + + +@OldAPIStack +class StochasticSampling(Exploration): + """An exploration that simply samples from a distribution. + + The sampling can be made deterministic by passing explore=False into + the call to `get_exploration_action`. + Also allows for scheduled parameters for the distributions, such as + lowering stddev, temperature, etc.. over time. + """ + + def __init__( + self, + action_space: gym.spaces.Space, + *, + framework: str, + model: ModelV2, + random_timesteps: int = 0, + **kwargs + ): + """Initializes a StochasticSampling Exploration object. + + Args: + action_space: The gym action space used by the environment. + framework: One of None, "tf", "torch". + model: The ModelV2 used by the owning Policy. + random_timesteps: The number of timesteps for which to act + completely randomly. Only after this number of timesteps, + actual samples will be drawn to get exploration actions. + """ + assert framework is not None + super().__init__(action_space, model=model, framework=framework, **kwargs) + + # Create the Random exploration module (used for the first n + # timesteps). + self.random_timesteps = random_timesteps + self.random_exploration = Random( + action_space, model=self.model, framework=self.framework, **kwargs + ) + + # The current timestep value (tf-var or python int). + self.last_timestep = get_variable( + np.array(0, np.int64), + framework=self.framework, + tf_name="timestep", + dtype=np.int64, + ) + + @override(Exploration) + def get_exploration_action( + self, + *, + action_distribution: ActionDistribution, + timestep: Optional[Union[int, TensorType]] = None, + explore: bool = True + ): + if self.framework == "torch": + return self._get_torch_exploration_action( + action_distribution, timestep, explore + ) + else: + return self._get_tf_exploration_action_op( + action_distribution, timestep, explore + ) + + def _get_tf_exploration_action_op(self, action_dist, timestep, explore): + ts = self.last_timestep + 1 + + stochastic_actions = tf.cond( + pred=tf.convert_to_tensor(ts < self.random_timesteps), + true_fn=lambda: ( + self.random_exploration.get_tf_exploration_action_op( + action_dist, explore=True + )[0] + ), + false_fn=lambda: action_dist.sample(), + ) + deterministic_actions = action_dist.deterministic_sample() + + action = tf.cond( + tf.constant(explore) if isinstance(explore, bool) else explore, + true_fn=lambda: stochastic_actions, + false_fn=lambda: deterministic_actions, + ) + + logp = tf.cond( + tf.math.logical_and( + explore, tf.convert_to_tensor(ts >= self.random_timesteps) + ), + true_fn=lambda: action_dist.sampled_action_logp(), + false_fn=functools.partial(zero_logps_from_actions, deterministic_actions), + ) + + # Increment `last_timestep` by 1 (or set to `timestep`). + if self.framework == "tf2": + self.last_timestep.assign_add(1) + return action, logp + else: + assign_op = ( + tf1.assign_add(self.last_timestep, 1) + if timestep is None + else tf1.assign(self.last_timestep, timestep) + ) + with tf1.control_dependencies([assign_op]): + return action, logp + + def _get_torch_exploration_action( + self, + action_dist: ActionDistribution, + timestep: Union[TensorType, int], + explore: Union[TensorType, bool], + ): + # Set last timestep or (if not given) increase by one. + self.last_timestep = ( + timestep if timestep is not None else self.last_timestep + 1 + ) + + # Apply exploration. + if explore: + # Random exploration phase. + if self.last_timestep < self.random_timesteps: + action, logp = self.random_exploration.get_torch_exploration_action( + action_dist, explore=True + ) + # Take a sample from our distribution. + else: + action = action_dist.sample() + logp = action_dist.sampled_action_logp() + + # No exploration -> Return deterministic actions. + else: + action = action_dist.deterministic_sample() + logp = torch.zeros_like(action_dist.sampled_action_logp()) + + return action, logp diff --git a/lib/python3.12/site-packages/ray/rllib/utils/exploration/thompson_sampling.py b/lib/python3.12/site-packages/ray/rllib/utils/exploration/thompson_sampling.py new file mode 100644 index 0000000000000000000000000000000000000000..3d4700790edb0adb100c28b259f210c231181d80 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/exploration/thompson_sampling.py @@ -0,0 +1,46 @@ +from typing import Union + +from ray.rllib.models.action_dist import ActionDistribution +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.exploration.exploration import Exploration +from ray.rllib.utils.framework import ( + TensorType, + try_import_tf, +) + +tf1, tf, tfv = try_import_tf() + + +@OldAPIStack +class ThompsonSampling(Exploration): + @override(Exploration) + def get_exploration_action( + self, + action_distribution: ActionDistribution, + timestep: Union[int, TensorType], + explore: bool = True, + ): + if self.framework == "torch": + return self._get_torch_exploration_action(action_distribution, explore) + elif self.framework == "tf2": + return self._get_tf_exploration_action(action_distribution, explore) + else: + raise NotImplementedError + + def _get_torch_exploration_action(self, action_dist, explore): + if explore: + return action_dist.inputs.argmax(dim=-1), None + else: + scores = self.model.predict(self.model.current_obs()) + return scores.argmax(dim=-1), None + + def _get_tf_exploration_action(self, action_dist, explore): + action = tf.argmax( + tf.cond( + pred=explore, + true_fn=lambda: action_dist.inputs, + false_fn=lambda: self.model.predict(self.model.current_obs()), + ), + axis=-1, + ) + return action, None diff --git a/lib/python3.12/site-packages/ray/rllib/utils/exploration/upper_confidence_bound.py b/lib/python3.12/site-packages/ray/rllib/utils/exploration/upper_confidence_bound.py new file mode 100644 index 0000000000000000000000000000000000000000..7e7e71efe187db0b1c42cd875acf03b70d949b6c --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/exploration/upper_confidence_bound.py @@ -0,0 +1,44 @@ +from typing import Union + +from ray.rllib.models.action_dist import ActionDistribution +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.exploration.exploration import Exploration +from ray.rllib.utils.framework import ( + TensorType, + try_import_tf, +) + +tf1, tf, tfv = try_import_tf() + + +@OldAPIStack +class UpperConfidenceBound(Exploration): + @override(Exploration) + def get_exploration_action( + self, + action_distribution: ActionDistribution, + timestep: Union[int, TensorType], + explore: bool = True, + ): + if self.framework == "torch": + return self._get_torch_exploration_action(action_distribution, explore) + elif self.framework == "tf2": + return self._get_tf_exploration_action(action_distribution, explore) + else: + raise NotImplementedError + + def _get_torch_exploration_action(self, action_dist, explore): + if explore: + return action_dist.inputs.argmax(dim=-1), None + else: + scores = self.model.value_function() + return scores.argmax(dim=-1), None + + def _get_tf_exploration_action(self, action_dist, explore): + action = tf.argmax( + tf.cond( + explore, lambda: action_dist.inputs, lambda: self.model.value_function() + ), + axis=-1, + ) + return action, None diff --git a/lib/python3.12/site-packages/ray/rllib/utils/filter.py b/lib/python3.12/site-packages/ray/rllib/utils/filter.py new file mode 100644 index 0000000000000000000000000000000000000000..ec22f25670ee0313c82e5eb49a8dfbce871b2145 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/filter.py @@ -0,0 +1,506 @@ +import logging +import threading + +import numpy as np +import tree # pip install dm_tree + +from ray._common.deprecation import Deprecated, deprecation_warning +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.numpy import ( + SMALL_NUMBER, +) # Assuming SMALL_NUMBER is a small float like 1e-8 +from ray.rllib.utils.serialization import _deserialize_ndarray, _serialize_ndarray +from ray.rllib.utils.typing import TensorStructType + +logger = logging.getLogger(__name__) + + +@OldAPIStack +class Filter: + """Processes input, possibly statefully.""" + + def apply_changes(self, other: "Filter", *args, **kwargs) -> None: + """Updates self with "new state" from other filter.""" + raise NotImplementedError + + def copy(self) -> "Filter": + """Creates a new object with same state as self. + + Returns: + A copy of self. + """ + raise NotImplementedError + + def sync(self, other: "Filter") -> None: + """Copies all state from other filter to self.""" + raise NotImplementedError + + def reset_buffer(self) -> None: + """Creates copy of current state and resets accumulated state""" + raise NotImplementedError + + def as_serializable(self) -> "Filter": + raise NotImplementedError + + @Deprecated(new="Filter.reset_buffer()", error=True) + def clear_buffer(self): + pass + + +@OldAPIStack +class NoFilter(Filter): + is_concurrent = True + + def __call__(self, x: TensorStructType, update=True): + # Process no further if already np.ndarray, dict, or tuple. + if isinstance(x, (np.ndarray, dict, tuple)): + return x + + try: + return np.asarray(x) + except Exception: + raise ValueError("Failed to convert to array", x) + + def apply_changes(self, other: "NoFilter", *args, **kwargs) -> None: + pass + + def copy(self) -> "NoFilter": + return self + + def sync(self, other: "NoFilter") -> None: + pass + + def reset_buffer(self) -> None: + pass + + def as_serializable(self) -> "NoFilter": + return self + + +# Based on Welford's algorithm for numerical stability +# http://www.johndcook.com/blog/standard_deviation/ [4] +@OldAPIStack +class RunningStat: + def __init__(self, shape=()): + """Initializes a `RunningStat` instance.""" + # Keep always a state and a delta from all attributes. Note, + # we use the state for filtering and the delta for updates. + # All deltas will be zero(s) after a state synchronization + # across different actors. + self.num_pushes = 0 + self.num_pushes_delta = 0 + # Stores the mean. + self.mean_array = np.zeros(shape) + self.mean_delta_array = np.zeros(shape) + # Stores the sum of squared demeaned observations. Note, this + # follows Wellington's algorithm. + self.sum_sq_diff_array = np.zeros(shape) + self.sum_sq_diff_delta_array = np.zeros(shape) + + def copy(self): + """Copies a `RunningStat`.""" + # Copy all attributes by creating a new `RunningStat` instance. + other = RunningStat(self.shape) + other.num_pushes = self.num_pushes + other.num_pushes_delta = self.num_pushes_delta + other.mean_array = np.copy(self.mean_array) + other.mean_delta_array = np.copy(self.mean_delta_array) + other.sum_sq_diff_array = np.copy(self.sum_sq_diff_array) + other.sum_sq_diff_delta_array = np.copy(self.sum_sq_diff_delta_array) + return other + + def push(self, x): + """Updates a `RunningStat` instance by a new value. + + Args: + x: A new value to update mean and sum of squares by. Must have the + same shape like the mean. + + Raises: + `ValueError` in case of a shape mismatch. + """ + x = np.asarray(x) + if x.shape != self.mean_array.shape: + raise ValueError( + "Unexpected input shape {}, expected {}, value = {}".format( + x.shape, self.mean_array.shape, x + ) + ) + + # Store old mean for Welford's sum of squares update. + old_mean = np.copy(self.mean_array) + self.num_pushes += 1 + # Also increase the delta counter since the last merge. + self.num_pushes_delta += 1 + + if self.num_pushes == 1: + self.mean_array[...] = x + self.mean_delta_array[...] = x + # sum_sq_diff_array remains 0 for the first element + else: + # Welford's update for mean + delta = x - old_mean + self.mean_array[...] += delta / self.num_pushes + # Update the mean delta. + self.mean_delta_array[...] += delta / self.num_pushes + + # Welford's update for sum of squared differences (S) + # S_k = S_{k-1} + (x_k - M_k)(x_k - M_{k-1}). + self.sum_sq_diff_array[...] += delta * (x - self.mean_array) + # Update the mean sum of squares. + self.sum_sq_diff_delta_array[...] += delta * (x - self.mean_array) + + def update(self, other): + """Update this `RunningStat` instance by another one. + + Args: + other: Another `RunningStat` instance whose state should me + merged with `self`. + """ + # Make this explicitly for future changes to avoid ever turning `num_pushes` into + # a float (this was a problem in earlier versions). + n1_int = self.num_pushes + # Note, we use only the delta for the updates, this reduces the risk of numerical + # instabilities significantly. + n2_int = other.num_pushes_delta + # For higher precision use float versions of the counters. + n1_flt = float(self.num_pushes) + n2_flt = float(other.num_pushes_delta) + n_flt = n1_flt + n2_flt + + # If none of the two `RunningStat`s has seen values, yet, return. + if n1_int + n2_int == 0: + # Avoid divide by zero, which creates nans + return + + # Numerically stable formula for combining means + # M_combined = (n1*M1 + n2*M2) / (n1+n2) + # This is equivalent to M1 + delta * n2 / n + delta_mean = other.mean_delta_array - self.mean_array + self.mean_array += delta_mean * n2_flt / n_flt + + # Numerically stable formula for combining sums of squared differences (S) + # S_combined = S1 + S2 + (n1*n2 / (n1+n2)) * (M1 - M2)^2 [6] + delta_mean_sq = delta_mean * delta_mean + self.sum_sq_diff_array += other.sum_sq_diff_delta_array + delta_mean_sq * ( + n1_flt * n2_flt / n_flt + ) + + # Update the counter with the interger versions of the two counters. + self.num_pushes = n1_int + n2_int + + def __repr__(self): + """Represents a `RunningStat` instance. + + Note, a `RunningStat` is represented by its mean, its standard deviation + and the number `n` of values used to compute the two statistics. + """ + return "(n={}, mean_mean={}, mean_std={})".format( + self.n, np.mean(self.mean), np.mean(self.std) + ) + + @property + def n(self): + """Returns the number of values seen by a `RunningStat` instance.""" + return self.num_pushes + + @property + def mean(self): + """Returns the (vector) mean estimate of a `RunningStat` instance.""" + return self.mean_array + + @property + def var(self): + """Returns the (unbiased vector) variance estimate of a `RunningStat` instance.""" + # For n=0 or n=1, variance is typically undefined or 0. + # Returning 0 for n <= 1 is a common convention for running variance. + if self.num_pushes <= 1: + return np.zeros_like(self.mean_array).astype(np.float32) + # Variance = S / (n-1) for sample variance + return (self.sum_sq_diff_array / (float(self.num_pushes) - 1)).astype( + np.float32 + ) + + @property + def std(self): + """Returns the (unbiased vector) std estimate of a `RunningStat` instance.ance.""" + # Ensure variance is non-negative before sqrt + return np.sqrt(np.maximum(0, self.var)) + + @property + def shape(self): + """Returns the shape of the `RunningStat` instance.""" + return self.mean_array.shape + + def to_state(self): + """Returns the pickable state of a `RunningStat` instance.""" + return { + "num_pushes": self.num_pushes, + "num_pushes_delta": self.num_pushes_delta, + "mean_array": _serialize_ndarray(self.mean_array), + "mean_delta_array": _serialize_ndarray(self.mean_delta_array), + "sum_sq_diff_array": _serialize_ndarray(self.sum_sq_diff_array), + "sum_sq_diff_delta_array": _serialize_ndarray(self.sum_sq_diff_delta_array), + } + + @staticmethod + def from_state(state): + """Builds a `RunningStat` instance from a pickable state.""" + # Need to pass shape to constructor for proper initialization + # Assuming shape can be inferred from mean_array in state + shape = _deserialize_ndarray(state["mean_array"]).shape + running_stats = RunningStat(shape) + running_stats.num_pushes = state["num_pushes"] + running_stats.num_pushes_delta = state["num_pushes_delta"] + running_stats.mean_array = _deserialize_ndarray(state["mean_array"]) + running_stats.mean_delta_array = _deserialize_ndarray(state["mean_delta_array"]) + running_stats.sum_sq_diff_array = _deserialize_ndarray( + state["sum_sq_diff_array"] + ) + running_stats.sum_sq_diff_delta_array = _deserialize_ndarray( + state["sum_sq_diff_delta_array"] + ) + return running_stats + + +@OldAPIStack +class MeanStdFilter(Filter): + """Keeps track of a running mean for seen states""" + + is_concurrent = False + + def __init__(self, shape, demean=True, destd=True, clip=10.0): + self.shape = shape + # We don't have a preprocessor, if shape is None (Discrete) or + # flat_shape is Tuple[np.ndarray] or Dict[str, np.ndarray] + # (complex inputs). + flat_shape = tree.flatten(self.shape) + self.no_preprocessor = shape is None or ( + isinstance(self.shape, (dict, tuple)) + and len(flat_shape) > 0 + and isinstance(flat_shape, np.ndarray) + ) + # If preprocessing (flattening dicts/tuples), make sure shape + # is an np.ndarray, so we don't confuse it with a complex Tuple + # space's shape structure (which is a Tuple[np.ndarray, ...]). + if not self.no_preprocessor: + self.shape = np.array(self.shape) + self.demean = demean + self.destd = destd + self.clip = clip + # Running stats. + self.running_stats = tree.map_structure(lambda s: RunningStat(s), self.shape) + + # In distributed rollouts, each worker sees different states. + # The buffer is used to keep track of deltas amongst all the + # observation filters. + self.buffer = None + self.reset_buffer() + + def reset_buffer(self) -> None: + self.buffer = tree.map_structure(lambda s: RunningStat(s), self.shape) + + def apply_changes( + self, other: "MeanStdFilter", with_buffer: bool = False, *args, **kwargs + ) -> None: + """Applies updates from the buffer of another filter. + + Args: + other: Other filter to apply info from + with_buffer: Flag for specifying if the buffer should be + copied from other. + + .. testcode:: + :skipif: True + + a = MeanStdFilter(()) + a(1) + a(2) + print([a.running_stats.n, a.running_stats.mean, a.buffer.n]) + + .. testoutput:: + + [2, 1.5, 2] + + .. testcode:: + :skipif: True + + b = MeanStdFilter(()) + b(10) + a.apply_changes(b, with_buffer=False) + print([a.running_stats.n, a.running_stats.mean, a.buffer.n]) + + .. testoutput:: + + [3, 4.333333333333333, 2] + + .. testcode:: + :skipif: True + + a.apply_changes(b, with_buffer=True) + print([a.running_stats.n, a.running_stats.mean, a.buffer.n]) + + .. testoutput:: + + [4, 5.75, 1] + """ + tree.map_structure( + lambda rs, other_rs: rs.update(other_rs), self.running_stats, other.buffer + ) + if with_buffer: + self.buffer = tree.map_structure(lambda b: b.copy(), other.buffer) + + def copy(self) -> "MeanStdFilter": + """Returns a copy of `self`.""" + other = MeanStdFilter(self.shape) + other.sync(self) + return other + + def as_serializable(self) -> "MeanStdFilter": + return self.copy() + + def sync(self, other: "MeanStdFilter") -> None: + """Syncs all fields together from other filter. + + .. testcode:: + :skipif: True + + a = MeanStdFilter(()) + a(1) + a(2) + print([a.running_stats.n, a.running_stats.mean, a.buffer.n]) + + .. testoutput:: + + [2, array(1.5), 2] + + .. testcode:: + :skipif: True + + b = MeanStdFilter(()) + b(10) + print([b.running_stats.n, b.running_stats.mean, b.buffer.n]) + + .. testoutput:: + + [1, array(10.0), 1] + + .. testcode:: + :skipif: True + + a.sync(b) + print([a.running_stats.n, a.running_stats.mean, a.buffer.n]) + + .. testoutput:: + + [1, array(10.0), 1] + """ + self.demean = other.demean + self.destd = other.destd + self.clip = other.clip + self.running_stats = tree.map_structure( + lambda rs: rs.copy(), other.running_stats + ) + self.buffer = tree.map_structure(lambda b: b.copy(), other.buffer) + + def __call__(self, x: TensorStructType, update: bool = True) -> TensorStructType: + if self.no_preprocessor: + x = tree.map_structure(lambda x_: np.asarray(x_), x) + else: + x = np.asarray(x) + + def _helper(x, rs, buffer, shape): + # Discrete|MultiDiscrete spaces -> No normalization. + if shape is None: + return x + + # Keep dtype as is througout this filter. + orig_dtype = x.dtype + + if update: + if len(x.shape) == len(rs.shape) + 1: + # The vectorized case. + for i in range(x.shape): + rs.push(x[i]) + buffer.push(x[i]) + else: + # The unvectorized case. + rs.push(x) + buffer.push(x) + if self.demean: + x = x - rs.mean + if self.destd: + x = x / (rs.std + SMALL_NUMBER) + if self.clip: + x = np.clip(x, -self.clip, self.clip) + return x.astype(orig_dtype) + + if self.no_preprocessor: + return tree.map_structure_up_to( + x, _helper, x, self.running_stats, self.buffer, self.shape + ) + else: + return _helper(x, self.running_stats, self.buffer, self.shape) + + +@OldAPIStack +class ConcurrentMeanStdFilter(MeanStdFilter): + is_concurrent = True + + def __init__(self, *args, **kwargs): + super(ConcurrentMeanStdFilter, self).__init__(*args, **kwargs) + deprecation_warning( + old="ConcurrentMeanStdFilter", + error=False, + help="ConcurrentMeanStd filters are only used for testing and will " + "therefore be deprecated in the course of moving to the " + "Connetors API, where testing of filters will be done by other " + "means.", + ) + + self._lock = threading.RLock() + + def lock_wrap(func): + def wrapper(*args, **kwargs): + with self._lock: + return func(*args, **kwargs) + + return wrapper + + self.__getattribute__ = lock_wrap(self.__getattribute__) + + def as_serializable(self) -> "MeanStdFilter": + """Returns non-concurrent version of current class""" + other = MeanStdFilter(self.shape) + other.sync(self) + return other + + def copy(self) -> "ConcurrentMeanStdFilter": + """Returns a copy of Filter.""" + other = ConcurrentMeanStdFilter(self.shape) + other.sync(self) + return other + + def __repr__(self) -> str: + return "ConcurrentMeanStdFilter({}, {}, {}, {}, {}, {})".format( + self.shape, + self.demean, + self.destd, + self.clip, + self.running_stats, + self.buffer, + ) + + +@OldAPIStack +def get_filter(filter_config, shape): + if filter_config == "MeanStdFilter": + return MeanStdFilter(shape, clip=None) + elif filter_config == "ConcurrentMeanStdFilter": + return ConcurrentMeanStdFilter(shape, clip=None) + elif filter_config == "NoFilter": + return NoFilter() + elif callable(filter_config): + return filter_config(shape) + else: + raise Exception("Unknown observation_filter: " + str(filter_config)) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/filter_manager.py b/lib/python3.12/site-packages/ray/rllib/utils/filter_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..b814ee163b84b0c1b61976a053d1e3f833a2376d --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/filter_manager.py @@ -0,0 +1,82 @@ +import logging +from typing import Optional + +import ray +from ray.rllib.utils.annotations import OldAPIStack + +logger = logging.getLogger(__name__) + + +@OldAPIStack +class FilterManager: + """Manages filters and coordination across remote evaluators that expose + `get_filters` and `sync_filters`. + """ + + @staticmethod + def synchronize( + local_filters, + worker_set, + update_remote=True, + timeout_seconds: Optional[float] = None, + use_remote_data_for_update: bool = True, + ): + """Aggregates filters from remote workers (if use_remote_data_for_update=True). + + Local copy is updated and then broadcasted to all remote evaluators + (if `update_remote` is True). + + Args: + local_filters: Filters to be synchronized. + worker_set: EnvRunnerGroup with remote EnvRunners with filters. + update_remote: Whether to push updates from the local filters to the remote + workers' filters. + timeout_seconds: How long to wait for filter to get or set filters + use_remote_data_for_update: Whether to use the `worker_set`'s remote workers + to update the local filters. If False, stats from the remote workers + will not be used and discarded. + """ + # No sync/update required in either direction -> Early out. + if not (update_remote or use_remote_data_for_update): + return + + logger.debug(f"Synchronizing filters: {local_filters}") + + # Get the filters from the remote workers. + remote_filters = worker_set.foreach_env_runner( + func=lambda worker: worker.get_filters(flush_after=True), + local_env_runner=False, + timeout_seconds=timeout_seconds, + ) + if len(remote_filters) != worker_set.num_healthy_remote_workers(): + logger.error( + "Failed to get remote filters from a rollout worker in " + "FilterManager! " + "Filtered metrics may be computed, but filtered wrong." + ) + + # Should we utilize the remote workers' filter stats to update the local + # filters? + if use_remote_data_for_update: + for rf in remote_filters: + for k in local_filters: + local_filters[k].apply_changes(rf[k], with_buffer=False) + + # Should we update the remote workers' filters from the (now possibly synched) + # local filters? + if update_remote: + copies = {k: v.as_serializable() for k, v in local_filters.items()} + remote_copy = ray.put(copies) + + logger.debug("Updating remote filters ...") + results = worker_set.foreach_env_runner( + func=lambda worker: worker.sync_filters(ray.get(remote_copy)), + local_env_runner=False, + timeout_seconds=timeout_seconds, + ) + if len(results) != worker_set.num_healthy_remote_workers(): + logger.error( + "Failed to set remote filters to a rollout worker in " + "FilterManager. " + "Filtered metrics may be computed, but filtered wrong." + ) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/framework.py b/lib/python3.12/site-packages/ray/rllib/utils/framework.py new file mode 100644 index 0000000000000000000000000000000000000000..1eb5674c8030cb5fb1e9abc859721713ee7acdd4 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/framework.py @@ -0,0 +1,408 @@ +import logging +import os +import sys +from typing import TYPE_CHECKING, Any, Optional + +import numpy as np +import tree # pip install dm_tree + +import ray +from ray._common.deprecation import Deprecated +from ray.rllib.utils.annotations import DeveloperAPI, PublicAPI +from ray.rllib.utils.typing import ( + TensorShape, + TensorStructType, + TensorType, +) + +if TYPE_CHECKING: + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + +logger = logging.getLogger(__name__) + + +@PublicAPI +def convert_to_tensor( + data: TensorStructType, + framework: str, + device: Optional[str] = None, +): + """Converts any nested numpy struct into framework-specific tensors. + + Args: + data: The input data (numpy) to convert to framework-specific tensors. + framework: The framework to convert to. Only "torch" and "tf2" allowed. + device: An optional device name (for torch only). + + Returns: + The converted tensor struct matching the input data. + """ + if framework == "torch": + from ray.rllib.utils.torch_utils import convert_to_torch_tensor + + return convert_to_torch_tensor(data, device=device) + elif framework == "tf2": + _, tf, _ = try_import_tf() + + return tree.map_structure(lambda s: tf.convert_to_tensor(s), data) + raise NotImplementedError( + f"framework={framework} not supported in `convert_to_tensor()`!" + ) + + +@PublicAPI +def get_device(config: "AlgorithmConfig", num_gpus_requested: int = 1): + """Returns a single device (CPU or some GPU) depending on a config. + + Args: + config: An AlgorithmConfig to extract information from about the device to use. + num_gpus_requested: The number of GPUs actually requested. This may be the value + of `config.num_gpus_per_env_runner` when for example calling this function + from an EnvRunner. + + Returns: + A single device (or name) given `config` and `num_gpus_requested`. + """ + if config.framework_str == "torch": + torch, _ = try_import_torch() + + # TODO (Kourosh): How do we handle model parallelism? + # TODO (Kourosh): Instead of using _TorchAccelerator, we should use the public + # API in ray.train but allow for session to be None without any errors raised. + if num_gpus_requested > 0: + from ray.air._internal.torch_utils import get_devices + + # `get_devices()` returns a list that contains the 0th device if + # it is called from outside a Ray Train session. It's necessary to give + # the user the option to run on the gpu of their choice, so we enable that + # option here through `config.local_gpu_idx`. + devices = get_devices() + # Note, if we have a single learner and we do not run on Ray Tune, the local + # learner is not an Ray actor and Ray does not manage devices for it. + if ( + len(devices) == 1 + and ray._private.worker._mode() == ray._private.worker.WORKER_MODE + ): + return devices[0] + else: + assert config.local_gpu_idx < torch.cuda.device_count(), ( + f"local_gpu_idx {config.local_gpu_idx} is not a valid GPU ID " + "or is not available." + ) + # This is an index into the available CUDA devices. For example, if + # `os.environ["CUDA_VISIBLE_DEVICES"] = "1"` then + # `torch.cuda.device_count() = 1` and torch.device(0) maps to that GPU + # with ID=1 on the node. + return torch.device(config.local_gpu_idx) + else: + return torch.device("cpu") + else: + raise NotImplementedError( + f"`framework_str` {config.framework_str} not supported!" + ) + + +@PublicAPI +def try_import_jax(error: bool = False): + """Tries importing JAX and FLAX and returns both modules (or Nones). + + Args: + error: Whether to raise an error if JAX/FLAX cannot be imported. + + Returns: + Tuple containing the jax- and the flax modules. + + Raises: + ImportError: If error=True and JAX is not installed. + """ + if "RLLIB_TEST_NO_JAX_IMPORT" in os.environ: + logger.warning("Not importing JAX for test purposes.") + return None, None + + try: + import flax + import jax + except ImportError: + if error: + raise ImportError( + "Could not import JAX! RLlib requires you to " + "install at least one deep-learning framework: " + "`pip install [torch|tensorflow|jax]`." + ) + return None, None + + return jax, flax + + +@PublicAPI +def try_import_tf(error: bool = False): + """Tries importing tf and returns the module (or None). + + Args: + error: Whether to raise an error if tf cannot be imported. + + Returns: + Tuple containing + 1) tf1.x module (either from tf2.x.compat.v1 OR as tf1.x). + 2) tf module (resulting from `import tensorflow`). Either tf1.x or + 2.x. 3) The actually installed tf version as int: 1 or 2. + + Raises: + ImportError: If error=True and tf is not installed. + """ + tf_stub = _TFStub() + # Make sure, these are reset after each test case + # that uses them: del os.environ["RLLIB_TEST_NO_TF_IMPORT"] + if "RLLIB_TEST_NO_TF_IMPORT" in os.environ: + logger.warning("Not importing TensorFlow for test purposes") + return None, tf_stub, None + + if "TF_CPP_MIN_LOG_LEVEL" not in os.environ: + os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3" + + # Try to reuse already imported tf module. This will avoid going through + # the initial import steps below and thereby switching off v2_behavior + # (switching off v2 behavior twice breaks all-framework tests for eager). + was_imported = False + if "tensorflow" in sys.modules: + tf_module = sys.modules["tensorflow"] + was_imported = True + + else: + try: + import tensorflow as tf_module + except ImportError: + if error: + raise ImportError( + "Could not import TensorFlow! RLlib requires you to " + "install at least one deep-learning framework: " + "`pip install [torch|tensorflow|jax]`." + ) + return None, tf_stub, None + + # Try "reducing" tf to tf.compat.v1. + try: + tf1_module = tf_module.compat.v1 + tf1_module.logging.set_verbosity(tf1_module.logging.ERROR) + if not was_imported: + tf1_module.disable_v2_behavior() + tf1_module.enable_resource_variables() + tf1_module.logging.set_verbosity(tf1_module.logging.WARN) + # No compat.v1 -> return tf as is. + except AttributeError: + tf1_module = tf_module + + if not hasattr(tf_module, "__version__"): + version = 1 # sphinx doc gen + else: + version = 2 if "2." in tf_module.__version__[:2] else 1 + + return tf1_module, tf_module, version + + +# Fake module for tf. +class _TFStub: + def __init__(self) -> None: + self.keras = _KerasStub() + + def __bool__(self): + # if tf should return False + return False + + +# Fake module for tf.keras. +class _KerasStub: + def __init__(self) -> None: + self.Model = _FakeTfClassStub + + +# Fake classes under keras (e.g for tf.keras.Model) +class _FakeTfClassStub: + def __init__(self, *a, **kw): + raise ImportError("Could not import `tensorflow`. Try pip install tensorflow.") + + +@DeveloperAPI +def tf_function(tf_module): + """Conditional decorator for @tf.function. + + Use @tf_function(tf) instead to avoid errors if tf is not installed.""" + + # The actual decorator to use (pass in `tf` (which could be None)). + def decorator(func): + # If tf not installed -> return function as is (won't be used anyways). + if tf_module is None or tf_module.executing_eagerly(): + return func + # If tf installed, return @tf.function-decorated function. + return tf_module.function(func) + + return decorator + + +@PublicAPI +def try_import_tfp(error: bool = False): + """Tries importing tfp and returns the module (or None). + + Args: + error: Whether to raise an error if tfp cannot be imported. + + Returns: + The tfp module. + + Raises: + ImportError: If error=True and tfp is not installed. + """ + if "RLLIB_TEST_NO_TF_IMPORT" in os.environ: + logger.warning("Not importing TensorFlow Probability for test purposes.") + return None + + try: + import tensorflow_probability as tfp + + return tfp + except ImportError as e: + if error: + raise e + return None + + +# Fake module for torch.nn. +class _NNStub: + def __init__(self, *a, **kw): + # Fake nn.functional module within torch.nn. + self.functional = None + self.Module = _FakeTorchClassStub + self.parallel = _ParallelStub() + + +# Fake class for e.g. torch.nn.Module to allow it to be inherited from. +class _FakeTorchClassStub: + def __init__(self, *a, **kw): + raise ImportError("Could not import `torch`. Try pip install torch.") + + +class _ParallelStub: + def __init__(self, *a, **kw): + self.DataParallel = _FakeTorchClassStub + self.DistributedDataParallel = _FakeTorchClassStub + + +@PublicAPI +def try_import_torch(error: bool = False): + """Tries importing torch and returns the module (or None). + + Args: + error: Whether to raise an error if torch cannot be imported. + + Returns: + Tuple consisting of the torch- AND torch.nn modules. + + Raises: + ImportError: If error=True and PyTorch is not installed. + """ + if "RLLIB_TEST_NO_TORCH_IMPORT" in os.environ: + logger.warning("Not importing PyTorch for test purposes.") + return _torch_stubs() + + try: + import torch + import torch.nn as nn + + return torch, nn + except ImportError: + if error: + raise ImportError( + "Could not import PyTorch! RLlib requires you to " + "install at least one deep-learning framework: " + "`pip install [torch|tensorflow|jax]`." + ) + return _torch_stubs() + + +def _torch_stubs(): + nn = _NNStub() + return None, nn + + +@DeveloperAPI +def get_variable( + value: Any, + framework: str = "tf", + trainable: bool = False, + tf_name: str = "unnamed-variable", + torch_tensor: bool = False, + device: Optional[str] = None, + shape: Optional[TensorShape] = None, + dtype: Optional[TensorType] = None, +) -> Any: + """Creates a tf variable, a torch tensor, or a python primitive. + + Args: + value: The initial value to use. In the non-tf case, this will + be returned as is. In the tf case, this could be a tf-Initializer + object. + framework: One of "tf", "torch", or None. + trainable: Whether the generated variable should be + trainable (tf)/require_grad (torch) or not (default: False). + tf_name: For framework="tf": An optional name for the + tf.Variable. + torch_tensor: For framework="torch": Whether to actually create + a torch.tensor, or just a python value (default). + device: An optional torch device to use for + the created torch tensor. + shape: An optional shape to use iff `value` + does not have any (e.g. if it's an initializer w/o explicit value). + dtype: An optional dtype to use iff `value` does + not have any (e.g. if it's an initializer w/o explicit value). + This should always be a numpy dtype (e.g. np.float32, np.int64). + + Returns: + A framework-specific variable (tf.Variable, torch.tensor, or + python primitive). + """ + if framework in ["tf2", "tf"]: + import tensorflow as tf + + dtype = dtype or getattr( + value, + "dtype", + tf.float32 + if isinstance(value, float) + else tf.int32 + if isinstance(value, int) + else None, + ) + return tf.compat.v1.get_variable( + tf_name, + initializer=value, + dtype=dtype, + trainable=trainable, + **({} if shape is None else {"shape": shape}), + ) + elif framework == "torch" and torch_tensor is True: + torch, _ = try_import_torch() + if not isinstance(value, np.ndarray): + value = np.array(value) + var_ = torch.from_numpy(value) + if dtype in [torch.float32, np.float32]: + var_ = var_.float() + elif dtype in [torch.int32, np.int32]: + var_ = var_.int() + elif dtype in [torch.float64, np.float64]: + var_ = var_.double() + + if device: + var_ = var_.to(device) + var_.requires_grad = trainable + return var_ + # torch or None: Return python primitive. + return value + + +@Deprecated( + old="rllib/utils/framework.py::get_activation_fn", + new="rllib/models/utils.py::get_activation_fn", + error=True, +) +def get_activation_fn(name: Optional[str] = None, framework: str = "tf"): + pass diff --git a/lib/python3.12/site-packages/ray/rllib/utils/from_config.py b/lib/python3.12/site-packages/ray/rllib/utils/from_config.py new file mode 100644 index 0000000000000000000000000000000000000000..3f80e785265b06fd00d3b372758ff4836823ad10 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/from_config.py @@ -0,0 +1,330 @@ +import importlib +import json +import os +import re +from copy import deepcopy +from functools import partial +from typing import TYPE_CHECKING, Optional + +import yaml + +from ray.rllib.utils import force_list, merge_dicts +from ray.rllib.utils.annotations import DeveloperAPI + +if TYPE_CHECKING: + from ray.rllib.utils.typing import FromConfigSpec + + +@DeveloperAPI +def from_config(cls, config: Optional["FromConfigSpec"] = None, **kwargs): + """Uses the given config to create an object. + + If `config` is a dict, an optional "type" key can be used as a + "constructor hint" to specify a certain class of the object. + If `config` is not a dict, `config`'s value is used directly as this + "constructor hint". + + The rest of `config` (if it's a dict) will be used as kwargs for the + constructor. Additional keys in **kwargs will always have precedence + (overwrite keys in `config` (if a dict)). + Also, if the config-dict or **kwargs contains the special key "_args", + it will be popped from the dict and used as *args list to be passed + separately to the constructor. + + The following constructor hints are valid: + - None: Use `cls` as constructor. + - An already instantiated object: Will be returned as is; no + constructor call. + - A string or an object that is a key in `cls`'s `__type_registry__` + dict: The value in `__type_registry__` for that key will be used + as the constructor. + - A python callable: Use that very callable as constructor. + - A string: Either a json/yaml filename or the name of a python + module+class (e.g. "ray.rllib. [...] .[some class name]") + + Args: + cls: The class to build an instance for (from `config`). + config (Optional[dict, str]): The config dict or type-string or + filename. + + Keyword Args: + kwargs: Optional possibility to pass the constructor arguments in + here and use `config` as the type-only info. Then we can call + this like: from_config([type]?, [**kwargs for constructor]) + If `config` is already a dict, then `kwargs` will be merged + with `config` (overwriting keys in `config`) after "type" has + been popped out of `config`. + If a constructor of a Configurable needs *args, the special + key `_args` can be passed inside `kwargs` with a list value + (e.g. kwargs={"_args": [arg1, arg2, arg3]}). + + Returns: + any: The object generated from the config. + """ + # `cls` is the config (config is None). + if config is None and isinstance(cls, (dict, str)): + config = cls + cls = None + # `config` is already a created object of this class -> + # Take it as is. + elif isinstance(cls, type) and isinstance(config, cls): + return config + + # `type_`: Indicator for the Configurable's constructor. + # `ctor_args`: *args arguments for the constructor. + # `ctor_kwargs`: **kwargs arguments for the constructor. + # Try to copy, so caller can reuse safely. + try: + config = deepcopy(config) + except Exception: + pass + if isinstance(config, dict): + type_ = config.pop("type", None) + if type_ is None and isinstance(cls, str): + type_ = cls + ctor_kwargs = config + # Give kwargs priority over things defined in config dict. + # This way, one can pass a generic `spec` and then override single + # constructor parameters via the kwargs in the call to `from_config`. + ctor_kwargs.update(kwargs) + else: + type_ = config + if type_ is None and "type" in kwargs: + type_ = kwargs.pop("type") + ctor_kwargs = kwargs + # Special `_args` field in kwargs for *args-utilizing constructors. + ctor_args = force_list(ctor_kwargs.pop("_args", [])) + + # Figure out the actual constructor (class) from `type_`. + # None: Try __default__object (if no args/kwargs), only then + # constructor of cls (using args/kwargs). + if type_ is None: + # We have a default constructor that was defined directly by cls + # (not by its children). + if ( + cls is not None + and hasattr(cls, "__default_constructor__") + and cls.__default_constructor__ is not None + and ctor_args == [] + and ( + not hasattr(cls.__bases__[0], "__default_constructor__") + or cls.__bases__[0].__default_constructor__ is None + or cls.__bases__[0].__default_constructor__ + is not cls.__default_constructor__ + ) + ): + constructor = cls.__default_constructor__ + # Default constructor's keywords into ctor_kwargs. + if isinstance(constructor, partial): + kwargs = merge_dicts(ctor_kwargs, constructor.keywords) + constructor = partial(constructor.func, **kwargs) + ctor_kwargs = {} # erase to avoid duplicate kwarg error + # No default constructor -> Try cls itself as constructor. + else: + constructor = cls + # Try the __type_registry__ of this class. + else: + constructor = _lookup_type(cls, type_) + + # Found in cls.__type_registry__. + if constructor is not None: + pass + # type_ is False or None (and this value is not registered) -> + # return value of type_. + elif type_ is False or type_ is None: + return type_ + # Python callable. + elif callable(type_): + constructor = type_ + # A string: Filename or a python module+class or a json/yaml str. + elif isinstance(type_, str): + if re.search("\\.(yaml|yml|json)$", type_): + return from_file(cls, type_, *ctor_args, **ctor_kwargs) + # Try un-json/un-yaml'ing the string into a dict. + obj = yaml.safe_load(type_) + if isinstance(obj, dict): + return from_config(cls, obj) + try: + obj = from_config(cls, json.loads(type_)) + except json.JSONDecodeError: + pass + else: + return obj + + # Test for absolute module.class path specifier. + if type_.find(".") != -1: + module_name, function_name = type_.rsplit(".", 1) + try: + module = importlib.import_module(module_name) + constructor = getattr(module, function_name) + # Module not found. + except (ModuleNotFoundError, ImportError, AttributeError): + pass + + # If constructor still not found, try attaching cls' module, + # then look for type_ in there. + if constructor is None: + if isinstance(cls, str): + # Module found, but doesn't have the specified + # c'tor/function. + raise ValueError( + f"Full classpath specifier ({type_}) must be a valid " + "full [module].[class] string! E.g.: " + "`my.cool.module.MyCoolClass`." + ) + + try: + module = importlib.import_module(cls.__module__) + constructor = getattr(module, type_) + except (ModuleNotFoundError, ImportError, AttributeError): + # Try the package as well. + try: + package_name = importlib.import_module( + cls.__module__ + ).__package__ + module = __import__(package_name, fromlist=[type_]) + constructor = getattr(module, type_) + except (ModuleNotFoundError, ImportError, AttributeError): + pass + + if constructor is None: + raise ValueError( + f"String specifier ({type_}) must be a valid filename, " + f"a [module].[class], a class within '{cls.__module__}', " + f"or a key into {cls.__name__}.__type_registry__!" + ) + + if not constructor: + raise TypeError("Invalid type '{}'. Cannot create `from_config`.".format(type_)) + + # Create object with inferred constructor. + try: + object_ = constructor(*ctor_args, **ctor_kwargs) + # Catch attempts to construct from an abstract class and return None. + except TypeError as e: + if re.match("Can't instantiate abstract class", e.args[0]): + return None + raise e # Re-raise + # No sanity check for fake (lambda)-"constructors". + if type(constructor).__name__ != "function": + assert isinstance( + object_, + constructor.func if isinstance(constructor, partial) else constructor, + ) + + return object_ + + +@DeveloperAPI +def from_file(cls, filename, *args, **kwargs): + """ + Create object from config saved in filename. Expects json or yaml file. + + Args: + filename: File containing the config (json or yaml). + + Returns: + any: The object generated from the file. + """ + path = os.path.join(os.getcwd(), filename) + if not os.path.isfile(path): + raise FileNotFoundError("File '{}' not found!".format(filename)) + + with open(path, "rt") as fp: + if path.endswith(".yaml") or path.endswith(".yml"): + config = yaml.safe_load(fp) + else: + config = json.load(fp) + + # Add possible *args. + config["_args"] = args + return from_config(cls, config=config, **kwargs) + + +def _lookup_type(cls, type_): + if ( + cls is not None + and hasattr(cls, "__type_registry__") + and isinstance(cls.__type_registry__, dict) + and ( + type_ in cls.__type_registry__ + or ( + isinstance(type_, str) + and re.sub("[\\W_]", "", type_.lower()) in cls.__type_registry__ + ) + ) + ): + available_class_for_type = cls.__type_registry__.get(type_) + if available_class_for_type is None: + available_class_for_type = cls.__type_registry__[ + re.sub("[\\W_]", "", type_.lower()) + ] + return available_class_for_type + return None + + +class _NotProvided: + """Singleton class to provide a "not provided" value for AlgorithmConfig signatures. + + Using the only instance of this class indicates that the user does NOT wish to + change the value of some property. + + .. testcode:: + :skipif: True + + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + config = AlgorithmConfig() + # Print out the default learning rate. + print(config.lr) + + .. testoutput:: + + 0.001 + + .. testcode:: + :skipif: True + + # Print out the default `preprocessor_pref`. + print(config.preprocessor_pref) + + .. testoutput:: + + "deepmind" + + .. testcode:: + :skipif: True + + # Will only set the `preprocessor_pref` property (to None) and leave + # all other properties at their default values. + config.training(preprocessor_pref=None) + config.preprocessor_pref is None + + .. testoutput:: + + True + + .. testcode:: + :skipif: True + + # Still the same value (didn't touch it in the call to `.training()`. + print(config.lr) + + .. testoutput:: + + 0.001 + """ + + class __NotProvided: + pass + + instance = None + + def __init__(self): + if _NotProvided.instance is None: + _NotProvided.instance = _NotProvided.__NotProvided() + + +# Use this object as default values in all method signatures of +# AlgorithmConfig, indicating that the respective property should NOT be touched +# in the call. +NotProvided = _NotProvided() diff --git a/lib/python3.12/site-packages/ray/rllib/utils/images.py b/lib/python3.12/site-packages/ray/rllib/utils/images.py new file mode 100644 index 0000000000000000000000000000000000000000..322d0e559c5c4402db29e90dcf82c8b374eae8c2 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/images.py @@ -0,0 +1,60 @@ +import importlib +import logging + +import numpy as np + +from ray.rllib.utils.annotations import DeveloperAPI + +logger = logging.getLogger(__name__) + + +@DeveloperAPI +def is_package_installed(package_name): + try: + importlib.metadata.version(package_name) + return True + except importlib.metadata.PackageNotFoundError: + return False + + +try: + import cv2 + + cv2.ocl.setUseOpenCL(False) + + logger.debug("CV2 found for image processing.") +except ImportError as e: + if is_package_installed("opencv-python"): + raise ImportError( + f"OpenCV is installed, but we failed to import it. This may be because " + f"you need to install `opencv-python-headless` instead of " + f"`opencv-python`. Error message: {e}", + ) + cv2 = None + + +@DeveloperAPI +def resize(img: np.ndarray, height: int, width: int) -> np.ndarray: + if not cv2: + raise ModuleNotFoundError( + "`opencv` not installed! Do `pip install opencv-python`" + ) + return cv2.resize(img, (width, height), interpolation=cv2.INTER_AREA) + + +@DeveloperAPI +def rgb2gray(img: np.ndarray) -> np.ndarray: + if not cv2: + raise ModuleNotFoundError( + "`opencv` not installed! Do `pip install opencv-python`" + ) + return cv2.cvtColor(img, cv2.COLOR_RGB2GRAY) + + +@DeveloperAPI +def imread(img_file: str) -> np.ndarray: + if not cv2: + raise ModuleNotFoundError( + "`opencv` not installed! Do `pip install opencv-python`" + ) + return cv2.imread(img_file).astype(np.float32) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/lambda_defaultdict.py b/lib/python3.12/site-packages/ray/rllib/utils/lambda_defaultdict.py new file mode 100644 index 0000000000000000000000000000000000000000..ce4653961c565304570da87eea9066dc366643e0 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/lambda_defaultdict.py @@ -0,0 +1,52 @@ +from collections import defaultdict +from typing import Any, Callable + + +class LambdaDefaultDict(defaultdict): + """A defaultdict that creates default values based on the associated key. + + Note that the standard defaultdict can only produce default values (via its factory) + that are independent of the key under which they are stored. + As opposed to that, the lambda functions used as factories for this + `LambdaDefaultDict` class do accept a single argument: The missing key. + If a missing key is accessed by the user, the provided lambda function is called + with this missing key as its argument. The returned value is stored in the + dictionary under that key and returned. + + Example: + + In this example, if you try to access a key that doesn't exist, it will call + the lambda function, passing it the missing key. The function will return a + string, which will be stored in the dictionary under that key. + + .. testcode:: + + from ray.rllib.utils.lambda_defaultdict import LambdaDefaultDict + + default_dict = LambdaDefaultDict(lambda missing_key: f"Value for {missing_key}") + print(default_dict["a"]) + + .. testoutput:: + + Value for a + """ # noqa: E501 + + def __init__(self, default_factory: Callable[[str], Any], *args, **kwargs): + """Initializes a LambdaDefaultDict instance. + + Args: + default_factory: The default factory callable, taking a string (key) + and returning the default value to use for that key. + """ + if not callable(default_factory): + raise TypeError("First argument must be a Callable!") + + # We will handle the factory in __missing__ method. + super().__init__(None, *args, **kwargs) + + self.default_factory = default_factory + + def __missing__(self, key): + # Call default factory with the key as argument. + self[key] = value = self.default_factory(key) + return value diff --git a/lib/python3.12/site-packages/ray/rllib/utils/memory.py b/lib/python3.12/site-packages/ray/rllib/utils/memory.py new file mode 100644 index 0000000000000000000000000000000000000000..323bec70c50f8144e296b4c94266dace76e094ac --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/memory.py @@ -0,0 +1,8 @@ +from ray._common.deprecation import deprecation_warning +from ray.rllib.utils.numpy import aligned_array, concat_aligned # noqa + +deprecation_warning( + old="ray.rllib.utils.memory.[...]", + new="ray.rllib.utils.numpy.[...]", + error=True, +) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/metrics/__init__.py b/lib/python3.12/site-packages/ray/rllib/utils/metrics/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2704918e32f0c83feb089c563520e2914620276b --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/metrics/__init__.py @@ -0,0 +1,189 @@ +from ray.rllib.core import ALL_MODULES # noqa + + +# Algorithm ResultDict keys. +AGGREGATOR_ACTOR_RESULTS = "aggregator_actors" +DIFFERENTIABLE_LEARNER_RESULTS = "differentiable_learners" +EVALUATION_RESULTS = "evaluation" +ENV_RUNNER_RESULTS = "env_runners" +OFFLINE_EVAL_RUNNER_RESULTS = "offline_eval_runners" +FAULT_TOLERANCE_STATS = "fault_tolerance" +LEARNER_GROUP = "learner_group" +LEARNER_RESULTS = "learners" +REPLAY_BUFFER_RESULTS = "replay_buffer" +TIMERS = "timers" + +# RLModule metrics. +NUM_TRAINABLE_PARAMETERS = "num_trainable_parameters" +NUM_NON_TRAINABLE_PARAMETERS = "num_non_trainable_parameters" + +# Number of times `training_step()` was called in one iteration. +NUM_TRAINING_STEP_CALLS_PER_ITERATION = "num_training_step_calls_per_iteration" + +# Counters for sampling, sampling (on eval workers) and +# training steps (env- and agent steps). +MEAN_NUM_EPISODE_LISTS_RECEIVED = "mean_num_episode_lists_received" +NUM_AGENT_STEPS_SAMPLED = "num_agent_steps_sampled" +NUM_AGENT_STEPS_SAMPLED_LIFETIME = "num_agent_steps_sampled_lifetime" +NUM_AGENT_STEPS_SAMPLED_THIS_ITER = "num_agent_steps_sampled_this_iter" # @OldAPIStack +NUM_ENV_STEPS_SAMPLED = "num_env_steps_sampled" +NUM_ENV_STEPS_SAMPLED_LIFETIME = "num_env_steps_sampled_lifetime" +NUM_ENV_STEPS_SAMPLED_PER_SECOND = "num_env_steps_sampled_per_second" # Deprecated +NUM_ENV_STEPS_SAMPLED_THIS_ITER = "num_env_steps_sampled_this_iter" # @OldAPIStack +NUM_ENV_STEPS_SAMPLED_FOR_EVALUATION_THIS_ITER = ( + "num_env_steps_sampled_for_evaluation_this_iter" +) +NUM_MODULE_STEPS_SAMPLED = "num_module_steps_sampled" +NUM_MODULE_STEPS_SAMPLED_LIFETIME = "num_module_steps_sampled_lifetime" +ENV_TO_MODULE_SUM_EPISODES_LENGTH_IN = "env_to_module_sum_episodes_length_in" +ENV_TO_MODULE_SUM_EPISODES_LENGTH_OUT = "env_to_module_sum_episodes_length_out" + +# Counters for adding and evicting in replay buffers. +ACTUAL_N_STEP = "actual_n_step" +AGENT_ACTUAL_N_STEP = "agent_actual_n_step" +AGENT_STEP_UTILIZATION = "agent_step_utilization" +MODULE_ACTUAL_N_STEP = "module_actual_n_step" +MODULE_STEP_UTILIZATION = "module_step_utilization" +ENV_STEP_UTILIZATION = "env_step_utilization" +NUM_AGENT_EPISODES_STORED = "num_agent_episodes" +NUM_AGENT_EPISODES_ADDED = "num_agent_episodes_added" +NUM_AGENT_EPISODES_ADDED_LIFETIME = "num_agent_episodes_added_lifetime" +NUM_AGENT_EPISODES_EVICTED = "num_agent_episodes_evicted" +NUM_AGENT_EPISODES_EVICTED_LIFETIME = "num_agent_episodes_evicted_lifetime" +NUM_AGENT_EPISODES_PER_SAMPLE = "num_agent_episodes_per_sample" +NUM_AGENT_RESAMPLES = "num_agent_resamples" +NUM_AGENT_STEPS_ADDED = "num_agent_steps_added" +NUM_AGENT_STEPS_ADDED_LIFETIME = "num_agent_steps_added_lifetime" +NUM_AGENT_STEPS_EVICTED = "num_agent_steps_evicted" +NUM_AGENT_STEPS_EVICTED_LIFETIME = "num_agent_steps_evicted_lifetime" +NUM_AGENT_STEPS_PER_SAMPLE = "num_agent_steps_per_sample" +NUM_AGENT_STEPS_PER_SAMPLE_LIFETIME = "num_agent_steps_per_sample_lifetime" +NUM_AGENT_STEPS_STORED = "num_agent_steps_stored" +NUM_ENV_STEPS_ADDED = "num_env_steps_added" +NUM_ENV_STEPS_ADDED_LIFETIME = "num_env_steps_added_lifetime" +NUM_ENV_STEPS_EVICTED = "num_env_steps_evicted" +NUM_ENV_STEPS_EVICTED_LIFETIME = "num_env_steps_evicted_lifetime" +NUM_ENV_STEPS_PER_SAMPLE = "num_env_steps_per_sample" +NUM_ENV_STEPS_PER_SAMPLE_LIFETIME = "num_env_steps_per_sample_lifetime" +NUM_ENV_STEPS_STORED = "num_env_steps_stored" +NUM_EPISODES_STORED = "num_episodes_stored" +NUM_EPISODES_ADDED = "num_episodes_added" +NUM_EPISODES_ADDED_LIFETIME = "num_episodes_added_lifetime" +NUM_EPISODES_EVICTED = "num_episodes_evicted" +NUM_EPISODES_EVICTED_LIFETIME = "num_episodes_evicted_lifetime" +NUM_EPISODES_PER_SAMPLE = "num_episodes_per_sample" +NUM_RESAMPLES = "num_resamples" +NUM_MODULE_EPISODES_STORED = "num_module_episodes" +NUM_MODULE_EPISODES_ADDED = "num_module_episodes_added" +NUM_MODULE_EPISODES_ADDED_LIFETIME = "num_module_episodes_added_lifetime" +NUM_MODULE_EPISODES_EVICTED = "num_module_episodes_evicted" +NUM_MODULE_EPISODES_EVICTED_LIFETIME = "num_module_episodes_evicted_lifetime" +NUM_MODULE_EPISODES_PER_SAMPLE = "num_module_episodes_per_sample" +NUM_MODULE_RESAMPLES = "num_module_resamples" +NUM_MODULE_STEPS_ADDED = "num_module_steps_added" +NUM_MODULE_STEPS_ADDED_LIFETIME = "num_module_steps_added_lifetime" +NUM_MODULE_STEPS_EVICTED = "num_module_steps_evicted" +NUM_MODULE_STEPS_EVICTED_LIFETIME = "num_module_steps_evicted_lifetime" +NUM_MODULE_STEPS_PER_SAMPLE = "num_module_steps_per_sample" +NUM_MODULE_STEPS_PER_SAMPLE_LIFETIME = "num_module_steps_per_sample_lifetime" +NUM_MODULE_STEPS_STORED = "num_module_steps_stored" + +EPISODE_DURATION_SEC_MEAN = "episode_duration_sec_mean" +EPISODE_LEN_MEAN = "episode_len_mean" +EPISODE_LEN_MAX = "episode_len_max" +EPISODE_LEN_MIN = "episode_len_min" +EPISODE_RETURN_MEAN = "episode_return_mean" +EPISODE_RETURN_MAX = "episode_return_max" +EPISODE_RETURN_MIN = "episode_return_min" +NUM_EPISODES = "num_episodes" +NUM_EPISODES_LIFETIME = "num_episodes_lifetime" +TIME_BETWEEN_SAMPLING = "time_between_sampling" + +DATASET_NUM_ITERS_TRAINED = "dataset_num_iters_trained" +DATASET_NUM_ITERS_TRAINED_LIFETIME = "dataset_num_iters_trained_lifetime" +DATASET_NUM_ITERS_EVALUATED = "dataset_num_iters_evaluated" +DATASET_NUM_ITERS_EVALUATED_LIFETIME = "dataset_num_iters_evaluated_lifetime" +MEAN_NUM_LEARNER_GROUP_UPDATE_CALLED = "mean_num_learner_group_update_called" +MEAN_NUM_LEARNER_RESULTS_RECEIVED = "mean_num_learner_results_received" +NUM_AGENT_STEPS_TRAINED = "num_agent_steps_trained" +NUM_AGENT_STEPS_TRAINED_LIFETIME = "num_agent_steps_trained_lifetime" +NUM_AGENT_STEPS_TRAINED_THIS_ITER = "num_agent_steps_trained_this_iter" # @OldAPIStack +NUM_ENV_STEPS_TRAINED = "num_env_steps_trained" +NUM_ENV_STEPS_TRAINED_LIFETIME = "num_env_steps_trained_lifetime" +NUM_ENV_STEPS_TRAINED_THIS_ITER = "num_env_steps_trained_this_iter" # @OldAPIStack +NUM_MODULE_STEPS_TRAINED = "num_module_steps_trained" +NUM_MODULE_STEPS_TRAINED_LIFETIME = "num_module_steps_trained_lifetime" +MODULE_SAMPLE_BATCH_SIZE_MEAN = "module_sample_batch_size_mean" +MODULE_TRAIN_BATCH_SIZE_MEAN = "module_train_batch_size_mean" +LEARNER_CONNECTOR_SUM_EPISODES_LENGTH_IN = "learner_connector_sum_episodes_length_in" +LEARNER_CONNECTOR_SUM_EPISODES_LENGTH_OUT = "learner_connector_sum_episodes_length_out" + +# Backward compatibility: Replace with num_env_steps_... or num_agent_steps_... +STEPS_TRAINED_THIS_ITER_COUNTER = "num_steps_trained_this_iter" + +# Counters for keeping track of worker weight updates (synchronization +# between local worker and remote workers). +NUM_SYNCH_WORKER_WEIGHTS = "num_weight_broadcasts" +NUM_TRAINING_STEP_CALLS_SINCE_LAST_SYNCH_WORKER_WEIGHTS = ( + "num_training_step_calls_since_last_synch_worker_weights" +) +# The running sequence number for a set of NN weights. If a worker's NN has a +# lower sequence number than some weights coming in for an update, the worker +# should perform the update, otherwise ignore the incoming weights (they are older +# or the same) as/than the ones it already has. +WEIGHTS_SEQ_NO = "weights_seq_no" +# Number of total gradient updates that have been performed on a policy. +NUM_GRAD_UPDATES_LIFETIME = "num_grad_updates_lifetime" +# Average difference between the number of grad-updates that the policy/ies had +# that collected the training batch vs the policy that was just updated (trained). +# Good measure for the off-policy'ness of training. Should be 0.0 for PPO and PG, +# small for IMPALA and APPO, and any (larger) value for DQN and other off-policy algos. +DIFF_NUM_GRAD_UPDATES_VS_SAMPLER_POLICY = "diff_num_grad_updates_vs_sampler_policy" + +# Counters to track target network updates. +LAST_TARGET_UPDATE_TS = "last_target_update_ts" +NUM_TARGET_UPDATES = "num_target_updates" + +# Performance timers +# ------------------ +# Duration of n `Algorithm.training_step()` calls making up one "iteration". +# Note that n may be >1 if the user has set up a min time (sec) or timesteps per +# iteration. +TRAINING_ITERATION_TIMER = "training_iteration" +# Duration of a `Algorithm.evaluate()` call. +EVALUATION_ITERATION_TIMER = "evaluation_iteration" +OFFLINE_EVALUATION_ITERATION_TIMER = "offline_evaluation_iteration_timer" +# Duration of a single `training_step()` call. +TRAINING_STEP_TIMER = "training_step" +APPLY_GRADS_TIMER = "apply_grad" +COMPUTE_GRADS_TIMER = "compute_grads" +GARBAGE_COLLECTION_TIMER = "garbage_collection" +RESTORE_ENV_RUNNERS_TIMER = "restore_env_runners" +RESTORE_EVAL_ENV_RUNNERS_TIMER = "restore_eval_env_runners" +RESTORE_OFFLINE_EVAL_RUNNERS_TIMER = "restore_offline_eval_runners" +SYNCH_WORKER_WEIGHTS_TIMER = "synch_weights" +SYNCH_ENV_CONNECTOR_STATES_TIMER = "synch_env_connectors" +SYNCH_EVAL_ENV_CONNECTOR_STATES_TIMER = "synch_eval_env_connectors" +GRAD_WAIT_TIMER = "grad_wait" +SAMPLE_TIMER = "sample" # @OldAPIStack +ENV_RUNNER_SAMPLING_TIMER = "env_runner_sampling_timer" +ENV_RESET_TIMER = "env_reset_timer" +ENV_STEP_TIMER = "env_step_timer" +ENV_TO_MODULE_CONNECTOR = "env_to_module_connector" +RLMODULE_INFERENCE_TIMER = "rlmodule_inference_timer" +MODULE_TO_ENV_CONNECTOR = "module_to_env_connector" +OFFLINE_SAMPLING_TIMER = "offline_sampling_timer" +REPLAY_BUFFER_ADD_DATA_TIMER = "replay_buffer_add_data_timer" +REPLAY_BUFFER_SAMPLE_TIMER = "replay_buffer_sampling_timer" +REPLAY_BUFFER_UPDATE_PRIOS_TIMER = "replay_buffer_update_prios_timer" +LEARNER_CONNECTOR = "learner_connector" +LEARNER_UPDATE_TIMER = "learner_update_timer" +LEARN_ON_BATCH_TIMER = "learn" # @OldAPIStack +LOAD_BATCH_TIMER = "load" +TARGET_NET_UPDATE_TIMER = "target_net_update" +CONNECTOR_PIPELINE_TIMER = "connector_pipeline_timer" +CONNECTOR_TIMERS = "connectors" + +# Learner. +LEARNER_STATS_KEY = "learner_stats" +TD_ERROR_KEY = "td_error" diff --git a/lib/python3.12/site-packages/ray/rllib/utils/metrics/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/metrics/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9900e9b1f83c5dbcfdfff2e7cd502bb565b8a499 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/metrics/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/metrics/__pycache__/learner_info.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/metrics/__pycache__/learner_info.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c7d7e2fbec16a4b553324f2a06dd1b82a9dca375 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/metrics/__pycache__/learner_info.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/metrics/__pycache__/metrics_logger.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/metrics/__pycache__/metrics_logger.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cea22b44235dfd87f3e82b1a889db916e793aed5 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/metrics/__pycache__/metrics_logger.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/metrics/__pycache__/ray_metrics.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/metrics/__pycache__/ray_metrics.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3c2d00a126992d6c43017503c7d469bab51e1c30 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/metrics/__pycache__/ray_metrics.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/metrics/__pycache__/stats.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/metrics/__pycache__/stats.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..07f2fb350a7748dabe3e791d4b61742c22e99114 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/metrics/__pycache__/stats.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/metrics/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/metrics/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0c52964742c8207d69386b84540fd1f22a2b51b0 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/metrics/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/metrics/__pycache__/window_stat.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/metrics/__pycache__/window_stat.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d30765f96c2f6449685740b540e8cdd088387ce2 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/metrics/__pycache__/window_stat.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/metrics/learner_info.py b/lib/python3.12/site-packages/ray/rllib/utils/metrics/learner_info.py new file mode 100644 index 0000000000000000000000000000000000000000..ad5a96b946cd5cf215a2668ef78a2476573ca818 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/metrics/learner_info.py @@ -0,0 +1,121 @@ +from collections import defaultdict +from typing import Dict + +import numpy as np +import tree # pip install dm_tree + +from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.typing import PolicyID + +# Instant metrics (keys for metrics.info). +LEARNER_INFO = "learner" +# By convention, metrics from optimizing the loss can be reported in the +# `grad_info` dict returned by learn_on_batch() / compute_grads() via this key. +LEARNER_STATS_KEY = "learner_stats" + + +@OldAPIStack +class LearnerInfoBuilder: + def __init__(self, num_devices: int = 1): + self.num_devices = num_devices + self.results_all_towers = defaultdict(list) + self.is_finalized = False + + def add_learn_on_batch_results( + self, + results: Dict, + policy_id: PolicyID = DEFAULT_POLICY_ID, + ) -> None: + """Adds a policy.learn_on_(loaded)?_batch() result to this builder. + + Args: + results: The results returned by Policy.learn_on_batch or + Policy.learn_on_loaded_batch. + policy_id: The policy's ID, whose learn_on_(loaded)_batch method + returned `results`. + """ + assert ( + not self.is_finalized + ), "LearnerInfo already finalized! Cannot add more results." + + # No towers: Single CPU. + if "tower_0" not in results: + self.results_all_towers[policy_id].append(results) + # Multi-GPU case: + else: + self.results_all_towers[policy_id].append( + tree.map_structure_with_path( + lambda p, *s: _all_tower_reduce(p, *s), + *( + results.pop("tower_{}".format(tower_num)) + for tower_num in range(self.num_devices) + ) + ) + ) + for k, v in results.items(): + if k == LEARNER_STATS_KEY: + for k1, v1 in results[k].items(): + self.results_all_towers[policy_id][-1][LEARNER_STATS_KEY][ + k1 + ] = v1 + else: + self.results_all_towers[policy_id][-1][k] = v + + def add_learn_on_batch_results_multi_agent( + self, + all_policies_results: Dict, + ) -> None: + """Adds multiple policy.learn_on_(loaded)?_batch() results to this builder. + + Args: + all_policies_results: The results returned by all Policy.learn_on_batch or + Policy.learn_on_loaded_batch wrapped as a dict mapping policy ID to + results. + """ + for pid, result in all_policies_results.items(): + if pid != "batch_count": + self.add_learn_on_batch_results(result, policy_id=pid) + + def finalize(self): + self.is_finalized = True + + info = {} + for policy_id, results_all_towers in self.results_all_towers.items(): + # Reduce mean across all minibatch SGD steps (axis=0 to keep + # all shapes as-is). + info[policy_id] = tree.map_structure_with_path( + _all_tower_reduce, *results_all_towers + ) + + return info + + +@OldAPIStack +def _all_tower_reduce(path, *tower_data): + """Reduces stats across towers based on their stats-dict paths.""" + # TD-errors: Need to stay per batch item in order to be able to update + # each item's weight in a prioritized replay buffer. + if len(path) == 1 and path[0] == "td_error": + return np.concatenate(tower_data, axis=0) + elif tower_data[0] is None: + return None + + if isinstance(path[-1], str): + # TODO(sven): We need to fix this terrible dependency on `str.starts_with` + # for determining, how to aggregate these stats! As "num_..." might + # be a good indicator for summing, it will fail if the stats is e.g. + # `num_samples_per_sec" :) + # Counter stats: Reduce sum. + # if path[-1].startswith("num_"): + # return np.nansum(tower_data) + # Min stats: Reduce min. + if path[-1].startswith("min_"): + return np.nanmin(tower_data) + # Max stats: Reduce max. + elif path[-1].startswith("max_"): + return np.nanmax(tower_data) + if np.isnan(tower_data).all(): + return np.nan + # Everything else: Reduce mean. + return np.nanmean(tower_data) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/metrics/metrics_logger.py b/lib/python3.12/site-packages/ray/rllib/utils/metrics/metrics_logger.py new file mode 100644 index 0000000000000000000000000000000000000000..0ae087254a8a38dc0cd706244b0c666fc9760877 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/metrics/metrics_logger.py @@ -0,0 +1,1367 @@ +import logging +from typing import Any, Dict, List, Optional, Tuple, Union + +import tree # pip install dm_tree + +from ray._common.deprecation import Deprecated, deprecation_warning +from ray.rllib.utils import deep_update, force_tuple +from ray.rllib.utils.framework import try_import_tf, try_import_torch +from ray.rllib.utils.metrics.stats import Stats, merge_stats +from ray.util import log_once +from ray.util.annotations import PublicAPI + +_, tf, _ = try_import_tf() +torch, _ = try_import_torch() +logger = logging.getLogger("ray.rllib") + + +@PublicAPI(stability="alpha") +class MetricsLogger: + """A generic class collecting and processing metrics in RL training and evaluation. + + This class represents the main API used by all of RLlib's components (internal and + user facing) in order to log, collect, and process (reduce) stats during training + and evaluation/inference. + + It supports: + - Logging of simple float/int values (for example a loss) over time or from + parallel runs (n Learner workers, each one reporting a loss from their respective + data shard). + - Logging of images, videos, or other more complex data structures over time. + - Reducing these collected values using a user specified reduction method (for + example "min" or "mean") and other settings controlling the reduction and internal + data, such as sliding windows or EMA coefficients. + - Optionally clearing all logged values after a `reduce()` call to make space for + new data. + - Tracking throughputs of logged values. + + .. testcode:: + + import time + from ray.rllib.utils.metrics.metrics_logger import MetricsLogger + from ray.rllib.utils.test_utils import check + + logger = MetricsLogger(root=True) + + # 1) Logging float values (mean over window): + # Log some loss under the "loss" key. By default, all logged values + # under that key are averaged and reported back, once `reduce()` is called. + logger.log_value("loss", 0.001, reduce="mean", window=10) + logger.log_value("loss", 0.002, reduce="mean", window=10) + # Peek at the current (reduced) value of "loss": + check(logger.peek("loss"), 0.0015) # <- expect average value + # Actually reduce the underlying Stats object(s). + results = logger.reduce() + check(results["loss"], 0.0015) + + # 2) Logging float values (minimum over window): + # Log the minimum of loss values under the "min_loss" key. + logger.log_value("min_loss", 0.1, reduce="min", window=2) + logger.log_value("min_loss", 0.01, reduce="min", window=2) + logger.log_value("min_loss", 0.1, reduce="min", window=2) + logger.log_value("min_loss", 0.02, reduce="min", window=2) + # Peek at the current (reduced) value of "min_loss": + check(logger.peek("min_loss"), 0.02) # <- expect min value (over window=2) + # Actually reduce the underlying Stats object(s). + results = logger.reduce() + check(results["min_loss"], 0.02) + + # 3) Log n counts in different (remote?) components and merge them on the + # controller side. + remote_logger_1 = MetricsLogger() + remote_logger_2 = MetricsLogger() + main_logger = MetricsLogger() + remote_logger_1.log_value("count", 2, reduce="sum", clear_on_reduce=True) + remote_logger_2.log_value("count", 3, reduce="sum", clear_on_reduce=True) + # Reduce the two remote loggers .. + remote_results_1 = remote_logger_1.reduce() + remote_results_2 = remote_logger_2.reduce() + # .. then merge the two results into the controller logger. + main_logger.aggregate([remote_results_1, remote_results_2]) + check(main_logger.peek("count"), 5) + + # 4) Time blocks of code using EMA (coeff=0.1). Note that the higher the coeff + # (the closer to 1.0), the more short term the EMA turns out. + logger = MetricsLogger() + + # First delta measurement: + with logger.log_time("my_block_to_be_timed", reduce="mean", ema_coeff=0.1): + time.sleep(1.0) + # EMA should be ~1sec. + assert 1.1 > logger.peek("my_block_to_be_timed") > 0.9 + # Second delta measurement (note that we don't have to repeat the args again, as + # the stats under that name have already been created above with the correct + # args). + with logger.log_time("my_block_to_be_timed"): + time.sleep(2.0) + # EMA should be ~1.1sec. + assert 1.15 > logger.peek("my_block_to_be_timed") > 1.05 + + # When calling `reduce()`, the internal values list gets cleaned up (reduced) + # and reduction results are returned. + results = logger.reduce() + # EMA should be ~1.1sec. + assert 1.15 > results["my_block_to_be_timed"] > 1.05 + + # 5) Keeping track of throughputs and compiling all metrics + logger = MetricsLogger() + logger.log_value("samples", 1.0, reduce="sum", with_throughput=True, throughput_ema_coeff=1.0) + time.sleep(1.0) + logger.log_value("samples", 2.0, reduce="sum", with_throughput=True, throughput_ema_coeff=1.0) + results = logger.compile() + check(results["samples"], 3.0) + # Since we have an ema_coeff of 1.0, the throughput should be the same as the last value we logged (after 1 second) + check(results["samples_throughput"], 2.0, rtol=0.1) + + """ + + def __init__(self, root=False): + """Initializes a MetricsLogger instance. + + Args: + root: Whether this logger is a root logger. If True, lifetime stats + (clear_on_reduce=False and reduction="sum") will not be cleared on reduce(). + """ + self.stats = {} + self._tensor_mode = False + # TODO (sven): We use a dummy RLock here for most RLlib algos, however, APPO + # and IMPALA require this to be an actual RLock (b/c of thread safety reasons). + # An actual RLock, however, breaks our current OfflineData and + # OfflinePreLearner logic, in which the Learner (which contains a + # MetricsLogger) is serialized and deserialized. We will have to fix this + # offline RL logic first, then can remove this hack here and return to always + # using the RLock. + self._threading_lock = _DummyRLock() + # Is this a root logger? + self._is_root_logger = root + + def __contains__(self, key: Union[str, Tuple[str, ...]]) -> bool: + """Returns True, if `key` can be found in self.stats. + + Args: + key: The key to find in self.stats. This must be either a str (single, + top-level key) or a tuple of str (nested key). + + Returns: + Whether `key` could be found in self.stats. + """ + return self._key_in_stats(key) + + def peek( + self, + key: Union[str, Tuple[str, ...], None] = None, + default=None, + compile: bool = True, + throughput: bool = False, + ) -> Any: + """Returns the reduced values found in this MetricsLogger. + + Note that calling this method does NOT cause an actual underlying value list + reduction, even though reduced values are being returned. It'll keep all + internal structures as-is. By default, this returns a single reduced value or, if + the Stats object has no reduce method, a list of values. When when compile is False, + the result is a list of one or more values. + + Args: + key: The key/key sequence of the sub-structure of `self`, whose (reduced) + values to return. + default: An optional default value in case `key` cannot be found in `self`. + If default is not provided and `key` cannot be found, throws a KeyError. + compile: If True, the result is compiled into a single value if possible. + throughput: If True, the throughput is returned instead of the + actual (reduced) value. + + .. testcode:: + from ray.rllib.utils.metrics.metrics_logger import MetricsLogger + from ray.rllib.utils.test_utils import check + + logger = MetricsLogger() + ema = 0.01 + + # Log some (EMA reduced) values. + key = ("some", "nested", "key", "sequence") + logger.log_value(key, 2.0, ema_coeff=ema) + logger.log_value(key, 3.0) + + # Expected reduced value: + expected_reduced = (1.0 - ema) * 2.0 + ema * 3.0 + + # Peek at the (reduced) value under `key`. + check(logger.peek(key), expected_reduced) + + # Peek at the (reduced) nested struct under ("some", "nested"). + check( + logger.peek(("some", "nested")), + {"key": {"sequence": expected_reduced}}, + ) + + # Log some more, check again. + logger.log_value(key, 4.0) + expected_reduced = (1.0 - ema) * expected_reduced + ema * 4.0 + check(logger.peek(key=key), expected_reduced) + + Returns: + The (reduced) values of the (possibly nested) sub-structure found under + the given key or key sequence. + """ + if throughput: + return self._get_throughputs(key=key, default=default) + + # Create a reduced view of the entire stats structure. + def _nested_peek(stats): + return tree.map_structure( + # If the Stats object has a reduce method, we need to convert the list to a single value + lambda s: ( + s.peek(compile=compile) + if s._reduce_method is not None + else s.peek(compile=compile)[0] + ) + if isinstance(s, Stats) + else s, + stats.copy(), + ) + + with self._threading_lock: + if key is None: + return _nested_peek(self.stats) + else: + if default is None: + stats = self._get_key(key, key_error=True) + else: + stats = self._get_key(key, key_error=False) + + if isinstance(stats, Stats): + # If the Stats object has a reduce method, we need to convert the list to a single value + return stats.peek(compile=compile) + + elif isinstance(stats, dict) and stats: + return _nested_peek(stats) + else: + return default + + @staticmethod + def peek_results(results: Any, compile: bool = True) -> Any: + """Performs `peek()` on any leaf element of an arbitrarily nested Stats struct. + + Args: + results: The nested structure of Stats-leafs to be peek'd and returned. + compile: If True, the result is compiled into a single value if possible. + + Returns: + A corresponding structure of the peek'd `results` (reduced float/int values; + no Stats objects). + """ + return tree.map_structure( + lambda s: s.peek(compile=compile) if isinstance(s, Stats) else s, results + ) + + def log_value( + self, + key: Union[str, Tuple[str, ...]], + value: Any, + *, + reduce: Optional[str] = "mean", + window: Optional[Union[int, float]] = None, + ema_coeff: Optional[float] = None, + percentiles: Union[List[int], bool] = False, + clear_on_reduce: bool = False, + with_throughput: bool = False, + throughput_ema_coeff: Optional[float] = None, + reduce_per_index_on_aggregate: bool = False, + ) -> None: + """Logs a new value under a (possibly nested) key to the logger. + + .. testcode:: + + from ray.rllib.utils.metrics.metrics_logger import MetricsLogger + from ray.rllib.utils.test_utils import check + + logger = MetricsLogger(root=True) + + # Log n simple float values under the "loss" key. By default, all logged + # values under that key are averaged, once `reduce()` is called. + logger.log_value("loss", 0.01, window=10) + logger.log_value("loss", 0.02, window=10) + logger.log_value("loss", 0.03, window=10) + + # Peek at the current (reduced) value. + # Note that in the underlying structure, the internal values list still + # contains all logged values (0.01, 0.02, and 0.03). + check(logger.peek("loss"), 0.02) + + # Log 10x (window size) the same value. + for _ in range(10): + logger.log_value("loss", 0.05, window=10) + check(logger.peek("loss"), 0.05) + + # Internals check (note that users should not be concerned with accessing + # these). Len should always be 10, since the underlying struct is a + # `deque(max_len=10)`. + check(len(logger.stats["loss"].values), 10) + + # Only, when we call `reduce` does the underlying structure get "cleaned + # up". In this case, the list is shortened to 10 items (window size). + results = logger.reduce() + check(results, {"loss": 0.05}) + check(len(logger.stats["loss"].values), 10) + + # Log a value under a deeper nested key. + logger.log_value(("some", "nested", "key"), -1.0) + check(logger.peek(("some", "nested", "key")), -1.0) + + # Log n values without reducing them (we want to just collect some items). + logger.log_value("some_items", 5.0, reduce=None) + logger.log_value("some_items", 6.0, reduce=None) + logger.log_value("some_items", 7.0, reduce=None) + # Peeking at these returns the full list of items (no reduction set up). + check(logger.peek("some_items"), [5.0, 6.0, 7.0]) + # If you don't want the internal list to grow indefinitely, you should set + # `clear_on_reduce=True`: + logger.log_value("some_more_items", -5.0, reduce=None, clear_on_reduce=True) + logger.log_value("some_more_items", -6.0, reduce=None, clear_on_reduce=True) + logger.log_value("some_more_items", -7.0, reduce=None, clear_on_reduce=True ) + # Peeking at these returns the full list of items (no reduction set up). + check(logger.peek("some_more_items"), [-5.0, -6.0, -7.0]) + # Reducing everything (and return plain values, not `Stats` objects). + results = logger.reduce() + check(results, { + "loss": 0.05, + "some": { + "nested": { + "key": -1.0, + }, + }, + "some_items": [5.0, 6.0, 7.0], # reduce=None; list as-is + "some_more_items": [-5.0, -6.0, -7.0], # reduce=None; list as-is + }) + # However, the `reduce()` call did empty the `some_more_items` list + # (b/c we set `clear_on_reduce=True`). + check(logger.peek("some_more_items"), []) + # ... but not the "some_items" list (b/c `clear_on_reduce=False`). + check(logger.peek("some_items"), []) + + Args: + key: The key (or nested key-tuple) to log the `value` under. + value: The value to log. This should be a numeric value. + reduce: The reduction method to apply, once `self.reduce()` is called. + If None, will collect all logged values under `key` in a list (and + also return that list upon calling `self.reduce()`). + window: An optional window size to reduce over. + If not None, then the reduction operation is only applied to the most + recent `window` items, and - after reduction - the internal values list + under `key` is shortened to hold at most `window` items (the most + recent ones). + Must be None if `ema_coeff` is provided. + If None (and `ema_coeff` is None), reduction must not be "mean". + ema_coeff: An optional EMA coefficient to use if `reduce` is "mean" + and no `window` is provided. Note that if both `window` and `ema_coeff` + are provided, an error is thrown. Also, if `ema_coeff` is provided, + `reduce` must be "mean". + The reduction formula for EMA is: + EMA(t1) = (1.0 - ema_coeff) * EMA(t0) + ema_coeff * new_value + percentiles: If reduce is `None`, we can compute the percentiles of the + values list given by `percentiles`. Defaults to [0, 0.5, 0.75, 0.9, 0.95, + 0.99, 1] if set to True. When using percentiles, a window must be provided. + This window should be chosen carefully. RLlib computes exact percentiles and + the computational complexity is O(m*n*log(n/m)) where n is the window size + and m is the number of parallel metrics loggers involved (for example, + m EnvRunners). + clear_on_reduce: If True, all values under `key` will be emptied after + `self.reduce()` is called. Setting this to True is useful for cases, + in which the internal values list would otherwise grow indefinitely, + for example if reduce is None and there is no `window` provided. + with_throughput: Whether to track a throughput estimate together with this + metric. This is only supported for `reduce=sum` and + `clear_on_reduce=False` metrics (aka. "lifetime counts"). The `Stats` + object under the logged key then keeps track of the time passed + between two consecutive calls to `reduce()` and update its throughput + estimate. The current throughput estimate of a key can be obtained + through: .peek(key, throughput=True). + throughput_ema_coeff: The EMA coefficient to use for throughput tracking. + Only used if with_throughput=True. Defaults to 0.05 if with_throughput is True. + reduce_per_index_on_aggregate: If True, when merging Stats objects in parallel, we reduce + incoming values per index such that the new value at index `n` will be + the reduced value of all incoming values at index `n`. + If False, when reducing `n` Stats, the first `n` merged values will be + the reduced value of all incoming values at index `0`, the next `n` merged + values will be the reduced values of all incoming values at index `1`, etc. + """ + # No reduction (continue appending to list) AND no window. + # -> We'll force-reset our values upon `reduce()`. + if reduce is None and (window is None or window == float("inf")): + clear_on_reduce = True + + # Set default ema_coeff to 0.01 if reduce is "mean" and no window is provided + if reduce == "mean" and window is None and ema_coeff is None: + ema_coeff = 0.01 + + if with_throughput and throughput_ema_coeff is None: + throughput_ema_coeff = 0.05 + + value = self._detach_tensor_if_necessary(value) + + with self._threading_lock: + # `key` doesn't exist -> Automatically create it. + if not self._key_in_stats(key): + self._set_key( + key, + ( + Stats( + value, + reduce=reduce, + percentiles=percentiles, + window=window, + ema_coeff=ema_coeff, + clear_on_reduce=clear_on_reduce, + throughput=with_throughput, + throughput_ema_coeff=throughput_ema_coeff, + reduce_per_index_on_aggregate=reduce_per_index_on_aggregate, + ) + ), + ) + else: + stats = self._get_key(key) + if reduce != stats._reduce_method and log_once(f"reduce_warning_{key}"): + logger.warning( + f"reduce should be the same for all logged values under the same key, " + f"but got argument reduce={reduce} while the existing Stats object {key} " + f"has reduce={stats._reduce_method}." + ) + if clear_on_reduce != stats._clear_on_reduce and log_once( + f"clear_on_reduce_warning_{key}" + ): + logger.warning( + f"clear_on_reduce should be the same for all logged values under the same key, " + f"but got argument clear_on_reduce={clear_on_reduce} while the existing Stats object {key} " + f"has clear_on_reduce={stats._clear_on_reduce}." + ) + if with_throughput != bool(stats.has_throughput) and log_once( + f"with_throughput_warning_{key}" + ): + logger.warning( + f"with_throughput should be the same for all logged values under the same key, " + f"but got argument with_throughput={with_throughput} while the existing Stats object {key} " + f"has has_throughput={stats.has_throughput}. This warning will always be shown if you are using an older checkpoint." + ) + if throughput_ema_coeff != stats._throughput_ema_coeff and log_once( + f"throughput_ema_coeff_warning_{key}" + ): + logger.warning( + f"throughput_ema_coeff should be the same for all logged values under the same key, " + f"but got argument throughput_ema_coeff={throughput_ema_coeff} while the existing Stats object {key} " + f"has throughput_ema_coeff={stats._throughput_ema_coeff}. This warning will always be shown if you are using an older checkpoint." + ) + if window != stats._window and log_once(f"window_warning_{key}"): + logger.warning( + f"window should be the same for all logged values under the same key, " + f"but got argument window={window} while the existing Stats object {key} " + f"has window={stats._window}." + ) + if percentiles != getattr(stats, "_percentiles", False) and log_once( + f"percentiles_warning_{key}" + ): + logger.warning( + "percentiles should be the same for all logged values under the same key, " + f"but got argument percentiles={percentiles} while the existing Stats object {key} " + f"has percentiles={getattr(stats, '_percentiles', False)}." + ) + + if ( + reduce_per_index_on_aggregate + != stats._reduce_per_index_on_aggregate + and log_once(f"reduce_per_index_on_aggregate_warning_{key}") + ): + logger.warning( + f"reduce_per_index_on_aggregate should be the same for all logged values under the same key, " + f"but got argument reduce_per_index_on_aggregate={reduce_per_index_on_aggregate} while the existing Stats object {key} " + f"has reduce_per_index_on_aggregate={stats._reduce_per_index_on_aggregate}." + ) + + stats.push(value) + + def log_dict( + self, + value_dict, + *, + key: Optional[Union[str, Tuple[str, ...]]] = None, + reduce: Optional[str] = "mean", + window: Optional[Union[int, float]] = None, + ema_coeff: Optional[float] = None, + percentiles: Union[List[int], bool] = False, + clear_on_reduce: bool = False, + with_throughput: bool = False, + throughput_ema_coeff: Optional[float] = None, + reduce_per_index_on_aggregate: bool = False, + ) -> None: + """Logs all leafs of a possibly nested dict of values to this logger. + + To aggregate logs from upstream components, use `aggregate`. + + This is a convinience function that is equivalent to: + ``` + tree.map_structure_with_path(lambda path, value: logger.log_value(path, value, ...), value_dict) + ``` + + Traverses through all leafs of `stats_dict` and - if a path cannot be found in + this logger yet, will add the `Stats` found at the leaf under that new key. + If a path already exists, will merge the found leaf (`Stats`) with the ones + already logged before. This way, `stats_dict` does NOT have to have + the same structure as what has already been logged to `self`, but can be used to + log values under new keys or nested key paths. + + .. testcode:: + from ray.rllib.utils.metrics.metrics_logger import MetricsLogger + from ray.rllib.utils.test_utils import check + + logger = MetricsLogger() + + # Log n dicts with keys "a" and (some) "b". By default, all logged values + # under that key are averaged, once `reduce()` is called. + logger.log_dict( + { + "a": 0.1, + "b": -0.1, + }, + window=10, + ) + logger.log_dict({ + "b": -0.2, + }) # don't have to repeat `window` arg if key already exists + logger.log_dict({ + "a": 0.2, + "c": {"d": 5.0}, # can also introduce an entirely new (nested) key + }) + + # Peek at the current (reduced) values under "a" and "b". + check(logger.peek("a"), 0.15) + check(logger.peek("b"), -0.15) + check(logger.peek(("c", "d")), 5.0) + + # Reduced all stats. + results = logger.reduce() + check(results, { + "a": 0.15, + "b": -0.15, + "c": {"d": 5.0}, + }) + + Args: + value_dict: The (possibly nested) dict with individual values as + leafs to be logged to this logger. + key: An additional key (or tuple of keys) to prepend to all the keys + (or tuples of keys in case of nesting) found inside `stats_dict`. + Useful to log the entire contents of `stats_dict` in a more organized + fashion under one new key, for example logging the results returned by + an EnvRunner under key + reduce: The reduction method to apply, once `self.reduce()` is called. + If None, will collect all logged values under `key` in a list (and + also return that list upon calling `self.reduce()`). + window: An optional window size to reduce over. + If not None, then the reduction operation is only applied to the most + recent `window` items, and - after reduction - the internal values list + under `key` is shortened to hold at most `window` items (the most + recent ones). + Must be None if `ema_coeff` is provided. + If None (and `ema_coeff` is None), reduction must not be "mean". + ema_coeff: An optional EMA coefficient to use if `reduce` is "mean" + and no `window` is provided. Note that if both `window` and `ema_coeff` + are provided, an error is thrown. Also, if `ema_coeff` is provided, + `reduce` must be "mean". + The reduction formula for EMA is: + EMA(t1) = (1.0 - ema_coeff) * EMA(t0) + ema_coeff * new_value + percentiles: If reduce is `None`, we can compute the percentiles of the + values list given by `percentiles`. Defaults to [0, 0.5, 0.75, 0.9, 0.95, + 0.99, 1] if set to True. When using percentiles, a window must be provided. + This window should be chosen carefully. RLlib computes exact percentiles and + the computational complexity is O(m*n*log(n/m)) where n is the window size + and m is the number of parallel metrics loggers involved (for example, + m EnvRunners). + clear_on_reduce: If True, all values under `key` will be emptied after + `self.reduce()` is called. Setting this to True is useful for cases, + in which the internal values list would otherwise grow indefinitely, + for example if reduce is None and there is no `window` provided. + with_throughput: Whether to track a throughput estimate together with this + metric. This is only supported for `reduce=sum` and + `clear_on_reduce=False` metrics (aka. "lifetime counts"). The `Stats` + object under the logged key then keeps track of the time passed + between two consecutive calls to `reduce()` and update its throughput + estimate. The current throughput estimate of a key can be obtained + through: .peek(key, throughput=True). + throughput_ema_coeff: The EMA coefficient to use for throughput tracking. + Only used if with_throughput=True. Defaults to 0.05 if with_throughput is True. + reduce_per_index_on_aggregate: If True, when merging Stats objects, we reduce + incoming values per index such that the new value at index `n` will be + the reduced value of all incoming values at index `n`. + If False, when reducing `n` Stats, the first `n` merged values will be + the reduced value of all incoming values at index `0`, the next `n` merged + values will be the reduced values of all incoming values at index `1`, etc. + """ + assert isinstance( + value_dict, dict + ), f"`stats_dict` ({value_dict}) must be dict!" + + prefix_key = force_tuple(key) + + def _map(path, stat_or_value): + extended_key = prefix_key + force_tuple(tree.flatten(path)) + + if isinstance(stat_or_value, Stats): + deprecation_warning( + old="MetricsLogger.log_dict() for Stats objects", + new="MetricsLogger.aggregate()", + error=True, + ) + + self.log_value( + extended_key, + value=stat_or_value, + reduce=reduce, + window=window, + ema_coeff=ema_coeff, + percentiles=percentiles, + clear_on_reduce=clear_on_reduce, + with_throughput=with_throughput, + throughput_ema_coeff=throughput_ema_coeff, + reduce_per_index_on_aggregate=reduce_per_index_on_aggregate, + ) + + with self._threading_lock: + tree.map_structure_with_path(_map, value_dict) + + @Deprecated(new="aggregate", error=False) + def merge_and_log_n_dicts(self, *args, **kwargs): + return self.aggregate(*args, **kwargs) + + def aggregate( + self, + stats_dicts: List[Dict[str, Any]], + *, + key: Optional[Union[str, Tuple[str, ...]]] = None, + ) -> None: + """Merges n stats_dicts and logs result by merging on the time axis with existing stats. + + The n stats_dicts should be generated by n parallel components such that merging their + respective stats in parallel is meaningful. This lets us aggregate stats in a tree structure of + MetricsLoggers. + + If you want to log a dictionary of values (not Stats objects), use `log_dict`. + + .. testcode:: + + from ray.rllib.utils.metrics.metrics_logger import MetricsLogger + from ray.rllib.utils.test_utils import check + + # Example: n Learners logging loss stats to be merged. + # Note that losses should usually be logged with a window=1 so they don't + # get smeared over time and instead provide an accurate picture of the + # current situation. + main_logger = MetricsLogger() + + logger_learner1 = MetricsLogger() + logger_learner1.log_value("loss", 0.1, window=1) + learner1_results = logger_learner1.reduce() + + logger_learner2 = MetricsLogger() + logger_learner2.log_value("loss", 0.2, window=1) + learner2_results = logger_learner2.reduce() + + # Merge the stats from both Learners. + main_logger.aggregate( + [learner1_results, learner2_results], + key="learners", + ) + check(main_logger.peek(("learners", "loss")), 0.15) + + # Example: m EnvRunners logging episode returns to be merged. + main_logger = MetricsLogger() + + logger_env_runner1 = MetricsLogger() + logger_env_runner1.log_value("mean_ret", 100.0, window=3) + logger_env_runner1.log_value("mean_ret", 200.0, window=3) + logger_env_runner1.log_value("mean_ret", 300.0, window=3) + logger_env_runner1.log_value("mean_ret", 400.0, window=3) + env_runner1_results = logger_env_runner1.reduce() + + logger_env_runner2 = MetricsLogger() + logger_env_runner2.log_value("mean_ret", 150.0, window=3) + logger_env_runner2.log_value("mean_ret", 250.0, window=3) + logger_env_runner2.log_value("mean_ret", 350.0, window=3) + logger_env_runner2.log_value("mean_ret", 450.0, window=3) + env_runner2_results = logger_env_runner2.reduce() + + # Merge the stats from both EnvRunners. + main_logger.aggregate( + [env_runner1_results, env_runner2_results], + key="env_runners", + ) + # The expected procedure is as follows: + # The individual internal values lists of the two loggers are as follows: + # env runner 1: [200, 300, 400] + # env runner 2: [250, 350, 450] + # Move backwards from index=-1 (each time, loop through both env runners) + # index=-1 -> [400, 450] -> mean -> [425] -> repeat 2 times (number + # of env runners) -> [425, 425] + # index=-2 -> [300, 350] -> mean -> [325] -> repeat 2 times + # -> append -> [425, 425, 325, 325] -> STOP b/c we have reached >= window. + # reverse the list -> [325, 325, 425, 425] + # deque(max_len=3) -> [325, 425, 425] + check( + main_logger.stats["env_runners"]["mean_ret"].values, + [325, 425, 425], + ) + check(main_logger.peek(("env_runners", "mean_ret")), (325 + 425 + 425) / 3) + + # Example: Lifetime sum over n parallel components' stats. + main_logger = MetricsLogger() + + logger1 = MetricsLogger() + logger1.log_value("some_stat", 50, reduce="sum", window=None) + logger1.log_value("some_stat", 25, reduce="sum", window=None) + logger1_results = logger1.reduce() + + logger2 = MetricsLogger() + logger2.log_value("some_stat", 75, reduce="sum", window=None) + logger2_results = logger2.reduce() + + # Merge the stats from both Learners. + main_logger.aggregate([logger1_results, logger2_results]) + check(main_logger.peek("some_stat"), 150) + + # Example: Sum over n parallel components' stats with a window of 3. + main_logger = MetricsLogger() + + logger1 = MetricsLogger() + logger1.log_value("some_stat", 50, reduce="sum", window=3) + logger1.log_value("some_stat", 25, reduce="sum", window=3) + logger1.log_value("some_stat", 10, reduce="sum", window=3) + logger1.log_value("some_stat", 5, reduce="sum", window=3) + logger1_results = logger1.reduce() + + logger2 = MetricsLogger() + logger2.log_value("some_stat", 75, reduce="sum", window=3) + logger2.log_value("some_stat", 100, reduce="sum", window=3) + logger2_results = logger2.reduce() + + # Merge the stats from both Learners. + main_logger.aggregate([logger1_results, logger2_results]) + # The expected procedure is as follows: + # The individual internal values lists of the two loggers are as follows: + # env runner 1: [50, 25, 10, 5] + # env runner 2: [75, 100] + # Move backwards from index=-1 (each time, loop through both loggers) + # index=-1 -> [5, 100] -> reduce over both two indices -> [(5 + 100) / 2, (5 + 100) / 2] = [52.5, 52.5] + # Result = [52.5, 52.5] + # len() = 2 < window = 3 + # index=-2 -> [10, 75] -> reduce over both two indices -> [(10 + 75) / 2, (10 + 75) / 2] = [42.5, 42.5] + # result = [42.5, 42.5, 52.5, 52.5] + # len() = 4 >= window = 3 + check(main_logger.peek("some_stat"), 147.5) # last 3 items (window) get sum'd + + Args: + stats_dicts: List of n stats dicts to be merged and then logged. + key: Optional top-level key under which to log all keys/key sequences + found in the n `stats_dicts`. + """ + all_keys = set() + + def traverse_and_add_paths(d, path=()): + if isinstance(d, dict): + new_dict = {} + for key, value in d.items(): + new_dict[key] = traverse_and_add_paths(value, path + (key,)) + return new_dict + elif isinstance(d, list): + all_keys.add(path) + if len(d) == 1: + return d[0] + return d + else: + # For lists and values, we add the path to the set of all keys + all_keys.add(path) + return d + + def build_nested_dict(stats_dict, key): + if isinstance(key, str): + return {key: stats_dict} + elif len(key) > 1: + # Key is tuple of keys so we build a nested dict recursively + return {key[0]: build_nested_dict(stats_dict, key[1:])} + else: + return {key[0]: stats_dict} + + # We do one pass over all the stats_dicts_or_loggers to 1. prepend the key if provided and 2. collect all the keys that lead to leaves (which may be lists or values). + incoming_stats_dicts_with_key = [] + for stats_dict in stats_dicts: + if key is not None: + stats_dict = build_nested_dict(stats_dict, key) + stats_dict = traverse_and_add_paths(stats_dict) + incoming_stats_dicts_with_key.append(stats_dict) + + tree.map_structure_with_path( + lambda path, _: all_keys.add(force_tuple(path)), + self.stats, + ) + + for key in all_keys: + # Get all incoming Stats objects for this key + incoming_stats = [ + self._get_key(key, stats=s) + for s in incoming_stats_dicts_with_key + if self._key_in_stats(key, stats=s) + ] + + structure_under_key = self._get_key(key, stats=self.stats, key_error=False) + # self._get_key returns {} if the key is not found + own_stats = ( + None if isinstance(structure_under_key, dict) else structure_under_key + ) + + merged_stats = merge_stats( + base_stats=own_stats, incoming_stats=incoming_stats + ) + + self._set_key(key, merged_stats) + + def log_time( + self, + key: Union[str, Tuple[str, ...]], + *, + reduce: str = "mean", + window: Optional[Union[int, float]] = None, + ema_coeff: Optional[float] = None, + percentiles: Union[List[int], bool] = False, + clear_on_reduce: bool = False, + with_throughput: bool = False, + throughput_ema_coeff: float = 0.05, + reduce_per_index_on_aggregate: bool = False, + ) -> Stats: + """Measures and logs a time delta value under `key` when used with a with-block. + + .. testcode:: + + import time + from ray.rllib.utils.metrics.metrics_logger import MetricsLogger + + logger = MetricsLogger() + + # First delta measurement: + with logger.log_time("my_block_to_be_timed", ema_coeff=0.1): + time.sleep(1.0) + + # EMA should be ~1sec. + assert 1.1 > logger.peek("my_block_to_be_timed") > 0.9 + + # Second delta measurement (note that we don't have to repeat the args + # again, as the stats under that name have already been created above with + # the correct args). + with logger.log_time("my_block_to_be_timed"): + time.sleep(2.0) + + # EMA should be ~1.1sec. + assert 1.15 > logger.peek("my_block_to_be_timed") > 1.05 + + # When calling `reduce()`, the latest, reduced value is returned. + results = logger.reduce() + # EMA should be ~1.1sec. + assert 1.15 > results["my_block_to_be_timed"] > 1.05 + + Args: + key: The key (or tuple of keys) to log the measured time delta under. + reduce: The reduction method to apply, once `self.reduce()` is called. + If None, will collect all logged values under `key` in a list (and + also return that list upon calling `self.reduce()`). + window: An optional window size to reduce over. + If not None, then the reduction operation is only applied to the most + recent `window` items, and - after reduction - the internal values list + under `key` is shortened to hold at most `window` items (the most + recent ones). + Must be None if `ema_coeff` is provided. + If None (and `ema_coeff` is None), reduction must not be "mean". + ema_coeff: An optional EMA coefficient to use if `reduce` is "mean" + and no `window` is provided. Note that if both `window` and `ema_coeff` + are provided, an error is thrown. Also, if `ema_coeff` is provided, + `reduce` must be "mean". + The reduction formula for EMA is: + EMA(t1) = (1.0 - ema_coeff) * EMA(t0) + ema_coeff * new_value + percentiles: If reduce is `None`, we can compute the percentiles of the + values list given by `percentiles`. Defaults to [0, 0.5, 0.75, 0.9, 0.95, + 0.99, 1] if set to True. When using percentiles, a window must be provided. + This window should be chosen carefully. RLlib computes exact percentiles and + the computational complexity is O(m*n*log(n/m)) where n is the window size + and m is the number of parallel metrics loggers involved (for example, + m EnvRunners). + clear_on_reduce: If True, all values under `key` will be emptied after + `self.reduce()` is called. Setting this to True is useful for cases, + in which the internal values list would otherwise grow indefinitely, + for example if reduce is None and there is no `window` provided. + with_throughput: Whether to track a throughput estimate together with this + metric. This is only supported for `reduce=sum` and + `clear_on_reduce=False` metrics (aka. "lifetime counts"). The `Stats` + object under the logged key then keeps track of the time passed + between two consecutive calls to `reduce()` and update its throughput + estimate. The current throughput estimate of a key can be obtained + through: .peek(key, throughput=True). + throughput_ema_coeff: The EMA coefficient to use for throughput tracking. + Only used if with_throughput=True. Defaults to 0.05. + reduce_per_index_on_aggregate: If True, when merging Stats objects, we reduce + incoming values per index such that the new value at index `n` will be + the reduced value of all incoming values at index `n`. + If False, when reducing `n` Stats, the first `n` merged values will be + the reduced value of all incoming values at index `0`, the next `n` merged + values will be the reduced values of all incoming values at index `1`, etc. + """ + # No reduction (continue appending to list) AND no window. + # -> We'll force-reset our values upon `reduce()`. + if reduce is None and (window is None or window == float("inf")): + clear_on_reduce = True + + # Set default ema_coeff to 0.01 if reduce is "mean" and no window is provided + if reduce == "mean" and window is None and ema_coeff is None: + ema_coeff = 0.01 + + if not self._key_in_stats(key): + self._set_key( + key, + Stats( + init_values=None, + reduce=reduce, + percentiles=percentiles, + window=window, + ema_coeff=ema_coeff, + clear_on_reduce=clear_on_reduce, + throughput=with_throughput, + throughput_ema_coeff=throughput_ema_coeff, + reduce_per_index_on_aggregate=reduce_per_index_on_aggregate, + ), + ) + + # Return the Stats object, so a `with` clause can enter and exit it. + return self._get_key(key) + + def reduce(self) -> Dict: + """Reduces all logged values based on their settings and returns a result dict. + + DO NOT CALL THIS METHOD under normal circumstances! RLlib's components call it + right before a distinct step has been completed and the (MetricsLogger-based) + results of that step need to be passed upstream to other components for further + processing. + + The returned result dict has the exact same structure as the logged keys (or + nested key sequences) combined. Values are Stats objects if this MetricsLogger + is not a root logger. If this MetricsLogger is a root logger, the values are + the actual reduced values. + + For example, imagine component A (e.g. an Algorithm) containing a MetricsLogger + and n remote components (e.g. n EnvRunners), each with their own + MetricsLogger object. Component A calls its n remote components, each of + which returns an equivalent, reduced dict with values at the leafs. + Component A can then further log these n result dicts through its own + MetricsLogger through: + `logger.aggregate([n returned result dicts from n subcomponents])`. + + .. testcode:: + from ray.rllib.utils.metrics.metrics_logger import MetricsLogger + from ray.rllib.utils.test_utils import check + logger = MetricsLogger(root=True) + + # Log some values under different keys. + logger.log_value("loss", 0.1, window=2) + logger.log_value("loss", 0.2, window=2) + logger.log_value("min_loss", 0.3, reduce="min", window=2) + logger.log_value("min_loss", 0.1, reduce="min", window=2) + # reduce() returns the reduced values. + results = logger.reduce() + check(results["loss"], 0.15) # mean of [0.1, 0.2] + check(results["min_loss"], 0.1) # min of [0.3, 0.1] + # We can also reduce a specific key using indexing. + check(logger.reduce()["loss"], 0.15) # mean of [0.1, 0.2] + # Or reduce a nested key structure. + logger.log_value(("nested", "key"), 1.0) + check(logger.reduce()["nested"]["key"], 1.0) + + Returns: + A (nested) dict matching the structure of `self.stats` (contains all ever + logged keys to this MetricsLogger) with the leafs being (reduced) Stats + objects if this MetricsLogger is not a root logger. If this MetricsLogger + is a root logger, the leafs are the actual (reduced) values. + """ + # For better error message, catch the last key-path (reducing of which might + # throw an error). + PATH = None + + def _reduce(path, stats: Stats): + nonlocal PATH + PATH = path + return stats.reduce(compile=self._is_root_logger) + + try: + with self._threading_lock: + reduced_stats_to_return = tree.map_structure_with_path( + _reduce, self.stats + ) + # Provide proper error message if reduction fails due to bad data. + except Exception as e: + raise ValueError( + "There was an error while reducing the Stats object under key=" + f"{PATH}! Check, whether you logged invalid or incompatible " + "values into this key over time in your custom code." + f"\nThe values under this key are: {self._get_key(PATH).values}." + f"\nThe original error was {str(e)}" + ) + + return reduced_stats_to_return + + def activate_tensor_mode(self): + self._tensor_mode = True + + def deactivate_tensor_mode(self): + self._tensor_mode = False + + @property + def tensor_mode(self): + return self._tensor_mode + + def _detach_tensor_if_necessary(self, value): + if self.tensor_mode: + if torch and torch.is_tensor(value): + return value.detach() + elif tf and tf.is_tensor(value): + return tf.stop_gradient(value) + return value + + def set_value( + self, + key: Union[str, Tuple[str, ...]], + value: Any, + *, + reduce: Optional[str] = "mean", + window: Optional[Union[int, float]] = None, + ema_coeff: Optional[float] = None, + percentiles: Union[List[int], bool] = False, + clear_on_reduce: bool = False, + with_throughput: bool = False, + throughput_ema_coeff: float = 0.05, + reduce_per_index_on_aggregate: bool = False, + ) -> None: + """Overrides the logged values under `key` with `value`. + + The internal values list under `key` is cleared and reset to [`value`]. If + `key` already exists, this method will NOT alter the reduce settings. Otherwise, + it will apply the provided reduce settings (`reduce`, `window`, `ema_coeff`, + and `clear_on_reduce`). + + Args: + key: The key to override. + value: The new value to set the internal values list to (will be set to + a list containing a single item `value`). + reduce: The reduction method to apply, once `self.reduce()` is called. + If None, will collect all logged values under `key` in a list (and + also return that list upon calling `self.reduce()`). + Note that this is only applied if `key` does not exist in `self` yet. + window: An optional window size to reduce over. + If not None, then the reduction operation is only applied to the most + recent `window` items, and - after reduction - the internal values list + under `key` is shortened to hold at most `window` items (the most + recent ones). + Must be None if `ema_coeff` is provided. + If None (and `ema_coeff` is None), reduction must not be "mean". + Note that this is only applied if `key` does not exist in `self` yet. + ema_coeff: An optional EMA coefficient to use if `reduce` is "mean" + and no `window` is provided. Note that if both `window` and `ema_coeff` + are provided, an error is thrown. Also, if `ema_coeff` is provided, + `reduce` must be "mean". + The reduction formula for EMA is: + EMA(t1) = (1.0 - ema_coeff) * EMA(t0) + ema_coeff * new_value + Note that this is only applied if `key` does not exist in `self` yet. + percentiles: If reduce is `None`, we can compute the percentiles of the + values list given by `percentiles`. Defaults to [0, 0.5, 0.75, 0.9, 0.95, + 0.99, 1] if set to True. When using percentiles, a window must be provided. + This window should be chosen carefully. RLlib computes exact percentiles and + the computational complexity is O(m*n*log(n/m)) where n is the window size + and m is the number of parallel metrics loggers involved (for example, + m EnvRunners). + clear_on_reduce: If True, all values under `key` will be emptied after + `self.reduce()` is called. Setting this to True is useful for cases, + in which the internal values list would otherwise grow indefinitely, + for example if reduce is None and there is no `window` provided. + Note that this is only applied if `key` does not exist in `self` yet. + with_throughput: Whether to track a throughput estimate together with this + metric. This is only supported for `reduce=sum` and + `clear_on_reduce=False` metrics (aka. "lifetime counts"). The `Stats` + object under the logged key then keeps track of the time passed + between two consecutive calls to `reduce()` and update its throughput + estimate. The current throughput estimate of a key can be obtained + through: peeked_value, throuthput_per_sec = + .peek([key], throughput=True). + throughput_ema_coeff: The EMA coefficient to use for throughput tracking. + Only used if with_throughput=True. Defaults to 0.05. + reduce_per_index_on_aggregate: If True, when merging Stats objects, we reduce + incoming values per index such that the new value at index `n` will be + the reduced value of all incoming values at index `n`. + If False, when reducing `n` Stats, the first `n` merged values will be + the reduced value of all incoming values at index `0`, the next `n` merged + values will be the reduced values of all incoming values at index `1`, etc. + Note that this is only applied if `key` does not exist in `self` yet. + """ + # Key already in self -> Erase internal values list with [`value`]. + if self._key_in_stats(key): + with self._threading_lock: + stats = self._get_key(key) + stats.values = [value] + # Key cannot be found in `self` -> Simply log as a (new) value. + else: + # Set default ema_coeff to 0.01 if reduce is "mean" and no window is provided + if reduce == "mean" and window is None and ema_coeff is None: + ema_coeff = 0.01 + self.log_value( + key, + value, + reduce=reduce, + window=window, + ema_coeff=ema_coeff, + percentiles=percentiles, + clear_on_reduce=clear_on_reduce, + with_throughput=with_throughput, + throughput_ema_coeff=throughput_ema_coeff, + reduce_per_index_on_aggregate=reduce_per_index_on_aggregate, + ) + + def reset(self) -> None: + """Resets all data stored in this MetricsLogger.""" + with self._threading_lock: + self.stats = {} + + def delete(self, *key: Tuple[str, ...], key_error: bool = True) -> None: + """Deletes the given `key` from this metrics logger's stats. + + Args: + key: The key or key sequence (for nested location within self.stats), + to delete from this MetricsLogger's stats. + key_error: Whether to throw a KeyError if `key` cannot be found in `self`. + + Raises: + KeyError: If `key` cannot be found in `self` AND `key_error` is True. + """ + self._del_key(key, key_error) + + def get_state(self) -> Dict[str, Any]: + """Returns the current state of `self` as a dict. + + Note that the state is merely the combination of all states of the individual + `Stats` objects stored under `self.stats`. + """ + stats_dict = {} + + def _map(path, stats): + # Convert keys to strings for msgpack-friendliness. + stats_dict["--".join(path)] = stats.get_state() + + with self._threading_lock: + tree.map_structure_with_path(_map, self.stats) + + return {"stats": stats_dict} + + def set_state(self, state: Dict[str, Any]) -> None: + """Sets the state of `self` to the given `state`. + + Args: + state: The state to set `self` to. + """ + with self._threading_lock: + # Reset all existing stats to ensure a clean state transition + self.stats = {} + for flat_key, stats_state in state["stats"].items(): + self._set_key(flat_key.split("--"), Stats.from_state(stats_state)) + + def _key_in_stats(self, flat_key, *, stats=None): + flat_key = force_tuple(tree.flatten(flat_key)) + _dict = stats if stats is not None else self.stats + for key in flat_key: + if key not in _dict: + return False + _dict = _dict[key] + return True + + def _get_key(self, flat_key, *, stats=None, key_error=True): + flat_key = force_tuple(tree.flatten(flat_key)) + _dict = stats if stats is not None else self.stats + for key in flat_key: + try: + _dict = _dict[key] + except KeyError as e: + if key_error: + raise e + else: + return {} + return _dict + + def _set_key(self, flat_key, stats): + flat_key = force_tuple(tree.flatten(flat_key)) + + with self._threading_lock: + _dict = self.stats + for i, key in enumerate(flat_key): + # If we are at the end of the key sequence, set + # the key, no matter, whether it already exists or not. + if i == len(flat_key) - 1: + _dict[key] = stats + return + # If an intermediary key in the sequence is missing, + # add a sub-dict under this key. + if key not in _dict: + _dict[key] = {} + _dict = _dict[key] + + def _del_key(self, flat_key, key_error=False): + flat_key = force_tuple(tree.flatten(flat_key)) + + with self._threading_lock: + # Erase the key from the (nested) `self.stats` dict. + _dict = self.stats + try: + for i, key in enumerate(flat_key): + if i == len(flat_key) - 1: + del _dict[key] + return + _dict = _dict[key] + except KeyError as e: + if key_error: + raise e + + def _get_throughputs( + self, key: Optional[Union[str, Tuple[str, ...]]] = None, default=None + ) -> Union[Dict, float]: + """Returns throughput values for Stats that have throughput tracking enabled. + + If no key is provided, returns a nested dict containing throughput values for all Stats + that have throughput tracking enabled. If a key is provided, returns the throughput value + for that specific key or nested structure. + + The throughput values represent the rate of change of the corresponding metrics per second. + For example, if a metric represents the number of steps taken, its throughput value would + represent steps per second. + + Args: + key: Optional key or nested key path to get throughput for. If provided, returns just + the throughput value for that key or nested structure. If None, returns a nested dict + with throughputs for all metrics. + default: Default value to return if no throughput values are found. + Returns: + If key is None: A nested dict with the same structure as self.stats but with "_throughput" + appended to leaf keys and throughput values as leaf values. Only includes entries for + Stats objects that have throughput tracking enabled. + + If key is provided: The throughput value for that specific key or nested structure. + """ + + def _nested_throughputs(stats): + """Helper function to calculate throughputs for a nested structure.""" + + def _transform(path, value): + if isinstance(value, Stats) and value.has_throughput: + # Convert path to tuple for consistent key handling + key = force_tuple(path) + # Add "_throughput" to the last key in the path + return key[:-1] + (key[-1] + "_throughput",), value.throughput + return path, value + + result = {} + for path, value in tree.flatten_with_path(stats): + new_path, new_value = _transform(path, value) + if isinstance(new_value, float): # Only include throughput values + _dict = result + for k in new_path[:-1]: + if k not in _dict: + _dict[k] = {} + _dict = _dict[k] + _dict[new_path[-1]] = new_value + return result + + with self._threading_lock: + if key is not None: + # Get the Stats object or nested structure for the key + stats = self._get_key(key, key_error=False) + + if isinstance(stats, Stats): + if not stats.has_throughput: + raise ValueError( + f"Key '{key}' does not have throughput tracking enabled" + ) + return stats.throughput + elif stats == {}: + # If the key is not found, return the default value + return default + else: + # stats is a non-empty dictionary + return _nested_throughputs(stats) + + throughputs = {} + + def _map(path, stats): + if isinstance(stats, Stats) and stats.has_throughput: + # Convert path to tuple for consistent key handling + key = force_tuple(path) + # Add "_throughput" to the last key in the path + key = key[:-1] + (key[-1] + "_throughput",) + # Set the throughput value in the nested structure + _dict = throughputs + for k in key[:-1]: + if k not in _dict: + _dict[k] = {} + _dict = _dict[k] + _dict[key[-1]] = stats.throughput + + tree.map_structure_with_path(_map, self.stats) + + return throughputs if throughputs else default + + def compile(self) -> Dict: + """Compiles all current values and throughputs into a single dictionary. + + This method combines the results of all stats and throughputs into a single + dictionary, with throughput values having a "_throughput" suffix. This is useful + for getting a complete snapshot of all metrics and their throughputs in one call. + + Returns: + A nested dictionary containing both the current values and throughputs for all + metrics. The structure matches self.stats, with throughput values having + "_throughput" suffix in their keys. + """ + # Get all current values + values = self.reduce() + + # Get all throughputs + throughputs = self._get_throughputs() + + deep_update(values, throughputs or {}, new_keys_allowed=True) + + def traverse_dict(d): + if isinstance(d, dict): + new_dict = {} + for key, value in d.items(): + new_dict[key] = traverse_dict(value) + return new_dict + elif isinstance(d, list): + if len(d) == 1: + return d[0] + # If value is a longer list, we should just return the list because there is no reduction method applied + return d + else: + # If the value is not a list, it is a single value and we can yield it + return d + + return traverse_dict(values) + + +class _DummyRLock: + def acquire(self, blocking=True, timeout=-1): + return True + + def release(self): + pass + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + pass diff --git a/lib/python3.12/site-packages/ray/rllib/utils/metrics/ray_metrics.py b/lib/python3.12/site-packages/ray/rllib/utils/metrics/ray_metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..a66598327d5635f7546fcf0ad525fecdd8736c88 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/metrics/ray_metrics.py @@ -0,0 +1,59 @@ +import time + +from ray.util.metrics import Histogram + +_num_buckets = 31 +_coeff = 4 +_short_event_min = 0.0001 # 0.1 ms +_short_event_max = 1.5 +_long_event_min = 0.1 +_long_event_max = 600.0 + + +def _create_buckets(coeff, event_min, event_max, num): + """Generates a list of `num` buckets between `event_min` and `event_max`. + `coeff` - specifies how much denser at the low end + """ + if num == 1: + return [event_min] + step = 1 / (num - 1) + return [ + (0 + step * i) ** coeff * (event_max - event_min) + event_min + for i in range(num) + ] + + +DEFAULT_HISTOGRAM_BOUNDARIES_SHORT_EVENTS = _create_buckets( + coeff=_coeff, + event_min=_short_event_min, + event_max=_short_event_max, + num=_num_buckets, +) +DEFAULT_HISTOGRAM_BOUNDARIES_LONG_EVENTS = _create_buckets( + coeff=_coeff, + event_min=_long_event_min, + event_max=_long_event_max, + num=_num_buckets, +) + + +class TimerAndPrometheusLogger: + """Context manager for timing code execution. + + Elapsed time is automatically logged to the provided Prometheus Histogram. + + Example: + with TimerAndPrometheusLogger(Histogram): + learner.update() + """ + + def __init__(self, histogram: Histogram): + self._histogram = histogram + + def __enter__(self): + self.start = time.perf_counter() + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.elapsed = time.perf_counter() - self.start + self._histogram.observe(self.elapsed) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/metrics/stats.py b/lib/python3.12/site-packages/ray/rllib/utils/metrics/stats.py new file mode 100644 index 0000000000000000000000000000000000000000..3a7826ced2dc455af70deb3f9c6f5044b54ee9e8 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/metrics/stats.py @@ -0,0 +1,1039 @@ +import copy +import heapq +import threading +import time +import uuid +from collections import defaultdict, deque +from typing import Any, Dict, List, Optional, Tuple, Union + +import numpy as np + +from ray.rllib.utils import force_list +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.numpy import convert_to_numpy +from ray.util.annotations import DeveloperAPI + +torch, _ = try_import_torch() + + +@DeveloperAPI +class Stats: + """A container class holding a number of values and executing reductions over them. + + The individual values in a Stats object may be of any type, for example python int + or float, numpy arrays, or more complex structured (tuple, dict) and are stored in + a list under `self.values`. This class is not meant to be interfaced with directly + from application code. Instead, use `MetricsLogger` to log and manipulate Stats. + + Stats can be used to store metrics of the same type over time, for example a loss + or a learning rate, and to reduce all stored values applying a certain reduction + mechanism (for example "mean" or "sum"). + + Available reduction mechanisms are: + - "mean" using EMA with a configurable EMA coefficient. + - "mean" using a sliding window (over the last n stored values). + - "max/min" with an optional sliding window (over the last n stored values). + - "sum" with an optional sliding window (over the last n stored values). + - None: Simply store all logged values to an ever-growing list. + + Through the `reduce()` API, one of the above-mentioned reduction mechanisms will + be executed on `self.values`. + """ + + def __init__( + self, + init_values: Optional[Any] = None, + reduce: Optional[str] = "mean", + percentiles: Union[List[int], bool] = False, + reduce_per_index_on_aggregate: bool = False, + window: Optional[Union[int, float]] = None, + ema_coeff: Optional[float] = None, + clear_on_reduce: bool = False, + throughput: Union[bool, float] = False, + throughput_ema_coeff: Optional[float] = None, + ): + """Initializes a Stats instance. + + Args: + init_values: Optional initial values to be placed into `self.values`. If None, + `self.values` will start empty. If percentiles is True, values must be ordered + if provided. + reduce: The name of the reduce method to be used. Allowed are "mean", "min", + "max", and "sum". Use None to apply no reduction method (leave + `self.values` as-is when reducing, except for shortening it to + `window`). Note that if both `reduce` and `window` are None, the user of + this Stats object needs to apply some caution over the values list not + growing infinitely. + percentiles: If reduce is `None`, we can compute the percentiles of the + values list given by `percentiles`. Defaults to [0, 50, 75, 90, 95, + 99, 100] if set to True. When using percentiles, a window must be provided. + This window should be chosen carfully. RLlib computes exact percentiles and + the computational complexity is O(m*n*log(n/m)) where n is the window size + and m is the number of parallel metrics loggers invovled (for example, + m EnvRunners). To be safe, choose a window < 1M and less than 1000 Stats + objects to aggregate. See #52963 for more details. + window: An optional window size to reduce over. + If `window` is not None, then the reduction operation is only applied to + the most recent `windows` items, and - after reduction - the values list + is shortened to hold at most `window` items (the most recent ones). + Must be None if `ema_coeff` is not None. + If `window` is None (and `ema_coeff` is None), reduction must not be + "mean". + reduce_per_index_on_aggregate: If True, when merging Stats objects, we reduce + incoming values per index such that the new value at index `n` will be + the reduced value of all incoming values at index `n`. + If False, when reducing `n` Stats, the first `n` merged values will be + the reduced value of all incoming values at index `0`, the next `n` merged + values will be the reduced values of all incoming values at index `1`, etc. + ema_coeff: An optional EMA coefficient to use if reduce is "mean" + and no `window` is provided. Note that if both `window` and `ema_coeff` + are provided, an error is thrown. Also, if `ema_coeff` is provided, + `reduce` must be "mean". + The reduction formula for EMA performed by Stats is: + EMA(t1) = (1.0 - ema_coeff) * EMA(t0) + ema_coeff * new_value + clear_on_reduce: If True, the Stats object will reset its entire values list + to an empty one after `self.reduce()` is called. However, it will then + return from the `self.reduce()` call a new Stats object with the + properly reduced (not completely emptied) new values. Setting this + to True is useful for cases, in which the internal values list would + otherwise grow indefinitely, for example if reduce is None and there + is no `window` provided. + throughput: If True, track a throughput estimate together with this + Stats. This is only supported for `reduce=sum` and + `clear_on_reduce=False` metrics (aka. "lifetime counts"). The `Stats` + then keeps track of the time passed between two consecutive calls to + `reduce()` and update its throughput estimate. The current throughput + estimate can be obtained through: + `throughput_per_sec = Stats.peek(throughput=True)`. + If a float, track throughput and also set current throughput estimate + to the given value. + throughput_ema_coeff: An optional EMA coefficient to use for throughput tracking. + Only used if throughput=True. + """ + # Thus far, we only support mean, max, min, and sum. + if reduce not in [None, "mean", "min", "max", "sum", "percentiles"]: + raise ValueError( + "`reduce` must be one of `mean|min|max|sum|percentiles` or None!" + ) + + # One or both window and ema_coeff must be None. + if window is not None and ema_coeff is not None: + raise ValueError("Only one of `window` or `ema_coeff` can be specified!") + # If `ema_coeff` is provided, `reduce` must be "mean". + if ema_coeff is not None and reduce != "mean": + raise ValueError( + "`ema_coeff` arg only allowed (not None) when `reduce=mean`!" + ) + + if percentiles is not False: + if reduce is not None: + raise ValueError( + "`reduce` must be `None` when `percentiles` is not `False`!" + ) + if window in (None, float("inf")): + raise ValueError( + "A window must be specified when reduce is 'percentiles'!" + ) + if reduce_per_index_on_aggregate is not False: + raise ValueError( + f"`reduce_per_index_on_aggregate` ({reduce_per_index_on_aggregate})" + f" must be `False` when `percentiles` is not `False`!" + ) + + if percentiles is True: + percentiles = [0, 50, 75, 90, 95, 99, 100] + else: + if type(percentiles) not in (bool, list): + raise ValueError("`percentiles` must be a list or bool!") + if isinstance(percentiles, list): + if not all(isinstance(p, (int, float)) for p in percentiles): + raise ValueError( + "`percentiles` must contain only ints or floats!" + ) + if not all(0 <= p <= 100 for p in percentiles): + raise ValueError( + "`percentiles` must contain only values between 0 and 100!" + ) + + self._percentiles = percentiles + + # If `window` is explicitly set to inf, `clear_on_reduce` must be True. + self._inf_window = window in [None, float("inf")] + + # If `window` is set to inf, `clear_on_reduce` must be True. + # Otherwise, we risk a memory leak. + if self._inf_window and not clear_on_reduce and reduce is None: + raise ValueError( + "When using an infinite window without reduction, `clear_on_reduce` must " + "be set to True!" + ) + + # If reduce=mean AND window=ema_coeff=None, we use EMA by default with a coeff + # of 0.01 (we do NOT support infinite window sizes for mean as that would mean + # to keep data in the cache forever). + if reduce == "mean" and self._inf_window and ema_coeff is None: + ema_coeff = 0.01 + + self._reduce_method = reduce + self._window = window + self._ema_coeff = ema_coeff + + if ( + self._reduce_method not in ["mean", "sum", "min", "max"] + and reduce_per_index_on_aggregate + ): + raise ValueError( + "reduce_per_index_on_aggregate is only supported for mean, sum, min, and max reduction!" + ) + + self._reduce_per_index_on_aggregate = reduce_per_index_on_aggregate + + # Timing functionality (keep start times per thread). + self._start_times = defaultdict(lambda: None) + + # Simply store ths flag for the user of this class. + self._clear_on_reduce = clear_on_reduce + + self._has_returned_zero = False + + # On each `.reduce()` call, we store the result of this call in + # self._last_reduce. + self._last_reduced = [np.nan] + # The ID of this Stats instance. + self.id_ = str(uuid.uuid4()) + self._prev_merge_values = defaultdict(int) + + self._throughput_ema_coeff = throughput_ema_coeff + self._throughput_stats = None + if throughput is not False: + self._throughput_stats = Stats( + # We have to check for bool here because in Python, bool is a subclass + # of int. + init_values=[throughput] + if ( + isinstance(throughput, (int, float)) + and not isinstance(throughput, bool) + ) + else None, + reduce="mean", + ema_coeff=throughput_ema_coeff, + window=None, + clear_on_reduce=False, + throughput=False, + throughput_ema_coeff=None, + ) + if init_values is not None: + self._last_throughput_measure_time = time.perf_counter() + else: + self._last_throughput_measure_time = ( + -1 + ) # Track last push time for throughput calculation + + # The actual, underlying data in this Stats object. + self.values: Union[List, deque.Deque] = None + self._set_values(force_list(init_values)) + + self._is_tensor = False + + # Track if new values were pushed since last reduce + if init_values is not None: + self._has_new_values = True + else: + self._has_new_values = False + + def check_value(self, value: Any) -> None: + # If we have a reduce method, value should always be a scalar + # If we don't reduce, we can keep track of value as it is + if self._reduce_method is not None: + if isinstance(value, np.ndarray) and value.shape == (): + return + elif torch and torch.is_tensor(value): + self._is_tensor = True + if tuple(value.shape) == (): + return + elif type(value) not in (list, tuple, deque): + return + raise ValueError( + f"Value ({value}) is required to be a scalar when using a reduce " + "method!" + ) + + def push(self, value: Any) -> None: + """Pushes a value into this Stats object. + + Args: + value: The value to be pushed. Can be of any type. + """ + self.check_value(value) + # If throughput tracking is enabled, calculate it based on time between pushes + if self.has_throughput: + self._recompute_throughput(value) + # Handle different reduction methods + if self._window is not None: + # For windowed operations, append to values and trim if needed + self.values.append(value) + if len(self.values) > self._window: + self.values.popleft() + else: + # For non-windowed operations, use _reduced_values + if len(self.values) == 0: + self._set_values([value]) + else: + self.values.append(value) + _, values = self._reduced_values() + self._set_values(values) + + # Mark that we have new values + self._has_new_values = True + + def __enter__(self) -> "Stats": + """Called when entering a context (with which users can measure a time delta). + + Returns: + This Stats instance (self), unless another thread has already entered (and + not exited yet), in which case a copy of `self` is returned. This way, the + second thread(s) cannot mess with the original Stat's (self) time-measuring. + This also means that only the first thread to __enter__ actually logs into + `self` and the following threads' measurements are discarded (logged into + a non-referenced shim-Stats object, which will simply be garbage collected). + """ + # In case another thread already is measuring this Stats (timing), simply ignore + # the "enter request" and return a clone of `self`. + thread_id = threading.get_ident() + self._start_times[thread_id] = time.perf_counter() + return self + + def __exit__(self, exc_type, exc_value, tb) -> None: + """Called when exiting a context (with which users can measure a time delta).""" + thread_id = threading.get_ident() + assert self._start_times[thread_id] is not None + time_delta_s = time.perf_counter() - self._start_times[thread_id] + self.push(time_delta_s) + + del self._start_times[thread_id] + + def peek(self, compile: bool = True) -> Union[Any, List[Any]]: + """Returns the result of reducing the internal values list. + + Note that this method does NOT alter the internal values list in this process. + Thus, users can call this method to get an accurate look at the reduced value(s) + given the current internal values list. + + Args: + compile: If True, the result is compiled into a single value if possible. + + Returns: + The result of reducing the internal values list. + """ + if self._has_new_values or (not compile and not self._inf_window): + reduced_value, reduced_values = self._reduced_values() + if not compile and not self._inf_window: + return reduced_values + if compile and self._reduce_method: + return reduced_value[0] + if compile and self._percentiles is not False: + return compute_percentiles(reduced_values, self._percentiles) + return reduced_value + else: + return_value = self._last_reduced + if compile: + # We don't need to check for self._reduce_method or percentiles here + # because we only store the reduced value if there is a reduce method. + return_value = return_value[0] + return return_value + + @property + def throughput(self) -> float: + """Returns the current throughput estimate per second. + + Raises: + ValueError: If throughput tracking is not enabled for this Stats object. + + Returns: + The current throughput estimate per second. + """ + if not self.has_throughput: + raise ValueError("Throughput tracking is not enabled for this Stats object") + # We can always return the first value here because throughput is a single value + return self._throughput_stats.peek() + + @property + def has_throughput(self) -> bool: + """Returns whether this Stats object tracks throughput. + + Returns: + True if this Stats object has throughput tracking enabled, False otherwise. + """ + return self._throughput_stats is not None + + def reduce(self, compile: bool = True) -> Union[Any, List[Any]]: + """Reduces the internal values list according to the constructor settings. + + Thereby, the internal values list is changed (note that this is different from + `peek()`, where the internal list is NOT changed). See the docstring of this + class for details on the reduction logic applied to the values list, based on + the constructor settings, such as `window`, `reduce`, etc.. + + Args: + compile: If True, the result is compiled into a single value if possible. + If it is not possible, the result is a list of values. + If False, the result is a list of one or more values. + + Returns: + The reduced value (can be of any type, depending on the input values and + reduction method). + """ + len_before_reduce = len(self) + if self._has_new_values: + # Only calculate and update history if there were new values pushed since + # last reduce + reduced, reduced_internal_values_list = self._reduced_values() + # `clear_on_reduce` -> Clear the values list. + if self._clear_on_reduce: + self._set_values([]) + else: + self._set_values(reduced_internal_values_list) + else: + reduced_internal_values_list = None + reduced = self._last_reduced + + reduced = self._numpy_if_necessary(reduced) + + # Shift historic reduced valued by one in our reduce_history. + if self._reduce_method is not None: + # It only makes sense to extend the history if we are reducing to a single + # value. We need to make a copy here because the new_values_list is a + # reference to the internal values list + self._last_reduced = force_list(reduced.copy()) + else: + # If there is a window and no reduce method, we don't want to use the reduce + # history to return reduced values in other methods + self._has_new_values = True + + if compile and self._reduce_method is not None: + assert ( + len(reduced) == 1 + ), f"Reduced values list must contain exactly one value, found {reduced}" + reduced = reduced[0] + + if not compile and not self._inf_window: + if reduced_internal_values_list is None: + _, reduced_internal_values_list = self._reduced_values() + return_values = self._numpy_if_necessary( + reduced_internal_values_list + ).copy() + elif compile and self._percentiles is not False: + if reduced_internal_values_list is None: + _, reduced_internal_values_list = self._reduced_values() + return_values = compute_percentiles( + reduced_internal_values_list, self._percentiles + ) + else: + return_values = reduced + + if compile: + return return_values + else: + if len_before_reduce == 0: + # return_values will be be 0 if we reduce a sum over zero elements + # But we don't want to create such a zero out of nothing for our new + # Stats object that we return here + return Stats.similar_to(self) + + return Stats.similar_to(self, init_values=return_values) + + def merge_on_time_axis(self, other: "Stats") -> None: + """Merges another Stats object's values into this one along the time axis. + + Args: + other: The other Stats object to merge values from. + """ + self.values.extend(other.values) + + # Mark that we have new values since we modified the values list + self._has_new_values = True + + def merge_in_parallel(self, *others: "Stats") -> None: + """Merges all internal values of `others` into `self`'s internal values list. + + Thereby, the newly incoming values of `others` are treated equally with respect + to each other as well as with respect to the internal values of self. + + Use this method to merge other `Stats` objects, which resulted from some + parallelly executed components, into this one. For example: n Learner workers + all returning a loss value in the form of `{"total_loss": [some value]}`. + + The following examples demonstrate the parallel merging logic for different + reduce- and window settings: + + Args: + others: One or more other Stats objects that need to be parallely merged + into `self, meaning with equal weighting as the existing values in + `self`. + """ + win = self._window or float("inf") + + # If any of the value lists have a length of 0 or if there is only one value and + # it is nan, we skip + stats_to_merge = [ + s + for s in [self, *others] + if not ( + len(s) == 0 + or ( + len(s) == 1 and np.all(np.isnan(self._numpy_if_necessary(s.values))) + ) + ) + ] + + # If there is only one stat to merge, and it is the same as self, return. + if len(stats_to_merge) == 0: + # If none of the stats have values, return. + return + elif len(stats_to_merge) == 1: + if stats_to_merge[0] == self: + # If no incoming stats have values, return. + return + else: + # If there is only one stat with values, and it's incoming, copy its + # values. + self.values = stats_to_merge[0].values + return + + # Take turns stepping through `self` and `*others` values, thereby moving + # backwards from last index to beginning and will up the resulting values list. + # Stop as soon as we reach the window size. + new_values = [] + tmp_values = [] + + if self._percentiles is not False: + # Use heapq to sort values (assumes that the values are already sorted) + # and then pick the correct percentiles + lists_to_merge = [list(self.values), *[list(o.values) for o in others]] + merged = list(heapq.merge(*lists_to_merge)) + self._set_values(merged) + else: + # Loop from index=-1 backward to index=start until our new_values list has + # at least a len of `win`. + for i in range(1, max(map(len, stats_to_merge)) + 1): + # Per index, loop through all involved stats, including `self` and add + # to `tmp_values`. + for stats in stats_to_merge: + if len(stats) < i: + continue + tmp_values.append(stats.values[-i]) + + # Now reduce across `tmp_values` based on the reduce-settings of this + # Stats. + if self._reduce_per_index_on_aggregate: + n_values = 1 + else: + n_values = len(tmp_values) + + if self._ema_coeff is not None: + new_values.extend([np.nanmean(tmp_values)] * n_values) + elif self._reduce_method is None: + new_values.extend(tmp_values) + elif self._reduce_method == "sum": + # We add [sum(tmp_values) / n_values] * n_values to the new values + # list instead of tmp_values, because every incoming element should + # have the same weight. + added_sum = self._reduced_values(values=tmp_values)[0][0] + new_values.extend([added_sum / n_values] * n_values) + if self.has_throughput: + self._recompute_throughput(added_sum) + else: + new_values.extend( + self._reduced_values(values=tmp_values)[0] * n_values + ) + + tmp_values.clear() + if len(new_values) >= win: + new_values = new_values[:win] + break + + self._set_values(list(reversed(new_values))) + + # Mark that we have new values since we modified the values list + self._has_new_values = True + + def clear_throughput(self) -> None: + """Clears the throughput Stats, if applicable and `self` has throughput. + + Also resets `self._last_throughput_measure_time` to -1 such that the Stats + object has to create a new timestamp first, before measuring any new throughput + values. + """ + if self.has_throughput: + self._throughput_stats._set_values([]) + self._last_throughput_measure_time = -1 + + def _recompute_throughput(self, value) -> None: + """Recomputes the current throughput value of this Stats instance.""" + # Make sure this Stats object does measure throughput. + assert self.has_throughput + # Take the current time stamp. + current_time = time.perf_counter() + # Check, whether we have a previous timestamp (non -1). + if self._last_throughput_measure_time >= 0: + # Compute the time delta. + time_diff = current_time - self._last_throughput_measure_time + # Avoid divisions by zero. + if time_diff > 0: + # Push new throughput value into our throughput stats object. + self._throughput_stats.push(value / time_diff) + # Update the time stamp of the most recent throughput computation (this one). + self._last_throughput_measure_time = current_time + + @staticmethod + def _numpy_if_necessary(values): + # Torch tensor handling. Convert to CPU/numpy first. + if torch and len(values) > 0 and torch.is_tensor(values[0]): + # Convert all tensors to numpy values. + values = [v.cpu().numpy() for v in values] + return values + + def __len__(self) -> int: + """Returns the length of the internal values list.""" + return len(self.values) + + def __repr__(self) -> str: + win_or_ema = ( + f"; win={self._window}" + if self._window + else f"; ema={self._ema_coeff}" + if self._ema_coeff + else "" + ) + return ( + f"Stats({self.peek()}; len={len(self)}; " + f"reduce={self._reduce_method}{win_or_ema})" + ) + + def __int__(self): + if self._reduce_method is None: + raise ValueError( + "Cannot convert Stats object with reduce method `None` to int because " + "it can not be reduced to a single value." + ) + else: + return int(self.peek()) + + def __float__(self): + if self._reduce_method is None: + raise ValueError( + "Cannot convert Stats object with reduce method `None` to float " + "because it can not be reduced to a single value." + ) + else: + return float(self.peek()) + + def __eq__(self, other): + if self._reduce_method is None: + self._comp_error("__eq__") + else: + return float(self) == float(other) + + def __le__(self, other): + if self._reduce_method is None: + self._comp_error("__le__") + else: + return float(self) <= float(other) + + def __ge__(self, other): + if self._reduce_method is None: + self._comp_error("__ge__") + else: + return float(self) >= float(other) + + def __lt__(self, other): + if self._reduce_method is None: + self._comp_error("__lt__") + else: + return float(self) < float(other) + + def __gt__(self, other): + if self._reduce_method is None: + self._comp_error("__gt__") + else: + return float(self) > float(other) + + def __add__(self, other): + if self._reduce_method is None: + self._comp_error("__add__") + else: + return float(self) + float(other) + + def __sub__(self, other): + if self._reduce_method is None: + self._comp_error("__sub__") + else: + return float(self) - float(other) + + def __mul__(self, other): + if self._reduce_method is None: + self._comp_error("__mul__") + else: + return float(self) * float(other) + + def __format__(self, fmt): + if self._reduce_method is None: + raise ValueError( + "Cannot format Stats object with reduce method `None` because it can " + "not be reduced to a single value." + ) + else: + return f"{float(self):{fmt}}" + + def _comp_error(self, comp): + raise ValueError( + f"Cannot {comp} Stats object with reduce method `None` to other " + "because it can not be reduced to a single value." + ) + + def get_state(self) -> Dict[str, Any]: + state = { + # Make sure we don't return any tensors here. + "values": convert_to_numpy(self.values), + "reduce": self._reduce_method, + "percentiles": self._percentiles, + "reduce_per_index_on_aggregate": self._reduce_per_index_on_aggregate, + "window": self._window, + "ema_coeff": self._ema_coeff, + "clear_on_reduce": self._clear_on_reduce, + "_last_reduced": self._last_reduced, + "_is_tensor": self._is_tensor, + } + if self._throughput_stats is not None: + state["throughput_stats"] = self._throughput_stats.get_state() + return state + + @staticmethod + def from_state(state: Dict[str, Any]) -> "Stats": + # If `values` could contain tensors, don't reinstate them (b/c we don't know + # whether we are on a supported device). + values = state["values"] + if "_is_tensor" in state and state["_is_tensor"]: + values = [] + + if "throughput_stats" in state: + throughput_stats = Stats.from_state(state["throughput_stats"]) + stats = Stats( + values, + reduce=state["reduce"], + percentiles=state.get("percentiles", False), + reduce_per_index_on_aggregate=state.get( + "reduce_per_index_on_aggregate", False + ), + window=state["window"], + ema_coeff=state["ema_coeff"], + clear_on_reduce=state["clear_on_reduce"], + throughput=throughput_stats.peek(), + throughput_ema_coeff=throughput_stats._ema_coeff, + ) + elif state.get("_throughput", False): + # Older checkpoints have a _throughput key that is boolean or + # a float (throughput value). They don't have a throughput_ema_coeff + # so we use a default of 0.05. + # TODO(Artur): Remove this after a few Ray releases. + stats = Stats( + values, + reduce=state["reduce"], + percentiles=state.get("percentiles", False), + window=state["window"], + ema_coeff=state["ema_coeff"], + clear_on_reduce=state["clear_on_reduce"], + throughput=state["_throughput"], + throughput_ema_coeff=0.05, + ) + else: + stats = Stats( + values, + reduce=state["reduce"], + percentiles=state.get("percentiles", False), + window=state["window"], + ema_coeff=state["ema_coeff"], + clear_on_reduce=state["clear_on_reduce"], + throughput=False, + throughput_ema_coeff=None, + ) + # Compatibility to old checkpoints where a reduce sometimes resulted in a single + # values instead of a list such that the history would be a list of integers + # instead of a list of lists. + if "_hist" in state: + # TODO(Artur): Remove this after a few Ray releases. + if not isinstance(state["_hist"][0], list): + state["_hist"] = list(map(lambda x: [x], state["_hist"])) + stats._last_reduced = state["_hist"][-1] + else: + stats._last_reduced = state.get("_last_reduced", [np.nan]) + return stats + + @staticmethod + def similar_to( + other: "Stats", + init_values: Optional[Any] = None, + ) -> "Stats": + """Returns a new Stats object that's similar to `other`. + + "Similar" here means it has the exact same settings (reduce, window, ema_coeff, + etc..). The initial values of the returned `Stats` are empty by default, but + can be set as well. + + Args: + other: The other Stats object to return a similar new Stats equivalent for. + init_value: The initial value to already push into the returned Stats. + + Returns: + A new Stats object similar to `other`, with the exact same settings and + maybe a custom initial value (if provided; otherwise empty). + """ + stats = Stats( + init_values=init_values, + reduce=other._reduce_method, + percentiles=other._percentiles, + reduce_per_index_on_aggregate=other._reduce_per_index_on_aggregate, + window=other._window, + ema_coeff=other._ema_coeff, + clear_on_reduce=other._clear_on_reduce, + throughput=other._throughput_stats.peek() + if other.has_throughput + else False, + throughput_ema_coeff=other._throughput_ema_coeff, + ) + stats.id_ = other.id_ + stats._last_reduced = other._last_reduced + return stats + + def _set_values(self, new_values): + # For stats with window, use a deque with maxlen=window. + # This way, we never store more values than absolutely necessary. + if not self._inf_window: + self.values = deque(new_values, maxlen=self._window) + # For infinite windows, use `new_values` as-is (a list). + else: + self.values = new_values + + self._has_new_values = True + + def _reduced_values(self, values=None) -> Tuple[Any, Any]: + """Runs a non-committed reduction procedure on given values (or `self.values`). + + Note that this method does NOT alter any state of `self` or the possibly + provided list of `values`. It only returns new values as they should be + adopted after a possible, actual reduction step. + + Args: + values: The list of values to reduce. If not None, use `self.values` + + Returns: + A tuple containing 1) the reduced values and 2) the new internal values list + to be used. If there is no reduciton method, the reduced values will be the same as the values. + """ + values = values if values is not None else self.values + + # No reduction method. Return list as-is OR reduce list to len=window. + if self._reduce_method is None: + if self._percentiles is not False: + # Sort values + values = list(values) + # (Artur): Numpy can sort faster than Python's built-in sort for large lists. Howoever, if we convert to an array here + # and then sort, this only slightly (<2x) improved the runtime of this method, even for an internal values list of 1M values. + values.sort() + return values, values + + # Special case: Internal values list is empty -> return NaN or 0.0 for sum. + elif len(values) == 0: + if self._reduce_method in ["min", "max", "mean"] or self._has_returned_zero: + # We also return np.nan if we have returned zero before. + # This helps with cases where stats are cleared on reduce, but we don't want to log 0's, except for the first time. + return [np.nan], [] + else: + return [0], [] + + # Do EMA (always a "mean" reduction; possibly using a window). + elif self._ema_coeff is not None: + # Perform EMA reduction over all values in internal values list. + mean_value = values[0] + for v in values[1:]: + mean_value = self._ema_coeff * v + (1.0 - self._ema_coeff) * mean_value + if self._inf_window: + return [mean_value], [mean_value] + else: + return [mean_value], values + # Non-EMA reduction (possibly using a window). + else: + # Use the numpy/torch "nan"-prefix to ignore NaN's in our value lists. + if torch and torch.is_tensor(values[0]): + self._is_tensor = True + # Only one item in the + if len(values[0].shape) == 0: + reduced = values[0] + else: + reduce_meth = getattr(torch, "nan" + self._reduce_method) + reduce_in = torch.stack(list(values)) + if self._reduce_method == "mean": + reduce_in = reduce_in.float() + reduced = reduce_meth(reduce_in) + else: + reduce_meth = getattr(np, "nan" + self._reduce_method) + + if np.all(np.isnan(values)): + # This avoids warnings for taking a mean of an empty array. + reduced = np.nan + else: + reduced = reduce_meth(values) + + def safe_isnan(value): + if torch and isinstance(value, torch.Tensor): + return torch.isnan(value) + return np.isnan(value) + + # Convert from numpy to primitive python types, if original `values` are + # python types. + if ( + not safe_isnan(reduced) + and reduced.shape == () + and isinstance(values[0], (int, float)) + ): + if reduced.dtype in [np.int32, np.int64, np.int8, np.int16]: + reduced = int(reduced) + else: + reduced = float(reduced) + + # For window=None|inf (infinite window) and reduce != mean, we don't have to + # keep any values, except the last (reduced) one. + if self._inf_window and self._reduce_method != "mean": + # TODO (sven): What if values are torch tensors? In this case, we + # would have to do reduction using `torch` above (not numpy) and only + # then return the python primitive AND put the reduced new torch + # tensor in the new `self.values`. + return [reduced], [reduced] + else: + # In all other cases, keep the values that were also used for the reduce + # operation. + return [reduced], values + + +@DeveloperAPI +def compute_percentiles(sorted_list, percentiles): + """Compute percentiles from an already sorted list. + + Note that this will not raise an error if the list is not sorted to avoid overhead. + + Args: + sorted_list: A list of numbers sorted in ascending order + percentiles: A list of percentile values (0-100) + + Returns: + A dictionary mapping percentile values to their corresponding data values + """ + n = len(sorted_list) + + if n == 0: + return {p: None for p in percentiles} + + results = {} + + for p in percentiles: + index = (p / 100) * (n - 1) + + if index.is_integer(): + results[p] = sorted_list[int(index)] + else: + lower_index = int(index) + upper_index = lower_index + 1 + weight = index - lower_index + results[p] = ( + sorted_list[lower_index] * (1 - weight) + + sorted_list[upper_index] * weight + ) + + return results + + +@DeveloperAPI +def merge_stats(base_stats: Optional[Stats], incoming_stats: List[Stats]) -> Stats: + """Merges Stats objects. + + If `base_stats` is None, we use the first incoming Stats object as the new base Stats object. + If `base_stats` is not None, we merge all incoming Stats objects into the base Stats object. + + Args: + base_stats: The base Stats object to merge into. + incoming_stats: The list of Stats objects to merge. + + Returns: + The merged Stats object. + """ + if base_stats is None: + new_root_stats = True + else: + new_root_stats = False + # Nothing to be merged + if len(incoming_stats) == 0: + return base_stats + + if new_root_stats: + # We need to deepcopy here first because stats from incoming_stats may be altered in the future + base_stats = copy.deepcopy(incoming_stats[0]) + base_stats.clear_throughput() + # Note that we may take a mean of means here, which is not the same as a + # mean of all values. In the future, we could implement a weighted mean + # of means here by introducing a new Stats object that counts samples + # for each mean Stats object. + if len(incoming_stats) > 1: + base_stats.merge_in_parallel(*incoming_stats[1:]) + if ( + base_stats._reduce_method == "sum" + and base_stats._inf_window + and base_stats._clear_on_reduce is False + ): + for stat in incoming_stats: + base_stats._prev_merge_values[stat.id_] = stat.peek() + + elif len(incoming_stats) > 0: + # Special case: `base_stats` is a lifetime sum (reduce=sum, + # clear_on_reduce=False) -> We subtract the previous value (from 2 + # `reduce()` calls ago) from all to-be-merged stats, so we don't count + # twice the older sum from before. + + # Also, for the merged, new throughput value, we need to find out what the + # actual value-delta is between before the last reduce and the current one. + + added_sum = 0.0 # Used in `base_stats._recompute_throughput` if applicable. + if ( + base_stats._reduce_method == "sum" + and base_stats._inf_window + and base_stats._clear_on_reduce is False + ): + for stat in incoming_stats: + # Subtract "lifetime counts" from the Stat's values to not count + # older "lifetime counts" more than once. + prev_reduction = base_stats._prev_merge_values[stat.id_] + new_reduction = stat.peek(compile=True) + base_stats.values[-1] -= prev_reduction + # Keep track of how many counts we actually gained (for throughput + # recomputation). + added_sum += new_reduction - prev_reduction + base_stats._prev_merge_values[stat.id_] = new_reduction + + parallel_merged_stat = copy.deepcopy(incoming_stats[0]) + if len(incoming_stats) > 1: + # There are more than one incoming parallel others -> Merge all of + # them in parallel (equal importance). + parallel_merged_stat.merge_in_parallel(*incoming_stats[1:]) + + # Merge incoming Stats object into base Stats object on time axis + # (giving incoming ones priority). + if base_stats._reduce_method == "mean" and not base_stats._clear_on_reduce: + # If we don't clear values, values that are not cleared would contribute + # to the mean multiple times. + base_stats._set_values(parallel_merged_stat.values.copy()) + else: + base_stats.merge_on_time_axis(parallel_merged_stat) + # Keep track of throughput through the sum of added counts. + if base_stats.has_throughput: + base_stats._recompute_throughput(added_sum) + + return base_stats diff --git a/lib/python3.12/site-packages/ray/rllib/utils/metrics/utils.py b/lib/python3.12/site-packages/ray/rllib/utils/metrics/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..f14fa8e1e12aa623b34642d73626a25509f9379d --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/metrics/utils.py @@ -0,0 +1,20 @@ +import re + + +def to_snake_case(class_name: str) -> str: + """Converts class name to snake case. + + This is used to unify metrics names when using class names within. + Args: + class_name: A string defining a class name usually in camel + case. + + Returns: + The class name in snake case. + """ + # Insert _ between a lower- or digit-char and an upper-char + name = re.sub(r"(?<=[a-z0-9])(?=[A-Z])", "_", class_name) + # Insert _ between an upper-char followed by upper+lower (to split + # "ABCReward" into "abc_reward") + name = re.sub(r"(?<=[A-Z])(?=[A-Z][a-z])", "_", name) + return name.lower() diff --git a/lib/python3.12/site-packages/ray/rllib/utils/metrics/window_stat.py b/lib/python3.12/site-packages/ray/rllib/utils/metrics/window_stat.py new file mode 100644 index 0000000000000000000000000000000000000000..1ea4fe35c956d83ab5cb26a174e2c25af62d148d --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/metrics/window_stat.py @@ -0,0 +1,79 @@ +import numpy as np + +from ray.rllib.utils.annotations import OldAPIStack + + +@OldAPIStack +class WindowStat: + """Handles/stores incoming dataset and provides window-based statistics. + + .. testcode:: + :skipif: True + + win_stats = WindowStat("level", 3) + win_stats.push(5.0) + win_stats.push(7.0) + win_stats.push(7.0) + win_stats.push(10.0) + # Expect 8.0 as the mean of the last 3 values: (7+7+10)/3=8.0 + print(win_stats.mean()) + + .. testoutput:: + + 8.0 + """ + + def __init__(self, name: str, n: int): + """Initializes a WindowStat instance. + + Args: + name: The name of the stats to collect and return stats for. + n: The window size. Statistics will be computed for the last n + items received from the stream. + """ + # The window-size. + self.window_size = n + # The name of the data (used for `self.stats()`). + self.name = name + # List of items to do calculations over (len=self.n). + self.items = [None] * self.window_size + # The current index to insert the next item into `self.items`. + self.idx = 0 + # How many items have been added over the lifetime of this object. + self.count = 0 + + def push(self, obj) -> None: + """Pushes a new value/object into the data buffer.""" + # Insert object at current index. + self.items[self.idx] = obj + # Increase insertion index by 1. + self.idx += 1 + # Increase lifetime count by 1. + self.count += 1 + # Fix index in case of rollover. + self.idx %= len(self.items) + + def mean(self) -> float: + """Returns the (NaN-)mean of the last `self.window_size` items.""" + return float(np.nanmean(self.items[: self.count])) + + def std(self) -> float: + """Returns the (NaN)-stddev of the last `self.window_size` items.""" + return float(np.nanstd(self.items[: self.count])) + + def quantiles(self) -> np.ndarray: + """Returns ndarray with 0, 10, 50, 90, and 100 percentiles.""" + if not self.count: + return np.ndarray([], dtype=np.float32) + else: + return np.nanpercentile( + self.items[: self.count], [0, 10, 50, 90, 100] + ).tolist() + + def stats(self): + return { + self.name + "_count": int(self.count), + self.name + "_mean": self.mean(), + self.name + "_std": self.std(), + self.name + "_quantiles": self.quantiles(), + } diff --git a/lib/python3.12/site-packages/ray/rllib/utils/minibatch_utils.py b/lib/python3.12/site-packages/ray/rllib/utils/minibatch_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..764e67c73a9508767e471878a3ba613c2c387d18 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/minibatch_utils.py @@ -0,0 +1,391 @@ +import math +from typing import List, Optional + +from ray.data import DataIterator +from ray.rllib.policy.sample_batch import MultiAgentBatch, SampleBatch, concat_samples +from ray.rllib.utils import unflatten_dict +from ray.rllib.utils.annotations import DeveloperAPI +from ray.rllib.utils.typing import DeviceType, EpisodeType + + +@DeveloperAPI +class MiniBatchIteratorBase: + """The base class for all minibatch iterators.""" + + def __init__( + self, + batch: MultiAgentBatch, + *, + num_epochs: int = 1, + shuffle_batch_per_epoch: bool = True, + minibatch_size: int, + num_total_minibatches: int = 0, + ) -> None: + """Initializes a MiniBatchIteratorBase instance. + + Args: + batch: The input multi-agent batch. + num_epochs: The number of complete passes over the entire train batch. Each + pass might be further split into n minibatches (if `minibatch_size` + provided). The train batch is generated from the given `episodes` + through the Learner connector pipeline. + minibatch_size: The size of minibatches to use to further split the train + batch into per epoch. The train batch is generated from the given + `episodes` through the Learner connector pipeline. + num_total_minibatches: The total number of minibatches to loop through + (over all `num_epochs` epochs). It's only required to set this to != 0 + in multi-agent + multi-GPU situations, in which the MultiAgentEpisodes + themselves are roughly sharded equally, however, they might contain + SingleAgentEpisodes with very lopsided length distributions. Thus, + without this fixed, pre-computed value, one Learner might go through a + different number of minibatche passes than others causing a deadlock. + """ + pass + + +@DeveloperAPI +class MiniBatchCyclicIterator(MiniBatchIteratorBase): + """This implements a simple multi-agent minibatch iterator. + + This iterator will split the input multi-agent batch into minibatches where the + size of batch for each module_id (aka policy_id) is equal to minibatch_size. If the + input batch is smaller than minibatch_size, then the iterator will cycle through + the batch until it has covered `num_epochs` epochs. + """ + + def __init__( + self, + batch: MultiAgentBatch, + *, + num_epochs: int = 1, + minibatch_size: int, + shuffle_batch_per_epoch: bool = True, + num_total_minibatches: int = 0, + ) -> None: + """Initializes a MiniBatchCyclicIterator instance.""" + super().__init__( + batch, + num_epochs=num_epochs, + minibatch_size=minibatch_size, + shuffle_batch_per_epoch=shuffle_batch_per_epoch, + ) + + self._batch = batch + self._minibatch_size = minibatch_size + self._num_epochs = num_epochs + self._shuffle_batch_per_epoch = shuffle_batch_per_epoch + + # mapping from module_id to the start index of the batch + self._start = {mid: 0 for mid in batch.policy_batches.keys()} + # mapping from module_id to the number of epochs covered for each module_id + self._num_covered_epochs = {mid: 0 for mid in batch.policy_batches.keys()} + + self._minibatch_count = 0 + self._num_total_minibatches = num_total_minibatches + + def __iter__(self): + while ( + # Make sure each item in the total batch gets at least iterated over + # `self._num_epochs` times. + ( + self._num_total_minibatches == 0 + and min(self._num_covered_epochs.values()) < self._num_epochs + ) + # Make sure we reach at least the given minimum number of mini-batches. + or ( + self._num_total_minibatches > 0 + and self._minibatch_count < self._num_total_minibatches + ) + ): + minibatch = {} + for module_id, module_batch in self._batch.policy_batches.items(): + + if len(module_batch) == 0: + raise ValueError( + f"The batch for module_id {module_id} is empty! " + "This will create an infinite loop because we need to cover " + "the same number of samples for each module_id." + ) + s = self._start[module_id] # start + + # TODO (sven): Fix this bug for LSTMs: + # In an RNN-setting, the Learner connector already has zero-padded + # and added a timerank to the batch. Thus, n_step would still be based + # on the BxT dimension, rather than the new B dimension (excluding T), + # which then leads to minibatches way too large. + # However, changing this already would break APPO/IMPALA w/o LSTMs as + # these setups require sequencing, BUT their batches are not yet time- + # ranked (this is done only in their loss functions via the + # `make_time_major` utility). + n_steps = self._minibatch_size + + samples_to_concat = [] + + # get_len is a function that returns the length of a batch + # if we are not slicing the batch in the batch dimension B, then + # the length of the batch is simply the length of the batch + # o.w the length of the batch is the length list of seq_lens. + if module_batch._slice_seq_lens_in_B: + assert module_batch.get(SampleBatch.SEQ_LENS) is not None, ( + "MiniBatchCyclicIterator requires SampleBatch.SEQ_LENS" + "to be present in the batch for slicing a batch in the batch " + "dimension B." + ) + + def get_len(b): + return len(b[SampleBatch.SEQ_LENS]) + + n_steps = int( + get_len(module_batch) + * (self._minibatch_size / len(module_batch)) + ) + + else: + + def get_len(b): + return len(b) + + # Cycle through the batch until we have enough samples. + while s + n_steps >= get_len(module_batch): + sample = module_batch[s:] + samples_to_concat.append(sample) + len_sample = get_len(sample) + assert len_sample > 0, "Length of a sample must be > 0!" + n_steps -= len_sample + s = 0 + self._num_covered_epochs[module_id] += 1 + # Shuffle the individual single-agent batch, if required. + # This should happen once per minibatch iteration in order to make + # each iteration go through a different set of minibatches. + if self._shuffle_batch_per_epoch: + module_batch.shuffle() + + e = s + n_steps # end + if e > s: + samples_to_concat.append(module_batch[s:e]) + + # concatenate all the samples, we should have minibatch_size of sample + # after this step + minibatch[module_id] = concat_samples(samples_to_concat) + # roll minibatch to zero when we reach the end of the batch + self._start[module_id] = e + + # Note (Kourosh): env_steps is the total number of env_steps that this + # multi-agent batch is covering. It should be simply inherited from the + # original multi-agent batch. + minibatch = MultiAgentBatch(minibatch, len(self._batch)) + yield minibatch + + self._minibatch_count += 1 + + +class MiniBatchDummyIterator(MiniBatchIteratorBase): + def __init__(self, batch: MultiAgentBatch, **kwargs): + super().__init__(batch, **kwargs) + self._batch = batch + + def __iter__(self): + yield self._batch + + +@DeveloperAPI +class MiniBatchRayDataIterator: + def __init__( + self, + *, + iterator: DataIterator, + device: DeviceType, + minibatch_size: int, + num_iters: Optional[int], + **kwargs, + ): + # A `ray.data.DataIterator` that can iterate in different ways over the data. + self._iterator = iterator + # Note, in multi-learner settings the `return_state` is in `kwargs`. + self._kwargs = {k: v for k, v in kwargs.items() if k != "return_state"} + + # Holds a batched_iterable over the dataset. + self._batched_iterable = self._iterator.iter_torch_batches( + batch_size=minibatch_size, + device=device, + **self._kwargs, + ) + # Create an iterator that can be stopped and resumed during an epoch. + self._epoch_iterator = iter(self._batched_iterable) + self._num_iters = num_iters + + def __iter__(self) -> MultiAgentBatch: + iteration = 0 + while self._num_iters is None or iteration < self._num_iters: + for batch in self._epoch_iterator: + # Update the iteration counter. + iteration += 1 + + batch = unflatten_dict(batch) + batch = MultiAgentBatch( + { + module_id: SampleBatch(module_data) + for module_id, module_data in batch.items() + }, + env_steps=sum( + len(next(iter(module_data.values()))) + for module_data in batch.values() + ), + ) + + yield (batch) + + # If `num_iters` is reached break and return. + if self._num_iters and iteration == self._num_iters: + break + else: + # Reinstantiate a new epoch iterator. + self._epoch_iterator = iter(self._batched_iterable) + # If a full epoch on the data should be run, stop. + if not self._num_iters: + # Exit the loop. + break + + +@DeveloperAPI +class ShardBatchIterator: + """Iterator for sharding batch into num_shards batches. + + Args: + batch: The input multi-agent batch. + num_shards: The number of shards to split the batch into. + + Yields: + A MultiAgentBatch of size len(batch) / num_shards. + """ + + def __init__(self, batch: MultiAgentBatch, num_shards: int): + self._batch = batch + self._num_shards = num_shards + + def __iter__(self): + for i in range(self._num_shards): + # TODO (sven): The following way of sharding a multi-agent batch destroys + # the relationship of the different agents' timesteps to each other. + # Thus, in case the algorithm requires agent-synchronized data (aka. + # "lockstep"), the `ShardBatchIterator` cannot be used. + batch_to_send = {} + for pid, sub_batch in self._batch.policy_batches.items(): + batch_size = math.ceil(len(sub_batch) / self._num_shards) + start = batch_size * i + end = min(start + batch_size, len(sub_batch)) + batch_to_send[pid] = sub_batch[int(start) : int(end)] + # TODO (Avnish): int(batch_size) ? How should we shard MA batches really? + new_batch = MultiAgentBatch(batch_to_send, int(batch_size)) + yield new_batch + + +@DeveloperAPI +class ShardEpisodesIterator: + """Iterator for sharding a list of Episodes into `num_shards` lists of Episodes.""" + + def __init__( + self, + episodes: List[EpisodeType], + num_shards: int, + len_lookback_buffer: Optional[int] = None, + ): + """Initializes a ShardEpisodesIterator instance. + + Args: + episodes: The input list of Episodes. + num_shards: The number of shards to split the episodes into. + len_lookback_buffer: An optional length of a lookback buffer to enforce + on the returned shards. When spitting an episode, the second piece + might need a lookback buffer (into the first piece) depending on the + user's settings. + """ + self._episodes = sorted(episodes, key=len, reverse=True) + self._num_shards = num_shards + self._len_lookback_buffer = len_lookback_buffer + self._total_length = sum(len(e) for e in episodes) + self._target_lengths = [0 for _ in range(self._num_shards)] + remaining_length = self._total_length + for s in range(self._num_shards): + len_ = remaining_length // (num_shards - s) + self._target_lengths[s] = len_ + remaining_length -= len_ + + def __iter__(self) -> List[EpisodeType]: + """Runs one iteration through this sharder. + + Yields: + A sub-list of Episodes of size roughly `len(episodes) / num_shards`. The + yielded sublists might have slightly different total sums of episode + lengths, in order to not have to drop even a single timestep. + """ + sublists = [[] for _ in range(self._num_shards)] + lengths = [0 for _ in range(self._num_shards)] + episode_index = 0 + + while episode_index < len(self._episodes): + episode = self._episodes[episode_index] + min_index = lengths.index(min(lengths)) + + # Add the whole episode if it fits within the target length + if lengths[min_index] + len(episode) <= self._target_lengths[min_index]: + sublists[min_index].append(episode) + lengths[min_index] += len(episode) + episode_index += 1 + # Otherwise, slice the episode + else: + remaining_length = self._target_lengths[min_index] - lengths[min_index] + if remaining_length > 0: + slice_part, remaining_part = ( + # Note that the first slice will automatically "inherit" the + # lookback buffer size of the episode. + episode[:remaining_length], + # However, the second slice might need a user defined lookback + # buffer (into the first slice). + episode.slice( + slice(remaining_length, None), + len_lookback_buffer=self._len_lookback_buffer, + ), + ) + sublists[min_index].append(slice_part) + lengths[min_index] += len(slice_part) + self._episodes[episode_index] = remaining_part + else: + assert remaining_length == 0 + sublists[min_index].append(episode) + episode_index += 1 + + for sublist in sublists: + yield sublist + + +@DeveloperAPI +class ShardObjectRefIterator: + """Iterator for sharding a list of ray ObjectRefs into num_shards sub-lists. + + Args: + object_refs: The input list of ray ObjectRefs. + num_shards: The number of shards to split the references into. + + Yields: + A sub-list of ray ObjectRefs with lengths as equal as possible. + """ + + def __init__(self, object_refs, num_shards: int): + self._object_refs = object_refs + self._num_shards = num_shards + + def __iter__(self): + # Calculate the size of each sublist + n = len(self._object_refs) + sublist_size = n // self._num_shards + remaining_elements = n % self._num_shards + + start = 0 + for i in range(self._num_shards): + # Determine the end index for the current sublist + end = start + sublist_size + (1 if i < remaining_elements else 0) + # Append the sublist to the result + yield self._object_refs[start:end] + # Update the start index for the next sublist + start = end diff --git a/lib/python3.12/site-packages/ray/rllib/utils/numpy.py b/lib/python3.12/site-packages/ray/rllib/utils/numpy.py new file mode 100644 index 0000000000000000000000000000000000000000..5e97a5a5c9629615d3345b4e1cbd615449bfec61 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/numpy.py @@ -0,0 +1,605 @@ +from collections import OrderedDict +from types import MappingProxyType +from typing import List, Optional + +import numpy as np +import tree # pip install dm_tree +from gymnasium.spaces import Discrete, MultiDiscrete + +from ray._common.deprecation import Deprecated +from ray.rllib.utils.annotations import PublicAPI +from ray.rllib.utils.framework import try_import_tf, try_import_torch +from ray.rllib.utils.typing import SpaceStruct, TensorStructType, TensorType, Union + +tf1, tf, tfv = try_import_tf() +torch, _ = try_import_torch() + +SMALL_NUMBER = 1e-6 +# Some large int number. May be increased here, if needed. +LARGE_INTEGER = 100000000 +# Min and Max outputs (clipped) from an NN-output layer interpreted as the +# log(x) of some x (e.g. a stddev of a normal +# distribution). +MIN_LOG_NN_OUTPUT = -5 +MAX_LOG_NN_OUTPUT = 2 + + +@PublicAPI +@Deprecated( + help="RLlib itself has no use for this anymore.", + error=False, +) +def aligned_array(size: int, dtype, align: int = 64) -> np.ndarray: + """Returns an array of a given size that is 64-byte aligned. + + The returned array can be efficiently copied into GPU memory by TensorFlow. + + Args: + size: The size (total number of items) of the array. For example, + array([[0.0, 1.0], [2.0, 3.0]]) would have size=4. + dtype: The numpy dtype of the array. + align: The alignment to use. + + Returns: + A np.ndarray with the given specifications. + """ + n = size * dtype.itemsize + empty = np.empty(n + (align - 1), dtype=np.uint8) + data_align = empty.ctypes.data % align + offset = 0 if data_align == 0 else (align - data_align) + if n == 0: + # stop np from optimising out empty slice reference + output = empty[offset : offset + 1][0:0].view(dtype) + else: + output = empty[offset : offset + n].view(dtype) + + assert len(output) == size, len(output) + assert output.ctypes.data % align == 0, output.ctypes.data + return output + + +@PublicAPI +@Deprecated( + help="RLlib itself has no use for this anymore.", + error=False, +) +def concat_aligned( + items: List[np.ndarray], time_major: Optional[bool] = None +) -> np.ndarray: + """Concatenate arrays, ensuring the output is 64-byte aligned. + + We only align float arrays; other arrays are concatenated as normal. + + This should be used instead of np.concatenate() to improve performance + when the output array is likely to be fed into TensorFlow. + + Args: + items: The list of items to concatenate and align. + time_major: Whether the data in items is time-major, in which + case, we will concatenate along axis=1. + + Returns: + The concat'd and aligned array. + """ + + if len(items) == 0: + return [] + elif len(items) == 1: + # we assume the input is aligned. In any case, it doesn't help + # performance to force align it since that incurs a needless copy. + return items[0] + elif isinstance(items[0], np.ndarray) and items[0].dtype in [ + np.float32, + np.float64, + np.uint8, + ]: + dtype = items[0].dtype + flat = aligned_array(sum(s.size for s in items), dtype) + if time_major is not None: + if time_major is True: + batch_dim = sum(s.shape[1] for s in items) + new_shape = (items[0].shape[0], batch_dim,) + items[ + 0 + ].shape[2:] + else: + batch_dim = sum(s.shape[0] for s in items) + new_shape = (batch_dim, items[0].shape[1],) + items[ + 0 + ].shape[2:] + else: + batch_dim = sum(s.shape[0] for s in items) + new_shape = (batch_dim,) + items[0].shape[1:] + output = flat.reshape(new_shape) + assert output.ctypes.data % 64 == 0, output.ctypes.data + np.concatenate(items, out=output, axis=1 if time_major else 0) + return output + else: + return np.concatenate(items, axis=1 if time_major else 0) + + +@PublicAPI +def convert_to_numpy(x: TensorStructType, reduce_type: bool = True) -> TensorStructType: + """Converts values in `stats` to non-Tensor numpy or python types. + + Args: + x: Any (possibly nested) struct, the values in which will be + converted and returned as a new struct with all torch/tf tensors + being converted to numpy types. + reduce_type: Whether to automatically reduce all float64 and int64 data + into float32 and int32 data, respectively. + + Returns: + A new struct with the same structure as `x`, but with all + values converted to numpy arrays (on CPU). + """ + + # The mapping function used to numpyize torch/tf Tensors (and move them + # to the CPU beforehand). + def mapping(item): + if torch and isinstance(item, torch.Tensor): + ret = ( + item.cpu().item() + if len(item.size()) == 0 + else item.detach().cpu().numpy() + ) + elif ( + tf and isinstance(item, (tf.Tensor, tf.Variable)) and hasattr(item, "numpy") + ): + assert tf.executing_eagerly() + ret = item.numpy() + else: + ret = item + if reduce_type and isinstance(ret, np.ndarray): + if np.issubdtype(ret.dtype, np.floating): + ret = ret.astype(np.float32) + elif np.issubdtype(ret.dtype, int): + ret = ret.astype(np.int32) + return ret + + return tree.map_structure(mapping, x) + + +@PublicAPI +def fc( + x: np.ndarray, + weights: np.ndarray, + biases: Optional[np.ndarray] = None, + framework: Optional[str] = None, +) -> np.ndarray: + """Calculates FC (dense) layer outputs given weights/biases and input. + + Args: + x: The input to the dense layer. + weights: The weights matrix. + biases: The biases vector. All 0s if None. + framework: An optional framework hint (to figure out, + e.g. whether to transpose torch weight matrices). + + Returns: + The dense layer's output. + """ + + def map_(data, transpose=False): + if torch: + if isinstance(data, torch.Tensor): + data = data.cpu().detach().numpy() + if tf and tf.executing_eagerly(): + if isinstance(data, tf.Variable): + data = data.numpy() + if transpose: + data = np.transpose(data) + return data + + x = map_(x) + # Torch stores matrices in transpose (faster for backprop). + transpose = framework == "torch" and ( + x.shape[1] != weights.shape[0] and x.shape[1] == weights.shape[1] + ) + weights = map_(weights, transpose=transpose) + biases = map_(biases) + + return np.matmul(x, weights) + (0.0 if biases is None else biases) + + +@PublicAPI +def flatten_inputs_to_1d_tensor( + inputs: TensorStructType, + spaces_struct: Optional[SpaceStruct] = None, + time_axis: bool = False, + batch_axis: bool = True, +) -> TensorType: + """Flattens arbitrary input structs according to the given spaces struct. + + Returns a single 1D tensor resulting from the different input + components' values. + + Thereby: + - Boxes (any shape) get flattened to (B, [T]?, -1). Note that image boxes + are not treated differently from other types of Boxes and get + flattened as well. + - Discrete (int) values are one-hot'd, e.g. a batch of [1, 0, 3] (B=3 with + Discrete(4) space) results in [[0, 1, 0, 0], [1, 0, 0, 0], [0, 0, 0, 1]]. + - MultiDiscrete values are multi-one-hot'd, e.g. a batch of + [[0, 2], [1, 4]] (B=2 with MultiDiscrete([2, 5]) space) results in + [[1, 0, 0, 0, 1, 0, 0], [0, 1, 0, 0, 0, 0, 1]]. + + Args: + inputs: The inputs to be flattened. + spaces_struct: The (possibly nested) structure of the spaces that `inputs` + belongs to. + time_axis: Whether all inputs have a time-axis (after the batch axis). + If True, will keep not only the batch axis (0th), but the time axis + (1st) as-is and flatten everything from the 2nd axis up. + batch_axis: Whether all inputs have a batch axis. + If True, will keep that batch axis as-is and flatten everything from the + other dims up. + + Returns: + A single 1D tensor resulting from concatenating all + flattened/one-hot'd input components. Depending on the time_axis flag, + the shape is (B, n) or (B, T, n). + + .. testcode:: + :skipif: True + + # B=2 + from ray.rllib.utils.tf_utils import flatten_inputs_to_1d_tensor + from gymnasium.spaces import Discrete, Box + out = flatten_inputs_to_1d_tensor( + {"a": [1, 0], "b": [[[0.0], [0.1]], [1.0], [1.1]]}, + spaces_struct=dict(a=Discrete(2), b=Box(shape=(2, 1))) + ) + print(out) + + # B=2; T=2 + out = flatten_inputs_to_1d_tensor( + ([[1, 0], [0, 1]], + [[[0.0, 0.1], [1.0, 1.1]], [[2.0, 2.1], [3.0, 3.1]]]), + spaces_struct=tuple([Discrete(2), Box(shape=(2, ))]), + time_axis=True + ) + print(out) + + .. testoutput:: + + [[0.0, 1.0, 0.0, 0.1], [1.0, 0.0, 1.0, 1.1]] # B=2 n=4 + [[[0.0, 1.0, 0.0, 0.1], [1.0, 0.0, 1.0, 1.1]], + [[1.0, 0.0, 2.0, 2.1], [0.0, 1.0, 3.0, 3.1]]] # B=2 T=2 n=4 + """ + # `time_axis` must not be True if `batch_axis` is False. + assert not (time_axis and not batch_axis) + + flat_inputs = tree.flatten(inputs) + flat_spaces = ( + tree.flatten(spaces_struct) + if spaces_struct is not None + else [None] * len(flat_inputs) + ) + + B = None + T = None + out = [] + for input_, space in zip(flat_inputs, flat_spaces): + # Store batch and (if applicable) time dimension. + if B is None and batch_axis: + B = input_.shape[0] + if time_axis: + T = input_.shape[1] + + # One-hot encoding. + if isinstance(space, Discrete): + if time_axis: + input_ = np.reshape(input_, [B * T]) + out.append(one_hot(input_, depth=space.n).astype(np.float32)) + # Multi one-hot encoding. + elif isinstance(space, MultiDiscrete): + if time_axis: + input_ = np.reshape(input_, [B * T, -1]) + if batch_axis: + out.append( + np.concatenate( + [ + one_hot(input_[:, i], depth=n).astype(np.float32) + for i, n in enumerate(space.nvec) + ], + axis=-1, + ) + ) + else: + out.append( + np.concatenate( + [ + one_hot(input_[i], depth=n).astype(np.float32) + for i, n in enumerate(space.nvec) + ], + axis=-1, + ) + ) + # Box: Flatten. + else: + # Special case for spaces: Box(.., shape=(), ..) + if isinstance(input_, float): + input_ = np.array([input_]) + + if time_axis: + input_ = np.reshape(input_, [B * T, -1]) + elif batch_axis: + input_ = np.reshape(input_, [B, -1]) + else: + input_ = np.reshape(input_, [-1]) + out.append(input_.astype(np.float32)) + + merged = np.concatenate(out, axis=-1) + # Restore the time-dimension, if applicable. + if time_axis: + merged = np.reshape(merged, [B, T, -1]) + return merged + + +@PublicAPI +def make_action_immutable(obj): + """Flags actions immutable to notify users when trying to change them. + + Can also be used with any tree-like structure containing either + dictionaries, numpy arrays or already immutable objects per se. + Note, however that `tree.map_structure()` will in general not + include the shallow object containing all others and therefore + immutability will hold only for all objects contained in it. + Use `tree.traverse(fun, action, top_down=False)` to include + also the containing object. + + Args: + obj: The object to be made immutable. + + Returns: + The immutable object. + + .. testcode:: + :skipif: True + + import tree + import numpy as np + from ray.rllib.utils.numpy import make_action_immutable + arr = np.arange(1,10) + d = dict(a = 1, b = (arr, arr)) + tree.traverse(make_action_immutable, d, top_down=False) + """ + if isinstance(obj, np.ndarray): + obj.setflags(write=False) + return obj + elif isinstance(obj, OrderedDict): + return MappingProxyType(dict(obj)) + elif isinstance(obj, dict): + return MappingProxyType(obj) + else: + return obj + + +@PublicAPI +def huber_loss(x: np.ndarray, delta: float = 1.0) -> np.ndarray: + """Reference: https://en.wikipedia.org/wiki/Huber_loss.""" + return np.where( + np.abs(x) < delta, np.power(x, 2.0) * 0.5, delta * (np.abs(x) - 0.5 * delta) + ) + + +@PublicAPI +def l2_loss(x: np.ndarray) -> np.ndarray: + """Computes half the L2 norm of a tensor (w/o the sqrt): sum(x**2) / 2. + + Args: + x: The input tensor. + + Returns: + The l2-loss output according to the above formula given `x`. + """ + return np.sum(np.square(x)) / 2.0 + + +@PublicAPI +def lstm( + x, + weights: np.ndarray, + biases: Optional[np.ndarray] = None, + initial_internal_states: Optional[np.ndarray] = None, + time_major: bool = False, + forget_bias: float = 1.0, +): + """Calculates LSTM layer output given weights/biases, states, and input. + + Args: + x: The inputs to the LSTM layer including time-rank + (0th if time-major, else 1st) and the batch-rank + (1st if time-major, else 0th). + weights: The weights matrix. + biases: The biases vector. All 0s if None. + initial_internal_states: The initial internal + states to pass into the layer. All 0s if None. + time_major: Whether to use time-major or not. Default: False. + forget_bias: Gets added to first sigmoid (forget gate) output. + Default: 1.0. + + Returns: + Tuple consisting of 1) The LSTM layer's output and + 2) Tuple: Last (c-state, h-state). + """ + sequence_length = x.shape[0 if time_major else 1] + batch_size = x.shape[1 if time_major else 0] + units = weights.shape[1] // 4 # 4 internal layers (3x sigmoid, 1x tanh) + + if initial_internal_states is None: + c_states = np.zeros(shape=(batch_size, units)) + h_states = np.zeros(shape=(batch_size, units)) + else: + c_states = initial_internal_states[0] + h_states = initial_internal_states[1] + + # Create a placeholder for all n-time step outputs. + if time_major: + unrolled_outputs = np.zeros(shape=(sequence_length, batch_size, units)) + else: + unrolled_outputs = np.zeros(shape=(batch_size, sequence_length, units)) + + # Push the batch 4 times through the LSTM cell and capture the outputs plus + # the final h- and c-states. + for t in range(sequence_length): + input_matrix = x[t, :, :] if time_major else x[:, t, :] + input_matrix = np.concatenate((input_matrix, h_states), axis=1) + input_matmul_matrix = np.matmul(input_matrix, weights) + biases + # Forget gate (3rd slot in tf output matrix). Add static forget bias. + sigmoid_1 = sigmoid(input_matmul_matrix[:, units * 2 : units * 3] + forget_bias) + c_states = np.multiply(c_states, sigmoid_1) + # Add gate (1st and 2nd slots in tf output matrix). + sigmoid_2 = sigmoid(input_matmul_matrix[:, 0:units]) + tanh_3 = np.tanh(input_matmul_matrix[:, units : units * 2]) + c_states = np.add(c_states, np.multiply(sigmoid_2, tanh_3)) + # Output gate (last slot in tf output matrix). + sigmoid_4 = sigmoid(input_matmul_matrix[:, units * 3 : units * 4]) + h_states = np.multiply(sigmoid_4, np.tanh(c_states)) + + # Store this output time-slice. + if time_major: + unrolled_outputs[t, :, :] = h_states + else: + unrolled_outputs[:, t, :] = h_states + + return unrolled_outputs, (c_states, h_states) + + +@PublicAPI +def one_hot( + x: Union[TensorType, int], + depth: int = 0, + on_value: float = 1.0, + off_value: float = 0.0, + dtype: type = np.float32, +) -> np.ndarray: + """One-hot utility function for numpy. + + Thanks to qianyizhang: + https://gist.github.com/qianyizhang/07ee1c15cad08afb03f5de69349efc30. + + Args: + x: The input to be one-hot encoded. + depth: The max. number to be one-hot encoded (size of last rank). + on_value: The value to use for on. Default: 1.0. + off_value: The value to use for off. Default: 0.0. + + Returns: + The one-hot encoded equivalent of the input array. + """ + + # Handle simple ints properly. + if isinstance(x, int): + x = np.array(x, dtype=np.int32) + # Handle torch arrays properly. + elif torch and isinstance(x, torch.Tensor): + x = x.numpy() + + # Handle bool arrays correctly. + if x.dtype == np.bool_: + x = x.astype(np.int_) + depth = 2 + + # If depth is not given, try to infer it from the values in the array. + if depth == 0: + depth = np.max(x) + 1 + assert ( + np.max(x) < depth + ), "ERROR: The max. index of `x` ({}) is larger than depth ({})!".format( + np.max(x), depth + ) + shape = x.shape + + out = np.ones(shape=(*shape, depth)) * off_value + indices = [] + for i in range(x.ndim): + tiles = [1] * x.ndim + s = [1] * x.ndim + s[i] = -1 + r = np.arange(shape[i]).reshape(s) + if i > 0: + tiles[i - 1] = shape[i - 1] + r = np.tile(r, tiles) + indices.append(r) + indices.append(x) + out[tuple(indices)] = on_value + return out.astype(dtype) + + +@PublicAPI +def one_hot_multidiscrete(x, depths=List[int]): + # Handle torch arrays properly. + if torch and isinstance(x, torch.Tensor): + x = x.numpy() + + shape = x.shape + return np.concatenate( + [ + one_hot(x[i] if len(shape) == 1 else x[:, i], depth=n).astype(np.float32) + for i, n in enumerate(depths) + ], + axis=-1, + ) + + +@PublicAPI +def relu(x: np.ndarray, alpha: float = 0.0) -> np.ndarray: + """Implementation of the leaky ReLU function. + + y = x * alpha if x < 0 else x + + Args: + x: The input values. + alpha: A scaling ("leak") factor to use for negative x. + + Returns: + The leaky ReLU output for x. + """ + return np.maximum(x, x * alpha, x) + + +@PublicAPI +def sigmoid(x: np.ndarray, derivative: bool = False) -> np.ndarray: + """ + Returns the sigmoid function applied to x. + Alternatively, can return the derivative or the sigmoid function. + + Args: + x: The input to the sigmoid function. + derivative: Whether to return the derivative or not. + Default: False. + + Returns: + The sigmoid function (or its derivative) applied to x. + """ + if derivative: + return x * (1 - x) + else: + return 1 / (1 + np.exp(-x)) + + +@PublicAPI +def softmax( + x: Union[np.ndarray, list], axis: int = -1, epsilon: Optional[float] = None +) -> np.ndarray: + """Returns the softmax values for x. + + The exact formula used is: + S(xi) = e^xi / SUMj(e^xj), where j goes over all elements in x. + + Args: + x: The input to the softmax function. + axis: The axis along which to softmax. + epsilon: Optional epsilon as a minimum value. If None, use + `SMALL_NUMBER`. + + Returns: + The softmax over x. + """ + epsilon = epsilon or SMALL_NUMBER + # x_exp = np.maximum(np.exp(x), SMALL_NUMBER) + x_exp = np.exp(x) + # return x_exp / + # np.maximum(np.sum(x_exp, axis, keepdims=True), SMALL_NUMBER) + return np.maximum(x_exp / np.sum(x_exp, axis, keepdims=True), epsilon) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/policy.py b/lib/python3.12/site-packages/ray/rllib/utils/policy.py new file mode 100644 index 0000000000000000000000000000000000000000..2a04b5866a6626d4e71140eb7ff342a164ce96b8 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/policy.py @@ -0,0 +1,303 @@ +import logging +from typing import ( + TYPE_CHECKING, + Callable, + Dict, + List, + Optional, + Tuple, + Type, + Union, +) + +import gymnasium as gym +import numpy as np +import tree # pip install dm_tree + +import ray.cloudpickle as pickle +from ray._common.deprecation import Deprecated +from ray.rllib.core.rl_module import validate_module_id +from ray.rllib.models.preprocessors import ATARI_OBS_SHAPE +from ray.rllib.policy.policy import PolicySpec +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.typing import ( + ActionConnectorDataType, + AgentConnectorDataType, + AgentConnectorsOutput, + PartialAlgorithmConfigDict, + PolicyState, + TensorStructType, + TensorType, +) +from ray.util import log_once + +if TYPE_CHECKING: + from ray.rllib.policy.policy import Policy + +logger = logging.getLogger(__name__) + +tf1, tf, tfv = try_import_tf() + + +@OldAPIStack +def create_policy_for_framework( + policy_id: str, + policy_class: Type["Policy"], + merged_config: PartialAlgorithmConfigDict, + observation_space: gym.Space, + action_space: gym.Space, + worker_index: int = 0, + session_creator: Optional[Callable[[], "tf1.Session"]] = None, + seed: Optional[int] = None, +): + """Framework-specific policy creation logics. + + Args: + policy_id: Policy ID. + policy_class: Policy class type. + merged_config: Complete policy config. + observation_space: Observation space of env. + action_space: Action space of env. + worker_index: Index of worker holding this policy. Default is 0. + session_creator: An optional tf1.Session creation callable. + seed: Optional random seed. + """ + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + + if isinstance(merged_config, AlgorithmConfig): + merged_config = merged_config.to_dict() + + # add policy_id to merged_config + merged_config["__policy_id"] = policy_id + + framework = merged_config.get("framework", "tf") + # Tf. + if framework in ["tf2", "tf"]: + var_scope = policy_id + (f"_wk{worker_index}" if worker_index else "") + # For tf static graph, build every policy in its own graph + # and create a new session for it. + if framework == "tf": + with tf1.Graph().as_default(): + # Session creator function provided manually -> Use this one to + # create the tf1 session. + if session_creator: + sess = session_creator() + # Use a default session creator, based only on our `tf_session_args` in + # the config. + else: + sess = tf1.Session( + config=tf1.ConfigProto(**merged_config["tf_session_args"]) + ) + + with sess.as_default(): + # Set graph-level seed. + if seed is not None: + tf1.set_random_seed(seed) + with tf1.variable_scope(var_scope): + return policy_class( + observation_space, action_space, merged_config + ) + # For tf-eager: no graph, no session. + else: + with tf1.variable_scope(var_scope): + return policy_class(observation_space, action_space, merged_config) + # Non-tf: No graph, no session. + else: + return policy_class(observation_space, action_space, merged_config) + + +@OldAPIStack +def parse_policy_specs_from_checkpoint( + path: str, +) -> Tuple[PartialAlgorithmConfigDict, Dict[str, PolicySpec], Dict[str, PolicyState]]: + """Read and parse policy specifications from a checkpoint file. + + Args: + path: Path to a policy checkpoint. + + Returns: + A tuple of: base policy config, dictionary of policy specs, and + dictionary of policy states. + """ + with open(path, "rb") as f: + checkpoint_dict = pickle.load(f) + # Policy data is contained as a serialized binary blob under their + # ID keys. + w = pickle.loads(checkpoint_dict["worker"]) + + policy_config = w["policy_config"] + policy_states = w.get("policy_states", w["state"]) + serialized_policy_specs = w["policy_specs"] + policy_specs = { + id: PolicySpec.deserialize(spec) for id, spec in serialized_policy_specs.items() + } + + return policy_config, policy_specs, policy_states + + +@OldAPIStack +def local_policy_inference( + policy: "Policy", + env_id: str, + agent_id: str, + obs: TensorStructType, + reward: Optional[float] = None, + terminated: Optional[bool] = None, + truncated: Optional[bool] = None, + info: Optional[Dict] = None, + explore: bool = None, + timestep: Optional[int] = None, +) -> TensorStructType: + """Run a connector enabled policy using environment observation. + + policy_inference manages policy and agent/action connectors, + so the user does not have to care about RNN state buffering or + extra fetch dictionaries. + Note that connectors are intentionally run separately from + compute_actions_from_input_dict(), so we can have the option + of running per-user connectors on the client side in a + server-client deployment. + + Args: + policy: Policy object used in inference. + env_id: Environment ID. RLlib builds environments' trajectories internally with + connectors based on this, i.e. one trajectory per (env_id, agent_id) tuple. + agent_id: Agent ID. RLlib builds agents' trajectories internally with connectors + based on this, i.e. one trajectory per (env_id, agent_id) tuple. + obs: Environment observation to base the action on. + reward: Reward that is potentially used during inference. If not required, + may be left empty. Some policies have ViewRequirements that require this. + This can be set to zero at the first inference step - for example after + calling gmy.Env.reset. + terminated: `Terminated` flag that is potentially used during inference. If not + required, may be left None. Some policies have ViewRequirements that + require this extra information. + truncated: `Truncated` flag that is potentially used during inference. If not + required, may be left None. Some policies have ViewRequirements that + require this extra information. + info: Info that is potentially used durin inference. If not required, + may be left empty. Some policies have ViewRequirements that require this. + explore: Whether to pick an exploitation or exploration action + (default: None -> use self.config["explore"]). + timestep: The current (sampling) time step. + + Returns: + List of outputs from policy forward pass. + """ + assert ( + policy.agent_connectors + ), "policy_inference only works with connector enabled policies." + + __check_atari_obs_space(obs) + + # Put policy in inference mode, so we don't spend time on training + # only transformations. + policy.agent_connectors.in_eval() + policy.action_connectors.in_eval() + + # TODO(jungong) : support multiple env, multiple agent inference. + input_dict = {SampleBatch.NEXT_OBS: obs} + if reward is not None: + input_dict[SampleBatch.REWARDS] = reward + if terminated is not None: + input_dict[SampleBatch.TERMINATEDS] = terminated + if truncated is not None: + input_dict[SampleBatch.TRUNCATEDS] = truncated + if info is not None: + input_dict[SampleBatch.INFOS] = info + + acd_list: List[AgentConnectorDataType] = [ + AgentConnectorDataType(env_id, agent_id, input_dict) + ] + ac_outputs: List[AgentConnectorsOutput] = policy.agent_connectors(acd_list) + outputs = [] + for ac in ac_outputs: + policy_output = policy.compute_actions_from_input_dict( + ac.data.sample_batch, + explore=explore, + timestep=timestep, + ) + + # Note (Kourosh): policy output is batched, the AgentConnectorDataType should + # not be batched during inference. This is the assumption made in AgentCollector + policy_output = tree.map_structure(lambda x: x[0], policy_output) + + action_connector_data = ActionConnectorDataType( + env_id, agent_id, ac.data.raw_dict, policy_output + ) + + if policy.action_connectors: + acd = policy.action_connectors(action_connector_data) + actions = acd.output + else: + actions = policy_output[0] + + outputs.append(actions) + + # Notify agent connectors with this new policy output. + # Necessary for state buffering agent connectors, for example. + policy.agent_connectors.on_policy_output(action_connector_data) + return outputs + + +@OldAPIStack +def compute_log_likelihoods_from_input_dict( + policy: "Policy", batch: Union[SampleBatch, Dict[str, TensorStructType]] +): + """Returns log likelihood for actions in given batch for policy. + + Computes likelihoods by passing the observations through the current + policy's `compute_log_likelihoods()` method + + Args: + batch: The SampleBatch or MultiAgentBatch to calculate action + log likelihoods from. This batch/batches must contain OBS + and ACTIONS keys. + + Returns: + The probabilities of the actions in the batch, given the + observations and the policy. + """ + num_state_inputs = 0 + for k in batch.keys(): + if k.startswith("state_in_"): + num_state_inputs += 1 + state_keys = ["state_in_{}".format(i) for i in range(num_state_inputs)] + log_likelihoods: TensorType = policy.compute_log_likelihoods( + actions=batch[SampleBatch.ACTIONS], + obs_batch=batch[SampleBatch.OBS], + state_batches=[batch[k] for k in state_keys], + prev_action_batch=batch.get(SampleBatch.PREV_ACTIONS), + prev_reward_batch=batch.get(SampleBatch.PREV_REWARDS), + actions_normalized=policy.config.get("actions_in_input_normalized", False), + ) + return log_likelihoods + + +@Deprecated(new="Policy.from_checkpoint([checkpoint path], [policy IDs]?)", error=True) +def load_policies_from_checkpoint(path, policy_ids=None): + pass + + +def __check_atari_obs_space(obs): + # TODO(Artur): Remove this after we have migrated deepmind style preprocessing into + # connectors (and don't auto-wrap in RW anymore) + if any( + o.shape == ATARI_OBS_SHAPE if isinstance(o, np.ndarray) else False + for o in tree.flatten(obs) + ): + if log_once("warn_about_possibly_non_wrapped_atari_env"): + logger.warning( + "The observation you fed into local_policy_inference() has " + "dimensions (210, 160, 3), which is the standard for atari " + "environments. If RLlib raises an error including a related " + "dimensionality mismatch, you may need to use " + "ray.rllib.env.wrappers.atari_wrappers.wrap_deepmind to wrap " + "you environment." + ) + + +# @OldAPIStack +validate_policy_id = validate_module_id diff --git a/lib/python3.12/site-packages/ray/rllib/utils/postprocessing/__init__.py b/lib/python3.12/site-packages/ray/rllib/utils/postprocessing/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/utils/postprocessing/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/postprocessing/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5213bd24560bbba3fc62421857847bc34e222772 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/postprocessing/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/postprocessing/__pycache__/episodes.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/postprocessing/__pycache__/episodes.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f4c26c66d8e76ee237e29991c3a33c28704ff183 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/postprocessing/__pycache__/episodes.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/postprocessing/__pycache__/value_predictions.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/postprocessing/__pycache__/value_predictions.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..97df17accc4b92f3a25e662ccce9e5f501fd1610 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/postprocessing/__pycache__/value_predictions.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/postprocessing/__pycache__/zero_padding.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/postprocessing/__pycache__/zero_padding.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..f1abc3a19982e98f7900511d7295375c4a01870a Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/postprocessing/__pycache__/zero_padding.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/postprocessing/episodes.py b/lib/python3.12/site-packages/ray/rllib/utils/postprocessing/episodes.py new file mode 100644 index 0000000000000000000000000000000000000000..fdc5b1321659366372b985116ad3cddcf8b2f899 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/postprocessing/episodes.py @@ -0,0 +1,144 @@ +from typing import List, Tuple + +import numpy as np + +from ray.rllib.env.single_agent_episode import SingleAgentEpisode +from ray.util.annotations import DeveloperAPI + + +@DeveloperAPI +def add_one_ts_to_episodes_and_truncate(episodes: List[SingleAgentEpisode]): + """Adds an artificial timestep to an episode at the end. + + In detail: The last observations, infos, actions, and all `extra_model_outputs` + will be duplicated and appended to each episode's data. An extra 0.0 reward + will be appended to the episode's rewards. The episode's timestep will be + increased by 1. Also, adds the truncated=True flag to each episode if the + episode is not already done (terminated or truncated). + + Useful for value function bootstrapping, where it is required to compute a + forward pass for the very last timestep within the episode, + i.e. using the following input dict: { + obs=[final obs], + state=[final state output], + prev. reward=[final reward], + etc.. + } + + Args: + episodes: The list of SingleAgentEpisode objects to extend by one timestep + and add a truncation flag if necessary. + + Returns: + A list of the original episodes' truncated values (so the episodes can be + properly restored later into their original states). + """ + orig_truncateds = [] + for episode in episodes: + orig_truncateds.append(episode.is_truncated) + + # Add timestep. + episode.t += 1 + # Use the episode API that allows appending (possibly complex) structs + # to the data. + episode.observations.append(episode.observations[-1]) + episode.infos.append(episode.infos[-1]) + episode.actions.append(episode.actions[-1]) + episode.rewards.append(0.0) + for v in episode.extra_model_outputs.values(): + v.append(v[-1]) + # Artificially make this episode truncated for the upcoming GAE + # computations. + if not episode.is_done: + episode.is_truncated = True + # Validate to make sure, everything is in order. + episode.validate() + + return orig_truncateds + + +@DeveloperAPI +def remove_last_ts_from_data( + episode_lens: List[int], + *data: Tuple[np._typing.NDArray], +) -> Tuple[np._typing.NDArray]: + """Removes the last timesteps from each given data item. + + Each item in data is a concatenated sequence of episodes data. + For example if `episode_lens` is [2, 4], then data is a shape=(6,) + ndarray. The returned corresponding value will have shape (4,), meaning + both episodes have been shortened by exactly one timestep to 1 and 3. + + ..testcode:: + + from ray.rllib.algorithms.ppo.ppo_learner import PPOLearner + import numpy as np + + unpadded = PPOLearner._remove_last_ts_from_data( + [5, 3], + np.array([0, 1, 2, 3, 4, 0, 1, 2]), + ) + assert (unpadded[0] == [0, 1, 2, 3, 0, 1]).all() + + unpadded = PPOLearner._remove_last_ts_from_data( + [4, 2, 3], + np.array([0, 1, 2, 3, 0, 1, 0, 1, 2]), + np.array([4, 5, 6, 7, 2, 3, 3, 4, 5]), + ) + assert (unpadded[0] == [0, 1, 2, 0, 0, 1]).all() + assert (unpadded[1] == [4, 5, 6, 2, 3, 4]).all() + + Args: + episode_lens: A list of current episode lengths. The returned + data will have the same lengths minus 1 timestep. + data: A tuple of data items (np.ndarrays) representing concatenated episodes + to be shortened by one timestep per episode. + Note that only arrays with `shape=(n,)` are supported! The + returned data will have `shape=(n-len(episode_lens),)` (each + episode gets shortened by one timestep). + + Returns: + A tuple of new data items shortened by one timestep. + """ + # Figure out the new slices to apply to each data item based on + # the given episode_lens. + slices = [] + sum = 0 + for len_ in episode_lens: + slices.append(slice(sum, sum + len_ - 1)) + sum += len_ + # Compiling return data by slicing off one timestep at the end of + # each episode. + ret = [] + for d in data: + ret.append(np.concatenate([d[s] for s in slices])) + return tuple(ret) if len(ret) > 1 else ret[0] + + +@DeveloperAPI +def remove_last_ts_from_episodes_and_restore_truncateds( + episodes: List[SingleAgentEpisode], + orig_truncateds: List[bool], +) -> None: + """Reverts the effects of `_add_ts_to_episodes_and_truncate`. + + Args: + episodes: The list of SingleAgentEpisode objects to extend by one timestep + and add a truncation flag if necessary. + orig_truncateds: A list of the original episodes' truncated values to be + applied to the `episodes`. + """ + + # Fix all episodes. + for episode, orig_truncated in zip(episodes, orig_truncateds): + # Reduce timesteps by 1. + episode.t -= 1 + # Remove all extra timestep data from the episode's buffers. + episode.observations.pop() + episode.infos.pop() + episode.actions.pop() + episode.rewards.pop() + for v in episode.extra_model_outputs.values(): + v.pop() + # Fix the truncateds flag again. + episode.is_truncated = orig_truncated diff --git a/lib/python3.12/site-packages/ray/rllib/utils/postprocessing/value_predictions.py b/lib/python3.12/site-packages/ray/rllib/utils/postprocessing/value_predictions.py new file mode 100644 index 0000000000000000000000000000000000000000..0c0f88e7f49e9a2969ebdc2fc2887fedcf1f3e32 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/postprocessing/value_predictions.py @@ -0,0 +1,108 @@ +import numpy as np + +from ray.util.annotations import DeveloperAPI + + +@DeveloperAPI +def compute_value_targets( + values, + rewards, + terminateds, + truncateds, + gamma: float, + lambda_: float, +): + """Computes value function (vf) targets given vf predictions and rewards. + + Note that advantages can then easily be computed via the formula: + advantages = targets - vf_predictions + """ + # Force-set all values at terminals (not at truncations!) to 0.0. + orig_values = flat_values = values * (1.0 - terminateds) + + flat_values = np.append(flat_values, 0.0) + intermediates = rewards + gamma * (1 - lambda_) * flat_values[1:] + continues = 1.0 - terminateds + + Rs = [] + last = flat_values[-1] + for t in reversed(range(intermediates.shape[0])): + last = intermediates[t] + continues[t] * gamma * lambda_ * last + Rs.append(last) + if truncateds[t]: + last = orig_values[t] + + # Reverse back to correct (time) direction. + value_targets = np.stack(list(reversed(Rs)), axis=0) + + return value_targets.astype(np.float32) + + +def extract_bootstrapped_values(vf_preds, episode_lengths, T): + """Returns a bootstrapped value batch given value predictions. + + Note that the incoming value predictions must have happened over (artificially) + elongated episodes (by 1 timestep at the end). This way, we can either extract the + `vf_preds` at these extra timesteps (as "bootstrap values") or skip over them + entirely if they lie in the middle of the T-slices. + + For example, given an episodes structure like this: + 01234a 0123456b 01c 012- 0123e 012- + where each episode is separated by a space and goes from 0 to n and ends in an + artificially elongated timestep (denoted by 'a', 'b', 'c', '-', or 'e'), where '-' + means that the episode was terminated and the bootstrap value at the end should be + zero and 'a', 'b', 'c', etc.. represent truncated episode ends with computed vf + estimates. + The output for the above sequence (and T=4) should then be: + 4 3 b 2 3 - + + Args: + vf_preds: The computed value function predictions over the artificially + elongated episodes (by one timestep at the end). + episode_lengths: The original (correct) episode lengths, NOT counting the + artificially added timestep at the end. + T: The size of the time dimension by which to slice the data. Note that the + sum of all episode lengths (`sum(episode_lengths)`) must be dividable by T. + + Returns: + The batch of bootstrapped values. + """ + bootstrapped_values = [] + if sum(episode_lengths) % T != 0: + raise ValueError( + "Can only extract bootstrapped values if the sum of episode lengths " + f"({sum(episode_lengths)}) is dividable by the given T ({T})!" + ) + + # Loop over all episode lengths and collect bootstrap values. + # Do not alter incoming `episode_lengths` list. + episode_lengths = episode_lengths[:] + i = -1 + while i < len(episode_lengths) - 1: + i += 1 + eps_len = episode_lengths[i] + # We can make another T-stride inside this episode -> + # - Use a vf prediction within the episode as bootstrapped value. + # - "Fix" the episode_lengths array and continue within the same episode. + if T < eps_len: + bootstrapped_values.append(vf_preds[T]) + vf_preds = vf_preds[T:] + episode_lengths[i] -= T + i -= 1 + # We can make another T-stride inside this episode, but will then be at the end + # of it -> + # - Use the value function prediction at the artificially added timestep + # as bootstrapped value. + # - Skip the additional timestep at the end and ,ove on with next episode. + elif T == eps_len: + bootstrapped_values.append(vf_preds[T]) + vf_preds = vf_preds[T + 1 :] + # The episode fits entirely into the T-stride -> + # - Move on to next episode ("fix" its length by make it seemingly longer). + else: + # Skip bootstrap value of current episode (not needed). + vf_preds = vf_preds[1:] + # Make next episode seem longer. + episode_lengths[i + 1] += eps_len + + return np.array(bootstrapped_values) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/postprocessing/zero_padding.py b/lib/python3.12/site-packages/ray/rllib/utils/postprocessing/zero_padding.py new file mode 100644 index 0000000000000000000000000000000000000000..15783668834f769d0e07d35c3fe9295b82166ce7 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/postprocessing/zero_padding.py @@ -0,0 +1,270 @@ +from collections import deque +from typing import List, Tuple, Union + +import numpy as np +import tree # pip install dm_tree + +from ray.rllib.utils.spaces.space_utils import BatchedNdArray, batch +from ray.util.annotations import DeveloperAPI + + +@DeveloperAPI +def create_mask_and_seq_lens(episode_len: int, T: int) -> Tuple[List, List]: + """Creates loss mask and a seq_lens array, given an episode length and T. + + Args: + episode_lens: A list of episode lengths to infer the loss mask and seq_lens + array from. + T: The maximum number of timesteps in each "row", also known as the maximum + sequence length (max_seq_len). Episodes are split into chunks that are at + most `T` long and remaining timesteps will be zero-padded (and masked out). + + Returns: + Tuple consisting of a) list of the loss masks to use (masking out areas that + are past the end of an episode (or rollout), but had to be zero-added due to + the added extra time rank (of length T) and b) the list of sequence lengths + resulting from splitting the given episodes into chunks of at most `T` + timesteps. + """ + mask = [] + seq_lens = [] + + len_ = min(episode_len, T) + seq_lens.append(len_) + row = np.array([1] * len_ + [0] * (T - len_), np.bool_) + mask.append(row) + + # Handle sequence lengths greater than T. + overflow = episode_len - T + while overflow > 0: + len_ = min(overflow, T) + seq_lens.append(len_) + extra_row = np.array([1] * len_ + [0] * (T - len_), np.bool_) + mask.append(extra_row) + overflow -= T + + return mask, seq_lens + + +@DeveloperAPI +def split_and_zero_pad( + item_list: List[Union[BatchedNdArray, np._typing.NDArray, float]], + max_seq_len: int, +) -> List[np._typing.NDArray]: + """Splits the contents of `item_list` into a new list of ndarrays and returns it. + + In the returned list, each item is one ndarray of len (axis=0) `max_seq_len`. + The last item in the returned list may be (right) zero-padded, if necessary, to + reach `max_seq_len`. + + If `item_list` contains one or more `BatchedNdArray` (instead of individual + items), these will be split accordingly along their axis=0 to yield the returned + structure described above. + + .. testcode:: + + from ray.rllib.utils.postprocessing.zero_padding import ( + BatchedNdArray, + split_and_zero_pad, + ) + from ray.rllib.utils.test_utils import check + + # Simple case: `item_list` contains individual floats. + check( + split_and_zero_pad([0, 1, 2, 3, 4, 5, 6, 7], 5), + [[0, 1, 2, 3, 4], [5, 6, 7, 0, 0]], + ) + + # `item_list` contains BatchedNdArray (ndarrays that explicitly declare they + # have a batch axis=0). + check( + split_and_zero_pad([ + BatchedNdArray([0, 1]), + BatchedNdArray([2, 3, 4, 5]), + BatchedNdArray([6, 7, 8]), + ], 5), + [[0, 1, 2, 3, 4], [5, 6, 7, 8, 0]], + ) + + Args: + item_list: A list of individual items or BatchedNdArrays to be split into + `max_seq_len` long pieces (the last of which may be zero-padded). + max_seq_len: The maximum length of each item in the returned list. + + Returns: + A list of np.ndarrays (all of length `max_seq_len`), which contains the same + data as `item_list`, but split into sub-chunks of size `max_seq_len`. + The last item in the returned list may be zero-padded, if necessary. + """ + zero_element = tree.map_structure( + lambda s: np.zeros_like([s[0]] if isinstance(s, BatchedNdArray) else s), + item_list[0], + ) + + # The replacement list (to be returned) for `items_list`. + # Items list contains n individual items. + # -> ret will contain m batched rows, where m == n // T and the last row + # may be zero padded (until T). + ret = [] + + # List of the T-axis item, collected to form the next row. + current_time_row = [] + current_t = 0 + + item_list = deque(item_list) + while len(item_list) > 0: + item = item_list.popleft() + t = max_seq_len - current_t + + # In case `item` is a complex struct. + item_flat = tree.flatten(item) + item_list_append = [] + current_time_row_flat_items = [] + add_to_current_t = 0 + + for itm in item_flat: + # `itm` is already a batched np.array: Split if necessary. + if isinstance(itm, BatchedNdArray): + current_time_row_flat_items.append(itm[:t]) + if len(itm) <= t: + add_to_current_t = len(itm) + else: + add_to_current_t = t + item_list_append.append(itm[t:]) + # `itm` is a single item (no batch axis): Append and continue with next + # item. + else: + current_time_row_flat_items.append(itm) + add_to_current_t = 1 + + current_t += add_to_current_t + current_time_row.append(tree.unflatten_as(item, current_time_row_flat_items)) + if item_list_append: + item_list.appendleft(tree.unflatten_as(item, item_list_append)) + + # `current_time_row` is "full" (max_seq_len): Append as ndarray (with batch + # axis) to `ret`. + if current_t == max_seq_len: + ret.append( + batch( + current_time_row, + individual_items_already_have_batch_dim="auto", + ) + ) + current_time_row = [] + current_t = 0 + + # `current_time_row` is unfinished: Pad, if necessary and append to `ret`. + if current_t > 0 and current_t < max_seq_len: + current_time_row.extend([zero_element] * (max_seq_len - current_t)) + ret.append( + batch(current_time_row, individual_items_already_have_batch_dim="auto") + ) + + return ret + + +@DeveloperAPI +def split_and_zero_pad_n_episodes( + nd_array: np._typing.NDArray, + episode_lens: List[int], + max_seq_len: int, +) -> List[np._typing.NDArray]: + """Splits and zero-pads a single np.ndarray based on episode lens and a maxlen. + + Args: + nd_array: The single np.ndarray to be split into n chunks, based on the given + `episode_lens` and the `max_seq_len` argument. For example, if `nd_array` + has a batch dimension (axis 0) of 21, `episode_lens` is [15, 3, 3], and + `max_seq_len` is 6, then the returned list would have np.ndarrays in it of + batch dimensions (axis 0): [6, 6, 6 (zero-padded), 6 (zero-padded), + 6 (zero-padded)]. + Note that this function doesn't work on nested data, such as dicts of + ndarrays. + episode_lens: A list of episode lengths along which to split and zero-pad the + given `nd_array`. + max_seq_len: The maximum sequence length to split at (and zero-pad). + + Returns: A list of n np.ndarrays, resulting from splitting and zero-padding the + given `nd_array`. + """ + ret = [] + + cursor = 0 + for episode_len in episode_lens: + items = BatchedNdArray(nd_array[cursor : cursor + episode_len]) + ret.extend(split_and_zero_pad([items], max_seq_len)) + cursor += episode_len + + return ret + + +@DeveloperAPI +def unpad_data_if_necessary( + episode_lens: List[int], + data: np._typing.NDArray, +) -> np._typing.NDArray: + """Removes right-side zero-padding from data based on `episode_lens`. + + ..testcode:: + + from ray.rllib.utils.postprocessing.zero_padding import unpad_data_if_necessary + import numpy as np + + unpadded = unpad_data_if_necessary( + episode_lens=[4, 2], + data=np.array([ + [2, 4, 5, 3, 0, 0, 0, 0], + [-1, 3, 0, 0, 0, 0, 0, 0], + ]), + ) + assert (unpadded == [2, 4, 5, 3, -1, 3]).all() + + unpadded = unpad_data_if_necessary( + episode_lens=[1, 5], + data=np.array([ + [2, 0, 0, 0, 0], + [-1, -2, -3, -4, -5], + ]), + ) + assert (unpadded == [2, -1, -2, -3, -4, -5]).all() + + Args: + episode_lens: A list of actual episode lengths. + data: A 2D np.ndarray with right-side zero-padded rows. + + Returns: + A 1D np.ndarray resulting from concatenation of the un-padded + input data along the 0-axis. + """ + # If data des NOT have time dimension, return right away. + if len(data.shape) == 1: + return data + + # Assert we only have B and T dimensions (meaning this function only operates + # on single-float data, such as value function predictions, advantages, or rewards). + assert len(data.shape) == 2 + + new_data = [] + row_idx = 0 + + T = data.shape[1] + for len_ in episode_lens: + # Calculate how many full rows this array occupies and how many elements are + # in the last, potentially partial row. + num_rows, col_idx = divmod(len_, T) + + # If the array spans multiple full rows, fully include these rows. + for i in range(num_rows): + new_data.append(data[row_idx]) + row_idx += 1 + + # If there are elements in the last, potentially partial row, add this + # partial row as well. + if col_idx > 0: + new_data.append(data[row_idx, :col_idx]) + + # Move to the next row for the next array (skip the zero-padding zone). + row_idx += 1 + + return np.concatenate(new_data) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/pre_checks/__init__.py b/lib/python3.12/site-packages/ray/rllib/utils/pre_checks/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/utils/pre_checks/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/pre_checks/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8ab45ba893195904dd34eb25b2c93da220bbf8c0 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/pre_checks/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/pre_checks/__pycache__/env.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/pre_checks/__pycache__/env.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5c031fe50ea5c3f84df188a1966c5157faed434f Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/pre_checks/__pycache__/env.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/pre_checks/env.py b/lib/python3.12/site-packages/ray/rllib/utils/pre_checks/env.py new file mode 100644 index 0000000000000000000000000000000000000000..7055f08695fb3b9a945def9556d226766baddad3 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/pre_checks/env.py @@ -0,0 +1,288 @@ +"""Common pre-checks for all RLlib experiments.""" +import logging +from typing import TYPE_CHECKING, Set + +import gymnasium as gym +import numpy as np +import tree # pip install dm_tree + +from ray.rllib.utils.annotations import DeveloperAPI +from ray.rllib.utils.error import ERR_MSG_OLD_GYM_API, UnsupportedSpaceException +from ray.rllib.utils.spaces.space_utils import get_base_struct_from_space +from ray.util import log_once + +if TYPE_CHECKING: + from ray.rllib.env import MultiAgentEnv + +logger = logging.getLogger(__name__) + + +@DeveloperAPI +def check_multiagent_environments(env: "MultiAgentEnv") -> None: + """Checking for common errors in RLlib MultiAgentEnvs. + + Args: + env: The env to be checked. + """ + from ray.rllib.env import MultiAgentEnv + + if not isinstance(env, MultiAgentEnv): + raise ValueError("The passed env is not a MultiAgentEnv.") + elif not ( + hasattr(env, "observation_space") + and hasattr(env, "action_space") + and hasattr(env, "_agent_ids") + ): + if log_once("ma_env_super_ctor_called"): + logger.warning( + f"Your MultiAgentEnv {env} does not have some or all of the needed " + "base-class attributes! Make sure you call `super().__init__()` from " + "within your MutiAgentEnv's constructor. " + "This will raise an error in the future." + ) + return + + try: + obs_and_infos = env.reset(seed=42, options={}) + except Exception as e: + raise ValueError( + ERR_MSG_OLD_GYM_API.format( + env, "In particular, the `reset()` method seems to be faulty." + ) + ) from e + reset_obs, reset_infos = obs_and_infos + + _check_if_element_multi_agent_dict(env, reset_obs, "reset()") + + sampled_action = { + aid: env.get_action_space(aid).sample() for aid in reset_obs.keys() + } + _check_if_element_multi_agent_dict( + env, sampled_action, "get_action_space(agent_id=..).sample()" + ) + + try: + results = env.step(sampled_action) + except Exception as e: + raise ValueError( + ERR_MSG_OLD_GYM_API.format( + env, "In particular, the `step()` method seems to be faulty." + ) + ) from e + next_obs, reward, done, truncated, info = results + + _check_if_element_multi_agent_dict(env, next_obs, "step, next_obs") + _check_if_element_multi_agent_dict(env, reward, "step, reward") + _check_if_element_multi_agent_dict(env, done, "step, done") + _check_if_element_multi_agent_dict(env, truncated, "step, truncated") + _check_if_element_multi_agent_dict(env, info, "step, info", allow_common=True) + _check_reward({"dummy_env_id": reward}, base_env=True, agent_ids=env.agents) + _check_done_and_truncated( + {"dummy_env_id": done}, + {"dummy_env_id": truncated}, + base_env=True, + agent_ids=env.agents, + ) + _check_info({"dummy_env_id": info}, base_env=True, agent_ids=env.agents) + + +def _check_reward(reward, base_env=False, agent_ids=None): + if base_env: + for _, multi_agent_dict in reward.items(): + for agent_id, rew in multi_agent_dict.items(): + if not ( + np.isreal(rew) + and not isinstance(rew, bool) + and ( + np.isscalar(rew) + or (isinstance(rew, np.ndarray) and rew.shape == ()) + ) + ): + error = ( + "Your step function must return rewards that are" + f" integer or float. reward: {rew}. Instead it was a " + f"{type(rew)}" + ) + raise ValueError(error) + if not (agent_id in agent_ids or agent_id == "__all__"): + error = ( + f"Your reward dictionary must have agent ids that belong to " + f"the environment. AgentIDs received from " + f"env.agents are: {agent_ids}" + ) + raise ValueError(error) + elif not ( + np.isreal(reward) + and not isinstance(reward, bool) + and ( + np.isscalar(reward) + or (isinstance(reward, np.ndarray) and reward.shape == ()) + ) + ): + error = ( + "Your step function must return a reward that is integer or float. " + "Instead it was a {}".format(type(reward)) + ) + raise ValueError(error) + + +def _check_done_and_truncated(done, truncated, base_env=False, agent_ids=None): + for what in ["done", "truncated"]: + data = done if what == "done" else truncated + if base_env: + for _, multi_agent_dict in data.items(): + for agent_id, done_ in multi_agent_dict.items(): + if not isinstance(done_, (bool, np.bool_)): + raise ValueError( + f"Your step function must return `{what}s` that are " + f"boolean. But instead was a {type(data)}" + ) + if not (agent_id in agent_ids or agent_id == "__all__"): + error = ( + f"Your `{what}s` dictionary must have agent ids that " + f"belong to the environment. AgentIDs received from " + f"env.agents are: {agent_ids}" + ) + raise ValueError(error) + elif not isinstance(data, (bool, np.bool_)): + error = ( + f"Your step function must return a `{what}` that is a boolean. But " + f"instead was a {type(data)}" + ) + raise ValueError(error) + + +def _check_info(info, base_env=False, agent_ids=None): + if base_env: + for _, multi_agent_dict in info.items(): + for agent_id, inf in multi_agent_dict.items(): + if not isinstance(inf, dict): + raise ValueError( + "Your step function must return infos that are a dict. " + f"instead was a {type(inf)}: element: {inf}" + ) + if not ( + agent_id in agent_ids + or agent_id == "__all__" + or agent_id == "__common__" + ): + error = ( + f"Your dones dictionary must have agent ids that belong to " + f"the environment. AgentIDs received from " + f"env.agents are: {agent_ids}" + ) + raise ValueError(error) + elif not isinstance(info, dict): + error = ( + "Your step function must return a info that " + f"is a dict. element type: {type(info)}. element: {info}" + ) + raise ValueError(error) + + +def _not_contained_error(func_name, _type): + _error = ( + f"The {_type} collected from {func_name} was not contained within" + f" your env's {_type} space. Its possible that there was a type" + f"mismatch (for example {_type}s of np.float32 and a space of" + f"np.float64 {_type}s), or that one of the sub-{_type}s was" + f"out of bounds" + ) + return _error + + +def _check_if_element_multi_agent_dict( + env, + element, + function_string, + base_env=False, + allow_common=False, +): + if not isinstance(element, dict): + if base_env: + error = ( + f"The element returned by {function_string} contains values " + f"that are not MultiAgentDicts. Instead, they are of " + f"type: {type(element)}" + ) + else: + error = ( + f"The element returned by {function_string} is not a " + f"MultiAgentDict. Instead, it is of type: " + f" {type(element)}" + ) + raise ValueError(error) + agent_ids: Set = set(env.agents) + agent_ids.add("__all__") + if allow_common: + agent_ids.add("__common__") + + if not all(k in agent_ids for k in element): + if base_env: + error = ( + f"The element returned by {function_string} has agent_ids" + f" that are not the names of the agents in the env." + f"agent_ids in this\nMultiEnvDict:" + f" {list(element.keys())}\nAgentIDs in this env: " + f"{env.agents}" + ) + else: + error = ( + f"The element returned by {function_string} has agent_ids" + f" that are not the names of the agents in the env. " + f"\nAgentIDs in this MultiAgentDict: " + f"{list(element.keys())}\nAgentIDs in this env: " + f"{env.agents}. You likely need to add the attribute `agents` to your " + f"env, which is a list containing the IDs of agents currently in your " + f"env/episode, as well as, `possible_agents`, which is a list of all " + f"possible agents that could ever show up in your env." + ) + raise ValueError(error) + + +def _find_offending_sub_space(space, value): + """Returns error, value, and space when offending `space.contains(value)` fails. + + Returns only the offending sub-value/sub-space in case `space` is a complex Tuple + or Dict space. + + Args: + space: The gym.Space to check. + value: The actual (numpy) value to check for matching `space`. + + Returns: + Tuple consisting of 1) key-sequence of the offending sub-space or the empty + string if `space` is not complex (Tuple or Dict), 2) the offending sub-space, + 3) the offending sub-space's dtype, 4) the offending sub-value, 5) the offending + sub-value's dtype. + + .. testcode:: + :skipif: True + + path, space, space_dtype, value, value_dtype = _find_offending_sub_space( + gym.spaces.Dict({ + -2.0, 1.5, (2, ), np.int8), np.array([-1.5, 3.0]) + ) + + """ + if not isinstance(space, (gym.spaces.Dict, gym.spaces.Tuple)): + return None, space, space.dtype, value, _get_type(value) + + structured_space = get_base_struct_from_space(space) + + def map_fn(p, s, v): + if not s.contains(v): + raise UnsupportedSpaceException((p, s, v)) + + try: + tree.map_structure_with_path(map_fn, structured_space, value) + except UnsupportedSpaceException as e: + space, value = e.args[0][1], e.args[0][2] + return "->".join(e.args[0][0]), space, space.dtype, value, _get_type(value) + + # This is actually an error. + return None, None, None, None, None + + +def _get_type(var): + return var.dtype if hasattr(var, "dtype") else type(var) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__init__.py b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c5f53f25e3e346e9eec23f958147d4577ef2b868 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__init__.py @@ -0,0 +1,45 @@ +from ray.rllib.utils.replay_buffers.episode_replay_buffer import EpisodeReplayBuffer +from ray.rllib.utils.replay_buffers.fifo_replay_buffer import FifoReplayBuffer +from ray.rllib.utils.replay_buffers.multi_agent_episode_buffer import ( + MultiAgentEpisodeReplayBuffer, +) +from ray.rllib.utils.replay_buffers.multi_agent_mixin_replay_buffer import ( + MultiAgentMixInReplayBuffer, +) +from ray.rllib.utils.replay_buffers.multi_agent_prioritized_episode_buffer import ( + MultiAgentPrioritizedEpisodeReplayBuffer, +) +from ray.rllib.utils.replay_buffers.multi_agent_prioritized_replay_buffer import ( + MultiAgentPrioritizedReplayBuffer, +) +from ray.rllib.utils.replay_buffers.multi_agent_replay_buffer import ( + MultiAgentReplayBuffer, + ReplayMode, +) +from ray.rllib.utils.replay_buffers.prioritized_episode_buffer import ( + PrioritizedEpisodeReplayBuffer, +) +from ray.rllib.utils.replay_buffers.prioritized_replay_buffer import ( + PrioritizedReplayBuffer, +) +from ray.rllib.utils.replay_buffers.replay_buffer import ReplayBuffer, StorageUnit +from ray.rllib.utils.replay_buffers.reservoir_replay_buffer import ReservoirReplayBuffer + +from ray.rllib.utils.replay_buffers import utils # isort: skip + +__all__ = [ + "EpisodeReplayBuffer", + "FifoReplayBuffer", + "MultiAgentEpisodeReplayBuffer", + "MultiAgentMixInReplayBuffer", + "MultiAgentPrioritizedEpisodeReplayBuffer", + "MultiAgentPrioritizedReplayBuffer", + "MultiAgentReplayBuffer", + "PrioritizedEpisodeReplayBuffer", + "PrioritizedReplayBuffer", + "ReplayMode", + "ReplayBuffer", + "ReservoirReplayBuffer", + "StorageUnit", + "utils", +] diff --git a/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..41633383a42177f34013e718471f60ba9befe089 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/base.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/base.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0518eb1754f2acfa8287fc53ddfe38e88a33eaa1 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/base.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/episode_replay_buffer.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/episode_replay_buffer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d9cdde944adcded72f78829ede0e8576f35def52 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/episode_replay_buffer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/fifo_replay_buffer.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/fifo_replay_buffer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b8d6560858f88f7f7a72ffbcfefc94fb88aec2d8 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/fifo_replay_buffer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/multi_agent_episode_buffer.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/multi_agent_episode_buffer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dd6a1aaa1ee7c09ff63634d92850b86459ca3669 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/multi_agent_episode_buffer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/multi_agent_mixin_replay_buffer.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/multi_agent_mixin_replay_buffer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..62c659390261dfab15b1614e3c6a01baf69b8564 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/multi_agent_mixin_replay_buffer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/multi_agent_prioritized_episode_buffer.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/multi_agent_prioritized_episode_buffer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..170a5761462b74e2c7f667d3f2ce7ce6dc81e6ba Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/multi_agent_prioritized_episode_buffer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/multi_agent_prioritized_replay_buffer.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/multi_agent_prioritized_replay_buffer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a009b44294b555d10d743ec0456ebd7d7f0d7d49 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/multi_agent_prioritized_replay_buffer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/multi_agent_replay_buffer.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/multi_agent_replay_buffer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e9b4f30c951e042e9dacda7e637cc76046d5224e Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/multi_agent_replay_buffer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/prioritized_episode_buffer.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/prioritized_episode_buffer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..14879db16b6468d5f5b146fd39dca702def59457 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/prioritized_episode_buffer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/prioritized_replay_buffer.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/prioritized_replay_buffer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5d63eea3413292858ca20930d91ee10cfea5d68c Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/prioritized_replay_buffer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/replay_buffer.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/replay_buffer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5b35d5a859f7ae8a26b64dc63dc4fe24efc95a67 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/replay_buffer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/reservoir_replay_buffer.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/reservoir_replay_buffer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4e1d796942aae74b6d3be3609fba7ee4cc5e1331 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/reservoir_replay_buffer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/simple_replay_buffer.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/simple_replay_buffer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..566fb8c9c40d7be3afc7122b2a3c6a202c2f9d42 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/simple_replay_buffer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dea040780d100049df65b779a0c7820faa9993b2 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/base.py b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/base.py new file mode 100644 index 0000000000000000000000000000000000000000..dc5fd7d6c35e17e482053830a09c0ccd2d1aa1d0 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/base.py @@ -0,0 +1,76 @@ +import platform +from abc import ABCMeta, abstractmethod +from typing import Any, Dict, Optional + +from ray.util.annotations import DeveloperAPI + + +@DeveloperAPI +class ReplayBufferInterface(metaclass=ABCMeta): + """Abstract base class for all of RLlib's replay buffers. + + Mainly defines the `add()` and `sample()` methods that every buffer class + must implement to be usable by an Algorithm. + Buffers may determine on all the implementation details themselves, e.g. + whether to store single timesteps, episodes, or episode fragments or whether + to return fixed batch sizes or per-call defined ones. + """ + + @abstractmethod + @DeveloperAPI + def __len__(self) -> int: + """Returns the number of items currently stored in this buffer.""" + + @abstractmethod + @DeveloperAPI + def add(self, batch: Any, **kwargs) -> None: + """Adds a batch of experiences or other data to this buffer. + + Args: + batch: Batch or data to add. + ``**kwargs``: Forward compatibility kwargs. + """ + + @abstractmethod + @DeveloperAPI + def sample(self, num_items: Optional[int] = None, **kwargs) -> Any: + """Samples `num_items` items from this buffer. + + The exact shape of the returned data depends on the buffer's implementation. + + Args: + num_items: Number of items to sample from this buffer. + ``**kwargs``: Forward compatibility kwargs. + + Returns: + A batch of items. + """ + + @abstractmethod + @DeveloperAPI + def get_state(self) -> Dict[str, Any]: + """Returns all local state in a dict. + + Returns: + The serializable local state. + """ + + @abstractmethod + @DeveloperAPI + def set_state(self, state: Dict[str, Any]) -> None: + """Restores all local state to the provided `state`. + + Args: + state: The new state to set this buffer. Can be obtained by calling + `self.get_state()`. + """ + + @DeveloperAPI + def get_host(self) -> str: + """Returns the computer's network name. + + Returns: + The computer's networks name or an empty string, if the network + name could not be determined. + """ + return platform.node() diff --git a/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/episode_replay_buffer.py b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/episode_replay_buffer.py new file mode 100644 index 0000000000000000000000000000000000000000..cc0a18a0128dc04390715c536ed5bc5666d19f50 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/episode_replay_buffer.py @@ -0,0 +1,1475 @@ +import copy +import hashlib +from collections import deque +from typing import Any, Dict, List, Optional, Tuple, Union + +import numpy as np +import scipy + +from ray.rllib.core import DEFAULT_AGENT_ID, DEFAULT_MODULE_ID +from ray.rllib.env.single_agent_episode import SingleAgentEpisode +from ray.rllib.env.utils.infinite_lookback_buffer import InfiniteLookbackBuffer +from ray.rllib.utils import force_list +from ray.rllib.utils.annotations import ( + OverrideToImplementCustomLogic_CallToSuperRecommended, + override, +) +from ray.rllib.utils.metrics import ( + ACTUAL_N_STEP, + AGENT_ACTUAL_N_STEP, + AGENT_STEP_UTILIZATION, + ENV_STEP_UTILIZATION, + MODULE_ACTUAL_N_STEP, + MODULE_STEP_UTILIZATION, + NUM_AGENT_EPISODES_ADDED, + NUM_AGENT_EPISODES_ADDED_LIFETIME, + NUM_AGENT_EPISODES_EVICTED, + NUM_AGENT_EPISODES_EVICTED_LIFETIME, + NUM_AGENT_EPISODES_PER_SAMPLE, + NUM_AGENT_EPISODES_STORED, + NUM_AGENT_RESAMPLES, + NUM_AGENT_STEPS_ADDED, + NUM_AGENT_STEPS_ADDED_LIFETIME, + NUM_AGENT_STEPS_EVICTED, + NUM_AGENT_STEPS_EVICTED_LIFETIME, + NUM_AGENT_STEPS_PER_SAMPLE, + NUM_AGENT_STEPS_PER_SAMPLE_LIFETIME, + NUM_AGENT_STEPS_SAMPLED, + NUM_AGENT_STEPS_SAMPLED_LIFETIME, + NUM_AGENT_STEPS_STORED, + NUM_ENV_STEPS_ADDED, + NUM_ENV_STEPS_ADDED_LIFETIME, + NUM_ENV_STEPS_EVICTED, + NUM_ENV_STEPS_EVICTED_LIFETIME, + NUM_ENV_STEPS_PER_SAMPLE, + NUM_ENV_STEPS_PER_SAMPLE_LIFETIME, + NUM_ENV_STEPS_SAMPLED, + NUM_ENV_STEPS_SAMPLED_LIFETIME, + NUM_ENV_STEPS_STORED, + NUM_EPISODES_ADDED, + NUM_EPISODES_ADDED_LIFETIME, + NUM_EPISODES_EVICTED, + NUM_EPISODES_EVICTED_LIFETIME, + NUM_EPISODES_PER_SAMPLE, + NUM_EPISODES_STORED, + NUM_MODULE_EPISODES_ADDED, + NUM_MODULE_EPISODES_ADDED_LIFETIME, + NUM_MODULE_EPISODES_EVICTED, + NUM_MODULE_EPISODES_EVICTED_LIFETIME, + NUM_MODULE_EPISODES_PER_SAMPLE, + NUM_MODULE_EPISODES_STORED, + NUM_MODULE_RESAMPLES, + NUM_MODULE_STEPS_ADDED, + NUM_MODULE_STEPS_ADDED_LIFETIME, + NUM_MODULE_STEPS_EVICTED, + NUM_MODULE_STEPS_EVICTED_LIFETIME, + NUM_MODULE_STEPS_PER_SAMPLE, + NUM_MODULE_STEPS_PER_SAMPLE_LIFETIME, + NUM_MODULE_STEPS_SAMPLED, + NUM_MODULE_STEPS_SAMPLED_LIFETIME, + NUM_RESAMPLES, +) +from ray.rllib.utils.metrics.metrics_logger import MetricsLogger +from ray.rllib.utils.replay_buffers.base import ReplayBufferInterface +from ray.rllib.utils.typing import AgentID, ModuleID, ResultDict, SampleBatchType + + +class EpisodeReplayBuffer(ReplayBufferInterface): + """Buffer that stores (completed or truncated) episodes by their ID. + + Each "row" (a slot in a deque) in the buffer is occupied by one episode. If an + incomplete episode is added to the buffer and then another chunk of that episode is + added at a later time, the buffer will automatically concatenate the new fragment to + the original episode. This way, episodes can be completed via subsequent `add` + calls. + + Sampling returns batches of size B (number of "rows"), where each row is a + trajectory of length T. Each trajectory contains consecutive timesteps from an + episode, but might not start at the beginning of that episode. Should an episode end + within such a trajectory, a random next episode (starting from its t0) will be + concatenated to that "row". Example: `sample(B=2, T=4)` -> + + 0 .. 1 .. 2 .. 3 <- T-axis + 0 e5 e6 e7 e8 + 1 f2 f3 h0 h2 + ^ B-axis + + .. where e, f, and h are different (randomly picked) episodes, the 0-index (e.g. h0) + indicates the start of an episode, and `f3` is an episode end (gym environment + returned terminated=True or truncated=True). + + 0-indexed returned timesteps contain the reset observation, a dummy 0.0 reward, as + well as the first action taken in the episode (action picked after observing + obs(0)). + The last index in an episode (e.g. f3 in the example above) contains the final + observation of the episode, the final reward received, a dummy action + (repeat the previous action), as well as either terminated=True or truncated=True. + """ + + __slots__ = ( + "capacity", + "batch_size_B", + "batch_length_T", + "episodes", + "episode_id_to_index", + "num_episodes_evicted", + "_indices", + "_num_timesteps", + "_num_timesteps_added", + "sampled_timesteps", + "rng", + ) + + def __init__( + self, + capacity: int = 10000, + *, + batch_size_B: int = 16, + batch_length_T: int = 64, + metrics_num_episodes_for_smoothing: int = 100, + **kwargs, + ): + """Initializes an EpisodeReplayBuffer instance. + + Args: + capacity: The total number of timesteps to be storable in this buffer. + Will start ejecting old episodes once this limit is reached. + batch_size_B: The number of rows in a SampleBatch returned from `sample()`. + batch_length_T: The length of each row in a SampleBatch returned from + `sample()`. + """ + self.capacity = capacity + self.batch_size_B = batch_size_B + self.batch_length_T = batch_length_T + + # The actual episode buffer. We are using a deque here for faster insertion + # (left side) and eviction (right side) of data. + self.episodes = deque() + # Maps (unique) episode IDs to the index under which to find this episode + # within our `self.episodes` deque. + # Note that even after eviction started, the indices in here will NOT be + # changed. We will therefore need to offset all indices in + # `self.episode_id_to_index` by the number of episodes that have already been + # evicted (self._num_episodes_evicted) in order to get the actual index to use + # on `self.episodes`. + self.episode_id_to_index = {} + # The number of episodes that have already been evicted from the buffer + # due to reaching capacity. + self._num_episodes_evicted = 0 + + # List storing all index tuples: (eps_idx, ts_in_eps_idx), where ... + # `eps_idx - self._num_episodes_evicted' is the index into self.episodes. + # `ts_in_eps_idx` is the timestep index within that episode + # (0 = 1st timestep, etc..). + # We sample uniformly from the set of these indices in a `sample()` + # call. + self._indices = [] + + # The size of the buffer in timesteps. + self._num_timesteps = 0 + # The number of timesteps added thus far. + self._num_timesteps_added = 0 + + # How many timesteps have been sampled from the buffer in total? + self.sampled_timesteps = 0 + + self.rng = np.random.default_rng(seed=None) + + # Initialize the metrics. + self.metrics = MetricsLogger() + self._metrics_num_episodes_for_smoothing = metrics_num_episodes_for_smoothing + + # Initialize the metrics. + self.metrics = MetricsLogger() + self._metrics_num_episodes_for_smoothing = metrics_num_episodes_for_smoothing + + @override(ReplayBufferInterface) + def __len__(self) -> int: + return self.get_num_timesteps() + + @override(ReplayBufferInterface) + def add(self, episodes: Union[List["SingleAgentEpisode"], "SingleAgentEpisode"]): + """Converts incoming SampleBatch into a number of SingleAgentEpisode objects. + + Then adds these episodes to the internal deque. + """ + episodes = force_list(episodes) + + # Set up some counters for metrics. + num_env_steps_added = 0 + agent_to_num_steps_added = {DEFAULT_AGENT_ID: 0} + module_to_num_steps_added = {DEFAULT_MODULE_ID: 0} + num_episodes_added = 0 + agent_to_num_episodes_added = {DEFAULT_AGENT_ID: 0} + module_to_num_episodes_added = {DEFAULT_MODULE_ID: 0} + num_episodes_evicted = 0 + agent_to_num_episodes_evicted = {DEFAULT_AGENT_ID: 0} + module_to_num_episodes_evicted = {DEFAULT_MODULE_ID: 0} + num_env_steps_evicted = 0 + agent_to_num_steps_evicted = {DEFAULT_AGENT_ID: 0} + module_to_num_steps_evicted = {DEFAULT_MODULE_ID: 0} + + for eps in episodes: + # Make sure we don't change what's coming in from the user. + # TODO (sven): It'd probably be better to make sure in the EnvRunner to not + # hold on to episodes (for metrics purposes only) that we are returning + # back to the user from `EnvRunner.sample()`. Then we wouldn't have to + # do any copying. Instead, either compile the metrics right away on the + # EnvRunner OR compile metrics entirely on the Algorithm side (this is + # actually preferred). + eps = copy.deepcopy(eps) + + eps_len = len(eps) + # TODO (simon): Check, if we can deprecate these two + # variables and instead peek into the metrics. + self._num_timesteps += eps_len + self._num_timesteps_added += eps_len + num_env_steps_added += eps_len + agent_to_num_steps_added[DEFAULT_AGENT_ID] += eps_len + module_to_num_steps_added[DEFAULT_MODULE_ID] += eps_len + + # Ongoing episode, concat to existing record. + if eps.id_ in self.episode_id_to_index: + eps_idx = self.episode_id_to_index[eps.id_] + existing_eps = self.episodes[eps_idx - self._num_episodes_evicted] + old_len = len(existing_eps) + self._indices.extend([(eps_idx, old_len + i) for i in range(len(eps))]) + existing_eps.concat_episode(eps) + # New episode. Add to end of our episodes deque. + else: + num_episodes_added += 1 + agent_to_num_episodes_added[DEFAULT_AGENT_ID] += 1 + module_to_num_episodes_added[DEFAULT_MODULE_ID] += 1 + self.episodes.append(eps) + eps_idx = len(self.episodes) - 1 + self._num_episodes_evicted + self.episode_id_to_index[eps.id_] = eps_idx + self._indices.extend([(eps_idx, i) for i in range(len(eps))]) + + # Eject old records from front of deque (only if we have more than 1 episode + # in the buffer). + while self._num_timesteps > self.capacity and self.get_num_episodes() > 1: + # Eject oldest episode. + evicted_eps = self.episodes.popleft() + evicted_eps_len = len(evicted_eps) + num_episodes_evicted += 1 + num_env_steps_evicted += evicted_eps_len + agent_to_num_episodes_evicted[DEFAULT_AGENT_ID] += 1 + module_to_num_episodes_evicted[DEFAULT_MODULE_ID] += 1 + agent_to_num_steps_evicted[ + DEFAULT_AGENT_ID + ] += evicted_eps.agent_steps() + module_to_num_steps_evicted[ + DEFAULT_MODULE_ID + ] += evicted_eps.agent_steps() + # Correct our size. + self._num_timesteps -= evicted_eps_len + + # Erase episode from all our indices: + # 1) Main episode index. + evicted_idx = self.episode_id_to_index[evicted_eps.id_] + del self.episode_id_to_index[evicted_eps.id_] + # 2) All timestep indices that this episode owned. + new_indices = [] # New indices that will replace self._indices. + idx_cursor = 0 + # Loop through all (eps_idx, ts_in_eps_idx)-tuples. + for i, idx_tuple in enumerate(self._indices): + # This tuple is part of the evicted episode -> Add everything + # up until here to `new_indices` (excluding this very index, b/c + # it's already part of the evicted episode). + if idx_cursor is not None and idx_tuple[0] == evicted_idx: + new_indices.extend(self._indices[idx_cursor:i]) + # Set to None to indicate we are in the eviction zone. + idx_cursor = None + # We are/have been in the eviction zone (i pointing/pointed to the + # evicted episode) .. + elif idx_cursor is None: + # ... but are now not anymore (i is now an index into a + # non-evicted episode) -> Set cursor to valid int again. + if idx_tuple[0] != evicted_idx: + idx_cursor = i + # But early-out if evicted episode was only 1 single + # timestep long. + if evicted_eps_len == 1: + break + # Early-out: We reached the end of the to-be-evicted episode. + # We can stop searching further here (all following tuples + # will NOT be in the evicted episode). + elif idx_tuple[1] == evicted_eps_len - 1: + assert self._indices[i + 1][0] != idx_tuple[0] + idx_cursor = i + 1 + break + + # Jump over (splice-out) the evicted episode if we are still in the + # eviction zone. + if idx_cursor is not None: + new_indices.extend(self._indices[idx_cursor:]) + + # Reset our `self._indices` to the newly compiled list. + self._indices = new_indices + + # Increase episode evicted counter. + self._num_episodes_evicted += 1 + + # Update the metrics. + self._update_add_metrics( + num_episodes_added=num_episodes_added, + num_env_steps_added=num_env_steps_added, + num_episodes_evicted=num_episodes_evicted, + num_env_steps_evicted=num_env_steps_evicted, + agent_to_num_episodes_added=agent_to_num_episodes_added, + agent_to_num_steps_added=agent_to_num_steps_added, + agent_to_num_episodes_evicted=agent_to_num_episodes_evicted, + agent_to_num_steps_evicted=agent_to_num_steps_evicted, + module_to_num_episodes_added=module_to_num_steps_added, + module_to_num_steps_added=module_to_num_episodes_added, + module_to_num_episodes_evicted=module_to_num_episodes_evicted, + module_to_num_steps_evicted=module_to_num_steps_evicted, + ) + + @OverrideToImplementCustomLogic_CallToSuperRecommended + def _update_add_metrics( + self, + *, + num_episodes_added: int, + num_env_steps_added: int, + num_episodes_evicted: int, + num_env_steps_evicted: int, + agent_to_num_episodes_added: Dict[AgentID, int], + agent_to_num_steps_added: Dict[AgentID, int], + agent_to_num_episodes_evicted: Dict[AgentID, int], + agent_to_num_steps_evicted: Dict[AgentID, int], + module_to_num_steps_added: Dict[ModuleID, int], + module_to_num_episodes_added: Dict[ModuleID, int], + module_to_num_episodes_evicted: Dict[ModuleID, int], + module_to_num_steps_evicted: Dict[ModuleID, int], + **kwargs, + ) -> None: + """Updates the replay buffer's adding metrics. + + Args: + num_episodes_added: The total number of episodes added to the + buffer in the `EpisodeReplayBuffer.add` call. + num_timesteps_added: The total number of environment steps added to the + buffer in the `EpisodeReplayBuffer.add` call. + num_episodes_evicted: The total number of environment steps evicted from + the buffer in the `EpisodeReplayBuffer.add` call. Note, this + does not include the number of episodes evicted before ever + added to the buffer (i.e. can happen in case a lot of episodes + were added and the buffer's capacity is not large enough). + num_env_steps_evicted: he total number of environment steps evicted from + the buffer in the `EpisodeReplayBuffer.add` call. Note, this + does not include the number of steps evicted before ever + added to the buffer (i.e. can happen in case a lot of episodes + were added and the buffer's capacity is not large enough). + agent_to_num_episodes_added: A dictionary with the number of episodes per + agent ID added to the buffer during the `EpisodeReplayBuffer.add` call. + agent_to_num_steps_added: A dictionary with the number of agent steps per + agent ID added to the buffer during the `EpisodeReplayBuffer.add` call. + agent_to_num_episodes_evicted: A dictionary with the number of episodes per + agent ID evicted to the buffer during the `EpisodeReplayBuffer.add` call. + agent_to_num_steps_evicted: A dictionary with the number of agent steps per + agent ID evicted to the buffer during the `EpisodeReplayBuffer.add` call. + module_to_num_episodes_added: A dictionary with the number of episodes per + module ID added to the buffer during the `EpisodeReplayBuffer.add` call. + module_to_num_steps_added: A dictionary with the number of agent steps per + module ID added to the buffer during the `EpisodeReplayBuffer.add` call. + module_to_num_episodes_evicted: A dictionary with the number of episodes per + module ID evicted to the buffer during the `EpisodeReplayBuffer.add` call. + module_to_num_steps_evicted: A dictionary with the number of agent steps per + module ID evicted to the buffer during the `EpisodeReplayBuffer.add` call. + """ + # Whole buffer episode metrics. + self.metrics.log_value( + NUM_EPISODES_STORED, + self.get_num_episodes(), + reduce="mean", + window=self._metrics_num_episodes_for_smoothing, + ) + # Number of new episodes added. Note, this metric could + # be zero when ongoing episodes were logged. + self.metrics.log_value( + NUM_EPISODES_ADDED, + num_episodes_added, + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + NUM_EPISODES_ADDED_LIFETIME, + num_episodes_added, + reduce="sum", + ) + self.metrics.log_value( + NUM_EPISODES_EVICTED, + num_episodes_evicted, + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + NUM_EPISODES_EVICTED_LIFETIME, + num_episodes_evicted, + reduce="sum", + ) + # Whole buffer step metrics. + self.metrics.log_value( + NUM_ENV_STEPS_STORED, + self.get_num_timesteps(), + reduce="mean", + window=self._metrics_num_episodes_for_smoothing, + ) + self.metrics.log_value( + NUM_ENV_STEPS_ADDED, + num_env_steps_added, + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + NUM_ENV_STEPS_ADDED_LIFETIME, + num_env_steps_added, + reduce="sum", + ) + self.metrics.log_value( + NUM_ENV_STEPS_EVICTED, + num_env_steps_evicted, + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + NUM_ENV_STEPS_EVICTED_LIFETIME, + num_env_steps_evicted, + reduce="sum", + ) + + # Log per-agent metrics. + for aid in agent_to_num_episodes_added: + # Number of new episodes added. Note, this metric could + # be zero. + self.metrics.log_value( + (NUM_AGENT_EPISODES_ADDED, aid), + agent_to_num_episodes_added[aid], + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + (NUM_AGENT_EPISODES_ADDED_LIFETIME, aid), + agent_to_num_episodes_added[aid], + reduce="sum", + ) + # Number of new agent steps added. Note, this metric could + # be zero, too. + self.metrics.log_value( + (NUM_AGENT_STEPS_ADDED, aid), + agent_to_num_steps_added[aid], + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + (NUM_AGENT_STEPS_ADDED_LIFETIME, aid), + agent_to_num_steps_added[aid], + reduce="sum", + ) + for aid in agent_to_num_episodes_evicted: + # Number of agent episodes evicted. Note, values could be zero. + self.metrics.log_value( + (NUM_AGENT_EPISODES_EVICTED, aid), + agent_to_num_episodes_evicted[aid], + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + (NUM_AGENT_EPISODES_EVICTED_LIFETIME, aid), + agent_to_num_episodes_evicted[aid], + reduce="sum", + ) + # Number of agent steps evicted. Note, values could be zero. + self.metrics.log_value( + (NUM_AGENT_STEPS_EVICTED, aid), + agent_to_num_steps_evicted[aid], + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + (NUM_AGENT_STEPS_EVICTED_LIFETIME, aid), + agent_to_num_steps_evicted[aid], + reduce="sum", + ) + # Note, we need to loop through the metrics here to receive + # metrics for all agents (not only the ones that steps). + for aid in self.metrics.stats[NUM_AGENT_STEPS_ADDED_LIFETIME]: + # Add default metrics for evicted episodes if not existent. + if aid not in agent_to_num_episodes_evicted: + self.metrics.log_value( + (NUM_AGENT_EPISODES_EVICTED, aid), + 0, + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + (NUM_AGENT_EPISODES_EVICTED_LIFETIME, aid), + 0, + reduce="sum", + ) + self.metrics.log_value( + (NUM_AGENT_STEPS_EVICTED, aid), + 0, + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + (NUM_AGENT_STEPS_EVICTED_LIFETIME, aid), + 0, + reduce="sum", + ) + # Number of episodes in the buffer. + agent_episodes_stored = self.metrics.peek( + (NUM_AGENT_EPISODES_ADDED_LIFETIME, aid) + ) - self.metrics.peek((NUM_AGENT_EPISODES_EVICTED_LIFETIME, aid)) + self.metrics.log_value( + (NUM_AGENT_EPISODES_STORED, aid), + agent_episodes_stored, + reduce="mean", + window=self._metrics_num_episodes_for_smoothing, + ) + # Number of agent steps in the buffer. + agent_steps_stored = self.metrics.peek( + (NUM_AGENT_STEPS_ADDED_LIFETIME, aid) + ) - self.metrics.peek((NUM_AGENT_EPISODES_EVICTED_LIFETIME, aid)) + + self.metrics.log_value( + (NUM_AGENT_STEPS_STORED, aid), + agent_steps_stored, + reduce="mean", + window=self._metrics_num_episodes_for_smoothing, + ) + + # Log per-module metrics. + for mid in module_to_num_episodes_added: + # Number of new episodes added. Note, this metric could + # be zero. + self.metrics.log_value( + (NUM_MODULE_EPISODES_ADDED, mid), + module_to_num_episodes_added[mid], + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + (NUM_MODULE_EPISODES_ADDED_LIFETIME, mid), + module_to_num_episodes_added[mid], + reduce="sum", + ) + # Number of new module steps added. Note, this metric could + # be zero, too. + self.metrics.log_value( + (NUM_MODULE_STEPS_ADDED, mid), + module_to_num_steps_added[mid], + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + (NUM_MODULE_STEPS_ADDED_LIFETIME, mid), + module_to_num_steps_added[mid], + reduce="sum", + ) + for mid in module_to_num_episodes_evicted: + # Number of module episodes evicted. Note, values could be zero. + self.metrics.log_value( + (NUM_MODULE_EPISODES_EVICTED, mid), + module_to_num_episodes_evicted[mid], + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + (NUM_MODULE_EPISODES_EVICTED_LIFETIME, mid), + module_to_num_episodes_evicted[mid], + reduce="sum", + ) + # Number of module steps evicted. Note, values could be zero. + self.metrics.log_value( + (NUM_MODULE_STEPS_EVICTED, mid), + module_to_num_steps_evicted[mid], + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + (NUM_MODULE_STEPS_EVICTED_LIFETIME, mid), + module_to_num_steps_evicted[mid], + reduce="sum", + ) + # Note, we need to loop through the metrics here to receive + # metrics for all agents (not only the ones that steps). + for mid in self.metrics.stats[NUM_MODULE_STEPS_ADDED_LIFETIME]: + # Number of episodes in the buffer. + if mid not in module_to_num_episodes_evicted: + self.metrics.log_value( + (NUM_MODULE_EPISODES_EVICTED, mid), + 0, + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + (NUM_MODULE_EPISODES_EVICTED_LIFETIME, mid), + 0, + reduce="sum", + ) + self.metrics.log_value( + (NUM_MODULE_STEPS_EVICTED, mid), + 0, + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + (NUM_MODULE_STEPS_EVICTED_LIFETIME, mid), + 0, + reduce="sum", + ) + module_episodes_stored = self.metrics.peek( + (NUM_MODULE_EPISODES_ADDED_LIFETIME, mid) + ) - self.metrics.peek((NUM_MODULE_EPISODES_EVICTED_LIFETIME, mid)) + self.metrics.log_value( + (NUM_MODULE_EPISODES_STORED, mid), + module_episodes_stored, + reduce="mean", + window=self._metrics_num_episodes_for_smoothing, + ) + + @override(ReplayBufferInterface) + def sample( + self, + num_items: Optional[int] = None, + *, + batch_size_B: Optional[int] = None, + batch_length_T: Optional[int] = None, + n_step: Optional[Union[int, Tuple]] = None, + beta: float = 0.0, + gamma: float = 0.99, + include_infos: bool = False, + include_extra_model_outputs: bool = False, + sample_episodes: Optional[bool] = False, + to_numpy: bool = False, + # TODO (simon): Check, if we need here 1 as default. + lookback: int = 0, + min_batch_length_T: int = 0, + **kwargs, + ) -> Union[SampleBatchType, SingleAgentEpisode]: + """Samples from a buffer in a randomized way. + + Each sampled item defines a transition of the form: + + `(o_t, a_t, sum(r_(t+1:t+n+1)), o_(t+n), terminated_(t+n), truncated_(t+n))` + + where `o_t` is drawn by randomized sampling.`n` is defined by the `n_step` + applied. + + If requested, `info`s of a transitions last timestep `t+n` and respective + extra model outputs (e.g. action log-probabilities) are added to + the batch. + + Args: + num_items: Number of items (transitions) to sample from this + buffer. + batch_size_B: The number of rows (transitions) to return in the + batch + batch_length_T: THe sequence length to sample. At this point in time + only sequences of length 1 are possible. + n_step: The n-step to apply. For the default the batch contains in + `"new_obs"` the observation and in `"obs"` the observation `n` + time steps before. The reward will be the sum of rewards + collected in between these two observations and the action will + be the one executed n steps before such that we always have the + state-action pair that triggered the rewards. + If `n_step` is a tuple, it is considered as a range to sample + from. If `None`, we use `n_step=1`. + gamma: The discount factor to be used when applying n-step calculations. + The default of `0.99` should be replaced by the `Algorithm`s + discount factor. + include_infos: A boolean indicating, if `info`s should be included in + the batch. This could be of advantage, if the `info` contains + values from the environment important for loss computation. If + `True`, the info at the `"new_obs"` in the batch is included. + include_extra_model_outputs: A boolean indicating, if + `extra_model_outputs` should be included in the batch. This could be + of advantage, if the `extra_mdoel_outputs` contain outputs from the + model important for loss computation and only able to compute with the + actual state of model e.g. action log-probabilities, etc.). If `True`, + the extra model outputs at the `"obs"` in the batch is included (the + timestep at which the action is computed). + to_numpy: If episodes should be numpy'ized. + lookback: A desired lookback. Any non-negative integer is valid. + min_batch_length_T: An optional minimal length when sampling sequences. It + ensures that sampled sequences are at least `min_batch_length_T` time + steps long. This can be used to prevent empty sequences during + learning, when using a burn-in period for stateful `RLModule`s. In rare + cases, such as when episodes are very short early in training, this may + result in longer sampling times. + + Returns: + Either a batch with transitions in each row or (if `return_episodes=True`) + a list of 1-step long episodes containing all basic episode data and if + requested infos and extra model outputs. + """ + + if sample_episodes: + return self._sample_episodes( + num_items=num_items, + batch_size_B=batch_size_B, + batch_length_T=batch_length_T, + n_step=n_step, + beta=beta, + gamma=gamma, + include_infos=include_infos, + include_extra_model_outputs=include_extra_model_outputs, + to_numpy=to_numpy, + lookback=lookback, + min_batch_length_T=min_batch_length_T, + ) + else: + return self._sample_batch( + num_items=num_items, + batch_size_B=batch_size_B, + batch_length_T=batch_length_T, + ) + + def _sample_batch( + self, + num_items: Optional[int] = None, + *, + batch_size_B: Optional[int] = None, + batch_length_T: Optional[int] = None, + ) -> SampleBatchType: + """Returns a batch of size B (number of "rows"), where each row has length T. + + Each row contains consecutive timesteps from an episode, but might not start + at the beginning of that episode. Should an episode end within such a + row (trajectory), a random next episode (starting from its t0) will be + concatenated to that row. For more details, see the docstring of the + EpisodeReplayBuffer class. + + Args: + num_items: See `batch_size_B`. For compatibility with the + `ReplayBufferInterface` abstract base class. + batch_size_B: The number of rows (trajectories) to return in the batch. + batch_length_T: The length of each row (in timesteps) to return in the + batch. + + Returns: + The sampled batch (observations, actions, rewards, terminateds, truncateds) + of dimensions [B, T, ...]. + """ + if num_items is not None: + assert batch_size_B is None, ( + "Cannot call `sample()` with both `num_items` and `batch_size_B` " + "provided! Use either one." + ) + batch_size_B = num_items + + # Use our default values if no sizes/lengths provided. + batch_size_B = batch_size_B or self.batch_size_B + batch_length_T = batch_length_T or self.batch_length_T + + # Rows to return. + observations = [[] for _ in range(batch_size_B)] + actions = [[] for _ in range(batch_size_B)] + rewards = [[] for _ in range(batch_size_B)] + is_first = [[False] * batch_length_T for _ in range(batch_size_B)] + is_last = [[False] * batch_length_T for _ in range(batch_size_B)] + is_terminated = [[False] * batch_length_T for _ in range(batch_size_B)] + is_truncated = [[False] * batch_length_T for _ in range(batch_size_B)] + + # Store the unique episode buffer indices to determine sample variation. + sampled_episode_idxs = set() + # Store the unique env step buffer indices to determine sample variation + sampled_env_step_idxs = set() + + B = 0 + T = 0 + while B < batch_size_B: + # Pull a new uniform random index tuple: (eps_idx, ts_in_eps_idx). + index_tuple = self._indices[self.rng.integers(len(self._indices))] + + # Compute the actual episode index (offset by the number of + # already evicted episodes). + episode_idx, episode_ts = ( + index_tuple[0] - self._num_episodes_evicted, + index_tuple[1], + ) + episode = self.episodes[episode_idx] + + # Starting a new chunk, set is_first to True. + is_first[B][T] = True + + # Begin of new batch item (row). + if len(rewards[B]) == 0: + # And we are at the start of an episode: Set reward to 0.0. + if episode_ts == 0: + rewards[B].append(0.0) + # We are in the middle of an episode: Set reward to the previous + # timestep's values. + else: + rewards[B].append(episode.rewards[episode_ts - 1]) + # We are in the middle of a batch item (row). Concat next episode to this + # row from the next episode's beginning. In other words, we never concat + # a middle of an episode to another truncated one. + else: + episode_ts = 0 + rewards[B].append(0.0) + + observations[B].extend(episode.observations[episode_ts:]) + # Repeat last action to have the same number of actions than observations. + actions[B].extend(episode.actions[episode_ts:]) + actions[B].append(episode.actions[-1]) + # Number of rewards are also the same as observations b/c we have the + # initial 0.0 one. + rewards[B].extend(episode.rewards[episode_ts:]) + assert len(observations[B]) == len(actions[B]) == len(rewards[B]) + + T = min(len(observations[B]), batch_length_T) + + # Set is_last=True. + is_last[B][T - 1] = True + # If episode is terminated and we have reached the end of it, set + # is_terminated=True. + if episode.is_terminated and T == len(observations[B]): + is_terminated[B][T - 1] = True + # If episode is truncated and we have reached the end of it, set + # is_truncated=True. + elif episode.is_truncated and T == len(observations[B]): + is_truncated[B][T - 1] = True + + # We are done with this batch row. + if T == batch_length_T: + # We may have overfilled this row: Clip trajectory at the end. + observations[B] = observations[B][:batch_length_T] + actions[B] = actions[B][:batch_length_T] + rewards[B] = rewards[B][:batch_length_T] + # Start filling the next row. + B += 1 + T = 0 + # Add the episode buffer index to the set of episode indexes. + sampled_episode_idxs.add(episode_idx) + # Add the unique hashcode for the timestep. + sampled_env_step_idxs.add( + hashlib.sha256(f"{episode.id_}-{episode_ts}".encode()).hexdigest() + ) + + # Update our sampled counter. + self.sampled_timesteps += batch_size_B * batch_length_T + + # Update the sample metrics. + num_env_steps_sampled = batch_size_B * batch_length_T + num_episodes_per_sample = len(sampled_episode_idxs) + num_env_steps_per_sample = len(sampled_env_step_idxs) + num_resamples = 0 + agent_to_sample_size = {DEFAULT_AGENT_ID: num_env_steps_sampled} + agent_to_num_episodes_per_sample = {DEFAULT_AGENT_ID: num_episodes_per_sample} + agent_to_num_steps_per_sample = {DEFAULT_AGENT_ID: num_env_steps_per_sample} + agent_to_num_resamples = {DEFAULT_AGENT_ID: num_resamples} + module_to_num_steps_sampled = {DEFAULT_MODULE_ID: num_env_steps_sampled} + module_to_num_episodes_per_sample = {DEFAULT_MODULE_ID: num_episodes_per_sample} + module_to_num_steps_per_sample = {DEFAULT_MODULE_ID: num_env_steps_per_sample} + module_to_num_resamples = {DEFAULT_MODULE_ID: num_resamples} + self._update_sample_metrics( + num_env_steps_sampled=num_env_steps_sampled, + num_episodes_per_sample=num_episodes_per_sample, + num_env_steps_per_sample=num_env_steps_per_sample, + num_resamples=num_resamples, + sampled_n_step=None, + agent_to_num_steps_sampled=agent_to_sample_size, + agent_to_num_episodes_per_sample=agent_to_num_episodes_per_sample, + agent_to_num_steps_per_sample=agent_to_num_steps_per_sample, + agent_to_num_resamples=agent_to_num_resamples, + agent_to_sampled_n_step=None, + module_to_num_steps_sampled=module_to_num_steps_sampled, + module_to_num_episodes_per_sample=module_to_num_episodes_per_sample, + module_to_num_steps_per_sample=module_to_num_steps_per_sample, + module_to_num_resamples=module_to_num_resamples, + module_to_sampled_n_step=None, + ) + + # TODO: Return SampleBatch instead of this simpler dict. + ret = { + "obs": np.array(observations), + "actions": np.array(actions), + "rewards": np.array(rewards), + "is_first": np.array(is_first), + "is_last": np.array(is_last), + "is_terminated": np.array(is_terminated), + "is_truncated": np.array(is_truncated), + } + + return ret + + def _sample_episodes( + self, + num_items: Optional[int] = None, + *, + batch_size_B: Optional[int] = None, + batch_length_T: Optional[int] = None, + n_step: Optional[Union[int, Tuple]] = None, + gamma: float = 0.99, + include_infos: bool = False, + include_extra_model_outputs: bool = False, + to_numpy: bool = False, + lookback: int = 1, + min_batch_length_T: int = 0, + **kwargs, + ) -> List[SingleAgentEpisode]: + """Samples episodes from a buffer in a randomized way. + + Each sampled item defines a transition of the form: + + `(o_t, a_t, sum(r_(t+1:t+n+1)), o_(t+n), terminated_(t+n), truncated_(t+n))` + + where `o_t` is drawn by randomized sampling.`n` is defined by the `n_step` + applied. + + If requested, `info`s of a transitions last timestep `t+n` and respective + extra model outputs (e.g. action log-probabilities) are added to + the batch. + + Args: + num_items: Number of items (transitions) to sample from this + buffer. + batch_size_B: The number of rows (transitions) to return in the + batch + batch_length_T: The sequence length to sample. Can be either `None` + (the default) or any positive integer. + n_step: The n-step to apply. For the default the batch contains in + `"new_obs"` the observation and in `"obs"` the observation `n` + time steps before. The reward will be the sum of rewards + collected in between these two observations and the action will + be the one executed n steps before such that we always have the + state-action pair that triggered the rewards. + If `n_step` is a tuple, it is considered as a range to sample + from. If `None`, we use `n_step=1`. + gamma: The discount factor to be used when applying n-step calculations. + The default of `0.99` should be replaced by the `Algorithm`s + discount factor. + include_infos: A boolean indicating, if `info`s should be included in + the batch. This could be of advantage, if the `info` contains + values from the environment important for loss computation. If + `True`, the info at the `"new_obs"` in the batch is included. + include_extra_model_outputs: A boolean indicating, if + `extra_model_outputs` should be included in the batch. This could be + of advantage, if the `extra_mdoel_outputs` contain outputs from the + model important for loss computation and only able to compute with the + actual state of model e.g. action log-probabilities, etc.). If `True`, + the extra model outputs at the `"obs"` in the batch is included (the + timestep at which the action is computed). + to_numpy: If episodes should be numpy'ized. + lookback: A desired lookback. Any non-negative integer is valid. + min_batch_length_T: An optional minimal length when sampling sequences. It + ensures that sampled sequences are at least `min_batch_length_T` time + steps long. This can be used to prevent empty sequences during + learning, when using a burn-in period for stateful `RLModule`s. In rare + cases, such as when episodes are very short early in training, this may + result in longer sampling times. + + Returns: + A list of 1-step long episodes containing all basic episode data and if + requested infos and extra model outputs. + """ + if num_items is not None: + assert batch_size_B is None, ( + "Cannot call `sample()` with both `num_items` and `batch_size_B` " + "provided! Use either one." + ) + batch_size_B = num_items + + # Use our default values if no sizes/lengths provided. + batch_size_B = batch_size_B or self.batch_size_B + + assert n_step is not None, ( + "When sampling episodes, `n_step` must be " + "provided, but `n_step` is `None`." + ) + # If no sequence should be sampled, we sample n-steps. + if not batch_length_T: + # Sample the `n_step`` itself, if necessary. + actual_n_step = n_step + random_n_step = isinstance(n_step, tuple) + # Otherwise we use an n-step of 1. + else: + assert ( + not isinstance(n_step, tuple) and n_step == 1 + ), "When sampling sequences n-step must be 1." + actual_n_step = n_step + + # Keep track of the indices that were sampled last for updating the + # weights later (see `ray.rllib.utils.replay_buffer.utils. + # update_priorities_in_episode_replay_buffer`). + self._last_sampled_indices = [] + + sampled_episodes = [] + # Record all the env step buffer indices that are contained in the sample. + sampled_env_step_idxs = set() + # Record all the episode buffer indices that are contained in the sample. + sampled_episode_idxs = set() + # Record all n-steps that have been used. + sampled_n_steps = [] + # Record all the env step buffer indices that are contained in the sample. + sampled_env_step_idxs = set() + # Record all the episode buffer indices that are contained in the sample. + sampled_episode_idxs = set() + # Record all n-steps that have been used. + sampled_n_steps = [] + + B = 0 + while B < batch_size_B: + # Pull a new uniform random index tuple: (eps_idx, ts_in_eps_idx). + index_tuple = self._indices[self.rng.integers(len(self._indices))] + + # Compute the actual episode index (offset by the number of + # already evicted episodes). + episode_idx, episode_ts = ( + index_tuple[0] - self._num_episodes_evicted, + index_tuple[1], + ) + episode = self.episodes[episode_idx] + + # If we use random n-step sampling, draw the n-step for this item. + if not batch_length_T and random_n_step: + actual_n_step = int(self.rng.integers(n_step[0], n_step[1])) + + # Skip, if we are too far to the end and `episode_ts` + n_step would go + # beyond the episode's end. + if min_batch_length_T > 0 and episode_ts + min_batch_length_T >= len( + episode + ): + continue + if episode_ts + (batch_length_T or 0) + (actual_n_step - 1) > len(episode): + actual_length = len(episode) + else: + actual_length = episode_ts + (batch_length_T or 0) + (actual_n_step - 1) + + # If no sequence should be sampled, we sample here the n-step. + if not batch_length_T: + sampled_episode = episode.slice( + slice( + episode_ts, + episode_ts + actual_n_step, + ) + ) + # Note, this will be the reward after executing action + # `a_(episode_ts-n_step+1)`. For `n_step>1` this will be the discounted + # sum of all discounted rewards that were collected over the last n + # steps. + raw_rewards = sampled_episode.get_rewards() + + rewards = scipy.signal.lfilter( + [1], [1, -gamma], raw_rewards[::-1], axis=0 + )[-1] + + sampled_episode = SingleAgentEpisode( + id_=sampled_episode.id_, + agent_id=sampled_episode.agent_id, + module_id=sampled_episode.module_id, + observation_space=sampled_episode.observation_space, + action_space=sampled_episode.action_space, + observations=[ + sampled_episode.get_observations(0), + sampled_episode.get_observations(-1), + ], + actions=[sampled_episode.get_actions(0)], + rewards=[rewards], + infos=[ + sampled_episode.get_infos(0), + sampled_episode.get_infos(-1), + ], + terminated=sampled_episode.is_terminated, + truncated=sampled_episode.is_truncated, + extra_model_outputs={ + **( + { + k: [episode.get_extra_model_outputs(k, 0)] + for k in episode.extra_model_outputs.keys() + } + if include_extra_model_outputs + else {} + ), + }, + t_started=episode_ts, + len_lookback_buffer=0, + ) + # Otherwise we simply slice the episode. + else: + sampled_episode = episode.slice( + slice( + episode_ts, + actual_length, + ), + len_lookback_buffer=lookback, + ) + # Record a has for the episode ID and timestep inside of the episode. + sampled_env_step_idxs.add( + hashlib.sha256(f"{episode.id_}-{episode_ts}".encode()).hexdigest() + ) + # Record a has for the episode ID and timestep inside of the episode. + sampled_env_step_idxs.add( + hashlib.sha256(f"{episode.id_}-{episode_ts}".encode()).hexdigest() + ) + # Remove reference to sampled episode. + del episode + + # Add the actually chosen n-step in this episode. + sampled_episode.extra_model_outputs["n_step"] = InfiniteLookbackBuffer( + np.full((len(sampled_episode) + lookback,), actual_n_step), + lookback=lookback, + ) + # Some loss functions need `weights` - which are only relevant when + # prioritizing. + sampled_episode.extra_model_outputs["weights"] = InfiniteLookbackBuffer( + np.ones((len(sampled_episode) + lookback,)), lookback=lookback + ) + + # Append the sampled episode. + sampled_episodes.append(sampled_episode) + sampled_episode_idxs.add(episode_idx) + sampled_n_steps.append(actual_n_step) + + # Increment counter. + B += (actual_length - episode_ts - (actual_n_step - 1) + 1) or 1 + + # Update the metric. + self.sampled_timesteps += batch_size_B + + # Update the sample metrics. + num_env_steps_sampled = batch_size_B + num_episodes_per_sample = len(sampled_episode_idxs) + num_env_steps_per_sample = len(sampled_env_step_idxs) + sampled_n_step = sum(sampled_n_steps) / batch_size_B + num_resamples = 0 + agent_to_num_steps_sampled = {DEFAULT_AGENT_ID: num_env_steps_sampled} + agent_to_num_episodes_per_sample = {DEFAULT_AGENT_ID: num_episodes_per_sample} + agent_to_num_steps_per_sample = {DEFAULT_AGENT_ID: num_env_steps_per_sample} + agent_to_sampled_n_step = {DEFAULT_AGENT_ID: sampled_n_step} + agent_to_num_resamples = {DEFAULT_AGENT_ID: num_resamples} + module_to_num_steps_sampled = {DEFAULT_MODULE_ID: num_env_steps_sampled} + module_to_num_episodes_per_sample = {DEFAULT_MODULE_ID: num_episodes_per_sample} + module_to_num_steps_per_sample = {DEFAULT_MODULE_ID: num_env_steps_per_sample} + module_to_sampled_n_step = {DEFAULT_MODULE_ID: sampled_n_step} + module_to_num_resamples = {DEFAULT_MODULE_ID: num_resamples} + self._update_sample_metrics( + num_env_steps_sampled=num_env_steps_sampled, + num_episodes_per_sample=num_episodes_per_sample, + num_env_steps_per_sample=num_env_steps_per_sample, + sampled_n_step=sampled_n_step, + num_resamples=num_resamples, + agent_to_num_steps_sampled=agent_to_num_steps_sampled, + agent_to_num_episodes_per_sample=agent_to_num_episodes_per_sample, + agent_to_num_steps_per_sample=agent_to_num_steps_per_sample, + agent_to_sampled_n_step=agent_to_sampled_n_step, + agent_to_num_resamples=agent_to_num_resamples, + module_to_num_steps_sampled=module_to_num_steps_sampled, + module_to_num_episodes_per_sample=module_to_num_episodes_per_sample, + module_to_num_steps_per_sample=module_to_num_steps_per_sample, + module_to_sampled_n_step=module_to_sampled_n_step, + module_to_num_resamples=module_to_num_resamples, + ) + + return sampled_episodes + + @OverrideToImplementCustomLogic_CallToSuperRecommended + def _update_sample_metrics( + self, + *, + num_env_steps_sampled: int, + num_episodes_per_sample: int, + num_env_steps_per_sample: int, + sampled_n_step: Optional[float], + num_resamples: int, + agent_to_num_steps_sampled: Dict[AgentID, int], + agent_to_num_episodes_per_sample: Dict[AgentID, int], + agent_to_num_steps_per_sample: Dict[AgentID, int], + agent_to_sampled_n_step: Dict[AgentID, float], + agent_to_num_resamples: Dict[AgentID, int], + module_to_num_steps_sampled: Dict[ModuleID, int], + module_to_num_episodes_per_sample: Dict[ModuleID, int], + module_to_num_steps_per_sample: Dict[ModuleID, int], + module_to_sampled_n_step: Dict[ModuleID, float], + module_to_num_resamples: Dict[ModuleID, int], + **kwargs: Dict[str, Any], + ) -> None: + """Updates the replay buffer's sample metrics. + + Args: + num_env_steps_sampled: The number of environment steps sampled + this iteration in the `sample` method. + num_episodes_per_sample: The number of unique episodes in the + sample. + num_env_steps_per_sample: The number of unique environment steps + in the sample. + sampled_n_step: The mean n-step used in the sample. Note, this + is constant, if the n-step is not sampled. + num_resamples: The number of resamples in a single call to + `PrioritizedEpisodeReplayBuffer.sample`. A resampling is triggered + when the sampled timestep is to near to the episode end to cover the + required n-step. + agent_to_num_steps_sampled: A dictionary with the number of agent + steps per agent ID sampled during the `EpisodeReplayBuffer.sample` + call. + agent_to_num_episodes_per_sample: A dictionary with the number of + unique episodes per agent ID contained in the sample returned by + the `EpisodeReplayBuffer.sample` call. + agent_to_num_steps_per_sample: A dictionary with the number of + unique agent steps per agent ID contained in the sample returned by + the `EpisodeReplayBuffer.sample` call. + agent_to_sampled_n_step: A dictionary with the mean n-step per agent ID + used in the sample returned by the `EpisodeReplayBuffer.sample` call. + agent_to_num_resamples: A dictionary with the number of resamples per + agent ID in a single call to `PrioritizedEpisodeReplayBuffer.sample`. + A resampling is triggered when the sampled timestep is to near to the + episode end to cover the required n-step. + module_to_num_steps_sampled: A dictionary with the number of module + steps per module ID sampled during the `EpisodeReplayBuffer.sample` + call. + module_to_num_episodes_per_sample: A dictionary with the number of + unique episodes per module ID contained in the sample returned by + the `EpisodeReplayBuffer.sample` call. + module_to_num_steps_per_sample: A dictionary with the number of + unique module steps per module ID contained in the sample returned by + the `EpisodeReplayBuffer.sample` call. + module_to_sampled_n_step: A dictionary with the mean n-step per module ID + used in the sample returned by the `EpisodeReplayBuffer.sample` call. + module_to_num_resamples: A dictionary with the number of resamples per + module ID in a single call to `PrioritizedEpisodeReplayBuffer.sample`. + A resampling is triggered when the sampled timestep is to near to the + episode end to cover the required n-step. + """ + # Whole buffer sampled env steps metrics. + self.metrics.log_value( + NUM_EPISODES_PER_SAMPLE, + num_episodes_per_sample, + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + NUM_ENV_STEPS_PER_SAMPLE, + num_env_steps_per_sample, + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + NUM_ENV_STEPS_PER_SAMPLE_LIFETIME, + num_env_steps_per_sample, + reduce="sum", + ) + self.metrics.log_value( + NUM_ENV_STEPS_SAMPLED, + num_env_steps_sampled, + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + NUM_ENV_STEPS_SAMPLED_LIFETIME, + num_env_steps_sampled, + reduce="sum", + ) + self.metrics.log_value( + ENV_STEP_UTILIZATION, + self.metrics.peek(NUM_ENV_STEPS_PER_SAMPLE_LIFETIME) + / self.metrics.peek(NUM_ENV_STEPS_SAMPLED_LIFETIME), + reduce="mean", + window=self._metrics_num_episodes_for_smoothing, + ) + if sampled_n_step: + self.metrics.log_value( + ACTUAL_N_STEP, + sampled_n_step, + reduce="mean", + window=self._metrics_num_episodes_for_smoothing, + ) + for aid in agent_to_sampled_n_step: + self.metrics.log_value( + (AGENT_ACTUAL_N_STEP, aid), + agent_to_sampled_n_step[aid], + reduce="mean", + window=self._metrics_num_episodes_for_smoothing, + ) + for mid in module_to_sampled_n_step: + self.metrics.log_value( + (MODULE_ACTUAL_N_STEP, mid), + module_to_sampled_n_step[mid], + reduce="mean", + window=self._metrics_num_episodes_for_smoothing, + ) + + self.metrics.log_value( + NUM_RESAMPLES, + num_resamples, + reduce="sum", + clear_on_reduce=True, + ) + + # Add per-agent metrics. + for aid in agent_to_num_steps_sampled: + self.metrics.log_value( + (NUM_AGENT_STEPS_SAMPLED, aid), + agent_to_num_steps_sampled[aid], + reduce="sum", + clear_on_reduce=True, + ) + # TODO (simon): Check, if we can then deprecate + # self.sampled_timesteps. + self.metrics.log_value( + (NUM_AGENT_STEPS_SAMPLED_LIFETIME, aid), + agent_to_num_steps_sampled[aid], + reduce="sum", + ) + self.metrics.log_value( + (NUM_AGENT_EPISODES_PER_SAMPLE, aid), + agent_to_num_episodes_per_sample[aid], + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + (NUM_AGENT_STEPS_PER_SAMPLE, aid), + agent_to_num_steps_per_sample[aid], + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + (NUM_AGENT_STEPS_PER_SAMPLE_LIFETIME, aid), + agent_to_num_steps_per_sample[aid], + reduce="sum", + ) + self.metrics.log_value( + (AGENT_STEP_UTILIZATION, aid), + self.metrics.peek((NUM_AGENT_STEPS_PER_SAMPLE_LIFETIME, aid)) + / self.metrics.peek((NUM_AGENT_STEPS_SAMPLED_LIFETIME, aid)), + reduce="mean", + window=self._metrics_num_episodes_for_smoothing, + ) + self.metrics.log_value( + (NUM_AGENT_RESAMPLES, aid), + agent_to_num_resamples[aid], + reduce="sum", + clear_on_reduce=True, + ) + # Add per-module metrics. + for mid in module_to_num_steps_sampled: + self.metrics.log_value( + (NUM_MODULE_STEPS_SAMPLED, mid), + module_to_num_steps_sampled[mid], + reduce="sum", + clear_on_reduce=True, + ) + # TODO (simon): Check, if we can then deprecate + # self.sampled_timesteps. + self.metrics.log_value( + (NUM_MODULE_STEPS_SAMPLED_LIFETIME, mid), + module_to_num_steps_sampled[mid], + reduce="sum", + ) + self.metrics.log_value( + (NUM_MODULE_EPISODES_PER_SAMPLE, mid), + module_to_num_episodes_per_sample[mid], + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + (NUM_MODULE_STEPS_PER_SAMPLE, mid), + module_to_num_steps_per_sample[mid], + reduce="sum", + clear_on_reduce=True, + ) + self.metrics.log_value( + (NUM_MODULE_STEPS_PER_SAMPLE_LIFETIME, mid), + module_to_num_steps_per_sample[mid], + reduce="sum", + ) + self.metrics.log_value( + (MODULE_STEP_UTILIZATION, mid), + self.metrics.peek((NUM_MODULE_STEPS_PER_SAMPLE_LIFETIME, mid)) + / self.metrics.peek((NUM_MODULE_STEPS_SAMPLED_LIFETIME, mid)), + reduce="mean", + window=self._metrics_num_episodes_for_smoothing, + ) + self.metrics.log_value( + (NUM_MODULE_RESAMPLES, mid), + module_to_num_resamples[mid], + reduce="sum", + clear_on_reduce=True, + ) + + # TODO (simon): Check, if we can instead peek into the metrics + # and deprecate all variables. + def get_num_episodes(self, module_id: Optional[ModuleID] = None) -> int: + """Returns number of episodes (completed or truncated) stored in the buffer.""" + return len(self.episodes) + + def get_num_episodes_evicted(self, module_id: Optional[ModuleID] = None) -> int: + """Returns number of episodes that have been evicted from the buffer.""" + return self._num_episodes_evicted + + def get_num_timesteps(self, module_id: Optional[ModuleID] = None) -> int: + """Returns number of individual timesteps stored in the buffer.""" + return len(self._indices) + + def get_sampled_timesteps(self, module_id: Optional[ModuleID] = None) -> int: + """Returns number of timesteps that have been sampled in buffer's lifetime.""" + return self.sampled_timesteps + + def get_added_timesteps(self, module_id: Optional[ModuleID] = None) -> int: + """Returns number of timesteps that have been added in buffer's lifetime.""" + return self._num_timesteps_added + + def get_metrics(self) -> ResultDict: + """Returns the metrics of the buffer and reduces them.""" + return self.metrics.reduce() + + @override(ReplayBufferInterface) + def get_state(self) -> Dict[str, Any]: + """Gets a pickable state of the buffer. + + This is used for checkpointing the buffer's state. It is specifically helpful, + for example, when a trial is paused and resumed later on. The buffer's state + can be saved to disk and reloaded when the trial is resumed. + + Returns: + A dict containing all necessary information to restore the buffer's state. + """ + return { + "episodes": [eps.get_state() for eps in self.episodes], + "episode_id_to_index": list(self.episode_id_to_index.items()), + "_num_episodes_evicted": self._num_episodes_evicted, + "_indices": self._indices, + "_num_timesteps": self._num_timesteps, + "_num_timesteps_added": self._num_timesteps_added, + "sampled_timesteps": self.sampled_timesteps, + } + + @override(ReplayBufferInterface) + def set_state(self, state) -> None: + """Sets the state of a buffer from a previously stored state. + + See `get_state()` for more information on what is stored in the state. This + method is used to restore the buffer's state from a previously stored state. + It is specifically helpful, for example, when a trial is paused and resumed + later on. The buffer's state can be saved to disk and reloaded when the trial + is resumed. + + Args: + state: The state to restore the buffer from. + """ + self._set_episodes(state) + self.episode_id_to_index = dict(state["episode_id_to_index"]) + self._num_episodes_evicted = state["_num_episodes_evicted"] + self._indices = state["_indices"] + self._num_timesteps = state["_num_timesteps"] + self._num_timesteps_added = state["_num_timesteps_added"] + self.sampled_timesteps = state["sampled_timesteps"] + + def _set_episodes(self, state) -> None: + """Sets the episodes from the state. + + Note, this method is used for class inheritance purposes. It is specifically + helpful when a subclass of this class wants to override the behavior of how + episodes are set from the state. By default, it sets `SingleAgentEpuisode`s, + but subclasses can override this method to set episodes of a different type. + """ + if not self.episodes: + self.episodes = deque( + [ + SingleAgentEpisode.from_state(eps_data) + for eps_data in state["episodes"] + ] + ) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/fifo_replay_buffer.py b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/fifo_replay_buffer.py new file mode 100644 index 0000000000000000000000000000000000000000..ca7cdbd6510238a88683ad279b9efe105113d821 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/fifo_replay_buffer.py @@ -0,0 +1,110 @@ +from typing import Any, Dict, Optional + +import numpy as np + +from ray.rllib.policy.sample_batch import MultiAgentBatch +from ray.rllib.utils.annotations import override +from ray.rllib.utils.replay_buffers.replay_buffer import ReplayBuffer, StorageUnit +from ray.rllib.utils.typing import SampleBatchType +from ray.util.annotations import DeveloperAPI + + +@DeveloperAPI +class FifoReplayBuffer(ReplayBuffer): + """This replay buffer implements a FIFO queue. + + Sometimes, e.g. for offline use cases, it may be desirable to use + off-policy algorithms without a Replay Buffer. + This FifoReplayBuffer can be used in-place to achieve the same effect + without having to introduce separate algorithm execution branches. + + For simplicity and efficiency reasons, this replay buffer stores incoming + sample batches as-is, and returns them one at time. + This is to avoid any additional load when this replay buffer is used. + """ + + def __init__(self, *args, **kwargs): + """Initializes a FifoReplayBuffer. + + Args: + ``*args`` : Forward compatibility args. + ``**kwargs``: Forward compatibility kwargs. + """ + # Completely by-passing underlying ReplayBuffer by setting its + # capacity to 1 (lowest allowed capacity). + ReplayBuffer.__init__(self, 1, StorageUnit.FRAGMENTS, **kwargs) + + self._queue = [] + + @DeveloperAPI + @override(ReplayBuffer) + def add(self, batch: SampleBatchType, **kwargs) -> None: + return self._queue.append(batch) + + @DeveloperAPI + @override(ReplayBuffer) + def sample(self, *args, **kwargs) -> Optional[SampleBatchType]: + """Sample a saved training batch from this buffer. + + Args: + ``*args`` : Forward compatibility args. + ``**kwargs``: Forward compatibility kwargs. + + Returns: + A single training batch from the queue. + """ + if len(self._queue) <= 0: + # Return empty SampleBatch if queue is empty. + return MultiAgentBatch({}, 0) + batch = self._queue.pop(0) + # Equal weights of 1.0. + batch["weights"] = np.ones(len(batch)) + return batch + + @DeveloperAPI + def update_priorities(self, *args, **kwargs) -> None: + """Update priorities of items at given indices. + + No-op for this replay buffer. + + Args: + ``*args`` : Forward compatibility args. + ``**kwargs``: Forward compatibility kwargs. + """ + pass + + @DeveloperAPI + @override(ReplayBuffer) + def stats(self, debug: bool = False) -> Dict: + """Returns the stats of this buffer. + + Args: + debug: If true, adds sample eviction statistics to the returned stats dict. + + Returns: + A dictionary of stats about this buffer. + """ + # As if this replay buffer has never existed. + return {} + + @DeveloperAPI + @override(ReplayBuffer) + def get_state(self) -> Dict[str, Any]: + """Returns all local state. + + Returns: + The serializable local state. + """ + # Pass through replay buffer does not save states. + return {} + + @DeveloperAPI + @override(ReplayBuffer) + def set_state(self, state: Dict[str, Any]) -> None: + """Restores all local state to the provided `state`. + + Args: + state: The new state to set this buffer. Can be obtained by calling + `self.get_state()`. + """ + pass diff --git a/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/multi_agent_episode_buffer.py b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/multi_agent_episode_buffer.py new file mode 100644 index 0000000000000000000000000000000000000000..f60804f4551a6f19edcb3992d485d87355e13a8a --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/multi_agent_episode_buffer.py @@ -0,0 +1,1180 @@ +import copy +import hashlib +from collections import defaultdict, deque +from typing import Any, Dict, List, Optional, Set, Tuple, Union + +import numpy as np +import scipy +from gymnasium.core import ActType, ObsType + +from ray.rllib.core.columns import Columns +from ray.rllib.env.multi_agent_episode import MultiAgentEpisode +from ray.rllib.env.single_agent_episode import SingleAgentEpisode +from ray.rllib.utils import force_list +from ray.rllib.utils.annotations import ( + DeveloperAPI, + override, +) +from ray.rllib.utils.replay_buffers.episode_replay_buffer import EpisodeReplayBuffer +from ray.rllib.utils.spaces.space_utils import batch +from ray.rllib.utils.typing import AgentID, ModuleID, SampleBatchType + + +@DeveloperAPI +class MultiAgentEpisodeReplayBuffer(EpisodeReplayBuffer): + """Multi-agent episode replay buffer that stores episodes by their IDs. + + This class implements a replay buffer as used in "playing Atari with Deep + Reinforcement Learning" (Mnih et al., 2013) for multi-agent reinforcement + learning, + + Each "row" (a slot in a deque) in the buffer is occupied by one episode. If an + incomplete episode is added to the buffer and then another chunk of that episode is + added at a later time, the buffer will automatically concatenate the new fragment to + the original episode. This way, episodes can be completed via subsequent `add` + calls. + + Sampling returns a size `B` episode list (number of 'rows'), where each episode + holds a tuple tuple of the form + + `(o_t, a_t, sum(r_t+1:t+n), o_t+n)` + + where `o_t` is the observation in `t`, `a_t` the action chosen at observation `o_t`, + `o_t+n` is the observation `n` timesteps later and `sum(r_t+1:t+n)` is the sum of + all rewards collected over the time steps between `t+1` and `t+n`. The `n`-step can + be chosen freely when sampling and defaults to `1`. If `n_step` is a tuple it is + sampled uniformly across the interval defined by the tuple (for each row in the + batch). + + Each episode contains - in addition to the data tuples presented above - two further + elements in its `extra_model_outputs`, namely `n_steps` and `weights`. The former + holds the `n_step` used for the sampled timesteps in the episode and the latter the + corresponding (importance sampling) weight for the transition. + + .. testcode:: + + import gymnasium as gym + + from ray.rllib.env.multi_agent_episode import MultiAgentEpisode + from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole + from ray.rllib.utils.replay_buffers import MultiAgentEpisodeReplayBuffer + + + # Create the environment. + env = MultiAgentCartPole({"num_agents": 2}) + + # Set up the loop variables + agent_ids = env.agents + agent_ids.append("__all__") + terminateds = {aid: False for aid in agent_ids} + truncateds = {aid: False for aid in agent_ids} + num_timesteps = 10000 + episodes = [] + + # Initialize the first episode entries. + eps = MultiAgentEpisode() + obs, infos = env.reset() + eps.add_env_reset(observations=obs, infos=infos) + + # Sample 10,000 env timesteps. + for i in range(num_timesteps): + # If terminated we create a new episode. + if eps.is_done: + episodes.append(eps.to_numpy()) + eps = MultiAgentEpisode() + terminateds = {aid: False for aid in agent_ids} + truncateds = {aid: False for aid in agent_ids} + obs, infos = env.reset() + eps.add_env_reset(observations=obs, infos=infos) + + # Sample a random action for all agents that should step in the episode + # next. + actions = { + aid: env.get_action_space(aid).sample() + for aid in eps.get_agents_to_act() + } + obs, rewards, terminateds, truncateds, infos = env.step(actions) + eps.add_env_step( + obs, + actions, + rewards, + infos, + terminateds=terminateds, + truncateds=truncateds + ) + + # Add the last (truncated) episode to the list of episodes. + if not eps.is_done: + episodes.append(eps) + + # Create the buffer. + buffer = MultiAgentEpisodeReplayBuffer() + # Add the list of episodes sampled. + buffer.add(episodes) + + # Pull a sample from the buffer using an `n-step` of 3. + sample = buffer.sample(num_items=256, gamma=0.95, n_step=3) + """ + + def __init__( + self, + capacity: int = 10000, + *, + batch_size_B: int = 16, + batch_length_T: int = 1, + metrics_num_episodes_for_smoothing: int = 100, + **kwargs, + ): + """Initializes a multi-agent episode replay buffer. + + Args: + capacity: The total number of timesteps to be storable in this buffer. + Will start ejecting old episodes once this limit is reached. + batch_size_B: The number of episodes returned from `sample()`. + batch_length_T: The length of each episode in the episode list returned from + `sample()`. + """ + # Initialize the base episode replay buffer. + super().__init__( + capacity=capacity, + batch_size_B=batch_size_B, + batch_length_T=batch_length_T, + metrics_num_episodes_for_smoothing=metrics_num_episodes_for_smoothing, + **kwargs, + ) + + # Stores indices of module (single-agent) timesteps. Each index is a tuple + # of the form: + # `(ma_episode_idx, agent_id, timestep)`. + # This information is stored for each timestep of an episode and is used in + # the `"independent"`` sampling process. The multi-agent episode index amd the + # agent ID are used to retrieve the single-agent episode. The timestep is then + # needed to retrieve the corresponding timestep data from that single-agent + # episode. + self._module_to_indices: Dict[ + ModuleID, List[Tuple[int, AgentID, int]] + ] = defaultdict(list) + + # Stores the number of single-agent timesteps in the buffer. + self._num_agent_timesteps: int = 0 + # Stores the number of single-agent timesteps per module. + self._num_module_timesteps: Dict[ModuleID, int] = defaultdict(int) + + # Stores the number of added single-agent timesteps over the + # lifetime of the buffer. + self._num_agent_timesteps_added: int = 0 + # Stores the number of added single-agent timesteps per module + # over the lifetime of the buffer. + self._num_module_timesteps_added: Dict[ModuleID, int] = defaultdict(int) + + self._num_module_episodes: Dict[ModuleID, int] = defaultdict(int) + # Stores the number of module episodes evicted. Note, this is + # important for indexing. + self._num_module_episodes_evicted: Dict[ModuleID, int] = defaultdict(int) + + # Stores hte number of module timesteps sampled. + self.sampled_timesteps_per_module: Dict[ModuleID, int] = defaultdict(int) + + @override(EpisodeReplayBuffer) + def add( + self, + episodes: Union[List["MultiAgentEpisode"], "MultiAgentEpisode"], + ) -> None: + """Adds episodes to the replay buffer. + + Note, if the incoming episodes' time steps cause the buffer to overflow, + older episodes are evicted. Because episodes usually come in chunks and + not complete, this could lead to edge cases (e.g. with very small capacity + or very long episode length) where the first part of an episode is evicted + while the next part just comes in. + To defend against such case, the complete episode is evicted, including + the new chunk, unless the episode is the only one in the buffer. In the + latter case the buffer will be allowed to overflow in a temporary fashion, + i.e. during the next addition of samples to the buffer an attempt is made + to fall below capacity again. + + The user is advised to select a large enough buffer with regard to the maximum + expected episode length. + + Args: + episodes: The multi-agent episodes to add to the replay buffer. Can be a + single episode or a list of episodes. + """ + episodes: List["MultiAgentEpisode"] = force_list(episodes) + + new_episode_ids: Set[str] = {eps.id_ for eps in episodes} + total_env_timesteps = sum([eps.env_steps() for eps in episodes]) + self._num_timesteps += total_env_timesteps + self._num_timesteps_added += total_env_timesteps + + # Set up some counters for metrics. + num_env_steps_added = 0 + agent_to_num_steps_added = defaultdict(int) + module_to_num_steps_added = defaultdict(int) + num_episodes_added = 0 + agent_to_num_episodes_added = defaultdict(int) + module_to_num_episodes_added = defaultdict(int) + num_episodes_evicted = 0 + agent_to_num_episodes_evicted = defaultdict(int) + module_to_num_episodes_evicted = defaultdict(int) + num_env_steps_evicted = 0 + agent_to_num_steps_evicted = defaultdict(int) + module_to_num_steps_evicted = defaultdict(int) + + # Evict old episodes. + eps_evicted_ids: Set[Union[str, int]] = set() + eps_evicted_idxs: Set[int] = set() + while ( + self._num_timesteps > self.capacity + and self._num_remaining_episodes(new_episode_ids, eps_evicted_ids) != 1 + ): + # Evict episode. + evicted_episode = self.episodes.popleft() + eps_evicted_ids.add(evicted_episode.id_) + eps_evicted_idxs.add(self.episode_id_to_index.pop(evicted_episode.id_)) + # If this episode has a new chunk in the new episodes added, + # we subtract it again. + # TODO (sven, simon): Should we just treat such an episode chunk + # as a new episode? + if evicted_episode.id_ in new_episode_ids: + idx = next( + i + for i, eps in enumerate(episodes) + if eps.id_ == evicted_episode.id_ + ) + new_eps_to_evict = episodes.pop(idx) + self._num_timesteps -= new_eps_to_evict.env_steps() + self._num_timesteps_added -= new_eps_to_evict.env_steps() + # Remove the timesteps of the evicted episode from the counter. + self._num_timesteps -= evicted_episode.env_steps() + self._num_agent_timesteps -= evicted_episode.agent_steps() + self._num_episodes_evicted += 1 + # Increase the counters. + num_episodes_evicted += 1 + num_env_steps_evicted += evicted_episode.env_steps() + for aid, a_eps in evicted_episode.agent_episodes.items(): + mid = evicted_episode._agent_to_module_mapping[aid] + agent_to_num_episodes_evicted[aid] += 1 + module_to_num_episodes_evicted[mid] += 1 + agent_to_num_steps_evicted[aid] += a_eps.agent_steps() + module_to_num_steps_evicted[mid] += a_eps.agent_steps() + # Remove the module timesteps of the evicted episode from the counters. + self._evict_module_episodes(evicted_episode) + del evicted_episode + + # Add agent and module steps. + for eps in episodes: + self._num_agent_timesteps += eps.agent_steps() + self._num_agent_timesteps_added += eps.agent_steps() + # Update the module counters by the module timesteps. + self._update_module_counters(eps) + + # Remove corresponding indices, if episodes were evicted. + if eps_evicted_idxs: + # If the episode is not evicted, we keep the index. + # Note, each index 2-tuple is of the form (ma_episode_idx, timestep) + # and refers to a certain environment timestep in a certain + # multi-agent episode. + self._indices = [ + idx_tuple + for idx_tuple in self._indices + if idx_tuple[0] not in eps_evicted_idxs + ] + # Also remove corresponding module indices. + for module_id, module_indices in self._module_to_indices.items(): + # Each index 3-tuple is of the form + # (ma_episode_idx, agent_id, timestep) and refers to a certain + # agent timestep in a certain multi-agent episode. + self._module_to_indices[module_id] = [ + idx_triplet + for idx_triplet in module_indices + if idx_triplet[0] not in eps_evicted_idxs + ] + + for eps in episodes: + eps = copy.deepcopy(eps) + # If the episode is part of an already existing episode, concatenate. + if eps.id_ in self.episode_id_to_index: + eps_idx = self.episode_id_to_index[eps.id_] + existing_eps = self.episodes[eps_idx - self._num_episodes_evicted] + existing_len = len(existing_eps) + self._indices.extend( + [ + ( + eps_idx, + existing_len + i, + ) + for i in range(len(eps)) + ] + ) + # Add new module indices. + self._add_new_module_indices(eps, eps_idx, True) + # Concatenate the episode chunk. + existing_eps.concat_episode(eps) + # Otherwise, create a new entry. + else: + # New episode. + self.episodes.append(eps) + # Update the counters + num_episodes_added += 1 + for aid, a_eps in eps.agent_episodes.items(): + mid = eps._agent_to_module_mapping[aid] + agent_to_num_episodes_added[aid] += 1 + module_to_num_episodes_added[mid] += 1 + eps_idx = len(self.episodes) - 1 + self._num_episodes_evicted + self.episode_id_to_index[eps.id_] = eps_idx + self._indices.extend([(eps_idx, i) for i in range(len(eps))]) + # Add new module indices. + self._add_new_module_indices(eps, eps_idx, False) + # Update the step counters. + num_env_steps_added += eps.env_steps() + for aid, e_eps in eps.agent_episodes.items(): + mid = eps._agent_to_module_mapping[aid] + agent_to_num_steps_added[aid] += e_eps.agent_steps() + module_to_num_steps_added[mid] += e_eps.agent_steps() + + # Update the adding metrics. + self._update_add_metrics( + num_episodes_added=num_episodes_added, + num_env_steps_added=num_env_steps_added, + num_episodes_evicted=num_episodes_evicted, + num_env_steps_evicted=num_env_steps_evicted, + agent_to_num_episodes_added=agent_to_num_episodes_added, + agent_to_num_steps_added=agent_to_num_steps_added, + agent_to_num_episodes_evicted=agent_to_num_episodes_evicted, + agent_to_num_steps_evicted=agent_to_num_steps_evicted, + module_to_num_episodes_added=module_to_num_steps_added, + module_to_num_steps_added=module_to_num_episodes_added, + module_to_num_episodes_evicted=module_to_num_episodes_evicted, + module_to_num_steps_evicted=module_to_num_steps_evicted, + ) + + @override(EpisodeReplayBuffer) + def sample( + self, + num_items: Optional[int] = None, + *, + batch_size_B: Optional[int] = None, + batch_length_T: Optional[int] = None, + n_step: Optional[Union[int, Tuple]] = 1, + gamma: float = 0.99, + include_infos: bool = False, + include_extra_model_outputs: bool = False, + replay_mode: str = "independent", + modules_to_sample: Optional[List[ModuleID]] = None, + **kwargs, + ) -> Union[List["MultiAgentEpisode"], List["SingleAgentEpisode"]]: + """Samples a batch of multi-agent transitions. + + Multi-agent transitions can be sampled either `"independent"` or + `"synchronized"` with the former sampling for each module independent agent + steps and the latter sampling agent transitions from the same environment step. + + The n-step parameter can be either a single integer or a tuple of two integers. + In the former case, the n-step is fixed to the given integer and in the latter + case, the n-step is sampled uniformly from the given range. Large n-steps could + potentially lead to a many retries because not all samples might have a full + n-step transition. + + Sampling returns batches of size B (number of 'rows'), where each row is a tuple + of the form + + `(o_t, a_t, sum(r_t+1:t+n), o_t+n)` + + where `o_t` is the observation in `t`, `a_t` the action chosen at observation + `o_t`, `o_t+n` is the observation `n` timesteps later and `sum(r_t+1:t+n)` is + the sum of all rewards collected over the time steps between `t+1` and `t+n`. + The n`-step can be chosen freely when sampling and defaults to `1`. If `n_step` + is a tuple it is sampled uniformly across the interval defined by the tuple (for + each row in the batch). + + Each batch contains - in addition to the data tuples presented above - two + further columns, namely `n_steps` and `weigths`. The former holds the `n_step` + used for each row in the batch and the latter a (default) weight of `1.0` for + each row in the batch. This weight is used for weighted loss calculations in + the training process. + + Args: + num_items: The number of items to sample. If provided, `batch_size_B` + should be `None`. + batch_size_B: The batch size to sample. If provided, `num_items` + should be `None`. + batch_length_T: The length of the sampled batch. If not provided, the + default batch length is used. This feature is not yet implemented. + n_step: The n-step to sample. If the n-step is a tuple, the n-step is + sampled uniformly from the given range. If not provided, the default + n-step of `1` is used. + gamma: The discount factor for the n-step reward calculation. + include_infos: Whether to include the infos in the sampled batch. + include_extra_model_outputs: Whether to include the extra model outputs + in the sampled batch. + replay_mode: The replay mode to use for sampling. Either `"independent"` + or `"synchronized"`. + modules_to_sample: A list of module IDs to sample from. If not provided, + transitions for aall modules are sampled. + + Returns: + A dictionary of the form `ModuleID -> SampleBatchType` containing the + sampled data for each module or each module in `modules_to_sample`, + if provided. + """ + if num_items is not None: + assert batch_size_B is None, ( + "Cannot call `sample()` with both `num_items` and `batch_size_B` " + "provided! Use either one." + ) + batch_size_B = num_items + + # Use our default values if no sizes/lengths provided. + batch_size_B = batch_size_B or self.batch_size_B + # TODO (simon): Implement trajectory sampling for RNNs. + batch_length_T = batch_length_T or self.batch_length_T + + # Sample for each module independently. + if replay_mode == "independent": + return self._sample_independent( + batch_size_B=batch_size_B, + batch_length_T=batch_length_T, + n_step=n_step, + gamma=gamma, + include_infos=include_infos, + include_extra_model_outputs=include_extra_model_outputs, + modules_to_sample=modules_to_sample, + ) + else: + return self._sample_synchonized( + batch_size_B=batch_size_B, + batch_length_T=batch_length_T, + n_step=n_step, + gamma=gamma, + include_infos=include_infos, + include_extra_model_outputs=include_extra_model_outputs, + modules_to_sample=modules_to_sample, + ) + + def get_added_agent_timesteps(self) -> int: + """Returns number of agent timesteps that have been added in buffer's lifetime. + + Note, this could be more than the `get_added_timesteps` returns as an + environment timestep could contain multiple agent timesteps (for eaxch agent + one). + """ + return self._num_agent_timesteps_added + + def get_module_ids(self) -> List[ModuleID]: + """Returns a list of module IDs stored in the buffer.""" + return list(self._module_to_indices.keys()) + + def get_num_agent_timesteps(self) -> int: + """Returns number of agent timesteps stored in the buffer. + + Note, this could be more than the `num_timesteps` as an environment timestep + could contain multiple agent timesteps (for eaxch agent one). + """ + return self._num_agent_timesteps + + @override(EpisodeReplayBuffer) + def get_num_episodes(self, module_id: Optional[ModuleID] = None) -> int: + """Returns number of episodes stored for a module in the buffer. + + Note, episodes could be either complete or truncated. + + Args: + module_id: The ID of the module to query. If not provided, the number of + episodes for all modules is returned. + + Returns: + The number of episodes stored for the module or all modules. + """ + return ( + self._num_module_episodes[module_id] + if module_id + else super().get_num_episodes() + ) + + @override(EpisodeReplayBuffer) + def get_num_episodes_evicted(self, module_id: Optional[ModuleID] = None) -> int: + """Returns number of episodes evicted for a module in the buffer.""" + return ( + self._num_module_episodes_evicted[module_id] + if module_id + else super().get_num_episodes_evicted() + ) + + @override(EpisodeReplayBuffer) + def get_num_timesteps(self, module_id: Optional[ModuleID] = None) -> int: + """Returns number of individual timesteps for a module stored in the buffer. + + Args: + module_id: The ID of the module to query. If not provided, the number of + timesteps for all modules are returned. + + Returns: + The number of timesteps stored for the module or all modules. + """ + return ( + self._num_module_timesteps[module_id] + if module_id + else super().get_num_timesteps() + ) + + @override(EpisodeReplayBuffer) + def get_sampled_timesteps(self, module_id: Optional[ModuleID] = None) -> int: + """Returns number of timesteps that have been sampled for a module. + + Args: + module_id: The ID of the module to query. If not provided, the number of + sampled timesteps for all modules are returned. + + Returns: + The number of timesteps sampled for the module or all modules. + """ + return ( + self.sampled_timesteps_per_module[module_id] + if module_id + else super().get_sampled_timesteps() + ) + + @override(EpisodeReplayBuffer) + def get_added_timesteps(self, module_id: Optional[ModuleID] = None) -> int: + """Returns the number of timesteps added in buffer's lifetime for given module. + + Args: + module_id: The ID of the module to query. If not provided, the total number + of timesteps ever added. + + Returns: + The number of timesteps added for `module_id` (or all modules if `module_id` + is None). + """ + return ( + self._num_module_timesteps_added[module_id] + if module_id + else super().get_added_timesteps() + ) + + @override(EpisodeReplayBuffer) + def get_state(self) -> Dict[str, Any]: + """Gets a pickable state of the buffer. + + This is used for checkpointing the buffer's state. It is specifically helpful, + for example, when a trial is paused and resumed later on. The buffer's state + can be saved to disk and reloaded when the trial is resumed. + + Returns: + A dict containing all necessary information to restore the buffer's state. + """ + return super().get_state() | { + "_module_to_indices": list(self._module_to_indices.items()), + "_num_agent_timesteps": self._num_agent_timesteps, + "_num_agent_timesteps_added": self._num_agent_timesteps_added, + "_num_module_timesteps": list(self._num_module_timesteps.items()), + "_num_module_timesteps_added": list( + self._num_module_timesteps_added.items() + ), + "_num_module_episodes": list(self._num_module_episodes.items()), + "_num_module_episodes_evicted": list( + self._num_module_episodes_evicted.items() + ), + "sampled_timesteps_per_module": list( + self.sampled_timesteps_per_module.items() + ), + } + + @override(EpisodeReplayBuffer) + def set_state(self, state) -> None: + """Sets the state of a buffer from a previously stored state. + + See `get_state()` for more information on what is stored in the state. This + method is used to restore the buffer's state from a previously stored state. + It is specifically helpful, for example, when a trial is paused and resumed + later on. The buffer's state can be saved to disk and reloaded when the trial + is resumed. + + Args: + state: The state to restore the buffer from. + """ + # Set the episodes. + self._set_episodes(state) + # Set the super's state. + super().set_state(state) + # Now set the remaining attributes. + self._module_to_indices = defaultdict(list, dict(state["_module_to_indices"])) + self._num_agent_timesteps = state["_num_agent_timesteps"] + self._num_agent_timesteps_added = state["_num_agent_timesteps_added"] + self._num_module_timesteps = defaultdict( + int, dict(state["_num_module_timesteps"]) + ) + self._num_module_timesteps_added = defaultdict( + int, dict(state["_num_module_timesteps_added"]) + ) + self._num_module_episodes = defaultdict( + int, dict(state["_num_module_episodes"]) + ) + self._num_module_episodes_evicted = defaultdict( + int, dict(state["_num_module_episodes_evicted"]) + ) + self.sampled_timesteps_per_module = defaultdict( + list, dict(state["sampled_timesteps_per_module"]) + ) + + def _set_episodes(self, state: Dict[str, Any]) -> None: + """Sets the episodes from the state.""" + if not self.episodes: + self.episodes = deque( + [ + MultiAgentEpisode.from_state(eps_data) + for eps_data in state["episodes"] + ] + ) + + def _sample_independent( + self, + batch_size_B: Optional[int], + batch_length_T: Optional[int], + n_step: Optional[Union[int, Tuple[int, int]]], + gamma: float, + include_infos: bool, + include_extra_model_outputs: bool, + modules_to_sample: Optional[Set[ModuleID]], + ) -> List["SingleAgentEpisode"]: + """Samples a batch of independent multi-agent transitions.""" + + actual_n_step = n_step or 1 + # Sample the n-step if necessary. + random_n_step = isinstance(n_step, (tuple, list)) + + sampled_episodes = [] + # Record the number of samples per module/agent/total. + num_env_steps_sampled = 0 + agent_to_num_steps_sampled = defaultdict(int) + module_to_num_steps_sampled = defaultdict(int) + # Record all the env step buffer indices that are contained in the sample. + sampled_env_step_idxs = set() + agent_to_sampled_env_step_idxs = defaultdict(set) + module_to_sampled_env_step_idxs = defaultdict(set) + # Record all the episode buffer indices that are contained in the sample. + sampled_episode_idxs = set() + agent_to_sampled_episode_idxs = defaultdict(set) + module_to_sampled_episode_idxs = defaultdict(set) + # Record all n-steps that have been used. + sampled_n_steps = [] + agent_to_sampled_n_steps = defaultdict(list) + module_to_sampled_n_steps = defaultdict(list) + # Record the number of times a sample needs to be resampled. + num_resamples = 0 + agent_to_num_resamples = defaultdict(int) + module_to_num_resamples = defaultdict(int) + + # TODO (simon): Ensure that the module has data and if not, skip it. + # TODO (sven): Should we then error out or skip? I think the Learner + # should handle this case when a module has no train data. + modules_to_sample = modules_to_sample or set(self._module_to_indices.keys()) + for module_id in modules_to_sample: + module_indices = self._module_to_indices[module_id] + B = 0 + while B < batch_size_B: + # Now sample from the single-agent timesteps. + index_tuple = module_indices[self.rng.integers(len(module_indices))] + + # This will be an agent timestep (not env timestep). + # TODO (simon, sven): Maybe deprecate sa_episode_idx (_) in the index + # quads. Is there any need for it? + ma_episode_idx, agent_id, sa_episode_ts = ( + index_tuple[0] - self._num_episodes_evicted, + index_tuple[1], + index_tuple[2], + ) + + # Get the multi-agent episode. + ma_episode = self.episodes[ma_episode_idx] + # Retrieve the single-agent episode for filtering. + sa_episode = ma_episode.agent_episodes[agent_id] + + # If we use random n-step sampling, draw the n-step for this item. + if random_n_step: + actual_n_step = int(self.rng.integers(n_step[0], n_step[1])) + # If we cannnot make the n-step, we resample. + if sa_episode_ts + actual_n_step > len(sa_episode): + num_resamples += 1 + agent_to_num_resamples[agent_id] += 1 + module_to_num_resamples[module_id] += 1 + continue + # Note, this will be the reward after executing action + # `a_(episode_ts)`. For `n_step>1` this will be the discounted sum + # of all rewards that were collected over the last n steps. + sa_raw_rewards = sa_episode.get_rewards( + slice(sa_episode_ts, sa_episode_ts + actual_n_step) + ) + sa_rewards = scipy.signal.lfilter( + [1], [1, -gamma], sa_raw_rewards[::-1], axis=0 + )[-1] + + sampled_sa_episode = SingleAgentEpisode( + id_=sa_episode.id_, + # Provide the IDs for the learner connector. + agent_id=sa_episode.agent_id, + module_id=sa_episode.module_id, + multi_agent_episode_id=ma_episode.id_, + # Ensure that each episode contains a tuple of the form: + # (o_t, a_t, sum(r_(t:t+n_step)), o_(t+n_step)) + # Two observations (t and t+n). + observations=[ + sa_episode.get_observations(sa_episode_ts), + sa_episode.get_observations(sa_episode_ts + actual_n_step), + ], + observation_space=sa_episode.observation_space, + infos=( + [ + sa_episode.get_infos(sa_episode_ts), + sa_episode.get_infos(sa_episode_ts + actual_n_step), + ] + if include_infos + else None + ), + actions=[sa_episode.get_actions(sa_episode_ts)], + action_space=sa_episode.action_space, + rewards=[sa_rewards], + # If the sampled single-agent episode is the single-agent episode's + # last time step, check, if the single-agent episode is terminated + # or truncated. + terminated=( + sa_episode_ts + actual_n_step >= len(sa_episode) + and sa_episode.is_terminated + ), + truncated=( + sa_episode_ts + actual_n_step >= len(sa_episode) + and sa_episode.is_truncated + ), + extra_model_outputs={ + "weights": [1.0], + "n_step": [actual_n_step], + **( + { + k: [ + sa_episode.get_extra_model_outputs(k, sa_episode_ts) + ] + for k in sa_episode.extra_model_outputs.keys() + } + if include_extra_model_outputs + else {} + ), + }, + # TODO (sven): Support lookback buffers. + len_lookback_buffer=0, + t_started=sa_episode_ts, + ) + # Append single-agent episode to the list of sampled episodes. + sampled_episodes.append(sampled_sa_episode) + # Add the episode indices. + sampled_episode_idxs.add(ma_episode_idx) + agent_to_sampled_episode_idxs[sa_episode.agent_id].add(sa_episode.id_) + module_to_sampled_episode_idxs[module_id].add(sa_episode.id_) + # Add the unique step hashes. + # Get the corresponding index in the `env_to_agent_t` mapping. + # TODO (simon, sven): This has complexity O(n) and could become + # expensive when the episode is large. Note, however, that conversion + # from list to `numpy.ndarray` is also complexity O(n) and we do this + # at many places - also in the `MultiAgentEpisode`s. + ma_episode_ts = ma_episode.env_t_to_agent_t[agent_id].data.index( + sa_episode_ts + ) + sampled_env_step_idxs.add( + hashlib.sha256( + f"{ma_episode.id_}-{ma_episode_ts}".encode() + ).hexdigest() + ) + hashed_agent_step = hashlib.sha256( + f"{sa_episode.id_}-{sa_episode_ts}".encode() + ).hexdigest() + agent_to_sampled_env_step_idxs[agent_id].add(hashed_agent_step) + module_to_sampled_env_step_idxs[module_id].add(hashed_agent_step) + # Add the actual n-step used in generating this sample. + sampled_n_steps.append(actual_n_step) + agent_to_sampled_n_steps[agent_id].append(actual_n_step) + module_to_sampled_n_steps[module_id].append(actual_n_step) + + # Increase counter. + B += 1 + + # Increase the per module timesteps counter. + self.sampled_timesteps_per_module[module_id] += B + # Increase the counter metrics. + num_env_steps_sampled += B + agent_to_num_steps_sampled[agent_id] += B + module_to_num_steps_sampled[module_id] += B + + # Increase the counter for environment timesteps. + self.sampled_timesteps += batch_size_B + + # Update the sample metrics. + num_episodes_per_sample = len(sampled_episode_idxs) + num_env_steps_per_sample = len(sampled_env_step_idxs) + sampled_n_step = sum(sampled_n_steps) / batch_size_B + agent_to_num_episodes_per_sample = { + aid: len(l) for aid, l in agent_to_sampled_episode_idxs.items() + } + module_to_num_episodes_per_sample = { + mid: len(l) for mid, l in module_to_sampled_episode_idxs.items() + } + agent_to_num_steps_per_sample = { + aid: len(l) for aid, l in agent_to_sampled_env_step_idxs.items() + } + module_to_num_steps_per_sample = { + mid: len(l) for mid, l in module_to_sampled_env_step_idxs.items() + } + agent_to_sampled_n_step = { + aid: sum(l) / len(l) for aid, l in agent_to_sampled_n_steps.items() + } + module_to_sampled_n_step = { + mid: sum(l) / len(l) for mid, l in module_to_sampled_n_steps.items() + } + self._update_sample_metrics( + num_env_steps_sampled=num_env_steps_sampled, + num_episodes_per_sample=num_episodes_per_sample, + num_env_steps_per_sample=num_env_steps_per_sample, + sampled_n_step=sampled_n_step, + num_resamples=num_resamples, + agent_to_num_steps_sampled=agent_to_num_steps_sampled, + agent_to_num_episodes_per_sample=agent_to_num_episodes_per_sample, + agent_to_num_steps_per_sample=agent_to_num_steps_per_sample, + agent_to_sampled_n_step=agent_to_sampled_n_step, + agent_to_num_resamples=agent_to_num_resamples, + module_to_num_steps_sampled=module_to_num_steps_sampled, + module_to_num_episodes_per_sample=module_to_num_episodes_per_sample, + module_to_num_steps_per_sample=module_to_num_steps_per_sample, + module_to_sampled_n_step=module_to_sampled_n_step, + module_to_num_resamples=module_to_num_resamples, + ) + + # Return multi-agent dictionary. + return sampled_episodes + + def _sample_synchonized( + self, + batch_size_B: Optional[int], + batch_length_T: Optional[int], + n_step: Optional[Union[int, Tuple]], + gamma: float, + include_infos: bool, + include_extra_model_outputs: bool, + modules_to_sample: Optional[List[ModuleID]], + ) -> SampleBatchType: + """Samples a batch of synchronized multi-agent transitions.""" + # Sample the n-step if necessary. + if isinstance(n_step, tuple): + # Use random n-step sampling. + random_n_step = True + else: + actual_n_step = n_step or 1 + random_n_step = False + + # Containers for the sampled data. + observations: Dict[ModuleID, List[ObsType]] = defaultdict(list) + next_observations: Dict[ModuleID, List[ObsType]] = defaultdict(list) + actions: Dict[ModuleID, List[ActType]] = defaultdict(list) + rewards: Dict[ModuleID, List[float]] = defaultdict(list) + is_terminated: Dict[ModuleID, List[bool]] = defaultdict(list) + is_truncated: Dict[ModuleID, List[bool]] = defaultdict(list) + weights: Dict[ModuleID, List[float]] = defaultdict(list) + n_steps: Dict[ModuleID, List[int]] = defaultdict(list) + # If `info` should be included, construct also a container for them. + if include_infos: + infos: Dict[ModuleID, List[Dict[str, Any]]] = defaultdict(list) + # If `extra_model_outputs` should be included, construct a container for them. + if include_extra_model_outputs: + extra_model_outputs: Dict[ModuleID, List[Dict[str, Any]]] = defaultdict( + list + ) + + B = 0 + while B < batch_size_B: + index_tuple = self._indices[self.rng.integers(len(self._indices))] + + # This will be an env timestep (not agent timestep) + ma_episode_idx, ma_episode_ts = ( + index_tuple[0] - self._num_episodes_evicted, + index_tuple[1], + ) + # If we use random n-step sampling, draw the n-step for this item. + if random_n_step: + actual_n_step = int(self.rng.integers(n_step[0], n_step[1])) + # If we are at the end of an episode, continue. + # Note, priority sampling got us `o_(t+n)` and we need for the loss + # calculation in addition `o_t`. + # TODO (simon): Maybe introduce a variable `num_retries` until the + # while loop should break when not enough samples have been collected + # to make n-step possible. + if ma_episode_ts - actual_n_step < 0: + continue + + # Retrieve the multi-agent episode. + ma_episode = self.episodes[ma_episode_idx] + + # Ensure that each row contains a tuple of the form: + # (o_t, a_t, sum(r_(t:t+n_step)), o_(t+n_step)) + # TODO (simon): Implement version for sequence sampling when using RNNs. + eps_observation = ma_episode.get_observations( + slice(ma_episode_ts - actual_n_step, ma_episode_ts + 1), + return_list=True, + ) + # Note, `MultiAgentEpisode` stores the action that followed + # `o_t` with `o_(t+1)`, therefore, we need the next one. + # TODO (simon): This gets the wrong action as long as the getters are not + # fixed. + eps_actions = ma_episode.get_actions(ma_episode_ts - actual_n_step) + # Make sure that at least a single agent should have full transition. + # TODO (simon): Filter for the `modules_to_sample`. + agents_to_sample = self._agents_with_full_transitions( + eps_observation, + eps_actions, + ) + # If not, we resample. + if not agents_to_sample: + continue + # TODO (simon, sven): Do we need to include the common agent rewards? + # Note, the reward that is collected by transitioning from `o_t` to + # `o_(t+1)` is stored in the next transition in `MultiAgentEpisode`. + eps_rewards = ma_episode.get_rewards( + slice(ma_episode_ts - actual_n_step, ma_episode_ts), + return_list=True, + ) + # TODO (simon, sven): Do we need to include the common infos? And are + # there common extra model outputs? + if include_infos: + # If infos are included we include the ones from the last timestep + # as usually the info contains additional values about the last state. + eps_infos = ma_episode.get_infos(ma_episode_ts) + if include_extra_model_outputs: + # If `extra_model_outputs` are included we include the ones from the + # first timestep as usually the `extra_model_outputs` contain additional + # values from the forward pass that produced the action at the first + # timestep. + # Note, we extract them into single row dictionaries similar to the + # infos, in a connector we can then extract these into single batch + # rows. + eps_extra_model_outputs = { + k: ma_episode.get_extra_model_outputs( + k, ma_episode_ts - actual_n_step + ) + for k in ma_episode.extra_model_outputs.keys() + } + # If the sampled time step is the episode's last time step check, if + # the episode is terminated or truncated. + episode_terminated = False + episode_truncated = False + if ma_episode_ts == ma_episode.env_t: + episode_terminated = ma_episode.is_terminated + episode_truncated = ma_episode.is_truncated + # TODO (simon): Filter for the `modules_to_sample`. + # TODO (sven, simon): We could here also sample for all agents in the + # `modules_to_sample` and then adapt the `n_step` for agents that + # have not a full transition. + for agent_id in agents_to_sample: + # Map our agent to the corresponding module we want to + # train. + module_id = ma_episode._agent_to_module_mapping[agent_id] + # Sample only for the modules in `modules_to_sample`. + if module_id not in ( + modules_to_sample or self._module_to_indices.keys() + ): + continue + # TODO (simon, sven): Here we could skip for modules not + # to be sampled in `modules_to_sample`. + observations[module_id].append(eps_observation[0][agent_id]) + next_observations[module_id].append(eps_observation[-1][agent_id]) + # Fill missing rewards with zeros. + agent_rewards = [r[agent_id] or 0.0 for r in eps_rewards] + rewards[module_id].append( + scipy.signal.lfilter([1], [1, -gamma], agent_rewards[::-1], axis=0)[ + -1 + ] + ) + # Note, this should exist, as we filtered for agents with full + # transitions. + actions[module_id].append(eps_actions[agent_id]) + if include_infos: + infos[module_id].append(eps_infos[agent_id]) + if include_extra_model_outputs: + extra_model_outputs[module_id].append( + { + k: eps_extra_model_outputs[agent_id][k] + for k in eps_extra_model_outputs[agent_id].keys() + } + ) + # If sampled observation is terminal for the agent. Either MAE + # episode is truncated/terminated or SAE episode is truncated/ + # terminated at this ts. + # TODO (simon, sven): Add method agent_alive(ts) to MAE. + # or add slicing to get_terminateds(). + agent_ts = ma_episode.env_t_to_agent_t[agent_id][ma_episode_ts] + agent_eps = ma_episode.agent_episodes[agent_id] + agent_terminated = agent_ts == agent_eps.t and agent_eps.is_terminated + agent_truncated = ( + agent_ts == agent_eps.t + and agent_eps.is_truncated + and not agent_eps.is_terminated + ) + if episode_terminated or agent_terminated: + is_terminated[module_id].append(True) + is_truncated[module_id].append(False) + elif episode_truncated or agent_truncated: + is_truncated[module_id].append(True) + is_terminated[module_id].append(False) + else: + is_terminated[module_id].append(False) + is_truncated[module_id].append(False) + # Increase the per module counter. + self.sampled_timesteps_per_module[module_id] += 1 + + # Increase counter. + B += 1 + # Increase the counter for environment timesteps. + self.sampled_timesteps += batch_size_B + + # Should be convertible to MultiAgentBatch. + ret = { + **{ + module_id: { + Columns.OBS: batch(observations[module_id]), + Columns.ACTIONS: batch(actions[module_id]), + Columns.REWARDS: np.array(rewards[module_id]), + Columns.NEXT_OBS: batch(next_observations[module_id]), + Columns.TERMINATEDS: np.array(is_terminated[module_id]), + Columns.TRUNCATEDS: np.array(is_truncated[module_id]), + "weights": np.array(weights[module_id]), + "n_step": np.array(n_steps[module_id]), + } + for module_id in observations.keys() + } + } + + # Return multi-agent dictionary. + return ret + + def _num_remaining_episodes(self, new_eps, evicted_eps): + """Calculates the number of remaining episodes. + + When adding episodes and evicting them in the `add()` method + this function calculates iteratively the number of remaining + episodes. + + Args: + new_eps: List of new episode IDs. + evicted_eps: List of evicted episode IDs. + + Returns: + Number of episodes remaining after evicting the episodes in + `evicted_eps` and adding the episode in `new_eps`. + """ + return len( + set(self.episode_id_to_index.keys()).union(set(new_eps)) - set(evicted_eps) + ) + + def _evict_module_episodes(self, ma_episode: MultiAgentEpisode) -> None: + """Evicts the module episodes from the buffer adn updates all counters. + + Args: + multi_agent_eps: The multi-agent episode to evict from the buffer. + """ + + # Note we need to take the agent ids from the evicted episode because + # different episodes can have different agents and module mappings. + for agent_id in ma_episode.agent_episodes: + # Retrieve the corresponding module ID and module episode. + module_id = ma_episode._agent_to_module_mapping[agent_id] + module_eps = ma_episode.agent_episodes[agent_id] + # Update all counters. + self._num_module_timesteps[module_id] -= module_eps.env_steps() + self._num_module_episodes[module_id] -= 1 + self._num_module_episodes_evicted[module_id] += 1 + + def _update_module_counters(self, ma_episode: MultiAgentEpisode) -> None: + """Updates the module counters after adding an episode. + + Args: + multi_agent_episode: The multi-agent episode to update the module counters + for. + """ + for agent_id in ma_episode.agent_ids: + agent_steps = ma_episode.agent_episodes[agent_id].env_steps() + # Only add if the agent has stepped in the episode (chunk). + if agent_steps > 0: + # Receive the corresponding module ID. + module_id = ma_episode.module_for(agent_id) + self._num_module_timesteps[module_id] += agent_steps + self._num_module_timesteps_added[module_id] += agent_steps + # if ma_episode.agent_episodes[agent_id].is_done: + # # TODO (simon): Check, if we do not count the same episode + # # multiple times. + # # Also add to the module episode counter. + # self._num_module_episodes[module_id] += 1 + + def _add_new_module_indices( + self, + ma_episode: MultiAgentEpisode, + episode_idx: int, + ma_episode_exists: bool = True, + ) -> None: + """Adds the module indices for new episode chunks. + + Args: + ma_episode: The multi-agent episode to add the module indices for. + episode_idx: The index of the episode in the `self.episodes`. + ma_episode_exists: Whether `ma_episode` is already in this buffer (with a + predecessor chunk to which we'll concatenate `ma_episode` later). + """ + existing_ma_episode = None + if ma_episode_exists: + existing_ma_episode = self.episodes[ + self.episode_id_to_index[ma_episode.id_] - self._num_episodes_evicted + ] + + # Note, we iterate through the agent episodes b/c we want to store records + # and some agents could not have entered the environment. + for agent_id in ma_episode.agent_episodes: + # Get the corresponding module id. + module_id = ma_episode.module_for(agent_id) + # Get the module episode. + module_eps = ma_episode.agent_episodes[agent_id] + + # Is the agent episode already in the buffer's existing `ma_episode`? + if ma_episode_exists and agent_id in existing_ma_episode.agent_episodes: + existing_sa_eps_len = len(existing_ma_episode.agent_episodes[agent_id]) + # Otherwise, it is a new single-agent episode and we increase the counter. + else: + existing_sa_eps_len = 0 + self._num_module_episodes[module_id] += 1 + + # Add new module indices. + self._module_to_indices[module_id].extend( + [ + ( + # Keep the MAE index for sampling + episode_idx, + agent_id, + existing_sa_eps_len + i, + ) + for i in range(len(module_eps)) + ] + ) + + def _agents_with_full_transitions( + self, observations: Dict[AgentID, ObsType], actions: Dict[AgentID, ActType] + ): + """Filters for agents that have full transitions. + + Args: + observations: The observations of the episode. + actions: The actions of the episode. + + Returns: + List of agent IDs that have full transitions. + """ + agents_to_sample = [] + for agent_id in observations[0].keys(): + # Only if the agent has an action at the first and an observation + # at the first and last timestep of the n-step transition, we can sample it. + if agent_id in actions and agent_id in observations[-1]: + agents_to_sample.append(agent_id) + return agents_to_sample diff --git a/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/multi_agent_mixin_replay_buffer.py b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/multi_agent_mixin_replay_buffer.py new file mode 100644 index 0000000000000000000000000000000000000000..4c406f36ed081a03b4df01db8edb48c114adb55f --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/multi_agent_mixin_replay_buffer.py @@ -0,0 +1,404 @@ +import collections +import logging +import random +from typing import Any, Dict, Optional + +import numpy as np + +from ray.rllib.policy.rnn_sequencing import timeslice_along_seq_lens_with_overlap +from ray.rllib.policy.sample_batch import ( + DEFAULT_POLICY_ID, + SampleBatch, + concat_samples_into_ma_batch, +) +from ray.rllib.utils.annotations import override +from ray.rllib.utils.replay_buffers.multi_agent_prioritized_replay_buffer import ( + MultiAgentPrioritizedReplayBuffer, +) +from ray.rllib.utils.replay_buffers.multi_agent_replay_buffer import ( + MultiAgentReplayBuffer, + ReplayMode, + merge_dicts_with_warning, +) +from ray.rllib.utils.replay_buffers.replay_buffer import _ALL_POLICIES, StorageUnit +from ray.rllib.utils.typing import PolicyID, SampleBatchType +from ray.util.annotations import DeveloperAPI +from ray.util.debug import log_once + +logger = logging.getLogger(__name__) + + +@DeveloperAPI +class MultiAgentMixInReplayBuffer(MultiAgentPrioritizedReplayBuffer): + """This buffer adds replayed samples to a stream of new experiences. + + - Any newly added batch (`add()`) is immediately returned upon + the next `sample` call (close to on-policy) as well as being moved + into the buffer. + - Additionally, a certain number of old samples is mixed into the + returned sample according to a given "replay ratio". + - If >1 calls to `add()` are made without any `sample()` calls + in between, all newly added batches are returned (plus some older samples + according to the "replay ratio"). + + .. testcode:: + :skipif: True + + # replay ratio 0.66 (2/3 replayed, 1/3 new samples): + buffer = MultiAgentMixInReplayBuffer(capacity=100, + replay_ratio=0.66) + buffer.add() + buffer.add() + buffer.sample(1) + + .. testoutput:: + + ..[, , ] + + .. testcode:: + :skipif: True + + buffer.add() + buffer.sample(1) + + .. testoutput:: + + [, , ] + or: [, , ], [, , ] or [, , ], + but always as it is the newest sample + + .. testcode:: + :skipif: True + + buffer.add() + buffer.sample(1) + + .. testoutput:: + + [, , ] + or [, , ], [, , ] or [, , ], etc.. + but always as it is the newest sample + + .. testcode:: + :skipif: True + + # replay proportion 0.0 -> replay disabled: + buffer = MixInReplay(capacity=100, replay_ratio=0.0) + buffer.add() + buffer.sample() + + .. testoutput:: + + [] + + .. testcode:: + :skipif: True + + buffer.add() + buffer.sample() + + .. testoutput:: + + [] + """ + + def __init__( + self, + capacity: int = 10000, + storage_unit: str = "timesteps", + num_shards: int = 1, + replay_mode: str = "independent", + replay_sequence_override: bool = True, + replay_sequence_length: int = 1, + replay_burn_in: int = 0, + replay_zero_init_states: bool = True, + replay_ratio: float = 0.66, + underlying_buffer_config: dict = None, + prioritized_replay_alpha: float = 0.6, + prioritized_replay_beta: float = 0.4, + prioritized_replay_eps: float = 1e-6, + **kwargs, + ): + """Initializes MultiAgentMixInReplayBuffer instance. + + Args: + capacity: The capacity of the buffer, measured in `storage_unit`. + storage_unit: Either 'timesteps', 'sequences' or + 'episodes'. Specifies how experiences are stored. If they + are stored in episodes, replay_sequence_length is ignored. + num_shards: The number of buffer shards that exist in total + (including this one). + replay_mode: One of "independent" or "lockstep". Determines, + whether batches are sampled independently or to an equal + amount. + replay_sequence_override: If True, ignore sequences found in incoming + batches, slicing them into sequences as specified by + `replay_sequence_length` and `replay_sequence_burn_in`. This only has + an effect if storage_unit is `sequences`. + replay_sequence_length: The sequence length (T) of a single + sample. If > 1, we will sample B x T from this buffer. This + only has an effect if storage_unit is 'timesteps'. + replay_burn_in: The burn-in length in case + `replay_sequence_length` > 0. This is the number of timesteps + each sequence overlaps with the previous one to generate a + better internal state (=state after the burn-in), instead of + starting from 0.0 each RNN rollout. + replay_zero_init_states: Whether the initial states in the + buffer (if replay_sequence_length > 0) are alwayas 0.0 or + should be updated with the previous train_batch state outputs. + replay_ratio: Ratio of replayed samples in the returned + batches. E.g. a ratio of 0.0 means only return new samples + (no replay), a ratio of 0.5 means always return newest sample + plus one old one (1:1), a ratio of 0.66 means always return + the newest sample plus 2 old (replayed) ones (1:2), etc... + underlying_buffer_config: A config that contains all necessary + constructor arguments and arguments for methods to call on + the underlying buffers. This replaces the standard behaviour + of the underlying PrioritizedReplayBuffer. The config + follows the conventions of the general + replay_buffer_config. kwargs for subsequent calls of methods + may also be included. Example: + "replay_buffer_config": {"type": PrioritizedReplayBuffer, + "capacity": 10, "storage_unit": "timesteps", + prioritized_replay_alpha: 0.5, prioritized_replay_beta: 0.5, + prioritized_replay_eps: 0.5} + prioritized_replay_alpha: Alpha parameter for a prioritized + replay buffer. Use 0.0 for no prioritization. + prioritized_replay_beta: Beta parameter for a prioritized + replay buffer. + prioritized_replay_eps: Epsilon parameter for a prioritized + replay buffer. + **kwargs: Forward compatibility kwargs. + """ + if not 0 <= replay_ratio <= 1: + raise ValueError("Replay ratio must be within [0, 1]") + + MultiAgentPrioritizedReplayBuffer.__init__( + self, + capacity=capacity, + storage_unit=storage_unit, + num_shards=num_shards, + replay_mode=replay_mode, + replay_sequence_override=replay_sequence_override, + replay_sequence_length=replay_sequence_length, + replay_burn_in=replay_burn_in, + replay_zero_init_states=replay_zero_init_states, + underlying_buffer_config=underlying_buffer_config, + prioritized_replay_alpha=prioritized_replay_alpha, + prioritized_replay_beta=prioritized_replay_beta, + prioritized_replay_eps=prioritized_replay_eps, + **kwargs, + ) + + self.replay_ratio = replay_ratio + + self.last_added_batches = collections.defaultdict(list) + + @DeveloperAPI + @override(MultiAgentPrioritizedReplayBuffer) + def add(self, batch: SampleBatchType, **kwargs) -> None: + """Adds a batch to the appropriate policy's replay buffer. + + Turns the batch into a MultiAgentBatch of the DEFAULT_POLICY_ID if + it is not a MultiAgentBatch. Subsequently, adds the individual policy + batches to the storage. + + Args: + batch: The batch to be added. + **kwargs: Forward compatibility kwargs. + """ + # Make a copy so the replay buffer doesn't pin plasma memory. + batch = batch.copy() + # Handle everything as if multi-agent. + batch = batch.as_multi_agent() + + kwargs = merge_dicts_with_warning(self.underlying_buffer_call_args, kwargs) + + pids_and_batches = self._maybe_split_into_policy_batches(batch) + + # We need to split batches into timesteps, sequences or episodes + # here already to properly keep track of self.last_added_batches + # underlying buffers should not split up the batch any further + with self.add_batch_timer: + if self.storage_unit == StorageUnit.TIMESTEPS: + for policy_id, sample_batch in pids_and_batches.items(): + timeslices = sample_batch.timeslices(1) + for time_slice in timeslices: + self.replay_buffers[policy_id].add(time_slice, **kwargs) + self.last_added_batches[policy_id].append(time_slice) + + elif self.storage_unit == StorageUnit.SEQUENCES: + for policy_id, sample_batch in pids_and_batches.items(): + timeslices = timeslice_along_seq_lens_with_overlap( + sample_batch=sample_batch, + seq_lens=sample_batch.get(SampleBatch.SEQ_LENS) + if self.replay_sequence_override + else None, + zero_pad_max_seq_len=self.replay_sequence_length, + pre_overlap=self.replay_burn_in, + zero_init_states=self.replay_zero_init_states, + ) + for slice in timeslices: + self.replay_buffers[policy_id].add(slice, **kwargs) + self.last_added_batches[policy_id].append(slice) + + elif self.storage_unit == StorageUnit.EPISODES: + for policy_id, sample_batch in pids_and_batches.items(): + for eps in sample_batch.split_by_episode(): + # Only add full episodes to the buffer + if eps.get(SampleBatch.T)[0] == 0 and ( + eps.get(SampleBatch.TERMINATEDS, [True])[-1] + or eps.get(SampleBatch.TRUNCATEDS, [False])[-1] + ): + self.replay_buffers[policy_id].add(eps, **kwargs) + self.last_added_batches[policy_id].append(eps) + else: + if log_once("only_full_episodes"): + logger.info( + "This buffer uses episodes as a storage " + "unit and thus allows only full episodes " + "to be added to it. Some samples may be " + "dropped." + ) + elif self.storage_unit == StorageUnit.FRAGMENTS: + for policy_id, sample_batch in pids_and_batches.items(): + self.replay_buffers[policy_id].add(sample_batch, **kwargs) + self.last_added_batches[policy_id].append(sample_batch) + + self._num_added += batch.count + + @DeveloperAPI + @override(MultiAgentReplayBuffer) + def sample( + self, num_items: int, policy_id: PolicyID = DEFAULT_POLICY_ID, **kwargs + ) -> Optional[SampleBatchType]: + """Samples a batch of size `num_items` from a specified buffer. + + Concatenates old samples to new ones according to + self.replay_ratio. If not enough new samples are available, mixes in + less old samples to retain self.replay_ratio on average. Returns + an empty batch if there are no items in the buffer. + + Args: + num_items: Number of items to sample from this buffer. + policy_id: ID of the policy that produced the experiences to be + sampled. + **kwargs: Forward compatibility kwargs. + + Returns: + Concatenated MultiAgentBatch of items. + """ + # Merge kwargs, overwriting standard call arguments + kwargs = merge_dicts_with_warning(self.underlying_buffer_call_args, kwargs) + + def mix_batches(_policy_id): + """Mixes old with new samples. + + Tries to mix according to self.replay_ratio on average. + If not enough new samples are available, mixes in less old samples + to retain self.replay_ratio on average. + """ + + def round_up_or_down(value, ratio): + """Returns an integer averaging to value*ratio.""" + product = value * ratio + ceil_prob = product % 1 + if random.uniform(0, 1) < ceil_prob: + return int(np.ceil(product)) + else: + return int(np.floor(product)) + + max_num_new = round_up_or_down(num_items, 1 - self.replay_ratio) + # if num_samples * self.replay_ratio is not round, + # we need one more sample with a probability of + # (num_items*self.replay_ratio) % 1 + + _buffer = self.replay_buffers[_policy_id] + output_batches = self.last_added_batches[_policy_id][:max_num_new] + self.last_added_batches[_policy_id] = self.last_added_batches[_policy_id][ + max_num_new: + ] + + # No replay desired + if self.replay_ratio == 0.0: + return concat_samples_into_ma_batch(output_batches) + # Only replay desired + elif self.replay_ratio == 1.0: + return _buffer.sample(num_items, **kwargs) + + num_new = len(output_batches) + + if np.isclose(num_new, num_items * (1 - self.replay_ratio)): + # The optimal case, we can mix in a round number of old + # samples on average + num_old = num_items - max_num_new + else: + # We never want to return more elements than num_items + num_old = min( + num_items - max_num_new, + round_up_or_down( + num_new, self.replay_ratio / (1 - self.replay_ratio) + ), + ) + + output_batches.append(_buffer.sample(num_old, **kwargs)) + # Depending on the implementation of underlying buffers, samples + # might be SampleBatches + output_batches = [batch.as_multi_agent() for batch in output_batches] + return concat_samples_into_ma_batch(output_batches) + + def check_buffer_is_ready(_policy_id): + if ( + (len(self.replay_buffers[policy_id]) == 0) and self.replay_ratio > 0.0 + ) or ( + len(self.last_added_batches[_policy_id]) == 0 + and self.replay_ratio < 1.0 + ): + return False + return True + + with self.replay_timer: + samples = [] + + if self.replay_mode == ReplayMode.LOCKSTEP: + assert ( + policy_id is None + ), "`policy_id` specifier not allowed in `lockstep` mode!" + if check_buffer_is_ready(_ALL_POLICIES): + samples.append(mix_batches(_ALL_POLICIES).as_multi_agent()) + elif policy_id is not None: + if check_buffer_is_ready(policy_id): + samples.append(mix_batches(policy_id).as_multi_agent()) + else: + for policy_id, replay_buffer in self.replay_buffers.items(): + if check_buffer_is_ready(policy_id): + samples.append(mix_batches(policy_id).as_multi_agent()) + + return concat_samples_into_ma_batch(samples) + + @DeveloperAPI + @override(MultiAgentPrioritizedReplayBuffer) + def get_state(self) -> Dict[str, Any]: + """Returns all local state. + + Returns: + The serializable local state. + """ + data = { + "last_added_batches": self.last_added_batches, + } + parent = MultiAgentPrioritizedReplayBuffer.get_state(self) + parent.update(data) + return parent + + @DeveloperAPI + @override(MultiAgentPrioritizedReplayBuffer) + def set_state(self, state: Dict[str, Any]) -> None: + """Restores all local state to the provided `state`. + + Args: + state: The new state to set this buffer. Can be obtained by + calling `self.get_state()`. + """ + self.last_added_batches = state["last_added_batches"] + MultiAgentPrioritizedReplayBuffer.set_state(state) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/multi_agent_prioritized_episode_buffer.py b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/multi_agent_prioritized_episode_buffer.py new file mode 100644 index 0000000000000000000000000000000000000000..933f262432cc21de449fa98fd211e2f685a63db2 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/multi_agent_prioritized_episode_buffer.py @@ -0,0 +1,1079 @@ +import copy +import hashlib +from collections import defaultdict, deque +from typing import Dict, List, Optional, Set, Tuple, Union + +import numpy as np +import scipy +from numpy.typing import NDArray + +from ray.rllib.env.multi_agent_episode import MultiAgentEpisode +from ray.rllib.env.single_agent_episode import SingleAgentEpisode +from ray.rllib.execution.segment_tree import MinSegmentTree, SumSegmentTree +from ray.rllib.utils import force_list +from ray.rllib.utils.annotations import override +from ray.rllib.utils.replay_buffers.multi_agent_episode_buffer import ( + MultiAgentEpisodeReplayBuffer, +) +from ray.rllib.utils.replay_buffers.prioritized_episode_buffer import ( + PrioritizedEpisodeReplayBuffer, +) +from ray.rllib.utils.typing import ModuleID + + +class MultiAgentPrioritizedEpisodeReplayBuffer( + MultiAgentEpisodeReplayBuffer, PrioritizedEpisodeReplayBuffer +): + """Multi-agent episode replay buffer that stores episodes by their IDs. + + This class implements a replay buffer as used in "Prioritized Experience + Replay" (Schaul et al., 2016) for multi-agent reinforcement learning, + + Each "row" (a slot in a deque) in the buffer is occupied by one episode. If an + incomplete episode is added to the buffer and then another chunk of that episode is + added at a later time, the buffer will automatically concatenate the new fragment to + the original episode. This way, episodes can be completed via subsequent `add` + calls. + + Sampling returns a size `B` episode list (number of 'rows'), where each episode + holds a tuple tuple of the form + + `(o_t, a_t, sum(r_t+1:t+n), o_t+n)` + + where `o_t` is the observation in `t`, `a_t` the action chosen at observation `o_t`, + `o_t+n` is the observation `n` timesteps later and `sum(r_t+1:t+n)` is the sum of + all rewards collected over the time steps between `t+1` and `t+n`. The `n`-step can + be chosen freely when sampling and defaults to `1`. If `n_step` is a tuple it is + sampled uniformly across the interval defined by the tuple (for each row in the + batch). + + Each episode contains - in addition to the data tuples presented above - two further + elements in its `extra_model_outputs`, namely `n_steps` and `weights`. The former + holds the `n_step` used for the sampled timesteps in the episode and the latter the + corresponding (importance sampling) weight for the transition. + + After sampling priorities can be updated (for the last sampled episode list) with + `self.update_priorities`. This method assigns the new priorities automatically to + the last sampled timesteps. Note, this implies that sampling timesteps and updating + their corresponding priorities needs to alternate (e.g. sampling several times and + then updating the priorities would not work because the buffer caches the last + sampled timestep indices). + + .. testcode:: + + import gymnasium as gym + + from ray.rllib.env.multi_agent_episode import MultiAgentEpisode + from ray.rllib.examples.envs.classes.multi_agent import MultiAgentCartPole + from ray.rllib.utils.replay_buffers import ( + MultiAgentPrioritizedEpisodeReplayBuffer, + ) + + + # Create the environment. + env = MultiAgentCartPole({"num_agents": 2}) + + # Set up the loop variables + agent_ids = env.agents + agent_ids.append("__all__") + terminateds = {aid: False for aid in agent_ids} + truncateds = {aid: False for aid in agent_ids} + num_timesteps = 10000 + episodes = [] + + # Initialize the first episode entries. + eps = MultiAgentEpisode() + obs, infos = env.reset() + eps.add_env_reset(observations=obs, infos=infos) + + # Sample 10,000 env timesteps. + for i in range(num_timesteps): + # If terminated we create a new episode. + if eps.is_done: + episodes.append(eps.to_numpy()) + eps = MultiAgentEpisode() + terminateds = {aid: False for aid in agent_ids} + truncateds = {aid: False for aid in agent_ids} + obs, infos = env.reset() + eps.add_env_reset(observations=obs, infos=infos) + + # Sample a random action for all agents that should step in the episode + # next. + actions = { + aid: env.get_action_space(aid).sample() + for aid in eps.get_agents_to_act() + } + obs, rewards, terminateds, truncateds, infos = env.step(actions) + eps.add_env_step( + obs, + actions, + rewards, + infos, + terminateds=terminateds, + truncateds=truncateds + ) + + # Add the last (truncated) episode to the list of episodes. + if not eps.is_done: + episodes.append(eps) + + # Create the buffer. + buffer = MultiAgentPrioritizedEpisodeReplayBuffer() + # Add the list of episodes sampled. + buffer.add(episodes) + + # Pull a sample from the buffer using an `n-step` of 3. + sample = buffer.sample(num_items=256, gamma=0.95, n_step=3, beta=0.5) + """ + + def __init__( + self, + capacity: int = 10000, + *, + batch_size_B: int = 16, + batch_length_T: int = 1, + alpha: float = 1.0, + metrics_num_episodes_for_smoothing: int = 100, + **kwargs, + ): + """Initializes a `MultiAgentPrioritizedEpisodeReplayBuffer` object + + Args: + capacity: The total number of timesteps to be storable in this buffer. + Will start ejecting old episodes once this limit is reached. + batch_size_B: The number of episodes returned from `sample()`. + batch_length_T: The length of each episode in the episode list returned from + `sample()`. + alpha: The amount of prioritization to be used: `alpha=1.0` means full + prioritization, `alpha=0.0` means no prioritization. + """ + # Initialize the parents. + MultiAgentEpisodeReplayBuffer.__init__( + self, + capacity=capacity, + batch_size_B=batch_size_B, + batch_length_T=batch_length_T, + metrics_num_episodes_for_smoothing=metrics_num_episodes_for_smoothing, + **kwargs, + ) + PrioritizedEpisodeReplayBuffer.__init__( + self, + capacity=capacity, + batch_size_B=batch_size_B, + batch_length_T=batch_length_T, + alpha=alpha, + metrics_num_episodes_for_smoothing=metrics_num_episodes_for_smoothing, + **kwargs, + ) + + # TODO (simon): If not needed in synchronized sampling, remove. + # Maps indices from samples to their corresponding tree index. + self._sample_idx_to_tree_idx = {} + # Initialize segment trees for the priority weights per module. Note, b/c + # the trees are binary we need for them a capacity that is an exponential + # of 2. Double it to enable temporary buffer overflow (we need then free + # nodes in the trees). + tree_capacity = int(2 ** np.ceil(np.log2(self.capacity))) + + # Each module receives its own segment trees for independent sampling. + self._module_to_max_priority: Dict[ModuleID, float] = defaultdict(lambda: 1.0) + self._module_to_sum_segment: Dict[ModuleID, "SumSegmentTree"] = defaultdict( + lambda: SumSegmentTree(2 * tree_capacity) + ) + self._module_to_min_segment: Dict[ModuleID, "MinSegmentTree"] = defaultdict( + lambda: MinSegmentTree(2 * tree_capacity) + ) + # At initialization all nodes are free. + self._module_to_free_nodes: Dict[ModuleID, "deque"] = defaultdict( + lambda: deque(list(range(2 * tree_capacity)), maxlen=2 * tree_capacity) + ) + # Keep track of the maximum index used from the trees. This helps + # to not traverse the complete trees. + self._module_to_max_idx: Dict[ModuleID, int] = defaultdict(lambda: 0) + # Map from tree indices to sample indices (i.e. `self._indices`). + self._module_to_tree_idx_to_sample_idx: Dict[ModuleID, dict] = defaultdict( + lambda: {} + ) + # Map from module ID to the last sampled indices to update priorities. + self._module_to_last_sampled_indices: Dict[ModuleID, list] = defaultdict( + lambda: [] + ) + + @override(MultiAgentEpisodeReplayBuffer) + def add( + self, + episodes: Union[List["MultiAgentEpisode"], "MultiAgentEpisode"], + weight: Optional[Union[float, Dict[ModuleID, float]]] = None, + ) -> None: + """Adds incoming episodes to the replay buffer. + + Note, if the incoming episodes' time steps cause the buffer to overflow, + older episodes are evicted. Because episodes usually come in chunks and + not complete, this could lead to edge cases (e.g. with very small capacity + or very long episode length) where the first part of an episode is evicted + while the next part just comes in. + To defend against such case, the complete episode is evicted, including + the new chunk, unless the episode is the only one in the buffer. In the + latter case the buffer will be allowed to overflow in a temporary fashion, + i.e. during the next addition of samples to the buffer an attempt is made + to fall below capacity again. + + The user is advised to select a large enough buffer with regard to the maximum + expected episode length. + + Args: + episodes: A list of `SingleAgentEpisode`s that contain the episode data. + weight: A starting priority for the time steps in `episodes`. If `None` + the maximum priority is used, i.e. 1.0 (as suggested in the original + paper we scale weights to the interval [0.0, 1.0]). If a dictionary + is provided, it must contain the weights for each module. + + """ + # Define the weights. + weight_per_module = {} + # If no weight is provided, use the maximum priority. + if weight is None: + weight = self._max_priority + # If `weight` is a dictionary, use the module weights. + elif isinstance(dict, weight): + weight_per_module = weight + # Define the weight as the mean of the module weights. + weight = np.mean(list(weight.values())) + + episodes: List["MultiAgentEpisode"] = force_list(episodes) + + new_episode_ids: List[str] = [eps.id_ for eps in episodes] + # Calculate the total number of environment timesteps in the new episodes. + # Note, we need the potential new sum of timesteps to decide whether to + # evict episodes. + total_env_timesteps = sum([eps.env_steps() for eps in episodes]) + self._num_timesteps += total_env_timesteps + self._num_timesteps_added += total_env_timesteps + + # Set up some counters for metrics. + num_env_steps_added = 0 + agent_to_num_steps_added = defaultdict(int) + module_to_num_steps_added = defaultdict(int) + num_episodes_added = 0 + agent_to_num_episodes_added = defaultdict(int) + module_to_num_episodes_added = defaultdict(int) + num_episodes_evicted = 0 + agent_to_num_episodes_evicted = defaultdict(int) + module_to_num_episodes_evicted = defaultdict(int) + num_env_steps_evicted = 0 + agent_to_num_steps_evicted = defaultdict(int) + module_to_num_steps_evicted = defaultdict(int) + + # Evict old episodes. + eps_evicted_ids: Set[Union[str, int]] = set() + eps_evicted_idxs: Set[int] = set() + # Only evict episodes if the buffer is full and there is more than one + + while ( + self._num_timesteps > self.capacity + and self._num_remaining_episodes(new_episode_ids, eps_evicted_ids) != 1 + ): + # Evict episode. + evicted_episode = self.episodes.popleft() + eps_evicted_ids.add(evicted_episode.id_) + eps_evicted_idxs.add(self.episode_id_to_index.pop(evicted_episode.id_)) + # If this episode has a new chunk in the new episodes added, + # we subtract it again. + # TODO (sven, simon): Should we just treat such an episode chunk + # as a new episode? + if evicted_episode.id_ in new_episode_ids: + idx = next( + i + for i, eps in enumerate(episodes) + if eps.id_ == evicted_episode.id_ + ) + new_eps_to_evict = episodes.pop(idx) + # Remove the timesteps of the evicted new episode from the counter. + self._num_timesteps -= new_eps_to_evict.env_steps() + self._num_timesteps_added -= new_eps_to_evict.env_steps() + # Remove the timesteps of the evicted old episode from the counter. + self._num_timesteps -= evicted_episode.env_steps() + self._num_agent_timesteps -= evicted_episode.agent_steps() + self._num_episodes_evicted += 1 + # Increase the counters. + num_episodes_evicted += 1 + num_env_steps_evicted += evicted_episode.env_steps() + for aid, a_eps in evicted_episode.agent_episodes.items(): + mid = evicted_episode._agent_to_module_mapping[aid] + agent_to_num_episodes_evicted[aid] += 1 + module_to_num_episodes_evicted[mid] += 1 + agent_to_num_steps_evicted[aid] += a_eps.agent_steps() + module_to_num_steps_evicted[mid] += a_eps.agent_steps() + # Remove the module timesteps of the evicted episode from the counters. + self._evict_module_episodes(evicted_episode) + del evicted_episode + + # Add agent and module steps. + for eps in episodes: + self._num_agent_timesteps += eps.agent_steps() + self._num_agent_timesteps_added += eps.agent_steps() + # Update the module counters by the module timesteps. + self._update_module_counters(eps) + + # Remove corresponding indices, if episodes were evicted. + if eps_evicted_idxs: + new_indices = [] + # Each index 2-tuple is of the form (ma_episode_idx, timestep) and + # refers to a certain environment timestep in a certain multi-agent + # episode. + i = 0 + for idx_tuple in self._indices: + # If episode index is from an evicted episode, remove it from the + # indices and clean up. + if idx_tuple[0] in eps_evicted_idxs: + # Here we need the index of a multi-agent sample in the segment + # tree. + self._free_nodes.appendleft(idx_tuple[2]) + # Remove also the potentially maximum index. + self._max_idx -= 1 if self._max_idx == idx_tuple[2] else 0 + # Reset to defaults. + self._sum_segment[idx_tuple[2]] = 0.0 + self._min_segment[idx_tuple[2]] = float("inf") + sample_idx = self._tree_idx_to_sample_idx[idx_tuple[2]] + self._tree_idx_to_sample_idx.pop(idx_tuple[2]) + self._sample_idx_to_tree_idx.pop(sample_idx) + # Otherwise, keep the index. + else: + new_indices.append(idx_tuple) + self._tree_idx_to_sample_idx[idx_tuple[2]] = i + self._sample_idx_to_tree_idx[i] = idx_tuple[2] + i += 1 + # Assign the new list of indices. + self._indices = new_indices + # Also remove corresponding module indices. + for module_id, module_indices in self._module_to_indices.items(): + new_module_indices = [] + # Each index 4-tuple is of the form + # (ma_episode_idx, agent_id, timestep, segtree_idx) and refers to a + # certain agent timestep in a certain multi-agent episode. + i = 0 + for idx_quadlet in module_indices: + # If episode index is from an evicted episode, remove it from the + # indices and clean up. + if idx_quadlet[0] in eps_evicted_idxs: + # Here we need the index of a multi-agent sample in the segment + # tree. + self._module_to_free_nodes[module_id].appendleft(idx_quadlet[3]) + # Remove also the potentially maximum index per module. + self._module_to_max_idx[module_id] -= ( + 1 + if self._module_to_max_idx[module_id] == idx_quadlet[3] + else 0 + ) + # Set to defaults. + self._module_to_sum_segment[module_id][idx_quadlet[3]] = 0.0 + self._module_to_min_segment[module_id][idx_quadlet[3]] = float( + "inf" + ) + self._module_to_tree_idx_to_sample_idx[module_id].pop( + idx_quadlet[3] + ) + # Otherwise, keep the index. + else: + new_module_indices.append(idx_quadlet) + self._module_to_tree_idx_to_sample_idx[module_id][ + idx_quadlet[3] + ] = i + i += 1 + # Assign the new list of indices for the module. + self._module_to_indices[module_id] = new_module_indices + + j = len(self._indices) + for eps in episodes: + eps = copy.deepcopy(eps) + # If the episode is part of an already existing episode, concatenate. + if eps.id_ in self.episode_id_to_index: + eps_idx = self.episode_id_to_index[eps.id_] + existing_eps = self.episodes[eps_idx - self._num_episodes_evicted] + existing_len = len(existing_eps) + self._indices.extend( + [ + ( + eps_idx, + existing_len + i, + # Get the index in the segment trees. + self._get_free_node_and_assign(j + i, weight), + ) + for i in range(len(eps)) + ] + ) + # Add new module indices. + self._add_new_module_indices(eps, eps_idx, True, weight_per_module) + # Concatenate the episode chunk. + existing_eps.concat_episode(eps) + # Otherwise, create a new entry. + else: + # New episode. + self.episodes.append(eps) + # Update the counters. + num_episodes_added += 1 + for aid, a_eps in eps.agent_episodes.items(): + mid = eps._agent_to_module_mapping[aid] + agent_to_num_episodes_added[aid] += 1 + module_to_num_episodes_added[mid] += 1 + eps_idx = len(self.episodes) - 1 + self._num_episodes_evicted + self.episode_id_to_index[eps.id_] = eps_idx + self._indices.extend( + [ + (eps_idx, i, self._get_free_node_and_assign(j + i, weight)) + for i in range(len(eps)) + ] + ) + # Add new module indices. + self._add_new_module_indices(eps, eps_idx, False, weight_per_module) + + # Update the step counters. + num_env_steps_added += eps.env_steps() + for aid, e_eps in eps.agent_episodes.items(): + mid = eps._agent_to_module_mapping[aid] + agent_to_num_steps_added[aid] += e_eps.agent_steps() + module_to_num_steps_added[mid] += e_eps.agent_steps() + # Increase index to the new length of `self._indices`. + j = len(self._indices) + + # Update the adding metrics. + # Update the adding metrics. + self._update_add_metrics( + num_episodes_added=num_episodes_added, + num_env_steps_added=num_env_steps_added, + num_episodes_evicted=num_episodes_evicted, + num_env_steps_evicted=num_env_steps_evicted, + agent_to_num_episodes_added=agent_to_num_episodes_added, + agent_to_num_steps_added=agent_to_num_steps_added, + agent_to_num_episodes_evicted=agent_to_num_episodes_evicted, + agent_to_num_steps_evicted=agent_to_num_steps_evicted, + module_to_num_episodes_added=module_to_num_steps_added, + module_to_num_steps_added=module_to_num_episodes_added, + module_to_num_episodes_evicted=module_to_num_episodes_evicted, + module_to_num_steps_evicted=module_to_num_steps_evicted, + ) + + @override(MultiAgentEpisodeReplayBuffer) + def sample( + self, + num_items: Optional[int] = None, + *, + batch_size_B: Optional[int] = None, + batch_length_T: Optional[int] = None, + n_step: Optional[Union[int, Tuple]] = 1, + gamma: float = 0.99, + include_infos: bool = False, + include_extra_model_outputs: bool = False, + replay_mode: str = "independent", + modules_to_sample: Optional[List[ModuleID]] = None, + beta: float = 0.0, + **kwargs, + ) -> Union[List["MultiAgentEpisode"], List["SingleAgentEpisode"]]: + """Samples a list of episodes with multi-agent transitions. + + This sampling method also adds (importance sampling) weights to the returned + batch. See for prioritized sampling Schaul et al. (2016). + + Multi-agent transitions can be sampled either `"independent"` or + `"synchronized"` with the former sampling for each module independent agent + steps and the latter sampling agent transitions from the same environment step. + + The n-step parameter can be either a single integer or a tuple of two integers. + In the former case, the n-step is fixed to the given integer and in the latter + case, the n-step is sampled uniformly from the given range. Large n-steps could + potentially lead to many retries because not all samples might have a full + n-step transition. + + Sampling returns episode lists of size B (number of 'rows'), where each episode + holds a transition of the form + + `(o_t, a_t, sum(r_t+1:t+n), o_t+n, terminated_t+n, truncated_t+n)` + + where `o_t` is the observation in `t`, `a_t` the action chosen at observation + `o_t`, `o_t+n` is the observation `n` timesteps later and `sum(r_t+1:t+n)` is + the sum of all rewards collected over the time steps between `t+1` and `t+n`. + The `n`-step can be chosen freely when sampling and defaults to `1`. If `n_step` + is a tuple it is sampled uniformly across the interval defined by the tuple (for + each row in the batch), i.e. from the interval `[n_step[0], n_step[1]]`. + + If requested, `info`s of a transition's first and last timestep `t+n` and/or + `extra_model_outputs` from the first timestep (e.g. log-probabilities, etc.) are + added to the batch. + + Each episode contains - in addition to the data tuples presented above - two + further entries in its `extra_model_outputs`, namely `n_steps` and `weigths`. + The former holds the `n_step` used for each transition and the latter the + (importance sampling) weight of `1.0` for each row in the batch. This weight + is used for weighted loss calculations in the training process. + + Args: + num_items: The number of items to sample. If provided, `batch_size_B` + should be `None`. + batch_size_B: The batch size to sample. If provided, `num_items` + should be `None`. + batch_length_T: The length of the sampled batch. If not provided, the + default batch length is used. This feature is not yet implemented. + n_step: The n-step to sample. If the n-step is a tuple, the n-step is + sampled uniformly from the given range. If not provided, the default + n-step of `1` is used. + gamma: The discount factor for the n-step reward calculation. + include_infos: Whether to include the infos in the sampled episodes. + include_extra_model_outputs: Whether to include the extra model outputs + in the sampled episodes. + replay_mode: The replay mode to use for sampling. Either `"independent"` + or `"synchronized"`. + modules_to_sample: A list of module IDs to sample from. If not provided, + transitions for aall modules are sampled. + beta: The exponent of the importance sampling weight (see Schaul et + al. (2016)). A `beta=0.0` does not correct for the bias introduced + by prioritized replay and `beta=1.0` fully corrects for it. + + Returns: + A list of 1-step long single-agent episodes containing all basic episode + data and if requested infos and extra model outputs. In addition extra model + outputs hold the (importance sampling) weights and the n-step used for each + transition. + """ + assert beta >= 0.0 + + if num_items is not None: + assert batch_size_B is None, ( + "Cannot call `sample()` with both `num_items` and `batch_size_B` " + "provided! Use either one." + ) + batch_size_B = num_items + + # Use our default values if no sizes/lengths provided. + batch_size_B = batch_size_B or self.batch_size_B + # TODO (simon): Implement trajectory sampling for RNNs. + batch_length_T = batch_length_T or self.batch_length_T + + # Sample for each module independently. + if replay_mode == "independent": + return self._sample_independent( + batch_size_B=batch_size_B, + batch_length_T=batch_length_T, + n_step=n_step, + gamma=gamma, + include_infos=include_infos, + include_extra_model_outputs=include_extra_model_outputs, + modules_to_sample=modules_to_sample, + beta=beta, + ) + # Otherwise, sample synchronized. + else: + return self._sample_synchonized( + batch_size_B=batch_size_B, + batch_length_T=batch_length_T, + n_step=n_step, + gamma=gamma, + include_infos=include_infos, + include_extra_model_outputs=include_extra_model_outputs, + modules_to_sample=modules_to_sample, + ) + + @override(PrioritizedEpisodeReplayBuffer) + def update_priorities( + self, + priorities: Union[NDArray, Dict[ModuleID, NDArray]], + module_id: ModuleID, + ) -> None: + """Update the priorities of items at corresponding indices. + + Usually, incoming priorities are TD-errors. + + Args: + priorities: Numpy array containing the new priorities to be used + in sampling for the items in the last sampled batch. + """ + + assert len(priorities) == len(self._module_to_last_sampled_indices[module_id]) + + for idx, priority in zip( + self._module_to_last_sampled_indices[module_id], priorities + ): + # sample_idx = self._module_to_tree_idx_to_sample_idx[module_id][idx] + # ma_episode_idx = ( + # self._module_to_indices[module_id][sample_idx][0] + # - self._num_episodes_evicted + # ) + + # ma_episode_indices.append(ma_episode_idx) + # Note, TD-errors come in as absolute values or results from + # cross-entropy loss calculations. + # assert priority > 0, f"priority was {priority}" + priority = max(priority, 1e-12) + assert 0 <= idx < self._module_to_sum_segment[module_id].capacity + # TODO (simon): Create metrics. + # delta = priority**self._alpha - self._sum_segment[idx] + # Update the priorities in the segment trees. + self._module_to_sum_segment[module_id][idx] = priority**self._alpha + self._module_to_min_segment[module_id][idx] = priority**self._alpha + # Update the maximal priority. + self._module_to_max_priority[module_id] = max( + self._module_to_max_priority[module_id], priority + ) + # Clear the corresponding index list for the module. + self._module_to_last_sampled_indices[module_id].clear() + + # TODO (simon): Use this later for synchronized sampling. + # for ma_episode_idx in ma_episode_indices: + # ma_episode_tree_idx = self._sample_idx_to_tree_idx(ma_episode_idx) + # ma_episode_idx = + + # # Update the weights + # self._sum_segment[tree_idx] = sum( + # self._module_to_sum_segment[module_id][idx] + # for module_id, idx in self._tree_idx_to_sample_idx[tree_idx] + # ) + # self._min_segment[tree_idx] = min( + # self._module_to_min_segment[module_id][idx] + # for module_id, idx in self._tree_idx_to_sample_idx[tree_idx] + # ) + + @override(MultiAgentEpisodeReplayBuffer) + def get_state(self): + return ( + MultiAgentEpisodeReplayBuffer.get_state(self) + | PrioritizedEpisodeReplayBuffer.get_state(self) + | { + "_module_to_max_priority": list(self._module_to_max_priority.items()), + "_module_to_sum_segment": list(self._module_to_sum_segment.items()), + "_module_to_min_segment": list(self._module_to_min_segment.items()), + "_module_to_free_nodes": list(self._module_to_free_nodes.items()), + "_module_to_max_idx": list(self._module_to_max_idx.items()), + "_module_to_tree_idx_to_sample_idx": list( + self._module_to_tree_idx_to_sample_idx.items() + ), + "_module_to_last_sampled_indices": list( + self._module_to_last_sampled_indices.items() + ), + } + ) + + @override(MultiAgentEpisodeReplayBuffer) + def set_state(self, state) -> None: + MultiAgentEpisodeReplayBuffer.set_state(self, state) + PrioritizedEpisodeReplayBuffer.set_state(self, state) + self._module_to_max_priority = defaultdict( + lambda: 1.0, dict(state["_module_to_max_priority"]) + ) + tree_capacity = int(2 ** np.ceil(np.log2(self.capacity))) + self._module_to_sum_segment = defaultdict( + lambda: SumSegmentTree(2 * tree_capacity), + dict(state["_module_to_sum_segment"]), + ) + self._module_to_min_segment = defaultdict( + lambda: SumSegmentTree(2 * tree_capacity), + dict(state["_module_to_min_segment"]), + ) + self._module_to_free_nodes = defaultdict( + lambda: deque(list(range(2 * tree_capacity)), maxlen=2 * tree_capacity), + dict(state["_module_to_free_nodes"]), + ) + self._module_to_max_idx = defaultdict( + lambda: 0, dict(state["_module_to_max_idx"]) + ) + self._module_to_tree_idx_to_sample_idx = defaultdict( + lambda: {}, dict(state["_module_to_tree_idx_to_sample_idx"]) + ) + self._module_to_last_sampled_indices = defaultdict( + lambda: [], dict(state["_module_to_last_sampled_indices"]) + ) + + @override(MultiAgentEpisodeReplayBuffer) + def _add_new_module_indices( + self, + ma_episode: MultiAgentEpisode, + ma_episode_idx: int, + ma_episode_exists: bool = True, + weight: Optional[Union[float, Dict[ModuleID, float]]] = None, + ) -> None: + """Adds the module indices for new episode chunks. + + Args: + multi_agent_episode: The multi-agent episode to add the module indices for. + episode_idx: The index of the episode in the `self.episodes`. + """ + existing_ma_episode = None + if ma_episode_exists: + existing_ma_episode = self.episodes[ + self.episode_id_to_index[ma_episode.id_] - self._num_episodes_evicted + ] + + for agent_id in ma_episode.agent_ids: + # Get the corresponding module id. + module_id = ma_episode.module_for(agent_id) + # Get the module episode. + module_eps = ma_episode.agent_episodes[agent_id] + + # Is the agent episode already in the buffer's existing `ma_episode`? + if ma_episode_exists and agent_id in existing_ma_episode.agent_episodes: + existing_sa_eps_len = len(existing_ma_episode.agent_episodes[agent_id]) + # Otherwise, it is a new single-agent episode and we increase the counter. + else: + existing_sa_eps_len = 0 + self._num_module_episodes[module_id] += 1 + + # Add new module indices. + module_weight = weight.get( + module_id, self._module_to_max_priority[module_id] + ) + j = len(self._module_to_indices[module_id]) + self._module_to_indices[module_id].extend( + [ + ( + # Keep the MAE index for sampling. + ma_episode_idx, + agent_id, + existing_sa_eps_len + i, + # Get the index in the segment trees. + self._get_free_node_per_module_and_assign( + module_id, + j + i, + module_weight, + ), + ) + for i in range(len(module_eps)) + ] + ) + + @override(PrioritizedEpisodeReplayBuffer) + def _get_free_node_and_assign(self, sample_index, weight: float = 1.0) -> int: + """Gets the next free node in the segment trees. + + In addition the initial priorities for a new transition are added + to the segment trees and the index of the nodes is added to the + index mapping. + + Args: + sample_index: The index of the sample in the `self._indices` list. + weight: The initial priority weight to be used in sampling for + the item at index `sample_index`. + + Returns: + The index in the segment trees `self._sum_segment` and + `self._min_segment` for the item at index `sample_index` in + ``self._indices`. + """ + # Get an index from the free nodes in the segment trees. + idx = self._free_nodes.popleft() + self._max_idx = idx if idx > self._max_idx else self._max_idx + # Add the weight to the segments. + self._sum_segment[idx] = weight**self._alpha + self._min_segment[idx] = weight**self._alpha + # Map the index in the trees to the index in `self._indices`. + self._tree_idx_to_sample_idx[idx] = sample_index + self._sample_idx_to_tree_idx[sample_index] = idx + # Return the index. + return idx + + def _get_free_node_per_module_and_assign( + self, module_id: ModuleID, sample_index, weight: float = 1.0 + ) -> int: + """Gets the next free node in the segment trees. + + In addition the initial priorities for a new transition are added + to the segment trees and the index of the nodes is added to the + index mapping. + + Args: + sample_index: The index of the sample in the `self._indices` list. + weight: The initial priority weight to be used in sampling for + the item at index `sample_index`. + + Returns: + The index in the segment trees `self._sum_segment` and + `self._min_segment` for the item at index `sample_index` in + ``self._indices`. + """ + # Get an index from the free nodes in the segment trees. + idx = self._module_to_free_nodes[module_id].popleft() + self._module_to_max_idx[module_id] = ( + idx + if idx > self._module_to_max_idx[module_id] + else self._module_to_max_idx[module_id] + ) + # Add the weight to the segments. + # TODO (simon): Allow alpha to be chosen per module. + self._module_to_sum_segment[module_id][idx] = weight**self._alpha + self._module_to_min_segment[module_id][idx] = weight**self._alpha + # Map the index in the trees to the index in `self._indices`. + self._module_to_tree_idx_to_sample_idx[module_id][idx] = sample_index + # Return the index. + return idx + + @override(MultiAgentEpisodeReplayBuffer) + def _sample_independent( + self, + batch_size_B: Optional[int], + batch_length_T: Optional[int], + n_step: Optional[Union[int, Tuple]], + gamma: float, + include_infos: bool, + include_extra_model_outputs: bool, + modules_to_sample: Optional[List[ModuleID]], + beta: Optional[float], + ) -> List["SingleAgentEpisode"]: + """Samples a single-agent episode list with independent transitions. + + Note, independent sampling means that each module samples its transitions + independently from the replay buffer. This is the default sampling mode. + In contrast, synchronized sampling samples transitions from the same + environment step. + """ + + actual_n_step = n_step or 1 + # Sample the n-step if necessary. + random_n_step = isinstance(n_step, tuple) + + # Keep track of the indices that were sampled last for updating the + # weights later (see `ray.rllib.utils.replay_buffer.utils. + # update_priorities_in_episode_replay_buffer`). + # self._last_sampled_indices = defaultdict(lambda: []) + + sampled_episodes = [] + # Record the number of samples per module/agent/total. + num_env_steps_sampled = 0 + agent_to_num_steps_sampled = defaultdict(int) + module_to_num_steps_sampled = defaultdict(int) + # Record all the env step buffer indices that are contained in the sample. + sampled_env_step_idxs = set() + agent_to_sampled_env_step_idxs = defaultdict(set) + module_to_sampled_env_step_idxs = defaultdict(set) + # Record all the episode buffer indices that are contained in the sample. + sampled_episode_idxs = set() + agent_to_sampled_episode_idxs = defaultdict(set) + module_to_sampled_episode_idxs = defaultdict(set) + # Record all n-steps that have been used. + sampled_n_steps = [] + agent_to_sampled_n_steps = defaultdict(list) + module_to_sampled_n_steps = defaultdict(list) + # Record the number of times a sample needs to be resampled. + num_resamples = 0 + agent_to_num_resamples = defaultdict(int) + module_to_num_resamples = defaultdict(int) + + # TODO (simon): Ensure that the module has data and if not, skip it. + # TODO (sven): Should we then error out or skip? I think the Learner + # should handle this case when a module has no train data. + modules_to_sample = modules_to_sample or set(self._module_to_indices.keys()) + for module_id in modules_to_sample: + # Sample proportionally from the replay buffer's module segments using the + # respective weights. + module_total_segment_sum = self._module_to_sum_segment[module_id].sum() + module_p_min = ( + self._module_to_min_segment[module_id].min() / module_total_segment_sum + ) + # TODO (simon): Allow individual betas per module. + module_max_weight = (module_p_min * self.get_num_timesteps(module_id)) ** ( + -beta + ) + B = 0 + while B < batch_size_B: + # First, draw a random sample from Uniform(0, sum over all weights). + # Note, transitions with higher weight get sampled more often (as + # more random draws fall into larger intervals). + module_random_sum = ( + self.rng.random() * self._module_to_sum_segment[module_id].sum() + ) + # Get the highest index in the sum-tree for which the sum is + # smaller or equal the random sum sample. + # Note, in contrast to Schaul et al. (2018) (who sample + # `o_(t + n_step)`, Algorithm 1) we sample `o_t`. + module_idx = self._module_to_sum_segment[module_id].find_prefixsum_idx( + module_random_sum + ) + # Get the theoretical probability mass for drawing this sample. + module_p_sample = ( + self._module_to_sum_segment[module_id][module_idx] + / module_total_segment_sum + ) + # Compute the importance sampling weight. + module_weight = ( + module_p_sample * self.get_num_timesteps(module_id) + ) ** (-beta) + # Now, get the transition stored at this index. + index_quadlet = self._module_to_indices[module_id][ + self._module_to_tree_idx_to_sample_idx[module_id][module_idx] + ] + + # This will be an agent timestep (not env timestep). + # TODO (simon, sven): Maybe deprecate sa_episode_idx (_) in the index + # quads. Is there any need for it? + ma_episode_idx, agent_id, sa_episode_ts = ( + index_quadlet[0] - self._num_episodes_evicted, + index_quadlet[1], + index_quadlet[2], + ) + + # Get the multi-agent episode. + ma_episode = self.episodes[ma_episode_idx] + # Retrieve the single-agent episode for filtering. + sa_episode = ma_episode.agent_episodes[agent_id] + + # If we use random n-step sampling, draw the n-step for this item. + if random_n_step: + actual_n_step = int(self.rng.integers(n_step[0], n_step[1])) + # If we cannnot make the n-step, we resample. + if sa_episode_ts + actual_n_step > len(sa_episode): + num_resamples += 1 + agent_to_num_resamples[agent_id] += 1 + module_to_num_resamples[module_id] += 1 + continue + # Note, this will be the reward after executing action + # `a_(episode_ts)`. For `n_step>1` this will be the discounted sum + # of all rewards that were collected over the last n steps. + sa_raw_rewards = sa_episode.get_rewards( + slice(sa_episode_ts, sa_episode_ts + actual_n_step) + ) + sa_rewards = scipy.signal.lfilter( + [1], [1, -gamma], sa_raw_rewards[::-1], axis=0 + )[-1] + + sampled_sa_episode = SingleAgentEpisode( + id_=sa_episode.id_, + # Provide the IDs for the learner connector. + agent_id=sa_episode.agent_id, + module_id=sa_episode.module_id, + multi_agent_episode_id=ma_episode.id_, + # Ensure that each episode contains a tuple of the form: + # (o_t, a_t, sum(r_(t:t+n_step)), o_(t+n_step)) + # Two observations (t and t+n). + observations=[ + sa_episode.get_observations(sa_episode_ts), + sa_episode.get_observations(sa_episode_ts + actual_n_step), + ], + observation_space=sa_episode.observation_space, + infos=( + [ + sa_episode.get_infos(sa_episode_ts), + sa_episode.get_infos(sa_episode_ts + actual_n_step), + ] + if include_infos + else None + ), + actions=[sa_episode.get_actions(sa_episode_ts)], + action_space=sa_episode.action_space, + rewards=[sa_rewards], + # If the sampled single-agent episode is the single-agent episode's + # last time step, check, if the single-agent episode is terminated + # or truncated. + terminated=( + sa_episode_ts + actual_n_step >= len(sa_episode) + and sa_episode.is_terminated + ), + truncated=( + sa_episode_ts + actual_n_step >= len(sa_episode) + and sa_episode.is_truncated + ), + extra_model_outputs={ + "weights": [ + module_weight / module_max_weight * 1 + ], # actual_size=1 + "n_step": [actual_n_step], + **( + { + k: [ + sa_episode.get_extra_model_outputs(k, sa_episode_ts) + ] + for k in sa_episode.extra_model_outputs.keys() + } + if include_extra_model_outputs + else {} + ), + }, + # TODO (sven): Support lookback buffers. + len_lookback_buffer=0, + t_started=sa_episode_ts, + ) + # Append single-agent episode to the list of sampled episodes. + sampled_episodes.append(sampled_sa_episode) + # Add the episode indices. + sampled_episode_idxs.add(ma_episode_idx) + agent_to_sampled_episode_idxs[sa_episode.agent_id].add(sa_episode.id_) + module_to_sampled_episode_idxs[module_id].add(sa_episode.id_) + # Get the corresponding index in the `env_to_agent_t` mapping. + # TODO (simon, sven): This has complexity O(n) and could become + # expensive when the episode is large. Note, however, that conversion + # from list to `numpy.ndarray` is also complexity O(n) and we do this + # at many places - also in the `MultiAgentEpisode`s. + ma_episode_ts = ma_episode.env_t_to_agent_t[agent_id].data.index( + sa_episode_ts + ) + # Add the unique step hashes. + sampled_env_step_idxs.add( + hashlib.sha256( + f"{ma_episode.id_}-{ma_episode_ts}".encode() + ).hexdigest() + ) + # Add the unique agent step hashes. + hashed_agent_step = hashlib.sha256( + f"{sa_episode.id_}-{sa_episode_ts}".encode() + ).hexdigest() + agent_to_sampled_env_step_idxs[agent_id].add(hashed_agent_step) + # Add the unique module step hashes. + module_to_sampled_env_step_idxs[module_id].add(hashed_agent_step) + # Add the actual n-step used in generating this sample. + sampled_n_steps.append(actual_n_step) + agent_to_sampled_n_steps[agent_id].append(actual_n_step) + module_to_sampled_n_steps[module_id].append(actual_n_step) + + # Increase counter. + B += 1 + # Keep track of sampled indices for updating priorities later for each + # module. + self._module_to_last_sampled_indices[module_id].append(module_idx) + + # Increase the per module timesteps counter. + self.sampled_timesteps_per_module[module_id] += B + # Increase the counter metrics. + num_env_steps_sampled += B + agent_to_num_steps_sampled[agent_id] += B + module_to_num_steps_sampled[module_id] += B + + # Increase the counter for environment timesteps. + self.sampled_timesteps += batch_size_B + + # Update the sample metrics. + num_episodes_per_sample = len(sampled_episode_idxs) + num_env_steps_per_sample = len(sampled_env_step_idxs) + sampled_n_step = sum(sampled_n_steps) / batch_size_B + agent_to_num_episodes_per_sample = { + aid: len(l) for aid, l in agent_to_sampled_episode_idxs.items() + } + module_to_num_episodes_per_sample = { + mid: len(l) for mid, l in module_to_sampled_episode_idxs.items() + } + agent_to_num_steps_per_sample = { + aid: len(l) for aid, l in agent_to_sampled_env_step_idxs.items() + } + module_to_num_steps_per_sample = { + mid: len(l) for mid, l in module_to_sampled_env_step_idxs.items() + } + agent_to_sampled_n_step = { + aid: sum(l) / len(l) for aid, l in agent_to_sampled_n_steps.items() + } + module_to_sampled_n_step = { + mid: sum(l) / len(l) for mid, l in module_to_sampled_n_steps.items() + } + self._update_sample_metrics( + num_env_steps_sampled=num_env_steps_sampled, + num_episodes_per_sample=num_episodes_per_sample, + num_env_steps_per_sample=num_env_steps_per_sample, + sampled_n_step=sampled_n_step, + num_resamples=num_resamples, + agent_to_num_steps_sampled=agent_to_num_steps_sampled, + agent_to_num_episodes_per_sample=agent_to_num_episodes_per_sample, + agent_to_num_steps_per_sample=agent_to_num_steps_per_sample, + agent_to_sampled_n_step=agent_to_sampled_n_step, + agent_to_num_resamples=agent_to_num_resamples, + module_to_num_steps_sampled=module_to_num_steps_sampled, + module_to_num_episodes_per_sample=module_to_num_episodes_per_sample, + module_to_num_steps_per_sample=module_to_num_steps_per_sample, + module_to_sampled_n_step=module_to_sampled_n_step, + module_to_num_resamples=module_to_num_resamples, + ) + + # Return multi-agent dictionary. + return sampled_episodes diff --git a/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/multi_agent_prioritized_replay_buffer.py b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/multi_agent_prioritized_replay_buffer.py new file mode 100644 index 0000000000000000000000000000000000000000..3ac5d8cbd30ac755d514f21272498da88f6db5ab --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/multi_agent_prioritized_replay_buffer.py @@ -0,0 +1,280 @@ +import logging +from typing import Dict + +import numpy as np + +from ray.rllib.policy.rnn_sequencing import timeslice_along_seq_lens_with_overlap +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils.annotations import override +from ray.rllib.utils.replay_buffers.multi_agent_replay_buffer import ( + MultiAgentReplayBuffer, + ReplayMode, + merge_dicts_with_warning, +) +from ray.rllib.utils.replay_buffers.prioritized_replay_buffer import ( + PrioritizedReplayBuffer, +) +from ray.rllib.utils.replay_buffers.replay_buffer import ( + StorageUnit, +) +from ray.rllib.utils.typing import PolicyID, SampleBatchType +from ray.util.annotations import DeveloperAPI +from ray.util.debug import log_once +from ray.util.timer import _Timer + +logger = logging.getLogger(__name__) + + +@DeveloperAPI +class MultiAgentPrioritizedReplayBuffer( + MultiAgentReplayBuffer, PrioritizedReplayBuffer +): + """A prioritized replay buffer shard for multiagent setups. + + This buffer is meant to be run in parallel to distribute experiences + across `num_shards` shards. Unlike simpler buffers, it holds a set of + buffers - one for each policy ID. + """ + + def __init__( + self, + capacity: int = 10000, + storage_unit: str = "timesteps", + num_shards: int = 1, + replay_mode: str = "independent", + replay_sequence_override: bool = True, + replay_sequence_length: int = 1, + replay_burn_in: int = 0, + replay_zero_init_states: bool = True, + underlying_buffer_config: dict = None, + prioritized_replay_alpha: float = 0.6, + prioritized_replay_beta: float = 0.4, + prioritized_replay_eps: float = 1e-6, + **kwargs, + ): + """Initializes a MultiAgentReplayBuffer instance. + + Args: + capacity: The capacity of the buffer, measured in `storage_unit`. + storage_unit: Either 'timesteps', 'sequences' or + 'episodes'. Specifies how experiences are stored. If they + are stored in episodes, replay_sequence_length is ignored. + If they are stored in episodes, replay_sequence_length is + ignored. + num_shards: The number of buffer shards that exist in total + (including this one). + replay_mode: One of "independent" or "lockstep". Determines, + whether batches are sampled independently or to an equal + amount. + replay_sequence_override: If True, ignore sequences found in incoming + batches, slicing them into sequences as specified by + `replay_sequence_length` and `replay_sequence_burn_in`. This only has + an effect if storage_unit is `sequences`. + replay_sequence_length: The sequence length (T) of a single + sample. If > 1, we will sample B x T from this buffer. + replay_burn_in: The burn-in length in case + `replay_sequence_length` > 0. This is the number of timesteps + each sequence overlaps with the previous one to generate a + better internal state (=state after the burn-in), instead of + starting from 0.0 each RNN rollout. + replay_zero_init_states: Whether the initial states in the + buffer (if replay_sequence_length > 0) are alwayas 0.0 or + should be updated with the previous train_batch state outputs. + underlying_buffer_config: A config that contains all necessary + constructor arguments and arguments for methods to call on + the underlying buffers. This replaces the standard behaviour + of the underlying PrioritizedReplayBuffer. The config + follows the conventions of the general + replay_buffer_config. kwargs for subsequent calls of methods + may also be included. Example: + "replay_buffer_config": {"type": PrioritizedReplayBuffer, + "capacity": 10, "storage_unit": "timesteps", + prioritized_replay_alpha: 0.5, prioritized_replay_beta: 0.5, + prioritized_replay_eps: 0.5} + prioritized_replay_alpha: Alpha parameter for a prioritized + replay buffer. Use 0.0 for no prioritization. + prioritized_replay_beta: Beta parameter for a prioritized + replay buffer. + prioritized_replay_eps: Epsilon parameter for a prioritized + replay buffer. + ``**kwargs``: Forward compatibility kwargs. + """ + if "replay_mode" in kwargs and ( + kwargs["replay_mode"] == "lockstep" + or kwargs["replay_mode"] == ReplayMode.LOCKSTEP + ): + if log_once("lockstep_mode_not_supported"): + logger.error( + "Replay mode `lockstep` is not supported for " + "MultiAgentPrioritizedReplayBuffer. " + "This buffer will run in `independent` mode." + ) + kwargs["replay_mode"] = "independent" + + if underlying_buffer_config is not None: + if log_once("underlying_buffer_config_not_supported"): + logger.info( + "PrioritizedMultiAgentReplayBuffer instantiated " + "with underlying_buffer_config. This will " + "overwrite the standard behaviour of the " + "underlying PrioritizedReplayBuffer." + ) + prioritized_replay_buffer_config = underlying_buffer_config + else: + prioritized_replay_buffer_config = { + "type": PrioritizedReplayBuffer, + "alpha": prioritized_replay_alpha, + "beta": prioritized_replay_beta, + } + + shard_capacity = capacity // num_shards + MultiAgentReplayBuffer.__init__( + self, + capacity=shard_capacity, + storage_unit=storage_unit, + replay_sequence_override=replay_sequence_override, + replay_mode=replay_mode, + replay_sequence_length=replay_sequence_length, + replay_burn_in=replay_burn_in, + replay_zero_init_states=replay_zero_init_states, + underlying_buffer_config=prioritized_replay_buffer_config, + **kwargs, + ) + + self.prioritized_replay_eps = prioritized_replay_eps + self.update_priorities_timer = _Timer() + + @DeveloperAPI + @override(MultiAgentReplayBuffer) + def _add_to_underlying_buffer( + self, policy_id: PolicyID, batch: SampleBatchType, **kwargs + ) -> None: + """Add a batch of experiences to the underlying buffer of a policy. + + If the storage unit is `timesteps`, cut the batch into timeslices + before adding them to the appropriate buffer. Otherwise, let the + underlying buffer decide how slice batches. + + Args: + policy_id: ID of the policy that corresponds to the underlying + buffer + batch: SampleBatch to add to the underlying buffer + ``**kwargs``: Forward compatibility kwargs. + """ + # Merge kwargs, overwriting standard call arguments + kwargs = merge_dicts_with_warning(self.underlying_buffer_call_args, kwargs) + + # For the storage unit `timesteps`, the underlying buffer will + # simply store the samples how they arrive. For sequences and + # episodes, the underlying buffer may split them itself. + if self.storage_unit is StorageUnit.TIMESTEPS: + timeslices = batch.timeslices(1) + elif self.storage_unit is StorageUnit.SEQUENCES: + timeslices = timeslice_along_seq_lens_with_overlap( + sample_batch=batch, + seq_lens=batch.get(SampleBatch.SEQ_LENS) + if self.replay_sequence_override + else None, + zero_pad_max_seq_len=self.replay_sequence_length, + pre_overlap=self.replay_burn_in, + zero_init_states=self.replay_zero_init_states, + ) + elif self.storage_unit == StorageUnit.EPISODES: + timeslices = [] + for eps in batch.split_by_episode(): + if eps.get(SampleBatch.T)[0] == 0 and ( + eps.get(SampleBatch.TERMINATEDS, [True])[-1] + or eps.get(SampleBatch.TRUNCATEDS, [False])[-1] + ): + # Only add full episodes to the buffer + timeslices.append(eps) + else: + if log_once("only_full_episodes"): + logger.info( + "This buffer uses episodes as a storage " + "unit and thus allows only full episodes " + "to be added to it. Some samples may be " + "dropped." + ) + elif self.storage_unit == StorageUnit.FRAGMENTS: + timeslices = [batch] + else: + raise ValueError("Unknown `storage_unit={}`".format(self.storage_unit)) + + for slice in timeslices: + # If SampleBatch has prio-replay weights, average + # over these to use as a weight for the entire + # sequence. + if self.replay_mode is ReplayMode.INDEPENDENT: + if "weights" in slice and len(slice["weights"]): + weight = np.mean(slice["weights"]) + else: + weight = None + + if "weight" in kwargs and weight is not None: + if log_once("overwrite_weight"): + logger.warning( + "Adding batches with column " + "`weights` to this buffer while " + "providing weights as a call argument " + "to the add method results in the " + "column being overwritten." + ) + + kwargs = {"weight": weight, **kwargs} + else: + if "weight" in kwargs: + if log_once("lockstep_no_weight_allowed"): + logger.warning( + "Settings weights for batches in " + "lockstep mode is not allowed." + "Weights are being ignored." + ) + + kwargs = {**kwargs, "weight": None} + self.replay_buffers[policy_id].add(slice, **kwargs) + + @DeveloperAPI + @override(PrioritizedReplayBuffer) + def update_priorities(self, prio_dict: Dict) -> None: + """Updates the priorities of underlying replay buffers. + + Computes new priorities from td_errors and prioritized_replay_eps. + These priorities are used to update underlying replay buffers per + policy_id. + + Args: + prio_dict: A dictionary containing td_errors for + batches saved in underlying replay buffers. + """ + with self.update_priorities_timer: + for policy_id, (batch_indexes, td_errors) in prio_dict.items(): + new_priorities = np.abs(td_errors) + self.prioritized_replay_eps + self.replay_buffers[policy_id].update_priorities( + batch_indexes, new_priorities + ) + + @DeveloperAPI + @override(MultiAgentReplayBuffer) + def stats(self, debug: bool = False) -> Dict: + """Returns the stats of this buffer and all underlying buffers. + + Args: + debug: If True, stats of underlying replay buffers are + fetched with debug=True. + + Returns: + stat: Dictionary of buffer stats. + """ + stat = { + "add_batch_time_ms": round(1000 * self.add_batch_timer.mean, 3), + "replay_time_ms": round(1000 * self.replay_timer.mean, 3), + "update_priorities_time_ms": round( + 1000 * self.update_priorities_timer.mean, 3 + ), + } + for policy_id, replay_buffer in self.replay_buffers.items(): + stat.update( + {"policy_{}".format(policy_id): replay_buffer.stats(debug=debug)} + ) + return stat diff --git a/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/multi_agent_replay_buffer.py b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/multi_agent_replay_buffer.py new file mode 100644 index 0000000000000000000000000000000000000000..776e99b4e237f33676e80b63fddda3fd5b43e6a7 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/multi_agent_replay_buffer.py @@ -0,0 +1,392 @@ +import collections +import logging +from enum import Enum +from typing import Any, Dict, Optional + +from ray._common.deprecation import Deprecated +from ray.rllib.policy.rnn_sequencing import timeslice_along_seq_lens_with_overlap +from ray.rllib.policy.sample_batch import MultiAgentBatch, SampleBatch +from ray.rllib.utils.annotations import override +from ray.rllib.utils.from_config import from_config +from ray.rllib.utils.replay_buffers.replay_buffer import ( + _ALL_POLICIES, + ReplayBuffer, + StorageUnit, +) +from ray.rllib.utils.typing import PolicyID, SampleBatchType +from ray.util.annotations import DeveloperAPI +from ray.util.debug import log_once +from ray.util.timer import _Timer + +logger = logging.getLogger(__name__) + + +@DeveloperAPI +class ReplayMode(Enum): + LOCKSTEP = "lockstep" + INDEPENDENT = "independent" + + +@DeveloperAPI +def merge_dicts_with_warning(args_on_init, args_on_call): + """Merge argument dicts, overwriting args_on_call with warning. + + The MultiAgentReplayBuffer supports setting standard arguments for calls + of methods of the underlying buffers. These arguments can be + overwritten. Such overwrites trigger a warning to the user. + """ + for arg_name, arg_value in args_on_call.items(): + if arg_name in args_on_init: + if log_once("overwrite_argument_{}".format((str(arg_name)))): + logger.warning( + "Replay Buffer was initialized to have " + "underlying buffers methods called with " + "argument `{}={}`, but was subsequently called " + "with `{}={}`.".format( + arg_name, + args_on_init[arg_name], + arg_name, + arg_value, + ) + ) + return {**args_on_init, **args_on_call} + + +@DeveloperAPI +class MultiAgentReplayBuffer(ReplayBuffer): + """A replay buffer shard for multiagent setups. + + This buffer is meant to be run in parallel to distribute experiences + across `num_shards` shards. Unlike simpler buffers, it holds a set of + buffers - one for each policy ID. + """ + + def __init__( + self, + capacity: int = 10000, + storage_unit: str = "timesteps", + num_shards: int = 1, + replay_mode: str = "independent", + replay_sequence_override: bool = True, + replay_sequence_length: int = 1, + replay_burn_in: int = 0, + replay_zero_init_states: bool = True, + underlying_buffer_config: dict = None, + **kwargs, + ): + """Initializes a MultiAgentReplayBuffer instance. + + Args: + capacity: The capacity of the buffer, measured in `storage_unit`. + storage_unit: Either 'timesteps', 'sequences' or + 'episodes'. Specifies how experiences are stored. If they + are stored in episodes, replay_sequence_length is ignored. + num_shards: The number of buffer shards that exist in total + (including this one). + replay_mode: One of "independent" or "lockstep". Determines, + whether batches are sampled independently or to an equal + amount. + replay_sequence_override: If True, ignore sequences found in incoming + batches, slicing them into sequences as specified by + `replay_sequence_length` and `replay_sequence_burn_in`. This only has + an effect if storage_unit is `sequences`. + replay_sequence_length: The sequence length (T) of a single + sample. If > 1, we will sample B x T from this buffer. This + only has an effect if storage_unit is 'timesteps'. + replay_burn_in: This is the number of timesteps + each sequence overlaps with the previous one to generate a + better internal state (=state after the burn-in), instead of + starting from 0.0 each RNN rollout. This only has an effect + if storage_unit is `sequences`. + replay_zero_init_states: Whether the initial states in the + buffer (if replay_sequence_length > 0) are alwayas 0.0 or + should be updated with the previous train_batch state outputs. + underlying_buffer_config: A config that contains all necessary + constructor arguments and arguments for methods to call on + the underlying buffers. + ``**kwargs``: Forward compatibility kwargs. + """ + shard_capacity = capacity // num_shards + ReplayBuffer.__init__(self, capacity, storage_unit) + + # If the user provides an underlying buffer config, we use to + # instantiate and interact with underlying buffers + self.underlying_buffer_config = underlying_buffer_config + if self.underlying_buffer_config is not None: + self.underlying_buffer_call_args = self.underlying_buffer_config + else: + self.underlying_buffer_call_args = {} + self.replay_sequence_override = replay_sequence_override + self.replay_mode = replay_mode + self.replay_sequence_length = replay_sequence_length + self.replay_burn_in = replay_burn_in + self.replay_zero_init_states = replay_zero_init_states + self.replay_sequence_override = replay_sequence_override + + if ( + replay_sequence_length > 1 + and self.storage_unit is not StorageUnit.SEQUENCES + ): + logger.warning( + "MultiAgentReplayBuffer configured with " + "`replay_sequence_length={}`, but `storage_unit={}`. " + "replay_sequence_length will be ignored and set to 1.".format( + replay_sequence_length, storage_unit + ) + ) + self.replay_sequence_length = 1 + + if replay_sequence_length == 1 and self.storage_unit is StorageUnit.SEQUENCES: + logger.warning( + "MultiAgentReplayBuffer configured with " + "`replay_sequence_length={}`, but `storage_unit={}`. " + "This will result in sequences equal to timesteps.".format( + replay_sequence_length, storage_unit + ) + ) + + if replay_mode in ["lockstep", ReplayMode.LOCKSTEP]: + self.replay_mode = ReplayMode.LOCKSTEP + if self.storage_unit in [StorageUnit.EPISODES, StorageUnit.SEQUENCES]: + raise ValueError( + "MultiAgentReplayBuffer does not support " + "lockstep mode with storage unit `episodes`" + "or `sequences`." + ) + elif replay_mode in ["independent", ReplayMode.INDEPENDENT]: + self.replay_mode = ReplayMode.INDEPENDENT + else: + raise ValueError("Unsupported replay mode: {}".format(replay_mode)) + + if self.underlying_buffer_config: + ctor_args = { + **{"capacity": shard_capacity, "storage_unit": StorageUnit.FRAGMENTS}, + **self.underlying_buffer_config, + } + + def new_buffer(): + return from_config(self.underlying_buffer_config["type"], ctor_args) + + else: + # Default case + def new_buffer(): + self.underlying_buffer_call_args = {} + return ReplayBuffer( + self.capacity, + storage_unit=StorageUnit.FRAGMENTS, + ) + + self.replay_buffers = collections.defaultdict(new_buffer) + + # Metrics. + self.add_batch_timer = _Timer() + self.replay_timer = _Timer() + self._num_added = 0 + + def __len__(self) -> int: + """Returns the number of items currently stored in this buffer.""" + return sum(len(buffer._storage) for buffer in self.replay_buffers.values()) + + @DeveloperAPI + @Deprecated( + old="ReplayBuffer.replay()", + new="ReplayBuffer.sample(num_items)", + error=True, + ) + def replay(self, num_items: int = None, **kwargs) -> Optional[SampleBatchType]: + """Deprecated in favor of new ReplayBuffer API.""" + pass + + @DeveloperAPI + @override(ReplayBuffer) + def add(self, batch: SampleBatchType, **kwargs) -> None: + """Adds a batch to the appropriate policy's replay buffer. + + Turns the batch into a MultiAgentBatch of the DEFAULT_POLICY_ID if + it is not a MultiAgentBatch. Subsequently, adds the individual policy + batches to the storage. + + Args: + batch : The batch to be added. + ``**kwargs``: Forward compatibility kwargs. + """ + if batch is None: + if log_once("empty_batch_added_to_buffer"): + logger.info( + "A batch that is `None` was added to {}. This can be " + "normal at the beginning of execution but might " + "indicate an issue.".format(type(self).__name__) + ) + return + # Make a copy so the replay buffer doesn't pin plasma memory. + batch = batch.copy() + # Handle everything as if multi-agent. + batch = batch.as_multi_agent() + + with self.add_batch_timer: + pids_and_batches = self._maybe_split_into_policy_batches(batch) + for policy_id, sample_batch in pids_and_batches.items(): + self._add_to_underlying_buffer(policy_id, sample_batch, **kwargs) + + self._num_added += batch.count + + @DeveloperAPI + def _add_to_underlying_buffer( + self, policy_id: PolicyID, batch: SampleBatchType, **kwargs + ) -> None: + """Add a batch of experiences to the underlying buffer of a policy. + + If the storage unit is `timesteps`, cut the batch into timeslices + before adding them to the appropriate buffer. Otherwise, let the + underlying buffer decide how slice batches. + + Args: + policy_id: ID of the policy that corresponds to the underlying + buffer + batch: SampleBatch to add to the underlying buffer + ``**kwargs``: Forward compatibility kwargs. + """ + # Merge kwargs, overwriting standard call arguments + kwargs = merge_dicts_with_warning(self.underlying_buffer_call_args, kwargs) + + # For the storage unit `timesteps`, the underlying buffer will + # simply store the samples how they arrive. For sequences and + # episodes, the underlying buffer may split them itself. + if self.storage_unit is StorageUnit.TIMESTEPS: + timeslices = batch.timeslices(1) + elif self.storage_unit is StorageUnit.SEQUENCES: + timeslices = timeslice_along_seq_lens_with_overlap( + sample_batch=batch, + seq_lens=batch.get(SampleBatch.SEQ_LENS) + if self.replay_sequence_override + else None, + zero_pad_max_seq_len=self.replay_sequence_length, + pre_overlap=self.replay_burn_in, + zero_init_states=self.replay_zero_init_states, + ) + elif self.storage_unit == StorageUnit.EPISODES: + timeslices = [] + for eps in batch.split_by_episode(): + if eps.get(SampleBatch.T)[0] == 0 and ( + eps.get(SampleBatch.TERMINATEDS, [True])[-1] + or eps.get(SampleBatch.TRUNCATEDS, [False])[-1] + ): + # Only add full episodes to the buffer + timeslices.append(eps) + else: + if log_once("only_full_episodes"): + logger.info( + "This buffer uses episodes as a storage " + "unit and thus allows only full episodes " + "to be added to it. Some samples may be " + "dropped." + ) + elif self.storage_unit == StorageUnit.FRAGMENTS: + timeslices = [batch] + else: + raise ValueError("Unknown `storage_unit={}`".format(self.storage_unit)) + + for slice in timeslices: + self.replay_buffers[policy_id].add(slice, **kwargs) + + @DeveloperAPI + @override(ReplayBuffer) + def sample( + self, num_items: int, policy_id: Optional[PolicyID] = None, **kwargs + ) -> Optional[SampleBatchType]: + """Samples a MultiAgentBatch of `num_items` per one policy's buffer. + + If less than `num_items` records are in the policy's buffer, + some samples in the results may be repeated to fulfil the batch size + `num_items` request. Returns an empty batch if there are no items in + the buffer. + + Args: + num_items: Number of items to sample from a policy's buffer. + policy_id: ID of the policy that created the experiences we sample. If + none is given, sample from all policies. + + Returns: + Concatenated MultiAgentBatch of items. + ``**kwargs``: Forward compatibility kwargs. + """ + # Merge kwargs, overwriting standard call arguments + kwargs = merge_dicts_with_warning(self.underlying_buffer_call_args, kwargs) + + with self.replay_timer: + # Lockstep mode: Sample from all policies at the same time an + # equal amount of steps. + if self.replay_mode == ReplayMode.LOCKSTEP: + assert ( + policy_id is None + ), "`policy_id` specifier not allowed in `lockstep` mode!" + # In lockstep mode we sample MultiAgentBatches + return self.replay_buffers[_ALL_POLICIES].sample(num_items, **kwargs) + elif policy_id is not None: + sample = self.replay_buffers[policy_id].sample(num_items, **kwargs) + return MultiAgentBatch({policy_id: sample}, sample.count) + else: + samples = {} + for policy_id, replay_buffer in self.replay_buffers.items(): + samples[policy_id] = replay_buffer.sample(num_items, **kwargs) + return MultiAgentBatch(samples, sum(s.count for s in samples.values())) + + @DeveloperAPI + @override(ReplayBuffer) + def stats(self, debug: bool = False) -> Dict: + """Returns the stats of this buffer and all underlying buffers. + + Args: + debug: If True, stats of underlying replay buffers + are fetched with debug=True. + + Returns: + stat: Dictionary of buffer stats. + """ + stat = { + "add_batch_time_ms": round(1000 * self.add_batch_timer.mean, 3), + "replay_time_ms": round(1000 * self.replay_timer.mean, 3), + } + for policy_id, replay_buffer in self.replay_buffers.items(): + stat.update( + {"policy_{}".format(policy_id): replay_buffer.stats(debug=debug)} + ) + return stat + + @DeveloperAPI + @override(ReplayBuffer) + def get_state(self) -> Dict[str, Any]: + """Returns all local state. + + Returns: + The serializable local state. + """ + state = {"num_added": self._num_added, "replay_buffers": {}} + for policy_id, replay_buffer in self.replay_buffers.items(): + state["replay_buffers"][policy_id] = replay_buffer.get_state() + return state + + @DeveloperAPI + @override(ReplayBuffer) + def set_state(self, state: Dict[str, Any]) -> None: + """Restores all local state to the provided `state`. + + Args: + state: The new state to set this buffer. Can be obtained by + calling `self.get_state()`. + """ + self._num_added = state["num_added"] + buffer_states = state["replay_buffers"] + for policy_id in buffer_states.keys(): + self.replay_buffers[policy_id].set_state(buffer_states[policy_id]) + + def _maybe_split_into_policy_batches(self, batch: SampleBatchType): + """Returns a dict of policy IDs and batches, depending on our replay mode. + + This method helps with splitting up MultiAgentBatches only if the + self.replay_mode requires it. + """ + if self.replay_mode == ReplayMode.LOCKSTEP: + return {_ALL_POLICIES: batch} + else: + return batch.policy_batches diff --git a/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/prioritized_episode_buffer.py b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/prioritized_episode_buffer.py new file mode 100644 index 0000000000000000000000000000000000000000..46e5663c5390e0e6b638f85bf66d67e806bddfbb --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/prioritized_episode_buffer.py @@ -0,0 +1,732 @@ +import copy +import hashlib +from collections import deque +from typing import Any, Dict, List, Optional, Tuple, Union + +import numpy as np +import scipy +from numpy.typing import NDArray + +from ray.rllib.core import DEFAULT_AGENT_ID, DEFAULT_MODULE_ID +from ray.rllib.env.single_agent_episode import SingleAgentEpisode +from ray.rllib.execution.segment_tree import MinSegmentTree, SumSegmentTree +from ray.rllib.utils import force_list +from ray.rllib.utils.annotations import ( + override, +) +from ray.rllib.utils.replay_buffers.episode_replay_buffer import EpisodeReplayBuffer +from ray.rllib.utils.typing import ModuleID, SampleBatchType + + +class PrioritizedEpisodeReplayBuffer(EpisodeReplayBuffer): + """Prioritized Replay Buffer that stores episodes by their ID. + + This replay buffer stores episode data (more specifically `SingleAgentEpisode` + objects) and implements prioritized experience replay first proposed + in the paper by Schaul et al. (2016, https://arxiv.org/abs/1511.05952). + + Implementation is based on segment trees as suggested by the authors of + the cited paper, i.e. we use proportional prioritization with an order + of O(log N) in updating and sampling. + + Each "row" (a slot in a deque) in the buffer is occupied by one episode. If an + incomplete episode is added to the buffer and then another chunk of that episode is + added at a later time, the buffer will automatically concatenate the new fragment to + the original episode. This way, episodes can be completed via subsequent `add` + calls. + + Sampling returns a size `B` episode list (number of 'rows'), where each episode + holds a tuple tuple of the form + + `(o_t, a_t, sum(r_t+1:t+n), o_t+n)` + + where `o_t` is the observation in `t`, `a_t` the action chosen at observation `o_t`, + `o_t+n` is the observation `n` timesteps later and `sum(r_t+1:t+n)` is the sum of + all rewards collected over the time steps between `t+1` and `t+n`. The `n`-step can + be chosen freely when sampling and defaults to `1`. If `n_step` is a tuple it is + sampled uniformly across the interval defined by the tuple (for each row in the + batch). + + Each episode contains - in addition to the data tuples presented above - two further + elements in its ` extra_model_outputs`, namely `n_steps` and `weights`. The former + holds the `n_step` used for the sampled timesteps in the episode and the latter the + corresponding (importance sampling) weight for the transition. + + After sampling priorities can be updated (for the last sampled episode list) with + `self.update_priorities`. This method assigns the new priorities automatically to + the last sampled timesteps. Note, this implies that sampling timesteps and updating + their corresponding priorities needs to alternate (e.g. sampling several times and + then updating the priorities would not work because the buffer caches the last + sampled timestep indices). + + .. testcode:: + + import gymnasium as gym + + from ray.rllib.env.single_agent_episode import SingleAgentEpisode + from ray.rllib.utils.replay_buffers import ( + PrioritizedEpisodeReplayBuffer + ) + + # Create the environment. + env = gym.make("CartPole-v1") + + # Set up the loop variables + terminated = False + truncated = False + num_timesteps = 10000 + episodes = [] + + # Initialize the first episode entries. + eps = SingleAgentEpisode() + obs, info = env.reset() + eps.add_env_reset(obs, info) + + # Sample 10,000 timesteps. + for i in range(num_timesteps): + # If terminated we create a new episode. + if terminated: + episodes.append(eps.to_numpy()) + eps = SingleAgentEpisode() + obs, info = env.reset() + eps.add_env_reset(obs, info) + + action = env.action_space.sample() + obs, reward, terminated, truncated, info = env.step(action) + eps.add_env_step( + obs, + action, + reward, + info, + terminated=terminated, + truncated=truncated + ) + + # Add the last (truncated) episode to the list of episodes. + if not terminated or truncated: + episodes.append(eps) + + # Create the buffer. + buffer = PrioritizedEpisodeReplayBuffer() + # Add the list of episodes sampled. + buffer.add(episodes) + + # Pull a sample from the buffer using an `n-step` of 3. + sample = buffer.sample(num_items=256, gamma=0.95, n_step=3) + """ + + def __init__( + self, + capacity: int = 10000, + *, + batch_size_B: int = 16, + batch_length_T: int = 1, + alpha: float = 1.0, + metrics_num_episodes_for_smoothing: int = 100, + **kwargs, + ): + """Initializes a `PrioritizedEpisodeReplayBuffer` object + + Args: + capacity: The total number of timesteps to be storable in this buffer. + Will start ejecting old episodes once this limit is reached. + batch_size_B: The number of episodes returned from `sample()`. + batch_length_T: The length of each episode in the episode list returned from + `sample()`. + alpha: The amount of prioritization to be used: `alpha=1.0` means full + prioritization, `alpha=0.0` means no prioritization. + """ + super().__init__( + capacity=capacity, + batch_size_B=batch_size_B, + batch_length_T=batch_length_T, + metrics_num_episodes_for_smoothing=metrics_num_episodes_for_smoothing, + ) + + # `alpha` should be non-negative. + assert alpha >= 0 + self._alpha = alpha + + # Initialize segment trees for the priority weights. Note, b/c the trees + # are binary we need for them a capacity that is an exponential of 2. + # Double it to enable temporary buffer overflow (we need then free nodes + # in the trees). + tree_capacity = int(2 ** np.ceil(np.log2(self.capacity))) + + self._max_priority = 1.0 + self._sum_segment = SumSegmentTree(2 * tree_capacity) + self._min_segment = MinSegmentTree(2 * tree_capacity) + # At initialization all nodes are free. + self._free_nodes = deque( + list(range(2 * tree_capacity)), maxlen=2 * tree_capacity + ) + # Keep track of the maximum index used from the trees. This helps + # to not traverse the complete trees. + self._max_idx = 0 + # Map from tree indices to sample indices (i.e. `self._indices`). + self._tree_idx_to_sample_idx = {} + # Keep track of the indices that were sampled last for updating the + # weights later. + self._last_sampled_indices = [] + + @override(EpisodeReplayBuffer) + def add( + self, + episodes: Union[List["SingleAgentEpisode"], "SingleAgentEpisode"], + weight: Optional[float] = None, + ) -> None: + """Adds incoming episodes to the replay buffer. + + Note, if the incoming episodes' time steps cause the buffer to overflow, + older episodes are evicted. Because episodes usually come in chunks and + not complete, this could lead to edge cases (e.g. with very small capacity + or very long episode length) where the first part of an episode is evicted + while the next part just comes in. + To defend against such case, the complete episode is evicted, including + the new chunk, unless the episode is the only one in the buffer. In the + latter case the buffer will be allowed to overflow in a temporary fashion, + i.e. during the next addition of samples to the buffer an attempt is made + to fall below capacity again. + + The user is advised to select a large enough buffer with regard to the maximum + expected episode length. + + Args: + episodes: A list of `SingleAgentEpisode`s that contain the episode data. + weight: A starting priority for the time steps in `episodes`. If `None` + the maximum priority is used, i.e. 1.0 (as suggested in the original + paper we scale weights to the interval [0.0, 1.0]).. + """ + + # TODO (sven, simon): Eventually allow here an array? + if weight is None: + weight = self._max_priority + + episodes = force_list(episodes) + + # Set up some counters for metrics. + num_env_steps_added = 0 + agent_to_num_steps_added = {DEFAULT_AGENT_ID: 0} + module_to_num_steps_added = {DEFAULT_MODULE_ID: 0} + num_episodes_added = 0 + agent_to_num_episodes_added = {DEFAULT_AGENT_ID: 0} + module_to_num_episodes_added = {DEFAULT_MODULE_ID: 0} + num_episodes_evicted = 0 + agent_to_num_episodes_evicted = {DEFAULT_AGENT_ID: 0} + module_to_num_episodes_evicted = {DEFAULT_MODULE_ID: 0} + num_env_steps_evicted = 0 + agent_to_num_steps_evicted = {DEFAULT_AGENT_ID: 0} + module_to_num_steps_evicted = {DEFAULT_MODULE_ID: 0} + + # Add first the timesteps of new episodes to have info about how many + # episodes should be evicted to stay below capacity. + new_episode_ids = [] + for eps in episodes: + new_episode_ids.append(eps.id_) + self._num_timesteps += len(eps) + self._num_timesteps_added += len(eps) + + eps_evicted = [] + eps_evicted_ids = [] + eps_evicted_idxs = [] + while ( + self._num_timesteps > self.capacity + and self._num_remaining_episodes(new_episode_ids, eps_evicted_ids) != 1 + ): + # Evict episode + eps_evicted.append(self.episodes.popleft()) + eps_evicted_ids.append(eps_evicted[-1].id_) + eps_evicted_idxs.append(self.episode_id_to_index.pop(eps_evicted_ids[-1])) + num_episodes_evicted += 1 + num_env_steps_evicted += len(eps_evicted[-1]) + agent_to_num_episodes_evicted[DEFAULT_AGENT_ID] += 1 + agent_to_num_steps_evicted[DEFAULT_AGENT_ID] += eps_evicted[ + -1 + ].agent_steps() + module_to_num_episodes_evicted[DEFAULT_MODULE_ID] += 1 + module_to_num_steps_evicted[DEFAULT_MODULE_ID] += eps_evicted[ + -1 + ].agent_steps() + # If this episode has a new chunk in the new episodes added, + # we subtract it again. + # TODO (sven, simon): Should we just treat such an episode chunk + # as a new episode? + if eps_evicted_ids[-1] in new_episode_ids: + # TODO (simon): Apply the same logic as in the MA-case. + len_to_subtract = len( + episodes[new_episode_ids.index(eps_evicted_idxs[-1])] + ) + self._num_timesteps -= len_to_subtract + self._num_timesteps_added -= len_to_subtract + # Remove the timesteps of the evicted episode from the counter. + self._num_timesteps -= len(eps_evicted[-1]) + self._num_episodes_evicted += 1 + + # Remove corresponding indices, if episodes were evicted. + # TODO (simon): Refactor into method such that MultiAgent + # version can inherit. + if eps_evicted_idxs: + new_indices = [] + i = 0 + for idx_triple in self._indices: + # If the index comes from an evicted episode free the nodes. + if idx_triple[0] in eps_evicted_idxs: + # Here we need the index of a sample in the segment tree. + self._free_nodes.appendleft(idx_triple[2]) + # Also remove the potentially maximum index. + self._max_idx -= 1 if self._max_idx == idx_triple[2] else 0 + self._sum_segment[idx_triple[2]] = 0.0 + self._min_segment[idx_triple[2]] = float("inf") + self._tree_idx_to_sample_idx.pop(idx_triple[2]) + # Otherwise update the index in the index mapping. + else: + new_indices.append(idx_triple) + self._tree_idx_to_sample_idx[idx_triple[2]] = i + i += 1 + # Assign the new list of indices. + self._indices = new_indices + + # Now append the indices for the new episodes. + j = len(self._indices) + for eps in episodes: + # If the episode chunk is part of an evicted episode continue. + if eps.id_ in eps_evicted_ids: + continue + # Otherwise, add the episode data to the buffer. + else: + eps = copy.deepcopy(eps) + # If the episode is part of an already existing episode, concatenate. + if eps.id_ in self.episode_id_to_index: + eps_idx = self.episode_id_to_index[eps.id_] + existing_eps = self.episodes[eps_idx - self._num_episodes_evicted] + old_len = len(existing_eps) + self._indices.extend( + [ + ( + eps_idx, + old_len + i, + # Get the index in the segment trees. + self._get_free_node_and_assign(j + i, weight), + ) + for i in range(len(eps)) + ] + ) + existing_eps.concat_episode(eps) + # Otherwise, create a new entry. + else: + num_episodes_added += 1 + agent_to_num_episodes_added[DEFAULT_AGENT_ID] += 1 + module_to_num_episodes_added[DEFAULT_MODULE_ID] += 1 + self.episodes.append(eps) + eps_idx = len(self.episodes) - 1 + self._num_episodes_evicted + self.episode_id_to_index[eps.id_] = eps_idx + self._indices.extend( + [ + ( + eps_idx, + i, + self._get_free_node_and_assign(j + i, weight), + ) + for i in range(len(eps)) + ] + ) + num_env_steps_added += len(eps) + agent_to_num_steps_added[DEFAULT_AGENT_ID] += eps.agent_steps() + module_to_num_steps_added[DEFAULT_MODULE_ID] += eps.agent_steps() + # Increase index to the new length of `self._indices`. + j = len(self._indices) + + # Increase metrics. + self._update_add_metrics( + num_episodes_added=num_episodes_added, + num_env_steps_added=num_env_steps_added, + num_episodes_evicted=num_episodes_evicted, + num_env_steps_evicted=num_env_steps_evicted, + agent_to_num_episodes_added=agent_to_num_episodes_added, + agent_to_num_steps_added=agent_to_num_steps_added, + agent_to_num_episodes_evicted=agent_to_num_episodes_evicted, + agent_to_num_steps_evicted=agent_to_num_steps_evicted, + module_to_num_episodes_added=module_to_num_steps_added, + module_to_num_steps_added=module_to_num_episodes_added, + module_to_num_episodes_evicted=module_to_num_episodes_evicted, + module_to_num_steps_evicted=module_to_num_steps_evicted, + ) + + @override(EpisodeReplayBuffer) + def sample( + self, + num_items: Optional[int] = None, + *, + batch_size_B: Optional[int] = None, + batch_length_T: Optional[int] = None, + n_step: Optional[Union[int, Tuple]] = None, + beta: float = 0.0, + gamma: float = 0.99, + include_infos: bool = False, + include_extra_model_outputs: bool = False, + to_numpy: bool = False, + **kwargs, + ) -> SampleBatchType: + """Samples from a buffer in a prioritized way. + + This sampling method also adds (importance sampling) weights to + the returned batch. See for prioritized sampling Schaul et al. + (2016). + + Each sampled item defines a transition of the form: + + `(o_t, a_t, sum(r_(t+1:t+n+1)), o_(t+n), terminated_(t+n), truncated_(t+n))` + + where `o_(t+n)` is drawn by prioritized sampling, i.e. the priority + of `o_(t+n)` led to the sample and defines the importance weight that + is returned in the sample batch. `n` is defined by the `n_step` applied. + + If requested, `info`s of a transitions last timestep `t+n` are added to + the batch. + + Args: + num_items: Number of items (transitions) to sample from this + buffer. + batch_size_B: The number of rows (transitions) to return in the + batch + batch_length_T: THe sequence length to sample. At this point in time + only sequences of length 1 are possible. + n_step: The n-step to apply. For the default the batch contains in + `"new_obs"` the observation and in `"obs"` the observation `n` + time steps before. The reward will be the sum of rewards + collected in between these two observations and the action will + be the one executed n steps before such that we always have the + state-action pair that triggered the rewards. + If `n_step` is a tuple, it is considered as a range to sample + from. If `None`, we use `n_step=1`. + beta: The exponent of the importance sampling weight (see Schaul et + al. (2016)). A `beta=0.0` does not correct for the bias introduced + by prioritized replay and `beta=1.0` fully corrects for it. + gamma: The discount factor to be used when applying n-step calculations. + The default of `0.99` should be replaced by the `Algorithm`s + discount factor. + include_infos: A boolean indicating, if `info`s should be included in + the batch. This could be of advantage, if the `info` contains + values from the environment important for loss computation. If + `True`, the info at the `"new_obs"` in the batch is included. + include_extra_model_outputs: A boolean indicating, if + `extra_model_outputs` should be included in the batch. This could be + of advantage, if the `extra_mdoel_outputs` contain outputs from the + model important for loss computation and only able to compute with the + actual state of model e.g. action log-probabilities, etc.). If `True`, + the extra model outputs at the `"obs"` in the batch is included (the + timestep at which the action is computed). + + Returns: + A list of 1-step long episodes containing all basic episode data and if + requested infos and extra model outputs. + """ + assert beta >= 0.0 + + if num_items is not None: + assert batch_size_B is None, ( + "Cannot call `sample()` with both `num_items` and `batch_size_B` " + "provided! Use either one." + ) + batch_size_B = num_items + + # Use our default values if no sizes/lengths provided. + batch_size_B = batch_size_B or self.batch_size_B + # TODO (simon): Implement trajectory sampling for RNNs. + batch_length_T = batch_length_T or self.batch_length_T + + # Sample the n-step if necessary. + actual_n_step = n_step or 1 + random_n_step = isinstance(n_step, tuple) + + # Keep track of the indices that were sampled last for updating the + # weights later (see `ray.rllib.utils.replay_buffer.utils. + # update_priorities_in_episode_replay_buffer`). + self._last_sampled_indices = [] + + sampled_episodes = [] + # Record all the env step buffer indices that are contained in the sample. + sampled_env_step_idxs = set() + # Record all the episode buffer indices that are contained in the sample. + sampled_episode_idxs = set() + # Record all n-steps that have been used. + sampled_n_steps = [] + # Record the number of times it needs to be resampled. + num_resamples = 0 + + # Sample proportionally from replay buffer's segments using the weights. + total_segment_sum = self._sum_segment.sum() + p_min = self._min_segment.min() / total_segment_sum + max_weight = (p_min * self.get_num_timesteps()) ** (-beta) + B = 0 + while B < batch_size_B: + # First, draw a random sample from Uniform(0, sum over all weights). + # Note, transitions with higher weight get sampled more often (as + # more random draws fall into larger intervals). + random_sum = self.rng.random() * self._sum_segment.sum() + # Get the highest index in the sum-tree for which the sum is + # smaller or equal the random sum sample. + # Note, in contrast to Schaul et al. (2018) (who sample `o_(t + n_step)`, + # Algorithm 1) we sample `o_t`. + idx = self._sum_segment.find_prefixsum_idx(random_sum) + # Get the theoretical probability mass for drawing this sample. + p_sample = self._sum_segment[idx] / total_segment_sum + # Compute the importance sampling weight. + weight = (p_sample * self.get_num_timesteps()) ** (-beta) + # Now, get the transition stored at this index. + index_triple = self._indices[self._tree_idx_to_sample_idx[idx]] + + # Compute the actual episode index (offset by the number of + # already evicted episodes) + episode_idx, episode_ts = ( + index_triple[0] - self._num_episodes_evicted, + index_triple[1], + ) + episode = self.episodes[episode_idx] + + # If we use random n-step sampling, draw the n-step for this item. + if random_n_step: + actual_n_step = int(self.rng.integers(n_step[0], n_step[1])) + + # Skip, if we are too far to the end and `episode_ts` + n_step would go + # beyond the episode's end. + if episode_ts + actual_n_step > len(episode): + num_resamples += 1 + continue + + # Note, this will be the reward after executing action + # `a_(episode_ts-n_step+1)`. For `n_step>1` this will be the discounted + # sum of all discounted rewards that were collected over the last n steps. + raw_rewards = episode.get_rewards( + slice(episode_ts, episode_ts + actual_n_step) + ) + rewards = scipy.signal.lfilter([1], [1, -gamma], raw_rewards[::-1], axis=0)[ + -1 + ] + + # Generate the episode to be returned. + sampled_episode = SingleAgentEpisode( + # Ensure that each episode contains a tuple of the form: + # (o_t, a_t, sum(r_(t:t+n_step)), o_(t+n_step)) + # Two observations (t and t+n). + observations=[ + episode.get_observations(episode_ts), + episode.get_observations(episode_ts + actual_n_step), + ], + observation_space=episode.observation_space, + infos=( + [ + episode.get_infos(episode_ts), + episode.get_infos(episode_ts + actual_n_step), + ] + if include_infos + else None + ), + actions=[episode.get_actions(episode_ts)], + action_space=episode.action_space, + rewards=[rewards], + # If the sampled time step is the episode's last time step check, if + # the episode is terminated or truncated. + terminated=( + False + if episode_ts + actual_n_step < len(episode) + else episode.is_terminated + ), + truncated=( + False + if episode_ts + actual_n_step < len(episode) + else episode.is_truncated + ), + extra_model_outputs={ + # TODO (simon): Check, if we have to correct here for sequences + # later. + "weights": [weight / max_weight * 1], # actual_size=1 + "n_step": [actual_n_step], + **( + { + k: [episode.get_extra_model_outputs(k, episode_ts)] + for k in episode.extra_model_outputs.keys() + } + if include_extra_model_outputs + else {} + ), + }, + # TODO (sven): Support lookback buffers. + len_lookback_buffer=0, + t_started=episode_ts, + ) + # Record here the episode time step via a hash code. + sampled_env_step_idxs.add( + hashlib.sha256(f"{episode.id_}-{episode_ts}".encode()).hexdigest() + ) + # Convert to numpy arrays, if required. + if to_numpy: + sampled_episode.to_numpy() + sampled_episodes.append(sampled_episode) + + # Add the episode buffer index to the sampled indices. + sampled_episode_idxs.add(episode_idx) + # Record the actual n-step for this sample. + sampled_n_steps.append(actual_n_step) + + # Increment counter. + B += 1 + + # Keep track of sampled indices for updating priorities later. + self._last_sampled_indices.append(idx) + + # Add to the sampled timesteps counter of the buffer. + self.sampled_timesteps += batch_size_B + + # Update the sample metrics. + num_env_steps_sampled = batch_size_B + num_episodes_per_sample = len(sampled_episode_idxs) + num_env_steps_per_sample = len(sampled_env_step_idxs) + sampled_n_step = sum(sampled_n_steps) / batch_size_B + agent_to_num_steps_sampled = {DEFAULT_AGENT_ID: num_env_steps_sampled} + agent_to_num_episodes_per_sample = {DEFAULT_AGENT_ID: num_episodes_per_sample} + agent_to_num_steps_per_sample = {DEFAULT_AGENT_ID: num_env_steps_per_sample} + agent_to_sampled_n_step = {DEFAULT_AGENT_ID: sampled_n_step} + agent_to_num_resamples = {DEFAULT_AGENT_ID: num_resamples} + module_to_num_steps_sampled = {DEFAULT_MODULE_ID: num_env_steps_sampled} + module_to_num_episodes_per_sample = {DEFAULT_MODULE_ID: num_episodes_per_sample} + module_to_num_steps_per_sample = {DEFAULT_MODULE_ID: num_env_steps_per_sample} + module_to_sampled_n_step = {DEFAULT_MODULE_ID: sampled_n_step} + module_to_num_resamples = {DEFAULT_MODULE_ID: num_resamples} + self._update_sample_metrics( + num_env_steps_sampled=num_env_steps_sampled, + num_episodes_per_sample=num_episodes_per_sample, + num_env_steps_per_sample=num_env_steps_per_sample, + sampled_n_step=sampled_n_step, + num_resamples=num_resamples, + agent_to_num_steps_sampled=agent_to_num_steps_sampled, + agent_to_num_episodes_per_sample=agent_to_num_episodes_per_sample, + agent_to_num_steps_per_sample=agent_to_num_steps_per_sample, + agent_to_sampled_n_step=agent_to_sampled_n_step, + agent_to_num_resamples=agent_to_num_resamples, + module_to_num_steps_sampled=module_to_num_steps_sampled, + module_to_num_episodes_per_sample=module_to_num_episodes_per_sample, + module_to_num_steps_per_sample=module_to_num_steps_per_sample, + module_to_sampled_n_step=module_to_sampled_n_step, + module_to_num_resamples=module_to_num_resamples, + ) + + return sampled_episodes + + @override(EpisodeReplayBuffer) + def get_state(self) -> Dict[str, Any]: + """Gets the state of a `PrioritizedEpisodeReplayBuffer`. + + Returns: + A state dict that can be stored in a checkpoint. + """ + # Get super's state. + state = super().get_state() + # Add additional attributes. + state.update( + { + "_sum_segment": self._sum_segment.get_state(), + "_min_segment": self._min_segment.get_state(), + "_free_nodes": list(self._free_nodes), + "_max_priority": self._max_priority, + "_max_idx": self._max_idx, + "_tree_idx_to_sample_idx": list(self._tree_idx_to_sample_idx.items()), + # TODO (sven, simon): Do we need these? + "_last_sampled_indices": self._last_sampled_indices, + } + ) + return state + + @override(EpisodeReplayBuffer) + def set_state(self, state) -> None: + """Sets the state of a `PrioritizedEpisodeReplayBuffer`. + + Args: + state: A buffer state stored (usually stored in a checkpoint). + """ + # Set super's state. + super().set_state(state) + # Set additional attributes. + self._sum_segment.set_state(state["_sum_segment"]) + self._min_segment.set_state(state["_min_segment"]) + self._free_nodes = deque(state["_free_nodes"]) + self._max_priority = state["_max_priority"] + self._max_idx = state["_max_idx"] + self._tree_idx_to_sample_idx = dict(state["_tree_idx_to_sample_idx"]) + # TODO (sven, simon): Do we need these? + self._last_sampled_indices = state["_last_sampled_indices"] + + def update_priorities( + self, priorities: NDArray, module_id: Optional[ModuleID] = None + ) -> None: + """Update the priorities of items at corresponding indices. + + Usually, incoming priorities are TD-errors. + + Args: + priorities: Numpy array containing the new priorities to be used + in sampling for the items in the last sampled batch. + """ + assert len(priorities) == len(self._last_sampled_indices) + + for idx, priority in zip(self._last_sampled_indices, priorities): + # Note, TD-errors come in as absolute values or results from + # cross-entropy loss calculations. + # assert priority > 0, f"priority was {priority}" + priority = max(priority, 1e-12) + assert 0 <= idx < self._sum_segment.capacity + # TODO (simon): Create metrics. + # delta = priority**self._alpha - self._sum_segment[idx] + # Update the priorities in the segment trees. + self._sum_segment[idx] = priority**self._alpha + self._min_segment[idx] = priority**self._alpha + # Update the maximal priority. + self._max_priority = max(self._max_priority, priority) + self._last_sampled_indices.clear() + + def _get_free_node_and_assign(self, sample_index, weight: float = 1.0) -> int: + """Gets the next free node in the segment trees. + + In addition the initial priorities for a new transition are added + to the segment trees and the index of the nodes is added to the + index mapping. + + Args: + sample_index: The index of the sample in the `self._indices` list. + weight: The initial priority weight to be used in sampling for + the item at index `sample_index`. + + Returns: + The index in the segment trees `self._sum_segment` and + `self._min_segment` for the item at index `sample_index` in + ``self._indices`. + """ + # Get an index from the free nodes in the segment trees. + idx = self._free_nodes.popleft() + self._max_idx = idx if idx > self._max_idx else self._max_idx + # Add the weight to the segments. + self._sum_segment[idx] = weight**self._alpha + self._min_segment[idx] = weight**self._alpha + # Map the index in the trees to the index in `self._indices`. + self._tree_idx_to_sample_idx[idx] = sample_index + # Return the index. + return idx + + def _num_remaining_episodes(self, new_eps, evicted_eps): + """Calculates the number of remaining episodes. + + When adding episodes and evicting them in the `add()` method + this function calculates iteratively the number of remaining + episodes. + + Args: + new_eps: List of new episode IDs. + evicted_eps: List of evicted episode IDs. + + Returns: + Number of episodes remaining after evicting the episodes in + `evicted_eps` and adding the episode in `new_eps`. + """ + return len( + set(self.episode_id_to_index.keys()).union(set(new_eps)) - set(evicted_eps) + ) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/prioritized_replay_buffer.py b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/prioritized_replay_buffer.py new file mode 100644 index 0000000000000000000000000000000000000000..77e13796777a634875f4d42633fb561110371ff7 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/prioritized_replay_buffer.py @@ -0,0 +1,241 @@ +import random +from typing import Any, Dict, List, Optional + +import numpy as np + +# Import ray before psutil will make sure we use psutil's bundled version +import ray # noqa F401 +from ray.rllib.execution.segment_tree import MinSegmentTree, SumSegmentTree +from ray.rllib.policy.sample_batch import SampleBatch +from ray.rllib.utils.annotations import override +from ray.rllib.utils.metrics.window_stat import WindowStat +from ray.rllib.utils.replay_buffers.replay_buffer import ReplayBuffer +from ray.rllib.utils.typing import SampleBatchType +from ray.util.annotations import DeveloperAPI + +import psutil # noqa E402 + + +@DeveloperAPI +class PrioritizedReplayBuffer(ReplayBuffer): + """This buffer implements Prioritized Experience Replay. + + The algorithm has been described by Tom Schaul et. al. in "Prioritized + Experience Replay". See https://arxiv.org/pdf/1511.05952.pdf for + the full paper. + """ + + def __init__( + self, + capacity: int = 10000, + storage_unit: str = "timesteps", + alpha: float = 1.0, + **kwargs, + ): + """Initializes a PrioritizedReplayBuffer instance. + + Args: + capacity: Max number of timesteps to store in the FIFO + buffer. After reaching this number, older samples will be + dropped to make space for new ones. + storage_unit: Either 'timesteps', 'sequences' or + 'episodes'. Specifies how experiences are stored. + alpha: How much prioritization is used + (0.0=no prioritization, 1.0=full prioritization). + ``**kwargs``: Forward compatibility kwargs. + """ + ReplayBuffer.__init__(self, capacity, storage_unit, **kwargs) + + assert alpha > 0 + self._alpha = alpha + + # Segment tree must have capacity that is a power of 2 + it_capacity = 1 + while it_capacity < self.capacity: + it_capacity *= 2 + + self._it_sum = SumSegmentTree(it_capacity) + self._it_min = MinSegmentTree(it_capacity) + self._max_priority = 1.0 + self._prio_change_stats = WindowStat("reprio", 1000) + + @DeveloperAPI + @override(ReplayBuffer) + def _add_single_batch(self, item: SampleBatchType, **kwargs) -> None: + """Add a batch of experiences to self._storage with weight. + + An item consists of either one or more timesteps, a sequence or an + episode. Differs from add() in that it does not consider the storage + unit or type of batch and simply stores it. + + Args: + item: The item to be added. + ``**kwargs``: Forward compatibility kwargs. + """ + weight = kwargs.get("weight", None) + + if weight is None: + weight = self._max_priority + + self._it_sum[self._next_idx] = weight**self._alpha + self._it_min[self._next_idx] = weight**self._alpha + + ReplayBuffer._add_single_batch(self, item) + + def _sample_proportional(self, num_items: int) -> List[int]: + res = [] + for _ in range(num_items): + # TODO(szymon): should we ensure no repeats? + mass = random.random() * self._it_sum.sum(0, len(self._storage)) + idx = self._it_sum.find_prefixsum_idx(mass) + res.append(idx) + return res + + @DeveloperAPI + @override(ReplayBuffer) + def sample( + self, num_items: int, beta: float, **kwargs + ) -> Optional[SampleBatchType]: + """Sample `num_items` items from this buffer, including prio. weights. + + Samples in the results may be repeated. + + Examples for storage of SamplesBatches: + - If storage unit `timesteps` has been chosen and batches of + size 5 have been added, sample(5) will yield a concatenated batch of + 15 timesteps. + - If storage unit 'sequences' has been chosen and sequences of + different lengths have been added, sample(5) will yield a concatenated + batch with a number of timesteps equal to the sum of timesteps in + the 5 sampled sequences. + - If storage unit 'episodes' has been chosen and episodes of + different lengths have been added, sample(5) will yield a concatenated + batch with a number of timesteps equal to the sum of timesteps in + the 5 sampled episodes. + + Args: + num_items: Number of items to sample from this buffer. + beta: To what degree to use importance weights (0 - no corrections, + 1 - full correction). + ``**kwargs``: Forward compatibility kwargs. + + Returns: + Concatenated SampleBatch of items including "weights" and + "batch_indexes" fields denoting IS of each sampled + transition and original idxes in buffer of sampled experiences. + """ + assert beta >= 0.0 + + if len(self) == 0: + raise ValueError("Trying to sample from an empty buffer.") + + idxes = self._sample_proportional(num_items) + + weights = [] + batch_indexes = [] + p_min = self._it_min.min() / self._it_sum.sum() + max_weight = (p_min * len(self)) ** (-beta) + + for idx in idxes: + p_sample = self._it_sum[idx] / self._it_sum.sum() + weight = (p_sample * len(self)) ** (-beta) + count = self._storage[idx].count + # If zero-padded, count will not be the actual batch size of the + # data. + if ( + isinstance(self._storage[idx], SampleBatch) + and self._storage[idx].zero_padded + ): + actual_size = self._storage[idx].max_seq_len + else: + actual_size = count + weights.extend([weight / max_weight] * actual_size) + batch_indexes.extend([idx] * actual_size) + self._num_timesteps_sampled += count + batch = self._encode_sample(idxes) + + # Note: prioritization is not supported in multi agent lockstep + if isinstance(batch, SampleBatch): + batch["weights"] = np.array(weights) + batch["batch_indexes"] = np.array(batch_indexes) + + return batch + + @DeveloperAPI + def update_priorities(self, idxes: List[int], priorities: List[float]) -> None: + """Update priorities of items at given indices. + + Sets priority of item at index idxes[i] in buffer + to priorities[i]. + + Args: + idxes: List of indices of items + priorities: List of updated priorities corresponding to items at the + idxes denoted by variable `idxes`. + """ + # Making sure we don't pass in e.g. a torch tensor. + assert isinstance( + idxes, (list, np.ndarray) + ), "ERROR: `idxes` is not a list or np.ndarray, but {}!".format( + type(idxes).__name__ + ) + assert len(idxes) == len(priorities) + for idx, priority in zip(idxes, priorities): + assert priority > 0 + assert 0 <= idx < len(self._storage) + delta = priority**self._alpha - self._it_sum[idx] + self._prio_change_stats.push(delta) + self._it_sum[idx] = priority**self._alpha + self._it_min[idx] = priority**self._alpha + + self._max_priority = max(self._max_priority, priority) + + @DeveloperAPI + @override(ReplayBuffer) + def stats(self, debug: bool = False) -> Dict: + """Returns the stats of this buffer. + + Args: + debug: If true, adds sample eviction statistics to the returned stats dict. + + Returns: + A dictionary of stats about this buffer. + """ + parent = ReplayBuffer.stats(self, debug) + if debug: + parent.update(self._prio_change_stats.stats()) + return parent + + @DeveloperAPI + @override(ReplayBuffer) + def get_state(self) -> Dict[str, Any]: + """Returns all local state. + + Returns: + The serializable local state. + """ + # Get parent state. + state = super().get_state() + # Add prio weights. + state.update( + { + "sum_segment_tree": self._it_sum.get_state(), + "min_segment_tree": self._it_min.get_state(), + "max_priority": self._max_priority, + } + ) + return state + + @DeveloperAPI + @override(ReplayBuffer) + def set_state(self, state: Dict[str, Any]) -> None: + """Restores all local state to the provided `state`. + + Args: + state: The new state to set this buffer. Can be obtained by calling + `self.get_state()`. + """ + super().set_state(state) + self._it_sum.set_state(state["sum_segment_tree"]) + self._it_min.set_state(state["min_segment_tree"]) + self._max_priority = state["max_priority"] diff --git a/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/replay_buffer.py b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/replay_buffer.py new file mode 100644 index 0000000000000000000000000000000000000000..7abb580ed70bca15926e5f71e6daa0ddd0b27927 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/replay_buffer.py @@ -0,0 +1,375 @@ +import logging +import random +from enum import Enum +from typing import Any, Dict, List, Optional, Union + +import numpy as np + +# Import ray before psutil will make sure we use psutil's bundled version +import ray # noqa F401 +from ray.rllib.policy.sample_batch import SampleBatch, concat_samples +from ray.rllib.utils.actor_manager import FaultAwareApply +from ray.rllib.utils.annotations import override +from ray.rllib.utils.metrics.window_stat import WindowStat +from ray.rllib.utils.replay_buffers.base import ReplayBufferInterface +from ray.rllib.utils.typing import SampleBatchType +from ray.util.annotations import DeveloperAPI +from ray.util.debug import log_once + +import psutil + +# Constant that represents all policies in lockstep replay mode. +_ALL_POLICIES = "__all__" + +logger = logging.getLogger(__name__) + + +@DeveloperAPI +class StorageUnit(Enum): + """Specifies how batches are structured in a ReplayBuffer. + + timesteps: One buffer slot per timestep. + sequences: One buffer slot per sequence. + episodes: One buffer slot per episode. + fragemts: One buffer slot per incoming batch. + """ + + TIMESTEPS = "timesteps" + SEQUENCES = "sequences" + EPISODES = "episodes" + FRAGMENTS = "fragments" + + +@DeveloperAPI +def warn_replay_capacity(*, item: SampleBatchType, num_items: int) -> None: + """Warn if the configured replay buffer capacity is too large.""" + if log_once("replay_capacity"): + item_size = item.size_bytes() + psutil_mem = psutil.virtual_memory() + total_gb = psutil_mem.total / 1e9 + mem_size = num_items * item_size / 1e9 + msg = ( + "Estimated max memory usage for replay buffer is {} GB " + "({} batches of size {}, {} bytes each), " + "available system memory is {} GB".format( + mem_size, num_items, item.count, item_size, total_gb + ) + ) + if mem_size > total_gb: + raise ValueError(msg) + elif mem_size > 0.2 * total_gb: + logger.warning(msg) + else: + logger.info(msg) + + +@DeveloperAPI +class ReplayBuffer(ReplayBufferInterface, FaultAwareApply): + """The lowest-level replay buffer interface used by RLlib. + + This class implements a basic ring-type of buffer with random sampling. + ReplayBuffer is the base class for advanced types that add functionality while + retaining compatibility through inheritance. + + The following examples show how buffers behave with different storage_units + and capacities. This behaviour is generally similar for other buffers, although + they might not implement all storage_units. + + Examples: + + .. testcode:: + + from ray.rllib.utils.replay_buffers.replay_buffer import ReplayBuffer + from ray.rllib.utils.replay_buffers.replay_buffer import StorageUnit + from ray.rllib.policy.sample_batch import SampleBatch + + # Store any batch as a whole + buffer = ReplayBuffer(capacity=10, storage_unit=StorageUnit.FRAGMENTS) + buffer.add(SampleBatch({"a": [1], "b": [2, 3, 4]})) + buffer.sample(1) + + # Store only complete episodes + buffer = ReplayBuffer(capacity=10, + storage_unit=StorageUnit.EPISODES) + buffer.add(SampleBatch({"c": [1, 2, 3, 4], + SampleBatch.T: [0, 1, 0, 1], + SampleBatch.TERMINATEDS: [False, True, False, True], + SampleBatch.EPS_ID: [0, 0, 1, 1]})) + buffer.sample(1) + + # Store single timesteps + buffer = ReplayBuffer(capacity=2, storage_unit=StorageUnit.TIMESTEPS) + buffer.add(SampleBatch({"a": [1, 2], SampleBatch.T: [0, 1]})) + buffer.sample(1) + + buffer.add(SampleBatch({"a": [3], SampleBatch.T: [2]})) + print(buffer._eviction_started) + buffer.sample(1) + + buffer = ReplayBuffer(capacity=10, storage_unit=StorageUnit.SEQUENCES) + buffer.add(SampleBatch({"c": [1, 2, 3], SampleBatch.SEQ_LENS: [1, 2]})) + buffer.sample(1) + + .. testoutput:: + + True + + `True` is not the output of the above testcode, but an artifact of unexpected + behaviour of sphinx doctests. + (see https://github.com/ray-project/ray/pull/32477#discussion_r1106776101) + """ + + def __init__( + self, + capacity: int = 10000, + storage_unit: Union[str, StorageUnit] = "timesteps", + **kwargs, + ): + """Initializes a (FIFO) ReplayBuffer instance. + + Args: + capacity: Max number of timesteps to store in this FIFO + buffer. After reaching this number, older samples will be + dropped to make space for new ones. + storage_unit: If not a StorageUnit, either 'timesteps', 'sequences' or + 'episodes'. Specifies how experiences are stored. + ``**kwargs``: Forward compatibility kwargs. + """ + + if storage_unit in ["timesteps", StorageUnit.TIMESTEPS]: + self.storage_unit = StorageUnit.TIMESTEPS + elif storage_unit in ["sequences", StorageUnit.SEQUENCES]: + self.storage_unit = StorageUnit.SEQUENCES + elif storage_unit in ["episodes", StorageUnit.EPISODES]: + self.storage_unit = StorageUnit.EPISODES + elif storage_unit in ["fragments", StorageUnit.FRAGMENTS]: + self.storage_unit = StorageUnit.FRAGMENTS + else: + raise ValueError( + f"storage_unit must be either '{StorageUnit.TIMESTEPS}', " + f"'{StorageUnit.SEQUENCES}', '{StorageUnit.EPISODES}' " + f"or '{StorageUnit.FRAGMENTS}', but is {storage_unit}" + ) + + # The actual storage (list of SampleBatches or MultiAgentBatches). + self._storage = [] + + # Caps the number of timesteps stored in this buffer + if capacity <= 0: + raise ValueError( + "Capacity of replay buffer has to be greater than zero " + "but was set to {}.".format(capacity) + ) + self.capacity = capacity + # The next index to override in the buffer. + self._next_idx = 0 + # len(self._hit_count) must always be less than len(capacity) + self._hit_count = np.zeros(self.capacity) + + # Whether we have already hit our capacity (and have therefore + # started to evict older samples). + self._eviction_started = False + + # Number of (single) timesteps that have been added to the buffer + # over its lifetime. Note that each added item (batch) may contain + # more than one timestep. + self._num_timesteps_added = 0 + self._num_timesteps_added_wrap = 0 + + # Number of (single) timesteps that have been sampled from the buffer + # over its lifetime. + self._num_timesteps_sampled = 0 + + self._evicted_hit_stats = WindowStat("evicted_hit", 1000) + self._est_size_bytes = 0 + + self.batch_size = None + + @override(ReplayBufferInterface) + def __len__(self) -> int: + return len(self._storage) + + @override(ReplayBufferInterface) + def add(self, batch: SampleBatchType, **kwargs) -> None: + """Adds a batch of experiences or other data to this buffer. + + Splits batch into chunks of timesteps, sequences or episodes, depending on + `self._storage_unit`. Calls `self._add_single_batch` to add resulting slices + to the buffer storage. + + Args: + batch: The batch to add. + ``**kwargs``: Forward compatibility kwargs. + """ + if not batch.count > 0: + return + + warn_replay_capacity(item=batch, num_items=self.capacity / batch.count) + + if self.storage_unit == StorageUnit.TIMESTEPS: + timeslices = batch.timeslices(1) + for t in timeslices: + self._add_single_batch(t, **kwargs) + + elif self.storage_unit == StorageUnit.SEQUENCES: + timestep_count = 0 + for seq_len in batch.get(SampleBatch.SEQ_LENS): + start_seq = timestep_count + end_seq = timestep_count + seq_len + self._add_single_batch(batch[start_seq:end_seq], **kwargs) + timestep_count = end_seq + + elif self.storage_unit == StorageUnit.EPISODES: + for eps in batch.split_by_episode(): + if eps.get(SampleBatch.T, [0])[0] == 0 and ( + eps.get(SampleBatch.TERMINATEDS, [True])[-1] + or eps.get(SampleBatch.TRUNCATEDS, [False])[-1] + ): + # Only add full episodes to the buffer + # Check only if info is available + self._add_single_batch(eps, **kwargs) + else: + if log_once("only_full_episodes"): + logger.info( + "This buffer uses episodes as a storage " + "unit and thus allows only full episodes " + "to be added to it (starting from T=0 and ending in " + "`terminateds=True` or `truncateds=True`. " + "Some samples may be dropped." + ) + + elif self.storage_unit == StorageUnit.FRAGMENTS: + self._add_single_batch(batch, **kwargs) + + @DeveloperAPI + def _add_single_batch(self, item: SampleBatchType, **kwargs) -> None: + """Add a SampleBatch of experiences to self._storage. + + An item consists of either one or more timesteps, a sequence or an + episode. Differs from add() in that it does not consider the storage + unit or type of batch and simply stores it. + + Args: + item: The batch to be added. + ``**kwargs``: Forward compatibility kwargs. + """ + self._num_timesteps_added += item.count + self._num_timesteps_added_wrap += item.count + + if self._next_idx >= len(self._storage): + self._storage.append(item) + self._est_size_bytes += item.size_bytes() + else: + item_to_be_removed = self._storage[self._next_idx] + self._est_size_bytes -= item_to_be_removed.size_bytes() + self._storage[self._next_idx] = item + self._est_size_bytes += item.size_bytes() + + # Eviction of older samples has already started (buffer is "full"). + if self._eviction_started: + self._evicted_hit_stats.push(self._hit_count[self._next_idx]) + self._hit_count[self._next_idx] = 0 + + # Wrap around storage as a circular buffer once we hit capacity. + if self._num_timesteps_added_wrap >= self.capacity: + self._eviction_started = True + self._num_timesteps_added_wrap = 0 + self._next_idx = 0 + else: + self._next_idx += 1 + + @override(ReplayBufferInterface) + def sample( + self, num_items: Optional[int] = None, **kwargs + ) -> Optional[SampleBatchType]: + """Samples `num_items` items from this buffer. + + The items depend on the buffer's storage_unit. + Samples in the results may be repeated. + + Examples for sampling results: + + 1) If storage unit 'timesteps' has been chosen and batches of + size 5 have been added, sample(5) will yield a concatenated batch of + 15 timesteps. + + 2) If storage unit 'sequences' has been chosen and sequences of + different lengths have been added, sample(5) will yield a concatenated + batch with a number of timesteps equal to the sum of timesteps in + the 5 sampled sequences. + + 3) If storage unit 'episodes' has been chosen and episodes of + different lengths have been added, sample(5) will yield a concatenated + batch with a number of timesteps equal to the sum of timesteps in + the 5 sampled episodes. + + Args: + num_items: Number of items to sample from this buffer. + ``**kwargs``: Forward compatibility kwargs. + + Returns: + Concatenated batch of items. + """ + if len(self) == 0: + raise ValueError("Trying to sample from an empty buffer.") + idxes = [random.randint(0, len(self) - 1) for _ in range(num_items)] + sample = self._encode_sample(idxes) + self._num_timesteps_sampled += sample.count + return sample + + @DeveloperAPI + def stats(self, debug: bool = False) -> dict: + """Returns the stats of this buffer. + + Args: + debug: If True, adds sample eviction statistics to the returned + stats dict. + + Returns: + A dictionary of stats about this buffer. + """ + data = { + "added_count": self._num_timesteps_added, + "added_count_wrapped": self._num_timesteps_added_wrap, + "eviction_started": self._eviction_started, + "sampled_count": self._num_timesteps_sampled, + "est_size_bytes": self._est_size_bytes, + "num_entries": len(self._storage), + } + if debug: + data.update(self._evicted_hit_stats.stats()) + return data + + @override(ReplayBufferInterface) + def get_state(self) -> Dict[str, Any]: + state = {"_storage": self._storage, "_next_idx": self._next_idx} + state.update(self.stats(debug=False)) + return state + + @override(ReplayBufferInterface) + def set_state(self, state: Dict[str, Any]) -> None: + # The actual storage. + self._storage = state["_storage"] + self._next_idx = state["_next_idx"] + # Stats and counts. + self._num_timesteps_added = state["added_count"] + self._num_timesteps_added_wrap = state["added_count_wrapped"] + self._eviction_started = state["eviction_started"] + self._num_timesteps_sampled = state["sampled_count"] + self._est_size_bytes = state["est_size_bytes"] + + @DeveloperAPI + def _encode_sample(self, idxes: List[int]) -> SampleBatchType: + """Fetches concatenated samples at given indices from the storage.""" + samples = [] + for i in idxes: + self._hit_count[i] += 1 + samples.append(self._storage[i]) + + if samples: + # We assume all samples are of same type + out = concat_samples(samples) + else: + out = SampleBatch() + out.decompress_if_needed() + return out diff --git a/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/reservoir_replay_buffer.py b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/reservoir_replay_buffer.py new file mode 100644 index 0000000000000000000000000000000000000000..cda00a1b2edd7a2af5a59557f9312d752eccbdd1 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/reservoir_replay_buffer.py @@ -0,0 +1,132 @@ +import random +from typing import Any, Dict + +# Import ray before psutil will make sure we use psutil's bundled version +import ray # noqa F401 +from ray.rllib.utils.annotations import ExperimentalAPI, override +from ray.rllib.utils.replay_buffers.replay_buffer import ( + ReplayBuffer, + warn_replay_capacity, +) +from ray.rllib.utils.typing import SampleBatchType + +import psutil # noqa E402 + + +# __sphinx_doc_reservoir_buffer__begin__ +@ExperimentalAPI +class ReservoirReplayBuffer(ReplayBuffer): + """This buffer implements reservoir sampling. + + The algorithm has been described by Jeffrey S. Vitter in "Random sampling + with a reservoir". + """ + + def __init__( + self, capacity: int = 10000, storage_unit: str = "timesteps", **kwargs + ): + """Initializes a ReservoirBuffer instance. + + Args: + capacity: Max number of timesteps to store in the FIFO + buffer. After reaching this number, older samples will be + dropped to make space for new ones. + storage_unit: Either 'timesteps', 'sequences' or + 'episodes'. Specifies how experiences are stored. + """ + ReplayBuffer.__init__(self, capacity, storage_unit) + self._num_add_calls = 0 + self._num_evicted = 0 + + @ExperimentalAPI + @override(ReplayBuffer) + def _add_single_batch(self, item: SampleBatchType, **kwargs) -> None: + """Add a SampleBatch of experiences to self._storage. + + An item consists of either one or more timesteps, a sequence or an + episode. Differs from add() in that it does not consider the storage + unit or type of batch and simply stores it. + + Args: + item: The batch to be added. + ``**kwargs``: Forward compatibility kwargs. + """ + self._num_timesteps_added += item.count + self._num_timesteps_added_wrap += item.count + + # Update add counts. + self._num_add_calls += 1 + # Update our timesteps counts. + + if self._num_timesteps_added < self.capacity: + self._storage.append(item) + self._est_size_bytes += item.size_bytes() + else: + # Eviction of older samples has already started (buffer is "full") + self._eviction_started = True + idx = random.randint(0, self._num_add_calls - 1) + if idx < len(self._storage): + self._num_evicted += 1 + self._evicted_hit_stats.push(self._hit_count[idx]) + self._hit_count[idx] = 0 + # This is a bit of a hack: ReplayBuffer always inserts at + # self._next_idx + self._next_idx = idx + self._evicted_hit_stats.push(self._hit_count[idx]) + self._hit_count[idx] = 0 + + item_to_be_removed = self._storage[idx] + self._est_size_bytes -= item_to_be_removed.size_bytes() + self._storage[idx] = item + self._est_size_bytes += item.size_bytes() + + assert item.count > 0, item + warn_replay_capacity(item=item, num_items=self.capacity / item.count) + + @ExperimentalAPI + @override(ReplayBuffer) + def stats(self, debug: bool = False) -> dict: + """Returns the stats of this buffer. + + Args: + debug: If True, adds sample eviction statistics to the returned + stats dict. + + Returns: + A dictionary of stats about this buffer. + """ + data = { + "num_evicted": self._num_evicted, + "num_add_calls": self._num_add_calls, + } + parent = ReplayBuffer.stats(self, debug) + parent.update(data) + return parent + + @ExperimentalAPI + @override(ReplayBuffer) + def get_state(self) -> Dict[str, Any]: + """Returns all local state. + + Returns: + The serializable local state. + """ + parent = ReplayBuffer.get_state(self) + parent.update(self.stats()) + return parent + + @ExperimentalAPI + @override(ReplayBuffer) + def set_state(self, state: Dict[str, Any]) -> None: + """Restores all local state to the provided `state`. + + Args: + state: The new state to set this buffer. Can be + obtained by calling `self.get_state()`. + """ + self._num_evicted = state["num_evicted"] + self._num_add_calls = state["num_add_calls"] + ReplayBuffer.set_state(self, state) + + +# __sphinx_doc_reservoir_buffer__end__ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/simple_replay_buffer.py b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/simple_replay_buffer.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/utils.py b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..13b90440da0e0047aa6d413e0ff0af748b45747f --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/replay_buffers/utils.py @@ -0,0 +1,440 @@ +import logging +from typing import Any, Dict, Optional + +import numpy as np + +from ray._common.deprecation import DEPRECATED_VALUE, deprecation_warning +from ray.rllib.policy.sample_batch import MultiAgentBatch, SampleBatch, concat_samples +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.from_config import from_config +from ray.rllib.utils.metrics import ALL_MODULES, TD_ERROR_KEY +from ray.rllib.utils.metrics.learner_info import LEARNER_STATS_KEY +from ray.rllib.utils.replay_buffers import ( + EpisodeReplayBuffer, + MultiAgentPrioritizedReplayBuffer, + MultiAgentReplayBuffer, + PrioritizedEpisodeReplayBuffer, + ReplayBuffer, +) +from ray.rllib.utils.typing import ( + AlgorithmConfigDict, + ModuleID, + ResultDict, + SampleBatchType, + TensorType, +) +from ray.util import log_once +from ray.util.annotations import DeveloperAPI + +import psutil + +logger = logging.getLogger(__name__) + + +@DeveloperAPI +def update_priorities_in_episode_replay_buffer( + *, + replay_buffer: EpisodeReplayBuffer, + td_errors: Dict[ModuleID, TensorType], +) -> None: + # Only update priorities, if the buffer supports them. + if isinstance(replay_buffer, PrioritizedEpisodeReplayBuffer): + + # The `ResultDict` will be multi-agent. + for module_id, td_error in td_errors.items(): + # Skip the `"__all__"` keys. + if module_id in ["__all__", ALL_MODULES]: + continue + + # Warn once, if we have no TD-errors to update priorities. + if TD_ERROR_KEY not in td_error or td_error[TD_ERROR_KEY] is None: + if log_once( + "no_td_error_in_train_results_from_module_{}".format(module_id) + ): + logger.warning( + "Trying to update priorities for module with ID " + f"`{module_id}` in prioritized episode replay buffer without " + "providing `td_errors` in train_results. Priority update for " + "this policy is being skipped." + ) + continue + # TODO (simon): Implement multi-agent version. Remove, happens in buffer. + # assert len(td_error[TD_ERROR_KEY]) == len( + # replay_buffer._last_sampled_indices + # ) + # TODO (simon): Implement for stateful modules. + + replay_buffer.update_priorities(td_error[TD_ERROR_KEY], module_id) + + +@OldAPIStack +def update_priorities_in_replay_buffer( + replay_buffer: ReplayBuffer, + config: AlgorithmConfigDict, + train_batch: SampleBatchType, + train_results: ResultDict, +) -> None: + """Updates the priorities in a prioritized replay buffer, given training results. + + The `abs(TD-error)` from the loss (inside `train_results`) is used as new + priorities for the row-indices that were sampled for the train batch. + + Don't do anything if the given buffer does not support prioritized replay. + + Args: + replay_buffer: The replay buffer, whose priority values to update. This may also + be a buffer that does not support priorities. + config: The Algorithm's config dict. + train_batch: The batch used for the training update. + train_results: A train results dict, generated by e.g. the `train_one_step()` + utility. + """ + # Only update priorities if buffer supports them. + if isinstance(replay_buffer, MultiAgentPrioritizedReplayBuffer): + # Go through training results for the different policies (maybe multi-agent). + prio_dict = {} + for policy_id, info in train_results.items(): + # TODO(sven): This is currently structured differently for + # torch/tf. Clean up these results/info dicts across + # policies (note: fixing this in torch_policy.py will + # break e.g. DDPPO!). + td_error = info.get("td_error", info[LEARNER_STATS_KEY].get("td_error")) + + policy_batch = train_batch.policy_batches[policy_id] + # Set the get_interceptor to None in order to be able to access the numpy + # arrays directly (instead of e.g. a torch array). + policy_batch.set_get_interceptor(None) + # Get the replay buffer row indices that make up the `train_batch`. + batch_indices = policy_batch.get("batch_indexes") + + if SampleBatch.SEQ_LENS in policy_batch: + # Batch_indices are represented per column, in order to update + # priorities, we need one index per td_error + _batch_indices = [] + + # Sequenced batches have been zero padded to max_seq_len. + # Depending on how batches are split during learning, not all + # sequences have an associated td_error (trailing ones missing). + if policy_batch.zero_padded: + seq_lens = len(td_error) * [policy_batch.max_seq_len] + else: + seq_lens = policy_batch[SampleBatch.SEQ_LENS][: len(td_error)] + + # Go through all indices by sequence that they represent and shrink + # them to one index per sequences + sequence_sum = 0 + for seq_len in seq_lens: + _batch_indices.append(batch_indices[sequence_sum]) + sequence_sum += seq_len + batch_indices = np.array(_batch_indices) + + if td_error is None: + if log_once( + "no_td_error_in_train_results_from_policy_{}".format(policy_id) + ): + logger.warning( + "Trying to update priorities for policy with id `{}` in " + "prioritized replay buffer without providing td_errors in " + "train_results. Priority update for this policy is being " + "skipped.".format(policy_id) + ) + continue + + if batch_indices is None: + if log_once( + "no_batch_indices_in_train_result_for_policy_{}".format(policy_id) + ): + logger.warning( + "Trying to update priorities for policy with id `{}` in " + "prioritized replay buffer without providing batch_indices in " + "train_batch. Priority update for this policy is being " + "skipped.".format(policy_id) + ) + continue + + # Try to transform batch_indices to td_error dimensions + if len(batch_indices) != len(td_error): + T = replay_buffer.replay_sequence_length + assert ( + len(batch_indices) > len(td_error) and len(batch_indices) % T == 0 + ) + batch_indices = batch_indices.reshape([-1, T])[:, 0] + assert len(batch_indices) == len(td_error) + prio_dict[policy_id] = (batch_indices, td_error) + + # Make the actual buffer API call to update the priority weights on all + # policies. + replay_buffer.update_priorities(prio_dict) + + +@DeveloperAPI +def sample_min_n_steps_from_buffer( + replay_buffer: ReplayBuffer, min_steps: int, count_by_agent_steps: bool +) -> Optional[SampleBatchType]: + """Samples a minimum of n timesteps from a given replay buffer. + + This utility method is primarily used by the QMIX algorithm and helps with + sampling a given number of time steps which has stored samples in units + of sequences or complete episodes. Samples n batches from replay buffer + until the total number of timesteps reaches `train_batch_size`. + + Args: + replay_buffer: The replay buffer to sample from + num_timesteps: The number of timesteps to sample + count_by_agent_steps: Whether to count agent steps or env steps + + Returns: + A concatenated SampleBatch or MultiAgentBatch with samples from the + buffer. + """ + train_batch_size = 0 + train_batches = [] + while train_batch_size < min_steps: + batch = replay_buffer.sample(num_items=1) + batch_len = batch.agent_steps() if count_by_agent_steps else batch.env_steps() + if batch_len == 0: + # Replay has not started, so we can't accumulate timesteps here + return batch + train_batches.append(batch) + train_batch_size += batch_len + # All batch types are the same type, hence we can use any concat_samples() + train_batch = concat_samples(train_batches) + return train_batch + + +@DeveloperAPI +def validate_buffer_config(config: dict) -> None: + """Checks and fixes values in the replay buffer config. + + Checks the replay buffer config for common misconfigurations, warns or raises + error in case validation fails. The type "key" is changed into the inferred + replay buffer class. + + Args: + config: The replay buffer config to be validated. + + Raises: + ValueError: When detecting severe misconfiguration. + """ + if config.get("replay_buffer_config", None) is None: + config["replay_buffer_config"] = {} + + if config.get("worker_side_prioritization", DEPRECATED_VALUE) != DEPRECATED_VALUE: + deprecation_warning( + old="config['worker_side_prioritization']", + new="config['replay_buffer_config']['worker_side_prioritization']", + error=True, + ) + + prioritized_replay = config.get("prioritized_replay", DEPRECATED_VALUE) + if prioritized_replay != DEPRECATED_VALUE: + deprecation_warning( + old="config['prioritized_replay'] or config['replay_buffer_config'][" + "'prioritized_replay']", + help="Replay prioritization specified by config key. RLlib's new replay " + "buffer API requires setting `config[" + "'replay_buffer_config']['type']`, e.g. `config[" + "'replay_buffer_config']['type'] = " + "'MultiAgentPrioritizedReplayBuffer'` to change the default " + "behaviour.", + error=True, + ) + + capacity = config.get("buffer_size", DEPRECATED_VALUE) + if capacity == DEPRECATED_VALUE: + capacity = config["replay_buffer_config"].get("buffer_size", DEPRECATED_VALUE) + if capacity != DEPRECATED_VALUE: + deprecation_warning( + old="config['buffer_size'] or config['replay_buffer_config'][" + "'buffer_size']", + new="config['replay_buffer_config']['capacity']", + error=True, + ) + + replay_burn_in = config.get("burn_in", DEPRECATED_VALUE) + if replay_burn_in != DEPRECATED_VALUE: + config["replay_buffer_config"]["replay_burn_in"] = replay_burn_in + deprecation_warning( + old="config['burn_in']", + help="config['replay_buffer_config']['replay_burn_in']", + ) + + replay_batch_size = config.get("replay_batch_size", DEPRECATED_VALUE) + if replay_batch_size == DEPRECATED_VALUE: + replay_batch_size = config["replay_buffer_config"].get( + "replay_batch_size", DEPRECATED_VALUE + ) + if replay_batch_size != DEPRECATED_VALUE: + deprecation_warning( + old="config['replay_batch_size'] or config['replay_buffer_config'][" + "'replay_batch_size']", + help="Specification of replay_batch_size is not supported anymore but is " + "derived from `train_batch_size`. Specify the number of " + "items you want to replay upon calling the sample() method of replay " + "buffers if this does not work for you.", + error=True, + ) + + # Deprecation of old-style replay buffer args + # Warnings before checking of we need local buffer so that algorithms + # Without local buffer also get warned + keys_with_deprecated_positions = [ + "prioritized_replay_alpha", + "prioritized_replay_beta", + "prioritized_replay_eps", + "no_local_replay_buffer", + "replay_zero_init_states", + "replay_buffer_shards_colocated_with_driver", + ] + for k in keys_with_deprecated_positions: + if config.get(k, DEPRECATED_VALUE) != DEPRECATED_VALUE: + deprecation_warning( + old="config['{}']".format(k), + help="config['replay_buffer_config']['{}']" "".format(k), + error=False, + ) + # Copy values over to new location in config to support new + # and old configuration style. + if config.get("replay_buffer_config") is not None: + config["replay_buffer_config"][k] = config[k] + + learning_starts = config.get( + "learning_starts", + config.get("replay_buffer_config", {}).get("learning_starts", DEPRECATED_VALUE), + ) + if learning_starts != DEPRECATED_VALUE: + deprecation_warning( + old="config['learning_starts'] or" + "config['replay_buffer_config']['learning_starts']", + help="config['num_steps_sampled_before_learning_starts']", + error=True, + ) + config["num_steps_sampled_before_learning_starts"] = learning_starts + + # Can't use DEPRECATED_VALUE here because this is also a deliberate + # value set for some algorithms + # TODO: (Artur): Compare to DEPRECATED_VALUE on deprecation + replay_sequence_length = config.get("replay_sequence_length", None) + if replay_sequence_length is not None: + config["replay_buffer_config"][ + "replay_sequence_length" + ] = replay_sequence_length + deprecation_warning( + old="config['replay_sequence_length']", + help="Replay sequence length specified at new " + "location config['replay_buffer_config'][" + "'replay_sequence_length'] will be overwritten.", + error=True, + ) + + replay_buffer_config = config["replay_buffer_config"] + assert ( + "type" in replay_buffer_config + ), "Can not instantiate ReplayBuffer from config without 'type' key." + + # Check if old replay buffer should be instantiated + buffer_type = config["replay_buffer_config"]["type"] + + if isinstance(buffer_type, str) and buffer_type.find(".") == -1: + # Create valid full [module].[class] string for from_config + config["replay_buffer_config"]["type"] = ( + "ray.rllib.utils.replay_buffers." + buffer_type + ) + + # Instantiate a dummy buffer to fail early on misconfiguration and find out about + # inferred buffer class + dummy_buffer = from_config(buffer_type, config["replay_buffer_config"]) + + config["replay_buffer_config"]["type"] = type(dummy_buffer) + + if hasattr(dummy_buffer, "update_priorities"): + if ( + config["replay_buffer_config"].get("replay_mode", "independent") + == "lockstep" + ): + raise ValueError( + "Prioritized replay is not supported when replay_mode=lockstep." + ) + elif config["replay_buffer_config"].get("replay_sequence_length", 0) > 1: + raise ValueError( + "Prioritized replay is not supported when " + "replay_sequence_length > 1." + ) + else: + if config["replay_buffer_config"].get("worker_side_prioritization"): + raise ValueError( + "Worker side prioritization is not supported when " + "prioritized_replay=False." + ) + + +@DeveloperAPI +def warn_replay_buffer_capacity(*, item: SampleBatchType, capacity: int) -> None: + """Warn if the configured replay buffer capacity is too large for machine's memory. + + Args: + item: A (example) item that's supposed to be added to the buffer. + This is used to compute the overall memory footprint estimate for the + buffer. + capacity: The capacity value of the buffer. This is interpreted as the + number of items (such as given `item`) that will eventually be stored in + the buffer. + + Raises: + ValueError: If computed memory footprint for the buffer exceeds the machine's + RAM. + """ + if log_once("warn_replay_buffer_capacity"): + item_size = item.size_bytes() + psutil_mem = psutil.virtual_memory() + total_gb = psutil_mem.total / 1e9 + mem_size = capacity * item_size / 1e9 + msg = ( + "Estimated max memory usage for replay buffer is {} GB " + "({} batches of size {}, {} bytes each), " + "available system memory is {} GB".format( + mem_size, capacity, item.count, item_size, total_gb + ) + ) + if mem_size > total_gb: + raise ValueError(msg) + elif mem_size > 0.2 * total_gb: + logger.warning(msg) + else: + logger.info(msg) + + +def patch_buffer_with_fake_sampling_method( + buffer: ReplayBuffer, fake_sample_output: SampleBatchType +) -> None: + """Patch a ReplayBuffer such that we always sample fake_sample_output. + + Transforms fake_sample_output into a MultiAgentBatch if it is not a + MultiAgentBatch and the buffer is a MultiAgentBuffer. This is useful for testing + purposes if we need deterministic sampling. + + Args: + buffer: The buffer to be patched + fake_sample_output: The output to be sampled + + """ + if isinstance(buffer, MultiAgentReplayBuffer) and not isinstance( + fake_sample_output, MultiAgentBatch + ): + fake_sample_output = SampleBatch(fake_sample_output).as_multi_agent() + + def fake_sample(_: Any = None, **kwargs) -> Optional[SampleBatchType]: + """Always returns a predefined batch. + + Args: + _: dummy arg to match signature of sample() method + __: dummy arg to match signature of sample() method + ``**kwargs``: dummy args to match signature of sample() method + + Returns: + Predefined MultiAgentBatch fake_sample_output + """ + + return fake_sample_output + + buffer.sample = fake_sample diff --git a/lib/python3.12/site-packages/ray/rllib/utils/runners/__init__.py b/lib/python3.12/site-packages/ray/rllib/utils/runners/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/utils/runners/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/runners/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..46e865775f9f5357d19c3253117854659a805608 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/runners/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/runners/__pycache__/runner.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/runners/__pycache__/runner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..70ec269b23ebe77af4479f71c01582789e9339a5 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/runners/__pycache__/runner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/runners/__pycache__/runner_group.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/runners/__pycache__/runner_group.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..869c2c21516633a09d95075d3c152b32dd34c1bf Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/runners/__pycache__/runner_group.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/runners/runner.py b/lib/python3.12/site-packages/ray/rllib/utils/runners/runner.py new file mode 100644 index 0000000000000000000000000000000000000000..fb3a8b61d278cca0af02ee61ab01afe9961e8db5 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/runners/runner.py @@ -0,0 +1,104 @@ +import abc +import logging +from typing import TYPE_CHECKING, Any, Union + +from ray.rllib.utils.actor_manager import FaultAwareApply +from ray.rllib.utils.metrics.metrics_logger import MetricsLogger +from ray.rllib.utils.typing import DeviceType, TensorType + +if TYPE_CHECKING: + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + +logger = logging.getLogger(__name__) + + +class Runner(FaultAwareApply, metaclass=abc.ABCMeta): + def __init__(self, *, config: "AlgorithmConfig", **kwargs): + """Initializes a `Runner` instance. + + Args: + config: The `AlgorithmConfig` to use to setup this `Runner`. + **kwargs: Forward compatibility `kwargs`. + """ + self.worker_index: int = kwargs.get("worker_index") + self.config: AlgorithmConfig = config.copy(copy_frozen=False) + # Set the device. + self.set_device() + # Generate the `RLModule`. + self.make_module() + self._weights_seq_no = 0 + + # Create a MetricsLogger object for logging custom stats. + self.metrics: MetricsLogger = MetricsLogger() + + # Initialize the `FaultAwareApply`. + super().__init__() + + @abc.abstractmethod + def assert_healthy(self): + """Checks that self.__init__() has been completed properly. + + Useful in case an `Runner` is run as @ray.remote (Actor) and the owner + would like to make sure the Ray Actor has been properly initialized. + + Raises: + AssertionError: If the `Runner` Actor has NOT been properly initialized. + """ + + @abc.abstractmethod + def make_module(self): + """Creates the `RLModule` for this `Runner` and assigns it to `self.module`. + + Note that users should be able to change the `Runner`'s config (e.g. change + `self.config.rl_module_spec`) and then call this method to create a new `RLModule` + with the updated configuration. + """ + pass + + @abc.abstractmethod + def run(self, **kwargs) -> Any: + """Runs the `Runner`. + + The exact logic of this method could have very different forms. + + Args: + **kwargs: Forward compatibility kwargs. + + Returns: + Anything. + """ + + @abc.abstractmethod + def get_metrics(self) -> Any: + """Returns metrics (in any form) of the logic run in this `Runner`. + + Returns: + Metrics of any form. + """ + + @abc.abstractmethod + def stop(self) -> None: + """Releases all resources used by this `Runner`. + + For example, when using a `gym.Env` in this `Runner`, you should make sure + that its `close()` method is called. + """ + + @property + @abc.abstractmethod + def _device(self) -> Union[DeviceType, None]: + """Returns the device of this `Runner`. None if framework is not supported.""" + pass + + @abc.abstractmethod + def set_device(self) -> None: + """Sets the device for this `Runner`.""" + pass + + @abc.abstractmethod + def __del__(self) -> None: + """If this Actor is deleted, clears all resources used by it.""" + + @abc.abstractmethod + def _convert_to_tensor(self, struct) -> TensorType: + """Converts structs to a framework-specific tensor.""" diff --git a/lib/python3.12/site-packages/ray/rllib/utils/runners/runner_group.py b/lib/python3.12/site-packages/ray/rllib/utils/runners/runner_group.py new file mode 100644 index 0000000000000000000000000000000000000000..eea39ff5d125cb79033fcef2086d6231468e33e2 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/runners/runner_group.py @@ -0,0 +1,751 @@ +import abc +import logging +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + List, + Optional, + Tuple, + TypeVar, + Union, +) + +import ray +from ray.actor import ActorHandle +from ray.exceptions import RayActorError +from ray.rllib.core import ( + COMPONENT_ENV_TO_MODULE_CONNECTOR, + COMPONENT_LEARNER, + COMPONENT_MODULE_TO_ENV_CONNECTOR, + COMPONENT_RL_MODULE, +) +from ray.rllib.core.learner.learner_group import LearnerGroup +from ray.rllib.utils.actor_manager import FaultTolerantActorManager +from ray.rllib.utils.metrics import NUM_ENV_STEPS_SAMPLED_LIFETIME, WEIGHTS_SEQ_NO +from ray.rllib.utils.runners.runner import Runner +from ray.rllib.utils.typing import PolicyID +from ray.util.annotations import DeveloperAPI + +if TYPE_CHECKING: + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + +logger = logging.getLogger(__name__) + +# Generic type var for `foreach_*` methods. +T = TypeVar("T") + + +@DeveloperAPI +class RunnerGroup(metaclass=abc.ABCMeta): + def __init__( + self, + config: "AlgorithmConfig", + # TODO (simon): Check, if this is needed. Derived classes could define + # this if needed. + # default_policy_class: Optional[Type[Policy]] + local_runner: Optional[bool] = False, + logdir: Optional[str] = None, + # TODO (simon): Check, if still needed. + tune_trial_id: Optional[str] = None, + pg_offset: int = 0, + _setup: bool = True, + **kwargs: Dict[str, Any], + ) -> None: + + # TODO (simon): Remove when old stack is deprecated. + self.config: AlgorithmConfig = ( + AlgorithmConfig.from_dict(config) + if isinstance(config, dict) + else (config or AlgorithmConfig()) + ) + + self._remote_config = config + self._remote_config_obj_ref = ray.put(self._remote_config) + + self._tune_trial_id = tune_trial_id + self._pg_offset = pg_offset + + self._logdir = logdir + + self._worker_manager = FaultTolerantActorManager( + max_remote_requests_in_flight_per_actor=self._max_requests_in_flight_per_runner, + init_id=1, + ) + + if _setup: + try: + self._setup( + config=config, + num_runners=self.num_runners, + local_runner=local_runner, + **kwargs, + ) + # `RunnerGroup` creation possibly fails, if some (remote) workers cannot + # be initialized properly (due to some errors in the `Runners`'s + # constructor). + except RayActorError as e: + # In case of an actor (remote worker) init failure, the remote worker + # may still exist and will be accessible, however, e.g. calling + # its `run.remote()` would result in strange "property not found" + # errors. + if e.actor_init_failed: + # Raise the original error here that the `Runners` raised + # during its construction process. This is to enforce transparency + # for the user (better to understand the real reason behind the + # failure). + # - e.args[0]: The `RayTaskError` (inside the caught `RayActorError`). + # - e.args[0].args[2]: The original `Exception` (e.g. a `ValueError` due + # to a config mismatch) thrown inside the actor. + raise e.args[0].args[2] + # In any other case, raise the `RayActorError` as-is. + else: + raise e + + def _setup( + self, + *, + config: Optional["AlgorithmConfig"] = None, + num_runners: int = 0, + local_runner: Optional[bool] = False, + validate: Optional[bool] = None, + **kwargs: Dict[str, Any], + ) -> None: + + # TODO (simon): Deprecate this as soon as we are deprecating the old stack. + self._local_runner = None + if num_runners == 0: + local_runner = True + + self.__local_config = config + + # Create a number of @ray.remote workers. + self.add_runners( + num_runners, + validate=validate + if validate is not None + else self._validate_runners_after_construction, + **kwargs, + ) + + if local_runner: + self._local_runner = self._make_runner( + runner_index=0, + num_runners=num_runners, + config=self._local_config, + **kwargs, + ) + + def add_runners(self, num_runners: int, validate: bool = False, **kwargs) -> None: + """Creates and adds a number of remote runners to this runner set.""" + + old_num_runners = self._worker_manager.num_actors() + new_runners = [ + self._make_runner( + runner_index=old_num_runners + i + 1, + num_runners=old_num_runners + num_runners, + # `self._remote_config` can be large and it's best practice to + # pass it by reference instead of value + # (https://docs.ray.io/en/latest/ray-core/patterns/pass-large-arg-by-value.html) # noqa + config=self._remote_config_obj_ref, + **kwargs, + ) + for i in range(num_runners) + ] + # Add the new workers to the worker manager. + self._worker_manager.add_actors(new_runners) + + # Validate here, whether all remote workers have been constructed properly + # and are "up and running". Establish initial states. + if validate: + self.validate() + + def validate(self) -> Exception: + for result in self._worker_manager.foreach_actor(lambda w: w.assert_healthy()): + # Simiply raise the error, which will get handled by the try-except + # clause around the _setup(). + if not result.ok: + e = result.get() + if self._ignore_ray_errors_on_runners: + logger.error( + f"Validation of {self.runner_cls.__name__} failed! Error={str(e)}" + ) + else: + raise e + + def _make_runner( + self, + *, + runner_index: int, + num_runners: int, + recreated_runner: bool = False, + config: "AlgorithmConfig", + **kwargs, + ) -> ActorHandle: + # TODO (simon): Change this in the `EnvRunner` API + # to `runner_*`. + kwargs = dict( + config=config, + worker_index=runner_index, + num_workers=num_runners, + recreated_worker=recreated_runner, + log_dir=self._logdir, + tune_trial_id=self._tune_trial_id, + **kwargs, + ) + + # If a local runner is requested just return a runner instance. + if runner_index == 0: + return self.runner_cls(**kwargs) + + # Otherwise define a bundle index and schedule the remote worker. + pg_bundle_idx = ( + -1 + if ray.util.get_current_placement_group() is None + else self._pg_offset + runner_index + ) + return ( + ray.remote(**self._remote_args)(self.runner_cls) + .options(placement_group_bundle_index=pg_bundle_idx) + .remote(**kwargs) + ) + + def sync_runner_states( + self, + *, + config: "AlgorithmConfig", + from_runner: Optional[Runner] = None, + env_steps_sampled: Optional[int] = None, + connector_states: Optional[List[Dict[str, Any]]] = None, + rl_module_state: Optional[Dict[str, Any]] = None, + runner_indices_to_update: Optional[List[int]] = None, + env_to_module=None, + module_to_env=None, + **kwargs, + ): + """Synchronizes the connectors of this `RunnerGroup`'s `Runner`s.""" + + # If no `Runner` is passed in synchronize through the local `Runner`. + from_runner = from_runner or self.local_runner + + merge = config.merge_runner_states or ( + config.merge_runner_states == "training_only" and config.in_evaluation + ) + + broadcast = config.broadcast_runner_states + + # Early out if the number of (healthy) remote workers is 0. In this case, the + # local worker is the only operating worker and thus of course always holds + # the reference connector state. + if self.num_healthy_remote_runners == 0 and self.local_runner: + self.local_runner.set_state( + { + **( + {NUM_ENV_STEPS_SAMPLED_LIFETIME: env_steps_sampled} + if env_steps_sampled is not None + else {} + ), + **(rl_module_state or {}), + } + ) + + # Also early out, if we don't merge AND don't broadcast. + if not merge and not broadcast: + return + + # Use states from all remote `Runner`s. + if merge: + if connector_states == []: + runner_states = {} + else: + if connector_states is None: + connector_states = self.foreach_runner( + lambda w: w.get_state( + components=[ + COMPONENT_ENV_TO_MODULE_CONNECTOR, + COMPONENT_MODULE_TO_ENV_CONNECTOR, + ] + ), + local_runner=False, + timeout_seconds=( + config.sync_filters_on_rollout_workers_timeout_s + ), + ) + env_to_module_states = [ + s[COMPONENT_ENV_TO_MODULE_CONNECTOR] + for s in connector_states + if COMPONENT_ENV_TO_MODULE_CONNECTOR in s + ] + module_to_env_states = [ + s[COMPONENT_MODULE_TO_ENV_CONNECTOR] + for s in connector_states + if COMPONENT_MODULE_TO_ENV_CONNECTOR in s + ] + + if ( + self.local_runner is not None + and hasattr(self.local_runner, "_env_to_module") + and hasattr(self.local_runner, "_module_to_env") + ): + assert env_to_module is None + env_to_module = self.local_runner._env_to_module + assert module_to_env is None + module_to_env = self.local_runner._module_to_env + + runner_states = {} + if env_to_module_states: + runner_states.update( + { + COMPONENT_ENV_TO_MODULE_CONNECTOR: ( + env_to_module.merge_states(env_to_module_states) + ), + } + ) + if module_to_env_states: + runner_states.update( + { + COMPONENT_MODULE_TO_ENV_CONNECTOR: ( + module_to_env.merge_states(module_to_env_states) + ), + } + ) + # Ignore states from remote `Runner`s (use the current `from_worker` states + # only). + else: + if from_runner is None: + runner_states = { + COMPONENT_ENV_TO_MODULE_CONNECTOR: env_to_module.get_state(), + COMPONENT_MODULE_TO_ENV_CONNECTOR: module_to_env.get_state(), + } + else: + runner_states = from_runner.get_state( + components=[ + COMPONENT_ENV_TO_MODULE_CONNECTOR, + COMPONENT_MODULE_TO_ENV_CONNECTOR, + ] + ) + + # Update the global number of environment steps, if necessary. + # Make sure to divide by the number of env runners (such that each `Runner` + # knows (roughly) its own(!) lifetime count and can infer the global lifetime + # count from it). + if env_steps_sampled is not None: + runner_states[NUM_ENV_STEPS_SAMPLED_LIFETIME] = env_steps_sampled // ( + config.num_runners or 1 + ) + + # If we do NOT want remote `Runner`s to get their Connector states updated, + # only update the local worker here (with all state components, except the model + # weights) and then remove the connector components. + if not broadcast: + if self.local_runner is not None: + self.local_runner.set_state(runner_states) + else: + env_to_module.set_state( + runner_states.get(COMPONENT_ENV_TO_MODULE_CONNECTOR), {} + ) + module_to_env.set_state( + runner_states.get(COMPONENT_MODULE_TO_ENV_CONNECTOR), {} + ) + runner_states.pop(COMPONENT_ENV_TO_MODULE_CONNECTOR, None) + runner_states.pop(COMPONENT_MODULE_TO_ENV_CONNECTOR, None) + + # If there are components in the state left -> Update remote workers with these + # state components (and maybe the local worker, if it hasn't been updated yet). + if runner_states: + # Update the local `Runner`, but NOT with the weights. If used at all for + # evaluation (through the user calling `self.evaluate`), RLlib would update + # the weights up front either way. + if self.local_runner is not None and broadcast: + self.local_runner.set_state(runner_states) + + # Send the model weights only to remote `Runner`s. + # In case the local `Runner` is ever needed for evaluation, + # RLlib updates its weight right before such an eval step. + if rl_module_state: + runner_states.update(rl_module_state) + + # Broadcast updated states back to all workers. + self.foreach_runner( + "set_state", # Call the `set_state()` remote method. + kwargs=dict(state=runner_states), + remote_worker_ids=runner_indices_to_update, + local_runner=False, + timeout_seconds=0.0, # This is a state update -> Fire-and-forget. + ) + + def sync_weights( + self, + policies: Optional[List[PolicyID]] = None, + from_worker_or_learner_group: Optional[Union[Runner, "LearnerGroup"]] = None, + to_worker_indices: Optional[List[int]] = None, + timeout_seconds: Optional[float] = 0.0, + inference_only: Optional[bool] = False, + **kwargs, + ) -> None: + """Syncs model weights from the given weight source to all remote workers. + + Weight source can be either a (local) rollout worker or a learner_group. It + should just implement a `get_weights` method. + + Args: + policies: Optional list of PolicyIDs to sync weights for. + If None (default), sync weights to/from all policies. + from_worker_or_learner_group: Optional (local) `Runner` instance or + LearnerGroup instance to sync from. If None (default), + sync from this `Runner`Group's local worker. + to_worker_indices: Optional list of worker indices to sync the + weights to. If None (default), sync to all remote workers. + global_vars: An optional global vars dict to set this + worker to. If None, do not update the global_vars. + timeout_seconds: Timeout in seconds to wait for the sync weights + calls to complete. Default is 0.0 (fire-and-forget, do not wait + for any sync calls to finish). Setting this to 0.0 might significantly + improve algorithm performance, depending on the algo's `training_step` + logic. + inference_only: Sync weights with workers that keep inference-only + modules. This is needed for algorithms in the new stack that + use inference-only modules. In this case only a part of the + parameters are synced to the workers. Default is False. + """ + if self.local_runner is None and from_worker_or_learner_group is None: + raise TypeError( + "No `local_runner` in `RunnerGroup`! Must provide " + "`from_worker_or_learner_group` arg in `sync_weights()`!" + ) + + # Only sync if we have remote workers or `from_worker_or_trainer` is provided. + rl_module_state = None + if self.num_remote_runners or from_worker_or_learner_group is not None: + weights_src = ( + from_worker_or_learner_group + if from_worker_or_learner_group is not None + else self.local_runner + ) + + if weights_src is None: + raise ValueError( + "`from_worker_or_trainer` is None. In this case, `RunnerGroup`^ " + "should have `local_runner`. But `local_runner` is also `None`." + ) + + modules = ( + [COMPONENT_RL_MODULE + "/" + p for p in policies] + if policies is not None + else [COMPONENT_RL_MODULE] + ) + # LearnerGroup has-a Learner, which has-a RLModule. + if isinstance(weights_src, LearnerGroup): + rl_module_state = weights_src.get_state( + components=[COMPONENT_LEARNER + "/" + m for m in modules], + inference_only=inference_only, + )[COMPONENT_LEARNER] + # `Runner` (new API stack). + else: + # Runner (remote) has a RLModule. + # TODO (sven): Replace this with a new ActorManager API: + # try_remote_request_till_success("get_state") -> tuple(int, + # remoteresult) + # `weights_src` could be the ActorManager, then. Then RLlib would know + # that it has to ping the manager to try all healthy actors until the + # first returns something. + if isinstance(weights_src, ActorHandle): + rl_module_state = ray.get( + weights_src.get_state.remote( + components=modules, + inference_only=inference_only, + ) + ) + # `Runner` (local) has an RLModule. + else: + rl_module_state = weights_src.get_state( + components=modules, + inference_only=inference_only, + ) + + # Make sure `rl_module_state` only contains the weights and the + # weight seq no, nothing else. + rl_module_state = { + k: v + for k, v in rl_module_state.items() + if k in [COMPONENT_RL_MODULE, WEIGHTS_SEQ_NO] + } + + # Move weights to the object store to avoid having to make n pickled + # copies of the weights dict for each worker. + rl_module_state_ref = ray.put(rl_module_state) + + # Sync to specified remote workers in this `Runner`Group. + self.foreach_runner( + func="set_state", + kwargs=dict(state=rl_module_state_ref), + local_runner=False, # Do not sync back to local worker. + remote_worker_ids=to_worker_indices, + timeout_seconds=timeout_seconds, + ) + + # If `from_worker_or_learner_group` is provided, also sync to this + # `RunnerGroup`'s local worker. + if self.local_runner is not None: + if from_worker_or_learner_group is not None: + self.local_runner.set_state(rl_module_state) + + def reset(self, new_remote_runners: List[ActorHandle]) -> None: + """Hard overrides the remote `Runner`s in this set with the provided ones. + + Args: + new_remote_workers: A list of new `Runner`s (as `ActorHandles`) to use as + new remote workers. + """ + self._worker_manager.clear() + self._worker_manager.add_actors(new_remote_runners) + + def stop(self) -> None: + """Calls `stop` on all `Runner`s (including the local one).""" + try: + # Make sure we stop all `Runner`s, include the ones that were just + # restarted / recovered or that are tagged unhealthy (at least, we should + # try). + self.foreach_runner( + lambda w: w.stop(), healthy_only=False, local_runner=True + ) + except Exception: + logger.exception("Failed to stop workers!") + finally: + self._worker_manager.clear() + + def foreach_runner( + self, + func: Union[Callable[[Runner], T], List[Callable[[Runner], T]], str, List[str]], + *, + kwargs=None, + local_runner: bool = True, + healthy_only: bool = True, + remote_worker_ids: List[int] = None, + timeout_seconds: Optional[float] = None, + return_obj_refs: bool = False, + mark_healthy: bool = False, + ) -> List[T]: + """Calls the given function with each `Runner` as its argument. + + Args: + func: The function to call for each `Runner`s. The only call argument is + the respective `Runner` instance. + local_env_runner: Whether to apply `func` to local `Runner`, too. + Default is True. + healthy_only: Apply `func` on known-to-be healthy `Runner`s only. + remote_worker_ids: Apply `func` on a selected set of remote `Runner`s. + Use None (default) for all remote `Runner`s. + timeout_seconds: Time to wait (in seconds) for results. Set this to 0.0 for + fire-and-forget. Set this to None (default) to wait infinitely (i.e. for + synchronous execution). + return_obj_refs: Whether to return `ObjectRef` instead of actual results. + Note, for fault tolerance reasons, these returned ObjectRefs should + never be resolved with ray.get() outside of this `RunnerGroup`. + mark_healthy: Whether to mark all those `Runner`s healthy again that are + currently marked unhealthy AND that returned results from the remote + call (within the given `timeout_seconds`). + Note that `Runner`s are NOT set unhealthy, if they simply time out + (only if they return a `RayActorError`). + Also note that this setting is ignored if `healthy_only=True` (b/c + `mark_healthy` only affects `Runner`s that are currently tagged as + unhealthy). + + Returns: + The list of return values of all calls to `func([worker])`. + """ + assert ( + not return_obj_refs or not local_runner + ), "Can not return `ObjectRef` from local worker." + + local_result = [] + if local_runner and self.local_runner is not None: + if kwargs: + local_kwargs = kwargs[0] + kwargs = kwargs[1:] + else: + local_kwargs = {} + kwargs = kwargs + if isinstance(func, str): + local_result = [getattr(self.local_runner, func)(**local_kwargs)] + else: + local_result = [func(self.local_runner, **local_kwargs)] + + if not self._worker_manager.actor_ids(): + return local_result + + remote_results = self._worker_manager.foreach_actor( + func, + kwargs=kwargs, + healthy_only=healthy_only, + remote_actor_ids=remote_worker_ids, + timeout_seconds=timeout_seconds, + return_obj_refs=return_obj_refs, + mark_healthy=mark_healthy, + ) + + FaultTolerantActorManager.handle_remote_call_result_errors( + remote_results, ignore_ray_errors=self._ignore_ray_errors_on_runners + ) + + # With application errors handled, return good results. + remote_results = [r.get() for r in remote_results.ignore_errors()] + + return local_result + remote_results + + def foreach_runner_async( + self, + func: Union[Callable[[Runner], T], List[Callable[[Runner], T]], str, List[str]], + *, + healthy_only: bool = True, + remote_worker_ids: List[int] = None, + ) -> int: + """Calls the given function asynchronously with each `Runner` as the argument. + + Does not return results directly. Instead, `fetch_ready_async_reqs()` can be + used to pull results in an async manner whenever they are available. + + Args: + func: The function to call for each `Runner`s. The only call argument is + the respective `Runner` instance. + healthy_only: Apply `func` on known-to-be healthy `Runner`s only. + remote_worker_ids: Apply `func` on a selected set of remote `Runner`s. + + Returns: + The number of async requests that have actually been made. This is the + length of `remote_worker_ids` (or self.num_remote_workers()` if + `remote_worker_ids` is None) minus the number of requests that were NOT + made b/c a remote `Runner` already had its + `max_remote_requests_in_flight_per_actor` counter reached. + """ + return self._worker_manager.foreach_actor_async( + func, + healthy_only=healthy_only, + remote_actor_ids=remote_worker_ids, + ) + + def fetch_ready_async_reqs( + self, + *, + timeout_seconds: Optional[float] = 0.0, + return_obj_refs: bool = False, + mark_healthy: bool = False, + ) -> List[Tuple[int, T]]: + """Get esults from outstanding asynchronous requests that are ready. + + Args: + timeout_seconds: Time to wait for results. Default is 0, meaning + those requests that are already ready. + return_obj_refs: Whether to return ObjectRef instead of actual results. + mark_healthy: Whether to mark all those workers healthy again that are + currently marked unhealthy AND that returned results from the remote + call (within the given `timeout_seconds`). + Note that workers are NOT set unhealthy, if they simply time out + (only if they return a RayActorError). + Also note that this setting is ignored if `healthy_only=True` (b/c + `mark_healthy` only affects workers that are currently tagged as + unhealthy). + + Returns: + A list of results successfully returned from outstanding remote calls, + paired with the indices of the callee workers. + """ + remote_results = self._worker_manager.fetch_ready_async_reqs( + timeout_seconds=timeout_seconds, + return_obj_refs=return_obj_refs, + mark_healthy=mark_healthy, + ) + + FaultTolerantActorManager.handle_remote_call_result_errors( + remote_results, + ignore_ray_errors=self._ignore_ray_errors_on_runners, + ) + + return [(r.actor_id, r.get()) for r in remote_results.ignore_errors()] + + def probe_unhealthy_runners(self) -> List[int]: + """Checks for unhealthy workers and tries restoring their states. + + Returns: + List of IDs of the workers that were restored. + """ + return self._worker_manager.probe_unhealthy_actors( + timeout_seconds=self.runner_health_probe_timeout_s, + mark_healthy=True, + ) + + @property + @abc.abstractmethod + def runner_health_probe_timeout_s(self): + """Number of seconds to wait for health probe calls to `Runner`s.""" + + @property + @abc.abstractmethod + def runner_cls(self) -> Callable: + """Class for each runner.""" + + @property + def _local_config(self) -> "AlgorithmConfig": + """Returns the config for a local `Runner`.""" + return self.__local_config + + @property + def local_runner(self) -> Runner: + """Returns the local `Runner`.""" + return self._local_runner + + @property + def healthy_runner_ids(self) -> List[int]: + """Returns the list of remote `Runner` IDs.""" + return self._worker_manager.healthy_actor_ids() + + @property + @abc.abstractmethod + def num_runners(self) -> int: + """Number of runners to schedule and manage.""" + + @property + def num_remote_runners(self) -> int: + """Number of remote `Runner`s.""" + return self._worker_manager.num_actors() + + @property + def num_healthy_remote_runners(self) -> int: + """Returns the number of healthy remote `Runner`s.""" + return self._worker_manager.num_healthy_actors() + + @property + def num_healthy_runners(self) -> int: + """Returns the number of healthy `Runner`s.""" + return int(bool(self._local_runner)) + self.num_healthy_remote_runners() + + @property + def num_in_flight_async_reqs(self) -> int: + """Returns the number of in-flight async requests.""" + return self._worker_manager.num_outstanding_async_reqs() + + @property + def num_remote_runner_restarts(self) -> int: + """Returns the number of times managed remote `Runner`s have been restarted.""" + return self._worker_manager.total_num_restarts() + + @property + @abc.abstractmethod + def _remote_args(self): + """Remote arguments for each runner.""" + + @property + @abc.abstractmethod + def _ignore_ray_errors_on_runners(self): + """If errors in runners should be ignored.""" + + @property + @abc.abstractmethod + def _max_requests_in_flight_per_runner(self): + """Maximum requests in flight per runner.""" + + @property + @abc.abstractmethod + def _validate_runners_after_construction(self): + """If runners should validated after constructed.""" diff --git a/lib/python3.12/site-packages/ray/rllib/utils/schedules/__init__.py b/lib/python3.12/site-packages/ray/rllib/utils/schedules/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..96a64c41a33a176477d806ec0c2e6cc4d0cf97ae --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/schedules/__init__.py @@ -0,0 +1,15 @@ +from ray.rllib.utils.schedules.constant_schedule import ConstantSchedule +from ray.rllib.utils.schedules.exponential_schedule import ExponentialSchedule +from ray.rllib.utils.schedules.linear_schedule import LinearSchedule +from ray.rllib.utils.schedules.piecewise_schedule import PiecewiseSchedule +from ray.rllib.utils.schedules.polynomial_schedule import PolynomialSchedule +from ray.rllib.utils.schedules.schedule import Schedule + +__all__ = [ + "ConstantSchedule", + "ExponentialSchedule", + "LinearSchedule", + "Schedule", + "PiecewiseSchedule", + "PolynomialSchedule", +] diff --git a/lib/python3.12/site-packages/ray/rllib/utils/schedules/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/schedules/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7514c7786bbe261ff9411a0701d067aca3d8a8eb Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/schedules/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/schedules/__pycache__/constant_schedule.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/schedules/__pycache__/constant_schedule.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b3fdbc0c1fb6fb4762d3958cba51c453dc81e996 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/schedules/__pycache__/constant_schedule.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/schedules/__pycache__/exponential_schedule.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/schedules/__pycache__/exponential_schedule.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4a170b1e23f15154a8d813b040210f0f97a905f8 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/schedules/__pycache__/exponential_schedule.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/schedules/__pycache__/linear_schedule.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/schedules/__pycache__/linear_schedule.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..33d9d8cbb8e65d8fe1e9624be38114dac4255fd7 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/schedules/__pycache__/linear_schedule.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/schedules/__pycache__/piecewise_schedule.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/schedules/__pycache__/piecewise_schedule.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..914197090471aabcacebbe5ab7a0724701956eb5 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/schedules/__pycache__/piecewise_schedule.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/schedules/__pycache__/polynomial_schedule.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/schedules/__pycache__/polynomial_schedule.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2908221c5285f4867fdf58d05e221dc9d530f5d6 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/schedules/__pycache__/polynomial_schedule.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/schedules/__pycache__/schedule.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/schedules/__pycache__/schedule.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..08b7044a79edf81e9ffcf29ea6280443f5162b88 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/schedules/__pycache__/schedule.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/schedules/__pycache__/scheduler.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/schedules/__pycache__/scheduler.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8452e4ea2584d0062e474c61730bc896bda1a585 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/schedules/__pycache__/scheduler.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/schedules/constant_schedule.py b/lib/python3.12/site-packages/ray/rllib/utils/schedules/constant_schedule.py new file mode 100644 index 0000000000000000000000000000000000000000..5ad58e266f8d3e21144daae1c251cd2333b85907 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/schedules/constant_schedule.py @@ -0,0 +1,32 @@ +from typing import Optional + +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.schedules.schedule import Schedule +from ray.rllib.utils.typing import TensorType + +tf1, tf, tfv = try_import_tf() + + +@OldAPIStack +class ConstantSchedule(Schedule): + """A Schedule where the value remains constant over time.""" + + def __init__(self, value: float, framework: Optional[str] = None): + """Initializes a ConstantSchedule instance. + + Args: + value: The constant value to return, independently of time. + framework: The framework descriptor string, e.g. "tf", + "torch", or None. + """ + super().__init__(framework=framework) + self._v = value + + @override(Schedule) + def _value(self, t: TensorType) -> TensorType: + return self._v + + @override(Schedule) + def _tf_value_op(self, t: TensorType) -> TensorType: + return tf.constant(self._v) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/schedules/exponential_schedule.py b/lib/python3.12/site-packages/ray/rllib/utils/schedules/exponential_schedule.py new file mode 100644 index 0000000000000000000000000000000000000000..4aafac19470baa3081d06a5a3bc4cbd18d76e384 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/schedules/exponential_schedule.py @@ -0,0 +1,50 @@ +from typing import Optional + +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.schedules.schedule import Schedule +from ray.rllib.utils.typing import TensorType + +torch, _ = try_import_torch() + + +@OldAPIStack +class ExponentialSchedule(Schedule): + """Exponential decay schedule from `initial_p` to `final_p`. + + Reduces output over `schedule_timesteps`. After this many time steps + always returns `final_p`. + """ + + def __init__( + self, + schedule_timesteps: int, + framework: Optional[str] = None, + initial_p: float = 1.0, + decay_rate: float = 0.1, + ): + """Initializes a ExponentialSchedule instance. + + Args: + schedule_timesteps: Number of time steps for which to + linearly anneal initial_p to final_p. + framework: The framework descriptor string, e.g. "tf", + "torch", or None. + initial_p: Initial output value. + decay_rate: The percentage of the original value after + 100% of the time has been reached (see formula above). + >0.0: The smaller the decay-rate, the stronger the decay. + 1.0: No decay at all. + """ + super().__init__(framework=framework) + assert schedule_timesteps > 0 + self.schedule_timesteps = schedule_timesteps + self.initial_p = initial_p + self.decay_rate = decay_rate + + @override(Schedule) + def _value(self, t: TensorType) -> TensorType: + """Returns the result of: initial_p * decay_rate ** (`t`/t_max).""" + if self.framework == "torch" and torch and isinstance(t, torch.Tensor): + t = t.float() + return self.initial_p * self.decay_rate ** (t / self.schedule_timesteps) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/schedules/linear_schedule.py b/lib/python3.12/site-packages/ray/rllib/utils/schedules/linear_schedule.py new file mode 100644 index 0000000000000000000000000000000000000000..d23647b9bbebf757f002444bb231a9b78d863a47 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/schedules/linear_schedule.py @@ -0,0 +1,18 @@ +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.schedules.polynomial_schedule import PolynomialSchedule + + +@OldAPIStack +class LinearSchedule(PolynomialSchedule): + """Linear interpolation between `initial_p` and `final_p`. + + Uses `PolynomialSchedule` with power=1.0. + + + The formula is: + value = `final_p` + (`initial_p` - `final_p`) * (1 - `t`/t_max) + """ + + def __init__(self, **kwargs): + """Initializes a LinearSchedule instance.""" + super().__init__(power=1.0, **kwargs) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/schedules/piecewise_schedule.py b/lib/python3.12/site-packages/ray/rllib/utils/schedules/piecewise_schedule.py new file mode 100644 index 0000000000000000000000000000000000000000..6c4b15478b3f21bfa8612bdd5e39f952d0801086 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/schedules/piecewise_schedule.py @@ -0,0 +1,105 @@ +from typing import Callable, List, Optional, Tuple + +from ray.rllib.utils.annotations import override +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.schedules.schedule import Schedule +from ray.rllib.utils.typing import TensorType +from ray.util.annotations import DeveloperAPI + +tf1, tf, tfv = try_import_tf() + + +def _linear_interpolation(left, right, alpha): + return left + alpha * (right - left) + + +@DeveloperAPI +class PiecewiseSchedule(Schedule): + """Implements a Piecewise Scheduler.""" + + def __init__( + self, + endpoints: List[Tuple[int, float]], + framework: Optional[str] = None, + interpolation: Callable[ + [TensorType, TensorType, TensorType], TensorType + ] = _linear_interpolation, + outside_value: Optional[float] = None, + ): + """Initializes a PiecewiseSchedule instance. + + Args: + endpoints: A list of tuples + `(t, value)` such that the output + is an interpolation (given by the `interpolation` callable) + between two values. + E.g. + t=400 and endpoints=[(0, 20.0),(500, 30.0)] + output=20.0 + 0.8 * (30.0 - 20.0) = 28.0 + NOTE: All the values for time must be sorted in an increasing + order. + framework: The framework descriptor string, e.g. "tf", + "torch", or None. + interpolation: A function that takes the left-value, + the right-value and an alpha interpolation parameter + (0.0=only left value, 1.0=only right value), which is the + fraction of distance from left endpoint to right endpoint. + outside_value: If t in call to `value` is + outside of all the intervals in `endpoints` this value is + returned. If None then an AssertionError is raised when outside + value is requested. + """ + super().__init__(framework=framework) + + idxes = [e[0] for e in endpoints] + assert idxes == sorted(idxes) + self.interpolation = interpolation + self.outside_value = outside_value + self.endpoints = [(int(e[0]), float(e[1])) for e in endpoints] + + @override(Schedule) + def _value(self, t: TensorType) -> TensorType: + # Find t in our list of endpoints. + for (l_t, l), (r_t, r) in zip(self.endpoints[:-1], self.endpoints[1:]): + # When found, return an interpolation (default: linear). + if l_t <= t < r_t: + alpha = float(t - l_t) / (r_t - l_t) + return self.interpolation(l, r, alpha) + + # t does not belong to any of the pieces, return `self.outside_value`. + assert self.outside_value is not None + return self.outside_value + + @override(Schedule) + def _tf_value_op(self, t: TensorType) -> TensorType: + assert self.outside_value is not None, ( + "tf-version of PiecewiseSchedule requires `outside_value` to be " + "provided!" + ) + + endpoints = tf.cast(tf.stack([e[0] for e in self.endpoints] + [-1]), tf.int64) + + # Create all possible interpolation results. + results_list = [] + for (l_t, l), (r_t, r) in zip(self.endpoints[:-1], self.endpoints[1:]): + alpha = tf.cast(t - l_t, tf.float32) / tf.cast(r_t - l_t, tf.float32) + results_list.append(self.interpolation(l, r, alpha)) + # If t does not belong to any of the pieces, return `outside_value`. + results_list.append(self.outside_value) + results_list = tf.stack(results_list) + + # Return correct results tensor depending on where we find t. + def _cond(i, x): + x = tf.cast(x, tf.int64) + return tf.logical_not( + tf.logical_or( + tf.equal(endpoints[i + 1], -1), + tf.logical_and(endpoints[i] <= x, x < endpoints[i + 1]), + ) + ) + + def _body(i, x): + return (i + 1, t) + + idx_and_t = tf.while_loop(_cond, _body, [tf.constant(0, dtype=tf.int64), t]) + return results_list[idx_and_t[0]] diff --git a/lib/python3.12/site-packages/ray/rllib/utils/schedules/polynomial_schedule.py b/lib/python3.12/site-packages/ray/rllib/utils/schedules/polynomial_schedule.py new file mode 100644 index 0000000000000000000000000000000000000000..17a2820d9af811d8c647f467501cd906f2bbdd9d --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/schedules/polynomial_schedule.py @@ -0,0 +1,67 @@ +from typing import Optional + +from ray.rllib.utils.annotations import OldAPIStack, override +from ray.rllib.utils.framework import try_import_tf, try_import_torch +from ray.rllib.utils.schedules.schedule import Schedule +from ray.rllib.utils.typing import TensorType + +tf1, tf, tfv = try_import_tf() +torch, _ = try_import_torch() + + +@OldAPIStack +class PolynomialSchedule(Schedule): + """Polynomial interpolation between `initial_p` and `final_p`. + + Over `schedule_timesteps`. After this many time steps, always returns + `final_p`. + """ + + def __init__( + self, + schedule_timesteps: int, + final_p: float, + framework: Optional[str], + initial_p: float = 1.0, + power: float = 2.0, + ): + """Initializes a PolynomialSchedule instance. + + Args: + schedule_timesteps: Number of time steps for which to + linearly anneal initial_p to final_p + final_p: Final output value. + framework: The framework descriptor string, e.g. "tf", + "torch", or None. + initial_p: Initial output value. + power: The exponent to use (default: quadratic). + """ + super().__init__(framework=framework) + assert schedule_timesteps > 0 + self.schedule_timesteps = schedule_timesteps + self.final_p = final_p + self.initial_p = initial_p + self.power = power + + @override(Schedule) + def _value(self, t: TensorType) -> TensorType: + """Returns the result of: + final_p + (initial_p - final_p) * (1 - `t`/t_max) ** power + """ + if self.framework == "torch" and torch and isinstance(t, torch.Tensor): + t = t.float() + t = min(t, self.schedule_timesteps) + return ( + self.final_p + + (self.initial_p - self.final_p) + * (1.0 - (t / self.schedule_timesteps)) ** self.power + ) + + @override(Schedule) + def _tf_value_op(self, t: TensorType) -> TensorType: + t = tf.math.minimum(t, self.schedule_timesteps) + return ( + self.final_p + + (self.initial_p - self.final_p) + * (1.0 - (t / self.schedule_timesteps)) ** self.power + ) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/schedules/schedule.py b/lib/python3.12/site-packages/ray/rllib/utils/schedules/schedule.py new file mode 100644 index 0000000000000000000000000000000000000000..fa5b113ff5a9c1d328cdc0db0ec8197d4efa8b0f --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/schedules/schedule.py @@ -0,0 +1,73 @@ +from abc import ABCMeta, abstractmethod +from typing import Any, Union + +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.typing import TensorType + +tf1, tf, tfv = try_import_tf() + + +@OldAPIStack +class Schedule(metaclass=ABCMeta): + """Schedule classes implement various time-dependent scheduling schemas. + + - Constant behavior. + - Linear decay. + - Piecewise decay. + - Exponential decay. + + Useful for backend-agnostic rate/weight changes for learning rates, + exploration epsilons, beta parameters for prioritized replay, loss weights + decay, etc.. + + Each schedule can be called directly with the `t` (absolute time step) + value and returns the value dependent on the Schedule and the passed time. + """ + + def __init__(self, framework): + self.framework = framework + + def value(self, t: Union[int, TensorType]) -> Any: + """Generates the value given a timestep (based on schedule's logic). + + Args: + t: The time step. This could be a tf.Tensor. + + Returns: + The calculated value depending on the schedule and `t`. + """ + if self.framework in ["tf2", "tf"]: + return self._tf_value_op(t) + return self._value(t) + + def __call__(self, t: Union[int, TensorType]) -> Any: + """Simply calls self.value(t). Implemented to make Schedules callable.""" + return self.value(t) + + @abstractmethod + def _value(self, t: Union[int, TensorType]) -> Any: + """ + Returns the value based on a time step input. + + Args: + t: The time step. This could be a tf.Tensor. + + Returns: + The calculated value depending on the schedule and `t`. + """ + raise NotImplementedError + + def _tf_value_op(self, t: TensorType) -> TensorType: + """ + Returns the tf-op that calculates the value based on a time step input. + + Args: + t: The time step op (int tf.Tensor). + + Returns: + The calculated value depending on the schedule and `t`. + """ + # By default (most of the time), tf should work with python code. + # Override only if necessary. + return self._value(t) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/schedules/scheduler.py b/lib/python3.12/site-packages/ray/rllib/utils/schedules/scheduler.py new file mode 100644 index 0000000000000000000000000000000000000000..d4de4e0ab5ba92778bd2b7e60efe06cf50620b44 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/schedules/scheduler.py @@ -0,0 +1,174 @@ +from typing import Optional + +from ray.rllib.utils.framework import try_import_tf, try_import_torch +from ray.rllib.utils.schedules.piecewise_schedule import PiecewiseSchedule +from ray.rllib.utils.typing import LearningRateOrSchedule, TensorType +from ray.util.annotations import DeveloperAPI + +_, tf, _ = try_import_tf() +torch, _ = try_import_torch() + + +@DeveloperAPI +class Scheduler: + """Class to manage a scheduled (framework-dependent) tensor variable. + + Uses the PiecewiseSchedule (for maximum configuration flexibility) + """ + + def __init__( + self, + fixed_value_or_schedule: LearningRateOrSchedule, + *, + framework: str = "torch", + device: Optional[str] = None, + ): + """Initializes a Scheduler instance. + + Args: + fixed_value_or_schedule: A fixed, constant value (in case no schedule should + be used) or a schedule configuration in the format of + [[timestep, value], [timestep, value], ...] + Intermediary timesteps will be assigned to linerarly interpolated + values. A schedule config's first entry must + start with timestep 0, i.e.: [[0, initial_value], [...]]. + framework: The framework string, for which to create the tensor variable + that hold the current value. This is the variable that can be used in + the graph, e.g. in a loss function. + device: Optional device (for torch) to place the tensor variable on. + """ + self.framework = framework + self.device = device + self.use_schedule = isinstance(fixed_value_or_schedule, (list, tuple)) + + if self.use_schedule: + # Custom schedule, based on list of + # ([ts], [value to be reached by ts])-tuples. + self._schedule = PiecewiseSchedule( + fixed_value_or_schedule, + outside_value=fixed_value_or_schedule[-1][-1], + framework=None, + ) + # As initial tensor valie, use the first timestep's (must be 0) value. + self._curr_value = self._create_tensor_variable( + initial_value=fixed_value_or_schedule[0][1] + ) + + # If no schedule, pin (fix) given value. + else: + self._curr_value = fixed_value_or_schedule + + @staticmethod + def validate( + *, + fixed_value_or_schedule: LearningRateOrSchedule, + setting_name: str, + description: str, + ) -> None: + """Performs checking of a certain schedule configuration. + + The first entry in `value_or_schedule` (if it's not a fixed value) must have a + timestep of 0. + + Args: + fixed_value_or_schedule: A fixed, constant value (in case no schedule should + be used) or a schedule configuration in the format of + [[timestep, value], [timestep, value], ...] + Intermediary timesteps will be assigned to linerarly interpolated + values. A schedule config's first entry must + start with timestep 0, i.e.: [[0, initial_value], [...]]. + setting_name: The property name of the schedule setting (within a config), + e.g. `lr` or `entropy_coeff`. + description: A full text description of the property that's being scheduled, + e.g. `learning rate`. + + Raises: + ValueError: In case, errors are found in the schedule's format. + """ + # Fixed (single) value. + if ( + isinstance(fixed_value_or_schedule, (int, float)) + or fixed_value_or_schedule is None + ): + return + + if not isinstance(fixed_value_or_schedule, (list, tuple)) or ( + len(fixed_value_or_schedule) < 2 + ): + raise ValueError( + f"Invalid `{setting_name}` ({fixed_value_or_schedule}) specified! " + f"Must be a list of 2 or more tuples, each of the form " + f"(`timestep`, `{description} to reach`), for example " + "`[(0, 0.001), (1e6, 0.0001), (2e6, 0.00005)]`." + ) + elif fixed_value_or_schedule[0][0] != 0: + raise ValueError( + f"When providing a `{setting_name}` schedule, the first timestep must " + f"be 0 and the corresponding lr value is the initial {description}! " + f"You provided ts={fixed_value_or_schedule[0][0]} {description}=" + f"{fixed_value_or_schedule[0][1]}." + ) + elif any(len(pair) != 2 for pair in fixed_value_or_schedule): + raise ValueError( + f"When providing a `{setting_name}` schedule, each tuple in the " + f"schedule list must have exctly 2 items of the form " + f"(`timestep`, `{description} to reach`), for example " + "`[(0, 0.001), (1e6, 0.0001), (2e6, 0.00005)]`." + ) + + def get_current_value(self) -> TensorType: + """Returns the current value (as a tensor variable). + + This method should be used in loss functions of other (in-graph) places + where the current value is needed. + + Returns: + The tensor variable (holding the current value to be used). + """ + return self._curr_value + + def update(self, timestep: int) -> float: + """Updates the underlying (framework specific) tensor variable. + + In case of a fixed value, this method does nothing and only returns the fixed + value as-is. + + Args: + timestep: The current timestep that the update might depend on. + + Returns: + The current value of the tensor variable as a python float. + """ + if self.use_schedule: + python_value = self._schedule.value(t=timestep) + if self.framework == "torch": + self._curr_value.data = torch.tensor(python_value) + else: + self._curr_value.assign(python_value) + else: + python_value = self._curr_value + + return python_value + + def _create_tensor_variable(self, initial_value: float) -> TensorType: + """Creates a framework-specific tensor variable to be scheduled. + + Args: + initial_value: The initial (float) value for the variable to hold. + + Returns: + The created framework-specific tensor variable. + """ + if self.framework == "torch": + return torch.tensor( + initial_value, + requires_grad=False, + dtype=torch.float32, + device=self.device, + ) + else: + return tf.Variable( + initial_value, + trainable=False, + dtype=tf.float32, + ) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/serialization.py b/lib/python3.12/site-packages/ray/rllib/utils/serialization.py new file mode 100644 index 0000000000000000000000000000000000000000..65698a80568b2d78bd6db2f753c6c9bd3fba93a5 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/serialization.py @@ -0,0 +1,418 @@ +import base64 +import importlib +import io +import zlib +from collections import OrderedDict +from typing import Any, Dict, Optional, Sequence, Type, Union + +import gymnasium as gym +import numpy as np + +import ray +from ray.rllib.utils.annotations import DeveloperAPI +from ray.rllib.utils.error import NotSerializable +from ray.rllib.utils.spaces.flexdict import FlexDict +from ray.rllib.utils.spaces.repeated import Repeated +from ray.rllib.utils.spaces.simplex import Simplex + +NOT_SERIALIZABLE = "__not_serializable__" + + +@DeveloperAPI +def convert_numpy_to_python_primitives(obj: Any): + """Convert an object that is a numpy type to a python type. + + If the object is not a numpy type, it is returned unchanged. + + Args: + obj: The object to convert. + """ + if isinstance(obj, np.integer): + return int(obj) + elif isinstance(obj, np.floating): + return float(obj) + elif isinstance(obj, np.bool_): + return bool(obj) + elif isinstance(obj, np.str_): + return str(obj) + elif isinstance(obj, np.ndarray): + ret = obj.tolist() + for i, v in enumerate(ret): + ret[i] = convert_numpy_to_python_primitives(v) + return ret + else: + return obj + + +def _serialize_ndarray(array: np.ndarray) -> str: + """Pack numpy ndarray into Base64 encoded strings for serialization. + + This function uses numpy.save() instead of pickling to ensure + compatibility. + + Args: + array: numpy ndarray. + + Returns: + b64 escaped string. + """ + buf = io.BytesIO() + np.save(buf, array) + return base64.b64encode(zlib.compress(buf.getvalue())).decode("ascii") + + +def _deserialize_ndarray(b64_string: str) -> np.ndarray: + """Unpack b64 escaped string into numpy ndarray. + + This function assumes the unescaped bytes are of npy format. + + Args: + b64_string: Base64 escaped string. + + Returns: + numpy ndarray. + """ + return np.load( + io.BytesIO(zlib.decompress(base64.b64decode(b64_string))), allow_pickle=True + ) + + +@DeveloperAPI +def gym_space_to_dict(space: gym.spaces.Space) -> Dict: + """Serialize a gym Space into a JSON-serializable dict. + + Args: + space: gym.spaces.Space + + Returns: + Serialized JSON string. + """ + if space is None: + return None + + def _box(sp: gym.spaces.Box) -> Dict: + return { + "space": "box", + "low": _serialize_ndarray(sp.low), + "high": _serialize_ndarray(sp.high), + "shape": sp._shape, # shape is a tuple. + "dtype": sp.dtype.str, + } + + def _discrete(sp: gym.spaces.Discrete) -> Dict: + d = { + "space": "discrete", + "n": int(sp.n), + } + # Offset is a relatively new Discrete space feature. + if hasattr(sp, "start"): + d["start"] = int(sp.start) + return d + + def _multi_binary(sp: gym.spaces.MultiBinary) -> Dict: + return { + "space": "multi-binary", + "n": sp.n, + } + + def _multi_discrete(sp: gym.spaces.MultiDiscrete) -> Dict: + return { + "space": "multi-discrete", + "nvec": _serialize_ndarray(sp.nvec), + "dtype": sp.dtype.str, + } + + def _tuple(sp: gym.spaces.Tuple) -> Dict: + return { + "space": "tuple", + "spaces": [gym_space_to_dict(sp) for sp in sp.spaces], + } + + def _dict(sp: gym.spaces.Dict) -> Dict: + return { + "space": "dict", + "spaces": {k: gym_space_to_dict(sp) for k, sp in sp.spaces.items()}, + } + + def _simplex(sp: Simplex) -> Dict: + return { + "space": "simplex", + "shape": sp._shape, # shape is a tuple. + "concentration": sp.concentration, + "dtype": sp.dtype.str, + } + + def _repeated(sp: Repeated) -> Dict: + return { + "space": "repeated", + "child_space": gym_space_to_dict(sp.child_space), + "max_len": sp.max_len, + } + + def _flex_dict(sp: FlexDict) -> Dict: + d = { + "space": "flex_dict", + } + for k, s in sp.spaces: + d[k] = gym_space_to_dict(s) + return d + + def _text(sp: "gym.spaces.Text") -> Dict: + # Note (Kourosh): This only works in gym >= 0.25.0 + charset = getattr(sp, "character_set", None) + if charset is None: + charset = getattr(sp, "charset", None) + if charset is None: + raise ValueError( + "Text space must have a character_set or charset attribute" + ) + return { + "space": "text", + "min_length": sp.min_length, + "max_length": sp.max_length, + "charset": charset, + } + + if isinstance(space, gym.spaces.Box): + return _box(space) + elif isinstance(space, gym.spaces.Discrete): + return _discrete(space) + elif isinstance(space, gym.spaces.MultiBinary): + return _multi_binary(space) + elif isinstance(space, gym.spaces.MultiDiscrete): + return _multi_discrete(space) + elif isinstance(space, gym.spaces.Tuple): + return _tuple(space) + elif isinstance(space, gym.spaces.Dict): + return _dict(space) + elif isinstance(space, gym.spaces.Text): + return _text(space) + elif isinstance(space, Simplex): + return _simplex(space) + elif isinstance(space, Repeated): + return _repeated(space) + elif isinstance(space, FlexDict): + return _flex_dict(space) + else: + raise ValueError("Unknown space type for serialization, ", type(space)) + + +@DeveloperAPI +def space_to_dict(space: gym.spaces.Space) -> Dict: + d = {"space": gym_space_to_dict(space)} + if "original_space" in space.__dict__: + d["original_space"] = space_to_dict(space.original_space) + return d + + +@DeveloperAPI +def gym_space_from_dict(d: Dict) -> gym.spaces.Space: + """De-serialize a dict into gym Space. + + Args: + str: serialized JSON str. + + Returns: + De-serialized gym space. + """ + if d is None: + return None + + def __common(d: Dict): + """Common updates to the dict before we use it to construct spaces""" + ret = d.copy() + del ret["space"] + if "dtype" in ret: + ret["dtype"] = np.dtype(ret["dtype"]) + return ret + + def _box(d: Dict) -> gym.spaces.Box: + ret = d.copy() + ret.update( + { + "low": _deserialize_ndarray(d["low"]), + "high": _deserialize_ndarray(d["high"]), + } + ) + return gym.spaces.Box(**__common(ret)) + + def _discrete(d: Dict) -> gym.spaces.Discrete: + return gym.spaces.Discrete(**__common(d)) + + def _multi_binary(d: Dict) -> gym.spaces.MultiBinary: + return gym.spaces.MultiBinary(**__common(d)) + + def _multi_discrete(d: Dict) -> gym.spaces.MultiDiscrete: + ret = d.copy() + ret.update( + { + "nvec": _deserialize_ndarray(ret["nvec"]), + } + ) + return gym.spaces.MultiDiscrete(**__common(ret)) + + def _tuple(d: Dict) -> gym.spaces.Discrete: + spaces = [gym_space_from_dict(sp) for sp in d["spaces"]] + return gym.spaces.Tuple(spaces=spaces) + + def _dict(d: Dict) -> gym.spaces.Discrete: + # We need to always use an OrderedDict here to cover the following two ways, by + # which a user might construct a Dict space originally. We need to restore this + # original Dict space with the exact order of keys the user intended to. + # - User provides an OrderedDict inside the gym.spaces.Dict constructor -> + # gymnasium should NOT further sort the keys. The same (user-provided) order + # must be restored. + # - User provides a simple dict inside the gym.spaces.Dict constructor -> + # By its API definition, gymnasium automatically sorts all keys alphabetically. + # The same (alphabetical) order must thus be restored. + spaces = OrderedDict( + {k: gym_space_from_dict(sp) for k, sp in d["spaces"].items()} + ) + return gym.spaces.Dict(spaces=spaces) + + def _simplex(d: Dict) -> Simplex: + return Simplex(**__common(d)) + + def _repeated(d: Dict) -> Repeated: + child_space = gym_space_from_dict(d["child_space"]) + return Repeated(child_space=child_space, max_len=d["max_len"]) + + def _flex_dict(d: Dict) -> FlexDict: + spaces = {k: gym_space_from_dict(s) for k, s in d.items() if k != "space"} + return FlexDict(spaces=spaces) + + def _text(d: Dict) -> "gym.spaces.Text": + return gym.spaces.Text(**__common(d)) + + space_map = { + "box": _box, + "discrete": _discrete, + "multi-binary": _multi_binary, + "multi-discrete": _multi_discrete, + "tuple": _tuple, + "dict": _dict, + "simplex": _simplex, + "repeated": _repeated, + "flex_dict": _flex_dict, + "text": _text, + } + + space_type = d["space"] + if space_type not in space_map: + raise ValueError("Unknown space type for de-serialization, ", space_type) + + return space_map[space_type](d) + + +@DeveloperAPI +def space_from_dict(d: Dict) -> gym.spaces.Space: + space = gym_space_from_dict(d["space"]) + if "original_space" in d: + assert "space" in d["original_space"] + if isinstance(d["original_space"]["space"], str): + # For backward compatibility reasons, if d["original_space"]["space"] + # is a string, this original space was serialized by gym_space_to_dict. + space.original_space = gym_space_from_dict(d["original_space"]) + else: + # Otherwise, this original space was serialized by space_to_dict. + space.original_space = space_from_dict(d["original_space"]) + return space + + +@DeveloperAPI +def check_if_args_kwargs_serializable(args: Sequence[Any], kwargs: Dict[str, Any]): + """Check if parameters to a function are serializable by ray. + + Args: + args: arguments to be checked. + kwargs: keyword arguments to be checked. + + Raises: + NoteSerializable if either args are kwargs are not serializable + by ray. + """ + for arg in args: + try: + # if the object is truly serializable we should be able to + # ray.put and ray.get it. + ray.get(ray.put(arg)) + except TypeError as e: + raise NotSerializable( + "RLModule constructor arguments must be serializable. " + f"Found non-serializable argument: {arg}.\n" + f"Original serialization error: {e}" + ) + for k, v in kwargs.items(): + try: + # if the object is truly serializable we should be able to + # ray.put and ray.get it. + ray.get(ray.put(v)) + except TypeError as e: + raise NotSerializable( + "RLModule constructor arguments must be serializable. " + f"Found non-serializable keyword argument: {k} = {v}.\n" + f"Original serialization error: {e}" + ) + + +@DeveloperAPI +def serialize_type(type_: Union[Type, str]) -> str: + """Converts a type into its full classpath ([module file] + "." + [class name]). + + Args: + type_: The type to convert. + + Returns: + The full classpath of the given type, e.g. "ray.rllib.algorithms.ppo.PPOConfig". + """ + # TODO (avnishn): find a way to incorporate the tune registry here. + # Already serialized. + if isinstance(type_, str): + return type_ + + return type_.__module__ + "." + type_.__qualname__ + + +@DeveloperAPI +def deserialize_type( + module: Union[str, Type], error: bool = False +) -> Optional[Union[str, Type]]: + """Resolves a class path to a class. + If the given module is already a class, it is returned as is. + If the given module is a string, it is imported and the class is returned. + + Args: + module: The classpath (str) or type to resolve. + error: Whether to throw a ValueError if `module` could not be resolved into + a class. If False and `module` is not resolvable, returns None. + + Returns: + The resolved class or `module` (if `error` is False and no resolution possible). + + Raises: + ValueError: If `error` is True and `module` cannot be resolved. + """ + # Already a class, return as-is. + if isinstance(module, type): + return module + # A string. + elif isinstance(module, str): + # Try interpreting (as classpath) and importing the given module. + try: + module_path, class_name = module.rsplit(".", 1) + module = importlib.import_module(module_path) + return getattr(module, class_name) + # Module not found OR not a module (but a registered string?). + except (ModuleNotFoundError, ImportError, AttributeError, ValueError) as e: + # Ignore if error=False. + if error: + raise ValueError( + f"Could not deserialize the given classpath `module={module}` into " + "a valid python class! Make sure you have all necessary pip " + "packages installed and all custom modules are in your " + "`PYTHONPATH` env variable." + ) from e + else: + raise ValueError(f"`module` ({module} must be type or string (classpath)!") + + return module diff --git a/lib/python3.12/site-packages/ray/rllib/utils/sgd.py b/lib/python3.12/site-packages/ray/rllib/utils/sgd.py new file mode 100644 index 0000000000000000000000000000000000000000..cc3a72506ee5bf1dd2ceec357190846601223afd --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/sgd.py @@ -0,0 +1,137 @@ +"""Utils for minibatch SGD across multiple RLlib policies.""" + +import logging +import random + +import numpy as np + +from ray.rllib.policy.sample_batch import MultiAgentBatch, SampleBatch +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.metrics.learner_info import LearnerInfoBuilder + +logger = logging.getLogger(__name__) + + +@OldAPIStack +def standardized(array: np.ndarray): + """Normalize the values in an array. + + Args: + array (np.ndarray): Array of values to normalize. + + Returns: + array with zero mean and unit standard deviation. + """ + return (array - array.mean()) / max(1e-4, array.std()) + + +@OldAPIStack +def minibatches(samples: SampleBatch, sgd_minibatch_size: int, shuffle: bool = True): + """Return a generator yielding minibatches from a sample batch. + + Args: + samples: SampleBatch to split up. + sgd_minibatch_size: Size of minibatches to return. + shuffle: Whether to shuffle the order of the generated minibatches. + Note that in case of a non-recurrent policy, the incoming batch + is globally shuffled first regardless of this setting, before + the minibatches are generated from it! + + Yields: + SampleBatch: Each of size `sgd_minibatch_size`. + """ + if not sgd_minibatch_size: + yield samples + return + + if isinstance(samples, MultiAgentBatch): + raise NotImplementedError( + "Minibatching not implemented for multi-agent in simple mode" + ) + + if "state_in_0" not in samples and "state_out_0" not in samples: + samples.shuffle() + + all_slices = samples._get_slice_indices(sgd_minibatch_size) + data_slices, state_slices = all_slices + + if len(state_slices) == 0: + if shuffle: + random.shuffle(data_slices) + for i, j in data_slices: + yield samples[i:j] + else: + all_slices = list(zip(data_slices, state_slices)) + if shuffle: + # Make sure to shuffle data and states while linked together. + random.shuffle(all_slices) + for (i, j), (si, sj) in all_slices: + yield samples.slice(i, j, si, sj) + + +@OldAPIStack +def do_minibatch_sgd( + samples, + policies, + local_worker, + num_sgd_iter, + sgd_minibatch_size, + standardize_fields, +): + """Execute minibatch SGD. + + Args: + samples: Batch of samples to optimize. + policies: Dictionary of policies to optimize. + local_worker: Master rollout worker instance. + num_sgd_iter: Number of epochs of optimization to take. + sgd_minibatch_size: Size of minibatches to use for optimization. + standardize_fields: List of sample field names that should be + normalized prior to optimization. + + Returns: + averaged info fetches over the last SGD epoch taken. + """ + + # Handle everything as if multi-agent. + samples = samples.as_multi_agent() + + # Use LearnerInfoBuilder as a unified way to build the final + # results dict from `learn_on_loaded_batch` call(s). + # This makes sure results dicts always have the same structure + # no matter the setup (multi-GPU, multi-agent, minibatch SGD, + # tf vs torch). + learner_info_builder = LearnerInfoBuilder(num_devices=1) + for policy_id, policy in policies.items(): + if policy_id not in samples.policy_batches: + continue + + batch = samples.policy_batches[policy_id] + for field in standardize_fields: + batch[field] = standardized(batch[field]) + + # Check to make sure that the sgd_minibatch_size is not smaller + # than max_seq_len otherwise this will cause indexing errors while + # performing sgd when using a RNN or Attention model + if ( + policy.is_recurrent() + and policy.config["model"]["max_seq_len"] > sgd_minibatch_size + ): + raise ValueError( + "`sgd_minibatch_size` ({}) cannot be smaller than" + "`max_seq_len` ({}).".format( + sgd_minibatch_size, policy.config["model"]["max_seq_len"] + ) + ) + + for i in range(num_sgd_iter): + for minibatch in minibatches(batch, sgd_minibatch_size): + results = ( + local_worker.learn_on_batch( + MultiAgentBatch({policy_id: minibatch}, minibatch.count) + ) + )[policy_id] + learner_info_builder.add_learn_on_batch_results(results, policy_id) + + learner_info = learner_info_builder.finalize() + return learner_info diff --git a/lib/python3.12/site-packages/ray/rllib/utils/spaces/__init__.py b/lib/python3.12/site-packages/ray/rllib/utils/spaces/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/ray/rllib/utils/spaces/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/spaces/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..23fbb9ba587c45a43dd8073176f41184bcf0e8b7 Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/spaces/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/spaces/__pycache__/flexdict.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/spaces/__pycache__/flexdict.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..0429d371f2e3aa2cb4e8784ba6ea26ecedae4c7e Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/spaces/__pycache__/flexdict.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/spaces/__pycache__/repeated.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/spaces/__pycache__/repeated.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cd3a3a6406d796db10accb72f3426e70e65c007c Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/spaces/__pycache__/repeated.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/spaces/__pycache__/simplex.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/spaces/__pycache__/simplex.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..42604c996cae15cefd77ce2f5ed5514be4d0ecbe Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/spaces/__pycache__/simplex.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/spaces/__pycache__/space_utils.cpython-312.pyc b/lib/python3.12/site-packages/ray/rllib/utils/spaces/__pycache__/space_utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..4a48b06387600fd4400dbe4b35b8709ba91a21cb Binary files /dev/null and b/lib/python3.12/site-packages/ray/rllib/utils/spaces/__pycache__/space_utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/ray/rllib/utils/spaces/flexdict.py b/lib/python3.12/site-packages/ray/rllib/utils/spaces/flexdict.py new file mode 100644 index 0000000000000000000000000000000000000000..905a05e2857b72856cee0bd142dd3af204029c07 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/spaces/flexdict.py @@ -0,0 +1,47 @@ +import gymnasium as gym + +from ray.rllib.utils.annotations import PublicAPI + + +@PublicAPI +class FlexDict(gym.spaces.Dict): + """Gym Dictionary with arbitrary keys updatable after instantiation + + Example: + space = FlexDict({}) + space['key'] = spaces.Box(4,) + See also: documentation for gym.spaces.Dict + """ + + def __init__(self, spaces=None, **spaces_kwargs): + err = "Use either Dict(spaces=dict(...)) or Dict(foo=x, bar=z)" + assert (spaces is None) or (not spaces_kwargs), err + + if spaces is None: + spaces = spaces_kwargs + + for space in spaces.values(): + self.assertSpace(space) + + super().__init__(spaces=spaces) + + def assertSpace(self, space): + err = "Values of the dict should be instances of gym.Space" + assert issubclass(type(space), gym.spaces.Space), err + + def sample(self): + return {k: space.sample() for k, space in self.spaces.items()} + + def __getitem__(self, key): + return self.spaces[key] + + def __setitem__(self, key, space): + self.assertSpace(space) + self.spaces[key] = space + + def __repr__(self): + return ( + "FlexDict(" + + ", ".join([str(k) + ":" + str(s) for k, s in self.spaces.items()]) + + ")" + ) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/spaces/repeated.py b/lib/python3.12/site-packages/ray/rllib/utils/spaces/repeated.py new file mode 100644 index 0000000000000000000000000000000000000000..77beaff288e749b01780c9e9bf6b25e3f123b2d8 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/spaces/repeated.py @@ -0,0 +1,38 @@ +import gymnasium as gym +import numpy as np + +from ray.rllib.utils.annotations import PublicAPI + + +@PublicAPI +class Repeated(gym.Space): + """Represents a variable-length list of child spaces. + + Example: + self.observation_space = spaces.Repeated(spaces.Box(4,), max_len=10) + --> from 0 to 10 boxes of shape (4,) + + See also: documentation for rllib.models.RepeatedValues, which shows how + the lists are represented as batched input for ModelV2 classes. + """ + + def __init__(self, child_space: gym.Space, max_len: int): + super().__init__() + self.child_space = child_space + self.max_len = max_len + + def sample(self): + return [ + self.child_space.sample() + for _ in range(self.np_random.integers(1, self.max_len + 1)) + ] + + def contains(self, x): + return ( + isinstance(x, (list, np.ndarray)) + and len(x) <= self.max_len + and all(self.child_space.contains(c) for c in x) + ) + + def __repr__(self): + return "Repeated({}, {})".format(self.child_space, self.max_len) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/spaces/simplex.py b/lib/python3.12/site-packages/ray/rllib/utils/spaces/simplex.py new file mode 100644 index 0000000000000000000000000000000000000000..7e3b4c843ad0e4d0303f106b5621f88428f8b931 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/spaces/simplex.py @@ -0,0 +1,61 @@ +import gymnasium as gym +import numpy as np + +from ray.rllib.utils.annotations import PublicAPI + + +@PublicAPI +class Simplex(gym.Space): + """Represents a d - 1 dimensional Simplex in R^d. + + That is, all coordinates are in [0, 1] and sum to 1. + The dimension d of the simplex is assumed to be shape[-1]. + + Additionally one can specify the underlying distribution of + the simplex as a Dirichlet distribution by providing concentration + parameters. By default, sampling is uniform, i.e. concentration is + all 1s. + + Example usage: + self.action_space = spaces.Simplex(shape=(3, 4)) + --> 3 independent 4d Dirichlet with uniform concentration + """ + + def __init__(self, shape, concentration=None, dtype=np.float32): + assert type(shape) in [tuple, list] + + super().__init__(shape, dtype) + self.dim = self.shape[-1] + + if concentration is not None: + assert ( + concentration.shape[0] == shape[-1] + ), f"{concentration.shape[0]} vs {shape[-1]}" + self.concentration = concentration + else: + self.concentration = np.array([1] * self.dim) + + def sample(self): + return np.random.dirichlet(self.concentration, size=self.shape[:-1]).astype( + self.dtype + ) + + def contains(self, x): + return x.shape == self.shape and np.allclose( + np.sum(x, axis=-1), np.ones_like(x[..., 0]) + ) + + def to_jsonable(self, sample_n): + return np.array(sample_n).tolist() + + def from_jsonable(self, sample_n): + return [np.asarray(sample) for sample in sample_n] + + def __repr__(self): + return "Simplex({}; {})".format(self.shape, self.concentration) + + def __eq__(self, other): + return ( + np.allclose(self.concentration, other.concentration) + and self.shape == other.shape + ) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/spaces/space_utils.py b/lib/python3.12/site-packages/ray/rllib/utils/spaces/space_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..d7afaeffb306a52e4ecf172efc617d38238d1e88 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/spaces/space_utils.py @@ -0,0 +1,513 @@ +from typing import Any, List, Optional, Union + +import gymnasium as gym +import numpy as np +import tree # pip install dm_tree + +from ray.rllib.utils.annotations import DeveloperAPI + + +@DeveloperAPI +class BatchedNdArray(np.ndarray): + """A ndarray-wrapper the usage of which indicates that there a batch dim exists. + + This is such that our `batch()` utility can distinguish between having to + stack n individual batch items (each one w/o any batch dim) vs having to + concatenate n already batched items (each one possibly with a different batch + dim, but definitely with some batch dim). + + TODO (sven): Maybe replace this by a list-override instead. + """ + + def __new__(cls, input_array): + # Use __new__ to create a new instance of our subclass. + obj = np.asarray(input_array).view(cls) + return obj + + +@DeveloperAPI +def get_original_space(space: gym.Space) -> gym.Space: + """Returns the original space of a space, if any. + + This function recursively traverses the given space and returns the original space + at the very end of the chain. + + Args: + space: The space to get the original space for. + + Returns: + The original space or the given space itself if no original space is found. + """ + if hasattr(space, "original_space"): + return get_original_space(space.original_space) + else: + return space + + +@DeveloperAPI +def is_composite_space(space: gym.Space) -> bool: + """Returns true, if the space is composite. + + Note, we follow here the glossary of `gymnasium` by which any spoace + that holds other spaces is defined as being 'composite'. + + Args: + space: The space to be checked for being composed of other spaces. + + Returns: + True, if the space is composed of other spaces, otherwise False. + """ + if type(space) in [ + gym.spaces.Dict, + gym.spaces.Graph, + gym.spaces.Sequence, + gym.spaces.Tuple, + ]: + return True + else: + return False + + +@DeveloperAPI +def flatten_space(space: gym.Space) -> List[gym.Space]: + """Flattens a gym.Space into its primitive components. + + Primitive components are any non Tuple/Dict spaces. + + Args: + space: The gym.Space to flatten. This may be any + supported type (including nested Tuples and Dicts). + + Returns: + List[gym.Space]: The flattened list of primitive Spaces. This list + does not contain Tuples or Dicts anymore. + """ + + def _helper_flatten(space_, return_list): + from ray.rllib.utils.spaces.flexdict import FlexDict + + if isinstance(space_, gym.spaces.Tuple): + for s in space_: + _helper_flatten(s, return_list) + elif isinstance(space_, (gym.spaces.Dict, FlexDict)): + for k in sorted(space_.spaces): + _helper_flatten(space_[k], return_list) + else: + return_list.append(space_) + + ret = [] + _helper_flatten(space, ret) + return ret + + +@DeveloperAPI +def get_base_struct_from_space(space): + """Returns a Tuple/Dict Space as native (equally structured) py tuple/dict. + + Args: + space: The Space to get the python struct for. + + Returns: + Union[dict,tuple,gym.Space]: The struct equivalent to the given Space. + Note that the returned struct still contains all original + "primitive" Spaces (e.g. Box, Discrete). + + .. testcode:: + :skipif: True + + get_base_struct_from_space(Dict({ + "a": Box(), + "b": Tuple([Discrete(2), Discrete(3)]) + })) + + .. testoutput:: + + dict(a=Box(), b=tuple(Discrete(2), Discrete(3))) + """ + + def _helper_struct(space_): + if isinstance(space_, gym.spaces.Tuple): + return tuple(_helper_struct(s) for s in space_) + elif isinstance(space_, gym.spaces.Dict): + return {k: _helper_struct(space_[k]) for k in space_.spaces} + else: + return space_ + + return _helper_struct(space) + + +@DeveloperAPI +def get_dummy_batch_for_space( + space: gym.Space, + batch_size: int = 32, + *, + fill_value: Union[float, int, str] = 0.0, + time_size: Optional[int] = None, + time_major: bool = False, + one_hot_discrete: bool = False, +) -> np.ndarray: + """Returns batched dummy data (using `batch_size`) for the given `space`. + + Note: The returned batch will not pass a `space.contains(batch)` test + as an additional batch dimension has to be added at axis 0, unless `batch_size` is + set to 0. + + Args: + space: The space to get a dummy batch for. + batch_size: The required batch size (B). Note that this can also + be 0 (only if `time_size` is None!), which will result in a + non-batched sample for the given space (no batch dim). + fill_value: The value to fill the batch with + or "random" for random values. + time_size: If not None, add an optional time axis + of `time_size` size to the returned batch. This time axis might either + be inserted at axis=1 (default) or axis=0, if `time_major` is True. + time_major: If True AND `time_size` is not None, return batch + as shape [T x B x ...], otherwise as [B x T x ...]. If `time_size` + if None, ignore this setting and return [B x ...]. + one_hot_discrete: If True, will return one-hot vectors (instead of + int-values) for those sub-components of a (possibly complex) `space` + that are Discrete or MultiDiscrete. Note that in case `fill_value` is 0.0, + this will result in zero-hot vectors (where all slots have a value of 0.0). + + Returns: + The dummy batch of size `bqtch_size` matching the given space. + """ + # Complex spaces. Perform recursive calls of this function. + if isinstance(space, (gym.spaces.Dict, gym.spaces.Tuple, dict, tuple)): + base_struct = space + if isinstance(space, (gym.spaces.Dict, gym.spaces.Tuple)): + base_struct = get_base_struct_from_space(space) + return tree.map_structure( + lambda s: get_dummy_batch_for_space( + space=s, + batch_size=batch_size, + fill_value=fill_value, + time_size=time_size, + time_major=time_major, + one_hot_discrete=one_hot_discrete, + ), + base_struct, + ) + + if one_hot_discrete: + if isinstance(space, gym.spaces.Discrete): + space = gym.spaces.Box(0.0, 1.0, (space.n,), np.float32) + elif isinstance(space, gym.spaces.MultiDiscrete): + space = gym.spaces.Box(0.0, 1.0, (np.sum(space.nvec),), np.float32) + + # Primitive spaces: Box, Discrete, MultiDiscrete. + # Random values: Use gym's sample() method. + if fill_value == "random": + if time_size is not None: + assert batch_size > 0 and time_size > 0 + if time_major: + return np.array( + [ + [space.sample() for _ in range(batch_size)] + for t in range(time_size) + ], + dtype=space.dtype, + ) + else: + return np.array( + [ + [space.sample() for t in range(time_size)] + for _ in range(batch_size) + ], + dtype=space.dtype, + ) + else: + return np.array( + [space.sample() for _ in range(batch_size)] + if batch_size > 0 + else space.sample(), + dtype=space.dtype, + ) + # Fill value given: Use np.full. + else: + if time_size is not None: + assert batch_size > 0 and time_size > 0 + if time_major: + shape = [time_size, batch_size] + else: + shape = [batch_size, time_size] + else: + shape = [batch_size] if batch_size > 0 else [] + return np.full( + shape + list(space.shape), fill_value=fill_value, dtype=space.dtype + ) + + +@DeveloperAPI +def flatten_to_single_ndarray(input_): + """Returns a single np.ndarray given a list/tuple of np.ndarrays. + + Args: + input_ (Union[List[np.ndarray], np.ndarray]): The list of ndarrays or + a single ndarray. + + Returns: + np.ndarray: The result after concatenating all single arrays in input_. + + .. testcode:: + :skipif: True + + flatten_to_single_ndarray([ + np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]), + np.array([7, 8, 9]), + ]) + + .. testoutput:: + + np.array([ + 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0 + ]) + """ + # Concatenate complex inputs. + if isinstance(input_, (list, tuple, dict)): + expanded = [] + for in_ in tree.flatten(input_): + expanded.append(np.reshape(in_, [-1])) + input_ = np.concatenate(expanded, axis=0).flatten() + return input_ + + +@DeveloperAPI +def batch( + list_of_structs: List[Any], + *, + individual_items_already_have_batch_dim: Union[bool, str] = False, +): + """Converts input from a list of (nested) structs to a (nested) struct of batches. + + Input: List of structs (each of these structs representing a single batch item). + [ + {"a": 1, "b": (4, 7.0)}, <- batch item 1 + {"a": 2, "b": (5, 8.0)}, <- batch item 2 + {"a": 3, "b": (6, 9.0)}, <- batch item 3 + ] + + Output: Struct of different batches (each batch has size=3 b/c there were 3 items + in the original list): + { + "a": np.array([1, 2, 3]), + "b": (np.array([4, 5, 6]), np.array([7.0, 8.0, 9.0])) + } + + Args: + list_of_structs: The list of (possibly nested) structs. Each item + in this list represents a single batch item. + individual_items_already_have_batch_dim: True, if the individual items in + `list_of_structs` already have a batch dim. In this case, we will + concatenate (instead of stack) at the end. In the example above, this would + look like this: Input: [{"a": [1], "b": ([4], [7.0])}, ...] -> Output: same + as in above example. + If the special value "auto" is used, + + Returns: + The struct of component batches. Each leaf item in this struct represents the + batch for a single component (in case struct is tuple/dict). If the input is a + simple list of primitive items, e.g. a list of floats, a np.array of floats + will be returned. + """ + if not list_of_structs: + raise ValueError("Input `list_of_structs` does not contain any items.") + + # TODO (sven): Maybe replace this by a list-override (usage of which indicated + # this method that concatenate should be used (not stack)). + if individual_items_already_have_batch_dim == "auto": + flat = tree.flatten(list_of_structs[0]) + individual_items_already_have_batch_dim = isinstance(flat[0], BatchedNdArray) + + np_func = np.concatenate if individual_items_already_have_batch_dim else np.stack + ret = tree.map_structure( + lambda *s: np.ascontiguousarray(np_func(s, axis=0)), *list_of_structs + ) + return ret + + +@DeveloperAPI +def unbatch(batches_struct): + """Converts input from (nested) struct of batches to batch of structs. + + Input: Struct of different batches (each batch has size=3): + { + "a": np.array([1, 2, 3]), + "b": (np.array([4, 5, 6]), np.array([7.0, 8.0, 9.0])) + } + Output: Batch (list) of structs (each of these structs representing a + single action): + [ + {"a": 1, "b": (4, 7.0)}, <- action 1 + {"a": 2, "b": (5, 8.0)}, <- action 2 + {"a": 3, "b": (6, 9.0)}, <- action 3 + ] + + Args: + batches_struct: The struct of component batches. Each leaf item + in this struct represents the batch for a single component + (in case struct is tuple/dict). + Alternatively, `batches_struct` may also simply be a batch of + primitives (non tuple/dict). + + Returns: + The list of individual structs. Each item in the returned list represents a + single (maybe complex) batch item. + """ + flat_batches = tree.flatten(batches_struct) + + out = [] + for batch_pos in range(len(flat_batches[0])): + out.append( + tree.unflatten_as( + batches_struct, + [flat_batches[i][batch_pos] for i in range(len(flat_batches))], + ) + ) + return out + + +@DeveloperAPI +def clip_action(action, action_space): + """Clips all components in `action` according to the given Space. + + Only applies to Box components within the action space. + + Args: + action: The action to be clipped. This could be any complex + action, e.g. a dict or tuple. + action_space: The action space struct, + e.g. `{"a": Distrete(2)}` for a space: Dict({"a": Discrete(2)}). + + Returns: + Any: The input action, but clipped by value according to the space's + bounds. + """ + + def map_(a, s): + if isinstance(s, gym.spaces.Box): + a = np.clip(a, s.low, s.high) + return a + + return tree.map_structure(map_, action, action_space) + + +@DeveloperAPI +def unsquash_action(action, action_space_struct): + """Unsquashes all components in `action` according to the given Space. + + Inverse of `normalize_action()`. Useful for mapping policy action + outputs (normalized between -1.0 and 1.0) to an env's action space. + Unsquashing results in cont. action component values between the + given Space's bounds (`low` and `high`). This only applies to Box + components within the action space, whose dtype is float32 or float64. + + Args: + action: The action to be unsquashed. This could be any complex + action, e.g. a dict or tuple. + action_space_struct: The action space struct, + e.g. `{"a": Box()}` for a space: Dict({"a": Box()}). + + Returns: + Any: The input action, but unsquashed, according to the space's + bounds. An unsquashed action is ready to be sent to the + environment (`BaseEnv.send_actions([unsquashed actions])`). + """ + + def map_(a, s): + if ( + isinstance(s, gym.spaces.Box) + and np.all(s.bounded_below) + and np.all(s.bounded_above) + ): + if s.dtype == np.float32 or s.dtype == np.float64: + # Assuming values are roughly between -1.0 and 1.0 -> + # unsquash them to the given bounds. + a = s.low + (a + 1.0) * (s.high - s.low) / 2.0 + # Clip to given bounds, just in case the squashed values were + # outside [-1.0, 1.0]. + a = np.clip(a, s.low, s.high) + elif np.issubdtype(s.dtype, np.integer): + # For Categorical and MultiCategorical actions, shift the selection + # into the proper range. + a = s.low + a + return a + + return tree.map_structure(map_, action, action_space_struct) + + +@DeveloperAPI +def normalize_action(action, action_space_struct): + """Normalizes all (Box) components in `action` to be in [-1.0, 1.0]. + + Inverse of `unsquash_action()`. Useful for mapping an env's action + (arbitrary bounded values) to a [-1.0, 1.0] interval. + This only applies to Box components within the action space, whose + dtype is float32 or float64. + + Args: + action: The action to be normalized. This could be any complex + action, e.g. a dict or tuple. + action_space_struct: The action space struct, + e.g. `{"a": Box()}` for a space: Dict({"a": Box()}). + + Returns: + Any: The input action, but normalized, according to the space's + bounds. + """ + + def map_(a, s): + if isinstance(s, gym.spaces.Box) and ( + s.dtype == np.float32 or s.dtype == np.float64 + ): + # Normalize values to be exactly between -1.0 and 1.0. + a = ((a - s.low) * 2.0) / (s.high - s.low) - 1.0 + return a + + return tree.map_structure(map_, action, action_space_struct) + + +@DeveloperAPI +def convert_element_to_space_type(element: Any, sampled_element: Any) -> Any: + """Convert all the components of the element to match the space dtypes. + + Args: + element: The element to be converted. + sampled_element: An element sampled from a space to be matched + to. + + Returns: + The input element, but with all its components converted to match + the space dtypes. + """ + + def map_(elem, s): + if isinstance(s, np.ndarray): + if not isinstance(elem, np.ndarray): + assert isinstance( + elem, (float, int) + ), f"ERROR: `elem` ({elem}) must be np.array, float or int!" + if s.shape == (): + elem = np.array(elem, dtype=s.dtype) + else: + raise ValueError( + "Element should be of type np.ndarray but is instead of \ + type {}".format( + type(elem) + ) + ) + elif s.dtype != elem.dtype: + elem = elem.astype(s.dtype) + + # Gymnasium now uses np.int_64 as the dtype of a Discrete action space + elif isinstance(s, int) or isinstance(s, np.int_): + if isinstance(elem, float) and elem.is_integer(): + elem = int(elem) + # Note: This does not check if the float element is actually an integer + if isinstance(elem, np.float_): + elem = np.int64(elem) + + return elem + + return tree.map_structure(map_, element, sampled_element, check_types=False) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/tensor_dtype.py b/lib/python3.12/site-packages/ray/rllib/utils/tensor_dtype.py new file mode 100644 index 0000000000000000000000000000000000000000..b590833c133a5dfca34ed6f753a5a4d38d2b6112 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/tensor_dtype.py @@ -0,0 +1,65 @@ +import numpy as np + +from ray.rllib.utils.framework import try_import_tf, try_import_torch +from ray.rllib.utils.typing import TensorType +from ray.util.annotations import DeveloperAPI + +torch, _ = try_import_torch() +_, tf, _ = try_import_tf() + + +# Dict of NumPy dtype -> torch dtype +if torch: + numpy_to_torch_dtype_dict = { + np.bool_: torch.bool, + np.uint8: torch.uint8, + np.int8: torch.int8, + np.int16: torch.int16, + np.int32: torch.int32, + np.int64: torch.int64, + np.float16: torch.float16, + np.float32: torch.float32, + np.float64: torch.float64, + np.complex64: torch.complex64, + np.complex128: torch.complex128, + } +else: + numpy_to_torch_dtype_dict = {} + +# Dict of NumPy dtype -> tf dtype +if tf: + numpy_to_tf_dtype_dict = { + np.bool_: tf.bool, + np.uint8: tf.uint8, + np.int8: tf.int8, + np.int16: tf.int16, + np.int32: tf.int32, + np.int64: tf.int64, + np.float16: tf.float16, + np.float32: tf.float32, + np.float64: tf.float64, + np.complex64: tf.complex64, + np.complex128: tf.complex128, + } +else: + numpy_to_tf_dtype_dict = {} + +# Dict of torch dtype -> NumPy dtype +torch_to_numpy_dtype_dict = { + value: key for (key, value) in numpy_to_torch_dtype_dict.items() +} +# Dict of tf dtype -> NumPy dtype +tf_to_numpy_dtype_dict = {value: key for (key, value) in numpy_to_tf_dtype_dict.items()} + + +@DeveloperAPI +def get_np_dtype(x: TensorType) -> np.dtype: + """Returns the NumPy dtype of the given tensor or array.""" + if torch and isinstance(x, torch.Tensor): + return torch_to_numpy_dtype_dict[x.dtype] + if tf and isinstance(x, tf.Tensor): + return tf_to_numpy_dtype_dict[x.dtype] + elif isinstance(x, np.ndarray): + return x.dtype + else: + raise TypeError("Unsupported type: {}".format(type(x))) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/test_utils.py b/lib/python3.12/site-packages/ray/rllib/utils/test_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..35e55f9fc9c9dec8e5b87894e9eefede3b5d99a2 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/test_utils.py @@ -0,0 +1,1928 @@ +import argparse +import json +import logging +import os +import pprint +import random +import re +import time +from typing import ( + TYPE_CHECKING, + Any, + Dict, + List, + Optional, + Tuple, + Type, + Union, +) + +import gymnasium as gym +import numpy as np +import tree # pip install dm_tree +from gymnasium.spaces import ( + Box, + Dict as GymDict, + Discrete, + MultiBinary, + MultiDiscrete, + Tuple as GymTuple, +) + +import ray +from ray import tune +from ray.air.integrations.wandb import WANDB_ENV_VAR, WandbLoggerCallback +from ray.rllib.core import DEFAULT_MODULE_ID, Columns +from ray.rllib.env.wrappers.atari_wrappers import is_atari, wrap_deepmind +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.error import UnsupportedSpaceException +from ray.rllib.utils.framework import try_import_jax, try_import_tf, try_import_torch +from ray.rllib.utils.metrics import ( + DIFF_NUM_GRAD_UPDATES_VS_SAMPLER_POLICY, + ENV_RUNNER_RESULTS, + EPISODE_RETURN_MEAN, + EVALUATION_RESULTS, + NUM_ENV_STEPS_SAMPLED_LIFETIME, + NUM_ENV_STEPS_TRAINED, +) +from ray.rllib.utils.typing import ResultDict +from ray.tune import CLIReporter +from ray.tune.result import TRAINING_ITERATION + +if TYPE_CHECKING: + from ray.rllib.algorithms import Algorithm, AlgorithmConfig + from ray.rllib.offline.dataset_reader import DatasetReader + +jax, _ = try_import_jax() +tf1, tf, tfv = try_import_tf() +torch, _ = try_import_torch() + +logger = logging.getLogger(__name__) + + +def add_rllib_example_script_args( + parser: Optional[argparse.ArgumentParser] = None, + default_reward: float = 100.0, + default_iters: int = 200, + default_timesteps: int = 100000, +) -> argparse.ArgumentParser: + """Adds RLlib-typical (and common) examples scripts command line args to a parser. + + TODO (sven): This function should be used by most of our examples scripts, which + already mostly have this logic in them (but written out). + + Args: + parser: The parser to add the arguments to. If None, create a new one. + default_reward: The default value for the --stop-reward option. + default_iters: The default value for the --stop-iters option. + default_timesteps: The default value for the --stop-timesteps option. + + Returns: + The altered (or newly created) parser object. + """ + if parser is None: + parser = argparse.ArgumentParser() + + # Algo and Algo config options. + parser.add_argument( + "--algo", type=str, default="PPO", help="The RLlib-registered algorithm to use." + ) + parser.add_argument( + "--framework", + choices=["tf", "tf2", "torch"], + default="torch", + help="The DL framework specifier.", + ) + parser.add_argument( + "--env", + type=str, + default=None, + help="The gym.Env identifier to run the experiment with.", + ) + parser.add_argument( + "--num-env-runners", + type=int, + default=None, + help="The number of (remote) EnvRunners to use for the experiment.", + ) + parser.add_argument( + "--num-envs-per-env-runner", + type=int, + default=None, + help="The number of (vectorized) environments per EnvRunner. Note that " + "this is identical to the batch size for (inference) action computations.", + ) + parser.add_argument( + "--num-agents", + type=int, + default=0, + help="If 0 (default), will run as single-agent. If > 0, will run as " + "multi-agent with the environment simply cloned n times and each agent acting " + "independently at every single timestep. The overall reward for this " + "experiment is then the sum over all individual agents' rewards.", + ) + + # Evaluation options. + parser.add_argument( + "--evaluation-num-env-runners", + type=int, + default=0, + help="The number of evaluation (remote) EnvRunners to use for the experiment.", + ) + parser.add_argument( + "--evaluation-interval", + type=int, + default=0, + help="Every how many iterations to run one round of evaluation. " + "Use 0 (default) to disable evaluation.", + ) + parser.add_argument( + "--evaluation-duration", + type=lambda v: v if v == "auto" else int(v), + default=10, + help="The number of evaluation units to run each evaluation round. " + "Use `--evaluation-duration-unit` to count either in 'episodes' " + "or 'timesteps'. If 'auto', will run as many as possible during train pass (" + "`--evaluation-parallel-to-training` must be set then).", + ) + parser.add_argument( + "--evaluation-duration-unit", + type=str, + default="episodes", + choices=["episodes", "timesteps"], + help="The evaluation duration unit to count by. One of 'episodes' or " + "'timesteps'. This unit will be run `--evaluation-duration` times in each " + "evaluation round. If `--evaluation-duration=auto`, this setting does not " + "matter.", + ) + parser.add_argument( + "--evaluation-parallel-to-training", + action="store_true", + help="Whether to run evaluation parallel to training. This might help speed up " + "your overall iteration time. Be aware that when using this option, your " + "reported evaluation results are referring to one iteration before the current " + "one.", + ) + + # RLlib logging options. + parser.add_argument( + "--output", + type=str, + default=None, + help="The output directory to write trajectories to, which are collected by " + "the algo's EnvRunners.", + ) + parser.add_argument( + "--log-level", + type=str, + default=None, # None -> use default + choices=["INFO", "DEBUG", "WARN", "ERROR"], + help="The log-level to be used by the RLlib logger.", + ) + + # tune.Tuner options. + parser.add_argument( + "--no-tune", + action="store_true", + help="Whether to NOT use tune.Tuner(), but rather a simple for-loop calling " + "`algo.train()` repeatedly until one of the stop criteria is met.", + ) + parser.add_argument( + "--num-samples", + type=int, + default=1, + help="How many (tune.Tuner.fit()) experiments to execute - if possible in " + "parallel.", + ) + parser.add_argument( + "--max-concurrent-trials", + type=int, + default=None, + help="How many (tune.Tuner) trials to run concurrently.", + ) + parser.add_argument( + "--verbose", + type=int, + default=2, + help="The verbosity level for the `tune.Tuner()` running the experiment.", + ) + parser.add_argument( + "--checkpoint-freq", + type=int, + default=0, + help=( + "The frequency (in training iterations) with which to create checkpoints. " + "Note that if --wandb-key is provided, all checkpoints will " + "automatically be uploaded to WandB." + ), + ) + parser.add_argument( + "--checkpoint-at-end", + action="store_true", + help=( + "Whether to create a checkpoint at the very end of the experiment. " + "Note that if --wandb-key is provided, all checkpoints will " + "automatically be uploaded to WandB." + ), + ) + + # WandB logging options. + parser.add_argument( + "--wandb-key", + type=str, + default=None, + help="The WandB API key to use for uploading results.", + ) + parser.add_argument( + "--wandb-project", + type=str, + default=None, + help="The WandB project name to use.", + ) + parser.add_argument( + "--wandb-run-name", + type=str, + default=None, + help="The WandB run name to use.", + ) + + # Experiment stopping and testing criteria. + parser.add_argument( + "--stop-reward", + type=float, + default=default_reward, + help="Reward at which the script should stop training.", + ) + parser.add_argument( + "--stop-iters", + type=int, + default=default_iters, + help="The number of iterations to train.", + ) + parser.add_argument( + "--stop-timesteps", + type=int, + default=default_timesteps, + help="The number of (environment sampling) timesteps to train.", + ) + parser.add_argument( + "--as-test", + action="store_true", + help="Whether this script should be run as a test. If set, --stop-reward must " + "be achieved within --stop-timesteps AND --stop-iters, otherwise this " + "script will throw an exception at the end.", + ) + parser.add_argument( + "--as-release-test", + action="store_true", + help="Whether this script should be run as a release test. If set, " + "all that applies to the --as-test option is true, plus, a short JSON summary " + "will be written into a results file whose location is given by the ENV " + "variable `TEST_OUTPUT_JSON`.", + ) + + # Learner scaling options. + parser.add_argument( + "--num-learners", + type=int, + default=None, + help="The number of Learners to use. If `None`, use the algorithm's default " + "value.", + ) + parser.add_argument( + "--num-cpus-per-learner", + type=float, + default=None, + help="The number of CPUs per Learner to use. If `None`, use the algorithm's " + "default value.", + ) + parser.add_argument( + "--num-gpus-per-learner", + type=float, + default=None, + help="The number of GPUs per Learner to use. If `None` and there are enough " + "GPUs for all required Learners (--num-learners), use a value of 1, " + "otherwise 0.", + ) + parser.add_argument( + "--num-aggregator-actors-per-learner", + type=int, + default=None, + help="The number of Aggregator actors to use per Learner. If `None`, use the " + "algorithm's default value.", + ) + + # Ray init options. + parser.add_argument("--num-cpus", type=int, default=0) + parser.add_argument( + "--local-mode", + action="store_true", + help="Init Ray in local mode for easier debugging.", + ) + + # Old API stack: config.num_gpus. + parser.add_argument( + "--num-gpus", + type=int, + default=None, + help="The number of GPUs to use (only on the old API stack).", + ) + parser.add_argument( + "--old-api-stack", + action="store_true", + help="Run this script on the old API stack of RLlib.", + ) + + # Deprecated options. Throws error when still used. Use `--old-api-stack` for + # disabling the new API stack. + parser.add_argument( + "--enable-new-api-stack", + action="store_true", + ) + + return parser + + +def check(x, y, decimals=5, atol=None, rtol=None, false=False): + """ + Checks two structures (dict, tuple, list, + np.array, float, int, etc..) for (almost) numeric identity. + All numbers in the two structures have to match up to `decimal` digits + after the floating point. Uses assertions. + + Args: + x: The value to be compared (to the expectation: `y`). This + may be a Tensor. + y: The expected value to be compared to `x`. This must not + be a tf-Tensor, but may be a tf/torch-Tensor. + decimals: The number of digits after the floating point up to + which all numeric values have to match. + atol: Absolute tolerance of the difference between x and y + (overrides `decimals` if given). + rtol: Relative tolerance of the difference between x and y + (overrides `decimals` if given). + false: Whether to check that x and y are NOT the same. + """ + # A dict type. + if isinstance(x, dict): + assert isinstance(y, dict), "ERROR: If x is dict, y needs to be a dict as well!" + y_keys = set(x.keys()) + for key, value in x.items(): + assert key in y, f"ERROR: y does not have x's key='{key}'! y={y}" + check(value, y[key], decimals=decimals, atol=atol, rtol=rtol, false=false) + y_keys.remove(key) + assert not y_keys, "ERROR: y contains keys ({}) that are not in x! y={}".format( + list(y_keys), y + ) + # A tuple type. + elif isinstance(x, (tuple, list)): + assert isinstance( + y, (tuple, list) + ), "ERROR: If x is tuple/list, y needs to be a tuple/list as well!" + assert len(y) == len( + x + ), "ERROR: y does not have the same length as x ({} vs {})!".format( + len(y), len(x) + ) + for i, value in enumerate(x): + check(value, y[i], decimals=decimals, atol=atol, rtol=rtol, false=false) + # Boolean comparison. + elif isinstance(x, (np.bool_, bool)): + if false is True: + assert bool(x) is not bool(y), f"ERROR: x ({x}) is y ({y})!" + else: + assert bool(x) is bool(y), f"ERROR: x ({x}) is not y ({y})!" + # Nones or primitives (excluding int vs float, which should be compared with + # tolerance/decimals as well). + elif ( + x is None + or y is None + or isinstance(x, str) + or (isinstance(x, int) and isinstance(y, int)) + ): + if false is True: + assert x != y, f"ERROR: x ({x}) is the same as y ({y})!" + else: + assert x == y, f"ERROR: x ({x}) is not the same as y ({y})!" + # String/byte comparisons. + elif ( + hasattr(x, "dtype") and (x.dtype == object or str(x.dtype).startswith(" raise error (not expected to be equal). + if false is True: + assert False, f"ERROR: x ({x}) is the same as y ({y})!" + + # Using atol/rtol. + else: + # Provide defaults for either one of atol/rtol. + if atol is None: + atol = 0 + if rtol is None: + rtol = 1e-7 + try: + np.testing.assert_allclose(x, y, atol=atol, rtol=rtol) + except AssertionError as e: + if false is False: + raise e + else: + if false is True: + assert False, f"ERROR: x ({x}) is the same as y ({y})!" + + +def check_compute_single_action( + algorithm, include_state=False, include_prev_action_reward=False +): + """Tests different combinations of args for algorithm.compute_single_action. + + Args: + algorithm: The Algorithm object to test. + include_state: Whether to include the initial state of the Policy's + Model in the `compute_single_action` call. + include_prev_action_reward: Whether to include the prev-action and + -reward in the `compute_single_action` call. + + Raises: + ValueError: If anything unexpected happens. + """ + # Have to import this here to avoid circular dependency. + from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID, SampleBatch + + # Some Algorithms may not abide to the standard API. + pid = DEFAULT_POLICY_ID + try: + # Multi-agent: Pick any learnable policy (or DEFAULT_POLICY if it's the only + # one). + pid = next(iter(algorithm.env_runner.get_policies_to_train())) + pol = algorithm.get_policy(pid) + except AttributeError: + pol = algorithm.policy + # Get the policy's model. + model = pol.model + + action_space = pol.action_space + + def _test( + what, method_to_test, obs_space, full_fetch, explore, timestep, unsquash, clip + ): + call_kwargs = {} + if what is algorithm: + call_kwargs["full_fetch"] = full_fetch + call_kwargs["policy_id"] = pid + + obs = obs_space.sample() + if isinstance(obs_space, Box): + obs = np.clip(obs, -1.0, 1.0) + state_in = None + if include_state: + state_in = model.get_initial_state() + if not state_in: + state_in = [] + i = 0 + while f"state_in_{i}" in model.view_requirements: + state_in.append( + model.view_requirements[f"state_in_{i}"].space.sample() + ) + i += 1 + action_in = action_space.sample() if include_prev_action_reward else None + reward_in = 1.0 if include_prev_action_reward else None + + if method_to_test == "input_dict": + assert what is pol + + input_dict = {SampleBatch.OBS: obs} + if include_prev_action_reward: + input_dict[SampleBatch.PREV_ACTIONS] = action_in + input_dict[SampleBatch.PREV_REWARDS] = reward_in + if state_in: + if what.config.get("enable_rl_module_and_learner", False): + input_dict["state_in"] = state_in + else: + for i, s in enumerate(state_in): + input_dict[f"state_in_{i}"] = s + input_dict_batched = SampleBatch( + tree.map_structure(lambda s: np.expand_dims(s, 0), input_dict) + ) + action = pol.compute_actions_from_input_dict( + input_dict=input_dict_batched, + explore=explore, + timestep=timestep, + **call_kwargs, + ) + # Unbatch everything to be able to compare against single + # action below. + # ARS and ES return action batches as lists. + if isinstance(action[0], list): + action = (np.array(action[0]), action[1], action[2]) + action = tree.map_structure(lambda s: s[0], action) + + try: + action2 = pol.compute_single_action( + input_dict=input_dict, + explore=explore, + timestep=timestep, + **call_kwargs, + ) + # Make sure these are the same, unless we have exploration + # switched on (or noisy layers). + if not explore and not pol.config.get("noisy"): + check(action, action2) + except TypeError: + pass + else: + action = what.compute_single_action( + obs, + state_in, + prev_action=action_in, + prev_reward=reward_in, + explore=explore, + timestep=timestep, + unsquash_action=unsquash, + clip_action=clip, + **call_kwargs, + ) + + state_out = None + if state_in or full_fetch or what is pol: + action, state_out, _ = action + if state_out: + for si, so in zip(tree.flatten(state_in), tree.flatten(state_out)): + if tf.is_tensor(si): + # If si is a tensor of Dimensions, we need to convert it + # We expect this to be the case for TF RLModules who's initial + # states are Tf Tensors. + si_shape = si.shape.as_list() + else: + si_shape = list(si.shape) + check(si_shape, so.shape) + + if unsquash is None: + unsquash = what.config["normalize_actions"] + if clip is None: + clip = what.config["clip_actions"] + + # Test whether unsquash/clipping works on the Algorithm's + # compute_single_action method: Both flags should force the action + # to be within the space's bounds. + if method_to_test == "single" and what == algorithm: + if not action_space.contains(action) and ( + clip or unsquash or not isinstance(action_space, Box) + ): + raise ValueError( + f"Returned action ({action}) of algorithm/policy {what} " + f"not in Env's action_space {action_space}" + ) + # We are operating in normalized space: Expect only smaller action + # values. + if ( + isinstance(action_space, Box) + and not unsquash + and what.config.get("normalize_actions") + and np.any(np.abs(action) > 15.0) + ): + raise ValueError( + f"Returned action ({action}) of algorithm/policy {what} " + "should be in normalized space, but seems too large/small " + "for that!" + ) + + # Loop through: Policy vs Algorithm; Different API methods to calculate + # actions; unsquash option; clip option; full fetch or not. + for what in [pol, algorithm]: + if what is algorithm: + # Get the obs-space from Workers.env (not Policy) due to possible + # pre-processor up front. + worker_set = getattr(algorithm, "env_runner_group", None) + assert worker_set + if not worker_set.local_env_runner: + obs_space = algorithm.get_policy(pid).observation_space + else: + obs_space = worker_set.local_env_runner.for_policy( + lambda p: p.observation_space, policy_id=pid + ) + obs_space = getattr(obs_space, "original_space", obs_space) + else: + obs_space = pol.observation_space + + for method_to_test in ["single"] + (["input_dict"] if what is pol else []): + for explore in [True, False]: + for full_fetch in [False, True] if what is algorithm else [False]: + timestep = random.randint(0, 100000) + for unsquash in [True, False, None]: + for clip in [False] if unsquash else [True, False, None]: + print("-" * 80) + print(f"what={what}") + print(f"method_to_test={method_to_test}") + print(f"explore={explore}") + print(f"full_fetch={full_fetch}") + print(f"unsquash={unsquash}") + print(f"clip={clip}") + _test( + what, + method_to_test, + obs_space, + full_fetch, + explore, + timestep, + unsquash, + clip, + ) + + +def check_inference_w_connectors(policy, env_name, max_steps: int = 100): + """Checks whether the given policy can infer actions from an env with connectors. + + Args: + policy: The policy to check. + env_name: Name of the environment to check + max_steps: The maximum number of steps to run the environment for. + + Raises: + ValueError: If the policy cannot infer actions from the environment. + """ + # Avoids circular import + from ray.rllib.utils.policy import local_policy_inference + + env = gym.make(env_name) + + # Potentially wrap the env like we do in RolloutWorker + if is_atari(env): + env = wrap_deepmind( + env, + dim=policy.config["model"]["dim"], + framestack=policy.config["model"].get("framestack"), + ) + + obs, info = env.reset() + reward, terminated, truncated = 0.0, False, False + ts = 0 + while not terminated and not truncated and ts < max_steps: + action_out = local_policy_inference( + policy, + env_id=0, + agent_id=0, + obs=obs, + reward=reward, + terminated=terminated, + truncated=truncated, + info=info, + ) + obs, reward, terminated, truncated, info = env.step(action_out[0][0]) + + ts += 1 + + +def check_learning_achieved( + tune_results: "tune.ResultGrid", + min_value: float, + evaluation: Optional[bool] = None, + metric: str = f"{ENV_RUNNER_RESULTS}/episode_return_mean", +): + """Throws an error if `min_reward` is not reached within tune_results. + + Checks the last iteration found in tune_results for its + "episode_return_mean" value and compares it to `min_reward`. + + Args: + tune_results: The tune.Tuner().fit() returned results object. + min_reward: The min reward that must be reached. + evaluation: If True, use `evaluation/env_runners/[metric]`, if False, use + `env_runners/[metric]`, if None, use evaluation sampler results if + available otherwise, use train sampler results. + + Raises: + ValueError: If `min_reward` not reached. + """ + # Get maximum value of `metrics` over all trials + # (check if at least one trial achieved some learning, not just the final one). + recorded_values = [] + for _, row in tune_results.get_dataframe().iterrows(): + if evaluation or ( + evaluation is None and f"{EVALUATION_RESULTS}/{metric}" in row + ): + recorded_values.append(row[f"{EVALUATION_RESULTS}/{metric}"]) + else: + recorded_values.append(row[metric]) + best_value = max(recorded_values) + if best_value < min_value: + raise ValueError(f"`{metric}` of {min_value} not reached!") + print(f"`{metric}` of {min_value} reached! ok") + + +def check_off_policyness( + results: ResultDict, + upper_limit: float, + lower_limit: float = 0.0, +) -> Optional[float]: + """Verifies that the off-policy'ness of some update is within some range. + + Off-policy'ness is defined as the average (across n workers) diff + between the number of gradient updates performed on the policy used + for sampling vs the number of gradient updates that have been performed + on the trained policy (usually the one on the local worker). + + Uses the published DIFF_NUM_GRAD_UPDATES_VS_SAMPLER_POLICY metric inside + a training results dict and compares to the given bounds. + + Note: Only works with single-agent results thus far. + + Args: + results: The training results dict. + upper_limit: The upper limit to for the off_policy_ness value. + lower_limit: The lower limit to for the off_policy_ness value. + + Returns: + The off-policy'ness value (described above). + + Raises: + AssertionError: If the value is out of bounds. + """ + + # Have to import this here to avoid circular dependency. + from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID + from ray.rllib.utils.metrics.learner_info import LEARNER_INFO + + # Assert that the off-policy'ness is within the given bounds. + learner_info = results["info"][LEARNER_INFO] + if DEFAULT_POLICY_ID not in learner_info: + return None + off_policy_ness = learner_info[DEFAULT_POLICY_ID][ + DIFF_NUM_GRAD_UPDATES_VS_SAMPLER_POLICY + ] + # Roughly: Reaches up to 0.4 for 2 rollout workers and up to 0.2 for + # 1 rollout worker. + if not (lower_limit <= off_policy_ness <= upper_limit): + raise AssertionError( + f"`off_policy_ness` ({off_policy_ness}) is outside the given bounds " + f"({lower_limit} - {upper_limit})!" + ) + + return off_policy_ness + + +def check_train_results_new_api_stack(train_results: ResultDict) -> None: + """Checks proper structure of a Algorithm.train() returned dict. + + Args: + train_results: The train results dict to check. + + Raises: + AssertionError: If `train_results` doesn't have the proper structure or + data in it. + """ + # Import these here to avoid circular dependencies. + from ray.rllib.utils.metrics import ( + ENV_RUNNER_RESULTS, + FAULT_TOLERANCE_STATS, + LEARNER_RESULTS, + TIMERS, + ) + + # Assert that some keys are where we would expect them. + for key in [ + ENV_RUNNER_RESULTS, + FAULT_TOLERANCE_STATS, + LEARNER_RESULTS, + TIMERS, + TRAINING_ITERATION, + "config", + ]: + assert ( + key in train_results + ), f"'{key}' not found in `train_results` ({train_results})!" + + # Make sure, `config` is an actual dict, not an AlgorithmConfig object. + assert isinstance( + train_results["config"], dict + ), "`config` in results not a python dict!" + + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + + is_multi_agent = ( + AlgorithmConfig() + .update_from_dict({"policies": train_results["config"]["policies"]}) + .is_multi_agent + ) + + # Check in particular the "info" dict. + learner_results = train_results[LEARNER_RESULTS] + + # Make sure we have a `DEFAULT_MODULE_ID key if we are not in a + # multi-agent setup. + if not is_multi_agent: + assert len(learner_results) == 0 or DEFAULT_MODULE_ID in learner_results, ( + f"'{DEFAULT_MODULE_ID}' not found in " + f"train_results['{LEARNER_RESULTS}']!" + ) + + for module_id, module_metrics in learner_results.items(): + # The ModuleID can be __all_modules__ in multi-agent case when the new learner + # stack is enabled. + if module_id == "__all_modules__": + continue + + # On the new API stack, policy has no LEARNER_STATS_KEY under it anymore. + for key, value in module_metrics.items(): + # Min- and max-stats should be single values. + if key.endswith("_min") or key.endswith("_max"): + assert np.isscalar(value), f"'key' value not a scalar ({value})!" + + return train_results + + +@OldAPIStack +def check_train_results(train_results: ResultDict): + """Checks proper structure of a Algorithm.train() returned dict. + + Args: + train_results: The train results dict to check. + + Raises: + AssertionError: If `train_results` doesn't have the proper structure or + data in it. + """ + # Import these here to avoid circular dependencies. + from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID + from ray.rllib.utils.metrics.learner_info import LEARNER_INFO, LEARNER_STATS_KEY + + # Assert that some keys are where we would expect them. + for key in [ + "config", + "custom_metrics", + ENV_RUNNER_RESULTS, + "info", + "iterations_since_restore", + "num_healthy_workers", + "perf", + "time_since_restore", + "time_this_iter_s", + "timers", + "time_total_s", + TRAINING_ITERATION, + ]: + assert ( + key in train_results + ), f"'{key}' not found in `train_results` ({train_results})!" + + for key in [ + "episode_len_mean", + "episode_reward_max", + "episode_reward_mean", + "episode_reward_min", + "hist_stats", + "policy_reward_max", + "policy_reward_mean", + "policy_reward_min", + "sampler_perf", + ]: + assert key in train_results[ENV_RUNNER_RESULTS], ( + f"'{key}' not found in `train_results[ENV_RUNNER_RESULTS]` " + f"({train_results[ENV_RUNNER_RESULTS]})!" + ) + + # Make sure, `config` is an actual dict, not an AlgorithmConfig object. + assert isinstance( + train_results["config"], dict + ), "`config` in results not a python dict!" + + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + + is_multi_agent = ( + AlgorithmConfig() + .update_from_dict({"policies": train_results["config"]["policies"]}) + .is_multi_agent + ) + + # Check in particular the "info" dict. + info = train_results["info"] + assert LEARNER_INFO in info, f"'learner' not in train_results['infos'] ({info})!" + assert ( + "num_steps_trained" in info or NUM_ENV_STEPS_TRAINED in info + ), f"'num_(env_)?steps_trained' not in train_results['infos'] ({info})!" + + learner_info = info[LEARNER_INFO] + + # Make sure we have a default_policy key if we are not in a + # multi-agent setup. + if not is_multi_agent: + # APEX algos sometimes have an empty learner info dict (no metrics + # collected yet). + assert len(learner_info) == 0 or DEFAULT_POLICY_ID in learner_info, ( + f"'{DEFAULT_POLICY_ID}' not found in " + f"train_results['infos']['learner'] ({learner_info})!" + ) + + for pid, policy_stats in learner_info.items(): + if pid == "batch_count": + continue + + # the pid can be __all__ in multi-agent case when the new learner stack is + # enabled. + if pid == "__all__": + continue + + # On the new API stack, policy has no LEARNER_STATS_KEY under it anymore. + if LEARNER_STATS_KEY in policy_stats: + learner_stats = policy_stats[LEARNER_STATS_KEY] + else: + learner_stats = policy_stats + for key, value in learner_stats.items(): + # Min- and max-stats should be single values. + if key.startswith("min_") or key.startswith("max_"): + assert np.isscalar(value), f"'key' value not a scalar ({value})!" + + return train_results + + +# TODO (simon): Use this function in the `run_rllib_example_experiment` when +# `no_tune` is `True`. +def should_stop( + stop: Dict[str, Any], results: ResultDict, keep_ray_up: bool = False +) -> bool: + """Checks stopping criteria on `ResultDict` + + Args: + stop: Dictionary of stopping criteria. Each criterium is a mapping of + a metric in the `ResultDict` of the algorithm to a certain criterium. + results: An RLlib `ResultDict` containing all results from a training step. + keep_ray_up: Optionally shutting down the runnin Ray instance. + + Returns: True, if any stopping criterium is fulfilled. Otherwise, False. + """ + for key, threshold in stop.items(): + val = results + for k in key.split("/"): + k = k.strip() + # If k exists in the current level, continue down; + # otherwise, set val to None and break out of this inner loop. + if isinstance(val, dict) and k in val: + val = val[k] + else: + val = None + break + + # If the key was not found, simply skip to the next criterion. + if val is None: + continue + + try: + # Check that val is numeric and meets the threshold. + if not np.isnan(val) and val >= threshold: + print(f"Stop criterion ({key}={threshold}) fulfilled!") + if not keep_ray_up: + ray.shutdown() + return True + except TypeError: + # If val isn't numeric, skip this criterion. + continue + + # If none of the criteria are fulfilled, return False. + return False + + +# TODO (sven): Make this the de-facto, well documented, and unified utility for most of +# our tests: +# - CI (label: "learning_tests") +# - release tests (benchmarks) +# - example scripts +def run_rllib_example_script_experiment( + base_config: "AlgorithmConfig", + args: Optional[argparse.Namespace] = None, + *, + stop: Optional[Dict] = None, + success_metric: Optional[Dict] = None, + trainable: Optional[Type] = None, + tune_callbacks: Optional[List] = None, + keep_config: bool = False, + keep_ray_up: bool = False, + scheduler=None, + progress_reporter=None, +) -> Union[ResultDict, tune.result_grid.ResultGrid]: + """Given an algorithm config and some command line args, runs an experiment. + + There are some constraints on what properties must be defined in `args`. + It should ideally be generated via calling + `args = add_rllib_example_script_args()`, which can be found in this very module + here. + + The function sets up an Algorithm object from the given config (altered by the + contents of `args`), then runs the Algorithm via Tune (or manually, if + `args.no_tune` is set to True) using the stopping criteria in `stop`. + + At the end of the experiment, if `args.as_test` is True, checks, whether the + Algorithm reached the `success_metric` (if None, use `env_runners/ + episode_return_mean` with a minimum value of `args.stop_reward`). + + See https://github.com/ray-project/ray/tree/master/rllib/examples for an overview + of all supported command line options. + + Args: + base_config: The AlgorithmConfig object to use for this experiment. This base + config will be automatically "extended" based on some of the provided + `args`. For example, `args.num_env_runners` is used to set + `config.num_env_runners`, etc.. + args: A argparse.Namespace object, ideally returned by calling + `args = add_rllib_example_script_args()`. It must have the following + properties defined: `stop_iters`, `stop_reward`, `stop_timesteps`, + `no_tune`, `verbose`, `checkpoint_freq`, `as_test`. Optionally, for WandB + logging: `wandb_key`, `wandb_project`, `wandb_run_name`. + stop: An optional dict mapping ResultDict key strings (using "/" in case of + nesting, e.g. "env_runners/episode_return_mean" for referring to + `result_dict['env_runners']['episode_return_mean']` to minimum + values, reaching of which will stop the experiment). Default is: + { + "env_runners/episode_return_mean": args.stop_reward, + "training_iteration": args.stop_iters, + "num_env_steps_sampled_lifetime": args.stop_timesteps, + } + success_metric: Only relevant if `args.as_test` is True. + A dict mapping a single(!) ResultDict key string (using "/" in + case of nesting, e.g. "env_runners/episode_return_mean" for referring + to `result_dict['env_runners']['episode_return_mean']` to a single(!) + minimum value to be reached in order for the experiment to count as + successful. If `args.as_test` is True AND this `success_metric` is not + reached with the bounds defined by `stop`, will raise an Exception. + trainable: The Trainable sub-class to run in the tune.Tuner. If None (default), + use the registered RLlib Algorithm class specified by args.algo. + tune_callbacks: A list of Tune callbacks to configure with the tune.Tuner. + In case `args.wandb_key` is provided, appends a WandB logger to this + list. + keep_config: Set this to True, if you don't want this utility to change the + given `base_config` in any way and leave it as-is. This is helpful + for those example scripts which demonstrate how to set config settings + that are otherwise taken care of automatically in this function (e.g. + `num_env_runners`). + + Returns: + The last ResultDict from a --no-tune run OR the tune.Tuner.fit() + results. + """ + if args is None: + parser = add_rllib_example_script_args() + args = parser.parse_args() + + # Deprecated args. + if args.enable_new_api_stack: + raise ValueError( + "`--enable-new-api-stack` flag no longer supported (it's the default " + "behavior now)! To switch back to the old API stack on your scripts, use " + "the `--old-api-stack` flag." + ) + + # If run --as-release-test, --as-test must also be set. + if args.as_release_test: + args.as_test = True + + # Initialize Ray. + ray.init( + num_cpus=args.num_cpus or None, + local_mode=args.local_mode, + ignore_reinit_error=True, + ) + + # Define one or more stopping criteria. + if stop is None: + stop = { + f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": args.stop_reward, + f"{ENV_RUNNER_RESULTS}/{NUM_ENV_STEPS_SAMPLED_LIFETIME}": ( + args.stop_timesteps + ), + TRAINING_ITERATION: args.stop_iters, + } + + config = base_config + + # Enhance the `base_config`, based on provided `args`. + if not keep_config: + # Set the framework. + config.framework(args.framework) + + # Add an env specifier (only if not already set in config)? + if args.env is not None and config.env is None: + config.environment(args.env) + + # Disable the new API stack? + if args.old_api_stack: + config.api_stack( + enable_rl_module_and_learner=False, + enable_env_runner_and_connector_v2=False, + ) + + # Define EnvRunner scaling and behavior. + if args.num_env_runners is not None: + config.env_runners(num_env_runners=args.num_env_runners) + if args.num_envs_per_env_runner is not None: + config.env_runners(num_envs_per_env_runner=args.num_envs_per_env_runner) + + # Define compute resources used automatically (only using the --num-learners + # and --num-gpus-per-learner args). + # New stack. + if config.enable_rl_module_and_learner: + if args.num_gpus is not None and args.num_gpus > 0: + raise ValueError( + "--num-gpus is not supported on the new API stack! To train on " + "GPUs, use the command line options `--num-gpus-per-learner=1` and " + "`--num-learners=[your number of available GPUs]`, instead." + ) + + # Do we have GPUs available in the cluster? + num_gpus_available = ray.cluster_resources().get("GPU", 0) + # Number of actual Learner instances (including the local Learner if + # `num_learners=0`). + num_actual_learners = ( + args.num_learners + if args.num_learners is not None + else config.num_learners + ) or 1 # 1: There is always a local Learner, if num_learners=0. + # How many were hard-requested by the user + # (through explicit `--num-gpus-per-learner >= 1`). + num_gpus_requested = (args.num_gpus_per_learner or 0) * num_actual_learners + # Number of GPUs needed, if `num_gpus_per_learner=None` (auto). + num_gpus_needed_if_available = ( + args.num_gpus_per_learner + if args.num_gpus_per_learner is not None + else 1 + ) * num_actual_learners + # Define compute resources used. + config.resources(num_gpus=0) # old API stack setting + if args.num_learners is not None: + config.learners(num_learners=args.num_learners) + + # User wants to use aggregator actors per Learner. + if args.num_aggregator_actors_per_learner is not None: + config.learners( + num_aggregator_actors_per_learner=( + args.num_aggregator_actors_per_learner + ) + ) + + # User wants to use GPUs if available, but doesn't hard-require them. + if args.num_gpus_per_learner is None: + if num_gpus_available >= num_gpus_needed_if_available: + config.learners(num_gpus_per_learner=1) + else: + config.learners(num_gpus_per_learner=0) + # User hard-requires n GPUs, but they are not available -> Error. + elif num_gpus_available < num_gpus_requested: + raise ValueError( + "You are running your script with --num-learners=" + f"{args.num_learners} and --num-gpus-per-learner=" + f"{args.num_gpus_per_learner}, but your cluster only has " + f"{num_gpus_available} GPUs!" + ) + + # All required GPUs are available -> Use them. + else: + config.learners(num_gpus_per_learner=args.num_gpus_per_learner) + + # Set CPUs per Learner. + if args.num_cpus_per_learner is not None: + config.learners(num_cpus_per_learner=args.num_cpus_per_learner) + + # Old stack (override only if arg was provided by user). + elif args.num_gpus is not None: + config.resources(num_gpus=args.num_gpus) + + # Evaluation setup. + if args.evaluation_interval > 0: + config.evaluation( + evaluation_num_env_runners=args.evaluation_num_env_runners, + evaluation_interval=args.evaluation_interval, + evaluation_duration=args.evaluation_duration, + evaluation_duration_unit=args.evaluation_duration_unit, + evaluation_parallel_to_training=args.evaluation_parallel_to_training, + ) + + # Set the log-level (if applicable). + if args.log_level is not None: + config.debugging(log_level=args.log_level) + + # Set the output dir (if applicable). + if args.output is not None: + config.offline_data(output=args.output) + + # Run the experiment w/o Tune (directly operate on the RLlib Algorithm object). + if args.no_tune: + assert not args.as_test and not args.as_release_test + algo = config.build() + for i in range(stop.get(TRAINING_ITERATION, args.stop_iters)): + results = algo.train() + if ENV_RUNNER_RESULTS in results: + mean_return = results[ENV_RUNNER_RESULTS].get( + EPISODE_RETURN_MEAN, np.nan + ) + print(f"iter={i} R={mean_return}", end="") + if ( + EVALUATION_RESULTS in results + and ENV_RUNNER_RESULTS in results[EVALUATION_RESULTS] + ): + Reval = results[EVALUATION_RESULTS][ENV_RUNNER_RESULTS][ + EPISODE_RETURN_MEAN + ] + print(f" R(eval)={Reval}", end="") + print() + for key, threshold in stop.items(): + val = results + for k in key.split("/"): + try: + val = val[k] + except KeyError: + val = None + break + if val is not None and not np.isnan(val) and val >= threshold: + print(f"Stop criterium ({key}={threshold}) fulfilled!") + if not keep_ray_up: + ray.shutdown() + return results + + if not keep_ray_up: + ray.shutdown() + return results + + # Run the experiment using Ray Tune. + + # Log results using WandB. + tune_callbacks = tune_callbacks or [] + if hasattr(args, "wandb_key") and ( + args.wandb_key is not None or WANDB_ENV_VAR in os.environ + ): + wandb_key = args.wandb_key or os.environ[WANDB_ENV_VAR] + project = args.wandb_project or ( + args.algo.lower() + "-" + re.sub("\\W+", "-", str(config.env).lower()) + ) + tune_callbacks.append( + WandbLoggerCallback( + api_key=wandb_key, + project=project, + upload_checkpoints=True, + **({"name": args.wandb_run_name} if args.wandb_run_name else {}), + ) + ) + # Auto-configure a CLIReporter (to log the results to the console). + # Use better ProgressReporter for multi-agent cases: List individual policy rewards. + if progress_reporter is None and args.num_agents > 0: + progress_reporter = CLIReporter( + metric_columns={ + **{ + TRAINING_ITERATION: "iter", + "time_total_s": "total time (s)", + NUM_ENV_STEPS_SAMPLED_LIFETIME: "ts", + f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": "combined return", + }, + **{ + ( + f"{ENV_RUNNER_RESULTS}/module_episode_returns_mean/" f"{pid}" + ): f"return {pid}" + for pid in config.policies + }, + }, + ) + + # Force Tuner to use old progress output as the new one silently ignores our custom + # `CLIReporter`. + os.environ["RAY_AIR_NEW_OUTPUT"] = "0" + + # Run the actual experiment (using Tune). + start_time = time.time() + results = tune.Tuner( + trainable or config.algo_class, + param_space=config, + run_config=tune.RunConfig( + stop=stop, + verbose=args.verbose, + callbacks=tune_callbacks, + checkpoint_config=tune.CheckpointConfig( + checkpoint_frequency=args.checkpoint_freq, + checkpoint_at_end=args.checkpoint_at_end, + ), + progress_reporter=progress_reporter, + ), + tune_config=tune.TuneConfig( + num_samples=args.num_samples, + max_concurrent_trials=args.max_concurrent_trials, + scheduler=scheduler, + ), + ).fit() + time_taken = time.time() - start_time + + if not keep_ray_up: + ray.shutdown() + + # Error out, if Tuner.fit() failed to run. Otherwise, erroneous examples might pass + # the CI tests w/o us knowing that they are broken (b/c some examples do not have + # a --as-test flag and/or any passing criteria). + if results.errors: + # Might cause an IndexError if the tuple is not long enough; in that case, use repr(e). + errors = [ + e.args[0].args[2] + if e.args and hasattr(e.args[0], "args") and len(e.args[0].args) > 2 + else repr(e) + for e in results.errors + ] + raise RuntimeError( + f"Running the example script resulted in one or more errors! {errors}" + ) + + # If run as a test, check whether we reached the specified success criteria. + test_passed = False + if args.as_test: + # Success metric not provided, try extracting it from `stop`. + if success_metric is None: + for try_it in [ + f"{EVALUATION_RESULTS}/{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}", + f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}", + ]: + if try_it in stop: + success_metric = {try_it: stop[try_it]} + break + if success_metric is None: + success_metric = { + f"{ENV_RUNNER_RESULTS}/{EPISODE_RETURN_MEAN}": args.stop_reward, + } + # TODO (sven): Make this work for more than one metric (AND-logic?). + # Get maximum value of `metric` over all trials + # (check if at least one trial achieved some learning, not just the final one). + success_metric_key, success_metric_value = next(iter(success_metric.items())) + best_value = max( + row[success_metric_key] for _, row in results.get_dataframe().iterrows() + ) + if best_value >= success_metric_value: + test_passed = True + print(f"`{success_metric_key}` of {success_metric_value} reached! ok") + + if args.as_release_test: + trial = results._experiment_analysis.trials[0] + stats = trial.last_result + stats.pop("config", None) + json_summary = { + "time_taken": float(time_taken), + "trial_states": [trial.status], + "last_update": float(time.time()), + "stats": stats, + "passed": [test_passed], + "not_passed": [not test_passed], + "failures": {str(trial): 1} if not test_passed else {}, + } + with open( + os.environ.get("TEST_OUTPUT_JSON", "/tmp/learning_test.json"), + "wt", + ) as f: + try: + json.dump(json_summary, f) + # Something went wrong writing json. Try again w/ simplified stats. + except Exception: + from ray.rllib.algorithms.algorithm import Algorithm + + simplified_stats = { + k: stats[k] for k in Algorithm._progress_metrics if k in stats + } + json_summary["stats"] = simplified_stats + json.dump(json_summary, f) + + if not test_passed: + raise ValueError( + f"`{success_metric_key}` of {success_metric_value} not reached!" + ) + + return results + + +def check_same_batch(batch1, batch2) -> None: + """Check if both batches are (almost) identical. + + For MultiAgentBatches, the step count and individual policy's + SampleBatches are checked for identity. For SampleBatches, identity is + checked as the almost numerical key-value-pair identity between batches + with ray.rllib.utils.test_utils.check(). unroll_id is compared only if + both batches have an unroll_id. + + Args: + batch1: Batch to compare against batch2 + batch2: Batch to compare against batch1 + """ + # Avoids circular import + from ray.rllib.policy.sample_batch import MultiAgentBatch, SampleBatch + + assert type(batch1) is type( + batch2 + ), "Input batches are of different types {} and {}".format( + str(type(batch1)), str(type(batch2)) + ) + + def check_sample_batches(_batch1, _batch2, _policy_id=None): + unroll_id_1 = _batch1.get("unroll_id", None) + unroll_id_2 = _batch2.get("unroll_id", None) + # unroll IDs only have to fit if both batches have them + if unroll_id_1 is not None and unroll_id_2 is not None: + assert unroll_id_1 == unroll_id_2 + + batch1_keys = set() + for k, v in _batch1.items(): + # unroll_id is compared above already + if k == "unroll_id": + continue + check(v, _batch2[k]) + batch1_keys.add(k) + + batch2_keys = set(_batch2.keys()) + # unroll_id is compared above already + batch2_keys.discard("unroll_id") + _difference = batch1_keys.symmetric_difference(batch2_keys) + + # Cases where one batch has info and the other has not + if _policy_id: + assert not _difference, ( + "SampleBatches for policy with ID {} " + "don't share information on the " + "following information: \n{}" + "".format(_policy_id, _difference) + ) + else: + assert not _difference, ( + "SampleBatches don't share information " + "on the following information: \n{}" + "".format(_difference) + ) + + if type(batch1) is SampleBatch: + check_sample_batches(batch1, batch2) + elif type(batch1) is MultiAgentBatch: + assert batch1.count == batch2.count + batch1_ids = set() + for policy_id, policy_batch in batch1.policy_batches.items(): + check_sample_batches( + policy_batch, batch2.policy_batches[policy_id], policy_id + ) + batch1_ids.add(policy_id) + + # Case where one ma batch has info on a policy the other has not + batch2_ids = set(batch2.policy_batches.keys()) + difference = batch1_ids.symmetric_difference(batch2_ids) + assert ( + not difference + ), f"MultiAgentBatches don't share the following information: \n{difference}." + else: + raise ValueError("Unsupported batch type " + str(type(batch1))) + + +def check_reproducibilty( + algo_class: Type["Algorithm"], + algo_config: "AlgorithmConfig", + *, + fw_kwargs: Dict[str, Any], + training_iteration: int = 1, +) -> None: + # TODO @kourosh: we can get rid of examples/deterministic_training.py once + # this is added to all algorithms + """Check if the algorithm is reproducible across different testing conditions: + + frameworks: all input frameworks + num_gpus: int(os.environ.get("RLLIB_NUM_GPUS", "0")) + num_workers: 0 (only local workers) or + 4 ((1) local workers + (4) remote workers) + num_envs_per_env_runner: 2 + + Args: + algo_class: Algorithm class to test. + algo_config: Base config to use for the algorithm. + fw_kwargs: Framework iterator keyword arguments. + training_iteration: Number of training iterations to run. + + Returns: + None + + Raises: + It raises an AssertionError if the algorithm is not reproducible. + """ + from ray.rllib.policy.sample_batch import DEFAULT_POLICY_ID + from ray.rllib.utils.metrics.learner_info import LEARNER_INFO + + stop_dict = {TRAINING_ITERATION: training_iteration} + # use 0 and 2 workers (for more that 4 workers we have to make sure the instance + # type in ci build has enough resources) + for num_workers in [0, 2]: + algo_config = ( + algo_config.debugging(seed=42).env_runners( + num_env_runners=num_workers, num_envs_per_env_runner=2 + ) + # new API + .learners( + num_gpus_per_learner=int(os.environ.get("RLLIB_NUM_GPUS", "0")), + ) + # old API + .resources( + num_gpus=int(os.environ.get("RLLIB_NUM_GPUS", "0")), + ) + ) + + print( + f"Testing reproducibility of {algo_class.__name__}" + f" with {num_workers} workers" + ) + print("/// config") + pprint.pprint(algo_config.to_dict()) + # test tune.Tuner().fit() reproducibility + results1 = tune.Tuner( + algo_class, + param_space=algo_config.to_dict(), + run_config=tune.RunConfig(stop=stop_dict, verbose=1), + ).fit() + results1 = results1.get_best_result().metrics + + results2 = tune.Tuner( + algo_class, + param_space=algo_config.to_dict(), + run_config=tune.RunConfig(stop=stop_dict, verbose=1), + ).fit() + results2 = results2.get_best_result().metrics + + # Test rollout behavior. + check( + results1[ENV_RUNNER_RESULTS]["hist_stats"], + results2[ENV_RUNNER_RESULTS]["hist_stats"], + ) + # As well as training behavior (minibatch sequence during SGD + # iterations). + # As well as training behavior (minibatch sequence during SGD + # iterations). + if algo_config.enable_rl_module_and_learner: + check( + results1["info"][LEARNER_INFO][DEFAULT_POLICY_ID], + results2["info"][LEARNER_INFO][DEFAULT_POLICY_ID], + ) + else: + check( + results1["info"][LEARNER_INFO][DEFAULT_POLICY_ID]["learner_stats"], + results2["info"][LEARNER_INFO][DEFAULT_POLICY_ID]["learner_stats"], + ) + + +def get_cartpole_dataset_reader(batch_size: int = 1) -> "DatasetReader": + """Returns a DatasetReader for the cartpole dataset. + Args: + batch_size: The batch size to use for the reader. + Returns: + A rllib DatasetReader for the cartpole dataset. + """ + from ray.rllib.algorithms import AlgorithmConfig + from ray.rllib.offline import IOContext + from ray.rllib.offline.dataset_reader import ( + DatasetReader, + get_dataset_and_shards, + ) + + path = "tests/data/cartpole/large.json" + input_config = {"format": "json", "paths": path} + dataset, _ = get_dataset_and_shards( + AlgorithmConfig().offline_data(input_="dataset", input_config=input_config) + ) + ioctx = IOContext( + config=( + AlgorithmConfig() + .training(train_batch_size=batch_size) + .offline_data(actions_in_input_normalized=True) + ), + worker_index=0, + ) + reader = DatasetReader(dataset, ioctx) + return reader + + +class ModelChecker: + """Helper class to compare architecturally identical Models across frameworks. + + Holds a ModelConfig, such that individual models can be added simply via their + framework string (by building them with config.build(framework=...). + A call to `check()` forces all added models to be compared in terms of their + number of trainable and non-trainable parameters, as well as, their + computation results given a common weights structure and values and identical + inputs to the models. + """ + + def __init__(self, config): + self.config = config + + # To compare number of params between frameworks. + self.param_counts = {} + # To compare computed outputs from fixed-weights-nets between frameworks. + self.output_values = {} + + # We will pass an observation filled with this one random value through + # all DL networks (after they have been set to fixed-weights) to compare + # the computed outputs. + self.random_fill_input_value = np.random.uniform(-0.01, 0.01) + + # Dict of models to check against each other. + self.models = {} + + def add(self, framework: str = "torch", obs=True, state=False) -> Any: + """Builds a new Model for the given framework.""" + model = self.models[framework] = self.config.build(framework=framework) + + # Pass a B=1 observation through the model. + inputs = np.full( + [1] + ([1] if state else []) + list(self.config.input_dims), + self.random_fill_input_value, + ) + if obs: + inputs = {Columns.OBS: inputs} + if state: + inputs[Columns.STATE_IN] = tree.map_structure( + lambda s: np.zeros(shape=[1] + list(s)), state + ) + if framework == "torch": + from ray.rllib.utils.torch_utils import convert_to_torch_tensor + + inputs = convert_to_torch_tensor(inputs) + # w/ old specs: inputs = model.input_specs.fill(self.random_fill_input_value) + + outputs = model(inputs) + + # Bring model into a reproducible, comparable state (so we can compare + # computations across frameworks). Use only a value-sequence of len=1 here + # as it could possibly be that the layers are stored in different order + # across the different frameworks. + model._set_to_dummy_weights(value_sequence=(self.random_fill_input_value,)) + + # Perform another forward pass. + comparable_outputs = model(inputs) + + # Store the number of parameters for this framework's net. + self.param_counts[framework] = model.get_num_parameters() + # Store the fixed-weights-net outputs for this framework's net. + if framework == "torch": + self.output_values[framework] = tree.map_structure( + lambda s: s.detach().numpy() if s is not None else None, + comparable_outputs, + ) + else: + self.output_values[framework] = tree.map_structure( + lambda s: s.numpy() if s is not None else None, comparable_outputs + ) + return outputs + + def check(self): + """Compares all added Models with each other and possibly raises errors.""" + + main_key = next(iter(self.models.keys())) + # Compare number of trainable and non-trainable params between all + # frameworks. + for c in self.param_counts.values(): + check(c, self.param_counts[main_key]) + + # Compare dummy outputs by exact values given that all nets received the + # same input and all nets have the same (dummy) weight values. + for v in self.output_values.values(): + check(v, self.output_values[main_key], atol=0.0005) + + +def _get_mean_action_from_algorithm(alg: "Algorithm", obs: np.ndarray) -> np.ndarray: + """Returns the mean action computed by the given algorithm. + + Note: This makes calls to `Algorithm.compute_single_action` + + Args: + alg: The constructed algorithm to run inference on. + obs: The observation to compute the action for. + + Returns: + The mean action computed by the algorithm over 5000 samples. + + """ + out = [] + for _ in range(5000): + out.append(float(alg.compute_single_action(obs))) + return np.mean(out) + + +def check_supported_spaces( + alg: str, + config: "AlgorithmConfig", + train: bool = True, + check_bounds: bool = False, + frameworks: Optional[Tuple[str, ...]] = None, + use_gpu: bool = False, +): + """Checks whether the given algorithm supports different action and obs spaces. + + Performs the checks by constructing an rllib algorithm from the config and + checking to see that the model inside the policy is the correct one given + the action and obs spaces. For example if the action space is discrete and + the obs space is an image, then the model should be a vision network with + a categorical action distribution. + + Args: + alg: The name of the algorithm to test. + config: The config to use for the algorithm. + train: Whether to train the algorithm for a few iterations. + check_bounds: Whether to check the bounds of the action space. + frameworks: The frameworks to test the algorithm with. + use_gpu: Whether to check support for training on a gpu. + + + """ + # Do these imports here because otherwise we have circular imports. + from ray.rllib.examples.envs.classes.random_env import RandomEnv + from ray.rllib.models.torch.complex_input_net import ( + ComplexInputNetwork as TorchComplexNet, + ) + from ray.rllib.models.torch.fcnet import FullyConnectedNetwork as TorchFCNet + from ray.rllib.models.torch.visionnet import VisionNetwork as TorchVisionNet + + action_spaces_to_test = { + # Test discrete twice here until we support multi_binary action spaces + "discrete": Discrete(5), + "continuous": Box(-1.0, 1.0, (5,), dtype=np.float32), + "int_actions": Box(0, 3, (2, 3), dtype=np.int32), + "multidiscrete": MultiDiscrete([1, 2, 3, 4]), + "tuple": GymTuple( + [Discrete(2), Discrete(3), Box(-1.0, 1.0, (5,), dtype=np.float32)] + ), + "dict": GymDict( + { + "action_choice": Discrete(3), + "parameters": Box(-1.0, 1.0, (1,), dtype=np.float32), + "yet_another_nested_dict": GymDict( + {"a": GymTuple([Discrete(2), Discrete(3)])} + ), + } + ), + } + + observation_spaces_to_test = { + "multi_binary": MultiBinary([3, 10, 10]), + "discrete": Discrete(5), + "continuous": Box(-1.0, 1.0, (5,), dtype=np.float32), + "vector2d": Box(-1.0, 1.0, (5, 5), dtype=np.float32), + "image": Box(-1.0, 1.0, (84, 84, 1), dtype=np.float32), + "tuple": GymTuple([Discrete(10), Box(-1.0, 1.0, (5,), dtype=np.float32)]), + "dict": GymDict( + { + "task": Discrete(10), + "position": Box(-1.0, 1.0, (5,), dtype=np.float32), + } + ), + } + + # The observation spaces that we test RLModules with + rlmodule_supported_observation_spaces = [ + "multi_binary", + "discrete", + "continuous", + "image", + "tuple", + "dict", + ] + + # The action spaces that we test RLModules with + rlmodule_supported_action_spaces = ["discrete", "continuous"] + + default_observation_space = default_action_space = "discrete" + + config["log_level"] = "ERROR" + config["env"] = RandomEnv + + def _do_check(alg, config, a_name, o_name): + # We need to copy here so that this validation does not affect the actual + # validation method call further down the line. + config_copy = config.copy() + config_copy.validate() + # If RLModules are enabled, we need to skip a few tests for now: + if config_copy.enable_rl_module_and_learner: + # Skip PPO cases in which RLModules don't support the given spaces yet. + if o_name not in rlmodule_supported_observation_spaces: + logger.warning( + "Skipping PPO test with RLModules for obs space {}".format(o_name) + ) + return + if a_name not in rlmodule_supported_action_spaces: + logger.warning( + "Skipping PPO test with RLModules for action space {}".format( + a_name + ) + ) + return + + fw = config["framework"] + action_space = action_spaces_to_test[a_name] + obs_space = observation_spaces_to_test[o_name] + print( + "=== Testing {} (fw={}) action_space={} obs_space={} ===".format( + alg, fw, action_space, obs_space + ) + ) + t0 = time.time() + config.update_from_dict( + dict( + env_config=dict( + action_space=action_space, + observation_space=obs_space, + reward_space=Box(1.0, 1.0, shape=(), dtype=np.float32), + p_terminated=1.0, + check_action_bounds=check_bounds, + ) + ) + ) + stat = "ok" + + try: + algo = config.build() + except ray.exceptions.RayActorError as e: + if len(e.args) >= 2 and isinstance(e.args[2], UnsupportedSpaceException): + stat = "unsupported" + elif isinstance(e.args[0].args[2], UnsupportedSpaceException): + stat = "unsupported" + else: + raise + except UnsupportedSpaceException: + stat = "unsupported" + else: + if alg not in ["SAC", "PPO"]: + # 2D (image) input: Expect VisionNet. + if o_name in ["atari", "image"]: + assert isinstance(algo.get_policy().model, TorchVisionNet) + # 1D input: Expect FCNet. + elif o_name == "continuous": + assert isinstance(algo.get_policy().model, TorchFCNet) + # Could be either one: ComplexNet (if disabled Preprocessor) + # or FCNet (w/ Preprocessor). + elif o_name == "vector2d": + assert isinstance( + algo.get_policy().model, (TorchComplexNet, TorchFCNet) + ) + if train: + algo.train() + algo.stop() + print("Test: {}, ran in {}s".format(stat, time.time() - t0)) + + if not frameworks: + frameworks = ("tf2", "tf", "torch") + + _do_check_remote = ray.remote(_do_check) + _do_check_remote = _do_check_remote.options(num_gpus=1 if use_gpu else 0) + # Test all action spaces first. + for a_name in action_spaces_to_test.keys(): + o_name = default_observation_space + ray.get(_do_check_remote.remote(alg, config, a_name, o_name)) + + # Now test all observation spaces. + for o_name in observation_spaces_to_test.keys(): + a_name = default_action_space + ray.get(_do_check_remote.remote(alg, config, a_name, o_name)) diff --git a/lib/python3.12/site-packages/ray/rllib/utils/tf_run_builder.py b/lib/python3.12/site-packages/ray/rllib/utils/tf_run_builder.py new file mode 100644 index 0000000000000000000000000000000000000000..26bf8207719e8026c0088fdf74f9fe03c462a9c3 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/tf_run_builder.py @@ -0,0 +1,115 @@ +import logging +import os +import time + +from ray.rllib.utils.annotations import OldAPIStack +from ray.rllib.utils.framework import try_import_tf +from ray.util.debug import log_once + +tf1, tf, tfv = try_import_tf() +logger = logging.getLogger(__name__) + + +@OldAPIStack +class _TFRunBuilder: + """Used to incrementally build up a TensorFlow run. + + This is particularly useful for batching ops from multiple different + policies in the multi-agent setting. + """ + + def __init__(self, session, debug_name): + self.session = session + self.debug_name = debug_name + self.feed_dict = {} + self.fetches = [] + self._executed = None + + def add_feed_dict(self, feed_dict): + assert not self._executed + for k in feed_dict: + if k in self.feed_dict: + raise ValueError("Key added twice: {}".format(k)) + self.feed_dict.update(feed_dict) + + def add_fetches(self, fetches): + assert not self._executed + base_index = len(self.fetches) + self.fetches.extend(fetches) + return list(range(base_index, len(self.fetches))) + + def get(self, to_fetch): + if self._executed is None: + try: + self._executed = _run_timeline( + self.session, + self.fetches, + self.debug_name, + self.feed_dict, + os.environ.get("TF_TIMELINE_DIR"), + ) + except Exception as e: + logger.exception( + "Error fetching: {}, feed_dict={}".format( + self.fetches, self.feed_dict + ) + ) + raise e + if isinstance(to_fetch, int): + return self._executed[to_fetch] + elif isinstance(to_fetch, list): + return [self.get(x) for x in to_fetch] + elif isinstance(to_fetch, tuple): + return tuple(self.get(x) for x in to_fetch) + else: + raise ValueError("Unsupported fetch type: {}".format(to_fetch)) + + +_count = 0 + + +def _run_timeline(sess, ops, debug_name, feed_dict=None, timeline_dir=None): + if feed_dict is None: + feed_dict = {} + + if timeline_dir: + from tensorflow.python.client import timeline + + try: + run_options = tf1.RunOptions(trace_level=tf.RunOptions.FULL_TRACE) + except AttributeError: + run_options = None + # In local mode, tf1.RunOptions is not available, see #26511 + if log_once("tf1.RunOptions_not_available"): + logger.exception( + "Can not access tf.RunOptions.FULL_TRACE. This may be because " + "you have used `ray.init(local_mode=True)`. RLlib will use " + "timeline without `options=tf.RunOptions.FULL_TRACE`." + ) + run_metadata = tf1.RunMetadata() + start = time.time() + fetches = sess.run( + ops, options=run_options, run_metadata=run_metadata, feed_dict=feed_dict + ) + trace = timeline.Timeline(step_stats=run_metadata.step_stats) + global _count + outf = os.path.join( + timeline_dir, + "timeline-{}-{}-{}.json".format(debug_name, os.getpid(), _count % 10), + ) + _count += 1 + trace_file = open(outf, "w") + logger.info( + "Wrote tf timeline ({} s) to {}".format( + time.time() - start, os.path.abspath(outf) + ) + ) + trace_file.write(trace.generate_chrome_trace_format()) + else: + if log_once("tf_timeline"): + logger.info( + "Executing TF run without tracing. To dump TF timeline traces " + "to disk, set the TF_TIMELINE_DIR environment variable." + ) + fetches = sess.run(ops, feed_dict=feed_dict) + return fetches diff --git a/lib/python3.12/site-packages/ray/rllib/utils/tf_utils.py b/lib/python3.12/site-packages/ray/rllib/utils/tf_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..7a5617add0b490ce8f5d8801ebc9cf7a0cee5c39 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/tf_utils.py @@ -0,0 +1,1041 @@ +import logging +from collections import OrderedDict, deque +from typing import TYPE_CHECKING, Any, Callable, List, Optional, Type, Union + +import gymnasium as gym +import numpy as np +import tree # pip install dm_tree +from gymnasium.spaces import Discrete, MultiDiscrete + +from ray.rllib.utils import force_list +from ray.rllib.utils.annotations import DeveloperAPI, PublicAPI +from ray.rllib.utils.framework import try_import_tf +from ray.rllib.utils.numpy import SMALL_NUMBER +from ray.rllib.utils.spaces.space_utils import get_base_struct_from_space +from ray.rllib.utils.typing import ( + LocalOptimizer, + ModelGradients, + NetworkType, + PartialAlgorithmConfigDict, + SpaceStruct, + TensorStructType, + TensorType, +) + +if TYPE_CHECKING: + from ray.rllib.algorithms.algorithm_config import AlgorithmConfig + from ray.rllib.core.learner.learner import ParamDict + from ray.rllib.policy.eager_tf_policy import EagerTFPolicy + from ray.rllib.policy.eager_tf_policy_v2 import EagerTFPolicyV2 + from ray.rllib.policy.tf_policy import TFPolicy + +logger = logging.getLogger(__name__) +tf1, tf, tfv = try_import_tf() + + +@PublicAPI +def clip_gradients( + gradients_dict: "ParamDict", + *, + grad_clip: Optional[float] = None, + grad_clip_by: str, +) -> Optional[float]: + """Performs gradient clipping on a grad-dict based on a clip value and clip mode. + + Changes the provided gradient dict in place. + + Args: + gradients_dict: The gradients dict, mapping str to gradient tensors. + grad_clip: The value to clip with. The way gradients are clipped is defined + by the `grad_clip_by` arg (see below). + grad_clip_by: One of 'value', 'norm', or 'global_norm'. + + Returns: + If `grad_clip_by`="global_norm" and `grad_clip` is not None, returns the global + norm of all tensors, otherwise returns None. + """ + # No clipping, return. + if grad_clip is None: + return + + # Clip by value (each gradient individually). + if grad_clip_by == "value": + for k, v in gradients_dict.copy().items(): + gradients_dict[k] = tf.clip_by_value(v, -grad_clip, grad_clip) + + # Clip by L2-norm (per gradient tensor). + elif grad_clip_by == "norm": + for k, v in gradients_dict.copy().items(): + gradients_dict[k] = tf.clip_by_norm(v, grad_clip) + + # Clip by global L2-norm (across all gradient tensors). + else: + assert grad_clip_by == "global_norm" + + clipped_grads, global_norm = tf.clip_by_global_norm( + list(gradients_dict.values()), grad_clip + ) + for k, v in zip(gradients_dict.copy().keys(), clipped_grads): + gradients_dict[k] = v + + # Return the computed global norm scalar. + return global_norm + + +@PublicAPI +def explained_variance(y: TensorType, pred: TensorType) -> TensorType: + """Computes the explained variance for a pair of labels and predictions. + + The formula used is: + max(-1.0, 1.0 - (std(y - pred)^2 / std(y)^2)) + + Args: + y: The labels. + pred: The predictions. + + Returns: + The explained variance given a pair of labels and predictions. + """ + _, y_var = tf.nn.moments(y, axes=[0]) + _, diff_var = tf.nn.moments(y - pred, axes=[0]) + return tf.maximum(-1.0, 1 - (diff_var / (y_var + SMALL_NUMBER))) + + +@PublicAPI +def flatten_inputs_to_1d_tensor( + inputs: TensorStructType, + spaces_struct: Optional[SpaceStruct] = None, + time_axis: bool = False, +) -> TensorType: + """Flattens arbitrary input structs according to the given spaces struct. + + Returns a single 1D tensor resulting from the different input + components' values. + + Thereby: + - Boxes (any shape) get flattened to (B, [T]?, -1). Note that image boxes + are not treated differently from other types of Boxes and get + flattened as well. + - Discrete (int) values are one-hot'd, e.g. a batch of [1, 0, 3] (B=3 with + Discrete(4) space) results in [[0, 1, 0, 0], [1, 0, 0, 0], [0, 0, 0, 1]]. + - MultiDiscrete values are multi-one-hot'd, e.g. a batch of + [[0, 2], [1, 4]] (B=2 with MultiDiscrete([2, 5]) space) results in + [[1, 0, 0, 0, 1, 0, 0], [0, 1, 0, 0, 0, 0, 1]]. + + Args: + inputs: The inputs to be flattened. + spaces_struct: The structure of the spaces that behind the input + time_axis: Whether all inputs have a time-axis (after the batch axis). + If True, will keep not only the batch axis (0th), but the time axis + (1st) as-is and flatten everything from the 2nd axis up. + + Returns: + A single 1D tensor resulting from concatenating all + flattened/one-hot'd input components. Depending on the time_axis flag, + the shape is (B, n) or (B, T, n). + + .. testcode:: + :skipif: True + + # B=2 + from ray.rllib.utils.tf_utils import flatten_inputs_to_1d_tensor + from gymnasium.spaces import Discrete, Box + out = flatten_inputs_to_1d_tensor( + {"a": [1, 0], "b": [[[0.0], [0.1]], [1.0], [1.1]]}, + spaces_struct=dict(a=Discrete(2), b=Box(shape=(2, 1))) + ) + print(out) + + # B=2; T=2 + out = flatten_inputs_to_1d_tensor( + ([[1, 0], [0, 1]], + [[[0.0, 0.1], [1.0, 1.1]], [[2.0, 2.1], [3.0, 3.1]]]), + spaces_struct=tuple([Discrete(2), Box(shape=(2, ))]), + time_axis=True + ) + print(out) + + .. testoutput:: + + [[0.0, 1.0, 0.0, 0.1], [1.0, 0.0, 1.0, 1.1]] # B=2 n=4 + [[[0.0, 1.0, 0.0, 0.1], [1.0, 0.0, 1.0, 1.1]], + [[1.0, 0.0, 2.0, 2.1], [0.0, 1.0, 3.0, 3.1]]] # B=2 T=2 n=4 + """ + + flat_inputs = tree.flatten(inputs) + flat_spaces = ( + tree.flatten(spaces_struct) + if spaces_struct is not None + else [None] * len(flat_inputs) + ) + + B = None + T = None + out = [] + for input_, space in zip(flat_inputs, flat_spaces): + input_ = tf.convert_to_tensor(input_) + shape = tf.shape(input_) + # Store batch and (if applicable) time dimension. + if B is None: + B = shape[0] + if time_axis: + T = shape[1] + + # One-hot encoding. + if isinstance(space, Discrete): + if time_axis: + input_ = tf.reshape(input_, [B * T]) + out.append(tf.cast(one_hot(input_, space), tf.float32)) + elif isinstance(space, MultiDiscrete): + if time_axis: + input_ = tf.reshape(input_, [B * T, -1]) + out.append(tf.cast(one_hot(input_, space), tf.float32)) + # Flatten. + else: + if time_axis: + input_ = tf.reshape(input_, [B * T, -1]) + else: + input_ = tf.reshape(input_, [B, -1]) + out.append(tf.cast(input_, tf.float32)) + + merged = tf.concat(out, axis=-1) + # Restore the time-dimension, if applicable. + if time_axis: + merged = tf.reshape(merged, [B, T, -1]) + + return merged + + +@PublicAPI +def get_gpu_devices() -> List[str]: + """Returns a list of GPU device names, e.g. ["/gpu:0", "/gpu:1"]. + + Supports both tf1.x and tf2.x. + + Returns: + List of GPU device names (str). + """ + if tfv == 1: + from tensorflow.python.client import device_lib + + devices = device_lib.list_local_devices() + else: + try: + devices = tf.config.list_physical_devices() + except Exception: + devices = tf.config.experimental.list_physical_devices() + + # Expect "GPU", but also stuff like: "XLA_GPU". + return [d.name for d in devices if "GPU" in d.device_type] + + +@PublicAPI +def get_placeholder( + *, + space: Optional[gym.Space] = None, + value: Optional[Any] = None, + name: Optional[str] = None, + time_axis: bool = False, + flatten: bool = True, +) -> "tf1.placeholder": + """Returns a tf1.placeholder object given optional hints, such as a space. + + Note that the returned placeholder will always have a leading batch + dimension (None). + + Args: + space: An optional gym.Space to hint the shape and dtype of the + placeholder. + value: An optional value to hint the shape and dtype of the + placeholder. + name: An optional name for the placeholder. + time_axis: Whether the placeholder should also receive a time + dimension (None). + flatten: Whether to flatten the given space into a plain Box space + and then create the placeholder from the resulting space. + + Returns: + The tf1 placeholder. + """ + from ray.rllib.models.catalog import ModelCatalog + + if space is not None: + if isinstance(space, (gym.spaces.Dict, gym.spaces.Tuple)): + if flatten: + return ModelCatalog.get_action_placeholder(space, None) + else: + return tree.map_structure_with_path( + lambda path, component: get_placeholder( + space=component, + name=name + "." + ".".join([str(p) for p in path]), + ), + get_base_struct_from_space(space), + ) + return tf1.placeholder( + shape=(None,) + ((None,) if time_axis else ()) + space.shape, + dtype=tf.float32 if space.dtype == np.float64 else space.dtype, + name=name, + ) + else: + assert value is not None + shape = value.shape[1:] + return tf1.placeholder( + shape=(None,) + + ((None,) if time_axis else ()) + + (shape if isinstance(shape, tuple) else tuple(shape.as_list())), + dtype=tf.float32 if value.dtype == np.float64 else value.dtype, + name=name, + ) + + +@PublicAPI +def get_tf_eager_cls_if_necessary( + orig_cls: Type["TFPolicy"], + config: Union["AlgorithmConfig", PartialAlgorithmConfigDict], +) -> Type[Union["TFPolicy", "EagerTFPolicy", "EagerTFPolicyV2"]]: + """Returns the corresponding tf-eager class for a given TFPolicy class. + + Args: + orig_cls: The original TFPolicy class to get the corresponding tf-eager + class for. + config: The Algorithm config dict or AlgorithmConfig object. + + Returns: + The tf eager policy class corresponding to the given TFPolicy class. + """ + cls = orig_cls + framework = config.get("framework", "tf") + + if framework in ["tf2", "tf"] and not tf1: + raise ImportError("Could not import tensorflow!") + + if framework == "tf2": + if not tf1.executing_eagerly(): + tf1.enable_eager_execution() + assert tf1.executing_eagerly() + + from ray.rllib.policy.eager_tf_policy import EagerTFPolicy + from ray.rllib.policy.eager_tf_policy_v2 import EagerTFPolicyV2 + from ray.rllib.policy.tf_policy import TFPolicy + + # Create eager-class (if not already one). + if hasattr(orig_cls, "as_eager") and not issubclass(orig_cls, EagerTFPolicy): + cls = orig_cls.as_eager() + # Could be some other type of policy or already + # eager-ized. + elif not issubclass(orig_cls, TFPolicy): + pass + else: + raise ValueError( + "This policy does not support eager execution: {}".format(orig_cls) + ) + + # Now that we know, policy is an eager one, add tracing, if necessary. + if config.get("eager_tracing") and issubclass( + cls, (EagerTFPolicy, EagerTFPolicyV2) + ): + cls = cls.with_tracing() + return cls + + +@PublicAPI +def huber_loss(x: TensorType, delta: float = 1.0) -> TensorType: + """Computes the huber loss for a given term and delta parameter. + + Reference: https://en.wikipedia.org/wiki/Huber_loss + Note that the factor of 0.5 is implicitly included in the calculation. + + Formula: + L = 0.5 * x^2 for small abs x (delta threshold) + L = delta * (abs(x) - 0.5*delta) for larger abs x (delta threshold) + + Args: + x: The input term, e.g. a TD error. + delta: The delta parmameter in the above formula. + + Returns: + The Huber loss resulting from `x` and `delta`. + """ + return tf.where( + tf.abs(x) < delta, # for small x -> apply the Huber correction + tf.math.square(x) * 0.5, + delta * (tf.abs(x) - 0.5 * delta), + ) + + +@PublicAPI +def l2_loss(x: TensorType) -> TensorType: + """Computes half the L2 norm over a tensor's values without the sqrt. + + output = 0.5 * sum(x ** 2) + + Args: + x: The input tensor. + + Returns: + 0.5 times the L2 norm over the given tensor's values (w/o sqrt). + """ + return 0.5 * tf.reduce_sum(tf.pow(x, 2.0)) + + +@PublicAPI +def make_tf_callable( + session_or_none: Optional["tf1.Session"], dynamic_shape: bool = False +) -> Callable: + """Returns a function that can be executed in either graph or eager mode. + + The function must take only positional args. + + If eager is enabled, this will act as just a function. Otherwise, it + will build a function that executes a session run with placeholders + internally. + + Args: + session_or_none: tf.Session if in graph mode, else None. + dynamic_shape: True if the placeholders should have a dynamic + batch dimension. Otherwise they will be fixed shape. + + Returns: + A function that can be called in either eager or static-graph mode. + """ + + if tf.executing_eagerly(): + assert session_or_none is None + else: + assert session_or_none is not None + + def make_wrapper(fn): + # Static-graph mode: Create placeholders and make a session call each + # time the wrapped function is called. Returns the output of this + # session call. + if session_or_none is not None: + args_placeholders = [] + kwargs_placeholders = {} + + symbolic_out = [None] + + def call(*args, **kwargs): + args_flat = [] + for a in args: + if type(a) is list: + args_flat.extend(a) + else: + args_flat.append(a) + args = args_flat + + # We have not built any placeholders yet: Do this once here, + # then reuse the same placeholders each time we call this + # function again. + if symbolic_out[0] is None: + with session_or_none.graph.as_default(): + + def _create_placeholders(path, value): + if dynamic_shape: + if len(value.shape) > 0: + shape = (None,) + value.shape[1:] + else: + shape = () + else: + shape = value.shape + return tf1.placeholder( + dtype=value.dtype, + shape=shape, + name=".".join([str(p) for p in path]), + ) + + placeholders = tree.map_structure_with_path( + _create_placeholders, args + ) + for ph in tree.flatten(placeholders): + args_placeholders.append(ph) + + placeholders = tree.map_structure_with_path( + _create_placeholders, kwargs + ) + for k, ph in placeholders.items(): + kwargs_placeholders[k] = ph + + symbolic_out[0] = fn(*args_placeholders, **kwargs_placeholders) + feed_dict = dict(zip(args_placeholders, tree.flatten(args))) + tree.map_structure( + lambda ph, v: feed_dict.__setitem__(ph, v), + kwargs_placeholders, + kwargs, + ) + ret = session_or_none.run(symbolic_out[0], feed_dict) + return ret + + return call + # Eager mode (call function as is). + else: + return fn + + return make_wrapper + + +# TODO (sven): Deprecate this function once we have moved completely to the Learner API. +# Replaced with `clip_gradients()`. +@PublicAPI +def minimize_and_clip( + optimizer: LocalOptimizer, + objective: TensorType, + var_list: List["tf.Variable"], + clip_val: float = 10.0, +) -> ModelGradients: + """Computes, then clips gradients using objective, optimizer and var list. + + Ensures the norm of the gradients for each variable is clipped to + `clip_val`. + + Args: + optimizer: Either a shim optimizer (tf eager) containing a + tf.GradientTape under `self.tape` or a tf1 local optimizer + object. + objective: The loss tensor to calculate gradients on. + var_list: The list of tf.Variables to compute gradients over. + clip_val: The global norm clip value. Will clip around -clip_val and + +clip_val. + + Returns: + The resulting model gradients (list or tuples of grads + vars) + corresponding to the input `var_list`. + """ + # Accidentally passing values < 0.0 will break all gradients. + assert clip_val is None or clip_val > 0.0, clip_val + + if tf.executing_eagerly(): + tape = optimizer.tape + grads_and_vars = list(zip(list(tape.gradient(objective, var_list)), var_list)) + else: + grads_and_vars = optimizer.compute_gradients(objective, var_list=var_list) + + return [ + (tf.clip_by_norm(g, clip_val) if clip_val is not None else g, v) + for (g, v) in grads_and_vars + if g is not None + ] + + +@PublicAPI +def one_hot(x: TensorType, space: gym.Space) -> TensorType: + """Returns a one-hot tensor, given and int tensor and a space. + + Handles the MultiDiscrete case as well. + + Args: + x: The input tensor. + space: The space to use for generating the one-hot tensor. + + Returns: + The resulting one-hot tensor. + + Raises: + ValueError: If the given space is not a discrete one. + + .. testcode:: + :skipif: True + + import gymnasium as gym + import tensorflow as tf + from ray.rllib.utils.tf_utils import one_hot + x = tf.Variable([0, 3], dtype=tf.int32) # batch-dim=2 + # Discrete space with 4 (one-hot) slots per batch item. + s = gym.spaces.Discrete(4) + one_hot(x, s) + + .. testoutput:: + + + + .. testcode:: + :skipif: True + + x = tf.Variable([[0, 1, 2, 3]], dtype=tf.int32) # batch-dim=1 + # MultiDiscrete space with 5 + 4 + 4 + 7 = 20 (one-hot) slots + # per batch item. + s = gym.spaces.MultiDiscrete([5, 4, 4, 7]) + one_hot(x, s) + + .. testoutput:: + + + """ + if isinstance(space, Discrete): + return tf.one_hot(x, space.n, dtype=tf.float32) + elif isinstance(space, MultiDiscrete): + if isinstance(space.nvec[0], np.ndarray): + nvec = np.ravel(space.nvec) + x = tf.reshape(x, (x.shape[0], -1)) + else: + nvec = space.nvec + return tf.concat( + [tf.one_hot(x[:, i], n, dtype=tf.float32) for i, n in enumerate(nvec)], + axis=-1, + ) + else: + raise ValueError("Unsupported space for `one_hot`: {}".format(space)) + + +@PublicAPI +def reduce_mean_ignore_inf(x: TensorType, axis: Optional[int] = None) -> TensorType: + """Same as tf.reduce_mean() but ignores -inf values. + + Args: + x: The input tensor to reduce mean over. + axis: The axis over which to reduce. None for all axes. + + Returns: + The mean reduced inputs, ignoring inf values. + """ + mask = tf.not_equal(x, tf.float32.min) + x_zeroed = tf.where(mask, x, tf.zeros_like(x)) + return tf.math.reduce_sum(x_zeroed, axis) / tf.math.reduce_sum( + tf.cast(mask, tf.float32), axis + ) + + +@PublicAPI +def scope_vars( + scope: Union[str, "tf1.VariableScope"], trainable_only: bool = False +) -> List["tf.Variable"]: + """Get variables inside a given scope. + + Args: + scope: Scope in which the variables reside. + trainable_only: Whether or not to return only the variables that were + marked as trainable. + + Returns: + The list of variables in the given `scope`. + """ + return tf1.get_collection( + tf1.GraphKeys.TRAINABLE_VARIABLES + if trainable_only + else tf1.GraphKeys.VARIABLES, + scope=scope if isinstance(scope, str) else scope.name, + ) + + +@PublicAPI +def symlog(x: "tf.Tensor") -> "tf.Tensor": + """The symlog function as described in [1]: + + [1] Mastering Diverse Domains through World Models - 2023 + D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap + https://arxiv.org/pdf/2301.04104v1.pdf + """ + return tf.math.sign(x) * tf.math.log(tf.math.abs(x) + 1) + + +@PublicAPI +def inverse_symlog(y: "tf.Tensor") -> "tf.Tensor": + """Inverse of the `symlog` function as desribed in [1]: + + [1] Mastering Diverse Domains through World Models - 2023 + D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap + https://arxiv.org/pdf/2301.04104v1.pdf + """ + # To get to symlog inverse, we solve the symlog equation for x: + # y = sign(x) * log(|x| + 1) + # <=> y / sign(x) = log(|x| + 1) + # <=> y = log( x + 1) V x >= 0 + # -y = log(-x + 1) V x < 0 + # <=> exp(y) = x + 1 V x >= 0 + # exp(-y) = -x + 1 V x < 0 + # <=> exp(y) - 1 = x V x >= 0 + # exp(-y) - 1 = -x V x < 0 + # <=> exp(y) - 1 = x V x >= 0 (if x >= 0, then y must also be >= 0) + # -exp(-y) - 1 = x V x < 0 (if x < 0, then y must also be < 0) + # <=> sign(y) * (exp(|y|) - 1) = x + return tf.math.sign(y) * (tf.math.exp(tf.math.abs(y)) - 1) + + +@PublicAPI +def two_hot( + value: "tf.Tensor", + num_buckets: int = 255, + lower_bound: float = -20.0, + upper_bound: float = 20.0, + dtype=None, +): + """Returns a two-hot vector of dim=num_buckets with two entries that are non-zero. + + See [1] for more details: + [1] Mastering Diverse Domains through World Models - 2023 + D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap + https://arxiv.org/pdf/2301.04104v1.pdf + + Entries in the vector represent equally sized buckets within some fixed range + (`lower_bound` to `upper_bound`). + Those entries not 0.0 at positions k and k+1 encode the actual `value` and sum + up to 1.0. They are the weights multiplied by the buckets values at k and k+1 for + retrieving `value`. + + Example: + num_buckets=11 + lower_bound=-5 + upper_bound=5 + value=2.5 + -> [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.5, 0.5, 0.0, 0.0] + -> [-5 -4 -3 -2 -1 0 1 2 3 4 5] (0.5*2 + 0.5*3=2.5) + + Example: + num_buckets=5 + lower_bound=-1 + upper_bound=1 + value=0.1 + -> [0.0, 0.0, 0.8, 0.2, 0.0] + -> [-1 -0.5 0 0.5 1] (0.2*0.5 + 0.8*0=0.1) + + Args: + value: The input tensor of shape (B,) to be two-hot encoded. + num_buckets: The number of buckets to two-hot encode into. + lower_bound: The lower bound value used for the encoding. If input values are + lower than this boundary, they will be encoded as `lower_bound`. + upper_bound: The upper bound value used for the encoding. If input values are + higher than this boundary, they will be encoded as `upper_bound`. + + Returns: + The two-hot encoded tensor of shape (B, num_buckets). + """ + # First make sure, values are clipped. + value = tf.clip_by_value(value, lower_bound, upper_bound) + # Tensor of batch indices: [0, B=batch size). + batch_indices = tf.cast( + tf.range(0, tf.shape(value)[0]), + dtype=dtype or tf.float32, + ) + # Calculate the step deltas (how much space between each bucket's central value?). + bucket_delta = (upper_bound - lower_bound) / (num_buckets - 1) + # Compute the float indices (might be non-int numbers: sitting between two buckets). + idx = (-lower_bound + value) / bucket_delta + # k + k = tf.math.floor(idx) + # k+1 + kp1 = tf.math.ceil(idx) + # In case k == kp1 (idx is exactly on the bucket boundary), move kp1 up by 1.0. + # Otherwise, this would result in a NaN in the returned two-hot tensor. + kp1 = tf.where(tf.equal(k, kp1), kp1 + 1.0, kp1) + # Iff `kp1` is one beyond our last index (because incoming value is larger than + # `upper_bound`), move it to one before k (kp1's weight is going to be 0.0 anyways, + # so it doesn't matter where it points to; we are just avoiding an index error + # with this). + kp1 = tf.where(tf.equal(kp1, num_buckets), kp1 - 2.0, kp1) + # The actual values found at k and k+1 inside the set of buckets. + values_k = lower_bound + k * bucket_delta + values_kp1 = lower_bound + kp1 * bucket_delta + # Compute the two-hot weights (adding up to 1.0) to use at index k and k+1. + weights_k = (value - values_kp1) / (values_k - values_kp1) + weights_kp1 = 1.0 - weights_k + # Compile a tensor of full paths (indices from batch index to feature index) to + # use for the scatter_nd op. + indices_k = tf.stack([batch_indices, k], -1) + indices_kp1 = tf.stack([batch_indices, kp1], -1) + indices = tf.concat([indices_k, indices_kp1], 0) + # The actual values (weights adding up to 1.0) to place at the computed indices. + updates = tf.concat([weights_k, weights_kp1], 0) + # Call the actual scatter update op, returning a zero-filled tensor, only changed + # at the given indices. + return tf.scatter_nd( + tf.cast(indices, tf.int32), + updates, + shape=(tf.shape(value)[0], num_buckets), + ) + + +@PublicAPI +def update_target_network( + main_net: NetworkType, + target_net: NetworkType, + tau: float, +) -> None: + """Updates a keras.Model target network using Polyak averaging. + + new_target_net_weight = ( + tau * main_net_weight + (1.0 - tau) * current_target_net_weight + ) + + Args: + main_net: The keras.Model to update from. + target_net: The target network to update. + tau: The tau value to use in the Polyak averaging formula. + """ + for old_var, current_var in zip(target_net.variables, main_net.variables): + updated_var = tau * current_var + (1.0 - tau) * old_var + old_var.assign(updated_var) + + +@PublicAPI +def zero_logps_from_actions(actions: TensorStructType) -> TensorType: + """Helper function useful for returning dummy logp's (0) for some actions. + + Args: + actions: The input actions. This can be any struct + of complex action components or a simple tensor of different + dimensions, e.g. [B], [B, 2], or {"a": [B, 4, 5], "b": [B]}. + + Returns: + A 1D tensor of 0.0 (dummy logp's) matching the batch + dim of `actions` (shape=[B]). + """ + # Need to flatten `actions` in case we have a complex action space. + # Take the 0th component to extract the batch dim. + action_component = tree.flatten(actions)[0] + logp_ = tf.zeros_like(action_component, dtype=tf.float32) + # Logp's should be single values (but with the same batch dim as + # `deterministic_actions` or `stochastic_actions`). In case + # actions are just [B], zeros_like works just fine here, but if + # actions are [B, ...], we have to reduce logp back to just [B]. + while len(logp_.shape) > 1: + logp_ = logp_[:, 0] + return logp_ + + +@DeveloperAPI +def warn_if_infinite_kl_divergence( + policy: Type["TFPolicy"], mean_kl: TensorType +) -> None: + def print_warning(): + logger.warning( + "KL divergence is non-finite, this will likely destabilize your model and" + " the training process. Action(s) in a specific state have near-zero" + " probability. This can happen naturally in deterministic environments" + " where the optimal policy has zero mass for a specific action. To fix this" + " issue, consider setting the coefficient for the KL loss term to zero or" + " increasing policy entropy." + ) + return tf.constant(0.0) + + if policy.loss_initialized(): + tf.cond( + tf.math.is_inf(mean_kl), + false_fn=lambda: tf.constant(0.0), + true_fn=lambda: print_warning(), + ) + + +def _unflatten(vector, shapes): + i = 0 + arrays = [] + for shape in shapes: + size = np.prod(shape, dtype=np.int_) + array = vector[i : (i + size)].reshape(shape) + arrays.append(array) + i += size + assert len(vector) == i, "Passed weight does not have the correct shape." + return arrays + + +@DeveloperAPI +class TensorFlowVariables: + """A class used to set and get weights for Tensorflow networks. + + Attributes: + sess (tf.Session): The tensorflow session used to run assignment. + variables (Dict[str, tf.Variable]): Extracted variables from the loss + or additional variables that are passed in. + placeholders (Dict[str, tf.placeholders]): Placeholders for weights. + assignment_nodes (Dict[str, tf.Tensor]): Nodes that assign weights. + """ + + def __init__(self, output, sess=None, input_variables=None): + """Creates TensorFlowVariables containing extracted variables. + + The variables are extracted by performing a BFS search on the + dependency graph with loss as the root node. After the tree is + traversed and those variables are collected, we append input_variables + to the collected variables. For each variable in the list, the + variable has a placeholder and assignment operation created for it. + + Args: + output (tf.Operation, List[tf.Operation]): The tensorflow + operation to extract all variables from. + sess (Optional[tf.Session]): Optional tf.Session used for running + the get and set methods in tf graph mode. + Use None for tf eager. + input_variables (List[tf.Variables]): Variables to include in the + list. + """ + self.sess = sess + output = force_list(output) + queue = deque(output) + variable_names = [] + explored_inputs = set(output) + + # We do a BFS on the dependency graph of the input function to find + # the variables. + while len(queue) != 0: + tf_obj = queue.popleft() + if tf_obj is None: + continue + # The object put into the queue is not necessarily an operation, + # so we want the op attribute to get the operation underlying the + # object. Only operations contain the inputs that we can explore. + if hasattr(tf_obj, "op"): + tf_obj = tf_obj.op + for input_op in tf_obj.inputs: + if input_op not in explored_inputs: + queue.append(input_op) + explored_inputs.add(input_op) + # Tensorflow control inputs can be circular, so we keep track of + # explored operations. + for control in tf_obj.control_inputs: + if control not in explored_inputs: + queue.append(control) + explored_inputs.add(control) + if "Variable" in tf_obj.node_def.op or "VarHandle" in tf_obj.node_def.op: + variable_names.append(tf_obj.node_def.name) + self.variables = OrderedDict() + variable_list = [ + v for v in tf1.global_variables() if v.op.node_def.name in variable_names + ] + if input_variables is not None: + variable_list += input_variables + + if not tf1.executing_eagerly(): + for v in variable_list: + self.variables[v.op.node_def.name] = v + + self.placeholders = {} + self.assignment_nodes = {} + + # Create new placeholders to put in custom weights. + for k, var in self.variables.items(): + self.placeholders[k] = tf1.placeholder( + var.value().dtype, + var.get_shape().as_list(), + name="Placeholder_" + k, + ) + self.assignment_nodes[k] = var.assign(self.placeholders[k]) + else: + for v in variable_list: + self.variables[v.name] = v + + def get_flat_size(self): + """Returns the total length of all of the flattened variables. + + Returns: + The length of all flattened variables concatenated. + """ + return sum(np.prod(v.get_shape().as_list()) for v in self.variables.values()) + + def get_flat(self): + """Gets the weights and returns them as a flat array. + + Returns: + 1D Array containing the flattened weights. + """ + # Eager mode. + if not self.sess: + return np.concatenate( + [v.numpy().flatten() for v in self.variables.values()] + ) + # Graph mode. + return np.concatenate( + [v.eval(session=self.sess).flatten() for v in self.variables.values()] + ) + + def set_flat(self, new_weights): + """Sets the weights to new_weights, converting from a flat array. + + Note: + You can only set all weights in the network using this function, + i.e., the length of the array must match get_flat_size. + + Args: + new_weights (np.ndarray): Flat array containing weights. + """ + shapes = [v.get_shape().as_list() for v in self.variables.values()] + arrays = _unflatten(new_weights, shapes) + if not self.sess: + for v, a in zip(self.variables.values(), arrays): + v.assign(a) + else: + placeholders = [self.placeholders[k] for k, v in self.variables.items()] + self.sess.run( + list(self.assignment_nodes.values()), + feed_dict=dict(zip(placeholders, arrays)), + ) + + def get_weights(self): + """Returns a dictionary containing the weights of the network. + + Returns: + Dictionary mapping variable names to their weights. + """ + # Eager mode. + if not self.sess: + return self.variables + # Graph mode. + return self.sess.run(self.variables) + + def set_weights(self, new_weights: dict): + """Sets the weights to new_weights. + + Note: + Can set subsets of variables as well, by only passing in the + variables you want to be set. + + Args: + new_weights: Dictionary mapping variable names to their + weights. + """ + if self.sess is None: + for name, var in self.variables.items(): + var.assign(new_weights[name]) + else: + assign_list, feed_dict = self._assign_weights(new_weights) + self.sess.run(assign_list, feed_dict=feed_dict) + + def _assign_weights(self, weights): + """Sets weigths using exact or closest assignable variable name + + Args: + weights: Dictionary mapping variable names to their + weights. + Returns: + Tuple[List, Dict]: assigned variables list, dict of + placeholders and weights + """ + + assigned = [] + feed_dict = {} + assignable = set(self.assignment_nodes.keys()) + + def nb_common_elem(l1, l2): + return len([e for e in l1 if e in l2]) + + def assign(name, value): + feed_dict[self.placeholders[name]] = value + assigned.append(name) + assignable.remove(name) + + for name, value in weights.items(): + if name in assignable: + assign(name, value) + else: + common = { + var: nb_common_elem(name.split("/"), var.split("/")) + for var in assignable + } + select = [ + close_var + for close_var, cn in sorted(common.items(), key=lambda i: -i[1]) + if cn > 0 and value.shape == self.assignment_nodes[close_var].shape + ] + if select: + assign(select[0], value) + + assert assigned, ( + "No variables in the input matched those in the network. " + "Possible cause: Two networks were defined in the same " + "TensorFlow graph. To fix this, place each network " + "definition in its own tf.Graph." + ) + + assert len(assigned) == len(weights), ( + "All weights couldn't be assigned because no variable " + "had an exact/close name or had same shape" + ) + + return [self.assignment_nodes[v] for v in assigned], feed_dict diff --git a/lib/python3.12/site-packages/ray/rllib/utils/threading.py b/lib/python3.12/site-packages/ray/rllib/utils/threading.py new file mode 100644 index 0000000000000000000000000000000000000000..a9a4461dadbf69afa24ed8e9007fe326640dc7de --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/threading.py @@ -0,0 +1,34 @@ +from typing import Callable + +from ray.rllib.utils.annotations import OldAPIStack + + +@OldAPIStack +def with_lock(func: Callable) -> Callable: + """Use as decorator (@withlock) around object methods that need locking. + + Note: The object must have a self._lock = threading.Lock() property. + Locking thus works on the object level (no two locked methods of the same + object can be called asynchronously). + + Args: + func: The function to decorate/wrap. + + Returns: + The wrapped (object-level locked) function. + """ + + def wrapper(self, *a, **k): + try: + with self._lock: + return func(self, *a, **k) + except AttributeError as e: + if "has no attribute '_lock'" in e.args[0]: + raise AttributeError( + "Object {} must have a `self._lock` property (assigned " + "to a threading.RLock() object in its " + "constructor)!".format(self) + ) + raise e + + return wrapper diff --git a/lib/python3.12/site-packages/ray/rllib/utils/torch_utils.py b/lib/python3.12/site-packages/ray/rllib/utils/torch_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a9238b59ba978c61a3a2258c3ff6e59409e6f9cf --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/torch_utils.py @@ -0,0 +1,886 @@ +import logging +import os +import warnings +from typing import TYPE_CHECKING, Dict, List, Optional, Union + +import gymnasium as gym +import numpy as np +import tree # pip install dm_tree +from gymnasium.spaces import Discrete, MultiDiscrete +from packaging import version + +from ray.rllib.models.repeated_values import RepeatedValues +from ray.rllib.utils.annotations import DeveloperAPI, OldAPIStack, PublicAPI +from ray.rllib.utils.framework import try_import_torch +from ray.rllib.utils.numpy import SMALL_NUMBER +from ray.rllib.utils.typing import ( + LocalOptimizer, + NetworkType, + SpaceStruct, + TensorStructType, + TensorType, +) + +if TYPE_CHECKING: + from ray.rllib.core.learner.learner import ParamDict, ParamList + from ray.rllib.policy.torch_policy import TorchPolicy + from ray.rllib.policy.torch_policy_v2 import TorchPolicyV2 + +logger = logging.getLogger(__name__) +torch, nn = try_import_torch() + +# Limit values suitable for use as close to a -inf logit. These are useful +# since -inf / inf cause NaNs during backprop. +FLOAT_MIN = -3.4e38 +FLOAT_MAX = 3.4e38 + +if torch: + TORCH_COMPILE_REQUIRED_VERSION = version.parse("2.0.0") +else: + TORCH_COMPILE_REQUIRED_VERSION = ValueError( + "torch is not installed. TORCH_COMPILE_REQUIRED_VERSION is not defined." + ) + + +@OldAPIStack +def apply_grad_clipping( + policy: "TorchPolicy", optimizer: LocalOptimizer, loss: TensorType +) -> Dict[str, TensorType]: + """Applies gradient clipping to already computed grads inside `optimizer`. + + Note: This function does NOT perform an analogous operation as + tf.clip_by_global_norm. It merely clips by norm (per gradient tensor) and + then computes the global norm across all given tensors (but without clipping + by that global norm). + + Args: + policy: The TorchPolicy, which calculated `loss`. + optimizer: A local torch optimizer object. + loss: The torch loss tensor. + + Returns: + An info dict containing the "grad_norm" key and the resulting clipped + gradients. + """ + grad_gnorm = 0 + if policy.config["grad_clip"] is not None: + clip_value = policy.config["grad_clip"] + else: + clip_value = np.inf + + num_none_grads = 0 + for param_group in optimizer.param_groups: + # Make sure we only pass params with grad != None into torch + # clip_grad_norm_. Would fail otherwise. + params = list(filter(lambda p: p.grad is not None, param_group["params"])) + if params: + # PyTorch clips gradients inplace and returns the norm before clipping + # We therefore need to compute grad_gnorm further down (fixes #4965) + global_norm = nn.utils.clip_grad_norm_(params, clip_value) + + if isinstance(global_norm, torch.Tensor): + global_norm = global_norm.cpu().numpy() + + grad_gnorm += min(global_norm, clip_value) + else: + num_none_grads += 1 + + # Note (Kourosh): grads could indeed be zero. This method should still return + # grad_gnorm in that case. + if num_none_grads == len(optimizer.param_groups): + # No grads available + return {} + return {"grad_gnorm": grad_gnorm} + + +@PublicAPI +def clip_gradients( + gradients_dict: "ParamDict", + *, + grad_clip: Optional[float] = None, + grad_clip_by: str = "value", +) -> TensorType: + """Performs gradient clipping on a grad-dict based on a clip value and clip mode. + + Changes the provided gradient dict in place. + + Args: + gradients_dict: The gradients dict, mapping str to gradient tensors. + grad_clip: The value to clip with. The way gradients are clipped is defined + by the `grad_clip_by` arg (see below). + grad_clip_by: One of 'value', 'norm', or 'global_norm'. + + Returns: + If `grad_clip_by`="global_norm" and `grad_clip` is not None, returns the global + norm of all tensors, otherwise returns None. + """ + # No clipping, return. + if grad_clip is None: + return + + if grad_clip_by not in ["value", "norm", "global_norm"]: + raise ValueError( + f"`grad_clip_by` ({grad_clip_by}) must be one of [value|norm|global_norm]!" + ) + + # Clip by value (each gradient individually). + if grad_clip_by == "value": + for k, v in gradients_dict.items(): + gradients_dict[k] = ( + None if v is None else torch.clip(v, -grad_clip, grad_clip) + ) + + # Clip by L2-norm (per gradient tensor). + elif grad_clip_by == "norm": + for k, v in gradients_dict.items(): + if v is not None: + # Compute the L2-norm of the gradient tensor. + norm = v.norm(2).nan_to_num(neginf=-10e8, posinf=10e8) + # Clip all the gradients. + if norm > grad_clip: + v.mul_(grad_clip / norm) + + # Clip by global L2-norm (across all gradient tensors). + else: + gradients_list = list(gradients_dict.values()) + total_norm = compute_global_norm(gradients_list) + if len(gradients_list) == 0: + return total_norm + # We do want the coefficient to be in between 0.0 and 1.0, therefore + # if the global_norm is smaller than the clip value, we use the clip value + # as normalization constant. + clip_coeff = grad_clip / torch.clamp(total_norm + 1e-6, min=grad_clip) + # Note: multiplying by the clamped coefficient is redundant when the coefficient + # is clamped to 1, but doing so avoids a `if clip_coeff < 1:` conditional which + # can require a CPU <=> device synchronization when the gradients reside in GPU + # memory. + clip_coeff_clamped = torch.clamp(clip_coeff, max=1.0) + for g in gradients_list: + if g is not None: + g.detach().mul_(clip_coeff_clamped.to(g.device)) + return total_norm + + +@PublicAPI +def compute_global_norm(gradients_list: "ParamList") -> TensorType: + """Computes the global norm for a gradients dict. + + Args: + gradients_list: The gradients list containing parameters. + + Returns: + Returns the global norm of all tensors in `gradients_list`. + """ + # Define the norm type to be L2. + norm_type = 2.0 + # If we have no grads, return zero. + if len(gradients_list) == 0: + return torch.tensor(0.0) + + # Compute the global norm. + total_norm = torch.norm( + torch.stack( + [ + torch.norm(g.detach(), norm_type) + # Note, we want to avoid overflow in the norm computation, this does + # not affect the gradients themselves as we clamp by multiplying and + # not by overriding tensor values. + .nan_to_num(neginf=-10e8, posinf=10e8) + for g in gradients_list + if g is not None + ] + ), + norm_type, + ).nan_to_num(neginf=-10e8, posinf=10e8) + + # Return the global norm. + return total_norm + + +@OldAPIStack +def concat_multi_gpu_td_errors( + policy: Union["TorchPolicy", "TorchPolicyV2"] +) -> Dict[str, TensorType]: + """Concatenates multi-GPU (per-tower) TD error tensors given TorchPolicy. + + TD-errors are extracted from the TorchPolicy via its tower_stats property. + + Args: + policy: The TorchPolicy to extract the TD-error values from. + + Returns: + A dict mapping strings "td_error" and "mean_td_error" to the + corresponding concatenated and mean-reduced values. + """ + td_error = torch.cat( + [ + t.tower_stats.get("td_error", torch.tensor([0.0])).to(policy.device) + for t in policy.model_gpu_towers + ], + dim=0, + ) + policy.td_error = td_error + return { + "td_error": td_error, + "mean_td_error": torch.mean(td_error), + } + + +@PublicAPI +def convert_to_torch_tensor( + x, + device: Optional[str] = None, + pin_memory: bool = False, + use_stream: bool = False, + stream: Optional[Union["torch.cuda.Stream", "torch.cuda.classes.Stream"]] = None, +): + """ + Converts any (possibly nested) structure to torch.Tensors. + + Args: + x: The input structure whose leaves will be converted. + device: The device to create the tensor on (e.g. "cuda:0" or "cpu"). + pin_memory: If True, calls `pin_memory()` on the created tensors. + use_stream: If True, uses a separate CUDA stream for `Tensor.to()`. + stream: An optional CUDA stream for the host-to-device copy in `Tensor.to()`. + + Returns: + A new structure with the same layout as `x` but with all leaves converted + to torch.Tensors. Leaves that are None are left unchanged. + """ + + # Convert the provided device (if any) to a torch.device; default to CPU. + device = torch.device(device) if device is not None else torch.device("cpu") + is_cuda = (device.type == "cuda") and torch.cuda.is_available() + + # Determine the appropriate stream. + if is_cuda: + if use_stream: + if stream is not None: + # Ensure the provided stream is of an acceptable type. + assert isinstance( + stream, (torch.cuda.Stream, torch.cuda.classes.Stream) + ), f"`stream` must be a torch.cuda.Stream but got {type(stream)}." + else: + stream = torch.cuda.Stream() + else: + stream = torch.cuda.default_stream(device=device) + else: + stream = None + + def mapping(item): + # Pass through None values. + if item is None: + return item + + # Special handling for "RepeatedValues" types. + if isinstance(item, RepeatedValues): + return RepeatedValues( + tree.map_structure(mapping, item.values), + item.lengths, + item.max_len, + ) + + # Convert to a tensor if not already one. + if torch.is_tensor(item): + tensor = item + elif isinstance(item, np.ndarray): + # Leave object or string arrays as is. + if item.dtype == object or item.dtype.type is np.str_: + return item + # If the numpy array is not writable, suppress warnings. + if not item.flags.writeable: + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + tensor = torch.from_numpy(item) + else: + tensor = torch.from_numpy(item) + else: + tensor = torch.from_numpy(np.asarray(item)) + + # Convert floating-point tensors from float64 to float32 (unless they are float16). + if tensor.is_floating_point() and tensor.dtype != torch.float16: + tensor = tensor.float() + + # Optionally pin memory for faster host-to-GPU copies. + if pin_memory and is_cuda: + tensor = tensor.pin_memory() + + # Move the tensor to the desired device. + # For CUDA devices, use the provided stream context if available. + if is_cuda: + if stream is not None: + with torch.cuda.stream(stream): + tensor = tensor.to(device, non_blocking=True) + else: + tensor = tensor.to(device, non_blocking=True) + else: + # For CPU (or non-CUDA), this is a no-op if already on the target device. + tensor = tensor.to(device) + + return tensor + + return tree.map_structure(mapping, x) + + +@PublicAPI +def copy_torch_tensors(x: TensorStructType, device: Optional[str] = None): + """Creates a copy of `x` and makes deep copies torch.Tensors in x. + + Also moves the copied tensors to the specified device (if not None). + + Note if an object in x is not a torch.Tensor, it will be shallow-copied. + + Args: + x : Any (possibly nested) struct possibly containing torch.Tensors. + device : The device to move the tensors to. + + Returns: + Any: A new struct with the same structure as `x`, but with all + torch.Tensors deep-copied and moved to the specified device. + + """ + + def mapping(item): + if isinstance(item, torch.Tensor): + return ( + torch.clone(item.detach()) + if device is None + else item.detach().to(device) + ) + else: + return item + + return tree.map_structure(mapping, x) + + +@PublicAPI +def explained_variance(y: TensorType, pred: TensorType) -> TensorType: + """Computes the explained variance for a pair of labels and predictions. + + The formula used is: + max(-1.0, 1.0 - (std(y - pred)^2 / std(y)^2)) + + Args: + y: The labels. + pred: The predictions. + + Returns: + The explained variance given a pair of labels and predictions. + """ + squeezed_y = y.squeeze() + y_var = torch.var(squeezed_y, dim=0) + diff_var = torch.var(squeezed_y - pred.squeeze(), dim=0) + min_ = torch.tensor([-1.0]).to(pred.device) + return torch.max(min_, 1 - (diff_var / (y_var + SMALL_NUMBER)))[0] + + +@PublicAPI +def flatten_inputs_to_1d_tensor( + inputs: TensorStructType, + spaces_struct: Optional[SpaceStruct] = None, + time_axis: bool = False, +) -> TensorType: + """Flattens arbitrary input structs according to the given spaces struct. + + Returns a single 1D tensor resulting from the different input + components' values. + + Thereby: + - Boxes (any shape) get flattened to (B, [T]?, -1). Note that image boxes + are not treated differently from other types of Boxes and get + flattened as well. + - Discrete (int) values are one-hot'd, e.g. a batch of [1, 0, 3] (B=3 with + Discrete(4) space) results in [[0, 1, 0, 0], [1, 0, 0, 0], [0, 0, 0, 1]]. + - MultiDiscrete values are multi-one-hot'd, e.g. a batch of + [[0, 2], [1, 4]] (B=2 with MultiDiscrete([2, 5]) space) results in + [[1, 0, 0, 0, 1, 0, 0], [0, 1, 0, 0, 0, 0, 1]]. + + Args: + inputs: The inputs to be flattened. + spaces_struct: The structure of the spaces that behind the input + time_axis: Whether all inputs have a time-axis (after the batch axis). + If True, will keep not only the batch axis (0th), but the time axis + (1st) as-is and flatten everything from the 2nd axis up. + + Returns: + A single 1D tensor resulting from concatenating all + flattened/one-hot'd input components. Depending on the time_axis flag, + the shape is (B, n) or (B, T, n). + + .. testcode:: + + from gymnasium.spaces import Discrete, Box + from ray.rllib.utils.torch_utils import flatten_inputs_to_1d_tensor + import torch + struct = { + "a": np.array([1, 3]), + "b": ( + np.array([[1.0, 2.0], [4.0, 5.0]]), + np.array( + [[[8.0], [7.0]], [[5.0], [4.0]]] + ), + ), + "c": { + "cb": np.array([1.0, 2.0]), + }, + } + struct_torch = tree.map_structure(lambda s: torch.from_numpy(s), struct) + spaces = dict( + { + "a": gym.spaces.Discrete(4), + "b": (gym.spaces.Box(-1.0, 10.0, (2,)), gym.spaces.Box(-1.0, 1.0, (2, + 1))), + "c": dict( + { + "cb": gym.spaces.Box(-1.0, 1.0, ()), + } + ), + } + ) + print(flatten_inputs_to_1d_tensor(struct_torch, spaces_struct=spaces)) + + .. testoutput:: + + tensor([[0., 1., 0., 0., 1., 2., 8., 7., 1.], + [0., 0., 0., 1., 4., 5., 5., 4., 2.]]) + + """ + + flat_inputs = tree.flatten(inputs) + flat_spaces = ( + tree.flatten(spaces_struct) + if spaces_struct is not None + else [None] * len(flat_inputs) + ) + + B = None + T = None + out = [] + for input_, space in zip(flat_inputs, flat_spaces): + # Store batch and (if applicable) time dimension. + if B is None: + B = input_.shape[0] + if time_axis: + T = input_.shape[1] + + # One-hot encoding. + if isinstance(space, Discrete): + if time_axis: + input_ = torch.reshape(input_, [B * T]) + out.append(one_hot(input_, space).float()) + # Multi one-hot encoding. + elif isinstance(space, MultiDiscrete): + if time_axis: + input_ = torch.reshape(input_, [B * T, -1]) + out.append(one_hot(input_, space).float()) + # Box: Flatten. + else: + if time_axis: + input_ = torch.reshape(input_, [B * T, -1]) + else: + input_ = torch.reshape(input_, [B, -1]) + out.append(input_.float()) + + merged = torch.cat(out, dim=-1) + # Restore the time-dimension, if applicable. + if time_axis: + merged = torch.reshape(merged, [B, T, -1]) + + return merged + + +@PublicAPI +def global_norm(tensors: List[TensorType]) -> TensorType: + """Returns the global L2 norm over a list of tensors. + + output = sqrt(SUM(t ** 2 for t in tensors)), + where SUM reduces over all tensors and over all elements in tensors. + + Args: + tensors: The list of tensors to calculate the global norm over. + + Returns: + The global L2 norm over the given tensor list. + """ + # List of single tensors' L2 norms: SQRT(SUM(xi^2)) over all xi in tensor. + single_l2s = [torch.pow(torch.sum(torch.pow(t, 2.0)), 0.5) for t in tensors] + # Compute global norm from all single tensors' L2 norms. + return torch.pow(sum(torch.pow(l2, 2.0) for l2 in single_l2s), 0.5) + + +@OldAPIStack +def huber_loss(x: TensorType, delta: float = 1.0) -> TensorType: + """Computes the huber loss for a given term and delta parameter. + + Reference: https://en.wikipedia.org/wiki/Huber_loss + Note that the factor of 0.5 is implicitly included in the calculation. + + Formula: + L = 0.5 * x^2 for small abs x (delta threshold) + L = delta * (abs(x) - 0.5*delta) for larger abs x (delta threshold) + + Args: + x: The input term, e.g. a TD error. + delta: The delta parmameter in the above formula. + + Returns: + The Huber loss resulting from `x` and `delta`. + """ + return torch.where( + torch.abs(x) < delta, + torch.pow(x, 2.0) * 0.5, + delta * (torch.abs(x) - 0.5 * delta), + ) + + +@OldAPIStack +def l2_loss(x: TensorType) -> TensorType: + """Computes half the L2 norm over a tensor's values without the sqrt. + + output = 0.5 * sum(x ** 2) + + Args: + x: The input tensor. + + Returns: + 0.5 times the L2 norm over the given tensor's values (w/o sqrt). + """ + return 0.5 * torch.sum(torch.pow(x, 2.0)) + + +@PublicAPI +def one_hot(x: TensorType, space: gym.Space) -> TensorType: + """Returns a one-hot tensor, given and int tensor and a space. + + Handles the MultiDiscrete case as well. + + Args: + x: The input tensor. + space: The space to use for generating the one-hot tensor. + + Returns: + The resulting one-hot tensor. + + Raises: + ValueError: If the given space is not a discrete one. + + .. testcode:: + + import torch + import gymnasium as gym + from ray.rllib.utils.torch_utils import one_hot + x = torch.IntTensor([0, 3]) # batch-dim=2 + # Discrete space with 4 (one-hot) slots per batch item. + s = gym.spaces.Discrete(4) + print(one_hot(x, s)) + x = torch.IntTensor([[0, 1, 2, 3]]) # batch-dim=1 + # MultiDiscrete space with 5 + 4 + 4 + 7 = 20 (one-hot) slots + # per batch item. + s = gym.spaces.MultiDiscrete([5, 4, 4, 7]) + print(one_hot(x, s)) + + .. testoutput:: + + tensor([[1, 0, 0, 0], + [0, 0, 0, 1]]) + tensor([[1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0, 0]]) + """ + if isinstance(space, Discrete): + return nn.functional.one_hot(x.long(), space.n) + elif isinstance(space, MultiDiscrete): + if isinstance(space.nvec[0], np.ndarray): + nvec = np.ravel(space.nvec) + x = x.reshape(x.shape[0], -1) + else: + nvec = space.nvec + return torch.cat( + [nn.functional.one_hot(x[:, i].long(), n) for i, n in enumerate(nvec)], + dim=-1, + ) + else: + raise ValueError("Unsupported space for `one_hot`: {}".format(space)) + + +@PublicAPI +def reduce_mean_ignore_inf(x: TensorType, axis: Optional[int] = None) -> TensorType: + """Same as torch.mean() but ignores -inf values. + + Args: + x: The input tensor to reduce mean over. + axis: The axis over which to reduce. None for all axes. + + Returns: + The mean reduced inputs, ignoring inf values. + """ + mask = torch.ne(x, float("-inf")) + x_zeroed = torch.where(mask, x, torch.zeros_like(x)) + return torch.sum(x_zeroed, axis) / torch.sum(mask.float(), axis) + + +@PublicAPI +def sequence_mask( + lengths: TensorType, + maxlen: Optional[int] = None, + dtype=None, + time_major: bool = False, +) -> TensorType: + """Offers same behavior as tf.sequence_mask for torch. + + Thanks to Dimitris Papatheodorou + (https://discuss.pytorch.org/t/pytorch-equivalent-for-tf-sequence-mask/ + 39036). + + Args: + lengths: The tensor of individual lengths to mask by. + maxlen: The maximum length to use for the time axis. If None, use + the max of `lengths`. + dtype: The torch dtype to use for the resulting mask. + time_major: Whether to return the mask as [B, T] (False; default) or + as [T, B] (True). + + Returns: + The sequence mask resulting from the given input and parameters. + """ + # If maxlen not given, use the longest lengths in the `lengths` tensor. + if maxlen is None: + maxlen = lengths.max() + + mask = torch.ones(tuple(lengths.shape) + (int(maxlen),)) + + mask = ~(mask.to(lengths.device).cumsum(dim=1).t() > lengths) + # Time major transformation. + if not time_major: + mask = mask.t() + + # By default, set the mask to be boolean. + mask.type(dtype or torch.bool) + + return mask + + +@PublicAPI +def update_target_network( + main_net: NetworkType, + target_net: NetworkType, + tau: float, +) -> None: + """Updates a torch.nn.Module target network using Polyak averaging. + + .. code-block:: text + + new_target_net_weight = ( + tau * main_net_weight + (1.0 - tau) * current_target_net_weight + ) + + Args: + main_net: The nn.Module to update from. + target_net: The target network to update. + tau: The tau value to use in the Polyak averaging formula. + """ + # Get the current parameters from the Q network. + state_dict = main_net.state_dict() + # Use here Polyak averaging. + new_state_dict = { + k: tau * state_dict[k] + (1 - tau) * v + for k, v in target_net.state_dict().items() + } + # Apply the new parameters to the target Q network. + target_net.load_state_dict(new_state_dict) + + +@DeveloperAPI +def warn_if_infinite_kl_divergence( + policy: "TorchPolicy", + kl_divergence: TensorType, +) -> None: + if policy.loss_initialized() and kl_divergence.isinf(): + logger.warning( + "KL divergence is non-finite, this will likely destabilize your model and" + " the training process. Action(s) in a specific state have near-zero" + " probability. This can happen naturally in deterministic environments" + " where the optimal policy has zero mass for a specific action. To fix this" + " issue, consider setting the coefficient for the KL loss term to zero or" + " increasing policy entropy." + ) + + +@PublicAPI +def set_torch_seed(seed: Optional[int] = None) -> None: + """Sets the torch random seed to the given value. + + Args: + seed: The seed to use or None for no seeding. + """ + if seed is not None and torch: + torch.manual_seed(seed) + # See https://github.com/pytorch/pytorch/issues/47672. + cuda_version = torch.version.cuda + if cuda_version is not None and float(torch.version.cuda) >= 10.2: + # See https://docs.nvidia.com/cuda/cublas/index.html#results-reproducibility. + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) # if using multi-GPU + else: + if version.Version(torch.__version__) >= version.Version("1.8.0"): + # Not all Operations support this. + torch.use_deterministic_algorithms(True) + else: + torch.set_deterministic(True) + # This is only for Convolution no problem. + torch.backends.cudnn.deterministic = True + # For benchmark=True, CuDNN may choose different algorithms depending on runtime + # conditions or slight differences in input sizes, even if the seed is fixed, + # which breaks determinism. + torch.backends.cudnn.benchmark = False + + +@PublicAPI +def softmax_cross_entropy_with_logits( + logits: TensorType, + labels: TensorType, +) -> TensorType: + """Same behavior as tf.nn.softmax_cross_entropy_with_logits. + + Args: + x: The input predictions. + labels: The labels corresponding to `x`. + + Returns: + The resulting softmax cross-entropy given predictions and labels. + """ + return torch.sum(-labels * nn.functional.log_softmax(logits, -1), -1) + + +@PublicAPI +def symlog(x: "torch.Tensor") -> "torch.Tensor": + """The symlog function as described in [1]: + + [1] Mastering Diverse Domains through World Models - 2023 + D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap + https://arxiv.org/pdf/2301.04104v1.pdf + """ + return torch.sign(x) * torch.log(torch.abs(x) + 1) + + +@PublicAPI +def inverse_symlog(y: "torch.Tensor") -> "torch.Tensor": + """Inverse of the `symlog` function as desribed in [1]: + + [1] Mastering Diverse Domains through World Models - 2023 + D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap + https://arxiv.org/pdf/2301.04104v1.pdf + """ + # To get to symlog inverse, we solve the symlog equation for x: + # y = sign(x) * log(|x| + 1) + # <=> y / sign(x) = log(|x| + 1) + # <=> y = log( x + 1) V x >= 0 + # -y = log(-x + 1) V x < 0 + # <=> exp(y) = x + 1 V x >= 0 + # exp(-y) = -x + 1 V x < 0 + # <=> exp(y) - 1 = x V x >= 0 + # exp(-y) - 1 = -x V x < 0 + # <=> exp(y) - 1 = x V x >= 0 (if x >= 0, then y must also be >= 0) + # -exp(-y) - 1 = x V x < 0 (if x < 0, then y must also be < 0) + # <=> sign(y) * (exp(|y|) - 1) = x + return torch.sign(y) * (torch.exp(torch.abs(y)) - 1) + + +@PublicAPI +def two_hot( + value: "torch.Tensor", + num_buckets: int = 255, + lower_bound: float = -20.0, + upper_bound: float = 20.0, + device: Optional[str] = None, +): + """Returns a two-hot vector of dim=num_buckets with two entries that are non-zero. + + See [1] for more details: + [1] Mastering Diverse Domains through World Models - 2023 + D. Hafner, J. Pasukonis, J. Ba, T. Lillicrap + https://arxiv.org/pdf/2301.04104v1.pdf + + Entries in the vector represent equally sized buckets within some fixed range + (`lower_bound` to `upper_bound`). + Those entries not 0.0 at positions k and k+1 encode the actual `value` and sum + up to 1.0. They are the weights multiplied by the buckets values at k and k+1 for + retrieving `value`. + + Example: + num_buckets=11 + lower_bound=-5 + upper_bound=5 + value=2.5 + -> [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.5, 0.5, 0.0, 0.0] + -> [-5 -4 -3 -2 -1 0 1 2 3 4 5] (0.5*2 + 0.5*3=2.5) + + Example: + num_buckets=5 + lower_bound=-1 + upper_bound=1 + value=0.1 + -> [0.0, 0.0, 0.8, 0.2, 0.0] + -> [-1 -0.5 0 0.5 1] (0.2*0.5 + 0.8*0=0.1) + + Args: + value: The input tensor of shape (B,) to be two-hot encoded. + num_buckets: The number of buckets to two-hot encode into. + lower_bound: The lower bound value used for the encoding. If input values are + lower than this boundary, they will be encoded as `lower_bound`. + upper_bound: The upper bound value used for the encoding. If input values are + higher than this boundary, they will be encoded as `upper_bound`. + + Returns: + The two-hot encoded tensor of shape (B, num_buckets). + """ + # First make sure, values are clipped. + value = torch.clamp(value, lower_bound, upper_bound) + # Tensor of batch indices: [0, B=batch size). + batch_indices = torch.arange(0, value.shape[0], device=device).float() + # Calculate the step deltas (how much space between each bucket's central value?). + bucket_delta = (upper_bound - lower_bound) / (num_buckets - 1) + # Compute the float indices (might be non-int numbers: sitting between two buckets). + idx = (-lower_bound + value) / bucket_delta + # k + k = torch.floor(idx) + # k+1 + kp1 = torch.ceil(idx) + # In case k == kp1 (idx is exactly on the bucket boundary), move kp1 up by 1.0. + # Otherwise, this would result in a NaN in the returned two-hot tensor. + kp1 = torch.where(k.eq(kp1), kp1 + 1.0, kp1) + # Iff `kp1` is one beyond our last index (because incoming value is larger than + # `upper_bound`), move it to one before k (kp1's weight is going to be 0.0 anyways, + # so it doesn't matter where it points to; we are just avoiding an index error + # with this). + kp1 = torch.where(kp1.eq(num_buckets), kp1 - 2.0, kp1) + # The actual values found at k and k+1 inside the set of buckets. + values_k = lower_bound + k * bucket_delta + values_kp1 = lower_bound + kp1 * bucket_delta + # Compute the two-hot weights (adding up to 1.0) to use at index k and k+1. + weights_k = (value - values_kp1) / (values_k - values_kp1) + weights_kp1 = 1.0 - weights_k + # Compile a tensor of full paths (indices from batch index to feature index) to + # use for the scatter_nd op. + indices_k = torch.stack([batch_indices, k], dim=-1) + indices_kp1 = torch.stack([batch_indices, kp1], dim=-1) + indices = torch.cat([indices_k, indices_kp1], dim=0).long() + # The actual values (weights adding up to 1.0) to place at the computed indices. + updates = torch.cat([weights_k, weights_kp1], dim=0) + # Call the actual scatter update op, returning a zero-filled tensor, only changed + # at the given indices. + output = torch.zeros(value.shape[0], num_buckets, device=device) + # Set our two-hot values at computed indices. + output[indices[:, 0], indices[:, 1]] = updates + return output + + +def _dynamo_is_available(): + # This only works if torch._dynamo is available + try: + # TODO(Artur): Remove this once torch._dynamo is available on CI + import torch._dynamo as dynamo # noqa: F401 + + return True + except ImportError: + return False diff --git a/lib/python3.12/site-packages/ray/rllib/utils/typing.py b/lib/python3.12/site-packages/ray/rllib/utils/typing.py new file mode 100644 index 0000000000000000000000000000000000000000..5c8bd00e91365a5024929d0ed15b7ffaba94e8f0 --- /dev/null +++ b/lib/python3.12/site-packages/ray/rllib/utils/typing.py @@ -0,0 +1,378 @@ +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + Hashable, + List, + Optional, + Sequence, + Tuple, + Type, + TypeVar, + Union, +) + +import gymnasium as gym + +from ray.rllib.utils.annotations import OldAPIStack + +if TYPE_CHECKING: + # Modules might be missing but supply users with type hints if they are installed. + import jax.numpy as jnp + import keras + import tensorflow as tf + import torch + from numpy.typing import NDArray + + from ray.rllib.core.rl_module.multi_rl_module import MultiRLModuleSpec + from ray.rllib.core.rl_module.rl_module import RLModuleSpec + from ray.rllib.env.env_context import EnvContext + from ray.rllib.env.multi_agent_episode import MultiAgentEpisode + from ray.rllib.env.single_agent_episode import SingleAgentEpisode + from ray.rllib.policy.dynamic_tf_policy_v2 import DynamicTFPolicyV2 + from ray.rllib.policy.eager_tf_policy_v2 import EagerTFPolicyV2 + from ray.rllib.policy.policy import PolicySpec + from ray.rllib.policy.sample_batch import MultiAgentBatch, SampleBatch + from ray.rllib.policy.view_requirement import ViewRequirement + + +TensorType = Union["NDArray[Any]", "jnp.ndarray", "tf.Tensor", "torch.Tensor"] +""" +Represents a generic tensor type. +This could be an np.ndarray, jnp.ndarray, tf.Tensor, or a torch.Tensor. +""" + +TensorStructType = Union[TensorType, dict, tuple] +"""Either a plain tensor, or a dict or tuple of tensors (or StructTensors).""" + +# A shape of a tensor. +TensorShape = Union[Tuple[int, ...], List[int]] + +NetworkType = Union["torch.nn.Module", "keras.Model"] +"""A neural network.""" + +DeviceType = Union[str, "torch.device", "int"] +""" +A device identifier, which can be a string (e.g. 'cpu', 'cuda:0'), +a torch.device object, or other types supported by torch. +""" + +RLModuleSpecType = Union["RLModuleSpec", "MultiRLModuleSpec"] +"""An RLModule spec (single-agent or multi-agent).""" + +StateDict = Dict[str, Any] +"""A state dict of an RLlib component (e.g. EnvRunner, Learner, RLModule).""" + +AlgorithmConfigDict = dict # @OldAPIStack +""" +Represents a fully filled out config of a Algorithm class. + +Note: + Policy config dicts are usually the same as AlgorithmConfigDict, but + parts of it may sometimes be altered in e.g. a multi-agent setup, + where we have >1 Policies in the same Algorithm. +""" + +PartialAlgorithmConfigDict = dict # @OldAPIStack +""" +An algorithm config dict that only has overrides. It needs to be combined with +the default algorithm config to be used. +""" + +ModelConfigDict = dict # @OldAPIStack +""" +Represents the model config sub-dict of the algo config that is passed to the +model catalog. +""" + +ConvFilterSpec = List[ + Tuple[int, Union[int, Tuple[int, int]], Union[int, Tuple[int, int]]] +] +""" +Conv2D configuration format. Each entry in the outer list represents one Conv2D +layer. Each inner list has the format: [num_output_filters, kernel, stride], where +kernel and stride may be single ints (width and height are the same) or 2-tuples +(int, int) for width and height (different values). +""" + +FromConfigSpec = Union[Dict[str, Union[Any, type, str]], type, str] +""" +Objects that can be created through the `from_config()` util method +need a config dict with a "type" key, a class path (str), or a type directly. +""" + +EnvConfigDict = dict +""" +Represents the env_config sub-dict of the algo config that is passed to +the env constructor. +""" + +EnvID = Union[int, str] +""" +Represents an environment id. These could be: +- An int index for a sub-env within a vectorized env. +- An external env ID (str), which changes(!) each episode. +""" + +# TODO (sven): Specify this type more strictly (it should just be gym.Env). +EnvType = Union[Any, gym.Env] +""" +Represents a BaseEnv, MultiAgentEnv, ExternalEnv, ExternalMultiAgentEnv, +VectorEnv, gym.Env, or ActorHandle. +""" + +EnvCreator = Callable[["EnvContext"], Optional[EnvType]] +""" +A callable, taking a EnvContext object +(config dict + properties: `worker_index`, `vector_index`, `num_workers`, +and `remote`) and returning an env object (or None if no env is used). +""" + +AgentID = Hashable +"""Represents a generic identifier for an agent (e.g., "agent1").""" + +PolicyID = str # @OldAPIStack +"""Represents a generic identifier for a policy (e.g., "pol1").""" + +ModuleID = str +"""Represents a generic identifier for a (single-agent) RLModule.""" + +MultiAgentPolicyConfigDict = Dict[PolicyID, "PolicySpec"] # @OldAPIStack +"""Type of the config.policies dict for multi-agent training.""" + +EpisodeType = Union["SingleAgentEpisode", "MultiAgentEpisode"] +"""A new stack Episode type: Either single-agent or multi-agent.""" + +# @ OldAPIStack +IsPolicyToTrain = Callable[[PolicyID, Optional["MultiAgentBatch"]], bool] +"""Is Policy to train callable.""" + +AgentToModuleMappingFn = Callable[[AgentID, EpisodeType], ModuleID] +"""Function describing an agent to module mapping.""" + +ShouldModuleBeUpdatedFn = Union[ + Sequence[ModuleID], + Callable[[ModuleID, Optional["MultiAgentBatch"]], bool], +] +""" +ModuleIDs that should be updated +or a callable to return whether a module should be updated. +""" + +PolicyState = Dict[str, TensorStructType] # @OldAPIStack +""" +State dict of a Policy, mapping strings (e.g. "weights") +to some state data (TensorStructType). +""" + +TFPolicyV2Type = Type[Union["DynamicTFPolicyV2", "EagerTFPolicyV2"]] # @OldAPIStack +"""Any tf Policy type (static-graph or eager Policy).""" + +EpisodeID = Union[int, str] +"""Represents an episode id (old and new API stack).""" + +UnrollID = int # @OldAPIStack +"""Represents an "unroll" (maybe across different sub-envs in a vector env).""" + +MultiAgentDict = Dict[AgentID, Any] +"""A dict keyed by agent ids, e.g. {"agent-1": value}.""" + +MultiEnvDict = Dict[EnvID, MultiAgentDict] +""" +A dict keyed by env ids that contain further nested dictionaries keyed by agent +ids. e.g., {"env-1": {"agent-1": value}}. +""" + +EnvObsType = Any +"""Represents an observation returned from the env. (Any alias)""" + +EnvActionType = Any +"""Represents an action passed to the env. (Any alias)""" + +EnvInfoDict = dict +""" +Info dictionary returned by calling `reset()` or `step()` on `gymnasium.Env` +instances. Might be an empty dict. +""" + +FileType = Any +"""Represents a File object. (Any alias)""" + +ViewRequirementsDict = Dict[str, "ViewRequirement"] # @OldAPIStack +""" +Represents a ViewRequirements dict mapping column names (str) to ViewRequirement +objects. +""" + +ResultDict = Dict +""" +Represents the result dict returned by Algorithm.train() and algorithm components, +such as EnvRunners, LearnerGroup, etc.. Also, the MetricsLogger used by all these +components returns this upon its `reduce()` method call, so a ResultDict can further +be accumulated (and reduced again) by downstream components. +""" + +LocalOptimizer = Union["torch.optim.Optimizer", "keras.optimizers.Optimizer"] +"""A tf or torch local optimizer object.""" + +Optimizer = LocalOptimizer +"""A tf or torch optimizer object.""" + +Param = Union["torch.Tensor", "tf.Variable"] +"""A parameter, either a torch.Tensor or tf.Variable.""" + +ParamRef = Hashable +"""A reference to a parameter. (Hashable alias)""" + +ParamDict = Dict[ParamRef, Param] +"""A dictionary mapping parameter references to parameters.""" + +ParamList = List[Param] +"""A list of parameters.""" + +NamedParamDict = Dict[str, Param] +"""A dictionary mapping parameter names to parameters.""" + +LearningRateOrSchedule = Union[ + float, + List[List[Union[int, float]]], + List[Tuple[int, Union[int, float]]], +] +""" +A single learning rate or a learning rate schedule (list of sub-lists, each of +the format: [ts (int), lr_to_reach_by_ts (float)]). +""" + +GradInfoDict = dict +""" +Dict of tensors returned by compute gradients on the policy, e.g., +{"td_error": [...], "learner_stats": {"vf_loss": ..., ...}}, +for multi-agent, {"policy1": {"learner_stats": ..., }, "policy2": ...}. +""" + +LearnerStatsDict = dict +""" +Dict of learner stats returned by compute gradients on the policy, e.g., +{"vf_loss": ..., ...}. This will always be nested under the "learner_stats" key(s) +of a GradInfoDict. In the multi-agent case, this will be keyed by policy id. +""" + +ModelGradients = Union[List[Tuple[TensorType, TensorType]], List[TensorType]] +""" +List of grads+var tuples (tf) or list of gradient tensors (torch) representing +model gradients and returned by compute_gradients(). +""" + +ModelWeights = dict +"""Type of dict returned by get_weights() representing model weights.""" + +ModelInputDict = Dict[str, TensorType] +"""An input dict used for direct ModelV2 calls.""" + +SampleBatchType = Union["SampleBatch", "MultiAgentBatch", Dict[str, Any]] +"""Some kind of sample batch.""" + +SpaceStruct = Union[ + gym.spaces.Space, Dict[str, gym.spaces.Space], Tuple[gym.spaces.Space, ...] +] +""" +A (possibly nested) space struct: Either a gym.spaces.Space or a (possibly +nested) dict|tuple of gym.space.Spaces. +""" + +StateBatches = List[List[Any]] # @OldAPIStack +""" +A list of batches of RNN states. +Each item in this list has dimension [B, S] (S=state vector size) +""" + +# __sphinx_doc_begin_policy_output_type__ +PolicyOutputType = Tuple[TensorStructType, StateBatches, Dict] # @OldAPIStack +"""Format of data output from policy forward pass.""" +# __sphinx_doc_end_policy_output_type__ + + +# __sphinx_doc_begin_agent_connector_data_type__ +@OldAPIStack +class AgentConnectorDataType: + """Data type that is fed into and yielded from agent connectors. + + Args: + env_id: ID of the environment. + agent_id: ID to help identify the agent from which the data is received. + data: A payload (``data``). With RLlib's default sampler, the payload + is a dictionary of arbitrary data columns (obs, rewards, terminateds, + truncateds, etc). + """ + + def __init__(self, env_id: str, agent_id: str, data: Any): + self.env_id = env_id + self.agent_id = agent_id + self.data = data + + +# __sphinx_doc_end_agent_connector_data_type__ + + +# __sphinx_doc_begin_action_connector_output__ +@OldAPIStack +class ActionConnectorDataType: + """Data type that is fed into and yielded from agent connectors. + + Args: + env_id: ID of the environment. + agent_id: ID to help identify the agent from which the data is received. + input_dict: Input data that was passed into the policy. + Sometimes output must be adapted based on the input, for example + action masking. So the entire input data structure is provided here. + output: An object of PolicyOutputType. It is is composed of the + action output, the internal state output, and additional data fetches. + + """ + + def __init__( + self, + env_id: str, + agent_id: str, + input_dict: TensorStructType, + output: PolicyOutputType, + ): + self.env_id = env_id + self.agent_id = agent_id + self.input_dict = input_dict + self.output = output + + +# __sphinx_doc_end_action_connector_output__ + + +# __sphinx_doc_begin_agent_connector_output__ +@OldAPIStack +class AgentConnectorsOutput: + """Final output data type of agent connectors. + + Args are populated depending on the AgentConnector settings. + The branching happens in ViewRequirementAgentConnector. + + Args: + raw_dict: The raw input dictionary that sampler can use to + build episodes and training batches. + This raw dict also gets passed into ActionConnectors in case + it contains data useful for action adaptation (e.g. action masks). + sample_batch: The SampleBatch that can be immediately used for + querying the policy for next action. + """ + + def __init__( + self, raw_dict: Dict[str, TensorStructType], sample_batch: "SampleBatch" + ): + self.raw_dict = raw_dict + self.sample_batch = sample_batch + + +# __sphinx_doc_end_agent_connector_output__ + + +# Generic type var. +T = TypeVar("T")