Add **kwargs to train_prm and train_rl to accept config parameters

src/reasoning/step_level_cot.py

@@ -353,6 +353,7 @@ class StepLevelCoTTrainer:
         learning_rate: float = 1e-5,
         save_steps: int = 500,
         eval_steps: int = 500,
+        **kwargs,  # Accept additional config like hidden_dim, num_layers, dropout, reward_scale
     ) -> None:
         """
         Train Process Reward Model.
@@ -362,9 +363,12 @@ class StepLevelCoTTrainer:
             learning_rate: Learning rate for PRM training
             save_steps: Save checkpoint every N steps
             eval_steps: Evaluate every N steps
+            **kwargs: Additional PRM configuration (hidden_dim, num_layers, dropout, reward_scale)
         """
         logger.info("Starting PRM training")
         logger.info(f"Max steps: {max_steps}, LR: {learning_rate}")
+        if kwargs:
+            logger.info(f"PRM config: {kwargs}")
         
         # Load datasets using the trainer's dataset loading methods
         train_dataset = self.load_step_dataset(split='train')
@@ -449,6 +453,7 @@ class StepLevelCoTTrainer:
         learning_rate: float = 5e-6,
         save_steps: int = 500,
         eval_steps: int = 500,
+        **kwargs,  # Accept additional config like gamma, lam, cliprange, etc.
     ) -> None:
         """
         Train policy with reinforcement learning.
@@ -458,9 +463,12 @@ class StepLevelCoTTrainer:
             learning_rate: Learning rate for RL training
             save_steps: Save checkpoint every N steps
             eval_steps: Evaluate every N steps
+            **kwargs: Additional PPO configuration (gamma, lam, cliprange, vf_coef, ent_coef, etc.)
         """
         logger.info("Starting RL training")
         logger.info(f"Max steps: {max_steps}, LR: {learning_rate}")
+        if kwargs:
+            logger.info(f"PPO config: {kwargs}")
         
         # Load dataset using the trainer's dataset loading methods
         train_dataset = self.load_step_dataset(split='train')