use agent trainer

Author: thwu1

examples/math_tinker/train_math_tinker.py

@@ -13,7 +13,7 @@
 from examples.math_tinker.math_reward import math_reward_fn
 from rllm.data.dataset import DatasetRegistry
 from rllm.environments.base.single_turn_env import SingleTurnEnvironment
-from rllm.trainer.tinker.tinker_agent_trainer import TinkerAgentTrainer
+from rllm.trainer import AgentTrainer
 
 
 @hydra.main(version_base=None, config_path="../../rllm/trainer/config", config_name="tinker_agent_trainer")
@@ -32,18 +32,19 @@ def main(config: DictConfig):
         raise ValueError("Datasets not found! Please run prepare_tinker_math_dataset.py first:\n  python -m examples.math_tinker.prepare_tinker_math_dataset")
 
     # Create trainer (uses separated components internally)
-    trainer = TinkerAgentTrainer(
+    trainer = AgentTrainer(
         config=config,
         agent_class=MathAgentWithFewshot,
         env_class=SingleTurnEnvironment,
         agent_args={"use_fewshot": True},
         env_args={"reward_fn": math_reward_fn},
         train_dataset=train_dataset,
         val_dataset=test_dataset,
+        backend="tinker",
     )
 
     # Train (all orchestration handled internally by TinkerAgentTrainer)
-    trainer.fit_agent()
+    trainer.train()
 
 
 if __name__ == "__main__":

examples/solver_judge_tinker/train_solver_judge_flow_tinker.py

@@ -3,15 +3,15 @@
 from examples.solver_judge.solver_judge_flow import SolverJudgeWorkflow
 from rllm.data.dataset import DatasetRegistry
 from rllm.rewards.countdown_reward import countdown_reward_fn
-from rllm.trainer.tinker.tinker_workflow_trainer import TinkerWorkflowTrainer
+from rllm.trainer import AgentTrainer
 
 
 @hydra.main(config_path="pkg://rllm.trainer.config", config_name="tinker_workflow_trainer", version_base=None)
 def main(config):
     train_dataset = DatasetRegistry.load_dataset("countdown", "train")
     test_dataset = DatasetRegistry.load_dataset("countdown", "test")
 
-    trainer = TinkerWorkflowTrainer(
+    trainer = AgentTrainer(
         workflow_class=SolverJudgeWorkflow,
         workflow_args={
             "n_solutions": 2,
@@ -20,8 +20,9 @@ def main(config):
         config=config,
         train_dataset=train_dataset,
         val_dataset=test_dataset,
+        backend="tinker",
     )
-    trainer.fit_agent()
+    trainer.train()
 
 
 if __name__ == "__main__":

rllm/trainer/agent_trainer.py

@@ -24,6 +24,7 @@ def __init__(
         config: dict[str, Any] | list[str] | None = None,
         train_dataset: Dataset | None = None,
         val_dataset: Dataset | None = None,
+        backend: str = "verl",
     ):
         """
         Initialize the AgentTrainer.
@@ -59,13 +60,48 @@ def __init__(
         self.env_args = env_args or {}
 
         self.config = config
+        self.train_dataset = train_dataset
+        self.val_dataset = val_dataset
+        self.backend = backend
+
+        assert self.backend in ["verl", "tinker"], f"Unsupported backend: {self.backend}, must be one of ['verl', 'tinker']"
 
         if train_dataset is not None and self.config is not None and hasattr(self.config, "data"):
             self.config.data.train_files = train_dataset.get_verl_data_path()
         if val_dataset is not None and self.config is not None and hasattr(self.config, "data"):
             self.config.data.val_files = val_dataset.get_verl_data_path()
 
     def train(self):
+        if self.backend == "verl":
+            self._train_verl()
+        elif self.backend == "tinker":
+            self._train_tinker()
+
+    def _train_tinker(self):
+        from rllm.trainer.tinker.tinker_agent_trainer import TinkerAgentTrainer
+        from rllm.trainer.tinker.tinker_workflow_trainer import TinkerWorkflowTrainer
+
+        if self.config.rllm.workflow.use_workflow:
+            trainer = TinkerWorkflowTrainer(
+                config=self.config,
+                workflow_class=self.workflow_class,
+                workflow_args=self.workflow_args,
+                train_dataset=self.train_dataset,
+                val_dataset=self.val_dataset,
+            )
+        else:
+            trainer = TinkerAgentTrainer(
+                config=self.config,
+                agent_class=self.agent_class,
+                env_class=self.env_class,
+                agent_args=self.agent_args,
+                env_args=self.env_args,
+                train_dataset=self.train_dataset,
+                val_dataset=self.val_dataset,
+            )
+        trainer.fit_agent()
+
+    def _train_verl(self):
         # Check if Ray is not initialized
         if not ray.is_initialized():
             # read off all the `ray_init` settings from the config

rllm/trainer/config/tinker_agent_trainer.yaml

@@ -63,6 +63,10 @@ trainer:
   val_before_train: true
   default_local_dir: '/tmp/rllm-tinker-checkpoints'
 
+rllm:
+  workflow:
+    use_workflow: false
+
 # Hydra configuration
 hydra:
   run:

rllm/trainer/config/tinker_workflow_trainer.yaml

@@ -60,6 +60,10 @@ trainer:
   val_before_train: true
   default_local_dir: '/tmp/rllm-tinker-checkpoints'
 
+rllm:
+  workflow:
+    use_workflow: true
+
 # Hydra configuration
 hydra:
   run: