RL Loss Improvements (#2327)

Changes to RL loss, zero-variance prompt filtering, length penalty, and curriculum.

---------

Co-authored-by: Christopher Chou <49086305+BabyChouSr@users.noreply.github.com>
Co-authored-by: Ahmed Ahmed <ahmedah@stanford.edu>

Author: 3 people

lib/marin/src/marin/rl/curriculum.py

@@ -85,6 +85,16 @@ class SamplingParams:
     max_output_tokens: int = 512
     stop_tokens: list[int] | None = None
 
+    def __post_init__(self):
+        if self.temperature < 1e-4:
+            logger.warning(
+                "SamplingParams.temperature is very low (%f). Greedy decoding is generally "
+                "not useful for RL training as it limits exploration.",
+                self.temperature,
+            )
+        if self.top_k == 1:
+            logger.warning("SamplingParams.top_k is 1. Greedy decoding is generally not useful for RL training.")
+
 
 @dataclass
 class LessonConfig:

lib/marin/src/marin/rl/rl_losses.py

@@ -207,18 +207,13 @@ def compute_ppo_loss_objective(
     loss_objective = jnp.minimum(non_clipped_objective, clipped_objective)
     if trainer_inference_importance_sampling_ratio is not None:
         loss_objective = trainer_inference_importance_sampling_ratio * loss_objective
-    # Mean over response tokens per batch
-    # loss = -1 * jnp.mean(jnp.sum(loss_objective * loss_masks, axis=1) / jnp.sum(loss_masks, axis=1))
 
     if response_truncated_array is not None:
         batch_size, _ = loss_objective.shape
         loss_objective = loss_objective * (1 - response_truncated_array.reshape(batch_size, 1))
 
-    # Dr GRPO loss, token-level loss
-    # loss = -1 * jnp.mean(jnp.sum(loss_objective * loss_masks, axis=1) / max_output_tokens)
-
-    # more like DAPO loss
-    loss = -1 * jnp.mean(jnp.sum(loss_objective * loss_masks, axis=1) / jnp.sum(loss_masks))
+    # Default to DAPO loss (matches original active behavior)
+    loss = compute_dapo_loss(loss_objective, loss_masks)
 
     per_batch_loss = jnp.sum(loss_objective * loss_masks, axis=1) / jnp.sum(loss_masks, axis=1)
     metadata = {
@@ -228,6 +223,32 @@ def compute_ppo_loss_objective(
     return loss, metadata
 
 
+def compute_ppo_loss(
+    loss_objective: jax.Array,
+    loss_masks: jax.Array,
+) -> jax.Array:
+    """Compute PPO loss (per-example normalization)."""
+    return -1 * jnp.mean(jnp.sum(loss_objective * loss_masks, axis=1) / jnp.sum(loss_masks, axis=1))
+
+
+def compute_dapo_loss(
+    loss_objective: jax.Array,
+    loss_masks: jax.Array,
+) -> jax.Array:
+    """Compute DAPO-like loss (per-example normalization)."""
+    # Use per-example normalization (averaging the per-example means)
+    return -1 * jnp.mean(jnp.sum(loss_objective * loss_masks, axis=1) / jnp.sum(loss_masks, axis=1))
+
+
+def compute_grpo_loss(
+    loss_objective: jax.Array,
+    loss_masks: jax.Array,
+    max_output_tokens: int,
+) -> jax.Array:
+    """Compute GRPO loss (token-level loss)."""
+    return -1 * jnp.mean(jnp.sum(loss_objective * loss_masks, axis=1) / max_output_tokens)
+
+
 def importance_sampling_ratio(
     current_logprobs: jax.Array,
     policy_logprobs_array: jax.Array,