feat(rl): plumb GRPO flags to the agentic learner

src/maxtext/configs/types.py

@@ -2596,6 +2596,29 @@ class RL(BaseModel):
           "If None, no chunking is applied, which may lead to OOM errors if tensors are too large."
       ),
   )
+  use_rollout_logps: bool = Field(
+      True,
+      description=(
+          "Use rollout engine's logprobs as old_per_token_logps. "
+          "False selects the step-0 re-forward path (trainer recomputes them)"
+      ),
+  )
+  force_on_policy_ratio: bool = Field(
+      False,
+      description=(
+          "Pin the PPO/GRPO surrogate ratio to exactly 1.0 by using "
+          "stop_gradient(current_logp) as old_per_token_logps. Valid only for "
+          "single-iteration on-policy training (num_iterations=1)."
+      ),
+  )
+  log_sampler_trainer_agreement: bool = Field(
+      False,
+      description=(
+          "Compute an extra trainer forward pass per step to log sampler-vs-trainer "
+          "logp agreement metrics. Costs ~1 forward pass; needed because "
+          "force_on_policy_ratio otherwise leaves no trainer logps to compare."
+      ),
+  )
 
 
 class RLDataset(BaseModel):

src/maxtext/trainers/post_train/rl/train_rl.py

@@ -48,6 +48,7 @@
 from functools import wraps
 from typing import Any, Callable, Optional, Sequence
 
+import dataclasses
 import datasets
 import grain
 import jax
@@ -619,11 +620,19 @@ def _reward_fn(**kwargs):
         beta=trainer_config.rl.grpo_beta,
         epsilon=trainer_config.rl.grpo_epsilon,
         loss_algo=trainer_config.rl.loss_algo,
+        loss_agg_mode=trainer_config.rl.loss_agg_mode,
         max_response_length=trainer_config.max_target_length - trainer_config.max_prefill_predict_length,
         max_concurrency=trainer_config.rl.max_concurrency,
         off_policy_steps=trainer_config.rl.off_policy_steps,
         system_prompt=trainer_config.rl.system_prompt,
         epsilon_high=trainer_config.rl.epsilon_high,
+        use_rollout_logps=trainer_config.rl.use_rollout_logps,
+        force_on_policy_ratio=trainer_config.rl.force_on_policy_ratio,
+        log_sampler_trainer_agreement=(trainer_config.rl.log_sampler_trainer_agreement),
+    )
+    max_logging.log(
+        "GRPO config resolved:\n"
+        + "\n".join(f"  {k} = {v!r}" for k, v in sorted(dataclasses.asdict(grpo_config).items()))
     )
     # Instantiate the custom MaxText chat parser
     template_config = load_data_template_from_file(trainer_config.data_template_path)