torchrl 0.11.0__cp314-cp314t-win_amd64.whl

torchrl/collectors/weight_update.py

@@ -0,0 +1,591 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+from __future__ import annotations
+
+import abc
+import weakref
+from collections.abc import Callable
+from typing import Any, TypeVar
+
+import torch
+from tensordict import TensorDict, TensorDictBase
+from tensordict.nn import TensorDictModuleBase
+from torchrl._utils import logger as torchrl_logger
+
+Policy = TypeVar("Policy", bound=TensorDictModuleBase)
+
+
+class WeightUpdaterBase(metaclass=abc.ABCMeta):
+    """A base class for updating remote policy weights on inference workers.
+
+    .. deprecated::
+        WeightUpdaterBase is deprecated and will be removed in a future version.
+        Please use WeightSyncScheme from torchrl.weight_update.weight_sync_schemes instead.
+
+    The weight updater is the central piece of the weight update scheme:
+
+    - In leaf collector nodes, it is responsible for sending the weights to the policy, which can be as simple as
+      updating a state-dict, or more complex if an inference server is being used.
+    - In server collector nodes, it is responsible for sending the weights to the leaf collectors.
+
+    In a collector, the updater is called within :meth:`~torchrl.collector.BaseCollector.update_policy_weights_`.`
+
+    The main method of this class is the :meth:`~._push_weights` method, which updates the policy weights in the worker /
+    policy. This method is called by :meth:`~.push_weights`, which also calls the post-hooks: only `_push_weights` should
+    be implemented by child classes.
+
+    To extend this class, implement the following abstract methods:
+
+    - `_get_server_weights` (optional): Define how to retrieve the weights from the server if they are not passed to
+        the updater directly. This method is only called if the weights (handle) is not passed directly.
+    - `_sync_weights_with_worker`: Define how to synchronize weights with a specific worker.
+        This method must be implemented by child classes.
+    - `_maybe_map_weights`: Optionally transform the server weights before distribution.
+        By default, this method returns the weights unchanged.
+    - `all_worker_ids`: Provide a list of all worker identifiers.
+        Returns `None` by default (no worker id).
+    - `from_policy` (optional classmethod): Define how to create an instance of the weight updater from a policy.
+        If implemented, this method will be called before falling back to the default constructor when initializing
+        a weight updater in a collector.
+
+    Attributes:
+        collector: The collector (or any container) of the weight receiver. The collector is registered via
+            :meth:`~torchrl.collectors.WeightUpdaterBase.register_collector`.
+
+    Methods:
+        push_weights: Updates the weights on specified or all remote workers.
+            The `__call__` method is a proxy to `push_weights`.
+        register_collector: Registers the collector (or any container) in the receiver through a weakref.
+            This will be called automatically by the collector upon registration of the updater.
+        from_policy: Optional classmethod to create an instance from a policy.
+
+    Post-hooks:
+        - `register_post_hook`: Registers a post-hook to be called after the weights are updated.
+            The post-hook must be a callable that takes no arguments.
+            The post-hook will be called after the weights are updated.
+            The post-hook will be called in the same process as the weight updater.
+            The post-hook will be called in the same order as the post-hooks were registered.
+
+    .. seealso:: :meth:`~torchrl.collectors.BaseCollector.update_policy_weights_`.
+
+    """
+
+    _collector_wrs: list[Any] = None
+    _post_hooks: list[Callable[[], Any]] | None = None
+
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        import warnings
+
+        warnings.warn(
+            f"Creating {cls.__name__} which inherits from WeightUpdaterBase is deprecated. "
+            "Please use WeightSyncScheme from torchrl.weight_update.weight_sync_schemes instead. "
+            "This will be removed in a future version.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
+    @property
+    def post_hooks(self) -> list[Callable[[], None]]:
+        """The list of post-hooks registered to the weight updater."""
+        if self._post_hooks is None:
+            self._post_hooks = []
+        return self._post_hooks
+
+    @classmethod
+    def from_policy(cls, policy: TensorDictModuleBase) -> WeightUpdaterBase | None:
+        """Optional classmethod to create a weight updater instance from a policy.
+
+        This method can be implemented by subclasses to provide custom initialization logic
+        based on the policy. If implemented, this method will be called before falling back
+        to the default constructor when initializing a weight updater in a collector.
+
+        Args:
+            policy (TensorDictModuleBase): The policy to create the weight updater from.
+
+        Returns:
+            WeightUpdaterBase | None: An instance of the weight updater, or None if the policy
+                cannot be used to create an instance.
+        """
+        return None
+
+    def register_collector(self, collector):  # noqa
+        """Register a collector in the updater.
+
+        Once registered, the updater will not accept another collector.
+
+        Args:
+            collector (BaseCollector): The collector to register.
+
+        """
+        if self._collector_wrs is None:
+            self._collector_wrs = []
+        self._collector_wrs.append(weakref.ref(collector))
+
+    @property
+    def collector(self) -> Any | None:  # noqa
+        """The collector or container of the receiver.
+
+        Returns `None` if the container is out-of-scope or not set.
+        """
+        if self._collector_wrs is None:
+            return None
+        if len(self._collector_wrs) > 1:
+            raise ValueError("Cannot access `collector` with multiple collectors.")
+        if self._collector_wrs:
+            collector = self._collector_wrs[0]()
+        else:
+            collector = None
+        return collector
+
+    @property
+    def collectors(self) -> list[Any] | None:
+        """The collectors or container of the receiver."""
+        if self._collector_wrs is None:
+            return None
+        if self._collector_wrs:
+            collectors = [
+                wr() if wr is not None else None for wr in self._collector_wrs
+            ]
+        else:
+            collectors = None
+        return collectors
+
+    def _push_weights(
+        self,
+        *,
+        policy_or_weights: TensorDictModuleBase | TensorDictBase | dict | None = None,
+        worker_ids: torch.device | int | list[int] | list[torch.device] | None = None,
+    ):
+        """Updates the weights of the policy, or on specified / all remote workers.
+
+        Args:
+            policy_or_weights: The source to get weights from. Can be:
+                - TensorDictModuleBase: A policy module whose weights will be extracted
+                - TensorDictBase: A TensorDict containing weights
+                - dict: A regular dict containing weights
+                - None: Will try to get weights from server using _get_server_weights()
+            worker_ids: An optional list of workers to update.
+
+        Returns: nothing.
+        """
+        if policy_or_weights is None:
+            # Get the weights on server (local)
+            server_weights = self._get_server_weights()
+        else:
+            server_weights = policy_or_weights
+
+        server_weights = self._maybe_map_weights(server_weights)
+
+        # Get the remote weights (inference workers)
+        if isinstance(worker_ids, (int, torch.device)):
+            worker_ids = [worker_ids]
+        elif worker_ids is None:
+            worker_ids = self.all_worker_ids()
+        if worker_ids is None:
+            self._sync_weights_with_worker(server_weights=server_weights)
+            return
+        for worker in worker_ids:
+            if self._skip_update(worker):
+                continue
+            self._sync_weights_with_worker(
+                worker_id=worker, server_weights=server_weights
+            )
+
+    def push_weights(
+        self,
+        policy_or_weights: TensorDictModuleBase | TensorDictBase | dict | None = None,
+        worker_ids: torch.device | int | list[int] | list[torch.device] | None = None,
+    ):
+        """Updates the weights of the policy, or on specified / all remote workers.
+
+        Args:
+            policy_or_weights: The source to get weights from. Can be:
+                - TensorDictModuleBase: A policy module whose weights will be extracted
+                - TensorDictBase: A TensorDict containing weights
+                - dict: A regular dict containing weights
+                - None: Will try to get weights from server using _get_server_weights()
+            worker_ids: An optional list of workers to update.
+
+        Returns: nothing.
+        """
+        self._push_weights(policy_or_weights=policy_or_weights, worker_ids=worker_ids)
+        self._call_post_hooks()
+
+    def init(self, *args, **kwargs):
+        """Initialize the weight updater with custom arguments.
+
+        This method can be overridden by subclasses to handle custom initialization.
+        By default, this is a no-op.
+
+        Args:
+            *args: Positional arguments for initialization
+            **kwargs: Keyword arguments for initialization
+        """
+        return
+
+    def register_post_hook(self, hook: Callable[[], None]):
+        """Registers a post-hook to be called after weights are updated.
+
+        Args:
+            hook (Callable[[], None]): The post-hook to register.
+        """
+        self.post_hooks.append(hook)
+
+    def _call_post_hooks(self):
+        """Calls all registered post-hooks in order."""
+        for hook in self.post_hooks:
+            hook()
+
+    def _skip_update(self, worker_id: int | torch.device) -> bool:
+        """Whether to skip updating weights for a worker.
+
+        By default, never skips updates. Subclasses can override this to implement
+        custom update frequency logic.
+
+        Args:
+            worker_id (int | torch.device): The worker ID to check.
+
+        Returns:
+            bool: Whether to skip the update.
+        """
+        return False
+
+    @abc.abstractmethod
+    def _sync_weights_with_worker(
+        self,
+        *,
+        worker_id: int | torch.device | None = None,
+        server_weights: TensorDictBase,
+    ) -> None:
+        """Synchronizes weights with a specific worker.
+
+        This method must be implemented by child classes to define how weights are
+        synchronized with workers.
+
+        Args:
+            worker_id (int | torch.device | None): The worker to sync with, if applicable.
+            server_weights (TensorDictBase): The weights from the server to sync.
+        """
+        raise NotImplementedError
+
+    def _get_server_weights(self) -> TensorDictBase | None:
+        """Gets the weights from the server.
+
+        This method is called when no weights are passed to push_weights().
+        By default returns None. Subclasses can override to implement custom
+        weight retrieval logic.
+
+        Returns:
+            TensorDictBase | None: The server weights, or None.
+        """
+        return None
+
+    def _maybe_map_weights(self, policy_or_weights: TensorDictBase) -> TensorDictBase:
+        """Optionally transforms server weights before distribution.
+
+        By default returns weights unchanged. Subclasses can override to implement
+        custom weight mapping logic.
+
+        Args:
+            policy_or_weights (Any): The weights - or any container or handler - to potentially transform, query or extract.
+
+        Returns:
+            TensorDictBase: The transformed weights.
+        """
+        if isinstance(policy_or_weights, TensorDictModuleBase):
+            # Extract weights from policy module
+            server_weights = TensorDict.from_module(policy_or_weights).data
+        elif isinstance(policy_or_weights, (TensorDictBase, dict)):
+            # Use weights directly
+            server_weights = policy_or_weights
+        else:
+            raise TypeError(
+                f"policy_or_weights must be None, TensorDictModuleBase, TensorDictBase or dict, got {type(policy_or_weights)}"
+            )
+        return server_weights
+
+    def all_worker_ids(self) -> list[int] | list[torch.device] | None:
+        """Gets list of all worker IDs.
+
+        Returns None by default. Subclasses should override to return actual worker IDs.
+
+        Returns:
+            list[int] | list[torch.device] | None: List of worker IDs or None.
+        """
+        return None
+
+    def __call__(
+        self,
+        policy_or_weights: TensorDictModuleBase | TensorDictBase | dict | None = None,
+        worker_ids: torch.device | int | list[int] | list[torch.device] | None = None,
+    ):
+        """Updates the weights of the policy, or on specified / all remote workers.
+
+        Args:
+            policy_or_weights: The source to get weights from. Can be:
+                - TensorDictModuleBase: A policy module whose weights will be extracted
+                - TensorDictBase: A TensorDict containing weights
+                - dict: A regular dict containing weights
+                - None: Will try to get weights from server using _get_server_weights()
+            worker_ids: An optional list of workers to update.
+
+        Returns: nothing.
+        """
+        return self.push_weights(
+            policy_or_weights=policy_or_weights, worker_ids=worker_ids
+        )
+
+    def increment_version(self):
+        """Increment the policy version."""
+        for collector in self.collectors:
+            collector.increment_version()
+
+
+# Specialized classes
+class VanillaWeightUpdater(WeightUpdaterBase):
+    """A simple implementation of :class:`~torchrl.collectors.WeightUpdaterBase` for updating local policy weights.
+
+    The `VanillaWeightSender` class provides a basic mechanism for updating the weights
+    of a local policy by directly fetching them from a specified source. It is typically used
+    in scenarios where the weight update logic is straightforward and does not require any
+    complex mapping or transformation.
+
+    This class is used by default in the `Collector` when no custom weight sender
+    is provided.
+
+    .. seealso:: :class:`~torchrl.collectors.WeightUpdaterBase` and :class:`~torchrl.collectors.Collector`.
+
+    Keyword Args:
+        weight_getter (Callable[[], TensorDictBase], optional): a callable that returns the weights from the server.
+            If not provided, the weights must be passed to ``push_weights`` directly.
+        policy_weights (TensorDictBase): a TensorDictBase containing the policy weights to be updated
+            in-place. Use ``push_weights`` to update the weights.
+    """
+
+    @classmethod
+    def from_policy(cls, policy: TensorDictModuleBase) -> WeightUpdaterBase | None:
+        """Creates a VanillaWeightUpdater instance from a policy.
+
+        This method creates a weight updater that will update the policy's weights directly
+        using its state dict.
+
+        Args:
+            policy (TensorDictModuleBase): The policy to create the weight updater from.
+
+        Returns:
+            VanillaWeightUpdater: An instance of the weight updater configured to update
+                the policy's weights.
+        """
+        policy_weights = TensorDict.from_module(policy)
+        return cls(policy_weights=policy_weights.lock_())
+
+    def __init__(
+        self,
+        *,
+        weight_getter: Callable[[], TensorDictBase] | None = None,
+        policy_weights: TensorDictBase,
+    ):
+        self.weight_getter = weight_getter
+        self.policy_weights = policy_weights
+
+    def _get_server_weights(self) -> TensorDictBase | None:
+        return self.weight_getter() if self.weight_getter is not None else None
+
+    def _get_local_weights(self) -> TensorDictBase:
+        return self.policy_weights
+
+    def _map_weights(self, server_weights: TensorDictBase) -> TensorDictBase:
+        return server_weights
+
+    def _maybe_map_weights(self, server_weights: TensorDictBase) -> TensorDictBase:
+        return server_weights
+
+    def _sync_weights_with_worker(
+        self, *, worker_id: None = None, server_weights: TensorDictBase
+    ) -> None:
+        if server_weights is None:
+            return
+        self.policy_weights.update_(server_weights)
+
+
+class MultiProcessedWeightUpdater(WeightUpdaterBase):
+    """A remote weight updater for synchronizing policy weights across multiple processes or devices.
+
+    The `MultiProcessedWeightUpdater` class provides a mechanism for updating the weights
+    of a policy across multiple inference workers in a multiprocessed environment. It is designed
+    to handle the distribution of weights from a central server to various devices or processes
+    that are running the policy.
+    This class is typically used in multiprocessed data collectors where each process or device
+    requires an up-to-date copy of the policy weights.
+
+    Keyword Args:
+        get_server_weights (Callable[[], TensorDictBase] | None): A callable that retrieves the
+            latest policy weights from the server or another centralized source.
+        policy_weights (Dict[torch.device, TensorDictBase]): A dictionary mapping each device or
+            process to its current policy weights, which will be updated.
+
+    .. note::
+        This class assumes that the server weights can be directly applied to the workers without
+        any additional processing. If your use case requires more complex weight mapping or synchronization
+        logic, consider extending `WeightUpdaterBase` with a custom implementation.
+
+    .. seealso:: :class:`~torchrl.collectors.WeightUpdaterBase` and
+        :class:`~torchrl.collectors.BaseCollector`.
+
+    """
+
+    def __init__(
+        self,
+        *,
+        get_server_weights: Callable[[], TensorDictBase] | None,
+        policy_weights: dict[torch.device, TensorDictBase],
+    ):
+        self.weights_getter = get_server_weights
+        self._policy_weights = policy_weights
+
+    def all_worker_ids(self) -> list[int] | list[torch.device]:
+        return list(self._policy_weights)
+
+    def _sync_weights_with_worker(
+        self, worker_id: int | torch.device, server_weights: TensorDictBase | None
+    ) -> None:
+        if server_weights is None:
+            return
+        self._policy_weights[worker_id].data.update_(server_weights)
+
+    def _get_server_weights(self) -> TensorDictBase | None:
+        # The weights getter can be none if no mapping is required
+        if self.weights_getter is None:
+            return
+        weights = self.weights_getter()
+        if weights is None:
+            return
+        return weights.data
+
+    def _maybe_map_weights(self, server_weights: TensorDictBase) -> TensorDictBase:
+        return server_weights
+
+
+class RemoteModuleWeightUpdater(WeightUpdaterBase):
+    """A weight updater for remote nn.Modules that requires explicit weight passing.
+
+    This weight updater is designed for scenarios where the master collector doesn't have
+    direct access to worker weights (e.g., when using policy_factory). It enforces that
+    weights must be passed explicitly when calling update_policy_weights_().
+
+    This updater does not try to retrieve weights from the server or workers automatically.
+    Instead, it raises an exception if weights are not provided, ensuring that the weight
+    synchronization is handled explicitly by the user.
+
+    Raises:
+        RuntimeError: If update_policy_weights_() is called without providing weights.
+
+    .. note::
+        This weight updater is primarily used to suppress warnings in tests and scenarios
+        where the weight synchronization is handled externally or the workers manage
+        their own weight updates.
+
+    .. seealso:: :class:`~torchrl.collectors.WeightUpdaterBase`
+    """
+
+    def __init__(self):
+        pass
+
+    def _get_server_weights(self) -> TensorDictBase | None:
+        """Returns None since this updater doesn't manage server weights."""
+        return None
+
+    def _sync_weights_with_worker(
+        self,
+        *,
+        worker_id: int | torch.device | None = None,
+        server_weights: TensorDictBase,
+    ) -> None:
+        """Raises an error if weights are not provided explicitly.
+
+        Since this updater is for remote modules where the master doesn't have access
+        to worker weights, it enforces that weights must be passed explicitly.
+        """
+        if server_weights is None:
+            raise RuntimeError(
+                "RemoteModuleWeightUpdater requires weights to be passed explicitly. "
+                "Call update_policy_weights_(weights) with the weights to be synchronized."
+            )
+        # If weights are provided, we assume the synchronization is handled elsewhere
+        # This is a no-op updater that just validates the weight passing pattern
+
+    def all_worker_ids(self) -> None:
+        """Returns None since this updater doesn't manage specific workers."""
+        return None
+
+
+class RayWeightUpdater(WeightUpdaterBase):
+    """A remote weight updater for synchronizing policy weights across remote workers using Ray.
+
+    The `RayWeightUpdater` class provides a mechanism for updating the weights of a policy
+    across remote inference workers managed by Ray. It leverages Ray's distributed computing
+    capabilities to efficiently distribute policy weights to remote collectors.
+    This class is typically used in distributed data collectors where each remote worker requires
+    an up-to-date copy of the policy weights.
+
+    Args:
+        policy_weights (TensorDictBase): The current weights of the policy that need to be distributed
+            to remote workers.
+        remote_collectors (List): A list of remote collectors that will receive the updated policy weights.
+        max_interval (int, optional): The maximum number of batches between weight updates for each worker.
+            Defaults to 0, meaning weights are updated every batch.
+
+    Methods:
+        all_worker_ids: Returns a list of all worker identifiers (indices of remote collectors).
+        _get_server_weights: Retrieves the latest weights from the server and stores them in Ray's object store.
+        _maybe_map_weights: Optionally maps server weights before distribution (no-op in this implementation).
+        _sync_weights_with_worker: Synchronizes the server weights with a specific remote worker using Ray.
+        _skip_update: Determines whether to skip the weight update for a specific worker based on the interval.
+
+    .. note::
+        This class assumes that the server weights can be directly applied to the remote workers without
+        any additional processing. If your use case requires more complex weight mapping or synchronization
+        logic, consider extending `WeightUpdaterBase` with a custom implementation.
+
+    .. seealso:: :class:`~torchrl.collectors.WeightUpdaterBase` and
+        :class:`~torchrl.collectors.distributed.RayCollector`.
+
+    """
+
+    def __init__(
+        self,
+        policy_weights: TensorDictBase,
+        remote_collectors: list,
+        max_interval: int = 0,
+    ):
+        self.policy_weights = policy_weights
+        self.remote_collectors = remote_collectors
+        self.max_interval = max(0, max_interval)
+        self._batches_since_weight_update = [0] * len(self.remote_collectors)
+
+    def all_worker_ids(self) -> list[int] | list[torch.device]:
+        return list(range(len(self.remote_collectors)))
+
+    def _get_server_weights(self) -> Any:
+        import ray
+
+        return ray.put(self.policy_weights.data)
+
+    def _maybe_map_weights(self, server_weights: Any) -> Any:
+        return server_weights
+
+    def _sync_weights_with_worker(self, worker_id: int, server_weights: Any) -> Any:
+        torchrl_logger.debug(f"syncing weights with worker {worker_id}")
+        c = self.remote_collectors[worker_id]
+        c.update_policy_weights_.remote(policy_weights=server_weights)
+        self._batches_since_weight_update[worker_id] = 0
+
+    def _skip_update(self, worker_id: int) -> bool:
+        self._batches_since_weight_update[worker_id] += 1
+        # Use gt because we just incremented it
+        if self._batches_since_weight_update[worker_id] > self.max_interval:
+            return False
+        return True

torchrl/csrc/numpy_utils.h

@@ -0,0 +1,38 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+//
+// This source code is licensed under the MIT license found in the
+// LICENSE file in the root directory of this source tree.
+
+#pragma once
+
+#include <pybind11/numpy.h>
+
+#include <algorithm>
+#include <cstdint>
+#include <vector>
+
+namespace py = pybind11;
+
+namespace torchrl {
+namespace utils {
+
+template <typename T>
+std::vector<int64_t> NumpyArrayShape(const py::array_t<T>& arr) {
+  const int64_t ndim = arr.ndim();
+  std::vector<int64_t> shape(ndim);
+  for (int64_t i = 0; i < ndim; ++i) {
+    shape[i] = static_cast<int64_t>(arr.shape(i));
+  }
+  return shape;
+}
+
+template <typename T_SRC, typename T_DST = T_SRC>
+py::array_t<T_DST> NumpyEmptyLike(const py::array_t<T_SRC>& src) {
+  py::array_t<T_DST> dst(src.size());
+  const std::vector<int64_t> shape = NumpyArrayShape(src);
+  dst.resize(shape);
+  return dst;
+}
+
+}  // namespace utils
+}  // namespace torchrl

torchrl/csrc/pybind.cpp

@@ -0,0 +1,27 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+//
+// This source code is licensed under the MIT license found in the
+// LICENSE file in the root directory of this source tree.
+
+#include <pybind11/numpy.h>
+#include <pybind11/pybind11.h>
+#include <torch/extension.h>
+#include <torch/torch.h>
+
+#include <memory>
+
+#include "segment_tree.h"
+#include "utils.h"
+
+namespace py = pybind11;
+
+PYBIND11_MODULE(_torchrl, m) {
+  torchrl::DefineSumSegmentTree<float>("Fp32", m);
+  torchrl::DefineSumSegmentTree<double>("Fp64", m);
+
+  torchrl::DefineMinSegmentTree<float>("Fp32", m);
+  torchrl::DefineMinSegmentTree<double>("Fp64", m);
+
+  m.def("safetanh", &safetanh, "Safe Tanh");
+  m.def("safeatanh", &safeatanh, "Safe Inverse Tanh");
+}