fix(sac): discount entropy backup in target

serl_launcher/serl_launcher/agents/continuous/sac.py

@@ -18,6 +18,26 @@
 from serl_launcher.networks.mlp import MLP
 
 
+def _compute_target_q(
+    rewards: jax.Array,
+    masks: jax.Array,
+    discount: float,
+    target_next_min_q: jax.Array,
+    entropy_term: Optional[jax.Array] = None,
+) -> jax.Array:
+    """Build the bootstrapped SAC target, including optional entropy backup.
+
+    The entropy correction is part of the next-state value and therefore must
+    be discounted and masked together with the target Q estimate.  Keeping the
+    arithmetic in one helper also makes the terminal-transition behaviour
+    explicit and easy to test.
+    """
+    target_next_value = (
+        target_next_min_q if entropy_term is None else target_next_min_q - entropy_term
+    )
+    return rewards + discount * masks * target_next_value
+
+
 class SACAgent(flax.struct.PyTreeNode):
     """
     Online actor-critic supporting several different algorithms depending on configuration:
@@ -161,15 +181,19 @@ def critic_loss_fn(self, batch, params: Params, rng: PRNGKey):
         target_next_min_q = target_next_qs.min(axis=0)
         chex.assert_shape(target_next_min_q, (batch_size,))
 
-        target_q = (
-            batch["rewards"]
-            + self.config["discount"] * batch["masks"] * target_next_min_q
-        )
-        chex.assert_shape(target_q, (batch_size,))
-
+        entropy_term = None
         if self.config["backup_entropy"]:
             temperature = self.forward_temperature()
-            target_q = target_q - temperature * next_actions_log_probs
+            entropy_term = temperature * next_actions_log_probs
+
+        target_q = _compute_target_q(
+            batch["rewards"],
+            batch["masks"],
+            self.config["discount"],
+            target_next_min_q,
+            entropy_term,
+        )
+        chex.assert_shape(target_q, (batch_size,))
 
         predicted_qs = self.forward_critic(
             batch["observations"], batch["actions"], rng=rng, grad_params=params

serl_launcher/tests/test_sac.py

@@ -0,0 +1,39 @@
+"""Regression tests for continuous-control SAC targets."""
+
+import jax.numpy as jnp
+
+from serl_launcher.agents.continuous.sac import _compute_target_q
+
+
+def test_entropy_backup_is_discounted_and_masked() -> None:
+    rewards = jnp.array([1.0, 2.0])
+    masks = jnp.array([1.0, 0.0])
+    target_next_min_q = jnp.array([4.0, 4.0])
+    entropy_term = jnp.array([0.5, 0.5])
+
+    target_q = _compute_target_q(
+        rewards,
+        masks,
+        discount=0.9,
+        target_next_min_q=target_next_min_q,
+        entropy_term=entropy_term,
+    )
+
+    expected = jnp.array([1.0 + 0.9 * (4.0 - 0.5), 2.0])
+    assert jnp.allclose(target_q, expected)
+
+
+def test_target_q_without_entropy_backup_is_unchanged() -> None:
+    rewards = jnp.array([1.0, 2.0])
+    masks = jnp.array([1.0, 0.0])
+    target_next_min_q = jnp.array([4.0, 4.0])
+
+    target_q = _compute_target_q(
+        rewards,
+        masks,
+        discount=0.9,
+        target_next_min_q=target_next_min_q,
+    )
+
+    expected = jnp.array([1.0 + 0.9 * 4.0, 2.0])
+    assert jnp.allclose(target_q, expected)