Configurable inference device, use MPS on mac

rlgym_ppo/learner.py

@@ -18,6 +18,7 @@
 
 import numpy as np
 import torch
+import torch.backends
 import wandb
 from wandb.wandb_run import Run
 from rlgym_sim import gym
@@ -76,7 +77,9 @@ def __init__(
         instance_launch_delay: Union[float,None] = None,
         random_seed: int = 123,
         n_checkpoints_to_keep: int = 5,
-        device: str = "auto"):
+
+        device: str = "auto",
+        inference_device: str = "auto"):
         assert (
             env_create_function is not None
         ), "MUST PROVIDE A FUNCTION TO CREATE RLGYM FUNCTIONS TO INITIALIZE RLGYM-PPO"
@@ -105,12 +108,31 @@ def __init__(
 
         if device in {"auto", "gpu"} and torch.cuda.is_available():
             self.device = "cuda:0"
-        elif device == "auto" and not torch.cuda.is_available():
+        elif device in {"auto", "gpu"} and hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+            self.device = "mps"
+        elif device == "auto":
             self.device = "cpu"
         else:
             self.device = device
+
+        if inference_device in {"auto", "gpu"} and torch.cuda.is_available():
+            # if multiple cuda devices are available, use the second one for inference
+            if torch.cuda.device_count() > 1:
+                self.inference_device = "cuda:1"
+            else:
+                self.inference_device = "cuda:0"
+        elif inference_device == "gpu" and hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+            # macos, apple m1, etc
+            # only go with mps for inference if gpu is explicitly requested.
+            # based on testing, it's quite a bit slower than using mps for
+            # backprop and CPU for inference
+            inference_device = "mps"
+        elif inference_device == "auto":
+            self.inference_device = "cpu"
+        else:
+            self.inference_device = inference_device
 
-        print(f"Using device {self.device}")
+        print(f"Using device {self.device} for backprop and {self.inference_device} for inference")
         self.exp_buffer_size = exp_buffer_size
         self.timestep_limit = timestep_limit
         self.ts_per_epoch = ts_per_iteration
@@ -160,7 +182,10 @@ def __init__(
             ent_coef=ppo_ent_coef,
         )
 
-        self.agent.policy = self.ppo_learner.policy
+        if self.device == self.inference_device:
+            self.agent.policy = self.ppo_learner.policy
+        else:
+            self.agent.policy = self.ppo_learner.policy.clone(to=self.inference_device)
 
         self.wandb_run = wandb_run
         wandb_loaded = checkpoint_load_folder is not None and self.load(
@@ -207,6 +232,12 @@ def _learn(self):
         # While the number of timesteps we have collected so far is less than the
         # amount we are allowed to collect.
         while self.agent.cumulative_timesteps < self.timestep_limit:
+
+            # if we're using different devices for inference and backprop, we need to
+            # copy the current policy to the inference device
+            if self.device != self.inference_device:
+                self.agent.policy = self.ppo_learner.policy.clone(to=self.inference_device)
+
             epoch_start = time.perf_counter()
             report = {}
 
@@ -255,7 +286,7 @@ def _learn(self):
             report.clear()
             ppo_report.clear()
 
-            if "cuda" in self.device:
+            if "cuda" in self.device or "cuda" in self.inference_device:
                 torch.cuda.empty_cache()
 
             # Check if keyboard press

rlgym_ppo/ppo/continuous_policy.py

@@ -23,6 +23,11 @@
 class ContinuousPolicy(nn.Module):
     def __init__(self, input_shape, output_shape, layer_sizes, device, var_min=0.1, var_max=1.0):
         super().__init__()
+        self.input_shape = input_shape
+        self.output_shape = output_shape
+        self.layer_sizes = layer_sizes
+        self.var_min = var_min
+        self.var_max = var_max
         self.device = device
         self.affine_map = torch_functions.MapContinuousToAction(range_min=var_min, range_max=var_max)
 
@@ -39,6 +44,15 @@ def __init__(self, input_shape, output_shape, layer_sizes, device, var_min=0.1,
         layers.append(nn.Linear(layer_sizes[-1], output_shape))
         layers.append(nn.Tanh())
         self.model = nn.Sequential(*layers).to(self.device)
+
+    def clone(self, to: str = None):
+        device = self.device if to is None else to
+        cloned_policy = ContinuousPolicy(self.input_shape, self.output_shape, self.layer_sizes, device, self.var_min, self.var_max)
+        cloned_policy.load_state_dict(self.state_dict())
+        if to is not None:
+            cloned_policy.to(to)
+
+        return cloned_policy
 
     @functools.lru_cache()
     def logpdf(self, x, mean, std):
@@ -118,4 +132,4 @@ def get_backprop_data(self, obs, acts, summed_probs=True):
         entropy = distribution.entropy()
         entropy = entropy.mean().to(self.device)
 
-        return log_probs, entropy
+        return log_probs, entropy

rlgym_ppo/ppo/discrete_policy.py

@@ -17,6 +17,8 @@ class DiscreteFF(nn.Module):
     def __init__(self, input_shape, n_actions, layer_sizes, device):
         super().__init__()
         self.device = device
+        self.input_shape = input_shape
+        self.layer_sizes = layer_sizes
 
         assert len(layer_sizes) != 0, "AT LEAST ONE LAYER MUST BE SPECIFIED TO BUILD THE NEURAL NETWORK!"
         layers = [nn.Linear(input_shape, layer_sizes[0]), nn.ReLU()]
@@ -32,6 +34,14 @@ def __init__(self, input_shape, n_actions, layer_sizes, device):
 
         self.n_actions = n_actions
 
+    def clone(self, to: str = None):
+        device = self.device if to is None else to
+        cloned_policy = DiscreteFF(self.input_shape, self.n_actions, self.layer_sizes, device)
+        cloned_policy.load_state_dict(self.state_dict())
+        if to is not None:
+            cloned_policy.to(to)
+
+        return cloned_policy
     def get_output(self, obs):
         t = type(obs)
         if t != torch.Tensor:
@@ -77,4 +87,4 @@ def get_backprop_data(self, obs, acts):
         action_log_probs = log_probs.gather(-1, acts)
         entropy = -(log_probs * probs).sum(dim=-1)
 
-        return action_log_probs.to(self.device), entropy.to(self.device).mean()
+        return action_log_probs.to(self.device), entropy.to(self.device).mean()

rlgym_ppo/ppo/multi_discrete_policy.py

@@ -17,6 +17,7 @@ class MultiDiscreteFF(nn.Module):
     def __init__(self, input_shape, layer_sizes, device):
         super().__init__()
         self.device = device
+        self.layer_sizes = layer_sizes
         bins = [3,3,3,3,3,2,2,2]
         n_output_nodes = sum(bins)
         assert len(layer_sizes) != 0, "AT LEAST ONE LAYER MUST BE SPECIFIED TO BUILD THE NEURAL NETWORK!"
@@ -33,6 +34,13 @@ def __init__(self, input_shape, layer_sizes, device):
         self.splits = bins
         self.multi_discrete = torch_functions.MultiDiscreteRolv(bins)
 
+    def clone(self, to: str = None):
+        device = self.device if to is None else to
+        cloned_policy = MultiDiscreteFF(self.input_shape, self.layer_sizes, device)
+        cloned_policy.load_state_dict(self.state_dict())
+        if to is not None:
+            cloned_policy.to(to)
+
     def get_output(self, obs):
         t = type(obs)
         if t != torch.Tensor:
@@ -86,4 +94,4 @@ def get_backprop_data(self, obs, acts):
         entropy = distribution.entropy().to(self.device)
         log_probs = distribution.log_prob(acts).to(self.device)
 
-        return log_probs, entropy.mean()
+        return log_probs, entropy.mean()

rlgym_ppo/ppo/value_estimator.py

@@ -13,6 +13,8 @@
 class ValueEstimator(nn.Module):
     def __init__(self, input_shape, layer_sizes, device):
         super().__init__()
+        self.input_shape = input_shape
+        self.layer_sizes = layer_sizes
         self.device = device
 
         assert len(layer_sizes) != 0, "AT LEAST ONE LAYER MUST BE SPECIFIED TO BUILD THE NEURAL NETWORK!"
@@ -26,6 +28,14 @@ def __init__(self, input_shape, layer_sizes, device):
 
         layers.append(nn.Linear(layer_sizes[-1], 1))
         self.model = nn.Sequential(*layers).to(self.device)
+
+
+    def clone(self, to: str = None):
+        device = self.device if to is None else to
+        cloned_value_estimator = ValueEstimator(self.input_shape, self.layer_sizes, device)
+        cloned_value_estimator.load_state_dict(self.state_dict())
+        if to is not None:
+            cloned_value_estimator.to(to)
 
     def forward(self, x):
         t = type(x)
@@ -34,4 +44,4 @@ def forward(self, x):
                 x = np.asarray(x)
             x = torch.as_tensor(x, dtype=torch.float32, device=self.device)
 
-        return self.model(x)
+        return self.model(x)