changing imports

Author: benja263

gbrl/__init__.py

@@ -8,10 +8,46 @@
 ##############################################################################
 __version__ = "1.0.7"
 
-from .ac_gbrl import (ActorCritic, GaussianActor, ContinuousCritic,
-                   DiscreteCritic, ParametricActor)
-from .gbt import GBRL
-from .gbrl_cpp import GBRL as GBRL_CPP
+import importlib.util
+import os
+import platform
+
+_loaded_cpp_module = None
+
+def load_cpp_module():
+    global _loaded_cpp_module
+    module_name = "gbrl_cpp"
+    if platform.system() == "Windows":
+        ext = ".pyd"
+    elif platform.system() == "Darwin":  # macOS
+        ext = ".dylib"
+    else:  # Assume Linux/Unix
+        ext = ".so"
+    possible_paths = [
+        os.path.join(os.path.dirname(__file__)),  # Current directory
+        os.path.join(os.path.dirname(__file__), "Release"),  # Release folder
+    ]
+    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}")
+
+
+# Load the C++ module dynamically
+_gbrl_cpp_module = load_cpp_module()
+
+# Create a global alias for the GBRL class
+GBRL_CPP = _gbrl_cpp_module.GBRL
 
 cuda_available = GBRL_CPP.cuda_available
 

gbrl/ac_gbrl.py

@@ -11,10 +11,9 @@
 import numpy as np
 import torch as th
 
-from .gbrl_wrapper import (GBTWrapper, SeparateActorCriticWrapper,
-                           SharedActorCriticWrapper, )
-from .gbt import GBRL
-from .utils import (setup_optimizer, clip_grad_norm, numerical_dtype, 
+from gbrl.gbrl_wrapper import GBTWrapper, SeparateActorCriticWrapper, SharedActorCriticWrapper, 
+from gbrl.gbt import GBRL
+from gbrl.utils import (setup_optimizer, clip_grad_norm, numerical_dtype, 
                     concatenate_arrays, validate_array, constant_like,
                     tensor_to_leaf)
 

gbrl/gbrl_wrapper.py

@@ -13,8 +13,8 @@
 import numpy as np
 import torch as th
 
-from .gbrl_cpp import GBRL as GBRL_CPP
-from .utils import (get_input_dim, get_poly_vectors, 
+from gbrl import GBRL_CPP
+from gbrl.utils import (get_input_dim, get_poly_vectors, 
                     to_numpy,
                     numerical_dtype, 
                     get_tensor_info,

gbrl/gbt.py

@@ -12,8 +12,8 @@
 import torch as th
 
 
-from .gbrl_wrapper import GBTWrapper
-from .utils import setup_optimizer, clip_grad_norm, validate_array
+from gbrl.gbrl_wrapper import GBTWrapper
+from gbrl.utils import setup_optimizer, clip_grad_norm, validate_array
 
 
 class GBRL:

gbrl/utils.py

@@ -11,7 +11,7 @@
 import torch as th
 from scipy.special import binom
 
-from .config import APPROVED_OPTIMIZERS, VALID_OPTIMIZER_ARGS
+from gbrl.config import APPROVED_OPTIMIZERS, VALID_OPTIMIZER_ARGS
 
 import numpy as np 
 # Define custom dtypes