Skip to content

Agent Trainer

Main training interface for RL algorithms and agent optimization.

rllm.trainer.agent_trainer

AgentTrainer

A wrapper class that allows users to easily train custom agents with custom environments without having to directly interact with the underlying training infrastructure.

Source code in rllm/trainer/agent_trainer.py 
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
class AgentTrainer:
    """
    A wrapper class that allows users to easily train custom agents with custom environments
    without having to directly interact with the underlying training infrastructure.
    """

    def __init__(
        self,
        workflow_class: type | None = None,
        workflow_args: dict[str, Any] | None = None,
        agent_class: type | None = None,
        env_class: type | None = None,
        agent_args: dict[str, Any] | None = None,
        env_args: dict[str, Any] | None = None,
        config: dict[str, Any] | list[str] | None = None,
        train_dataset: Dataset | None = None,
        val_dataset: Dataset | None = None,
    ):
        """
        Initialize the AgentTrainer.

        Args:
            agent_class: The custom agent class to use for training
            env_class: The custom environment class to use for training
            config: Configuration overrides to apply to the default config
                   Can be a dictionary with dot notation keys (e.g., {"data.train_batch_size": 8})
                   or a list of strings in the format "key=value" (e.g., ["data.train_batch_size=8"])
            train_dataset: Optional train dataset to use
            val_dataset: Optional validation dataset to use
            agent_args: Optional arguments to pass to the agent class
            env_args: Optional arguments to pass to the environment class
        """

        if workflow_class is not None and config.rllm.workflow.use_workflow:
            if agent_class is not None:
                raise ValueError("agent_class is not supported when using workflow, instead use workflow_args['agent_cls']")
            if agent_args is not None:
                raise ValueError("agent_args is not supported when using workflow, instead use workflow_args['agent_args']")
            if env_class is not None:
                raise ValueError("env_class is not supported when using workflow, instead use workflow_args['env_cls']")
            if env_args is not None:
                raise ValueError("env_args is not supported when using workflow, instead use workflow_args['env_args']")

        self.workflow_class = workflow_class
        self.workflow_args = workflow_args or {}

        self.agent_class = agent_class
        self.env_class = env_class
        self.agent_args = agent_args or {}
        self.env_args = env_args or {}

        self.config = config

        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):
        # Check if Ray is not initialized
        if not ray.is_initialized():
            # read off all the `ray_init` settings from the config
            if self.config is not None and hasattr(self.config, "ray_init"):
                ray_init_settings = {k: v for k, v in self.config.ray_init.items() if v is not None}
            else:
                ray_init_settings = {}
            ray.init(runtime_env=get_ppo_ray_runtime_env(), **ray_init_settings)

        runner = TaskRunner.remote()

        ray.get(
            runner.run.remote(
                config=self.config,
                workflow_class=self.workflow_class,
                workflow_args=self.workflow_args,
                agent_class=self.agent_class,
                env_class=self.env_class,
                agent_args=self.agent_args,
                env_args=self.env_args,
            )
        )

__init__ 

__init__(workflow_class: type | None = None, workflow_args: dict[str, Any] | None = None, agent_class: type | None = None, env_class: type | None = None, agent_args: dict[str, Any] | None = None, env_args: dict[str, Any] | None = None, config: dict[str, Any] | list[str] | None = None, train_dataset: Dataset | None = None, val_dataset: Dataset | None = None)

Initialize the AgentTrainer.

Parameters:

Name Type Description Default
agent_class type | None

The custom agent class to use for training

 None
env_class type | None

The custom environment class to use for training

 None
config dict[str, Any] | list[str] | None

Configuration overrides to apply to the default config Can be a dictionary with dot notation keys (e.g., {"data.train_batch_size": 8}) or a list of strings in the format "key=value" (e.g., ["data.train_batch_size=8"])

 None
train_dataset Dataset | None

Optional train dataset to use

 None
val_dataset Dataset | None

Optional validation dataset to use

 None
agent_args dict[str, Any] | None

Optional arguments to pass to the agent class

 None
env_args dict[str, Any] | None

Optional arguments to pass to the environment class

 None
Source code in rllm/trainer/agent_trainer.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
def __init__(
    self,
    workflow_class: type | None = None,
    workflow_args: dict[str, Any] | None = None,
    agent_class: type | None = None,
    env_class: type | None = None,
    agent_args: dict[str, Any] | None = None,
    env_args: dict[str, Any] | None = None,
    config: dict[str, Any] | list[str] | None = None,
    train_dataset: Dataset | None = None,
    val_dataset: Dataset | None = None,
):
    """
    Initialize the AgentTrainer.

    Args:
        agent_class: The custom agent class to use for training
        env_class: The custom environment class to use for training
        config: Configuration overrides to apply to the default config
               Can be a dictionary with dot notation keys (e.g., {"data.train_batch_size": 8})
               or a list of strings in the format "key=value" (e.g., ["data.train_batch_size=8"])
        train_dataset: Optional train dataset to use
        val_dataset: Optional validation dataset to use
        agent_args: Optional arguments to pass to the agent class
        env_args: Optional arguments to pass to the environment class
    """

    if workflow_class is not None and config.rllm.workflow.use_workflow:
        if agent_class is not None:
            raise ValueError("agent_class is not supported when using workflow, instead use workflow_args['agent_cls']")
        if agent_args is not None:
            raise ValueError("agent_args is not supported when using workflow, instead use workflow_args['agent_args']")
        if env_class is not None:
            raise ValueError("env_class is not supported when using workflow, instead use workflow_args['env_cls']")
        if env_args is not None:
            raise ValueError("env_args is not supported when using workflow, instead use workflow_args['env_args']")

    self.workflow_class = workflow_class
    self.workflow_args = workflow_args or {}

    self.agent_class = agent_class
    self.env_class = env_class
    self.agent_args = agent_args or {}
    self.env_args = env_args or {}

    self.config = config

    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()