Fix autotune for new ray tune

src/pytorch_mppi/autotune.py

@@ -108,6 +108,10 @@ def ensure_valid_value(self, value):
     def apply_parameter_value(self, value):
         """Apply the parameter value to the underlying object"""
 
+    @abc.abstractmethod
+    def attach_to_state(self, state: dict):
+        """Reattach/reinitialize the parameter to a new internal state. This should be similar to a call to __init__"""
+
     def get_parameter_value_from_config(self, config):
         """Get the serialized value of the parameter from a config dictionary, where each name is a scalar"""
         return config[self.name()]
@@ -118,10 +122,19 @@ def get_config_from_parameter_value(self, value):
 
 
 class MPPIParameter(TunableParameter, abc.ABC):
-    def __init__(self, mppi: MPPI):
+    def __init__(self, mppi: MPPI, dim=None):
         self.mppi = mppi
-        self.d = mppi.d
-        self.dtype = mppi.dtype
+        self._dim = dim
+        if self.mppi is not None:
+            self.d = self.mppi.d
+            self.dtype = self.mppi.dtype
+            if dim is None:
+                self._dim = self.mppi.nu
+
+    def attach_to_state(self, state: dict):
+        self.mppi = state['mppi']
+        self.d = self.mppi.d
+        self.dtype = self.mppi.dtype
 
 
 class SigmaParameter(MPPIParameter):
@@ -132,10 +145,10 @@ def name():
         return 'sigma'
 
     def dim(self):
-        return self.mppi.nu
+        return self._dim
 
     def get_current_parameter_value(self):
-        return torch.cat([self.mppi.noise_sigma[i][i].view(1) for i in range(self.mppi.nu)])
+        return torch.cat([self.mppi.noise_sigma[i][i].view(1) for i in range(self.dim())])
 
     def ensure_valid_value(self, value):
         sigma = ensure_tensor(self.d, self.dtype, value)
@@ -149,10 +162,10 @@ def apply_parameter_value(self, value):
         self.mppi.noise_sigma_inv = torch.inverse(self.mppi.noise_sigma.detach())
 
     def get_parameter_value_from_config(self, config):
-        return torch.tensor([config[f'{self.name()}{i}'] for i in range(self.mppi.nu)], dtype=self.dtype, device=self.d)
+        return torch.tensor([config[f'{self.name()}{i}'] for i in range(self.dim())], dtype=self.dtype, device=self.d)
 
     def get_config_from_parameter_value(self, value):
-        return {f'{self.name()}{i}': value[i].item() for i in range(self.mppi.nu)}
+        return {f'{self.name()}{i}': value[i].item() for i in range(self.dim())}
 
 
 class MuParameter(MPPIParameter):
@@ -161,7 +174,7 @@ def name():
         return 'mu'
 
     def dim(self):
-        return self.mppi.nu
+        return self._dim
 
     def get_current_parameter_value(self):
         return self.mppi.noise_mu.clone()
@@ -176,10 +189,10 @@ def apply_parameter_value(self, value):
         self.mppi.noise_sigma_inv = torch.inverse(self.mppi.noise_sigma.detach())
 
     def get_parameter_value_from_config(self, config):
-        return torch.tensor([config[f'{self.name()}{i}'] for i in range(self.mppi.nu)], dtype=self.dtype, device=self.d)
+        return torch.tensor([config[f'{self.name()}{i}'] for i in range(self.dim())], dtype=self.dtype, device=self.d)
 
     def get_config_from_parameter_value(self, value):
-        return {f'{self.name()}{i}': value[i].item() for i in range(self.mppi.nu)}
+        return {f'{self.name()}{i}': value[i].item() for i in range(self.dim())}
 
 
 class LambdaParameter(MPPIParameter):
@@ -236,15 +249,25 @@ class Autotune:
     eps = 0.0001
 
     def __init__(self, params_to_tune: typing.Sequence[TunableParameter],
-                 evaluate_fn: typing.Callable[[], EvaluationResult], optimizer=CMAESOpt()):
+                 evaluate_fn: typing.Callable[[], EvaluationResult],
+                 reload_state_fn: typing.Callable[[], dict] = None,
+                 optimizer=CMAESOpt()):
+        """
+
+        :param params_to_tune: sequence of tunable parameters
+        :param evaluate_fn: function that returns an EvaluationResult that we want to minimize
+        :param reload_state_fn: function that returns a dictionary of state to reattach to the parameters
+        :param optimizer: optimizer that searches in the parameter space
+        """
         self.evaluate_fn = evaluate_fn
+        self.reload_state_fn = reload_state_fn
 
         self.params = params_to_tune
         self.optim = optimizer
         self.optim.tuner = self
         self.results = []
 
-        self.get_parameter_values(self.params)
+        self.attach_parameters()
         self.optim.setup_optimization()
 
     def optimize_step(self) -> EvaluationResult:
@@ -303,6 +326,17 @@ def apply_parameters(self, param_values):
         for p in self.params:
             p.apply_parameter_value(param_values[p.name()])
 
+    def attach_parameters(self):
+        """Attach parameters to any underlying state they require In most cases the parameters are defined already
+        attached to whatever state it needs, e.g. the MPPI controller object for changing the parameter values.
+        However, there are cases where the full state is not serializable, e.g. when using a multiprocessing pool
+        and so we pass only the information required to load the state. We then must load the state and reattach
+        the parameters to the state each training iteration."""
+        if self.reload_state_fn is not None:
+            state = self.reload_state_fn()
+            for p in self.params:
+                p.attach_to_state(state)
+
     def config_to_params(self, config):
         """Configs are param dictionaries where each must be a scalar"""
         return {p.name(): p.get_parameter_value_from_config(config) for p in self.params}

src/pytorch_mppi/autotune_global.py

@@ -1,8 +1,10 @@
 import abc
 import numpy as np
+import torch.cuda
 
 # pip install "ray[tune]" bayesian-optimization hyperopt
 from ray import tune
+from ray import train
 
 from pytorch_mppi import autotune
 from ray.tune.search.hyperopt import HyperOptSearch
@@ -47,35 +49,35 @@ def _linearize_space_value(space, v):
 
 
 class SigmaGlobalParameter(autotune.SigmaParameter, GlobalTunableParameter):
-    def __init__(self, *args, search_space=tune.loguniform(1e-4, 1e2)):
-        super().__init__(*args)
+    def __init__(self, *args, search_space=tune.loguniform(1e-4, 1e2), **kwargs):
+        super().__init__(*args, **kwargs)
         GlobalTunableParameter.__init__(self, search_space)
 
     def total_search_space(self) -> dict:
-        return {f"{self.name()}{i}": self.search_space for i in range(self.mppi.nu)}
+        return {f"{self.name()}{i}": self.search_space for i in range(self.dim())}
 
 
 class MuGlobalParameter(autotune.MuParameter, GlobalTunableParameter):
-    def __init__(self, *args, search_space=tune.uniform(-1, 1)):
-        super().__init__(*args)
+    def __init__(self, *args, search_space=tune.uniform(-1, 1), **kwargs):
+        super().__init__(*args, **kwargs)
         GlobalTunableParameter.__init__(self, search_space)
 
     def total_search_space(self) -> dict:
-        return {f"{self.name()}{i}": self.search_space for i in range(self.mppi.nu)}
+        return {f"{self.name()}{i}": self.search_space for i in range(self.dim())}
 
 
 class LambdaGlobalParameter(autotune.LambdaParameter, GlobalTunableParameter):
-    def __init__(self, *args, search_space=tune.loguniform(1e-5, 1e3)):
-        super().__init__(*args)
+    def __init__(self, *args, search_space=tune.loguniform(1e-5, 1e3), **kwargs):
+        super().__init__(*args, **kwargs)
         GlobalTunableParameter.__init__(self, search_space)
 
     def total_search_space(self) -> dict:
         return {self.name(): self.search_space}
 
 
 class HorizonGlobalParameter(autotune.HorizonParameter, GlobalTunableParameter):
-    def __init__(self, *args, search_space=tune.randint(1, 50)):
-        super().__init__(*args)
+    def __init__(self, *args, search_space=tune.randint(1, 50), **kwargs):
+        super().__init__(*args, **kwargs)
         GlobalTunableParameter.__init__(self, search_space)
 
     def total_search_space(self) -> dict:
@@ -124,8 +126,10 @@ def setup_optimization(self):
         init = self.tuner.initial_value()
 
         hyperopt_search = self.search_alg(points_to_evaluate=[init], metric="cost", mode="min")
+
+        trainable_with_resources = tune.with_resources(self.trainable, {"gpu": 1 if torch.cuda.is_available() else 0})
         self.optim = tune.Tuner(
-            self.trainable,
+            trainable_with_resources,
             tune_config=tune.TuneConfig(
                 num_samples=self.iterations,
                 search_alg=hyperopt_search,
@@ -136,9 +140,10 @@ def setup_optimization(self):
         )
 
     def trainable(self, config):
+        self.tuner.attach_parameters()
         self.tuner.apply_parameters(self.tuner.config_to_params(config))
         res = self.tuner.evaluate_fn()
-        tune.report(cost=res.costs.mean().item())
+        train.report({'cost': res.costs.mean().item()})
 
     def optimize_step(self):
         raise RuntimeError("Ray optimizers only allow tuning of all iterations at once")