Make AgenticGRPOLearner import optional.

PiperOrigin-RevId: 896861388

Author: abhinavclemson

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

@@ -67,7 +67,6 @@
 from tunix.rl import rl_cluster as rl_cluster_lib
 from tunix.rl.rollout import base_rollout
 from tunix.rl.grpo.grpo_learner import GrpoConfig, GrpoLearner
-from tunix.rl.agentic.agentic_grpo_learner import GrpoConfig as AgenticGrpoConfig, GrpoLearner as AgenticGrpoLearner
 from tunix.sft import metrics_logger, profiler
 
 # for vLLM we can skip JAX precompilation with this flag, it makes startup faster
@@ -601,6 +600,15 @@ def _reward_fn(**kwargs):
   max_logging.log("Setting up RL trainer...")
   if trainer_config.rl.use_agentic_rollout:
     max_logging.log("Using AgenticGRPOLearner with async online rollouts.")
+    # TODO: Remove this try-except once the dependency on tunix is fixed.
+    try:
+      from tunix.rl.agentic.agentic_grpo_learner import GrpoConfig as AgenticGrpoConfig  # pylint: disable=import-outside-toplevel
+      from tunix.rl.agentic.agentic_grpo_learner import GrpoLearner as AgenticGrpoLearner  # pylint: disable=import-outside-toplevel
+    except ImportError as e:
+      raise ValueError(
+          "tunix.rl.agentic dependencies are not installed! "
+          "Please install tunix with agentic support to use 'use_agentic_rollout'."
+      ) from e
     grpo_config = AgenticGrpoConfig(
         num_generations=trainer_config.rl.num_generations,
         num_iterations=trainer_config.rl.num_iterations,

tests/post_training/unit/train_rl_test.py

@@ -319,6 +319,7 @@ def get_dataset_side_effect(model_tokenizer, config, data_dir, split, data_files
     # Configs
     trainer_config = SimpleNamespace(
         debug=SimpleNamespace(rl=False),
+        rl=SimpleNamespace(use_agentic_rollout=False),
         tokenizer_path="dummy_path",
         dataset_name="dummy_dataset",
         train_split="train",