Allow configurable micro-batching for compute_logps in agentic RL.

PiperOrigin-RevId: 914005102

Author: hgao327

tests/rl/agentic/agentic_rl_learner_test.py

@@ -14,10 +14,15 @@
 
 """Tests for agentic_rl_learner."""
 
+import asyncio
+from typing import Any
 from unittest import mock
 
+from absl import logging
 from absl.testing import absltest
 from absl.testing import parameterized
+from tunix.rl import rl_cluster as rl_cluster_lib
+from tunix.rl import utils as rl_utils
 from tunix.rl.agentic import agentic_rl_learner
 from tunix.rl.rollout import base_rollout
 
@@ -135,6 +140,44 @@ def test_validate_rollout_config_vllm_missing_server_mode(self):
           algo_config=algo_config,
       )
 
+  def test_train_batch_size_mismatch_raises_error(self):
+    with mock.patch.object(
+        rl_utils, "is_sharing_weights", return_value=False
+    ):
+      rl_cluster = mock.Mock()
+      rl_cluster.cluster_config = mock.Mock()
+      rl_cluster.cluster_config.role_to_mesh = {
+          rl_cluster_lib.Role.ACTOR: mock.Mock(),
+          rl_cluster_lib.Role.ROLLOUT: mock.Mock(),
+      }
+      training_config = mock.Mock()
+      training_config.compute_logps_micro_batch_size = 2
+      training_config.train_micro_batch_size = 1
+      training_config.mini_batch_size = None
+      rl_cluster.cluster_config.training_config = training_config
+      rl_cluster.cluster_config.rollout_config = base_rollout.RolloutConfig(
+          max_tokens_to_generate=10, return_logprobs=True
+      )
+      rl_cluster.cluster_config.rollout_engine = 'generic'
+      rl_cluster.actor_trainer = mock.Mock()
+      rl_cluster.actor_trainer.restored_global_step.return_value = 0
+      rl_cluster.actor_trainer.iter_steps = 0
+      rl_cluster.rollout = mock.Mock()
+      rl_cluster.tokenizer = mock.Mock()
+      algo_config = agentic_rl_learner.AgenticRLConfig(max_response_length=10)
+      learner = DummyLearner(
+          rl_cluster=rl_cluster,
+          reward_fns=mock.Mock(),
+          algo_config=algo_config,
+      )
+      train_dataset = [{'prompt': ['p1']}]
+      with self.assertRaisesRegex(
+          ValueError,
+          r'compute_logps_micro_batch_size \(2\) must be equal to'
+          r' train_micro_batch_size \(1\)',
+      ):
+        learner.train(train_dataset)
+
 
 if __name__ == "__main__":
   absltest.main()

tunix/rl/agentic/agentic_rl_learner.py

@@ -211,7 +211,7 @@ def __init__(
         self._training_config.rollout_micro_batch_size
     )
     self._compute_logps_micro_batch_size = (
-        self._training_config.compute_logps_micro_batch_size
+        self._training_config.compute_logps_micro_batch_size or 1
     )
     sft_utils.show_hbm_usage(title="AgenticRLLearner init")
 
@@ -413,8 +413,6 @@ def _model_call(
       if "pair_index" in env.extra_kwargs:
         tags[perf_constants.PAIR_INDEX] = env.extra_kwargs["pair_index"]
 
-
-
     result = self.rl_cluster.generate(
         prompts=chat_lists,
         apply_chat_template=False if self.chat_parser else True,
@@ -694,22 +692,20 @@ def train(
     train_micro_batch_size = (
         self._training_config.train_micro_batch_size or mini_batch_size
     )
-    # Rollout and compute_logps micro batch sizes have to be 1 since we only
-    # process inidividual prompts.
+    # Rollout micro batch size has to be 1 since we only process individual
+    # prompts.
     self._rollout_micro_batch_size = 1
-
-    compute_logps_mb = self._training_config.compute_logps_micro_batch_size
     self._process_in_consumer = False
-    if compute_logps_mb is not None and compute_logps_mb > 1:
-      if compute_logps_mb != train_micro_batch_size:
+
+    if self._compute_logps_micro_batch_size > 1:
+      if self._compute_logps_micro_batch_size != train_micro_batch_size:
         raise ValueError(
-            f"compute_logps_micro_batch_size ({compute_logps_mb}) must be"
-            f" equal to train_micro_batch_size ({train_micro_batch_size})"
+            "compute_logps_micro_batch_size"
+            f" ({self._compute_logps_micro_batch_size}) must be equal to"
+            f" train_micro_batch_size ({train_micro_batch_size})"
         )
       self._process_in_consumer = True
-      self._compute_logps_micro_batch_size = compute_logps_mb
-    else:
-      self._compute_logps_micro_batch_size = 1
+
     for v, n in [
         (self._rollout_micro_batch_size, f"{self._rollout_micro_batch_size=}"),
         (