decouple force_on_policy_ratio from step0 reforward.

tests/rl/agentic/agentic_grpo_learner_test.py

@@ -109,10 +109,13 @@ def _mock_generate(
     else:
       prompt_tokens.append(tokenizer.encode(" ".join(m["content"] for m in p)))
   max_p_len = max(len(pt) for pt in prompt_tokens)
-  padded_prompts = np.array([
-      np.pad(pt, (max(0, max_p_len - len(pt)), 0), constant_values=0)
-      for pt in prompt_tokens
-  ], dtype=np.int32)
+  padded_prompts = np.array(
+      [
+          np.pad(pt, (max(0, max_p_len - len(pt)), 0), constant_values=0)
+          for pt in prompt_tokens
+      ],
+      dtype=np.int32,
+  )
   return base_rollout.RolloutOutput(
       text=text,
       tokens=tokens,
@@ -585,9 +588,7 @@ def test_compute_logps_chunk_size(self):
         chat_parser=MockChatParser(),
     )
 
-    train_ds = _dummy_dataset(
-        MySource(data=["1", "2"], repeat=1), batch_size=2
-    )
+    train_ds = _dummy_dataset(MySource(data=["1", "2"], repeat=1), batch_size=2)
 
     with (
         mock.patch.object(
@@ -639,9 +640,7 @@ class MockModel(nnx.Module):
       def __init__(self, *, rngs: nnx.Rngs):
         self.lm_head = 1
 
-      def __call__(
-          self, inputs, positions, cache, attention_mask, **kwargs
-      ):
+      def __call__(self, inputs, positions, cache, attention_mask, **kwargs):
         del kwargs
         return (
             jnp.full(
@@ -703,9 +702,7 @@ class MockModel(nnx.Module):
       def __init__(self, *, rngs: nnx.Rngs):
         self.lm_head = 1
 
-      def __call__(
-          self, inputs, positions, cache, attention_mask, **kwargs
-      ):
+      def __call__(self, inputs, positions, cache, attention_mask, **kwargs):
         del kwargs
         return (
             jnp.full(
@@ -2375,6 +2372,122 @@ def _patch_process_results(
         decoded_completion.count("Assistant:"), 2
     )  # 3 turns but terminal env obs does not append generation msg
 
+  def test_force_on_policy_ratio_bypasses_actor_recompute(self):
+    """Verifies force_on_policy_ratio=True sets old_logps to None with 0 extra passes."""
+    vocab = _mock_vocab()
+    tokenizer = tokenizer_adapter.TokenizerAdapter(vocab)
+    model = test_common.ToyTransformer(
+        config=test_common.ModelConfig(vocab_size=vocab.GetPieceSize()),
+        rngs=nnx.Rngs(0),
+    )
+    ref_model = test_common.ToyTransformer(
+        config=test_common.ModelConfig(vocab_size=vocab.GetPieceSize()),
+        rngs=nnx.Rngs(0),
+    )
+    mesh = pxla.thread_resources.env.physical_mesh
+    cluster_config = rl_engine_lib.ClusterConfig(
+        role_to_mesh={
+            rl_engine_lib.Role.ACTOR: mesh,
+            rl_engine_lib.Role.REFERENCE: mesh,
+            rl_engine_lib.Role.ROLLOUT: mesh,
+        },
+        rollout_engine="vanilla",
+        offload_to_cpu=False,
+        training_config=rl_engine_lib.RLTrainingConfig(
+            actor_optimizer=optax.sgd(1e-3),
+            eval_every_n_steps=10,
+            max_steps=10,
+        ),
+        rollout_config=base_rollout.RolloutConfig(
+            max_prompt_length=32,
+            max_tokens_to_generate=10,
+            return_logprobs=True,
+        ),
+    )
+    rl_engine = rl_engine_lib.RLEngine(
+        actor=model,
+        reference=ref_model,
+        tokenizer=tokenizer,
+        cluster_config=cluster_config,
+    )
+    grpo_config = agentic_grpo_learner.GRPOConfig(
+        beta=0.0,
+        force_compute_kl=False,
+        max_response_length=10,
+        num_generations=2,
+        num_iterations=1,
+        use_rollout_logps=True,
+        force_on_policy_ratio=True,
+    )
+    learner = agentic_grpo_learner.GRPOLearner(
+        rl_engine=rl_engine,
+        reward_fns=reward_fn_1,
+        algo_config=grpo_config,
+        chat_parser=MockChatParser(),
+    )
+
+    class MockTraj:
+
+      def __init__(self, index):
+        self.traj = {
+            "conversation_text": [
+                {"role": "assistant", "content": f"msg {index}"}
+            ],
+            "conversation_tokens": np.array([1, 2, 3]),
+            "conversation_masks": np.array([1, 1, 1]),
+            "old_logprobs": np.full(3, 1.0, dtype=np.float32),
+            "policy_version": 0,
+            "trajectory_reward": 1.0,
+            "prompt_tokens": np.array([4, 5]),
+            "original_input": {"prompts": "hello"},
+            "group_id": "test_group",
+        }
+
+    trajectories = [MockTraj(0), MockTraj(1)]
+
+    with mock.patch.object(
+        rl_engine,
+        "get_actor_per_token_logps",
+        return_value=jnp.full((2, 10), -1.0),
+        autospec=True,
+    ) as mock_get_actor_logps:
+      results = learner._process_results(trajectories, expected_step=1)
+      self.assertLen(results, 1)
+      train_example = results[0]
+
+      # 1. Asserts 0 extra trainer forward passes!
+      mock_get_actor_logps.assert_not_called()
+
+      # 2. Asserts old_per_token_logps is None (forces ratio to 1.0 via stop_gradient in loss)
+      self.assertIsNone(train_example.old_per_token_logps)
+
+  def test_force_on_policy_ratio_config_validation(self):
+    """force_on_policy_ratio rejects multi-iteration, allows stale trajectories."""
+    # num_iterations > 1 raises ValueError: old_logp is re-derived from the current
+    # policy on every inner epoch, so the trust region vanishes after the first.
+    with self.assertRaisesRegex(
+        ValueError, "can only be True when num_iterations == 1"
+    ):
+      agentic_grpo_learner.GRPOConfig(
+          num_generations=2,
+          num_iterations=2,
+          force_on_policy_ratio=True,
+      )
+
+    # off_policy_steps > 0 is a supported trade-off (near-on-policy training on
+    # slightly stale trajectories), so it warns rather than raising.
+    with mock.patch.object(
+        agentic_grpo_learner.logging, "warning"
+    ) as mock_warn:
+      config = agentic_grpo_learner.GRPOConfig(
+          num_generations=2,
+          num_iterations=1,
+          off_policy_steps=1,
+          force_on_policy_ratio=True,
+      )
+    mock_warn.assert_called_once()
+    self.assertIn("off_policy_steps", mock_warn.call_args[0][0])
+
 
 if __name__ == "__main__":
   absltest.main()

tunix/rl/agentic/agentic_grpo_learner.py

@@ -37,19 +37,19 @@
 import jax
 import jax.numpy as jnp
 import numpy as np
-from tunix.rl import algo_core  # pylint: disable=unused-import
 from tunix.perf.experimental import constants as perf_constants
+from tunix.rl import algo_core  # pylint: disable=unused-import
 from tunix.rl import common
 from tunix.rl import function_registry
 from tunix.rl import rl_cluster as rl_engine_lib
 from tunix.rl import utils as rl_utils
 from tunix.rl.agentic import agentic_rl_learner
 from tunix.rl.agentic import utils as agentic_utils
-from tunix.utils import compat
 from tunix.rl.agentic.agents import base_agent
 from tunix.rl.agentic.agents import model_agent
 from tunix.rl.agentic.environments import base_environment
 from tunix.rl.agentic.environments import task_environment
+from tunix.utils import compat
 from tunix.utils import trajectory_logger
 
 TrainingInputT = agentic_rl_learner.TrainingInputT
@@ -83,6 +83,17 @@ class GRPOConfig(agentic_rl_learner.AgenticRLConfig):
     max_concurrency: Maximum number of concurrent rollout engines.
     off_policy_steps: Number of off-policy steps can be accepted before a policy
       update.
+    use_rollout_logps: Use the rollout engine's log-probabilities as
+      old_per_token_logps. False makes the trainer recompute them.
+    force_on_policy_ratio: When num_iterations == 1, use
+      stop_gradients(current_logp) as old_per_token_logps instead of recomputing
+      or using rollout logps. Pins the surrogate ratio to 1.0, so clipping never
+      fires and sampler-vs-trainer numerical noise is removed from the ratio.
+    log_sampler_trainer_agreement: Optionally spend one extra trainer forward
+      pass per step to log sampler-vs-trainer log-probability agreement metrics.
+      Without force_on_policy_ratio these metrics come for free from the logps
+      already being computed; with it, no trainer logps exist, so this flag
+      pays for them explicitly. Default False
     degenerate_group_masking: Whether to mask out degenerate groups with all-0
       advantages. Deprecated. Will remove in the next release.
   """
@@ -112,6 +123,12 @@ class GRPOConfig(agentic_rl_learner.AgenticRLConfig):
       False  # Whether to mask out degenerate groups with all-0 advantages.
   )
   use_rollout_logps: bool = True
+  # Pin the surrogate ratio to 1.0 (old_logp := stop_gradient(current_logp)).
+  # Valid for single-iteration on-policy training only.
+  force_on_policy_ratio: bool = False
+  # Costs one trainer forward pass; keeps the sampler/trainer agreement metrics
+  # alive when force_on_policy_ratio would otherwise leave nothing to compare.
+  log_sampler_trainer_agreement: bool = False
   # Truncated importance-sampling (TIS) correction for the residual mismatch
   # between the rollout sampler and the trainer's recomputed log-probabilities.
   # Set to ``"token"`` to enable per-token TIS weights. When enabled, the loss
@@ -140,6 +157,28 @@ def __post_init__(self):
           "loss_algo should be either grpo or gspo-token. Received: "
           f"{self.loss_algo}"
       )
+    if self.force_on_policy_ratio:
+      if self.num_iterations > 1:
+        raise ValueError(
+            "force_on_policy_ratio can only be True when num_iterations == 1."
+            " With num_iterations > 1 the policy is updated several times on"
+            " the same batch, so the surrogate ratio is genuinely != 1 after"
+            " the first inner epoch; pinning it to 1 removes the trust region"
+            " for every subsequent epoch. Got"
+            f" num_iterations={self.num_iterations}"
+        )
+
+      if self.off_policy_steps > 0:
+        logging.warning(
+            "force_on_policy_ratio=True with off_policy_steps=%d: trajectories "
+            "may be up to %d policy updates stale, but the surrogate ratio is "
+            "pinned to 1.0, so the off-policy correction is discarded. This is "
+            "a deliberate trade for near-on-policy training; pair it with an "
+            "importance-sampling correction if the behavior and target "
+            "policies can drift apart.",
+            self.off_policy_steps,
+            self.off_policy_steps,
+        )
 
 
 TGrpoConfig = TypeVar("TGrpoConfig", bound=GRPOConfig)
@@ -228,9 +267,11 @@ def __init__(
     else:
       logging.warning("Metrics log dir is None, skipping trajectory logging.")
 
-    self.algo_config.temperature = self.rl_engine.get_rollout_config(  # pyrefly: ignore[missing-attribute]
-        mode=rl_engine_lib.Mode.TRAIN
-    ).temperature
+    self.algo_config.temperature = (
+        self.rl_engine.get_rollout_config(  # pyrefly: ignore[missing-attribute]
+            mode=rl_engine_lib.Mode.TRAIN
+        ).temperature
+    )
 
     # Workaround to pass loss fn with algorithm flag
     policy_loss_fn = function_registry.get_policy_loss_fn(
@@ -414,6 +455,7 @@ def _compute_packed_logps(self, example: TrainExample) -> TrainExample:
     if (
         example.old_per_token_logps is None
         and not self.algo_config.use_rollout_logps
+        and not self.algo_config.force_on_policy_ratio
     ):
       updates["old_per_token_logps"] = self.rl_engine.get_actor_per_token_logps(
           prompt_tokens=prompt_tokens,
@@ -439,7 +481,11 @@ def _compute_packed_logps(self, example: TrainExample) -> TrainExample:
     # The rollout-logps path defers its trainer recompute here too. Not just
     # diagnostics: sampler_is="token" consumes it as old_per_token_logps.
     need_trainer_logps = (
-        self.algo_config.use_rollout_logps
+        (
+            not self.algo_config.force_on_policy_ratio
+            or self.algo_config.log_sampler_trainer_agreement
+        )
+        and self.algo_config.use_rollout_logps
         and example.old_per_token_logps is not None
         and (self._have_actor_mesh() or self.algo_config.sampler_is == "token")
     )
@@ -465,7 +511,10 @@ def _compute_packed_logps(self, example: TrainExample) -> TrainExample:
         )
       if sampler_is_weights is not None:
         updates["sampler_is_weights"] = sampler_is_weights
-      if self.algo_config.sampler_is == "token":
+      if (
+          self.algo_config.sampler_is == "token"
+          and not self.algo_config.force_on_policy_ratio
+      ):
         updates["old_per_token_logps"] = trainer_logps
 
     if updates:
@@ -573,7 +622,7 @@ def _process_results(
       )
       if (
           len(completion_tokens) >= max_response_length
-          and completion_mask[-1] != eos_value
+          and completion_tokens[-1] != eos_value
       ):
         clipped_completion_count += 1
       padded_prompt, padded_completion, _ = (
@@ -646,7 +695,24 @@ def _process_results(
 
     rollout_per_token_logps = None
     trainer_per_token_logps = None
-    if self.algo_config.use_rollout_logps and padded_old_logprobs:
+    if self.algo_config.force_on_policy_ratio:
+      # The loss derives old_per_token_logps from the actor's own forward pass.
+      old_per_token_logps = None
+      if padded_old_logprobs:
+        rollout_per_token_logps = jnp.asarray(padded_old_logprobs)
+        if (
+            self.algo_config.log_sampler_trainer_agreement
+            and have_actor_mesh
+            and not is_packed
+        ):
+          trainer_per_token_logps = self.rl_engine.get_actor_per_token_logps(
+              prompt_tokens=prompt_ids,
+              completion_tokens=completion_ids,
+              pad_id=pad_value,
+              eos_id=eos_value,
+              micro_batch_size=compute_logps_micro_batch_size,
+          )
+    elif self.algo_config.use_rollout_logps and padded_old_logprobs:
       rollout_per_token_logps = jnp.asarray(padded_old_logprobs)
       old_per_token_logps = rollout_per_token_logps
       # The diagnostic pass (and the sampler-IS ``token`` path, which needs the
@@ -845,11 +911,11 @@ def _process_results(
             f"{prefix}/{sub_key}/max": (np.max(flat_vals), np.max),
             f"{prefix}/{sub_key}/min": (np.min(flat_vals), np.min),
         })
-        self.rl_engine.buffer_metrics_async(
-            metrics_to_log,  # pyrefly: ignore[bad-argument-type]
-            mode=mode,
-            step=expected_step,  # pyrefly: ignore[bad-argument-type]
-        )
+      self.rl_engine.buffer_metrics_async(
+          metrics_to_log,  # pyrefly: ignore[bad-argument-type]
+          mode=mode,
+          step=expected_step,  # pyrefly: ignore[bad-argument-type]
+      )
 
     for metric_fn in self.metric_fns:
       user_defined_metric = metric_fn(
@@ -864,7 +930,9 @@ def _process_results(
           },
       )
       self.rl_engine.buffer_metrics_async(
-          user_defined_metric, mode=mode, step=expected_step  # pyrefly: ignore[bad-argument-type]
+          user_defined_metric,
+          mode=mode,
+          step=expected_step,  # pyrefly: ignore[bad-argument-type]
       )
 
     combined_batch = TrainExample(