Add reward_functions_path + reward_functions CLI knobs for custom rewards

Currently the reward stack is hardcoded to a 3-fn list:
[match_format_exactly, match_format_approximately, check_numbers].
Replacing it requires editing train_rl.py.

Add two new config fields:
  - reward_functions_path: path to a user Python file
  - reward_functions: comma-separated list of function names to import

When both are set, the built-in reward stack is REPLACED entirely by the
user-provided functions (so users can pin a single VTC-style partial-credit
reward, swap in a math_verify-based scorer, etc., without editing maxtext).

Each user function must accept (prompts, completions, tmvp_config, **kwargs)
and return a list of floats.

Default (both empty) keeps existing behavior unchanged.

Reuses `_load_custom_callable` helper added in the previous commit.

Author: Pooya Moradi

src/maxtext/configs/post_train/rl.yml

@@ -238,6 +238,13 @@ eval_dataset_name: 'openai/gsm8k'
 # tmvp_config, x) -> dict with keys {prompts, question, answer}. When empty
 # (default), the built-in utils_rl.process_data is used.
 dataset_processor_path: ''
+# Optional: path to a user-provided Python file with custom reward functions,
+# AND a comma-separated list of function names to import from that file.
+# Each function signature: fn(prompts, completions, tmvp_config, **kwargs) -> list[float].
+# When both are set, the built-in [match_format_exactly, match_format_approximately,
+# check_numbers] reward list is REPLACED entirely.
+reward_functions_path: ''
+reward_functions: ''
 train_split: 'train'
 eval_split: 'test'
 hf_name: 'main' # subset of Hugging Face dataset

src/maxtext/configs/types.py

@@ -2056,6 +2056,20 @@ class RLDataset(BaseModel):
           "When set, replaces the built-in dataset processor for custom datasets."
       ),
   )
+  reward_functions_path: str = Field(
+      "",
+      description=(
+          "Optional path to a user Python file containing custom reward functions. "
+          "Used with `reward_functions` to fully replace the built-in reward stack."
+      ),
+  )
+  reward_functions: str = Field(
+      "",
+      description=(
+          "Comma-separated names of reward functions to import from `reward_functions_path`. "
+          "Each function signature: (prompts, completions, tmvp_config, **kwargs) -> list[float]."
+      ),
+  )
 
 
 class RLEvaluation(BaseModel):

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

@@ -526,11 +526,28 @@ def _reward_fn(**kwargs):
 
     return _reward_fn
 
-  reward_fns = [  # type: ignore
-      make_reward_fn(utils_rl.match_format_exactly),
-      make_reward_fn(utils_rl.match_format_approximately),
-      make_reward_fn(utils_rl.check_numbers),
-  ]
+  # Optional user-provided reward functions. `reward_functions_path` is a
+  # filesystem path to a Python file. `reward_functions` is a comma-separated
+  # list of function names to import from that file. Each function must accept
+  # `prompts`, `completions`, `tmvp_config`, and `**kwargs` and return a list
+  # of floats.
+  # When both are set, the built-in reward_fns list below is REPLACED entirely
+  # by the user-provided functions (so users have full control over their
+  # reward stack — match the GPU recipe by passing your own vtc reward fn).
+  _custom_rewards_path = getattr(trainer_config, "reward_functions_path", "") or ""
+  _custom_rewards_names = getattr(trainer_config, "reward_functions", "") or ""
+  if _custom_rewards_path and _custom_rewards_names:
+    _names = [n.strip() for n in _custom_rewards_names.split(",") if n.strip()]
+    reward_fns = [make_reward_fn(utils_rl.load_custom_callable(_custom_rewards_path, n)) for n in _names]
+    max_logging.log(
+        f"reward_fns: using {len(reward_fns)} custom reward function(s) " f"{_names} from {_custom_rewards_path}"
+    )
+  else:
+    reward_fns = [  # type: ignore
+        make_reward_fn(utils_rl.match_format_exactly),
+        make_reward_fn(utils_rl.match_format_approximately),
+        make_reward_fn(utils_rl.check_numbers),
+    ]
 
   # Create RL trainer
   max_logging.log("Setting up RL trainer...")