[Versioning] Import implement_for from pyvers (#3166)

Author: vmoens

pyproject.toml

@@ -29,6 +29,7 @@ classifiers = [
 requires-python = ">=3.9"
 dependencies = [
     "torch>=2.1.0",
+    "pyvers",
     "numpy",
     "packaging",
     "cloudpickle",

torchrl/_utils.py

@@ -18,15 +18,14 @@
 import warnings
 from collections.abc import Callable
 from contextlib import nullcontext
-from copy import copy
 from functools import wraps
-from importlib import import_module
 from textwrap import indent
 from typing import Any, cast, TypeVar
 
 import numpy as np
 import torch
-from packaging.version import parse
+
+from pyvers import implement_for  # noqa: F401
 from tensordict import unravel_key
 from tensordict.utils import NestedKey
 from torch import multiprocessing as mp, Tensor
@@ -390,274 +389,6 @@ def __repr__(self):
 _CKPT_BACKEND = _Dynamic_CKPT_BACKEND()
 
 
-class implement_for:
-    """A version decorator that checks the version in the environment and implements a function with the fitting one.
-
-    If specified module is missing or there is no fitting implementation, call of the decorated function
-    will lead to the explicit error.
-    In case of intersected ranges, last fitting implementation is used.
-
-    This wrapper also works to implement different backends for a same function (eg. gym vs gymnasium,
-    numpy vs jax-numpy etc).
-
-    Args:
-        module_name (str or callable): version is checked for the module with this
-            name (e.g. "gym"). If a callable is provided, it should return the
-            module.
-        from_version: version from which implementation is compatible. Can be open (None).
-        to_version: version from which implementation is no longer compatible. Can be open (None).
-
-    Keyword Args:
-        class_method (bool, optional): if ``True``, the function will be written as a class method.
-            Defaults to ``False``.
-        compilable (bool, optional): If ``False``, the module import happens
-            only on the first call to the wrapped function. If ``True``, the
-            module import happens when the wrapped function is initialized. This
-            allows the wrapped function to work well with ``torch.compile``.
-            Defaults to ``False``.
-
-    Examples:
-        >>> @implement_for("gym", "0.13", "0.14")
-        >>> def fun(self, x):
-        ...     # Older gym versions will return x + 1
-        ...     return x + 1
-        ...
-        >>> @implement_for("gym", "0.14", "0.23")
-        >>> def fun(self, x):
-        ...     # More recent gym versions will return x + 2
-        ...     return x + 2
-        ...
-        >>> @implement_for(lambda: import_module("gym"), "0.23", None)
-        >>> def fun(self, x):
-        ...     # More recent gym versions will return x + 2
-        ...     return x + 2
-        ...
-        >>> @implement_for("gymnasium", None, "1.0.0")
-        >>> def fun(self, x):
-        ...     # If gymnasium is to be used instead of gym, x+3 will be returned
-        ...     return x + 3
-        ...
-
-        This indicates that the function is compatible with gym 0.13+, but doesn't with gym 0.14+.
-    """
-
-    # Stores pointers to fitting implementations: dict[func_name] = func_pointer
-    _implementations = {}
-    _setters = []
-    _cache_modules = {}
-
-    def __init__(
-        self,
-        module_name: str | Callable,
-        from_version: str | None = None,
-        to_version: str | None = None,
-        *,
-        class_method: bool = False,
-        compilable: bool = False,
-    ):
-        self.module_name = module_name
-        self.from_version = from_version
-        self.to_version = to_version
-        self.class_method = class_method
-        self._compilable = compilable
-        implement_for._setters.append(self)
-
-    @staticmethod
-    def check_version(version: str, from_version: str | None, to_version: str | None):
-        version = parse(".".join([str(v) for v in parse(version).release]))
-        return (from_version is None or version >= parse(from_version)) and (
-            to_version is None or version < parse(to_version)
-        )
-
-    @staticmethod
-    def get_class_that_defined_method(f):
-        """Returns the class of a method, if it is defined, and None otherwise."""
-        out = f.__globals__.get(f.__qualname__.split(".")[0], None)
-        return out
-
-    @classmethod
-    def get_func_name(cls, fn):
-        # produces a name like torchrl.module.Class.method or torchrl.module.function
-        fn_str = str(fn).split(".")
-        if fn_str[0].startswith("<bound method "):
-            first = fn_str[0][len("<bound method ") :]
-        elif fn_str[0].startswith("<function "):
-            first = fn_str[0][len("<function ") :]
-        else:
-            raise RuntimeError(f"Unknown func representation {fn}")
-        last = fn_str[1:]
-        if last:
-            first = [first]
-            last[-1] = last[-1].split(" ")[0]
-        else:
-            last = [first.split(" ")[0]]
-            first = []
-        return ".".join([fn.__module__] + first + last)
-
-    def _get_cls(self, fn):
-        cls = self.get_class_that_defined_method(fn)
-        if cls is None:
-            # class not yet defined
-            return
-        if cls.__class__.__name__ == "function":
-            cls = inspect.getmodule(fn)
-        return cls
-
-    def module_set(self):
-        """Sets the function in its module, if it exists already."""
-        prev_setter = type(self)._implementations.get(self.get_func_name(self.fn), None)
-        if prev_setter is not None:
-            prev_setter.do_set = False
-        type(self)._implementations[self.get_func_name(self.fn)] = self
-        cls = self.get_class_that_defined_method(self.fn)
-        if cls is not None:
-            if cls.__class__.__name__ == "function":
-                cls = inspect.getmodule(self.fn)
-        else:
-            # class not yet defined
-            return
-        try:
-            delattr(cls, self.fn.__name__)
-        except AttributeError:
-            pass
-
-        name = self.fn.__name__
-        if self.class_method:
-            fn = classmethod(self.fn)
-        else:
-            fn = self.fn
-        setattr(cls, name, fn)
-
-    @classmethod
-    def import_module(cls, module_name: Callable | str) -> str:
-        """Imports module and returns its version."""
-        if not callable(module_name):
-            module = cls._cache_modules.get(module_name, None)
-            if module is None:
-                if module_name in sys.modules:
-                    sys.modules[module_name] = module = import_module(module_name)
-                else:
-                    cls._cache_modules[module_name] = module = import_module(
-                        module_name
-                    )
-        else:
-            module = module_name()
-        return module.__version__
-
-    _lazy_impl = collections.defaultdict(list)
-
-    def _delazify(self, func_name):
-        out = None
-        # Make a copy of the list to avoid issues when clearing during iteration
-        lazy_calls = implement_for._lazy_impl[func_name][:]
-        for local_call in lazy_calls:
-            out = local_call()
-        # Only clear for compilable decorators, since non-compilable decorators
-        # need to keep the list to allow multiple lazy calls
-        # Check if any of the decorators are compilable
-        any_compilable = any(
-            hasattr(call, "__self__") and call.__self__._compilable
-            for call in lazy_calls
-        )
-        if any_compilable:
-            implement_for._lazy_impl[func_name].clear()
-        return out
-
-    def __call__(self, fn):
-        # function names are unique
-        self.func_name = self.get_func_name(fn)
-        self.fn = fn
-        implement_for._lazy_impl[self.func_name].append(self._call)
-
-        if self._compilable:
-            _call_fn = self._delazify(self.func_name)
-
-            if self.class_method:
-                return classmethod(_call_fn)
-
-            return _call_fn
-        else:
-
-            @wraps(fn)
-            def _lazy_call_fn(*args, **kwargs):
-                # first time we call the function, we also do the replacement.
-                # This will cause the imports to occur only during the first call to fn
-
-                result = self._delazify(self.func_name)(*args, **kwargs)
-                return result
-
-            if self.class_method:
-                return classmethod(_lazy_call_fn)
-
-            return _lazy_call_fn
-
-    def _call(self):
-
-        # If the module is missing replace the function with the mock.
-        fn = self.fn
-        func_name = self.func_name
-        implementations = implement_for._implementations
-
-        @wraps(fn)
-        def unsupported(*args, **kwargs):
-            raise ModuleNotFoundError(
-                f"Supported version of '{func_name}' has not been found."
-            )
-
-        self.do_set = False
-        # Return fitting implementation if it was encountered before.
-        if func_name in implementations:
-            try:
-                # check that backends don't conflict
-                version = self.import_module(self.module_name)
-                if self.check_version(version, self.from_version, self.to_version):
-                    if VERBOSE:
-                        module = import_module(self.module_name)
-                        warnings.warn(
-                            f"Got multiple backends for {func_name}. "
-                            f"Using the last queried ({module} with version {version})."
-                        )
-                    self.do_set = True
-                if not self.do_set:
-                    return implementations[func_name].fn
-            except ModuleNotFoundError:
-                # then it's ok, there is no conflict
-                return implementations[func_name].fn
-        else:
-            try:
-                version = self.import_module(self.module_name)
-                if self.check_version(version, self.from_version, self.to_version):
-                    self.do_set = True
-            except ModuleNotFoundError:
-                return unsupported
-        if self.do_set:
-            self.module_set()
-            return fn
-        return unsupported
-
-    @classmethod
-    def reset(cls, setters_dict: dict[str, implement_for] = None):
-        """Resets the setters in setter_dict.
-
-        ``setter_dict`` is a copy of implementations. We just need to iterate through its
-        values and call :meth:`module_set` for each.
-
-        """
-        if VERBOSE:
-            logger.info("resetting implement_for")
-        if setters_dict is None:
-            setters_dict = copy(cls._implementations)
-        for setter in setters_dict.values():
-            setter.module_set()
-
-    def __repr__(self):
-        return (
-            f"{self.__class__.__name__}("
-            f"module_name={self.module_name}({self.from_version, self.to_version}), "
-            f"fn_name={self.fn.__name__}, cls={self._get_cls(self.fn)})"
-        )
-
-
 def accept_remote_rref_invocation(func):
     """Decorator that allows a method to be invoked remotely.
 

torchrl/trainers/trainers.py

@@ -1051,6 +1051,7 @@ class LogScalar(TrainerHookBase):
     Args:
         key (NestedKey): the key where to find the value in the input batch.
             Can be a string for simple keys or a tuple for nested keys.
+            Default is `torchrl.trainers.trainers.REWARD_KEY` (= `("next", "reward")`).
         logname (str, optional): name of the metric to be logged. If None, will use
             the key as the log name. Default is None.
         log_pbar (bool, optional): if ``True``, the value will be logged on
@@ -1077,7 +1078,7 @@ class LogScalar(TrainerHookBase):
 
     def __init__(
         self,
-        key: NestedKey,
+        key: NestedKey = REWARD_KEY,
         logname: str | None = None,
         log_pbar: bool = False,
         include_std: bool = True,