rlmodule 0.2.0__tar.gz → 0.2.1__tar.gz

{rlmodule-0.2.0/rlmodule.egg-info → rlmodule-0.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rlmodule
-Version: 0.2.0
+Version: 0.2.1
 Summary: Flexible reinforcement learning models instantiators library
 Author: Lopatovsky
 Author-email: lopatovsky@gmail.com

{rlmodule-0.2.0 → rlmodule-0.2.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "rlmodule"
-version = "0.2.0"
+version = "0.2.1"
 description = "Flexible reinforcement learning models instantiators library"
 readme = "README.md"
 requires-python = ">=3.10"

{rlmodule-0.2.0 → rlmodule-0.2.1}/rlmodule/source/output_layer.py

@@ -120,6 +120,11 @@ class GaussianLayer(OutputLayer):
             >>> print(actions.shape, log_prob.shape, outputs["mean_actions"].shape)
             torch.Size([4096, 8]) torch.Size([4096, 1]) torch.Size([4096, 8])
         """
+        # ``RLModel.forward`` passes ``inputs["taken_pre_squash"]`` here (``_squashed_forward``).
+        taken_pre_squash = outputs_dict.pop("taken_pre_squash", None)
+        if taken_pre_squash is not None and not self._squash:
+            raise ValueError("taken_pre_squash is the pre-tanh sample of a squashed head, but squash=False")
+
         # Optional per-head pre-projection (see OutputLayerCfg.head).
         head_out = self._head(input)
 
@@ -134,7 +139,7 @@ class GaussianLayer(OutputLayer):
         self._distribution = Normal(mean_actions, self._cached_log_std.exp())
 
         if self._squash:
-            return self._squashed_forward(mean_actions, taken_actions, outputs_dict)
+            return self._squashed_forward(mean_actions, taken_actions, outputs_dict, taken_pre_squash)
 
         # sample using the reparametrization trick
         actions = self._distribution.rsample()
@@ -180,7 +185,7 @@ class GaussianLayer(OutputLayer):
     # Keeps atanh finite for actions stored at the bound (float32 tanh saturates to 1).
     _SQUASH_EPS = 1e-6
 
-    def _squashed_forward(self, mean_actions, taken_actions, outputs_dict):
+    def _squashed_forward(self, mean_actions, taken_actions, outputs_dict, taken_pre_squash=None):
         """Tanh-squashed Gaussian: ``u ~ N(mean, std)``, ``a = output_scale * tanh(u)``.
 
         ``self._distribution`` stays the pre-squash Normal over ``u``. ``mean_actions``
@@ -190,11 +195,20 @@ class GaussianLayer(OutputLayer):
         ``log_prob`` applies the change of variables per action dimension before the
         reduction: ``log p(a) = log N(u) - log(output_scale) - log(1 - tanh(u)^2)``,
         with ``log(1 - tanh(u)^2) = 2 * (log 2 - u - softplus(-2u))`` for stability.
+
+        A sample's ``u`` is returned as ``outputs_dict["pre_squash"]``. Re-evaluating a
+        taken action recovers ``u = atanh(a / output_scale)``, which float32 loses once
+        ``tanh(u)`` rounds to 1 (``|u|`` above about 8): at a std of 0.08 a ``u`` of 12
+        comes back 4.9 off, thousands of nats in ``log_prob``. ``taken_pre_squash``
+        (the stored ``u``) skips the inversion; PPO stores it (``PPOSquash``).
         """
         scale = self._output_scale
         self._distribution = Normal(mean_actions / scale, self._cached_log_std.exp())
 
-        if taken_actions is None:
+        if taken_pre_squash is not None:
+            pre_squash = taken_pre_squash
+            actions = scale * torch.tanh(pre_squash)
+        elif taken_actions is None:
             pre_squash = self._distribution.rsample()
             actions = scale * torch.tanh(pre_squash)
         else:
@@ -211,6 +225,7 @@ class GaussianLayer(OutputLayer):
 
         outputs_dict["mean_actions"] = scale * torch.tanh(mean_actions / scale)
         outputs_dict["log_prob"] = log_prob
+        outputs_dict["pre_squash"] = pre_squash
         return actions, outputs_dict
 
     def get_entropy(self, role: str = "") -> torch.Tensor:

{rlmodule-0.2.0 → rlmodule-0.2.1}/rlmodule/source/rlmodel.py

@@ -45,6 +45,21 @@ def _select_trunk_input(inputs: Mapping[str, Union[torch.Tensor, Any]], role: st
     return trunk_input
 
 
+def _episode_ends(inputs: Mapping[str, Union[torch.Tensor, Any]]) -> Optional[torch.Tensor]:
+    """Where a sampled sequence's episodes end, for the recurrent trunk's resets.
+
+    skrl's RNN agents zero the hidden state after a ``terminated`` *or* a
+    ``truncated`` step while collecting, and pass both with every sampled
+    sequence; resetting on ``terminated`` alone would carry a timed-out
+    episode's hidden state into the next episode of the same sequence.
+    """
+    terminated = inputs.get("terminated", None)
+    truncated = inputs.get("truncated", None)
+    if terminated is None or truncated is None:
+        return terminated if terminated is not None else truncated
+    return terminated | truncated
+
+
 class RLModel(Model):
     def __init__(
         self,
@@ -75,11 +90,13 @@ class RLModel(Model):
         if self._input_actions:
             trunk_input = torch.cat((trunk_input, inputs["taken_actions"]), dim=1)
         if self._rnn:
-            output, output_dict = self._net(trunk_input, inputs.get("terminated", None), inputs["rnn"])
+            output, output_dict = self._net(trunk_input, _episode_ends(inputs), inputs["rnn"])
         else:
             output = self._net(trunk_input)
             output_dict = {}
 
+        if "taken_pre_squash" in inputs:
+            output_dict = {**output_dict, "taken_pre_squash": inputs["taken_pre_squash"]}
         return self._output_layer(output, inputs.get("taken_actions", None), output_dict)
 
 
@@ -156,7 +173,7 @@ class SharedRLModel(Model):
             self._reset_cache()
 
             if self._rnn:
-                output, output_dict = self._net(states, inputs.get("terminated", None), inputs["rnn"])
+                output, output_dict = self._net(states, _episode_ends(inputs), inputs["rnn"])
             else:
                 output = self._net(states)
                 output_dict = {}
@@ -167,6 +184,8 @@ class SharedRLModel(Model):
         output, output_dict = self._cached_output
 
         if role == "policy":
+            if "taken_pre_squash" in inputs:
+                output_dict = {**output_dict, "taken_pre_squash": inputs["taken_pre_squash"]}
             return self._policy_output_layer(output, inputs.get("taken_actions", None), output_dict)
         elif role == "value":
             return self._value_output_layer(output, inputs.get("taken_actions", None), output_dict)

{rlmodule-0.2.0 → rlmodule-0.2.1/rlmodule.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rlmodule
-Version: 0.2.0
+Version: 0.2.1
 Summary: Flexible reinforcement learning models instantiators library
 Author: Lopatovsky
 Author-email: lopatovsky@gmail.com

{rlmodule-0.2.0 → rlmodule-0.2.1}/tests/test_rnn_memory_and_correctness.py

@@ -442,3 +442,34 @@ def test_num_envs_happy_path_returns_expected_shape():
     spec = model.get_specification()
     assert spec["rnn"]["sequence_length"] == L
     assert spec["rnn"]["sizes"] == [(NUM_LAYERS, N, HOUT), (NUM_LAYERS, N, HOUT)]
+
+
+def test_a_truncated_step_resets_the_hidden_state_like_a_terminated_one():
+    """skrl's RNN agents zero the hidden state after a truncated step too, so
+    the model's training forward must reset there: a time-out in the middle
+    of a sampled sequence starts a new episode.
+    """
+    torch.manual_seed(0)
+    model = _build_shared_lstm_model()
+    model.train(True)
+    states = torch.randn(N * L, HIN)
+    h, c = (torch.randn(NUM_LAYERS, N * L, HOUT) for _ in range(2))
+    none = torch.zeros(N * L, 1, dtype=torch.bool)
+    at_1 = none.view(N, L).clone()
+    at_1[:, 1] = True
+    at_1 = at_1.view(N * L, 1)
+
+    def mean_actions(terminated, truncated):
+        model._reset_cache()
+        inputs = {
+            "states": states,
+            "rnn": [h, c],
+            "terminated": terminated,
+            "truncated": truncated,
+        }
+        with torch.no_grad():
+            return model(inputs, role="policy")[1]["mean_actions"]
+
+    truncated = mean_actions(none, at_1)
+    assert torch.allclose(truncated, mean_actions(at_1, none))
+    assert not torch.allclose(truncated, mean_actions(none, none))

{rlmodule-0.2.0 → rlmodule-0.2.1}/tests/test_squashed_gaussian.py

@@ -68,6 +68,40 @@ def test_squashed_log_prob_is_finite_at_the_bound() -> None:
     assert torch.isfinite(out["log_prob"]).all()
 
 
+def _saturated(mean: float, std: float) -> GaussianLayer:
+    """A squashed layer whose pre-squash mean is *mean* everywhere, at *std*."""
+    layer = _squashed(log_std=ParameterLogStdCfg(initial_log_std=math.log(std)))
+    with torch.no_grad():
+        layer._net[0].weight.zero_()
+        layer._net[0].bias.fill_(mean)
+    return layer
+
+
+def test_the_sample_returns_its_pre_squash_value() -> None:
+    layer = _squashed(output_scale=2.0)
+    actions, out = layer(torch.randn(64, N_IN), None, {})
+    torch.testing.assert_close(2.0 * torch.tanh(out["pre_squash"]), actions)
+
+
+def test_taken_pre_squash_gives_the_sampled_log_prob_where_tanh_saturates() -> None:
+    # Run 89b's failure: at |u| ~ 12 float32 tanh is 1, so atanh(a) cannot recover u,
+    # and at a std of 0.08 the re-evaluated log_prob is thousands of nats off.
+    layer = _saturated(mean=12.0, std=0.08)
+    x = torch.randn(256, N_IN)
+    actions, out = layer(x, None, {})
+    assert out["pre_squash"].abs().min() > 11.0
+    _, exact = layer(x, actions, {"taken_pre_squash": out["pre_squash"]})
+    torch.testing.assert_close(exact["log_prob"], out["log_prob"])
+    _, inverted = layer(x, actions, {})
+    assert (inverted["log_prob"] - out["log_prob"]).abs().max() > 1000.0
+
+
+def test_taken_pre_squash_on_an_unsquashed_layer_is_refused() -> None:
+    layer = _layer()
+    with pytest.raises(ValueError, match="squash=False"):
+        layer(torch.randn(3, N_IN), torch.zeros(3, N_ACT), {"taken_pre_squash": torch.zeros(3, N_ACT)})
+
+
 def test_squashed_mean_action_is_tanh_of_the_mean() -> None:
     layer = _squashed(output_scale=2.0)
     x = torch.randn(8, N_IN)