fix: initialisation weirdness

Author: hysmio

pufferlib/models.py

@@ -117,7 +117,14 @@ def __init__(self, env, policy, input_size=128, hidden_size=128):
             if "bias" in name:
                 nn.init.constant_(param, 0)
             elif "weight" in name and param.ndim >= 2:
-                nn.init.orthogonal_(param, 1.0)
+                if param.device.type == 'mps':
+                    # Apple MPS does not support orthogonal
+                    
+                    param.to(device='cpu')
+                    nn.init.orthogonal_(param, 1.0)
+                    param.to(device=param.device)
+                else:
+                    nn.init.orthogonal_(param, 1.0)
 
         self.lstm = nn.LSTM(input_size, hidden_size)
 

pufferlib/pytorch.py

@@ -164,7 +164,13 @@ def _flattened_tensor_size(native_dtype):
 
 def layer_init(layer, std=np.sqrt(2), bias_const=0.0):
     """CleanRL's default layer initialization"""
-    torch.nn.init.orthogonal_(layer.weight, std)
+    if layer.weight.device.type == 'mps':
+        # Apple MPS does not support orthogonal
+        layer.weight.to(device='cpu')
+        nn.init.orthogonal_(layer.weight, std)
+        layer.weight.to(device=layer.device)
+    else:
+        nn.init.orthogonal_(layer.weight, std)
     torch.nn.init.constant_(layer.bias, bias_const)
     return layer