Cast initial_log_std parameter to float in PyTorch (#297)

Toni-SM · web-flow · commit 4e622dcf691c · 2025-03-18T16:06:22.000-04:00
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -9,6 +9,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 - Allow `None` type spaces and samples/values in spaces utilities
 
 ### Fixed
+- Cast model instantiator's `initial_log_std` parameter to `float` in PyTorch
 - Fix common property overwriting (e.g. `clip_actions`) in shared models composed of different mixin types
 
 ## [1.4.1] - 2025-01-27
diff --git a/skrl/utils/model_instantiators/torch/gaussian.py b/skrl/utils/model_instantiators/torch/gaussian.py
@@ -102,7 +102,7 @@ def __init__(self, observation_space, action_space, device, clip_actions,
         GaussianMixin.__init__(self, clip_actions, clip_log_std, min_log_std, max_log_std, reduction)
 
         {networks}
-        self.log_std_parameter = nn.Parameter(torch.full(size=({output["size"]},), fill_value={initial_log_std}), requires_grad={not fixed_log_std})
+        self.log_std_parameter = nn.Parameter(torch.full(size=({output["size"]},), fill_value={float(initial_log_std)}), requires_grad={not fixed_log_std})
 
     def compute(self, inputs, role=""):
         states = unflatten_tensorized_space(self.observation_space, inputs.get("states"))
diff --git a/skrl/utils/model_instantiators/torch/multivariate_gaussian.py b/skrl/utils/model_instantiators/torch/multivariate_gaussian.py
@@ -97,7 +97,7 @@ def __init__(self, observation_space, action_space, device, clip_actions,
         MultivariateGaussianMixin.__init__(self, clip_actions, clip_log_std, min_log_std, max_log_std)
 
         {networks}
-        self.log_std_parameter = nn.Parameter(torch.full(size=({output["size"]},), fill_value={initial_log_std}), requires_grad={not fixed_log_std})
+        self.log_std_parameter = nn.Parameter(torch.full(size=({output["size"]},), fill_value={float(initial_log_std)}), requires_grad={not fixed_log_std})
 
     def compute(self, inputs, role=""):
         states = unflatten_tensorized_space(self.observation_space, inputs.get("states"))
diff --git a/tests/jax/test_jax_model_instantiators.py b/tests/jax/test_jax_model_instantiators.py
@@ -3,7 +3,12 @@
 import yaml
 from gymnasium import spaces
 
-from skrl.utils.model_instantiators.jax import categorical_model, deterministic_model, gaussian_model
+from skrl.utils.model_instantiators.jax import (
+    categorical_model,
+    deterministic_model,
+    gaussian_model,
+    multicategorical_model,
+)
 from skrl.utils.spaces.jax import flatten_tensorized_space, sample_space
 
 
@@ -87,10 +92,42 @@ def test_categorical_model(capsys, device):
         )
         model.init_state_dict("model")
 
-        output = model.act({"states": flatten_tensorized_space(sample_space(observation_space, 10, "jax", device))})
+        output = model.act(
+            {
+                "states": flatten_tensorized_space(
+                    sample_space(observation_space, batch_size=10, backend="native", device=device)
+                )
+            }
+        )
         assert output[0].shape == (10, 1)
 
 
+@pytest.mark.parametrize("device", [None, "cpu", "cuda:0"])
+def test_multicategorical_model(capsys, device):
+    # observation
+    action_space = spaces.MultiDiscrete([2, 3])
+    for observation_space_type in [spaces.Box, spaces.Tuple, spaces.Dict]:
+        observation_space = NETWORK_SPEC_OBSERVATION[observation_space_type][1]
+        model = multicategorical_model(
+            observation_space=observation_space,
+            action_space=action_space,
+            device=device,
+            unnormalized_log_prob=True,
+            network=yaml.safe_load(NETWORK_SPEC_OBSERVATION[observation_space_type][0])["network"],
+            output="ACTIONS",
+        )
+        model.init_state_dict("model")
+
+        output = model.act(
+            {
+                "states": flatten_tensorized_space(
+                    sample_space(observation_space, batch_size=10, backend="native", device=device)
+                )
+            }
+        )
+        assert output[0].shape == (10, 2)
+
+
 @pytest.mark.parametrize("device", [None, "cpu", "cuda:0"])
 def test_deterministic_model(capsys, device):
     # observation
@@ -107,7 +144,13 @@ def test_deterministic_model(capsys, device):
         )
         model.init_state_dict("model")
 
-        output = model.act({"states": flatten_tensorized_space(sample_space(observation_space, 10, "jax", device))})
+        output = model.act(
+            {
+                "states": flatten_tensorized_space(
+                    sample_space(observation_space, batch_size=10, backend="native", device=device)
+                )
+            }
+        )
         assert output[0].shape == (10, 2)
 
 
@@ -131,5 +174,11 @@ def test_gaussian_model(capsys, device):
         )
         model.init_state_dict("model")
 
-        output = model.act({"states": flatten_tensorized_space(sample_space(observation_space, 10, "jax", device))})
+        output = model.act(
+            {
+                "states": flatten_tensorized_space(
+                    sample_space(observation_space, batch_size=10, backend="native", device=device)
+                )
+            }
+        )
         assert output[0].shape == (10, 2)
diff --git a/tests/jax/test_jax_model_instantiators_definition.py b/tests/jax/test_jax_model_instantiators_definition.py
@@ -9,7 +9,13 @@
 import jax.numpy as jnp
 import numpy as np
 
-from skrl.utils.model_instantiators.jax import Shape, categorical_model, deterministic_model, gaussian_model
+from skrl.utils.model_instantiators.jax import (
+    Shape,
+    categorical_model,
+    deterministic_model,
+    gaussian_model,
+    multicategorical_model,
+)
 from skrl.utils.model_instantiators.jax.common import _generate_modules, _get_activation_function, _parse_input
 
 
@@ -255,3 +261,40 @@ def test_categorical_model(capsys):
     observations = jnp.ones((10, model.num_observations))
     output = model.act({"states": observations})
     assert output[0].shape == (10, 1)
+
+
+def test_multicategorical_model(capsys):
+    device = "cpu"
+    observation_space = gym.spaces.Box(np.array([-1] * 5), np.array([1] * 5))
+    action_space = gym.spaces.MultiDiscrete([2, 3])
+
+    content = r"""
+    unnormalized_log_prob: True
+    network:
+      - name: net
+        input: OBSERVATIONS
+        layers:
+          - linear: 32
+          - linear: [32]
+          - linear: {out_features: 32}
+        activations: elu
+    output: ACTIONS
+    """
+    content = yaml.safe_load(content)
+    # source
+    model = multicategorical_model(
+        observation_space=observation_space, action_space=action_space, device=device, return_source=True, **content
+    )
+    with capsys.disabled():
+        print(model)
+    # instance
+    model = multicategorical_model(
+        observation_space=observation_space, action_space=action_space, device=device, return_source=False, **content
+    )
+    model.init_state_dict("model")
+    with capsys.disabled():
+        print(model)
+
+    observations = jnp.ones((10, model.num_observations))
+    output = model.act({"states": observations})
+    assert output[0].shape == (10, 2)
diff --git a/tests/torch/test_torch_model_instantiators.py b/tests/torch/test_torch_model_instantiators.py
@@ -7,6 +7,7 @@
     categorical_model,
     deterministic_model,
     gaussian_model,
+    multicategorical_model,
     multivariate_gaussian_model,
     shared_model,
 )
@@ -91,12 +92,44 @@ def test_categorical_model(capsys, device):
             network=yaml.safe_load(NETWORK_SPEC_OBSERVATION[observation_space_type][0])["network"],
             output="ACTIONS",
         )
-        model.to(device=device)
+        model.to(device=model.device)
 
-        output = model.act({"states": flatten_tensorized_space(sample_space(observation_space, 10, "torch", device))})
+        output = model.act(
+            {
+                "states": flatten_tensorized_space(
+                    sample_space(observation_space, batch_size=10, backend="native", device=device)
+                )
+            }
+        )
         assert output[0].shape == (10, 1)
 
 
+@pytest.mark.parametrize("device", [None, "cpu", "cuda:0"])
+def test_multicategorical_model(capsys, device):
+    # observation
+    action_space = spaces.MultiDiscrete([2, 3])
+    for observation_space_type in [spaces.Box, spaces.Tuple, spaces.Dict]:
+        observation_space = NETWORK_SPEC_OBSERVATION[observation_space_type][1]
+        model = multicategorical_model(
+            observation_space=observation_space,
+            action_space=action_space,
+            device=device,
+            unnormalized_log_prob=True,
+            network=yaml.safe_load(NETWORK_SPEC_OBSERVATION[observation_space_type][0])["network"],
+            output="ACTIONS",
+        )
+        model.to(device=model.device)
+
+        output = model.act(
+            {
+                "states": flatten_tensorized_space(
+                    sample_space(observation_space, batch_size=10, backend="native", device=device)
+                )
+            }
+        )
+        assert output[0].shape == (10, 2)
+
+
 @pytest.mark.parametrize("device", [None, "cpu", "cuda:0"])
 def test_deterministic_model(capsys, device):
     # observation
@@ -111,9 +144,15 @@ def test_deterministic_model(capsys, device):
             network=yaml.safe_load(NETWORK_SPEC_OBSERVATION[observation_space_type][0])["network"],
             output="ACTIONS",
         )
-        model.to(device=device)
+        model.to(device=model.device)
 
-        output = model.act({"states": flatten_tensorized_space(sample_space(observation_space, 10, "torch", device))})
+        output = model.act(
+            {
+                "states": flatten_tensorized_space(
+                    sample_space(observation_space, batch_size=10, backend="native", device=device)
+                )
+            }
+        )
         assert output[0].shape == (10, 2)
 
 
@@ -135,9 +174,15 @@ def test_gaussian_model(capsys, device):
             network=yaml.safe_load(NETWORK_SPEC_OBSERVATION[observation_space_type][0])["network"],
             output="ACTIONS",
         )
-        model.to(device=device)
+        model.to(device=model.device)
 
-        output = model.act({"states": flatten_tensorized_space(sample_space(observation_space, 10, "torch", device))})
+        output = model.act(
+            {
+                "states": flatten_tensorized_space(
+                    sample_space(observation_space, batch_size=10, backend="native", device=device)
+                )
+            }
+        )
         assert output[0].shape == (10, 2)
 
 
@@ -159,9 +204,15 @@ def test_multivariate_gaussian_model(capsys, device):
             network=yaml.safe_load(NETWORK_SPEC_OBSERVATION[observation_space_type][0])["network"],
             output="ACTIONS",
         )
-        model.to(device=device)
+        model.to(device=model.device)
 
-        output = model.act({"states": flatten_tensorized_space(sample_space(observation_space, 10, "torch", device))})
+        output = model.act(
+            {
+                "states": flatten_tensorized_space(
+                    sample_space(observation_space, batch_size=10, backend="native", device=device)
+                )
+            }
+        )
         assert output[0].shape == (10, 2)
 
 
@@ -196,9 +247,13 @@ def test_shared_gaussian_deterministic_model(capsys, device, single_forward_pass
             ],
             single_forward_pass=single_forward_pass,
         )
-        model.to(device=device)
+        model.to(device=model.device)
 
-        inputs = {"states": flatten_tensorized_space(sample_space(observation_space, 10, "torch", device))}
+        inputs = {
+            "states": flatten_tensorized_space(
+                sample_space(observation_space, batch_size=10, backend="native", device=device)
+            )
+        }
         output = model.act(inputs, role="role_0")
         assert output[0].shape == (10, 2)
         output = model.act(inputs, role="role_1")
@@ -236,9 +291,13 @@ def test_shared_multivariate_gaussian_deterministic_model(capsys, device, single
             ],
             single_forward_pass=single_forward_pass,
         )
-        model.to(device=device)
+        model.to(device=model.device)
 
-        inputs = {"states": flatten_tensorized_space(sample_space(observation_space, 10, "torch", device))}
+        inputs = {
+            "states": flatten_tensorized_space(
+                sample_space(observation_space, batch_size=10, backend="native", device=device)
+            )
+        }
         output = model.act(inputs, role="role_0")
         assert output[0].shape == (10, 2)
         output = model.act(inputs, role="role_1")
@@ -249,7 +308,7 @@ def test_shared_multivariate_gaussian_deterministic_model(capsys, device, single
 @pytest.mark.parametrize("device", [None, "cpu", "cuda:0"])
 def test_shared_categorical_deterministic_model(capsys, device, single_forward_pass):
     # observation
-    action_space = spaces.Box(low=-1, high=1, shape=(2,))
+    action_space = spaces.Discrete(2)
     for observation_space_type in [spaces.Box, spaces.Tuple, spaces.Dict]:
         observation_space = NETWORK_SPEC_OBSERVATION[observation_space_type][1]
         model = shared_model(
@@ -272,10 +331,54 @@ def test_shared_categorical_deterministic_model(capsys, device, single_forward_p
             ],
             single_forward_pass=single_forward_pass,
         )
-        model.to(device=device)
+        model.to(device=model.device)
 
-        inputs = {"states": flatten_tensorized_space(sample_space(observation_space, 10, "torch", device))}
+        inputs = {
+            "states": flatten_tensorized_space(
+                sample_space(observation_space, batch_size=10, backend="native", device=device)
+            )
+        }
         output = model.act(inputs, role="role_0")
         assert output[0].shape == (10, 1)
         output = model.act(inputs, role="role_1")
         assert output[0].shape == (10, 1)
+
+
+@pytest.mark.parametrize("single_forward_pass", [True, False])
+@pytest.mark.parametrize("device", [None, "cpu", "cuda:0"])
+def test_shared_multicategorical_deterministic_model(capsys, device, single_forward_pass):
+    # observation
+    action_space = spaces.MultiDiscrete([2, 3])
+    for observation_space_type in [spaces.Box, spaces.Tuple, spaces.Dict]:
+        observation_space = NETWORK_SPEC_OBSERVATION[observation_space_type][1]
+        model = shared_model(
+            observation_space=observation_space,
+            action_space=action_space,
+            device=device,
+            structure=["MultiCategoricalMixin", "DeterministicMixin"],
+            roles=["role_0", "role_1"],
+            parameters=[
+                {
+                    "unnormalized_log_prob": True,
+                    "network": yaml.safe_load(NETWORK_SPEC_OBSERVATION[observation_space_type][0])["network"],
+                    "output": "ACTIONS",
+                },
+                {
+                    "clip_actions": False,
+                    "network": yaml.safe_load(NETWORK_SPEC_OBSERVATION[observation_space_type][0])["network"],
+                    "output": "ONE",
+                },
+            ],
+            single_forward_pass=single_forward_pass,
+        )
+        model.to(device=model.device)
+
+        inputs = {
+            "states": flatten_tensorized_space(
+                sample_space(observation_space, batch_size=10, backend="native", device=device)
+            )
+        }
+        output = model.act(inputs, role="role_0")
+        assert output[0].shape == (10, 2)
+        output = model.act(inputs, role="role_1")
+        assert output[0].shape == (10, 1)
diff --git a/tests/torch/test_torch_model_instantiators_definition.py b/tests/torch/test_torch_model_instantiators_definition.py
