Merge pull request #234 from KraHsu/hotfix/skrl-offpolicy-stability

fix(rl/skrl): off-policy training stability (obs norm, critic LayerNorm, SAC clip)

Author: KraHsu

src/genelab/rl/backends/skrl/__init__.py

@@ -120,6 +120,18 @@ def _build_agent(agent_cfg: SkrlAgentCfg, wrapper: Any, device: Any, log_dir: Pa
     models = build_models(agent_cfg.algorithm, wrapper, agent_cfg.models, device)
     agent_cls, default_cfg = _agent_class(agent_cfg.algorithm)
     cfg = _build_skrl_cfg(agent_cfg, default_cfg, log_dir)
+    if agent_cfg.models.obs_normalization and "state_preprocessor" in cfg:
+        # Empirical obs normalization (analogue of rsl_rl's obs_normalization). The
+        # scaler is owned by the agent and applied to states before policy/critic
+        # forward passes (see skrl ``_update``); ``size`` is taken from the observation
+        # space so it matches whichever ``obs_group`` feeds the agent.
+        from skrl.resources.preprocessors.torch import RunningStandardScaler
+
+        cfg["state_preprocessor"] = RunningStandardScaler
+        cfg["state_preprocessor_kwargs"] = {
+            "size": wrapper.observation_space,
+            "device": device,
+        }
     return agent_cls(
         models=models,
         memory=memory,

src/genelab/rl/backends/skrl/config.py

@@ -32,6 +32,11 @@ class SkrlModelCfg:
     min_log_std: float = -20.0
     max_log_std: float = 2.0
     clip_actions: bool = False
+    # Empirical observation normalization (skrl ``RunningStandardScaler`` state
+    # preprocessor), the analogue of ``RslRlModelCfg.obs_normalization``. Off by
+    # default; load-bearing for off-policy stability (SAC/TD3 regress a bootstrapped
+    # Q and diverge without standardized critic inputs).
+    obs_normalization: bool = False
 
 
 @dataclass

src/genelab/rl/backends/skrl/models.py

@@ -128,8 +128,28 @@ def compute(self, inputs: Mapping[str, Any], role: str = "") -> tuple[Any, dict[
         return self.net(inputs["states"]), {}
 
 
+def _mlp_layernorm(
+    in_dim: int, hidden_dims: tuple[int, ...], out_dim: int, activation: str
+) -> nn.Sequential:
+    """MLP with LayerNorm after every hidden linear: ``Linear -> LayerNorm -> act``.
+
+    LayerNorm in the value network bounds the critic's output extrapolation, the key
+    ingredient that prevents the bootstrapped-Q overestimation divergence off-policy
+    algorithms suffer on long-horizon / contact-rich tasks (cf. CrossQ, SAC-LN). The
+    final linear has no norm/activation so the Q range stays unconstrained.
+    """
+    act = _activation(activation)
+    layers: list[nn.Module] = []
+    prev = in_dim
+    for width in hidden_dims:
+        layers += [nn.Linear(prev, width), nn.LayerNorm(width), act()]
+        prev = width
+    layers.append(nn.Linear(prev, out_dim))
+    return nn.Sequential(*layers)
+
+
 class QCritic(DeterministicMixin, Model):
-    """Action-value Q(s, a) for SAC / TD3 / DDPG."""
+    """Action-value Q(s, a) for SAC / TD3 / DDPG (LayerNorm-stabilized)."""
 
     def __init__(
         self,
@@ -142,7 +162,7 @@ def __init__(
         DeterministicMixin.__init__(self, clip_actions=False)
         num_obs = cast(int, self.num_observations)
         num_act = cast(int, self.num_actions)
-        self.net = _mlp(num_obs + num_act, model_cfg.hidden_dims, 1, model_cfg.activation)
+        self.net = _mlp_layernorm(num_obs + num_act, model_cfg.hidden_dims, 1, model_cfg.activation)
         self.to(device)
 
     def compute(self, inputs: Mapping[str, Any], role: str = "") -> tuple[Any, dict[str, Any]]:
@@ -197,8 +217,14 @@ def q() -> Model:
             "target_critic_2": q(),
         }
     if algorithm == "SAC":
+        # clip_actions MUST be False here: skrl's GaussianMixin clamps the sampled
+        # action *before* computing its log-prob, so with clipping on, a policy mean
+        # that drifts outside the action bounds yields log_prob(±1 | N(μ,σ)) → -inf,
+        # and SAC's entropy term (-α·log_prob) then blows up the Q-target. Left
+        # unclipped the log-prob is in-distribution (bounded); the env clamps the
+        # action downstream. (TD3/DDPG use a deterministic tanh policy, no such issue.)
         return {
-            "policy": gaussian(clip_actions=True),
+            "policy": gaussian(clip_actions=False),
             "critic_1": q(),
             "critic_2": q(),
             "target_critic_1": q(),

tests/test_skrl_pipeline.py

@@ -12,9 +12,12 @@
 skrl = pytest.importorskip("skrl")  # noqa: F841
 
 import gymnasium  # noqa: E402
+import torch.nn as nn  # noqa: E402
 
 from genelab.rl import RslRlOnPolicyRunnerCfg, SkrlAgentCfg  # noqa: E402
 from genelab.rl.backends import select_backend  # noqa: E402
+from genelab.rl.backends.skrl import _build_agent  # noqa: E402  # pyright: ignore[reportPrivateUsage]
+from genelab.rl.backends.skrl.config import SkrlModelCfg  # noqa: E402
 from genelab.rl.backends.skrl.models import build_models  # noqa: E402
 from genelab.rl.vecenvs.skrl import GenelabSkrlWrapper  # noqa: E402
 
@@ -124,3 +127,54 @@ def test_build_models_per_algorithm(algorithm: str) -> None:
     wrapped = GenelabSkrlWrapper(_FakeEnv())
     models = build_models(algorithm, wrapped, SkrlAgentCfg().models, device="cpu")
     assert set(models) == _EXPECTED_MODEL_KEYS[algorithm]
+
+
+# --- off-policy stability fixes ------------------------------------------------
+# Three changes make off-policy (SAC/TD3/DDPG) actually trainable instead of
+# diverging: empirical obs normalization, a LayerNorm-stabilized Q-critic, and a
+# non-clipping SAC policy (so the entropy term can't blow up the Q-target).
+
+
+def test_obs_normalization_defaults_off() -> None:
+    assert SkrlModelCfg().obs_normalization is False
+
+
+def test_qcritic_is_layernorm_stabilized() -> None:
+    # The Q-critic (SAC/TD3/DDPG only) gets LayerNorm to bound value extrapolation.
+    wrapped = GenelabSkrlWrapper(_FakeEnv())
+    models = build_models("SAC", wrapped, SkrlAgentCfg().models, device="cpu")
+    q_has_ln = any(isinstance(m, nn.LayerNorm) for m in models["critic_1"].net.modules())
+    assert q_has_ln
+    # PPO uses the state-value critic, which must stay a plain MLP (no behaviour change).
+    ppo = build_models("PPO", wrapped, SkrlAgentCfg().models, device="cpu")
+    assert not any(isinstance(m, nn.LayerNorm) for m in ppo["value"].net.modules())
+
+
+def test_sac_policy_does_not_clip_actions() -> None:
+    # skrl clamps the action before taking its log-prob; clipping on would let a
+    # drifting policy mean send -alpha*log_prob (and the Q-target) to infinity.
+    wrapped = GenelabSkrlWrapper(_FakeEnv())
+    policy = build_models("SAC", wrapped, SkrlAgentCfg().models, device="cpu")["policy"]
+    assert policy._g_clip_actions is False  # skrl GaussianMixin flag
+
+
+def test_obs_normalization_wires_state_preprocessor() -> None:
+    from skrl.resources.preprocessors.torch import RunningStandardScaler
+
+    wrapped = GenelabSkrlWrapper(_FakeEnv(), state_group=None)
+    off = _build_agent(
+        SkrlAgentCfg(algorithm="SAC", state_group=None, learning_starts=0), wrapped, "cpu", None
+    )
+    assert not isinstance(off._state_preprocessor, RunningStandardScaler)
+    on = _build_agent(
+        SkrlAgentCfg(
+            algorithm="SAC",
+            state_group=None,
+            learning_starts=0,
+            models=SkrlModelCfg(obs_normalization=True),
+        ),
+        wrapped,
+        "cpu",
+        None,
+    )
+    assert isinstance(on._state_preprocessor, RunningStandardScaler)