Add support for custom policy in yaml file (#303)

* Add support for custom policy in yaml file

* Update changelog

* Add test for custom policies

Co-authored-by: Antonin Raffin <antonin.raffin@dlr.de>

Author: Rick-v-E

CHANGELOG.md

@@ -1,3 +1,17 @@
+## Release 1.7.0a0 (WIP)
+
+### Breaking Changes
+
+### New Features
+- Specifying custom policies in yaml file is now supported (@Rick-v-E)
+
+### Bug fixes
+
+### Documentation
+
+### Other
+
+
 ## Release 1.6.3 (2022-10-13)
 
 ### Breaking Changes
@@ -13,6 +27,7 @@
 ### Other
 - Used issue forms instead of issue templates
 
+
 ## Release 1.6.2.post2 (2022-10-10)
 
 ### Breaking Changes

README.md

@@ -206,6 +206,13 @@ Specify a different activation function for the network:
   policy_kwargs: "dict(activation_fn=nn.ReLU)"
 ```
 
+For a custom policy:
+
+```yaml
+  policy: my_package.MyCustomPolicy  # for instance stable_baselines3.ppo.MlpPolicy
+```
+
+
 ## Hyperparameter Tuning
 
 We use [Optuna](https://optuna.org/) for optimizing the hyperparameters.

rl_zoo3/exp_manager.py

@@ -47,7 +47,7 @@
 import rl_zoo3.import_envs  # noqa: F401 pytype: disable=import-error
 from rl_zoo3.callbacks import SaveVecNormalizeCallback, TrialEvalCallback
 from rl_zoo3.hyperparams_opt import HYPERPARAMS_SAMPLER
-from rl_zoo3.utils import ALGOS, get_callback_list, get_latest_run_id, get_wrapper_class, linear_schedule
+from rl_zoo3.utils import ALGOS, get_callback_list, get_class_by_name, get_latest_run_id, get_wrapper_class, linear_schedule
 
 
 class ExperimentManager:
@@ -390,6 +390,10 @@ def _preprocess_hyperparams(
             self.frame_stack = hyperparams["frame_stack"]
             del hyperparams["frame_stack"]
 
+        # import the policy when using a custom policy
+        if "policy" in hyperparams and "." in hyperparams["policy"]:
+            hyperparams["policy"] = get_class_by_name(hyperparams["policy"])
+
         # obtain a class object from a wrapper name string in hyperparams
         # and delete the entry
         env_wrapper = get_wrapper_class(hyperparams)

rl_zoo3/utils.py

@@ -2,7 +2,7 @@
 import glob
 import importlib
 import os
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
 
 import gym
 import stable_baselines3 as sb3  # noqa: F401
@@ -118,6 +118,26 @@ def wrap_env(env: gym.Env) -> gym.Env:
         return None
 
 
+def get_class_by_name(name: str) -> Type:
+    """
+    Imports and returns a class given the name, e.g. passing
+    'stable_baselines3.common.callbacks.CheckpointCallback' returns the
+    CheckpointCallback class.
+
+    :param name:
+    :return:
+    """
+
+    def get_module_name(name: str) -> str:
+        return ".".join(name.split(".")[:-1])
+
+    def get_class_name(name: str) -> str:
+        return name.split(".")[-1]
+
+    module = importlib.import_module(get_module_name(name))
+    return getattr(module, get_class_name(name))
+
+
 def get_callback_list(hyperparams: Dict[str, Any]) -> List[BaseCallback]:
     """
     Get one or more Callback class specified as a hyper-parameter
@@ -135,12 +155,6 @@ def get_callback_list(hyperparams: Dict[str, Any]) -> List[BaseCallback]:
     :return:
     """
 
-    def get_module_name(callback_name):
-        return ".".join(callback_name.split(".")[:-1])
-
-    def get_class_name(callback_name):
-        return callback_name.split(".")[-1]
-
     callbacks = []
 
     if "callback" in hyperparams.keys():
@@ -168,8 +182,8 @@ def get_class_name(callback_name):
                 kwargs = callback_dict[callback_name]
             else:
                 kwargs = {}
-            callback_module = importlib.import_module(get_module_name(callback_name))
-            callback_class = getattr(callback_module, get_class_name(callback_name))
+
+            callback_class = get_class_by_name(callback_name)
             callbacks.append(callback_class(**kwargs))
 
     return callbacks

rl_zoo3/version.txt

@@ -1 +1 @@
-1.6.3
+1.7.0a0

tests/test_train.py

@@ -121,6 +121,8 @@ def test_custom_yaml(tmp_path):
         "n_steps:50",
         "n_epochs:2",
         "batch_size:4",
+        # Test custom policy
+        "policy:'stable_baselines3.ppo.MlpPolicy'",
     ]
 
     return_code = subprocess.call(["python", "train.py"] + args)