Update model instantiators implementations

Author: Toni-SM

skrl/utils/model_instantiators/torch/__init__.py

@@ -1,20 +1,5 @@
-from enum import Enum
-
 from skrl.utils.model_instantiators.torch.categorical import categorical_model
 from skrl.utils.model_instantiators.torch.deterministic import deterministic_model
 from skrl.utils.model_instantiators.torch.gaussian import gaussian_model
 from skrl.utils.model_instantiators.torch.multivariate_gaussian import multivariate_gaussian_model
 from skrl.utils.model_instantiators.torch.shared import shared_model
-
-
-# keep for compatibility with versions prior to 1.3.0
-class Shape(Enum):
-    """
-    Enum to select the shape of the model's inputs and outputs
-    """
-
-    ONE = 1
-    STATES = 0
-    OBSERVATIONS = 0
-    ACTIONS = -1
-    STATES_ACTIONS = -2

skrl/utils/model_instantiators/torch/categorical.py

@@ -8,52 +8,37 @@
 
 from skrl.models.torch import CategoricalMixin  # noqa
 from skrl.models.torch import Model
-from skrl.utils.model_instantiators.torch.common import convert_deprecated_parameters, generate_containers
+from skrl.utils.model_instantiators.torch.common import generate_containers
 from skrl.utils.spaces.torch import unflatten_tensorized_space  # noqa
 
 
 def categorical_model(
-    observation_space: Optional[Union[int, Tuple[int], gymnasium.Space]] = None,
-    action_space: Optional[Union[int, Tuple[int], gymnasium.Space]] = None,
+    *,
+    observation_space: Optional[gymnasium.Space] = None,
+    state_space: Optional[gymnasium.Space] = None,
+    action_space: Optional[gymnasium.Space] = None,
     device: Optional[Union[str, torch.device]] = None,
     unnormalized_log_prob: bool = True,
     network: Sequence[Mapping[str, Any]] = [],
     output: Union[str, Sequence[str]] = "",
     return_source: bool = False,
-    *args,
-    **kwargs,
 ) -> Union[Model, str]:
-    """Instantiate a categorical model
+    """Instantiate a :class:`~skrl.models.torch.categorical.CategoricalMixin`-based model.
 
-    :param observation_space: Observation/state space or shape (default: None).
-                              If it is not None, the num_observations property will contain the size of that space
-    :type observation_space: int, tuple or list of integers, gymnasium.Space or None, optional
-    :param action_space: Action space or shape (default: None).
-                         If it is not None, the num_actions property will contain the size of that space
-    :type action_space: int, tuple or list of integers, gymnasium.Space or None, optional
-    :param device: Device on which a tensor/array is or will be allocated (default: ``None``).
-                   If None, the device will be either ``"cuda"`` if available or ``"cpu"``
-    :type device: str or torch.device, optional
-    :param unnormalized_log_prob: Flag to indicate how to be interpreted the model's output (default: True).
-                                  If True, the model's output is interpreted as unnormalized log probabilities
-                                  (it can be any real number), otherwise as normalized probabilities
-                                  (the output must be non-negative, finite and have a non-zero sum)
-    :type unnormalized_log_prob: bool, optional
-    :param network: Network definition (default: [])
-    :type network: list of dict, optional
-    :param output: Output expression (default: "")
-    :type output: list or str, optional
+    :param observation_space: Observation space. The ``num_observations`` property will contain the size of the space.
+    :param state_space: State space. The ``num_states`` property will contain the size of the space.
+    :param action_space: Action space. The ``num_actions`` property will contain the size of the space.
+    :param device: Data allocation and computation device. If not specified, the default device will be used.
+    :param unnormalized_log_prob: Flag to indicate how to the model's output will be interpreted.
+        If True, the model's output is interpreted as unnormalized log probabilities (it can be any real number),
+        otherwise as normalized probabilities (the output must be non-negative, finite and have a non-zero sum).
+    :param network: Network definition.
+    :param output: Output expression.
     :param return_source: Whether to return the source string containing the model class used to
-                          instantiate the model rather than the model instance (default: False).
-    :type return_source: bool, optional
+        instantiate the model rather than the model instance.
 
-    :return: Categorical model instance or definition source
-    :rtype: Model
+    :return: Categorical model instance or definition source (if ``return_source`` is True).
     """
-    # compatibility with versions prior to 1.3.0
-    if not network and kwargs:
-        network, output = convert_deprecated_parameters(kwargs)
-
     # parse model definition
     containers, output = generate_containers(network, output, embed_output=True, indent=1)
 
@@ -77,14 +62,21 @@ def categorical_model(
     forward = textwrap.indent("\n".join(forward), prefix=" " * 8)[8:]
 
     template = f"""class CategoricalModel(CategoricalMixin, Model):
-    def __init__(self, observation_space, action_space, device, unnormalized_log_prob):
-        Model.__init__(self, observation_space, action_space, device)
-        CategoricalMixin.__init__(self, unnormalized_log_prob)
+    def __init__(self, observation_space, state_space, action_space, device=None, unnormalized_log_prob=True, role=""):
+        Model.__init__(
+            self,
+            observation_space=observation_space,
+            state_space=state_space,
+            action_space=action_space,
+            device=device,
+        )
+        CategoricalMixin.__init__(self, unnormalized_log_prob=unnormalized_log_prob, role=role)
 
         {networks}
 
     def compute(self, inputs, role=""):
-        states = unflatten_tensorized_space(self.observation_space, inputs.get("states"))
+        observations = unflatten_tensorized_space(self.observation_space, inputs.get("observations"))
+        states = unflatten_tensorized_space(self.state_space, inputs.get("states"))
         taken_actions = unflatten_tensorized_space(self.action_space, inputs.get("taken_actions"))
         {forward}
         return output, {{}}
@@ -98,6 +90,7 @@ def compute(self, inputs, role=""):
     exec(template, globals(), _locals)
     return _locals["CategoricalModel"](
         observation_space=observation_space,
+        state_space=state_space,
         action_space=action_space,
         device=device,
         unnormalized_log_prob=unnormalized_log_prob,

skrl/utils/model_instantiators/torch/common.py

@@ -2,11 +2,9 @@
 
 import ast
 
-from skrl import logger
-
 
 def _get_activation_function(activation: Union[str, None], as_module: bool = True) -> Union[str, None]:
-    """Get the activation function
+    """Get the activation function.
 
     Supported activation functions:
 
@@ -20,10 +18,10 @@ def _get_activation_function(activation: Union[str, None], as_module: bool = Tru
     - "softsign"
     - "tanh"
 
-    :param activation: Activation function name
-    :param as_module: Whether to return a PyTorch module instance rather than a functional method
+    :param activation: Activation function name.
+    :param as_module: Whether to return a PyTorch module instance rather than a functional method.
 
-    :return: Activation function or None if the activation is not supported
+    :return: Activation function or ``None`` if the activation is not supported.
     """
     activations = {
         "elu": "nn.ELU()" if as_module else "functional.elu",
@@ -40,11 +38,11 @@ def _get_activation_function(activation: Union[str, None], as_module: bool = Tru
 
 
 def _parse_input(source: str) -> str:
-    """Parse a network input expression by replacing substitutions and applying operations
+    """Parse a network input expression by replacing substitutions and applying operations.
 
-    :param source: Input expression
+    :param source: Input expression.
 
-    :return: Parsed network input
+    :return: Parsed network input.
     """
 
     class NodeTransformer(ast.NodeTransformer):
@@ -64,24 +62,20 @@ def visit_Call(self, node: ast.Call):
     NodeTransformer().visit(tree)
     source = ast.unparse(tree)
     # enum substitutions
-    source = source.replace("Shape.STATES_ACTIONS", "STATES_ACTIONS").replace(
-        "STATES_ACTIONS", "torch.cat([states, taken_actions], dim=1)"
-    )
-    source = source.replace("Shape.OBSERVATIONS_ACTIONS", "OBSERVATIONS_ACTIONS").replace(
-        "OBSERVATIONS_ACTIONS", "torch.cat([states, taken_actions], dim=1)"
-    )
-    source = source.replace("Shape.STATES", "STATES").replace("STATES", "states")
-    source = source.replace("Shape.OBSERVATIONS", "OBSERVATIONS").replace("OBSERVATIONS", "states")
-    source = source.replace("Shape.ACTIONS", "ACTIONS").replace("ACTIONS", "taken_actions")
+    source = source.replace("OBSERVATIONS_ACTIONS", "torch.cat([observations, taken_actions], dim=1)")
+    source = source.replace("STATES_ACTIONS", "torch.cat([states, taken_actions], dim=1)")
+    source = source.replace("OBSERVATIONS", "observations")
+    source = source.replace("STATES", "states")
+    source = source.replace("ACTIONS", "taken_actions")
     return source
 
 
 def _parse_output(source: Union[str, Sequence[str]]) -> Tuple[Union[str, Sequence[str]], Sequence[str], int]:
-    """Parse the network output expression by replacing substitutions and applying operations
+    """Parse the network output expression by replacing substitutions and applying operations.
 
-    :param source: Output expression
+    :param source: Output expression.
 
-    :return: Tuple with the parsed network output, generated modules and output size/shape
+    :return: Tuple with the parsed network output, generated modules and output size/shape.
     """
 
     class NodeTransformer(ast.NodeTransformer):
@@ -101,7 +95,6 @@ def visit_Call(self, node: ast.Call):
     modules = []
     if type(source) is str:
         # enum substitutions
-        source = source.replace("Shape.ACTIONS", "ACTIONS").replace("Shape.ONE", "ONE")
         token = "ACTIONS" if "ACTIONS" in source else None
         token = "ONE" if "ONE" in source else token
         if token:
@@ -120,13 +113,13 @@ def visit_Call(self, node: ast.Call):
 
 
 def _generate_modules(layers: Sequence[str], activations: Union[Sequence[str], str]) -> Sequence[str]:
-    """Generate network modules
+    """Generate network modules.
 
     :param layers: Layer definitions
     :param activations: Activation function definitions applied after each layer (except ``flatten`` layers).
-                        If a single activation function is specified (str or lis), it will be applied after each layer
+        If a single activation function is specified (str or list), it will be applied after each layer.
 
-    :return: A list of generated modules
+    :return: A list of generated modules.
     """
     # expand activations
     if type(activations) is str:
@@ -224,21 +217,24 @@ def _generate_modules(layers: Sequence[str], activations: Union[Sequence[str], s
 
 
 def get_num_units(token: Union[str, Any]) -> Union[str, Any]:
-    """Get the number of units/features a token represent
+    """Get the number of units/features a token represents.
 
-    :param token: Token
+    :param token: Token.
 
-    :return: Number of units/features a token represent. If the token is unknown, its value will be returned as it
+    :return: Number of units/features a token represents. If the token is unknown, its value will be returned as it.
     """
     num_units = {
         "ONE": "1",
-        "STATES": "self.num_observations",
+        "NUM_OBSERVATIONS": "self.num_observations",
+        "NUM_STATES": "self.num_states",
+        "NUM_ACTIONS": "self.num_actions",
         "OBSERVATIONS": "self.num_observations",
+        "STATES": "self.num_states",
         "ACTIONS": "self.num_actions",
-        "STATES_ACTIONS": "self.num_observations + self.num_actions",
         "OBSERVATIONS_ACTIONS": "self.num_observations + self.num_actions",
+        "STATES_ACTIONS": "self.num_states + self.num_actions",
     }
-    token_as_str = str(token).replace("Shape.", "")
+    token_as_str = str(token)
     if token_as_str in num_units:
         return num_units[token_as_str]
     return token
@@ -247,16 +243,16 @@ def get_num_units(token: Union[str, Any]) -> Union[str, Any]:
 def generate_containers(
     network: Sequence[Mapping[str, Any]], output: Union[str, Sequence[str]], embed_output: bool = True, indent: int = -1
 ) -> Tuple[Sequence[Mapping[str, Any]], Mapping[str, Any]]:
-    """Generate network containers
+    """Generate network containers.
 
-    :param network: Network definition
-    :param output: Network's output expression
+    :param network: Network definition.
+    :param output: Network's output expression.
     :param embed_output: Whether to embed the output modules (if any) in the container definition.
-                         If True, the output modules will be append to the last container module
+        If True, the output modules will be append to the last container module.
     :param indent: Indentation level used to generate the Sequential definition.
-                   If negative, no indentation will be applied
+        If negative, no indentation will be applied.
 
-    :return: Network containers and output
+    :return: Network containers and output.
     """
     # parse output
     output, output_modules, output_size = _parse_output(output)
@@ -290,37 +286,3 @@ def generate_containers(
             output = output.replace("PLACEHOLDER", container["name"] if embed_output else "output")
     output = {"output": output, "modules": output_modules, "size": output_size}
     return containers, output
-
-
-def convert_deprecated_parameters(parameters: Mapping[str, Any]) -> Tuple[Mapping[str, Any], str]:
-    """Function to convert deprecated parameters to network-output format
-
-    :param parameters: Deprecated parameters and their values.
-
-    :return: Network and output definitions
-    """
-    logger.warning(
-        f'The following parameters ({", ".join(list(parameters.keys()))}) are deprecated. '
-        "See https://skrl.readthedocs.io/en/latest/api/utils/model_instantiators.html"
-    )
-    # network definition
-    activations = parameters.get("hidden_activation", [])
-    if type(activations) in [list, tuple] and len(set(activations)) == 1:
-        activations = activations[0]
-    network = [
-        {
-            "name": "net",
-            "input": str(parameters.get("input_shape", "STATES")),
-            "layers": parameters.get("hiddens", []),
-            "activations": activations,
-        }
-    ]
-    # output
-    output_scale = parameters.get("output_scale", 1.0)
-    scale_operation = f"{output_scale} * " if output_scale != 1.0 else ""
-    if parameters.get("output_activation", None):
-        output = f'{scale_operation}{parameters["output_activation"]}({str(parameters.get("output_shape", "ACTIONS"))})'
-    else:
-        output = f'{scale_operation}{str(parameters.get("output_shape", "ACTIONS"))}'
-
-    return network, output