update/refactor parameter updaters

pyrobolearn/parameters/updater.py

@@ -20,31 +20,85 @@ __status__ = "Development"
 
 
 class ParameterUpdater(object):
+    r"""Parameter updater (abstract) class.
 
-    def __init__(self, sleep_count=1):
+    The parameter updater represents a scheme on how to update certain parameters (int, float, torch.nn.Module,
+    torch.tensor, np.array, etc).
+    """
+
+    def __init__(self, current=None, target=None, sleep_count=1):
         """
         Initialize the parameter updater.
 
         Args:
+            current (int, float, torch.tensor, np.array, torch.nn.Module, None): current parameter(s).
+            target (int, float, torch.tensor, np.array, torch.nn.Module, None): target parameter(s) to be modified
+                based on the current parameter(s).
             sleep_count (int): number of ticks to sleep
         """
         self.sleep_count = 1 if sleep_count <= 0 else int(sleep_count)
         self.counter = 0
-        self.target = None
+        self.current = current
+        self.target = target
 
     def _compute(self, current, target):
+        """
+        Inner update that needs to be inherited in all the chid classes. It updates the target parameter(s) based on
+        the current parameter(s). This is an in-place operation; it modifies the given target.
+
+        Args:
+            current (int, float, torch.tensor, np.array, torch.nn.Module): current parameter(s).
+            target (int, float, torch.tensor, np.array, torch.nn.Module): target parameter(s) to be modified based on
+            the current parameter(s).
+
+        Returns:
+            int, float, torch.tensor, np.array, torch.nn.Module: updated target parameter(s).
+        """
         pass
 
-    def compute(self, current, target=None):
+    def compute(self, current=None, target=None):
+        """
+        Update the target parameter(s) based on the current parameter(s).
+
+        Args:
+            current (int, float, torch.tensor, np.array, torch.nn.Module, None): current parameter(s). If None, it
+                will be set to the current parameter(s) given at the initialization.
+            target (int, float, torch.tensor, np.array, torch.nn.Module, None): target parameter(s) to be modified
+                based on the current parameter(s). If None, it will be set to the current parameters.
+
+        Returns:
+            int, float, torch.tensor, np.array, torch.nn.Module: updated target parameter(s).
+        """
+        # if time to update the parameters
         if (self.counter % self.sleep_count) == 0:
+
+            # if the current parameters is None, take the ones given at the initialization
+            if current is None:
+                current = self.current
+                if current is None:
+                    raise ValueError("Expecting to be given current parameters, instead got None.")
+
+            # if the target is None, set it to be the current parameters.
             if target is None:
                 target = current
-            self._compute(current, target)
+
+            # the target and current needs to be of the same type (except if current is a float or int)
+            if not isinstance(current, (int, float)) and type(target) != type(current):
+                raise TypeError("The given target and current parameters are of different types: "
+                                "type(target)={} and type(current)={}".format(type(target), type(current)))
+
+            # inner computation
+            self.target = self._compute(current, target)
+
+        # increment counter
         self.counter += 1
+
+        # return the new target parameter(s).
         return self.target
 
-    def __call__(self, current, target):
-        return self.compute(current, target)
+    def __call__(self, current=None, target=None):
+        """Alias to :func:`compute` method. See doc for :func:`compute`."""
+        return self.compute(current=current, target=target)
 
 
 class CopyParameter(ParameterUpdater):
@@ -53,23 +107,41 @@ class CopyParameter(ParameterUpdater):
     Copy the parameters every `sleep_count` times it is called.
     """
 
-    def __init__(self, sleep_count=1):
+    def __init__(self, current=None, target=None, sleep_count=1):
         """
         Initialize the copy parameter updater.
 
         Args:
+            current (int, float, torch.tensor, np.array, torch.nn.Module, None): current parameter(s).
+            target (int, float, torch.tensor, np.array, torch.nn.Module, None): target parameter(s) to be modified
+                based on the current parameter(s).
             sleep_count (int): number of ticks to sleep
         """
-        super(CopyParameter, self).__init__(sleep_count)
+        super(CopyParameter, self).__init__(current, target, sleep_count)
 
     def _compute(self, current, target):
+        """
+        Updates the target parameter(s) based on the current parameter(s).
+
+        Args:
+            current (int, float, torch.tensor, np.array, torch.nn.Module): current parameter(s).
+            target (int, float, torch.tensor, np.array, torch.nn.Module): target parameter(s) to be modified based on
+            the current parameter(s).
+
+        Returns:
+            int, float, torch.tensor, np.array, torch.nn.Module: updated target parameter(s).
+        """
         if isinstance(target, torch.nn.Module):
             if isinstance(current, torch.nn.Module):
                 for p1, p2 in zip(target.parameters(), current.parameters()):
                     p1.data.copy_(p2.data)
+        elif isinstance(target, torch.Tensor):
+            target.data.copy_(current.data)
+        elif isinstance(target, np.ndarray):
+            target[:] = current
         else:
             target = current
-        self.target = target
+        return target
 
 
 class LinearDecay(ParameterUpdater):
@@ -80,7 +152,7 @@ class LinearDecay(ParameterUpdater):
     step size.
     """
 
-    def __init__(self, slope, step=0.01, end=None, sleep_count=1):
+    def __init__(self, slope, step=0.01, end=None, current=None, target=None, sleep_count=1):
         """
         Initialize the linear decay parameter updater.
 
@@ -89,9 +161,12 @@ class LinearDecay(ParameterUpdater):
             step (float): integration step size.
             end (float, None): end value. If the slope is negative (resp. positive), this is the minimum (resp.
                 maximum) value it can take.
+            current (int, float, torch.tensor, np.array, torch.nn.Module, None): current parameter(s).
+            target (int, float, torch.tensor, np.array, torch.nn.Module, None): target parameter(s) to be modified
+                based on the current parameter(s).
             sleep_count (int): number of ticks to sleep
         """
-        super(LinearDecay, self).__init__(sleep_count)
+        super(LinearDecay, self).__init__(current, target, sleep_count)
         self.slope = slope
         self.dt = step
         if end is None:
@@ -99,20 +174,42 @@ class LinearDecay(ParameterUpdater):
         self.end = end
 
     def _compute(self, current, target):
+        """
+        Updates the target parameter(s) based on the current parameter(s).
+
+        Args:
+            current (int, float, torch.tensor, np.array, torch.nn.Module): current parameter(s).
+            target (int, float, torch.tensor, np.array, torch.nn.Module): target parameter(s) to be modified based on
+            the current parameter(s).
+
+        Returns:
+            int, float, torch.tensor, np.array, torch.nn.Module: updated target parameter(s).
+        """
         if isinstance(target, torch.nn.Module):
             if isinstance(current, torch.nn.Module):
                 for p1, p2 in zip(target.parameters(), current.parameters()):
-                    data = p2.data + self.slope * torch.ones_like(p2.data) * self.dt
+                    data = p2.data + self.slope * self.dt
                     if self.slope < 0:
                         torch.clamp_min_(data, self.end)
                     else:
                         torch.clamp_max_(data, self.end)
                     p1.data.copy_(data)
+        elif isinstance(target, torch.Tensor):
+            data = current.data + self.slope * self.dt
+            if self.slope < 0:
+                torch.clamp_min_(data, self.end)
+            else:
+                torch.clamp_max_(data, self.end)
+            target.data.copy_(data)
+        elif isinstance(target, np.ndarray):
+            target[:] = current + self.slope * self.dt
+            if (self.slope < 0 and target < self.end) or (self.slope > 0 and target > self.end):
+                target[:] = self.end
         else:
             target = current + self.slope * self.dt
             if (self.slope < 0 and target < self.end) or (self.slope > 0 and target > self.end):
                 target = self.end
-        self.target = target
+        return target
 
 
 class ExponentialDecay(ParameterUpdater):
@@ -122,7 +219,7 @@ class ExponentialDecay(ParameterUpdater):
     at which the exponential converges to 0 if negative, and diverges to infinity if positive.
     """
 
-    def __int__(self, speed, step=0.01, end=None, sleep_count=1):
+    def __init__(self, speed, step=0.01, end=None, current=None, target=None, sleep_count=1):
         """
         Initialize the exponential decay parameter updater.
 
@@ -131,8 +228,12 @@ class ExponentialDecay(ParameterUpdater):
             step (float): integration step size.
             end (float, None): end value. If the speed is negative (resp. positive), this is the minimum (resp.
                 maximum) value it can take.
+            current (int, float, torch.tensor, np.array, torch.nn.Module, None): current parameter(s).
+            target (int, float, torch.tensor, np.array, torch.nn.Module, None): target parameter(s) to be modified
+                based on the current parameter(s).
             sleep_count (int): number of ticks to sleep
         """
+        super(ExponentialDecay, self).__init__(current, target, sleep_count)
         self.speed = speed
         self.dt = step
         if end is None:
@@ -140,6 +241,17 @@ class ExponentialDecay(ParameterUpdater):
         self.end = end
 
     def _compute(self, current, target):
+        """
+        Updates the target parameter(s) based on the current parameter(s).
+
+        Args:
+            current (int, float, torch.tensor, np.array, torch.nn.Module): current parameter(s).
+            target (int, float, torch.tensor, np.array, torch.nn.Module): target parameter(s) to be modified based on
+            the current parameter(s).
+
+        Returns:
+            int, float, torch.tensor, np.array, torch.nn.Module: updated target parameter(s).
+        """
         if isinstance(target, torch.nn.Module):
             if isinstance(current, torch.nn.Module):
                 for p1, p2 in zip(target.parameters(), current.parameters()):
@@ -149,11 +261,22 @@ class ExponentialDecay(ParameterUpdater):
                     else:
                         torch.clamp_max_(data, self.end)
                     p1.data.copy_(data)
+        elif isinstance(target, torch.Tensor):
+            data = current.data + self.speed * current.data * self.dt
+            if self.speed < 0:
+                torch.clamp_min_(data, self.end)
+            else:
+                torch.clamp_max_(data, self.end)
+            target.data.copy_(data)
+        elif isinstance(target, np.ndarray):
+            target[:] = current + self.speed * current * self.dt
+            if (self.speed < 0 and target < self.end) or (self.speed > 0 and target > self.end):
+                target[:] = self.end
         else:
             target = current + self.speed * current * self.dt
             if (self.speed < 0 and target < self.end) or (self.speed > 0 and target > self.end):
                 target = self.end
-        self.target = target
+        return target
 
 
 class PolyakAveraging(ParameterUpdater):
@@ -163,12 +286,16 @@ class PolyakAveraging(ParameterUpdater):
     :math:`y^*` is the target, and :math:`y` is the current value.
     """
 
-    def __init__(self, rho=0., sleep_count=1):
+    def __init__(self, rho=0., current=None, target=None, sleep_count=1):
         """
         Initialize the polyak averaging parameter updater.
 
         Args:
-            rho (float): float value between 0 and 1. If 1, it won't do anything, if 0 it will
+            rho (float): float value between 0 and 1. If 1, it will let the target parameter(s) unchanged, if 0 it
+                will just copy the current parameter(s).
+            current (int, float, torch.tensor, np.array, torch.nn.Module, None): current parameter(s).
+            target (int, float, torch.tensor, np.array, torch.nn.Module, None): target parameter(s) to be modified
+                based on the current parameter(s).
             sleep_count (int): number of ticks to sleep
         """
         super(PolyakAveraging, self).__init__(sleep_count)
@@ -177,10 +304,25 @@ class PolyakAveraging(ParameterUpdater):
         self.rho = rho
 
     def _compute(self, current, target):
+        """
+        Updates the target parameter(s) based on the current parameter(s).
+
+        Args:
+            current (int, float, torch.tensor, np.array, torch.nn.Module): current parameter(s).
+            target (int, float, torch.tensor, np.array, torch.nn.Module): target parameter(s) to be modified based on
+            the current parameter(s).
+
+        Returns:
+            int, float, torch.tensor, np.array, torch.nn.Module: updated target parameter(s).
+        """
         if isinstance(target, torch.nn.Module):
             if isinstance(current, torch.nn.Module):
                 for p1, p2 in zip(target.parameters(), current.parameters()):
                     p1.data.copy_(self.rho * p1.data + (1 - self.rho) * p2.data)
+        elif isinstance(target, torch.Tensor):
+            target.data.copy_(self.rho * target.data + (1 - self.rho) * current.data)
+        elif isinstance(target, np.ndarray):
+            target[:] = self.rho * target + (1 - self.rho) * current
         else:
             target = self.rho * target + (1 - self.rho) * current
-        self.target = target
+        return target