Fixed training continuation

.github/workflows/ci.yml

@@ -17,7 +17,7 @@ jobs:
 
     strategy:
       matrix:
-        os: [ubuntu-latest, windows-latest]
+        os: [ubuntu-latest, windows-latest, macos-latest]
         python-version: ['3.10']
 
     steps:
@@ -36,10 +36,19 @@ jobs:
         # sudo apt-get install -y lcov g++ gcc
         sudo apt-get install -y g++ gcc
 
-    - name: Install dependencies on Windows
-      if: matrix.os == 'windows-latest'
+    - name: Install LLVM via Homebrew
+      if: matrix.os == 'macos-latest'
+      run: brew install libomp llvm
+
+    - name: Configure Environment Variables
+      if: matrix.os == 'macos-latest'
       run: |
-        python -m pip install --upgrade pip
+        echo "PATH=$(brew --prefix llvm)/bin:$PATH" >> $GITHUB_ENV
+        echo "LDFLAGS=-L$(brew --prefix libomp)/lib -L$(brew --prefix llvm)/lib -L$(brew --prefix llvm)/lib/c++ -Wl,-rpath,$(brew --prefix llvm)/lib/c++" >> $GITHUB_ENV
+        echo "CPPFLAGS=-I$(brew --prefix libomp)/include -I$(brew --prefix llvm)/include" >> $GITHUB_ENV
+        echo "CC=$(brew --prefix llvm)/bin/clang" >> $GITHUB_ENV
+        echo "CXX=$(brew --prefix llvm)/bin/clang++" >> $GITHUB_ENV
+        echo "DYLD_LIBRARY_PATH=$(brew --prefix llvm)/lib:$(brew --prefix libomp)/lib" >> $GITHUB_ENV
 
     - name: Set up MSVC environment
       if: matrix.os == 'windows-latest'

CMakeLists.txt

@@ -236,7 +236,6 @@ if (USE_CUDA)
 endif()
 # Platform-specific settings and linking
 if (CMAKE_SYSTEM_NAME STREQUAL "Linux")
-
     target_link_libraries(gbrl_cpp PRIVATE OpenMP::OpenMP_CXX)
     if (USE_CUDA)
         target_link_libraries(gbrl_cpp PRIVATE CUDA::cudart)

gbrl/__init__.py

@@ -6,7 +6,7 @@
 # https://nvlabs.github.io/gbrl/license.html 
 #
 ##############################################################################
-__version__ = "1.0.8"
+__version__ = "1.0.9"
 
 import importlib.util
 import os
@@ -39,7 +39,21 @@ def load_cpp_module():
                     spec.loader.exec_module(module)
                     _loaded_cpp_module = module = module
                     return module
-
+
+    if platform.system() == "Darwin":  # check for .so on Darwin
+        ext = ".so"
+        for dir_path  in possible_paths:
+            if os.path.exists(dir_path):
+            # Scan for files that match the module name and extension
+                for file_name in os.listdir(dir_path):
+                    if file_name.startswith(module_name) and file_name.endswith(ext):
+                        # Dynamically load the matching shared library
+                        file_path = os.path.join(dir_path, file_name)
+                        spec = importlib.util.spec_from_file_location(module_name, file_path)
+                        module = importlib.util.module_from_spec(spec)
+                        spec.loader.exec_module(module)
+                        _loaded_cpp_module = module = module
+                        return module
     raise ImportError(f"Could not find {module_name}{ext} in any of the expected locations: {possible_paths}")
 
 

gbrl/ac_gbrl.py

@@ -21,6 +21,7 @@
 class ActorCritic(GBRL):
     def __init__(self, 
                  tree_struct: Dict,
+                 input_dim: int,
                  output_dim: int, 
                  policy_optimizer: Dict,
                  value_optimizer: Dict= None,
@@ -56,6 +57,7 @@ def __init__(self,
         if value_optimizer is not None:
             value_optimizer = setup_optimizer(value_optimizer, prefix='value_')
         super().__init__(tree_struct,
+                         input_dim,
                          output_dim,
                          None,
                          gbrl_params,
@@ -68,11 +70,13 @@ def __init__(self,
         self.bias = bias if bias is not None else np.zeros(self.output_dim if shared_tree_struct else self.output_dim - 1, dtype=numerical_dtype)
         # init model
         if self.shared_tree_struct:
-            self._model = SharedActorCriticWrapper(self.output_dim, self.tree_struct, self.policy_optimizer, self.value_optimizer, self.gbrl_params, self.verbose, self.device) 
+            self._model = SharedActorCriticWrapper(self.input_dim, self.output_dim, self.tree_struct, self.policy_optimizer, self.value_optimizer, self.gbrl_params, self.verbose, self.device) 
+            self._model.reset()
+            self._model.set_bias(self.bias)
         else:
-            self._model = SeparateActorCriticWrapper(self.output_dim, self.tree_struct, self.policy_optimizer, self.value_optimizer, self.gbrl_params, self.verbose, self.device)
-        self._model.reset()
-        self._model.set_bias(self.bias)
+            self._model = SeparateActorCriticWrapper(self.input_dim, self.output_dim, self.tree_struct, self.policy_optimizer, self.value_optimizer, self.gbrl_params, self.verbose, self.device)
+            self._model.reset()
+            self._model.set_policy_bias(self.bias)
         self.policy_grad = None 
         self.value_grad = None
 
@@ -100,6 +104,7 @@ def load_model(cls, load_name: str, device: str) -> "ActorCritic":
             instance.bias = instance._model.get_bias()
         instance.value_optimizer = instance._model.value_optimizer
         instance.policy_optimizer = instance._model.policy_optimizer
+        instance.input_dim = instance._model.input_dim
         instance.output_dim = instance._model.output_dim
         instance.verbose = instance._model.verbose
         instance.tree_struct = instance._model.tree_struct
@@ -257,14 +262,15 @@ def copy(self) -> "ActorCritic":
 
     def __copy__(self) -> "ActorCritic":
         value_optimizer = None if self.value_optimizer is None else self.value_optimizer.copy()
-        copy_ = ActorCritic(self.tree_struct.copy(), self.output_dim, self.policy_optimizer.copy(), value_optimizer, self.shared_tree_struct, self.gbrl_params, self.bias, self.verbose, self.device)
+        copy_ = ActorCritic(self.tree_struct.copy(), self.input_dim, self.output_dim, self.policy_optimizer.copy(), value_optimizer, self.shared_tree_struct, self.gbrl_params, self.bias, self.verbose, self.device)
         if self._model is not None:
             copy_._model = self._model.copy()
         return copy_
 
 class ParametricActor(GBRL):
     def __init__(self, 
                  tree_struct: Dict,
+                 input_dim: int, 
                  output_dim: int, 
                  policy_optimizer: Dict,
                  gbrl_params: Dict=dict(),
@@ -294,6 +300,7 @@ def __init__(self,
         """
         policy_optimizer = setup_optimizer(policy_optimizer, prefix='policy_')
         super().__init__(tree_struct,
+                         input_dim,
                          output_dim,
                          None,
                          gbrl_params,
@@ -341,6 +348,7 @@ def load_model(cls, load_name: str, device: str) -> "ParametricActor":
         instance._model = GBTWrapper.load(load_name, device)
         instance.bias = instance._model.get_bias()
         instance.policy_optimizer = instance._model.optimizer
+        instance.input_dim = instance._model.input_dim
         instance.output_dim = instance._model.output_dim
         instance.verbose = instance._model.verbose
         instance.tree_struct = instance._model.tree_struct
@@ -378,14 +386,15 @@ def __call__(self, observations: Union[np.ndarray, th.Tensor], requires_grad : b
         return params
 
     def __copy__(self) -> "ParametricActor":
-        copy_ = ParametricActor(self.tree_struct.copy(), self.output_dim, self.policy_optimizer.copy(),  self.gbrl_params, self.bias, self.verbose, self.device)
+        copy_ = ParametricActor(self.tree_struct.copy(), self.input_dim, self.output_dim, self.policy_optimizer.copy(),  self.gbrl_params, self.bias, self.verbose, self.device)
         if self._model is not None:
             copy_._model = self._model.copy()
         return copy_
 
 class GaussianActor(GBRL):
     def __init__(self, 
                  tree_struct: Dict,
+                 input_dim: int,
                  output_dim: int, 
                  mu_optimizer: Dict,
                  std_optimizer: Dict = None,
@@ -433,6 +442,7 @@ def __init__(self,
         self.log_std_init = log_std_init
 
         super().__init__(tree_struct,
+                         input_dim,
                          output_dim,
                          None, 
                          gbrl_params,
@@ -442,7 +452,7 @@ def __init__(self,
 
         self.policy_dim = policy_dim
         # init model
-        self._model = GBTWrapper(self.output_dim, self.tree_struct, [mu_optimizer, std_optimizer], self.gbrl_params, self.verbose, self.device)
+        self._model = GBTWrapper(self.input_dim, self.output_dim, self.tree_struct, [mu_optimizer, std_optimizer], self.gbrl_params, self.verbose, self.device)
         self._model.reset()
         self._model.set_bias(self.bias)
 
@@ -513,7 +523,7 @@ def __call__(self, observations: Union[np.ndarray, th.Tensor], requires_grad : b
 
     def __copy__(self) -> "GaussianActor":
         std_optimizer = None if self.std_optimizer is None else self.std_optimizer.copy()
-        copy_ = GaussianActor(self.tree_struct.copy(), self.output_dim, self.mu_optimizer.copy(), std_optimizer, self.gbrl_params, self.bias, self.verbose, self.device)
+        copy_ = GaussianActor(self.tree_struct.copy(), self.input_dim, self.output_dim, self.mu_optimizer.copy(), std_optimizer, self.gbrl_params, self.bias, self.verbose, self.device)
         if self._model is not None:
             copy_._model = self._model.copy()
         return copy_
@@ -522,6 +532,7 @@ def __copy__(self) -> "GaussianActor":
 class ContinuousCritic(GBRL):
     def __init__(self, 
                  tree_struct: Dict,
+                 input_dim: int,
                  output_dim: int, 
                  weights_optimizer: Dict,
                  bias_optimizer: Dict =  None,
@@ -564,6 +575,7 @@ def __init__(self,
         bias_optimizer = setup_optimizer(bias_optimizer, prefix='bias_')
 
         super().__init__(tree_struct,
+                         input_dim,
                          output_dim,
                          None, 
                          gbrl_params,
@@ -575,7 +587,7 @@ def __init__(self,
         self.bias = bias if bias is not None else np.zeros(self.output_dim, dtype=numerical_dtype)
         self.target_update_interval = target_update_interval
         # init model
-        self._model = GBTWrapper(self.output_dim, self.tree_struct, [weights_optimizer, bias_optimizer], self.gbrl_params, self.verbose, self.device)
+        self._model = GBTWrapper(self.input_dim, self.output_dim, self.tree_struct, [weights_optimizer, bias_optimizer], self.gbrl_params, self.verbose, self.device)
         self._model.reset()
         self._model.set_bias(self.bias)
 
@@ -657,7 +669,7 @@ def __call__(self, observations: Union[np.ndarray, th.Tensor], requires_grad: bo
         return weights, bias
 
     def __copy__(self) -> "ContinuousCritic":
-        copy_ = ContinuousCritic(self.tree_struct.copy(), self.output_dim, self.weights_optimizer.copy(), self.bias_optimizer.copy() if isinstance(self.critic_optimizer, dict) else {'weights_optimizer': self.critic_optimizer[0], 'bias_optimizer': self.critic_optimizer[1]}, self.gbrl_params, self.target_update_interval, self.bias, self.verbose, self.device)
+        copy_ = ContinuousCritic(self.tree_struct.copy(), self.input_dim, self.output_dim, self.weights_optimizer.copy(), self.bias_optimizer.copy() if isinstance(self.critic_optimizer, dict) else {'weights_optimizer': self.critic_optimizer[0], 'bias_optimizer': self.critic_optimizer[1]}, self.gbrl_params, self.target_update_interval, self.bias, self.verbose, self.device)
         if self._model is not None:
             copy_._model = self._model.copy()
         return copy_
@@ -666,6 +678,7 @@ def __copy__(self) -> "ContinuousCritic":
 class DiscreteCritic(GBRL):
     def __init__(self, 
                  tree_struct: Dict,
+                 input_dim: int,
                  output_dim: int, 
                  critic_optimizer: Dict,
                  gbrl_params: Dict=dict(),
@@ -695,6 +708,7 @@ def __init__(self,
         """
         critic_optimizer = setup_optimizer(critic_optimizer, prefix='critic_')
         super().__init__(tree_struct,
+                         input_dim,
                          output_dim,
                          None, 
                          gbrl_params,
@@ -704,7 +718,7 @@ def __init__(self,
         self.target_update_interval = target_update_interval
         self.bias = bias if bias is not None else np.zeros(self.output_dim, dtype=numerical_dtype)
         # init model
-        self._model = GBTWrapper(self.output_dim, self.tree_struct, self.critic_optimizer, self.gbrl_params, self.verbose, self.device)
+        self._model = GBTWrapper(self.input_dim, self.output_dim, self.tree_struct, self.critic_optimizer, self.gbrl_params, self.verbose, self.device)
         self._model.reset()
         self._model.set_bias(self.bias)
 
@@ -764,7 +778,7 @@ def predict_target(self, observations: Union[np.ndarray, th.Tensor], tensor: boo
         return self._model.predict(observations, requires_grad=False, stop_idx=max(n_trees - self.target_update_interval, 1), tensor=tensor)
 
     def __copy__(self) -> "DiscreteCritic":
-        copy_ = DiscreteCritic(self.tree_struct.copy(), self.output_dim, self.critic_optimizer.copy(), self.gbrl_params, self.target_update_interval, self.bias, self.verbose, self.device)
+        copy_ = DiscreteCritic(self.tree_struct.copy(), self.input_dim, self.output_dim, self.critic_optimizer.copy(), self.gbrl_params, self.target_update_interval, self.bias, self.verbose, self.device)
         if self._model is not None:
             copy_._model = self._model.copy()
         return copy_