Add class mapping of categorical model (#216)

* Add class mapping for Categorical model * Apply format --------- Co-authored-by: Toni-SM <[email protected]>
Toni-SM · Nov 3, 2024 · eff7295 · eff7295
1 parent 9252ec9
commit eff7295
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/skrl/utils/runner/jax/runner.py b/skrl/utils/runner/jax/runner.py
@@ -14,7 +14,7 @@
 from skrl.resources.schedulers.jax import KLAdaptiveLR  # noqa
 from skrl.trainers.jax import SequentialTrainer, Trainer
 from skrl.utils import set_seed
-from skrl.utils.model_instantiators.jax import deterministic_model, gaussian_model
+from skrl.utils.model_instantiators.jax import categorical_model, deterministic_model, gaussian_model
 
 
 class Runner:
@@ -35,6 +35,7 @@ def __init__(self, env: Union[Wrapper, MultiAgentEnvWrapper], cfg: Mapping[str,
         self._class_mapping = {
             # model
             "gaussianmixin": gaussian_model,
+            "categoricalmixin": categorical_model,
             "deterministicmixin": deterministic_model,
             "shared": None,
             # memory

diff --git a/skrl/utils/runner/torch/runner.py b/skrl/utils/runner/torch/runner.py
@@ -14,7 +14,7 @@
 from skrl.resources.schedulers.torch import KLAdaptiveLR  # noqa
 from skrl.trainers.torch import SequentialTrainer, Trainer
 from skrl.utils import set_seed
-from skrl.utils.model_instantiators.torch import deterministic_model, gaussian_model, shared_model
+from skrl.utils.model_instantiators.torch import categorical_model, deterministic_model, gaussian_model, shared_model
 
 
 class Runner:
@@ -35,6 +35,7 @@ def __init__(self, env: Union[Wrapper, MultiAgentEnvWrapper], cfg: Mapping[str,
         self._class_mapping = {
             # model
             "gaussianmixin": gaussian_model,
+            "categoricalmixin": categorical_model,
             "deterministicmixin": deterministic_model,
             "shared": shared_model,
             # memory