Update annotation and doctring

Toni-SM · Aug 13, 2024 · afadb50 · afadb50
1 parent 4d55d04
commit afadb50
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 16 deletions.
diff --git a/skrl/utils/model_instantiators/jax/__init__.py b/skrl/utils/model_instantiators/jax/__init__.py
@@ -152,7 +152,7 @@ def gaussian_model(observation_space: Optional[Union[int, Tuple[int], gym.Space,
                    output_shape: Shape = Shape.ACTIONS,
                    output_activation: Optional[str] = "tanh",
                    output_scale: float = 1.0,
-                   return_source: bool = False) -> Model:
+                   return_source: bool = False) -> Union[Model, str]:
     """Instantiate a Gaussian model
 
     :param observation_space: Observation/state space or shape (default: None).
@@ -191,7 +191,7 @@ def gaussian_model(observation_space: Optional[Union[int, Tuple[int], gym.Space,
                           instantiate the model rather than the model instance (default: False).
     :type return_source: bool, optional
 
-    :return: Gaussian model instance
+    :return: Gaussian model instance or definition source
     :rtype: Model
     """
     # network
@@ -242,7 +242,7 @@ def deterministic_model(observation_space: Optional[Union[int, Tuple[int], gym.S
                         output_shape: Shape = Shape.ACTIONS,
                         output_activation: Optional[str] = "tanh",
                         output_scale: float = 1.0,
-                        return_source: bool = False) -> Model:
+                        return_source: bool = False) -> Union[Model, str]:
     """Instantiate a deterministic model
 
     :param observation_space: Observation/state space or shape (default: None).
@@ -273,7 +273,7 @@ def deterministic_model(observation_space: Optional[Union[int, Tuple[int], gym.S
                           instantiate the model rather than the model instance (default: False).
     :type return_source: bool, optional
 
-    :return: Deterministic model instance
+    :return: Deterministic model instance or definition source
     :rtype: Model
     """
     # network
@@ -318,7 +318,7 @@ def categorical_model(observation_space: Optional[Union[int, Tuple[int], gym.Spa
                       hidden_activation: list = ["relu", "relu"],
                       output_shape: Shape = Shape.ACTIONS,
                       output_activation: Optional[str] = None,
-                      return_source: bool = False) -> Model:
+                      return_source: bool = False) -> Union[Model, str]:
     """Instantiate a categorical model
 
     :param observation_space: Observation/state space or shape (default: None).
@@ -349,7 +349,7 @@ def categorical_model(observation_space: Optional[Union[int, Tuple[int], gym.Spa
                           instantiate the model rather than the model instance (default: False).
     :type return_source: bool, optional
 
-    :return: Categorical model instance
+    :return: Categorical model instance or definition source
     :rtype: Model
     """
     # network

diff --git a/skrl/utils/model_instantiators/torch/__init__.py b/skrl/utils/model_instantiators/torch/__init__.py
@@ -161,7 +161,7 @@ def gaussian_model(observation_space: Optional[Union[int, Tuple[int], gym.Space,
                    output_shape: Shape = Shape.ACTIONS,
                    output_activation: Optional[str] = "tanh",
                    output_scale: float = 1.0,
-                   return_source: bool = False) -> Model:
+                   return_source: bool = False) -> Union[Model, str]:
     """Instantiate a Gaussian model
 
     :param observation_space: Observation/state space or shape (default: None).
@@ -200,7 +200,7 @@ def gaussian_model(observation_space: Optional[Union[int, Tuple[int], gym.Space,
                           instantiate the model rather than the model instance (default: False).
     :type return_source: bool, optional
 
-    :return: Gaussian model instance
+    :return: Gaussian model instance or definition source
     :rtype: Model
     """
     # network
@@ -254,7 +254,7 @@ def multivariate_gaussian_model(observation_space: Optional[Union[int, Tuple[int
                                 output_shape: Shape = Shape.ACTIONS,
                                 output_activation: Optional[str] = "tanh",
                                 output_scale: float = 1.0,
-                                return_source: bool = False) -> Model:
+                                return_source: bool = False) -> Union[Model, str]:
     """Instantiate a multivariate Gaussian model
 
     :param observation_space: Observation/state space or shape (default: None).
@@ -293,7 +293,7 @@ def multivariate_gaussian_model(observation_space: Optional[Union[int, Tuple[int
                           instantiate the model rather than the model instance (default: False).
     :type return_source: bool, optional
 
-    :return: Multivariate Gaussian model instance
+    :return: Multivariate Gaussian model instance or definition source
     :rtype: Model
     """
     # network
@@ -343,7 +343,7 @@ def deterministic_model(observation_space: Optional[Union[int, Tuple[int], gym.S
                         output_shape: Shape = Shape.ACTIONS,
                         output_activation: Optional[str] = "tanh",
                         output_scale: float = 1.0,
-                        return_source: bool = False) -> Model:
+                        return_source: bool = False) -> Union[Model, str]:
     """Instantiate a deterministic model
 
     :param observation_space: Observation/state space or shape (default: None).
@@ -374,7 +374,7 @@ def deterministic_model(observation_space: Optional[Union[int, Tuple[int], gym.S
                           instantiate the model rather than the model instance (default: False).
     :type return_source: bool, optional
 
-    :return: Deterministic model instance
+    :return: Deterministic model instance or definition source
     :rtype: Model
     """
     # network
@@ -418,7 +418,7 @@ def categorical_model(observation_space: Optional[Union[int, Tuple[int], gym.Spa
                       hidden_activation: list = ["relu", "relu"],
                       output_shape: Shape = Shape.ACTIONS,
                       output_activation: Optional[str] = None,
-                      return_source: bool = False) -> Model:
+                      return_source: bool = False) -> Union[Model, str]:
     """Instantiate a categorical model
 
     :param observation_space: Observation/state space or shape (default: None).
@@ -449,7 +449,7 @@ def categorical_model(observation_space: Optional[Union[int, Tuple[int], gym.Spa
                           instantiate the model rather than the model instance (default: False).
     :type return_source: bool, optional
 
-    :return: Categorical model instance
+    :return: Categorical model instance or definition source
     :rtype: Model
     """
     # network
@@ -489,7 +489,7 @@ def shared_model(observation_space: Optional[Union[int, Tuple[int], gym.Space, g
                  roles: Sequence[str] = [],
                  parameters: Sequence[Mapping[str, Any]] = [],
                  single_forward_pass: bool = True,
-                 return_source: bool = False) -> Model:
+                 return_source: bool = False) -> Union[Model, str]:
     """Instantiate a shared model
 
     :param observation_space: Observation/state space or shape (default: None).
@@ -514,7 +514,7 @@ def shared_model(observation_space: Optional[Union[int, Tuple[int], gym.Space, g
                           instantiate the model rather than the model instance (default: False).
     :type return_source: bool, optional
 
-    :return: Shared model instance
+    :return: Shared model instance or definition source
     :rtype: Model
     """
     # network